diff --git a/src/server/responses-tool-search-repair.ts b/src/server/responses-tool-search-repair.ts index 6c18708406..05673ee8f1 100644 --- a/src/server/responses-tool-search-repair.ts +++ b/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(); const ordinaryItemIds = new Set(); + 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,13 +208,7 @@ 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 }; @@ -184,7 +216,7 @@ export function createRoutedToolSearchRestoreBlockRewrite( ? 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 diff --git a/tests/responses-tool-search-repair.test.ts b/tests/responses-tool-search-repair.test.ts index a080574a3d..db9bee12af 100644 --- a/tests/responses-tool-search-repair.test.ts +++ b/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); + }); });