From 7d92aa7527206cdd47df29874da01384dfa91f8c Mon Sep 17 00:00:00 2001 From: Jan Leon Date: Tue, 18 Aug 2026 15:53:05 +0200 Subject: [PATCH] fix(streaming): preserve completed Codex tool handoffs (#10608) --- open-sse/handlers/chatCore.ts | 2 + open-sse/utils/responsesToolHandoff.ts | 132 ++++++ open-sse/utils/streamHandler.ts | 71 +++- .../codex-tool-handoff-disconnect-499.test.ts | 377 ++++++++++++++++++ 4 files changed, 581 insertions(+), 1 deletion(-) create mode 100644 open-sse/utils/responsesToolHandoff.ts create mode 100644 tests/unit/codex-tool-handoff-disconnect-499.test.ts diff --git a/open-sse/handlers/chatCore.ts b/open-sse/handlers/chatCore.ts index 3c8b286d90..a8d60759ee 100644 --- a/open-sse/handlers/chatCore.ts +++ b/open-sse/handlers/chatCore.ts @@ -2877,6 +2877,8 @@ export async function handleChatCore({ connectionId, clientResponseFormat, clientAbortSignal: clientRawRequest?.signal, + allowCompletedToolHandoffGrace: isCodexResponsesEcho, + clientDisconnectGracePeriodMs: STREAM_DISCONNECT_GRACE_PERIOD_MS, }); const dedupRequestBody = { ...translatedBody, model: `${provider}/${model}`, stream }; diff --git a/open-sse/utils/responsesToolHandoff.ts b/open-sse/utils/responsesToolHandoff.ts new file mode 100644 index 0000000000..e9d80a954c --- /dev/null +++ b/open-sse/utils/responsesToolHandoff.ts @@ -0,0 +1,132 @@ +type CompletedToolItem = { + keys: string[]; + type: "function_call" | "custom_tool_call"; + value: string; +}; + +function getResponsesEventKeys( + payload: Record, + item?: Record +): string[] { + const keys = new Set(); + const addStringKey = (prefix: string, value: unknown) => { + if (typeof value === "string" && value.trim()) keys.add(`${prefix}:${value.trim()}`); + }; + const addIndexKey = (value: unknown) => { + if (typeof value === "number" && Number.isInteger(value) && value >= 0) { + keys.add(`index:${value}`); + } + }; + + addStringKey("item", payload.item_id); + addStringKey("call", payload.call_id); + addIndexKey(payload.output_index); + if (item) { + addStringKey("item", item.id); + addStringKey("call", item.call_id); + } + return [...keys]; +} + +/** + * Codex can start its next turn as soon as it receives a complete client-side + * tool call, closing the current HTTP response before response.completed. This + * watcher accepts only a matching done-payload plus a completed tool item; + * ordinary message/reasoning items and partial calls never qualify. + */ +export function createCompletedResponsesToolHandoffWatcher() { + let buffer = ""; + let completed = false; + const functionArgumentsDone = new Map(); + const customToolInputDone = new Map(); + const completedToolItems: CompletedToolItem[] = []; + + const matchesDonePayload = (item: CompletedToolItem): boolean => { + const doneValues = item.type === "function_call" ? functionArgumentsDone : customToolInputDone; + return item.keys.some((key) => doneValues.get(key) === item.value); + }; + + const evaluate = () => { + completed = completed || completedToolItems.some(matchesDonePayload); + }; + + const notePayload = (payload: Record, eventType: string) => { + if ( + eventType === "response.function_call_arguments.done" && + typeof payload.arguments === "string" + ) { + for (const key of getResponsesEventKeys(payload)) { + functionArgumentsDone.set(key, payload.arguments); + } + evaluate(); + return; + } + + if (eventType === "response.custom_tool_call_input.done" && typeof payload.input === "string") { + for (const key of getResponsesEventKeys(payload)) { + customToolInputDone.set(key, payload.input); + } + evaluate(); + return; + } + + if (eventType !== "response.output_item.done") return; + const item = + payload.item && typeof payload.item === "object" && !Array.isArray(payload.item) + ? (payload.item as Record) + : null; + if (!item) return; + if (item.type !== "function_call" && item.type !== "custom_tool_call") return; + if (typeof item.call_id !== "string" || !item.call_id.trim()) return; + if (typeof item.name !== "string" || !item.name.trim()) return; + if (item.status !== undefined && item.status !== "completed") return; + + const valueKey = item.type === "function_call" ? "arguments" : "input"; + const value = item[valueKey]; + if (typeof value !== "string") return; + const keys = getResponsesEventKeys(payload, item); + if (keys.length === 0) return; + + completedToolItems.push({ keys, type: item.type, value }); + if (completedToolItems.length > 32) completedToolItems.shift(); + evaluate(); + }; + + const noteFrame = (frame: string) => { + let eventType = ""; + const dataLines: string[] = []; + for (const rawLine of frame.split(/\r?\n/)) { + const line = rawLine.trimStart(); + if (line.startsWith("event:")) { + eventType = line.slice("event:".length).trim(); + } else if (line.startsWith("data:")) { + dataLines.push(line.slice("data:".length).trimStart()); + } + } + if (dataLines.length === 0) return; + + try { + const parsed = JSON.parse(dataLines.join("\n")); + if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) return; + const payload = parsed as Record; + notePayload(payload, typeof payload.type === "string" ? payload.type : eventType); + } catch { + // A partial/malformed frame is not evidence of a completed tool handoff. + } + }; + + return { + note(text: string): boolean { + if (completed) return true; + buffer += text; + let boundary = /\r?\n\r?\n/.exec(buffer); + while (boundary) { + noteFrame(buffer.slice(0, boundary.index)); + buffer = buffer.slice(boundary.index + boundary[0].length); + boundary = /\r?\n\r?\n/.exec(buffer); + } + if (buffer.length > 65_536) buffer = buffer.slice(-65_536); + return completed; + }, + }; +} diff --git a/open-sse/utils/streamHandler.ts b/open-sse/utils/streamHandler.ts index 4e7bf382dd..e79c858b13 100644 --- a/open-sse/utils/streamHandler.ts +++ b/open-sse/utils/streamHandler.ts @@ -2,6 +2,7 @@ import { trackPendingRequest } from "@/lib/usageDb"; import { STREAM_IDLE_TIMEOUT_MS } from "../config/constants.ts"; import { FORMATS } from "../translator/formats.ts"; import { PENDING_REQUEST_CLEARED_MARKER } from "./stream.ts"; +import { createCompletedResponsesToolHandoffWatcher } from "./responsesToolHandoff.ts"; import { createStreamContentWatcher, type StreamContentWatcher } from "./streamReadiness.ts"; // Stream handler with disconnect detection - shared for all providers @@ -36,6 +37,8 @@ type StreamControllerOptions = { connectionId?: string | null; clientResponseFormat?: string | null; clientAbortSignal?: AbortSignal | null; + allowCompletedToolHandoffGrace?: boolean; + clientDisconnectGracePeriodMs?: number; }; type StreamController = ReturnType; @@ -238,11 +241,15 @@ export function createStreamController({ connectionId, clientResponseFormat, clientAbortSignal, + allowCompletedToolHandoffGrace = false, + clientDisconnectGracePeriodMs = 0, }: StreamControllerOptions = {}) { const abortController = new AbortController(); const startTime = Date.now(); let disconnected = false; let clientTerminalSeen = false; + let completedToolHandoffSeen = false; + let completedToolHandoffDrain: (() => void) | null = null; let pendingRequestCleared = false; let cleanupClientAbortSignal: (() => void) | null = null; @@ -316,7 +323,16 @@ export function createStreamController({ // fire when the client aborts mid-stream, so we must clean up here. clearPendingRequest(); - abortController.abort(reason); + const deferUpstreamAbort = + allowCompletedToolHandoffGrace && + clientDisconnectGracePeriodMs > 0 && + completedToolHandoffSeen && + completedToolHandoffDrain !== null; + if (deferUpstreamAbort) { + completedToolHandoffDrain?.(); + } else { + abortController.abort(reason); + } onDisconnect?.({ reason, duration: Date.now() - startTime }); }, @@ -334,6 +350,20 @@ export function createStreamController({ clientTerminalSeen = true; }, + markCompletedToolHandoffSeen: () => { + completedToolHandoffSeen = true; + }, + + registerCompletedToolHandoffDrain: (drain: () => void) => { + completedToolHandoffDrain = drain; + }, + + shouldDeferCompletedToolHandoff: () => + allowCompletedToolHandoffGrace && + clientDisconnectGracePeriodMs > 0 && + completedToolHandoffSeen && + completedToolHandoffDrain !== null, + // Call on error handleError: (error: unknown) => { cleanupClientAbortListener(); @@ -387,6 +417,7 @@ export function createStreamController({ abortController.abort(); }, clientResponseFormat, + clientDisconnectGracePeriodMs, }; if (clientAbortSignal && typeof clientAbortSignal.addEventListener === "function") { @@ -556,9 +587,38 @@ export function createDisconnectAwareStream(transformStream, streamController) { const terminalDecoder = new TextDecoder(); const contentDecoder = new TextDecoder(); const contentWatcher = createStreamContentWatcher(); + const completedToolHandoffWatcher = createCompletedResponsesToolHandoffWatcher(); + const toolHandoffDecoder = new TextDecoder(); let terminalTail = ""; let clientTerminalSeen = false; let bytesWereForwarded = false; + let completedToolHandoffDrainStarted = false; + + const drainCompletedToolHandoff = () => { + if (completedToolHandoffDrainStarted) return; + completedToolHandoffDrainStarted = true; + const gracePeriodMs = Math.max(0, Number(streamController.clientDisconnectGracePeriodMs) || 0); + const timeoutReason = "completed_tool_handoff_grace_expired"; + const timeout = setTimeout(() => { + streamController.abort(); + void Promise.allSettled([reader.cancel(timeoutReason), writer.abort(timeoutReason)]); + }, gracePeriodMs); + + void (async () => { + try { + while (true) { + const { done } = await reader.read(); + if (done) break; + } + streamController.handleComplete(); + } catch (error) { + streamController.handleError(error); + } finally { + clearTimeout(timeout); + } + })(); + }; + streamController.registerCompletedToolHandoffDrain?.(drainCompletedToolHandoff); const noteClientChunk = (chunk: unknown) => { if (!(chunk instanceof Uint8Array)) return; @@ -566,6 +626,12 @@ export function createDisconnectAwareStream(transformStream, streamController) { // Runs past clientTerminalSeen: the frame that carries the terminal marker // can carry the only content too, and #8649 needs the whole stream scanned. contentWatcher.note(contentDecoder.decode(chunk, { stream: true })); + if ( + isResponsesClientFormat(streamController.clientResponseFormat) && + completedToolHandoffWatcher.note(toolHandoffDecoder.decode(chunk, { stream: true })) + ) { + streamController.markCompletedToolHandoffSeen?.(); + } if (clientTerminalSeen) return; terminalTail += terminalDecoder.decode(chunk, { stream: true }); @@ -676,11 +742,14 @@ export function createDisconnectAwareStream(transformStream, streamController) { }, async cancel(reason) { + const deferCompletedToolHandoff = + streamController.shouldDeferCompletedToolHandoff?.() === true; if (clientTerminalSeen) { streamController.handleComplete(); } else { streamController.handleDisconnect(reason || "cancelled"); } + if (deferCompletedToolHandoff) return; await Promise.allSettled([reader.cancel(reason), writer.abort(reason)]); }, }, diff --git a/tests/unit/codex-tool-handoff-disconnect-499.test.ts b/tests/unit/codex-tool-handoff-disconnect-499.test.ts new file mode 100644 index 0000000000..9514345e92 --- /dev/null +++ b/tests/unit/codex-tool-handoff-disconnect-499.test.ts @@ -0,0 +1,377 @@ +/** + * Regression coverage for Codex Responses tool handoffs: Codex can close the + * current HTTP response immediately after receiving a complete tool-call item, + * before the trailing response.completed frame reaches the client. OmniRoute + * must keep the upstream transform alive briefly so its normal completion and + * usage bookkeeping can still win over the delayed 499 finalizer. + */ +import test from "node:test"; +import assert from "node:assert/strict"; + +import { FORMATS } from "../../open-sse/translator/formats.ts"; +import { createCompletedResponsesToolHandoffWatcher } from "../../open-sse/utils/responsesToolHandoff.ts"; +import { createPassthroughStreamWithLogger } from "../../open-sse/utils/stream.ts"; +import { + createDisconnectAwareStream, + createNoopAbortWritable, + createStreamController, +} from "../../open-sse/utils/streamHandler.ts"; +import { createClientDisconnectGraceHandler } from "../../open-sse/utils/streamFailureFinalization.ts"; + +const encoder = new TextEncoder(); +const decoder = new TextDecoder(); + +function sse(event: string, data: Record): string { + return `event: ${event}\ndata: ${JSON.stringify({ type: event, ...data })}\n\n`; +} + +function wait(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + +async function cancelAfterFirstChunk({ + sseText, + allowCompletedToolHandoffGrace = true, +}: { + sseText: string; + allowCompletedToolHandoffGrace?: boolean; +}): Promise<{ upstreamCancelled: boolean; disconnects: number; signalAborted: boolean }> { + let upstreamCancelled = false; + let disconnects = 0; + const transformStream = { + readable: new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode(sseText)); + }, + cancel() { + upstreamCancelled = true; + }, + }), + writable: createNoopAbortWritable(), + }; + const streamController = createStreamController({ + clientResponseFormat: FORMATS.OPENAI_RESPONSES, + allowCompletedToolHandoffGrace, + clientDisconnectGracePeriodMs: 50, + onDisconnect: () => { + disconnects++; + }, + }); + const clientStream = createDisconnectAwareStream(transformStream, streamController); + const reader = clientStream.getReader(); + assert.equal((await reader.read()).done, false); + await reader.cancel("request_signal_aborted"); + + return { + upstreamCancelled, + disconnects, + signalAborted: streamController.signal.aborted, + }; +} + +test("Codex tool handoff drains the trailing Responses completion instead of persisting 499", async () => { + let upstreamController: ReadableStreamDefaultController | null = null; + let upstreamCancelled = false; + let completionRecorded = false; + let completionStatus: number | null = null; + let disconnectFinalizedAs499 = false; + const clientAbortController = new AbortController(); + + const providerStream = new ReadableStream({ + start(controller) { + upstreamController = controller; + controller.enqueue( + encoder.encode( + sse("response.output_item.added", { + output_index: 0, + item: { + id: "ctc_1", + type: "custom_tool_call", + call_id: "call_1", + name: "apply_patch", + input: "", + status: "in_progress", + }, + }) + + sse("response.custom_tool_call_input.done", { + item_id: "ctc_1", + output_index: 0, + input: "*** Begin Patch\n*** End Patch", + }) + + sse("response.output_item.done", { + output_index: 0, + item: { + id: "ctc_1", + type: "custom_tool_call", + call_id: "call_1", + name: "apply_patch", + input: "*** Begin Patch\n*** End Patch", + status: "completed", + }, + }) + ) + ); + }, + cancel() { + upstreamCancelled = true; + }, + }); + + const disconnectGraceHandler = createClientDisconnectGraceHandler({ + isStreamCompletionRecorded: () => completionRecorded, + gracePeriodMs: 50, + pollIntervalMs: 5, + finalize: () => { + disconnectFinalizedAs499 = true; + completionStatus = 499; + }, + }); + const streamController = createStreamController({ + clientResponseFormat: FORMATS.OPENAI_RESPONSES, + allowCompletedToolHandoffGrace: true, + clientDisconnectGracePeriodMs: 50, + clientAbortSignal: clientAbortController.signal, + onDisconnect: disconnectGraceHandler, + }); + const transformStream = createPassthroughStreamWithLogger( + "codex", + null, + null, + "gpt-5.6-sol", + "connection-1", + { model: "gpt-5.6-sol", stream: true }, + (payload) => { + completionRecorded = true; + completionStatus = payload.status; + }, + null, + null, + FORMATS.OPENAI_RESPONSES + ); + const transformedBody = providerStream.pipeThrough(transformStream); + const clientStream = createDisconnectAwareStream( + { readable: transformedBody, writable: createNoopAbortWritable() }, + streamController + ); + const reader = clientStream.getReader(); + + let received = ""; + while (!received.includes("response.output_item.done")) { + const chunk = await reader.read(); + assert.equal(chunk.done, false); + received += decoder.decode(chunk.value, { stream: true }); + } + + clientAbortController.abort("request_signal_aborted"); + const cancelPromise = reader.cancel("request_signal_aborted"); + setTimeout(() => { + try { + upstreamController?.enqueue( + encoder.encode( + sse("response.completed", { + response: { + id: "resp_1", + status: "completed", + output: [ + { + id: "ctc_1", + type: "custom_tool_call", + call_id: "call_1", + name: "apply_patch", + input: "*** Begin Patch\n*** End Patch", + status: "completed", + }, + ], + usage: { input_tokens: 10, output_tokens: 2, total_tokens: 12 }, + }, + }) + ) + ); + upstreamController?.close(); + } catch { + // The unchanged implementation cancels the upstream before this trailing + // completion can arrive; the assertions below expose that regression. + } + }, 0); + + await cancelPromise; + await wait(70); + + assert.equal(upstreamCancelled, false, "the completed tool handoff must be drained, not aborted"); + assert.equal(disconnectFinalizedAs499, false, "the real completion must beat the 499 finalizer"); + assert.equal(completionRecorded, true); + assert.equal(completionStatus, 200); +}); + +test("Codex handoff grace does not apply to an incomplete custom tool call", async () => { + const result = await cancelAfterFirstChunk({ + sseText: sse("response.output_item.done", { + output_index: 0, + item: { + id: "ctc_1", + type: "custom_tool_call", + call_id: "call_1", + name: "apply_patch", + input: "partial", + status: "completed", + }, + }), + }); + + assert.deepEqual(result, { upstreamCancelled: true, disconnects: 1, signalAborted: true }); +}); + +test("Codex handoff detection accepts a complete function call split across SSE chunks", () => { + const watcher = createCompletedResponsesToolHandoffWatcher(); + const frames = + sse("response.function_call_arguments.done", { + item_id: "fc_1", + output_index: 0, + arguments: '{"path":"README.md"}', + }) + + sse("response.output_item.done", { + output_index: 0, + item: { + id: "fc_1", + type: "function_call", + call_id: "call_1", + name: "read_file", + arguments: '{"path":"README.md"}', + status: "completed", + }, + }); + const splitAt = frames.indexOf("response.output_item.done") + 9; + + assert.equal(watcher.note(frames.slice(0, splitAt)), false); + assert.equal(watcher.note(frames.slice(splitAt)), true); +}); + +test("Codex handoff grace requires matching done input and completed item payloads", async () => { + const result = await cancelAfterFirstChunk({ + sseText: + sse("response.custom_tool_call_input.done", { + item_id: "ctc_1", + output_index: 0, + input: "complete input", + }) + + sse("response.output_item.done", { + output_index: 0, + item: { + id: "ctc_1", + type: "custom_tool_call", + call_id: "call_1", + name: "apply_patch", + input: "different input", + status: "completed", + }, + }), + }); + + assert.deepEqual(result, { upstreamCancelled: true, disconnects: 1, signalAborted: true }); +}); + +test("completed tool calls from non-Codex Responses clients keep normal abort behavior", async () => { + const result = await cancelAfterFirstChunk({ + allowCompletedToolHandoffGrace: false, + sseText: + sse("response.function_call_arguments.done", { + item_id: "fc_1", + output_index: 0, + arguments: "{}", + }) + + sse("response.output_item.done", { + output_index: 0, + item: { + id: "fc_1", + type: "function_call", + call_id: "call_1", + name: "read_file", + arguments: "{}", + status: "completed", + }, + }), + }); + + assert.deepEqual(result, { upstreamCancelled: true, disconnects: 1, signalAborted: true }); +}); + +test("Codex handoff grace still finalizes 499 and aborts when no completion arrives", async () => { + let upstreamCancelled = false; + let finalizedAs499 = false; + let completionRecorded = false; + const clientAbortController = new AbortController(); + const providerStream = new ReadableStream({ + start(controller) { + controller.enqueue( + encoder.encode( + sse("response.custom_tool_call_input.done", { + item_id: "ctc_1", + output_index: 0, + input: "complete input", + }) + + sse("response.output_item.done", { + output_index: 0, + item: { + id: "ctc_1", + type: "custom_tool_call", + call_id: "call_1", + name: "apply_patch", + input: "complete input", + status: "completed", + }, + }) + ) + ); + }, + cancel() { + upstreamCancelled = true; + }, + }); + const disconnectGraceHandler = createClientDisconnectGraceHandler({ + isStreamCompletionRecorded: () => completionRecorded, + gracePeriodMs: 25, + pollIntervalMs: 5, + finalize: () => { + finalizedAs499 = true; + }, + }); + const streamController = createStreamController({ + clientResponseFormat: FORMATS.OPENAI_RESPONSES, + allowCompletedToolHandoffGrace: true, + clientDisconnectGracePeriodMs: 25, + clientAbortSignal: clientAbortController.signal, + onDisconnect: disconnectGraceHandler, + }); + const transformStream = createPassthroughStreamWithLogger( + "codex", + null, + null, + "gpt-5.6-sol", + "connection-1", + { model: "gpt-5.6-sol", stream: true }, + () => { + completionRecorded = true; + }, + null, + null, + FORMATS.OPENAI_RESPONSES + ); + const clientStream = createDisconnectAwareStream( + { + readable: providerStream.pipeThrough(transformStream), + writable: createNoopAbortWritable(), + }, + streamController + ); + const reader = clientStream.getReader(); + assert.equal((await reader.read()).done, false); + + clientAbortController.abort("request_signal_aborted"); + await reader.cancel("request_signal_aborted"); + await wait(60); + + assert.equal(completionRecorded, false); + assert.equal(finalizedAs499, true); + assert.equal(upstreamCancelled, true); + assert.equal(streamController.signal.aborted, true); +});