From 1d8889001fcbe5a5352575f92cace27127932618 Mon Sep 17 00:00:00 2001 From: Cita Date: Sat, 21 Mar 2026 19:07:14 +0800 Subject: [PATCH] fix: harden codex runtime routing and oauth recovery --- src/server/config.test.ts | 9 + src/server/config.ts | 1 + .../providers/codexProviderProfile.ts | 3 +- .../runtime/codexWebsocketHeaders.ts | 7 +- .../runtime/codexWebsocketRuntime.test.ts | 93 +++++- .../runtime/codexWebsocketRuntime.ts | 38 ++- src/server/proxy-core/surfaces/chatSurface.ts | 4 +- .../surfaces/openAiResponsesSurface.ts | 11 +- src/server/routes/api/oauth.test.ts | 147 ++++++++ .../routes/proxy/chat.codex-oauth.test.ts | 43 +++ .../proxy/responses.codex-oauth.test.ts | 47 ++- .../routes/proxy/responses.websocket.test.ts | 163 +++++++-- src/server/routes/proxy/responsesWebsocket.ts | 314 +++++++++++------- .../routes/proxy/upstreamEndpoint.test.ts | 7 +- src/server/services/oauth/codexProvider.ts | 8 +- src/server/services/oauth/service.ts | 21 +- .../services/tokenRouter.siteStatus.test.ts | 45 +++ src/server/services/tokenRouter.ts | 9 +- 18 files changed, 775 insertions(+), 195 deletions(-) diff --git a/src/server/config.test.ts b/src/server/config.test.ts index 85db8ae..ecaf6f1 100644 --- a/src/server/config.test.ts +++ b/src/server/config.test.ts @@ -51,12 +51,21 @@ describe('buildConfig', () => { const config = buildConfig({}); expect(config.codexClientId).toBe('app_EMoamEEZ73f0CkXaXp7hrann'); + expect(config.codexResponsesWebsocketBeta).toBe('responses_websockets=2026-02-06'); expect(config.claudeClientId).toBe('9d1c250a-e61b-44d9-88ed-5944d1962f5e'); expect(config.claudeClientSecret).toBe(''); expect(config.geminiCliClientId).toBe('681255809395-oo8ft2oprdrnp9e3aqf6av3hmdib135j.apps.googleusercontent.com'); expect(config.geminiCliClientSecret).toBe('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'); }); + it('allows overriding the codex websocket beta gate from environment', () => { + const config = buildConfig({ + CODEX_RESPONSES_WEBSOCKET_BETA: 'responses_websockets=2099-01-01', + }); + + expect(config.codexResponsesWebsocketBeta).toBe('responses_websockets=2099-01-01'); + }); + it('accepts JSON request bodies larger than Fastify default 1 MiB', async () => { const app = Fastify(buildFastifyOptions(buildConfig({}))); const largeText = 'a'.repeat(2 * 1024 * 1024); diff --git a/src/server/config.ts b/src/server/config.ts index 7f1f083..2b85533 100644 --- a/src/server/config.ts +++ b/src/server/config.ts @@ -114,6 +114,7 @@ export function buildConfig(env: NodeJS.ProcessEnv) { proxyFileRetentionPruneIntervalMinutes: Math.max(1, Math.trunc(parseNumber(env.PROXY_FILE_RETENTION_PRUNE_INTERVAL_MINUTES, 60))), proxyErrorKeywords: parseCsvList(env.PROXY_ERROR_KEYWORDS), proxyEmptyContentFailEnabled: parseBoolean(env.PROXY_EMPTY_CONTENT_FAIL, false), + codexResponsesWebsocketBeta: parseOptionalSecret(env.CODEX_RESPONSES_WEBSOCKET_BETA) || 'responses_websockets=2026-02-06', codexHeaderDefaults: { userAgent: parseOptionalSecret(env.CODEX_HEADER_DEFAULTS_USER_AGENT), betaFeatures: parseOptionalSecret(env.CODEX_HEADER_DEFAULTS_BETA_FEATURES), diff --git a/src/server/proxy-core/providers/codexProviderProfile.ts b/src/server/proxy-core/providers/codexProviderProfile.ts index 948f03e..547449b 100644 --- a/src/server/proxy-core/providers/codexProviderProfile.ts +++ b/src/server/proxy-core/providers/codexProviderProfile.ts @@ -4,7 +4,6 @@ import { config } from '../../config.js'; const CODEX_CLIENT_VERSION = '0.101.0'; const CODEX_DEFAULT_USER_AGENT = 'codex_cli_rs/0.101.0 (Mac OS 26.0.1; arm64) Apple_Terminal/464'; -const CODEX_RESPONSES_WEBSOCKET_BETA = 'responses_websockets=2026-02-06'; function asTrimmedString(value: unknown): string { return typeof value === 'string' ? value.trim() : ''; @@ -81,7 +80,7 @@ function buildCodexRuntimeHeaders(input: { const codexTurnMetadata = getInputHeader(input.baseHeaders, 'x-codex-turn-metadata'); const timingMetrics = getInputHeader(input.baseHeaders, 'x-responsesapi-include-timing-metrics'); const openAiBeta = getInputHeader(input.baseHeaders, 'openai-beta') - || (websocketTransport ? CODEX_RESPONSES_WEBSOCKET_BETA : null); + || (websocketTransport ? asTrimmedString(config.codexResponsesWebsocketBeta) : null); const explicitSessionId = asTrimmedString(input.explicitSessionId); const continuityKey = asTrimmedString(input.continuityKey); const sessionId = ( diff --git a/src/server/proxy-core/runtime/codexWebsocketHeaders.ts b/src/server/proxy-core/runtime/codexWebsocketHeaders.ts index 308f6cc..fbae37f 100644 --- a/src/server/proxy-core/runtime/codexWebsocketHeaders.ts +++ b/src/server/proxy-core/runtime/codexWebsocketHeaders.ts @@ -1,4 +1,4 @@ -const CODEX_RESPONSES_WEBSOCKET_BETA = 'responses_websockets=2026-02-06'; +import { config } from '../../config.js'; function getHeaderValue(headers: Record, key: string): string { const expected = key.trim().toLowerCase(); @@ -11,13 +11,14 @@ function getHeaderValue(headers: Record, key: string): string { export function buildCodexWebsocketHandshakeHeaders(headers: Record): Record { const next = { ...headers }; + const websocketBeta = (config.codexResponsesWebsocketBeta || '').trim() || 'responses_websockets=2026-02-06'; const openAiBeta = getHeaderValue(next, 'openai-beta').trim(); if (!openAiBeta) { - next['OpenAI-Beta'] = CODEX_RESPONSES_WEBSOCKET_BETA; + next['OpenAI-Beta'] = websocketBeta; return next; } if (!openAiBeta.includes('responses_websockets=')) { - next['OpenAI-Beta'] = `${openAiBeta},${CODEX_RESPONSES_WEBSOCKET_BETA}`; + next['OpenAI-Beta'] = `${openAiBeta},${websocketBeta}`; } return next; } diff --git a/src/server/proxy-core/runtime/codexWebsocketRuntime.test.ts b/src/server/proxy-core/runtime/codexWebsocketRuntime.test.ts index 318def5..27f415c 100644 --- a/src/server/proxy-core/runtime/codexWebsocketRuntime.test.ts +++ b/src/server/proxy-core/runtime/codexWebsocketRuntime.test.ts @@ -7,6 +7,7 @@ describe('codexWebsocketRuntime', () => { let upstreamWsUrl: string; let upstreamConnectionCount = 0; let upstreamRequests: Record[] = []; + let upstreamMessageHandler: (socket: import('ws').WebSocket, parsed: Record, requestIndex: number) => void; beforeAll(async () => { upstreamServer = new WebSocketServer({ port: 0 }); @@ -15,22 +16,7 @@ describe('codexWebsocketRuntime', () => { socket.on('message', (payload) => { const parsed = JSON.parse(String(payload)) as Record; upstreamRequests.push(parsed); - const responseId = `resp-${upstreamRequests.length}`; - socket.send(JSON.stringify({ - type: 'response.completed', - response: { - id: responseId, - object: 'response', - model: parsed.model || 'gpt-5.4', - status: 'completed', - output: [], - usage: { - input_tokens: 1, - output_tokens: 1, - total_tokens: 2, - }, - }, - })); + upstreamMessageHandler(socket, parsed, upstreamRequests.length); }); }); await new Promise((resolve) => upstreamServer.once('listening', () => resolve())); @@ -41,6 +27,24 @@ describe('codexWebsocketRuntime', () => { beforeEach(() => { upstreamConnectionCount = 0; upstreamRequests = []; + upstreamMessageHandler = (socket, parsed, requestIndex) => { + const responseId = `resp-${requestIndex}`; + socket.send(JSON.stringify({ + type: 'response.completed', + response: { + id: responseId, + object: 'response', + model: parsed.model || 'gpt-5.4', + status: 'completed', + output: [], + usage: { + input_tokens: 1, + output_tokens: 1, + total_tokens: 2, + }, + }, + })); + }; }); afterAll(async () => { @@ -134,4 +138,61 @@ describe('codexWebsocketRuntime', () => { expect(upstreamConnectionCount).toBe(2); await runtime.closeSession('exec-session-close'); }); + + it('treats response.incomplete as a terminal websocket failure instead of waiting for socket close', async () => { + upstreamMessageHandler = (socket, parsed) => { + socket.send(JSON.stringify({ + type: 'response.incomplete', + response: { + id: 'resp-incomplete', + model: parsed.model || 'gpt-5.4', + status: 'incomplete', + }, + })); + }; + + const { createCodexWebsocketRuntime, CodexWebsocketRuntimeError } = await import('./codexWebsocketRuntime.js'); + const runtime = createCodexWebsocketRuntime(); + + await expect(runtime.sendRequest({ + sessionId: 'exec-session-incomplete', + requestUrl: upstreamWsUrl, + headers: { + Authorization: 'Bearer oauth-access-token', + 'OpenAI-Beta': 'responses_websockets=2026-02-06', + }, + body: { + model: 'gpt-5.4', + input: [], + }, + })).rejects.toBeInstanceOf(CodexWebsocketRuntimeError); + }); + + it('treats top-level error frames as terminal websocket failures', async () => { + upstreamMessageHandler = (socket) => { + socket.send(JSON.stringify({ + type: 'error', + error: { + message: 'account mismatch', + type: 'invalid_request_error', + }, + })); + }; + + const { createCodexWebsocketRuntime, CodexWebsocketRuntimeError } = await import('./codexWebsocketRuntime.js'); + const runtime = createCodexWebsocketRuntime(); + + await expect(runtime.sendRequest({ + sessionId: 'exec-session-error', + requestUrl: upstreamWsUrl, + headers: { + Authorization: 'Bearer oauth-access-token', + 'OpenAI-Beta': 'responses_websockets=2026-02-06', + }, + body: { + model: 'gpt-5.4', + input: [], + }, + })).rejects.toBeInstanceOf(CodexWebsocketRuntimeError); + }); }); diff --git a/src/server/proxy-core/runtime/codexWebsocketRuntime.ts b/src/server/proxy-core/runtime/codexWebsocketRuntime.ts index 560e76f..9bedbc7 100644 --- a/src/server/proxy-core/runtime/codexWebsocketRuntime.ts +++ b/src/server/proxy-core/runtime/codexWebsocketRuntime.ts @@ -23,7 +23,31 @@ function isRecord(value: unknown): value is Record { function isTerminalEvent(payload: Record): boolean { const type = asTrimmedString(payload.type); - return type === 'response.completed' || type === 'response.failed'; + return type === 'response.completed' + || type === 'response.failed' + || type === 'response.incomplete' + || type === 'error'; +} + +function isFailureTerminalEvent(payload: Record): boolean { + const type = asTrimmedString(payload.type); + return type === 'response.failed' || type === 'response.incomplete' || type === 'error'; +} + +function extractTerminalErrorMessage(payload: Record): string { + const type = asTrimmedString(payload.type); + if (type === 'error' && isRecord(payload.error)) { + return asTrimmedString(payload.error.message) || 'upstream websocket error'; + } + if ((type === 'response.failed' || type === 'response.incomplete') && isRecord(payload.response)) { + if (isRecord(payload.response.error)) { + return asTrimmedString(payload.response.error.message) || `upstream ${type}`; + } + if (isRecord(payload.response.incomplete_details)) { + return asTrimmedString(payload.response.incomplete_details.reason) || `upstream ${type}`; + } + } + return `upstream ${type || 'websocket error'}`; } export class CodexWebsocketRuntimeError extends Error { @@ -212,6 +236,18 @@ async function sendSessionRequest( events.push(parsed); if (!isTerminalEvent(parsed)) return; if (settled) return; + if (isFailureTerminalEvent(parsed)) { + settled = true; + cleanup(); + clearSessionSocket(session, socket); + void closeSocket(socket); + reject(new CodexWebsocketRuntimeError(extractTerminalErrorMessage(parsed), { + events: [...events], + status: typeof parsed.status === 'number' && Number.isFinite(parsed.status) ? parsed.status : undefined, + payload: parsed, + })); + return; + } settled = true; cleanup(); resolve({ diff --git a/src/server/proxy-core/surfaces/chatSurface.ts b/src/server/proxy-core/surfaces/chatSurface.ts index b1809ef..b6a7c07 100644 --- a/src/server/proxy-core/surfaces/chatSurface.ts +++ b/src/server/proxy-core/surfaces/chatSurface.ts @@ -234,7 +234,7 @@ export async function handleChatSurfaceRequest( dispatchRequest, }); const tryRecover = async (ctx: Parameters>[0]) => { - if (ctx.response.status === 401 && oauth) { + if ((ctx.response.status === 401 || ctx.response.status === 403) && oauth) { try { const refreshed = await refreshOauthAccessTokenSingleflight(selected.account.id); selected.tokenValue = refreshed.accessToken; @@ -916,7 +916,7 @@ export async function handleClaudeCountTokensSurfaceRequest( }, resolveChannelProxyUrl(selected.site, selected.account.extraConfig)), }); - if (upstream.status === 401 && oauth) { + if ((upstream.status === 401 || upstream.status === 403) && oauth) { try { const refreshed = await refreshOauthAccessTokenSingleflight(selected.account.id); selected.tokenValue = refreshed.accessToken; diff --git a/src/server/proxy-core/surfaces/openAiResponsesSurface.ts b/src/server/proxy-core/surfaces/openAiResponsesSurface.ts index 006284e..598aa67 100644 --- a/src/server/proxy-core/surfaces/openAiResponsesSurface.ts +++ b/src/server/proxy-core/surfaces/openAiResponsesSurface.ts @@ -137,13 +137,6 @@ function carriesResponsesFileUrlInput(value: unknown): boolean { type UsageSummary = ReturnType; -function deriveCodexExplicitSessionId(body: Record): string | null { - const promptCacheKey = typeof body.prompt_cache_key === 'string' - ? body.prompt_cache_key.trim() - : ''; - return promptCacheKey || null; -} - export async function handleOpenAiResponsesSurfaceRequest( request: FastifyRequest, reply: FastifyReply, @@ -274,7 +267,6 @@ export async function handleOpenAiResponsesSurfaceRequest( ); const buildEndpointRequest = (endpoint: 'chat' | 'messages' | 'responses') => { const upstreamStream = isStream || (isCodexSite && endpoint === 'responses'); - const codexExplicitSessionId = deriveCodexExplicitSessionId(normalizedResponsesBody); const endpointRequest = buildUpstreamEndpointRequest({ endpoint, modelName, @@ -289,7 +281,6 @@ export async function handleOpenAiResponsesSurfaceRequest( responsesOriginalBody: normalizedResponsesBody, downstreamHeaders: request.headers as Record, providerHeaders: buildProviderHeaders(), - codexExplicitSessionId, }); const upstreamPath = ( isCompactRequest && endpoint === 'responses' @@ -323,7 +314,7 @@ export async function handleOpenAiResponsesSurfaceRequest( dispatchRequest, }); const tryRecover = async (ctx: Parameters>[0]) => { - if (ctx.response.status === 401 && oauth) { + if ((ctx.response.status === 401 || ctx.response.status === 403) && oauth) { try { const refreshed = await refreshOauthAccessTokenSingleflight(selected.account.id); selected.tokenValue = refreshed.accessToken; diff --git a/src/server/routes/api/oauth.test.ts b/src/server/routes/api/oauth.test.ts index d00d588..f873b3f 100644 --- a/src/server/routes/api/oauth.test.ts +++ b/src/server/routes/api/oauth.test.ts @@ -500,6 +500,153 @@ describe('oauth routes', { timeout: 15_000 }, () => { }); }); + it('keeps the existing codex connection intact when a rebind callback fails model discovery', async () => { + const originalJwt = buildJwt({ + email: 'codex-existing@example.com', + 'https://api.openai.com/auth': { + chatgpt_account_id: 'chatgpt-account-existing', + chatgpt_plan_type: 'team', + }, + }); + const site = await db.insert(schema.sites).values({ + name: 'ChatGPT Codex OAuth', + url: 'https://chatgpt.com/backend-api/codex', + platform: 'codex', + status: 'active', + }).returning().get(); + const existing = await db.insert(schema.accounts).values({ + siteId: site.id, + username: 'codex-existing@example.com', + accessToken: 'stable-access-token', + apiToken: null, + status: 'active', + oauthProvider: 'codex', + oauthAccountKey: 'chatgpt-account-existing', + extraConfig: JSON.stringify({ + credentialMode: 'session', + oauth: { + provider: 'codex', + accountId: 'chatgpt-account-existing', + accountKey: 'chatgpt-account-existing', + email: 'codex-existing@example.com', + planType: 'team', + refreshToken: 'stable-refresh-token', + idToken: originalJwt, + }, + }), + }).returning().get(); + + const reboundJwt = buildJwt({ + email: 'codex-existing@example.com', + 'https://api.openai.com/auth': { + chatgpt_account_id: 'chatgpt-account-rebound', + chatgpt_plan_type: 'plus', + }, + }); + fetchMock + .mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => ({ + access_token: 'rebound-access-token', + refresh_token: 'rebound-refresh-token', + id_token: reboundJwt, + expires_in: 3600, + token_type: 'Bearer', + }), + text: async () => JSON.stringify({ ok: true }), + }) + .mockResolvedValueOnce({ + ok: false, + status: 403, + json: async () => ({ error: 'forbidden' }), + text: async () => 'forbidden', + }); + + const startResponse = await app.inject({ + method: 'POST', + url: '/api/oauth/providers/codex/start', + headers: { + host: 'metapi.example', + 'x-forwarded-proto': 'https', + }, + payload: { + accountId: existing.id, + }, + }); + const startBody = startResponse.json() as { state: string }; + + const callbackResponse = await app.inject({ + method: 'POST', + url: `/api/oauth/sessions/${encodeURIComponent(startBody.state)}/manual-callback`, + payload: { + callbackUrl: `http://localhost:1455/auth/callback?state=${encodeURIComponent(startBody.state)}&code=oauth-code-rebind-fail`, + }, + }); + expect(callbackResponse.statusCode).toBe(500); + + const stored = await db.select().from(schema.accounts).where(eq(schema.accounts.id, existing.id)).get(); + expect(stored).toMatchObject({ + id: existing.id, + username: 'codex-existing@example.com', + accessToken: 'stable-access-token', + oauthAccountKey: 'chatgpt-account-existing', + }); + expect(JSON.parse(stored?.extraConfig || '{}')).toMatchObject({ + oauth: { + accountId: 'chatgpt-account-existing', + refreshToken: 'stable-refresh-token', + idToken: originalJwt, + }, + }); + }); + + it('fails codex oauth onboarding when token exchange does not expose chatgpt_account_id', async () => { + const jwt = buildJwt({ + email: 'codex-no-account@example.com', + 'https://api.openai.com/auth': { + chatgpt_plan_type: 'plus', + }, + }); + fetchMock.mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => ({ + access_token: 'oauth-access-token', + refresh_token: 'oauth-refresh-token', + id_token: jwt, + expires_in: 3600, + token_type: 'Bearer', + }), + text: async () => JSON.stringify({ ok: true }), + }); + + const startResponse = await app.inject({ + method: 'POST', + url: '/api/oauth/providers/codex/start', + headers: { + host: 'metapi.example', + 'x-forwarded-proto': 'https', + }, + }); + const startBody = startResponse.json() as { state: string }; + + const callbackResponse = await app.inject({ + method: 'POST', + url: `/api/oauth/sessions/${encodeURIComponent(startBody.state)}/manual-callback`, + payload: { + callbackUrl: `http://localhost:1455/auth/callback?state=${encodeURIComponent(startBody.state)}&code=oauth-code-no-account-id`, + }, + }); + expect(callbackResponse.statusCode).toBe(500); + expect(callbackResponse.json()).toMatchObject({ + message: expect.stringContaining('chatgpt_account_id'), + }); + + const accounts = await db.select().from(schema.accounts).all(); + expect(accounts).toEqual([]); + }); + it('marks gemini oauth session as error when token exchange fails before account persistence', async () => { fetchMock.mockResolvedValueOnce({ ok: false, diff --git a/src/server/routes/proxy/chat.codex-oauth.test.ts b/src/server/routes/proxy/chat.codex-oauth.test.ts index f561f87..ac80e7e 100644 --- a/src/server/routes/proxy/chat.codex-oauth.test.ts +++ b/src/server/routes/proxy/chat.codex-oauth.test.ts @@ -361,4 +361,47 @@ describe('chat proxy codex oauth compatibility', () => { expect(secondOptions.headers.Authorization).toBe('Bearer fresh-access-token'); expect(response.json()?.choices?.[0]?.message?.content).toBe('ok after refresh'); }); + + it('retries oauth chat requests after a 403 auth failure', async () => { + fetchMock + .mockResolvedValueOnce(new Response(JSON.stringify({ + error: { message: 'forbidden account mismatch', type: 'invalid_request_error' }, + }), { + status: 403, + headers: { 'content-type': 'application/json' }, + })) + .mockResolvedValueOnce(new Response(JSON.stringify({ + id: 'chatcmpl_refreshed_403', + object: 'chat.completion', + created: 1706000000, + model: 'gpt-4o-mini', + choices: [ + { + index: 0, + message: { role: 'assistant', content: 'ok after forbidden refresh' }, + finish_reason: 'stop', + }, + ], + usage: { prompt_tokens: 4, completion_tokens: 2, total_tokens: 6 }, + }), { + status: 200, + headers: { 'content-type': 'application/json' }, + })); + + const response = await app.inject({ + method: 'POST', + url: '/v1/chat/completions', + payload: { + model: 'gpt-4o-mini', + messages: [{ role: 'user', content: 'hello oauth' }], + }, + }); + + expect(response.statusCode).toBe(200); + expect(refreshOauthAccessTokenSingleflightMock).toHaveBeenCalledWith(33); + expect(fetchMock).toHaveBeenCalledTimes(2); + const [, secondOptions] = fetchMock.mock.calls[1] as [string, any]; + expect(secondOptions.headers.Authorization).toBe('Bearer fresh-access-token'); + expect(response.json()?.choices?.[0]?.message?.content).toBe('ok after forbidden refresh'); + }); }); diff --git a/src/server/routes/proxy/responses.codex-oauth.test.ts b/src/server/routes/proxy/responses.codex-oauth.test.ts index 13187a6..463cd30 100644 --- a/src/server/routes/proxy/responses.codex-oauth.test.ts +++ b/src/server/routes/proxy/responses.codex-oauth.test.ts @@ -186,6 +186,7 @@ describe('responses proxy codex oauth refresh', () => { url: '/v1/responses', headers: { 'user-agent': 'CodexClient/1.0', + 'Chatgpt-Account-Id': 'spoofed-account', }, payload: { model: 'gpt-5.2-codex', @@ -214,6 +215,46 @@ describe('responses proxy codex oauth refresh', () => { expect(response.json()?.output_text).toContain('ok after codex token refresh'); }); + it('refreshes codex oauth token and retries the same responses request on 403', async () => { + fetchMock + .mockResolvedValueOnce(new Response(JSON.stringify({ + error: { message: 'forbidden account mismatch', type: 'invalid_request_error' }, + }), { + status: 403, + headers: { 'content-type': 'application/json' }, + })) + .mockResolvedValueOnce(new Response(JSON.stringify({ + id: 'resp_codex_refreshed_403', + object: 'response', + model: 'gpt-5.2-codex', + status: 'completed', + output_text: 'ok after codex forbidden refresh', + usage: { input_tokens: 4, output_tokens: 2, total_tokens: 6 }, + }), { + status: 200, + headers: { 'content-type': 'application/json' }, + })); + + const response = await app.inject({ + method: 'POST', + url: '/v1/responses', + headers: { + 'user-agent': 'CodexClient/1.0', + }, + payload: { + model: 'gpt-5.2-codex', + input: 'hello codex', + }, + }); + + expect(response.statusCode).toBe(200); + expect(refreshOauthAccessTokenSingleflightMock).toHaveBeenCalledWith(33); + expect(fetchMock).toHaveBeenCalledTimes(2); + const [, secondOptions] = fetchMock.mock.calls[1] as [string, any]; + expect(secondOptions.headers.Authorization).toBe('Bearer fresh-access-token'); + expect(response.json()?.output_text).toContain('ok after codex forbidden refresh'); + }); + it('retries oauth responses requests with a normalized upstream URL after refresh', async () => { selectChannelMock.mockReturnValue({ channel: { id: 11, routeId: 22 }, @@ -316,7 +357,7 @@ describe('responses proxy codex oauth refresh', () => { ]); }); - it('preserves explicit prompt_cache_key for codex responses requests', async () => { + it('preserves explicit prompt_cache_key for codex responses requests without converting it into codex session headers', async () => { fetchMock.mockResolvedValue(new Response(JSON.stringify({ id: 'resp_codex_with_cache_key', object: 'response', @@ -344,8 +385,8 @@ describe('responses proxy codex oauth refresh', () => { const [, options] = fetchMock.mock.calls[0] as [string, any]; const forwardedBody = JSON.parse(options.body); - expect(options.headers.Session_id || options.headers.session_id).toBe('codex-cache-123'); - expect(options.headers.Conversation_id || options.headers.conversation_id).toBe('codex-cache-123'); + expect(String(options.headers.Session_id || options.headers.session_id || '')).toMatch(/^[0-9a-f-]{36}$/i); + expect(options.headers.Conversation_id || options.headers.conversation_id).toBeUndefined(); expect(forwardedBody.prompt_cache_key).toBe('codex-cache-123'); }); diff --git a/src/server/routes/proxy/responses.websocket.test.ts b/src/server/routes/proxy/responses.websocket.test.ts index 62bc119..a11cd32 100644 --- a/src/server/routes/proxy/responses.websocket.test.ts +++ b/src/server/routes/proxy/responses.websocket.test.ts @@ -3,6 +3,7 @@ import { createServer, type Server } from 'node:http'; import { AddressInfo } from 'node:net'; import WebSocket, { WebSocketServer } from 'ws'; import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'; +import { config } from '../../config.js'; const fetchMock = vi.fn(); const selectChannelMock = vi.fn(); @@ -10,6 +11,8 @@ const selectNextChannelMock = vi.fn(); const previewSelectedChannelMock = vi.fn(); const recordSuccessMock = vi.fn(); const recordFailureMock = vi.fn(); +const authorizeDownstreamTokenMock = vi.fn(); +const consumeManagedKeyRequestMock = vi.fn(); const refreshModelsAndRebuildRoutesMock = vi.fn(); const reportProxyAllFailedMock = vi.fn(); const reportTokenExpiredMock = vi.fn(); @@ -47,6 +50,22 @@ vi.mock('../../services/alertService.js', () => ({ reportTokenExpired: (...args: unknown[]) => reportTokenExpiredMock(...args), })); +vi.mock('../../services/downstreamApiKeyService.js', () => ({ + authorizeDownstreamToken: (...args: unknown[]) => authorizeDownstreamTokenMock(...args), + consumeManagedKeyRequest: (...args: unknown[]) => consumeManagedKeyRequestMock(...args), + isModelAllowedByPolicyOrAllowedRoutes: async ( + model: string, + policy: { supportedModels?: string[]; allowedRouteIds?: number[]; denyAllWhenEmpty?: boolean }, + ) => { + const supportedModels = Array.isArray(policy?.supportedModels) ? policy.supportedModels : []; + const allowedRouteIds = Array.isArray(policy?.allowedRouteIds) ? policy.allowedRouteIds : []; + if (supportedModels.length === 0 && allowedRouteIds.length === 0) { + return policy?.denyAllWhenEmpty === true ? false : true; + } + return supportedModels.includes(model); + }, +})); + vi.mock('../../services/alertRules.js', () => ({ isTokenExpiredError: () => false, })); @@ -175,7 +194,17 @@ function waitForSocketMessages(socket: WebSocket, count: number, timeoutMs = 100 }); } +function createClientSocket(baseUrl: string, headers: Record = {}) { + return new WebSocket(`${baseUrl}/v1/responses`, { + headers: { + Authorization: 'Bearer sk-global-proxy-token', + ...headers, + }, + }); +} + describe('responses websocket transport', () => { + const originalCodexResponsesWebsocketBeta = config.codexResponsesWebsocketBeta; let app: FastifyInstance; let baseUrl: string; let upstreamServer: WebSocketServer; @@ -186,6 +215,9 @@ describe('responses websocket transport', () => { let upstreamMessageHandler: (socket: WebSocket, parsed: Record, requestIndex: number) => void; let rejectedUpgradeServer: Server; let rejectedUpgradeSiteUrl: string; + let rejectedUpgradeStatus: number; + let rejectedUpgradeStatusText: string; + let rejectedUpgradeBody: string; beforeAll(async () => { const { responsesProxyRoute } = await import('./responses.js'); @@ -214,13 +246,14 @@ describe('responses websocket transport', () => { rejectedUpgradeServer = createServer(); rejectedUpgradeServer.on('upgrade', (_request, socket) => { + const body = rejectedUpgradeBody; socket.write( - 'HTTP/1.1 426 Upgrade Required\r\n' + `HTTP/1.1 ${rejectedUpgradeStatus} ${rejectedUpgradeStatusText}\r\n` + 'Content-Type: text/plain\r\n' - + 'Content-Length: 16\r\n' + + `Content-Length: ${Buffer.byteLength(body)}\r\n` + 'Connection: close\r\n' + '\r\n' - + 'Upgrade Required', + + body, ); socket.destroy(); }); @@ -236,6 +269,8 @@ describe('responses websocket transport', () => { previewSelectedChannelMock.mockReset(); recordSuccessMock.mockReset(); recordFailureMock.mockReset(); + authorizeDownstreamTokenMock.mockReset(); + consumeManagedKeyRequestMock.mockReset(); refreshModelsAndRebuildRoutesMock.mockReset(); reportProxyAllFailedMock.mockReset(); reportTokenExpiredMock.mockReset(); @@ -249,6 +284,21 @@ describe('responses websocket transport', () => { upstreamConnectionCount = 0; upstreamUpgradeHeaders = {}; upstreamRequests = []; + (config as any).codexResponsesWebsocketBeta = originalCodexResponsesWebsocketBeta; + rejectedUpgradeStatus = 426; + rejectedUpgradeStatusText = 'Upgrade Required'; + rejectedUpgradeBody = 'Upgrade Required'; + authorizeDownstreamTokenMock.mockResolvedValue({ + ok: true, + source: 'global', + token: 'sk-global-proxy-token', + key: null, + policy: { + supportedModels: [], + allowedRouteIds: [], + siteWeightMultipliers: {}, + }, + }); upstreamMessageHandler = (socket, parsed, requestIndex) => { const responseId = `resp_upstream_${requestIndex}`; socket.send(JSON.stringify({ @@ -331,7 +381,7 @@ describe('responses websocket transport', () => { })); }; - const socket = new WebSocket(`${baseUrl}/v1/responses`); + const socket = createClientSocket(baseUrl); await waitForSocketOpen(socket); const messagesPromise = waitForSocketMessages(socket, 4); @@ -362,10 +412,8 @@ describe('responses websocket transport', () => { }); it('echoes x-codex-turn-state on websocket upgrade responses', async () => { - const socket = new WebSocket(`${baseUrl}/v1/responses`, { - headers: { - 'x-codex-turn-state': 'turn-state-123', - }, + const socket = createClientSocket(baseUrl, { + 'x-codex-turn-state': 'turn-state-123', }); const [upgrade] = await Promise.all([ @@ -378,17 +426,16 @@ describe('responses websocket transport', () => { }); it('reuses one upstream codex websocket session across sequential websocket turns', async () => { + (config as any).codexResponsesWebsocketBeta = 'responses_websockets=2099-01-01'; const selectedChannel = createSelectedChannel({ siteUrl: upstreamSiteUrl, }); selectChannelMock.mockReturnValue(selectedChannel); previewSelectedChannelMock.mockResolvedValue(selectedChannel); - const socket = new WebSocket(`${baseUrl}/v1/responses`, { - headers: { - 'x-codex-turn-state': 'turn-state-123', - 'x-codex-beta-features': 'feature-a,feature-b', - }, + const socket = createClientSocket(baseUrl, { + 'x-codex-turn-state': 'turn-state-123', + 'x-codex-beta-features': 'feature-a,feature-b', }); await waitForSocketOpen(socket); const firstMessagesPromise = waitForSocketMessages(socket, 1); @@ -427,7 +474,7 @@ describe('responses websocket transport', () => { }); expect(upstreamUpgradeHeaders['x-codex-turn-state']).toBe('turn-state-123'); expect(upstreamUpgradeHeaders['x-codex-beta-features']).toBe('feature-a,feature-b'); - expect(upstreamUpgradeHeaders['openai-beta']).toContain('responses_websockets='); + expect(upstreamUpgradeHeaders['openai-beta']).toContain('responses_websockets=2099-01-01'); }); it('falls back to the HTTP responses executor when the upstream codex websocket upgrade returns 426', async () => { @@ -442,7 +489,7 @@ describe('responses websocket transport', () => { 'data: [DONE]\n\n', ])); - const socket = new WebSocket(`${baseUrl}/v1/responses`); + const socket = createClientSocket(baseUrl); await waitForSocketOpen(socket); const messagesPromise = waitForSocketMessages(socket, 1); @@ -460,6 +507,80 @@ describe('responses websocket transport', () => { expect(messages[0]?.response?.id).toBe('resp_http_fallback'); }); + it('falls back to the HTTP responses executor when the upstream codex websocket upgrade returns 401', async () => { + rejectedUpgradeStatus = 401; + rejectedUpgradeStatusText = 'Unauthorized'; + rejectedUpgradeBody = JSON.stringify({ + error: { + message: 'expired token', + type: 'invalid_request_error', + }, + }); + const selectedChannel = createSelectedChannel({ + siteUrl: rejectedUpgradeSiteUrl, + }); + selectChannelMock.mockReturnValue(selectedChannel); + previewSelectedChannelMock.mockResolvedValue(selectedChannel); + fetchMock.mockResolvedValueOnce(createSseResponse([ + 'event: response.completed\n', + 'data: {"type":"response.completed","response":{"id":"resp_http_fallback_401","model":"gpt-5.4","status":"completed","output":[],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}\n\n', + 'data: [DONE]\n\n', + ])); + + const socket = createClientSocket(baseUrl); + await waitForSocketOpen(socket); + const messagesPromise = waitForSocketMessages(socket, 1); + + socket.send(JSON.stringify({ + type: 'response.create', + model: 'gpt-5.4', + input: [], + })); + + const messages = await messagesPromise; + socket.close(); + + expect(fetchMock).toHaveBeenCalledTimes(1); + expect(messages.map((message) => message.type)).toEqual(['response.completed']); + expect(messages[0]?.response?.id).toBe('resp_http_fallback_401'); + }); + + it('rejects websocket turns whose model is blocked by the downstream key policy before channel selection', async () => { + authorizeDownstreamTokenMock.mockResolvedValueOnce({ + ok: true, + source: 'managed', + token: 'sk-managed-denied', + key: { + id: 99, + name: 'limited-key', + }, + policy: { + supportedModels: ['gpt-4.1'], + allowedRouteIds: [], + siteWeightMultipliers: {}, + }, + }); + + const socket = createClientSocket(baseUrl); + await waitForSocketOpen(socket); + const messagesPromise = waitForSocketMessages(socket, 1); + + socket.send(JSON.stringify({ + type: 'response.create', + model: 'gpt-5.4', + input: [], + })); + + const messages = await messagesPromise; + socket.close(); + + expect(messages[0]).toMatchObject({ + type: 'error', + status: 403, + }); + expect(selectChannelMock).not.toHaveBeenCalled(); + }); + it('merges follow-up response.create payloads when the selected upstream does not support incremental mode', async () => { const selectedChannel = createSelectedChannel({ sitePlatform: 'openai', @@ -483,7 +604,7 @@ describe('responses websocket transport', () => { 'data: [DONE]\n\n', ])); - const socket = new WebSocket(`${baseUrl}/v1/responses`); + const socket = createClientSocket(baseUrl); await waitForSocketOpen(socket); const firstResponsePromise = waitForSocketMessages(socket, 1); @@ -609,7 +730,7 @@ describe('responses websocket transport', () => { })); }; - const socket = new WebSocket(`${baseUrl}/v1/responses`); + const socket = createClientSocket(baseUrl); await waitForSocketOpen(socket); const firstResponsePromise = waitForSocketMessages(socket, 1); @@ -694,7 +815,7 @@ describe('responses websocket transport', () => { 'data: [DONE]\n\n', ])); - const socket = new WebSocket(`${baseUrl}/v1/responses`); + const socket = createClientSocket(baseUrl); await waitForSocketOpen(socket); const firstResponsePromise = waitForSocketMessages(socket, 1); @@ -753,7 +874,7 @@ describe('responses websocket transport', () => { 'data: [DONE]\n\n', ])); - const socket = new WebSocket(`${baseUrl}/v1/responses`); + const socket = createClientSocket(baseUrl); await waitForSocketOpen(socket); const prewarmMessagesPromise = waitForSocketMessages(socket, 2); @@ -805,7 +926,7 @@ describe('responses websocket transport', () => { selectChannelMock.mockReturnValue(selectedChannel); previewSelectedChannelMock.mockResolvedValue(selectedChannel); - const socket = new WebSocket(`${baseUrl}/v1/responses`); + const socket = createClientSocket(baseUrl); await waitForSocketOpen(socket); const messagesPromise = waitForSocketMessages(socket, 1); @@ -852,7 +973,7 @@ describe('responses websocket transport', () => { socket.close(); }; - const socket = new WebSocket(`${baseUrl}/v1/responses`); + const socket = createClientSocket(baseUrl); await waitForSocketOpen(socket); const messagesPromise = waitForSocketMessages(socket, 3, 400); diff --git a/src/server/routes/proxy/responsesWebsocket.ts b/src/server/routes/proxy/responsesWebsocket.ts index abfa389..a7c2c0f 100644 --- a/src/server/routes/proxy/responsesWebsocket.ts +++ b/src/server/routes/proxy/responsesWebsocket.ts @@ -1,8 +1,15 @@ import { randomUUID } from 'node:crypto'; import type { FastifyInstance } from 'fastify'; import type { IncomingMessage } from 'node:http'; +import type { Duplex } from 'node:stream'; import { WebSocketServer, type RawData, type WebSocket } from 'ws'; import { createCodexWebsocketRuntime, CodexWebsocketRuntimeError } from '../../proxy-core/runtime/codexWebsocketRuntime.js'; +import { + authorizeDownstreamToken, + consumeManagedKeyRequest, + isModelAllowedByPolicyOrAllowedRoutes, + type DownstreamTokenAuthSuccess, +} from '../../services/downstreamApiKeyService.js'; import { tokenRouter } from '../../services/tokenRouter.js'; import { buildOauthProviderHeaders } from '../../services/oauth/service.js'; import { openAiResponsesTransformer } from '../../transformers/openai/responses/index.js'; @@ -15,6 +22,7 @@ const RESPONSES_WEBSOCKET_TRANSPORT_HEADER = 'x-metapi-responses-websocket-trans const codexWebsocketRuntime = createCodexWebsocketRuntime(); type SelectedChannel = NonNullable>>; +type ResponsesWebsocketAuthContext = DownstreamTokenAuthSuccess; type NormalizedResponsesWebsocketRequest = | { @@ -131,8 +139,8 @@ function shouldReuseSelectedChannel( } function deriveCodexExplicitSessionId(body: Record, sessionId: string): string { - const promptCacheKey = asTrimmedString(body.prompt_cache_key); - return promptCacheKey || sessionId; + void body; + return sessionId; } function parseJsonObject(raw: RawData): Record | null { @@ -417,15 +425,46 @@ function buildInjectHeaders(request: IncomingMessage): Record, lastRequest: Record | null, + authContext: ResponsesWebsocketAuthContext, ): Promise { const requestModel = asTrimmedString(parsed.model) || asTrimmedString(lastRequest?.model); if (!requestModel) return false; try { - const selected = await tokenRouter.previewSelectedChannel(requestModel); + const selected = await tokenRouter.previewSelectedChannel(requestModel, authContext.policy); return selectedChannelSupportsIncrementalInput(selected, requestModel); } catch { return false; @@ -436,6 +475,7 @@ async function handleResponsesWebsocketConnection( app: FastifyInstance, socket: WebSocket, request: IncomingMessage, + authContext: ResponsesWebsocketAuthContext, ) { const websocketSessionId = headerValueToTrimmedString(request.headers['session_id']) || headerValueToTrimmedString(request.headers['session-id']) @@ -449,129 +489,141 @@ async function handleResponsesWebsocketConnection( }); socket.on('message', async (raw) => { - const parsed = parseJsonObject(raw); - if (!parsed) { - writeResponsesWebsocketError(socket, 400, 'Invalid websocket JSON payload'); - return; - } - - const requestModel = asTrimmedString(parsed.model) || asTrimmedString(lastRequest?.model); - const supportsIncrementalInput = selectedChannelSupportsIncrementalInput(selectedChannel, requestModel) - || await supportsResponsesWebsocketIncrementalInput(parsed, lastRequest); - const shouldHandleLocalPrewarm = shouldHandleResponsesWebsocketPrewarmLocally( - parsed, - lastRequest, - supportsIncrementalInput, - ); - const normalized = normalizeResponsesWebsocketRequest( - parsed, - lastRequest, - lastResponseOutput, - supportsIncrementalInput, - ); - if (!normalized.ok) { - writeResponsesWebsocketError(socket, normalized.status, normalized.message); - return; - } - - lastRequest = normalized.nextRequestSnapshot; - if (shouldHandleLocalPrewarm) { - lastResponseOutput = []; - for (const payload of synthesizePrewarmResponsePayloads(normalized.request)) { - socket.send(JSON.stringify(payload)); + try { + const parsed = parseJsonObject(raw); + if (!parsed) { + writeResponsesWebsocketError(socket, 400, 'Invalid websocket JSON payload'); + return; } - return; - } - if (!shouldReuseSelectedChannel(selectedChannel, requestModel)) { - selectedChannel = requestModel - ? await tokenRouter.selectChannel(requestModel) + const requestModel = asTrimmedString(parsed.model) || asTrimmedString(lastRequest?.model); + if (requestModel && !await isModelAllowedByPolicyOrAllowedRoutes(requestModel, authContext.policy)) { + writeResponsesWebsocketError(socket, 403, 'model is not allowed for this downstream key'); + return; + } + const supportsIncrementalInput = selectedChannelSupportsIncrementalInput(selectedChannel, requestModel) + || await supportsResponsesWebsocketIncrementalInput(parsed, lastRequest, authContext); + const shouldHandleLocalPrewarm = shouldHandleResponsesWebsocketPrewarmLocally( + parsed, + lastRequest, + supportsIncrementalInput, + ); + const normalized = normalizeResponsesWebsocketRequest( + parsed, + lastRequest, + lastResponseOutput, + supportsIncrementalInput, + ); + if (!normalized.ok) { + writeResponsesWebsocketError(socket, normalized.status, normalized.message); + return; + } + + if (authContext.source === 'managed' && authContext.key?.id) { + await consumeManagedKeyRequest(authContext.key.id); + } + + lastRequest = normalized.nextRequestSnapshot; + if (shouldHandleLocalPrewarm) { + lastResponseOutput = []; + for (const payload of synthesizePrewarmResponsePayloads(normalized.request)) { + socket.send(JSON.stringify(payload)); + } + return; + } + + if (!shouldReuseSelectedChannel(selectedChannel, requestModel)) { + selectedChannel = requestModel + ? await tokenRouter.selectChannel(requestModel, authContext.policy) + : null; + } + + const codexWebsocketChannel = selectedChannelSupportsCodexWebsocketTransport(selectedChannel, requestModel) + ? selectedChannel : null; - } - const codexWebsocketChannel = selectedChannelSupportsCodexWebsocketTransport(selectedChannel, requestModel) - ? selectedChannel - : null; - - if (codexWebsocketChannel) { - const downstreamHeaders: Record = { - ...(request.headers as Record), - [RESPONSES_WEBSOCKET_TRANSPORT_HEADER]: '1', - ...(supportsIncrementalInput ? { [RESPONSES_WEBSOCKET_MODE_HEADER]: 'incremental' } : {}), - }; - const providerHeaders = buildOauthProviderHeaders({ - extraConfig: typeof codexWebsocketChannel.account.extraConfig === 'string' - ? codexWebsocketChannel.account.extraConfig - : null, - downstreamHeaders, - }); - const prepared = buildUpstreamEndpointRequest({ - endpoint: 'responses', - modelName: asTrimmedString(codexWebsocketChannel.actualModel) || requestModel, - stream: true, - tokenValue: codexWebsocketChannel.tokenValue, - sitePlatform: codexWebsocketChannel.site.platform, - siteUrl: codexWebsocketChannel.site.url, - openaiBody: normalized.request, - downstreamFormat: 'responses', - responsesOriginalBody: normalized.request, - downstreamHeaders, - providerHeaders, - codexExplicitSessionId: deriveCodexExplicitSessionId(normalized.request, websocketSessionId), - }); - const requestUrl = `${codexWebsocketChannel.site.url.replace(/\/+$/, '')}${prepared.path}`; - - try { - const runtimeResult = await codexWebsocketRuntime.sendRequest({ - sessionId: websocketSessionId, - requestUrl, - headers: prepared.headers, - body: prepared.body, + if (codexWebsocketChannel) { + const downstreamHeaders: Record = { + ...(request.headers as Record), + [RESPONSES_WEBSOCKET_TRANSPORT_HEADER]: '1', + ...(supportsIncrementalInput ? { [RESPONSES_WEBSOCKET_MODE_HEADER]: 'incremental' } : {}), + }; + const providerHeaders = buildOauthProviderHeaders({ + extraConfig: typeof codexWebsocketChannel.account.extraConfig === 'string' + ? codexWebsocketChannel.account.extraConfig + : null, + downstreamHeaders, }); - lastResponseOutput = collectResponsesOutput(runtimeResult.events); - for (const payload of runtimeResult.events) { - socket.send(JSON.stringify(payload)); - } - } catch (error) { - const runtimeError = error instanceof CodexWebsocketRuntimeError - ? error - : new CodexWebsocketRuntimeError('upstream websocket request failed'); - if (runtimeError.status === 426) { - const fallbackOutput = await forwardResponsesRequestViaHttp({ - app, - socket, - request, - payload: normalized.request, - preserveIncrementalMode: false, - }); - if (fallbackOutput) { - lastResponseOutput = fallbackOutput; - } - return; - } - lastResponseOutput = collectResponsesOutput(runtimeError.events); - for (const payload of runtimeError.events) { - socket.send(JSON.stringify(payload)); - } - writeResponsesWebsocketError( - socket, - runtimeError.status || 408, - runtimeError.message, - runtimeError.payload, - ); - } - return; - } + const prepared = buildUpstreamEndpointRequest({ + endpoint: 'responses', + modelName: asTrimmedString(codexWebsocketChannel.actualModel) || requestModel, + stream: true, + tokenValue: codexWebsocketChannel.tokenValue, + sitePlatform: codexWebsocketChannel.site.platform, + siteUrl: codexWebsocketChannel.site.url, + openaiBody: normalized.request, + downstreamFormat: 'responses', + responsesOriginalBody: normalized.request, + downstreamHeaders, + providerHeaders, + codexExplicitSessionId: deriveCodexExplicitSessionId(normalized.request, websocketSessionId), + }); + const requestUrl = `${codexWebsocketChannel.site.url.replace(/\/+$/, '')}${prepared.path}`; - const forwardedOutput = await forwardResponsesRequestViaHttp({ - app, - socket, - request, - payload: normalized.request, - preserveIncrementalMode: supportsIncrementalInput, - }); - if (forwardedOutput) { - lastResponseOutput = forwardedOutput; + try { + const runtimeResult = await codexWebsocketRuntime.sendRequest({ + sessionId: websocketSessionId, + requestUrl, + headers: prepared.headers, + body: prepared.body, + }); + lastResponseOutput = collectResponsesOutput(runtimeResult.events); + for (const payload of runtimeResult.events) { + socket.send(JSON.stringify(payload)); + } + } catch (error) { + const runtimeError = error instanceof CodexWebsocketRuntimeError + ? error + : new CodexWebsocketRuntimeError('upstream websocket request failed'); + if (runtimeError.status && runtimeError.events.length === 0) { + const fallbackOutput = await forwardResponsesRequestViaHttp({ + app, + socket, + request, + payload: normalized.request, + preserveIncrementalMode: false, + }); + if (fallbackOutput) { + lastResponseOutput = fallbackOutput; + } + return; + } + lastResponseOutput = collectResponsesOutput(runtimeError.events); + for (const payload of runtimeError.events) { + socket.send(JSON.stringify(payload)); + } + writeResponsesWebsocketError( + socket, + runtimeError.status || 408, + runtimeError.message, + runtimeError.payload, + ); + } + return; + } + + const forwardedOutput = await forwardResponsesRequestViaHttp({ + app, + socket, + request, + payload: normalized.request, + preserveIncrementalMode: supportsIncrementalInput, + }); + if (forwardedOutput) { + lastResponseOutput = forwardedOutput; + } + } catch { + writeResponsesWebsocketError(socket, 500, 'internal websocket proxy error'); } }); } @@ -588,10 +640,24 @@ export function ensureResponsesWebsocketTransport(app: FastifyInstance) { }); app.server.on('upgrade', (request, socket, head) => { - const url = new URL(request.url || '/', 'http://localhost'); - if (url.pathname !== '/v1/responses') return; - websocketServer.handleUpgrade(request, socket, head, (client) => { - void handleResponsesWebsocketConnection(app, client, request); + void (async () => { + const url = new URL(request.url || '/', 'http://localhost'); + if (url.pathname !== '/v1/responses') return; + const token = extractWebsocketAuthToken(request, url); + if (!token) { + writeUpgradeHttpError(socket, 401, 'Missing Authorization, x-api-key, x-goog-api-key, or key query parameter'); + return; + } + const authResult = await authorizeDownstreamToken(token); + if (!authResult.ok) { + writeUpgradeHttpError(socket, authResult.statusCode, authResult.error); + return; + } + websocketServer.handleUpgrade(request, socket, head, (client) => { + void handleResponsesWebsocketConnection(app, client, request, authResult); + }); + })().catch(() => { + writeUpgradeHttpError(socket, 500, 'internal websocket proxy error'); }); }); diff --git a/src/server/routes/proxy/upstreamEndpoint.test.ts b/src/server/routes/proxy/upstreamEndpoint.test.ts index 7ca1a04..b3f47a9 100644 --- a/src/server/routes/proxy/upstreamEndpoint.test.ts +++ b/src/server/routes/proxy/upstreamEndpoint.test.ts @@ -581,7 +581,7 @@ describe('buildUpstreamEndpointRequest', () => { expect(request.body.prompt_cache_key).toBeUndefined(); }); - it('preserves explicit prompt_cache_key for native codex responses requests', () => { + it('preserves explicit prompt_cache_key for native codex responses requests without mapping it into codex session headers', () => { const request = buildUpstreamEndpointRequest({ endpoint: 'responses', modelName: 'gpt-5.4', @@ -599,11 +599,10 @@ describe('buildUpstreamEndpointRequest', () => { providerHeaders: { Originator: 'codex_cli_rs', }, - codexExplicitSessionId: 'codex-cache-123', } as any); - expect(request.headers.Session_id).toBe('codex-cache-123'); - expect(request.headers.Conversation_id).toBe('codex-cache-123'); + expect(request.headers.Session_id).toMatch(/^[0-9a-f-]{36}$/i); + expect(request.headers.Conversation_id).toBeUndefined(); expect(request.body.prompt_cache_key).toBe('codex-cache-123'); }); diff --git a/src/server/services/oauth/codexProvider.ts b/src/server/services/oauth/codexProvider.ts index 8895abd..b57a323 100644 --- a/src/server/services/oauth/codexProvider.ts +++ b/src/server/services/oauth/codexProvider.ts @@ -91,11 +91,15 @@ function parseTokenResponsePayload(payload: unknown): CodexTokenExchangeResult { throw new Error('codex token exchange response missing required fields'); } const claims = parseJwtClaims(idToken); + const accountId = asTrimmedString(claims?.['https://api.openai.com/auth']?.chatgpt_account_id); + if (!accountId) { + throw new Error('codex token exchange response missing chatgpt_account_id'); + } return { accessToken, refreshToken, idToken, - accountId: asTrimmedString(claims?.['https://api.openai.com/auth']?.chatgpt_account_id), + accountId, email: asTrimmedString(claims?.email), planType: asTrimmedString(claims?.['https://api.openai.com/auth']?.chatgpt_plan_type), tokenExpiresAt: Date.now() + expiresIn * 1000, @@ -220,7 +224,7 @@ export const codexOauthProvider: OAuthProviderDefinition = { }; }, buildProxyHeaders: ({ oauth, downstreamHeaders }) => { - const accountId = getHeaderValue(downstreamHeaders, 'chatgpt-account-id') || oauth.accountId || oauth.accountKey; + const accountId = oauth.accountId || oauth.accountKey; const originator = getHeaderValue(downstreamHeaders, 'originator') || 'codex_cli_rs'; const headers: Record = { Originator: originator, diff --git a/src/server/services/oauth/service.ts b/src/server/services/oauth/service.ts index 9a73d70..d6367ea 100644 --- a/src/server/services/oauth/service.ts +++ b/src/server/services/oauth/service.ts @@ -178,7 +178,7 @@ async function findExistingOauthAccount(input: { if (byKey) return byKey; } - if (!accountKey && email) { + if (!accountKey && email && input.provider !== 'codex') { const byEmail = await db.select().from(schema.accounts).where(and( eq(schema.accounts.oauthProvider, input.provider), eq(schema.accounts.username, email), @@ -253,6 +253,7 @@ async function upsertOauthAccount(input: { account: await db.select().from(schema.accounts).where(eq(schema.accounts.id, existing.id)).get(), site, created: false, + previousAccount: existing, }; } @@ -270,7 +271,7 @@ async function upsertOauthAccount(input: { isPinned: false, sortOrder: await getNextAccountSortOrder(), }).returning().get(); - return { account: created, site, created: true }; + return { account: created, site, created: true, previousAccount: null }; } export function listOauthProviders() { @@ -365,7 +366,7 @@ export async function handleOauthCallback(input: { codeVerifier: session.codeVerifier, projectId: session.projectId, }); - const { account, site, created } = await upsertOauthAccount({ + const { account, site, created, previousAccount } = await upsertOauthAccount({ definition, exchange, rebindAccountId: session.rebindAccountId, @@ -379,6 +380,20 @@ export async function handleOauthCallback(input: { if (refreshResult.status !== 'success') { if (created) { await db.delete(schema.accounts).where(eq(schema.accounts.id, account.id)).run(); + } else if (previousAccount) { + await db.update(schema.accounts).set({ + siteId: previousAccount.siteId, + username: previousAccount.username, + accessToken: previousAccount.accessToken, + apiToken: previousAccount.apiToken, + checkinEnabled: previousAccount.checkinEnabled, + status: previousAccount.status, + oauthProvider: previousAccount.oauthProvider, + oauthAccountKey: previousAccount.oauthAccountKey, + oauthProjectId: previousAccount.oauthProjectId, + extraConfig: previousAccount.extraConfig, + updatedAt: previousAccount.updatedAt, + }).where(eq(schema.accounts.id, previousAccount.id)).run(); } await rebuildTokenRoutesFromAvailability(); const errorMessage = refreshResult.errorMessage || `${input.provider} model discovery failed`; diff --git a/src/server/services/tokenRouter.siteStatus.test.ts b/src/server/services/tokenRouter.siteStatus.test.ts index 0784b3b..e785252 100644 --- a/src/server/services/tokenRouter.siteStatus.test.ts +++ b/src/server/services/tokenRouter.siteStatus.test.ts @@ -164,4 +164,49 @@ describe('TokenRouter site status guard', () => { expect(selected).not.toBeNull(); expect(selected?.tokenValue).toBe('oauth-access-token'); }); + + it('prefers codex oauth access token over a stale legacy api token on oauth-managed rows', async () => { + const site = await db.insert(schema.sites).values({ + name: 'codex-site', + url: 'https://chatgpt.com/backend-api/codex', + platform: 'codex', + status: 'active', + }).returning().get(); + + const account = await db.insert(schema.accounts).values({ + siteId: site.id, + username: 'codex-user@example.com', + accessToken: 'oauth-access-token', + apiToken: 'stale-direct-token', + status: 'active', + extraConfig: JSON.stringify({ + credentialMode: 'session', + oauth: { + provider: 'codex', + accountId: 'chatgpt-account-123', + email: 'codex-user@example.com', + planType: 'plus', + }, + }), + }).returning().get(); + + const route = await db.insert(schema.tokenRoutes).values({ + modelPattern: 'gpt-5.2-codex', + enabled: true, + }).returning().get(); + + await db.insert(schema.routeChannels).values({ + routeId: route.id, + accountId: account.id, + tokenId: null, + priority: 0, + weight: 10, + enabled: true, + manualOverride: false, + }).run(); + + const selected = await new TokenRouter().selectChannel('gpt-5.2-codex'); + expect(selected).not.toBeNull(); + expect(selected?.tokenValue).toBe('oauth-access-token'); + }); }); diff --git a/src/server/services/tokenRouter.ts b/src/server/services/tokenRouter.ts index b0cf1e4..5c4e389 100644 --- a/src/server/services/tokenRouter.ts +++ b/src/server/services/tokenRouter.ts @@ -2037,14 +2037,15 @@ export class TokenRouter { return token ? token : null; } - const fallback = candidate.account.apiToken?.trim(); - if (fallback) return fallback; - if (getOauthInfoFromExtraConfig(candidate.account.extraConfig)) { const accessToken = candidate.account.accessToken?.trim(); - return accessToken || null; + if (accessToken) return accessToken; + return null; } + const fallback = candidate.account.apiToken?.trim(); + if (fallback) return fallback; + return null; }