fix(streaming): preserve completed Codex tool handoffs (#10608)

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 };

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<string, unknown>,
+  item?: Record<string, unknown>
+): string[] {
+  const keys = new Set<string>();
+  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<string, string>();
+  const customToolInputDone = new Map<string, string>();
+  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<string, unknown>, 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<string, unknown>)
+        : 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<string, unknown>;
+      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;
+    },
+  };
+}

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<typeof createStreamController>;
@@ -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)]);
       },
     },

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, unknown>): string {
+  return `event: ${event}\ndata: ${JSON.stringify({ type: event, ...data })}\n\n`;
+}
+
+function wait(ms: number): Promise<void> {
+  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<Uint8Array>({
+      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<Uint8Array> | null = null;
+  let upstreamCancelled = false;
+  let completionRecorded = false;
+  let completionStatus: number | null = null;
+  let disconnectFinalizedAs499 = false;
+  const clientAbortController = new AbortController();
+
+  const providerStream = new ReadableStream<Uint8Array>({
+    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<Uint8Array>({
+    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);
+});