pi-openai-codex-compat 0.0.2 → 0.0.3

package/extensions/openai-codex-compat/codex-transport.ts

@@ -6,9 +6,11 @@ import {
   type OpenAICodexResponsesOptions,
   type ProviderEnv,
   type ProviderHeaders,
+  uuidv7,
 } from "@earendil-works/pi-ai";
+import { codexCacheKey } from "./codex-cache-key.ts";
 import { isObject, type JsonRecord } from "./codex-protocol.ts";
-import { normalizeReplayItem, replayItemsEqual, stableResponsesJson } from "./responses-replay.ts";
+import { normalizeReplayItem } from "./responses-replay.ts";
 
 /**
  * Focused adaptation of @earendil-works/pi-ai@0.83.0
@@ -19,10 +21,13 @@ const DEFAULT_CODEX_BASE_URL = "https://chatgpt.com/backend-api";
 const OPENAI_BETA_RESPONSES_WEBSOCKETS = "responses_websockets=2026-02-06";
 const DEFAULT_MAX_RETRIES = 0;
 const BASE_DELAY_MS = 1_000;
+const DEFAULT_MAX_RETRY_DELAY_MS = 60_000;
 const REQUEST_COMPRESSION_ZSTD_LEVEL = 3;
 const DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS = 15_000;
 const SESSION_WEBSOCKET_CACHE_TTL_MS = 5 * 60 * 1_000;
 const SESSION_WEBSOCKET_MAX_AGE_MS = 55 * 60 * 1_000;
+const WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE = "websocket_connection_limit_reached";
+const PREVIOUS_RESPONSE_NOT_FOUND_CODE = "previous_response_not_found";
 
 type ProcessWithBuiltinModules = typeof process & {
   getBuiltinModule?: {
@@ -31,10 +36,6 @@ type ProcessWithBuiltinModules = typeof process & {
   };
 };
 
-type CodexTransportOptions = OpenAICodexResponsesOptions & {
-  env?: ProviderEnv;
-};
-
 export type CodexJsonRequestOptions = {
   apiKey: string;
   headers?: ProviderHeaders;
@@ -43,6 +44,60 @@ export type CodexJsonRequestOptions = {
   fetch?: typeof fetch;
 };
 
+export type CodexTransportDiagnostic = {
+  type: "provider_transport_failure";
+  timestamp: number;
+  error: {
+    name?: string;
+    message: string;
+    stack?: string;
+    code?: string | number;
+  };
+  details: {
+    configuredTransport: string;
+    fallbackTransport?: "sse";
+    eventsEmitted: boolean;
+    phase: "before_message_stream_start" | "after_message_stream_start";
+    requestBytes: number;
+  };
+};
+
+export type CodexContinuationHandle = {
+  readonly responseId: string;
+  replaceResponseItems(items: readonly JsonRecord[]): boolean;
+};
+
+export type CodexWebSocketResponseHandle = {
+  discard(): boolean;
+  failParsing(error: unknown): boolean;
+};
+
+export interface OpenAICodexWebSocketDebugStats {
+  requests: number;
+  connectionsCreated: number;
+  connectionsReused: number;
+  cachedContextRequests: number;
+  storeTrueRequests: number;
+  fullContextRequests: number;
+  deltaRequests: number;
+  lastInputItems: number;
+  lastDeltaInputItems?: number;
+  lastPreviousResponseId?: string;
+  websocketFailures: number;
+  sseFallbacks: number;
+  websocketFallbackActive?: boolean;
+  lastWebSocketError?: string;
+}
+
+type CodexTransportOptions = OpenAICodexResponsesOptions & {
+  accountId?: string;
+  env?: ProviderEnv;
+  onContinuationReady?(handle: CodexContinuationHandle): void;
+  onWebSocketResponseHandle?(handle: CodexWebSocketResponseHandle): void;
+  onTransportStart?(): void;
+  onTransportDiagnostic?(diagnostic: CodexTransportDiagnostic): void;
+};
+
 type WebSocketEventType = "open" | "message" | "error" | "close";
 type WebSocketListener = (event: unknown) => void;
 
@@ -71,10 +126,132 @@ type CachedWebSocket = {
   };
 };
 
-const websocketSessions = new Map<string, CachedWebSocket>();
+const websocketSessions = new Map<string, Map<string, CachedWebSocket>>();
 const websocketFallbackSessions = new Set<string>();
+const websocketDebugStats = new Map<string, OpenAICodexWebSocketDebugStats>();
+
+function getOrCreateWebSocketDebugStats(sessionId: string): OpenAICodexWebSocketDebugStats {
+  let stats = websocketDebugStats.get(sessionId);
+  if (!stats) {
+    stats = {
+      requests: 0,
+      connectionsCreated: 0,
+      connectionsReused: 0,
+      cachedContextRequests: 0,
+      storeTrueRequests: 0,
+      fullContextRequests: 0,
+      deltaRequests: 0,
+      lastInputItems: 0,
+      websocketFailures: 0,
+      sseFallbacks: 0,
+    };
+    websocketDebugStats.set(sessionId, stats);
+  }
+  return stats;
+}
+
+export function getOpenAICodexWebSocketDebugStats(
+  sessionId: string,
+): OpenAICodexWebSocketDebugStats | undefined {
+  const stats = websocketDebugStats.get(sessionId);
+  return stats ? { ...stats } : undefined;
+}
+
+export function resetOpenAICodexWebSocketDebugStats(sessionId?: string): void {
+  if (sessionId) {
+    websocketDebugStats.delete(sessionId);
+    websocketFallbackSessions.delete(sessionId);
+    return;
+  }
+  websocketDebugStats.clear();
+  websocketFallbackSessions.clear();
+}
 
-class CodexResponseError extends Error {}
+export function closeOpenAICodexWebSocketSessions(sessionId?: string): void {
+  const closeEntry = (entry: CachedWebSocket): void => {
+    if (entry.idleTimer) clearTimeout(entry.idleTimer);
+    closeSocket(entry.socket, "debug_close");
+  };
+  if (sessionId) {
+    for (const entry of websocketSessions.get(sessionId)?.values() ?? []) closeEntry(entry);
+    websocketSessions.delete(sessionId);
+    websocketFallbackSessions.delete(sessionId);
+    const stats = websocketDebugStats.get(sessionId);
+    if (stats?.websocketFallbackActive !== undefined) stats.websocketFallbackActive = false;
+    return;
+  }
+  for (const accountEntries of websocketSessions.values()) {
+    for (const entry of accountEntries.values()) closeEntry(entry);
+  }
+  websocketSessions.clear();
+  websocketFallbackSessions.clear();
+  for (const stats of websocketDebugStats.values()) {
+    if (stats.websocketFallbackActive !== undefined) stats.websocketFallbackActive = false;
+  }
+}
+
+function isWebSocketSseFallbackActive(sessionId: string | undefined): boolean {
+  return sessionId ? websocketFallbackSessions.has(sessionId) : false;
+}
+
+function recordWebSocketSseFallback(sessionId: string | undefined): void {
+  if (!sessionId) return;
+  const stats = getOrCreateWebSocketDebugStats(sessionId);
+  stats.sseFallbacks += 1;
+  stats.websocketFallbackActive = isWebSocketSseFallbackActive(sessionId);
+}
+
+function recordWebSocketFailure(sessionId: string | undefined, error: unknown): void {
+  if (!sessionId) return;
+  websocketFallbackSessions.add(sessionId);
+  const stats = getOrCreateWebSocketDebugStats(sessionId);
+  stats.websocketFailures += 1;
+  stats.lastWebSocketError = thrownMessage(error);
+  stats.websocketFallbackActive = true;
+}
+
+class CodexApiError extends Error {
+  readonly code: string | undefined;
+  readonly payload: JsonRecord;
+
+  constructor(message: string, code: string | undefined, payload: JsonRecord) {
+    super(message);
+    this.name = "CodexApiError";
+    this.code = code;
+    this.payload = payload;
+  }
+}
+
+class CodexProtocolError extends Error {
+  readonly payload: unknown;
+
+  constructor(message: string, payload: unknown, cause: unknown) {
+    super(message, { cause });
+    this.name = "CodexProtocolError";
+    this.payload = payload;
+  }
+}
+
+class WebSocketCloseError extends Error {
+  readonly code: number | undefined;
+  readonly reason: string | undefined;
+  readonly wasClean: boolean | undefined;
+
+  constructor(
+    message: string,
+    options: {
+      code: number | undefined;
+      reason: string | undefined;
+      wasClean: boolean | undefined;
+    },
+  ) {
+    super(message);
+    this.name = "WebSocketCloseError";
+    this.code = options.code;
+    this.reason = options.reason;
+    this.wasClean = options.wasClean;
+  }
+}
 
 function nodeOs(): typeof NodeOs | undefined {
   const currentProcess = process as ProcessWithBuiltinModules;
@@ -86,8 +263,94 @@ function nodeZlib(): typeof NodeZlib | undefined {
   return currentProcess.getBuiltinModule?.("node:zlib");
 }
 
-function explain(error: unknown): string {
-  return error instanceof Error ? error.message : String(error);
+function extractWebSocketError(event: unknown): Error {
+  if (isObject(event)) {
+    if (typeof event["message"] === "string" && event["message"].length > 0) {
+      return new Error(event["message"]);
+    }
+    const nestedError = event["error"];
+    if (nestedError instanceof Error && nestedError.message.length > 0) return nestedError;
+    if (
+      isObject(nestedError) &&
+      typeof nestedError["message"] === "string" &&
+      nestedError["message"].length > 0
+    ) {
+      return new Error(nestedError["message"]);
+    }
+  }
+  return new Error("WebSocket error");
+}
+
+function extractWebSocketCloseError(
+  event: unknown,
+  context = "WebSocket closed",
+): WebSocketCloseError {
+  const code = isObject(event) && typeof event["code"] === "number" ? event["code"] : undefined;
+  let reason =
+    isObject(event) && typeof event["reason"] === "string" && event["reason"].length > 0
+      ? event["reason"]
+      : undefined;
+  if (reason === undefined && code === 1_009) reason = "message too big";
+  const wasClean =
+    isObject(event) && typeof event["wasClean"] === "boolean" ? event["wasClean"] : undefined;
+  const details = [
+    code === undefined ? undefined : `code ${code}`,
+    reason === undefined ? undefined : `reason: ${reason}`,
+    wasClean === undefined ? undefined : `wasClean: ${String(wasClean)}`,
+  ].filter((detail): detail is string => detail !== undefined);
+  return new WebSocketCloseError(
+    details.length > 0 ? `${context} (${details.join(", ")})` : context,
+    { code, reason, wasClean },
+  );
+}
+
+function thrownMessage(error: unknown): string {
+  return error instanceof Error ? error.message || error.name : String(error);
+}
+
+function isCodexNonTransportError(error: unknown): boolean {
+  return error instanceof CodexApiError || error instanceof CodexProtocolError;
+}
+
+function isWebSocketConnectionLimitReachedError(error: unknown): boolean {
+  return error instanceof CodexApiError && error.code === WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE;
+}
+
+function isPreviousResponseNotFoundError(error: unknown): boolean {
+  return error instanceof CodexApiError && error.code === PREVIOUS_RESPONSE_NOT_FOUND_CODE;
+}
+
+function diagnosticError(error: unknown): CodexTransportDiagnostic["error"] {
+  if (!(error instanceof Error)) {
+    return { name: "ThrownValue", message: thrownMessage(error) };
+  }
+  const code = (error as Error & { code?: unknown }).code;
+  return {
+    ...(error.name ? { name: error.name } : {}),
+    message: error.message || error.name,
+    ...(error.stack ? { stack: error.stack } : {}),
+    ...(typeof code === "string" || typeof code === "number" ? { code } : {}),
+  };
+}
+
+function transportDiagnostic(
+  error: unknown,
+  transport: string,
+  emitted: boolean,
+  requestBytes: number,
+): CodexTransportDiagnostic {
+  return {
+    type: "provider_transport_failure",
+    timestamp: Date.now(),
+    error: diagnosticError(error),
+    details: {
+      configuredTransport: transport,
+      ...(!emitted ? { fallbackTransport: "sse" as const } : {}),
+      eventsEmitted: emitted,
+      phase: emitted ? "after_message_stream_start" : "before_message_stream_start",
+      requestBytes,
+    },
+  };
 }
 
 function sleep(milliseconds: number, signal?: AbortSignal): Promise<void> {
@@ -108,6 +371,79 @@ function sleep(milliseconds: number, signal?: AbortSignal): Promise<void> {
   });
 }
 
+function normalizeTimeoutMs(value: number | undefined): number | undefined {
+  if (value === undefined) return undefined;
+  if (!Number.isFinite(value) || value < 0) {
+    throw new Error(`Invalid timeoutMs: ${String(value)}`);
+  }
+  return Math.floor(value);
+}
+
+function isTerminalRateLimitError(errorText: string): boolean {
+  return /GoUsageLimitError|FreeUsageLimitError|Monthly usage limit reached|available balance|insufficient_quota|out of budget|quota exceeded|billing/i.test(
+    errorText,
+  );
+}
+
+function retryDelayMs(headers: Headers): number | undefined {
+  const retryAfterMs = headers.get("retry-after-ms");
+  if (retryAfterMs !== null) {
+    const milliseconds = Number(retryAfterMs);
+    if (Number.isFinite(milliseconds)) return Math.max(0, milliseconds);
+  }
+
+  const retryAfter = headers.get("retry-after");
+  if (!retryAfter) return undefined;
+  const seconds = Number(retryAfter);
+  if (Number.isFinite(seconds)) return Math.max(0, seconds * 1_000);
+  const date = Date.parse(retryAfter);
+  return Number.isNaN(date) ? undefined : Math.max(0, date - Date.now());
+}
+
+class RetryDelayExceededError extends Error {}
+
+function validateRetryDelay(delayMs: number, options: CodexTransportOptions): number {
+  const maximum = options.maxRetryDelayMs ?? DEFAULT_MAX_RETRY_DELAY_MS;
+  if (maximum > 0 && delayMs > maximum) {
+    throw new RetryDelayExceededError(
+      `Server requested ${Math.ceil(delayMs / 1_000)}s retry delay (max: ${Math.ceil(maximum / 1_000)}s)`,
+    );
+  }
+  return delayMs;
+}
+
+function codexHttpError(status: number, statusText: string, raw: string): Error {
+  let message = raw || statusText || "Request failed";
+  let friendlyMessage: string | undefined;
+  try {
+    const parsed = JSON.parse(raw) as {
+      error?: {
+        code?: string;
+        type?: string;
+        message?: string;
+        plan_type?: string;
+        resets_at?: number;
+      };
+    };
+    const error = parsed?.error;
+    if (!error) return new Error(message);
+    const code = error.code || error.type || "";
+    if (
+      status === 429 ||
+      /usage_limit_reached|usage_not_included|rate_limit_exceeded/i.test(code)
+    ) {
+      const plan = error.plan_type ? ` (${error.plan_type.toLowerCase()} plan)` : "";
+      const resetMinutes = error.resets_at
+        ? Math.max(0, Math.round((error.resets_at * 1_000 - Date.now()) / 60_000))
+        : undefined;
+      const reset = resetMinutes === undefined ? "" : ` Try again in ~${String(resetMinutes)} min.`;
+      friendlyMessage = `You have hit your ChatGPT usage limit${plan}.${reset}`.trim();
+    }
+    message = error.message || friendlyMessage || message;
+  } catch {}
+  return new Error(friendlyMessage || message);
+}
+
 function combineAbortSignals(signals: readonly (AbortSignal | undefined)[]): {
   signal?: AbortSignal;
   cleanup(): void;
@@ -150,7 +486,11 @@ function extractAccountId(token: string): string {
     if (parts.length !== 3) throw new Error("Invalid token");
     const payload = JSON.parse(Buffer.from(parts[1]!, "base64url").toString("utf8")) as JsonRecord;
     const authentication = payload["https://api.openai.com/auth"];
-    if (!isObject(authentication) || typeof authentication["chatgpt_account_id"] !== "string") {
+    if (
+      !isObject(authentication) ||
+      typeof authentication["chatgpt_account_id"] !== "string" ||
+      authentication["chatgpt_account_id"].length === 0
+    ) {
       throw new Error("No account ID");
     }
     return authentication["chatgpt_account_id"];
@@ -159,6 +499,11 @@ function extractAccountId(token: string): string {
   }
 }
 
+export function validateCodexAuthentication(model: Model<any>, apiKey: string | undefined): string {
+  if (!apiKey) throw new Error(`No API key for provider: ${model.provider}`);
+  return extractAccountId(apiKey);
+}
+
 function resolveCodexUrl(baseUrl?: string): string {
   const raw = baseUrl?.trim() || DEFAULT_CODEX_BASE_URL;
   const normalized = raw.replace(/\/+$/, "");
@@ -268,13 +613,13 @@ function compressBody(body: string): Uint8Array | undefined {
 }
 
 function isRetryable(status: number, text: string): boolean {
-  if (
-    status === 429 &&
-    /usage limit|insufficient_quota|out of budget|quota exceeded|billing/i.test(text)
-  ) {
-    return false;
+  if (status === 429 && isTerminalRateLimitError(text)) return false;
+  if (status === 429 || status === 500 || status === 502 || status === 503 || status === 504) {
+    return true;
   }
-  return status === 429 || status === 500 || status === 502 || status === 503 || status === 504;
+  return /rate.?limit|overloaded|service.?unavailable|upstream.?connect|connection.?refused/i.test(
+    text,
+  );
 }
 
 async function* parseSse(response: Response, signal?: AbortSignal): AsyncGenerator<JsonRecord> {
@@ -289,6 +634,7 @@ async function* parseSse(response: Response, signal?: AbortSignal): AsyncGenerat
     while (true) {
       if (signal?.aborted) throw new Error("Request was aborted");
       const { done, value } = await reader.read();
+      if (signal?.aborted) throw new Error("Request was aborted");
       if (done) break;
       buffer += decoder.decode(value, { stream: true }).replaceAll("\r\n", "\n");
       let boundary = buffer.indexOf("\n\n");
@@ -302,7 +648,16 @@ async function* parseSse(response: Response, signal?: AbortSignal): AsyncGenerat
           .join("\n")
           .trim();
         if (data && data !== "[DONE]") {
-          const parsed = JSON.parse(data) as unknown;
+          let parsed: unknown;
+          try {
+            parsed = JSON.parse(data) as unknown;
+          } catch (error) {
+            throw new CodexProtocolError(
+              `Invalid Codex SSE JSON: ${thrownMessage(error)}`,
+              data,
+              error,
+            );
+          }
           if (!isObject(parsed)) throw new Error("Invalid Codex SSE event");
           yield parsed;
         }
@@ -312,7 +667,9 @@ async function* parseSse(response: Response, signal?: AbortSignal): AsyncGenerat
   } finally {
     signal?.removeEventListener("abort", onAbort);
     await reader.cancel().catch(() => {});
-    reader.releaseLock();
+    try {
+      reader.releaseLock();
+    } catch {}
   }
 }
 
@@ -331,12 +688,14 @@ function socketReusable(socket: WebSocketLike): boolean {
   return socket.readyState === undefined || socket.readyState === 1;
 }
 
-function scheduleSocketExpiry(sessionId: string, entry: CachedWebSocket): void {
+function scheduleSocketExpiry(sessionId: string, accountId: string, entry: CachedWebSocket): void {
   if (entry.idleTimer) clearTimeout(entry.idleTimer);
   entry.idleTimer = setTimeout(() => {
     if (entry.busy) return;
     closeSocket(entry.socket, "idle_timeout");
-    websocketSessions.delete(sessionId);
+    const accountEntries = websocketSessions.get(sessionId);
+    if (accountEntries?.get(accountId) === entry) accountEntries.delete(accountId);
+    if (accountEntries?.size === 0) websocketSessions.delete(sessionId);
   }, SESSION_WEBSOCKET_CACHE_TTL_MS);
 }
 
@@ -347,7 +706,9 @@ async function connectWebSocket(
   timeoutMs: number,
 ): Promise<WebSocketLike> {
   const WebSocketClass = websocketConstructor();
-  if (!WebSocketClass) throw new Error("WebSocket transport is unavailable");
+  if (!WebSocketClass) {
+    throw new Error("WebSocket transport is not available in this runtime");
+  }
   const requestHeaders = headersToRecord(headers);
   delete requestHeaders["OpenAI-Beta"];
 
@@ -363,11 +724,11 @@ async function connectWebSocket(
       socket.removeEventListener("close", onClose);
       signal?.removeEventListener("abort", onAbort);
     };
-    const fail = (error: Error) => {
+    const fail = (error: Error, closeReason: string) => {
       if (settled) return;
       settled = true;
       cleanup();
-      closeSocket(socket, "connect_failure");
+      closeSocket(socket, closeReason);
       reject(error);
     };
     const onOpen = () => {
@@ -376,25 +737,27 @@ async function connectWebSocket(
       cleanup();
       resolve(socket);
     };
-    const onError = (event: unknown) => fail(new Error(`WebSocket error: ${explain(event)}`));
+    const onError = (event: unknown) => fail(extractWebSocketError(event), "connect_failure");
     const onClose = (event: unknown) =>
-      fail(new Error(`WebSocket closed during connect: ${explain(event)}`));
-    const onAbort = () => fail(new Error("Request was aborted"));
+      fail(extractWebSocketCloseError(event, "WebSocket closed during connect"), "connect_failure");
+    const onAbort = () => fail(new Error("Request was aborted"), "aborted");
 
     try {
       socket = new WebSocketClass(url, { headers: requestHeaders });
     } catch (error) {
-      reject(error);
+      reject(error instanceof Error ? error : new Error(String(error)));
       return;
     }
     socket.addEventListener("open", onOpen);
     socket.addEventListener("error", onError);
     socket.addEventListener("close", onClose);
     signal?.addEventListener("abort", onAbort, { once: true });
-    timer = setTimeout(
-      () => fail(new Error(`WebSocket connect timeout after ${timeoutMs}ms`)),
-      timeoutMs,
-    );
+    if (timeoutMs > 0) {
+      timer = setTimeout(
+        () => fail(new Error(`WebSocket connect timeout after ${timeoutMs}ms`), "connect_timeout"),
+        timeoutMs,
+      );
+    }
     if (signal?.aborted) onAbort();
   });
 }
@@ -403,11 +766,19 @@ async function acquireWebSocket(
   url: string,
   headers: Headers,
   sessionId: string | undefined,
+  accountId: string,
   signal: AbortSignal | undefined,
   timeoutMs: number,
-): Promise<{ socket: WebSocketLike; entry?: CachedWebSocket; release(keep: boolean): void }> {
+): Promise<{
+  socket: WebSocketLike;
+  entry?: CachedWebSocket;
+  reused: boolean;
+  release(keep: boolean): void;
+}> {
+  if (signal?.aborted) throw new Error("Request was aborted");
   if (sessionId) {
-    const cached = websocketSessions.get(sessionId);
+    let accountEntries = websocketSessions.get(sessionId);
+    const cached = accountEntries?.get(accountId);
     if (cached?.idleTimer) {
       clearTimeout(cached.idleTimer);
       delete cached.idleTimer;
@@ -418,70 +789,101 @@ async function acquireWebSocket(
       return {
         socket: cached.socket,
         entry: cached,
+        reused: true,
         release(keep) {
           if (!keep || !socketReusable(cached.socket)) {
+            if (cached.idleTimer) {
+              clearTimeout(cached.idleTimer);
+              delete cached.idleTimer;
+            }
             closeSocket(cached.socket);
-            websocketSessions.delete(sessionId);
+            const currentEntries = websocketSessions.get(sessionId);
+            if (currentEntries?.get(accountId) === cached) currentEntries.delete(accountId);
+            if (currentEntries?.size === 0) websocketSessions.delete(sessionId);
             return;
           }
           cached.busy = false;
-          scheduleSocketExpiry(sessionId, cached);
+          scheduleSocketExpiry(sessionId, accountId, cached);
         },
       };
     }
     if (cached && !cached.busy) {
-      closeSocket(cached.socket);
-      websocketSessions.delete(sessionId);
+      closeSocket(cached.socket, expired ? "connection_age_limit" : "done");
+      accountEntries?.delete(accountId);
+      if (accountEntries?.size === 0) websocketSessions.delete(sessionId);
+    }
+    if (cached?.busy) {
+      const socket = await connectWebSocket(url, headers, signal, timeoutMs);
+      return { socket, reused: false, release: () => closeSocket(socket) };
+    }
+
+    const socket = await connectWebSocket(url, headers, signal, timeoutMs);
+    const entry: CachedWebSocket = { socket, busy: true, createdAt: Date.now() };
+    accountEntries = websocketSessions.get(sessionId);
+    if (!accountEntries) {
+      accountEntries = new Map();
+      websocketSessions.set(sessionId, accountEntries);
     }
+    accountEntries.set(accountId, entry);
+    return {
+      socket,
+      entry,
+      reused: false,
+      release(keep) {
+        if (!keep || !socketReusable(socket)) {
+          closeSocket(socket);
+          if (entry.idleTimer) clearTimeout(entry.idleTimer);
+          const currentEntries = websocketSessions.get(sessionId);
+          if (currentEntries?.get(accountId) === entry) currentEntries.delete(accountId);
+          if (currentEntries?.size === 0) websocketSessions.delete(sessionId);
+          return;
+        }
+        entry.busy = false;
+        scheduleSocketExpiry(sessionId, accountId, entry);
+      },
+    };
   }
 
   const socket = await connectWebSocket(url, headers, signal, timeoutMs);
-  if (!sessionId) {
-    return { socket, release: () => closeSocket(socket) };
-  }
-  const entry: CachedWebSocket = { socket, busy: true, createdAt: Date.now() };
-  websocketSessions.set(sessionId, entry);
-  return {
-    socket,
-    entry,
-    release(keep) {
-      if (!keep || !socketReusable(socket)) {
-        closeSocket(socket);
-        websocketSessions.delete(sessionId);
-        return;
-      }
-      entry.busy = false;
-      scheduleSocketExpiry(sessionId, entry);
-    },
-  };
+  return { socket, reused: false, release: () => closeSocket(socket) };
 }
 
 function requestWithoutHistory(body: JsonRecord): JsonRecord {
-  const result = structuredClone(body);
+  const result = { ...body };
   delete result.input;
   delete result.previous_response_id;
   return result;
 }
 
+function jsonWireRequestBody(body: JsonRecord): JsonRecord {
+  const snapshot = JSON.parse(JSON.stringify(body)) as unknown;
+  if (!isObject(snapshot)) {
+    throw new Error("Codex request body must serialize to a JSON object");
+  }
+  return snapshot;
+}
+
 function cachedRequestBody(entry: CachedWebSocket, body: JsonRecord): JsonRecord {
   const continuation = entry.continuation;
   if (!continuation) return body;
   if (
-    stableResponsesJson(requestWithoutHistory(body)) !==
-    stableResponsesJson(requestWithoutHistory(continuation.lastRequestBody))
+    JSON.stringify(requestWithoutHistory(body)) !==
+    JSON.stringify(requestWithoutHistory(continuation.lastRequestBody))
   ) {
     delete entry.continuation;
     return body;
   }
 
-  const currentInput = Array.isArray(body.input) ? body.input.filter(isObject) : [];
-  const previousInput = Array.isArray(continuation.lastRequestBody.input)
-    ? continuation.lastRequestBody.input.filter(isObject)
-    : [];
+  const currentInput = body.input ?? [];
+  const previousInput = continuation.lastRequestBody.input ?? [];
+  if (!Array.isArray(currentInput) || !Array.isArray(previousInput)) {
+    delete entry.continuation;
+    return body;
+  }
   const baseline = [...previousInput, ...continuation.lastResponseItems];
   if (
     currentInput.length < baseline.length ||
-    !replayItemsEqual(currentInput.slice(0, baseline.length), baseline)
+    JSON.stringify(currentInput.slice(0, baseline.length)) !== JSON.stringify(baseline)
   ) {
     delete entry.continuation;
     return body;
@@ -494,6 +896,10 @@ function cachedRequestBody(entry: CachedWebSocket, body: JsonRecord): JsonRecord
   };
 }
 
+function requestInputLength(body: JsonRecord): number {
+  return typeof body.input === "string" || Array.isArray(body.input) ? body.input.length : 0;
+}
+
 async function decodeWebSocketData(data: unknown): Promise<string | undefined> {
   if (typeof data === "string") return data;
   if (data instanceof ArrayBuffer) return new TextDecoder().decode(new Uint8Array(data));
@@ -525,9 +931,10 @@ async function* parseWebSocket(
   };
   const onMessage = (event: unknown) => {
     void (async () => {
+      let text: string | undefined;
       try {
         if (!isObject(event)) return;
-        const text = await decodeWebSocketData(event["data"]);
+        text = await decodeWebSocketData(event["data"]);
         if (!text) return;
         const parsed = JSON.parse(text) as unknown;
         if (!isObject(parsed)) throw new Error("Invalid WebSocket event");
@@ -543,19 +950,23 @@ async function* parseWebSocket(
         queue.push(parsed);
         notify();
       } catch (error) {
-        failure = error instanceof Error ? error : new Error(String(error));
+        failure = new CodexProtocolError(
+          `Invalid Codex WebSocket JSON: ${thrownMessage(error)}`,
+          text,
+          error,
+        );
         done = true;
         notify();
       }
     })();
   };
   const onError = (event: unknown) => {
-    failure = new Error(`WebSocket error: ${explain(event)}`);
+    if (!failure) failure = extractWebSocketError(event);
     done = true;
     notify();
   };
   const onClose = (event: unknown) => {
-    if (!terminal) failure = new Error(`WebSocket closed: ${explain(event)}`);
+    if (!terminal && !failure) failure = extractWebSocketCloseError(event);
     done = true;
     notify();
   };
@@ -571,6 +982,7 @@ async function* parseWebSocket(
   signal?.addEventListener("abort", onAbort, { once: true });
   try {
     while (true) {
+      if (signal?.aborted) throw new Error("Request was aborted");
       if (queue.length > 0) {
         yield queue.shift()!;
         continue;
@@ -579,10 +991,14 @@ async function* parseWebSocket(
       await new Promise<void>((resolve, reject) => {
         wake = resolve;
         if (timeoutMs !== undefined && timeoutMs > 0) {
-          const timer = setTimeout(
-            () => reject(new Error(`WebSocket idle timeout after ${timeoutMs}ms`)),
-            timeoutMs,
-          );
+          const timer = setTimeout(() => {
+            const error = new Error(`WebSocket idle timeout after ${timeoutMs}ms`);
+            failure = error;
+            done = true;
+            wake = undefined;
+            closeSocket(socket, "idle_timeout");
+            reject(error);
+          }, timeoutMs);
           const priorWake = wake;
           wake = () => {
             clearTimeout(timer);
@@ -592,7 +1008,7 @@ async function* parseWebSocket(
       });
     }
     if (failure) throw failure;
-    if (!terminal) throw new Error("WebSocket ended without a terminal response");
+    if (!terminal) throw new Error("WebSocket stream closed before response.completed");
   } finally {
     socket.removeEventListener("message", onMessage);
     socket.removeEventListener("error", onError);
@@ -601,24 +1017,35 @@ async function* parseWebSocket(
   }
 }
 
-function normalizeEvent(event: JsonRecord): JsonRecord {
-  const type = event.type;
+function normalizeEvent(event: JsonRecord): JsonRecord | undefined {
+  const type = typeof event.type === "string" ? event.type : undefined;
+  if (!type) return undefined;
   if (type === "error") {
     const nested = isObject(event["error"]) ? event["error"] : undefined;
-    throw new CodexResponseError(
+    const code =
+      typeof event["code"] === "string"
+        ? event["code"]
+        : typeof nested?.["code"] === "string"
+          ? nested["code"]
+          : undefined;
+    const message =
       typeof event["message"] === "string"
         ? event["message"]
         : typeof nested?.["message"] === "string"
           ? nested["message"]
-          : "Codex request failed",
+          : undefined;
+    throw new CodexApiError(
+      `Codex error: ${message || code || JSON.stringify(event)}`,
+      code,
+      event,
     );
   }
   if (type === "response.failed") {
     const response = isObject(event.response) ? event.response : undefined;
     const error = isObject(response?.["error"]) ? response["error"] : undefined;
-    throw new CodexResponseError(
-      typeof error?.["message"] === "string" ? error["message"] : "Codex response failed",
-    );
+    const message = typeof error?.["message"] === "string" ? error["message"] : undefined;
+    const code = typeof error?.["code"] === "string" ? error["code"] : undefined;
+    throw new CodexApiError(message || "Codex response failed", code, event);
   }
   if (type === "response.done") {
     return { ...event, type: "response.completed" };
@@ -626,59 +1053,93 @@ function normalizeEvent(event: JsonRecord): JsonRecord {
   return event;
 }
 
+function isTerminalEvent(event: JsonRecord): boolean {
+  return event.type === "response.completed" || event.type === "response.incomplete";
+}
+
 async function* requestSse(
   model: Model<any>,
-  body: JsonRecord,
+  bodyJson: string,
   options: CodexTransportOptions,
   headers: Headers,
+  timeoutMs: number | undefined,
 ): AsyncGenerator<JsonRecord> {
-  const bodyJson = JSON.stringify(body);
   const compressed = compressBody(bodyJson);
   if (compressed) headers.set("content-encoding", "zstd");
   const requestBody = compressed ?? bodyJson;
   const maxRetries = options.maxRetries ?? DEFAULT_MAX_RETRIES;
+  let response: Response | undefined;
+  let lastError: Error | undefined;
 
   for (let attempt = 0; attempt <= maxRetries; attempt++) {
-    const timeoutSignal =
-      options.timeoutMs !== undefined && options.timeoutMs > 0
-        ? AbortSignal.timeout(options.timeoutMs)
-        : undefined;
-    const combined = combineAbortSignals([options.signal, timeoutSignal]);
-    let response: Response;
+    if (options.signal?.aborted) throw new Error("Request was aborted");
     try {
+      const timeoutSignal =
+        timeoutMs !== undefined && timeoutMs > 0 ? AbortSignal.timeout(timeoutMs) : undefined;
+      const combined = combineAbortSignals([options.signal, timeoutSignal]);
       try {
-        response = await (options.fetch ?? globalThis.fetch)(resolveCodexUrl(model.baseUrl), {
-          method: "POST",
-          headers,
-          body: requestBody,
-          ...(combined.signal ? { signal: combined.signal } : {}),
-        });
-      } catch (error) {
-        if (attempt < maxRetries && !options.signal?.aborted) {
-          await sleep(BASE_DELAY_MS * 2 ** attempt, options.signal);
-          continue;
+        try {
+          response = await (options.fetch ?? globalThis.fetch)(resolveCodexUrl(model.baseUrl), {
+            method: "POST",
+            headers,
+            body: requestBody,
+            ...(combined.signal ? { signal: combined.signal } : {}),
+          });
+        } catch (error) {
+          if (timeoutSignal?.aborted && !options.signal?.aborted) {
+            throw new Error(`Codex SSE response headers timed out after ${String(timeoutMs)}ms`);
+          }
+          throw error;
         }
-        throw error;
+      } finally {
+        combined.cleanup();
       }
-    } finally {
-      combined.cleanup();
-    }
-    await options.onResponse?.(
-      { status: response.status, headers: headersToRecord(response.headers) },
-      model,
-    );
-    if (!response.ok) {
+      await options.onResponse?.(
+        { status: response.status, headers: headersToRecord(response.headers) },
+        model,
+      );
+      if (response.ok) break;
+
       const errorText = await response.text();
       if (attempt < maxRetries && isRetryable(response.status, errorText)) {
+        const requestedDelay = retryDelayMs(response.headers);
+        const delay =
+          requestedDelay === undefined
+            ? BASE_DELAY_MS * 2 ** attempt
+            : validateRetryDelay(requestedDelay, options);
+        await sleep(delay, options.signal);
+        continue;
+      }
+      throw codexHttpError(response.status, response.statusText, errorText);
+    } catch (error) {
+      if (
+        options.signal?.aborted ||
+        (error instanceof Error &&
+          (error.name === "AbortError" || error.message === "Request was aborted"))
+      ) {
+        throw new Error("Request was aborted");
+      }
+      lastError = error instanceof Error ? error : new Error(String(error));
+      if (
+        attempt < maxRetries &&
+        !(lastError instanceof RetryDelayExceededError) &&
+        !lastError.message.includes("usage limit")
+      ) {
         await sleep(BASE_DELAY_MS * 2 ** attempt, options.signal);
         continue;
       }
-      throw new Error(errorText || `Codex request failed with status ${response.status}`);
+      throw lastError;
     }
-    for await (const event of parseSse(response, options.signal)) {
-      yield normalizeEvent(event);
-    }
-    return;
+  }
+
+  if (!response?.ok) throw lastError ?? new Error("Failed after retries");
+  if (!response.body) throw new Error("No response body");
+  options.onTransportStart?.();
+  for await (const event of parseSse(response, options.signal)) {
+    const normalized = normalizeEvent(event);
+    if (!normalized) continue;
+    yield normalized;
+    if (isTerminalEvent(normalized)) return;
   }
 }
 
@@ -688,24 +1149,74 @@ async function* requestWebSocket(
   options: CodexTransportOptions,
   headers: Headers,
   sessionId: string | undefined,
+  accountId: string,
+  timeoutMs: number | undefined,
+  connectTimeoutMs: number,
+  requestBytes: number,
 ): AsyncGenerator<JsonRecord> {
   const acquired = await acquireWebSocket(
     resolveCodexWebSocketUrl(model.baseUrl),
     headers,
     sessionId,
+    accountId,
     options.signal,
-    options.websocketConnectTimeoutMs ?? DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS,
+    connectTimeoutMs,
   );
-  let keep = true;
+  let keep = false;
+  let responseHandleActive = true;
+  const discard = (): boolean => {
+    if (!responseHandleActive) return false;
+    responseHandleActive = false;
+    if (acquired.entry) delete acquired.entry.continuation;
+    acquired.release(false);
+    return true;
+  };
   try {
+    if (options.signal?.aborted) throw new Error("Request was aborted");
     const useContinuation =
       options.transport === "auto" || options.transport === "websocket-cached";
+    const fullBody = useContinuation && acquired.entry ? jsonWireRequestBody(body) : body;
     const requestBody =
-      useContinuation && acquired.entry ? cachedRequestBody(acquired.entry, body) : body;
+      useContinuation && acquired.entry ? cachedRequestBody(acquired.entry, fullBody) : fullBody;
+    const stats = sessionId ? getOrCreateWebSocketDebugStats(sessionId) : undefined;
+    if (stats) {
+      stats.requests += 1;
+      if (acquired.reused) stats.connectionsReused += 1;
+      else stats.connectionsCreated += 1;
+      if (useContinuation) stats.cachedContextRequests += 1;
+      if (requestBody.store === true) stats.storeTrueRequests += 1;
+      stats.lastInputItems = requestInputLength(requestBody);
+      if (requestBody.previous_response_id) {
+        stats.deltaRequests += 1;
+        stats.lastDeltaInputItems = requestInputLength(requestBody);
+        stats.lastPreviousResponseId = requestBody.previous_response_id as string;
+      } else {
+        stats.fullContextRequests += 1;
+        delete stats.lastDeltaInputItems;
+        delete stats.lastPreviousResponseId;
+      }
+    }
+    options.onWebSocketResponseHandle?.({
+      discard,
+      failParsing(error) {
+        if (!discard()) return false;
+        options.onTransportDiagnostic?.(
+          transportDiagnostic(error, options.transport ?? "auto", true, requestBytes),
+        );
+        return true;
+      },
+    });
     acquired.socket.send(JSON.stringify({ type: "response.create", ...requestBody }));
     const responseItems: JsonRecord[] = [];
     let responseId: string | undefined;
-    for await (const event of parseWebSocket(acquired.socket, options.signal, options.timeoutMs)) {
+    for await (const event of parseWebSocket(acquired.socket, options.signal, timeoutMs)) {
+      if (
+        event.type === "response.created" &&
+        isObject(event.response) &&
+        typeof event.response.id === "string"
+      ) {
+        responseId = event.response.id;
+      }
       if (event.type === "response.output_item.done" && isObject(event.item)) {
         responseItems.push(structuredClone(event.item));
       }
@@ -727,20 +1238,43 @@ async function* requestWebSocket(
           }
         }
       }
-      yield normalizeEvent(event);
+      const normalized = normalizeEvent(event);
+      if (!normalized) continue;
+      yield normalized;
+      if (isTerminalEvent(normalized)) break;
     }
+    if (options.signal?.aborted) throw new Error("Request was aborted");
     if (useContinuation && acquired.entry && responseId) {
-      acquired.entry.continuation = {
-        lastRequestBody: structuredClone(body),
+      const entry = acquired.entry;
+      const continuation = {
+        lastRequestBody: fullBody,
         lastResponseId: responseId,
         lastResponseItems: responseItems.map(normalizeReplayItem),
       };
+      entry.continuation = continuation;
+      if (sessionId) {
+        options.onContinuationReady?.({
+          responseId,
+          replaceResponseItems(items) {
+            if (
+              websocketSessions.get(sessionId)?.get(accountId) !== entry ||
+              entry.continuation !== continuation ||
+              continuation.lastResponseId !== responseId
+            ) {
+              return false;
+            }
+            continuation.lastResponseItems = items.map((item) => structuredClone(item));
+            return true;
+          },
+        });
+      }
     }
+    keep = true;
   } catch (error) {
-    if (acquired.entry) delete acquired.entry.continuation;
-    keep = false;
+    responseHandleActive = false;
     throw error;
   } finally {
+    if (!keep && acquired.entry) delete acquired.entry.continuation;
     acquired.release(keep);
   }
 }
@@ -755,7 +1289,7 @@ export async function requestCodexJson(
     model.headers,
     options.headers,
     options.extraHeaders,
-    extractAccountId(options.apiKey),
+    validateCodexAuthentication(model, options.apiKey),
     options.apiKey,
   );
   const response = await (options.fetch ?? globalThis.fetch)(
@@ -769,7 +1303,7 @@ export async function requestCodexJson(
   );
   const responseText = await response.text();
   if (!response.ok) {
-    throw new Error(responseText || `Codex request failed with status ${response.status}`);
+    throw codexHttpError(response.status, response.statusText, responseText);
   }
   try {
     return JSON.parse(responseText) as unknown;
@@ -789,37 +1323,70 @@ export class CodexTransport {
     options: CodexTransportOptions,
   ): AsyncGenerator<JsonRecord> {
     if (!options.apiKey) throw new Error(`No API key for provider: ${model.provider}`);
-    const accountId = extractAccountId(options.apiKey);
-    const sessionId = options.cacheRetention === "none" ? undefined : options.sessionId;
+    const timeoutMs = normalizeTimeoutMs(options.timeoutMs);
+    const connectTimeoutMs =
+      normalizeTimeoutMs(options.websocketConnectTimeoutMs) ?? DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS;
+    const bodyJson = JSON.stringify(body);
+    const requestBytes = new TextEncoder().encode(bodyJson).byteLength;
+    const accountId = options.accountId ?? validateCodexAuthentication(model, options.apiKey);
+    const cacheSessionId = options.cacheRetention === "none" ? undefined : options.sessionId;
+    const requestId = codexCacheKey(cacheSessionId);
     const transport = options.transport ?? "auto";
 
-    const websocketDisabled =
-      transport === "auto" && sessionId !== undefined && websocketFallbackSessions.has(sessionId);
+    const websocketDisabled = transport !== "sse" && isWebSocketSseFallbackActive(cacheSessionId);
+    if (websocketDisabled) recordWebSocketSseFallback(cacheSessionId);
     if (transport !== "sse" && !websocketDisabled) {
       const headers = websocketHeaders(
         model.headers,
         options.headers,
         accountId,
         options.apiKey,
-        sessionId ?? crypto.randomUUID(),
+        requestId || uuidv7(),
       );
-      let emitted = false;
-      try {
-        for await (const event of requestWebSocket(model, body, options, headers, sessionId)) {
-          emitted = true;
-          yield event;
-        }
-        return;
-      } catch (error) {
-        if (
-          emitted ||
-          error instanceof CodexResponseError ||
-          transport === "websocket" ||
-          transport === "websocket-cached"
-        ) {
-          throw error;
+      let retriedConnectionLimit = false;
+      let retriedMissingContinuation = false;
+      while (true) {
+        let emitted = false;
+        try {
+          for await (const event of requestWebSocket(
+            model,
+            body,
+            options,
+            headers,
+            cacheSessionId,
+            accountId,
+            timeoutMs,
+            connectTimeoutMs,
+            requestBytes,
+          )) {
+            if (!emitted) options.onTransportStart?.();
+            emitted = true;
+            yield event;
+          }
+          return;
+        } catch (error) {
+          const aborted = options.signal?.aborted;
+          const connectionLimitBeforeStart =
+            !emitted && isWebSocketConnectionLimitReachedError(error);
+          if (!aborted && isPreviousResponseNotFoundError(error) && !retriedMissingContinuation) {
+            retriedMissingContinuation = true;
+            continue;
+          }
+          if (!aborted && connectionLimitBeforeStart && !retriedConnectionLimit) {
+            retriedConnectionLimit = true;
+            continue;
+          }
+          if (aborted || (isCodexNonTransportError(error) && !connectionLimitBeforeStart)) {
+            throw error;
+          }
+          options.onTransportDiagnostic?.(
+            transportDiagnostic(error, transport, emitted, requestBytes),
+          );
+          recordWebSocketFailure(cacheSessionId, error);
+          if (emitted) throw error;
+          recordWebSocketSseFallback(cacheSessionId);
+          break;
         }
-        if (sessionId) websocketFallbackSessions.add(sessionId);
       }
     }
 
@@ -828,28 +1395,14 @@ export class CodexTransport {
       options.headers,
       accountId,
       options.apiKey,
-      sessionId,
+      requestId,
     );
-    yield* requestSse(model, body, options, headers);
+    yield* requestSse(model, bodyJson, options, headers, timeoutMs);
   }
 
   close(sessionId?: string): void {
-    if (sessionId) {
-      const entry = websocketSessions.get(sessionId);
-      if (entry?.idleTimer) clearTimeout(entry.idleTimer);
-      if (entry) closeSocket(entry.socket, "session_shutdown");
-      websocketSessions.delete(sessionId);
-      websocketFallbackSessions.delete(sessionId);
-      return;
-    }
-    for (const entry of websocketSessions.values()) {
-      if (entry.idleTimer) clearTimeout(entry.idleTimer);
-      closeSocket(entry.socket, "shutdown");
-    }
-    websocketSessions.clear();
-    websocketFallbackSessions.clear();
+    closeOpenAICodexWebSocketSessions(sessionId);
   }
 }
 
-const transportCleanup = new CodexTransport();
-registerSessionResourceCleanup((sessionId) => transportCleanup.close(sessionId));
+registerSessionResourceCleanup(closeOpenAICodexWebSocketSessions);