fix: address codex review findings

2026-03-21 19:45:29 +08:00
parent 1d8889001f
commit d9e550c279
10 changed files with 638 additions and 167 deletions
@@ -255,6 +255,33 @@ describe('responses proxy codex oauth refresh', () => {
    expect(response.json()?.output_text).toContain('ok after codex forbidden refresh');
  });

+  it('does not refresh codex oauth token on non-auth 403 responses', async () => {
+    fetchMock.mockResolvedValueOnce(new Response(JSON.stringify({
+      error: { message: 'quota exceeded for workspace', type: 'usage_limit_reached' },
+    }), {
+      status: 403,
+      headers: { 'content-type': 'application/json' },
+    }));
+
+    const response = await app.inject({
+      method: 'POST',
+      url: '/v1/responses',
+      payload: {
+        model: 'gpt-5.2-codex',
+        input: 'hello codex',
+      },
+    });
+
+    expect(response.statusCode).toBe(403);
+    expect(refreshOauthAccessTokenSingleflightMock).not.toHaveBeenCalled();
+    expect(fetchMock).toHaveBeenCalledTimes(1);
+    expect(response.json()).toMatchObject({
+      error: {
+        message: expect.stringContaining('quota exceeded for workspace'),
+      },
+    });
+  });
+
  it('retries oauth responses requests with a normalized upstream URL after refresh', async () => {
    selectChannelMock.mockReturnValue({
      channel: { id: 11, routeId: 22 },
@@ -194,6 +194,16 @@ function waitForSocketMessages(socket: WebSocket, count: number, timeoutMs = 100
  });
 }

+function createDeferred<T>() {
+  let resolve!: (value: T | PromiseLike<T>) => void;
+  let reject!: (reason?: unknown) => void;
+  const promise = new Promise<T>((res, rej) => {
+    resolve = res;
+    reject = rej;
+  });
+  return { promise, resolve, reject };
+}
+
 function createClientSocket(baseUrl: string, headers: Record<string, string> = {}) {
  return new WebSocket(`${baseUrl}/v1/responses`, {
    headers: {
@@ -203,6 +213,10 @@ function createClientSocket(baseUrl: string, headers: Record<string, string> = {
  });
 }

+function createClientSocketForPath(path: string, headers: Record<string, string> = {}) {
+  return new WebSocket(path, { headers });
+}
+
 describe('responses websocket transport', () => {
  const originalCodexResponsesWebsocketBeta = config.codexResponsesWebsocketBeta;
  let app: FastifyInstance;
@@ -545,6 +559,54 @@ describe('responses websocket transport', () => {
    expect(messages[0]?.response?.id).toBe('resp_http_fallback_401');
  });

+  it('preserves query parameter auth when websocket transport falls back to the HTTP responses route', 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);
+    authorizeDownstreamTokenMock.mockResolvedValueOnce({
+      ok: true,
+      source: 'global',
+      token: 'sk-query-auth',
+      key: null,
+      policy: {
+        supportedModels: [],
+        allowedRouteIds: [],
+        siteWeightMultipliers: {},
+      },
+    });
+    fetchMock.mockResolvedValueOnce(createSseResponse([
+      'event: response.completed\n',
+      'data: {"type":"response.completed","response":{"id":"resp_http_fallback_query","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 = createClientSocketForPath(`${baseUrl}/v1/responses?key=sk-query-auth`);
+    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[0]?.response?.id).toBe('resp_http_fallback_query');
+  });
+
  it('rejects websocket turns whose model is blocked by the downstream key policy before channel selection', async () => {
    authorizeDownstreamTokenMock.mockResolvedValueOnce({
      ok: true,
@@ -678,6 +740,111 @@ describe('responses websocket transport', () => {
    ]);
  });

+  it('serializes websocket messages per connection so follow-up turns wait for the previous HTTP fallback to finish', async () => {
+    const selectedChannel = createSelectedChannel({
+      sitePlatform: 'openai',
+      actualModel: 'gpt-4.1',
+    });
+    selectChannelMock.mockReturnValue(selectedChannel);
+    previewSelectedChannelMock.mockResolvedValue(selectedChannel);
+    const firstResponseGate = createDeferred<Response>();
+    fetchMock
+      .mockImplementationOnce(() => firstResponseGate.promise)
+      .mockResolvedValueOnce(createSseResponse([
+        'event: response.completed\n',
+        'data: {"type":"response.completed","response":{"id":"resp_ws_2","model":"gpt-4.1","status":"completed","output":[{"id":"msg_2","type":"message","role":"assistant","status":"completed","content":[{"type":"output_text","text":"done"}]}],"usage":{"input_tokens":5,"output_tokens":1,"total_tokens":6}}}\n\n',
+        'data: [DONE]\n\n',
+      ]));
+
+    const socket = createClientSocket(baseUrl);
+    await waitForSocketOpen(socket);
+
+    socket.send(JSON.stringify({
+      type: 'response.create',
+      model: 'gpt-4.1',
+      instructions: 'be helpful',
+      input: [
+        {
+          id: 'msg_user_1',
+          type: 'message',
+          role: 'user',
+          content: [{ type: 'input_text', text: 'call the tool' }],
+        },
+      ],
+    }));
+
+    while (fetchMock.mock.calls.length < 1) {
+      await new Promise((resolve) => setTimeout(resolve, 0));
+    }
+
+    socket.send(JSON.stringify({
+      type: 'response.create',
+      previous_response_id: 'resp_ws_1',
+      input: [
+        {
+          id: 'tool_out_1',
+          type: 'function_call_output',
+          call_id: 'call_1',
+          output: 'tool result',
+        },
+      ],
+    }));
+
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    expect(fetchMock).toHaveBeenCalledTimes(1);
+
+    const messagesPromise = waitForSocketMessages(socket, 4);
+    firstResponseGate.resolve(createSseResponse([
+      'event: response.output_item.done\n',
+      'data: {"type":"response.output_item.done","output_index":0,"item":{"id":"fc_1","type":"function_call","call_id":"call_1"}}\n\n',
+      'event: response.output_item.done\n',
+      'data: {"type":"response.output_item.done","output_index":1,"item":{"id":"msg_1","type":"message","role":"assistant","status":"completed","content":[{"type":"output_text","text":"call tool"}]}}\n\n',
+      'event: response.completed\n',
+      'data: {"type":"response.completed","response":{"id":"resp_ws_1","model":"gpt-4.1","status":"completed","output":[{"id":"fc_1","type":"function_call","call_id":"call_1"},{"id":"msg_1","type":"message","role":"assistant","status":"completed","content":[{"type":"output_text","text":"call tool"}]}],"usage":{"input_tokens":3,"output_tokens":1,"total_tokens":4}}}\n\n',
+      'data: [DONE]\n\n',
+    ]));
+
+    const messages = await messagesPromise;
+    socket.close();
+
+    expect(fetchMock).toHaveBeenCalledTimes(2);
+    expect(messages.map((message) => message.type)).toEqual([
+      'response.output_item.done',
+      'response.output_item.done',
+      'response.completed',
+      'response.completed',
+    ]);
+    const [, secondOptions] = fetchMock.mock.calls[1] as [string, RequestInit];
+    const secondBody = JSON.parse(String(secondOptions.body));
+    expect(secondBody.input).toEqual([
+      {
+        id: 'msg_user_1',
+        type: 'message',
+        role: 'user',
+        content: [{ type: 'input_text', text: 'call the tool' }],
+      },
+      {
+        id: 'fc_1',
+        type: 'function_call',
+        call_id: 'call_1',
+        status: 'completed',
+      },
+      {
+        id: 'msg_1',
+        type: 'message',
+        role: 'assistant',
+        status: 'completed',
+        content: [{ type: 'output_text', text: 'call tool' }],
+      },
+      {
+        id: 'tool_out_1',
+        type: 'function_call_output',
+        call_id: 'call_1',
+        output: 'tool result',
+      },
+    ]);
+  });
+
  it('preserves incremental response.create payloads with previous_response_id for websocket-capable upstreams', async () => {
    const selectedChannel = createSelectedChannel({
      siteUrl: upstreamSiteUrl,
@@ -339,15 +339,25 @@ async function forwardResponsesRequestViaHttp(input: {
  request: IncomingMessage;
  payload: Record<string, unknown>;
  preserveIncrementalMode: boolean;
+  authToken: string;
 }): Promise<unknown[] | null> {
+  const injectHeaders = {
+    ...buildInjectHeaders(input.request),
+    [RESPONSES_WEBSOCKET_TRANSPORT_HEADER]: '1',
+    ...(input.preserveIncrementalMode ? { [RESPONSES_WEBSOCKET_MODE_HEADER]: 'incremental' } : {}),
+  };
+  if (
+    !headerValueToTrimmedString(injectHeaders.authorization)
+    && !headerValueToTrimmedString(injectHeaders['x-api-key'])
+    && !headerValueToTrimmedString(injectHeaders['x-goog-api-key'])
+  ) {
+    injectHeaders.authorization = `Bearer ${input.authToken}`;
+  }
+
  const response = await input.app.inject({
    method: 'POST',
    url: '/v1/responses',
-    headers: {
-      ...buildInjectHeaders(input.request),
-      [RESPONSES_WEBSOCKET_TRANSPORT_HEADER]: '1',
-      ...(input.preserveIncrementalMode ? { [RESPONSES_WEBSOCKET_MODE_HEADER]: 'incremental' } : {}),
-    },
+    headers: injectHeaders,
    payload: input.payload,
  });

@@ -444,7 +454,7 @@ function writeUpgradeHttpError(socket: Duplex, status: number, message: string):
        ? 'Bad Request'
        : 'Error';
  const body = JSON.stringify({ error: message });
-  socket.write(
+  socket.end(
    `HTTP/1.1 ${status} ${statusText}\r\n`
    + 'Content-Type: application/json\r\n'
    + `Content-Length: ${Buffer.byteLength(body)}\r\n`
@@ -452,7 +462,6 @@ function writeUpgradeHttpError(socket: Duplex, status: number, message: string):
    + '\r\n'
    + body,
  );
-  socket.destroy();
 }

 async function supportsResponsesWebsocketIncrementalInput(
@@ -483,148 +492,155 @@ async function handleResponsesWebsocketConnection(
  let lastRequest: Record<string, unknown> | null = null;
  let lastResponseOutput: unknown[] = [];
  let selectedChannel: SelectedChannel | null = null;
+  let messageQueue = Promise.resolve();

  socket.once('close', () => {
    void codexWebsocketRuntime.closeSession(websocketSessionId);
  });

-  socket.on('message', async (raw) => {
-    try {
-      const parsed = parseJsonObject(raw);
-      if (!parsed) {
-        writeResponsesWebsocketError(socket, 400, 'Invalid websocket JSON payload');
-        return;
-      }
-
-      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;
-
-      if (codexWebsocketChannel) {
-        const downstreamHeaders: Record<string, unknown> = {
-          ...(request.headers as Record<string, unknown>),
-          [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}`;
-
+  socket.on('message', (raw) => {
+    messageQueue = messageQueue
+      .catch(() => undefined)
+      .then(async () => {
        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));
+          const parsed = parseJsonObject(raw);
+          if (!parsed) {
+            writeResponsesWebsocketError(socket, 400, 'Invalid websocket JSON payload');
+            return;
          }
-        } 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;
+
+          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;
          }
-          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 (!shouldReuseSelectedChannel(selectedChannel, requestModel)) {
+            selectedChannel = requestModel
+              ? await tokenRouter.selectChannel(requestModel, authContext.policy)
+              : null;
+          }
+
+          const codexWebsocketChannel = selectedChannelSupportsCodexWebsocketTransport(selectedChannel, requestModel)
+            ? selectedChannel
+            : null;
+
+          if (codexWebsocketChannel) {
+            const downstreamHeaders: Record<string, unknown> = {
+              ...(request.headers as Record<string, unknown>),
+              [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,
+              });
+              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,
+                  authToken: authContext.token,
+                });
+                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,
+            authToken: authContext.token,
+          });
+          if (forwardedOutput) {
+            lastResponseOutput = forwardedOutput;
+          }
+        } catch {
+          writeResponsesWebsocketError(socket, 500, 'internal websocket proxy error');
+        }
      });
-      if (forwardedOutput) {
-        lastResponseOutput = forwardedOutput;
-      }
-    } catch {
-      writeResponsesWebsocketError(socket, 500, 'internal websocket proxy error');
-    }
  });
 }