diff --git a/packages/templates/src/templates/bb-guide-plugins.md b/packages/templates/src/templates/bb-guide-plugins.md index 8f0e61307b..a1dbbd8956 100644 --- a/packages/templates/src/templates/bb-guide-plugins.md +++ b/packages/templates/src/templates/bb-guide-plugins.md @@ -95,6 +95,20 @@ Authentication failures and temporary upstream failures can fail over to another eligible account before a response reaches the client. The hub never replays a response after it starts streaming to the client. +Conversation affinity keeps requests on the same eligible account, scoped by +provider and host. Claude uses the session ID in JSON or legacy +`metadata.user_id`. Codex uses its native `session-id` header, then the existing +`session_id` form, then the body `prompt_cache_key`. Session IDs and cache keys +have separate namespaces. Requests without a usable identifier follow the +priority, in-flight count, and weekly-reset order. + +The hub rebinds a conversation when its account becomes ineligible. It keeps +up to 4,096 bindings in memory, expires them after 30 minutes of inactivity, +and resets them on restart. Codex's native `session-id` and `thread-id` headers, +cache fields, and Responses-based compaction payloads pass through, including +encrypted items. Affinity avoids unnecessary account switches but does not +guarantee an upstream prompt-cache hit. + The builtin Keep Awake plugin prevents macOS idle sleep while bb is running. Its settings page lets you target all hosts or selected hosts. The CLI equivalents are: diff --git a/plugins/account-pool/PLUGIN_OVERVIEW.md b/plugins/account-pool/PLUGIN_OVERVIEW.md index e294f248d1..036ced570c 100644 --- a/plugins/account-pool/PLUGIN_OVERVIEW.md +++ b/plugins/account-pool/PLUGIN_OVERVIEW.md @@ -3,7 +3,7 @@ Keep a Claude Code or Codex thread running when one account hits its limit. The ## What you get - A pool of Claude and Codex accounts, added by importing the login already on the machine, signing in through the browser, or pasting an Anthropic API key. -- Per-request selection that follows your priority order, then the account with the fewest requests in flight, then the account whose weekly window resets first. +- New conversations follow your priority order, then the account with the fewest requests in flight, then the account whose weekly window resets first. A conversation stays on its account while that account remains eligible. - Live limit windows per account and model family in the plugin's settings page, and the same numbers from `bb pool status`. - A routing switch per provider and a bypass per thread, so one thread can go straight to its own credentials. @@ -15,6 +15,8 @@ Concurrent requests share one OAuth refresh per account. During a temporary refr If an OAuth request receives HTTP 401, the hub refreshes its credential once and retries. It reuses a token already refreshed by another request. Authentication failures and temporary upstream failures can move the request to another eligible account before a response reaches the client. Once a response starts, the hub does not replay it on another account. +Conversation affinity uses Claude's session metadata or Codex's session header, with its prompt cache key as a fallback. Bindings are separate for each provider and host. The hub chooses another account when the bound account becomes ineligible. Bindings expire after 30 minutes without a request and reset when the hub restarts. Requests without a usable identifier follow ordinary account selection. + ## Requirements Accounts you own and are permitted to use this way. diff --git a/plugins/account-pool/src/claude-adapter.ts b/plugins/account-pool/src/claude-adapter.ts index 133cd0bdcf..e366f2e949 100644 --- a/plugins/account-pool/src/claude-adapter.ts +++ b/plugins/account-pool/src/claude-adapter.ts @@ -73,9 +73,14 @@ export function createClaudeAdapter(options: { }, }; }, - modelFamily: (body) => parseRequestBody(body).family, - prepareBody: (body, account) => - parseRequestBody(body).forAccount(account.accountUuid), + parseRequest(body) { + const parsed = parseRequestBody(body); + return { + family: parsed.family, + affinityId: parsed.affinityId, + forAccount: (account) => parsed.forAccount(account.accountUuid), + }; + }, upstreamUrl: (request, settings) => mountedUpstreamUrl(request, settings.anthropicUpstreamBaseUrl), requestHeaders(inbound, _account, secret) { diff --git a/plugins/account-pool/src/codex-adapter.ts b/plugins/account-pool/src/codex-adapter.ts index 6dd556c7d6..09260e2e6f 100644 --- a/plugins/account-pool/src/codex-adapter.ts +++ b/plugins/account-pool/src/codex-adapter.ts @@ -1,4 +1,5 @@ import { z } from "zod"; +import { parseCodexRequestBody } from "./request-body.js"; import type { AccountQuota, AccountSecret, @@ -30,7 +31,9 @@ const ALLOWED_REQUEST_HEADERS = new Set([ "content-type", "openai-beta", "originator", + "session-id", "session_id", + "thread-id", "user-agent", ]); const ALLOWED_REQUEST_HEADER_PREFIXES = ["x-codex-", "x-stainless-"]; @@ -257,8 +260,14 @@ export function createCodexAdapter(options: { }, }; }, - modelFamily: () => "other", - prepareBody: (body) => body, + parseRequest(body, headers) { + const parsed = parseCodexRequestBody(body, headers); + return { + family: parsed.family, + affinityId: parsed.affinityId, + forAccount: () => body, + }; + }, upstreamUrl: (request, settings) => mountedUpstreamUrl(request, settings.codexUpstreamBaseUrl, "v1/"), requestHeaders(inbound, account, secret) { diff --git a/plugins/account-pool/src/hub.ts b/plugins/account-pool/src/hub.ts index c5ac772076..2ea4a16346 100644 --- a/plugins/account-pool/src/hub.ts +++ b/plugins/account-pool/src/hub.ts @@ -37,6 +37,8 @@ const MAX_REFRESH_BACKOFF_MS = 60_000; const MAX_REFRESH_BACKOFFS = 1_024; const MAX_FAILURE_DETAIL_BYTES = 1_024; const FAILURE_DISPOSAL_TIMEOUT_MS = 250; +const AFFINITY_IDLE_TTL_MS = 30 * 60 * 1_000; +const MAX_AFFINITY_BINDINGS = 4_096; const DROPPED_RESPONSE_HEADERS = new Set([ "content-encoding", "content-length", @@ -108,6 +110,10 @@ export class AccountPoolHub { private readonly activeControllers = new Set(); private readonly refreshes = new Map(); private readonly refreshBackoffs = new Map(); + private readonly affinityBindings = new Map< + string, + { accountId: string; lastUsedAt: number } + >(); private readonly usageRefreshes = new Map>(); private readonly lastUsageRefreshAt = new Map(); private readonly drainWaiters = new Set<() => void>(); @@ -115,6 +121,7 @@ export class AccountPoolHub { constructor(private readonly options: HubOptions) {} async start(signal: AbortSignal): Promise { + this.affinityBindings.clear(); this.stopped = new AbortController(); this.accepting = true; while (!signal.aborted) { @@ -280,12 +287,18 @@ export class AccountPoolHub { ): Promise { const signal = AbortSignal.any([request.signal, this.stopped.signal]); const attempted = new Set(); + let previousAccountId: string | null = null; let failure: FailureSummary | null = null; const accounts = (await this.options.accounts.list()).filter( (account) => account.provider === adapter.provider, ); const candidateIds = new Set(accounts.map((account) => account.id)); - const family = adapter.modelFamily(body); + const parsed = adapter.parseRequest(body, request.headers); + const family = parsed.family; + const affinityKey = + hostId === null || parsed.affinityId === null + ? null + : JSON.stringify([adapter.provider, hostId, parsed.affinityId]); try { while (attempted.size < candidateIds.size) { signal.throwIfAborted(); @@ -294,8 +307,12 @@ export class AccountPoolHub { candidateIds, attempted, family, + affinityKey, + previousAccountId, + signal, ); if (selected === null) break; + previousAccountId = selected.account.id; attempted.add(selected.account.id); if (hostId !== null) { const changed = await this.options.accounts.recordUsed( @@ -329,7 +346,7 @@ export class AccountPoolHub { try { upstream = await this.fetchUpstream( request, - adapter.prepareBody(body, selected.account), + parsed.forAccount(selected.account), selected.account, secret, adapter, @@ -567,17 +584,16 @@ export class AccountPoolHub { candidateIds: ReadonlySet, attempted: ReadonlySet, family: ModelFamily, + affinityKey: string | null, + previousAccountId: string | null, + signal: AbortSignal, ): Promise { + const accounts = await this.options.accounts.list(); + signal.throwIfAborted(); const now = this.options.now(); const threshold = this.options.getSettings().switchThreshold; - const candidates = (await this.options.accounts.list()) - .filter( - (account) => - account.provider === provider && - candidateIds.has(account.id) && - account.enabled && - !attempted.has(account.id), - ) + const eligible = accounts + .filter((account) => account.provider === provider && account.enabled) .map((account) => ({ account, quota: this.options.quotas.get(account.id), @@ -585,6 +601,10 @@ export class AccountPoolHub { .filter(({ quota }) => quota.error === null) .filter(({ quota }) => quota.heldUntil === null || quota.heldUntil <= now) .filter(({ quota }) => !isQuotaExhausted(quota, family, threshold, now)); + const candidates = eligible.filter( + ({ account }) => + candidateIds.has(account.id) && !attempted.has(account.id), + ); candidates.sort((left, right) => { const priority = left.account.priority - right.account.priority; if (priority !== 0) return priority; @@ -598,7 +618,34 @@ export class AccountPoolHub { (governingWeeklyResetAt(right.quota, family) ?? Number.MAX_SAFE_INTEGER) ); }); - return candidates[0] ?? null; + const binding = + affinityKey === null ? undefined : this.affinityBindings.get(affinityKey); + const bound = + binding !== undefined && now - binding.lastUsedAt < AFFINITY_IDLE_TTL_MS + ? eligible.find(({ account }) => account.id === binding.accountId) + : undefined; + const selected = + bound !== undefined && candidates.includes(bound) + ? bound + : (candidates[0] ?? null); + if ( + affinityKey !== null && + selected !== null && + (bound === undefined || + bound.account.id === selected.account.id || + bound.account.id === previousAccountId) + ) { + this.affinityBindings.delete(affinityKey); + this.affinityBindings.set(affinityKey, { + accountId: selected.account.id, + lastUsedAt: now, + }); + while (this.affinityBindings.size > MAX_AFFINITY_BINDINGS) { + const oldest = this.affinityBindings.keys().next(); + if (!oldest.done) this.affinityBindings.delete(oldest.value); + } + } + return selected; } private async freshSecret( diff --git a/plugins/account-pool/src/provider-adapter.ts b/plugins/account-pool/src/provider-adapter.ts index d5dd56a52f..a7abef493d 100644 --- a/plugins/account-pool/src/provider-adapter.ts +++ b/plugins/account-pool/src/provider-adapter.ts @@ -53,8 +53,14 @@ export interface ProviderAdapter { provider: PoolProvider; upstreamName: string; importAccount(): Promise; - modelFamily(body: Uint8Array): ModelFamily; - prepareBody(body: Uint8Array, account: Account): Uint8Array; + parseRequest( + body: Uint8Array, + headers: Headers, + ): { + family: ModelFamily; + affinityId: string | null; + forAccount: (account: Account) => Uint8Array; + }; upstreamUrl(request: Request, settings: HubSettings): URL; requestHeaders( inbound: Headers, diff --git a/plugins/account-pool/src/request-body.test.ts b/plugins/account-pool/src/request-body.test.ts new file mode 100644 index 0000000000..473ec8993b --- /dev/null +++ b/plugins/account-pool/src/request-body.test.ts @@ -0,0 +1,191 @@ +import { describe, expect, it } from "vitest"; +import { parseCodexRequestBody, parseRequestBody } from "./request-body.js"; + +const accountUuid = "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa"; +const nextAccountUuid = "bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb"; +const sessionId = "cccccccc-cccc-4ccc-8ccc-cccccccccccc"; +const encode = (value: object) => + new TextEncoder().encode(JSON.stringify(value)); + +describe("Claude request parsing", () => { + it.each([accountUuid, "invalid-account-uuid", 123, null])( + "extracts the own session independently of account UUID %s", + (account) => { + const body = encode({ + model: "claude-fable-5", + metadata: { + user_id: JSON.stringify({ + account_uuid: account, + session_id: sessionId, + parent_session_id: "parent", + device_id: "device", + }), + }, + }); + const parsed = parseRequestBody(body); + expect(parsed.family).toBe("fable"); + expect(parsed.affinityId).toBe(`session:${sessionId}`); + if (account === "invalid-account-uuid" || account === 123) + expect(parsed.forAccount(nextAccountUuid)).toBe(body); + }, + ); + + it.each([ + undefined, + null, + "", + " ", + 123, + [], + { nested: "session" }, + "a\nb", + "a\u007fb", + "a\u0085b", + "x".repeat(513), + ])( + "rejects invalid session %j while preserving it and extra metadata during account rewrite", + (session) => { + const user = { + account_uuid: accountUuid, + session_id: session, + parent_session_id: sessionId, + device_id: sessionId, + extension: { keep: true }, + }; + const request = { + model: "claude-fable-5", + metadata: { user_id: JSON.stringify(user), extra: "keep" }, + messages: [{ role: "user", content: "message" }], + }; + const parsed = parseRequestBody(encode(request)); + expect(parsed.affinityId).toBeNull(); + const rewritten = JSON.parse( + new TextDecoder().decode(parsed.forAccount(nextAccountUuid)), + ); + expect(rewritten).toEqual({ + ...request, + metadata: { + ...request.metadata, + user_id: JSON.stringify({ ...user, account_uuid: nextAccountUuid }), + }, + }); + }, + ); + + it("accepts a 512-character session and preserves legacy metadata except the account UUID", () => { + const id = "x".repeat(512); + expect( + parseRequestBody( + encode({ metadata: { user_id: JSON.stringify({ session_id: id }) } }), + ).affinityId, + ).toBe(`session:${id}`); + const body = encode({ + metadata: { + user_id: `user_hash_account_${accountUuid}_session_${sessionId}`, + extra: "keep", + }, + }); + const parsed = parseRequestBody(body); + expect(parsed.affinityId).toBe(`session:${sessionId}`); + expect( + JSON.parse(new TextDecoder().decode(parsed.forAccount(nextAccountUuid))), + ).toEqual({ + metadata: { + user_id: `user_hash_account_${nextAccountUuid}_session_${sessionId}`, + extra: "keep", + }, + }); + }); + + it.each([ + "not-json", + "[]", + '{"metadata":{"user_id":"malformed"}}', + '{"metadata":{"user_id":123}}', + ])("leaves malformed request %s unbound and byte-identical", (raw) => { + const body = new TextEncoder().encode(raw); + const parsed = parseRequestBody(body); + expect(parsed.affinityId).toBeNull(); + expect(parsed.forAccount(nextAccountUuid)).toBe(body); + }); +}); + +describe("Codex request parsing", () => { + it.each<{ + headers: Record; + cacheKey: string | null; + expected: string | null; + }>([ + { + headers: { "session-id": "native", session_id: "legacy" }, + cacheKey: "cache", + expected: "session:native", + }, + { + headers: { session_id: "legacy" }, + cacheKey: "cache", + expected: "session:legacy", + }, + { + headers: { "session-id": "", session_id: "legacy" }, + cacheKey: "cache", + expected: "session:legacy", + }, + { + headers: { "session-id": "x".repeat(513), session_id: "legacy" }, + cacheKey: "cache", + expected: "session:legacy", + }, + { + headers: { "session-id": "bad\u007fid" }, + cacheKey: "cache", + expected: "cache:cache", + }, + { headers: { "thread-id": "thread" }, cacheKey: null, expected: null }, + { headers: {}, cacheKey: "same", expected: "cache:same" }, + { + headers: { "session-id": "same" }, + cacheKey: "same", + expected: "session:same", + }, + { headers: {}, cacheKey: "", expected: null }, + { headers: {}, cacheKey: "bad\nid", expected: null }, + { headers: {}, cacheKey: "x".repeat(513), expected: null }, + { + headers: {}, + cacheKey: "x".repeat(512), + expected: `cache:${"x".repeat(512)}`, + }, + ])( + "resolves affinity $expected without changing request bytes", + ({ headers, cacheKey, expected }) => { + const body = new TextEncoder().encode( + JSON.stringify( + { + prompt_cache_key: cacheKey, + client_metadata: { session_id: sessionId }, + input: [ + { type: "compaction_trigger" }, + { type: "compaction", encrypted_content: "fixture" }, + ], + }, + null, + 2, + ), + ); + const parsed = parseCodexRequestBody(body, new Headers(headers)); + expect(parsed.affinityId).toBe(expected); + expect(parsed.forAccount(nextAccountUuid)).toBe(body); + }, + ); + + it.each(["not-json", '{"prompt_cache_key":123}', '{"prompt_cache_key":[]}'])( + "does not infer cache affinity from malformed payload %s", + (raw) => { + const body = new TextEncoder().encode(raw); + const parsed = parseCodexRequestBody(body, new Headers()); + expect(parsed.affinityId).toBeNull(); + expect(parsed.forAccount(null)).toBe(body); + }, + ); +}); diff --git a/plugins/account-pool/src/request-body.ts b/plugins/account-pool/src/request-body.ts index fed450c832..699aacdba9 100644 --- a/plugins/account-pool/src/request-body.ts +++ b/plugins/account-pool/src/request-body.ts @@ -12,32 +12,69 @@ const requestSchema = z }) .passthrough(); -const encodedUserSchema = z - .object({ account_uuid: z.string().uuid().nullish() }) - .passthrough(); +const encodedUserSchema = z.object({}).passthrough(); +const accountUuidSchema = z.string().uuid().nullish(); const ACCOUNT_COMPONENT = /(^|_)account_([0-9a-f]{8}(?:-[0-9a-f]{4}){3}-[0-9a-f]{12})(?=_|$)/iu; +const SESSION_COMPONENT = + /(?:^|_)session_([0-9a-f]{8}(?:-[0-9a-f]{4}){3}-[0-9a-f]{12})$/iu; +const codexRequestSchema = z.object({ prompt_cache_key: z.string().nullish() }); export interface ParsedRequestBody { family: ModelFamily; + affinityId: string | null; forAccount: (accountUuid: string | null) => Uint8Array; } -function rewriteUserId(userId: string, accountUuid: string): string | null { +function affinityIdentifier(value: string | null | undefined): string | null { + return value !== undefined && + value !== null && + value.length <= 512 && + value.trim().length > 0 && + !/[\u0000-\u001f\u007f-\u009f]/u.test(value) + ? value + : null; +} + +function parseUserId(userId: string): { + sessionId: string | null; + forAccount: (accountUuid: string) => string | null; +} { try { const encoded = encodedUserSchema.safeParse(JSON.parse(userId)); - if (encoded.success && encoded.data.account_uuid !== undefined) { - if (encoded.data.account_uuid === accountUuid) return null; - return JSON.stringify({ ...encoded.data, account_uuid: accountUuid }); + if (encoded.success) { + const originalAccount = accountUuidSchema.safeParse( + encoded.data.account_uuid, + ); + return { + sessionId: + typeof encoded.data.session_id === "string" + ? affinityIdentifier(encoded.data.session_id) + : null, + forAccount(accountUuid) { + if ( + !originalAccount.success || + originalAccount.data === undefined || + originalAccount.data === accountUuid + ) + return null; + return JSON.stringify({ ...encoded.data, account_uuid: accountUuid }); + }, + }; } } catch {} - if (!ACCOUNT_COMPONENT.test(userId)) return null; - const rewritten = userId.replace( - ACCOUNT_COMPONENT, - (_match, prefix: string) => `${prefix}account_${accountUuid}`, - ); - return rewritten === userId ? null : rewritten; + return { + sessionId: affinityIdentifier(SESSION_COMPONENT.exec(userId)?.[1]), + forAccount(accountUuid) { + if (!ACCOUNT_COMPONENT.test(userId)) return null; + const rewritten = userId.replace( + ACCOUNT_COMPONENT, + (_match, prefix: string) => `${prefix}account_${accountUuid}`, + ); + return rewritten === userId ? null : rewritten; + }, + }; } export function parseRequestBody(body: Uint8Array): ParsedRequestBody { @@ -46,17 +83,21 @@ export function parseRequestBody(body: Uint8Array): ParsedRequestBody { const parsed = requestSchema.safeParse( JSON.parse(new TextDecoder().decode(body)), ); - if (!parsed.success) { - return { family: "other", forAccount: () => original }; - } + if (!parsed.success) + return { family: "other", affinityId: null, forAccount: () => original }; const request = parsed.data; + const userId = request.metadata?.user_id; + const user = + userId === undefined || userId === null ? null : parseUserId(userId); return { family: modelFamily(request.model ?? null), + affinityId: + user?.sessionId === undefined || user.sessionId === null + ? null + : `session:${user.sessionId}`, forAccount(accountUuid) { - if (accountUuid === null) return original; - const userId = request.metadata?.user_id; - if (userId === undefined || userId === null) return original; - const rewritten = rewriteUserId(userId, accountUuid); + if (accountUuid === null || user === null) return original; + const rewritten = user.forAccount(accountUuid); if (rewritten === null) return original; return new TextEncoder().encode( JSON.stringify({ @@ -67,6 +108,29 @@ export function parseRequestBody(body: Uint8Array): ParsedRequestBody { }, }; } catch { - return { family: "other", forAccount: () => original }; + return { family: "other", affinityId: null, forAccount: () => original }; + } +} + +export function parseCodexRequestBody( + body: Uint8Array, + headers: Headers, +): ParsedRequestBody { + let affinityId: string | null = null; + const sessionId = + affinityIdentifier(headers.get("session-id")) ?? + affinityIdentifier(headers.get("session_id")); + if (sessionId !== null) affinityId = `session:${sessionId}`; + else { + try { + const parsed = codexRequestSchema.safeParse( + JSON.parse(new TextDecoder().decode(body)), + ); + const cacheKey = parsed.success + ? affinityIdentifier(parsed.data.prompt_cache_key) + : null; + if (cacheKey !== null) affinityId = `cache:${cacheKey}`; + } catch {} } + return { family: "other", affinityId, forAccount: () => body }; } diff --git a/plugins/account-pool/src/server.test.ts b/plugins/account-pool/src/server.test.ts index 966b1156e2..3aa7161548 100644 --- a/plugins/account-pool/src/server.test.ts +++ b/plugins/account-pool/src/server.test.ts @@ -93,11 +93,12 @@ async function resolveToken( async function resolveCodexToken( host: ReturnType, + hostId = "host-one", ): Promise<{ token: string; baseUrl: string }> { const entries = await host.harness.behavior.resolveProviderEnv("codex", { threadId: "thread-codex", projectId: "project-one", - hostId: "host-one", + hostId, }); const token = entries.find((entry) => entry.name === "CODEX_POOL_AUTH_TOKEN"); const baseUrl = entries.find( @@ -3294,6 +3295,715 @@ describe("Account Pool plugin", () => { expect(attempts).toBe(1); }); + describe("session affinity", () => { + const sessionId = "cccccccc-cccc-4ccc-8ccc-cccccccccccc"; + const claudeBody = (id: string, model = "claude-fable-5") => + JSON.stringify({ + model, + metadata: { + user_id: JSON.stringify({ + account_uuid: "invalid-account-uuid", + device_id: "device", + parent_session_id: "parent", + session_id: id, + }), + }, + }); + const openStream = () => + new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode("data: started\n\n")); + }, + }), + { headers: { "content-type": "text/event-stream" } }, + ); + + async function affinityFixture( + provider: "claude" | "codex", + upstreamFetch: typeof fetch, + now = () => 1_800_000_000_000, + ): Promise { + let imported = 0; + const fixture = await createFixture({ + upstreamUrl: "https://upstream.example", + provider, + apiKey: "sk-first", + source: provider === "codex" ? "import" : "api-key", + options: { + now, + codexUsageUrl: EMPTY_USAGE_URL, + fetch: (input, init) => + String(input) === EMPTY_USAGE_URL + ? Promise.resolve(Response.json({})) + : upstreamFetch(input, init), + importCodexCredentials: async () => ({ + accessToken: imported++ === 0 ? "sk-first" : "sk-second", + refreshToken: "refresh", + idToken: null, + accountId: `codex-account-${imported}`, + email: null, + expiresAt: now() + 24 * 60 * 60 * 1_000, + }), + }, + }); + if (provider === "claude") await addApiAccount(fixture, "sk-second"); + else + await fixture.host.harness.behavior.callRpc("account.add", { + provider, + source: { kind: "import" }, + label: null, + priority: 100, + }); + return fixture; + } + + it.each<{ + name: string; + provider: "claude" | "codex"; + headers: Record; + body: string; + }>([ + { + name: "Claude JSON", + provider: "claude", + headers: {}, + body: claudeBody(sessionId), + }, + { + name: "Claude legacy", + provider: "claude", + headers: {}, + body: JSON.stringify({ + metadata: { + user_id: `user_hash_account_aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa_session_${sessionId}`, + }, + }), + }, + { + name: "Codex native", + provider: "codex", + headers: { "session-id": sessionId, "thread-id": "native-thread" }, + body: "{}", + }, + { + name: "Codex legacy", + provider: "codex", + headers: { session_id: sessionId }, + body: "{}", + }, + { + name: "Codex cache", + provider: "codex", + headers: {}, + body: JSON.stringify({ prompt_cache_key: sessionId }), + }, + ])( + "keeps $name sessions sticky with host isolation and idle expiry", + async ({ provider, headers, body }) => { + let now = 1_800_000_000_000; + const attempts: Array = []; + const fixture = await affinityFixture( + provider, + async (_input, init) => { + const requestHeaders = new Headers(init?.headers); + attempts.push( + provider === "claude" + ? requestHeaders.get("x-api-key") + : requestHeaders.get("authorization"), + ); + return attempts.length === 1 ? openStream() : Response.json({}); + }, + () => now, + ); + const route = provider === "claude" ? "/v1/messages" : "/v1/responses"; + const keyFor = (key: string) => + provider === "claude" ? key : `Bearer ${key}`; + const send = async ( + key: string, + requestBody: string, + sessionHeaders = headers, + ) => { + const response = await fixture.host.harness.behavior.fetchHttp( + "POST", + route, + { + headers: { ...authHeaders(key), ...sessionHeaders }, + body: requestBody, + }, + ); + await response.text(); + expect(response.status).toBe(200); + return attempts.at(-1); + }; + const held = await fixture.host.harness.behavior.fetchHttp( + "POST", + route, + { + headers: { ...authHeaders(fixture.key), ...headers }, + body, + }, + ); + try { + const otherHost = + provider === "codex" + ? (await resolveCodexToken(fixture.host, "host-two")).token + : await resolveToken(fixture.host, "host-two"); + expect(await send(fixture.key, body)).toBe(keyFor("sk-first")); + expect(await send(otherHost, body)).toBe(keyFor("sk-second")); + expect(await send(fixture.key, "{}", {})).toBe(keyFor("sk-second")); + now += 29 * 60 * 1_000; + expect(await send(fixture.key, body)).toBe(keyFor("sk-first")); + now += 2 * 60 * 1_000; + expect(await send(fixture.key, body)).toBe(keyFor("sk-first")); + now += 31 * 60 * 1_000; + expect(await send(fixture.key, body)).toBe(keyFor("sk-second")); + } finally { + await held.body?.cancel(); + } + }, + ); + + it("separates Codex session and cache namespaces and prefers the native header", async () => { + const attempts: Array = []; + const fixture = await affinityFixture("codex", async (_input, init) => { + attempts.push(new Headers(init?.headers).get("authorization")); + return attempts.length === 1 ? openStream() : Response.json({}); + }); + const held = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/responses", + { + headers: { ...authHeaders(fixture.key), "session-id": sessionId }, + body: "{}", + }, + ); + try { + const variants: Array> = [ + {}, + { session_id: sessionId }, + { "session-id": sessionId, session_id: "different-session" }, + ]; + for (const headers of variants) { + const response = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/responses", + { + headers: { ...authHeaders(fixture.key), ...headers }, + body: JSON.stringify({ prompt_cache_key: sessionId }), + }, + ); + await response.text(); + } + expect(attempts).toEqual([ + "Bearer sk-first", + "Bearer sk-second", + "Bearer sk-first", + "Bearer sk-first", + ]); + } finally { + await held.body?.cancel(); + } + }); + + it.each([ + "family quota", + "auth error", + "disabled account", + "network error", + ])( + "rebinds after %s without letting an older response completion restore the binding", + async (reason) => { + const attempts: Array = []; + let finishOld = () => {}; + const fixture = await createFixture({ + upstreamUrl: "https://upstream.example", + apiKey: "sk-first", + options: { + fetch: async (_input, init) => { + const key = new Headers(init?.headers).get("x-api-key"); + attempts.push(key); + if (attempts.length === 1) + return new Response( + new ReadableStream({ + start(controller) { + controller.enqueue( + new TextEncoder().encode("data: old\n\n"), + ); + finishOld = () => { + controller.close(); + finishOld = () => {}; + }; + }, + }), + { + headers: + reason === "family quota" + ? { + "anthropic-ratelimit-unified-7d_fable-status": + "rejected", + "anthropic-ratelimit-unified-7d_fable-reset": + "4102444800", + } + : {}, + }, + ); + if ( + reason === "network error" && + key === "sk-first" && + attempts.length === 2 + ) + throw new TypeError("network failed"); + return Response.json( + {}, + { + status: + reason === "auth error" && + key === "sk-first" && + attempts.length === 2 + ? 403 + : 200, + }, + ); + }, + }, + }); + await addApiAccount(fixture, "sk-second"); + const send = (model: string) => + fixture.host.harness.behavior.fetchHttp("POST", "/v1/messages", { + headers: authHeaders(fixture.key), + body: claudeBody(sessionId, model), + }); + const held = await send( + reason === "family quota" ? "claude-fable-5" : "claude-opus-4-1", + ); + try { + if (reason === "disabled account") + await fixture.host.harness.behavior.callRpc("account.disable", { + id: fixture.account.id, + }); + const rebound = await send("claude-fable-5"); + expect(rebound.status).toBe(200); + await rebound.text(); + if (reason !== "family quota") + await fixture.host.harness.behavior.callRpc("account.enable", { + id: fixture.account.id, + }); + finishOld(); + await held.text(); + const afterCompletion = await send("claude-opus-4-1"); + await afterCompletion.text(); + expect(attempts).toEqual( + reason === "auth error" || reason === "network error" + ? ["sk-first", "sk-first", "sk-second", "sk-second"] + : ["sk-first", "sk-second", "sk-second"], + ); + } finally { + finishOld(); + if (!held.bodyUsed) await held.body?.cancel(); + } + }, + ); + + it("shares the first binding when simultaneous account listings resume under different load", async () => { + const attempts: Array = []; + const fixture = await createFixture({ + upstreamUrl: "https://upstream.example", + apiKey: "sk-first", + options: { + fetch: async (_input, init) => { + attempts.push(new Headers(init?.headers).get("x-api-key")); + return attempts.length === 1 ? openStream() : Response.json({}); + }, + }, + }); + await addApiAccount(fixture, "sk-second"); + const gates = [deferred(), deferred()]; + const originalList = AccountStore.prototype.list; + let listings = 0; + const list = vi + .spyOn(AccountStore.prototype, "list") + .mockImplementation(async function (this: AccountStore) { + const index = listings++; + const accounts = await originalList.call(this); + if (index < 2) await gates[index]?.promise; + return accounts; + }); + const requests = [1, 2].map(() => + fixture.host.harness.behavior.fetchHttp("POST", "/v1/messages", { + headers: authHeaders(fixture.key), + body: claudeBody(sessionId), + }), + ); + try { + await vi.waitFor(() => expect(listings).toBe(2)); + gates[0]?.resolve(); + await vi.waitFor(() => expect(attempts).toHaveLength(1)); + gates[1]?.resolve(); + const responses = await Promise.all(requests); + await responses[1]?.text(); + await responses[0]?.body?.cancel(); + expect(attempts).toEqual(["sk-first", "sk-first"]); + } finally { + for (const gate of gates) gate.resolve(); + await Promise.allSettled( + requests.map(async (request) => { + const response = await request; + if (!response.bodyUsed) await response.body?.cancel(); + }), + ); + list.mockRestore(); + } + }); + + it("keeps native Codex HTTP and WebSocket compaction continuations on the same account", async () => { + const seen: Array<{ headers: Headers; body: string }> = []; + const compacted = { + type: "compaction", + encrypted_content: "encrypted-compaction-fixture", + }; + const fixture = await affinityFixture("codex", async (_input, init) => { + seen.push({ + headers: new Headers(init?.headers), + body: new TextDecoder().decode( + init?.body instanceof ArrayBuffer ? init.body : new ArrayBuffer(0), + ), + }); + if (seen.length === 1) return openStream(); + return new Response( + `data: ${JSON.stringify({ type: "response.completed", response: { id: `response-${seen.length}`, output: [compacted] } })}\n\n`, + { headers: { "content-type": "text/event-stream" } }, + ); + }); + const headers = { + ...authHeaders(fixture.key), + "session-id": sessionId, + "thread-id": "native-thread", + }; + const held = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/responses", + { headers, body: "{}" }, + ); + const socket = await fixture.host.harness.experimental_openWebSocket( + "/v1/responses", + { headers }, + ); + try { + const fields = { + model: "gpt-5", + prompt_cache_key: "cache-key", + client_metadata: { + session_id: "body-session", + thread_id: "body-thread", + "x-codex-turn-metadata": "fixture", + }, + include: ["reasoning.encrypted_content"], + reasoning: { effort: "high" }, + }; + const input = [ + { + type: "reasoning", + encrypted_content: "encrypted-reasoning-fixture", + }, + { type: "compaction_trigger" }, + ]; + await socket.receive( + JSON.stringify({ type: "response.create", ...fields, input }), + ); + await vi.waitFor(() => expect(socket.sent).toHaveLength(1)); + const first = completedResponse(socket.sent[0]); + const delta = { type: "message", role: "user", content: "next" }; + await socket.receive( + JSON.stringify({ + type: "response.create", + ...fields, + previous_response_id: first.response.id, + input: [delta], + }), + ); + await vi.waitFor(() => expect(socket.sent).toHaveLength(2)); + expect(JSON.parse(seen[1]?.body ?? "{}")).toMatchObject({ + ...fields, + input, + }); + expect(JSON.parse(seen[2]?.body ?? "{}")).toMatchObject({ + ...fields, + input: [...input, compacted, delta], + }); + expect(seen.map(({ headers }) => headers.get("session-id"))).toEqual([ + sessionId, + sessionId, + sessionId, + ]); + expect(seen.map(({ headers }) => headers.get("thread-id"))).toEqual([ + "native-thread", + "native-thread", + "native-thread", + ]); + expect(seen.map(({ headers }) => headers.get("authorization"))).toEqual( + ["Bearer sk-first", "Bearer sk-first", "Bearer sk-first"], + ); + } finally { + await socket.close(1000, "done"); + await held.body?.cancel(); + } + }); + + it.each([false, true])( + "preserves a newer session binding outside an older candidate snapshot with a fallback: %s", + async (hasFallback) => { + const gate = deferred(); + const attempts: Array = []; + const fixture = await createFixture({ + upstreamUrl: "https://upstream.example", + apiKey: "sk-first", + priority: 0, + options: { + fetch: async (_input, init) => { + attempts.push(new Headers(init?.headers).get("x-api-key")); + if (attempts.length === 1) { + await gate.promise; + return Response.json({}, { status: 503 }); + } + return Response.json({}); + }, + }, + }); + const fallback = await addApiAccount(fixture, "sk-fallback", 20); + if (!hasFallback) + await fixture.host.harness.behavior.callRpc("account.disable", { + id: fallback.id, + }); + const send = () => + fixture.host.harness.behavior.fetchHttp("POST", "/v1/messages", { + headers: authHeaders(fixture.key), + body: claudeBody(sessionId), + }); + const old = send(); + try { + await vi.waitFor(() => expect(attempts).toEqual(["sk-first"])); + await addApiAccount(fixture, "sk-new", 10); + await fixture.host.harness.behavior.callRpc("account.disable", { + id: fixture.account.id, + }); + const rebound = await send(); + expect(rebound.status).toBe(200); + await rebound.text(); + gate.resolve(); + const exhausted = await old; + expect(exhausted.status).toBe(hasFallback ? 200 : 503); + await exhausted.text(); + await fixture.host.harness.behavior.callRpc("account.enable", { + id: fixture.account.id, + }); + const next = await send(); + expect(next.status).toBe(200); + await next.text(); + expect(attempts).toEqual( + hasFallback + ? ["sk-first", "sk-new", "sk-fallback", "sk-new"] + : ["sk-first", "sk-new", "sk-new"], + ); + } finally { + gate.resolve(); + const response = await old; + if (!response.bodyUsed) await response.text(); + } + }, + ); + + it("preserves a newer session binding to an already-attempted account", async () => { + const gate = deferred(); + const attempts: Array = []; + const fixture = await createFixture({ + upstreamUrl: "https://upstream.example", + apiKey: "sk-first", + priority: 0, + options: { + fetch: async (_input, init) => { + const key = new Headers(init?.headers).get("x-api-key"); + attempts.push(key); + if (attempts.length === 1) + return Response.json({}, { status: 503 }); + if (key === "sk-second") { + await gate.promise; + return Response.json({}, { status: 503 }); + } + return Response.json({}); + }, + }, + }); + const second = await addApiAccount(fixture, "sk-second", 10); + await addApiAccount(fixture, "sk-fallback", 20); + const send = () => + fixture.host.harness.behavior.fetchHttp("POST", "/v1/messages", { + headers: authHeaders(fixture.key), + body: claudeBody(sessionId), + }); + const old = send(); + try { + await vi.waitFor(() => + expect(attempts).toEqual(["sk-first", "sk-second"]), + ); + await fixture.host.harness.behavior.callRpc("account.disable", { + id: second.id, + }); + const rebound = await send(); + expect(rebound.status).toBe(200); + await rebound.text(); + gate.resolve(); + const fallback = await old; + expect(fallback.status).toBe(200); + await fallback.text(); + const next = await send(); + expect(next.status).toBe(200); + await next.text(); + expect(attempts).toEqual([ + "sk-first", + "sk-second", + "sk-first", + "sk-fallback", + "sk-first", + ]); + } finally { + gate.resolve(); + const response = await old; + if (!response.bodyUsed) await response.text(); + } + }); + + it("isolates provider bindings and clears them on hub restart", async () => { + const attempts: Array = []; + const started = new Set(); + const fixture = await affinityFixture("codex", async (_input, init) => { + const headers = new Headers(init?.headers); + const provider = headers.has("x-api-key") ? "claude" : "codex"; + attempts.push(headers.get("x-api-key") ?? headers.get("authorization")); + if (!started.has(provider)) { + started.add(provider); + return openStream(); + } + return Response.json({}); + }); + await addApiAccount(fixture, "sk-claude-first"); + const claudeSecond = await addApiAccount(fixture, "sk-claude-second"); + const accounts = z + .array(accountSummarySchema) + .parse( + await fixture.host.harness.behavior.callRpc("account.list", null), + ); + const codexSecond = accounts.find( + (account) => account.codexAccountId === "codex-account-2", + ); + if (codexSecond === undefined) + throw new Error("Missing second Codex account."); + const sendClaude = () => + fixture.host.harness.behavior.fetchHttp("POST", "/v1/messages", { + headers: authHeaders(fixture.key), + body: claudeBody(sessionId), + }); + const sendCodex = () => + fixture.host.harness.behavior.fetchHttp("POST", "/v1/responses", { + headers: { ...authHeaders(fixture.key), "session-id": sessionId }, + body: "{}", + }); + const held = [await sendClaude(), await sendCodex()]; + try { + for (const account of [claudeSecond, codexSecond]) + await fixture.host.harness.behavior.callRpc("account.setPriority", { + accountId: account.id, + priority: 0, + }); + await (await sendClaude()).text(); + await (await sendCodex()).text(); + for (const response of held) await response.body?.cancel(); + fixture.service.controller.abort(); + await fixture.service.done; + const restarted = fixture.host.harness.behavior.runService("hub"); + cleanups.push(async () => { + restarted.controller.abort(); + await restarted.done; + }); + await vi.waitFor(async () => { + const status = statusSchema.parse( + await fixture.host.harness.behavior.callRpc("status.get", null), + ); + expect(status.accepting).toBe(true); + }); + await (await sendClaude()).text(); + await (await sendCodex()).text(); + expect(attempts).toEqual([ + "sk-claude-first", + "Bearer sk-first", + "sk-claude-first", + "Bearer sk-first", + "sk-claude-second", + "Bearer sk-second", + ]); + } finally { + for (const response of held) + if (!response.bodyUsed) await response.body?.cancel(); + } + }); + + it("evicts the least recently used binding after 4096 sessions", async () => { + let lastKey: string | null = null; + let attempts = 0; + const fixture = await createFixture({ + upstreamUrl: "https://upstream.example", + apiKey: "sk-first", + options: { + now: () => 1_800_000_000_000, + fetch: async (_input, init) => { + lastKey = new Headers(init?.headers).get("x-api-key"); + attempts += 1; + return attempts === 1 ? openStream() : Response.json({}); + }, + }, + }); + await addApiAccount(fixture, "sk-second"); + const send = async (id: string) => { + const response = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { + headers: authHeaders(fixture.key), + body: claudeBody(id), + }, + ); + await response.text(); + return lastKey; + }; + const held = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { + headers: authHeaders(fixture.key), + body: claudeBody("oldest"), + }, + ); + try { + for (let index = 1; index < 4096; index += 1) + await send(`session-${index}`); + const touched = await send("oldest"); + await send("newest"); + const retained = await send("oldest"); + await held.body?.cancel(); + const nextOldest = await send("session-2"); + const evicted = await send("session-1"); + expect([touched, retained, nextOldest, evicted]).toEqual([ + "sk-first", + "sk-first", + "sk-second", + "sk-first", + ]); + } finally { + if (!held.bodyUsed) await held.body?.cancel(); + } + }, 20_000); + }); + it("serializes refresh, writes new tokens with 0600 mode, and uses them", async () => { let now = 1_800_000_000_000; let refreshCalls = 0;