From f5b27ebac9443ceeda8a3b4f2dc9bcf8920fa624 Mon Sep 17 00:00:00 2001
From: luvs01 <luvs01@hanmail.net>
Date: Mon, 10 Aug 2026 11:09:55 +0900
Subject: [PATCH] fix(kiro): bound non-streaming event collection

---
 src/adapters/kiro.ts         | 36 +++++++++++++++++++++++-------------
 src/lib/translator-budget.ts | 15 +++++++++++++++
 tests/kiro-stream.test.ts    | 18 ++++++++++++++++++
 3 files changed, 56 insertions(+), 13 deletions(-)

diff --git a/src/adapters/kiro.ts b/src/adapters/kiro.ts
index 2ea1bc84d..f859dc5b3 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 {
@@ -1906,19 +1908,27 @@ export function createKiroAdapter(provider: OcxProviderConfig): ProviderAdapter
     // tool failed with "web-search sidecar requires a non-streaming adapter" (kiro-only).
     async parseResponse(response: Response, budget: TranslatorBudget): Promise<AdapterEvent[]> {
       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.length === 0);
+          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 1bba4512d..076cdfd72 100644
--- a/src/lib/translator-budget.ts
+++ b/src/lib/translator-budget.ts
@@ -71,6 +71,21 @@ export interface TranslatorBudget {
 
 const retainedEventOwnership = new WeakMap<object, { budget: TranslatorBudget; bytes: number }>();
 
+/**
+ * Charge one event while an adapter incrementally materializes a response batch.
+ * The first event owns both JSON array brackets; later events own their comma, so
+ * the aggregate charge is identical to `retainTranslatedEventBatch`.
+ */
+export function retainTranslatedEvent<T extends object>(
+  event: T,
+  budget: TranslatorBudget,
+  first: boolean,
+): void {
+  const bytes = Buffer.byteLength(JSON.stringify(event)) + (first ? 2 : 1);
+  budget.chargeRetained(bytes, { kind: "retained_collectors" });
+  retainedEventOwnership.set(event, { budget, bytes });
+}
+
 /**
  * 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 744fa8ab1..388f4349c 100644
--- a/tests/kiro-stream.test.ts
+++ b/tests/kiro-stream.test.ts
@@ -1751,6 +1751,24 @@ 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 });
+    try {
+      const adapter = createKiroAdapterProduction(provider);
+      const response = new Response(streamOf(
+        eventFrame({ content: "a".repeat(80) }),
+        eventFrame({ content: "b".repeat(80) }),
+      ));
+
+      await expect(adapter.parseResponse!(response, budget)).rejects.toMatchObject({
+        code: "translation_buffer_limit",
+      });
+      expect(budget.snapshot().currentBytes).toBe(0);
+    } 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