mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-07-31 04:12:10 +03:00
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:
committed by
GitHub
parent
8a47bf2f5f
commit
41b761b666
@@ -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;
|
||||
|
||||
167
tests/unit/kiro-tool-args-streaming.test.ts
Normal file
167
tests/unit/kiro-tool-args-streaming.test.ts
Normal 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");
|
||||
});
|
||||
Reference in New Issue
Block a user