diff --git a/src/adapters/kiro-thinking.ts b/src/adapters/kiro-thinking.ts index e88c62ffc5..ee144e2878 100644 --- a/src/adapters/kiro-thinking.ts +++ b/src/adapters/kiro-thinking.ts @@ -75,6 +75,14 @@ export class KiroThinkingParser { return []; } + /** Release any partial tag/content carry when the owning stream stops early. */ + dispose(): void { + this.replaceCarry("preBuffer", ""); + this.replaceCarry("thinkingBuffer", ""); + this.closeTag = ""; + this.state = "streaming"; + } + private drainThinking(): AdapterEvent[] { const close = this.closeTag; const idx = this.thinkingBuffer.indexOf(close); diff --git a/src/adapters/kiro.ts b/src/adapters/kiro.ts index 0b6951f7c8..7fdfcc1db8 100644 --- a/src/adapters/kiro.ts +++ b/src/adapters/kiro.ts @@ -19,6 +19,8 @@ import { createKiroToolNameRegistry, fallbackToolUseId, fingerprint, invocationI import { namespacedToolName } from "../types"; import { isTranslatorBudgetExceededError, + releaseTranslatedEvent, + retainTranslatedEvent, type TranslatorBudget, } from "../lib/translator-budget"; import type { @@ -636,8 +638,8 @@ export function buildKiroPayload( // Stream parsing (shared by parseStream + parseResponse) // CodeWhisperer GenerateAssistantResponse ALWAYS returns an AWS eventstream body (there is no -// non-streaming mode), so both the streaming bridge and the non-streaming web-search sidecar loop -// decode the same way — parseResponse just collects what parseStream yields. +// non-streaming wire mode), so the streaming bridge and non-streaming Responses path decode the +// same way — parseResponse just collects what parseStream yields. interface KiroAttemptParseResult { terminal?: AdapterEvent; needsFallback?: boolean; @@ -859,16 +861,12 @@ async function* parseKiroAttempt( ); let handedOff = false; try { - let next = await attempt.next(); - while (!next.done) { - yield next.value; - next = await attempt.next(); - } + const result = yield* attempt; for (const event of deferred.splice(0)) { try { yield event; } finally { retention.releaseEvent(event); } } handedOff = true; - return { ...next.value, releaseRetained: () => retention.releaseAll() }; + return { ...result, releaseRetained: () => retention.releaseAll() }; } finally { if (!handedOff) retention.releaseAll(); } @@ -897,6 +895,12 @@ async function* parseKiroAttemptEvents( } let open: { id: string; name: string; chunks: string[]; completion: boolean } | null = null; + let openCallId: string | undefined; + const closeOpenCall = () => { + if (!openCallId) return; + budget.closeCall(openCallId); + openCallId = undefined; + }; let outputChars = ""; let outputCharsBytes = 0; let contextUsagePercentage: number | undefined; @@ -1109,7 +1113,7 @@ async function* parseKiroAttemptEvents( if (!open) return { events: [] }; const tool = open; open = null; - budget.closeCall(tool.id); + closeOpenCall(); const input = tool.chunks.join(""); if (!isCompleteKiroToolInput(input)) { return { events: [], terminal: protocolTerminal(kiroTruncationErrorMessage("incomplete tool input JSON"), tool.completion) }; @@ -1221,11 +1225,12 @@ async function* parseKiroAttemptEvents( if (started.terminal) return { assistantText, sawReasoning, terminal: started.terminal }; open = started.tool!; budget.openCall(open.id); + openCallId = open.id; } else if ( (ev.toolUseId && ev.toolUseId !== open.id) || (ev.name && open.name !== "unknown" && ev.name !== open.name) ) { - budget.closeCall(open.id); + closeOpenCall(); open = null; return { assistantText, sawReasoning, terminal: protocolTerminal(kiroTruncationErrorMessage("tool input changed identity before stop")) }; } @@ -1489,7 +1494,7 @@ async function* parseKiroAttemptEvents( }; } catch (err) { if (isTranslatorBudgetExceededError(err)) { - if (open) budget.closeCall(open.id); + closeOpenCall(); return { assistantText, sawReasoning, @@ -1531,6 +1536,9 @@ async function* parseKiroAttemptEvents( usage: usage(), }, }; + } finally { + thinking.dispose(); + closeOpenCall(); } } @@ -1547,7 +1555,7 @@ export async function* parseKiroStream( contextInputEstimate?: number, ): AsyncGenerator { const contextWindowState: KiroContextWindowState = { value: contextWindow }; - const first = parseKiroAttempt( + const firstResult = yield* parseKiroAttempt( response, budget, completionMode, @@ -1559,12 +1567,6 @@ export async function* parseKiroStream( contextInputEstimate, false, ); - let firstNext = await first.next(); - while (!firstNext.done) { - yield firstNext.value; - firstNext = await first.next(); - } - const firstResult = firstNext.value; try { if (!firstResult.needsFallback) { if (firstResult.terminal) yield firstResult.terminal; @@ -1642,7 +1644,7 @@ export async function* parseKiroStream( return; } - const second = parseKiroAttempt( + const secondResult = yield* parseKiroAttempt( fallback.response, budget, "text_fallback", @@ -1656,12 +1658,6 @@ export async function* parseKiroStream( // A zero-output transport failure here must stay non-retryable to avoid duplicating that text. priorEmittedOutput, ); - let secondNext = await second.next(); - while (!secondNext.done) { - yield secondNext.value; - secondNext = await second.next(); - } - const secondResult = secondNext.value; try { if (!secondResult.terminal) { yield retryableKiroIncomplete( @@ -1909,25 +1905,32 @@ export function createKiroAdapter(provider: OcxProviderConfig): ProviderAdapter return safeKiroHttpErrorMessage(status, headers, payloadText); }, - // Non-streaming path used by the web-search sidecar loop (loop.ts runs each iteration - // non-streamed so it can inspect tool calls). CW only ever event-streams, so we drain the - // same decoder into an array. Without this, any Codex request that includes the web_search - // tool failed with "web-search sidecar requires a non-streaming adapter" (kiro-only). + // Kiro always returns an event stream, including for non-streaming Responses requests. Drain + // the decoder into a budget-owned batch so an upstream stream cannot grow this array without + // bound while the caller waits for the complete JSON response. async parseResponse(response: Response, budget: TranslatorBudget): Promise { const events: AdapterEvent[] = []; - for await (const e of parseKiroStream( - response, - budget, - modelId, - inputTokens, - contextWindow, - toolNameMap, - conversationId, - completionMode, - completionMode === "required" ? fallbackFactory : undefined, - contextInputEstimate, - )) events.push(e); - return events; + try { + for await (const e of parseKiroStream( + response, + budget, + modelId, + inputTokens, + contextWindow, + toolNameMap, + conversationId, + completionMode, + completionMode === "required" ? fallbackFactory : undefined, + contextInputEstimate, + )) { + retainTranslatedEvent(e, budget, events.at(-1)); + events.push(e); + } + return events; + } catch (error) { + for (const event of events) releaseTranslatedEvent(event, budget); + throw error; + } }, }; } diff --git a/src/lib/translator-budget.ts b/src/lib/translator-budget.ts index 1bba4512d7..18400eebf7 100644 --- a/src/lib/translator-budget.ts +++ b/src/lib/translator-budget.ts @@ -71,6 +71,40 @@ export interface TranslatorBudget { const retainedEventOwnership = new WeakMap(); +/** + * Charge one event appended to an incrementally materialized adapter-event batch. + * The newest event owns the closing array bracket; moving that byte from the old + * tail keeps in-order release accounting equal to the still-retained JSON array. + */ +export function retainTranslatedEvent( + event: T, + budget: TranslatorBudget, + previousTail?: object, +): void { + if (retainedEventOwnership.has(event)) { + throw new Error("translated event is already retained"); + } + if (previousTail === event) { + throw new Error("incremental translated event tail must be a distinct object"); + } + const previousOwnership = previousTail === undefined + ? undefined + : retainedEventOwnership.get(previousTail); + if ( + previousTail !== undefined + && (!previousOwnership || previousOwnership.budget !== budget || previousOwnership.bytes < 2) + ) { + throw new Error("incremental translated event tail is not retained by this budget"); + } + + const serializedBytes = Buffer.byteLength(JSON.stringify(event)); + budget.chargeRetained(serializedBytes + (previousTail === undefined ? 2 : 1), { + kind: "retained_collectors", + }); + if (previousOwnership) previousOwnership.bytes -= 1; + retainedEventOwnership.set(event, { budget, bytes: serializedBytes + 2 }); +} + /** * Charge a materialized adapter-event batch and attach its lease to the events themselves. * A copied event array (for example terminal-guard collection) preserves the event objects, so diff --git a/tests/kiro-stream.test.ts b/tests/kiro-stream.test.ts index 744fa8ab1f..f38b0d1b2d 100644 --- a/tests/kiro-stream.test.ts +++ b/tests/kiro-stream.test.ts @@ -14,6 +14,7 @@ import { } from "../src/adapters/kiro-constants"; import { parseKiroEvent } from "../src/adapters/kiro-events"; import { resetKiroThrottleStateForTests } from "../src/adapters/kiro-retry"; +import { buildResponseJSON } from "../src/bridge"; import { encodeMessage } from "../src/lib/eventstream-decoder"; import { estimateTokens } from "../src/lib/token-estimate"; import { createTranslatorBudget } from "../src/lib/translator-budget"; @@ -1730,11 +1731,35 @@ describe("kiro adapter — parseStream", () => { }); }); -describe("kiro adapter — parseResponse (web-search sidecar non-streaming path)", () => { - test("adapter exposes parseResponse so the web_search sidecar accepts kiro", async () => { +describe("kiro adapter — non-streaming parseResponse", () => { + test("adapter exposes parseResponse for non-streaming Responses requests", async () => { expect(typeof createKiroAdapter(provider).parseResponse).toBe("function"); }); + test("returning the outer parser cancels its active attempt and releases retained state", async () => { + const budget = createTranslatorBudget(); + let bodyCancelled = false; + try { + const response = new Response(new ReadableStream({ + start(controller) { + controller.enqueue(eventFrame({ content: `${"x".repeat(30)}` })); + }, + cancel() { + bodyCancelled = true; + }, + })); + const events = parseKiroStream(response, budget); + expect((await events.next()).done).toBe(false); + await events.return(undefined); + + expect(bodyCancelled).toBe(true); + expect(budget.snapshot().currentBytes).toBe(0); + expect(budget.snapshot().activeCalls).toBe(0); + } finally { + budget.dispose(); + } + }); + test("drains the same CW eventstream into an AdapterEvent[] (parity with parseStream)", async () => { const frames = [ eventFrame({ content: "Hi " }), @@ -1751,6 +1776,53 @@ describe("kiro adapter — parseResponse (web-search sidecar non-streaming path) expect(start).toMatchObject({ id: "t1", name: "bash" }); }); + test("bounds events while collecting a non-streaming response", async () => { + const budget = createTranslatorBudget({ maxTurnBytes: 500 }); + let bodyCancelled = false; + try { + const adapter = createKiroAdapterProduction(provider); + const response = new Response(new ReadableStream({ + start(controller) { + for (let index = 0; index < 40; index++) { + controller.enqueue(eventFrame({ name: "bash", toolUseId: "pending-tool" })); + } + }, + cancel() { + bodyCancelled = true; + }, + })); + + await expect(adapter.parseResponse!(response, budget)).rejects.toMatchObject({ + code: "translation_buffer_limit", + }); + expect(bodyCancelled).toBe(true); + expect(budget.snapshot().currentBytes).toBe(0); + expect(budget.snapshot().activeCalls).toBe(0); + } finally { + budget.dispose(); + } + }); + + test("transfers collected event ownership to the non-streaming response builder", async () => { + const budget = createTranslatorBudget(); + try { + const adapter = createKiroAdapterProduction(provider); + const events = await adapter.parseResponse!( + new Response(streamOf(eventFrame({ content: "bounded" }))), + budget, + ); + expect(budget.snapshot().currentBytes).toBeGreaterThan(0); + + const json = buildResponseJSON(events, "kiro/test", { translatorBudget: budget }); + expect(json.status).toBe("completed"); + const output = json.output as Array>; + const outputBytes = output.reduce((total, item) => total + Buffer.byteLength(JSON.stringify(item)), 0); + expect(budget.snapshot().currentBytes).toBe(outputBytes); + } finally { + budget.dispose(); + } + }); + // The parity test above never calls buildRequest(), so the contextInputEstimate closure that // buildRequest() installs is never activated on the non-streaming path. Build a long-history // request first, then assert the terminal usage carries the absolute checkpoint rather than only diff --git a/tests/translator-budget.test.ts b/tests/translator-budget.test.ts index 48e25a3b9a..149f08afdc 100644 --- a/tests/translator-budget.test.ts +++ b/tests/translator-budget.test.ts @@ -7,6 +7,8 @@ import { TRANSLATOR_MAX_CALL_ARGUMENT_BYTES, TRANSLATOR_MAX_TURN_BYTES, createTranslatorBudget, + releaseTranslatedEvent, + retainTranslatedEvent, translatorObservedBufferSnapshot, } from "../src/lib/translator-budget"; import type { AdapterEvent } from "../src/types"; @@ -20,6 +22,41 @@ async function textWithin(stream: ReadableStream, timeoutMs = 2_000) } describe("translator budget", () => { + test("incremental event retention transfers array-tail ownership during in-order release", () => { + const budget = createTranslatorBudget({ maxTurnBytes: 4_096 }); + const first = { type: "text_delta", text: "first" }; + const second = { type: "done" }; + try { + retainTranslatedEvent(first, budget); + expect(budget.snapshot().currentBytes).toBe(Buffer.byteLength(JSON.stringify([first]))); + + retainTranslatedEvent(second, budget, first); + expect(budget.snapshot().currentBytes).toBe(Buffer.byteLength(JSON.stringify([first, second]))); + + releaseTranslatedEvent(first, budget); + expect(budget.snapshot().currentBytes).toBe(Buffer.byteLength(JSON.stringify([second]))); + releaseTranslatedEvent(second, budget); + expect(budget.snapshot().currentBytes).toBe(0); + } finally { + budget.dispose(); + } + }); + + test("incremental event retention rejects an event object that already owns a lease", () => { + const budget = createTranslatorBudget({ maxTurnBytes: 4_096 }); + const event = { type: "heartbeat" }; + try { + retainTranslatedEvent(event, budget); + const retainedBytes = budget.snapshot().currentBytes; + expect(() => retainTranslatedEvent(event, budget)).toThrow(/already retained/); + expect(budget.snapshot().currentBytes).toBe(retainedBytes); + releaseTranslatedEvent(event, budget); + expect(budget.snapshot().currentBytes).toBe(0); + } finally { + budget.dispose(); + } + }); + test("one one-shot tool call admits exactly 2 MiB and rejects one byte over", () => { const exact = createTranslatorBudget(); exact.openCall("call");