@bitkyc08/opencodex 2.86.0 → 2.87.0

package/src/server/responses/adapter-dispatch.ts

@@ -54,7 +54,7 @@ import { describeUpstreamConnectFailure } from "./upstream-error";
 import type { OpaqueBlobRecoveryGuard } from "./core-opaque-recovery";
 import type { AttemptRecoveryKind } from "../../usage/log";
 import type { OAuthAccessSnapshot } from "../../oauth";
-import { OAuthAccountPausedError, OAuthLoginRequiredError, publicOAuthAuthenticationErrorMessage } from "../../oauth";
+import { OAuthAccountPausedError, OAuthLoginRequiredError, OAuthTokenRefreshBusyError, publicOAuthAuthenticationErrorMessage } from "../../oauth";
 import { getAccountSet, markAccountNeedsReauthIfGeneration } from "../../oauth/store";
 import { tryKiroAlternateAfterTerminalRefresh } from "../../oauth/kiro-terminal-failover";
 import { classifyKiroRefusal } from "../../adapters/kiro-refusal";
@@ -932,6 +932,15 @@ export async function prepareAdapterExchange(
         try {
           refreshed = await refreshResolvedOAuthSelection(transportState.sentOAuthSnapshot);
         } catch (err) {
+          const cancelled = options.abortSignal?.aborted || req.signal.aborted;
+          if (cancelled || err instanceof OAuthTokenRefreshBusyError) {
+            cleanupUpstreamAbort();
+            upstream.abort();
+            try { void upstreamResponse.body?.cancel().catch(() => {}); } catch { /* already closed */ }
+            releaseCodexAuthContextProbeLease(admissionState.authCtx);
+            return cancelled ? clientCancelledResponse()
+              : formatErrorResponse(503, "server_error", publicOAuthAuthenticationErrorMessage(err), { retryAfter: "1" });
+          }
           const failed = transportState.sentOAuthSnapshot;
           if (route.providerName === "google-antigravity" && !antigravityAuthRotationAttempted
             && err instanceof OAuthLoginRequiredError && failed
@@ -1337,9 +1346,17 @@ export async function prepareAdapterExchange(
             resolveWireProtocolOverride(route.providerName, route.modelId, route.provider, inboundWire, route.staticPolicy),
             config.cacheRetention,
           );
+          bindRouteReasoningReplayScope({
+            parsed,
+            providerName: route.providerName,
+            provider: route.provider,
+            adapterName: transportState.activeAdapter.name,
+            oauthCredentialSnapshot: transportState.replayOAuthCredentialSnapshot,
+          });
           sealRequestAttemptIdentity(logCtx.activeAttempt, logCtx.provider, transportState.activeAdapter.name, logCtx.accountLogLabel);
           recordAttemptCredentialSource(logCtx.activeAttempt, route.providerName, route.provider, transportState.activeAdapter.name);
-         const result = await rebuildAndRefetch(upstreamResponse.status === 401 ? "oauth-401" : "anthropic-oauth-429");
+         const result = await rebuildAndRefetch(upstreamResponse.status === 401 ? "oauth-401"
+           : upstreamResponse.status === 403 ? "oauth-account-403" : "anthropic-oauth-429");
          if ("failed" in result) return result.failed;
          upstreamResponse = result;
           if (isNonReplayableResponse(upstreamResponse)) continue recovery;

package/src/server/responses/advisor-retention.ts

@@ -0,0 +1,117 @@
+import { jsonUtf8Bytes } from "../../lib/json-byte-size";
+import {
+  TRANSLATOR_MAX_TURN_BYTES,
+  TranslatorBudgetExceededError,
+  type TranslatorBudget,
+} from "../../lib/translator-budget";
+import type { AdapterEvent, OcxMessage, OcxUsage } from "../../types";
+import { mergeUsage } from "./terminal-guard";
+
+export const MAX_ADVISOR_RETAINED_EVENTS = 4_096;
+export const MAX_ADVISOR_RETAINED_CALLS = 1_024;
+
+/** Independent ownership: never borrow the bridge's event or tool-call leases. */
+export class AdvisorRetention {
+  private legBytes = 0;
+  private historyBytes = 0;
+  private copyBytes = 0;
+  private usageBytes = 0;
+  private events = 0;
+  private calls = 0;
+
+  constructor(private readonly shared?: TranslatorBudget) {}
+
+  private charge(bytes: number, kind: "retained_collectors" | "request_copies"): void {
+    if (bytes > TRANSLATOR_MAX_TURN_BYTES - this.legBytes - this.historyBytes - this.copyBytes - this.usageBytes) {
+      throw new TranslatorBudgetExceededError(kind, TRANSLATOR_MAX_TURN_BYTES);
+    }
+    this.shared?.chargeRetained(bytes, { kind });
+  }
+
+  /** Usage can alias a terminal payload. Keep its own owner after the raw leg ends. */
+  usage(first: OcxUsage | undefined, second: OcxUsage | undefined): OcxUsage | undefined {
+    if (!second) return first;
+    // Two-input merge materializes a new object. Reserve before merging, including
+    // a conservative bound for concatenation if a provider supplied malformed scalars.
+    const bytes = first
+      ? (jsonUtf8Bytes(first) + jsonUtf8Bytes(second)) * 2 + 256
+      : jsonUtf8Bytes(second) + 1;
+    this.charge(bytes, "request_copies");
+    this.copyBytes += bytes;
+    let retained = 0;
+    try {
+      const merged = mergeUsage(first, second);
+      const measured = jsonUtf8Bytes(merged, bytes) + 1;
+      if (measured > bytes) throw new TranslatorBudgetExceededError("request_copies", TRANSLATOR_MAX_TURN_BYTES);
+      retained = measured;
+      const previous = this.usageBytes;
+      this.usageBytes = retained;
+      this.shared?.releaseRetained(previous, { kind: "request_copies" });
+      return merged;
+    } finally {
+      this.copyBytes -= bytes;
+      this.shared?.releaseRetained(bytes - retained, { kind: "request_copies" });
+    }
+  }
+
+  /**
+   * `counted: false` is only for non-empty text/thinking deltas: the leg's 16 KiB character
+   * caps already bound how many of those can be held, and an ordinary long answer streams
+   * thousands of them. They still consume bytes.
+   */
+  retain(event: AdapterEvent, call = false, counted = true): void {
+    if ((counted && this.events >= MAX_ADVISOR_RETAINED_EVENTS) || (call && this.calls >= MAX_ADVISOR_RETAINED_CALLS)) {
+      throw new TranslatorBudgetExceededError("retained_collectors", TRANSLATOR_MAX_TURN_BYTES);
+    }
+    const bytes = jsonUtf8Bytes(event) + 1;
+    this.charge(bytes, "retained_collectors");
+    this.legBytes += bytes;
+    if (counted) this.events++;
+    if (call) this.calls++;
+  }
+
+  /** Hold a bounded materialization while its raw sources are still retained. */
+  copy<T>(bytes: number, build: () => T): T {
+    this.charge(bytes, "request_copies");
+    this.copyBytes += bytes;
+    try { return build(); }
+    finally {
+      this.copyBytes -= bytes;
+      this.shared?.releaseRetained(bytes, { kind: "request_copies" });
+    }
+  }
+
+  /** Reserve before constructing replay JSON; transfer the reservation to history. */
+  history<T extends OcxMessage | undefined>(bytes: number, build: () => T): T {
+    this.charge(bytes, "request_copies");
+    this.copyBytes += bytes;
+    let retained = 0;
+    try {
+      const message = build();
+      const measured = message ? jsonUtf8Bytes(message, bytes) + 1 : 0;
+      if (measured > bytes) throw new TranslatorBudgetExceededError("request_copies", TRANSLATOR_MAX_TURN_BYTES);
+      retained = measured;
+      this.historyBytes += retained;
+      return message;
+    } finally {
+      this.copyBytes -= bytes;
+      this.shared?.releaseRetained(bytes - retained, { kind: "request_copies" });
+    }
+  }
+
+  /** Event envelopes dominate the replay content; allow new message/call framing too. */
+  replayReservation(): number { return this.legBytes + (this.calls + 1) * 256; }
+
+  releaseLeg(): void {
+    this.shared?.releaseRetained(this.legBytes, { kind: "retained_collectors" });
+    this.legBytes = this.events = this.calls = 0;
+  }
+
+  release(): void {
+    this.releaseLeg();
+    this.shared?.releaseRetained(this.historyBytes, { kind: "request_copies" });
+    this.historyBytes = 0;
+    this.shared?.releaseRetained(this.usageBytes, { kind: "request_copies" });
+    this.usageBytes = 0;
+  }
+}

package/src/server/responses/advisor-slot.ts

@@ -1,3 +1,4 @@
+import { assistantReplayContent } from "../../responses/assistant-replay";
 /**
  * Core-owned registration slot for the optional advisor subsystem (src/advisor).
  *
@@ -28,7 +29,9 @@ import type {
   OcxToolResultMessage,
   OcxUsage,
 } from "../../types";
-import { mergeUsage } from "./terminal-guard";
+import { jsonStringPartsUtf8Bytes, jsonUtf8Bytes } from "../../lib/json-byte-size";
+import { isTranslatorBudgetExceededError, TRANSLATOR_MAX_CALL_ARGUMENT_BYTES, TranslatorBudgetExceededError, type TranslatorBudget } from "../../lib/translator-budget";
+import { AdvisorRetention } from "./advisor-retention";
 
 export const ADVISOR_TOOL_NAME = "advisor";
 
@@ -53,7 +56,7 @@ export interface AdvisorConsultOutcome {
    * Operator config blocked the consultation before any outbound call. Not a provider
    * failure: preflight must release a claim instead of recording a cooldown.
    */
-  blocked?: "consent" | "settings";
+  blocked?: "consent" | "settings" | "budget" | "limit";
 }
 
 /** The structural plan the optional advisor subsystem registers per request. */
@@ -78,12 +81,19 @@ export interface AdvisorPlan {
 interface HeldAdvisorCall {
   id: string;
   name: string;
-  argsBuf: string;
+  args: string[];
+  argumentBytes: number;
+  question?: string;
   closed: boolean;
   providerMetadata?: import("../../types").OcxProviderOpaqueToolCallMetadata;
 }
 
-function parseAdvisorArgs(argsBuf: string): { question?: string } {
+function parseAdvisorArgs(call: HeldAdvisorCall, retention: AdvisorRetention): { question?: string } {
+  // Reserve both the joined input and parsed JSON before either is materialized.
+  return retention.copy(call.argumentBytes * 2 + 4, () => parseArgs(call.args.join("")));
+}
+
+function parseArgs(argsBuf: string): { question?: string } {
   if (argsBuf.trim() === "") return {};
   try {
     const parsed: unknown = JSON.parse(argsBuf);
@@ -99,31 +109,18 @@ function assistantMessageFromLeg(
   events: readonly AdapterEvent[],
   held: readonly HeldAdvisorCall[],
   timestamp: number,
+  retention: AdvisorRetention,
 ): OcxAssistantMessage | undefined {
-  let text = "";
-  let thinking = "";
-  let signature: string | undefined;
-  const redacted: string[] = [];
-  for (const event of events) {
-    if (event.type === "text_delta") text += event.text;
-    else if (event.type === "thinking_delta") thinking += event.thinking;
-    else if (event.type === "thinking_signature") signature = event.signature;
-    else if (event.type === "redacted_thinking") redacted.push(event.data);
-  }
-  if (text.length > MAX_LEG_TEXT_CHARS) text = text.slice(0, MAX_LEG_TEXT_CHARS);
-  if (thinking.length > MAX_LEG_THINKING_CHARS) thinking = thinking.slice(0, MAX_LEG_THINKING_CHARS);
-  const content: OcxAssistantContentPart[] = [];
-  if (thinking || signature || redacted.length > 0) {
-    content.push({ type: "thinking", thinking, ...(signature ? { signature } : {}), ...(redacted.length > 0 ? { redacted } : {}) });
-  }
-  if (text) content.push({ type: "text", text });
+  const content = assistantReplayContent(events);
   for (const call of held) {
     if (!call.closed) continue;
+    const args = parseAdvisorArgs(call, retention);
+    call.question = args.question;
     content.push({
       type: "toolCall",
       id: call.id,
       name: call.name,
-      arguments: parseAdvisorArgs(call.argsBuf),
+      arguments: args,
       ...(call.providerMetadata ? { providerMetadata: call.providerMetadata } : {}),
     });
   }
@@ -150,6 +147,8 @@ export interface AdvisorGuardOptions {
   parsed: OcxParsedRequest;
   plan: AdvisorPlan;
   firstEvents: AsyncIterable<AdapterEvent>;
+  /** Independent guard ownership within the optional shared request budget. */
+  translatorBudget?: TranslatorBudget;
   /** One bounded worker continuation re-dispatch (same machinery as the terminal guard). */
   continuation: (parsed: OcxParsedRequest) => AsyncIterable<AdapterEvent> | Promise<AsyncIterable<AdapterEvent>>;
 }
@@ -172,7 +171,8 @@ export function createAdvisorGuard(plan: AdvisorPlan): NonNullable<OcxParsedRequ
     let parsed = options.parsed;
     let accumulatedUsage: OcxUsage | undefined;
     let source: AsyncIterable<AdapterEvent> = options.firstEvents;
-
+    const retention = new AdvisorRetention(options.translatorBudget);
+    try {
     while (true) {
       const held: HeldAdvisorCall[] = [];
       const legEvents: AdapterEvent[] = [];
@@ -189,15 +189,33 @@ export function createAdvisorGuard(plan: AdvisorPlan): NonNullable<OcxParsedRequ
       for await (const event of source) {
         if (event.type === "tool_call_start") {
           if (pending) { held.push(pending); pending = null; }
-          pending = { id: event.id, name: event.name, argsBuf: "", closed: false, ...(event.providerMetadata ? { providerMetadata: event.providerMetadata } : {}) };
           holdingCurrent = event.name === ADVISOR_TOOL_NAME;
-          if (!holdingCurrent) hasRealToolCall = true;
-          else continue;
+          if (!holdingCurrent) {
+            hasRealToolCall = true;
+            // No continuation can use this leg. Real-tool arguments are never copied.
+            held.length = legEvents.length = 0;
+            retention.releaseLeg();
+          } else {
+            if (!hasRealToolCall) {
+              retention.retain(event, true);
+              pending = { id: event.id, name: event.name, args: [], argumentBytes: 0, closed: false, ...(event.providerMetadata ? { providerMetadata: event.providerMetadata } : {}) };
+            }
+            continue;
+          }
         } else if (event.type === "tool_call_delta") {
-          if (pending) pending.argsBuf += event.arguments;
+          if (pending) {
+            const bytes = jsonStringPartsUtf8Bytes([event.arguments], TRANSLATOR_MAX_CALL_ARGUMENT_BYTES + 2) - 2;
+            if (bytes > TRANSLATOR_MAX_CALL_ARGUMENT_BYTES - pending.argumentBytes) {
+              throw new TranslatorBudgetExceededError("tool_args", TRANSLATOR_MAX_CALL_ARGUMENT_BYTES);
+            }
+            retention.retain(event);
+            pending.argumentBytes += bytes;
+            pending.args.push(event.arguments);
+          }
           if (holdingCurrent) continue;
         } else if (event.type === "tool_call_end") {
           if (pending) {
+            retention.retain(event);
             pending.closed = true;
             if (pending.name === ADVISOR_TOOL_NAME) held.push(pending);
             pending = null;
@@ -207,6 +225,7 @@ export function createAdvisorGuard(plan: AdvisorPlan): NonNullable<OcxParsedRequ
             continue;
           }
         } else if (event.type === "done") {
+          if (!hasRealToolCall) retention.retain(event);
           if (pending) { held.push(pending); pending = null; }
           terminalEvent = event;
           terminalConsumed = true;
@@ -214,22 +233,23 @@ export function createAdvisorGuard(plan: AdvisorPlan): NonNullable<OcxParsedRequ
         } else if (event.type === "incomplete" || event.type === "error") {
           // A broken leg cannot safely anchor a continuation: surface the terminal as-is.
           if (pending) { pending = null; }
-          const usage = mergeUsage(accumulatedUsage, event.usage);
-          yield usage ? { ...event, usage } : event;
+          accumulatedUsage = retention.usage(accumulatedUsage, event.usage);
+          yield accumulatedUsage ? { ...event, usage: accumulatedUsage } : event;
           return;
         } else {
           // Retain (bounded) the leg's visible content so the rebuilt assistant message is complete.
           if (event.type === "text_delta") {
             legTextChars += event.text.length;
-            if (legTextChars <= MAX_LEG_TEXT_CHARS) legEvents.push(event);
+            // Empty deltas add nothing to the rebuilt message; non-empty ones are bounded by the cap.
+            if (!hasRealToolCall && event.text !== "" && legTextChars <= MAX_LEG_TEXT_CHARS) { retention.retain(event, false, false); legEvents.push(event); }
           } else if (event.type === "thinking_delta") {
             legThinkingChars += event.thinking.length;
-            if (legThinkingChars <= MAX_LEG_THINKING_CHARS) legEvents.push(event);
+            if (!hasRealToolCall && event.thinking !== "" && legThinkingChars <= MAX_LEG_THINKING_CHARS) { retention.retain(event, false, false); legEvents.push(event); }
           } else if (
             event.type === "thinking_signature"
             || event.type === "redacted_thinking"
           ) {
-            legEvents.push(event);
+            if (!hasRealToolCall) { retention.retain(event); legEvents.push(event); }
           }
         }
         yield event;
@@ -245,14 +265,14 @@ export function createAdvisorGuard(plan: AdvisorPlan): NonNullable<OcxParsedRequ
       if (!shouldIntercept) {
         // Plain leg: surface the terminal with merged usage from any earlier intercepted legs.
         if (terminalEvent) {
-          const usage = mergeUsage(accumulatedUsage, terminalEvent.usage);
-          yield usage ? { ...terminalEvent, usage } : terminalEvent;
+          accumulatedUsage = retention.usage(accumulatedUsage, terminalEvent.usage);
+          yield accumulatedUsage ? { ...terminalEvent, usage: accumulatedUsage } : terminalEvent;
         }
         return;
       }
 
       // Consume this leg's done — the continuation replaces it.
-      accumulatedUsage = mergeUsage(accumulatedUsage, terminalEvent?.usage);
+      accumulatedUsage = retention.usage(accumulatedUsage, terminalEvent?.usage);
       if (continuations >= MAX_ADVISOR_CONTINUATIONS_PER_REQUEST || finalContinuationUsed) {
         yield {
           type: "error",
@@ -269,22 +289,22 @@ export function createAdvisorGuard(plan: AdvisorPlan): NonNullable<OcxParsedRequ
       continuations += 1;
       finalContinuationUsed ||= isFinalContinuation;
       const timestamp = Date.now();
-      const assistant = assistantMessageFromLeg(legEvents, advisorCalls, timestamp);
+      const assistant = retention.history(retention.replayReservation(), () => assistantMessageFromLeg(legEvents, advisorCalls, timestamp, retention));
       const messages: OcxMessage[] = [...parsed.context.messages];
       if (assistant) messages.push(assistant);
 
       for (const call of advisorCalls) {
         if (consultations >= maxConsultations) {
           // Runtime-owned limit text: not advice, and never a genuine advice wrapper.
-          messages.push(advisorToolResult(call, {
+          const outcome = {
             ok: false,
             isError: true,
             content: plan.formatUnavailable("limit", "consultation limit reached for this request"),
-          }, timestamp));
+          };
+          messages.push(retention.history(jsonUtf8Bytes(outcome) + jsonUtf8Bytes(call.id) + 256, () => advisorToolResult(call, outcome, timestamp))!);
           continue;
         }
         consultations += 1;
-        const args = parseAdvisorArgs(call.argsBuf);
         // Later calls in the same leg must see what came before them: the rebuilt worker turn
         // and every earlier advisor result in this leg's accumulated messages.
         const consultParsed: OcxParsedRequest = {
@@ -293,8 +313,9 @@ export function createAdvisorGuard(plan: AdvisorPlan): NonNullable<OcxParsedRequ
         };
         let outcome: AdvisorConsultOutcome;
         try {
-          outcome = await plan.consult(consultParsed, "manual", args.question);
+          outcome = await plan.consult(consultParsed, "manual", call.question);
         } catch (error) {
+          if (isTranslatorBudgetExceededError(error)) throw error;
           // Runtime-owned failure text only: the implementation neutralizes untrusted exception
           // text so it can never forge a genuine advice wrapper.
           outcome = {
@@ -303,7 +324,7 @@ export function createAdvisorGuard(plan: AdvisorPlan): NonNullable<OcxParsedRequ
             content: plan.formatUnavailable("manual", error instanceof Error ? error.message : String(error)),
           };
         }
-        messages.push(advisorToolResult(call, outcome, timestamp));
+        messages.push(retention.history(jsonUtf8Bytes(outcome) + jsonUtf8Bytes(call.id) + 256, () => advisorToolResult(call, outcome, timestamp))!);
       }
 
       const exhausted = consultations >= maxConsultations;
@@ -320,10 +341,16 @@ export function createAdvisorGuard(plan: AdvisorPlan): NonNullable<OcxParsedRequ
           : parsed.options,
       };
       parsed = nextParsed;
+      // All replay owners were admitted while the raw leg was still live.
+      held.length = legEvents.length = advisorCalls.length = 0;
+      pending = null;
+      terminalEvent = undefined;
+      retention.releaseLeg();
       yield { type: "assistant_boundary" };
       try {
         source = await options.continuation(parsed);
       } catch (error) {
+        if (isTranslatorBudgetExceededError(error)) throw error;
         yield {
           type: "error",
           message: error instanceof Error ? error.message : String(error),
@@ -332,5 +359,10 @@ export function createAdvisorGuard(plan: AdvisorPlan): NonNullable<OcxParsedRequ
         return;
       }
     }
+    } catch (error) {
+      if (!isTranslatorBudgetExceededError(error)) throw error;
+      yield { type: "error", status: 502, errorType: "upstream_error", code: "translation_buffer_limit",
+        message: "upstream translation buffer exceeded the safe limit", ...(accumulatedUsage ? { usage: accumulatedUsage } : {}) };
+    } finally { retention.release(); }
   };
 }

package/src/server/responses/codex-ws-wire.ts

@@ -1,3 +1,4 @@
+import { normalizeResponsesIncompleteReason } from "../../responses/incomplete-reason";
 import { MAX_CLIENT_SSE_FRAME_BYTES } from "../sse-frame-buffer";
 import {
   isTransientUpstreamStatus,
@@ -8,6 +9,7 @@ import {
 import { isRateLimitOrQuotaFailureMessage, recognizedUpstreamError } from "../../lib/errors";
 import type { RequestFailureCause, RequestFailureStage } from "../../lib/request-failure-model";
 import { packageVersion } from "../../lib/package-version";
+import { registerProviderResponseProvenance } from "../../providers/request-pacing";
 // If the 101 never arrives (network black hole), give SSE a chance well before
 // the caller's connect timeout (default 200s) would fire.
 export const UPGRADE_DEADLINE_MS = 10_000;
@@ -49,9 +51,12 @@ export const CODEX_WS_CREATE_FRAME_LIMIT_BYTES =
 /** Close code the backend uses for an oversized message (RFC 6455 "message too big"). */
 const WS_CLOSE_MESSAGE_TOO_BIG = 1009;
 
-const codexWsUpstreamResponses = new WeakSet<Response>();
-const quotaObservedResponses = new WeakSet<Response>();
-const codexWsPreludeProjections = new WeakSet<Response>();
+const codexWsUpstreamResponses = new WeakMap<Response, true>();
+const quotaObservedResponses = new WeakMap<Response, true>();
+const codexWsPreludeProjections = new WeakMap<Response, true>();
+registerProviderResponseProvenance(codexWsUpstreamResponses);
+registerProviderResponseProvenance(quotaObservedResponses);
+registerProviderResponseProvenance(codexWsPreludeProjections);
 
 /** Quota arrived directly at its captured account; do not replay old HTTP prelude headers. */
 export function isCodexWsQuotaObservedResponse(response: Response): boolean {
@@ -65,7 +70,7 @@ export function isCodexWsUpstreamResponse(response: Response): boolean {
 
 /** Pre-response projection; its headers retain the exchange's prelude snapshot. */
 export function markCodexWsPreludeProjection(response: Response): void {
-  codexWsPreludeProjections.add(response);
+  codexWsPreludeProjections.set(response, true);
 }
 
 export function isCodexWsPreludeProjection(response: Response): boolean {
@@ -73,8 +78,8 @@ export function isCodexWsPreludeProjection(response: Response): boolean {
 }
 
 export function markCodexWsResponse(response: Response, observed: boolean): void {
-  codexWsUpstreamResponses.add(response);
-  if (observed) quotaObservedResponses.add(response);
+  codexWsUpstreamResponses.set(response, true);
+  if (observed) quotaObservedResponses.set(response, true);
 }
 
 /**
@@ -313,6 +318,7 @@ export function restoreCodexWsPreOutputRefusal(response: Response): Response {
 }
 
 const codexWsStageByResponse = new WeakMap<Response, CodexWsStageRecord>();
+registerProviderResponseProvenance(codexWsStageByResponse);
 
 export function markCodexWsStage(response: Response, record: CodexWsStageRecord): void {
   const current = codexWsStageByResponse.get(response);
@@ -457,6 +463,7 @@ export function projectCodexWsFailure(
 }
 
 const socketDeathStages = new WeakMap<Response, RequestFailureStage>();
+registerProviderResponseProvenance(socketDeathStages);
 
 /**
  * Record that a pre-response settle came from the socket closing or failing under the send (#4191),
@@ -517,12 +524,12 @@ export function normalizeResponsesWsRelayEvent(text: string): ResponsesWsRelayEv
     return null;
   }
   if (!payload || typeof payload !== "object" || Array.isArray(payload)) return null;
-  const record = payload as Record<string, unknown>;
+  const record = normalizeResponsesIncompleteReason(payload as Record<string, unknown>);
   if (typeof record.type !== "string") return null;
   // A native error may be pretty-printed JSON. SSE prefixes one data line;
   // embedded physical newlines would otherwise truncate the JSON for readers.
   if (record.type !== "response.done") return {
-    type: record.type, text: /[\r\n]/.test(text) ? JSON.stringify(record) : text, payload: record,
+    type: record.type, text: record !== payload || /[\r\n]/.test(text) ? JSON.stringify(record) : text, payload: record,
   };
 
   const response = record.response;