diff --git a/docs-site/src/content/docs/reference/configuration/providers.md b/docs-site/src/content/docs/reference/configuration/providers.md index 082c0dd688f..7190882103a 100644 --- a/docs-site/src/content/docs/reference/configuration/providers.md +++ b/docs-site/src/content/docs/reference/configuration/providers.md @@ -125,7 +125,7 @@ separate. Full request URLs such as `/api/v1/responses` are not provider base UR | `codexAccountPriorityFailback?` | `boolean` | `false` | Opt-in return of ongoing Codex Pool tasks to a usable higher-priority account under the `quota` strategy. Requests trigger a background quota refresh at most once every five minutes, including failed refresh attempts. When recent quota observations show headroom, the next request may rebind even if its current account is below the threshold. This explicit preference can discard a warm prompt cache; it does not change default cache affinity. No requests means no polling. Manual pins, model access, cooldowns, pauses and reauthentication exclusions still apply; in-flight requests keep their captured account. Requires the current source account's effective threshold to be positive (an override of 0 disables it; a positive override works with global 0); exact account routes and Direct mode are unchanged. | | `autoSwitchThreshold?` | `number` | `80` | Usage threshold for placing new/unbound work. `quota` can re-evaluate unbound tasks on their next request once usage crosses this threshold. Bound tasks keep their account past the threshold by default (`pool.cacheAffinity`); they leave only when that account is exhausted or otherwise cannot serve, and then only for an account with genuine quota headroom and strictly lower usage. Set `pool.cacheAffinity: false` to re-evaluate bound tasks at this threshold, still only onto such a destination. `fill-first` uses it only as the drain point for unbound assignment; `round-robin` uses it only when placing new/unbound work: it skips accounts at or over their effective threshold while any eligible account is still under it, and never moves a bound task for it. The score uses the hottest known 5h, weekly, or 30d quota window. `0` disables usage-based proactive switching only, not unbound assignment or failure recovery. | | `codexAccountAutoSwitchThresholds?` | `Record` | — | Per-account overrides for `autoSwitchThreshold`: account id → integer from `0` to `100`. An absent entry inherits the global value; `0` disables usage-driven switching only when that account is the source. `quota` uses the current source account's effective threshold for bound-task re-evaluation and unbound switching. Fill-first drain and selection-order tier headroom are evaluated against each account's own effective threshold; main-account pins and subagent fallback also use the relevant account's effective threshold. Supports the main `__main__` account. A malformed map is ignored with a warning. Managed by each account card on the Codex Auth page; disabling an override removes its entry. Enabling the override copies the current global threshold into a fixed account value. The override, including `0`, takes precedence over later global changes. Disabling it sends `threshold: null`, removes the entry, and restores inheritance of the current global threshold and future global changes. | -| `accountPoolStrategy?` | `"quota" \| "round-robin" \| "fill-first" \| "reset-first"` | `"quota"` | Assignment strategy for new/unbound Codex requests. A request is unbound when it has no live (parent thread id, quota scope) affinity; a visible existing task can become unbound after proxy restart or affinity reset. `quota` picks the lowest-usage eligible account when no active account exists, keeps an eligible active account below `autoSwitchThreshold`, and after the threshold may move an unbound request. Bound tasks follow `pool.cacheAffinity` (on by default): they stay until the account is exhausted (known usage at 100%) or otherwise cannot serve, and then may rebind only to an account with genuine quota headroom and strictly lower usage. Set the flag `false` to proactively rebind a bound task at the threshold, still only onto such a destination. `round-robin` distributes unbound requests evenly across eligible accounts still under their effective `autoSwitchThreshold`, and across all eligible accounts when every one is at or over it; `fill-first` keeps assigning unbound requests to the active account until cooldown, unavailability, or the configured drain threshold. `reset-first`: Prefer the nearest future 5-hour or weekly reset among accounts below the usage threshold. Bound tasks follow the configured affinity policy. Independent model quotas use quota ordering. Monthly resets do not determine this ordering. | +| `accountPoolStrategy?` | `"quota" \| "round-robin" \| "fill-first" \| "reset-first"` | `"quota"` | Assignment strategy for new/unbound Codex requests. A request is unbound when it has no live (parent thread id, quota scope) affinity; a visible existing task can become unbound after proxy restart or affinity reset. `quota` picks the lowest-usage eligible account when no active account exists, keeps an eligible active account below `autoSwitchThreshold`, and after the threshold may move an unbound request. Bound tasks follow `pool.cacheAffinity` (on by default): they stay until the account is exhausted (known usage at 100%) or otherwise cannot serve, and then may rebind only to an account with genuine quota headroom and strictly lower usage. Set the flag `false` to proactively rebind a bound task at the threshold, still only onto such a destination. `round-robin` distributes unbound requests evenly across eligible accounts still under their effective `autoSwitchThreshold`, and across all eligible accounts when every one is at or over it; `fill-first` keeps assigning unbound requests to the active account until cooldown, unavailability, or the configured drain threshold. `reset-first`: Prefer the nearest future 5-hour or weekly reset among accounts below the usage threshold. Bound tasks follow the configured affinity policy. Independent model quotas use quota ordering. Monthly resets do not determine this ordering. Under every strategy, a Claude Code turn that reaches a Codex account route through `/v1/messages` is bound by its Claude session: a caller `session_id`/`session-id` alias if present, else the metadata session, else `x-claude-code-session-id`; the binding is keyed to the admitted credential (or to local loopback use) with a server-held secret, so callers cannot share, guess or reuse another credential's binding. The binding is used only to choose the account; the headers sent upstream are unchanged. Spawn turns and turns that already carry `x-codex-parent-thread-id`, `thread-id` or a complete session/thread pair keep their own identity. A bound session normally stays on one account; the affinity and failure-recovery rules above can still move it. | | `pool.cacheAffinity?` | `boolean` | `true` | Cache-affinity ordering for bound Codex threads, independent of `pool.kernel`. On by default; omitting the key or setting `true` keeps a bound task on its account until that account genuinely cannot serve. Only an explicit `false` restores threshold-based rebinding of bound tasks. A live binding outranks quota headroom: `quota` does not move the thread merely because usage crossed `autoSwitchThreshold`. The thread still leaves if that account cannot serve — paused, unusable, or genuinely exhausted (known usage at 100%) — and then only to an account with genuine quota headroom and strictly lower usage. Under either setting, an account with unknown usage is never chosen as a destination for a bound task, so when every account sits above the threshold the task stays put. Affinity is a reordering, not a pin. | | `pool.credentialGroups?` | `Array<{ id: string; credentials: string[]; note?: string }>` | `[]` | Accepted and validated, but not yet consumed by routing: declaring a group changes no routing decision until a consuming layer lands. Operator-declared quota domains: groups of credentials that demonstrably share one upstream usage limit. Members of one group count once toward available capacity, and a quota refusal inside a group is not answered by rotating to another member — the limit is the same, so the move would pay a cold prefix for zero new capacity. Declared groups speak only to quota; sharing a limit says nothing about prompt-cache compatibility, which is classified separately. Each member is written provider-qualified as `":"`, because a credential id means something only inside its provider; the provider segment accepts the usual aliases (`chatgpt:` and `codex:` both mean OpenAI). Group ids must be unique, `credentials` must be non-empty, and a credential may appear in at most one group — an ambiguous declaration is rejected on write and dropped with a warning on load rather than resolved by whichever group is listed first, since that would merge two unrelated quota domains. A malformed list costs only the grouping: `pool.kernel` and `pool.cacheAffinity` are preserved. Absent or empty means no declared grouping, so an unconfigured install behaves exactly as before. | | `accountPoolStickyLimit?` | `number` | `1` | New/unbound task assignments retained on one round-robin selection before advancing; the counter advances when a task is bound, not after an upstream success. Range 1–100. | diff --git a/scripts/test-layout/layout-server-inbound.json b/scripts/test-layout/layout-server-inbound.json index f38f6530ecd..41002edc674 100644 --- a/scripts/test-layout/layout-server-inbound.json +++ b/scripts/test-layout/layout-server-inbound.json @@ -1,5 +1,7 @@ { "config-concurrent-inbound-body.test.ts": "config", "inbound-body-admission.test.ts": "server", - "inbound-body-budget.test.ts": "server" + "inbound-body-budget.test.ts": "server", + "claude-bridge-pool-session.test.ts": "claude-integration", + "claude-bridge-pool-dispatch.test.ts": "claude-integration" } diff --git a/src/server/claude-bridge-pool-session.ts b/src/server/claude-bridge-pool-session.ts new file mode 100644 index 00000000000..6fb34e375c4 --- /dev/null +++ b/src/server/claude-bridge-pool-session.ts @@ -0,0 +1,59 @@ +/** Private pool identity for translated Claude turns; replay headers remain unchanged. */ +import { createHmac, randomBytes } from "node:crypto"; +import { codexPoolAffinityKey } from "../codex/auth-context"; +import { contextPrincipalIdOf, type DataPlaneAdmission } from "./auth-cors"; +import { safeCallerSessionId } from "./caller-session-identity"; +import { isThreadSpawnRequest } from "./effort-policy"; +import type { RouteResult } from "../router"; +import type { HandleResponsesOptions } from "./responses/core-options"; + +/** + * Process-local, like the lineage key: a synthesized identity never leaves the process (replay and + * upstream headers stay uncompleted), so no caller can name it, not even as a complete + * session-id + thread-id pair, which reaches the pool unscoped. + */ +const BRIDGE_POOL_KEY = randomBytes(32); + +export function resolveClaudeBridgePoolSession(args: { + ingress: Headers; replay: Headers; metadataSessionId?: string; admission?: DataPlaneAdmission; +}): string | undefined { + const { ingress, replay, metadataSessionId, admission } = args; + if (!admission || isThreadSpawnRequest(ingress) || isThreadSpawnRequest(replay)) return undefined; + if (ingress.has("x-codex-parent-thread-id") || ingress.has("thread-id")) return undefined; + if (codexPoolAffinityKey(replay) !== undefined) return undefined; + let value: string | undefined; + if (ingress.has("session_id") || ingress.has("session-id")) { + const first = ingress.get("session_id"); + const second = ingress.get("session-id"); + if (first && second && first !== second) return undefined; + value = safeCallerSessionId(first || second); + } else { + value = safeCallerSessionId(metadataSessionId ?? ingress.get("x-claude-code-session-id")); + } + if (!value) return undefined; + // Every identity this bridge names, caller aliases included, is bound to its admission scope: + // a shared or guessed value must never select another principal's conversation key. + let scope: unknown = "loopback"; + if (admission.kind !== "loopback") { + const principal = contextPrincipalIdOf(admission); + if (!principal) return undefined; + scope = ["principal", principal]; + } + return createHmac("sha256", BRIDGE_POOL_KEY) + .update(JSON.stringify(["opencodex-claude-bridge-pool-v1", scope, value])).digest("hex"); +} + +export function completeClaudeBridgePoolIdentity(headers: Headers, options: Pick): Headers { + const session = options.claudeBridgePoolSession; + if (!session || headers.get("x-codex-parent-thread-id") || headers.get("thread-id") + || codexPoolAffinityKey(headers) !== undefined) return headers; + const completed = new Headers(headers); + completed.set("session-id", session); + completed.set("thread-id", session); + return completed; +} + +export function applyClaudeBridgePoolSession(headers: Headers, route: Pick, + options: Pick): Headers { + return route.codexAccountMode !== undefined ? completeClaudeBridgePoolIdentity(headers, options) : headers; +} diff --git a/src/server/claude-messages.ts b/src/server/claude-messages.ts index 45170d2f755..344d15d8462 100644 --- a/src/server/claude-messages.ts +++ b/src/server/claude-messages.ts @@ -8,6 +8,7 @@ import { CodexInstallationIdUnavailableError } from "../codex/installation-id"; * unchanged. The Responses output (SSE or JSON) is converted back to Anthropic shape. */ import { FORWARD_HEADERS } from "../adapters/openai-responses"; +import { resolveClaudeBridgePoolSession } from "./claude-bridge-pool-session"; import { admissionModelDeniedResponse, AdmissionModelDeniedError, @@ -1176,6 +1177,9 @@ async function handleClaudeMessagesWithBudget( ?? normalizeLogConversationId(req.headers.get("x-opencode-session")) ?? metadataGoLane ?? getOrAllocateRequestSessionLane(req); + const claudeBridgePoolSession = resolveClaudeBridgePoolSession({ + ingress: req.headers, replay: headers, metadataSessionId: claudeNativeSessionId, admission: logIds?.admission, + }); let internalReq: Request; try { // The UTF-16 JSON string and the Request's UTF-8 body coexist until dispatch. @@ -1239,6 +1243,7 @@ async function handleClaudeMessagesWithBudget( ? { tools: anthropicBody.tools, messages: anthropicBody.messages, tool_choice: anthropicBody.tool_choice } : undefined, claudeGoAffinity: { sessionLane: claudeGoSessionLane }, claudeNativeSessionId, + ...(claudeBridgePoolSession ? { claudeBridgePoolSession } : {}), stripClaudeMainAuthForNoncanonicalForward: true, ...(trustedClaudeMainAuth ? { trustedClaudeMainAuth } : {}), // Claude's internal stored-main enrichment is not an original caller credential. diff --git a/src/server/responses/core-auth.ts b/src/server/responses/core-auth.ts index 6f3990d4cc7..d5ba2c19b99 100644 --- a/src/server/responses/core-auth.ts +++ b/src/server/responses/core-auth.ts @@ -1,3 +1,4 @@ +import { applyClaudeBridgePoolSession } from "../claude-bridge-pool-session"; import type { OcxProviderConfig, OcxConfig } from "../../types"; import { isCanonicalOpenAiForwardProvider } from "../../providers/openai-tiers"; import type { CodexAuthContext } from "../../codex/auth-context"; @@ -111,7 +112,7 @@ function validatedCallerDirectHeaders( return captureCallerDirectAuth(headers, config) ? headers : undefined; } -export function codexRouteCredentialDomainHeaders( +function codexRouteCredentialDomainHeadersFor( req: Request, route: RouteResult, options: HandleResponsesOptions, @@ -156,6 +157,20 @@ export function codexRouteCredentialDomainHeaders( } +/** + * Credential-domain headers for Codex account SELECTION (previews, recovery, final resolution). + * A Claude bridge pool identity is completed here only; upstream headers are always materialized + * from the uncompleted form, so the synthesized pair never reaches the wire. + */ +export function codexRouteCredentialDomainHeaders( + req: Request, route: RouteResult, options: HandleResponsesOptions, + credentialDomainWasRewritten: boolean, config: OcxConfig, +): Headers { + return applyClaudeBridgePoolSession( + codexRouteCredentialDomainHeadersFor(req, route, options, credentialDomainWasRewritten, config), route, options, + ); +} + /** * Does this route substitute OUR stored main credential, and does the caller own the credential * this request will authenticate with? @@ -210,7 +225,7 @@ export async function resolveResponsesCodexAuth( // routed credential with it either, because that would be a silent credential swap. if (hasAgentAssertionAuthorization(req.headers)) return { ok: false, response: agentAssertionUnsupportedResponse() }; try { - let authInputHeaders = codexRouteCredentialDomainHeaders( + let authInputHeaders = codexRouteCredentialDomainHeadersFor( req, route, options, @@ -270,7 +285,9 @@ export async function resolveResponsesCodexAuth( validateForwardAdmissionCredential(authInputHeaders, config); } if (route.codexAccountMode) { - authCtx = await resolveCodexAuthContext(authInputHeaders, config, route.codexAccountMode, { + // Only selection sees the Claude bridge pool identity; materialization below keeps the + // uncompleted headers, so the wire matches the original request. + authCtx = await resolveCodexAuthContext(applyClaudeBridgePoolSession(authInputHeaders, route, options), config, route.codexAccountMode, { admission: options.admission, codexAuthPolicy: options.codexAuthPolicy, accountId: route.codexAccountId, diff --git a/src/server/responses/core-codex-account.ts b/src/server/responses/core-codex-account.ts index f48f5cf7b82..736c2c1f1e7 100644 --- a/src/server/responses/core-codex-account.ts +++ b/src/server/responses/core-codex-account.ts @@ -27,6 +27,7 @@ import { formatErrorResponse } from "../../bridge"; import { readBoundedResponseBody } from "../../lib/bounded-body"; import { upstreamErrorMessageFromPayload, isRateLimitOrQuotaFailureMessage } from "../../lib/errors"; import { isNonReplayableResponse, isTransientUpstreamStatus, SendBudgetExhaustedError } from "../../lib/upstream-retry"; +import { applyClaudeBridgePoolSession } from "../claude-bridge-pool-session"; import { codexWsPreOutputRefusalHold, codexWsPreOutputRefusalOutcomeRecorded, @@ -490,6 +491,8 @@ export interface CodexPoolAccountRetryArgs { abortSignal?: AbortSignal; onCodexAuthContextResolved?: (ctx: CodexAuthContext) => void; deferCodexResetDerivedCooldown?: boolean; + /** Private Claude bridge pool identity; alternate selection completes it like the first pick. */ + claudeBridgePoolSession?: string; // Narrowed subset of HandleResponsesOptions: the retry rebuilds the adapter, so it // needs the inbound scope or the retry could land on a different wire than the // first attempt. @@ -767,7 +770,9 @@ export async function retryCodexPoolOnAlternateAccount( } try { retryAuthCtx ??= await resolveCodexAuthContext( - callerAuthHeaders, + // Alternate selection keeps the initial Claude bridge pool identity; the retry's upstream + // headers are still built from uncompleted callerAuthHeaders below. + applyClaudeBridgePoolSession(callerAuthHeaders, { codexAccountMode: "pool" }, options), config, "pool", { diff --git a/src/server/responses/core-options.ts b/src/server/responses/core-options.ts index 36473e8ed7b..049463cff9b 100644 --- a/src/server/responses/core-options.ts +++ b/src/server/responses/core-options.ts @@ -118,6 +118,8 @@ export interface HandleResponsesOptions { claudeGoAffinity?: { sessionLane?: string }; /** Validated Claude metadata identity; projected only into final canonical attempt headers. */ claudeNativeSessionId?: string; + /** Private Claude bridge pool identity (validated, principal-scoped when derived); completes session-id/thread-id only on Codex account-selection headers. */ + claudeBridgePoolSession?: string; /** Original live policy owner; separate from caller-specific routing/sidecar snapshots. */ codexAuthPolicy?: CodexAuthPolicyConfig; turnAdmissionLease?: AdmissionLease; diff --git a/src/server/responses/request-prepare.ts b/src/server/responses/request-prepare.ts index 065c3f66357..7985abb8b60 100644 --- a/src/server/responses/request-prepare.ts +++ b/src/server/responses/request-prepare.ts @@ -1,3 +1,4 @@ +import { completeClaudeBridgePoolIdentity } from "../claude-bridge-pool-session"; import type { HandleResponsesOptions, ResponsesRequestContext, ResponsesAdmissionState, ResponsesDispatchers } from "./core-options"; import type { PolicyRequestEvidence } from "../../routing/evaluator"; import type { RequestLogContext } from "../request-log"; @@ -952,7 +953,7 @@ export async function prepareResponsesRequest( let subagentFallbackModelEligibleAccountIdsForModel: SubagentModelEligibleAccountIds | undefined; let subagentQuotaFailureModel = parsed.modelId; const parentThreadId = req.headers.get("x-codex-parent-thread-id")?.trim() ?? null; - const poolAffinityKey = codexPoolAffinityKey(req.headers) ?? null; + const poolAffinityKey = codexPoolAffinityKey(completeClaudeBridgePoolIdentity(req.headers, options)) ?? null; // Preview has to see the same lineage resolve does. Without it, a child's first turn is // previewed as a cold pick and resolved onto the family account, and the subagent fallback // then decides model eligibility against an account the request will never use. diff --git a/structure/data-planes/inbound-compat.md b/structure/data-planes/inbound-compat.md index 00135946d9d..b00ba62b047 100644 --- a/structure/data-planes/inbound-compat.md +++ b/structure/data-planes/inbound-compat.md @@ -389,6 +389,23 @@ canonical Go transport, including combo selection and failover. No Go-only repla reaches non-Go destinations. Shared system cache keys never become conversation identity. The request-scoped fallback is stable across retries and distinct across client requests. +The Messages bridge carries a separate private Codex pool identity without changing replay headers. +It requires admission and excludes spawn markers, any ingress parent or thread header (even empty) +and a replay with a complete pool key. Explicit session_id then session-id aliases win; conflicting, +empty-only or invalid aliases block derivation. Otherwise the metadata session (the uuid form of the +metadata-derived prompt_cache_key), then the Claude session header, supplies the value; a system-derived +cohort key never does. Every value, aliases included, is HMAC-keyed with a process-local secret over its +admission scope (loopback, or the principal a remote caller must have), so no caller can name another +scope's pool key, not even as a complete pair. A caller-supplied complete pair stays unscoped. + +The identity is selection-only: session-id and thread-id are completed for Codex account selection +(final and alternate-account), selection previews, recovery and the route-independent pool key used +as state-binding fallback. Every upstream header set (first send, 401 refresh, alternate-account retry) +is built from the uncompleted headers, so the wire is unchanged and no synthesized pair or digest is +forwarded. Non-Codex routes, Anthropic passthrough, request logs, reasoning replay, Go lanes and adapter +fingerprints are unchanged. claude-bridge-pool-session.test.ts and claude-bridge-pool-dispatch.test.ts +cover the rules, the wire contract, 401 refresh and alternate-account failover. + Claude metadata also supplies a private native-session value. Only the final canonical ChatGPT attempt receives it, in copied forwarding headers; an explicit underscore session, hyphenated session or thread header suppresses synthesis. Refresh and alternate-account retries retain diff --git a/structure/providers/codex-installation-identity.md b/structure/providers/codex-installation-identity.md index 0e16b361d64..d16f82cf067 100644 --- a/structure/providers/codex-installation-identity.md +++ b/structure/providers/codex-installation-identity.md @@ -70,3 +70,5 @@ Authenticated live and frame contracts are covered by `tests/codex-integration/codex-installation-live-ws.test.ts`; compact empty-header and turn-size contracts are covered by `tests/responses/codex-installation-compact-empty.test.ts` and `tests/responses/codex-installation-turn-limit.test.ts`. + +The Claude bridge carries its private pool identity through Codex selection; see [Claude affinity](../data-planes/inbound-compat.md#claude-affinity-at-final-go-dispatch). diff --git a/structure/transports/response-turn-membership.md b/structure/transports/response-turn-membership.md index 894a20b21e6..67b27165ac0 100644 --- a/structure/transports/response-turn-membership.md +++ b/structure/transports/response-turn-membership.md @@ -36,3 +36,5 @@ retry and queued rebuilds through `src/server/responses/adapter-dispatch.ts` and `src/server/responses/request-transport.ts`; missing metadata keeps earlier tool delivery off. This structural mechanism does not alter supplier bodies, reasoning effort, retry policy, settlement or provider-opaque replay scope. + +The route-independent Claude bridge pool key agrees with final Codex selection; see [Claude affinity](../data-planes/inbound-compat.md#claude-affinity-at-final-go-dispatch). diff --git a/tests/claude-integration/claude-bridge-pool-dispatch.test.ts b/tests/claude-integration/claude-bridge-pool-dispatch.test.ts new file mode 100644 index 00000000000..5b900ee7fe9 --- /dev/null +++ b/tests/claude-integration/claude-bridge-pool-dispatch.test.ts @@ -0,0 +1,258 @@ +import { afterEach, beforeEach, expect, spyOn, test } from "bun:test"; +import { mkdtempSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { MAIN_CODEX_ACCOUNT_ID } from "../../src/codex/main-account"; +import { clearMainAccountInfoCache } from "../../src/codex/auth-api"; +import { fakeChatGptJwt } from "../helpers/fake-chatgpt-jwt"; +import { clearCodexThreadLineageForTests, recordCodexThreadLineage, resolveCodexThreadLineage } from "../../src/codex/lineage"; +import { saveCodexAccountCredential } from "../../src/codex/account-store"; +import { clearThreadAccountMap, clearCodexUpstreamHealth } from "../../src/codex/routing"; +import { codexPoolAffinityKey } from "../../src/codex/auth-context"; +import * as authContext from "../../src/codex/auth-context"; +import { clearConversationStateIssuerMap, peekConversationStateIssuer } from "../../src/codex/routing/thread-affinity"; +import { clearAccountQuota, updateAccountQuota } from "../../src/codex/auth-api"; +import { handleClaudeMessages } from "../../src/server/claude-messages"; +import { resolveClaudeBridgePoolSession } from "../../src/server/claude-bridge-pool-session"; +import { tryAdmitTurn } from "../../src/server/lifecycle"; +import type { DataPlaneAdmission } from "../../src/server/auth-cors"; +import type { OcxConfig } from "../../src/types"; +import { installIsolatedCodexHome, type IsolatedCodexHome } from "../helpers/isolated-codex-home"; +import { acquireOwnedSpendHome } from "../helpers/owned-spend-home"; +import { removeTreeWithRetry } from "../helpers/remove-tree"; + +const originalFetch = globalThis.fetch; +const local: DataPlaneAdmission = { kind: "loopback", source: "loopback" }; +const metadataUuid = "9745d86c-d579-494a-8d0e-f69a5214cf96"; +const principal = (id: string): DataPlaneAdmission => ({ kind: "configured", keyId: "fixture", source: "dedicated", contextPrincipalId: id }); +// What a caller could compute from public inputs; the keyed bridge identity must never equal it. +const publicDigest = (principalId: string, value: string) => new Bun.CryptoHasher("sha256") + .update(JSON.stringify(["opencodex-caller-session-v1", principalId, value])).digest("hex"); +// The bridge's own identity for these inputs (same process key), used to read its bookkeeping. +const bridgeIdentity = (admission: DataPlaneAdmission, ingress: Record = {}, metadataSessionId?: string) => + resolveClaudeBridgePoolSession({ ingress: new Headers(ingress), replay: new Headers(ingress), admission, metadataSessionId })!; +const poolKey = (identity: string) => codexPoolAffinityKey(new Headers({ "session-id": identity, "thread-id": identity }))!; +const account = (wire: Headers | undefined) => wire?.get("chatgpt-account-id"); +/** The pool identity is selection-only: no attempt may carry a synthesized pair upstream. */ +function expectNoSynthesizedPair(wire: Headers) { + expect(wire.has("session-id")).toBe(false); + expect(wire.has("thread-id")).toBe(false); + for (const name of ["session_id", "session-id", "thread-id"]) expect(wire.get(name) ?? "").not.toMatch(/^[0-9a-f]{64}$/); +} +let home: string, previous: string | undefined, isolated: IsolatedCodexHome, release: () => void; +let wires: Headers[]; +let cfg: OcxConfig; + +beforeEach(() => { + previous = process.env.OPENCODEX_HOME; + home = mkdtempSync(join(tmpdir(), "ocx-bridge-pool-")); + process.env.OPENCODEX_HOME = home; + isolated = installIsolatedCodexHome("ocx-bridge-codex-"); + release = acquireOwnedSpendHome(); + clearThreadAccountMap(); clearCodexUpstreamHealth(); clearAccountQuota(); clearCodexThreadLineageForTests(); clearConversationStateIssuerMap(); + for (const id of ["fixture-a", "fixture-b"]) { + saveCodexAccountCredential(id, { accessToken: `fixture-access-${id}`, refreshToken: `fixture-refresh-${id}`, expiresAt: Date.now() + 3_600_000, chatgptAccountId: `account-${id}` }); + updateAccountQuota(id, 10, 10); + } + cfg = { defaultProvider: "openai", activeCodexAccountId: "fixture-a", autoSwitchThreshold: 0, + accountPoolStrategy: "round-robin", accountPoolStickyLimit: 1, + codexAccounts: [{ id: "fixture-a", label: "a" }, { id: "fixture-b", label: "b" }], + providers: { openai: { adapter: "openai-responses", baseUrl: "https://chatgpt.com/backend-api/codex", authMode: "forward", codexAccountMode: "pool", upstreamWebsocket: false } }, + } as unknown as OcxConfig; + wires = []; + globalThis.fetch = (async (_input: RequestInfo | URL, init?: RequestInit) => { + wires.push(new Headers(init?.headers)); + return Response.json({ id: "resp_fixture", object: "response", status: "completed", output: [], usage: { input_tokens: 1, output_tokens: 0, total_tokens: 1 } }); + }) as typeof fetch; +}); +afterEach(() => { + globalThis.fetch = originalFetch; clearMainAccountInfoCache(); release(); isolated.restore(); + if (previous === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = previous; + removeTreeWithRetry(home); clearThreadAccountMap(); clearCodexUpstreamHealth(); clearAccountQuota(); clearCodexThreadLineageForTests(); clearConversationStateIssuerMap(); +}); + +async function turn(session: string | undefined, admission: DataPlaneAdmission = local, headers: Record = {}, model = "openai/gpt-5.4") { + const lease = tryAdmitTurn(); + expect(lease).not.toBeNull(); + try { + const response = await handleClaudeMessages(new Request("http://localhost/v1/messages", { method: "POST", + headers: { "content-type": "application/json", ...headers }, + body: JSON.stringify({ model, max_tokens: 32, stream: false, ...(session ? { metadata: { user_id: session } } : {}), messages: [{ role: "user", content: "ping" }] }), + }), cfg, { model: "", provider: "" }, { requestId: "fixture", start: Date.now(), turnAdmissionLease: lease!, admission }); + await response.text(); + expect(response.status).toBe(200); + } finally { lease?.release(); } +} + +test("pool binds multiple turns and separates sessions while preserving the native session", async () => { + for (const session of ["user_test_account__session_conversation-native", "user_test_account__session_conversation-native", "user_test_account__session_other", "user_test_account__session_other"]) await turn(session); + expect(wires).toHaveLength(4); + expect(wires[0]!.get("chatgpt-account-id")).toBe(wires[1]!.get("chatgpt-account-id")); + expect(wires[2]!.get("chatgpt-account-id")).toBe(wires[3]!.get("chatgpt-account-id")); + expect(wires[0]!.get("chatgpt-account-id")).not.toBe(wires[2]!.get("chatgpt-account-id")); + expect(wires[0]!.get("session_id")).toBe(metadataUuid); + for (const wire of wires) expectNoSynthesizedPair(wire); +}); + +test("a synthesized identity with a recorded parent follows the parent's account", async () => { + const child = bridgeIdentity(local, {}, metadataUuid); + // A complete caller pair keeps dev behavior, so this binds the root to the first round-robin pick. + await turn(undefined, local, { "session-id": "fixture-root", "thread-id": "fixture-root" }); + const childHeaders = new Headers({ "session-id": child, "thread-id": child, "x-codex-parent-thread-id": "fixture-root" }); + recordCodexThreadLineage(childHeaders); + const rootPair = new Headers({ "session-id": child, "thread-id": child }); + expect(resolveCodexThreadLineage(rootPair)?.parentThreadId).toBe("fixture-root"); + await turn("user_test_account__session_conversation-native"); + await turn("user_test_account__session_conversation-native"); + // Control: an unrelated cold session lands on the other account, so the child's placement + // above came from the recorded parent rather than from a cold pick. + await turn("user_test_account__session_unrelated"); + const rootAccount = wires[0]!.get("chatgpt-account-id"); + expect(rootAccount).toBeTruthy(); + expect(wires[1]!.get("chatgpt-account-id")).toBe(rootAccount); + expect(wires[2]!.get("chatgpt-account-id")).toBe(rootAccount); + expect(wires[3]!.get("chatgpt-account-id")).not.toBe(rootAccount); + expect(resolveCodexThreadLineage(rootPair)?.parentThreadId).toBe("fixture-root"); +}); + +test("remote principals with identical metadata pin separately and keep the dev wire", async () => { + for (const id of ["a", "a", "b", "b", "a"]) await turn("user_test_account__session_conversation-native", principal(id)); + expect(account(wires[1])).toBe(account(wires[0])); + expect(account(wires[2])).not.toBe(account(wires[0])); + expect(account(wires[3])).toBe(account(wires[2])); + expect(account(wires[4])).toBe(account(wires[0])); + for (const wire of wires) { + expect(wire.get("session_id")).toBe(metadataUuid); + expectNoSynthesizedPair(wire); + } +}); + +test("remote principals sending the same session alias do not share a pool binding", async () => { + for (const id of ["a", "a", "b", "b", "a"]) await turn(undefined, principal(id), { session_id: "fixture-alias" }); + expect(account(wires[1])).toBe(account(wires[0])); + expect(account(wires[2])).not.toBe(account(wires[0])); + expect(account(wires[3])).toBe(account(wires[2])); + expect(account(wires[4])).toBe(account(wires[0])); + for (const wire of wires) { + expect(wire.get("session_id")).toBe("fixture-alias"); + expectNoSynthesizedPair(wire); + } +}); + +test("an alias equal to another principal's derived identity does not join that binding", async () => { + await turn("user_test_account__session_conversation-native", principal("b")); + const victim = bridgeIdentity(principal("b"), {}, metadataUuid); + await turn(undefined, principal("a"), { session_id: victim }); + await turn(undefined, principal("a"), { session_id: victim }); + expect(account(wires[1])).not.toBe(account(wires[0])); + expect(account(wires[2])).toBe(account(wires[1])); + expect(peekConversationStateIssuer(poolKey(victim))).toBe(account(wires[0])!.replace(/^account-/, "")); +}); + +test("a complete caller pair built from public inputs cannot address a synthesized identity", async () => { + await turn("user_test_account__session_conversation-native", principal("b")); + const victim = bridgeIdentity(principal("b"), {}, metadataUuid); + const victimAccount = account(wires[0])!.replace(/^account-/, ""); + // Raw metadata session and the publicly computable digest are what an attacker could know. + for (const guess of [metadataUuid, publicDigest("b", metadataUuid)]) { + expect(guess).not.toBe(victim); + expect(poolKey(guess)).not.toBe(poolKey(victim)); + await turn(undefined, principal("a"), { "session-id": guess, "thread-id": guess }); + } + expect(account(wires[1])).not.toBe(account(wires[0])); + expect(peekConversationStateIssuer(poolKey(victim))).toBe(victimAccount); +}); + +test("caller session-id aliases and complete pairs reach the wire exactly as sent", async () => { + await turn(undefined, local, { "session-id": "fixture-hyphen" }); + expect(wires[0]!.get("session-id")).toBe("fixture-hyphen"); + expect(wires[0]!.get("session_id")).toBe("fixture-hyphen"); + expect(wires[0]!.has("thread-id")).toBe(false); + await turn(undefined, principal("a"), { "session-id": "fixture-pair", "thread-id": "fixture-pair" }); + expect(wires[1]!.get("session-id")).toBe("fixture-pair"); + expect(wires[1]!.get("thread-id")).toBe("fixture-pair"); +}); + +test("spawn markers do not synthesize thread identity", async () => { + for (const marker of [{ "x-openai-subagent": "collab_spawn" }, { "x-codex-turn-metadata": '{"subagent_kind":"thread_spawn"}' }]) { + await turn("user_test_account__session_conversation-native", local, marker); + } + expect(wires).toHaveLength(2); + for (const wire of wires) expect(wire.has("thread-id")).toBe(false); +}); + +test("API-key Responses routes do not receive pool identity", async () => { + cfg.providers.other = { adapter: "openai-responses", authMode: "key", apiKey: "fixture-key", baseUrl: "https://example.invalid/v1", models: ["m"] }; + await turn("user_test_account__session_conversation-native", local, {}, "other/m"); + expect(wires[0]!.has("thread-id")).toBe(false); + expect(wires[0]!.has("session-id")).toBe(false); +}); + +for (const mode of ["stored", "main"] as const) for (const variant of ["local", "remote", "header"] as const) test(`${mode} pool 401 refresh keeps the dev wire for ${variant} identity on retry`, async () => { + cfg.codexAccounts = mode === "stored" ? [{ id: "fixture-a", label: "a" }] : []; + if (mode === "main") { + cfg.activeCodexAccountId = MAIN_CODEX_ACCOUNT_ID; + clearMainAccountInfoCache(); + const token = fakeChatGptJwt({ exp: Math.floor(Date.now() / 1000) + 86400, chatgpt_account_id: "fixture-main" }); + writeFileSync(join(isolated.path, "auth.json"), JSON.stringify({ tokens: { access_token: token, refresh_token: "fixture-main-grant", account_id: "fixture-main" } })); + } + let refreshes = 0; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + const url = new URL(input instanceof Request ? input.url : String(input)); + if (url.hostname === "auth.openai.com") { + refreshes++; + return Response.json({ access_token: mode === "main" ? fakeChatGptJwt({ exp: Math.floor(Date.now() / 1000) + 86400, chatgpt_account_id: "fixture-main" }) : "fixture-refreshed-access", refresh_token: "fixture-refreshed-grant", expires_in: 3600 }); + } + if (!url.pathname.endsWith("/responses")) return Response.json({ rate_limit: { primary_window: { used_percent: 10 } } }); + wires.push(new Headers(init?.headers)); + return wires.length === 1 ? Response.json({ error: { message: "expired" } }, { status: 401 }) + : Response.json({ id: "resp_fixture", object: "response", status: "completed", output: [], usage: { input_tokens: 1, output_tokens: 0, total_tokens: 1 } }); + }) as typeof fetch; + const admission: DataPlaneAdmission = variant === "remote" ? { kind: "configured", keyId: "fixture", source: "dedicated", contextPrincipalId: "fixture-principal" } : local; + await turn(variant === "header" ? undefined : "user_test_account__session_conversation-native", admission, + variant === "header" ? { "x-claude-code-session-id": "fixture-header" } : {}); + expect(refreshes).toBe(1); + expect(wires).toHaveLength(2); + for (const name of ["session_id", "session-id", "thread-id"]) { + expect(wires[1]!.get(name)).toBe(wires[0]!.get(name)); + } + expect(wires[0]!.get("session_id")).toBe(variant === "header" ? null : metadataUuid); + for (const wire of wires) expectNoSynthesizedPair(wire); +}); + +for (const variant of ["local-metadata", "remote-alias"] as const) test(`alternate-account failover binds ${variant} under the initial selection key`, async () => { + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + const url = new URL(input instanceof Request ? input.url : String(input)); + if (!url.pathname.endsWith("/responses")) return Response.json({ rate_limit: { primary_window: { used_percent: 10 } } }); + wires.push(new Headers(init?.headers)); + return wires.length === 1 ? Response.json({ error: { message: "rate limited" } }, { status: 429 }) + : Response.json({ id: "resp_fixture", object: "response", status: "completed", output: [], usage: { input_tokens: 1, output_tokens: 0, total_tokens: 1 } }); + }) as typeof fetch; + const send = () => variant === "local-metadata" + ? turn("user_test_account__session_conversation-native") + : turn(undefined, principal("fixture-principal"), { session_id: "fixture-alias" }); + // Observe the alternate pick itself: its selection headers must carry the initial pool key. + const alternatePickKeys: Array = []; + const resolveSpy = spyOn(authContext, "resolveCodexAuthContext"); + try { + await send(); + for (const call of resolveSpy.mock.calls) { + const [headers, , , options] = call as Parameters; + if (options?.excludeAccountId) alternatePickKeys.push(codexPoolAffinityKey(headers)); + } + } finally { resolveSpy.mockRestore(); } + expect(wires).toHaveLength(2); + const alternate = account(wires[1]); + expect(alternate).toBeTruthy(); + expect(alternate).not.toBe(account(wires[0])); + const identity = variant === "local-metadata" + ? bridgeIdentity(local, {}, metadataUuid) + : bridgeIdentity(principal("fixture-principal"), { session_id: "fixture-alias" }); + const key = poolKey(identity); + expect(alternatePickKeys).toEqual([key]); + // The alternate account's conversation-state issuer is recorded under the initial selection key. + expect(peekConversationStateIssuer(key)).toBe(alternate!.replace(/^account-/, "")); + await send(); + expect(account(wires[2])).toBe(alternate); + for (const wire of wires) expectNoSynthesizedPair(wire); +}); diff --git a/tests/claude-integration/claude-bridge-pool-session.test.ts b/tests/claude-integration/claude-bridge-pool-session.test.ts new file mode 100644 index 00000000000..532711467ec --- /dev/null +++ b/tests/claude-integration/claude-bridge-pool-session.test.ts @@ -0,0 +1,83 @@ +import { expect, test } from "bun:test"; +import { codexPoolAffinityKey } from "../../src/codex/auth-context"; +import { getDefaultConfig } from "../../src/config"; +import { routeConcreteModel } from "../../src/router"; +import { codexRouteCredentialDomainHeaders } from "../../src/server/responses/core-auth"; +import { applyClaudeBridgePoolSession, completeClaudeBridgePoolIdentity, resolveClaudeBridgePoolSession } from "../../src/server/claude-bridge-pool-session"; +import type { DataPlaneAdmission } from "../../src/server/auth-cors"; + +const local: DataPlaneAdmission = { kind: "loopback", source: "loopback" }; +const remote = (principal?: string): DataPlaneAdmission => ({ kind: "configured", keyId: "fixture", source: "bearer", contextPrincipalId: principal }); +const resolve = (headers: Record = {}, admission: DataPlaneAdmission | undefined = local, metadataSessionId?: string, replay = new Headers(headers)) => + resolveClaudeBridgePoolSession({ ingress: new Headers(headers), replay, admission, metadataSessionId }); + +test("admission, spawn and explicit lineage gate synthesis", () => { + expect(resolveClaudeBridgePoolSession({ ingress: new Headers({ session_id: "alias" }), replay: new Headers() })).toBeUndefined(); + expect(resolveClaudeBridgePoolSession({ ingress: new Headers(), replay: new Headers(), metadataSessionId: "metadata" })).toBeUndefined(); + for (const marker of [{ "x-openai-subagent": "collab_spawn" }, { "x-codex-turn-metadata": '{"subagent_kind":"thread_spawn"}' }]) { + expect(resolve(marker, local, "metadata")).toBeUndefined(); + expect(resolve({}, local, "metadata", new Headers(marker))).toBeUndefined(); + } + for (const name of ["thread-id", "x-codex-parent-thread-id"]) for (const value of ["", "parent"]) { + expect(resolve({ [name]: value }, local, "metadata")).toBeUndefined(); + } + expect(resolve({}, local, "metadata", new Headers({ "session-id": "existing", "thread-id": "existing" }))).toBeUndefined(); +}); + +test("aliases take precedence and conflicts or invalid values block derivation", () => { + // An alias names the same keyed identity a derived value of the same text would in that scope. + const alias = resolve({}, local, "alias"); + for (const headers of [{ session_id: "alias" }, { "session-id": "alias" }, { session_id: "alias", "session-id": "alias" }, { session_id: "", "session-id": "alias" }]) expect(resolve(headers, local, "metadata")).toBe(alias); + for (const headers of [{ session_id: "a", "session-id": "b" }, { session_id: "" }, { session_id: "bad value", "session-id": "valid" }, { session_id: "bad value" }]) expect(resolve(headers, local, "metadata")).toBeUndefined(); +}); + +test("derived metadata and header identities are keyed per admission scope", () => { + const loopback = resolve({}, local, "metadata"); + expect(loopback).toMatch(/^[0-9a-f]{64}$/); + expect(loopback).not.toBe("metadata"); + expect(resolve({}, local, "metadata")).toBe(loopback); + const a = resolve({}, remote("a"), "metadata"); + expect(a).toMatch(/^[0-9a-f]{64}$/); + expect(resolve({}, remote("a"), "metadata")).toBe(a); + expect(resolve({}, remote("b"), "metadata")).not.toBe(a); + expect(a).not.toBe(loopback); + expect(resolve({}, remote(), "metadata")).toBeUndefined(); + expect(resolve({ "x-claude-code-session-id": "header" })).toBe(resolve({}, local, "header")); + expect(resolve({ "x-claude-code-session-id": "header" }, remote("a"))).toMatch(/^[0-9a-f]{64}$/); + expect(resolve({ "x-claude-code-session-id": "bad value" })).toBeUndefined(); + expect(resolve({ "prompt_cache_key": "shared-system" })).toBeUndefined(); +}); + +test("caller aliases are keyed to their principal and never computable from public inputs", () => { + const publicDigest = (principal: string, value: string) => new Bun.CryptoHasher("sha256") + .update(JSON.stringify(["opencodex-caller-session-v1", principal, value])).digest("hex"); + for (const name of ["session_id", "session-id"]) { + const a = resolve({ [name]: "fixture-alias" }, remote("a"), "metadata"); + expect(a).toMatch(/^[0-9a-f]{64}$/); + expect(a).not.toBe(publicDigest("a", "fixture-alias")); + expect(resolve({ [name]: "fixture-alias" }, remote("a"), "metadata")).toBe(a); + expect(resolve({ [name]: "fixture-alias" }, remote("b"), "metadata")).not.toBe(a); + expect(resolve({ [name]: "fixture-alias" }, remote(), "metadata")).toBeUndefined(); + expect(resolve({ [name]: "fixture-alias" }, local, "metadata")).not.toBe("fixture-alias"); + } + // A caller that knows another principal's derived identity cannot reuse it as an alias. + const victim = resolve({}, remote("b"), "metadata")!; + expect(resolve({ session_id: victim }, remote("a"))).not.toBe(victim); +}); + +test("completion agrees with final credential-domain headers for both Codex routes", () => { + const config = getDefaultConfig(); + const route = routeConcreteModel(config, "openai/gpt-5.4")!; + const headers = new Headers(); + const options = { claudeBridgePoolSession: "conversation" }; + expect(applyClaudeBridgePoolSession(headers, { codexAccountMode: undefined }, options)).toBe(headers); + const complete = completeClaudeBridgePoolIdentity(headers, options); + const key = codexPoolAffinityKey(complete); + expect(key).toBeDefined(); + const req = new Request("http://localhost/v1/responses", { headers }); + for (const [mode, rewrite] of [["pool", false], ["direct", true]] as const) { + const codexRoute = { ...route, codexAccountMode: mode }; + expect(codexPoolAffinityKey(codexRouteCredentialDomainHeaders(req, codexRoute, options, rewrite, config))).toBe(key); + } + expect(completeClaudeBridgePoolIdentity(complete, options)).toBe(complete); +}); diff --git a/tests/fixtures/test-layout-expected.json b/tests/fixtures/test-layout-expected.json index 053e11343bf..a839b41d9d3 100644 --- a/tests/fixtures/test-layout-expected.json +++ b/tests/fixtures/test-layout-expected.json @@ -508,6 +508,7 @@ "messages-native-decline-trace.test.ts": "claude-integration", "messages-native-oauth.test.ts": "claude-integration", "messages-native-opaque-state.test.ts": "claude-integration", "messages-surface-matrix.test.ts": "claude-integration", "claude-model-info.test.ts": "claude-integration", "claude-models-discovery.test.ts": "claude-integration", + "claude-bridge-pool-session.test.ts": "claude-integration", "claude-bridge-pool-dispatch.test.ts": "claude-integration", "claude-native-affinity.test.ts": "claude-integration", "claude-native-passthrough.test.ts": "claude-integration", "messages-shared-sse.test.ts": "claude-integration", "claude-native-rate-limit-headers.test.ts": "claude-integration", "claude-outbound.test.ts": "claude-integration",