diff --git a/src/web-search/passthrough-bridge.ts b/src/web-search/passthrough-bridge.ts index 6212e25c82..0e5afdc7ce 100644 --- a/src/web-search/passthrough-bridge.ts +++ b/src/web-search/passthrough-bridge.ts @@ -53,6 +53,9 @@ const MAX_QUERIES_PER_CALL = 3; const MAX_RETAINED_OUTPUT_ITEMS = 500; /** Refuse to buffer an unbounded partial SSE event from a misbehaving upstream. */ const MAX_SSE_BUFFER_CHARS = 8 * 1024 * 1024; +/** Bound client-tool events withheld while the bridge determines whether a leg can succeed. */ +const MAX_HELD_CALL_EVENTS = 1_000; +const MAX_HELD_CALL_CHARS = 8 * 1024 * 1024; export const WEB_SEARCH_BRIDGE_MIXED_TOOLS_ERROR_CODE = "web_search_bridge_mixed_tools"; export const WEB_SEARCH_BRIDGE_ERROR_CODE = "web_search_bridge_failed"; @@ -302,6 +305,7 @@ class BridgeStreamState { * failing the turn would let Codex start running a tool for a turn that never completes. */ private heldCalls: HeldCallEvent[] = []; + private heldCallChars = 0; private heldIndexes = new Set(); private heldItemIds = new Set(); private terminalPayload: Record | undefined; @@ -312,6 +316,7 @@ class BridgeStreamState { this.suppressedItemIds = new Map(); this.searches = []; this.heldCalls = []; + this.heldCallChars = 0; this.heldIndexes = new Set(); this.heldItemIds = new Set(); this.terminalPayload = undefined; @@ -357,6 +362,15 @@ class BridgeStreamState { + "data: " + JSON.stringify({ ...data, type, sequence_number: this.sequence++ }); } + private holdCall(payload: Record, dataChars: number, upstreamIndex?: number): void { + if (this.heldCalls.length >= MAX_HELD_CALL_EVENTS + || dataChars > MAX_HELD_CALL_CHARS - this.heldCallChars) { + throw new Error("upstream client tool events exceeded the web-search bridge buffer bound"); + } + this.heldCalls.push({ payload, ...(upstreamIndex === undefined ? {} : { upstreamIndex }) }); + this.heldCallChars += dataChars; + } + failureFrames(code: string, message: string): string[] { const failure = { type: "upstream_error", code, message }; return [ @@ -440,7 +454,7 @@ class BridgeStreamState { if (isClientExecutedItem(item)) { if (upstreamIndex !== undefined) this.heldIndexes.add(upstreamIndex); if (typeof item.id === "string") this.heldItemIds.add(item.id); - this.heldCalls.push({ payload, ...(upstreamIndex === undefined ? {} : { upstreamIndex }) }); + this.holdCall(payload, data.length, upstreamIndex); return []; } } @@ -462,7 +476,7 @@ class BridgeStreamState { if ((upstreamIndex !== undefined && this.heldIndexes.has(upstreamIndex)) || (itemId !== undefined && this.heldItemIds.has(itemId))) { - this.heldCalls.push({ payload, ...(upstreamIndex === undefined ? {} : { upstreamIndex }) }); + this.holdCall(payload, data.length, upstreamIndex); return []; } @@ -473,18 +487,20 @@ class BridgeStreamState { } /** Release the withheld client tool calls once the turn is known to end here. */ - flushHeldCalls(): string[] { - const blocks: string[] = []; - for (const held of this.heldCalls) { - const rewritten: Record = { ...held.payload }; - if (held.upstreamIndex !== undefined) { - rewritten.output_index = this.clientIndexFor(held.upstreamIndex); + *flushHeldCalls(): Generator { + try { + for (const held of this.heldCalls) { + const rewritten: Record = { ...held.payload }; + if (held.upstreamIndex !== undefined) { + rewritten.output_index = this.clientIndexFor(held.upstreamIndex); + } + if (held.payload.type === "response.output_item.done") this.retain(held.payload.item); + yield this.render(String(held.payload.type), rewritten); } - if (held.payload.type === "response.output_item.done") this.retain(held.payload.item); - blocks.push(this.render(String(held.payload.type), rewritten)); + } finally { + this.heldCalls = []; + this.heldCallChars = 0; } - this.heldCalls = []; - return blocks; } /** Decide what the leg's terminal means once the whole leg has been read. */ @@ -613,7 +629,7 @@ async function* bridgeStreamBlocks( // One continuation leg per allowed search, plus one final leg for the answer itself. let legsRemaining = options.plan.maxSearches + 1; - const emit = function* (blocks: readonly string[]): Generator { + const emit = function* (blocks: Iterable): Generator { for (const block of blocks) yield block + "\n\n"; }; diff --git a/tests/web-search/web-search-passthrough-bridge.test.ts b/tests/web-search/web-search-passthrough-bridge.test.ts index 9e7029bb53..398860afac 100644 --- a/tests/web-search/web-search-passthrough-bridge.test.ts +++ b/tests/web-search/web-search-passthrough-bridge.test.ts @@ -385,6 +385,41 @@ describe("the bridged client stream", () => { expect((cell!.item as Record).status).toBe("failed"); }); + test("bounds client tool events withheld from a hostile upstream", async () => { + const clientCall = { + type: "function_call", + id: "fc_attack", + call_id: "call_attack", + name: "exec", + arguments: "", + }; + const blocks = [ + frame("response.output_item.added", { output_index: 0, item: clientCall }), + ]; + for (let index = 0; index < 1_000; index += 1) { + blocks.push(frame("response.function_call_arguments.delta", { + output_index: 0, + item_id: "fc_attack", + delta: "x", + })); + } + + const stream = createPassthroughWebSearchBridgeStream({ + plan, + firstLeg: streamFromText(sseBody(...blocks)), + requestBody: initialBody, + send: async () => new Response(null, { status: 500 }), + execute: async () => ({ text: "unused", sources: [] }), + }); + + const body = await new Response(stream).text(); + expect(body).not.toContain("\"name\":\"exec\""); + const failed = clientEvents(body).find(event => event.type === "response.failed"); + const error = (failed!.response as { error: Record }).error; + expect(error.code).toBe(WEB_SEARCH_BRIDGE_ERROR_CODE); + expect(String(error.message)).toContain("client tool events exceeded"); + }); + test("a search that is not the last item keeps its streamed position", async () => { // The model searches first and keeps talking; the hosted cell must open where the call stood. const leg = sseBody(