diff --git a/packages/ai/src/api/openai-codex-responses.ts b/packages/ai/src/api/openai-codex-responses.ts index 36078b2bd..52e1becfb 100644 --- a/packages/ai/src/api/openai-codex-responses.ts +++ b/packages/ai/src/api/openai-codex-responses.ts @@ -328,6 +328,7 @@ export const stream: StreamFunction<"openai-codex-responses", OpenAICodexRespons httpTimeoutMs, websocketConnectTimeoutMs, cacheSessionId, + accountId, grammarToolInputProperties, options, ); @@ -868,7 +869,7 @@ export interface OpenAICodexWebSocketDebugStats { lastWebSocketError?: string; } -const websocketSessionCache = new Map(); +const websocketSessionCache = new Map>(); const websocketDebugStats = new Map(); const websocketSseFallbackSessions = new Set(); @@ -913,13 +914,12 @@ export function closeOpenAICodexWebSocketSessions(sessionId?: string): void { closeWebSocketSilently(entry.socket, 1000, "debug_close"); }; if (sessionId) { - const entry = websocketSessionCache.get(sessionId); - if (entry) closeEntry(entry); + for (const entry of websocketSessionCache.get(sessionId)?.values() ?? []) closeEntry(entry); websocketSessionCache.delete(sessionId); return; } - for (const entry of websocketSessionCache.values()) { - closeEntry(entry); + for (const accountEntries of websocketSessionCache.values()) { + for (const entry of accountEntries.values()) closeEntry(entry); } websocketSessionCache.clear(); } @@ -1021,14 +1021,16 @@ function closeWebSocketSilently(socket: WebSocketLike, code = 1000, reason = "do } catch {} } -function scheduleSessionWebSocketExpiry(sessionId: string, entry: CachedWebSocketConnection): void { +function scheduleSessionWebSocketExpiry(sessionId: string, accountId: string, entry: CachedWebSocketConnection): void { if (entry.idleTimer) { clearTimeout(entry.idleTimer); } entry.idleTimer = setTimeout(() => { if (entry.busy) return; closeWebSocketSilently(entry.socket, 1000, "idle_timeout"); - websocketSessionCache.delete(sessionId); + const accountEntries = websocketSessionCache.get(sessionId); + if (accountEntries?.get(accountId) === entry) accountEntries.delete(accountId); + if (accountEntries?.size === 0) websocketSessionCache.delete(sessionId); }, SESSION_WEBSOCKET_CACHE_TTL_MS); } @@ -1114,6 +1116,7 @@ async function acquireWebSocket( url: string, headers: Headers, sessionId: string | undefined, + accountId: string, signal?: AbortSignal, connectTimeoutMs?: number, env?: ProviderEnv, @@ -1132,7 +1135,8 @@ async function acquireWebSocket( }; } - const cached = websocketSessionCache.get(sessionId); + let accountEntries = websocketSessionCache.get(sessionId); + const cached = accountEntries?.get(accountId); if (cached) { if (cached.idleTimer) { clearTimeout(cached.idleTimer); @@ -1140,7 +1144,8 @@ async function acquireWebSocket( } if (!cached.busy && isWebSocketSessionExpired(cached)) { closeWebSocketSilently(cached.socket, 1000, "connection_age_limit"); - websocketSessionCache.delete(sessionId); + accountEntries?.delete(accountId); + if (accountEntries?.size === 0) websocketSessionCache.delete(sessionId); } else if (!cached.busy && isWebSocketReusable(cached.socket)) { cached.busy = true; return { @@ -1150,11 +1155,13 @@ async function acquireWebSocket( release: ({ keep } = {}) => { if (!keep || !isWebSocketReusable(cached.socket)) { closeWebSocketSilently(cached.socket); - websocketSessionCache.delete(sessionId); + const currentEntries = websocketSessionCache.get(sessionId); + if (currentEntries?.get(accountId) === cached) currentEntries.delete(accountId); + if (currentEntries?.size === 0) websocketSessionCache.delete(sessionId); return; } cached.busy = false; - scheduleSessionWebSocketExpiry(sessionId, cached); + scheduleSessionWebSocketExpiry(sessionId, accountId, cached); }, }; } @@ -1170,13 +1177,19 @@ async function acquireWebSocket( } if (!isWebSocketReusable(cached.socket)) { closeWebSocketSilently(cached.socket); - websocketSessionCache.delete(sessionId); + accountEntries?.delete(accountId); + if (accountEntries?.size === 0) websocketSessionCache.delete(sessionId); } } const socket = await connectWebSocket(url, headers, signal, connectTimeoutMs, env); const entry: CachedWebSocketConnection = { socket, busy: true, createdAt: Date.now() }; - websocketSessionCache.set(sessionId, entry); + accountEntries = websocketSessionCache.get(sessionId); + if (!accountEntries) { + accountEntries = new Map(); + websocketSessionCache.set(sessionId, accountEntries); + } + accountEntries.set(accountId, entry); return { socket, entry, @@ -1185,13 +1198,13 @@ async function acquireWebSocket( if (!keep || !isWebSocketReusable(entry.socket)) { closeWebSocketSilently(entry.socket); if (entry.idleTimer) clearTimeout(entry.idleTimer); - if (websocketSessionCache.get(sessionId) === entry) { - websocketSessionCache.delete(sessionId); - } + const currentEntries = websocketSessionCache.get(sessionId); + if (currentEntries?.get(accountId) === entry) currentEntries.delete(accountId); + if (currentEntries?.size === 0) websocketSessionCache.delete(sessionId); return; } entry.busy = false; - scheduleSessionWebSocketExpiry(sessionId, entry); + scheduleSessionWebSocketExpiry(sessionId, accountId, entry); }, }; } @@ -1450,6 +1463,7 @@ async function processWebSocketStream( idleTimeoutMs: number | undefined, websocketConnectTimeoutMs: number | undefined, cacheSessionId: string | undefined, + accountId: string, grammarToolInputProperties: ReadonlyMap, options?: OpenAICodexResponsesOptions, ): Promise { @@ -1457,6 +1471,7 @@ async function processWebSocketStream( url, headers, cacheSessionId, + accountId, options?.signal, websocketConnectTimeoutMs, options?.env, diff --git a/packages/ai/test/openai-codex-stream.test.ts b/packages/ai/test/openai-codex-stream.test.ts index 582ec9a1c..a8423c887 100644 --- a/packages/ai/test/openai-codex-stream.test.ts +++ b/packages/ai/test/openai-codex-stream.test.ts @@ -28,9 +28,9 @@ afterEach(() => { vi.restoreAllMocks(); }); -function mockToken(): string { +function mockToken(accountId = "acc_test"): string { const payload = Buffer.from( - JSON.stringify({ "https://api.openai.com/auth": { chatgpt_account_id: "acc_test" } }), + JSON.stringify({ "https://api.openai.com/auth": { chatgpt_account_id: accountId } }), "utf8", ).toString("base64"); return `aaa.${payload}.bbb`; @@ -1334,6 +1334,97 @@ describe("openai-codex streaming", () => { }); }); + it("scopes cached websockets to the authenticated account", async () => { + // Regression for #7284: rotating accounts must not reuse a socket authenticated by another account. + const connectedHeaders: Record[] = []; + let responseId = 0; + + class MockWebSocket { + static OPEN = 1; + readyState = MockWebSocket.OPEN; + private listeners = new Map void>>(); + + constructor(_url: string, protocols?: string | string[] | { headers?: Record }) { + const headers = + protocols && typeof protocols === "object" && !Array.isArray(protocols) ? protocols.headers : undefined; + connectedHeaders.push(headers ?? {}); + queueMicrotask(() => this.dispatch("open", {})); + } + + addEventListener(type: string, listener: (event: unknown) => void): void { + let listeners = this.listeners.get(type); + if (!listeners) { + listeners = new Set(); + this.listeners.set(type, listeners); + } + listeners.add(listener); + } + + removeEventListener(type: string, listener: (event: unknown) => void): void { + this.listeners.get(type)?.delete(listener); + } + + send(): void { + queueMicrotask(() => { + this.dispatch("message", { + data: JSON.stringify({ + type: "response.completed", + response: { + id: `resp_${++responseId}`, + status: "completed", + usage: { input_tokens: 1, output_tokens: 1, total_tokens: 2 }, + }, + }), + }); + }); + } + + close(): void { + this.readyState = 3; + } + + private dispatch(type: string, event: unknown): void { + for (const listener of this.listeners.get(type) ?? []) listener(event); + } + } + + vi.stubGlobal("WebSocket", MockWebSocket); + vi.stubGlobal( + "fetch", + vi.fn(async () => new Response("unexpected fetch", { status: 500 })), + ); + + const model: Model<"openai-codex-responses"> = { + id: "gpt-5.1-codex", + name: "GPT-5.1 Codex", + api: "openai-codex-responses", + provider: "openai-codex", + baseUrl: "https://chatgpt.com/backend-api", + reasoning: true, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 400000, + maxTokens: 128000, + }; + const context: Context = { systemPrompt: "", messages: [] }; + const options = { sessionId: "shared-session", transport: "websocket-cached" as const }; + + await streamOpenAICodexResponses(model, context, { ...options, apiKey: mockToken("account-a") }).result(); + await streamOpenAICodexResponses(model, context, { ...options, apiKey: mockToken("account-b") }).result(); + await streamOpenAICodexResponses(model, context, { ...options, apiKey: mockToken("account-a") }).result(); + + expect(connectedHeaders.map((headers) => headers["chatgpt-account-id"])).toEqual(["account-a", "account-b"]); + expect(connectedHeaders.map((headers) => headers.authorization)).toEqual([ + `Bearer ${mockToken("account-a")}`, + `Bearer ${mockToken("account-b")}`, + ]); + expect(global.fetch).not.toHaveBeenCalled(); + expect(getOpenAICodexWebSocketDebugStats("shared-session")).toMatchObject({ + connectionsCreated: 2, + connectionsReused: 1, + }); + }); + it("closes one-shot websockets when cacheRetention is none", async () => { const token = mockToken(); const sentBodies: Array<{ prompt_cache_key?: string }> = [];