From a7e050cdd701ae8c0165d525046e67c2d150b6e6 Mon Sep 17 00:00:00 2001 From: iceteaSA <171169159+iceteaSA@users.noreply.github.com> Date: Sun, 19 Jul 2026 21:36:09 +0200 Subject: [PATCH] fix(ws): mark and reroute on admission-time usage limit Admission-time usage_limit_reached frames bypassed the response.failed handler, leaving the exhausted account unmarked so OpenCode's retry loop selected it again. Raw upgrade 429s also discarded the status and provider reset body. Mark the connection's captured account before returning a ResponseStreamError. Its retryable APICallError contract is what OpenCode honors when reissuing the turn, while the explicit provider reset keeps the routing mark authoritative. --- packages/opencode/src/core/quota-manager.ts | 24 +- packages/opencode/src/index.ts | 25 +- packages/opencode/src/raw-ws-bun.ts | 43 +- packages/opencode/src/raw-ws-node.ts | 39 +- packages/opencode/src/raw-ws-upgrade.ts | 60 +++ .../opencode/src/tests/integration.test.ts | 117 ++++ .../src/tests/quota-rate-limit.test.ts | 32 ++ packages/opencode/src/tests/raw-ws.test.ts | 500 ++++++++++++++++++ packages/opencode/src/tests/ws-pool.test.ts | 387 ++++++++++++++ packages/opencode/src/ws-pool.ts | 56 +- packages/opencode/src/ws.ts | 93 +++- 11 files changed, 1317 insertions(+), 59 deletions(-) create mode 100644 packages/opencode/src/raw-ws-upgrade.ts diff --git a/packages/opencode/src/core/quota-manager.ts b/packages/opencode/src/core/quota-manager.ts index 19ca16d..8a8cd1a 100644 --- a/packages/opencode/src/core/quota-manager.ts +++ b/packages/opencode/src/core/quota-manager.ts @@ -53,23 +53,27 @@ export function quotaWindowResetIsPast( } /** - * Resolve when a mid-stream-exhausted account's window actually resets, from - * its own last-known cached quota snapshot (a response.failed frame carries no - * reset time itself). The frame always names the exhausted window (primary or - * secondary), so use only THAT window's cached reset; if it's unknown, fall - * back to a conservative bounded default rather than borrowing the OTHER - * window's reset — an unrelated window can reset far in the future and would - * blackhole the account long past its actual rate-limit. The default is - * self-correcting: if the account is still exhausted, the next response.failed - * frame re-marks it. Pure — single source of truth shared by the fetch-override - * reroute and its unit tests. + * Resolve when a WS-exhausted account becomes usable. Admission errors can + * carry the provider's authoritative reset; response.failed frames do not, so + * those fall back to the named window in the last-known quota snapshot. Never + * borrow the other window's reset — an unrelated window can blackhole the + * account long past its actual limit. The bounded default self-corrects when a + * later request reports exhaustion again. */ export function resolveMidStreamRateLimitResetAt( quota: OAuthQuotaSnapshot | undefined, window: string, now: number, defaultMs: number, + explicitResetAt?: number, ): number { + if ( + explicitResetAt !== undefined && + Number.isFinite(explicitResetAt) && + explicitResetAt > now + ) { + return explicitResetAt + } const named = window === PRIMARY ? quota?.primary diff --git a/packages/opencode/src/index.ts b/packages/opencode/src/index.ts index a206a57..0c32df8 100644 --- a/packages/opencode/src/index.ts +++ b/packages/opencode/src/index.ts @@ -1277,13 +1277,12 @@ export async function CodexAuthPlugin( await writeMachineSidebarState(quotaManager, latestStorage) } - // Resolve when a mid-stream-exhausted account's window actually resets, - // from its own last-known cached quota. The resolution logic itself - // lives in quota-manager.ts (single source of truth, unit-tested - // there) — this wrapper only supplies the account's cached snapshot. + // The pure resolver prefers an admission error's explicit reset and + // otherwise uses this account's cached named-window quota. function midStreamRateLimitResetAt( accountKey: string, window: string, + explicitResetAt?: number, ): number { const quota = accountKey === 'main' @@ -1294,6 +1293,7 @@ export async function CodexAuthPlugin( window, Date.now(), DEFAULT_MID_STREAM_RATE_LIMIT_RESET_MS, + explicitResetAt, ) } @@ -1314,13 +1314,10 @@ export async function CodexAuthPlugin( true, ) }, - // Mid-stream quota exhaustion (response.failed carrying - // rate_limit_reached_type) — the frame itself is the authority, - // so this marks the account rate-limited without any quota-API - // call. The fetch override then reroutes away from it: a - // rate-limited main is treated like a killswitch block, and a - // rate-limited fallback is dropped from candidate selection. - onRateLimitReached: (window, accountId) => { + // WS quota exhaustion is authoritative without a quota-API call. + // Admission errors supply their own reset; response.failed uses + // the cached named-window reset or the bounded default. + onRateLimitReached: (window, accountId, explicitResetAt) => { if (!accountId) { // The loader always sets the internal quota-account header, // so this should never fire today — but a future call site @@ -1332,7 +1329,11 @@ export async function CodexAuthPlugin( ) } const accountKey = accountId ?? 'main' - const resetAt = midStreamRateLimitResetAt(accountKey, window) + const resetAt = midStreamRateLimitResetAt( + accountKey, + window, + explicitResetAt, + ) quotaManager.markRateLimited(accountKey, resetAt) logQ.debug('mid-stream rate limit mark', { pid: process.pid, diff --git a/packages/opencode/src/raw-ws-bun.ts b/packages/opencode/src/raw-ws-bun.ts index 47874d8..052178a 100644 --- a/packages/opencode/src/raw-ws-bun.ts +++ b/packages/opencode/src/raw-ws-bun.ts @@ -15,6 +15,7 @@ // send(string), close(), and a `url` field. Text frames only on receive. import { dumpDiagnostic } from './dump' +import { rejectedUpgradeEvent, rejectedUpgradeStatus } from './raw-ws-upgrade' const WS_GUID = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11' @@ -65,6 +66,7 @@ export class RawWebSocket { } private rxBuffer: Uint8Array = new Uint8Array(0) private handshakeDone = false + private upgradeFailureEmitted = false private expectedAccept = '' // message reassembly across fragmented frames private fragOpcode = 0 @@ -166,12 +168,14 @@ export class RawWebSocket { data: (_s: unknown, d: Uint8Array) => this.onData(d), drain: () => this.flushWrites(), error: (_s: unknown, e: unknown) => { + if (this.emitRejectedUpgrade(true)) return this.readyState = 3 const message = e instanceof Error ? e.message : String(e) void this.log('socket_error', { message }) this.emit('error', { message }) }, close: () => { + if (this.emitRejectedUpgrade(true)) return this.readyState = 3 void this.log('socket_close') this.emit('close', { code: 1006, reason: 'socket closed' }) @@ -195,15 +199,25 @@ export class RawWebSocket { const headerText = text.slice(0, idx) const statusLine = headerText.split('\r\n')[0] ?? '' const acceptMatch = headerText.match(/sec-websocket-accept:\s*(\S+)/i) - if (!statusLine.includes(' 101')) { - this.readyState = 3 + const status = rejectedUpgradeStatus(statusLine) + if (status !== 101) { + if (status === undefined) { + this.readyState = 3 + try { + this.socket?.end() + } catch { + /* ignore */ + } + void this.log('upgrade_failed', { statusLine }) + this.emit('error', { message: `WS upgrade failed: ${statusLine}` }) + return + } + if (!this.emitRejectedUpgrade()) return try { this.socket?.end() } catch { /* ignore */ } - void this.log('upgrade_failed', { statusLine }) - this.emit('error', { message: `WS upgrade failed: ${statusLine}` }) return } if (!acceptMatch || acceptMatch[1] !== this.expectedAccept) { @@ -228,6 +242,27 @@ export class RawWebSocket { this.drainFrames() } + private emitRejectedUpgrade(finalizePartial = false): boolean { + if (this.upgradeFailureEmitted) return true + if (this.handshakeDone) return false + const headerEnd = Buffer.from(this.rxBuffer).indexOf('\r\n\r\n') + if (headerEnd === -1) return false + const errorEvent = rejectedUpgradeEvent( + this.rxBuffer, + headerEnd, + finalizePartial, + ) + if (!errorEvent) return false + this.upgradeFailureEmitted = true + this.readyState = 3 + void this.log('upgrade_failed', { + status: errorEvent.status, + partial: finalizePartial, + }) + this.emit('error', errorEvent) + return true + } + private drainFrames(): void { // RFC 6455 frame parsing for server->client (never masked). for (;;) { diff --git a/packages/opencode/src/raw-ws-node.ts b/packages/opencode/src/raw-ws-node.ts index e285dad..1f07e34 100644 --- a/packages/opencode/src/raw-ws-node.ts +++ b/packages/opencode/src/raw-ws-node.ts @@ -18,6 +18,7 @@ import { connect as netConnect } from 'node:net' import { connect as tlsConnect } from 'node:tls' import { dumpDiagnostic } from './dump' +import { rejectedUpgradeEvent, rejectedUpgradeStatus } from './raw-ws-upgrade' const WS_GUID = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11' @@ -61,6 +62,7 @@ export class RawWebSocket { } private rxBuffer: Uint8Array = new Uint8Array(0) private handshakeDone = false + private upgradeFailureEmitted = false private expectedAccept = '' // message reassembly across fragmented frames private fragOpcode = 0 @@ -159,12 +161,14 @@ export class RawWebSocket { ) }) socket.on('error', (e) => { + if (this.emitRejectedUpgrade(true)) return this.readyState = 3 const message = e instanceof Error ? e.message : String(e) void this.log('socket_error', { message }) this.emit('error', { message }) }) socket.on('close', (hadError) => { + if (this.emitRejectedUpgrade(true)) return this.readyState = 3 void this.log('socket_close', { hadError }) this.emit('close', { code: 1006, reason: 'socket closed' }) @@ -189,11 +193,17 @@ export class RawWebSocket { const headerText = text.slice(0, idx) const statusLine = headerText.split('\r\n')[0] ?? '' const acceptMatch = headerText.match(/sec-websocket-accept:\s*(\S+)/i) - if (!statusLine.includes(' 101')) { - this.readyState = 3 + const status = rejectedUpgradeStatus(statusLine) + if (status !== 101) { + if (status === undefined) { + this.readyState = 3 + this.socket?.destroy?.() + void this.log('upgrade_failed', { statusLine }) + this.emit('error', { message: `WS upgrade failed: ${statusLine}` }) + return + } + if (!this.emitRejectedUpgrade()) return this.socket?.destroy?.() - void this.log('upgrade_failed', { statusLine }) - this.emit('error', { message: `WS upgrade failed: ${statusLine}` }) return } if (!acceptMatch || acceptMatch[1] !== this.expectedAccept) { @@ -214,6 +224,27 @@ export class RawWebSocket { this.drainFrames() } + private emitRejectedUpgrade(finalizePartial = false): boolean { + if (this.upgradeFailureEmitted) return true + if (this.handshakeDone) return false + const headerEnd = Buffer.from(this.rxBuffer).indexOf('\r\n\r\n') + if (headerEnd === -1) return false + const errorEvent = rejectedUpgradeEvent( + this.rxBuffer, + headerEnd, + finalizePartial, + ) + if (!errorEvent) return false + this.upgradeFailureEmitted = true + this.readyState = 3 + void this.log('upgrade_failed', { + status: errorEvent.status, + partial: finalizePartial, + }) + this.emit('error', errorEvent) + return true + } + private drainFrames(): void { // RFC 6455 frame parsing for server->client (never masked). for (;;) { diff --git a/packages/opencode/src/raw-ws-upgrade.ts b/packages/opencode/src/raw-ws-upgrade.ts new file mode 100644 index 0000000..b4fee39 --- /dev/null +++ b/packages/opencode/src/raw-ws-upgrade.ts @@ -0,0 +1,60 @@ +import { isRecord } from './util/record' + +export function rejectedUpgradeStatus(statusLine: string): number | undefined { + const status = Number(statusLine.match(/^HTTP\/\d(?:\.\d)?\s+(\d{3})/)?.[1]) + return Number.isFinite(status) ? status : undefined +} + +export function rejectedUpgradeEvent( + buffer: Uint8Array, + headerEnd: number, + finalizePartial = false, +): Record | undefined { + const headerText = Buffer.from(buffer.slice(0, headerEnd)).toString('latin1') + const lines = headerText.split('\r\n') + const statusLine = lines[0] ?? '' + const status = rejectedUpgradeStatus(statusLine) + if (status === undefined || status === 101) return undefined + + const headers = Object.fromEntries( + lines.slice(1).flatMap((line) => { + const separator = line.indexOf(':') + return separator === -1 + ? [] + : [ + [ + line.slice(0, separator).trim().toLowerCase(), + line.slice(separator + 1).trim(), + ], + ] + }), + ) + const bodyBytes = buffer.slice(headerEnd + 4) + const contentLengthHeader = headers['content-length'] + const contentLength = Number(contentLengthHeader) + if ( + !finalizePartial && + (contentLengthHeader === undefined || + (Number.isFinite(contentLength) && bodyBytes.length < contentLength)) + ) { + return undefined + } + + const body = Buffer.from(bodyBytes).toString('utf8') + const parsed = (() => { + try { + const value = JSON.parse(body) + return isRecord(value) ? value : undefined + } catch { + return undefined + } + })() + return { + message: `WS upgrade failed: ${statusLine}`, + status, + status_code: status, + headers, + body, + ...(isRecord(parsed?.error) ? { error: parsed.error } : {}), + } +} diff --git a/packages/opencode/src/tests/integration.test.ts b/packages/opencode/src/tests/integration.test.ts index 8d15364..9a532d5 100644 --- a/packages/opencode/src/tests/integration.test.ts +++ b/packages/opencode/src/tests/integration.test.ts @@ -1483,6 +1483,123 @@ describe('integration: killswitch enforcement', () => { ) }) + it('reroutes to a healthy fallback after main rejects WebSocket admission with usage_limit_reached', async () => { + const fallback: OAuthAccount = { + id: 'fb-admission-healthy', + type: 'oauth', + access: 'sk-fb-admission-access', + refresh: 'sk-fb-admission-refresh', + expires: Date.now() + 24 * 3600_000, + enabled: true, + addedAt: Date.now(), + lastUsed: Date.now(), + } + writeFileSync( + configFile, + JSON.stringify({ + version: 1, + main: { type: 'opencode', provider: 'openai' }, + accounts: [fallback], + }), + ) + + const originalFetch = globalThis.fetch + globalThis.fetch = (async () => + new Response('{}', { status: 200 })) as unknown as typeof globalThis.fetch + + let mainSends = 0 + let fallbackSends = 0 + let hooks: Hooks | undefined + await withFakeWebSocket( + ({ message, authorization }) => ({ + send() { + if (authorization === 'Bearer access-main-admission-rl') { + mainSends++ + message( + JSON.stringify({ + type: 'error', + error: { + type: 'usage_limit_reached', + message: 'The usage limit has been reached', + plan_type: 'team', + resets_at: 1_784_958_366, + eligible_promo: null, + resets_in_seconds: 514_504, + }, + }), + ) + return + } + fallbackSends++ + message( + JSON.stringify({ + type: 'response.completed', + response: { id: `resp_fb_admission_${fallbackSends}` }, + }), + ) + }, + }), + async () => { + try { + hooks = await CodexAuthPlugin(createMockPluginInput(), { + experimentalWebSockets: true, + }) + const authHook = hooks.auth + if (!authHook?.loader) throw new Error('No auth loader') + const loaderResult = await authHook.loader( + async () => ({ + type: 'oauth' as const, + provider: 'openai', + access: 'access-main-admission-rl', + refresh: refreshToken, + expires: Date.now() + 24 * 3600_000, + accountId: 'chatgpt-main-admission-rl', + }), + { + id: 'openai', + label: 'OpenAI', + models: [], + } as unknown as Parameters< + NonNullable<(typeof authHook)['loader']> + >[1], + ) + const fetchOverride = (loaderResult as Record) + .fetch as (url: string, init?: RequestInit) => Promise + const wsRequest: RequestInit = { + method: 'POST', + headers: { + 'content-type': 'application/json', + 'session-id': 'ws-admission-rate-limit-reroute-session', + }, + body: JSON.stringify({ model: 'gpt-5.5', input: [], stream: true }), + } + + const first = await fetchOverride( + 'https://api.openai.com/v1/responses', + wsRequest, + ) + expect(first.status).toBe(200) + await expect(first.text()).rejects.toMatchObject({ + isRetryable: true, + }) + + const second = await fetchOverride( + 'https://api.openai.com/v1/responses', + wsRequest, + ) + expect(second.status).toBe(200) + await second.text() + + expect(mainSends).toBe(1) + expect(fallbackSends).toBe(1) + } finally { + globalThis.fetch = originalFetch + await hooks?.dispose?.() + } + }, + ) + }) + it('excludes a rate-limited fallback from candidate selection after a mid-stream mark', async () => { // fallback-first: the single fallback is tried before main. Its own // mid-stream rate_limit_reached_type mark must make usableFallbackCandidates diff --git a/packages/opencode/src/tests/quota-rate-limit.test.ts b/packages/opencode/src/tests/quota-rate-limit.test.ts index 2471f55..6d16650 100644 --- a/packages/opencode/src/tests/quota-rate-limit.test.ts +++ b/packages/opencode/src/tests/quota-rate-limit.test.ts @@ -354,6 +354,38 @@ describe('resolveMidStreamRateLimitResetAt', () => { ).toBe(now + 300_000) }) + it('prefers an explicit provider reset over cached quota reset math', async () => { + const { resolveMidStreamRateLimitResetAt } = await import( + '../core/quota-manager.ts' + ) + const providerResetAt = now + 900_000 + const quota = { + primary: { + usedPercent: 100, + remainingPercent: 0, + resetsAt: new Date(now + 300_000).toISOString(), + checkedAt: now, + }, + } + const resolveWithExplicitReset = resolveMidStreamRateLimitResetAt as ( + snapshot: typeof quota, + window: string, + now: number, + defaultMs: number, + explicitResetAt?: number, + ) => number + + expect( + resolveWithExplicitReset( + quota, + 'usage_limit_reached', + now, + defaultMs, + providerResetAt, + ), + ).toBe(providerResetAt) + }) + it('uses the bounded default (NOT the other window) when the named window is absent', async () => { // The frame named `primary`, but only `secondary` has a cached reset. We // must NOT borrow secondary's reset — an unrelated window can reset far in diff --git a/packages/opencode/src/tests/raw-ws.test.ts b/packages/opencode/src/tests/raw-ws.test.ts index 4ad874e..9abe84f 100644 --- a/packages/opencode/src/tests/raw-ws.test.ts +++ b/packages/opencode/src/tests/raw-ws.test.ts @@ -1,5 +1,7 @@ import { describe, expect, it } from 'bun:test' +import { createServer, type Socket } from 'node:net' import { RawWebSocket } from '../raw-ws-bun' +import { RawWebSocket as NodeRawWebSocket } from '../raw-ws-node' describe('RawWebSocket Bun', () => { it('queues handshake write and handles partial writes', async () => { @@ -55,4 +57,502 @@ describe('RawWebSocket Bun', () => { } } }) + + it('preserves status and structured usage-limit details from a rejected upgrade', async () => { + let dataCallback: ((socket: unknown, data: Uint8Array) => void) | undefined + const mockSocket = { + write(data: Uint8Array) { + return data.length + }, + end() {}, + } + const mockConnect = async (opts: any) => { + dataCallback = opts.socket.data + opts.socket.open(mockSocket) + return mockSocket + } + const originalConnect = (globalThis as any).Bun?.connect + if ((globalThis as any).Bun) { + ;(globalThis as any).Bun.connect = mockConnect + } + + try { + const ws = new RawWebSocket('ws://localhost:8080', {}) + const errorEvent = new Promise>((resolve) => { + ws.addEventListener('error', (event) => + resolve(event as Record), + ) + }) + await new Promise((resolve) => setTimeout(resolve, 0)) + const body = JSON.stringify({ + error: { + type: 'usage_limit_reached', + message: 'The usage limit has been reached', + plan_type: 'team', + resets_at: 1_784_958_366, + eligible_promo: null, + resets_in_seconds: 514_504, + }, + }) + dataCallback?.( + mockSocket, + new TextEncoder().encode( + `HTTP/1.1 429 Too Many Requests\r\nContent-Type: application/json\r\nContent-Length: ${body.length}\r\n\r\n${body}`, + ), + ) + + expect(await errorEvent).toMatchObject({ + status: 429, + error: { + type: 'usage_limit_reached', + resets_at: 1_784_958_366, + }, + }) + } finally { + if ((globalThis as any).Bun) { + ;(globalThis as any).Bun.connect = originalConnect + } + } + }) + + it('fails a malformed rejected-upgrade status line instead of remaining CONNECTING', async () => { + let dataCallback: ((socket: unknown, data: Uint8Array) => void) | undefined + let ended = false + const mockSocket = { + write(data: Uint8Array) { + return data.length + }, + end() { + ended = true + }, + } + const mockConnect = async (opts: any) => { + dataCallback = opts.socket.data + opts.socket.open(mockSocket) + return mockSocket + } + const originalConnect = (globalThis as any).Bun?.connect + if ((globalThis as any).Bun) { + ;(globalThis as any).Bun.connect = mockConnect + } + + try { + const ws = new RawWebSocket('ws://localhost:8080', {}) + const errorEvent = new Promise>((resolve) => { + ws.addEventListener('error', (event) => + resolve(event as Record), + ) + }) + await new Promise((resolve) => setTimeout(resolve, 0)) + dataCallback?.( + mockSocket, + new TextEncoder().encode('NOT HTTP\r\nHeader: value\r\n\r\n'), + ) + + expect( + await Promise.race([ + errorEvent, + new Promise((_, reject) => + setTimeout( + () => reject(new Error('timed out waiting for upgrade error')), + 100, + ), + ), + ]), + ).toMatchObject({ message: 'WS upgrade failed: NOT HTTP' }) + expect(ws.readyState).toBe(3) + expect(ended).toBe(true) + } finally { + if ((globalThis as any).Bun) { + ;(globalThis as any).Bun.connect = originalConnect + } + } + }) +}) + +describe('RawWebSocket Node', () => { + it('fails a malformed rejected-upgrade status line instead of remaining CONNECTING', async () => { + let peer: Socket | undefined + const server = createServer((socket) => { + peer = socket + socket.once('data', () => { + socket.write('NOT HTTP\r\nHeader: value\r\n\r\n') + }) + }) + await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve)) + const address = server.address() + if (!address || typeof address === 'string') { + throw new Error('test server did not expose a TCP port') + } + const ws = new NodeRawWebSocket(`ws://127.0.0.1:${address.port}`, {}) + + try { + const event = await Promise.race([ + new Promise>((resolve) => { + ws.addEventListener('error', (error) => + resolve(error as Record), + ) + }), + new Promise((_, reject) => + setTimeout( + () => reject(new Error('timed out waiting for upgrade error')), + 500, + ), + ), + ]) + expect(event).toMatchObject({ message: 'WS upgrade failed: NOT HTTP' }) + expect(ws.readyState).toBe(3) + } finally { + ws.close() + peer?.destroy() + await new Promise((resolve) => server.close(() => resolve())) + } + }) +}) + +describe('rejected upgrade finalization', () => { + it('finalizes a truncated rejected-upgrade body when the Bun socket closes', async () => { + let dataCallback: ((socket: unknown, data: Uint8Array) => void) | undefined + let closeCallback: (() => void) | undefined + const mockSocket = { + write(data: Uint8Array) { + return data.length + }, + end() {}, + } + const mockConnect = async (opts: any) => { + dataCallback = opts.socket.data + closeCallback = opts.socket.close + opts.socket.open(mockSocket) + return mockSocket + } + const originalConnect = (globalThis as any).Bun?.connect + if ((globalThis as any).Bun) { + ;(globalThis as any).Bun.connect = mockConnect + } + + const partialBody = '{"error":{"type":"usage_limit_reached"' + try { + const ws = new RawWebSocket('ws://localhost:8080', {}) + let errorEmitted = false + let closeEvents = 0 + const errorEvent = new Promise>((resolve) => { + ws.addEventListener('error', (event) => { + errorEmitted = true + resolve(event as Record) + }) + }) + ws.addEventListener('close', () => closeEvents++) + await new Promise((resolve) => setTimeout(resolve, 0)) + dataCallback?.( + mockSocket, + new TextEncoder().encode( + `HTTP/1.1 429 Too Many Requests\r\nContent-Type: application/json\r\nContent-Length: ${partialBody.length + 100}\r\n\r\n${partialBody}`, + ), + ) + await Promise.resolve() + expect(errorEmitted).toBe(false) + closeCallback?.() + + expect( + await Promise.race([ + errorEvent, + new Promise((_, reject) => + setTimeout( + () => reject(new Error('timed out waiting for upgrade error')), + 100, + ), + ), + ]), + ).toMatchObject({ status: 429, body: partialBody }) + expect(closeEvents).toBe(0) + } finally { + if ((globalThis as any).Bun) { + ;(globalThis as any).Bun.connect = originalConnect + } + } + }) + + it('accumulates a fragmented close-delimited body until the Bun socket closes', async () => { + let dataCallback: ((socket: unknown, data: Uint8Array) => void) | undefined + let closeCallback: (() => void) | undefined + const mockSocket = { + write(data: Uint8Array) { + return data.length + }, + end() {}, + } + const mockConnect = async (opts: any) => { + dataCallback = opts.socket.data + closeCallback = opts.socket.close + opts.socket.open(mockSocket) + return mockSocket + } + const originalConnect = (globalThis as any).Bun?.connect + if ((globalThis as any).Bun) { + ;(globalThis as any).Bun.connect = mockConnect + } + + const firstBody = '{"error":{"type":"usage_limit_reached",' + const secondBody = '"resets_at":1784958366}}' + try { + const ws = new RawWebSocket('ws://localhost:8080', {}) + let errorEmitted = false + const errorEvent = new Promise>((resolve) => { + ws.addEventListener('error', (event) => { + errorEmitted = true + resolve(event as Record) + }) + }) + await new Promise((resolve) => setTimeout(resolve, 0)) + dataCallback?.( + mockSocket, + new TextEncoder().encode( + `HTTP/1.1 429 Too Many Requests\r\nContent-Type: application/json\r\n\r\n${firstBody}`, + ), + ) + await Promise.resolve() + expect(errorEmitted).toBe(false) + + dataCallback?.(mockSocket, new TextEncoder().encode(secondBody)) + await Promise.resolve() + expect(errorEmitted).toBe(false) + closeCallback?.() + + expect( + await Promise.race([ + errorEvent, + new Promise((_, reject) => + setTimeout( + () => reject(new Error('timed out waiting for upgrade error')), + 100, + ), + ), + ]), + ).toMatchObject({ + status: 429, + body: `${firstBody}${secondBody}`, + error: { + type: 'usage_limit_reached', + resets_at: 1_784_958_366, + }, + }) + } finally { + if ((globalThis as any).Bun) { + ;(globalThis as any).Bun.connect = originalConnect + } + } + }) + + it('finalizes an empty rejected-upgrade body (no Content-Length) when the Bun socket closes', async () => { + let dataCallback: ((socket: unknown, data: Uint8Array) => void) | undefined + let closeCallback: (() => void) | undefined + const mockSocket = { + write(data: Uint8Array) { + return data.length + }, + end() {}, + } + const mockConnect = async (opts: any) => { + dataCallback = opts.socket.data + closeCallback = opts.socket.close + opts.socket.open(mockSocket) + return mockSocket + } + const originalConnect = (globalThis as any).Bun?.connect + if ((globalThis as any).Bun) { + ;(globalThis as any).Bun.connect = mockConnect + } + + try { + const ws = new RawWebSocket('ws://localhost:8080', {}) + let errorEmitted = false + const errorEvent = new Promise>((resolve) => { + ws.addEventListener('error', (event) => { + errorEmitted = true + resolve(event as Record) + }) + }) + await new Promise((resolve) => setTimeout(resolve, 0)) + // Status line + headers only: no Content-Length and no body. The client + // cannot know the body is complete until the socket closes. + dataCallback?.( + mockSocket, + new TextEncoder().encode( + 'HTTP/1.1 429 Too Many Requests\r\nContent-Type: application/json\r\n\r\n', + ), + ) + await Promise.resolve() + expect(errorEmitted).toBe(false) + closeCallback?.() + + expect( + await Promise.race([ + errorEvent, + new Promise((_, reject) => + setTimeout( + () => reject(new Error('timed out waiting for upgrade error')), + 100, + ), + ), + ]), + ).toMatchObject({ status: 429, body: '' }) + } finally { + if ((globalThis as any).Bun) { + ;(globalThis as any).Bun.connect = originalConnect + } + } + }) +}) + +describe('RawWebSocket Node rejected-upgrade matrix', () => { + const usageLimitBody = JSON.stringify({ + error: { + type: 'usage_limit_reached', + message: 'The usage limit has been reached', + resets_at: 1_784_958_366, + }, + }) + + async function withUpgradeServer( + handler: (socket: Socket) => void, + run: (url: string) => Promise, + ) { + const server = createServer(handler) + await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve)) + const address = server.address() + if (!address || typeof address === 'string') { + throw new Error('test server did not expose a TCP port') + } + try { + await run(`ws://127.0.0.1:${address.port}`) + } finally { + await new Promise((resolve) => server.close(() => resolve())) + } + } + + function awaitErrorEvent(ws: NodeRawWebSocket, timeoutMs = 1000) { + return Promise.race([ + new Promise>((resolve) => { + ws.addEventListener('error', (error) => + resolve(error as Record), + ) + }), + new Promise((_, reject) => + setTimeout( + () => reject(new Error('timed out waiting for upgrade error')), + timeoutMs, + ), + ), + ]) + } + + it('preserves status and body from a complete 429 (Content-Length)', async () => { + await withUpgradeServer( + (socket) => { + socket.once('data', () => { + socket.write( + `HTTP/1.1 429 Too Many Requests\r\nContent-Type: application/json\r\nContent-Length: ${usageLimitBody.length}\r\n\r\n${usageLimitBody}`, + () => socket.destroy(), + ) + }) + }, + async (url) => { + const ws = new NodeRawWebSocket(url, {}) + try { + expect(await awaitErrorEvent(ws)).toMatchObject({ + status: 429, + error: { + type: 'usage_limit_reached', + resets_at: 1_784_958_366, + }, + }) + } finally { + ws.close() + } + }, + ) + }) + + it('finalizes a partial 429 body when the socket closes early', async () => { + const partialBody = '{"error":{"type":"usage_limit_reached"' + await withUpgradeServer( + (socket) => { + socket.once('data', () => { + // Advertise more bytes than we actually send, then drop the connection. + socket.write( + `HTTP/1.1 429 Too Many Requests\r\nContent-Type: application/json\r\nContent-Length: ${partialBody.length + 100}\r\n\r\n${partialBody}`, + () => socket.destroy(), + ) + }) + }, + async (url) => { + const ws = new NodeRawWebSocket(url, {}) + try { + expect(await awaitErrorEvent(ws)).toMatchObject({ + status: 429, + body: partialBody, + }) + } finally { + ws.close() + } + }, + ) + }) + + it('accumulates a chunked 429 body (no Content-Length) until close', async () => { + const firstBody = '{"error":{"type":"usage_limit_reached",' + const secondBody = '"resets_at":1784958366}}' + await withUpgradeServer( + (socket) => { + socket.once('data', () => { + socket.write( + `HTTP/1.1 429 Too Many Requests\r\nContent-Type: application/json\r\n\r\n${firstBody}`, + ) + setTimeout(() => { + socket.write(secondBody, () => socket.destroy()) + }, 10) + }) + }, + async (url) => { + const ws = new NodeRawWebSocket(url, {}) + try { + expect(await awaitErrorEvent(ws)).toMatchObject({ + status: 429, + body: `${firstBody}${secondBody}`, + error: { + type: 'usage_limit_reached', + resets_at: 1_784_958_366, + }, + }) + } finally { + ws.close() + } + }, + ) + }) + + it('finalizes an empty 429 (no body) when the socket closes', async () => { + await withUpgradeServer( + (socket) => { + socket.once('data', () => { + socket.write( + 'HTTP/1.1 429 Too Many Requests\r\nContent-Type: application/json\r\n\r\n', + () => socket.destroy(), + ) + }) + }, + async (url) => { + const ws = new NodeRawWebSocket(url, {}) + try { + expect(await awaitErrorEvent(ws)).toMatchObject({ + status: 429, + body: '', + }) + } finally { + ws.close() + } + }, + ) + }) }) diff --git a/packages/opencode/src/tests/ws-pool.test.ts b/packages/opencode/src/tests/ws-pool.test.ts index 764b1b1..85d547b 100644 --- a/packages/opencode/src/tests/ws-pool.test.ts +++ b/packages/opencode/src/tests/ws-pool.test.ts @@ -277,6 +277,393 @@ describe('createWebSocketFetch', () => { ) }) + test('marks an admission-time usage limit with the provider reset before rejecting the stream', async () => { + const resetAtSeconds = 1_784_958_366 + const rateLimitCalls: Array<{ + window: string + accountId?: string + resetAt?: number + }> = [] + await withFakeWebSocket( + ({ message }) => ({ + send() { + message( + JSON.stringify({ + type: 'error', + error: { + type: 'usage_limit_reached', + message: 'The usage limit has been reached', + plan_type: 'team', + resets_at: resetAtSeconds, + eligible_promo: null, + resets_in_seconds: 514_504, + }, + }), + ) + }, + }), + async () => { + const websocketFetch = createWebSocketFetch({ + url: 'https://example.test/backend-api/codex/responses', + onRateLimitReached: (( + window: string, + accountId: string | undefined, + resetAt: number | undefined, + ) => { + rateLimitCalls.push({ window, accountId, resetAt }) + }) as NonNullable< + Parameters[0] + >['onRateLimitReached'], + }) + + const response = await websocketFetch( + 'https://example.test/backend-api/codex/responses', + { + method: 'POST', + headers: { + 'session-id': 'sess-admission-rl', + authorization: 'Bearer tok-main', + 'x-openai-auth-quota-account': 'main', + }, + body: JSON.stringify({ stream: true, input: [] }), + }, + ) + + await expect(response.text()).rejects.toMatchObject({ + isRetryable: true, + }) + expect(rateLimitCalls).toEqual([ + { + window: 'usage_limit_reached', + accountId: 'main', + resetAt: resetAtSeconds * 1000, + }, + ]) + websocketFetch.close() + }, + ) + }) + + test('does not classify a 503 circuit-open admission error as a rate limit', async () => { + const rateLimitCalls: string[] = [] + await withFakeWebSocket( + ({ message }) => ({ + send() { + message( + JSON.stringify({ + error: { + message: 'Service Unavailable', + type: null, + code: 'biscuit_baker_service_me_circuit_open', + param: null, + }, + status: 503, + type: 'error', + }), + ) + }, + }), + async () => { + const websocketFetch = createWebSocketFetch({ + url: 'https://example.test/backend-api/codex/responses', + onRateLimitReached: (window) => rateLimitCalls.push(window), + }) + const response = await websocketFetch( + 'https://example.test/backend-api/codex/responses', + streamRequest({ input: [] }), + ) + + expect(response.status).toBe(503) + expect(await response.json()).toMatchObject({ status: 503 }) + expect(rateLimitCalls).toEqual([]) + websocketFetch.close() + }, + ) + }) + + test('classifies status 429 with another error type but leaves reset undefined for the bounded default', async () => { + const rateLimitCalls: Array<{ window: string; resetAt?: number }> = [] + await withFakeWebSocket( + ({ message }) => ({ + send() { + message( + JSON.stringify({ + type: 'error', + status: 429, + error: { + type: 'rate_limit_exceeded', + message: 'Rate limit exceeded', + }, + }), + ) + }, + }), + async () => { + const websocketFetch = createWebSocketFetch({ + url: 'https://example.test/backend-api/codex/responses', + onRateLimitReached: (window, _accountId, resetAt) => + rateLimitCalls.push({ window, resetAt }), + }) + const response = await websocketFetch( + 'https://example.test/backend-api/codex/responses', + streamRequest({ input: [] }), + ) + + await expect(response.text()).rejects.toBeInstanceOf( + ResponseStreamError, + ) + expect(rateLimitCalls).toEqual([ + { window: 'rate_limit_exceeded', resetAt: undefined }, + ]) + const { resolveMidStreamRateLimitResetAt } = await import( + '../core/quota-manager' + ) + expect( + resolveMidStreamRateLimitResetAt( + undefined, + rateLimitCalls[0]!.window, + 1_000, + 60_000, + rateLimitCalls[0]!.resetAt, + ), + ).toBe(61_000) + websocketFetch.close() + }, + ) + }) + + test('keeps websocket connection-limit admission errors on HTTP fallback without a rate-limit mark', async () => { + const rateLimitCalls: string[] = [] + let httpRequests = 0 + await withFakeWebSocket( + ({ message }) => ({ + send() { + message( + JSON.stringify({ + type: 'error', + status: 400, + error: { + type: 'invalid_request_error', + code: 'websocket_connection_limit_reached', + message: 'Responses websocket connection limit reached', + }, + }), + ) + }, + }), + async () => { + const httpFetch: typeof globalThis.fetch = Object.assign( + async () => { + httpRequests++ + return new Response('http') + }, + { preconnect: () => {} }, + ) + const websocketFetch = createWebSocketFetch({ + url: 'https://example.test/backend-api/codex/responses', + httpFetch, + onRateLimitReached: (window) => rateLimitCalls.push(window), + }) + const response = await websocketFetch( + 'https://example.test/backend-api/codex/responses', + streamRequest({ input: [] }), + ) + + expect(await response.text()).toBe('http') + expect(httpRequests).toBe(1) + expect(rateLimitCalls).toEqual([]) + websocketFetch.close() + }, + ) + }) + + test('leaves an untyped statusless error on the benign terminal path without a rate-limit mark', async () => { + const rateLimitCalls: string[] = [] + await withFakeWebSocket( + ({ message }) => ({ + send() { + message( + JSON.stringify({ + type: 'error', + error: { message: 'provider rejected request' }, + }), + ) + }, + }), + async () => { + const websocketFetch = createWebSocketFetch({ + url: 'https://example.test/backend-api/codex/responses', + onRateLimitReached: (window) => rateLimitCalls.push(window), + }) + const response = await websocketFetch( + 'https://example.test/backend-api/codex/responses', + streamRequest({ input: [] }), + ) + const text = await response.text() + + expect(text).toContain('provider rejected request') + expect(text).toContain('data: [DONE]') + expect(rateLimitCalls).toEqual([]) + websocketFetch.close() + }, + ) + }) + + test('keeps admission rate-limit reroute alive after a filtered control frame (no emitted output)', async () => { + // A hosted-web-search lifecycle frame is filtered from the SSE output + // (translateHostedWebSearchEvent returns undefined), so it must NOT mark the + // stream as having emitted. If it did, a following admission-time usage limit + // would skip the retryable reroute and silently enqueue the error frame. + const rateLimitCalls: Array<{ window: string; accountId?: string }> = [] + await withFakeWebSocket( + ({ message }) => ({ + send() { + message( + JSON.stringify({ + type: 'response.web_search_call.searching', + item_id: 'ws_1', + }), + ) + message( + JSON.stringify({ + type: 'error', + error: { + type: 'usage_limit_reached', + message: 'The usage limit has been reached', + resets_at: 1_784_958_366, + }, + }), + ) + }, + }), + async () => { + const websocketFetch = createWebSocketFetch({ + url: 'https://example.test/backend-api/codex/responses', + onRateLimitReached: (window, accountId) => { + rateLimitCalls.push({ window, accountId }) + }, + }) + const response = await websocketFetch( + 'https://example.test/backend-api/codex/responses', + { + method: 'POST', + headers: { + 'session-id': 'sess-control-frame-admission', + authorization: 'Bearer tok-main', + 'x-openai-auth-quota-account': 'main', + }, + body: JSON.stringify({ stream: true, input: [] }), + }, + ) + await expect(response.text()).rejects.toMatchObject({ + isRetryable: true, + }) + expect(rateLimitCalls).toEqual([ + { window: 'usage_limit_reached', accountId: 'main' }, + ]) + websocketFetch.close() + }, + ) + }) + + test('falls back to HTTP when a native WebSocket upgrade fails before open', async () => { + // The standard (non-raw) WebSocket API exposes no HTTP status/body on a + // failed upgrade, so an admission 429 there is invisible to + // parseRateLimitSignal. Any upgrade failure on the native client must fall + // back to HTTP, which surfaces the real status for classification. + let httpRequests = 0 + await withFakeWebSocket( + ({ close }) => { + queueMicrotask(() => close(1006, 'upgrade rejected')) + return { autoOpen: false } + }, + async () => { + const httpFetch: typeof globalThis.fetch = Object.assign( + async () => { + httpRequests++ + return new Response('http') + }, + { preconnect: () => {} }, + ) + const websocketFetch = createWebSocketFetch({ + url: 'https://example.test/backend-api/codex/responses', + httpFetch, + }) + const response = await websocketFetch( + 'https://example.test/backend-api/codex/responses', + streamRequest({ input: [] }), + ) + expect(await response.text()).toBe('http') + expect(httpRequests).toBe(1) + websocketFetch.close() + }, + ) + }) + + test('marks and reroutes when the prewarm connection hits an admission rate limit', async () => { + // The first request on a fresh connection prewarms (generate:false) before + // the main turn. A usage limit on the PREWARM connection must mark the + // account and surface a retryable error too — not only on the main + // connection — so the re-issued turn reroutes off the exhausted account. + const rateLimitCalls: Array<{ window: string; accountId?: string }> = [] + const sent: Array> = [] + await withFakeWebSocket( + ({ message }) => ({ + send(data) { + const parsed = JSON.parse(data) as Record + sent.push(parsed) + if (parsed.generate === false) { + message( + JSON.stringify({ + type: 'error', + error: { + type: 'usage_limit_reached', + message: 'The usage limit has been reached', + resets_at: 1_784_958_366, + }, + }), + ) + } + }, + }), + async () => { + const websocketFetch = createWebSocketFetch({ + url: 'https://example.test/backend-api/codex/responses', + onRateLimitReached: (window, accountId) => { + rateLimitCalls.push({ window, accountId }) + }, + }) + const response = await websocketFetch( + 'https://example.test/backend-api/codex/responses', + { + method: 'POST', + headers: { + 'session-id': 'sess-prewarm-admission', + authorization: 'Bearer tok-main', + 'x-openai-auth-quota-account': 'main', + }, + body: JSON.stringify({ + stream: true, + input: [ + { role: 'user', content: [{ type: 'input_text', text: 'hi' }] }, + ], + }), + }, + ) + await expect(response.text()).rejects.toMatchObject({ + isRetryable: true, + }) + expect(rateLimitCalls).toEqual([ + { window: 'usage_limit_reached', accountId: 'main' }, + ]) + // Only the prewarm went out on this account; the main turn never sent. + expect(sent).toHaveLength(1) + expect(sent[0]).toMatchObject({ generate: false }) + websocketFetch.close() + }, + ) + }) + test('does NOT force a retry when output already streamed before a rate-limit failure (no duplicate replay)', async () => { // A rate_limit_reached_type arriving AFTER meaningful output has streamed // must still MARK the account (steer the next turn away) but must NOT error diff --git a/packages/opencode/src/ws-pool.ts b/packages/opencode/src/ws-pool.ts index b1c6e90..0572a5e 100644 --- a/packages/opencode/src/ws-pool.ts +++ b/packages/opencode/src/ws-pool.ts @@ -46,17 +46,12 @@ export interface CreateWebSocketFetchOptions { accountId: string | undefined, servedChatgptAccountId: string | undefined, ) => void - /** - * Fires when a response.failed frame carries a rate_limit_reached_type — - * mid-stream quota exhaustion the reactive HTTP-status fallback cannot see - * on the WebSocket transport (ws.ts synthesizes status:200 at upgrade, - * before generation runs). - * - * Receives the window label plus the per-request internal quota account key - * captured at send time, same as onQuota, so the signal is attributed to - * the connection's own account rather than a shared mutable global. - */ - onRateLimitReached?: (window: string, accountId: string | undefined) => void + /** Receives quota exhaustion plus the account identity captured at send time. */ + onRateLimitReached?: ( + window: string, + accountId: string | undefined, + resetAt?: number, + ) => void } interface PoolEntry { @@ -198,6 +193,14 @@ export function createWebSocketFetch(options?: CreateWebSocketFetchOptions) { ) const acctDisc = accountDiscriminator(sourceHeadersForKey) const key = `${sessionID}:${acctDisc}:conversation` + const requestAccountId = + typeof internalHeaders[QUOTA_ACCOUNT_HEADER] === 'string' + ? internalHeaders[QUOTA_ACCOUNT_HEADER] + : undefined + const requestOnRateLimitReached = onRateLimitReached + ? (window: string, resetAt?: number) => + onRateLimitReached(window, requestAccountId, resetAt) + : undefined const entry = pool.get(key) ?? { lastUsedAt: Date.now(), @@ -233,10 +236,6 @@ export function createWebSocketFetch(options?: CreateWebSocketFetchOptions) { // (the account this request was actually sent on), NOT the wire // chatgpt-account-id header. The internal key was read from the raw init // before internal-header stripping. - const requestAccountId = - typeof internalHeaders[QUOTA_ACCOUNT_HEADER] === 'string' - ? internalHeaders[QUOTA_ACCOUNT_HEADER] - : undefined const requestServedChatgptAccountId = typeof sourceHeaders['chatgpt-account-id'] === 'string' ? sourceHeaders['chatgpt-account-id'] @@ -250,14 +249,6 @@ export function createWebSocketFetch(options?: CreateWebSocketFetchOptions) { requestServedChatgptAccountId, ) : undefined - // Same capture-at-send-time discipline as requestOnQuota: the - // response.failed frame arrives asynchronously, so it must be - // attributed to THIS connection's account, never a shared mutable - // global a concurrent request could have overwritten. - const requestOnRateLimitReached = onRateLimitReached - ? (window: string) => onRateLimitReached(window, requestAccountId) - : undefined - const socketHeaders = !entry.socket && !entry.continuation ? prewarmHeaders(sourceHeaders) @@ -390,6 +381,23 @@ export function createWebSocketFetch(options?: CreateWebSocketFetchOptions) { throw error } + const admissionRateLimit = OpenAIWebSocket.parseRateLimitSignal(error) + if (admissionRateLimit) { + requestOnRateLimitReached?.( + admissionRateLimit.window, + admissionRateLimit.resetAt, + ) + } + + // A native (non-raw) WebSocket exposes no HTTP status/body on a failed + // upgrade, so an admission 429 there is invisible to parseRateLimitSignal. + // Fall back to HTTP — which surfaces the real status — on any upgrade + // failure for the standard client. The raw client classifies its own + // upgrade 429s above and keeps its mark-and-reroute path. + if (OpenAIWebSocket.isUpgradeFailure(error) && !rawWebSocket) { + entry.fallback = true + } + recordStreamFailure(entry) entry.continuation = undefined invalidate(entry) @@ -537,7 +545,7 @@ async function prewarm( sessionID?: string url?: string headers?: HeadersInit - onRateLimitReached?: (window: string) => void + onRateLimitReached?: (window: string, resetAt?: number) => void }, ) { if (!entry.socket || !Array.isArray(body.input) || body.input.length === 0) diff --git a/packages/opencode/src/ws.ts b/packages/opencode/src/ws.ts index 39e420c..27c6232 100644 --- a/packages/opencode/src/ws.ts +++ b/packages/opencode/src/ws.ts @@ -105,8 +105,8 @@ export interface StreamResponsesWebSocketOptions { onAbort?: (error: Error) => void /** Push per-turn quota from a codex.rate_limits in-band frame. */ onQuota?: (s: Record) => void - /** Called when response.failed carries a rate_limit_reached_type — the only mid-stream quota-exhaustion signal on this transport. */ - onRateLimitReached?: (window: string) => void + /** Called when the transport reports quota exhaustion for the current connection. */ + onRateLimitReached?: (window: string, resetAt?: number) => void } export interface WrappedError { @@ -158,6 +158,24 @@ export function isAbortError(error: unknown): error is DOMException { return error instanceof DOMException && error.name === 'AbortError' } +// Errors rejected during the connect/upgrade phase (before the socket opens), +// as opposed to failures mid-stream. The standard WebSocket API exposes no HTTP +// status/body on a failed upgrade, so the pool keys off this marker to fall back +// to HTTP (which surfaces the real status) rather than treat it as a stream error. +const upgradeFailures = new WeakSet() + +export function isUpgradeFailure(error: unknown): boolean { + return ( + typeof error === 'object' && error !== null && upgradeFailures.has(error) + ) +} + +function upgradeFailure(message: string, cause?: unknown): Error { + const error = new Error(message, cause === undefined ? undefined : { cause }) + upgradeFailures.add(error) + return error +} + export function connectResponsesWebSocket( options: ConnectResponsesWebSocketOptions, ) { @@ -222,13 +240,13 @@ export function connectResponsesWebSocket( function onError(error: Event) { cleanup() - reject(new Error(errorMessage(error), { cause: error })) + reject(upgradeFailure(errorMessage(error), error)) } function onClose(event: CloseEvent) { cleanup() reject( - new Error( + upgradeFailure( closeMessage( 'WebSocket closed before open', event.code, @@ -332,6 +350,26 @@ export function streamResponsesWebSocket( return } + const admissionRateLimit = !emitted + ? parseRateLimitSignal(event) + : undefined + if (admissionRateLimit && event) { + completed = true + cleanup() + options.onRateLimitReached?.( + admissionRateLimit.window, + admissionRateLimit.resetAt, + ) + options.onTerminal?.(event) + options.onFirstEvent?.() + controller?.error( + new ResponseStreamError( + `OpenAI account rate limit reached at admission (${admissionRateLimit.window})`, + ), + ) + return + } + if (event?.type === 'error' && options.onRetryableTerminal) { cleanupSocket() if (idleTimer) clearTimeout(idleTimer) @@ -444,8 +482,12 @@ export function streamResponsesWebSocket( const translatedEvent = event ? translateHostedWebSearchEvent(event) : event if (!translatedEvent) { + // Filtered/control frame (e.g. a hosted-web-search lifecycle event) — no + // SSE output is enqueued, so it must NOT set `emitted`. Setting it here + // would block the no-replay reroute for a later admission-time rate limit + // (the `!emitted` gate above). onFirstEvent still fires so the idle timer + // and the pool's first-event gate register the activity. if (!emitted) options.onFirstEvent?.() - emitted = true resetIdleTimeout('idle timeout waiting for websocket') return } @@ -679,6 +721,47 @@ function parseWrappedError( } } +export function parseRateLimitSignal(value: unknown): + | { + window: string + resetAt?: number + } + | undefined { + if (typeof value === 'string') { + try { + return parseRateLimitSignal(JSON.parse(value)) + } catch { + return undefined + } + } + if (!isRecord(value)) return undefined + + const causeSignal = parseRateLimitSignal(value.cause) + if (causeSignal) return causeSignal + const bodySignal = parseRateLimitSignal(value.body) + if (bodySignal) return bodySignal + + const detail = isRecord(value.error) ? value.error : value + const type = typeof detail.type === 'string' ? detail.type : undefined + const status = value.status ?? value.status_code + if (type !== 'usage_limit_reached' && status !== 429) return undefined + + const resetsAtSeconds = detail.resets_at + const resetsInSeconds = detail.resets_in_seconds + const resetAt = + typeof resetsAtSeconds === 'number' && + Number.isFinite(resetsAtSeconds) && + resetsAtSeconds > 0 + ? resetsAtSeconds * 1000 + : typeof resetsInSeconds === 'number' && + Number.isFinite(resetsInSeconds) && + resetsInSeconds > 0 + ? Date.now() + resetsInSeconds * 1000 + : undefined + + return { window: type ?? 'primary', resetAt } +} + function cancelError(reason: unknown) { if (isAbortError(reason)) return reason if (reason instanceof Error) return reason