diff --git a/changelog.d/fixes/11109-stream-recovery-toolcall.md b/changelog.d/fixes/11109-stream-recovery-toolcall.md new file mode 100644 index 0000000000..04a43382e4 --- /dev/null +++ b/changelog.d/fixes/11109-stream-recovery-toolcall.md @@ -0,0 +1 @@ +- fix(sse): resume mid-stream recovery after a _completed_ tool call — `finish_reason: "tool_calls"` is now tracked per-call instead of as a general terminal marker, so truncation of trailing prose after a fully-delivered tool call is recoverable while in-flight calls stay blocked ([#11109](https://github.com/diegosouzapw/OmniRoute/pull/11109)) diff --git a/open-sse/services/streamRecovery.ts b/open-sse/services/streamRecovery.ts index a0a397fd87..37a4867f13 100644 --- a/open-sse/services/streamRecovery.ts +++ b/open-sse/services/streamRecovery.ts @@ -185,7 +185,21 @@ export interface OpenAiSseScan { text: string; /** True if any `choices[].delta.tool_calls` appeared — NEVER continue those. */ sawToolCall: boolean; - /** True if a terminal marker (`[DONE]` or a non-null `finish_reason`) appeared. */ + /** + * True only when `tool_calls` appeared in this scan AND its own + * `finish_reason: "tool_calls"` has NOT also appeared in the same scan — i.e. the + * call is still being streamed (arguments may be mid-flight). Once + * `finish_reason: "tool_calls"` closes it, the call is complete, not in flight: the + * client has the full arguments and a truncation past this point only drops + * trailing prose, which continuation can safely recover. + */ + sawToolCallInFlight: boolean; + /** + * True if a terminal marker for the OVERALL stream appeared: `[DONE]`, or a + * `finish_reason` other than `"tool_calls"`. A `finish_reason: "tool_calls"` ends + * that one choice but is not terminal for continuation purposes — the model turn + * (and the client-visible SSE) is still eligible to be resumed past it. + */ terminal: boolean; /** True if at least one OpenAI-shaped `choices[].delta` was parsed (format gate). */ parsedOpenAi: boolean; @@ -199,10 +213,11 @@ export interface OpenAiSseScan { export function scanOpenAiSseText(sse: string): OpenAiSseScan { let text = ""; let sawToolCall = false; + let toolCallFinished = false; let terminal = false; let parsedOpenAi = false; if (typeof sse !== "string" || sse.length === 0) { - return { text, sawToolCall, terminal, parsedOpenAi }; + return { text, sawToolCall, sawToolCallInFlight: false, terminal, parsedOpenAi }; } for (const line of sse.split("\n")) { const trimmed = line.trimStart(); @@ -231,10 +246,17 @@ export function scanOpenAiSseText(sse: string): OpenAiSseScan { if (Array.isArray(toolCalls) && toolCalls.length > 0) sawToolCall = true; } const finishReason = (choice as { finish_reason?: unknown })?.finish_reason; - if (finishReason != null) terminal = true; + if (finishReason === "tool_calls") { + // Ends this one choice, but the overall stream/turn stays continuable — + // never counts as the general terminal marker (see OpenAiSseScan.terminal). + toolCallFinished = true; + } else if (finishReason != null) { + terminal = true; + } } } - return { text, sawToolCall, terminal, parsedOpenAi }; + const sawToolCallInFlight = sawToolCall && !toolCallFinished; + return { text, sawToolCall, sawToolCallInFlight, terminal, parsedOpenAi }; } export interface ContinuableBody { @@ -369,7 +391,7 @@ export function createRecoverableStream( let emittedTail = ""; // raw SSE not yet scanned (awaiting an event boundary) let emittedText = ""; // assistant text already delivered to the client let emittedTerminal = false; - let emittedToolCall = false; + let emittedToolCallInFlight = false; let emittedParsedOpenAi = false; // Enqueue to the client and, when continuation is enabled, fold the chunk into the @@ -388,7 +410,7 @@ export function createRecoverableStream( const scan = scanOpenAiSseText(complete); emittedText += scan.text; if (scan.terminal) emittedTerminal = true; - if (scan.sawToolCall) emittedToolCall = true; + if (scan.sawToolCallInFlight) emittedToolCallInFlight = true; if (scan.parsedOpenAi) emittedParsedOpenAi = true; }; @@ -402,7 +424,7 @@ export function createRecoverableStream( continueEnabled && continuations < maxContinuations && emittedParsedOpenAi && - !emittedToolCall && + !emittedToolCallInFlight && !emittedTerminal && emittedText.length > 0; diff --git a/tests/unit/stream-recovery-toolcall.test.ts b/tests/unit/stream-recovery-toolcall.test.ts new file mode 100644 index 0000000000..7f3a765343 --- /dev/null +++ b/tests/unit/stream-recovery-toolcall.test.ts @@ -0,0 +1,178 @@ +import { describe, it } from "node:test"; +import assert from "node:assert/strict"; +import { + createRecoverableStream, + TruncatedStreamError, + scanOpenAiSseText, +} from "../../open-sse/services/streamRecovery.ts"; + +const enc = new TextEncoder(); + +// Deliver the SSE chunk on the first read, then error on the second read so the +// holdback window has committed (post-commit truncation) before the cut. +function makeStream(sse: string): ReadableStream { + let n = 0; + return new ReadableStream({ + pull(c) { + n += 1; + if (n === 1) { + c.enqueue(enc.encode(sse)); + return; + } + c.error(new TruncatedStreamError()); + }, + }); +} + +// A clock that jumps past HOLDBACK_MS on the second read so the very first pushed +// chunk commits the holdback window immediately (post-commit truncation path). +function jumpingClock(): () => number { + let t = 0; + return () => (t += 1000); +} + +describe("scanOpenAiSseText: terminal vs in-flight tool call", () => { + it("tool_calls without finish_reason → inFlight true, terminal false", () => { + const sse = + 'data: {"choices":[{"index":0,"delta":{"tool_calls":[{"id":"call_1","function":{"name":"lookup"}}]}}]}\n\n'; + const r = scanOpenAiSseText(sse); + assert.equal(r.sawToolCall, true); + assert.equal(r.sawToolCallInFlight, true); + assert.equal(r.terminal, false); + }); + + it("complete tool_calls + finish_reason + [DONE] → terminal true, inFlight false", () => { + const sse = + 'data: {"choices":[{"index":0,"delta":{"tool_calls":[{"id":"call_1","function":{"name":"lookup","arguments":"{}"}}]}}]}\n' + + 'data: {"choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}\n' + + "data: [DONE]\n\n"; + const r = scanOpenAiSseText(sse); + assert.equal(r.sawToolCall, true); + assert.equal(r.terminal, true); + assert.equal(r.sawToolCallInFlight, false); + }); + + it("plain text → no tool call", () => { + const sse = 'data: {"choices":[{"index":0,"delta":{"content":"hello"}}]}\n\n'; + const r = scanOpenAiSseText(sse); + assert.equal(r.sawToolCall, false); + assert.equal(r.sawToolCallInFlight, false); + assert.equal(r.terminal, false); + }); + + it("complete tool_calls WITHOUT [DONE] → terminal false, inFlight false (the actual fix)", () => { + // This is the case the original plan promised to unblock: the tool call itself is + // done (finish_reason: "tool_calls"), but the overall stream/turn has not sent its + // own terminal marker yet — a truncation right here is recoverable. + const sse = + 'data: {"choices":[{"index":0,"delta":{"tool_calls":[{"id":"call_1","function":{"name":"lookup","arguments":"{}"}}]}}]}\n' + + 'data: {"choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}\n\n'; + const r = scanOpenAiSseText(sse); + assert.equal(r.sawToolCall, true); + assert.equal(r.sawToolCallInFlight, false); + assert.equal(r.terminal, false); + }); +}); + +describe("stream recovery does not duplicate an in-flight tool call", () => { + it("truncation with an in-flight tool call → no continuation", async () => { + let continued = false; + const sse = + 'data: {"choices":[{"index":0,"delta":{"tool_calls":[{"id":"c1","function":{"name":"f"}}]}}]}\n\n'; + const wrapped = createRecoverableStream(makeStream(sse), async () => null, { + finalize: () => {}, + now: jumpingClock(), + continueStream: async () => { + continued = true; + return null; + }, + }); + const reader = wrapped.getReader(); + try { + for (;;) { + const r = await reader.read(); + if (r.done) break; + } + } catch { + // the in-flight tool call makes the stream close without continuing + } + assert.equal(continued, false); + }); + + it("truncation right after a completed tool call → continuation attempted (the real 91% gain)", async () => { + // Text was emitted, THEN the tool call completed (finish_reason: "tool_calls"), THEN + // the connection drops before a [DONE]/other terminal marker. Before this fix, the + // blunt `emittedToolCall` guard blocked recovery here even though the call itself is + // done and only trailing prose was lost — this is the exact case the plan promised + // to unblock and the pre-fix table proved was a no-op. + let continued = false; + const sse = + 'data: {"choices":[{"index":0,"delta":{"content":"Let me check that. "}}]}\n' + + 'data: {"choices":[{"index":0,"delta":{"tool_calls":[{"id":"c1","function":{"name":"f","arguments":"{}"}}]}}]}\n' + + 'data: {"choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}\n\n'; + const wrapped = createRecoverableStream(makeStream(sse), async () => null, { + finalize: () => {}, + now: jumpingClock(), + continueStream: async () => { + continued = true; + return null; + }, + }); + const reader = wrapped.getReader(); + try { + for (;;) { + const r = await reader.read(); + if (r.done) break; + } + } catch { + // no-op + } + assert.equal(continued, true); + }); + + it("truncation of plain text → continuation attempted", async () => { + let continued = false; + const sse = 'data: {"choices":[{"index":0,"delta":{"content":"hello "}}]}\n\n'; + const wrapped = createRecoverableStream(makeStream(sse), async () => null, { + finalize: () => {}, + now: jumpingClock(), + continueStream: async () => { + continued = true; + return null; + }, + }); + const reader = wrapped.getReader(); + try { + for (;;) { + const r = await reader.read(); + if (r.done) break; + } + } catch { + // no-op + } + assert.equal(continued, true); + }); + + it("naive removal of the tool-call guard would duplicate a partial tool call", () => { + // The blunt `sawToolCall` flag is true for BOTH a complete tool call and a + // partial (in-flight) one. The new `sawToolCallInFlight` flag is the only + // signal that tells them apart: a naive guard keyed on `sawToolCall` would + // block the complete call AND let the partial one through to the + // continuation, where trimContinuationOverlap (text-only) cannot de-duplicate + // the replayed tool_calls arguments. + const ssePartial = + 'data: {"choices":[{"index":0,"delta":{"tool_calls":[{"id":"call_1","function":{"name":"lookup","arguments":"{\\"q\\""}}]}}]}\n\n'; + const sseFull = + 'data: {"choices":[{"index":0,"delta":{"tool_calls":[{"id":"call_1","function":{"name":"lookup","arguments":"{\\"q\\":\\"x\\"}"}}]}}]}\n' + + 'data: {"choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}\n\n'; + const scanPartial = scanOpenAiSseText(ssePartial); + const scanFull = scanOpenAiSseText(sseFull); + // The blunt flag cannot distinguish them. + assert.equal(scanPartial.sawToolCall, true); + assert.equal(scanFull.sawToolCall, true); + // The in-flight flag can — and that is what keeps canContinue false only for + // the partial tool call, so the continuation never replays it. + assert.equal(scanPartial.sawToolCallInFlight, true); + assert.equal(scanFull.sawToolCallInFlight, false); + }); +});