From 41b761b666b404f57adccffae6bbcdbf5c332d77 Mon Sep 17 00:00:00 2001 From: Diego Rodrigues de Sa e Souza <8016841+diegosouzapw@users.noreply.github.com> Date: Fri, 26 Jun 2026 00:16:06 -0300 Subject: [PATCH] fix(sse): emit valid concatenable kiro tool_calls.arguments deltas (#4855) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Integrated into release/v3.8.37 — cherry-picked defining commit onto release tip; CHANGELOG re-merged; tests green. --- open-sse/executors/kiro.ts | 128 +++++++++++---- tests/unit/kiro-tool-args-streaming.test.ts | 167 ++++++++++++++++++++ 2 files changed, 265 insertions(+), 30 deletions(-) create mode 100644 tests/unit/kiro-tool-args-streaming.test.ts diff --git a/open-sse/executors/kiro.ts b/open-sse/executors/kiro.ts index 44f7b32c99..382b11919a 100644 --- a/open-sse/executors/kiro.ts +++ b/open-sse/executors/kiro.ts @@ -27,6 +27,8 @@ type KiroStreamState = { hasToolCalls: boolean; toolCallIndex: number; seenToolIds: Map; + toolArgsEmitted: Map; + toolArgsBuffered: Map; totalContentLength?: number; contextUsagePercentage?: number; hasContextUsage?: boolean; @@ -125,6 +127,56 @@ function crc32(buf: Uint8Array) { return (crc ^ 0xffffffff) >>> 0; } +/** + * Flush buffered tool arguments at finish boundaries. + * + * Kiro/CodeWhisperer streams toolUseEvent.input as PARTIAL OBJECTS that grow over time + * (e.g. {command:"cat /home"} then {command:"cat /home/wxsys"}). Re-stringifying each one + * and emitting it as an OpenAI argument delta produces overlapping prefixes that + * concatenate into unparseable garbage downstream ("Unterminated string"). + * + * Fix: defer object-form payloads into state.toolArgsBuffered keyed by toolCallId, keep + * only the latest canonical, and emit ONCE here as the complete arguments string (the + * final object is the source of truth — intermediate states are noise). String-form + * payloads are already concatenable deltas and are emitted incrementally. + */ +export function flushBufferedToolArgs( + state: Pick, + controller: { enqueue: (chunk: Uint8Array) => void }, + ctx: { responseId: string; created: number; model: string } +): void { + if (!state.toolArgsBuffered || state.toolArgsBuffered.size === 0) return; + const { responseId, created, model } = ctx; + for (const [toolCallId, info] of state.toolArgsBuffered) { + const alreadyEmitted = state.toolArgsEmitted.get(toolCallId) || ""; + if (info.canonical && info.canonical !== alreadyEmitted) { + const argsChunk: JsonRecord = { + id: responseId, + object: "chat.completion.chunk", + created, + model, + choices: [ + { + index: 0, + delta: { + tool_calls: [ + { + index: info.toolIndex, + function: { arguments: info.canonical }, + }, + ], + }, + finish_reason: null, + }, + ], + }; + controller.enqueue(TEXT_ENCODER.encode(`data: ${JSON.stringify(argsChunk)}\n\n`)); + state.toolArgsEmitted.set(toolCallId, info.canonical); + } + } + state.toolArgsBuffered.clear(); +} + function buildKiroFinishChunk( state: KiroStreamState, responseId: string, @@ -310,6 +362,8 @@ export class KiroExecutor extends BaseExecutor { hasToolCalls: false, toolCallIndex: 0, seenToolIds: new Map(), + toolArgsEmitted: new Map(), + toolArgsBuffered: new Map(), }; const transformStream = new TransformStream( @@ -469,46 +523,56 @@ export class KiroExecutor extends BaseExecutor { } if (toolInput !== undefined) { - let argumentsStr; - if (typeof toolInput === "string") { - argumentsStr = toolInput; - } else if (typeof toolInput === "object") { - argumentsStr = JSON.stringify(toolInput); - } else { - continue; - } + // String-form payloads are already concatenable incremental deltas — + // emit immediately and track what we've sent. + state.toolArgsEmitted.set( + toolCallId, + (state.toolArgsEmitted.get(toolCallId) || "") + toolInput + ); - const argsChunk = { - id: responseId, - object: "chat.completion.chunk", - created, - model, - choices: [ - { - index: 0, - delta: { - tool_calls: [ - { - index: toolIndex, - function: { - arguments: argumentsStr, + const argsChunk = { + id: responseId, + object: "chat.completion.chunk", + created, + model, + choices: [ + { + index: 0, + delta: { + tool_calls: [ + { + index: toolIndex, + function: { + arguments: toolInput, + }, }, - }, - ], + ], + }, + finish_reason: null, }, - finish_reason: null, - }, - ], - }; - chunkIndex++; - controller.enqueue(TEXT_ENCODER.encode(`data: ${JSON.stringify(argsChunk)}\n\n`)); + ], + }; + chunkIndex++; + controller.enqueue( + TEXT_ENCODER.encode(`data: ${JSON.stringify(argsChunk)}\n\n`) + ); + } else if (typeof toolInput === "object" && toolInput !== null) { + // Object-form payloads are PARTIAL OBJECTS that grow over time. Buffer + // the latest canonical and flush once at a finish boundary, otherwise the + // overlapping JSON prefixes concatenate into unparseable garbage. + state.toolArgsBuffered.set(toolCallId, { + toolIndex, + canonical: JSON.stringify(toolInput), + }); + } } } } // Handle messageStopEvent if (eventType === "messageStopEvent") { + flushBufferedToolArgs(state, controller, { responseId, created, model }); state.stopSeen = true; } @@ -576,6 +640,10 @@ export class KiroExecutor extends BaseExecutor { }, flush(controller) { + // Flush any buffered tool arguments (partial-object payloads) before finishing — + // idempotent against toolArgsEmitted if messageStopEvent already flushed them. + flushBufferedToolArgs(state, controller, { responseId, created, model }); + // Emit finish chunk if not already sent if (!state.finishEmitted) { state.finishEmitted = true; diff --git a/tests/unit/kiro-tool-args-streaming.test.ts b/tests/unit/kiro-tool-args-streaming.test.ts new file mode 100644 index 0000000000..52a10c7cb6 --- /dev/null +++ b/tests/unit/kiro-tool-args-streaming.test.ts @@ -0,0 +1,167 @@ +/** + * Regression test for Kiro/CodeWhisperer streaming tool_calls.arguments deltas. + * + * Kiro streams toolUseEvent.input as PARTIAL OBJECTS that grow over time + * (e.g. {command:"cat /home"} then {command:"cat /home/wxsys"}). The old executor + * re-stringified each partial object and emitted it as an OpenAI argument delta, so the + * overlapping JSON prefixes concatenated downstream into unparseable garbage + * ("Unterminated string"). The fix buffers object-form payloads keyed by toolCallId, keeps + * only the latest canonical, and flushes ONCE at a finish boundary. String-form payloads + * remain concatenable incremental deltas. + */ +import test from "node:test"; +import assert from "node:assert/strict"; + +const { flushBufferedToolArgs } = await import("../../open-sse/executors/kiro.ts"); + +type ArgsBuffered = Map; +type State = { toolArgsEmitted: Map; toolArgsBuffered: ArgsBuffered }; + +function makeMockController() { + const enqueued: string[] = []; + const decoder = new TextDecoder(); + return { + enqueued, + enqueue(bytes: Uint8Array) { + enqueued.push(decoder.decode(bytes)); + }, + }; +} + +function parseSSEData(sseLines: string[]) { + return sseLines + .map((line) => { + const m = line.match(/^data: (.+)\n\n$/); + return m ? JSON.parse(m[1]) : null; + }) + .filter(Boolean); +} + +function reconstructArguments(chunks: any[]) { + const perIndex = new Map(); + for (const chunk of chunks) { + const tc = chunk.choices?.[0]?.delta?.tool_calls?.[0]; + if (!tc) continue; + const idx = tc.index as number; + const args = (tc.function?.arguments as string) || ""; + perIndex.set(idx, (perIndex.get(idx) || "") + args); + } + return perIndex; +} + +const CTX = { responseId: "resp_test", created: 1, model: "kiro" }; + +test("flushes a single buffered tool call with valid concatenable JSON", () => { + const state: State = { + toolArgsEmitted: new Map(), + toolArgsBuffered: new Map([ + [ + "tool_abc", + { + toolIndex: 0, + canonical: '{"command":"cat /home/wxsys/Project/naskin/.impeccable.md"}', + }, + ], + ]), + }; + const controller = makeMockController(); + flushBufferedToolArgs(state, controller, CTX); + + const chunks = parseSSEData(controller.enqueued); + assert.equal(chunks.length, 1); + + const args = reconstructArguments(chunks); + const final = args.get(0)!; + assert.doesNotThrow(() => JSON.parse(final)); + assert.deepEqual(JSON.parse(final), { + command: "cat /home/wxsys/Project/naskin/.impeccable.md", + }); + + assert.equal(state.toolArgsBuffered.size, 0); + assert.equal(state.toolArgsEmitted.get("tool_abc"), final); +}); + +test("4 partial-object events produce a single valid flush (reported bug)", () => { + const state: State = { toolArgsEmitted: new Map(), toolArgsBuffered: new Map() }; + + const partialPayloads = [ + { command: "cat /home" }, + { command: "cat /home/wxsys" }, + { command: "cat /home/wxsys/Project" }, + { command: "cat /home/wxsys/Project/naskin/.impeccable.md" }, + ]; + // Each event overwrites the buffered canonical for the same toolCallId. + for (const input of partialPayloads) { + state.toolArgsBuffered.set("tool_abc", { toolIndex: 0, canonical: JSON.stringify(input) }); + } + + const controller = makeMockController(); + flushBufferedToolArgs(state, controller, CTX); + + const chunks = parseSSEData(controller.enqueued); + assert.equal(chunks.length, 1); + + const final = reconstructArguments(chunks).get(0)!; + assert.doesNotThrow(() => JSON.parse(final)); + assert.equal(JSON.parse(final).command, "cat /home/wxsys/Project/naskin/.impeccable.md"); +}); + +test("flushes multiple concurrent tool calls preserving toolIndex", () => { + const state: State = { + toolArgsEmitted: new Map(), + toolArgsBuffered: new Map([ + ["T1", { toolIndex: 0, canonical: '{"filePath":"/a/b/c.txt"}' }], + ["T2", { toolIndex: 1, canonical: '{"command":"ls"}' }], + ]), + }; + const controller = makeMockController(); + flushBufferedToolArgs(state, controller, CTX); + + const chunks = parseSSEData(controller.enqueued); + assert.equal(chunks.length, 2); + + const args = reconstructArguments(chunks); + assert.deepEqual(JSON.parse(args.get(0)!), { filePath: "/a/b/c.txt" }); + assert.deepEqual(JSON.parse(args.get(1)!), { command: "ls" }); +}); + +test("does not re-emit when canonical equals already-emitted (idempotent flush)", () => { + const state: State = { + toolArgsEmitted: new Map([["tool_abc", '{"command":"ls"}']]), + toolArgsBuffered: new Map([["tool_abc", { toolIndex: 0, canonical: '{"command":"ls"}' }]]), + }; + const controller = makeMockController(); + flushBufferedToolArgs(state, controller, CTX); + + assert.equal(controller.enqueued.length, 0); + assert.equal(state.toolArgsBuffered.size, 0); +}); + +test("no-op when buffer is empty", () => { + const state: State = { toolArgsEmitted: new Map(), toolArgsBuffered: new Map() }; + const controller = makeMockController(); + flushBufferedToolArgs(state, controller, CTX); + + assert.equal(controller.enqueued.length, 0); +}); + +test("OLD BUG: re-stringifying each partial object produces unparseable concat", () => { + const partials = [ + { command: "cat /home" }, + { command: "cat /home/wxsys" }, + { command: "cat /home/wxsys/Project/naskin/.impeccable.md" }, + ]; + const buggyConcat = partials.map((p) => JSON.stringify(p)).join(""); + assert.throws(() => JSON.parse(buggyConcat)); +}); + +test("FIX: emitting only the final canonical produces parseable JSON", () => { + const partials = [ + { command: "cat /home" }, + { command: "cat /home/wxsys" }, + { command: "cat /home/wxsys/Project/naskin/.impeccable.md" }, + ]; + const fixedConcat = JSON.stringify(partials[partials.length - 1]); + assert.doesNotThrow(() => JSON.parse(fixedConcat)); + assert.equal(JSON.parse(fixedConcat).command, "cat /home/wxsys/Project/naskin/.impeccable.md"); +});