lmin/zmo-cli
1
0
Fork 0
Code Issues Pull Requests Actions 6 Packages Projects Releases Wiki Activity

fix(openai): retry websocket stream failures (#29673)

Browse Source
This commit is contained in:
Aiden Cline
2026-05-28 01:16:18 -05:00
committed by GitHub
parent e5524f5bf9
commit 14e0b9b17f
7 changed files with 305 additions and 186 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
packages/opencode/src/plugin/openai/README.md
View File

@@ -20,9 +20,9 @@ Enabled by default on `local`, `dev`, and `beta`. On `latest` and `prod`, set `O
## Retries
- If WebSocket setup fails or it fails before its first event, replay over HTTP and keep that session on HTTP until idle-pruned.
- If the server returns `websocket_connection_limit_reached` before output, reconnect up to 5 times, then follow the same HTTP fallback.
- If a WebSocket fails after its first event, fail the stream. Do not replay partial output.
- Retry WebSocket stream/setup failures up to 5 times, then use HTTP for that session until the pool entry is idle-pruned.
- `websocket_connection_limit_reached` consumes the same retry budget and HTTP fallback.
- If a WebSocket fails after its first event, fail it as retryable rather than replaying partial output in transport.
- Abort or cancel closes the socket.
## Next Steps

67
packages/opencode/src/plugin/openai/ws-pool.ts
View File

@@ -1,5 +1,6 @@
import WebSocket from "ws"
import * as Log from "@opencode-ai/core/util/log"
import { ProviderError } from "@/provider/error"
import { isRecord } from "@/util/record"
import { OpenAIWebSocket } from "./ws"
@@ -13,7 +14,7 @@ export interface CreateWebSocketFetchOptions {
connectTimeout?: number
idleTimeout?: number
maxConnectionAge?: number
connectionLimitRetries?: number
streamRetries?: number
}
interface PoolEntry {
@@ -22,6 +23,7 @@ interface PoolEntry {
lastUsedAt: number
busy: boolean
fallback: boolean
streamFailures: number
}
const DEFAULT_CONNECT_TIMEOUT = 15_000
@@ -35,7 +37,7 @@ export function createWebSocketFetch(options?: CreateWebSocketFetchOptions) {
const connectTimeout = options?.connectTimeout ?? DEFAULT_CONNECT_TIMEOUT
const idleTimeout = options?.idleTimeout ?? DEFAULT_IDLE_TIMEOUT
const maxConnectionAge = options?.maxConnectionAge ?? DEFAULT_MAX_CONNECTION_AGE
const connectionLimitRetries = options?.connectionLimitRetries ?? 5
const streamRetries = options?.streamRetries ?? 5
const pruneTimer = setInterval(() => prune(), Math.min(idleTimeout, 60_000))
if (typeof pruneTimer === "object" && "unref" in pruneTimer && typeof pruneTimer.unref === "function") {
pruneTimer.unref()
@@ -72,7 +74,7 @@ export function createWebSocketFetch(options?: CreateWebSocketFetchOptions) {
}
const key = `${sessionID}:conversation`
const entry = pool.get(key) ?? { lastUsedAt: Date.now(), busy: false, fallback: false }
const entry = pool.get(key) ?? { lastUsedAt: Date.now(), busy: false, fallback: false, streamFailures: 0 }
pool.set(key, entry)
if (entry.fallback) {
@@ -87,7 +89,6 @@ export function createWebSocketFetch(options?: CreateWebSocketFetchOptions) {
entry.busy = true
entry.lastUsedAt = Date.now()
try {
let connectionLimitAttempts = 0
entry.socket = await socket(
entry,
options?.url ?? url,
@@ -111,15 +112,16 @@ export function createWebSocketFetch(options?: CreateWebSocketFetchOptions) {
onTerminal: (event) => {
entry.busy = false
entry.lastUsedAt = Date.now()
entry.streamFailures = 0
if (event.type !== "response.completed" && event.type !== "response.done") {
log.warn("websocket terminal failure", { key, type: event.type })
invalidate(entry)
}
},
onConnectionInvalid: (error) => {
log.warn("websocket invalidated", { key, error: error instanceof Error ? error.message : String(error) })
log.warn("websocket invalidated", { key, error: error.message })
entry.busy = false
entry.fallback = true
if (!entry.fallback) recordStreamFailure(entry)
invalidate(entry)
resolveFirstEvent(false)
},
@@ -127,51 +129,52 @@ export function createWebSocketFetch(options?: CreateWebSocketFetchOptions) {
log.debug("websocket aborted", { key })
entry.busy = false
entry.lastUsedAt = Date.now()
entry.streamFailures = 0
invalidate(entry)
rejectFirstEvent(error)
},
onRetryableTerminal: async (event) => {
const error = connectionLimitError(event)
if (!error) return undefined
if (connectionLimitAttempts >= connectionLimitRetries) throw error
connectionLimitAttempts++
log.warn("websocket connection limit reached", { key, attempt: connectionLimitAttempts })
invalidate(entry)
entry.socket = await socket(
entry,
options?.url ?? url,
OpenAIWebSocket.normalizeHeaders(httpInit?.headers),
connectTimeout,
maxConnectionAge,
init?.signal,
)
entry.lastUsedAt = Date.now()
return entry.socket
log.warn("websocket connection limit reached", { key })
throw error
},
})
if (await firstEvent) return response
log.debug("http fallback", { key, reason: "websocket_failed_before_first_event" })
if (!entry.fallback) return response
log.debug("http fallback", { key, reason: "websocket_retries_exhausted" })
return httpFetch(input, httpInit)
} catch (error) {
entry.busy = false
entry.lastUsedAt = Date.now()
if (OpenAIWebSocket.isAbortError(error)) {
entry.streamFailures = 0
invalidate(entry)
throw error
}
entry.fallback = true
recordStreamFailure(entry)
log.warn("websocket setup failed", {
key,
error: error instanceof Error ? error.message : String(error),
fallback: "http",
fallback: entry.fallback ? "http" : undefined,
})
invalidate(entry)
return httpFetch(input, httpInit)
if (entry.fallback) return httpFetch(input, httpInit)
return failedResponse(
new ProviderError.ResponseStreamError(error instanceof Error ? error.message : String(error), {
cause: error,
}),
)
}
}
function recordStreamFailure(entry: PoolEntry) {
entry.streamFailures++
// Codex counts retries after the initial failed WebSocket attempt.
if (entry.streamFailures > streamRetries) entry.fallback = true
}
function prune() {
const now = Date.now()
for (const [key, entry] of pool) {
@@ -198,6 +201,20 @@ function connectionLimitError(event: Record<string, unknown>) {
return new Error(typeof event.error.message === "string" ? event.error.message : CONNECTION_LIMIT_REACHED_CODE)
}
function failedResponse(error: ProviderError.ResponseStreamError) {
return new Response(
new ReadableStream({
start(controller) {
controller.error(error)
},
}),
{
status: 200,
headers: { "content-type": "text/event-stream" },
},
)
}
async function socket(
entry: PoolEntry,
url: string,

29
packages/opencode/src/plugin/openai/ws.ts
View File

@@ -2,6 +2,7 @@
// fallback, and continuation state intentionally live above this file.
import WebSocket from "ws"
import { ProviderError } from "@/provider/error"
export const PROTOCOL_HEADER = "responses_websockets=2026-02-06"
@@ -21,7 +22,7 @@ export interface StreamResponsesWebSocketOptions {
onComplete?: (event: Record<string, unknown>) => void
onTerminal?: (event: Record<string, unknown>) => void
onRetryableTerminal?: (event: Record<string, unknown>) => Promise<WebSocket | undefined>
onConnectionInvalid?: (error: Error) => void
onConnectionInvalid?: (error: ProviderError.ResponseStreamError) => void
onAbort?: (error: Error) => void
}
@@ -101,7 +102,7 @@ export function connectResponsesWebSocket(options: ConnectResponsesWebSocketOpti
function onClose(code: number, reason: Buffer) {
cleanup()
reject(closeError("WebSocket closed before open", code, reason))
reject(new Error(closeMessage("WebSocket closed before open", code, reason)))
}
function onAbort() {
@@ -145,7 +146,7 @@ export function streamResponsesWebSocket(options: StreamResponsesWebSocketOption
controller?.close()
}
function invalidate(error: Error) {
function invalidate(error: ProviderError.ResponseStreamError) {
if (completed) return
completed = true
cleanup()
@@ -157,7 +158,7 @@ export function streamResponsesWebSocket(options: StreamResponsesWebSocketOption
if (completed) return
if (!options.idleTimeout) return
if (idleTimer) clearTimeout(idleTimer)
idleTimer = setTimeout(() => invalidate(new Error(message)), options.idleTimeout)
idleTimer = setTimeout(() => invalidate(new ProviderError.ResponseStreamError(message)), options.idleTimeout)
if (typeof idleTimer === "object" && "unref" in idleTimer && typeof idleTimer.unref === "function") {
idleTimer.unref()
}
@@ -166,7 +167,7 @@ export function streamResponsesWebSocket(options: StreamResponsesWebSocketOption
async function onMessage(data: WebSocket.RawData, isBinary: boolean) {
if (completed) return
if (isBinary) {
invalidate(new Error("Unexpected binary WebSocket frame"))
invalidate(new ProviderError.ResponseStreamError("Unexpected binary WebSocket frame"))
return
}
@@ -195,7 +196,11 @@ export function streamResponsesWebSocket(options: StreamResponsesWebSocketOption
return
}
} catch (error) {
invalidate(error instanceof Error ? error : new Error(String(error)))
invalidate(
new ProviderError.ResponseStreamError(error instanceof Error ? error.message : String(error), {
cause: error,
}),
)
return
}
}
@@ -230,12 +235,14 @@ export function streamResponsesWebSocket(options: StreamResponsesWebSocketOption
}
function onError(error: Error) {
invalidate(error)
invalidate(new ProviderError.ResponseStreamError(error.message, { cause: error }))
}
function onClose(code: number, reason: Buffer) {
if (completed) return
invalidate(closeError("WebSocket closed before response.completed", code, reason))
invalidate(
new ProviderError.ResponseStreamError(closeMessage("WebSocket closed before response.completed", code, reason)),
)
}
function onAbort() {
@@ -272,7 +279,7 @@ export function streamResponsesWebSocket(options: StreamResponsesWebSocketOption
socket.send(JSON.stringify({ type: "response.create", ...payload }), (error) => {
if (completed) return
resetIdleTimeout("idle timeout waiting for websocket")
if (error) invalidate(error)
if (error) invalidate(new ProviderError.ResponseStreamError(error.message, { cause: error }))
})
}
@@ -312,11 +319,11 @@ function abortError(signal: AbortSignal | undefined) {
return new DOMException(reason instanceof Error ? reason.message : "Aborted", "AbortError")
}
function closeError(message: string, code: number, reason: Buffer) {
function closeMessage(message: string, code: number, reason: Buffer) {
const details = [`code ${code}`]
if (code === 1009) details.push("message too big")
if (reason.length > 0) details.push(reason.toString())
return new Error(`${message} (${details.join(": ")})`)
return `${message} (${details.join(": ")})`
}
export * as OpenAIWebSocket from "./ws"