[WRONG BRANCH] fix(responses): bound retained tool-search item IDs

src/server/responses-tool-search-repair.ts

@@ -1,5 +1,6 @@
 import {
   isTranslatorBudgetExceededError,
+  TranslatorBudgetExceededError,
   type TranslatorBudget,
 } from "../lib/translator-budget";
 import {
@@ -24,6 +25,8 @@ type PendingArgumentBlock = {
 
 const MAX_PENDING_ARGUMENT_FRAMES = 256;
 const MAX_PENDING_ARGUMENT_BYTES = 1024 * 1024;
+const MAX_CLASSIFIED_ITEM_IDS = 256;
+const MAX_CLASSIFIED_ITEM_ID_BYTES = 256 * 1024;
 
 /**
  * Public Responses gateways stream a lowered search as a normal function lifecycle. Codex expects
@@ -36,6 +39,7 @@ export function createRoutedToolSearchRestoreBlockRewrite(
 ): SseBlockRewrite {
   const routedItemIds = new Set<string>();
   const ordinaryItemIds = new Set<string>();
+  let classifiedItemIdBytes = 0;
   let pendingArguments: PendingArgumentBlock[] = [];
   let pendingArgumentBytes = 0;
   let passthrough = false;
@@ -49,10 +53,44 @@ export function createRoutedToolSearchRestoreBlockRewrite(
     }
     pendingArguments = [];
     pendingArgumentBytes = 0;
+    if (classifiedItemIdBytes > 0) {
+      budget?.releaseRetained(classifiedItemIdBytes, { kind: "item_ids" });
+    }
+    classifiedItemIdBytes = 0;
     routedItemIds.clear();
     ordinaryItemIds.clear();
   };
 
+  const clearOrdinaryItemIds = (): void => {
+    let releasedBytes = 0;
+    for (const itemId of ordinaryItemIds) {
+      releasedBytes += Buffer.byteLength(JSON.stringify(itemId), "utf8");
+    }
+    ordinaryItemIds.clear();
+    classifiedItemIdBytes = Math.max(0, classifiedItemIdBytes - releasedBytes);
+    if (releasedBytes > 0) budget?.releaseRetained(releasedBytes, { kind: "item_ids" });
+  };
+
+  const classifyItemId = (itemId: string, routed: boolean): void => {
+    const target = routed ? routedItemIds : ordinaryItemIds;
+    const previous = routed ? ordinaryItemIds : routedItemIds;
+    if (target.has(itemId)) return;
+    if (previous.delete(itemId)) {
+      target.add(itemId);
+      return;
+    }
+    const retainedBytes = Buffer.byteLength(JSON.stringify(itemId), "utf8");
+    if (
+      routedItemIds.size + ordinaryItemIds.size >= MAX_CLASSIFIED_ITEM_IDS
+      || classifiedItemIdBytes + retainedBytes > MAX_CLASSIFIED_ITEM_ID_BYTES
+    ) {
+      throw new TranslatorBudgetExceededError("item_ids", MAX_CLASSIFIED_ITEM_ID_BYTES);
+    }
+    budget?.chargeRetained(retainedBytes, { kind: "item_ids" });
+    target.add(itemId);
+    classifiedItemIdBytes += retainedBytes;
+  };
+
   const retainPending = (
     block: string,
     itemId: string | undefined,
@@ -73,7 +111,7 @@ export function createRoutedToolSearchRestoreBlockRewrite(
       // that we forget what we already classified: an item restored to `tool_search_call`
       // upstream of here would otherwise start emitting `function_call_arguments.*` again and
       // the client would see a mixed private/public lifecycle for one call.
-      ordinaryItemIds.clear();
+      clearOrdinaryItemIds();
       return flushed;
     }
     if (retainedBytes > 0) {
@@ -90,7 +128,7 @@ export function createRoutedToolSearchRestoreBlockRewrite(
         passthrough = true;
         // Same reasoning as the frame/byte overflow above: an already-restored routed item
         // must keep its frames suppressed even once buffering stops.
-        ordinaryItemIds.clear();
+        clearOrdinaryItemIds();
         return flushed;
       }
     }
@@ -170,21 +208,15 @@ export function createRoutedToolSearchRestoreBlockRewrite(
       const itemId = typeof parsed.item.id === "string" ? parsed.item.id : undefined;
       const routed = names.has(parsed.item.name);
       if (itemId) {
-        if (routed) {
-          routedItemIds.add(itemId);
-          ordinaryItemIds.delete(itemId);
-        } else {
-          ordinaryItemIds.add(itemId);
-          routedItemIds.delete(itemId);
-        }
+        classifyItemId(itemId, routed);
       }
       const pending = takePending(itemId, outputIndex);
       const restored = routed ? restoreRoutedToolSearchCalls(parsed, names) : { value: parsed, changed: false };
       const restoredBlock = restored.changed
         ? replaceSseDataPayload(block, JSON.stringify(restored.value))
         : block;
       // Classification is retained past `output_item.done` for BOTH kinds, until the terminal
-      // event releases everything.
+      // event releases the bounded, budgeted state.
       //
       // `done` ends the item, not the id's relevance. Forgetting a ROUTED id let a trailing
       // `function_call_arguments.*` — which some upstreams emit after done — fall through to

tests/responses-tool-search-repair.test.ts

@@ -237,7 +237,7 @@ describe("routed Responses tool-search compatibility", () => {
       arguments: {},
       status: "in_progress",
     });
-    expect(budget.snapshot().currentBytes).toBe(0);
+    expect(budget.snapshot().currentBytes).toBeGreaterThan(0);
 
     expect(rewrite(frame("response.function_call_arguments.done", {
       output_index: 0,
@@ -628,4 +628,25 @@ describe("routed Responses tool-search compatibility", () => {
     }));
     expect(trailing).toHaveLength(1);
   });
+
+  test("bounds and budgets item ids retained for trailing argument frames", () => {
+    const budget = createTestTranslatorBudget();
+    const rewrite = createRoutedToolSearchRestoreBlockRewrite(new Set(["tool_search"]), budget);
+
+    for (let index = 0; index < 256; index += 1) {
+      rewrite(frame("response.output_item.done", {
+        output_index: index,
+        item: { type: "function_call", id: `fc_${index}`, name: "tool_search", arguments: "{}" },
+      }));
+    }
+
+    expect(budget.snapshot().currentBytes).toBeGreaterThan(0);
+    expect(() => rewrite(frame("response.output_item.done", {
+      output_index: 256,
+      item: { type: "function_call", id: "fc_overflow", name: "tool_search", arguments: "{}" },
+    }))).toThrow(/translator item_ids buffer exceeded/);
+
+    rewrite.dispose?.();
+    expect(budget.snapshot().currentBytes).toBe(0);
+  });
 });