shumengya/sproutclaw
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(ai): fall back from codex websocket to sse (#4133)

Browse Source
This commit is contained in:
Armin Ronacher
2026-05-03 22:51:42 +02:00
committed by Mario Zechner
parent 0959d8ec2b
commit ce377fc4a1
5 changed files with 233 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
packages/ai/src/index.ts
View File

@@ -26,6 +26,7 @@ export type { OpenAIResponsesOptions } from "./providers/openai-responses.js";
export * from "./providers/register-builtins.js";
export * from "./stream.js";
export * from "./types.js";
export * from "./utils/diagnostics.js";
export * from "./utils/event-stream.js";
export * from "./utils/json-parse.js";
export type {

205
packages/ai/src/providers/openai-codex-responses.ts
View File

@@ -32,6 +32,11 @@ import type {
StreamOptions,
Usage,
} from "../types.js";
import {
appendAssistantMessageDiagnostic,
createAssistantMessageDiagnostic,
formatThrownValue,
} from "../utils/diagnostics.js";
import { AssistantMessageEventStream } from "../utils/event-stream.js";
import { headersToRecord } from "../utils/headers.js";
import { convertResponsesMessages, convertResponsesTools, processResponsesStream } from "./openai-responses-shared.js";
@@ -46,6 +51,7 @@ const JWT_CLAIM_PATH = "https://api.openai.com/auth" as const;
const MAX_RETRIES = 3;
const BASE_DELAY_MS = 1000;
const CODEX_TOOL_CALL_PROVIDERS = new Set(["openai", "openai-codex", "opencode"]);
const WEBSOCKET_MESSAGE_TOO_BIG_CLOSE_CODE = 1009;
const CODEX_RESPONSE_STATUSES = new Set<CodexResponseStatus>([
"completed",
@@ -166,8 +172,12 @@ export const streamOpenAICodexResponses: StreamFunction<"openai-codex-responses"
);
const bodyJson = JSON.stringify(body);
const transport = options?.transport || "auto";
const websocketDisabledForSession = transport !== "sse" && isWebSocketSseFallbackActive(options?.sessionId);
if (websocketDisabledForSession) {
recordWebSocketSseFallback(options?.sessionId);
}
if (transport !== "sse") {
if (transport !== "sse" && !websocketDisabledForSession) {
let websocketStarted = false;
try {
await processWebSocketStream(
@@ -194,9 +204,25 @@ export const streamOpenAICodexResponses: StreamFunction<"openai-codex-responses"
stream.end();
return;
} catch (error) {
if (transport === "websocket" || transport === "websocket-cached" || websocketStarted) {
const aborted = options?.signal?.aborted;
if (aborted || isCodexNonTransportError(error)) {
throw error;
}
appendAssistantMessageDiagnostic(
output,
createAssistantMessageDiagnostic("provider_transport_failure", error, {
configuredTransport: transport,
fallbackTransport: websocketStarted ? undefined : "sse",
eventsEmitted: websocketStarted,
phase: websocketStarted ? "after_message_stream_start" : "before_message_stream_start",
requestBytes: new TextEncoder().encode(bodyJson).byteLength,
}),
);
recordWebSocketFailure(options?.sessionId, error);
if (websocketStarted) {
throw error;
}
recordWebSocketSseFallback(options?.sessionId);
}
}
@@ -434,6 +460,34 @@ async function processStream(
});
}
class CodexApiError extends Error {
readonly code?: string;
readonly payload?: Record<string, unknown>;
constructor(message: string, options?: { code?: string; payload?: Record<string, unknown>; cause?: unknown }) {
super(message);
this.name = "CodexApiError";
this.code = options?.code;
this.payload = options?.payload;
this.cause = options?.cause;
}
}
class CodexProtocolError extends Error {
readonly payload?: unknown;
constructor(message: string, options?: { payload?: unknown; cause?: unknown }) {
super(message);
this.name = "CodexProtocolError";
this.payload = options?.payload;
this.cause = options?.cause;
}
}
function isCodexNonTransportError(error: unknown): boolean {
return error instanceof CodexApiError || error instanceof CodexProtocolError;
}
async function* mapCodexEvents(events: AsyncIterable<Record<string, unknown>>): AsyncGenerator<ResponseStreamEvent> {
for await (const event of events) {
const type = typeof event.type === "string" ? event.type : undefined;
@@ -442,12 +496,17 @@ async function* mapCodexEvents(events: AsyncIterable<Record<string, unknown>>):
if (type === "error") {
const code = (event as { code?: string }).code || "";
const message = (event as { message?: string }).message || "";
throw new Error(`Codex error: ${message || code || JSON.stringify(event)}`);
throw new CodexApiError(`Codex error: ${message || code || JSON.stringify(event)}`, {
code: code || undefined,
payload: event,
});
}
if (type === "response.failed") {
const msg = (event as { response?: { error?: { message?: string } } }).response?.error?.message;
throw new Error(msg || "Codex response failed");
const response = (event as { response?: { error?: { code?: string; message?: string } } }).response;
const code = response?.error?.code;
const message = response?.error?.message;
throw new CodexApiError(message || "Codex response failed", { code, payload: event });
}
if (type === "response.done" || type === "response.completed" || type === "response.incomplete") {
@@ -498,8 +557,13 @@ async function* parseSSE(response: Response): AsyncGenerator<Record<string, unkn
const data = dataLines.join("\n").trim();
if (data && data !== "[DONE]") {
try {
yield JSON.parse(data);
} catch {}
yield JSON.parse(data) as Record<string, unknown>;
} catch (cause) {
throw new CodexProtocolError(`Invalid Codex SSE JSON: ${formatThrownValue(cause)}`, {
cause,
payload: data,
});
}
}
}
idx = buffer.indexOf("\n\n");
@@ -556,10 +620,15 @@ export interface OpenAICodexWebSocketDebugStats {
lastInputItems: number;
lastDeltaInputItems?: number;
lastPreviousResponseId?: string;
websocketFailures: number;
sseFallbacks: number;
websocketFallbackActive?: boolean;
lastWebSocketError?: string;
}
const websocketSessionCache = new Map<string, CachedWebSocketConnection>();
const websocketDebugStats = new Map<string, OpenAICodexWebSocketDebugStats>();
const websocketSseFallbackSessions = new Set<string>();
function getOrCreateWebSocketDebugStats(sessionId: string): OpenAICodexWebSocketDebugStats {
let stats = websocketDebugStats.get(sessionId);
@@ -573,6 +642,8 @@ function getOrCreateWebSocketDebugStats(sessionId: string): OpenAICodexWebSocket
fullContextRequests: 0,
deltaRequests: 0,
lastInputItems: 0,
websocketFailures: 0,
sseFallbacks: 0,
};
websocketDebugStats.set(sessionId, stats);
}
@@ -587,9 +658,11 @@ export function getOpenAICodexWebSocketDebugStats(sessionId: string): OpenAICode
export function resetOpenAICodexWebSocketDebugStats(sessionId?: string): void {
if (sessionId) {
websocketDebugStats.delete(sessionId);
websocketSseFallbackSessions.delete(sessionId);
return;
}
websocketDebugStats.clear();
websocketSseFallbackSessions.clear();
}
export function closeOpenAICodexWebSocketSessions(sessionId?: string): void {
@@ -609,6 +682,27 @@ export function closeOpenAICodexWebSocketSessions(sessionId?: string): void {
websocketSessionCache.clear();
}
function isWebSocketSseFallbackActive(sessionId: string | undefined): boolean {
return sessionId ? websocketSseFallbackSessions.has(sessionId) : false;
}
function recordWebSocketSseFallback(sessionId: string | undefined): void {
if (!sessionId) return;
const stats = getOrCreateWebSocketDebugStats(sessionId);
stats.sseFallbacks++;
stats.websocketFallbackActive = isWebSocketSseFallbackActive(sessionId);
}
function recordWebSocketFailure(sessionId: string | undefined, error: unknown): void {
if (!sessionId) return;
websocketSseFallbackSessions.add(sessionId);
const stats = getOrCreateWebSocketDebugStats(sessionId);
stats.websocketFailures++;
stats.lastWebSocketError = formatThrownValue(error);
stats.websocketFallbackActive = true;
}
type WebSocketConstructor = new (
url: string,
protocols?: string | string[] | { headers?: Record<string, string> },
@@ -620,6 +714,20 @@ function getWebSocketConstructor(): WebSocketConstructor | null {
return ctor as unknown as WebSocketConstructor;
}
class WebSocketCloseError extends Error {
readonly code?: number;
readonly reason?: string;
readonly wasClean?: boolean;
constructor(message: string, options?: { code?: number; reason?: string; wasClean?: boolean }) {
super(message);
this.name = "WebSocketCloseError";
this.code = options?.code;
this.reason = options?.reason;
this.wasClean = options?.wasClean;
}
}
function getWebSocketReadyState(socket: WebSocketLike): number | undefined {
const readyState = (socket as { readyState?: unknown }).readyState;
return typeof readyState === "number" ? readyState : undefined;
@@ -798,11 +906,22 @@ async function acquireWebSocket(
}
function extractWebSocketError(event: unknown): Error {
if (event && typeof event === "object" && "message" in event) {
const message = (event as { message?: unknown }).message;
if (event && typeof event === "object") {
const message = "message" in event ? (event as { message?: unknown }).message : undefined;
if (typeof message === "string" && message.length > 0) {
return new Error(message);
}
const nestedError = "error" in event ? (event as { error?: unknown }).error : undefined;
if (nestedError instanceof Error && nestedError.message.length > 0) {
return nestedError;
}
if (nestedError && typeof nestedError === "object" && "message" in nestedError) {
const nestedMessage = (nestedError as { message?: unknown }).message;
if (typeof nestedMessage === "string" && nestedMessage.length > 0) {
return new Error(nestedMessage);
}
}
}
return new Error("WebSocket error");
}
@@ -811,9 +930,17 @@ function extractWebSocketCloseError(event: unknown): Error {
if (event && typeof event === "object") {
const code = "code" in event ? (event as { code?: unknown }).code : undefined;
const reason = "reason" in event ? (event as { reason?: unknown }).reason : undefined;
const wasClean = "wasClean" in event ? (event as { wasClean?: unknown }).wasClean : undefined;
const codeText = typeof code === "number" ? ` ${code}` : "";
const reasonText = typeof reason === "string" && reason.length > 0 ? ` ${reason}` : "";
return new Error(`WebSocket closed${codeText}${reasonText}`.trim());
let reasonText = typeof reason === "string" && reason.length > 0 ? ` ${reason}` : "";
if (!reasonText && code === WEBSOCKET_MESSAGE_TOO_BIG_CLOSE_CODE) {
reasonText = " message too big";
}
return new WebSocketCloseError(`WebSocket closed${codeText}${reasonText}`.trim(), {
code: typeof code === "number" ? code : undefined,
reason: typeof reason === "string" && reason.length > 0 ? reason : undefined,
wasClean: typeof wasClean === "boolean" ? wasClean : undefined,
});
}
return new Error("WebSocket closed");
}
@@ -851,10 +978,11 @@ async function* parseWebSocket(socket: WebSocketLike, signal?: AbortSignal): Asy
const onMessage: WebSocketListener = (event) => {
void (async () => {
if (!event || typeof event !== "object" || !("data" in event)) return;
const text = await decodeWebSocketData((event as { data?: unknown }).data);
if (!text) return;
let text: string | null = null;
try {
if (!event || typeof event !== "object" || !("data" in event)) return;
text = await decodeWebSocketData((event as { data?: unknown }).data);
if (!text) return;
const parsed = JSON.parse(text) as Record<string, unknown>;
const type = typeof parsed.type === "string" ? parsed.type : "";
if (type === "response.completed" || type === "response.done" || type === "response.incomplete") {
@@ -863,7 +991,14 @@ async function* parseWebSocket(socket: WebSocketLike, signal?: AbortSignal): Asy
}
queue.push(parsed);
wake();
} catch {}
} catch (cause) {
failed = new CodexProtocolError(`Invalid Codex WebSocket JSON: ${formatThrownValue(cause)}`, {
cause,
payload: text,
});
done = true;
wake();
}
})();
};
@@ -980,6 +1115,23 @@ function buildCachedWebSocketRequestBody(entry: CachedWebSocketConnection, body:
};
}
async function* startWebSocketOutputOnFirstEvent(
events: AsyncIterable<ResponseStreamEvent>,
output: AssistantMessage,
stream: AssistantMessageEventStream,
onStart: () => void,
): AsyncGenerator<ResponseStreamEvent> {
let started = false;
for await (const event of events) {
if (!started) {
started = true;
onStart();
stream.push({ type: "start", partial: output });
}
yield event;
}
}
async function processWebSocketStream(
url: string,
body: RequestBody,
@@ -1017,13 +1169,22 @@ async function processWebSocketStream(
}
try {
socket.send(JSON.stringify({ type: "response.create", ...requestBody }));
onStart();
stream.push({ type: "start", partial: output });
await processResponsesStream(mapCodexEvents(parseWebSocket(socket, options?.signal)), output, stream, model, {
serviceTier: options?.serviceTier,
resolveServiceTier: resolveCodexServiceTier,
applyServiceTierPricing: (usage, serviceTier) => applyServiceTierPricing(usage, serviceTier, model),
});
await processResponsesStream(
startWebSocketOutputOnFirstEvent(
mapCodexEvents(parseWebSocket(socket, options?.signal)),
output,
stream,
onStart,
),
output,
stream,
model,
{
serviceTier: options?.serviceTier,
resolveServiceTier: resolveCodexServiceTier,
applyServiceTierPricing: (usage, serviceTier) => applyServiceTierPricing(usage, serviceTier, model),
},
);
if (options?.signal?.aborted) {
keepConnection = false;
} else if (useCachedContext && entry && output.responseId) {

2
packages/ai/src/types.ts
View File

@@ -1,3 +1,4 @@
import type { AssistantMessageDiagnostic } from "./utils/diagnostics.js";
import type { AssistantMessageEventStream } from "./utils/event-stream.js";
export type { AssistantMessageEventStream } from "./utils/event-stream.js";
@@ -224,6 +225,7 @@ export interface AssistantMessage {
model: string;
responseModel?: string; // Concrete `chunk.model` when different from the requested `model` (e.g. OpenRouter `auto` -> `anthropic/...`)
responseId?: string; // Provider-specific response/message identifier when the upstream API exposes one
diagnostics?: AssistantMessageDiagnostic[]; // Redacted provider/runtime diagnostics for failures and recoveries.
usage: Usage;
stopReason: StopReason;
errorMessage?: string;

45
packages/ai/src/utils/diagnostics.ts Normal file
View File

@@ -0,0 +1,45 @@
export interface DiagnosticErrorInfo {
name?: string;
message: string;
stack?: string;
code?: string | number;
}
export interface AssistantMessageDiagnostic {
type: string;
timestamp: number;
error?: DiagnosticErrorInfo;
details?: Record<string, unknown>;
}
export function formatThrownValue(value: unknown): string {
if (value instanceof Error) return value.message || value.name;
if (typeof value === "string") return value;
return String(value);
}
export function extractDiagnosticError(error: unknown): DiagnosticErrorInfo {
if (!(error instanceof Error)) return { name: "ThrownValue", message: formatThrownValue(error) };
const code = (error as Error & { code?: unknown }).code;
return {
name: error.name || undefined,
message: error.message || error.name,
stack: error.stack,
code: typeof code === "string" || typeof code === "number" ? code : undefined,
};
}
export function createAssistantMessageDiagnostic(
type: string,
error: unknown,
details?: Record<string, unknown>,
): AssistantMessageDiagnostic {
return { type, timestamp: Date.now(), error: extractDiagnosticError(error), details };
}
export function appendAssistantMessageDiagnostic<T extends { diagnostics?: AssistantMessageDiagnostic[] }>(
message: T,
diagnostic: AssistantMessageDiagnostic,
): void {
message.diagnostics = [...(message.diagnostics ?? []), diagnostic];
}

4
packages/coding-agent/src/core/agent-session.ts
View File

@@ -2417,8 +2417,8 @@ export class AgentSession {
if (isContextOverflow(message, contextWindow)) return false;
const err = message.errorMessage;
// Match: overloaded_error, provider returned error, rate limit, 429, 500, 502, 503, 504, service unavailable, network/connection errors (including connection lost), fetch failed, request ended without sending chunks, HTTP/2 closed before response, terminated, retry delay exceeded
return /overloaded|provider.?returned.?error|rate.?limit|too many requests|429|500|502|503|504|service.?unavailable|server.?error|internal.?error|network.?error|connection.?error|connection.?refused|connection.?lost|other side closed|fetch failed|upstream.?connect|reset before headers|socket hang up|ended without|http2 request did not get a response|timed? out|timeout|terminated|retry delay/i.test(
// Match: overloaded_error, provider returned error, rate limit, 429, 500, 502, 503, 504, service unavailable, network/connection errors (including connection lost), WebSocket transport closes/errors, fetch failed, request ended without sending chunks, HTTP/2 closed before response, terminated, retry delay exceeded
return /overloaded|provider.?returned.?error|rate.?limit|too many requests|429|500|502|503|504|service.?unavailable|server.?error|internal.?error|network.?error|connection.?error|connection.?refused|connection.?lost|websocket.?closed|websocket.?error|other side closed|fetch failed|upstream.?connect|reset before headers|socket hang up|ended without|http2 request did not get a response|timed? out|timeout|terminated|retry delay/i.test(
err,
);
}