fix(sse): emit valid concatenable kiro tool_calls.arguments deltas (#4855)

Integrated into release/v3.8.37 — cherry-picked defining commit onto release tip; CHANGELOG re-merged; tests green.
This commit is contained in:
Diego Rodrigues de Sa e Souza
2026-06-26 00:16:06 -03:00
committed by GitHub
parent 8a47bf2f5f
commit 41b761b666
2 changed files with 265 additions and 30 deletions

View File

@@ -27,6 +27,8 @@ type KiroStreamState = {
hasToolCalls: boolean;
toolCallIndex: number;
seenToolIds: Map<string, number>;
toolArgsEmitted: Map<string, string>;
toolArgsBuffered: Map<string, { toolIndex: number; canonical: string }>;
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<KiroStreamState, "toolArgsBuffered" | "toolArgsEmitted">,
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;

View File

@@ -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<string, { toolIndex: number; canonical: string }>;
type State = { toolArgsEmitted: Map<string, string>; 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<number, string>();
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");
});