From 4d13efac5ff4894a2af80821983348bd33d8ccb2 Mon Sep 17 00:00:00 2001 From: 0disoft Date: Wed, 4 Nov 2026 01:00:07 +0900 Subject: [PATCH 1/2] perf(kiro): count streamed tool argument bytes incrementally --- scripts/test-layout/layout.json | 2 +- src/adapters/kiro/stream.ts | 9 ++- structure/transports/byte-accounting.md | 6 ++ tests/fixtures/test-layout-expected.json | 2 +- .../kiro/kiro-tool-byte-accounting.test.ts | 76 +++++++++++++++++++ 5 files changed, 89 insertions(+), 6 deletions(-) create mode 100644 tests/providers/kiro/kiro-tool-byte-accounting.test.ts diff --git a/scripts/test-layout/layout.json b/scripts/test-layout/layout.json index 77a11f28fe2..7f6148cf72a 100644 --- a/scripts/test-layout/layout.json +++ b/scripts/test-layout/layout.json @@ -1100,7 +1100,7 @@ "kiro-metering-events.test.ts": "providers/kiro", "kiro-metering-usage.test.ts": "providers/kiro", "kiro-single-final.test.ts": "providers/kiro", "kiro-structured-output.test.ts": "providers/kiro", "kiro-structured-output-strict.test.ts": "providers/kiro", - "kiro-stream.test.ts": "providers/kiro", + "kiro-stream.test.ts": "providers/kiro", "kiro-tool-byte-accounting.test.ts": "providers/kiro", "kiro-transport-parity.test.ts": "providers/kiro", "kiro-usage-quota.test.ts": "providers/kiro", "kiro-usage-restart.test.ts": "providers/kiro", diff --git a/src/adapters/kiro/stream.ts b/src/adapters/kiro/stream.ts index 90c8124bb78..e49a782a1f2 100644 --- a/src/adapters/kiro/stream.ts +++ b/src/adapters/kiro/stream.ts @@ -385,7 +385,7 @@ async function* parseKiroAttemptEvents( }; } - let open: { id: string; name: string; chunks: string[]; completion: boolean } | null = null; + let open: { id: string; name: string; chunks: string[]; inputBytes: number; completion: boolean } | null = null; let openCallId: string | undefined; const closeOpenCall = () => { if (!openCallId) return; @@ -571,8 +571,8 @@ async function* parseKiroAttemptEvents( const beginTool = ( id: string, name: string, - ): { tool?: { id: string; name: string; chunks: string[]; completion: boolean }; terminal?: AdapterEvent } => { - const next = { id, name, chunks: [], completion: false }; + ): { tool?: { id: string; name: string; chunks: string[]; inputBytes: number; completion: boolean }; terminal?: AdapterEvent } => { + const next = { id, name, chunks: [], inputBytes: 0, completion: false }; const terminal = classifyTool(next); return terminal ? { terminal } : { tool: next }; }; @@ -831,12 +831,13 @@ async function* parseKiroAttemptEvents( } } if (open && ev.input !== undefined) { - const previousCallBytes = open.chunks.reduce((total, chunk) => total + Buffer.byteLength(chunk), 0); + const previousCallBytes = open.inputBytes; const nextCallBytes = previousCallBytes + Buffer.byteLength(ev.input); const callReservation = budget.reserveTransient(nextCallBytes, { kind: "tool_args", callId: open.id }); open.chunks.push(ev.input); callReservation.commitRetained(); budget.releaseRetained(previousCallBytes, { kind: "tool_args", callId: open.id }); + open.inputBytes = nextCallBytes; const nextOutputCharsBytes = appendedUtf8Bytes(outputChars, outputCharsBytes, ev.input); const toolOutputReservation = budget.reserveTransient(nextOutputCharsBytes, { kind: "retained_collectors" }); outputChars += ev.input; diff --git a/structure/transports/byte-accounting.md b/structure/transports/byte-accounting.md index 40d7f722caf..eed43592c37 100644 --- a/structure/transports/byte-accounting.md +++ b/structure/transports/byte-accounting.md @@ -143,6 +143,12 @@ envelopes, and real HTTP streams; `tests/server/inbound-body-budget.test.ts` cov ## Stream-buffer accounting +Kiro tool-input collection in `src/adapters/kiro/stream.ts` keeps a per-call UTF-8 byte +total, measuring only the arriving fragment instead of remeasuring retained chunks. +Each new call starts at zero; transient replacement admission, per-call limits and +close/cancellation release remain unchanged. `tests/providers/kiro/kiro-tool-byte-accounting.test.ts` +bounds measurement work and covers Unicode byte limits and consecutive-call reset. + Command Code NDJSON in `src/adapters/command-code.ts` holds one `live_transient` charge for its decoded buffer. Growth reserves only the positive byte delta; draining releases only the removed bytes. Recounting an unchanged residual does not reserve it twice. Failed growth keeps diff --git a/tests/fixtures/test-layout-expected.json b/tests/fixtures/test-layout-expected.json index c9204b18355..aa318e96f87 100644 --- a/tests/fixtures/test-layout-expected.json +++ b/tests/fixtures/test-layout-expected.json @@ -963,7 +963,7 @@ "kiro-review-regressions.test.ts": "providers/kiro", "kiro-metering-events.test.ts": "providers/kiro", "kiro-metering-usage.test.ts": "providers/kiro", "kiro-completion-with-real-tools.test.ts": "providers/kiro", "kiro-single-final.test.ts": "providers/kiro", "kiro-structured-output.test.ts": "providers/kiro", "kiro-structured-output-strict.test.ts": "providers/kiro", - "kiro-stream.test.ts": "providers/kiro", "kiro-transport-parity.test.ts": "providers/kiro", + "kiro-stream.test.ts": "providers/kiro", "kiro-tool-byte-accounting.test.ts": "providers/kiro", "kiro-transport-parity.test.ts": "providers/kiro", "kiro-usage-quota.test.ts": "providers/kiro", "kiro-usage-restart.test.ts": "providers/kiro", "kiro-windows-cli-db-path.test.ts": "providers/kiro", "kiro-windows-cli-executable-path.test.ts": "providers/kiro", "kiro-wire-estimate.test.ts": "providers/kiro", "lab-activation.test.ts": "lab", diff --git a/tests/providers/kiro/kiro-tool-byte-accounting.test.ts b/tests/providers/kiro/kiro-tool-byte-accounting.test.ts new file mode 100644 index 00000000000..8fef03cb2df --- /dev/null +++ b/tests/providers/kiro/kiro-tool-byte-accounting.test.ts @@ -0,0 +1,76 @@ +import { expect, spyOn, test } from "bun:test"; +import { parseKiroStream } from "../../../src/adapters/kiro/stream"; +import { encodeMessage } from "../../../src/lib/eventstream-decoder"; +import { createTranslatorBudget } from "../../../src/lib/translator-budget"; +import type { AdapterEvent } from "../../../src/types"; + +/** Exercise the real Smithy decoder and Kiro parser without a provider request. */ +function responseFor(calls: string[][]): Response { + const encoder = new TextEncoder(); + const frames = calls.flatMap((chunks, index) => [ + { name: "write_file", toolUseId: `call-${index}` }, + ...chunks.map(input => ({ input })), + { stop: true }, + ]).map(event => encodeMessage( + { ":message-type": "event", ":event-type": "toolUseEvent" }, + encoder.encode(JSON.stringify(event)), + )); + let index = 0; + return new Response(new ReadableStream({ + pull(controller) { + if (index < frames.length) controller.enqueue(frames[index++]); + else controller.close(); + }, + })); +} + +test("fragmented Kiro tool arguments require linear UTF-8 measurement work", async () => { + const fragment = "x".repeat(256); + const count = 256; + const chunks = ['{"value":"', ...Array(count).fill(fragment), '"}']; + const response = responseFor([chunks]); + const byteLength = Buffer.byteLength; + let measurements = 0; + const measuring = spyOn(Buffer, "byteLength").mockImplementation((...args: Parameters) => { + if (args[0] === fragment) measurements++; + return Reflect.apply(byteLength, Buffer, args); + }); + const budget = createTranslatorBudget(); + const events: AdapterEvent[] = []; + try { + for await (const event of parseKiroStream(response, budget)) events.push(event); + expect(events.filter(event => event.type === "tool_call_delta").map(event => event.arguments).join("")).toBe(chunks.join("")); + expect(events.at(-1)?.type).toBe("done"); + expect(budget.snapshot().currentBytes).toBe(0); + expect(budget.snapshot().activeCalls).toBe(0); + // One measurement for call retention and one for output accounting; leave bounded headroom. + expect(measurements).toBeLessThanOrEqual(3 * count); + } finally { + measuring.mockRestore(); + budget.dispose(); + } +}); + +test("Kiro argument byte totals reset per call and retain exact UTF-8 limits", async () => { + const chunks = ['{"value":"', "ν•œκΈ€πŸ˜€", '"}']; + const bytes = chunks.reduce((total, chunk) => total + Buffer.byteLength(chunk), 0); + for (const limit of [bytes - 1, bytes, bytes + 1]) { + const budget = createTranslatorBudget({ maxCallArgumentBytes: limit }); + const events: AdapterEvent[] = []; + try { + for await (const event of parseKiroStream(responseFor([chunks, chunks]), budget)) events.push(event); + if (limit < bytes) { + expect(events.some(event => event.type === "error")).toBe(true); + expect(events.some(event => event.type === "tool_call_start")).toBe(false); + } else { + expect(events.filter(event => event.type === "tool_call_start").map(event => event.id)).toEqual(["call-0", "call-1"]); + expect(events.filter(event => event.type === "tool_call_delta").map(event => event.arguments).join("")).toBe(chunks.join("").repeat(2)); + expect(events.at(-1)?.type).toBe("done"); + } + expect(budget.snapshot().currentBytes).toBe(0); + expect(budget.snapshot().activeCalls).toBe(0); + } finally { + budget.dispose(); + } + } +}); From 0c882a631ba72ae5ed1f030bde3a2791d0063190 Mon Sep 17 00:00:00 2001 From: 0disoft Date: Wed, 4 Nov 2026 02:06:49 +0900 Subject: [PATCH 2/2] test(kiro): bound total byte-measurement input size --- tests/providers/kiro/kiro-tool-byte-accounting.test.ts | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/providers/kiro/kiro-tool-byte-accounting.test.ts b/tests/providers/kiro/kiro-tool-byte-accounting.test.ts index 8fef03cb2df..08ba068ba40 100644 --- a/tests/providers/kiro/kiro-tool-byte-accounting.test.ts +++ b/tests/providers/kiro/kiro-tool-byte-accounting.test.ts @@ -31,8 +31,10 @@ test("fragmented Kiro tool arguments require linear UTF-8 measurement work", asy const response = responseFor([chunks]); const byteLength = Buffer.byteLength; let measurements = 0; + let measuredChars = 0; const measuring = spyOn(Buffer, "byteLength").mockImplementation((...args: Parameters) => { if (args[0] === fragment) measurements++; + if (typeof args[0] === "string") measuredChars += args[0].length; return Reflect.apply(byteLength, Buffer, args); }); const budget = createTranslatorBudget(); @@ -45,6 +47,8 @@ test("fragmented Kiro tool arguments require linear UTF-8 measurement work", asy expect(budget.snapshot().activeCalls).toBe(0); // One measurement for call retention and one for output accounting; leave bounded headroom. expect(measurements).toBeLessThanOrEqual(3 * count); + // Include growing-prefix measurements, with headroom for serialized adapter events. + expect(measuredChars).toBeLessThanOrEqual(16 * chunks.join("").length); } finally { measuring.mockRestore(); budget.dispose();