fix(openai): split websocket and HTTP timeouts

This commit is contained in:
Aiden Cline 2026-06-03 22:44:55 -05:00
parent 2a33addd29
commit f64319366f
4 changed files with 100 additions and 4 deletions

View file

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

View file

@ -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[] = []

View file

@ -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,

View file

@ -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) => {