Procházet zdrojové kódy

fix(openai): split websocket and HTTP timeouts

Aiden Cline před 2 měsíci
rodič
revize
f64319366f

+ 29 - 3
packages/opencode/src/plugin/openai/codex.ts

@@ -6,6 +6,7 @@ import os from "os"
 import { setTimeout as sleep } from "node:timers/promises"
 import { createServer } from "http"
 import { OpenAIWebSocketPool } from "./ws-pool"
+import { ProviderError } from "../../provider/error"
 
 const log = Log.create({ service: "plugin.codex" })
 
@@ -14,8 +15,24 @@ const ISSUER = "https://auth.openai.com"
 const CODEX_API_ENDPOINT = "https://chatgpt.com/backend-api/codex/responses"
 const OAUTH_PORT = 1455
 const OAUTH_POLLING_SAFETY_MARGIN_MS = 3000
+const DEFAULT_HTTP_HEADER_TIMEOUT = 10_000
+// This adapter applies a fresh header timeout only when it actually uses HTTP.
+const HEADER_TIMEOUT = Symbol.for("opencode.provider.header-timeout")
 const ALLOWED_MODELS = new Set(["gpt-5.5", "gpt-5.3-codex-spark", "gpt-5.4", "gpt-5.4-mini"])
 
+export function fetchWithHeaderTimeout(
+  fetch: typeof globalThis.fetch,
+  input: RequestInfo | URL,
+  init: RequestInit | undefined,
+  timeout: number | false,
+) {
+  if (timeout === false) return fetch(input, init)
+  const controller = new AbortController()
+  const timer = setTimeout(() => controller.abort(new ProviderError.HeaderTimeoutError(timeout)), timeout)
+  const signal = init?.signal ? AbortSignal.any([init.signal, controller.signal]) : controller.signal
+  return fetch(input, { ...init, signal }).finally(() => clearTimeout(timer))
+}
+
 interface PkceCodes {
   verifier: string
   challenge: string
@@ -354,10 +371,17 @@ function waitForOAuthCallback(pkce: PkceCodes, state: string): Promise<TokenResp
 export async function CodexAuthPlugin(input: PluginInput, options: CodexAuthPluginOptions = {}): Promise<Hooks> {
   const issuer = options.issuer ?? ISSUER
   const codexApiEndpoint = options.codexApiEndpoint ?? CODEX_API_ENDPOINT
+  let httpHeaderTimeout: number | false = DEFAULT_HTTP_HEADER_TIMEOUT
   let websocketFetchInstalled = false
   const websocketFetches: Array<ReturnType<typeof OpenAIWebSocketPool.createWebSocketFetch>> = []
+  const httpFetch: typeof globalThis.fetch = (requestInput, init) =>
+    fetchWithHeaderTimeout(fetch, requestInput, init, httpHeaderTimeout)
 
   return {
+    async config(config) {
+      const value = config.provider?.openai?.options?.headerTimeout
+      httpHeaderTimeout = typeof value === "number" || value === false ? value : DEFAULT_HTTP_HEADER_TIMEOUT
+    },
     async dispose() {
       for (const websocketFetch of websocketFetches) websocketFetch.close()
       websocketFetches.length = 0
@@ -404,7 +428,7 @@ export async function CodexAuthPlugin(input: PluginInput, options: CodexAuthPlug
       async loader(getAuth) {
         const auth = await getAuth()
         const websocketFetch = options.experimentalWebSockets
-          ? OpenAIWebSocketPool.createWebSocketFetch({ httpFetch: fetch })
+          ? Object.assign(OpenAIWebSocketPool.createWebSocketFetch({ httpFetch }), { [HEADER_TIMEOUT]: false })
           : undefined
         if (websocketFetch) {
           websocketFetches.push(websocketFetch)
@@ -419,7 +443,7 @@ export async function CodexAuthPlugin(input: PluginInput, options: CodexAuthPlug
             }>
           | undefined
 
-        return {
+        const result = {
           apiKey: OAUTH_DUMMY_KEY,
           async fetch(requestInput: RequestInfo | URL, init?: RequestInit) {
             if (init?.headers) {
@@ -504,9 +528,11 @@ export async function CodexAuthPlugin(input: PluginInput, options: CodexAuthPlug
               headers,
             }
             if (websocketFetch && parsed.pathname.endsWith("/responses")) return websocketFetch(url, requestInit)
-            return fetch(url, OpenAIWebSocketPool.withoutInternalHeaders(requestInit))
+            return (websocketFetch ? httpFetch : fetch)(url, OpenAIWebSocketPool.withoutInternalHeaders(requestInit))
           },
         }
+        if (websocketFetch) Object.assign(result.fetch, { [HEADER_TIMEOUT]: false })
+        return result
       },
       methods: [
         {

+ 4 - 1
packages/opencode/src/provider/provider.ts

@@ -33,6 +33,8 @@ import { ProviderError } from "./error"
 
 const log = Log.create({ service: "provider" })
 const OPENAI_HEADER_TIMEOUT_DEFAULT = 10_000
+// Custom fetch adapters can opt out when they apply route-specific header timing internally.
+const HEADER_TIMEOUT = Symbol.for("opencode.provider.header-timeout")
 
 function wrapSSE(res: Response, ms: number, ctl: AbortController) {
   if (typeof ms !== "number" || ms <= 0) return res
@@ -1603,7 +1605,8 @@ export const layer = Layer.effect(
           const fetchFn = customFetch ?? fetch
           const opts = init ?? {}
           const chunkAbortCtl = typeof chunkTimeout === "number" && chunkTimeout > 0 ? new AbortController() : undefined
-          const headerTimeoutMs = headerTimeout === false ? undefined : headerTimeout
+          const headerTimeoutMs =
+            headerTimeout === false || customFetch?.[HEADER_TIMEOUT] === false ? undefined : headerTimeout
           const headerTimeoutCtl = typeof headerTimeoutMs === "number" ? timeoutController(headerTimeoutMs) : undefined
           const signals: AbortSignal[] = []
 

+ 46 - 0
packages/opencode/test/plugin/codex.test.ts

@@ -1,11 +1,13 @@
 import { describe, expect, test } from "bun:test"
 import {
   CodexAuthPlugin,
+  fetchWithHeaderTimeout,
   parseJwtClaims,
   extractAccountIdFromClaims,
   extractAccountId,
   type IdTokenClaims,
 } from "../../src/plugin/openai/codex"
+import { ProviderError } from "../../src/provider/error"
 
 function createTestJwt(payload: object): string {
   const header = Buffer.from(JSON.stringify({ alg: "none" })).toString("base64url")
@@ -137,9 +139,53 @@ describe("plugin.codex", () => {
 
     expect(disabledOptions.fetch).toBeUndefined()
     expect(enabledOptions.fetch).toBeFunction()
+    expect(enabledOptions.fetch?.[Symbol.for("opencode.provider.header-timeout")]).toBe(false)
     await enabled.dispose?.()
   })
 
+  test("applies configured header timeout to websocket HTTP fallback", async () => {
+    using server = Bun.serve({
+      port: 0,
+      async fetch() {
+        await Bun.sleep(50)
+        return new Response("http")
+      },
+    })
+    const hooks = await CodexAuthPlugin({} as never, { experimentalWebSockets: true })
+    await hooks.config!({ provider: { openai: { options: { headerTimeout: 20 } } } } as never)
+    const loaded = await hooks.auth!.loader!(async () => ({ type: "api", key: "sk-test" }) as never, {} as never)
+
+    await expect(
+      loaded.fetch!(new URL("/v1/responses", server.url), {
+        method: "POST",
+        body: JSON.stringify({ stream: true }),
+      }),
+    ).rejects.toBeInstanceOf(ProviderError.HeaderTimeoutError)
+    await hooks.dispose?.()
+  })
+
+  test("marks websocket OAuth transport as managing its own header timeout", async () => {
+    const hooks = await CodexAuthPlugin({} as never, { experimentalWebSockets: true })
+    const loaded = await hooks.auth!.loader!(
+      async () => ({ type: "oauth", refresh: "refresh", access: "access", expires: Date.now() + 60_000 }) as never,
+      {} as never,
+    )
+
+    expect(loaded.fetch?.[Symbol.for("opencode.provider.header-timeout")]).toBe(false)
+    await hooks.dispose?.()
+  })
+
+  test("can disable websocket HTTP fallback header timeout", async () => {
+    const response = await fetchWithHeaderTimeout(
+      async () => new Response("http"),
+      "https://example.com/v1/responses",
+      undefined,
+      false,
+    )
+
+    expect(await response.text()).toBe("http")
+  })
+
   test("deduplicates concurrent Codex token refreshes", async () => {
     let auth = {
       type: "oauth" as const,

+ 21 - 0
packages/opencode/test/plugin/openai-ws.test.ts

@@ -7,6 +7,7 @@ import { APICallError } from "ai"
 import { ProviderError } from "../../src/provider/error"
 import { OpenAIWebSocket } from "../../src/plugin/openai/ws"
 import { OpenAIWebSocketPool, TITLE_HEADER } from "../../src/plugin/openai/ws-pool"
+import { fetchWithHeaderTimeout } from "../../src/plugin/openai/codex"
 
 describe("plugin.openai.ws", () => {
   test("derives websocket URLs and sends auth plus protocol headers", async () => {
@@ -166,6 +167,26 @@ describe("plugin.openai.ws-pool", () => {
     fetch.close()
   })
 
+  test("does not apply HTTP header timeout while waiting for the first websocket event", async () => {
+    await using server = await createWebSocketServer((socket) => {
+      socket.once("message", () => {
+        setTimeout(() => {
+          socket.send(JSON.stringify({ type: "response.completed", response: { id: "resp_delayed" } }))
+        }, 50)
+      })
+    })
+    const fetch = OpenAIWebSocketPool.createWebSocketFetch({
+      url: server.url,
+      httpFetch: (input, init) => fetchWithHeaderTimeout(globalThis.fetch, input, init, 20),
+      idleTimeout: 100,
+    })
+
+    const response = await fetch(server.url, streamRequest())
+
+    expect(await response.text()).toContain("data: [DONE]")
+    fetch.close()
+  })
+
   test("rotates a socket that exceeds max connection age", async () => {
     let connections = 0
     await using server = await createWebSocketServer((socket) => {