@bitkyc08/opencodex 2.10.0 → 2.10.1

package/src/bridge.ts

@@ -2,11 +2,13 @@ import type { AdapterEvent, OcxMessagePhase, OcxProviderContinuationState, OcxUs
 import { adapterFailureFromMessage, classifyError, CYBER_POLICY_ERROR_CODE, isCyberPolicyCode, type OcxErrorPayload } from "./lib/errors";
 import { encodeCompactionSummary } from "./responses/compaction";
 import { encodeReasoningEnvelope, type ReasoningEnvelope } from "./responses/reasoning-envelope";
+import { rememberReasoningForCall } from "./responses/reasoning-replay-cache";
 import { resolveStallTimeoutSec } from "./stall-timeout";
 import { usageDisplayTotalTokens } from "./usage/totals";
 import {
   isTranslatorBudgetExceededError,
   releaseTranslatedEvent,
+  createTranslatorBudget,
   type TranslatorBudget,
   type TranslatorBufferKind,
 } from "./lib/translator-budget";
@@ -15,6 +17,13 @@ function uuid(): string {
   return crypto.randomUUID().replace(/-/g, "");
 }
 
+/** Test-only: bound the abandoned-owned-budget watchdog delay (null restores). */
+let ownedBudgetAbandonedMs = 10 * 60 * 1000;
+const OWNED_BUDGET_ABANDONED_DEFAULT_MS = ownedBudgetAbandonedMs;
+export function setOwnedBudgetAbandonedMsForTests(ms: number | null): void {
+  ownedBudgetAbandonedMs = ms ?? OWNED_BUDGET_ABANDONED_DEFAULT_MS;
+}
+
 function sseEvent(name: string, data: Record<string, unknown>): string {
   return `event: ${name}\ndata: ${JSON.stringify(data)}\n\n`;
 }
@@ -107,13 +116,28 @@ function adapterFailureFromEvent(event: Extract<AdapterEvent, { type: "error" }>
 export { adapterFailureFromMessage } from "./lib/errors";
 
 /**
- * Build the native `WebSearchAction::Search` payload from the queries that ran. codex-rs prefers a
- * non-empty `query` over `queries` for the cell label, and only renders "<first> ..." when `query`
- * is absent and `queries.len() > 1`. So a single query → `{ query }`; multiple → `{ queries }` with
- * no singular `query`, so Codex shows the native plural ellipsis. Empty → `{ query: "" }`.
+ * Build the native `WebSearchAction::Search` payload from the queries that ran.
+ *
+ * Single query → `{ query, queries: [query] }`. Batch → `{ queries }` with NO singular
+ * `query`. Empty → `{ query: "", queries: [""] }`.
+ *
+ * The asymmetry is load-bearing in both directions. codex-rs prefers a non-empty `query`
+ * for the cell label and renders "<first> ..." only when `query` is ABSENT and
+ * `queries.len() > 1`, so adding `query` to a batch would collapse the plural ellipsis.
+ * Meanwhile DeepSeek's native Responses parser makes `queries` a required field, so a
+ * replayed one-term `web_search_call` — carried in the history of every subsequent turn
+ * — fails deserialization with `missing field 'queries'` and 400s the rest of the
+ * conversation (#930). Carrying both keys in the single case satisfies the strict parser
+ * without changing what codex-rs displays.
+ *
+ * This fixes items created from here on. History recorded before it is repaired at the
+ * replay boundary by `backfillWebSearchQueries()` in the Responses adapter.
  */
 function webSearchAction(queries: string[]): Record<string, unknown> {
-  if (queries.length <= 1) return { type: "search", query: queries[0] ?? "" };
+  if (queries.length <= 1) {
+    const query = queries[0] ?? "";
+    return { type: "search", query, queries: [query] };
+  }
   return { type: "search", queries };
 }
 
@@ -157,6 +181,12 @@ export function bridgeToResponsesSSE(
      */
     onUsage?: (usage: OcxUsage | undefined) => void;
     translatorBudget?: TranslatorBudget;
+    /**
+     * Conversation identity for the reasoning replay cache (issue #950).
+     * Provider call ids are not globally unique; scoping by thread keeps one
+     * conversation's reasoning out of another's continuations.
+     */
+    replayCacheScope?: string;
     /**
      * Test seam for the wire/stall beat loop. Production omits this and uses the
      * global timers; injecting here must not change scheduling semantics.
@@ -167,6 +197,7 @@ export function bridgeToResponsesSSE(
     };
   },
 ): ReadableStream<Uint8Array> {
+  const replayCacheScope = options?.replayCacheScope ?? "global";
   const setBeatInterval = options?.timers?.setInterval ?? ((handler: () => void, ms: number) => setInterval(handler, ms));
   const clearBeatInterval = options?.timers?.clearInterval ?? ((id: unknown) => clearInterval(id as ReturnType<typeof setInterval>));
   // Freeform/custom tools (apply_patch) carry their body in `input`; the model is given a
@@ -208,7 +239,25 @@ export function bridgeToResponsesSSE(
     try { const o = JSON.parse(args); return o && typeof o === "object" ? o : {}; } catch { return {}; }
   };
   const encoder = new TextEncoder();
-  const budget = options?.translatorBudget;
+  // Default-budget safety net: omission is SAFE (default turn limits), never
+  // unbounded. Production callers always pass one; an owned default is disposed
+  // at terminal/cancel below.
+  const ownsBudget = !options?.translatorBudget;
+  const budget = options?.translatorBudget ?? createTranslatorBudget();
+  // Idempotent: safe to call at every stream-death path; disposal must come
+  // AFTER the final charges (emitDone), never inside reportTerminal.
+  const disposeOwnedBudget = () => { if (ownsBudget) budget.dispose(); };
+  // A dropped stream (never read, never cancelled) reaches no terminal path,
+  // so the owned budget would sit in liveBudgets for the process lifetime.
+  // One unref'd watchdog per owned budget bounds that to a timeout and clears
+  // itself on any settle (the delay is test-overridable).
+  const ownedWatchdog = ownsBudget
+    ? setTimeout(() => disposeOwnedBudget(), ownedBudgetAbandonedMs)
+    : undefined;
+  ownedWatchdog?.unref?.();
+  const clearOwnedWatchdog = () => {
+    if (ownedWatchdog !== undefined) clearTimeout(ownedWatchdog);
+  };
   const bytesOf = (value: string): number => Buffer.byteLength(value);
   const appendString = (
     previous: string,
@@ -219,7 +268,6 @@ export function bridgeToResponsesSSE(
   ): { value: string; bytes: number } => {
     const fragmentBytes = bytesOf(fragment);
     const nextBytes = previousBytes + fragmentBytes;
-    if (!budget) return { value: previous + fragment, bytes: nextBytes };
     const scope = { kind, ...(callId ? { callId } : {}) };
     const reservation = budget.reserveTransient(nextBytes, scope);
     try {
@@ -257,6 +305,7 @@ export function bridgeToResponsesSSE(
     if (terminalReported || clientCancelled || closed) return;
     terminalReported = true;
     try { options?.onTerminal?.(status); } catch { /* terminal metrics must not break the stream */ }
+    clearOwnedWatchdog();
   };
   // RC3 keep-alive: Codex's idle timer is timeout(idle_timeout, stream.next()) over an
   // eventsource_stream; ANY received event re-arms it, while an unknown type is ignored
@@ -289,6 +338,7 @@ export function bridgeToResponsesSSE(
             return;
           }
           closed = true;
+          disposeOwnedBudget();
         }
       };
       const emitDone = () => {
@@ -384,8 +434,15 @@ export function bridgeToResponsesSSE(
       // encodeReasoningEnvelope: takeReasoningEnvelope's sig/red guard would drop txt-only.
       let hiddenRawReasoningText = "";
       let hiddenRawReasoningBytes = 0;
+      // Raw reasoning text flushed most recently, waiting for the tool call it
+      // preceded. Recorded into the replay cache on tool_call_start so a later
+      // continuation can re-attach it when history lost the reasoning item
+      // (issue #950). Kept until new reasoning/text arrives: parallel tool
+      // calls share the same preceding reasoning block.
+      let rawReasoningForNextToolCall = "";
       const flushHiddenRawReasoning = () => {
         if (!hiddenRawReasoningText) return;
+        rawReasoningForNextToolCall = hiddenRawReasoningText;
         const previousBytes = hiddenRawReasoningBytes;
         const encrypted = encodeReasoningEnvelope({ txt: hiddenRawReasoningText });
         const reservation = budget?.reserveTransient(bytesOf(encrypted), { kind: "reasoning" });
@@ -400,6 +457,29 @@ export function bridgeToResponsesSSE(
         retainFinishedItem(item as OutputItem, bytesOf(encrypted), "reasoning");
         outputIndex++;
       };
+      // Kiro reasoning round-trip. Kiro sends its encrypted blob at the END of a turn, while the
+      // assistant message is still open, so this CANNOT emit on arrival: the open message still
+      // owns `outputIndex` (it only advances on close), and an item emitted here would both reuse
+      // that index and land BEFORE the message — where the parser's backwards pairing drops it as
+      // orphaned. Stash it and flush after `done` has closed every open item instead.
+      let pendingKiroRedacted: string | undefined;
+      let pendingKiroRedactedBytes = 0;
+      const flushKiroRedactedReasoning = () => {
+        if (!pendingKiroRedacted) return;
+        const previousBytes = pendingKiroRedactedBytes;
+        const encrypted = encodeReasoningEnvelope({ krc: pendingKiroRedacted });
+        const reservation = budget?.reserveTransient(bytesOf(encrypted), { kind: "reasoning" });
+        pendingKiroRedacted = undefined;
+        pendingKiroRedactedBytes = 0;
+        reservation?.commitRetained();
+        budget?.releaseRetained(previousBytes, { kind: "reasoning" });
+        const itemId = `rs_${uuid()}`;
+        const item = { type: "reasoning", id: itemId, summary: [] as never[], encrypted_content: encrypted };
+        emit("response.output_item.added", { output_index: outputIndex, item });
+        emit("response.output_item.done", { output_index: outputIndex, item });
+        retainFinishedItem(item as OutputItem, bytesOf(encrypted), "reasoning");
+        outputIndex++;
+      };
       // Full assistant text of a compaction turn (across message boundaries) — becomes the
       // synthetic compaction item's payload on done.
       let compactionText = "";
@@ -477,6 +557,7 @@ export function bridgeToResponsesSSE(
 
       const closeCurrentRawReasoning = () => {
         if (!currentRawReasoning) return;
+        rawReasoningForNextToolCall = currentRawReasoning.text;
         const item = {
           type: "reasoning", id: currentRawReasoning.itemId, summary: [],
           content: [{ type: "reasoning_text", text: currentRawReasoning.text }],
@@ -682,6 +763,7 @@ export function bridgeToResponsesSSE(
         beat = undefined;
         try { controller.close(); } catch { /* already closed */ }
         closed = true;
+        disposeOwnedBudget();
         gated = true;
         stepping = false;
       };
@@ -694,6 +776,15 @@ export function bridgeToResponsesSSE(
         while (!terminated && !closed && emittedFrames === emittedAtStart) {
           iteratorStarted = true;
           const next = await it.next();
+          // A cancel during this await disposes the owned budget; a late event
+          // must never be processed or charged against it. Exit step() outright:
+          // falling into EOF synthesis would let closeCurrentMessage() charge
+          // finished-item retention against the disposed budget.
+          if (closed || clientCancelled) {
+            gated = true;
+            stepping = false;
+            return;
+          }
           if (next.done) { upstreamDone = true; break; }
           const event = next.value;
           let terminalEvent = false;
@@ -727,6 +818,7 @@ export function bridgeToResponsesSSE(
               if (currentReasoning) closeCurrentReasoning();
               if (currentRawReasoning) closeCurrentRawReasoning();
               flushHiddenRawReasoning();
+              rawReasoningForNextToolCall = "";
               if (currentToolCall) closeCurrentToolCall();
               flushHiddenReasoningEnvelope();
               break;
@@ -735,6 +827,8 @@ export function bridgeToResponsesSSE(
               if (currentReasoning) closeCurrentReasoning();
               if (currentRawReasoning) closeCurrentRawReasoning();
               flushHiddenRawReasoning();
+              // Reasoning consumed by a text turn, not a tool call: no cache target.
+              rawReasoningForNextToolCall = "";
               if (currentToolCall) closeCurrentToolCall();
               // Only flush on an explicit phase change. A later delta that omits `phase` must
               // keep appending to the current message rather than wiping the earlier phase.
@@ -769,6 +863,12 @@ export function bridgeToResponsesSSE(
             }
             case "thinking_delta": {
               if (options?.hideThinkingSummary) {
+                // The hidden branch returns early, so flush any raw reasoning
+                // that preceded the thinking block and clear the replay-cache
+                // candidate — otherwise a stale reasoning_raw_delta would be
+                // recorded for a LATER tool call (CodeRabbit on #971).
+                flushHiddenRawReasoning();
+                rawReasoningForNextToolCall = "";
                 ({ value: hiddenThinkingText, bytes: hiddenThinkingBytes } = appendString(
                   hiddenThinkingText,
                   hiddenThinkingBytes,
@@ -780,6 +880,7 @@ export function bridgeToResponsesSSE(
               if (currentMsg) closeCurrentMessage("commentary");
               if (currentRawReasoning) closeCurrentRawReasoning();
               flushHiddenRawReasoning();
+              rawReasoningForNextToolCall = "";
               if (currentToolCall) closeCurrentToolCall();
               if (!currentReasoning) {
                 const itemId = `rs_${uuid()}`;
@@ -817,6 +918,12 @@ export function bridgeToResponsesSSE(
               pendingRedacted.push(event.data);
               break;
             }
+            case "kiro_redacted_reasoning": {
+              // Stash only — see flushKiroRedactedReasoning. One blob per turn, so last wins.
+              pendingKiroRedactedBytes = replaceRetainedString(pendingKiroRedactedBytes, event.data, "reasoning");
+              pendingKiroRedacted = event.data;
+              break;
+            }
             case "reasoning_raw_delta": {
               if (options?.hideThinkingSummary) {
                 ({ value: hiddenRawReasoningText, bytes: hiddenRawReasoningBytes } = appendString(
@@ -853,6 +960,9 @@ export function bridgeToResponsesSSE(
               if (currentReasoning) closeCurrentReasoning();
               if (currentRawReasoning) closeCurrentRawReasoning();
               flushHiddenRawReasoning();
+              if (rawReasoningForNextToolCall) {
+                rememberReasoningForCall(event.id, rawReasoningForNextToolCall, replayCacheScope);
+              }
               if (currentToolCall) closeCurrentToolCall();
               const mapped = toolNsMap?.get(event.name);
               const realName = mapped?.name ?? event.name;
@@ -987,6 +1097,9 @@ export function bridgeToResponsesSSE(
               // Redacted-only turns (or hidden thinking without a trailing signature event) still
               // need their envelope-only reasoning item so the blocks replay next turn.
               flushHiddenReasoningEnvelope();
+              // After every close above, so the blob lands AFTER the assistant message it belongs
+              // to and the parser's backwards pairing finds it.
+              flushKiroRedactedReasoning();
               if (options?.compaction) {
                 // Exactly one compaction item per turn; codex-rs takes the first and fatals on 0.
                 const item = {
@@ -1141,6 +1254,7 @@ export function bridgeToResponsesSSE(
         /* already closed (e.g. client cancelled) */
       }
       closed = true;
+      disposeOwnedBudget();
       gated = true;
       stepping = false;
       };
@@ -1176,6 +1290,7 @@ export function bridgeToResponsesSSE(
             beat = undefined;
             try { controller.close(); } catch { /* already closed */ }
             closed = true;
+            disposeOwnedBudget();
             return;
           }
           // Wire silence is independent of upstream adapter heartbeats.
@@ -1188,6 +1303,7 @@ export function bridgeToResponsesSSE(
             emittedFrames++;
           } catch {
             closed = true;
+            disposeOwnedBudget();
           }
         }, heartbeatMs);
       };
@@ -1205,13 +1321,31 @@ export function bridgeToResponsesSSE(
       // cancelled turn does not leak the upstream stream or keep draining tokens (RC2).
         clientCancelled = true;
         closed = true;
+        clearOwnedWatchdog();
         if (beat !== undefined) clearBeatInterval(beat);
         cancelUpstreamOnce();
+        disposeOwnedBudget();
       },
     });
   }
 
 export function buildResponseJSON(
+  events: AdapterEvent[],
+  modelId: string,
+  options?: Parameters<typeof buildResponseJSONWithBudget>[2],
+): Record<string, unknown> {
+  // Default-budget safety net: a caller that omits the budget gets a bounded
+  // default (disposed with the call), never the unbounded append path.
+  if (options?.translatorBudget) return buildResponseJSONWithBudget(events, modelId, options);
+  const budget = createTranslatorBudget();
+  try {
+    return buildResponseJSONWithBudget(events, modelId, { ...options, translatorBudget: budget });
+  } finally {
+    budget.dispose();
+  }
+}
+
+function buildResponseJSONWithBudget(
   events: AdapterEvent[],
   modelId: string,
   options?: {
@@ -1225,9 +1359,12 @@ export function buildResponseJSON(
     /** Raw adapter-reported usage before wire normalization (see bridgeToResponsesSSE onUsage). */
     onUsage?: (usage: OcxUsage | undefined) => void;
     translatorBudget?: TranslatorBudget;
+    /** Conversation identity for the reasoning replay cache (issue #950). */
+    replayCacheScope?: string;
   },
 ): Record<string, unknown> {
   const responseId = `resp_${uuid()}`;
+  const replayCacheScope = options?.replayCacheScope ?? "global";
   const output: OutputItem[] = [];
   const budget = options?.translatorBudget;
   const encoder = new TextEncoder();
@@ -1283,11 +1420,18 @@ export function buildResponseJSON(
   let currentSummaryReasoningBytes = 0;
   let currentRawReasoning = "";
   let currentRawReasoningBytes = 0;
+  // Same replay-cache handoff as the streaming path (issue #950): the most
+  // recently flushed raw reasoning waits for the tool call it preceded.
+  let rawReasoningForNextToolCall = "";
   // Anthropic extended-thinking round-trip (batch): see bridgeToResponsesSSE counterpart.
   let batchSignature: string | undefined;
   let batchSignatureBytes = 0;
   let batchRedacted: string[] = [];
   let batchRedactedBytes = 0;
+  // Kiro reasoning blob, held until after the trailing flushes so it lands AFTER the assistant
+  // message (see the streaming path). Retained because it outlives releaseTranslatedEvent.
+  let batchKiroRedacted: string | undefined;
+  let batchKiroRedactedBytes = 0;
   let currentToolCallId = "";
   let currentToolCallName = "";
   let currentToolCallArgs = "";
@@ -1350,6 +1494,7 @@ export function buildResponseJSON(
   };
   const flushRawReasoning = () => {
     if (!currentRawReasoning) return;
+    rawReasoningForNextToolCall = currentRawReasoning;
     if (options?.hideThinkingSummary === true) {
       // Same contract as the streaming path: no visible reasoning, txt-only envelope round-trip.
       pushOutput({
@@ -1407,6 +1552,7 @@ export function buildResponseJSON(
         flushText("commentary");
         flushSummaryReasoning();
         flushRawReasoning();
+        rawReasoningForNextToolCall = "";
         flushToolCall();
         break;
       case "text_delta":
@@ -1415,6 +1561,7 @@ export function buildResponseJSON(
         if (currentText && e.phase !== undefined && currentTextPhase !== e.phase) flushText("commentary");
         if (currentSummaryReasoning) flushSummaryReasoning();
         if (currentRawReasoning) flushRawReasoning();
+        rawReasoningForNextToolCall = "";
         if (currentToolCallId) flushToolCall();
         // Compaction turns keep the summary out of normal message output (replay dedup — see
         // bridgeToResponsesSSE); it ships only inside the synthetic compaction item below.
@@ -1433,6 +1580,7 @@ export function buildResponseJSON(
       case "thinking_delta":
         if (currentText) flushText("commentary");
         if (currentRawReasoning) flushRawReasoning();
+        rawReasoningForNextToolCall = "";
         if (currentToolCallId) flushToolCall();
         {
           ({ value: currentSummaryReasoning, bytes: currentSummaryReasoningBytes } = appendBatchString(
@@ -1455,6 +1603,16 @@ export function buildResponseJSON(
         }
         batchRedacted.push(e.data);
         break;
+      case "kiro_redacted_reasoning":
+        // Stash only — pushed after the trailing flushes. One blob per turn, so last wins.
+        {
+          const dataBytes = bytesOf(e.data);
+          budget?.chargeRetained(dataBytes, { kind: "reasoning" });
+          if (batchKiroRedactedBytes > 0) budget?.releaseRetained(batchKiroRedactedBytes, { kind: "reasoning" });
+          batchKiroRedactedBytes = dataBytes;
+        }
+        batchKiroRedacted = e.data;
+        break;
       case "reasoning_raw_delta":
         if (currentText) flushText("commentary");
         if (currentSummaryReasoning) flushSummaryReasoning();
@@ -1469,6 +1627,9 @@ export function buildResponseJSON(
         if (currentText) flushText("commentary");
         if (currentSummaryReasoning) flushSummaryReasoning();
         if (currentRawReasoning) flushRawReasoning();
+        if (rawReasoningForNextToolCall) {
+          rememberReasoningForCall(e.id, rawReasoningForNextToolCall, replayCacheScope);
+        }
         flushToolCall();
         currentToolCallId = e.id;
         budget?.openCall(e.id);
@@ -1553,6 +1714,15 @@ export function buildResponseJSON(
   flushRawReasoning();
   // Open tool call on a failed/incomplete turn must not land as status:"completed".
   if (currentToolCallId) flushToolCall(errorEvent || incompleteEvent ? "incomplete" : "completed");
+  if (batchKiroRedacted) {
+    // pushOutput reserves the item itself and releases the retained raw blob it replaces.
+    pushOutput({
+      type: "reasoning", id: `rs_${uuid()}`, summary: [],
+      encrypted_content: encodeReasoningEnvelope({ krc: batchKiroRedacted }),
+    }, batchKiroRedactedBytes, "reasoning");
+    batchKiroRedacted = undefined;
+    batchKiroRedactedBytes = 0;
+  }
   // A truncated turn must never be installed as replacement history: emit the
   // compaction item only when the turn actually completed (#422).
   if (

package/src/chat/outbound.ts

@@ -319,7 +319,9 @@ export function responsesSseToChatCompletionsSse(
           closeToolCalls();
           try { void sseIterator?.return(undefined).catch(() => {}); } catch { /* already closed */ }
           classified.code = "translation_buffer_limit";
-          classified.type = "invalid_request_error";
+          // Provider-controlled overflow is an upstream failure on every path:
+          // streaming frame, collector, and defensive JSON agree on 502.
+          classified.type = "upstream_error";
         } else if (isCyberPolicyCode(details?.code) || classified.code === CYBER_POLICY_ERROR_CODE) {
           classified.code = CYBER_POLICY_ERROR_CODE;
           classified.type = "invalid_request_error";
@@ -468,7 +470,7 @@ export function responsesSseToChatCompletionsSse(
             fail(message, {
               code,
               ...(code === "translation_buffer_limit"
-                ? { status: 413, type: "invalid_request_error" }
+                ? { status: 502, type: "upstream_error" }
                 : { type, ...(code === CYBER_POLICY_ERROR_CODE ? { status: 400 } : {}) }),
             });
             break;
@@ -514,7 +516,7 @@ export function responsesSseToChatCompletionsSse(
           if (isTranslatorBudgetExceededError(err)) {
             upstreamAbort.abort(err);
             closeToolCalls();
-            fail(err.message, { status: 413, type: "invalid_request_error", code: err.code });
+            fail(err.message, { status: 502, type: "upstream_error", code: err.code });
           } else {
             fail(err instanceof Error ? err.message : String(err));
           }
@@ -629,6 +631,9 @@ export async function collectChatCompletion(
   let content = "";
   let reasoning = "";
   const toolCalls = new Map<number, { id: string; name: string; arguments: string; argumentBytes: number }>();
+  // Per-call budget scopes (2 MiB/call enforced by the budget): the map key is the
+  // wire index, which is stable across deltas and present before the call id.
+  const callScope = (index: number) => `chat_collect_${index}`;
   let finishReason = "stop";
   let usage: unknown;
   let streamError: ChatCompletionsStreamError | null = null;
@@ -648,9 +653,11 @@ export async function collectChatCompletion(
       } catch (err) {
         if (isChatCompletionsStreamError(err)) throw err;
         if (isTranslatorBudgetExceededError(err)) {
+          // Provider-controlled overflow is an upstream failure, not a client
+          // request error: match the adapter/bridge contract (502 upstream_error).
           throw new ChatCompletionsStreamError(err.message, {
-            status: 413,
-            type: "invalid_request_error",
+            status: 502,
+            type: "upstream_error",
             code: err.code,
           });
         }
@@ -677,11 +684,15 @@ export async function collectChatCompletion(
             const type = typeof parsed.error.type === "string" ? parsed.error.type : "server_error";
             const code = typeof parsed.error.code === "string" ? parsed.error.code : null;
             const status = code === "translation_buffer_limit"
-              ? 413
+              ? 502
               : code === CYBER_POLICY_ERROR_CODE || isCyberPolicyMessage(message)
                 ? 400
                 : streamErrorStatus(message);
-            streamError = new ChatCompletionsStreamError(message, { status, type, code });
+            streamError = new ChatCompletionsStreamError(message, {
+              status,
+              type: code === "translation_buffer_limit" ? "upstream_error" : type,
+              code,
+            });
             continue;
           }
           if (parsed.usage) usage = parsed.usage;
@@ -697,7 +708,12 @@ export async function collectChatCompletion(
             for (const tc of delta.tool_calls) {
               if (!isRec(tc)) continue;
               const index = typeof tc.index === "number" ? tc.index : 0;
-              const current = toolCalls.get(index) ?? { id: "", name: "", arguments: "", argumentBytes: 0 };
+              let current = toolCalls.get(index);
+              if (!current) {
+                current = { id: "", name: "", arguments: "", argumentBytes: 0 };
+                toolCalls.set(index, current);
+                translatorBudget.openCall(callScope(index));
+              }
               if (typeof tc.id === "string") current.id = tc.id;
               const fn = isRec(tc.function) ? tc.function : {};
               // Done-frame final arguments are authoritative last-write-wins snapshots.
@@ -707,27 +723,43 @@ export async function collectChatCompletion(
                 const nextBytes = replace
                   ? Buffer.byteLength(fn.arguments)
                   : appendedUtf8Bytes(current.arguments, current.argumentBytes, fn.arguments);
-                const reservation = translatorBudget.reserveTransient(nextBytes, { kind: "retained_collectors" });
+                const reservation = translatorBudget.reserveTransient(nextBytes, { kind: "tool_args", callId: callScope(index) });
                 try {
                   current.arguments = replace ? fn.arguments : current.arguments + fn.arguments;
                   reservation.commitRetained();
-                  translatorBudget.releaseRetained(current.argumentBytes, { kind: "retained_collectors" });
+                  translatorBudget.releaseRetained(current.argumentBytes, { kind: "tool_args", callId: callScope(index) });
                   current.argumentBytes = nextBytes;
                 } catch (error) {
                   reservation.release();
                   throw error;
                 }
               }
-              toolCalls.set(index, current);
             }
           }
         }
       }
     }
+  } catch (error) {
+    // Never leak an open call scope on the error path; the turn budget's
+    // dispose is a backstop, not the owner of this transfer.
+    for (const index of toolCalls.keys()) translatorBudget.closeCall(callScope(index));
+    // Processing-time overflow (per-call or turn cap) gets the same typed
+    // contract as read-time overflow: 502 upstream_error.
+    if (isTranslatorBudgetExceededError(error)) {
+      throw new ChatCompletionsStreamError(error.message, {
+        status: 502,
+        type: "upstream_error",
+        code: error.code,
+      });
+    }
+    throw error;
   } finally {
     reader.releaseLock();
   }
-  if (streamError) throw streamError;
+  if (streamError) {
+    for (const index of toolCalls.keys()) translatorBudget.closeCall(callScope(index));
+    throw streamError;
+  }
 
   const message: Rec = {
     role: "assistant",
@@ -735,17 +767,40 @@ export async function collectChatCompletion(
   };
   if (reasoning) message.reasoning_content = reasoning;
   if (toolCalls.size > 0) {
-    message.tool_calls = [...toolCalls.entries()]
-      .sort((a, b) => a[0] - b[0])
-      .map(([, tc]) => {
-        const copy = {
-        id: tc.id || `call_${uuid().slice(0, 16)}`,
-        type: "function",
-        function: { name: tc.name, arguments: tc.arguments },
-        };
-        translatorBudget.chargeRetained(Buffer.byteLength(JSON.stringify(copy)), { kind: "retained_collectors" });
-        return copy;
-      });
+    // Final owner transfer is itself a charging operation: if it overflows,
+    // release every copy already charged in this loop and close every scope
+    // still open — the error must not escape as a raw budget exception.
+    const chargedCopies: number[] = [];
+    try {
+      message.tool_calls = [...toolCalls.entries()]
+        .sort((a, b) => a[0] - b[0])
+        .map(([index, tc]) => {
+          const copy = {
+          id: tc.id || `call_${uuid().slice(0, 16)}`,
+          type: "function",
+          function: { name: tc.name, arguments: tc.arguments },
+          };
+          const copyBytes = Buffer.byteLength(JSON.stringify(copy));
+          translatorBudget.chargeRetained(copyBytes, { kind: "retained_collectors" });
+          chargedCopies.push(copyBytes);
+          // The serialized owner is charged; release the per-call accumulation.
+          translatorBudget.closeCall(callScope(index));
+          return copy;
+        });
+    } catch (error) {
+      for (const copyBytes of chargedCopies) {
+        translatorBudget.releaseRetained(copyBytes, { kind: "retained_collectors" });
+      }
+      for (const index of toolCalls.keys()) translatorBudget.closeCall(callScope(index));
+      if (isTranslatorBudgetExceededError(error)) {
+        throw new ChatCompletionsStreamError(error.message, {
+          status: 502,
+          type: "upstream_error",
+          code: error.code,
+        });
+      }
+      throw error;
+    }
     if (finishReason === "stop") finishReason = "tool_calls";
   }