[daily] route codex websocket through site api endpoints (#386)

* fix: route codex websocket through site api endpoints

* fix: preserve websocket endpoint exhaustion errors

* test: align proxy route mocks with first-byte timeout helpers

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,

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' },

src/server/routes/proxy/chat.codex-oauth.test.ts

@@ -21,9 +21,13 @@ const dbInsertMock = vi.fn((_arg?: any) => ({
   }),
 }));
 
-vi.mock('undici', () => ({
+vi.mock('undici', async () => {
+  const actual = await vi.importActual<typeof import('undici')>('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: {

src/server/routes/proxy/chat.count-tokens.test.ts

@@ -25,9 +25,13 @@ const dbInsertMock = vi.fn((_arg?: any) => ({
   }),
 }));
 
-vi.mock('undici', () => ({
+vi.mock('undici', async () => {
+  const actual = await vi.importActual<typeof import('undici')>('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: {

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', () => ({
+vi.mock('undici', async () => {
+  const actual = await vi.importActual<typeof import('undici')>('undici');
+  return {
+    ...actual,
     fetch: (...args: unknown[]) => fetchMock(...args),
-}));
+  };
+});
 
 vi.mock('../../services/tokenRouter.js', () => ({
   tokenRouter: {

src/server/routes/proxy/chat.stream.test.ts

@@ -26,9 +26,13 @@ const dbInsertMock = vi.fn((_arg?: any) => ({
   }),
 }));
 
-vi.mock('undici', () => ({
+vi.mock('undici', async () => {
+  const actual = await vi.importActual<typeof import('undici')>('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: {

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', () => ({
+vi.mock('undici', async () => {
+  const actual = await vi.importActual<typeof import('undici')>('undici');
+  return {
+    ...actual,
     fetch: (...args: unknown[]) => fetchMock(...args),
-}));
+  };
+});
 
 vi.mock('../../services/tokenRouter.js', () => ({
   tokenRouter: {

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', () => ({
+vi.mock('undici', async () => {
+  const actual = await vi.importActual<typeof import('undici')>('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', () => {

src/server/routes/proxy/gemini.test.ts

@@ -48,10 +48,13 @@ function createDbSelectChain() {
   };
 }
 
-vi.mock('undici', () => ({
+vi.mock('undici', async () => {
+  const actual = await vi.importActual<typeof import('undici')>('undici');
+  return {
+    ...actual,
     fetch: (...args: unknown[]) => fetchMock(...args),
-  Response: globalThis.Response,
-}));
+  };
+});
 
 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: {

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: {

src/server/routes/proxy/responses.codex-oauth.test.ts

@@ -35,9 +35,13 @@ const dbInsertMock = vi.fn((_arg?: any) => ({
   },
 }));
 
-vi.mock('undici', () => ({
+vi.mock('undici', async () => {
+  const actual = await vi.importActual<typeof import('undici')>('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: {

src/server/routes/proxy/responses.compact-upstream.test.ts

@@ -23,9 +23,13 @@ const dbInsertMock = vi.fn((_arg?: any) => ({
   }),
 }));
 
-vi.mock('undici', () => ({
+vi.mock('undici', async () => {
+  const actual = await vi.importActual<typeof import('undici')>('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: {

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<WebSocket>();
+let siteApiEndpointRows: Array<Record<string, unknown>> = [];
 const dbInsertMock = vi.fn((_arg?: any) => ({
   values: () => ({
     run: () => undefined,
   }),
 }));
 
-vi.mock('undici', () => ({
+vi.mock('undici', async () => {
+  const actual = await vi.importActual<typeof import('undici')>('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',

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,23 +641,7 @@ 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}`;
 
-            try {
             const websocketRuntimeSessionKey = buildCodexSessionResponseStoreKey({
               sessionId: websocketSessionId,
               siteId: codexWebsocketChannel.site.id,
@@ -652,20 +649,51 @@ async function handleResponsesWebsocketConnection(
               channelId: codexWebsocketChannel.channel.id,
             }) || websocketSessionId;
             runtimeSessionKeys.add(websocketRuntimeSessionKey);
-              const runtimeResult = await codexWebsocketRuntime.sendRequest({
+
+            try {
+              const runtimeResult = await runWithSiteApiEndpointPool(
+                codexWebsocketChannel.site as Parameters<typeof runWithSiteApiEndpointPool>[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,

src/server/routes/proxy/search.test.ts

@@ -16,9 +16,13 @@ const dbInsertMock = vi.fn((_arg?: any) => ({
   }),
 }));
 
-vi.mock('undici', () => ({
+vi.mock('undici', async () => {
+  const actual = await vi.importActual<typeof import('undici')>('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: {

src/server/routes/proxy/videos.test.ts

@@ -85,6 +85,7 @@ vi.mock('../../db/index.js', () => ({
       }),
     }),
   },
+  hasProxyLogStreamTimingColumns: async () => false,
   schema: {
     siteApiEndpoints: {
       id: {},