diff --git a/src/codex/auth-context.ts b/src/codex/auth-context.ts index 660e4f61026..fb6d40e5259 100644 --- a/src/codex/auth-context.ts +++ b/src/codex/auth-context.ts @@ -322,8 +322,7 @@ export type CodexAuthContext = transientProbe?: TransientProbeGrant; } | { - // Main Codex account participating in rotation: token injected from ~/.codex/auth.json - // (Option A). Distinct from "main" (request-owned passthrough or Direct mode). + // Main account rotating in the Pool (Option A), token from ~/.codex/auth.json; never request-owned "main". kind: "main-pool"; accountId: string; writerGeneration: number; @@ -336,14 +335,14 @@ export type CodexAuthContext = chatgptAccountId: string; /** Bypass Pool selection and suppress quota/transient failover for an exact selector. */ fixedAccount?: boolean; - /** See `pool.affinityKey`. */ - affinityKey?: string; + affinityKey?: string; // See `pool.affinityKey`. /** See `pool.probeLeaseId`. */ probeLeaseId?: string; quotaScope?: CodexQuotaScope; probeQuotaScope?: CodexQuotaScope; /** See `pool.transientProbe`. */ transientProbe?: TransientProbeGrant; + affinityDecision?: CodexAffinityDecision; }; /** Probe lease carried by this context, when it holds one. */ @@ -1733,6 +1732,7 @@ async function resolveCodexAuthContextUnannotated( ...(quotaScope ? { quotaScope } : {}), ...(probeLeaseId ? { probeLeaseId } : {}), ...(probeQuotaScope ? { probeQuotaScope } : {}), + ...(affinityDecision ? { affinityDecision } : {}), ...(transientProbe ? { transientProbe } : {}), }; } diff --git a/src/server/request-log.ts b/src/server/request-log.ts index cf8feddc5b9..aeedf279e40 100644 --- a/src/server/request-log.ts +++ b/src/server/request-log.ts @@ -19,6 +19,7 @@ import { } from "../lib/errors"; import { CODEX_CONFIG_PATH, readRootTomlString } from "../codex/paths"; import type { CodexAffinityMove, CodexAffinityReason } from "../codex/routing"; +import type { CodexAuthContext } from "../codex/auth-context"; import { readCodexCatalogPath } from "../codex/catalog"; import type { AttemptTierOutcome, OcxProviderConfig, OcxUsage } from "../types"; import { normalizeRouteDecisionTrace, type RouteDecisionTraceV1 } from "../routing/trace"; @@ -1263,6 +1264,18 @@ export function noteAffinityMove( (logCtx.affinityMoveReasons ??= []).push(reason); } +/** + * Copy the selection's own affinity decision onto the request log. Both Pool-served kinds carry + * one: a stored Pool account and the main login rotating as `main-pool` (#4546, #7799). A move + * then reads as an event with its cause instead of something inferred from account labels. + */ +export function noteCodexSelectionAffinity(logCtx: RequestLogContext, authCtx: CodexAuthContext): void { + const decision = authCtx.kind === "pool" || authCtx.kind === "main-pool" ? authCtx.affinityDecision : undefined; + if (!decision) return; + logCtx.affinity = decision.move; + logCtx.affinityReason = decision.reason; +} + /** * The affinity scope a released binding belonged to: one thread, one model lane. * diff --git a/src/server/responses/request-prepare.ts b/src/server/responses/request-prepare.ts index 353823b43c2..637c3a74262 100644 --- a/src/server/responses/request-prepare.ts +++ b/src/server/responses/request-prepare.ts @@ -76,7 +76,7 @@ import { anthropicSessionKeyFromParts } from "../../oauth/anthropic-routing"; import { configuredAnthropicInstance } from "../../providers/anthropic-instance"; import { isTranslatorBudgetExceededError } from "../../lib/translator-budget"; import { bindTurnTerminationScope, rememberDeliveredFinalAnswer } from "../../responses/turn-termination"; -import { observeCacheDiagnosticInbound, rebindCacheDiagnosticBody, requestLogSpeedLabel, readConfiguredCodexServiceTier } from "../request-log"; +import { noteCodexSelectionAffinity, observeCacheDiagnosticInbound, rebindCacheDiagnosticBody, requestLogSpeedLabel, readConfiguredCodexServiceTier } from "../request-log"; import type { RouteResult } from "../../router"; import { captureRouteStaticPolicy, @@ -1585,10 +1585,7 @@ export async function prepareResponsesRequest( // A move is the expensive event: it discards the prefix warmed on the previous account. Record // it as an event with its cause, so the operator reads it off one line instead of inferring it // from account labels across many (#4546). - if (admissionState.authCtx.kind === "pool" && admissionState.authCtx.affinityDecision) { - logCtx.affinity = admissionState.authCtx.affinityDecision.move; - logCtx.affinityReason = admissionState.authCtx.affinityDecision.reason; - } + noteCodexSelectionAffinity(logCtx, admissionState.authCtx); { const binding = conversationStateBindingFromAuth(admissionState.authCtx, poolAffinityKey); if (binding) { diff --git a/tests/codex-integration/codex-auth-context.test.ts b/tests/codex-integration/codex-auth-context.test.ts index b26014ff727..f6bdc2fb2eb 100644 --- a/tests/codex-integration/codex-auth-context.test.ts +++ b/tests/codex-integration/codex-auth-context.test.ts @@ -696,6 +696,32 @@ describe("Codex auth context", () => { }); }); + test("main-pool selection carries the same affinity decision as a stored-pool selection (#7799)", async () => { + const cfg = config(); + cfg.activeCodexAccountId = MAIN_CODEX_ACCOUNT_ID; + writeFileSync(join(testDir, "auth.json"), JSON.stringify({ + tokens: { access_token: "main-token", account_id: "main-account" }, + })); + // Thread and session ids give the request a bindable conversation identity. + const headers = { "thread-id": "affinity-7799", "session-id": "affinity-7799" }; + const resolve = () => resolveCodexAuthContext(new Headers(headers), cfg, "pool", { + isMainAccountTokenLive: () => true, + getMainAccountToken: () => ({ accessToken: "main-token", chatgptAccountId: "main-account" }), + primeCodexPoolQuotas: async () => {}, + }); + + await expect(resolve()).resolves.toMatchObject({ + kind: "main-pool", + accountId: MAIN_CODEX_ACCOUNT_ID, + affinityDecision: { move: "new_bind", reason: "healthy" }, + }); + await expect(resolve()).resolves.toMatchObject({ + kind: "main-pool", + accountId: MAIN_CODEX_ACCOUNT_ID, + affinityDecision: { move: "reused", reason: "healthy" }, + }); + }); + test("account-gated routing distinguishes an unavailable grant from no grant", async () => { const cfg = config(); saveCodexAccountCredential("pool-a", { diff --git a/tests/codex-integration/codex-main-rotation.test.ts b/tests/codex-integration/codex-main-rotation.test.ts index bfc114c935d..4c9c86665e3 100644 --- a/tests/codex-integration/codex-main-rotation.test.ts +++ b/tests/codex-integration/codex-main-rotation.test.ts @@ -170,6 +170,7 @@ describe("main account rotation (Option A)", () => { identityKey: expect.stringMatching(/^[a-f0-9]{64}$/), identityGeneration: expect.any(Number), }, + affinityDecision: { move: "new_bind", reason: "healthy" }, }); expect(isCodexAuthContextUsable(ctx, config)).toBe(true); const headers = headersForCodexAuthContext(new Headers(), ctx); @@ -202,6 +203,7 @@ describe("main account rotation (Option A)", () => { identityKey: expect.stringMatching(/^[a-f0-9]{64}$/), identityGeneration: expect.any(Number), }, + affinityDecision: { move: "new_bind", reason: "healthy" }, }); expect(isCodexAccountInCooldown(MAIN_CODEX_ACCOUNT_ID)).toBe(false); expect(isAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID)).toBe(false); diff --git a/tests/server/spend-instrumentation-log.test.ts b/tests/server/spend-instrumentation-log.test.ts index acf7c3eaa67..eeaaa2bdcaa 100644 --- a/tests/server/spend-instrumentation-log.test.ts +++ b/tests/server/spend-instrumentation-log.test.ts @@ -11,6 +11,7 @@ import { finishRequestAttempt, noteAffinityMove, noteAttemptSend, + noteCodexSelectionAffinity, noteNoAccountAffinityReason, recordNoAccountAffinityFailure, requestLogEntryFromPersistedUsage, @@ -201,3 +202,31 @@ describe("no-account failures explain themselves", () => { clearNoAccountAffinityReasonsForTests(); }); }); + +describe("Pool selections report their affinity decision on the request log (#7799)", () => { + const decision = { move: "rebound", reason: "quota_refusal" } as const; + const base = { accountId: "acct", writerGeneration: 0, accessToken: "token", chatgptAccountId: "chatgpt" }; + + test("a main-pool selection reports the same decision fields as a stored-pool selection", () => { + const mainPool: RequestLogContext = { provider: "openai", model: "gpt-5.6-sol" }; + noteCodexSelectionAffinity(mainPool, { kind: "main-pool", ...base, affinityDecision: decision }); + expect(mainPool.affinity).toBe("rebound"); + expect(mainPool.affinityReason).toBe("quota_refusal"); + + const pool: RequestLogContext = { provider: "openai", model: "gpt-5.6-sol" }; + noteCodexSelectionAffinity(pool, { kind: "pool", ...base, generation: 1, affinityDecision: decision }); + expect({ affinity: pool.affinity, affinityReason: pool.affinityReason }) + .toEqual({ affinity: mainPool.affinity, affinityReason: mainPool.affinityReason }); + }); + + test("a selection without a decision, or a request-owned main context, leaves the fields unset", () => { + const noDecision: RequestLogContext = { provider: "openai", model: "gpt-5.6-sol" }; + noteCodexSelectionAffinity(noDecision, { kind: "main-pool", ...base }); + expect(noDecision.affinity).toBeUndefined(); + expect(noDecision.affinityReason).toBeUndefined(); + + const requestOwned: RequestLogContext = { provider: "openai", model: "gpt-5.6-sol" }; + noteCodexSelectionAffinity(requestOwned, { kind: "main", accountId: null }); + expect(requestOwned.affinity).toBeUndefined(); + }); +});