diff --git a/src/server/proxy-core/surfaces/sharedSurface.test.ts b/src/server/proxy-core/surfaces/sharedSurface.test.ts index 018ceb7..e5be236 100644 --- a/src/server/proxy-core/surfaces/sharedSurface.test.ts +++ b/src/server/proxy-core/surfaces/sharedSurface.test.ts @@ -384,6 +384,8 @@ describe('selectSurfaceChannelForAttempt', () => { modelActual: 'upstream-model', status: 'failed', httpStatus: 502, + isStream: null, + firstByteLatencyMs: null, latencyMs: 1200, promptTokens: 10, completionTokens: 5, @@ -995,6 +997,8 @@ describe('selectSurfaceChannelForAttempt', () => { modelRequested: 'gpt-5.2', status: 'success', httpStatus: 200, + isStream: null, + firstByteLatencyMs: null, latencyMs: 250, errorMessage: null, retryCount: 1, diff --git a/src/server/routes/api/checkin.lock.test.ts b/src/server/routes/api/checkin.lock.test.ts index ca22ffb..13ac91f 100644 --- a/src/server/routes/api/checkin.lock.test.ts +++ b/src/server/routes/api/checkin.lock.test.ts @@ -35,6 +35,7 @@ vi.mock('../../db/index.js', () => { insert: () => insertChain, select: () => queryChain, }, + hasProxyLogStreamTimingColumns: async () => false, schema: { settings: { key: 'key' }, checkinLogs: { accountId: 'accountId', createdAt: 'createdAt' }, diff --git a/src/server/routes/proxy/chat.codex-oauth.test.ts b/src/server/routes/proxy/chat.codex-oauth.test.ts index cd47ddb..d3690b4 100644 --- a/src/server/routes/proxy/chat.codex-oauth.test.ts +++ b/src/server/routes/proxy/chat.codex-oauth.test.ts @@ -21,9 +21,13 @@ const dbInsertMock = vi.fn((_arg?: any) => ({ }), })); -vi.mock('undici', () => ({ - fetch: (...args: unknown[]) => fetchMock(...args), -})); +vi.mock('undici', async () => { + const actual = await vi.importActual('undici'); + return { + ...actual, + fetch: (...args: unknown[]) => fetchMock(...args), + }; +}); vi.mock('../../services/tokenRouter.js', () => ({ tokenRouter: { @@ -90,6 +94,7 @@ vi.mock('../../db/index.js', () => ({ hasProxyLogBillingDetailsColumn: async () => false, hasProxyLogClientColumns: async () => false, hasProxyLogDownstreamApiKeyIdColumn: async () => false, + hasProxyLogStreamTimingColumns: async () => false, schema: { proxyLogs: {}, siteApiEndpoints: { diff --git a/src/server/routes/proxy/chat.count-tokens.test.ts b/src/server/routes/proxy/chat.count-tokens.test.ts index 86d9403..181f810 100644 --- a/src/server/routes/proxy/chat.count-tokens.test.ts +++ b/src/server/routes/proxy/chat.count-tokens.test.ts @@ -25,9 +25,13 @@ const dbInsertMock = vi.fn((_arg?: any) => ({ }), })); -vi.mock('undici', () => ({ - fetch: (...args: unknown[]) => fetchMock(...args), -})); +vi.mock('undici', async () => { + const actual = await vi.importActual('undici'); + return { + ...actual, + fetch: (...args: unknown[]) => fetchMock(...args), + }; +}); vi.mock('../../services/tokenRouter.js', () => ({ tokenRouter: { @@ -107,6 +111,7 @@ vi.mock('../../db/index.js', () => ({ hasProxyLogBillingDetailsColumn: async () => false, hasProxyLogClientColumns: async () => false, hasProxyLogDownstreamApiKeyIdColumn: async () => false, + hasProxyLogStreamTimingColumns: async () => false, schema: { proxyLogs: {}, siteApiEndpoints: { diff --git a/src/server/routes/proxy/chat.siteApiEndpoint.test.ts b/src/server/routes/proxy/chat.siteApiEndpoint.test.ts index 420882f..c2eafa8 100644 --- a/src/server/routes/proxy/chat.siteApiEndpoint.test.ts +++ b/src/server/routes/proxy/chat.siteApiEndpoint.test.ts @@ -25,9 +25,13 @@ const resolveProxyUsageWithSelfLogFallbackMock = vi.fn(async ({ usage }: any) => })); const insertProxyLogMock = vi.fn(); -vi.mock('undici', () => ({ - fetch: (...args: unknown[]) => fetchMock(...args), -})); +vi.mock('undici', async () => { + const actual = await vi.importActual('undici'); + return { + ...actual, + fetch: (...args: unknown[]) => fetchMock(...args), + }; +}); vi.mock('../../services/tokenRouter.js', () => ({ tokenRouter: { diff --git a/src/server/routes/proxy/chat.stream.test.ts b/src/server/routes/proxy/chat.stream.test.ts index 9be8ba5..c8aad7b 100644 --- a/src/server/routes/proxy/chat.stream.test.ts +++ b/src/server/routes/proxy/chat.stream.test.ts @@ -26,9 +26,13 @@ const dbInsertMock = vi.fn((_arg?: any) => ({ }), })); -vi.mock('undici', () => ({ - fetch: (...args: unknown[]) => fetchMock(...args), -})); +vi.mock('undici', async () => { + const actual = await vi.importActual('undici'); + return { + ...actual, + fetch: (...args: unknown[]) => fetchMock(...args), + }; +}); vi.mock('../../services/tokenRouter.js', () => ({ tokenRouter: { @@ -91,6 +95,7 @@ vi.mock('../../db/index.js', () => ({ hasProxyLogBillingDetailsColumn: async () => false, hasProxyLogClientColumns: async () => false, hasProxyLogDownstreamApiKeyIdColumn: async () => false, + hasProxyLogStreamTimingColumns: async () => false, schema: { proxyLogs: {}, siteApiEndpoints: { diff --git a/src/server/routes/proxy/completions.siteApiEndpoint.test.ts b/src/server/routes/proxy/completions.siteApiEndpoint.test.ts index f919a21..d08961b 100644 --- a/src/server/routes/proxy/completions.siteApiEndpoint.test.ts +++ b/src/server/routes/proxy/completions.siteApiEndpoint.test.ts @@ -17,9 +17,13 @@ const insertProxyLogMock = vi.fn(); const resolveProxyUsageWithSelfLogFallbackMock = vi.fn(); const resolveProxyLogBillingMock = vi.fn(); -vi.mock('undici', () => ({ - fetch: (...args: unknown[]) => fetchMock(...args), -})); +vi.mock('undici', async () => { + const actual = await vi.importActual('undici'); + return { + ...actual, + fetch: (...args: unknown[]) => fetchMock(...args), + }; +}); vi.mock('../../services/tokenRouter.js', () => ({ tokenRouter: { diff --git a/src/server/routes/proxy/downstreamClientContext.routes.test.ts b/src/server/routes/proxy/downstreamClientContext.routes.test.ts index fb44848..f661c57 100644 --- a/src/server/routes/proxy/downstreamClientContext.routes.test.ts +++ b/src/server/routes/proxy/downstreamClientContext.routes.test.ts @@ -24,9 +24,13 @@ const dbInsertMock = vi.fn((_arg?: any) => ({ values: (arg: any) => dbValuesMock(arg), })); -vi.mock('undici', () => ({ - fetch: (...args: unknown[]) => fetchMock(...args), -})); +vi.mock('undici', async () => { + const actual = await vi.importActual('undici'); + return { + ...actual, + fetch: (...args: unknown[]) => fetchMock(...args), + }; +}); vi.mock('../../services/tokenRouter.js', () => ({ tokenRouter: { @@ -97,6 +101,7 @@ vi.mock('../../db/index.js', () => ({ hasProxyLogBillingDetailsColumn: async () => false, hasProxyLogClientColumns: async () => false, hasProxyLogDownstreamApiKeyIdColumn: async () => false, + hasProxyLogStreamTimingColumns: async () => false, })); describe('downstream client context route logging', () => { diff --git a/src/server/routes/proxy/gemini.test.ts b/src/server/routes/proxy/gemini.test.ts index 1c9b1ee..ca9696f 100644 --- a/src/server/routes/proxy/gemini.test.ts +++ b/src/server/routes/proxy/gemini.test.ts @@ -48,10 +48,13 @@ function createDbSelectChain() { }; } -vi.mock('undici', () => ({ - fetch: (...args: unknown[]) => fetchMock(...args), - Response: globalThis.Response, -})); +vi.mock('undici', async () => { + const actual = await vi.importActual('undici'); + return { + ...actual, + fetch: (...args: unknown[]) => fetchMock(...args), + }; +}); vi.mock('../../services/modelPricingService.js', () => ({ fetchModelPricingCatalog: (...args: unknown[]) => fetchModelPricingCatalogMock(...args), @@ -91,6 +94,7 @@ vi.mock('../../db/index.js', () => ({ hasProxyLogBillingDetailsColumn: async () => false, hasProxyLogClientColumns: async () => false, hasProxyLogDownstreamApiKeyIdColumn: async () => false, + hasProxyLogStreamTimingColumns: async () => false, schema: { proxyLogs: {}, modelAvailability: { diff --git a/src/server/routes/proxy/images.edits.test.ts b/src/server/routes/proxy/images.edits.test.ts index 061b5e3..a494cea 100644 --- a/src/server/routes/proxy/images.edits.test.ts +++ b/src/server/routes/proxy/images.edits.test.ts @@ -79,6 +79,7 @@ vi.mock('../../db/index.js', () => ({ hasProxyLogBillingDetailsColumn: async () => false, hasProxyLogClientColumns: async () => false, hasProxyLogDownstreamApiKeyIdColumn: async () => false, + hasProxyLogStreamTimingColumns: async () => false, schema: { proxyLogs: {}, siteApiEndpoints: { diff --git a/src/server/routes/proxy/responses.codex-oauth.test.ts b/src/server/routes/proxy/responses.codex-oauth.test.ts index f7c0fd2..32e462e 100644 --- a/src/server/routes/proxy/responses.codex-oauth.test.ts +++ b/src/server/routes/proxy/responses.codex-oauth.test.ts @@ -35,9 +35,13 @@ const dbInsertMock = vi.fn((_arg?: any) => ({ }, })); -vi.mock('undici', () => ({ - fetch: (...args: unknown[]) => fetchMock(...args), -})); +vi.mock('undici', async () => { + const actual = await vi.importActual('undici'); + return { + ...actual, + fetch: (...args: unknown[]) => fetchMock(...args), + }; +}); vi.mock('../../services/tokenRouter.js', () => ({ tokenRouter: { @@ -109,6 +113,7 @@ vi.mock('../../db/index.js', () => ({ hasProxyLogBillingDetailsColumn: async () => false, hasProxyLogClientColumns: async () => false, hasProxyLogDownstreamApiKeyIdColumn: async () => false, + hasProxyLogStreamTimingColumns: async () => false, schema: { proxyLogs: {}, siteApiEndpoints: { diff --git a/src/server/routes/proxy/responses.compact-upstream.test.ts b/src/server/routes/proxy/responses.compact-upstream.test.ts index a34c6b8..12569d9 100644 --- a/src/server/routes/proxy/responses.compact-upstream.test.ts +++ b/src/server/routes/proxy/responses.compact-upstream.test.ts @@ -23,9 +23,13 @@ const dbInsertMock = vi.fn((_arg?: any) => ({ }), })); -vi.mock('undici', () => ({ - fetch: (...args: unknown[]) => fetchMock(...args), -})); +vi.mock('undici', async () => { + const actual = await vi.importActual('undici'); + return { + ...actual, + fetch: (...args: unknown[]) => fetchMock(...args), + }; +}); vi.mock('../../services/tokenRouter.js', () => ({ tokenRouter: { @@ -88,6 +92,7 @@ vi.mock('../../db/index.js', () => ({ hasProxyLogBillingDetailsColumn: async () => false, hasProxyLogClientColumns: async () => false, hasProxyLogDownstreamApiKeyIdColumn: async () => false, + hasProxyLogStreamTimingColumns: async () => false, schema: { proxyLogs: {}, siteApiEndpoints: { diff --git a/src/server/routes/proxy/responses.websocket.test.ts b/src/server/routes/proxy/responses.websocket.test.ts index 7bc12da..92fb080 100644 --- a/src/server/routes/proxy/responses.websocket.test.ts +++ b/src/server/routes/proxy/responses.websocket.test.ts @@ -24,15 +24,20 @@ const resolveProxyUsageWithSelfLogFallbackMock = vi.fn(async ({ usage }: any) => recoveredFromSelfLog: false, })); const trackedClientSockets = new Set(); +let siteApiEndpointRows: Array> = []; const dbInsertMock = vi.fn((_arg?: any) => ({ values: () => ({ run: () => undefined, }), })); -vi.mock('undici', () => ({ - fetch: (...args: unknown[]) => fetchMock(...args), -})); +vi.mock('undici', async () => { + const actual = await vi.importActual('undici'); + return { + ...actual, + fetch: (...args: unknown[]) => fetchMock(...args), + }; +}); vi.mock('../../services/tokenRouter.js', () => ({ tokenRouter: { @@ -101,7 +106,7 @@ vi.mock('../../db/index.js', () => ({ from: () => ({ where: () => ({ orderBy: () => ({ - all: async () => [], + all: async () => siteApiEndpointRows, }), }), }), @@ -117,6 +122,7 @@ vi.mock('../../db/index.js', () => ({ hasProxyLogBillingDetailsColumn: async () => false, hasProxyLogClientColumns: async () => false, hasProxyLogDownstreamApiKeyIdColumn: async () => false, + hasProxyLogStreamTimingColumns: async () => false, schema: { proxyLogs: {}, siteApiEndpoints: { @@ -156,6 +162,7 @@ function createSelectedChannel(options?: { return { channel: { id: 11, routeId: 22 }, site: { + id: 44, name: options?.siteName ?? (isCodex ? 'codex-site' : 'openai-site'), url: options?.siteUrl ?? (isCodex ? 'https://chatgpt.com/backend-api/codex' : 'https://api.openai.com'), platform: sitePlatform, @@ -365,6 +372,7 @@ describe('responses websocket transport', () => { reportTokenExpiredMock.mockReset(); resolveProxyUsageWithSelfLogFallbackMock.mockClear(); dbInsertMock.mockClear(); + siteApiEndpointRows = []; const selectedChannel = createSelectedChannel(); selectChannelMock.mockReturnValue(selectedChannel); @@ -530,6 +538,92 @@ describe('responses websocket transport', () => { expect(upstreamConnectionCount).toBe(1); }); + it('uses the configured site api endpoint pool for codex websocket transport', async () => { + siteApiEndpointRows = [{ + id: 901, + siteId: 44, + url: upstreamSiteUrl, + enabled: true, + sortOrder: 0, + cooldownUntil: null, + lastSelectedAt: null, + lastFailedAt: null, + lastFailureReason: null, + updatedAt: null, + }]; + const selectedChannel = createSelectedChannel({ + siteUrl: rejectedUpgradeSiteUrl, + }); + selectChannelMock.mockReturnValue(selectedChannel); + previewSelectedChannelMock.mockResolvedValue(selectedChannel); + + const socket = createClientSocket(baseUrl); + await waitForSocketOpen(socket); + const messagePromise = waitForSocketMessageMatching( + socket, + (message) => message?.type === 'response.completed', + ); + + socket.send(JSON.stringify({ + type: 'response.create', + model: 'gpt-5.4', + input: [], + })); + + const message = await messagePromise; + socket.close(); + + expect(message?.type).toBe('response.completed'); + expect(fetchMock).toHaveBeenCalledTimes(0); + expect(upstreamConnectionCount).toBe(1); + }); + + it('preserves the site endpoint exhaustion message on websocket errors', async () => { + siteApiEndpointRows = [{ + id: 902, + siteId: 44, + url: upstreamSiteUrl, + enabled: false, + sortOrder: 0, + cooldownUntil: null, + lastSelectedAt: null, + lastFailedAt: null, + lastFailureReason: null, + updatedAt: null, + }]; + const selectedChannel = createSelectedChannel({ + siteUrl: rejectedUpgradeSiteUrl, + }); + selectChannelMock.mockReturnValue(selectedChannel); + previewSelectedChannelMock.mockResolvedValue(selectedChannel); + + const socket = createClientSocket(baseUrl); + await waitForSocketOpen(socket); + const errorPromise = waitForSocketMessageMatching( + socket, + (message) => message?.type === 'error', + ); + + socket.send(JSON.stringify({ + type: 'response.create', + model: 'gpt-5.4', + input: [], + })); + + const errorMessage = await errorPromise; + socket.close(); + + expect(errorMessage).toMatchObject({ + type: 'error', + status: 408, + error: { + message: '当前站点的 AI 请求地址均不可用', + }, + }); + expect(fetchMock).toHaveBeenCalledTimes(0); + expect(upstreamConnectionCount).toBe(0); + }); + it('echoes x-codex-turn-state on websocket upgrade responses', async () => { const socket = createClientSocket(baseUrl, { 'x-codex-turn-state': 'turn-state-123', diff --git a/src/server/routes/proxy/responsesWebsocket.ts b/src/server/routes/proxy/responsesWebsocket.ts index e0e6b6b..dde198c 100644 --- a/src/server/routes/proxy/responsesWebsocket.ts +++ b/src/server/routes/proxy/responsesWebsocket.ts @@ -11,6 +11,7 @@ import { isModelAllowedByPolicyOrAllowedRoutes, type DownstreamTokenAuthSuccess, } from '../../services/downstreamApiKeyService.js'; +import { runWithSiteApiEndpointPool, SiteApiEndpointRequestError } from '../../services/siteApiEndpointService.js'; import { tokenRouter } from '../../services/tokenRouter.js'; import { buildOauthProviderHeaders } from '../../services/oauth/service.js'; import { openAiResponsesTransformer } from '../../transformers/openai/responses/index.js'; @@ -130,6 +131,18 @@ function selectedChannelSupportsIncrementalInput( return selectedChannelSupportsCodexWebsocketTransport(selectedChannel, requestModel); } +function unwrapCodexWebsocketRuntimeError(error: unknown): CodexWebsocketRuntimeError { + if (error instanceof CodexWebsocketRuntimeError) return error; + if (error instanceof SiteApiEndpointRequestError && error.cause instanceof CodexWebsocketRuntimeError) { + return error.cause; + } + return new CodexWebsocketRuntimeError( + error instanceof Error && error.message.trim() + ? error.message + : 'upstream websocket request failed', + ); +} + function shouldReuseSelectedChannel( selectedChannel: SelectedChannel | null, requestModel: string, @@ -628,44 +641,59 @@ async function handleResponsesWebsocketConnection( account: codexWebsocketChannel.account, 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}`; + + const websocketRuntimeSessionKey = buildCodexSessionResponseStoreKey({ + sessionId: websocketSessionId, + siteId: codexWebsocketChannel.site.id, + accountId: codexWebsocketChannel.account.id, + channelId: codexWebsocketChannel.channel.id, + }) || websocketSessionId; + runtimeSessionKeys.add(websocketRuntimeSessionKey); try { - const websocketRuntimeSessionKey = buildCodexSessionResponseStoreKey({ - sessionId: websocketSessionId, - siteId: codexWebsocketChannel.site.id, - accountId: codexWebsocketChannel.account.id, - channelId: codexWebsocketChannel.channel.id, - }) || websocketSessionId; - runtimeSessionKeys.add(websocketRuntimeSessionKey); - const runtimeResult = await codexWebsocketRuntime.sendRequest({ - sessionId: websocketRuntimeSessionKey, - requestUrl, - headers: prepared.headers, - body: prepared.body, - }); + const runtimeResult = await runWithSiteApiEndpointPool( + codexWebsocketChannel.site as Parameters[0], + async (target) => { + const prepared = buildUpstreamEndpointRequest({ + endpoint: 'responses', + modelName: asTrimmedString(codexWebsocketChannel.actualModel) || requestModel, + stream: true, + tokenValue: codexWebsocketChannel.tokenValue, + sitePlatform: codexWebsocketChannel.site.platform, + siteUrl: target.baseUrl, + openaiBody: normalized.request, + downstreamFormat: 'responses', + responsesOriginalBody: normalized.request, + downstreamHeaders, + providerHeaders, + codexExplicitSessionId: deriveCodexExplicitSessionId(normalized.request, websocketSessionId), + }); + const requestUrl = `${target.baseUrl.replace(/\/+$/, '')}${prepared.path}`; + + try { + return await codexWebsocketRuntime.sendRequest({ + sessionId: websocketRuntimeSessionKey, + requestUrl, + headers: prepared.headers, + body: prepared.body, + }); + } catch (error) { + const runtimeError = error instanceof CodexWebsocketRuntimeError + ? error + : new CodexWebsocketRuntimeError('upstream websocket request failed'); + throw new SiteApiEndpointRequestError(runtimeError.message, { + status: runtimeError.status, + cause: runtimeError, + }); + } + }, + ); 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'); + const runtimeError = unwrapCodexWebsocketRuntimeError(error); if (runtimeError.status && runtimeError.events.length === 0) { const forwarded = await forwardResponsesRequestViaHttp({ app, diff --git a/src/server/routes/proxy/search.test.ts b/src/server/routes/proxy/search.test.ts index eb38164..3cef5a9 100644 --- a/src/server/routes/proxy/search.test.ts +++ b/src/server/routes/proxy/search.test.ts @@ -16,9 +16,13 @@ const dbInsertMock = vi.fn((_arg?: any) => ({ }), })); -vi.mock('undici', () => ({ - fetch: (...args: unknown[]) => fetchMock(...args), -})); +vi.mock('undici', async () => { + const actual = await vi.importActual('undici'); + return { + ...actual, + fetch: (...args: unknown[]) => fetchMock(...args), + }; +}); vi.mock('../../services/tokenRouter.js', () => ({ tokenRouter: { @@ -75,6 +79,7 @@ vi.mock('../../db/index.js', () => ({ hasProxyLogBillingDetailsColumn: async () => false, hasProxyLogClientColumns: async () => false, hasProxyLogDownstreamApiKeyIdColumn: async () => false, + hasProxyLogStreamTimingColumns: async () => false, schema: { proxyLogs: {}, siteApiEndpoints: { diff --git a/src/server/routes/proxy/videos.test.ts b/src/server/routes/proxy/videos.test.ts index 4a0d0bb..9654d9c 100644 --- a/src/server/routes/proxy/videos.test.ts +++ b/src/server/routes/proxy/videos.test.ts @@ -85,6 +85,7 @@ vi.mock('../../db/index.js', () => ({ }), }), }, + hasProxyLogStreamTimingColumns: async () => false, schema: { siteApiEndpoints: { id: {},