diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index 7281bcd305..11ed2284b8 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -15,6 +15,7 @@ import { nativeContextLimits } from "../../codex/catalog"; import { describeUpstreamConnectFailure } from "./upstream-error"; import type { CodexWsQuotaObserver } from "./codex-ws-metadata"; import { applyAccountQuotaFromUpstreamHeaders as applyCapturedCodexQuota } from "../../codex/quota"; +import { isCodexAccountGenerationLive } from "../../codex/account-store"; import { isCodexWsQuotaObservedResponse } from "./ws-upstream"; import { multiAgentGuidanceEnabled, @@ -1004,8 +1005,12 @@ export function usesCodexForwardPoolAuth( function codexWsQuotaObserver(authCtx: CodexAuthContext, provider: OcxProviderConfig): CodexWsQuotaObserver | undefined { if (!isCanonicalOpenAiForwardProvider(provider) || !usesCodexForwardPoolAuth(authCtx, provider)) return undefined; const { accountId, writerGeneration } = authCtx; + const credentialGeneration = authCtx.kind === "pool" ? authCtx.generation : undefined; const mainWriter = authCtx.kind === "main-pool" ? authCtx.mainQuotaWriter : undefined; - return headers => applyCapturedCodexQuota(accountId, headers, writerGeneration, mainWriter); + return headers => { + if (credentialGeneration !== undefined && !isCodexAccountGenerationLive(accountId, credentialGeneration)) return; + applyCapturedCodexQuota(accountId, headers, writerGeneration, mainWriter); + }; } export function preAuthUpstreamHostCircuitKey( diff --git a/tests/responses/responses-account-label.test.ts b/tests/responses/responses-account-label.test.ts index b7be5c8d7f..e96c4556d9 100644 --- a/tests/responses/responses-account-label.test.ts +++ b/tests/responses/responses-account-label.test.ts @@ -190,6 +190,65 @@ describe("Responses account usage attribution", () => { } }); + test("late WS quota from a replaced pool credential cannot repopulate cleared state", async () => { + const originalWebSocket = globalThis.WebSocket; + let releaseFinalQuota!: () => void; + const finalQuotaAllowed = new Promise(resolve => { releaseFinalQuota = resolve; }); + try { + await withPoolHome(async () => { + savePoolCredential("pool-ws-replaced"); + class MetadataSocket { + listeners = new Map void>>(); + constructor() { queueMicrotask(() => this.emit("open", {})); } + addEventListener(type: string, listener: (event: unknown) => void) { + this.listeners.set(type, [...(this.listeners.get(type) ?? []), listener]); + } + removeEventListener(type: string, listener: (event: unknown) => void) { + this.listeners.set(type, (this.listeners.get(type) ?? []).filter(value => value !== listener)); + } + emit(type: string, event: unknown) { + for (const listener of this.listeners.get(type) ?? []) listener(event); + } + send() { + const payload = (value: unknown) => this.emit("message", { data: JSON.stringify(value) }); + queueMicrotask(() => { + payload({ type: "codex.rate_limits", rate_limits: { + primary: { used_percent: 10, window_minutes: 10080 }, + } }); + payload({ type: "response.created", response: { id: "quota-response" } }); + void finalQuotaAllowed.then(() => { + payload({ type: "codex.rate_limits", rate_limits: { + primary: { used_percent: 100, window_minutes: 10080 }, + } }); + payload({ type: "response.completed", response: { id: "quota-response", status: "completed", output: [] } }); + }); + }); + } + close() { this.emit("close", {}); } + } + globalThis.WebSocket = MetadataSocket as unknown as typeof WebSocket; + globalThis.fetch = (async () => { throw new Error("unexpected HTTP request"); }) as typeof fetch; + const response = await handleResponses(new Request("http://localhost/v1/responses", { + method: "POST", headers: { "content-type": "application/json" }, + body: JSON.stringify({ model: "gpt-5.5", input: "hello", stream: true }), + }), poolConfig(["pool-ws-replaced"]), { model: "", provider: "" }, { + codexWsRuntimeIdentity: "1.4.0", + }); + expect(getAccountQuota("pool-ws-replaced")?.weeklyPercent).toBe(10); + + savePoolCredential("pool-ws-replaced"); + clearAccountQuota("pool-ws-replaced"); + releaseFinalQuota(); + await response.text(); + + expect(getAccountQuota("pool-ws-replaced")).toBeNull(); + }); + } finally { + releaseFinalQuota(); + globalThis.WebSocket = originalWebSocket; + } + }); + test("main-pool and legacy added accounts carry their effective labels", async () => { await withPoolHome(async home => { writeFileSync(join(home, "auth.json"), JSON.stringify({