Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 23 additions & 13 deletions src/adapters/kiro.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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;
}
},
};
}
15 changes: 15 additions & 0 deletions src/lib/translator-budget.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
18 changes: 18 additions & 0 deletions tests/kiro-stream.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading