From b82295fe331d8ee4196dc39c8da0eda7f6b10a90 Mon Sep 17 00:00:00 2001 From: timo <44401485+Timo972@users.noreply.github.com> Date: Mon, 20 Jul 2026 17:59:22 +0200 Subject: [PATCH 01/17] docs: design concurrent SSE terminal guard --- ...20-concurrent-sse-terminal-guard-design.md | 91 +++++++++++++++++++ 1 file changed, 91 insertions(+) create mode 100644 docs/superpowers/specs/2026-07-20-concurrent-sse-terminal-guard-design.md diff --git a/docs/superpowers/specs/2026-07-20-concurrent-sse-terminal-guard-design.md b/docs/superpowers/specs/2026-07-20-concurrent-sse-terminal-guard-design.md new file mode 100644 index 0000000..b3ed446 --- /dev/null +++ b/docs/superpowers/specs/2026-07-20-concurrent-sse-terminal-guard-design.md @@ -0,0 +1,91 @@ +# Concurrent SSE Terminal Guard Design + +## Problem + +CC-Router forwards Anthropic Messages SSE responses transparently. It records a +request as successful as soon as upstream response headers arrive, but it does +not verify that the SSE body contains the required terminal `message_stop` +event. + +The affected machine shows repeated responses that deliver content and a +terminal `message_delta` with `stop_reason`, then produce no further event. +Claude Code waits for `message_stop` and reports `Response stalled mid-stream` +at its five-minute watchdog boundary. Multiple concurrent sessions expose the +failure more frequently. Raising the timeout would only delay the same error. + +## Chosen Approach + +Add a streaming terminal-event guard to the existing Anthropic proxy response +path. It will observe SSE bytes while preserving immediate passthrough and keep +independent state for every response. + +For each successful Anthropic SSE response, the guard will: + +1. Parse event frames incrementally across arbitrary chunk boundaries. +2. Remember whether it has seen a terminal `message_delta` with a non-null + `stop_reason` and whether it has seen `message_stop`. +3. Forward every upstream byte immediately and unchanged. +4. When a terminal `message_delta` arrives, allow a short one-second grace + period for the normal `message_stop`. If it does not arrive, append a + canonical `message_stop` SSE frame, end the downstream response, and release + the stalled upstream stream. +5. If upstream ends before the grace timer after a terminal `message_delta`, + append `message_stop` before ending the downstream response. +6. If upstream ends without either terminal signal, do not fabricate success; + record the stream as incomplete and let the downstream client handle the + truncated response. + +The repair is intentionally narrow: a terminal `message_delta` proves that the +model finished and that only the protocol terminator is missing. The router +will not invent content, stop reasons, tool results, or successful completion +for an otherwise incomplete response. + +## Integration + +The guard will be a small, independently tested stream component rather than +additional ad hoc listeners inside `server.ts`. The Anthropic proxy will use it +only for successful `text/event-stream` responses. JSON responses, compressed +responses, OpenAI routes, and error status bodies retain their current paths. + +Existing token-usage accounting will continue to observe the stream. Stream +completion will update the pending request log after the body finishes, so an +HTTP 200 with a truncated body is no longer indistinguishable from a complete +response. + +Each guard instance owns its parser buffer and terminal flags. There is no +module-level stream state, so concurrent requests cannot affect one another. + +## Error Handling + +- Complete stream: forward unchanged and record normal completion. +- Terminal delta without `message_stop`: wait up to one second, append + `message_stop`, record that the terminator was repaired, and end normally. +- Truncated before a terminal delta: forward what was received, mark the log as + an incomplete-stream error, and do not claim completion. +- Upstream stream error: preserve the existing proxy error behavior and mark + the request log with the connection error. +- Client disconnect: stop forwarding and release listeners/state without + attempting a repair for a client that is no longer present. + +## Testing + +Regression tests will cover: + +- A complete SSE stream remains byte-for-byte unchanged. +- A terminal `message_delta` followed by either EOF or silence receives exactly + one synthesized `message_stop`. +- A stream already containing `message_stop` never receives a duplicate. +- SSE frames split across chunks are parsed correctly. +- A stream ending before any terminal delta is reported as incomplete and is + not repaired. +- Several concurrent streams maintain isolated state, including a mixture of + complete, repairable, and genuinely truncated responses. +- The focused tests, full Vitest suite, typecheck, and build all pass. + +## Out of Scope + +- Retrying or failing over a response after partial content has reached the + client. +- Buffering full model responses. +- Changing account selection or rate-limit policy. +- Increasing timeouts as the primary fix. From d5ab64e35f222bc993e72bfdb3d31080602e432c Mon Sep 17 00:00:00 2001 From: timo <44401485+Timo972@users.noreply.github.com> Date: Mon, 20 Jul 2026 18:23:24 +0200 Subject: [PATCH 02/17] docs: redesign routing around session affinity --- ...20-concurrent-sse-terminal-guard-design.md | 91 --------- ...7-20-intelligent-session-routing-design.md | 180 ++++++++++++++++++ 2 files changed, 180 insertions(+), 91 deletions(-) delete mode 100644 docs/superpowers/specs/2026-07-20-concurrent-sse-terminal-guard-design.md create mode 100644 docs/superpowers/specs/2026-07-20-intelligent-session-routing-design.md diff --git a/docs/superpowers/specs/2026-07-20-concurrent-sse-terminal-guard-design.md b/docs/superpowers/specs/2026-07-20-concurrent-sse-terminal-guard-design.md deleted file mode 100644 index b3ed446..0000000 --- a/docs/superpowers/specs/2026-07-20-concurrent-sse-terminal-guard-design.md +++ /dev/null @@ -1,91 +0,0 @@ -# Concurrent SSE Terminal Guard Design - -## Problem - -CC-Router forwards Anthropic Messages SSE responses transparently. It records a -request as successful as soon as upstream response headers arrive, but it does -not verify that the SSE body contains the required terminal `message_stop` -event. - -The affected machine shows repeated responses that deliver content and a -terminal `message_delta` with `stop_reason`, then produce no further event. -Claude Code waits for `message_stop` and reports `Response stalled mid-stream` -at its five-minute watchdog boundary. Multiple concurrent sessions expose the -failure more frequently. Raising the timeout would only delay the same error. - -## Chosen Approach - -Add a streaming terminal-event guard to the existing Anthropic proxy response -path. It will observe SSE bytes while preserving immediate passthrough and keep -independent state for every response. - -For each successful Anthropic SSE response, the guard will: - -1. Parse event frames incrementally across arbitrary chunk boundaries. -2. Remember whether it has seen a terminal `message_delta` with a non-null - `stop_reason` and whether it has seen `message_stop`. -3. Forward every upstream byte immediately and unchanged. -4. When a terminal `message_delta` arrives, allow a short one-second grace - period for the normal `message_stop`. If it does not arrive, append a - canonical `message_stop` SSE frame, end the downstream response, and release - the stalled upstream stream. -5. If upstream ends before the grace timer after a terminal `message_delta`, - append `message_stop` before ending the downstream response. -6. If upstream ends without either terminal signal, do not fabricate success; - record the stream as incomplete and let the downstream client handle the - truncated response. - -The repair is intentionally narrow: a terminal `message_delta` proves that the -model finished and that only the protocol terminator is missing. The router -will not invent content, stop reasons, tool results, or successful completion -for an otherwise incomplete response. - -## Integration - -The guard will be a small, independently tested stream component rather than -additional ad hoc listeners inside `server.ts`. The Anthropic proxy will use it -only for successful `text/event-stream` responses. JSON responses, compressed -responses, OpenAI routes, and error status bodies retain their current paths. - -Existing token-usage accounting will continue to observe the stream. Stream -completion will update the pending request log after the body finishes, so an -HTTP 200 with a truncated body is no longer indistinguishable from a complete -response. - -Each guard instance owns its parser buffer and terminal flags. There is no -module-level stream state, so concurrent requests cannot affect one another. - -## Error Handling - -- Complete stream: forward unchanged and record normal completion. -- Terminal delta without `message_stop`: wait up to one second, append - `message_stop`, record that the terminator was repaired, and end normally. -- Truncated before a terminal delta: forward what was received, mark the log as - an incomplete-stream error, and do not claim completion. -- Upstream stream error: preserve the existing proxy error behavior and mark - the request log with the connection error. -- Client disconnect: stop forwarding and release listeners/state without - attempting a repair for a client that is no longer present. - -## Testing - -Regression tests will cover: - -- A complete SSE stream remains byte-for-byte unchanged. -- A terminal `message_delta` followed by either EOF or silence receives exactly - one synthesized `message_stop`. -- A stream already containing `message_stop` never receives a duplicate. -- SSE frames split across chunks are parsed correctly. -- A stream ending before any terminal delta is reported as incomplete and is - not repaired. -- Several concurrent streams maintain isolated state, including a mixture of - complete, repairable, and genuinely truncated responses. -- The focused tests, full Vitest suite, typecheck, and build all pass. - -## Out of Scope - -- Retrying or failing over a response after partial content has reached the - client. -- Buffering full model responses. -- Changing account selection or rate-limit policy. -- Increasing timeouts as the primary fix. diff --git a/docs/superpowers/specs/2026-07-20-intelligent-session-routing-design.md b/docs/superpowers/specs/2026-07-20-intelligent-session-routing-design.md new file mode 100644 index 0000000..908f264 --- /dev/null +++ b/docs/superpowers/specs/2026-07-20-intelligent-session-routing-design.md @@ -0,0 +1,180 @@ +# Intelligent Session Routing Design + +## Problem + +CC-Router currently calls `TokenPool.getNext()` for every `/v1/*` request. The +`X-Claude-Code-Session-Id` header is used only to label traffic as coming from +Claude Code; it does not influence account selection. + +This request-level round robin has two harmful effects: + +1. A conversation moves between subscription accounts. Prompt caches are + isolated between accounts, so each account repeatedly creates overlapping + cache prefixes instead of extending one warm conversation cache. +2. The pool does not track in-flight requests. Concurrent sessions can select + an account that is already serving a long stream, increasing upstream + throttling and stalled-stream risk. + +The router must make account assignment at the session level while preserving +the upstream and downstream protocols unchanged. + +## Goals + +- Keep every healthy Claude Code session on one account for cache locality. +- Distribute new sessions across usable accounts according to live load and + rate-limit headroom. +- Move a session only when its account cannot serve it. +- Track request leases until the downstream response actually completes. +- Preserve Anthropic response status, headers, SSE bytes, ordering, and timing. +- Keep routing state bounded and avoid persisting session identifiers. + +## Non-Goals + +- Modifying, normalizing, or synthesizing Anthropic SSE events. +- Retrying a partially delivered response on another account. +- Sharing prompt caches across accounts; Anthropic isolates those caches. +- Persisting affinity mappings across router restarts. + +## Routing Model + +### Session identity + +For Claude Code traffic, the normalized `X-Claude-Code-Session-Id` header is +the affinity key. Empty values and values longer than 256 UTF-8 bytes are +ignored. Requests +without a usable session ID use load-aware request routing without persistent +affinity. + +The router stores only an in-memory mapping: + +```text +session ID -> account ID + last-seen timestamp +``` + +Mappings expire after one hour of inactivity. This bounds stale state and +matches the longest prompt-cache lifetime relevant to normal Claude Code use. +The map is capped at 10,000 entries. Expired entries are swept first; if it is +still full, the least recently used mapping is evicted. Session IDs are never +logged or written to disk. + +### Existing sessions + +If a session has a binding and the account remains enabled, healthy, below its +configured caps, and outside cooldown, the router reuses that account. Cache +affinity takes priority over global load balancing for an established session. + +An account's current in-flight count does not break an existing binding. This +prevents related requests within one Claude Code session from bouncing to a +different cache domain. + +### New sessions and unscoped requests + +The router filters accounts using the existing health, enablement, cap, and +cooldown rules. It then ranks eligible accounts by: + +1. Lowest in-flight request count. +2. Fewest active session bindings. +3. Greatest remaining rate-limit headroom, based on the worse of the 5-hour + and 7-day utilization ratios relative to configured caps. +4. Round-robin order as the deterministic tie-breaker. + +The headroom score is the lower-is-better value +`max(fiveHourUtil / sessionCap, sevenDayUtil / weeklyCap)`, with percentages +normalized to fractions. A zero cap makes the account ineligible; missing +upstream utilization starts at zero. The rotating tie-break cursor advances +only after an account is selected. + +The chosen account is bound to the session before forwarding begins, ensuring +simultaneous new sessions cannot all observe and select the same idle account. +Requests without a session ID use the same ranking but do not create a binding. + +### Failover and rebinding + +A binding is invalidated when its account is disabled, unhealthy, over a user +cap, in cooldown, or returns an account-specific 401, 429, or 529 response. +The failed response is still relayed unchanged. Claude Code's next retry is +then assigned to the best eligible account and establishes a new binding. + +The router never retries after response bytes have been sent. This avoids +duplicating text or tool calls. + +If every account is unavailable, the existing fallback policy remains, but it +selects the least-loaded account and logs the fallback reason. User caps remain +advisory when bypassing them is the only way to avoid an empty pool. + +## Request Leases + +Selection returns an account lease rather than a bare account. Acquiring a +lease increments that account's in-flight count synchronously. Releasing it is +idempotent and decrements the count exactly once. + +The HTTP layer releases the lease on all terminal paths: + +- downstream response `finish`; +- downstream/client `close`; +- proxy or upstream error; +- refresh failure or another pre-forward rejection. + +This lifecycle represents the complete body, not merely receipt of upstream +headers. It therefore works for long SSE responses and ordinary JSON bodies. + +Cooldown is represented by an expiry timestamp rather than conflating active +work with the existing `busy` boolean. Expired cooldowns are swept during +selection and health polling. + +## Components + +### `TokenPool` + +- Owns per-account in-flight counts. +- Exposes lease acquisition and idempotent release. +- Preserves existing health, cap, cooldown, mutation, and fallback behavior. +- Provides the load and headroom values needed by the session router. + +### `SessionRouter` + +- Owns bounded, expiring session-to-account mappings. +- Reuses valid bindings. +- Chooses and binds accounts for new sessions. +- Invalidates bindings after account-specific failures or management changes. +- Does not depend on Express or inspect message content. + +### Proxy integration + +- Extracts and validates the session header. +- Acquires a routed lease before token refresh and forwarding. +- Attaches one idempotent cleanup function to every response path. +- Invalidates affinity after 401, 429, and 529 responses. +- Leaves `http-proxy-middleware` response passthrough unchanged. + +## Observability + +Health/account views add an in-flight count and active-session count. Route +logs may record the non-sensitive selection reason (`sticky`, `new-session`, +`unscoped`, or `failover`) but never the session ID. A response remains logged +with its actual upstream status. + +## Testing + +Unit tests will verify: + +- Repeated requests from one session retain the same account. +- New sessions distribute across idle accounts before reusing one. +- In-flight load, active bindings, rate-limit headroom, and round-robin order + are applied in the documented priority. +- A valid sticky binding survives unrelated load changes. +- Disabled, unhealthy, capped, cooling-down, and failed accounts cause a + controlled rebind on the next request. +- Lease release is idempotent and never produces a negative count. +- Expired and excess mappings are evicted without logging session IDs. +- Requests without a session ID are load-aware but not sticky. + +Integration tests will hold several SSE responses open concurrently and verify: + +- Different sessions use separate idle accounts when available. +- A session's follow-up request retains its account. +- Leases remain active until the full response finishes or disconnects. +- Complete upstream SSE bodies remain byte-for-byte identical downstream. +- No `message_stop` or other SSE frame is inserted, removed, or rewritten. + +The focused tests, full Vitest suite, typecheck, and production build must pass. From a7cfbd42e2f363929b79342cffe3cf00526ce896 Mon Sep 17 00:00:00 2001 From: timo <44401485+Timo972@users.noreply.github.com> Date: Mon, 20 Jul 2026 18:39:02 +0200 Subject: [PATCH 03/17] docs: plan intelligent session routing --- .../2026-07-20-intelligent-session-routing.md | 661 ++++++++++++++++++ 1 file changed, 661 insertions(+) create mode 100644 docs/superpowers/plans/2026-07-20-intelligent-session-routing.md diff --git a/docs/superpowers/plans/2026-07-20-intelligent-session-routing.md b/docs/superpowers/plans/2026-07-20-intelligent-session-routing.md new file mode 100644 index 0000000..13f5798 --- /dev/null +++ b/docs/superpowers/plans/2026-07-20-intelligent-session-routing.md @@ -0,0 +1,661 @@ +# Intelligent Session Routing Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Replace request-level round robin for Anthropic subscription traffic with sticky, load-aware Claude Code session routing that preserves prompt-cache affinity and tracks concurrent streams until they really finish. + +**Architecture:** `TokenPool` owns eligibility, cooldown timestamps, per-account in-flight counters, load/headroom ranking, and idempotent request leases. A new Express-independent `SessionRouter` owns a bounded in-memory session-to-account map and chooses sticky or load-aware routes. The existing Anthropic proxy acquires one routed lease before refresh/forwarding, invalidates bindings after account-specific failures, and continues to let `http-proxy-middleware` relay upstream responses without handling or rewriting SSE bodies. + +**Tech Stack:** TypeScript, Node.js 20+, Express 4, `http-proxy-middleware` 3, Vitest 4, native Node HTTP test servers. + +## Global Constraints + +- Use normalized `X-Claude-Code-Session-Id` as the affinity key. Trim surrounding whitespace; ignore non-string/duplicate values, empty values, and values longer than 256 UTF-8 bytes. +- Store affinity only in memory as `session ID -> account ID + lastSeen`. Never log or persist session IDs. +- Expire bindings after 1 hour of inactivity. Cap the map at 10,000 entries; sweep expired entries first, then evict the least recently used entry before inserting. +- A valid existing binding wins even when its account already has in-flight work. It remains valid only while the account is enabled, healthy, below both user caps, not rate-limited, and outside cooldown. +- Rank new sessions and unscoped requests lexicographically by: lowest in-flight requests, fewest active session bindings, lowest `max(fiveHourUtil / sessionCap, sevenDayUtil / weeklyCap)`, then a rotating account-order tie-break. Normalize percent caps to fractions; a zero cap is ineligible; missing utilization is zero. Advance the tie-break cursor only after selection. +- Bind a new session synchronously before the HTTP request is forwarded. Requests without a valid session ID use the same ranking but create no binding. +- Invalidate a binding after its account is disabled, unhealthy, capped, cooling down, removed, or returns 401, 429, or 529. Relay the failed upstream response unchanged and reassign only the next client retry. +- Never retry after upstream response bytes have been received. Never synthesize, remove, reorder, parse-and-reencode, or otherwise rewrite SSE events, including `message_stop`. +- If all accounts are unavailable, keep caps advisory and select the least-loaded fallback account while invoking the existing fallback observability hook. +- Acquiring a lease increments in-flight state synchronously. Releasing is idempotent, decrements exactly once, and never makes a counter negative. +- Release leases on downstream `finish`, downstream/client `close`, proxy/upstream error, refresh failure, and other pre-forward terminal responses. +- Represent new 429/529 cooldown state with an expiry timestamp owned by `TokenPool`; do not use active-work state as cooldown state. Keep the existing public `busy` field only for backward compatibility while routing and health metrics use explicit cooldown/in-flight APIs. +- OpenAI subscription routing, Anthropic request bodies, authentication header behavior, rate-limit extraction, and usage accounting stay behaviorally unchanged. +- Health/account JSON adds numeric `inFlightRequests` and `activeSessions` fields and does not expose tokens or session identifiers. +- Every production change starts with a focused failing test. Focused tests, the full Vitest suite, `npm run lint`, and `npm run build` must pass with pristine output before completion. + +--- + +### Task 1: Lease-Aware, Load-Aware Token Pool + +**Files:** + +- Modify: `src/proxy/token-pool.ts` +- Modify: `src/__tests__/token-pool.test.ts` + +- [ ] **Step 1: Add failing lease lifecycle tests** + +Extend `src/__tests__/token-pool.test.ts` with deterministic tests for synchronous acquisition and idempotent release: + +```ts +describe("TokenPool — request leases", () => { + it("tracks a request until its lease is released exactly once", () => { + const pool = new TokenPool([makeAccount("a")]); + const lease = pool.acquireBest(new Map()); + + expect(lease.account.id).toBe("a"); + expect(pool.getInFlight("a")).toBe(1); + + lease.release(); + lease.release(); + expect(pool.getInFlight("a")).toBe(0); + }); + + it("prefers the account with fewer in-flight requests", () => { + const pool = new TokenPool([makeAccount("a"), makeAccount("b")]); + const first = pool.tryAcquire("a"); + expect(first).not.toBeNull(); + + const next = pool.acquireBest(new Map()); + expect(next.account.id).toBe("b"); + + first!.release(); + next.release(); + }); +}); +``` + +- [ ] **Step 2: Run the focused test and record RED evidence** + +Run: + +```bash +npm test -- src/__tests__/token-pool.test.ts +``` + +Expected: FAIL because `acquireBest`, `tryAcquire`, and `getInFlight` do not exist. + +- [ ] **Step 3: Add the lease and ranking contracts** + +In `src/proxy/token-pool.ts`, add these public contracts and inject a clock for deterministic cooldown tests: + +```ts +export interface AccountLease { + readonly account: Account; + readonly fallback: boolean; + release(): void; +} + +export interface TokenPoolOptions { + now?: () => number; +} + +export class TokenPool { + private readonly inFlight = new Map(); + private readonly cooldownUntil = new Map(); + private readonly now: () => number; + private currentIndex = 0; + + constructor(private readonly accounts: Account[], options: TokenPoolOptions = {}) { + this.now = options.now ?? Date.now; + } +} +``` + +Implement the following API: + +```ts +acquireBest(activeSessions: ReadonlyMap): AccountLease; +tryAcquire(accountId: string): AccountLease | null; +isEligible(accountId: string): boolean; +getInFlight(accountId: string): number; +setCooldown(accountId: string, durationMs: number): void; +isCoolingDown(accountId: string): boolean; +``` + +`tryAcquire` must reject missing or currently ineligible accounts. Both acquisition paths must increment `requestCount`, set `lastUsed`, increment the in-flight map, and return an idempotent closure that decrements only its own lease once. + +- [ ] **Step 4: Add failing ranking and cooldown tests** + +Cover the complete priority order in separate tests: + +```ts +it("uses active session count after in-flight load ties", () => { + const pool = new TokenPool([makeAccount("a"), makeAccount("b")]); + const lease = pool.acquireBest(new Map([["a", 2], ["b", 1]])); + expect(lease.account.id).toBe("b"); + lease.release(); +}); + +it("uses relative rate-limit headroom after load and binding ties", () => { + const a = makeAccount("a"); + const b = makeAccount("b"); + a.sessionLimitPercent = 80; + a.rateLimits.fiveHourUtil = 0.6; + b.sessionLimitPercent = 80; + b.rateLimits.fiveHourUtil = 0.2; + const pool = new TokenPool([a, b]); + + const lease = pool.acquireBest(new Map([["a", 1], ["b", 1]])); + expect(lease.account.id).toBe("b"); + lease.release(); +}); + +it("keeps an account unavailable until its timestamp cooldown expires", () => { + let now = 1_000; + const pool = new TokenPool( + [makeAccount("a"), makeAccount("b")], + { now: () => now }, + ); + pool.setCooldown("a", 500); + expect(pool.isEligible("a")).toBe(false); + now = 1_500; + expect(pool.isEligible("a")).toBe(true); +}); +``` + +Also verify: zero-percent caps are ineligible; five-hour and seven-day ratios use the worse value; the rotating tie-break distributes exact ties; a valid `tryAcquire` does not reject solely due to existing in-flight work; and all-unavailable fallback chooses the lowest in-flight count. + +- [ ] **Step 5: Implement lexicographic selection and timestamp cooldown** + +Keep the existing rate-window rollover helpers. Eligible selection must compare an account tuple equivalent to: + +```ts +[ + getInFlight(account.id), + activeSessions.get(account.id) ?? 0, + Math.max( + account.rateLimits.fiveHourUtil / (account.sessionLimitPercent / 100), + account.rateLimits.sevenDayUtil / (account.weeklyLimitPercent / 100), + ), + circularDistanceFromCursor, +] +``` + +Do not evaluate the ratio for zero caps because those accounts are ineligible. When there are no eligible accounts, choose lexicographically by in-flight count and existing reset preference from progressively broader fallback sets, and set `fallback: true`. Invoke `onCapBypass` when the selected fallback bypasses user caps. + +Replace the old server-facing use of `busy` with `cooldownUntil`, but preserve `getNext()` as a compatibility wrapper: + +```ts +getNext(): Account { + const lease = this.acquireBest(new Map()); + lease.release(); + return lease.account; +} +``` + +The wrapper preserves existing callers/tests without leaking an in-flight lease. Include `inFlightRequests` and derived `coolingDown` in `getStats()`. + +- [ ] **Step 6: Run focused and full regression tests** + +Run: + +```bash +npm test -- src/__tests__/token-pool.test.ts +npm test +npm run lint +``` + +Expected: all pass with no warnings or unhandled timer output. + +- [ ] **Step 7: Commit Task 1** + +```bash +git add src/proxy/token-pool.ts src/__tests__/token-pool.test.ts +git commit -m "feat: add load-aware account leases" +``` + +--- + +### Task 2: Bounded Sticky Session Router + +**Files:** + +- Create: `src/proxy/session-router.ts` +- Create: `src/__tests__/session-router.test.ts` + +- [ ] **Step 1: Write failing normalization and affinity tests** + +Create `src/__tests__/session-router.test.ts` with a local `makeAccount` fixture and these first behaviors: + +```ts +import { describe, expect, it } from "vitest"; +import { SessionRouter, normalizeSessionId } from "../proxy/session-router.js"; +import { TokenPool } from "../proxy/token-pool.js"; + +it("normalizes one bounded string header", () => { + expect(normalizeSessionId(" session-a ")).toBe("session-a"); + expect(normalizeSessionId(" ")).toBeUndefined(); + expect(normalizeSessionId(["a", "b"])).toBeUndefined(); + expect(normalizeSessionId("é".repeat(129))).toBeUndefined(); +}); + +it("keeps repeated requests from one session on its account", () => { + const pool = new TokenPool([makeAccount("a"), makeAccount("b")]); + const router = new SessionRouter(pool); + + const first = router.acquire("session-a"); + const second = router.acquire("session-a"); + + expect(first.account.id).toBe(second.account.id); + expect(first.reason).toBe("new-session"); + expect(second.reason).toBe("sticky"); + first.release(); + second.release(); +}); +``` + +- [ ] **Step 2: Run the focused test and record RED evidence** + +Run: + +```bash +npm test -- src/__tests__/session-router.test.ts +``` + +Expected: FAIL because the session router module does not exist. + +- [ ] **Step 3: Implement the public routing contract** + +Create `src/proxy/session-router.ts` with: + +```ts +import type { AccountLease } from "./token-pool.js"; +import { TokenPool } from "./token-pool.js"; + +export type RouteReason = "sticky" | "new-session" | "unscoped" | "failover"; + +export interface RoutedAccountLease extends AccountLease { + readonly reason: RouteReason; + readonly sessionId?: string; +} + +export interface SessionRouterOptions { + now?: () => number; + ttlMs?: number; + maxEntries?: number; +} + +export function normalizeSessionId(value: unknown): string | undefined; + +export class SessionRouter { + constructor(pool: TokenPool, options?: SessionRouterOptions); + acquire(sessionHeader: unknown): RoutedAccountLease; + invalidate(sessionHeader: unknown, expectedAccountId?: string): boolean; + invalidateAccount(accountId: string): number; + getActiveSessionCount(accountId: string): number; + getBindingCount(): number; +} +``` + +Use defaults `ttlMs = 60 * 60 * 1000` and `maxEntries = 10_000`. Keep both a binding map and per-account active-binding counts so the pool can rank without rescanning the whole map. + +- [ ] **Step 4: Add failing distribution, failover, and bound-state tests** + +Add tests proving: + +- two simultaneous new sessions bind separate idle accounts; +- unscoped requests are load-aware but do not increase `getBindingCount()`; +- a sticky binding stays put when another account later has lower load; +- disabling, marking unhealthy, capping, cooling down, or removing the bound account yields `reason: "failover"` on the next acquire; +- `invalidate(sessionId, accountId)` only removes the matching current binding, preventing an old response from deleting a newer rebind; +- `invalidateAccount` removes every binding for one account and fixes counts; +- bindings update `lastSeen` on access and expire at one hour; +- insertion at 10,000 entries evicts the least-recently-used binding after sweeping expired entries; +- no public result or diagnostic callback logs/exposes any session ID other than the transient `sessionId` field required by the HTTP layer. + +Use small injected `maxEntries` and a mutable injected clock for eviction tests rather than allocating 10,001 test sessions. + +- [ ] **Step 5: Implement binding, expiry, LRU eviction, and failover** + +`acquire` must run synchronously from lookup through lease acquisition and binding insertion. Its control flow is: + +```ts +const sessionId = normalizeSessionId(sessionHeader); +sweepExpiredBindings(); + +if (!sessionId) { + return wrap(pool.acquireBest(activeSessionCounts), "unscoped"); +} + +const existing = bindings.get(sessionId); +if (existing) { + const stickyLease = pool.tryAcquire(existing.accountId); + if (stickyLease) { + existing.lastSeen = now(); + return wrap(stickyLease, "sticky", sessionId); + } + removeBinding(sessionId); +} + +const lease = pool.acquireBest(activeSessionCounts); +insertBinding(sessionId, lease.account.id, now()); +return wrap(lease, existing ? "failover" : "new-session", sessionId); +``` + +The wrapper delegates the original idempotent `release`; it must not introduce a second in-flight increment. LRU eviction compares `lastSeen` and uses map insertion order only as a deterministic tie-break. + +- [ ] **Step 6: Run focused and dependency tests** + +Run: + +```bash +npm test -- src/__tests__/session-router.test.ts src/__tests__/token-pool.test.ts +npm test +npm run lint +``` + +Expected: all pass and no session ID appears in emitted console output. + +- [ ] **Step 7: Commit Task 2** + +```bash +git add src/proxy/session-router.ts src/__tests__/session-router.test.ts +git commit -m "feat: add sticky Claude session routing" +``` + +--- + +### Task 3: Proxy Lifecycle, Failure Invalidation, and Transparent SSE + +**Files:** + +- Create: `src/proxy/anthropic-proxy.ts` +- Create: `src/proxy/lease-lifecycle.ts` +- Create: `src/__tests__/anthropic-proxy.test.ts` +- Create: `src/__tests__/lease-lifecycle.test.ts` +- Modify: `src/proxy/server.ts` + +- [ ] **Step 1: Write failing lease terminal-path tests** + +Create `src/__tests__/lease-lifecycle.test.ts` around a Node `EventEmitter`-backed response double: + +```ts +import { EventEmitter } from "node:events"; +import { describe, expect, it, vi } from "vitest"; +import { attachLeaseLifecycle } from "../proxy/lease-lifecycle.js"; + +it.each(["finish", "close"])("releases once on downstream %s", (event) => { + const response = new EventEmitter(); + const release = vi.fn(); + attachLeaseLifecycle(response, { release }); + + response.emit(event); + response.emit("finish"); + response.emit("close"); + expect(release).toHaveBeenCalledTimes(1); +}); +``` + +Add a test for the returned explicit cleanup function so proxy errors and pre-forward failures share the same one-shot release. + +- [ ] **Step 2: Run the lifecycle test and record RED evidence** + +Run: + +```bash +npm test -- src/__tests__/lease-lifecycle.test.ts +``` + +Expected: FAIL because `lease-lifecycle.ts` does not exist. + +- [ ] **Step 3: Implement one-shot HTTP lease cleanup** + +Create `src/proxy/lease-lifecycle.ts`: + +```ts +export interface ResponseLifecycle { + once(event: "finish" | "close", listener: () => void): unknown; +} + +export interface Releasable { + release(): void; +} + +export function attachLeaseLifecycle( + response: ResponseLifecycle, + lease: Releasable, +): () => void { + let released = false; + const release = () => { + if (released) return; + released = true; + lease.release(); + }; + response.once("finish", release); + response.once("close", release); + return release; +} +``` + +- [ ] **Step 4: Write a failing byte-exact proxy integration test** + +Create `src/__tests__/anthropic-proxy.test.ts`. Start a native Node upstream server that writes deliberately split SSE chunks, including a final `message_stop`, with small delays. Mount the production proxy factory on an ephemeral Express server, collect the downstream body as `Buffer`, and assert: + +```ts +expect(downstream.status).toBe(200); +expect(downstream.headers.get("content-type")).toContain("text/event-stream"); +expect(Buffer.compare(downstreamBody, upstreamBody)).toBe(0); +expect(downstreamBody.toString("utf8").match(/event: message_stop/g)).toHaveLength(1); +``` + +Add a concurrent case holding two streams open and asserting neither body completes before its own upstream terminates. Always close both ephemeral servers in `finally` blocks. + +- [ ] **Step 5: Extract the transparent Anthropic proxy factory** + +Create `src/proxy/anthropic-proxy.ts` as the single place that configures `http-proxy-middleware` transport behavior: + +```ts +import type { RequestHandler, Request } from "express"; +import type { ServerResponse } from "node:http"; +import { createProxyMiddleware } from "http-proxy-middleware"; +import type { Options } from "http-proxy-middleware"; + +export interface AnthropicProxyOptions { + target: string; + timeoutMs: number; + on: NonNullable["on"]>; +} + +export function createAnthropicProxy(options: AnthropicProxyOptions): RequestHandler { + return createProxyMiddleware({ + target: options.target, + changeOrigin: true, + pathRewrite: path => `/v1${path}`, + proxyTimeout: options.timeoutMs, + timeout: options.timeoutMs, + on: options.on, + }); +} +``` + +Do not set `selfHandleResponse`, buffer/replace the response stream, or add data-transform pipes. Move the existing `proxyReq`, `proxyRes`, and `error` callbacks from the inline factory in `server.ts` into the `on` argument without changing auth, rate-limit extraction, logging, or usage listeners. + +- [ ] **Step 6: Add failing server-routing assertions** + +Extend the lifecycle test or add focused exported-helper tests that prove: + +- a request with a session header obtains a `RoutedAccountLease` and records only its `reason` in `_pendingLog.details`; +- a refresh-failure response invokes explicit cleanup before returning; +- proxy error cleanup is safe after response `close` already fired; +- 401 invalidates the matching session binding; +- 429 invalidates the binding and calls `pool.setCooldown` using sanitized `Retry-After` seconds with a 60-second fallback; +- 529 invalidates the binding and applies a 30-second cooldown; +- no handler automatically replays the request or writes an SSE event. + +Export narrowly scoped pure helpers from `lease-lifecycle.ts` if needed; do not export the Express app or add test-only branches to production code. + +- [ ] **Step 7: Integrate `SessionRouter` into `server.ts`** + +Instantiate one router beside the pool: + +```ts +const pool = new TokenPool(accounts); +const sessionRouter = new SessionRouter(pool); +``` + +Extend the request augmentation with `_ccRoute?: RoutedAccountLease` and `_ccReleaseLease?: () => void`. In `/v1`, replace `pool.getNext()` with: + +```ts +const route = sessionRouter.acquire(req.headers["x-claude-code-session-id"]); +req._ccRoute = route; +req._ccReleaseLease = attachLeaseLifecycle(res, route); +const account = route.account; +``` + +Acquire before token refresh. On every refresh/pre-forward error, invoke `_ccReleaseLease()` before ending the response; `finish`/`close` may invoke it again safely. In the proxy error callback, invoke it before writing the 502. + +In `proxyRes`, for 401/429/529 call: + +```ts +sessionRouter.invalidate(req._ccRoute?.sessionId, account.id); +``` + +For 429 and 529 call `pool.setCooldown` instead of mutating `account.busy` or creating `setTimeout` callbacks. Sanitize `Retry-After` as a finite non-negative number of seconds. Keep relaying the original response; do not retry from `proxyRes`. + +On management disable and successful removal, call `sessionRouter.invalidateAccount(account.id)`. Other invalid states are rejected lazily by `TokenPool.isEligible` on the next routed request. + +- [ ] **Step 8: Run focused proxy tests and full regression tests** + +Run: + +```bash +npm test -- src/__tests__/lease-lifecycle.test.ts src/__tests__/anthropic-proxy.test.ts src/__tests__/session-router.test.ts +npm test +npm run lint +``` + +Expected: all pass; SSE buffers match byte-for-byte; tests close all sockets/timers without hanging. + +- [ ] **Step 9: Commit Task 3** + +```bash +git add src/proxy/anthropic-proxy.ts src/proxy/lease-lifecycle.ts src/proxy/server.ts src/__tests__/anthropic-proxy.test.ts src/__tests__/lease-lifecycle.test.ts +git commit -m "fix: route concurrent streams by Claude session" +``` + +--- + +### Task 4: Routing Observability and User Documentation + +**Files:** + +- Modify: `src/proxy/server.ts` +- Modify: `src/ui/Dashboard.tsx` +- Modify: `src/__tests__/server-health-accounts.test.ts` +- Modify: `README.md` +- Modify: `package.json` + +- [ ] **Step 1: Add failing health-view tests** + +Extend `src/__tests__/server-health-accounts.test.ts` to pass a routing-metrics resolver and assert both providers return numeric counters: + +```ts +const views = createHealthAccountViews( + [makeAnthropicAccount()], + [openAIAccount], + accountId => accountId === "max-account-1" + ? { inFlightRequests: 2, activeSessions: 3, coolingDown: true } + : { inFlightRequests: 0, activeSessions: 0, coolingDown: false }, +); + +expect(views[0]).toMatchObject({ + busy: true, + inFlightRequests: 2, + activeSessions: 3, +}); +expect(views[1]).toMatchObject({ + inFlightRequests: 0, + activeSessions: 0, +}); +expect(JSON.stringify(views)).not.toContain("session-a"); +``` + +- [ ] **Step 2: Run the health test and record RED evidence** + +Run: + +```bash +npm test -- src/__tests__/server-health-accounts.test.ts +``` + +Expected: FAIL because the health view lacks routing counters and the resolver argument. + +- [ ] **Step 3: Add routing metrics to health/account JSON** + +Add to `HealthAccountView`: + +```ts +inFlightRequests: number; +activeSessions: number; +``` + +Define the resolver contract: + +```ts +export interface AccountRoutingMetrics { + inFlightRequests: number; + activeSessions: number; + coolingDown: boolean; +} + +type RoutingMetricsResolver = (accountId: string) => AccountRoutingMetrics; +``` + +Give `createHealthAccountViews` a backward-compatible default resolver returning zeros. At the live health and accounts endpoints, pass a resolver using `pool.getInFlight`, `sessionRouter.getActiveSessionCount`, and `pool.isCoolingDown`. Derive Anthropic `busy` as `account.busy || metrics.coolingDown`; OpenAI counters remain zero. + +Update the local `AccountStat` UI type in `src/ui/Dashboard.tsx` and show compact `N active / M streams` text for Anthropic accounts without altering account-management controls. + +- [ ] **Step 4: Document cache-aware routing and failure semantics** + +Update `README.md` to explain: + +- one Claude Code session stays on one Anthropic subscription account for prompt-cache locality; +- new sessions prefer idle accounts, then fewer bound sessions, then more rate-limit headroom; +- 401/429/529 invalidates affinity for the next retry while the current response is passed through unchanged; +- mappings are memory-only, expire after one hour, and session IDs are never logged; +- `proxyRequestTimeoutMs` is a transport safety limit, not a fix for missing upstream terminal events; +- the router never appends a synthetic `message_stop`. + +Change the package description from "Round-robin proxy" to "Cache-aware session router" and replace the `round-robin` keyword with `session-routing`. + +- [ ] **Step 5: Run health, UI, and full verification** + +Run: + +```bash +npm test -- src/__tests__/server-health-accounts.test.ts +npm test +npm run lint +npm run build +``` + +Expected: all tests and both TypeScript compilation commands pass with pristine output. + +- [ ] **Step 6: Commit Task 4** + +```bash +git add src/proxy/server.ts src/ui/Dashboard.tsx src/__tests__/server-health-accounts.test.ts README.md package.json +git commit -m "docs: explain cache-aware session routing" +``` + +--- + +## Final Verification and Review + +- [ ] Generate a whole-branch review package from the implementation base to the final task commit. +- [ ] Dispatch a fresh whole-branch reviewer against the approved design and this plan. +- [ ] Resolve every Critical or Important finding and re-run the affected focused tests. +- [ ] Run final verification from a clean working tree: + +```bash +npm test +npm run lint +npm run build +git status --short +``` + +- [ ] Confirm the full suite includes the concurrent byte-exact SSE test and exits without open handles. +- [ ] Use superpowers:finishing-a-development-branch to present the verified integration options. From c5d8f5aaf13813861b546f1fca5814e1816461e2 Mon Sep 17 00:00:00 2001 From: timo <44401485+Timo972@users.noreply.github.com> Date: Mon, 20 Jul 2026 18:39:26 +0200 Subject: [PATCH 04/17] chore: ignore local worktrees --- .gitignore | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.gitignore b/.gitignore index 0aa58bd..2042519 100644 --- a/.gitignore +++ b/.gitignore @@ -44,3 +44,6 @@ coverage/ # Temp files *.tmp + +# Isolated implementation worktrees +.worktrees/ From 18842b8ed19406ebf25f0f10d3d4177b0c6a3c22 Mon Sep 17 00:00:00 2001 From: timo <44401485+Timo972@users.noreply.github.com> Date: Mon, 20 Jul 2026 18:44:12 +0200 Subject: [PATCH 05/17] feat: add load-aware account leases --- src/__tests__/token-pool.test.ts | 148 ++++++++++++++++++++ src/proxy/token-pool.ts | 232 +++++++++++++++++++++++++------ 2 files changed, 335 insertions(+), 45 deletions(-) diff --git a/src/__tests__/token-pool.test.ts b/src/__tests__/token-pool.test.ts index e00e5a5..42909c3 100644 --- a/src/__tests__/token-pool.test.ts +++ b/src/__tests__/token-pool.test.ts @@ -56,6 +56,142 @@ describe("TokenPool — round-robin", () => { }); }); +describe("TokenPool — request leases", () => { + it("tracks a request until its lease is released exactly once", () => { + const pool = new TokenPool([makeAccount("a")]); + const lease = pool.acquireBest(new Map()); + + expect(lease.account.id).toBe("a"); + expect(pool.getInFlight("a")).toBe(1); + + lease.release(); + lease.release(); + expect(pool.getInFlight("a")).toBe(0); + }); + + it("prefers the account with fewer in-flight requests", () => { + const pool = new TokenPool([makeAccount("a"), makeAccount("b")]); + const first = pool.tryAcquire("a"); + expect(first).not.toBeNull(); + + const next = pool.acquireBest(new Map()); + expect(next.account.id).toBe("b"); + + first!.release(); + next.release(); + }); + + it("uses active session count after in-flight load ties", () => { + const pool = new TokenPool([makeAccount("a"), makeAccount("b")]); + const lease = pool.acquireBest(new Map([["a", 2], ["b", 1]])); + expect(lease.account.id).toBe("b"); + lease.release(); + }); + + it("uses relative rate-limit headroom after load and binding ties", () => { + const a = makeAccount("a"); + const b = makeAccount("b"); + a.sessionLimitPercent = 80; + a.rateLimits.fiveHourUtil = 0.6; + b.sessionLimitPercent = 80; + b.rateLimits.fiveHourUtil = 0.2; + const pool = new TokenPool([a, b]); + + const lease = pool.acquireBest(new Map([["a", 1], ["b", 1]])); + expect(lease.account.id).toBe("b"); + lease.release(); + }); + + it("uses the worse of five-hour and seven-day headroom ratios", () => { + const a = makeAccount("a"); + const b = makeAccount("b"); + a.rateLimits.fiveHourUtil = 0.1; + a.rateLimits.sevenDayUtil = 0.7; + b.rateLimits.fiveHourUtil = 0.5; + b.rateLimits.sevenDayUtil = 0.2; + const pool = new TokenPool([a, b]); + + const lease = pool.acquireBest(new Map()); + expect(lease.account.id).toBe("b"); + lease.release(); + }); + + it("rotates exact ties only after selection", () => { + const pool = new TokenPool([ + makeAccount("a"), + makeAccount("b"), + makeAccount("c"), + ]); + + const ids = Array.from({ length: 4 }, () => { + const lease = pool.acquireBest(new Map()); + lease.release(); + return lease.account.id; + }); + + expect(ids).toEqual(["a", "b", "c", "a"]); + }); + + it("allows a valid sticky acquisition despite existing in-flight work", () => { + let now = 2_000; + const pool = new TokenPool([makeAccount("a")], { now: () => now }); + const first = pool.tryAcquire("a"); + now = 2_500; + const second = pool.tryAcquire("a"); + + expect(first).not.toBeNull(); + expect(second).not.toBeNull(); + expect(pool.getInFlight("a")).toBe(2); + expect(pool.getAll()[0].requestCount).toBe(2); + expect(pool.getAll()[0].lastUsed).toBe(2_500); + + first!.release(); + second!.release(); + }); + + it("chooses the lowest in-flight account for all-unavailable fallback", () => { + const a = makeAccount("a"); + const b = makeAccount("b"); + const pool = new TokenPool([a, b]); + const first = pool.tryAcquire("a"); + a.enabled = false; + b.enabled = false; + + const fallback = pool.acquireBest(new Map()); + expect(fallback.account.id).toBe("b"); + expect(fallback.fallback).toBe(true); + + first!.release(); + fallback.release(); + }); +}); + +describe("TokenPool — timestamp cooldown", () => { + it("keeps an account unavailable until its timestamp cooldown expires", () => { + let now = 1_000; + const pool = new TokenPool( + [makeAccount("a"), makeAccount("b")], + { now: () => now }, + ); + pool.setCooldown("a", 500); + expect(pool.isEligible("a")).toBe(false); + expect(pool.isCoolingDown("a")).toBe(true); + + now = 1_500; + expect(pool.isEligible("a")).toBe(true); + expect(pool.isCoolingDown("a")).toBe(false); + }); + + it("treats zero-percent caps as ineligible", () => { + const a = makeAccount("a"); + a.sessionLimitPercent = 0; + const pool = new TokenPool([a, makeAccount("b")]); + + expect(pool.isEligible("a")).toBe(false); + expect(pool.tryAcquire("a")).toBeNull(); + }); +}); + describe("TokenPool — unhealthy accounts", () => { it("skips unhealthy accounts", () => { const pool = new TokenPool([ @@ -129,6 +265,18 @@ describe("TokenPool — stats", () => { expect(s.sessionLimitPercent).toBe(100); expect(s.weeklyLimitPercent).toBe(100); }); + + it("getStats() includes current in-flight and cooldown state", () => { + const pool = new TokenPool([makeAccount("a")]); + const lease = pool.acquireBest(new Map()); + pool.setCooldown("a", 1_000); + + const s = pool.getStats()[0]; + expect(s.inFlightRequests).toBe(1); + expect(s.coolingDown).toBe(true); + + lease.release(); + }); }); describe("TokenPool — mutation API", () => { diff --git a/src/proxy/token-pool.ts b/src/proxy/token-pool.ts index 66da310..6fb7e97 100644 --- a/src/proxy/token-pool.ts +++ b/src/proxy/token-pool.ts @@ -47,8 +47,12 @@ function limitingReset(a: Account): number { * window expires, `status` flips back to `"allowed"`. The callback fires * once per recovery so the dashboard can surface it. */ -function clearExpiredCooldown(a: Account, onExpired?: (a: Account) => void): void { - const nowSec = Math.floor(Date.now() / 1000); +function clearExpiredRateLimitWindows( + a: Account, + nowMs: number, + onExpired?: (a: Account) => void, +): void { + const nowSec = Math.floor(nowMs / 1000); const r = a.rateLimits; let changed = false; let recovered = false; @@ -79,7 +83,7 @@ function clearExpiredCooldown(a: Account, onExpired?: (a: Account) => void): voi } } - if (changed) r.lastUpdated = Date.now(); + if (changed) r.lastUpdated = nowMs; if (recovered && onExpired) onExpired(a); } @@ -102,16 +106,29 @@ export interface AccountPatch { weeklyLimitPercent?: number; } +export interface AccountLease { + readonly account: Account; + readonly fallback: boolean; + release(): void; +} + +export interface TokenPoolOptions { + now?: () => number; +} + export class TokenPool { - private accounts: Account[]; + private readonly inFlight = new Map(); + private readonly cooldownUntil = new Map(); + private readonly now: () => number; private currentIndex = 0; - constructor(accounts: Account[]) { - this.accounts = accounts; + constructor(private readonly accounts: Account[], options: TokenPoolOptions = {}) { + this.now = options.now ?? Date.now; } /** - * Round-robin selection among accounts that are: + * Compatibility wrapper for request sites that do not yet retain leases. + * Selection considers accounts that are: * • healthy * • not busy * • not rate-limited by Anthropic @@ -125,57 +142,174 @@ export class TokenPool { * ban that would leave Claude Code with no working account. The fallback * is logged via the optional onCapBypass callback so the dashboard can * surface it instead of silently exceeding the cap. - * 3. accounts[0] as a last resort (only if every account is unhealthy). + * 3. Any account as a last resort (only if every account is unhealthy). + * + * Fallback sets prefer the lowest in-flight load, then the earliest reset. * * Throws `EmptyPoolError` when there are no accounts at all — callers in * the request path should map this to a 503. The DELETE endpoint guards * against this state by refusing to remove the last account. */ getNext(): Account { + const lease = this.acquireBest(new Map()); + lease.release(); + return lease.account; + } + + /** + * Acquire the best eligible account using load, session affinity pressure, + * rate-limit headroom, and a rotating tie-break, in that order. + */ + acquireBest(activeSessions: ReadonlyMap): AccountLease { if (this.accounts.length === 0) { throw new EmptyPoolError("token pool is empty — add an account first"); } - // Sweep expired cooldowns before filtering — otherwise rate_limited accounts - // whose reset window has passed would never re-enter rotation. - for (const a of this.accounts) clearExpiredCooldown(a, this.onCooldownExpired); - - const available = this.accounts.filter(a => - a.healthy && - !a.busy && - a.rateLimits.status !== "rate_limited" && - isUsable(a) - ); - - if (available.length === 0) { - const healthyUsable = this.accounts.filter(a => a.healthy && isUsable(a)); - if (healthyUsable.length > 0) { - return healthyUsable.reduce((best, a) => - earliestReset(a) < earliestReset(best) ? a : best - ); - } - const healthy = this.accounts.filter(a => a.healthy); - if (healthy.length === 0) { - return this.accounts[0]; - } - // All healthy accounts are either busy, rate-limited, or over user caps. - // Fall back to the one that'll reset soonest — see docstring. Notify - // the listener so the bypass becomes visible in the dashboard. - const fallback = healthy.reduce((best, a) => - earliestReset(a) < earliestReset(best) ? a : best + this.sweepExpiredCooldowns(); + const eligible = this.accounts.filter(account => this.isEligibleWithoutSweep(account)); + + if (eligible.length > 0) { + const account = this.selectEligible(eligible, activeSessions); + this.advanceCursor(account); + return this.createLease(account, false); + } + + const healthyUsable = this.accounts.filter(account => account.healthy && isUsable(account)); + const healthy = this.accounts.filter(account => account.healthy); + const fallbackCandidates = healthyUsable.length > 0 + ? healthyUsable + : healthy.length > 0 + ? healthy + : this.accounts; + const account = this.selectFallback(fallbackCandidates); + this.advanceCursor(account); + if (overUserCap(account)) this.onCapBypass?.(account); + return this.createLease(account, true); + } + + /** Acquire a specific account for an existing sticky session. */ + tryAcquire(accountId: string): AccountLease | null { + const account = this.findById(accountId); + if (!account) return null; + clearExpiredRateLimitWindows(account, this.now(), this.onCooldownExpired); + if (!this.isEligibleWithoutSweep(account)) return null; + return this.createLease(account, false); + } + + isEligible(accountId: string): boolean { + const account = this.findById(accountId); + if (!account) return false; + clearExpiredRateLimitWindows(account, this.now(), this.onCooldownExpired); + return this.isEligibleWithoutSweep(account); + } + + getInFlight(accountId: string): number { + return this.inFlight.get(accountId) ?? 0; + } + + setCooldown(accountId: string, durationMs: number): void { + if (!this.findById(accountId)) return; + if (durationMs <= 0) { + this.cooldownUntil.delete(accountId); + return; + } + this.cooldownUntil.set(accountId, this.now() + durationMs); + } + + isCoolingDown(accountId: string): boolean { + const until = this.cooldownUntil.get(accountId); + if (until === undefined) return false; + if (this.now() < until) return true; + this.cooldownUntil.delete(accountId); + return false; + } + + private isEligibleWithoutSweep(account: Account): boolean { + return account.healthy && + !account.busy && + !this.isCoolingDown(account.id) && + account.rateLimits.status !== "rate_limited" && + isUsable(account); + } + + private selectEligible( + candidates: Account[], + activeSessions: ReadonlyMap, + ): Account { + return candidates.reduce((best, account) => { + const comparison = this.compareTuple( + [ + this.getInFlight(account.id), + activeSessions.get(account.id) ?? 0, + this.headroomScore(account), + this.circularDistance(account), + ], + [ + this.getInFlight(best.id), + activeSessions.get(best.id) ?? 0, + this.headroomScore(best), + this.circularDistance(best), + ], ); - const someCapped = this.accounts.some(a => a.healthy && overUserCap(a)); - if (someCapped && this.onCapBypass) { - this.onCapBypass(fallback); - } - return fallback; + return comparison < 0 ? account : best; + }); + } + + private selectFallback(candidates: Account[]): Account { + return candidates.reduce((best, account) => { + const comparison = this.compareTuple( + [this.getInFlight(account.id), earliestReset(account), this.circularDistance(account)], + [this.getInFlight(best.id), earliestReset(best), this.circularDistance(best)], + ); + return comparison < 0 ? account : best; + }); + } + + private headroomScore(account: Account): number { + const fiveHourCap = account.sessionLimitPercent / 100; + const sevenDayCap = account.weeklyLimitPercent / 100; + const fiveHourUtil = Number.isFinite(account.rateLimits.fiveHourUtil) + ? Math.max(0, account.rateLimits.fiveHourUtil) + : 0; + const sevenDayUtil = Number.isFinite(account.rateLimits.sevenDayUtil) + ? Math.max(0, account.rateLimits.sevenDayUtil) + : 0; + return Math.max(fiveHourUtil / fiveHourCap, sevenDayUtil / sevenDayCap); + } + + private circularDistance(account: Account): number { + const index = this.accounts.indexOf(account); + return (index - this.currentIndex + this.accounts.length) % this.accounts.length; + } + + private compareTuple(left: number[], right: number[]): number { + for (let i = 0; i < left.length; i++) { + if (left[i] !== right[i]) return left[i] - right[i]; } + return 0; + } + + private advanceCursor(account: Account): void { + const index = this.accounts.indexOf(account); + this.currentIndex = (index + 1) % this.accounts.length; + } - const account = available[this.currentIndex % available.length]; - this.currentIndex = (this.currentIndex + 1) % available.length; + private createLease(account: Account, fallback: boolean): AccountLease { + this.inFlight.set(account.id, this.getInFlight(account.id) + 1); account.requestCount++; - account.lastUsed = Date.now(); - return account; + account.lastUsed = this.now(); + let released = false; + return { + account, + fallback, + release: () => { + if (released) return; + released = true; + const remaining = Math.max(0, this.getInFlight(account.id) - 1); + if (remaining === 0) this.inFlight.delete(account.id); + else this.inFlight.set(account.id, remaining); + }, + }; } /** Optional listener fired when a request is routed to a capped account @@ -192,7 +326,11 @@ export class TokenPool { * poll loop so the UI reflects recovery without waiting for a new request. */ sweepExpiredCooldowns(): void { - for (const a of this.accounts) clearExpiredCooldown(a, this.onCooldownExpired); + const now = this.now(); + for (const a of this.accounts) { + clearExpiredRateLimitWindows(a, now, this.onCooldownExpired); + this.isCoolingDown(a.id); + } } getAll(): Account[] { @@ -208,6 +346,8 @@ export class TokenPool { id: a.id, healthy: a.healthy, busy: a.busy, + inFlightRequests: this.getInFlight(a.id), + coolingDown: this.isCoolingDown(a.id), requestCount: a.requestCount, errorCount: a.errorCount, expiresInMs: a.tokens.expiresAt - Date.now(), @@ -294,6 +434,8 @@ export class TokenPool { const idx = this.accounts.findIndex(a => a.id === id); if (idx === -1) return false; this.accounts.splice(idx, 1); + this.inFlight.delete(id); + this.cooldownUntil.delete(id); if (this.accounts.length > 0) { this.currentIndex = this.currentIndex % this.accounts.length; } else { From a19d9029b1de10a48183d110c3cb51e5297e9589 Mon Sep 17 00:00:00 2001 From: timo <44401485+Timo972@users.noreply.github.com> Date: Mon, 20 Jul 2026 18:48:44 +0200 Subject: [PATCH 06/17] fix: harden lease and cooldown ownership --- src/__tests__/token-pool.test.ts | 37 ++++++++++++++++++++++++++++++++ src/proxy/token-pool.ts | 13 ++++++----- 2 files changed, 45 insertions(+), 5 deletions(-) diff --git a/src/__tests__/token-pool.test.ts b/src/__tests__/token-pool.test.ts index 42909c3..c88307a 100644 --- a/src/__tests__/token-pool.test.ts +++ b/src/__tests__/token-pool.test.ts @@ -164,6 +164,30 @@ describe("TokenPool — request leases", () => { first!.release(); fallback.release(); }); + + it("does not let an old lease release a replacement account incarnation", () => { + const pool = new TokenPool([makeAccount("a")]); + const oldLease = pool.tryAcquire("a"); + expect(oldLease).not.toBeNull(); + + pool.removeAccount("a"); + pool.addAccount({ + id: "a", + accessToken: "sk-ant-oat01-replacement", + refreshToken: "sk-ant-ort01-replacement", + expiresAt: Date.now() + 3_600_000, + scopes: ["user:inference"], + }); + const replacementLease = pool.tryAcquire("a"); + expect(replacementLease).not.toBeNull(); + expect(pool.getInFlight("a")).toBe(1); + + oldLease!.release(); + expect(pool.getInFlight("a")).toBe(1); + + replacementLease!.release(); + expect(pool.getInFlight("a")).toBe(0); + }); }); describe("TokenPool — timestamp cooldown", () => { @@ -190,6 +214,19 @@ describe("TokenPool — timestamp cooldown", () => { expect(pool.isEligible("a")).toBe(false); expect(pool.tryAcquire("a")).toBeNull(); }); + + it("does not shorten an existing cooldown with a later shorter cooldown", () => { + let now = 1_000; + const pool = new TokenPool([makeAccount("a")], { now: () => now }); + pool.setCooldown("a", 60_000); + pool.setCooldown("a", 30_000); + + now = 32_000; + expect(pool.isCoolingDown("a")).toBe(true); + + now = 61_000; + expect(pool.isCoolingDown("a")).toBe(false); + }); }); describe("TokenPool — unhealthy accounts", () => { diff --git a/src/proxy/token-pool.ts b/src/proxy/token-pool.ts index 6fb7e97..c721846 100644 --- a/src/proxy/token-pool.ts +++ b/src/proxy/token-pool.ts @@ -209,11 +209,10 @@ export class TokenPool { setCooldown(accountId: string, durationMs: number): void { if (!this.findById(accountId)) return; - if (durationMs <= 0) { - this.cooldownUntil.delete(accountId); - return; - } - this.cooldownUntil.set(accountId, this.now() + durationMs); + if (!Number.isFinite(durationMs) || durationMs <= 0) return; + const proposedExpiry = this.now() + durationMs; + const existingExpiry = this.cooldownUntil.get(accountId) ?? 0; + this.cooldownUntil.set(accountId, Math.max(existingExpiry, proposedExpiry)); } isCoolingDown(accountId: string): boolean { @@ -305,6 +304,10 @@ export class TokenPool { release: () => { if (released) return; released = true; + // IDs may be reused after an account is removed. A lease belongs to + // the exact Account instance it acquired and must never decrement a + // replacement account's load counter. + if (this.findById(account.id) !== account) return; const remaining = Math.max(0, this.getInFlight(account.id) - 1); if (remaining === 0) this.inFlight.delete(account.id); else this.inFlight.set(account.id, remaining); From 689ab74cc7625c4e803f5d31fcbdb3358c8ecaa3 Mon Sep 17 00:00:00 2001 From: timo <44401485+Timo972@users.noreply.github.com> Date: Mon, 20 Jul 2026 18:54:08 +0200 Subject: [PATCH 07/17] feat: add sticky Claude session routing --- src/__tests__/session-router.test.ts | 257 +++++++++++++++++++++++++++ src/proxy/session-router.ts | 175 ++++++++++++++++++ 2 files changed, 432 insertions(+) create mode 100644 src/__tests__/session-router.test.ts create mode 100644 src/proxy/session-router.ts diff --git a/src/__tests__/session-router.test.ts b/src/__tests__/session-router.test.ts new file mode 100644 index 0000000..c7ba428 --- /dev/null +++ b/src/__tests__/session-router.test.ts @@ -0,0 +1,257 @@ +import { describe, expect, it, vi } from "vitest"; +import { SessionRouter, normalizeSessionId } from "../proxy/session-router.js"; +import { TokenPool } from "../proxy/token-pool.js"; +import type { Account } from "../proxy/types.js"; +import { DEFAULT_RATE_LIMITS } from "../proxy/types.js"; + +function makeAccount(id: string): Account { + return { + id, + tokens: { + accessToken: `sk-ant-oat01-${id}`, + refreshToken: `sk-ant-ort01-${id}`, + expiresAt: Date.now() + 3_600_000, + scopes: ["user:inference", "user:profile"], + }, + healthy: true, + busy: false, + requestCount: 0, + errorCount: 0, + lastUsed: 0, + lastRefresh: 0, + consecutiveErrors: 0, + rateLimits: { ...DEFAULT_RATE_LIMITS }, + enabled: true, + sessionLimitPercent: 100, + weeklyLimitPercent: 100, + }; +} + +describe("session ID normalization", () => { + it("normalizes one bounded string header", () => { + expect(normalizeSessionId(" session-a ")).toBe("session-a"); + expect(normalizeSessionId(" ")).toBeUndefined(); + expect(normalizeSessionId(["a", "b"])).toBeUndefined(); + expect(normalizeSessionId(42)).toBeUndefined(); + expect(normalizeSessionId(` ${"é".repeat(128)} `)).toBe("é".repeat(128)); + expect(normalizeSessionId("é".repeat(129))).toBeUndefined(); + expect(normalizeSessionId("a".repeat(257))).toBeUndefined(); + }); +}); + +describe("SessionRouter", () => { + it("keeps repeated requests from one session on its account", () => { + const pool = new TokenPool([makeAccount("a"), makeAccount("b")]); + const router = new SessionRouter(pool); + + const first = router.acquire("session-a"); + const second = router.acquire("session-a"); + + expect(first.account.id).toBe(second.account.id); + expect(first.reason).toBe("new-session"); + expect(second.reason).toBe("sticky"); + first.release(); + second.release(); + }); + + it("binds simultaneous new sessions to separate idle accounts", () => { + const pool = new TokenPool([makeAccount("a"), makeAccount("b")]); + const router = new SessionRouter(pool); + + const first = router.acquire("session-a"); + const second = router.acquire("session-b"); + + expect(first.account.id).toBe("a"); + expect(second.account.id).toBe("b"); + expect(router.getActiveSessionCount("a")).toBe(1); + expect(router.getActiveSessionCount("b")).toBe(1); + expect(router.getBindingCount()).toBe(2); + first.release(); + second.release(); + }); + + it("load-balances unscoped requests without creating bindings", () => { + const pool = new TokenPool([makeAccount("a"), makeAccount("b")]); + const router = new SessionRouter(pool); + + const first = router.acquire(undefined); + const second = router.acquire(["invalid"]); + + expect(first.account.id).toBe("a"); + expect(second.account.id).toBe("b"); + expect(first.reason).toBe("unscoped"); + expect(second.reason).toBe("unscoped"); + expect(first.sessionId).toBeUndefined(); + expect(router.getBindingCount()).toBe(0); + first.release(); + second.release(); + }); + + it("keeps a valid sticky binding despite lower load elsewhere", () => { + const pool = new TokenPool([makeAccount("a"), makeAccount("b")]); + const router = new SessionRouter(pool); + const first = router.acquire("session-a"); + + expect(pool.getInFlight("a")).toBe(1); + expect(pool.getInFlight("b")).toBe(0); + + const sticky = router.acquire("session-a"); + expect(sticky.account.id).toBe("a"); + expect(sticky.reason).toBe("sticky"); + + first.release(); + sticky.release(); + }); + + it.each([ + ["disabled", (account: Account, pool: TokenPool) => { account.enabled = false; }], + ["unhealthy", (account: Account, pool: TokenPool) => { account.healthy = false; }], + ["capped", (account: Account, pool: TokenPool) => { + account.sessionLimitPercent = 50; + account.rateLimits.fiveHourUtil = 0.5; + }], + ["cooling down", (account: Account, pool: TokenPool) => { pool.setCooldown(account.id, 60_000); }], + ["removed", (account: Account, pool: TokenPool) => { pool.removeAccount(account.id); }], + ])("fails over when the bound account is %s", (_label, makeUnavailable) => { + const a = makeAccount("a"); + const pool = new TokenPool([a, makeAccount("b")]); + const router = new SessionRouter(pool); + router.acquire("session-a").release(); + + makeUnavailable(a, pool); + const failover = router.acquire("session-a"); + + expect(failover.account.id).toBe("b"); + expect(failover.reason).toBe("failover"); + expect(router.getActiveSessionCount("a")).toBe(0); + expect(router.getActiveSessionCount("b")).toBe(1); + failover.release(); + }); + + it("does not let an old response invalidate a newer binding", () => { + const a = makeAccount("a"); + const pool = new TokenPool([a, makeAccount("b")]); + const router = new SessionRouter(pool); + router.acquire("session-a").release(); + a.enabled = false; + const rebound = router.acquire("session-a"); + rebound.release(); + + expect(router.invalidate("session-a", "a")).toBe(false); + expect(router.getActiveSessionCount("b")).toBe(1); + + const sticky = router.acquire("session-a"); + expect(sticky.account.id).toBe("b"); + expect(sticky.reason).toBe("sticky"); + sticky.release(); + + expect(router.invalidate("session-a", "b")).toBe(true); + expect(router.getBindingCount()).toBe(0); + }); + + it("invalidates every binding for one account and repairs counts", () => { + const pool = new TokenPool([makeAccount("a"), makeAccount("b")]); + const router = new SessionRouter(pool); + router.acquire("session-1").release(); + router.acquire("session-2").release(); + router.acquire("session-3").release(); + + expect(router.getActiveSessionCount("a")).toBe(2); + expect(router.getActiveSessionCount("b")).toBe(1); + expect(router.invalidateAccount("a")).toBe(2); + expect(router.getActiveSessionCount("a")).toBe(0); + expect(router.getActiveSessionCount("b")).toBe(1); + expect(router.getBindingCount()).toBe(1); + }); + + it("refreshes last-seen on access and expires at one hour of inactivity", () => { + const hour = 60 * 60 * 1000; + let now = 100; + const pool = new TokenPool([makeAccount("a"), makeAccount("b")], { now: () => now }); + const router = new SessionRouter(pool, { now: () => now }); + router.acquire("session-a").release(); + + now += hour - 1; + const refreshed = router.acquire("session-a"); + expect(refreshed.reason).toBe("sticky"); + refreshed.release(); + + now += hour - 1; + const stillSticky = router.acquire("session-a"); + expect(stillSticky.reason).toBe("sticky"); + stillSticky.release(); + + now += hour; + const expired = router.acquire("session-a"); + expect(expired.reason).toBe("new-session"); + expect(router.getBindingCount()).toBe(1); + expired.release(); + }); + + it("sweeps expired bindings before applying the capacity limit", () => { + let now = 0; + const pool = new TokenPool([makeAccount("a"), makeAccount("b")], { now: () => now }); + const router = new SessionRouter(pool, { now: () => now, ttlMs: 10, maxEntries: 2 }); + router.acquire("expired").release(); + now = 5; + router.acquire("live").release(); + + now = 10; + router.acquire("new").release(); + + expect(router.getBindingCount()).toBe(2); + expect(router.invalidate("expired")).toBe(false); + expect(router.invalidate("live")).toBe(true); + expect(router.invalidate("new")).toBe(true); + }); + + it("evicts the least recently used binding with insertion order as tie-break", () => { + let now = 0; + const pool = new TokenPool([makeAccount("a"), makeAccount("b")], { now: () => now }); + const router = new SessionRouter(pool, { now: () => now, maxEntries: 2 }); + router.acquire("oldest-tie").release(); + router.acquire("newer-tie").release(); + + router.acquire("third").release(); + + expect(router.getBindingCount()).toBe(2); + expect(router.invalidate("oldest-tie")).toBe(false); + expect(router.invalidate("newer-tie")).toBe(true); + expect(router.invalidate("third")).toBe(true); + }); + + it("updates recency so a recently used binding survives LRU eviction", () => { + let now = 0; + const pool = new TokenPool([makeAccount("a"), makeAccount("b")], { now: () => now }); + const router = new SessionRouter(pool, { now: () => now, maxEntries: 2 }); + router.acquire("first").release(); + now = 1; + router.acquire("second").release(); + now = 2; + router.acquire("first").release(); + now = 3; + router.acquire("third").release(); + + expect(router.invalidate("first")).toBe(true); + expect(router.invalidate("second")).toBe(false); + expect(router.invalidate("third")).toBe(true); + }); + + it("does not log session IDs or expose bindings through its public results", () => { + const log = vi.spyOn(console, "log").mockImplementation(() => undefined); + const warn = vi.spyOn(console, "warn").mockImplementation(() => undefined); + const router = new SessionRouter(new TokenPool([makeAccount("a")])); + + const lease = router.acquire("private-session-id"); + + expect(lease.sessionId).toBe("private-session-id"); + expect(Object.keys(lease).sort()).toEqual( + ["account", "fallback", "reason", "release", "sessionId"].sort(), + ); + expect(log).not.toHaveBeenCalled(); + expect(warn).not.toHaveBeenCalled(); + lease.release(); + log.mockRestore(); + warn.mockRestore(); + }); +}); diff --git a/src/proxy/session-router.ts b/src/proxy/session-router.ts new file mode 100644 index 0000000..154747e --- /dev/null +++ b/src/proxy/session-router.ts @@ -0,0 +1,175 @@ +import type { AccountLease } from "./token-pool.js"; +import { TokenPool } from "./token-pool.js"; + +const DEFAULT_TTL_MS = 60 * 60 * 1000; +const DEFAULT_MAX_ENTRIES = 10_000; +const MAX_SESSION_ID_BYTES = 256; + +export type RouteReason = "sticky" | "new-session" | "unscoped" | "failover"; + +export interface RoutedAccountLease extends AccountLease { + readonly reason: RouteReason; + readonly sessionId?: string; +} + +export interface SessionRouterOptions { + now?: () => number; + ttlMs?: number; + maxEntries?: number; +} + +interface SessionBinding { + accountId: string; + lastSeen: number; +} + +/** + * Normalize Claude Code's session header without retaining malformed or + * unexpectedly large values. Arrays are rejected rather than choosing one + * value because a session must have exactly one unambiguous identity. + */ +export function normalizeSessionId(value: unknown): string | undefined { + if (typeof value !== "string") return undefined; + const normalized = value.trim(); + if (normalized.length === 0) return undefined; + if (Buffer.byteLength(normalized, "utf8") > MAX_SESSION_ID_BYTES) return undefined; + return normalized; +} + +/** + * In-memory session affinity for Claude requests. Bindings and their counts + * are intentionally process-local and are never exposed for persistence or + * diagnostics. + */ +export class SessionRouter { + private readonly bindings = new Map(); + private readonly activeSessionCounts = new Map(); + private readonly now: () => number; + private readonly ttlMs: number; + private readonly maxEntries: number; + + constructor( + private readonly pool: TokenPool, + options: SessionRouterOptions = {}, + ) { + this.now = options.now ?? Date.now; + this.ttlMs = options.ttlMs ?? DEFAULT_TTL_MS; + this.maxEntries = options.maxEntries ?? DEFAULT_MAX_ENTRIES; + + if (!Number.isFinite(this.ttlMs) || this.ttlMs <= 0) { + throw new RangeError("session binding TTL must be a positive finite number"); + } + if (!Number.isInteger(this.maxEntries) || this.maxEntries <= 0) { + throw new RangeError("session binding capacity must be a positive integer"); + } + } + + acquire(sessionHeader: unknown): RoutedAccountLease { + const sessionId = normalizeSessionId(sessionHeader); + const now = this.now(); + this.sweepExpiredBindings(now); + + if (!sessionId) { + return this.wrap(this.pool.acquireBest(this.activeSessionCounts), "unscoped"); + } + + const existing = this.bindings.get(sessionId); + if (existing) { + const stickyLease = this.pool.tryAcquire(existing.accountId); + if (stickyLease) { + existing.lastSeen = now; + return this.wrap(stickyLease, "sticky", sessionId); + } + this.removeBinding(sessionId); + } + + const lease = this.pool.acquireBest(this.activeSessionCounts); + this.insertBinding(sessionId, lease.account.id, now); + return this.wrap(lease, existing ? "failover" : "new-session", sessionId); + } + + /** + * Remove a binding only if it still belongs to the expected account. This + * protects a new failover binding from a late response on the old account. + */ + invalidate(sessionHeader: unknown, expectedAccountId?: string): boolean { + const sessionId = normalizeSessionId(sessionHeader); + if (!sessionId) return false; + const binding = this.bindings.get(sessionId); + if (!binding) return false; + if (expectedAccountId !== undefined && binding.accountId !== expectedAccountId) { + return false; + } + return this.removeBinding(sessionId); + } + + invalidateAccount(accountId: string): number { + let removed = 0; + for (const [sessionId, binding] of this.bindings) { + if (binding.accountId !== accountId) continue; + if (this.removeBinding(sessionId)) removed++; + } + return removed; + } + + getActiveSessionCount(accountId: string): number { + return this.activeSessionCounts.get(accountId) ?? 0; + } + + getBindingCount(): number { + return this.bindings.size; + } + + private wrap( + lease: AccountLease, + reason: RouteReason, + sessionId?: string, + ): RoutedAccountLease { + return { + account: lease.account, + fallback: lease.fallback, + release: lease.release, + reason, + ...(sessionId === undefined ? {} : { sessionId }), + }; + } + + private insertBinding(sessionId: string, accountId: string, lastSeen: number): void { + if (this.bindings.size >= this.maxEntries) this.evictLeastRecentlyUsed(); + this.bindings.set(sessionId, { accountId, lastSeen }); + this.activeSessionCounts.set(accountId, this.getActiveSessionCount(accountId) + 1); + } + + private removeBinding(sessionId: string): boolean { + const binding = this.bindings.get(sessionId); + if (!binding) return false; + this.bindings.delete(sessionId); + + const remaining = this.getActiveSessionCount(binding.accountId) - 1; + if (remaining <= 0) this.activeSessionCounts.delete(binding.accountId); + else this.activeSessionCounts.set(binding.accountId, remaining); + return true; + } + + private sweepExpiredBindings(now: number): void { + for (const [sessionId, binding] of this.bindings) { + if (now - binding.lastSeen >= this.ttlMs) this.removeBinding(sessionId); + } + } + + private evictLeastRecentlyUsed(): void { + let oldestSessionId: string | undefined; + let oldestLastSeen = Infinity; + + // Map iteration order is insertion order, so equal timestamps retain the + // first entry as the deterministic eviction candidate. + for (const [sessionId, binding] of this.bindings) { + if (binding.lastSeen < oldestLastSeen) { + oldestSessionId = sessionId; + oldestLastSeen = binding.lastSeen; + } + } + + if (oldestSessionId !== undefined) this.removeBinding(oldestSessionId); + } +} From 2ec9e2311c0d56d4cd05c8908bcc9f90adaa3039 Mon Sep 17 00:00:00 2001 From: timo <44401485+Timo972@users.noreply.github.com> Date: Mon, 20 Jul 2026 18:57:38 +0200 Subject: [PATCH 08/17] fix: expire idle session counts on read --- src/__tests__/session-router.test.ts | 22 ++++++++++++++++++++++ src/proxy/session-router.ts | 12 +++++++++--- 2 files changed, 31 insertions(+), 3 deletions(-) diff --git a/src/__tests__/session-router.test.ts b/src/__tests__/session-router.test.ts index c7ba428..c4e2a3a 100644 --- a/src/__tests__/session-router.test.ts +++ b/src/__tests__/session-router.test.ts @@ -188,6 +188,28 @@ describe("SessionRouter", () => { expired.release(); }); + it("sweeps expired bindings when reading the total binding count", () => { + let now = 0; + const pool = new TokenPool([makeAccount("a")], { now: () => now }); + const router = new SessionRouter(pool, { now: () => now, ttlMs: 10 }); + router.acquire("session-a").release(); + + now = 10; + + expect(router.getBindingCount()).toBe(0); + }); + + it("sweeps expired bindings when reading an account's active session count", () => { + let now = 0; + const pool = new TokenPool([makeAccount("a")], { now: () => now }); + const router = new SessionRouter(pool, { now: () => now, ttlMs: 10 }); + router.acquire("session-a").release(); + + now = 10; + + expect(router.getActiveSessionCount("a")).toBe(0); + }); + it("sweeps expired bindings before applying the capacity limit", () => { let now = 0; const pool = new TokenPool([makeAccount("a"), makeAccount("b")], { now: () => now }); diff --git a/src/proxy/session-router.ts b/src/proxy/session-router.ts index 154747e..195ad52 100644 --- a/src/proxy/session-router.ts +++ b/src/proxy/session-router.ts @@ -113,10 +113,12 @@ export class SessionRouter { } getActiveSessionCount(accountId: string): number { - return this.activeSessionCounts.get(accountId) ?? 0; + this.sweepExpiredBindings(this.now()); + return this.getRawActiveSessionCount(accountId); } getBindingCount(): number { + this.sweepExpiredBindings(this.now()); return this.bindings.size; } @@ -137,7 +139,7 @@ export class SessionRouter { private insertBinding(sessionId: string, accountId: string, lastSeen: number): void { if (this.bindings.size >= this.maxEntries) this.evictLeastRecentlyUsed(); this.bindings.set(sessionId, { accountId, lastSeen }); - this.activeSessionCounts.set(accountId, this.getActiveSessionCount(accountId) + 1); + this.activeSessionCounts.set(accountId, this.getRawActiveSessionCount(accountId) + 1); } private removeBinding(sessionId: string): boolean { @@ -145,12 +147,16 @@ export class SessionRouter { if (!binding) return false; this.bindings.delete(sessionId); - const remaining = this.getActiveSessionCount(binding.accountId) - 1; + const remaining = this.getRawActiveSessionCount(binding.accountId) - 1; if (remaining <= 0) this.activeSessionCounts.delete(binding.accountId); else this.activeSessionCounts.set(binding.accountId, remaining); return true; } + private getRawActiveSessionCount(accountId: string): number { + return this.activeSessionCounts.get(accountId) ?? 0; + } + private sweepExpiredBindings(now: number): void { for (const [sessionId, binding] of this.bindings) { if (now - binding.lastSeen >= this.ttlMs) this.removeBinding(sessionId); From ce9631c395d12d250a7dd86db2d6d1b90281522a Mon Sep 17 00:00:00 2001 From: timo <44401485+Timo972@users.noreply.github.com> Date: Mon, 20 Jul 2026 19:07:47 +0200 Subject: [PATCH 09/17] fix: route concurrent streams by Claude session --- src/__tests__/anthropic-proxy.test.ts | 240 ++++++++++++++++++++++++++ src/__tests__/lease-lifecycle.test.ts | 155 +++++++++++++++++ src/proxy/anthropic-proxy.ts | 26 +++ src/proxy/lease-lifecycle.ts | 104 +++++++++++ src/proxy/server.ts | 106 ++++++++---- 5 files changed, 594 insertions(+), 37 deletions(-) create mode 100644 src/__tests__/anthropic-proxy.test.ts create mode 100644 src/__tests__/lease-lifecycle.test.ts create mode 100644 src/proxy/anthropic-proxy.ts create mode 100644 src/proxy/lease-lifecycle.ts diff --git a/src/__tests__/anthropic-proxy.test.ts b/src/__tests__/anthropic-proxy.test.ts new file mode 100644 index 0000000..ff59e99 --- /dev/null +++ b/src/__tests__/anthropic-proxy.test.ts @@ -0,0 +1,240 @@ +import { createServer, request, type Server } from "node:http"; +import type { AddressInfo } from "node:net"; +import express from "express"; +import { describe, expect, it, vi } from "vitest"; +import { createAnthropicProxy } from "../proxy/anthropic-proxy.js"; +import { applyUpstreamFailureRouting } from "../proxy/lease-lifecycle.js"; + +async function listen(server: Server): Promise { + await new Promise((resolve, reject) => { + server.once("error", reject); + server.listen(0, "127.0.0.1", () => { + server.off("error", reject); + resolve(); + }); + }); + return (server.address() as AddressInfo).port; +} + +async function close(server: Server): Promise { + await new Promise((resolve, reject) => { + server.close(error => error ? reject(error) : resolve()); + server.closeAllConnections(); + }); +} + +function deferred(): { promise: Promise; resolve: () => void } { + let resolve!: () => void; + const promise = new Promise(done => { resolve = done; }); + return { promise, resolve }; +} + +function collect(url: URL): Promise<{ + status: number; + headers: Record; + body: Buffer; +}> { + return new Promise((resolve, reject) => { + const req = request(url, response => { + const chunks: Buffer[] = []; + response.on("data", chunk => chunks.push(Buffer.from(chunk))); + response.on("end", () => resolve({ + status: response.statusCode ?? 0, + headers: response.headers, + body: Buffer.concat(chunks), + })); + response.on("error", reject); + }); + req.on("error", reject); + req.end(); + }); +} + +function startCollecting(url: URL): { + firstChunk: Promise; + completed: Promise; + hasCompleted: () => boolean; +} { + const first = deferred(); + let completed = false; + const body = new Promise((resolve, reject) => { + const req = request(url, response => { + const chunks: Buffer[] = []; + response.once("data", () => first.resolve()); + response.on("data", chunk => chunks.push(Buffer.from(chunk))); + response.on("end", () => { + completed = true; + resolve(Buffer.concat(chunks)); + }); + response.on("error", reject); + }); + req.on("error", reject); + req.end(); + }); + return { firstChunk: first.promise, completed: body, hasCompleted: () => completed }; +} + +describe("createAnthropicProxy", () => { + it("forwards deliberately split SSE bytes without inserting or removing events", async () => { + const chunks = [ + Buffer.from("event: message_start\nda"), + Buffer.from("ta: {\"type\":\"message_start\"}\n\n"), + Buffer.from("event: content_block_delta\ndata: {\"delta\":{\"text\":\"hello\"}}\n\n"), + Buffer.from("event: message_"), + Buffer.from("stop\ndata: {\"type\":\"message_stop\"}\n\n"), + ]; + const upstreamBody = Buffer.concat(chunks); + const observedChunks: Buffer[] = []; + let upstreamPath = ""; + const upstream = createServer((_req, res) => { + upstreamPath = _req.url ?? ""; + res.writeHead(200, { "content-type": "text/event-stream" }); + let index = 0; + const writeNext = () => { + if (index === chunks.length) { + res.end(); + return; + } + res.write(chunks[index++]); + setImmediate(writeNext); + }; + writeNext(); + }); + const upstreamPort = await listen(upstream); + + const app = express(); + app.use("/v1", createAnthropicProxy({ + target: `http://127.0.0.1:${upstreamPort}`, + timeoutMs: 2_000, + on: { + proxyRes: proxyResponse => { + proxyResponse.on("data", chunk => observedChunks.push(Buffer.from(chunk))); + }, + }, + })); + const downstream = createServer(app); + const downstreamPort = await listen(downstream); + + try { + const response = await collect(new URL(`http://127.0.0.1:${downstreamPort}/v1/messages`)); + + expect(response.status).toBe(200); + expect(response.headers["content-type"]).toContain("text/event-stream"); + expect(upstreamPath).toBe("/v1/messages"); + expect(Buffer.compare(response.body, upstreamBody)).toBe(0); + expect(Buffer.compare(Buffer.concat(observedChunks), upstreamBody)).toBe(0); + expect(response.body.toString("utf8").match(/event: message_stop/g)).toHaveLength(1); + } finally { + await close(downstream); + await close(upstream); + } + }); + + it("keeps concurrent SSE responses open until each upstream stream terminates", async () => { + const gates = new Map([ + ["/v1/one", deferred()], + ["/v1/two", deferred()], + ]); + const expected = new Map(); + const upstream = createServer(async (req, res) => { + const path = req.url ?? ""; + const gate = gates.get(path); + if (!gate) { + res.writeHead(404).end(); + return; + } + const prefix = Buffer.from(`event: ping\ndata: {\"stream\":\"${path}\"}\n\n`); + const suffix = Buffer.from("event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"); + expected.set(path, Buffer.concat([prefix, suffix])); + res.writeHead(200, { "content-type": "text/event-stream" }); + res.write(prefix); + await gate.promise; + res.end(suffix); + }); + const upstreamPort = await listen(upstream); + + const app = express(); + app.use("/v1", createAnthropicProxy({ + target: `http://127.0.0.1:${upstreamPort}`, + timeoutMs: 2_000, + on: {}, + })); + const downstream = createServer(app); + const downstreamPort = await listen(downstream); + + const one = startCollecting(new URL(`http://127.0.0.1:${downstreamPort}/v1/one`)); + const two = startCollecting(new URL(`http://127.0.0.1:${downstreamPort}/v1/two`)); + try { + await Promise.all([one.firstChunk, two.firstChunk]); + expect(one.hasCompleted()).toBe(false); + expect(two.hasCompleted()).toBe(false); + + gates.get("/v1/one")!.resolve(); + expect(await one.completed).toEqual(expected.get("/v1/one")); + expect(two.hasCompleted()).toBe(false); + + gates.get("/v1/two")!.resolve(); + expect(await two.completed).toEqual(expected.get("/v1/two")); + } finally { + gates.get("/v1/one")!.resolve(); + gates.get("/v1/two")!.resolve(); + await Promise.allSettled([one.completed, two.completed]); + await close(downstream); + await close(upstream); + } + }); + + it("relays a failed response unchanged and mutates only the next-request routing state", async () => { + const failureBody = Buffer.from("{\"type\":\"error\",\"error\":{\"type\":\"rate_limit_error\"}}\n"); + let upstreamRequests = 0; + const upstream = createServer((_req, res) => { + upstreamRequests++; + res.writeHead(429, { + "content-type": "application/json", + "retry-after": "invalid", + }); + res.write(failureBody.subarray(0, 11)); + res.end(failureBody.subarray(11)); + }); + const upstreamPort = await listen(upstream); + const invalidate = vi.fn(); + const setCooldown = vi.fn(); + const route = { + account: { id: "account-a" }, + sessionId: "session-a", + }; + + const app = express(); + app.use("/v1", createAnthropicProxy({ + target: `http://127.0.0.1:${upstreamPort}`, + timeoutMs: 2_000, + on: { + proxyRes: proxyResponse => { + applyUpstreamFailureRouting( + proxyResponse.statusCode ?? 0, + proxyResponse.headers["retry-after"], + route, + { invalidate }, + { setCooldown }, + ); + }, + }, + })); + const downstream = createServer(app); + const downstreamPort = await listen(downstream); + + try { + const response = await collect(new URL(`http://127.0.0.1:${downstreamPort}/v1/messages`)); + + expect(response.status).toBe(429); + expect(response.body).toEqual(failureBody); + expect(upstreamRequests).toBe(1); + expect(invalidate).toHaveBeenCalledWith("session-a", "account-a"); + expect(setCooldown).toHaveBeenCalledWith("account-a", 60_000); + expect(response.body.toString("utf8")).not.toContain("message_stop"); + } finally { + await close(downstream); + await close(upstream); + } + }); +}); diff --git a/src/__tests__/lease-lifecycle.test.ts b/src/__tests__/lease-lifecycle.test.ts new file mode 100644 index 0000000..4ebacab --- /dev/null +++ b/src/__tests__/lease-lifecycle.test.ts @@ -0,0 +1,155 @@ +import { EventEmitter } from "node:events"; +import { describe, expect, it, vi } from "vitest"; +import { + acquireRequestRoute, + applyUpstreamFailureRouting, + attachLeaseLifecycle, + routeReasonDetails, +} from "../proxy/lease-lifecycle.js"; + +describe("attachLeaseLifecycle", () => { + it.each(["finish", "close"] as const)("releases once on downstream %s", (event) => { + const response = new EventEmitter(); + const release = vi.fn(); + attachLeaseLifecycle(response, { release }); + + response.emit(event); + response.emit("finish"); + response.emit("close"); + + expect(release).toHaveBeenCalledTimes(1); + }); + + it.each(["proxy error", "refresh/pre-forward failure"])( + "shares one-shot release with explicit %s cleanup", + () => { + const response = new EventEmitter(); + const release = vi.fn(); + const cleanup = attachLeaseLifecycle(response, { release }); + + cleanup(); + cleanup(); + response.emit("finish"); + response.emit("close"); + + expect(release).toHaveBeenCalledTimes(1); + }, + ); + + it("keeps explicit proxy-error cleanup safe after downstream close", () => { + const response = new EventEmitter(); + const release = vi.fn(); + const cleanup = attachLeaseLifecycle(response, { release }); + + response.emit("close"); + cleanup(); + response.emit("finish"); + + expect(release).toHaveBeenCalledTimes(1); + }); +}); + +describe("routeReasonDetails", () => { + it("records only the bounded routing reason, never the session identifier", () => { + const details = routeReasonDetails({ + reason: "sticky", + sessionId: "private-session-value", + }); + + expect(details).toBe("sticky"); + expect(details).not.toContain("private-session-value"); + }); + + it("acquires the session route, attaches cleanup, and returns only reason details", () => { + const response = new EventEmitter(); + const route = { + account: { id: "account-a" }, + reason: "new-session" as const, + sessionId: "private-session-value", + release: vi.fn(), + }; + const acquire = vi.fn().mockReturnValue(route); + + const selected = acquireRequestRoute( + "private-session-value", + response, + { acquire }, + ); + + expect(acquire).toHaveBeenCalledWith("private-session-value"); + expect(selected.route).toBe(route); + expect(selected.details).toBe("new-session"); + expect(selected.details).not.toContain("private-session-value"); + selected.release(); + response.emit("close"); + expect(route.release).toHaveBeenCalledTimes(1); + }); +}); + +describe("applyUpstreamFailureRouting", () => { + const route = { + account: { id: "account-a" }, + sessionId: "session-a", + }; + + it("invalidates the matching binding on 401 without applying a cooldown", () => { + const invalidate = vi.fn(); + const setCooldown = vi.fn(); + + expect(applyUpstreamFailureRouting(401, undefined, route, { invalidate }, { setCooldown })) + .toBeUndefined(); + expect(invalidate).toHaveBeenCalledWith("session-a", "account-a"); + expect(setCooldown).not.toHaveBeenCalled(); + }); + + it("invalidates and applies a numeric Retry-After cooldown on 429", () => { + const invalidate = vi.fn(); + const setCooldown = vi.fn(); + + expect(applyUpstreamFailureRouting(429, "12.5", route, { invalidate }, { setCooldown })) + .toBe(12.5); + expect(invalidate).toHaveBeenCalledWith("session-a", "account-a"); + expect(setCooldown).toHaveBeenCalledWith("account-a", 12_500); + }); + + it.each([undefined, "not-a-number", "-1", "Infinity", ["1", "2"]])( + "uses the 60-second 429 fallback for an unsafe Retry-After value %j", + (retryAfter) => { + const invalidate = vi.fn(); + const setCooldown = vi.fn(); + + expect(applyUpstreamFailureRouting(429, retryAfter, route, { invalidate }, { setCooldown })) + .toBe(60); + expect(setCooldown).toHaveBeenCalledWith("account-a", 60_000); + }, + ); + + it("accepts a finite zero-second Retry-After value", () => { + const invalidate = vi.fn(); + const setCooldown = vi.fn(); + + expect(applyUpstreamFailureRouting(429, "0", route, { invalidate }, { setCooldown })) + .toBe(0); + expect(setCooldown).toHaveBeenCalledWith("account-a", 0); + }); + + it("invalidates and applies the fixed 30-second cooldown on 529", () => { + const invalidate = vi.fn(); + const setCooldown = vi.fn(); + + expect(applyUpstreamFailureRouting(529, "900", route, { invalidate }, { setCooldown })) + .toBe(30); + expect(invalidate).toHaveBeenCalledWith("session-a", "account-a"); + expect(setCooldown).toHaveBeenCalledWith("account-a", 30_000); + }); + + it("does not mutate routing for successful responses", () => { + const invalidate = vi.fn(); + const setCooldown = vi.fn(); + + expect(applyUpstreamFailureRouting(200, undefined, route, { invalidate }, { setCooldown })) + .toBeUndefined(); + expect(invalidate).not.toHaveBeenCalled(); + expect(setCooldown).not.toHaveBeenCalled(); + }); +}); diff --git a/src/proxy/anthropic-proxy.ts b/src/proxy/anthropic-proxy.ts new file mode 100644 index 0000000..fcd1f3f --- /dev/null +++ b/src/proxy/anthropic-proxy.ts @@ -0,0 +1,26 @@ +import type { ServerResponse } from "node:http"; +import type { Request, RequestHandler } from "express"; +import { createProxyMiddleware } from "http-proxy-middleware"; +import type { Options } from "http-proxy-middleware"; + +export interface AnthropicProxyOptions { + target: string; + timeoutMs: number; + on: NonNullable["on"]>; +} + +/** + * Construct the Anthropic transport with http-proxy-middleware's native + * response piping. In particular, this deliberately does not self-handle, + * buffer, transform, or synthesize any response bytes. + */ +export function createAnthropicProxy(options: AnthropicProxyOptions): RequestHandler { + return createProxyMiddleware({ + target: options.target, + changeOrigin: true, + pathRewrite: path => `/v1${path}`, + proxyTimeout: options.timeoutMs, + timeout: options.timeoutMs, + on: options.on, + }); +} diff --git a/src/proxy/lease-lifecycle.ts b/src/proxy/lease-lifecycle.ts new file mode 100644 index 0000000..a1406c4 --- /dev/null +++ b/src/proxy/lease-lifecycle.ts @@ -0,0 +1,104 @@ +export interface ResponseLifecycle { + once(event: "finish" | "close", listener: () => void): unknown; +} + +export interface Releasable { + release(): void; +} + +export interface RouteSummary { + readonly reason: "sticky" | "new-session" | "unscoped" | "failover"; + readonly sessionId?: string; +} + +export interface FailureRoute { + readonly account: { readonly id: string }; + readonly sessionId?: string; +} + +export interface BindingInvalidator { + invalidate(sessionHeader: unknown, expectedAccountId?: string): boolean; +} + +export interface CooldownSetter { + setCooldown(accountId: string, durationMs: number): void; +} + +export interface RoutedRequestLease extends Releasable, RouteSummary, FailureRoute {} + +export interface RouteAcquirer { + acquire(sessionHeader: unknown): T; +} + +/** + * Tie an account lease to every terminal HTTP response path while retaining + * one explicit cleanup callback for failures that happen before forwarding. + */ +export function attachLeaseLifecycle( + response: ResponseLifecycle, + lease: Releasable, +): () => void { + let released = false; + const release = () => { + if (released) return; + released = true; + lease.release(); + }; + + response.once("finish", release); + response.once("close", release); + return release; +} + +/** Keep route diagnostics useful without ever copying a session ID to logs. */ +export function routeReasonDetails(route: RouteSummary): string { + return route.reason; +} + +/** Acquire and immediately bind a routed lease to its response lifecycle. */ +export function acquireRequestRoute( + sessionHeader: unknown, + response: ResponseLifecycle, + router: RouteAcquirer, +): { route: T; release: () => void; details: string } { + const route = router.acquire(sessionHeader); + return { + route, + release: attachLeaseLifecycle(response, route), + details: routeReasonDetails(route), + }; +} + +function retryAfterSeconds(value: unknown): number { + const candidate = Array.isArray(value) + ? value.length === 1 ? Number(value[0]) : Number.NaN + : Number(value); + return Number.isFinite(candidate) && candidate >= 0 ? candidate : 60; +} + +/** + * Apply only routing state changes implied by an upstream failure. The + * current response remains owned by the proxy's native byte stream; callers + * use the returned seconds solely for status logging. + */ +export function applyUpstreamFailureRouting( + status: number, + retryAfterHeader: unknown, + route: FailureRoute, + router: BindingInvalidator, + pool: CooldownSetter, +): number | undefined { + if (status !== 401 && status !== 429 && status !== 529) return undefined; + + router.invalidate(route.sessionId, route.account.id); + if (status === 429) { + const seconds = retryAfterSeconds(retryAfterHeader); + pool.setCooldown(route.account.id, seconds * 1_000); + return seconds; + } + if (status === 529) { + pool.setCooldown(route.account.id, 30_000); + return 30; + } + return undefined; +} diff --git a/src/proxy/server.ts b/src/proxy/server.ts index 30e5101..ff1398d 100644 --- a/src/proxy/server.ts +++ b/src/proxy/server.ts @@ -25,11 +25,20 @@ import { mountMessagesCrossProviderRoute } from "./messages-cross-route.js"; import { mountModelsRoute } from "./models-server.js"; import type { ModelRoutingConfig } from "../protocol/model-ref.js"; import chalk from "chalk"; +import { SessionRouter } from "./session-router.js"; +import type { RoutedAccountLease } from "./session-router.js"; +import { createAnthropicProxy } from "./anthropic-proxy.js"; +import { + acquireRequestRoute, + applyUpstreamFailureRouting, +} from "./lease-lifecycle.js"; // Augment Request to carry the selected account and pending log entry declare module "express-serve-static-core" { interface Request { _ccAccount?: Account; + _ccRoute?: RoutedAccountLease; + _ccReleaseLease?: () => void; _startTime?: number; _pendingLog?: Partial; } @@ -264,6 +273,7 @@ export async function startServer(opts: ServerOptions = {}): Promise { } const pool = new TokenPool(accounts); + const sessionRouter = new SessionRouter(pool); const pickOpenAIAccount = createOpenAIAccountPicker(openAIAccounts); const initialConfig = readConfig(); const modelRouting = initialConfig.modelRouting ?? {}; @@ -386,6 +396,7 @@ export async function startServer(opts: ServerOptions = {}): Promise { if (provider === "anthropic_subscription") { for (const account of pool.getAll()) { pool.updateAccount(account.id, { enabled }); + if (!enabled) sessionRouter.invalidateAccount(account.id); } } else { for (const account of openAIAccounts) { @@ -485,6 +496,7 @@ export async function startServer(opts: ServerOptions = {}): Promise { res.status(500).json({ error: `Failed to persist accounts.json: ${result.message}` }); return; } + if (patch.enabled === false) sessionRouter.invalidateAccount(id); res.json({ account: publicAnthropicAccountView(updated) }); }); @@ -565,6 +577,7 @@ export async function startServer(opts: ServerOptions = {}): Promise { res.status(500).json({ error: `Failed to persist accounts.json: ${result.message}` }); return; } + sessionRouter.invalidateAccount(id); res.json({ ok: true, id }); }); @@ -599,15 +612,9 @@ export async function startServer(opts: ServerOptions = {}): Promise { // ─── Proxy middleware ────────────────────────────────────────────────────── // IMPORTANT: selfHandleResponse must be false (default) for SSE streaming to // work transparently. Setting it to true breaks streaming. - const proxy = createProxyMiddleware({ + const proxy = createAnthropicProxy({ target, - changeOrigin: true, - // Express strips the /v1 mount prefix from req.url before passing it to middleware. - // pathRewrite restores it so the proxy forwards /v1/messages, not /messages. - pathRewrite: (path) => `/v1${path}`, - // Long timeouts — Claude Code requests can be >5min (thinking, agents) - proxyTimeout: proxyRequestTimeoutMs, - timeout: proxyRequestTimeoutMs, + timeoutMs: proxyRequestTimeoutMs, on: { proxyReq: (proxyReq, req) => { const account = (req as Request)._ccAccount; @@ -647,6 +654,7 @@ export async function startServer(opts: ServerOptions = {}): Promise { proxyRes: (proxyRes, req) => { const account = (req as Request)._ccAccount; + const route = (req as Request)._ccRoute; if (!account) return; const status = proxyRes.statusCode ?? 0; @@ -664,6 +672,16 @@ export async function startServer(opts: ServerOptions = {}): Promise { pendingLog.statusCode = status; if (durationMs !== undefined) pendingLog.durationMs = durationMs; + const cooldownSeconds = route + ? applyUpstreamFailureRouting( + status, + proxyRes.headers["retry-after"], + route, + sessionRouter, + pool, + ) + : undefined; + if (status === 401) { // Token invalid or expired mid-request. // Forward the 401 to the client (Claude Code will retry on 401). @@ -681,13 +699,10 @@ export async function startServer(opts: ServerOptions = {}): Promise { // Rate limited — put account on cooldown for Retry-After seconds. stats.totalErrors++; account.errorCount++; - const retryAfter = Number(proxyRes.headers["retry-after"] ?? 60); + const retryAfter = cooldownSeconds ?? 60; pendingLog.type = "error"; pendingLog.details = `rate limited — cooldown ${retryAfter}s`; logError(account.id, 429, `Rate limited — cooldown ${retryAfter}s`); - - account.busy = true; - setTimeout(() => { account.busy = false; }, retryAfter * 1_000); } else if (status === 529) { // Anthropic service overloaded — short cooldown on this account. stats.totalErrors++; @@ -695,9 +710,6 @@ export async function startServer(opts: ServerOptions = {}): Promise { pendingLog.type = "error"; pendingLog.details = "service overloaded — cooldown 30s"; logError(account.id, 529, "Service overloaded — cooldown 30s"); - - account.busy = true; - setTimeout(() => { account.busy = false; }, 30_000); } // ── Capture rate limit utilization from response headers ──────────── @@ -769,17 +781,19 @@ export async function startServer(opts: ServerOptions = {}): Promise { }, error: (err: Error, _req: IncomingMessage, res: ServerResponse | Socket) => { + const request = _req as Request; + request._ccReleaseLease?.(); stats.totalErrors++; logError("proxy", 0, err.message); // Complete the pending log entry for connection-level errors - const pendingLog = (_req as Request)._pendingLog; + const pendingLog = request._pendingLog; if (pendingLog) { pendingLog.type = "error"; pendingLog.statusCode = 0; pendingLog.details = err.message; - if ((_req as Request)._startTime) { - pendingLog.durationMs = Date.now() - (_req as Request)._startTime!; + if (request._startTime) { + pendingLog.durationMs = Date.now() - request._startTime; } stats.addLog(pendingLog as LogEntry); } @@ -801,9 +815,13 @@ export async function startServer(opts: ServerOptions = {}): Promise { // CRITICAL: Do NOT use express.json() here — it consumes the body stream // and breaks SSE streaming passthrough. app.use("/v1", async (req, res, next) => { - let account: Account; + let selected: { route: RoutedAccountLease; release: () => void; details: string }; try { - account = pool.getNext(); + selected = acquireRequestRoute( + req.headers["x-claude-code-session-id"], + res, + sessionRouter, + ); } catch (err) { if (err instanceof EmptyPoolError) { stats.totalErrors++; @@ -814,28 +832,41 @@ export async function startServer(opts: ServerOptions = {}): Promise { }); return; } - throw err; + next(err); + return; } - // Synchronous refresh if token expires within the buffer window - if (needsRefresh(account)) { - const ok = await refreshAccountToken(account); - if (ok) saveAccounts(pool.getAll()); - if (!ok) { - stats.totalErrors++; - logError(account.id, 401, "Token refresh failed"); - res.status(401).json({ - type: "error", - error: { - type: "authentication_error", - message: "Anthropic subscription token refresh failed", - }, - }); - return; + const { route } = selected; + req._ccRoute = route; + req._ccReleaseLease = selected.release; + const account = route.account; + req._ccAccount = account; + + try { + // Synchronous refresh if token expires within the buffer window + if (needsRefresh(account)) { + const ok = await refreshAccountToken(account); + if (ok) saveAccounts(pool.getAll()); + if (!ok) { + req._ccReleaseLease(); + stats.totalErrors++; + logError(account.id, 401, "Token refresh failed"); + res.status(401).json({ + type: "error", + error: { + type: "authentication_error", + message: "Anthropic subscription token refresh failed", + }, + }); + return; + } } + } catch (err) { + req._ccReleaseLease(); + next(err); + return; } - req._ccAccount = account; req._startTime = Date.now(); const source = req.headers["x-claude-code-session-id"] ? "cli" as const @@ -851,6 +882,7 @@ export async function startServer(opts: ServerOptions = {}): Promise { method: req.method, path: req.path, source, + details: selected.details, }; stats.totalRequests++; From b9e7449eb54a30d32245686fda427cfd7ac273db Mon Sep 17 00:00:00 2001 From: timo <44401485+Timo972@users.noreply.github.com> Date: Mon, 20 Jul 2026 19:13:39 +0200 Subject: [PATCH 10/17] fix: harden routing failure state --- src/__tests__/lease-lifecycle.test.ts | 12 +++++- src/__tests__/provider-routing.test.ts | 58 ++++++++++++++++++++++++++ src/proxy/lease-lifecycle.ts | 14 +++++-- src/proxy/provider-routing.ts | 28 +++++++++++++ src/proxy/server.ts | 15 +++++-- 5 files changed, 118 insertions(+), 9 deletions(-) create mode 100644 src/__tests__/provider-routing.test.ts create mode 100644 src/proxy/provider-routing.ts diff --git a/src/__tests__/lease-lifecycle.test.ts b/src/__tests__/lease-lifecycle.test.ts index 4ebacab..ba9b9ba 100644 --- a/src/__tests__/lease-lifecycle.test.ts +++ b/src/__tests__/lease-lifecycle.test.ts @@ -112,7 +112,17 @@ describe("applyUpstreamFailureRouting", () => { expect(setCooldown).toHaveBeenCalledWith("account-a", 12_500); }); - it.each([undefined, "not-a-number", "-1", "Infinity", ["1", "2"]])( + it.each([ + undefined, + "", + " ", + "not-a-number", + "-1", + "Infinity", + "1e308", + ["1"], + ["1", "2"], + ])( "uses the 60-second 429 fallback for an unsafe Retry-After value %j", (retryAfter) => { const invalidate = vi.fn(); diff --git a/src/__tests__/provider-routing.test.ts b/src/__tests__/provider-routing.test.ts new file mode 100644 index 0000000..ddbd242 --- /dev/null +++ b/src/__tests__/provider-routing.test.ts @@ -0,0 +1,58 @@ +import { describe, expect, it, vi } from "vitest"; +import { persistProviderEnabledState } from "../proxy/provider-routing.js"; + +describe("persistProviderEnabledState", () => { + it("does not invalidate Anthropic affinity when persistence fails", () => { + const persist = vi.fn(() => { throw new Error("disk full"); }); + const invalidateAccount = vi.fn(); + + expect(() => persistProviderEnabledState({ + provider: "anthropic_subscription", + enabled: false, + accountIds: ["account-a", "account-b"], + persist, + invalidateAccount, + })).toThrow("disk full"); + + expect(invalidateAccount).not.toHaveBeenCalled(); + }); + + it("invalidates every Anthropic binding only after a successful disable", () => { + const order: string[] = []; + const persist = vi.fn(() => { + order.push("persist"); + return 2; + }); + const invalidateAccount = vi.fn((accountId: string) => { + order.push(`invalidate:${accountId}`); + }); + + expect(persistProviderEnabledState({ + provider: "anthropic_subscription", + enabled: false, + accountIds: ["account-a", "account-b"], + persist, + invalidateAccount, + })).toBe(2); + + expect(order).toEqual(["persist", "invalidate:account-a", "invalidate:account-b"]); + }); + + it.each([ + ["anthropic_subscription", true], + ["openai_subscription", false], + ["openai_subscription", true], + ] as const)("does not invalidate for provider %s with enabled=%s", (provider, enabled) => { + const invalidateAccount = vi.fn(); + + persistProviderEnabledState({ + provider, + enabled, + accountIds: ["account-a"], + persist: () => 1, + invalidateAccount, + }); + + expect(invalidateAccount).not.toHaveBeenCalled(); + }); +}); diff --git a/src/proxy/lease-lifecycle.ts b/src/proxy/lease-lifecycle.ts index a1406c4..f85b14e 100644 --- a/src/proxy/lease-lifecycle.ts +++ b/src/proxy/lease-lifecycle.ts @@ -70,10 +70,16 @@ export function acquireRequestRoute( } function retryAfterSeconds(value: unknown): number { - const candidate = Array.isArray(value) - ? value.length === 1 ? Number(value[0]) : Number.NaN - : Number(value); - return Number.isFinite(candidate) && candidate >= 0 ? candidate : 60; + if (typeof value !== "string" && typeof value !== "number") return 60; + if (typeof value === "string" && value.trim().length === 0) return 60; + + const candidate = Number(value); + const milliseconds = candidate * 1_000; + return Number.isFinite(candidate) && + candidate >= 0 && + Number.isFinite(milliseconds) + ? candidate + : 60; } /** diff --git a/src/proxy/provider-routing.ts b/src/proxy/provider-routing.ts new file mode 100644 index 0000000..f62e1c1 --- /dev/null +++ b/src/proxy/provider-routing.ts @@ -0,0 +1,28 @@ +export type ManagedProvider = "anthropic_subscription" | "openai_subscription"; + +export interface PersistProviderEnabledStateOptions { + provider: ManagedProvider; + enabled: boolean; + accountIds: Iterable; + persist(): T; + invalidateAccount(accountId: string): unknown; +} + +/** + * Persist a provider toggle before discarding any Anthropic affinity. If + * persistence throws, the caller can roll back runtime enablement without + * losing bindings that still point at valid accounts. + */ +export function persistProviderEnabledState( + options: PersistProviderEnabledStateOptions, +): T { + const result = options.persist(); + + if (options.provider === "anthropic_subscription" && !options.enabled) { + for (const accountId of options.accountIds) { + options.invalidateAccount(accountId); + } + } + + return result; +} diff --git a/src/proxy/server.ts b/src/proxy/server.ts index ff1398d..84f9763 100644 --- a/src/proxy/server.ts +++ b/src/proxy/server.ts @@ -32,6 +32,7 @@ import { acquireRequestRoute, applyUpstreamFailureRouting, } from "./lease-lifecycle.js"; +import { persistProviderEnabledState } from "./provider-routing.js"; // Augment Request to carry the selected account and pending log entry declare module "express-serve-static-core" { @@ -385,6 +386,7 @@ export async function startServer(opts: ServerOptions = {}): Promise { res.status(400).json({ error: "enabled must be boolean" }); return; } + const enabled = body.enabled; const provider = providerParam; const snapshots = { @@ -396,7 +398,6 @@ export async function startServer(opts: ServerOptions = {}): Promise { if (provider === "anthropic_subscription") { for (const account of pool.getAll()) { pool.updateAccount(account.id, { enabled }); - if (!enabled) sessionRouter.invalidateAccount(account.id); } } else { for (const account of openAIAccounts) { @@ -415,10 +416,16 @@ export async function startServer(opts: ServerOptions = {}): Promise { } }; - applyRuntime(body.enabled); + applyRuntime(enabled); try { - const changed = setProviderAccountsEnabled(provider, body.enabled, accountsPath); - res.json({ provider, enabled: body.enabled, changed }); + const changed = persistProviderEnabledState({ + provider, + enabled, + accountIds: pool.getAll().map(account => account.id), + persist: () => setProviderAccountsEnabled(provider, enabled, accountsPath), + invalidateAccount: accountId => sessionRouter.invalidateAccount(accountId), + }); + res.json({ provider, enabled, changed }); } catch (err) { rollback(); const message = err instanceof Error ? err.message : String(err); From 04e9f56f13e4f8dd731091133f91dcdb7a44b08f Mon Sep 17 00:00:00 2001 From: timo <44401485+Timo972@users.noreply.github.com> Date: Mon, 20 Jul 2026 19:19:39 +0200 Subject: [PATCH 11/17] docs: explain cache-aware session routing --- README.md | 16 ++++-- package.json | 4 +- src/__tests__/server-health-accounts.test.ts | 36 ++++++++++++ src/proxy/server.ts | 59 +++++++++++++++++--- src/ui/Dashboard.tsx | 5 ++ 5 files changed, 107 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index 621d994..7c27106 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,7 @@ Distribute Claude Code requests across Claude subscriptions, and expose an OpenA ### Features -- **Round-robin token rotation** — distribute requests across 2-20 Claude Max accounts automatically +- **Cache-aware session routing** — keep each Claude Code session on one account while distributing new sessions across 2-20 Claude Max accounts - **Multi-provider routing** — route `openai/*` models to OpenAI ChatGPT/Codex subscription accounts and Claude models to Claude subscriptions - **Transparent Claude proxy** — Claude Code works normally; streaming, thinking, tool use, prompt caching all pass through - **Codex CLI support** — configure Codex to use CC-Router as a Responses-compatible provider @@ -62,6 +62,14 @@ Claude Desktop ─[mitmproxy]─┐ (optional — intercepts api.anthropic.com All standard Claude Code features work transparently on the Claude route: streaming, extended thinking, tool use, prompt caching. OpenAI subscription routing is available for Codex-compatible Responses requests and Claude Code cross-routing with the limitations documented below. +### Cache-aware Claude account routing + +CC-Router keeps requests from one Claude Code session on the same healthy Anthropic subscription account. This session affinity preserves prompt-cache locality instead of scattering a conversation's shared prefix across account-specific caches. New sessions prefer the account with the fewest in-flight requests, then the fewest bound sessions, then the most rate-limit headroom. + +If an upstream account returns 401, 429, or 529, CC-Router passes that response through unchanged and invalidates the session's affinity. The client's next retry can then select another usable account; the router never retries after response bytes have started. Affinity mappings exist only in process memory, expire after one hour of inactivity, and are capped in size. Session IDs are never persisted or logged. + +Streaming remains byte-transparent. In particular, CC-Router never appends a synthetic `message_stop` event. `proxyRequestTimeoutMs` is only a transport safety limit for an upstream request that stops making progress; increasing it may allow a genuinely long request to finish, but does not fix an upstream or forwarding path that failed to deliver its terminal event. + **Claude Desktop support** is opt-in and requires a small interceptor (mitmproxy) because Claude Desktop doesn't expose a custom API endpoint setting. See [Claude Desktop support](#claude-desktop-support). --- @@ -76,7 +84,7 @@ With two accounts you double your effective rate limit. With three, you triple i ```text 1 account → hit limit, wait 60s, continue -3 accounts → request rotates across all three, limit effectively tripled +3 accounts → new sessions spread across all three; each session stays cache-local ``` --- @@ -170,7 +178,7 @@ server { } ``` -For longer requests, set `proxyRequestTimeoutMs` in `~/.cc-router/config.json` (milliseconds) and keep `proxy_read_timeout` at least as high. +For longer requests, set `proxyRequestTimeoutMs` in `~/.cc-router/config.json` (milliseconds) and keep `proxy_read_timeout` at least as high. This timeout is a transport safety limit, not a repair for a missing upstream terminal event; CC-Router does not synthesize `message_stop` when it expires. Each developer then points to: ```json @@ -539,7 +547,7 @@ Then start the interceptor: cc-router client start-desktop ``` -Open Claude Desktop and send a message. The request will be intercepted, redirected to CC-Router, and round-robinned across your accounts just like Claude Code traffic. +Open Claude Desktop and send a message. The request will be intercepted and redirected to CC-Router, which applies the same cache-aware account routing used for Claude Code traffic. ### Stopping / removing Desktop interception diff --git a/package.json b/package.json index 32390b2..93f757f 100644 --- a/package.json +++ b/package.json @@ -1,7 +1,7 @@ { "name": "ai-cc-router", "version": "0.6.2", - "description": "Round-robin proxy for Claude Max OAuth tokens — use multiple Claude Max accounts with Claude Code", + "description": "Cache-aware session router for Claude Max OAuth tokens — use multiple Claude Max accounts with Claude Code", "type": "module", "bin": { "cc-router": "dist/cli/index.js" @@ -19,7 +19,7 @@ "anthropic", "proxy", "oauth", - "round-robin", + "session-routing", "claude-code", "claude-max" ], diff --git a/src/__tests__/server-health-accounts.test.ts b/src/__tests__/server-health-accounts.test.ts index 22a8156..7dba740 100644 --- a/src/__tests__/server-health-accounts.test.ts +++ b/src/__tests__/server-health-accounts.test.ts @@ -56,13 +56,49 @@ describe("createHealthAccountViews", () => { expect(views[1]).toMatchObject({ healthy: true, busy: false, + inFlightRequests: 0, + activeSessions: 0, requestCount: 0, errorCount: 0, enabled: true, }); + expect(views[0]).toMatchObject({ + inFlightRequests: 0, + activeSessions: 0, + }); expect(views[1].rateLimits).toBeUndefined(); }); + it("includes safe routing counters without exposing session identifiers", () => { + const openAIAccount: OpenAISubscriptionAccount = { + id: "openai-primary", + provider: "openai_subscription", + accessToken: "openai-access", + refreshToken: "openai-refresh", + expiresAt: Date.now() + 120_000, + enabled: true, + }; + + const views = createHealthAccountViews( + [makeAnthropicAccount()], + [openAIAccount], + accountId => accountId === "max-account-1" + ? { inFlightRequests: 2, activeSessions: 3, coolingDown: true } + : { inFlightRequests: 0, activeSessions: 0, coolingDown: false }, + ); + + expect(views[0]).toMatchObject({ + busy: true, + inFlightRequests: 2, + activeSessions: 3, + }); + expect(views[1]).toMatchObject({ + inFlightRequests: 0, + activeSessions: 0, + }); + expect(JSON.stringify(views)).not.toContain("session-a"); + }); + it("does not count disabled Anthropic accounts as healthy", () => { const disabled = { ...makeAnthropicAccount(), enabled: false }; diff --git a/src/proxy/server.ts b/src/proxy/server.ts index 84f9763..2c9d4cc 100644 --- a/src/proxy/server.ts +++ b/src/proxy/server.ts @@ -58,6 +58,8 @@ export interface HealthAccountView { enabled: boolean; healthy: boolean; busy: boolean; + inFlightRequests: number; + activeSessions: number; requestCount: number; errorCount: number; expiresInMs: number; @@ -68,6 +70,20 @@ export interface HealthAccountView { weeklyLimitPercent?: number; } +export interface AccountRoutingMetrics { + inFlightRequests: number; + activeSessions: number; + coolingDown: boolean; +} + +type RoutingMetricsResolver = (accountId: string) => AccountRoutingMetrics; + +const zeroRoutingMetrics: RoutingMetricsResolver = () => ({ + inFlightRequests: 0, + activeSessions: 0, + coolingDown: false, +}); + export interface OperationalStatus { mode: string; target: string; @@ -150,14 +166,20 @@ export function createOperationalStatus(opts: { export function createHealthAccountViews( anthropicAccounts: Account[], openAIAccounts: OpenAISubscriptionAccount[], + resolveRoutingMetrics: RoutingMetricsResolver = zeroRoutingMetrics, ): HealthAccountView[] { return [ - ...anthropicAccounts.map(publicAnthropicAccountView), + ...anthropicAccounts.map(account => ( + publicAnthropicAccountView(account, resolveRoutingMetrics(account.id)) + )), ...openAIAccounts.map(publicOpenAIAccountView), ]; } -function publicAnthropicAccountView(a: Account): HealthAccountView { +function publicAnthropicAccountView( + a: Account, + metrics: AccountRoutingMetrics, +): HealthAccountView { return { id: a.id, provider: "anthropic_subscription", @@ -165,7 +187,9 @@ function publicAnthropicAccountView(a: Account): HealthAccountView { sessionLimitPercent: a.sessionLimitPercent, weeklyLimitPercent: a.weeklyLimitPercent, healthy: a.enabled !== false && a.healthy, - busy: a.busy, + busy: a.busy || metrics.coolingDown, + inFlightRequests: metrics.inFlightRequests, + activeSessions: metrics.activeSessions, requestCount: a.requestCount, errorCount: a.errorCount, expiresInMs: a.tokens.expiresAt - Date.now(), @@ -183,6 +207,8 @@ function publicOpenAIAccountView(a: OpenAISubscriptionAccount): HealthAccountVie enabled: a.enabled !== false, healthy: a.enabled !== false && expiresInMs > 0, busy: false, + inFlightRequests: 0, + activeSessions: 0, requestCount: 0, errorCount: 0, expiresInMs, @@ -275,6 +301,11 @@ export async function startServer(opts: ServerOptions = {}): Promise { const pool = new TokenPool(accounts); const sessionRouter = new SessionRouter(pool); + const resolveRoutingMetrics: RoutingMetricsResolver = accountId => ({ + inFlightRequests: pool.getInFlight(accountId), + activeSessions: sessionRouter.getActiveSessionCount(accountId), + coolingDown: pool.isCoolingDown(accountId), + }); const pickOpenAIAccount = createOpenAIAccountPicker(openAIAccounts); const initialConfig = readConfig(); const modelRouting = initialConfig.modelRouting ?? {}; @@ -337,7 +368,11 @@ export async function startServer(opts: ServerOptions = {}): Promise { // Sweep expired cooldowns on each poll so the dashboard reflects recovery // even during idle periods when no /v1 request would trigger getNext(). pool.sweepExpiredCooldowns(); - const accountViews = createHealthAccountViews(pool.getAll(), openAIAccounts); + const accountViews = createHealthAccountViews( + pool.getAll(), + openAIAccounts, + resolveRoutingMetrics, + ); res.json({ status: accountViews.some(a => a.healthy) ? "ok" : "degraded", mode, @@ -371,7 +406,13 @@ export async function startServer(opts: ServerOptions = {}): Promise { // Shape returned to clients — NEVER includes access/refresh tokens. accountsRouter.get("/", (_req, res) => { - res.json({ accounts: createHealthAccountViews(pool.getAll(), openAIAccounts) }); + res.json({ + accounts: createHealthAccountViews( + pool.getAll(), + openAIAccounts, + resolveRoutingMetrics, + ), + }); }); accountsRouter.patch("/providers/:provider", (req, res) => { @@ -504,7 +545,9 @@ export async function startServer(opts: ServerOptions = {}): Promise { return; } if (patch.enabled === false) sessionRouter.invalidateAccount(id); - res.json({ account: publicAnthropicAccountView(updated) }); + res.json({ + account: publicAnthropicAccountView(updated, resolveRoutingMetrics(updated.id)), + }); }); accountsRouter.post("/", (req, res) => { @@ -552,7 +595,9 @@ export async function startServer(opts: ServerOptions = {}): Promise { res.status(500).json({ error: `Failed to persist accounts.json: ${result.message}` }); return; } - res.status(201).json({ account: publicAnthropicAccountView(added) }); + res.status(201).json({ + account: publicAnthropicAccountView(added, resolveRoutingMetrics(added.id)), + }); }); accountsRouter.delete("/:id", (req, res) => { diff --git a/src/ui/Dashboard.tsx b/src/ui/Dashboard.tsx index b32466b..831c590 100644 --- a/src/ui/Dashboard.tsx +++ b/src/ui/Dashboard.tsx @@ -29,6 +29,8 @@ interface AccountStat { provider?: "anthropic_subscription" | "openai_subscription"; healthy: boolean; busy: boolean; + inFlightRequests?: number; + activeSessions?: number; requestCount: number; errorCount: number; expiresInMs: number; @@ -878,6 +880,9 @@ function AccountRow({ account: a, selected }: { account: AccountStat; selected: {expiryLabel.padEnd(8)} last {formatAgo(a.lastUsedMs)} + {a.provider !== "openai_subscription" && ( + {a.activeSessions ?? 0} active / {a.inFlightRequests ?? 0} streams + )} {capsHint && {capsHint}} {rl.lastUpdated > 0 && ( From cee1d15c2c15019d05abe2671bc7504db8833df7 Mon Sep 17 00:00:00 2001 From: timo <44401485+Timo972@users.noreply.github.com> Date: Mon, 20 Jul 2026 19:23:25 +0200 Subject: [PATCH 12/17] docs: clarify routing tie break --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 7c27106..cb7c97d 100644 --- a/README.md +++ b/README.md @@ -64,7 +64,7 @@ All standard Claude Code features work transparently on the Claude route: stream ### Cache-aware Claude account routing -CC-Router keeps requests from one Claude Code session on the same healthy Anthropic subscription account. This session affinity preserves prompt-cache locality instead of scattering a conversation's shared prefix across account-specific caches. New sessions prefer the account with the fewest in-flight requests, then the fewest bound sessions, then the most rate-limit headroom. +CC-Router keeps requests from one Claude Code session on the same healthy Anthropic subscription account. This session affinity preserves prompt-cache locality instead of scattering a conversation's shared prefix across account-specific caches. New sessions prefer the account with the fewest in-flight requests, then the fewest bound sessions, then the most rate-limit headroom; exact ties use a rotating round-robin order. If an upstream account returns 401, 429, or 529, CC-Router passes that response through unchanged and invalidates the session's affinity. The client's next retry can then select another usable account; the router never retries after response bytes have started. Affinity mappings exist only in process memory, expire after one hour of inactivity, and are capped in size. Session IDs are never persisted or logged. From 19838204cdda72cfe89a2bab6401abc3df4ee19c Mon Sep 17 00:00:00 2001 From: timo <44401485+Timo972@users.noreply.github.com> Date: Mon, 20 Jul 2026 19:49:23 +0200 Subject: [PATCH 13/17] fix: harden intelligent session routing --- src/__tests__/account-deletion.test.ts | 89 ++++++++++++ src/__tests__/anthropic-proxy.test.ts | 176 ++++++++++++++++++++++-- src/__tests__/anthropic-routing.test.ts | 171 +++++++++++++++++++++++ src/__tests__/lease-lifecycle.test.ts | 109 ++++++++++++--- src/__tests__/session-router.test.ts | 53 ++++++- src/__tests__/token-pool.test.ts | 18 +++ src/proxy/account-deletion.ts | 27 ++++ src/proxy/anthropic-routing.ts | 121 ++++++++++++++++ src/proxy/lease-lifecycle.ts | 28 ++-- src/proxy/server.ts | 133 ++++++++---------- src/proxy/session-router.ts | 45 +++++- src/proxy/token-pool.ts | 13 +- 12 files changed, 849 insertions(+), 134 deletions(-) create mode 100644 src/__tests__/account-deletion.test.ts create mode 100644 src/__tests__/anthropic-routing.test.ts create mode 100644 src/proxy/account-deletion.ts create mode 100644 src/proxy/anthropic-routing.ts diff --git a/src/__tests__/account-deletion.test.ts b/src/__tests__/account-deletion.test.ts new file mode 100644 index 0000000..dd51b9c --- /dev/null +++ b/src/__tests__/account-deletion.test.ts @@ -0,0 +1,89 @@ +import { describe, expect, it, vi } from "vitest"; +import { deleteAnthropicAccountTransaction } from "../proxy/account-deletion.js"; +import { SessionRouter } from "../proxy/session-router.js"; +import { TokenPool } from "../proxy/token-pool.js"; +import type { Account } from "../proxy/types.js"; +import { DEFAULT_RATE_LIMITS } from "../proxy/types.js"; + +function makeAccount(id: string): Account { + return { + id, + tokens: { + accessToken: `access-${id}`, + refreshToken: `refresh-${id}`, + expiresAt: Date.now() + 60_000, + scopes: ["user:inference"], + }, + healthy: true, + busy: false, + requestCount: 0, + errorCount: 0, + lastUsed: 0, + lastRefresh: 0, + consecutiveErrors: 0, + rateLimits: { ...DEFAULT_RATE_LIMITS }, + enabled: true, + sessionLimitPercent: 100, + weeklyLimitPercent: 100, + }; +} + +describe("deleteAnthropicAccountTransaction", () => { + it("leaves the exact runtime state untouched when prospective persistence fails", () => { + let now = 0; + const first = makeAccount("a"); + const second = makeAccount("b"); + const pool = new TokenPool([first, second], { now: () => now }); + const router = new SessionRouter(pool, { now: () => now }); + const openLease = router.acquire("session-a"); + expect(openLease.account).toBe(first); + pool.setCooldownForAccount(first, 60_000); + const persist = vi.fn(() => { throw new Error("disk full"); }); + + expect(() => deleteAnthropicAccountTransaction({ + id: "a", + pool, + sessionRouter: router, + persist, + })).toThrow("disk full"); + + expect(persist).toHaveBeenCalledTimes(1); + expect(persist.mock.calls[0][0]).toEqual([second]); + expect(pool.getAll()).toEqual([first, second]); + expect(pool.findById("a")).toBe(first); + expect(pool.getInFlight("a")).toBe(1); + expect(pool.isCoolingDown("a")).toBe(true); + expect(router.getActiveSessionCount("a")).toBe(1); + now = 60_000; + const sticky = router.acquire("session-a"); + expect(sticky.account).toBe(first); + expect(sticky.reason).toBe("sticky"); + expect(sticky.bindingGeneration).toBe(openLease.bindingGeneration); + sticky.release(); + openLease.release(); + }); + + it("persists prospective state before removing runtime state and bindings", () => { + const first = makeAccount("a"); + const second = makeAccount("b"); + const pool = new TokenPool([first, second]); + const router = new SessionRouter(pool); + router.acquire("session-a").release(); + const persist = vi.fn((prospective: Account[]) => { + expect(prospective).toEqual([second]); + expect(pool.findById("a")).toBe(first); + expect(router.getBindingCount()).toBe(1); + }); + + const removed = deleteAnthropicAccountTransaction({ + id: "a", + pool, + sessionRouter: router, + persist, + }); + + expect(removed).toBe(first); + expect(pool.getAll()).toEqual([second]); + expect(router.getBindingCount()).toBe(0); + }); +}); diff --git a/src/__tests__/anthropic-proxy.test.ts b/src/__tests__/anthropic-proxy.test.ts index ff59e99..3643288 100644 --- a/src/__tests__/anthropic-proxy.test.ts +++ b/src/__tests__/anthropic-proxy.test.ts @@ -1,9 +1,40 @@ -import { createServer, request, type Server } from "node:http"; +import { createServer, request, type ClientRequest, type Server } from "node:http"; import type { AddressInfo } from "node:net"; import express from "express"; import { describe, expect, it, vi } from "vitest"; import { createAnthropicProxy } from "../proxy/anthropic-proxy.js"; import { applyUpstreamFailureRouting } from "../proxy/lease-lifecycle.js"; +import { + createAnthropicRefreshMiddleware, + createAnthropicRoutingMiddleware, +} from "../proxy/anthropic-routing.js"; +import { SessionRouter } from "../proxy/session-router.js"; +import { TokenPool } from "../proxy/token-pool.js"; +import type { Account } from "../proxy/types.js"; +import { DEFAULT_RATE_LIMITS } from "../proxy/types.js"; + +function makeAccount(id: string): Account { + return { + id, + tokens: { + accessToken: `access-${id}`, + refreshToken: `refresh-${id}`, + expiresAt: Date.now() + 60_000, + scopes: ["user:inference"], + }, + healthy: true, + busy: false, + requestCount: 0, + errorCount: 0, + lastUsed: 0, + lastRefresh: 0, + consecutiveErrors: 0, + rateLimits: { ...DEFAULT_RATE_LIMITS }, + enabled: true, + sessionLimitPercent: 100, + weeklyLimitPercent: 100, + }; +} async function listen(server: Server): Promise { await new Promise((resolve, reject) => { @@ -29,13 +60,16 @@ function deferred(): { promise: Promise; resolve: () => void } { return { promise, resolve }; } -function collect(url: URL): Promise<{ +function collect( + url: URL, + headers: Record = {}, +): Promise<{ status: number; headers: Record; body: Buffer; }> { return new Promise((resolve, reject) => { - const req = request(url, response => { + const req = request(url, { headers }, response => { const chunks: Buffer[] = []; response.on("data", chunk => chunks.push(Buffer.from(chunk))); response.on("end", () => resolve({ @@ -50,15 +84,20 @@ function collect(url: URL): Promise<{ }); } -function startCollecting(url: URL): { +function startCollecting( + url: URL, + headers: Record = {}, +): { firstChunk: Promise; completed: Promise; hasCompleted: () => boolean; + request: ClientRequest; } { const first = deferred(); let completed = false; + let clientRequest!: ClientRequest; const body = new Promise((resolve, reject) => { - const req = request(url, response => { + clientRequest = request(url, { headers }, response => { const chunks: Buffer[] = []; response.once("data", () => first.resolve()); response.on("data", chunk => chunks.push(Buffer.from(chunk))); @@ -68,10 +107,15 @@ function startCollecting(url: URL): { }); response.on("error", reject); }); - req.on("error", reject); - req.end(); + clientRequest.on("error", reject); + clientRequest.end(); }); - return { firstChunk: first.promise, completed: body, hasCompleted: () => completed }; + return { + firstChunk: first.promise, + completed: body, + hasCompleted: () => completed, + request: clientRequest, + }; } describe("createAnthropicProxy", () => { @@ -184,6 +228,111 @@ describe("createAnthropicProxy", () => { } }); + it("routes concurrent production-stack SSE streams with sticky leases and abort cleanup", async () => { + const gates = new Map([ + ["/v1/one", deferred()], + ["/v1/two", deferred()], + ["/v1/abort", deferred()], + ]); + const expected = new Map(); + const authorization = new Map(); + const upstream = createServer(async (req, res) => { + const path = req.url ?? ""; + authorization.set(path, String(req.headers.authorization ?? "")); + const prefix = Buffer.from(`event: ping\ndata: {"stream":"${path}"}\n\n`); + const suffix = Buffer.from("event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"); + expected.set(path, Buffer.concat([prefix, suffix])); + res.writeHead(200, { "content-type": "text/event-stream" }); + res.write(prefix); + const gate = gates.get(path); + if (gate) await gate.promise; + if (!res.destroyed) res.end(suffix); + }); + const upstreamPort = await listen(upstream); + const accounts = [makeAccount("a"), makeAccount("b")]; + const pool = new TokenPool(accounts); + const sessionRouter = new SessionRouter(pool); + const app = express(); + app.use("/v1", createAnthropicRoutingMiddleware({ sessionRouter })); + app.use("/v1", createAnthropicRefreshMiddleware({ + needsRefresh: () => false, + refresh: async () => true, + onRefreshFailure: vi.fn(), + })); + app.use("/v1", createAnthropicProxy({ + target: `http://127.0.0.1:${upstreamPort}`, + timeoutMs: 2_000, + on: { + proxyReq: (proxyRequest, req) => { + const account = (req as express.Request)._ccAccount!; + proxyRequest.setHeader("authorization", `Bearer ${account.tokens.accessToken}`); + }, + }, + })); + const downstream = createServer(app); + const downstreamPort = await listen(downstream); + const base = `http://127.0.0.1:${downstreamPort}`; + const one = startCollecting(new URL(`${base}/v1/one`), { + "X-Claude-Code-Session-Id": "session-one", + }); + const two = startCollecting(new URL(`${base}/v1/two`), { + "X-Claude-Code-Session-Id": "session-two", + }); + let aborted: ReturnType | undefined; + + try { + await Promise.all([one.firstChunk, two.firstChunk]); + expect(authorization.get("/v1/one")).toBe("Bearer access-a"); + expect(authorization.get("/v1/two")).toBe("Bearer access-b"); + expect(pool.getInFlight("a")).toBe(1); + expect(pool.getInFlight("b")).toBe(1); + expect([...sessionRouter.getActiveSessionCountsSnapshot()]).toEqual([ + ["a", 1], + ["b", 1], + ]); + + const follow = await collect(new URL(`${base}/v1/follow`), { + "X-Claude-Code-Session-Id": "session-one", + }); + expect(authorization.get("/v1/follow")).toBe("Bearer access-a"); + expect(follow.body).toEqual(expected.get("/v1/follow")); + expect(follow.body.toString("utf8").match(/event: message_stop/g)).toHaveLength(1); + expect(pool.getInFlight("a")).toBe(1); + + aborted = startCollecting(new URL(`${base}/v1/abort`), { + "X-Claude-Code-Session-Id": "session-two", + }); + await aborted.firstChunk; + expect(authorization.get("/v1/abort")).toBe("Bearer access-b"); + expect(pool.getInFlight("b")).toBe(2); + aborted.request.destroy(); + await Promise.allSettled([aborted.completed]); + await vi.waitFor(() => expect(pool.getInFlight("b")).toBe(1)); + + gates.get("/v1/one")!.resolve(); + gates.get("/v1/two")!.resolve(); + const [oneBody, twoBody] = await Promise.all([one.completed, two.completed]); + expect(oneBody).toEqual(expected.get("/v1/one")); + expect(twoBody).toEqual(expected.get("/v1/two")); + expect(oneBody.toString("utf8").match(/event: message_stop/g)).toHaveLength(1); + expect(twoBody.toString("utf8").match(/event: message_stop/g)).toHaveLength(1); + expect(pool.getInFlight("a")).toBe(0); + expect(pool.getInFlight("b")).toBe(0); + } finally { + for (const gate of gates.values()) gate.resolve(); + one.request.destroy(); + two.request.destroy(); + aborted?.request.destroy(); + await Promise.allSettled([ + one.completed, + two.completed, + ...(aborted ? [aborted.completed] : []), + ]); + await close(downstream); + await close(upstream); + } + }); + it("relays a failed response unchanged and mutates only the next-request routing state", async () => { const failureBody = Buffer.from("{\"type\":\"error\",\"error\":{\"type\":\"rate_limit_error\"}}\n"); let upstreamRequests = 0; @@ -198,10 +347,11 @@ describe("createAnthropicProxy", () => { }); const upstreamPort = await listen(upstream); const invalidate = vi.fn(); - const setCooldown = vi.fn(); + const setCooldownForAccount = vi.fn(); const route = { - account: { id: "account-a" }, + account: makeAccount("account-a"), sessionId: "session-a", + bindingGeneration: 1, }; const app = express(); @@ -215,7 +365,7 @@ describe("createAnthropicProxy", () => { proxyResponse.headers["retry-after"], route, { invalidate }, - { setCooldown }, + { setCooldownForAccount }, ); }, }, @@ -229,8 +379,8 @@ describe("createAnthropicProxy", () => { expect(response.status).toBe(429); expect(response.body).toEqual(failureBody); expect(upstreamRequests).toBe(1); - expect(invalidate).toHaveBeenCalledWith("session-a", "account-a"); - expect(setCooldown).toHaveBeenCalledWith("account-a", 60_000); + expect(invalidate).toHaveBeenCalledWith("session-a", "account-a", 1); + expect(setCooldownForAccount).toHaveBeenCalledWith(route.account, 60_000); expect(response.body.toString("utf8")).not.toContain("message_stop"); } finally { await close(downstream); diff --git a/src/__tests__/anthropic-routing.test.ts b/src/__tests__/anthropic-routing.test.ts new file mode 100644 index 0000000..8dd63e3 --- /dev/null +++ b/src/__tests__/anthropic-routing.test.ts @@ -0,0 +1,171 @@ +import { createServer, request, type ClientRequest, type Server } from "node:http"; +import type { AddressInfo } from "node:net"; +import express from "express"; +import { describe, expect, it, vi } from "vitest"; +import { + createAnthropicRefreshMiddleware, + createAnthropicRoutingMiddleware, +} from "../proxy/anthropic-routing.js"; +import { SessionRouter } from "../proxy/session-router.js"; +import { TokenPool } from "../proxy/token-pool.js"; +import type { Account } from "../proxy/types.js"; +import { DEFAULT_RATE_LIMITS } from "../proxy/types.js"; + +function makeAccount(id: string): Account { + return { + id, + tokens: { + accessToken: `access-${id}`, + refreshToken: `refresh-${id}`, + expiresAt: Date.now() + 60_000, + scopes: ["user:inference"], + }, + healthy: true, + busy: false, + requestCount: 0, + errorCount: 0, + lastUsed: 0, + lastRefresh: 0, + consecutiveErrors: 0, + rateLimits: { ...DEFAULT_RATE_LIMITS }, + enabled: true, + sessionLimitPercent: 100, + weeklyLimitPercent: 100, + }; +} + +async function listen(server: Server): Promise { + await new Promise((resolve, reject) => { + server.once("error", reject); + server.listen(0, "127.0.0.1", () => { + server.off("error", reject); + resolve(); + }); + }); + return (server.address() as AddressInfo).port; +} + +async function close(server: Server): Promise { + await new Promise((resolve, reject) => { + server.close(error => error ? reject(error) : resolve()); + server.closeAllConnections(); + }); +} + +function deferred(): { + promise: Promise; + resolve: (value: T) => void; +} { + let resolve!: (value: T) => void; + const promise = new Promise(done => { resolve = done; }); + return { promise, resolve }; +} + +function send(options: Parameters[0]): Promise<{ + status: number; + body: string; +}> { + return new Promise((resolve, reject) => { + const req = request(options, response => { + const chunks: Buffer[] = []; + response.on("data", chunk => chunks.push(Buffer.from(chunk))); + response.on("end", () => resolve({ + status: response.statusCode ?? 0, + body: Buffer.concat(chunks).toString("utf8"), + })); + response.on("error", reject); + }); + req.on("error", reject); + req.end(); + }); +} + +describe("production Anthropic routing middleware", () => { + it("rejects duplicate native HTTP session fields as unscoped", async () => { + const pool = new TokenPool([makeAccount("a")]); + const sessionRouter = new SessionRouter(pool); + let observedHeaderFields = 0; + const app = express(); + app.use(createAnthropicRoutingMiddleware({ sessionRouter })); + app.use((req, res) => { + observedHeaderFields = req.rawHeaders.filter( + value => value.toLowerCase() === "x-claude-code-session-id", + ).length; + res.json({ reason: req._ccRoute?.reason }); + }); + const server = createServer(app); + const port = await listen(server); + + try { + const response = await send({ + host: "127.0.0.1", + port, + path: "/v1/messages", + headers: { + "X-Claude-Code-Session-Id": ["session-a", "session-b"], + }, + }); + + expect(response.status).toBe(200); + expect(observedHeaderFields).toBe(2); + expect(JSON.parse(response.body)).toEqual({ reason: "unscoped" }); + expect(sessionRouter.getBindingCount()).toBe(0); + } finally { + await close(server); + } + }); + + it("does not continue after the client disconnects during deferred refresh", async () => { + const pool = new TokenPool([makeAccount("a")]); + const sessionRouter = new SessionRouter(pool); + const refreshStarted = deferred(); + const refreshResult = deferred(); + const forwarded = vi.fn(); + const app = express(); + app.use(createAnthropicRoutingMiddleware({ sessionRouter })); + app.use(createAnthropicRefreshMiddleware({ + needsRefresh: () => true, + refresh: async () => { + refreshStarted.resolve(); + return refreshResult.promise; + }, + onRefreshFailure: vi.fn(), + })); + app.use((_req, res) => { + forwarded(); + res.end("forwarded"); + }); + const server = createServer(app); + const port = await listen(server); + let client: ClientRequest | undefined; + + try { + const clientClosed = new Promise(resolve => { + client = request({ + host: "127.0.0.1", + port, + path: "/v1/messages", + headers: { "X-Claude-Code-Session-Id": "session-a" }, + }); + client.on("error", () => resolve()); + client.end(); + }); + + await refreshStarted.promise; + expect(pool.getInFlight("a")).toBe(1); + client.destroy(); + await clientClosed; + await vi.waitFor(() => expect(pool.getInFlight("a")).toBe(0)); + + refreshResult.resolve(true); + await new Promise(resolve => setImmediate(resolve)); + + expect(forwarded).not.toHaveBeenCalled(); + expect(pool.getInFlight("a")).toBe(0); + } finally { + refreshResult.resolve(true); + client?.destroy(); + await close(server); + } + }); +}); diff --git a/src/__tests__/lease-lifecycle.test.ts b/src/__tests__/lease-lifecycle.test.ts index ba9b9ba..e665510 100644 --- a/src/__tests__/lease-lifecycle.test.ts +++ b/src/__tests__/lease-lifecycle.test.ts @@ -6,6 +6,33 @@ import { attachLeaseLifecycle, routeReasonDetails, } from "../proxy/lease-lifecycle.js"; +import { SessionRouter } from "../proxy/session-router.js"; +import { TokenPool } from "../proxy/token-pool.js"; +import type { Account } from "../proxy/types.js"; +import { DEFAULT_RATE_LIMITS } from "../proxy/types.js"; + +function makeAccount(id: string): Account { + return { + id, + tokens: { + accessToken: `access-${id}`, + refreshToken: `refresh-${id}`, + expiresAt: Date.now() + 60_000, + scopes: ["user:inference"], + }, + healthy: true, + busy: false, + requestCount: 0, + errorCount: 0, + lastUsed: 0, + lastRefresh: 0, + consecutiveErrors: 0, + rateLimits: { ...DEFAULT_RATE_LIMITS }, + enabled: true, + sessionLimitPercent: 100, + weeklyLimitPercent: 100, + }; +} describe("attachLeaseLifecycle", () => { it.each(["finish", "close"] as const)("releases once on downstream %s", (event) => { @@ -88,28 +115,29 @@ describe("routeReasonDetails", () => { describe("applyUpstreamFailureRouting", () => { const route = { - account: { id: "account-a" }, + account: makeAccount("account-a"), sessionId: "session-a", + bindingGeneration: 7, }; it("invalidates the matching binding on 401 without applying a cooldown", () => { const invalidate = vi.fn(); - const setCooldown = vi.fn(); + const setCooldownForAccount = vi.fn(); - expect(applyUpstreamFailureRouting(401, undefined, route, { invalidate }, { setCooldown })) + expect(applyUpstreamFailureRouting(401, undefined, route, { invalidate }, { setCooldownForAccount })) .toBeUndefined(); - expect(invalidate).toHaveBeenCalledWith("session-a", "account-a"); - expect(setCooldown).not.toHaveBeenCalled(); + expect(invalidate).toHaveBeenCalledWith("session-a", "account-a", 7); + expect(setCooldownForAccount).not.toHaveBeenCalled(); }); it("invalidates and applies a numeric Retry-After cooldown on 429", () => { const invalidate = vi.fn(); - const setCooldown = vi.fn(); + const setCooldownForAccount = vi.fn(); - expect(applyUpstreamFailureRouting(429, "12.5", route, { invalidate }, { setCooldown })) + expect(applyUpstreamFailureRouting(429, "12.5", route, { invalidate }, { setCooldownForAccount })) .toBe(12.5); - expect(invalidate).toHaveBeenCalledWith("session-a", "account-a"); - expect(setCooldown).toHaveBeenCalledWith("account-a", 12_500); + expect(invalidate).toHaveBeenCalledWith("session-a", "account-a", 7); + expect(setCooldownForAccount).toHaveBeenCalledWith(route.account, 12_500); }); it.each([ @@ -126,40 +154,77 @@ describe("applyUpstreamFailureRouting", () => { "uses the 60-second 429 fallback for an unsafe Retry-After value %j", (retryAfter) => { const invalidate = vi.fn(); - const setCooldown = vi.fn(); + const setCooldownForAccount = vi.fn(); - expect(applyUpstreamFailureRouting(429, retryAfter, route, { invalidate }, { setCooldown })) + expect(applyUpstreamFailureRouting(429, retryAfter, route, { invalidate }, { setCooldownForAccount })) .toBe(60); - expect(setCooldown).toHaveBeenCalledWith("account-a", 60_000); + expect(setCooldownForAccount).toHaveBeenCalledWith(route.account, 60_000); }, ); it("accepts a finite zero-second Retry-After value", () => { const invalidate = vi.fn(); - const setCooldown = vi.fn(); + const setCooldownForAccount = vi.fn(); - expect(applyUpstreamFailureRouting(429, "0", route, { invalidate }, { setCooldown })) + expect(applyUpstreamFailureRouting(429, "0", route, { invalidate }, { setCooldownForAccount })) .toBe(0); - expect(setCooldown).toHaveBeenCalledWith("account-a", 0); + expect(setCooldownForAccount).toHaveBeenCalledWith(route.account, 0); }); it("invalidates and applies the fixed 30-second cooldown on 529", () => { const invalidate = vi.fn(); - const setCooldown = vi.fn(); + const setCooldownForAccount = vi.fn(); - expect(applyUpstreamFailureRouting(529, "900", route, { invalidate }, { setCooldown })) + expect(applyUpstreamFailureRouting(529, "900", route, { invalidate }, { setCooldownForAccount })) .toBe(30); - expect(invalidate).toHaveBeenCalledWith("session-a", "account-a"); - expect(setCooldown).toHaveBeenCalledWith("account-a", 30_000); + expect(invalidate).toHaveBeenCalledWith("session-a", "account-a", 7); + expect(setCooldownForAccount).toHaveBeenCalledWith(route.account, 30_000); }); it("does not mutate routing for successful responses", () => { const invalidate = vi.fn(); - const setCooldown = vi.fn(); + const setCooldownForAccount = vi.fn(); - expect(applyUpstreamFailureRouting(200, undefined, route, { invalidate }, { setCooldown })) + expect(applyUpstreamFailureRouting(200, undefined, route, { invalidate }, { setCooldownForAccount })) .toBeUndefined(); expect(invalidate).not.toHaveBeenCalled(); - expect(setCooldown).not.toHaveBeenCalled(); + expect(setCooldownForAccount).not.toHaveBeenCalled(); + }); + + it("ignores an old same-account failure after the session was rebound", () => { + const pool = new TokenPool([makeAccount("a")]); + const router = new SessionRouter(pool); + const oldRoute = router.acquire("session-a"); + oldRoute.release(); + router.invalidate(oldRoute.sessionId, oldRoute.account.id, oldRoute.bindingGeneration); + const rebound = router.acquire("session-a"); + rebound.release(); + + applyUpstreamFailureRouting(401, undefined, oldRoute, router, pool); + + const sticky = router.acquire("session-a"); + expect(sticky.reason).toBe("sticky"); + expect(sticky.bindingGeneration).toBe(rebound.bindingGeneration); + sticky.release(); + }); + + it("does not cool down a replacement account from an old incarnation's failure", () => { + const oldAccount = makeAccount("a"); + const pool = new TokenPool([oldAccount]); + const router = new SessionRouter(pool); + const oldRoute = router.acquire("session-a"); + oldRoute.release(); + pool.removeAccount("a"); + pool.addAccount({ + id: "a", + accessToken: "replacement-access", + refreshToken: "replacement-refresh", + expiresAt: Date.now() + 60_000, + scopes: ["user:inference"], + }); + + applyUpstreamFailureRouting(429, "60", oldRoute, router, pool); + + expect(pool.isCoolingDown("a")).toBe(false); }); }); diff --git a/src/__tests__/session-router.test.ts b/src/__tests__/session-router.test.ts index c4e2a3a..ed8b2b5 100644 --- a/src/__tests__/session-router.test.ts +++ b/src/__tests__/session-router.test.ts @@ -149,6 +149,34 @@ describe("SessionRouter", () => { expect(router.getBindingCount()).toBe(0); }); + it("does not let an old same-account response invalidate a rebound generation", () => { + const pool = new TokenPool([makeAccount("a")]); + const router = new SessionRouter(pool); + const oldRoute = router.acquire("session-a"); + oldRoute.release(); + + expect(router.invalidate( + oldRoute.sessionId, + oldRoute.account.id, + oldRoute.bindingGeneration, + )).toBe(true); + const rebound = router.acquire("session-a"); + rebound.release(); + + expect(rebound.account).toBe(oldRoute.account); + expect(rebound.bindingGeneration).not.toBe(oldRoute.bindingGeneration); + expect(router.invalidate( + oldRoute.sessionId, + oldRoute.account.id, + oldRoute.bindingGeneration, + )).toBe(false); + + const sticky = router.acquire("session-a"); + expect(sticky.reason).toBe("sticky"); + expect(sticky.bindingGeneration).toBe(rebound.bindingGeneration); + sticky.release(); + }); + it("invalidates every binding for one account and repairs counts", () => { const pool = new TokenPool([makeAccount("a"), makeAccount("b")]); const router = new SessionRouter(pool); @@ -210,6 +238,29 @@ describe("SessionRouter", () => { expect(router.getActiveSessionCount("a")).toBe(0); }); + it("returns one aggregate account snapshot after a single expiry sweep", () => { + let now = 0; + let clockReads = 0; + const clock = () => { + clockReads++; + return now; + }; + const pool = new TokenPool([makeAccount("a"), makeAccount("b")], { now: clock }); + const router = new SessionRouter(pool, { now: clock, ttlMs: 10 }); + router.acquire("expired-session").release(); + now = 5; + router.acquire("live-session").release(); + now = 10; + clockReads = 0; + + const snapshot = router.getActiveSessionCountsSnapshot(); + + expect(clockReads).toBe(1); + expect([...snapshot]).toEqual([["b", 1]]); + expect([...snapshot.keys()]).not.toContain("expired-session"); + expect([...snapshot.keys()]).not.toContain("live-session"); + }); + it("sweeps expired bindings before applying the capacity limit", () => { let now = 0; const pool = new TokenPool([makeAccount("a"), makeAccount("b")], { now: () => now }); @@ -268,7 +319,7 @@ describe("SessionRouter", () => { expect(lease.sessionId).toBe("private-session-id"); expect(Object.keys(lease).sort()).toEqual( - ["account", "fallback", "reason", "release", "sessionId"].sort(), + ["account", "bindingGeneration", "fallback", "reason", "release", "sessionId"].sort(), ); expect(log).not.toHaveBeenCalled(); expect(warn).not.toHaveBeenCalled(); diff --git a/src/__tests__/token-pool.test.ts b/src/__tests__/token-pool.test.ts index c88307a..688fec6 100644 --- a/src/__tests__/token-pool.test.ts +++ b/src/__tests__/token-pool.test.ts @@ -227,6 +227,24 @@ describe("TokenPool — timestamp cooldown", () => { now = 61_000; expect(pool.isCoolingDown("a")).toBe(false); }); + + it("does not let an old account incarnation cool down a replacement with the same ID", () => { + const oldAccount = makeAccount("a"); + const pool = new TokenPool([oldAccount]); + pool.removeAccount("a"); + const replacement = pool.addAccount({ + id: "a", + accessToken: "sk-ant-oat01-replacement", + refreshToken: "sk-ant-ort01-replacement", + expiresAt: Date.now() + 3_600_000, + scopes: ["user:inference"], + }); + + pool.setCooldownForAccount(oldAccount, 60_000); + + expect(pool.findById("a")).toBe(replacement); + expect(pool.isCoolingDown("a")).toBe(false); + }); }); describe("TokenPool — unhealthy accounts", () => { diff --git a/src/proxy/account-deletion.ts b/src/proxy/account-deletion.ts new file mode 100644 index 0000000..8b25327 --- /dev/null +++ b/src/proxy/account-deletion.ts @@ -0,0 +1,27 @@ +import type { SessionRouter } from "./session-router.js"; +import type { TokenPool } from "./token-pool.js"; +import type { Account } from "./types.js"; + +export interface DeleteAnthropicAccountOptions { + id: string; + pool: TokenPool; + sessionRouter: SessionRouter; + persist(accounts: Account[]): void; +} + +/** Persist prospective state before irreversibly removing runtime routing state. */ +export function deleteAnthropicAccountTransaction( + options: DeleteAnthropicAccountOptions, +): Account { + const account = options.pool.findById(options.id); + if (!account) throw new Error(`Account "${options.id}" not found`); + + const prospective = options.pool.getAll().filter(candidate => candidate !== account); + options.persist(prospective); + + if (!options.pool.removeAccount(options.id)) { + throw new Error(`Account "${options.id}" disappeared during deletion`); + } + options.sessionRouter.invalidateAccount(options.id); + return account; +} diff --git a/src/proxy/anthropic-routing.ts b/src/proxy/anthropic-routing.ts new file mode 100644 index 0000000..d57f3f6 --- /dev/null +++ b/src/proxy/anthropic-routing.ts @@ -0,0 +1,121 @@ +import type { IncomingMessage } from "node:http"; +import type { NextFunction, Request, RequestHandler, Response } from "express"; +import { acquireRequestRoute } from "./lease-lifecycle.js"; +import { normalizeSessionId, type RoutedAccountLease, type SessionRouter } from "./session-router.js"; +import { EmptyPoolError } from "./token-pool.js"; +import type { Account } from "./types.js"; + +const SESSION_HEADER = "x-claude-code-session-id"; + +type RoutedRequest = Request & { + _ccAccount?: Account; + _ccRoute?: RoutedAccountLease; + _ccReleaseLease?: () => void; +}; + +/** Extract exactly one native HTTP session header field without joined duplicates. */ +export function extractClaudeSessionId(request: IncomingMessage): string | undefined { + const distinct = request.headersDistinct; + if (distinct !== undefined) { + const values = distinct[SESSION_HEADER]; + if (!values || values.length !== 1) return undefined; + return normalizeSessionId(values[0]); + } + + const values: string[] = []; + for (let index = 0; index < request.rawHeaders.length; index += 2) { + if (request.rawHeaders[index]?.toLowerCase() !== SESSION_HEADER) continue; + values.push(request.rawHeaders[index + 1] ?? ""); + } + if (values.length !== 1) return undefined; + return normalizeSessionId(values[0]); +} + +export interface AnthropicRoutingMiddlewareOptions { + sessionRouter: SessionRouter; + onEmptyPool?: (error: EmptyPoolError, request: Request, response: Response) => void; +} + +/** Acquire the production route and bind its lease to downstream termination. */ +export function createAnthropicRoutingMiddleware( + options: AnthropicRoutingMiddlewareOptions, +): RequestHandler { + return (request, response, next) => { + const routedRequest = request as RoutedRequest; + try { + const selected = acquireRequestRoute( + extractClaudeSessionId(request), + response, + options.sessionRouter, + ); + routedRequest._ccRoute = selected.route; + routedRequest._ccReleaseLease = selected.release; + routedRequest._ccAccount = selected.route.account; + next(); + } catch (error) { + if (error instanceof EmptyPoolError && options.onEmptyPool) { + options.onEmptyPool(error, request, response); + return; + } + next(error); + } + }; +} + +export interface AnthropicRefreshMiddlewareOptions { + needsRefresh(account: Account): boolean; + /** Refresh and durably persist rotated credentials before resolving true. */ + refresh(account: Account): Promise; + onRefreshFailure(account: Account): void; +} + +function requestTerminated(request: Request, response: Response): boolean { + return request.aborted || response.destroyed || response.writableEnded; +} + +/** Prepare the selected account, but never continue after downstream termination. */ +export function createAnthropicRefreshMiddleware( + options: AnthropicRefreshMiddlewareOptions, +): RequestHandler { + return async (request: Request, response: Response, next: NextFunction) => { + const routedRequest = request as RoutedRequest; + const account = routedRequest._ccAccount; + const release = routedRequest._ccReleaseLease; + if (!account || !release) { + next(new Error("Anthropic route missing before refresh")); + return; + } + + try { + if (options.needsRefresh(account)) { + const ok = await options.refresh(account); + if (requestTerminated(request, response)) { + release(); + return; + } + if (!ok) { + release(); + options.onRefreshFailure(account); + response.status(401).json({ + type: "error", + error: { + type: "authentication_error", + message: "Anthropic subscription token refresh failed", + }, + }); + return; + } + } + + if (requestTerminated(request, response)) { + release(); + return; + } + next(); + } catch (error) { + release(); + if (requestTerminated(request, response)) return; + next(error); + } + }; +} diff --git a/src/proxy/lease-lifecycle.ts b/src/proxy/lease-lifecycle.ts index f85b14e..285b952 100644 --- a/src/proxy/lease-lifecycle.ts +++ b/src/proxy/lease-lifecycle.ts @@ -9,19 +9,25 @@ export interface Releasable { export interface RouteSummary { readonly reason: "sticky" | "new-session" | "unscoped" | "failover"; readonly sessionId?: string; + readonly bindingGeneration?: number; } -export interface FailureRoute { - readonly account: { readonly id: string }; +export interface FailureRoute { + readonly account: TAccount; readonly sessionId?: string; + readonly bindingGeneration?: number; } export interface BindingInvalidator { - invalidate(sessionHeader: unknown, expectedAccountId?: string): boolean; + invalidate( + sessionHeader: unknown, + expectedAccountId?: string, + expectedGeneration?: number, + ): boolean; } -export interface CooldownSetter { - setCooldown(accountId: string, durationMs: number): void; +export interface CooldownSetter { + setCooldownForAccount(account: TAccount, durationMs: number): void; } export interface RoutedRequestLease extends Releasable, RouteSummary, FailureRoute {} @@ -87,23 +93,23 @@ function retryAfterSeconds(value: unknown): number { * current response remains owned by the proxy's native byte stream; callers * use the returned seconds solely for status logging. */ -export function applyUpstreamFailureRouting( +export function applyUpstreamFailureRouting( status: number, retryAfterHeader: unknown, - route: FailureRoute, + route: FailureRoute, router: BindingInvalidator, - pool: CooldownSetter, + pool: CooldownSetter, ): number | undefined { if (status !== 401 && status !== 429 && status !== 529) return undefined; - router.invalidate(route.sessionId, route.account.id); + router.invalidate(route.sessionId, route.account.id, route.bindingGeneration); if (status === 429) { const seconds = retryAfterSeconds(retryAfterHeader); - pool.setCooldown(route.account.id, seconds * 1_000); + pool.setCooldownForAccount(route.account, seconds * 1_000); return seconds; } if (status === 529) { - pool.setCooldown(route.account.id, 30_000); + pool.setCooldownForAccount(route.account, 30_000); return 30; } return undefined; diff --git a/src/proxy/server.ts b/src/proxy/server.ts index 2c9d4cc..7cbf43c 100644 --- a/src/proxy/server.ts +++ b/src/proxy/server.ts @@ -5,9 +5,9 @@ import { timingSafeEqual } from "crypto"; import type { IncomingMessage } from "http"; import type { Socket } from "net"; import type { Request } from "express"; -import { TokenPool, EmptyPoolError } from "./token-pool.js"; +import { TokenPool } from "./token-pool.js"; import { needsRefresh, refreshAccountToken, saveAccounts, startRefreshLoop } from "./token-refresher.js"; -import { loadAccounts, loadOpenAIAccounts, saveOpenAIAccounts, accountsFileExists, readAccountsFromPath, readConfig, writeConfig, serialize, getProxyRequestTimeoutMs, migrateLegacyAccountProviders, setProviderAccountsEnabled } from "../config/manager.js"; +import { loadAccounts, loadOpenAIAccounts, saveOpenAIAccounts, accountsFileExists, readAccountsFromPath, readConfig, writeConfig, getProxyRequestTimeoutMs, migrateLegacyAccountProviders, setProviderAccountsEnabled } from "../config/manager.js"; import { checkForUpdate, performUpdate, restartSelf } from "../utils/self-update.js"; import { trackEvent, startHeartbeat } from "../utils/telemetry.js"; import { loadTelemetryState } from "../config/telemetry.js"; @@ -29,10 +29,14 @@ import { SessionRouter } from "./session-router.js"; import type { RoutedAccountLease } from "./session-router.js"; import { createAnthropicProxy } from "./anthropic-proxy.js"; import { - acquireRequestRoute, applyUpstreamFailureRouting, } from "./lease-lifecycle.js"; import { persistProviderEnabledState } from "./provider-routing.js"; +import { deleteAnthropicAccountTransaction } from "./account-deletion.js"; +import { + createAnthropicRefreshMiddleware, + createAnthropicRoutingMiddleware, +} from "./anthropic-routing.js"; // Augment Request to carry the selected account and pending log entry declare module "express-serve-static-core" { @@ -301,11 +305,14 @@ export async function startServer(opts: ServerOptions = {}): Promise { const pool = new TokenPool(accounts); const sessionRouter = new SessionRouter(pool); - const resolveRoutingMetrics: RoutingMetricsResolver = accountId => ({ - inFlightRequests: pool.getInFlight(accountId), - activeSessions: sessionRouter.getActiveSessionCount(accountId), - coolingDown: pool.isCoolingDown(accountId), - }); + const createRoutingMetricsResolver = (): RoutingMetricsResolver => { + const activeSessionCounts = sessionRouter.getActiveSessionCountsSnapshot(); + return accountId => ({ + inFlightRequests: pool.getInFlight(accountId), + activeSessions: activeSessionCounts.get(accountId) ?? 0, + coolingDown: pool.isCoolingDown(accountId), + }); + }; const pickOpenAIAccount = createOpenAIAccountPicker(openAIAccounts); const initialConfig = readConfig(); const modelRouting = initialConfig.modelRouting ?? {}; @@ -368,6 +375,7 @@ export async function startServer(opts: ServerOptions = {}): Promise { // Sweep expired cooldowns on each poll so the dashboard reflects recovery // even during idle periods when no /v1 request would trigger getNext(). pool.sweepExpiredCooldowns(); + const resolveRoutingMetrics = createRoutingMetricsResolver(); const accountViews = createHealthAccountViews( pool.getAll(), openAIAccounts, @@ -406,6 +414,7 @@ export async function startServer(opts: ServerOptions = {}): Promise { // Shape returned to clients — NEVER includes access/refresh tokens. accountsRouter.get("/", (_req, res) => { + const resolveRoutingMetrics = createRoutingMetricsResolver(); res.json({ accounts: createHealthAccountViews( pool.getAll(), @@ -546,7 +555,7 @@ export async function startServer(opts: ServerOptions = {}): Promise { } if (patch.enabled === false) sessionRouter.invalidateAccount(id); res.json({ - account: publicAnthropicAccountView(updated, resolveRoutingMetrics(updated.id)), + account: publicAnthropicAccountView(updated, createRoutingMetricsResolver()(updated.id)), }); }); @@ -596,7 +605,7 @@ export async function startServer(opts: ServerOptions = {}): Promise { return; } res.status(201).json({ - account: publicAnthropicAccountView(added, resolveRoutingMetrics(added.id)), + account: publicAnthropicAccountView(added, createRoutingMetricsResolver()(added.id)), }); }); @@ -614,22 +623,19 @@ export async function startServer(opts: ServerOptions = {}): Promise { res.status(404).json({ error: `Account "${id}" not found` }); return; } - // Snapshot for rollback. serialize() gives us a persistable AccountRecord. - const snapshot = serialize([existing])[0]; - const removed = pool.removeAccount(id); - if (!removed) { - res.status(404).json({ error: `Account "${id}" not found` }); - return; - } - - const result = tryPersist(() => { - pool.addAccount(snapshot); - }); - if (!result.ok) { - res.status(500).json({ error: `Failed to persist accounts.json: ${result.message}` }); + try { + deleteAnthropicAccountTransaction({ + id, + pool, + sessionRouter, + persist: saveAccounts, + }); + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + logError("accounts", 0, `Failed to persist accounts.json: ${message}`); + res.status(500).json({ error: `Failed to persist accounts.json: ${message}` }); return; } - sessionRouter.invalidateAccount(id); res.json({ ok: true, id }); }); @@ -866,61 +872,32 @@ export async function startServer(opts: ServerOptions = {}): Promise { // ─── /v1/* — select account, refresh if needed, then proxy ─────────────── // CRITICAL: Do NOT use express.json() here — it consumes the body stream // and breaks SSE streaming passthrough. - app.use("/v1", async (req, res, next) => { - let selected: { route: RoutedAccountLease; release: () => void; details: string }; - try { - selected = acquireRequestRoute( - req.headers["x-claude-code-session-id"], - res, - sessionRouter, - ); - } catch (err) { - if (err instanceof EmptyPoolError) { - stats.totalErrors++; - logError("proxy", 503, err.message); - res.status(503).json({ - type: "error", - error: { type: "no_accounts", message: err.message }, - }); - return; - } - next(err); - return; - } - - const { route } = selected; - req._ccRoute = route; - req._ccReleaseLease = selected.release; + app.use("/v1", createAnthropicRoutingMiddleware({ + sessionRouter, + onEmptyPool: (err, _req, res) => { + stats.totalErrors++; + logError("proxy", 503, err.message); + res.status(503).json({ + type: "error", + error: { type: "no_accounts", message: err.message }, + }); + }, + }), createAnthropicRefreshMiddleware({ + needsRefresh, + refresh: async (account) => { + const ok = await refreshAccountToken(account); + if (ok) saveAccounts(pool.getAll()); + return ok; + }, + onRefreshFailure: (account) => { + stats.totalErrors++; + logError(account.id, 401, "Token refresh failed"); + }, + }), (req, _res, next) => { + const route = req._ccRoute!; const account = route.account; - req._ccAccount = account; - - try { - // Synchronous refresh if token expires within the buffer window - if (needsRefresh(account)) { - const ok = await refreshAccountToken(account); - if (ok) saveAccounts(pool.getAll()); - if (!ok) { - req._ccReleaseLease(); - stats.totalErrors++; - logError(account.id, 401, "Token refresh failed"); - res.status(401).json({ - type: "error", - error: { - type: "authentication_error", - message: "Anthropic subscription token refresh failed", - }, - }); - return; - } - } - } catch (err) { - req._ccReleaseLease(); - next(err); - return; - } - req._startTime = Date.now(); - const source = req.headers["x-claude-code-session-id"] + const source = route.sessionId !== undefined ? "cli" as const : req.headers["x-api-key"] ? "desktop" as const @@ -934,7 +911,7 @@ export async function startServer(opts: ServerOptions = {}): Promise { method: req.method, path: req.path, source, - details: selected.details, + details: route.reason, }; stats.totalRequests++; diff --git a/src/proxy/session-router.ts b/src/proxy/session-router.ts index 195ad52..c57ab30 100644 --- a/src/proxy/session-router.ts +++ b/src/proxy/session-router.ts @@ -10,6 +10,7 @@ export type RouteReason = "sticky" | "new-session" | "unscoped" | "failover"; export interface RoutedAccountLease extends AccountLease { readonly reason: RouteReason; readonly sessionId?: string; + readonly bindingGeneration?: number; } export interface SessionRouterOptions { @@ -21,6 +22,7 @@ export interface SessionRouterOptions { interface SessionBinding { accountId: string; lastSeen: number; + generation: number; } /** @@ -47,6 +49,7 @@ export class SessionRouter { private readonly now: () => number; private readonly ttlMs: number; private readonly maxEntries: number; + private nextBindingGeneration = 1; constructor( private readonly pool: TokenPool, @@ -78,21 +81,30 @@ export class SessionRouter { const stickyLease = this.pool.tryAcquire(existing.accountId); if (stickyLease) { existing.lastSeen = now; - return this.wrap(stickyLease, "sticky", sessionId); + return this.wrap(stickyLease, "sticky", sessionId, existing.generation); } this.removeBinding(sessionId); } const lease = this.pool.acquireBest(this.activeSessionCounts); - this.insertBinding(sessionId, lease.account.id, now); - return this.wrap(lease, existing ? "failover" : "new-session", sessionId); + const binding = this.insertBinding(sessionId, lease.account.id, now); + return this.wrap( + lease, + existing ? "failover" : "new-session", + sessionId, + binding.generation, + ); } /** * Remove a binding only if it still belongs to the expected account. This * protects a new failover binding from a late response on the old account. */ - invalidate(sessionHeader: unknown, expectedAccountId?: string): boolean { + invalidate( + sessionHeader: unknown, + expectedAccountId?: string, + expectedGeneration?: number, + ): boolean { const sessionId = normalizeSessionId(sessionHeader); if (!sessionId) return false; const binding = this.bindings.get(sessionId); @@ -100,6 +112,9 @@ export class SessionRouter { if (expectedAccountId !== undefined && binding.accountId !== expectedAccountId) { return false; } + if (expectedGeneration !== undefined && binding.generation !== expectedGeneration) { + return false; + } return this.removeBinding(sessionId); } @@ -122,10 +137,17 @@ export class SessionRouter { return this.bindings.size; } + /** Sweep once and expose only aggregate account IDs/counts. */ + getActiveSessionCountsSnapshot(): ReadonlyMap { + this.sweepExpiredBindings(this.now()); + return new Map(this.activeSessionCounts); + } + private wrap( lease: AccountLease, reason: RouteReason, sessionId?: string, + bindingGeneration?: number, ): RoutedAccountLease { return { account: lease.account, @@ -133,13 +155,24 @@ export class SessionRouter { release: lease.release, reason, ...(sessionId === undefined ? {} : { sessionId }), + ...(bindingGeneration === undefined ? {} : { bindingGeneration }), }; } - private insertBinding(sessionId: string, accountId: string, lastSeen: number): void { + private insertBinding( + sessionId: string, + accountId: string, + lastSeen: number, + ): SessionBinding { if (this.bindings.size >= this.maxEntries) this.evictLeastRecentlyUsed(); - this.bindings.set(sessionId, { accountId, lastSeen }); + const binding = { + accountId, + lastSeen, + generation: this.nextBindingGeneration++, + }; + this.bindings.set(sessionId, binding); this.activeSessionCounts.set(accountId, this.getRawActiveSessionCount(accountId) + 1); + return binding; } private removeBinding(sessionId: string): boolean { diff --git a/src/proxy/token-pool.ts b/src/proxy/token-pool.ts index c721846..6569089 100644 --- a/src/proxy/token-pool.ts +++ b/src/proxy/token-pool.ts @@ -208,11 +208,18 @@ export class TokenPool { } setCooldown(accountId: string, durationMs: number): void { - if (!this.findById(accountId)) return; + const account = this.findById(accountId); + if (!account) return; + this.setCooldownForAccount(account, durationMs); + } + + /** Apply cooldown only to the exact account incarnation that was routed. */ + setCooldownForAccount(account: Account, durationMs: number): void { + if (this.findById(account.id) !== account) return; if (!Number.isFinite(durationMs) || durationMs <= 0) return; const proposedExpiry = this.now() + durationMs; - const existingExpiry = this.cooldownUntil.get(accountId) ?? 0; - this.cooldownUntil.set(accountId, Math.max(existingExpiry, proposedExpiry)); + const existingExpiry = this.cooldownUntil.get(account.id) ?? 0; + this.cooldownUntil.set(account.id, Math.max(existingExpiry, proposedExpiry)); } isCoolingDown(accountId: string): boolean { From d0bf1867fe2a77b16ff91146dcce059c6e1d51df Mon Sep 17 00:00:00 2001 From: timo <44401485+Timo972@users.noreply.github.com> Date: Mon, 20 Jul 2026 20:01:47 +0200 Subject: [PATCH 14/17] fix: enforce scoped route identity --- src/__tests__/anthropic-proxy.test.ts | 31 ++++++++++-------- src/__tests__/lease-lifecycle.test.ts | 22 +++++++++++++ src/__tests__/session-router.test.ts | 37 +++++++++++++++++++++ src/proxy/lease-lifecycle.ts | 4 ++- src/proxy/session-router.ts | 46 +++++++++++++++++++-------- 5 files changed, 112 insertions(+), 28 deletions(-) diff --git a/src/__tests__/anthropic-proxy.test.ts b/src/__tests__/anthropic-proxy.test.ts index 3643288..32cb8e3 100644 --- a/src/__tests__/anthropic-proxy.test.ts +++ b/src/__tests__/anthropic-proxy.test.ts @@ -200,7 +200,7 @@ describe("createAnthropicProxy", () => { const app = express(); app.use("/v1", createAnthropicProxy({ target: `http://127.0.0.1:${upstreamPort}`, - timeoutMs: 2_000, + timeoutMs: 10_000, on: {}, })); const downstream = createServer(app); @@ -261,7 +261,7 @@ describe("createAnthropicProxy", () => { })); app.use("/v1", createAnthropicProxy({ target: `http://127.0.0.1:${upstreamPort}`, - timeoutMs: 2_000, + timeoutMs: 10_000, on: { proxyReq: (proxyRequest, req) => { const account = (req as express.Request)._ccAccount!; @@ -282,32 +282,37 @@ describe("createAnthropicProxy", () => { try { await Promise.all([one.firstChunk, two.firstChunk]); - expect(authorization.get("/v1/one")).toBe("Bearer access-a"); - expect(authorization.get("/v1/two")).toBe("Bearer access-b"); - expect(pool.getInFlight("a")).toBe(1); - expect(pool.getInFlight("b")).toBe(1); - expect([...sessionRouter.getActiveSessionCountsSnapshot()]).toEqual([ + const sessionOneAuthorization = authorization.get("/v1/one"); + const sessionTwoAuthorization = authorization.get("/v1/two"); + expect(sessionOneAuthorization).toMatch(/^Bearer access-[ab]$/); + expect(sessionTwoAuthorization).toMatch(/^Bearer access-[ab]$/); + expect(sessionOneAuthorization).not.toBe(sessionTwoAuthorization); + const sessionOneAccountId = sessionOneAuthorization!.endsWith("-a") ? "a" : "b"; + const sessionTwoAccountId = sessionTwoAuthorization!.endsWith("-a") ? "a" : "b"; + expect(pool.getInFlight(sessionOneAccountId)).toBe(1); + expect(pool.getInFlight(sessionTwoAccountId)).toBe(1); + expect(sessionRouter.getActiveSessionCountsSnapshot()).toEqual(new Map([ ["a", 1], ["b", 1], - ]); + ])); const follow = await collect(new URL(`${base}/v1/follow`), { "X-Claude-Code-Session-Id": "session-one", }); - expect(authorization.get("/v1/follow")).toBe("Bearer access-a"); + expect(authorization.get("/v1/follow")).toBe(sessionOneAuthorization); expect(follow.body).toEqual(expected.get("/v1/follow")); expect(follow.body.toString("utf8").match(/event: message_stop/g)).toHaveLength(1); - expect(pool.getInFlight("a")).toBe(1); + expect(pool.getInFlight(sessionOneAccountId)).toBe(1); aborted = startCollecting(new URL(`${base}/v1/abort`), { "X-Claude-Code-Session-Id": "session-two", }); await aborted.firstChunk; - expect(authorization.get("/v1/abort")).toBe("Bearer access-b"); - expect(pool.getInFlight("b")).toBe(2); + expect(authorization.get("/v1/abort")).toBe(sessionTwoAuthorization); + expect(pool.getInFlight(sessionTwoAccountId)).toBe(2); aborted.request.destroy(); await Promise.allSettled([aborted.completed]); - await vi.waitFor(() => expect(pool.getInFlight("b")).toBe(1)); + await vi.waitFor(() => expect(pool.getInFlight(sessionTwoAccountId)).toBe(1)); gates.get("/v1/one")!.resolve(); gates.get("/v1/two")!.resolve(); diff --git a/src/__tests__/lease-lifecycle.test.ts b/src/__tests__/lease-lifecycle.test.ts index e665510..a67226c 100644 --- a/src/__tests__/lease-lifecycle.test.ts +++ b/src/__tests__/lease-lifecycle.test.ts @@ -7,6 +7,7 @@ import { routeReasonDetails, } from "../proxy/lease-lifecycle.js"; import { SessionRouter } from "../proxy/session-router.js"; +import type { RoutedAccountLease } from "../proxy/session-router.js"; import { TokenPool } from "../proxy/token-pool.js"; import type { Account } from "../proxy/types.js"; import { DEFAULT_RATE_LIMITS } from "../proxy/types.js"; @@ -208,6 +209,27 @@ describe("applyUpstreamFailureRouting", () => { sticky.release(); }); + it("does not invalidate a binding when a malformed scoped route omits its generation", () => { + const pool = new TokenPool([makeAccount("a")]); + const router = new SessionRouter(pool); + const current = router.acquire("session-a"); + current.release(); + const malformedRoute = { + account: current.account, + reason: "sticky", + sessionId: current.sessionId, + fallback: false, + release: vi.fn(), + } as unknown as RoutedAccountLease; + + applyUpstreamFailureRouting(401, undefined, malformedRoute, router, pool); + + const sticky = router.acquire("session-a"); + expect(sticky.reason).toBe("sticky"); + expect(sticky.bindingGeneration).toBe(current.bindingGeneration); + sticky.release(); + }); + it("does not cool down a replacement account from an old incarnation's failure", () => { const oldAccount = makeAccount("a"); const pool = new TokenPool([oldAccount]); diff --git a/src/__tests__/session-router.test.ts b/src/__tests__/session-router.test.ts index ed8b2b5..97400c8 100644 --- a/src/__tests__/session-router.test.ts +++ b/src/__tests__/session-router.test.ts @@ -1,5 +1,6 @@ import { describe, expect, it, vi } from "vitest"; import { SessionRouter, normalizeSessionId } from "../proxy/session-router.js"; +import type { RoutedAccountLease } from "../proxy/session-router.js"; import { TokenPool } from "../proxy/token-pool.js"; import type { Account } from "../proxy/types.js"; import { DEFAULT_RATE_LIMITS } from "../proxy/types.js"; @@ -27,6 +28,18 @@ function makeAccount(id: string): Account { }; } +if (false) { + // @ts-expect-error A scoped route cannot compile without a binding generation. + const scopedRouteMissingGeneration: RoutedAccountLease = { + account: makeAccount("compile-only"), + fallback: false, + release: () => undefined, + reason: "sticky", + sessionId: "session-a", + }; + void scopedRouteMissingGeneration; +} + describe("session ID normalization", () => { it("normalizes one bounded string header", () => { expect(normalizeSessionId(" session-a ")).toBe("session-a"); @@ -54,6 +67,30 @@ describe("SessionRouter", () => { second.release(); }); + it("returns both identity fields on every scoped route and neither on unscoped routes", () => { + const a = makeAccount("a"); + const pool = new TokenPool([a, makeAccount("b")]); + const router = new SessionRouter(pool); + const created = router.acquire("session-a"); + const sticky = router.acquire("session-a"); + a.enabled = false; + const failover = router.acquire("session-a"); + const unscoped = router.acquire(undefined); + + for (const route of [created, sticky, failover]) { + expect(route.sessionId).toBe("session-a"); + expect(typeof route.bindingGeneration).toBe("number"); + } + expect(unscoped.reason).toBe("unscoped"); + expect("sessionId" in unscoped).toBe(false); + expect("bindingGeneration" in unscoped).toBe(false); + + created.release(); + sticky.release(); + failover.release(); + unscoped.release(); + }); + it("binds simultaneous new sessions to separate idle accounts", () => { const pool = new TokenPool([makeAccount("a"), makeAccount("b")]); const router = new SessionRouter(pool); diff --git a/src/proxy/lease-lifecycle.ts b/src/proxy/lease-lifecycle.ts index 285b952..d2ffbdf 100644 --- a/src/proxy/lease-lifecycle.ts +++ b/src/proxy/lease-lifecycle.ts @@ -102,7 +102,9 @@ export function applyUpstreamFailureRouting; -export interface RoutedAccountLease extends AccountLease { - readonly reason: RouteReason; - readonly sessionId?: string; - readonly bindingGeneration?: number; +export interface UnscopedRoutedAccountLease extends AccountLease { + readonly reason: "unscoped"; + readonly sessionId?: never; + readonly bindingGeneration?: never; } +export interface ScopedRoutedAccountLease extends AccountLease { + readonly reason: ScopedRouteReason; + readonly sessionId: string; + readonly bindingGeneration: number; +} + +export type RoutedAccountLease = UnscopedRoutedAccountLease | ScopedRoutedAccountLease; + export interface SessionRouterOptions { now?: () => number; ttlMs?: number; @@ -73,7 +82,7 @@ export class SessionRouter { this.sweepExpiredBindings(now); if (!sessionId) { - return this.wrap(this.pool.acquireBest(this.activeSessionCounts), "unscoped"); + return this.wrapUnscoped(this.pool.acquireBest(this.activeSessionCounts)); } const existing = this.bindings.get(sessionId); @@ -81,14 +90,14 @@ export class SessionRouter { const stickyLease = this.pool.tryAcquire(existing.accountId); if (stickyLease) { existing.lastSeen = now; - return this.wrap(stickyLease, "sticky", sessionId, existing.generation); + return this.wrapScoped(stickyLease, "sticky", sessionId, existing.generation); } this.removeBinding(sessionId); } const lease = this.pool.acquireBest(this.activeSessionCounts); const binding = this.insertBinding(sessionId, lease.account.id, now); - return this.wrap( + return this.wrapScoped( lease, existing ? "failover" : "new-session", sessionId, @@ -143,19 +152,28 @@ export class SessionRouter { return new Map(this.activeSessionCounts); } - private wrap( + private wrapUnscoped(lease: AccountLease): UnscopedRoutedAccountLease { + return { + account: lease.account, + fallback: lease.fallback, + release: lease.release, + reason: "unscoped", + }; + } + + private wrapScoped( lease: AccountLease, - reason: RouteReason, - sessionId?: string, - bindingGeneration?: number, - ): RoutedAccountLease { + reason: ScopedRouteReason, + sessionId: string, + bindingGeneration: number, + ): ScopedRoutedAccountLease { return { account: lease.account, fallback: lease.fallback, release: lease.release, reason, - ...(sessionId === undefined ? {} : { sessionId }), - ...(bindingGeneration === undefined ? {} : { bindingGeneration }), + sessionId, + bindingGeneration, }; } From d1b876ff979ef5bd576a424d8047a51857f1c2e6 Mon Sep 17 00:00:00 2001 From: timo <44401485+Timo972@users.noreply.github.com> Date: Mon, 20 Jul 2026 20:22:05 +0200 Subject: [PATCH 15/17] fix: serialize refresh account ownership --- src/__tests__/account-deletion.test.ts | 139 +++++++++++++++++++++++-- src/__tests__/lease-lifecycle.test.ts | 14 +++ src/__tests__/token-refresher.test.ts | 123 ++++++++++++++++++---- src/proxy/account-deletion.ts | 20 +++- src/proxy/lease-lifecycle.ts | 3 +- src/proxy/server.ts | 28 ++--- src/proxy/token-refresher.ts | 39 ++++++- 7 files changed, 319 insertions(+), 47 deletions(-) diff --git a/src/__tests__/account-deletion.test.ts b/src/__tests__/account-deletion.test.ts index dd51b9c..cc20360 100644 --- a/src/__tests__/account-deletion.test.ts +++ b/src/__tests__/account-deletion.test.ts @@ -1,7 +1,11 @@ import { describe, expect, it, vi } from "vitest"; -import { deleteAnthropicAccountTransaction } from "../proxy/account-deletion.js"; +import { + deleteAnthropicAccountTransaction, + LastAccountDeletionError, +} from "../proxy/account-deletion.js"; import { SessionRouter } from "../proxy/session-router.js"; import { TokenPool } from "../proxy/token-pool.js"; +import { refreshAccountToken } from "../proxy/token-refresher.js"; import type { Account } from "../proxy/types.js"; import { DEFAULT_RATE_LIMITS } from "../proxy/types.js"; @@ -28,8 +32,27 @@ function makeAccount(id: string): Account { }; } +function deferred(): { promise: Promise; resolve: (value: T) => void } { + let resolve!: (value: T) => void; + const promise = new Promise(done => { resolve = done; }); + return { promise, resolve }; +} + +function successfulRefreshResponse(): Response { + return { + ok: true, + json: async () => ({ + access_token: "rotated-access", + refresh_token: "rotated-refresh", + expires_in: 28800, + scope: "user:inference", + token_type: "Bearer", + }), + } as Response; +} + describe("deleteAnthropicAccountTransaction", () => { - it("leaves the exact runtime state untouched when prospective persistence fails", () => { + it("leaves the exact runtime state untouched when prospective persistence fails", async () => { let now = 0; const first = makeAccount("a"); const second = makeAccount("b"); @@ -40,12 +63,12 @@ describe("deleteAnthropicAccountTransaction", () => { pool.setCooldownForAccount(first, 60_000); const persist = vi.fn(() => { throw new Error("disk full"); }); - expect(() => deleteAnthropicAccountTransaction({ + await expect(deleteAnthropicAccountTransaction({ id: "a", pool, sessionRouter: router, persist, - })).toThrow("disk full"); + })).rejects.toThrow("disk full"); expect(persist).toHaveBeenCalledTimes(1); expect(persist.mock.calls[0][0]).toEqual([second]); @@ -63,7 +86,7 @@ describe("deleteAnthropicAccountTransaction", () => { openLease.release(); }); - it("persists prospective state before removing runtime state and bindings", () => { + it("persists prospective state before removing runtime state and bindings", async () => { const first = makeAccount("a"); const second = makeAccount("b"); const pool = new TokenPool([first, second]); @@ -75,7 +98,7 @@ describe("deleteAnthropicAccountTransaction", () => { expect(router.getBindingCount()).toBe(1); }); - const removed = deleteAnthropicAccountTransaction({ + const removed = await deleteAnthropicAccountTransaction({ id: "a", pool, sessionRouter: router, @@ -86,4 +109,108 @@ describe("deleteAnthropicAccountTransaction", () => { expect(pool.getAll()).toEqual([second]); expect(router.getBindingCount()).toBe(0); }); + + it("waits for an exact-account refresh before persisting and removing", async () => { + const first = makeAccount("a"); + const second = makeAccount("b"); + const pool = new TokenPool([first, second]); + const router = new SessionRouter(pool); + const response = deferred(); + vi.stubGlobal("fetch", vi.fn(() => response.promise)); + const persist = vi.fn(); + + try { + const refresh = refreshAccountToken(first); + const deletion = deleteAnthropicAccountTransaction({ + id: "a", + pool, + sessionRouter: router, + persist, + }); + await Promise.resolve(); + const persistedBeforeRefresh = persist.mock.calls.length; + + response.resolve(successfulRefreshResponse()); + expect(await refresh).toBe(true); + const removed = await deletion; + + expect(persistedBeforeRefresh).toBe(0); + expect(first.tokens.refreshToken).toBe("rotated-refresh"); + expect(persist).toHaveBeenCalledWith([second]); + expect(removed).toBe(first); + expect(pool.findById("a")).toBeNull(); + } finally { + response.resolve(successfulRefreshResponse()); + vi.unstubAllGlobals(); + } + }); + + it("does not persist or delete a replacement that appears while refresh is pending", async () => { + const oldAccount = makeAccount("a"); + const second = makeAccount("b"); + const pool = new TokenPool([oldAccount, second]); + const router = new SessionRouter(pool); + const response = deferred(); + vi.stubGlobal("fetch", vi.fn(() => response.promise)); + const persist = vi.fn(); + + try { + const refresh = refreshAccountToken(oldAccount); + const deletion = deleteAnthropicAccountTransaction({ + id: "a", + pool, + sessionRouter: router, + persist, + }); + await Promise.resolve(); + pool.removeAccount("a"); + const replacement = pool.addAccount({ + id: "a", + accessToken: "replacement-access", + refreshToken: "replacement-refresh", + expiresAt: Date.now() + 60_000, + scopes: ["user:inference"], + }); + response.resolve(successfulRefreshResponse()); + await refresh; + + await expect(deletion).rejects.toThrow(/changed during deletion/); + expect(persist).not.toHaveBeenCalled(); + expect(pool.findById("a")).toBe(replacement); + expect(replacement.tokens.refreshToken).toBe("replacement-refresh"); + } finally { + response.resolve(successfulRefreshResponse()); + vi.unstubAllGlobals(); + } + }); + + it("does not let concurrent async deletions remove the final account", async () => { + const first = makeAccount("a"); + const second = makeAccount("b"); + const pool = new TokenPool([first, second]); + const router = new SessionRouter(pool); + const persist = vi.fn(); + + const results = await Promise.allSettled([ + deleteAnthropicAccountTransaction({ + id: "a", + pool, + sessionRouter: router, + persist, + }), + deleteAnthropicAccountTransaction({ + id: "b", + pool, + sessionRouter: router, + persist, + }), + ]); + + expect(results.filter(result => result.status === "fulfilled")).toHaveLength(1); + expect(results.filter(result => result.status === "rejected")).toHaveLength(1); + const rejected = results.find(result => result.status === "rejected"); + expect(rejected?.status === "rejected" ? rejected.reason : undefined) + .toBeInstanceOf(LastAccountDeletionError); + expect(pool.getAll()).toHaveLength(1); + }); }); diff --git a/src/__tests__/lease-lifecycle.test.ts b/src/__tests__/lease-lifecycle.test.ts index a67226c..7625b70 100644 --- a/src/__tests__/lease-lifecycle.test.ts +++ b/src/__tests__/lease-lifecycle.test.ts @@ -112,6 +112,20 @@ describe("routeReasonDetails", () => { response.emit("close"); expect(route.release).toHaveBeenCalledTimes(1); }); + + it("marks emergency fallback in bounded route details without exposing the session", () => { + const account = makeAccount("account-a"); + account.enabled = false; + const router = new SessionRouter(new TokenPool([account])); + const route = router.acquire("private-fallback-session"); + + const details = routeReasonDetails(route); + + expect(route.fallback).toBe(true); + expect(details).toBe("new-session:fallback"); + expect(details).not.toContain("private-fallback-session"); + route.release(); + }); }); describe("applyUpstreamFailureRouting", () => { diff --git a/src/__tests__/token-refresher.test.ts b/src/__tests__/token-refresher.test.ts index e1c0665..65cd613 100644 --- a/src/__tests__/token-refresher.test.ts +++ b/src/__tests__/token-refresher.test.ts @@ -1,5 +1,10 @@ import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; -import { needsRefresh, refreshAccountToken } from "../proxy/token-refresher.js"; +import { + needsRefresh, + refreshAccountIfCurrent, + refreshAccountToken, +} from "../proxy/token-refresher.js"; +import { TokenPool } from "../proxy/token-pool.js"; import type { Account } from "../proxy/types.js"; function makeAccount(expiresAt: number): Account { @@ -21,6 +26,25 @@ function makeAccount(expiresAt: number): Account { }; } +function deferred(): { promise: Promise; resolve: (value: T) => void } { + let resolve!: (value: T) => void; + const promise = new Promise(done => { resolve = done; }); + return { promise, resolve }; +} + +function successfulRefreshResponse(suffix: string): Response { + return { + ok: true, + json: async () => ({ + access_token: `sk-ant-oat01-${suffix}`, + refresh_token: `sk-ant-ort01-${suffix}`, + expires_in: 28800, + scope: "user:inference user:profile", + token_type: "Bearer", + }), + } as Response; +} + // ─── needsRefresh ───────────────────────────────────────────────────────────── describe("needsRefresh", () => { @@ -163,34 +187,93 @@ describe("refreshAccountToken", () => { it("deduplicates concurrent refresh calls for the same account", async () => { const account = makeAccount(Date.now() + 5 * 60 * 1000); - let callCount = 0; - - vi.mocked(fetch).mockImplementation(async () => { - callCount++; - await new Promise(r => setTimeout(r, 10)); - return { - ok: true, - json: async () => ({ - access_token: "sk-ant-oat01-NEW", - refresh_token: "sk-ant-ort01-NEW", - expires_in: 28800, - scope: "user:inference user:profile", - token_type: "Bearer", - }), - } as Response; - }); + const response = deferred(); + vi.mocked(fetch).mockImplementation(() => response.promise); // Fire 3 concurrent refreshes for the same account - const [r1, r2, r3] = await Promise.all([ + const pending = [ refreshAccountToken(account), refreshAccountToken(account), refreshAccountToken(account), - ]); + ]; + expect(fetch).toHaveBeenCalledTimes(1); + response.resolve(successfulRefreshResponse("NEW")); + const [r1, r2, r3] = await Promise.all(pending); // All return the same result but fetch was only called once expect(r1).toBe(true); expect(r2).toBe(true); expect(r3).toBe(true); - expect(callCount).toBe(1); + expect(fetch).toHaveBeenCalledTimes(1); + }); + + it("does not coalesce different account objects that reuse the same ID", async () => { + const oldAccount = makeAccount(Date.now() + 5 * 60 * 1000); + const replacement = makeAccount(Date.now() + 5 * 60 * 1000); + const oldResponse = deferred(); + const replacementResponse = deferred(); + vi.mocked(fetch) + .mockImplementationOnce(() => oldResponse.promise) + .mockImplementationOnce(() => replacementResponse.promise); + + const oldRefresh = refreshAccountToken(oldAccount); + const replacementRefresh = refreshAccountToken(replacement); + oldResponse.resolve(successfulRefreshResponse("OLD")); + replacementResponse.resolve(successfulRefreshResponse("REPLACEMENT")); + const results = await Promise.all([oldRefresh, replacementRefresh]); + + expect(fetch).toHaveBeenCalledTimes(2); + expect(results).toEqual([true, true]); + expect(oldAccount.tokens.refreshToken).toBe("sk-ant-ort01-OLD"); + expect(replacement.tokens.refreshToken).toBe("sk-ant-ort01-REPLACEMENT"); + }); + + it("does not start a stale 401 refresh after the account ID is reused", async () => { + const oldAccount = makeAccount(Date.now() + 5 * 60 * 1000); + const pool = new TokenPool([oldAccount]); + pool.removeAccount(oldAccount.id); + const replacement = pool.addAccount({ + id: oldAccount.id, + accessToken: "replacement-access", + refreshToken: "replacement-refresh", + expiresAt: Date.now() + 60_000, + scopes: ["user:inference"], + }); + const refresh = vi.fn(async () => true); + const persist = vi.fn(); + + const result = await refreshAccountIfCurrent(oldAccount, pool, { refresh, persist }); + + expect(result).toBe(false); + expect(refresh).not.toHaveBeenCalled(); + expect(persist).not.toHaveBeenCalled(); + expect(replacement.tokens).toMatchObject({ + accessToken: "replacement-access", + refreshToken: "replacement-refresh", + }); + }); + + it("does not persist a refresh that loses exact account ownership while pending", async () => { + const oldAccount = makeAccount(Date.now() + 5 * 60 * 1000); + const pool = new TokenPool([oldAccount]); + const refreshResult = deferred(); + const refresh = vi.fn(() => refreshResult.promise); + const persist = vi.fn(); + + const pending = refreshAccountIfCurrent(oldAccount, pool, { refresh, persist }); + expect(refresh).toHaveBeenCalledWith(oldAccount); + pool.removeAccount(oldAccount.id); + const replacement = pool.addAccount({ + id: oldAccount.id, + accessToken: "replacement-access", + refreshToken: "replacement-refresh", + expiresAt: Date.now() + 60_000, + scopes: ["user:inference"], + }); + refreshResult.resolve(true); + + expect(await pending).toBe(false); + expect(persist).not.toHaveBeenCalled(); + expect(replacement.tokens.refreshToken).toBe("replacement-refresh"); }); }); diff --git a/src/proxy/account-deletion.ts b/src/proxy/account-deletion.ts index 8b25327..1469c11 100644 --- a/src/proxy/account-deletion.ts +++ b/src/proxy/account-deletion.ts @@ -1,6 +1,14 @@ import type { SessionRouter } from "./session-router.js"; import type { TokenPool } from "./token-pool.js"; import type { Account } from "./types.js"; +import { waitForAccountRefresh } from "./token-refresher.js"; + +export class LastAccountDeletionError extends Error { + constructor() { + super("Cannot remove the last account — at least one must remain"); + this.name = "LastAccountDeletionError"; + } +} export interface DeleteAnthropicAccountOptions { id: string; @@ -10,12 +18,20 @@ export interface DeleteAnthropicAccountOptions { } /** Persist prospective state before irreversibly removing runtime routing state. */ -export function deleteAnthropicAccountTransaction( +export async function deleteAnthropicAccountTransaction( options: DeleteAnthropicAccountOptions, -): Account { +): Promise { const account = options.pool.findById(options.id); if (!account) throw new Error(`Account "${options.id}" not found`); + await waitForAccountRefresh(account); + if (options.pool.findById(options.id) !== account) { + throw new Error(`Account "${options.id}" changed during deletion`); + } + if (options.pool.getAll().length <= 1) { + throw new LastAccountDeletionError(); + } + const prospective = options.pool.getAll().filter(candidate => candidate !== account); options.persist(prospective); diff --git a/src/proxy/lease-lifecycle.ts b/src/proxy/lease-lifecycle.ts index d2ffbdf..3b168b0 100644 --- a/src/proxy/lease-lifecycle.ts +++ b/src/proxy/lease-lifecycle.ts @@ -8,6 +8,7 @@ export interface Releasable { export interface RouteSummary { readonly reason: "sticky" | "new-session" | "unscoped" | "failover"; + readonly fallback?: boolean; readonly sessionId?: string; readonly bindingGeneration?: number; } @@ -58,7 +59,7 @@ export function attachLeaseLifecycle( /** Keep route diagnostics useful without ever copying a session ID to logs. */ export function routeReasonDetails(route: RouteSummary): string { - return route.reason; + return route.fallback ? `${route.reason}:fallback` : route.reason; } /** Acquire and immediately bind a routed lease to its response lifecycle. */ diff --git a/src/proxy/server.ts b/src/proxy/server.ts index 7cbf43c..0618c8e 100644 --- a/src/proxy/server.ts +++ b/src/proxy/server.ts @@ -6,7 +6,7 @@ import type { IncomingMessage } from "http"; import type { Socket } from "net"; import type { Request } from "express"; import { TokenPool } from "./token-pool.js"; -import { needsRefresh, refreshAccountToken, saveAccounts, startRefreshLoop } from "./token-refresher.js"; +import { needsRefresh, refreshAccountIfCurrent, saveAccounts, startRefreshLoop } from "./token-refresher.js"; import { loadAccounts, loadOpenAIAccounts, saveOpenAIAccounts, accountsFileExists, readAccountsFromPath, readConfig, writeConfig, getProxyRequestTimeoutMs, migrateLegacyAccountProviders, setProviderAccountsEnabled } from "../config/manager.js"; import { checkForUpdate, performUpdate, restartSelf } from "../utils/self-update.js"; import { trackEvent, startHeartbeat } from "../utils/telemetry.js"; @@ -30,9 +30,13 @@ import type { RoutedAccountLease } from "./session-router.js"; import { createAnthropicProxy } from "./anthropic-proxy.js"; import { applyUpstreamFailureRouting, + routeReasonDetails, } from "./lease-lifecycle.js"; import { persistProviderEnabledState } from "./provider-routing.js"; -import { deleteAnthropicAccountTransaction } from "./account-deletion.js"; +import { + deleteAnthropicAccountTransaction, + LastAccountDeletionError, +} from "./account-deletion.js"; import { createAnthropicRefreshMiddleware, createAnthropicRoutingMiddleware, @@ -609,7 +613,7 @@ export async function startServer(opts: ServerOptions = {}): Promise { }); }); - accountsRouter.delete("/:id", (req, res) => { + accountsRouter.delete("/:id", async (req, res) => { const { id } = req.params; // Refuse to remove the last account — downstream /v1/* would have no // token to route with and the pool would throw EmptyPoolError on the @@ -624,7 +628,7 @@ export async function startServer(opts: ServerOptions = {}): Promise { return; } try { - deleteAnthropicAccountTransaction({ + await deleteAnthropicAccountTransaction({ id, pool, sessionRouter, @@ -632,6 +636,10 @@ export async function startServer(opts: ServerOptions = {}): Promise { }); } catch (err) { const message = err instanceof Error ? err.message : String(err); + if (err instanceof LastAccountDeletionError) { + res.status(409).json({ error: message }); + return; + } logError("accounts", 0, `Failed to persist accounts.json: ${message}`); res.status(500).json({ error: `Failed to persist accounts.json: ${message}` }); return; @@ -750,9 +758,7 @@ export async function startServer(opts: ServerOptions = {}): Promise { pendingLog.details = "token invalid"; logError(account.id, 401, "Token invalid — scheduling background refresh"); - refreshAccountToken(account).then(ok => { - if (ok) saveAccounts(pool.getAll()); - }).catch(console.error); + void refreshAccountIfCurrent(account, pool).catch(console.error); } else if (status === 429) { // Rate limited — put account on cooldown for Retry-After seconds. stats.totalErrors++; @@ -884,11 +890,7 @@ export async function startServer(opts: ServerOptions = {}): Promise { }, }), createAnthropicRefreshMiddleware({ needsRefresh, - refresh: async (account) => { - const ok = await refreshAccountToken(account); - if (ok) saveAccounts(pool.getAll()); - return ok; - }, + refresh: account => refreshAccountIfCurrent(account, pool), onRefreshFailure: (account) => { stats.totalErrors++; logError(account.id, 401, "Token refresh failed"); @@ -911,7 +913,7 @@ export async function startServer(opts: ServerOptions = {}): Promise { method: req.method, path: req.path, source, - details: route.reason, + details: routeReasonDetails(route), }; stats.totalRequests++; diff --git a/src/proxy/token-refresher.ts b/src/proxy/token-refresher.ts index 144be7b..deb3c6d 100644 --- a/src/proxy/token-refresher.ts +++ b/src/proxy/token-refresher.ts @@ -22,8 +22,8 @@ const REFRESH_BUFFER_MS = 10 * 60 * 1000; /** Check every 5 minutes */ const CHECK_INTERVAL_MS = 5 * 60 * 1000; -/** Per-account refresh locks — prevent concurrent refreshes for the same account */ -const refreshLocks = new Map>(); +/** Exact-object refresh locks prevent stale account incarnations from sharing work. */ +const refreshLocks = new Map>(); export function needsRefresh(account: Account): boolean { return (account.tokens.expiresAt - Date.now()) < REFRESH_BUFFER_MS; @@ -31,18 +31,47 @@ export function needsRefresh(account: Account): boolean { export async function refreshAccountToken(account: Account): Promise { // Deduplicate concurrent refresh calls for the same account - const existing = refreshLocks.get(account.id); + const existing = refreshLocks.get(account); if (existing) return existing; const promise = _doRefresh(account); - refreshLocks.set(account.id, promise); + refreshLocks.set(account, promise); try { return await promise; } finally { - refreshLocks.delete(account.id); + if (refreshLocks.get(account) === promise) refreshLocks.delete(account); } } +/** Wait for refresh work owned by this exact account incarnation, if any. */ +export async function waitForAccountRefresh(account: Account): Promise { + const activeRefresh = refreshLocks.get(account); + if (activeRefresh) await activeRefresh; +} + +export interface AccountOwnershipView { + findById(id: string): Account | null; + getAll(): Account[]; +} + +export interface RefreshAccountIfCurrentOptions { + refresh?: (account: Account) => Promise; + persist?: (accounts: Account[]) => void; +} + +/** Refresh and persist only while the pool still owns this exact object. */ +export async function refreshAccountIfCurrent( + account: Account, + pool: AccountOwnershipView, + options: RefreshAccountIfCurrentOptions = {}, +): Promise { + if (pool.findById(account.id) !== account) return false; + const ok = await (options.refresh ?? refreshAccountToken)(account); + if (!ok || pool.findById(account.id) !== account) return false; + (options.persist ?? saveAccounts)(pool.getAll()); + return true; +} + async function _doRefresh(account: Account): Promise { try { const body = new URLSearchParams({ From 18558881b6f8cb72f5eb106c8bfd8985ccf3c1e5 Mon Sep 17 00:00:00 2001 From: timo <44401485+Timo972@users.noreply.github.com> Date: Mon, 20 Jul 2026 20:42:33 +0200 Subject: [PATCH 16/17] fix: reserve refresh ownership during deletion --- src/__tests__/account-deletion.test.ts | 74 ++++++++++++- src/__tests__/lease-lifecycle.test.ts | 17 +++ src/__tests__/token-refresher.test.ts | 73 +++++++++++++ src/proxy/account-deletion.ts | 46 +++++--- src/proxy/lease-lifecycle.ts | 14 +++ src/proxy/server.ts | 22 +++- src/proxy/token-refresher.ts | 145 +++++++++++++++++++++---- 7 files changed, 349 insertions(+), 42 deletions(-) diff --git a/src/__tests__/account-deletion.test.ts b/src/__tests__/account-deletion.test.ts index cc20360..a88c972 100644 --- a/src/__tests__/account-deletion.test.ts +++ b/src/__tests__/account-deletion.test.ts @@ -1,11 +1,16 @@ import { describe, expect, it, vi } from "vitest"; import { + AccountDeletionConflictError, + accountDeletionStatusCode, deleteAnthropicAccountTransaction, LastAccountDeletionError, } from "../proxy/account-deletion.js"; import { SessionRouter } from "../proxy/session-router.js"; import { TokenPool } from "../proxy/token-pool.js"; -import { refreshAccountToken } from "../proxy/token-refresher.js"; +import { + refreshAccountIfCurrent, + refreshAccountToken, +} from "../proxy/token-refresher.js"; import type { Account } from "../proxy/types.js"; import { DEFAULT_RATE_LIMITS } from "../proxy/types.js"; @@ -52,6 +57,12 @@ function successfulRefreshResponse(): Response { } describe("deleteAnthropicAccountTransaction", () => { + it("maps typed deletion conflicts to 409 and persistence failures to 500", () => { + expect(accountDeletionStatusCode(new AccountDeletionConflictError("a"))).toBe(409); + expect(accountDeletionStatusCode(new LastAccountDeletionError())).toBe(409); + expect(accountDeletionStatusCode(new Error("disk full"))).toBe(500); + }); + it("leaves the exact runtime state untouched when prospective persistence fails", async () => { let now = 0; const first = makeAccount("a"); @@ -77,6 +88,16 @@ describe("deleteAnthropicAccountTransaction", () => { expect(pool.getInFlight("a")).toBe(1); expect(pool.isCoolingDown("a")).toBe(true); expect(router.getActiveSessionCount("a")).toBe(1); + + const refreshAfterFailure = vi.fn(async () => true); + const persistAfterFailure = vi.fn(); + expect(await refreshAccountIfCurrent(first, pool, { + refresh: refreshAfterFailure, + persist: persistAfterFailure, + })).toBe(true); + expect(refreshAfterFailure).toHaveBeenCalledWith(first); + expect(persistAfterFailure).toHaveBeenCalledWith([first, second]); + now = 60_000; const sticky = router.acquire("session-a"); expect(sticky.account).toBe(first); @@ -174,7 +195,7 @@ describe("deleteAnthropicAccountTransaction", () => { response.resolve(successfulRefreshResponse()); await refresh; - await expect(deletion).rejects.toThrow(/changed during deletion/); + await expect(deletion).rejects.toBeInstanceOf(AccountDeletionConflictError); expect(persist).not.toHaveBeenCalled(); expect(pool.findById("a")).toBe(replacement); expect(replacement.tokens.refreshToken).toBe("replacement-refresh"); @@ -213,4 +234,53 @@ describe("deleteAnthropicAccountTransaction", () => { .toBeInstanceOf(LastAccountDeletionError); expect(pool.getAll()).toHaveLength(1); }); + + it("blocks a P2 refresh reaction until P1 persistence and deletion complete", async () => { + const first = makeAccount("a"); + const second = makeAccount("b"); + const pool = new TokenPool([first, second]); + const router = new SessionRouter(pool); + const p1Gate = deferred(); + const p2Refresh = vi.fn(async () => true); + const events: string[] = []; + let persistedRotatedToken = ""; + + const p1 = refreshAccountIfCurrent(first, pool, { + refresh: async (account) => { + await p1Gate.promise; + account.tokens.refreshToken = "rotated-refresh"; + return true; + }, + persist: (accounts) => { + expect(pool.findById("a")).toBe(first); + persistedRotatedToken = accounts[0].tokens.refreshToken; + events.push("p1-persist"); + }, + }); + const p2 = p1.then(() => refreshAccountIfCurrent(first, pool, { + refresh: p2Refresh, + persist: () => events.push("p2-persist"), + })); + const deletion = deleteAnthropicAccountTransaction({ + id: "a", + pool, + sessionRouter: router, + persist: () => events.push("delete-persist"), + }); + await Promise.resolve(); + await Promise.resolve(); + const removedBeforeP1Completed = pool.findById("a") === null; + + p1Gate.resolve(true); + const [p1Result, p2Result, removed] = await Promise.all([p1, p2, deletion]); + + expect(removedBeforeP1Completed).toBe(false); + expect(p1Result).toBe(true); + expect(p2Result).toBe(false); + expect(p2Refresh).not.toHaveBeenCalled(); + expect(removed).toBe(first); + expect(persistedRotatedToken).toBe("rotated-refresh"); + expect(events).toEqual(["p1-persist", "delete-persist"]); + expect(pool.findById("a")).toBeNull(); + }); }); diff --git a/src/__tests__/lease-lifecycle.test.ts b/src/__tests__/lease-lifecycle.test.ts index 7625b70..47409b2 100644 --- a/src/__tests__/lease-lifecycle.test.ts +++ b/src/__tests__/lease-lifecycle.test.ts @@ -4,6 +4,7 @@ import { acquireRequestRoute, applyUpstreamFailureRouting, attachLeaseLifecycle, + routeFailureDetails, routeReasonDetails, } from "../proxy/lease-lifecycle.js"; import { SessionRouter } from "../proxy/session-router.js"; @@ -126,6 +127,22 @@ describe("routeReasonDetails", () => { expect(details).not.toContain("private-fallback-session"); route.release(); }); + + it.each(["rate-limited", "proxy-error"] as const)( + "retains fallback routing details after a %s failure without exposing the session", + (failure) => { + const account = makeAccount("account-a"); + account.healthy = false; + const router = new SessionRouter(new TokenPool([account])); + const route = router.acquire("private-fallback-session"); + + const details = routeFailureDetails(route, failure); + + expect(details).toBe(`new-session:fallback:${failure}`); + expect(details).not.toContain("private-fallback-session"); + route.release(); + }, + ); }); describe("applyUpstreamFailureRouting", () => { diff --git a/src/__tests__/token-refresher.test.ts b/src/__tests__/token-refresher.test.ts index 65cd613..27337a5 100644 --- a/src/__tests__/token-refresher.test.ts +++ b/src/__tests__/token-refresher.test.ts @@ -1,8 +1,10 @@ import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; import { needsRefresh, + refreshAccountsOnce, refreshAccountIfCurrent, refreshAccountToken, + reserveAccountForDeletion, } from "../proxy/token-refresher.js"; import { TokenPool } from "../proxy/token-pool.js"; import type { Account } from "../proxy/types.js"; @@ -276,4 +278,75 @@ describe("refreshAccountToken", () => { expect(persist).not.toHaveBeenCalled(); expect(replacement.tokens.refreshToken).toBe("replacement-refresh"); }); + + it("retries durability without another upstream refresh after persistence fails", async () => { + const account = makeAccount(Date.now() + 5 * 60 * 1000); + const pool = new TokenPool([account]); + vi.mocked(fetch).mockResolvedValue(successfulRefreshResponse("ROTATED")); + const persist = vi.fn() + .mockImplementationOnce(() => { throw new Error("disk full"); }) + .mockImplementationOnce(() => undefined); + + await expect(refreshAccountIfCurrent(account, pool, { persist })) + .rejects.toThrow("disk full"); + + expect(fetch).toHaveBeenCalledTimes(1); + expect(account.tokens.refreshToken).toBe("sk-ant-ort01-ROTATED"); + expect(needsRefresh(account)).toBe(true); + + expect(await refreshAccountIfCurrent(account, pool, { persist })).toBe(true); + expect(fetch).toHaveBeenCalledTimes(1); + expect(persist).toHaveBeenCalledTimes(2); + expect(account.tokens.refreshToken).toBe("sk-ant-ort01-ROTATED"); + expect(needsRefresh(account)).toBe(false); + }); + + it("scheduled and foreground refresh share ownership and refuse stale persistence", async () => { + const oldAccount = makeAccount(Date.now() + 5 * 60 * 1000); + const accounts = [oldAccount]; + const pool = new TokenPool(accounts); + const response = deferred(); + vi.mocked(fetch).mockImplementation(() => response.promise); + const persist = vi.fn(); + + const scheduled = refreshAccountsOnce(accounts, { persist }); + const foreground = refreshAccountIfCurrent(oldAccount, pool, { persist }); + expect(fetch).toHaveBeenCalledTimes(1); + pool.removeAccount(oldAccount.id); + const replacement = pool.addAccount({ + id: oldAccount.id, + accessToken: "replacement-access", + refreshToken: "replacement-refresh", + expiresAt: Date.now() + 60_000, + scopes: ["user:inference"], + }); + response.resolve(successfulRefreshResponse("STALE")); + + await scheduled; + expect(await foreground).toBe(false); + expect(fetch).toHaveBeenCalledTimes(1); + expect(persist).not.toHaveBeenCalled(); + expect(replacement.tokens.refreshToken).toBe("replacement-refresh"); + }); + + it("keeps an exact account blocked until every deletion reservation releases", async () => { + const account = makeAccount(Date.now() + 5 * 60 * 1000); + const pool = new TokenPool([account]); + vi.mocked(fetch).mockResolvedValue(successfulRefreshResponse("AFTER")); + const persist = vi.fn(); + + const [releaseFirst, releaseSecond] = await Promise.all([ + reserveAccountForDeletion(account), + reserveAccountForDeletion(account), + ]); + + expect(await refreshAccountToken(account)).toBe(false); + expect(await refreshAccountIfCurrent(account, pool, { persist })).toBe(false); + releaseFirst(); + releaseFirst(); + expect(await refreshAccountToken(account)).toBe(false); + releaseSecond(); + expect(await refreshAccountToken(account)).toBe(true); + expect(fetch).toHaveBeenCalledTimes(1); + }); }); diff --git a/src/proxy/account-deletion.ts b/src/proxy/account-deletion.ts index 1469c11..9d93807 100644 --- a/src/proxy/account-deletion.ts +++ b/src/proxy/account-deletion.ts @@ -1,7 +1,14 @@ import type { SessionRouter } from "./session-router.js"; import type { TokenPool } from "./token-pool.js"; import type { Account } from "./types.js"; -import { waitForAccountRefresh } from "./token-refresher.js"; +import { reserveAccountForDeletion } from "./token-refresher.js"; + +export class AccountDeletionConflictError extends Error { + constructor(id: string) { + super(`Account "${id}" changed during deletion`); + this.name = "AccountDeletionConflictError"; + } +} export class LastAccountDeletionError extends Error { constructor() { @@ -10,6 +17,13 @@ export class LastAccountDeletionError extends Error { } } +export function accountDeletionStatusCode(error: unknown): 409 | 500 { + return error instanceof AccountDeletionConflictError || + error instanceof LastAccountDeletionError + ? 409 + : 500; +} + export interface DeleteAnthropicAccountOptions { id: string; pool: TokenPool; @@ -24,20 +38,24 @@ export async function deleteAnthropicAccountTransaction( const account = options.pool.findById(options.id); if (!account) throw new Error(`Account "${options.id}" not found`); - await waitForAccountRefresh(account); - if (options.pool.findById(options.id) !== account) { - throw new Error(`Account "${options.id}" changed during deletion`); - } - if (options.pool.getAll().length <= 1) { - throw new LastAccountDeletionError(); - } + const releaseReservation = await reserveAccountForDeletion(account); + try { + if (options.pool.findById(options.id) !== account) { + throw new AccountDeletionConflictError(options.id); + } + if (options.pool.getAll().length <= 1) { + throw new LastAccountDeletionError(); + } - const prospective = options.pool.getAll().filter(candidate => candidate !== account); - options.persist(prospective); + const prospective = options.pool.getAll().filter(candidate => candidate !== account); + options.persist(prospective); - if (!options.pool.removeAccount(options.id)) { - throw new Error(`Account "${options.id}" disappeared during deletion`); + if (!options.pool.removeAccount(options.id)) { + throw new AccountDeletionConflictError(options.id); + } + options.sessionRouter.invalidateAccount(options.id); + return account; + } finally { + releaseReservation(); } - options.sessionRouter.invalidateAccount(options.id); - return account; } diff --git a/src/proxy/lease-lifecycle.ts b/src/proxy/lease-lifecycle.ts index 3b168b0..a3371c2 100644 --- a/src/proxy/lease-lifecycle.ts +++ b/src/proxy/lease-lifecycle.ts @@ -62,6 +62,20 @@ export function routeReasonDetails(route: RouteSummary): string { return route.fallback ? `${route.reason}:fallback` : route.reason; } +export type RouteFailureReason = + | "token-invalid" + | "rate-limited" + | "service-overloaded" + | "proxy-error"; + +/** Retain bounded routing context when a later failure updates the log. */ +export function routeFailureDetails( + route: RouteSummary, + failure: RouteFailureReason, +): string { + return `${routeReasonDetails(route)}:${failure}`; +} + /** Acquire and immediately bind a routed lease to its response lifecycle. */ export function acquireRequestRoute( sessionHeader: unknown, diff --git a/src/proxy/server.ts b/src/proxy/server.ts index 0618c8e..dcc16a3 100644 --- a/src/proxy/server.ts +++ b/src/proxy/server.ts @@ -30,12 +30,13 @@ import type { RoutedAccountLease } from "./session-router.js"; import { createAnthropicProxy } from "./anthropic-proxy.js"; import { applyUpstreamFailureRouting, + routeFailureDetails, routeReasonDetails, } from "./lease-lifecycle.js"; import { persistProviderEnabledState } from "./provider-routing.js"; import { + accountDeletionStatusCode, deleteAnthropicAccountTransaction, - LastAccountDeletionError, } from "./account-deletion.js"; import { createAnthropicRefreshMiddleware, @@ -636,7 +637,8 @@ export async function startServer(opts: ServerOptions = {}): Promise { }); } catch (err) { const message = err instanceof Error ? err.message : String(err); - if (err instanceof LastAccountDeletionError) { + const status = accountDeletionStatusCode(err); + if (status === 409) { res.status(409).json({ error: message }); return; } @@ -755,7 +757,9 @@ export async function startServer(opts: ServerOptions = {}): Promise { stats.totalErrors++; account.errorCount++; pendingLog.type = "error"; - pendingLog.details = "token invalid"; + pendingLog.details = route + ? routeFailureDetails(route, "token-invalid") + : "token-invalid"; logError(account.id, 401, "Token invalid — scheduling background refresh"); void refreshAccountIfCurrent(account, pool).catch(console.error); @@ -765,14 +769,18 @@ export async function startServer(opts: ServerOptions = {}): Promise { account.errorCount++; const retryAfter = cooldownSeconds ?? 60; pendingLog.type = "error"; - pendingLog.details = `rate limited — cooldown ${retryAfter}s`; + pendingLog.details = route + ? routeFailureDetails(route, "rate-limited") + : "rate-limited"; logError(account.id, 429, `Rate limited — cooldown ${retryAfter}s`); } else if (status === 529) { // Anthropic service overloaded — short cooldown on this account. stats.totalErrors++; account.errorCount++; pendingLog.type = "error"; - pendingLog.details = "service overloaded — cooldown 30s"; + pendingLog.details = route + ? routeFailureDetails(route, "service-overloaded") + : "service-overloaded"; logError(account.id, 529, "Service overloaded — cooldown 30s"); } @@ -855,7 +863,9 @@ export async function startServer(opts: ServerOptions = {}): Promise { if (pendingLog) { pendingLog.type = "error"; pendingLog.statusCode = 0; - pendingLog.details = err.message; + pendingLog.details = request._ccRoute + ? routeFailureDetails(request._ccRoute, "proxy-error") + : "proxy-error"; if (request._startTime) { pendingLog.durationMs = Date.now() - request._startTime; } diff --git a/src/proxy/token-refresher.ts b/src/proxy/token-refresher.ts index deb3c6d..54f2b87 100644 --- a/src/proxy/token-refresher.ts +++ b/src/proxy/token-refresher.ts @@ -22,31 +22,76 @@ const REFRESH_BUFFER_MS = 10 * 60 * 1000; /** Check every 5 minutes */ const CHECK_INTERVAL_MS = 5 * 60 * 1000; -/** Exact-object refresh locks prevent stale account incarnations from sharing work. */ -const refreshLocks = new Map>(); +/** Exact-object locks prevent stale account incarnations from sharing work. */ +const rawRefreshLocks = new Map>(); +const ownedRefreshLocks = new Map>(); + +/** A count is required because concurrent deletion attempts may reserve the same object. */ +const deletionReservations = new Map(); + +/** Rotated credentials that still need to be durably written must not rotate again. */ +const pendingDurability = new WeakSet(); + +function isReservedForDeletion(account: Account): boolean { + return (deletionReservations.get(account) ?? 0) > 0; +} export function needsRefresh(account: Account): boolean { - return (account.tokens.expiresAt - Date.now()) < REFRESH_BUFFER_MS; + return pendingDurability.has(account) || + (account.tokens.expiresAt - Date.now()) < REFRESH_BUFFER_MS; } export async function refreshAccountToken(account: Account): Promise { + // A deletion reservation rejects every new caller, including callers that + // would otherwise attach themselves to already-running raw refresh work. + if (isReservedForDeletion(account)) return false; + // Deduplicate concurrent refresh calls for the same account - const existing = refreshLocks.get(account); + const existing = rawRefreshLocks.get(account); if (existing) return existing; const promise = _doRefresh(account); - refreshLocks.set(account, promise); + rawRefreshLocks.set(account, promise); try { return await promise; } finally { - if (refreshLocks.get(account) === promise) refreshLocks.delete(account); + if (rawRefreshLocks.get(account) === promise) rawRefreshLocks.delete(account); } } /** Wait for refresh work owned by this exact account incarnation, if any. */ export async function waitForAccountRefresh(account: Account): Promise { - const activeRefresh = refreshLocks.get(account); - if (activeRefresh) await activeRefresh; + const activeRefresh = rawRefreshLocks.get(account); + const activeOwnedRefresh = ownedRefreshLocks.get(account); + await Promise.allSettled( + [activeRefresh, activeOwnedRefresh].filter( + (promise): promise is Promise => promise !== undefined, + ), + ); +} + +/** + * Prevent new refreshes for an exact account object and wait for all work that + * started before the reservation, including owned persistence, to settle. + */ +export async function reserveAccountForDeletion(account: Account): Promise<() => void> { + deletionReservations.set(account, (deletionReservations.get(account) ?? 0) + 1); + let released = false; + const release = () => { + if (released) return; + released = true; + const remaining = (deletionReservations.get(account) ?? 1) - 1; + if (remaining > 0) deletionReservations.set(account, remaining); + else deletionReservations.delete(account); + }; + + try { + await waitForAccountRefresh(account); + return release; + } catch (error) { + release(); + throw error; + } } export interface AccountOwnershipView { @@ -59,19 +104,59 @@ export interface RefreshAccountIfCurrentOptions { persist?: (accounts: Account[]) => void; } -/** Refresh and persist only while the pool still owns this exact object. */ -export async function refreshAccountIfCurrent( +async function performOwnedRefresh( account: Account, pool: AccountOwnershipView, - options: RefreshAccountIfCurrentOptions = {}, + options: RefreshAccountIfCurrentOptions, ): Promise { if (pool.findById(account.id) !== account) return false; + + if (pendingDurability.has(account)) { + (options.persist ?? saveAccounts)(pool.getAll()); + pendingDurability.delete(account); + return true; + } + const ok = await (options.refresh ?? refreshAccountToken)(account); if (!ok || pool.findById(account.id) !== account) return false; - (options.persist ?? saveAccounts)(pool.getAll()); + try { + (options.persist ?? saveAccounts)(pool.getAll()); + } catch (error) { + pendingDurability.add(account); + throw error; + } + pendingDurability.delete(account); return true; } +/** + * Refresh and persist only while the pool still owns this exact object. + * Production callers for the same object coalesce through persistence. + */ +export function refreshAccountIfCurrent( + account: Account, + pool: AccountOwnershipView, + options: RefreshAccountIfCurrentOptions = {}, +): Promise { + if (isReservedForDeletion(account)) return Promise.resolve(false); + + const existing = ownedRefreshLocks.get(account); + if (existing) return existing; + + let operation!: Promise; + operation = (async () => { + try { + return await performOwnedRefresh(account, pool, options); + } finally { + if (ownedRefreshLocks.get(account) === operation) { + ownedRefreshLocks.delete(account); + } + } + })(); + ownedRefreshLocks.set(account, operation); + return operation; +} + async function _doRefresh(account: Account): Promise { try { const body = new URLSearchParams({ @@ -130,19 +215,39 @@ export function saveAccounts(accounts: Account[]): void { writeAnthropicAccountsPreservingOtherProviders(serialize(accounts)); } +export interface RefreshAccountsOnceOptions { + persist?: (accounts: Account[]) => void; + onError?: (error: unknown) => void; +} + +/** Run one ownership-aware scheduled refresh pass. */ +export async function refreshAccountsOnce( + accounts: Account[], + options: RefreshAccountsOnceOptions = {}, +): Promise { + const ownershipView: AccountOwnershipView = { + findById: id => accounts.find(account => account.id === id) ?? null, + getAll: () => accounts, + }; + + for (const account of [...accounts]) { + if (!needsRefresh(account)) continue; + try { + await refreshAccountIfCurrent(account, ownershipView, { + persist: options.persist, + }); + } catch (error) { + (options.onError ?? console.error)(error); + } + } +} + /** * Background refresh loop: checks every 5 minutes and refreshes any * token expiring within the REFRESH_BUFFER_MS window. */ export function startRefreshLoop(accounts: Account[]): void { - const check = async () => { - for (const account of accounts) { - if (needsRefresh(account)) { - const ok = await refreshAccountToken(account); - if (ok) saveAccounts(accounts); - } - } - }; + const check = () => refreshAccountsOnce(accounts); // Run immediately on startup (catches already-expired tokens) check().catch(console.error); From 84825384f44f6c5cb41c75d308df622b5801ee0d Mon Sep 17 00:00:00 2001 From: timo <44401485+Timo972@users.noreply.github.com> Date: Mon, 20 Jul 2026 21:04:08 +0200 Subject: [PATCH 17/17] fix: join active refresh during request preparation --- README.md | 2 +- src/__tests__/anthropic-routing.test.ts | 87 +++++++++++++++++++++++++ src/proxy/token-refresher.ts | 3 +- 3 files changed, 90 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index cb7c97d..a8637cf 100644 --- a/README.md +++ b/README.md @@ -547,7 +547,7 @@ Then start the interceptor: cc-router client start-desktop ``` -Open Claude Desktop and send a message. The request will be intercepted and redirected to CC-Router, which applies the same cache-aware account routing used for Claude Code traffic. +Open Claude Desktop and send a message. The request will be intercepted and redirected to CC-Router. Requests carrying exactly one valid `X-Claude-Code-Session-Id` receive cache-aware sticky affinity; requests without one valid session header use load-aware **unscoped** routing and do not receive sticky affinity. Claude Desktop traffic normally follows the unscoped path. ### Stopping / removing Desktop interception diff --git a/src/__tests__/anthropic-routing.test.ts b/src/__tests__/anthropic-routing.test.ts index 8dd63e3..0c5f2d7 100644 --- a/src/__tests__/anthropic-routing.test.ts +++ b/src/__tests__/anthropic-routing.test.ts @@ -8,6 +8,10 @@ import { } from "../proxy/anthropic-routing.js"; import { SessionRouter } from "../proxy/session-router.js"; import { TokenPool } from "../proxy/token-pool.js"; +import { + needsRefresh, + refreshAccountIfCurrent, +} from "../proxy/token-refresher.js"; import type { Account } from "../proxy/types.js"; import { DEFAULT_RATE_LIMITS } from "../proxy/types.js"; @@ -61,6 +65,19 @@ function deferred(): { return { promise, resolve }; } +function successfulRefreshResponse(): Response { + return { + ok: true, + json: async () => ({ + access_token: "rotated-access", + refresh_token: "rotated-refresh", + expires_in: 28800, + scope: "user:inference", + token_type: "Bearer", + }), + } as Response; +} + function send(options: Parameters[0]): Promise<{ status: number; body: string; @@ -168,4 +185,74 @@ describe("production Anthropic routing middleware", () => { await close(server); } }); + + it("joins an active forced refresh before forwarding a far-future token retry", async () => { + const account = makeAccount("a"); + account.tokens.expiresAt = Date.now() + 2 * 60 * 60 * 1000; + const pool = new TokenPool([account]); + const sessionRouter = new SessionRouter(pool); + const refreshResponse = deferred(); + vi.stubGlobal("fetch", vi.fn(() => refreshResponse.promise)); + let persisted = false; + const persist = vi.fn(() => { persisted = true; }); + const forcedRefresh = refreshAccountIfCurrent(account, pool, { persist }); + const preparationCheck = deferred(); + let joinedRefresh: Promise | undefined; + let tokenAtForward = ""; + let persistedAtForward = false; + const forwarded = vi.fn(); + + const app = express(); + app.use(createAnthropicRoutingMiddleware({ sessionRouter })); + app.use(createAnthropicRefreshMiddleware({ + needsRefresh: selected => { + const required = needsRefresh(selected); + preparationCheck.resolve(required); + return required; + }, + refresh: selected => { + joinedRefresh = refreshAccountIfCurrent(selected, pool, { persist }); + return joinedRefresh; + }, + onRefreshFailure: vi.fn(), + })); + app.use((_req, res) => { + forwarded(); + tokenAtForward = account.tokens.accessToken; + persistedAtForward = persisted; + res.end("forwarded"); + }); + const server = createServer(app); + const port = await listen(server); + const retry = send({ + host: "127.0.0.1", + port, + path: "/v1/messages", + headers: { "X-Claude-Code-Session-Id": "session-a" }, + }); + + try { + const preparationNeeded = await preparationCheck.promise; + await new Promise(resolve => setImmediate(resolve)); + + expect(preparationNeeded).toBe(true); + expect(forwarded).not.toHaveBeenCalled(); + expect(fetch).toHaveBeenCalledTimes(1); + + refreshResponse.resolve(successfulRefreshResponse()); + expect(await forcedRefresh).toBe(true); + expect((await retry).status).toBe(200); + + expect(joinedRefresh).toBe(forcedRefresh); + expect(fetch).toHaveBeenCalledTimes(1); + expect(persist).toHaveBeenCalledTimes(1); + expect(tokenAtForward).toBe("rotated-access"); + expect(persistedAtForward).toBe(true); + } finally { + refreshResponse.resolve(successfulRefreshResponse()); + await Promise.allSettled([forcedRefresh, retry]); + vi.unstubAllGlobals(); + await close(server); + } + }); }); diff --git a/src/proxy/token-refresher.ts b/src/proxy/token-refresher.ts index 54f2b87..6b0eb5b 100644 --- a/src/proxy/token-refresher.ts +++ b/src/proxy/token-refresher.ts @@ -37,7 +37,8 @@ function isReservedForDeletion(account: Account): boolean { } export function needsRefresh(account: Account): boolean { - return pendingDurability.has(account) || + return ownedRefreshLocks.has(account) || + pendingDurability.has(account) || (account.tokens.expiresAt - Date.now()) < REFRESH_BUFFER_MS; }