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

packages/ai/src/index.ts

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

packages/ai/src/providers/openai-codex-responses.ts

@@ -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;
+			const response = (event as { response?: { error?: { code?: string; message?: string } } }).response;
-			throw new Error(msg || "Codex response failed");
+			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);
+							yield JSON.parse(data) as Record<string, unknown>;
-						} catch {}
+						} 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) {
+	if (event && typeof event === "object") {
-		const message = (event as { message?: unknown }).message;
+		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}` : "";
+		let reasonText = typeof reason === "string" && reason.length > 0 ? ` ${reason}` : "";
-		return new Error(`WebSocket closed${codeText}${reasonText}`.trim());
+		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;
+			let text: string | null = null;
-			const text = await decodeWebSocketData((event as { data?: unknown }).data);
-			if (!text) return;
 			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();
+		await processResponsesStream(
-		stream.push({ type: "start", partial: output });
+			startWebSocketOutputOnFirstEvent(
-		await processResponsesStream(mapCodexEvents(parseWebSocket(socket, options?.signal)), output, stream, model, {
+				mapCodexEvents(parseWebSocket(socket, options?.signal)),
-			serviceTier: options?.serviceTier,
+				output,
-			resolveServiceTier: resolveCodexServiceTier,
+				stream,
-			applyServiceTierPricing: (usage, serviceTier) => applyServiceTierPricing(usage, serviceTier, model),
+				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) {

packages/ai/src/types.ts

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

packages/ai/src/utils/diagnostics.ts

@@ -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];
+}

packages/coding-agent/src/core/agent-session.ts

@@ -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
+		// 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|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(
+		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,
 		);
 	}