From 43e1c28f3fd409eb71096b80f06b72e8c3a2fa6d Mon Sep 17 00:00:00 2001 From: Dohyun Jung Date: Thu, 6 Aug 2026 09:42:52 +0900 Subject: [PATCH] fix(kiro): read usage from the frames Kiro actually sends (#9035) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Validated in local merge-train (devbox-vm-06-dev002) @ combined-tip (FAST gates green: static + changed tests + vitest — only pre-existing audit.test.ts flake). Evidence: /home/diegosouzapw/dev/proxys/OmniRoute/.claude/worktrees/merge-train-20260805-213228-suite.log --- open-sse/executors/kiro.ts | 143 ++++++++++++--- open-sse/services/kiroModels.ts | 31 +++- tests/unit/executor-kiro.test.ts | 212 +++++++++++++++++++++++ tests/unit/kiro-available-models.test.ts | 76 ++++++++ 4 files changed, 434 insertions(+), 28 deletions(-) diff --git a/open-sse/executors/kiro.ts b/open-sse/executors/kiro.ts index 3dab64c395..eb0a93eacc 100644 --- a/open-sse/executors/kiro.ts +++ b/open-sse/executors/kiro.ts @@ -6,6 +6,7 @@ import { type ProviderCredentials, } from "./base.ts"; import { PROVIDERS } from "../config/constants.ts"; +import { getRegistryEntry } from "../config/providerRegistry.ts"; import { v4 as uuidv4 } from "uuid"; import { refreshKiroToken } from "../services/tokenRefresh.ts"; import { @@ -130,7 +131,45 @@ function buildKiroFinishChunk( return finishChunk; } -function ensureKiroUsage(state: KiroStreamState) { +/** + * Kiro's fallback input-token budget when the model is absent from the registry. + * Mirrors the registry's own `defaultContextLength` and kiro-gateway's + * DEFAULT_MAX_INPUT_TOKENS. + */ +const KIRO_DEFAULT_MAX_INPUT_TOKENS = 200000; + +/** + * Input-token budget for a Kiro model, used to turn `contextUsagePercentage` + * into an absolute token count. + * + * Kiro reports only a percentage, so the budget it is a percentage OF decides the + * result. A fixed 200000 undercounts every model with a larger window by the + * ratio of the two windows — claude-sonnet-5 (1M) by 5x, gpt-5.6-* (272k) by + * ~26% — and those numbers land in usage_history and the API-key token-limit + * counters. + */ +function resolveKiroMaxInputTokens(model: string): number { + const entry = getRegistryEntry("kiro"); + const modelEntry = entry?.models?.find((m) => m.id === model); + return modelEntry?.contextLength || entry?.defaultContextLength || KIRO_DEFAULT_MAX_INPUT_TOKENS; +} + +/** + * Synthesize a usage block when Kiro sent no token counts of its own. + * + * Live `generateAssistantResponse` traffic carries no token counts at all — only + * `contextUsageEvent.contextUsagePercentage` and a `meteringEvent` credit figure + * (verified against the live API: frames are assistantResponseEvent / + * metadataEvent / contextUsageEvent / meteringEvent). So these numbers are + * ESTIMATES, derived the same way kiro-gateway derives them: the percentage + * yields the total, the response text yields the completion, and the prompt is + * the remainder. + * + * Subtracting matters: the percentage already covers the whole context, so + * adding a separately-estimated completion on top would double-count it and + * inflate `total_tokens`. + */ +function ensureKiroUsage(state: KiroStreamState, model: string) { if (state.usage) return; const estimatedOutputTokens = @@ -138,17 +177,30 @@ function ensureKiroUsage(state: KiroStreamState) { ? Math.max(1, Math.floor(state.totalContentLength / 4)) : 0; - const estimatedInputTokens = + const estimatedTotalTokens = state.contextUsagePercentage && state.contextUsagePercentage > 0 - ? Math.floor((state.contextUsagePercentage * 200000) / 100) + ? Math.floor((state.contextUsagePercentage * resolveKiroMaxInputTokens(model)) / 100) : 0; - if (estimatedInputTokens <= 0 && estimatedOutputTokens <= 0) return; + if (estimatedTotalTokens <= 0 && estimatedOutputTokens <= 0) return; + + // Without a percentage there is no total to split, so the output estimate is + // all that is known and stands on its own. + if (estimatedTotalTokens <= 0) { + state.usage = { + prompt_tokens: 0, + completion_tokens: estimatedOutputTokens, + total_tokens: estimatedOutputTokens, + }; + return; + } + + const promptTokens = Math.max(0, estimatedTotalTokens - estimatedOutputTokens); state.usage = { - prompt_tokens: estimatedInputTokens, + prompt_tokens: promptTokens, completion_tokens: estimatedOutputTokens, - total_tokens: estimatedInputTokens + estimatedOutputTokens, + total_tokens: promptTokens + estimatedOutputTokens, }; } @@ -685,37 +737,74 @@ export class KiroExecutor extends BaseExecutor { state.hasMeteringEvent = true; } - // Handle metricsEvent for token usage - if (eventType === "metricsEvent") { - // Extract usage data from metricsEvent payload - const metrics = event.payload?.metricsEvent || event.payload; + // Handle token usage. Kiro reports it under more than one frame: the + // `metricsEvent` shape covered by unit tests, and a `metadataEvent` + // carrying a nested `usage` object — the shape observed on live + // API-key traffic (see tests/unit/executor-kiro.test.ts, the + // "live API-key event shape" case, whose frames are + // assistantResponseEvent / metadataEvent / contextUsageEvent / + // meteringEvent with no metricsEvent at all). Reading only + // `metricsEvent` meant cache tokens were never picked up in + // production even after their field names were corrected, because + // the branch holding that code never ran. + if (eventType === "metricsEvent" || eventType === "metadataEvent") { + const metrics = + event.payload?.metricsEvent || + event.payload?.usage || + (event.payload?.metadataEvent as JsonRecord)?.usage || + event.payload; if (metrics && typeof metrics === "object") { + const readNumber = (...candidates: unknown[]) => + candidates.find((value) => typeof value === "number") as number | undefined; + + // Bedrock-style (`inputTokens`) and OpenAI-style + // (`prompt_tokens`) spellings both appear across Kiro frames. const inputTokens = - typeof (metrics as JsonRecord).inputTokens === "number" - ? ((metrics as JsonRecord).inputTokens as number) - : 0; + readNumber( + (metrics as JsonRecord).inputTokens, + (metrics as JsonRecord).prompt_tokens + ) || 0; const outputTokens = - typeof (metrics as JsonRecord).outputTokens === "number" - ? ((metrics as JsonRecord).outputTokens as number) - : 0; + readNumber( + (metrics as JsonRecord).outputTokens, + (metrics as JsonRecord).completion_tokens + ) || 0; - const cacheReadTokens = - typeof (metrics as JsonRecord).cacheReadTokens === "number" - ? ((metrics as JsonRecord).cacheReadTokens as number) - : 0; + const cacheReadTokens = readNumber( + (metrics as JsonRecord).cacheReadInputTokens, + (metrics as JsonRecord).cacheReadTokens, + (metrics as JsonRecord).cache_read_input_tokens + ); - const cacheCreationTokens = - typeof (metrics as JsonRecord).cacheCreationTokens === "number" - ? ((metrics as JsonRecord).cacheCreationTokens as number) - : 0; + const cacheCreationTokens = readNumber( + (metrics as JsonRecord).cacheWriteInputTokens, + (metrics as JsonRecord).cacheCreationTokens, + (metrics as JsonRecord).cache_creation_input_tokens + ); if (inputTokens > 0 || outputTokens > 0) { state.usage = { prompt_tokens: inputTokens, completion_tokens: outputTokens, total_tokens: inputTokens + outputTokens, - ...(cacheReadTokens > 0 && { cache_read_input_tokens: cacheReadTokens }), - ...(cacheCreationTokens > 0 && { + ...((cacheReadTokens || 0) > 0 && { + cache_read_input_tokens: cacheReadTokens, + }), + ...((cacheCreationTokens || 0) > 0 && { + cache_creation_input_tokens: cacheCreationTokens, + }), + }; + } else if ((cacheReadTokens || 0) > 0 || (cacheCreationTokens || 0) > 0) { + // Cache counts can arrive on a frame that carries no + // input/output totals. Preserve them instead of dropping the + // whole frame, and let ensureKiroUsage() fill the totals from + // contextUsagePercentage. + state.usage = { + ...(state.usage || {}), + ...((cacheReadTokens || 0) > 0 && { + cache_read_input_tokens: cacheReadTokens, + }), + ...((cacheCreationTokens || 0) > 0 && { cache_creation_input_tokens: cacheCreationTokens, }), }; @@ -772,7 +861,7 @@ export class KiroExecutor extends BaseExecutor { // Emit finish chunk if not already sent if (!state.finishEmitted) { state.finishEmitted = true; - ensureKiroUsage(state); + ensureKiroUsage(state, model); const finishChunk = buildKiroFinishChunk(state, responseId, created, model, true); controller.enqueue(TEXT_ENCODER.encode(`data: ${JSON.stringify(finishChunk)}\n\n`)); } diff --git a/open-sse/services/kiroModels.ts b/open-sse/services/kiroModels.ts index dade388504..5a64410fc4 100644 --- a/open-sse/services/kiroModels.ts +++ b/open-sse/services/kiroModels.ts @@ -56,6 +56,12 @@ function toNonEmptyString(value: unknown): string | null { return trimmed.length > 0 ? trimmed : null; } +export type KiroPromptCaching = { + supportsPromptCaching: boolean; + minimumTokensPerCacheCheckpoint: number | null; + maximumCacheCheckpointsPerRequest: number | null; +}; + export type KiroModel = { id: string; name: string; @@ -68,8 +74,28 @@ export type KiroModel = { rateMultiplier?: number; upstreamModelId?: string; description?: string; + promptCaching?: KiroPromptCaching; }; +function toNonNegativeInteger(value: unknown): number | null { + return typeof value === "number" && Number.isInteger(value) && value >= 0 ? value : null; +} + +function parsePromptCaching(value: unknown): KiroPromptCaching | undefined { + const promptCaching = asRecord(value); + if (typeof promptCaching.supportsPromptCaching !== "boolean") return undefined; + + return { + supportsPromptCaching: promptCaching.supportsPromptCaching, + minimumTokensPerCacheCheckpoint: toNonNegativeInteger( + promptCaching.minimumTokensPerCacheCheckpoint + ), + maximumCacheCheckpointsPerRequest: toNonNegativeInteger( + promptCaching.maximumCacheCheckpointsPerRequest + ), + }; +} + export type KiroModelsResult = { models: KiroModel[]; /** "api" = live discovery; "fallback" = static catalog (offline/unauthed/error). */ @@ -98,7 +124,8 @@ export function parseKiroModels(data: unknown): KiroModel[] { if (!id || seen.has(id)) continue; seen.add(id); const name = toNonEmptyString(item.modelName) || toNonEmptyString(item.name) || id; - models.push({ id, name, owned_by: "kiro" }); + const promptCaching = parsePromptCaching(item.promptCaching); + models.push({ id, name, owned_by: "kiro", ...(promptCaching && { promptCaching }) }); } return models; @@ -162,6 +189,7 @@ function expandKiroModels(data: unknown): KiroModel[] { const tokenLimits = asRecord(item.tokenLimits); const contextLength = Number(tokenLimits.maxInputTokens) || 200000; const rateMultiplier = Number(item.rateMultiplier); + const promptCaching = parsePromptCaching(item.promptCaching); for (const variant of buildVariants(upstreamId, display)) { if (seen.has(variant.id)) continue; @@ -172,6 +200,7 @@ function expandKiroModels(data: unknown): KiroModel[] { rateMultiplier: Number.isFinite(rateMultiplier) ? rateMultiplier : 1.0, upstreamModelId: upstreamId, description: toNonEmptyString(item.description) || "", + ...(promptCaching && { promptCaching }), }); } } diff --git a/tests/unit/executor-kiro.test.ts b/tests/unit/executor-kiro.test.ts index 8fde25b3e3..38856c07be 100644 --- a/tests/unit/executor-kiro.test.ts +++ b/tests/unit/executor-kiro.test.ts @@ -241,6 +241,218 @@ test("KiroExecutor.transformEventStreamToSSE converts text, tool calls, usage an assert.match(text, /\[DONE\]/); }); +test("KiroExecutor normalizes Bedrock cache-token fields from a metricsEvent", async () => { + const executor = new KiroExecutor(); + const response = buildEventStreamResponse([ + buildEventFrame("assistantResponseEvent", { content: "cached" }), + buildEventFrame("metricsEvent", { + inputTokens: 7, + outputTokens: 2, + cacheReadInputTokens: 1024, + cacheWriteInputTokens: 256, + }), + ]); + + const transformed = executor.transformEventStreamToSSE(response, "kiro-model"); + const chunks = parseSSEJsonChunks(await transformed.text()); + const finish = chunks.find((chunk) => chunk.choices?.[0]?.finish_reason); + + assert.deepEqual(finish.usage, { + prompt_tokens: 7, + completion_tokens: 2, + total_tokens: 9, + cache_read_input_tokens: 1024, + cache_creation_input_tokens: 256, + }); +}); + +test("KiroExecutor does not invent cache tokens for the live API-key event shape", async () => { + const executor = new KiroExecutor(); + const response = buildEventStreamResponse([ + buildEventFrame("assistantResponseEvent", { + content: "live-shaped response", + modelId: "claude-sonnet-4.5", + }), + buildEventFrame("metadataEvent", { stopReason: "END_TURN" }), + buildEventFrame("contextUsageEvent", { contextUsagePercentage: 4.93 }), + buildEventFrame("meteringEvent", { + unit: "credit", + unitPlural: "credits", + usage: 0.022, + }), + ]); + + const transformed = executor.transformEventStreamToSSE(response, "kiro-model"); + const chunks = parseSSEJsonChunks(await transformed.text()); + const finish = chunks.find((chunk) => chunk.choices?.[0]?.finish_reason); + + assert.equal(finish.usage.cache_read_input_tokens, undefined); + assert.equal(finish.usage.cache_creation_input_tokens, undefined); +}); + +// The cache-field fix in d205650a7 landed inside the `metricsEvent` branch, but +// live API-key traffic (the test above) sends no `metricsEvent` at all — it sends +// a `metadataEvent`. The corrected code was therefore unreachable in production +// and cache stats stayed empty. Usage extraction now accepts both frames. +test("KiroExecutor reads usage and cache tokens from a metadataEvent frame", async () => { + const executor = new KiroExecutor(); + const response = buildEventStreamResponse([ + buildEventFrame("assistantResponseEvent", { content: "cached reply" }), + buildEventFrame("metadataEvent", { + stopReason: "END_TURN", + usage: { + inputTokens: 12, + outputTokens: 3, + cacheReadInputTokens: 2048, + cacheWriteInputTokens: 512, + }, + }), + ]); + + const transformed = executor.transformEventStreamToSSE(response, "kiro-model"); + const chunks = parseSSEJsonChunks(await transformed.text()); + const finish = chunks.find((chunk) => chunk.choices?.[0]?.finish_reason); + + assert.deepEqual(finish.usage, { + prompt_tokens: 12, + completion_tokens: 3, + total_tokens: 15, + cache_read_input_tokens: 2048, + cache_creation_input_tokens: 512, + }); +}); + +// Cache counts can arrive on a frame carrying no input/output totals. Dropping +// the frame (the old behavior, gated on `inputTokens > 0 || outputTokens > 0`) +// lost the cache accounting entirely. +test("KiroExecutor keeps cache tokens that arrive without input/output totals", async () => { + const executor = new KiroExecutor(); + const response = buildEventStreamResponse([ + buildEventFrame("assistantResponseEvent", { content: "hi" }), + buildEventFrame("contextUsageEvent", { contextUsagePercentage: 10 }), + buildEventFrame("metadataEvent", { + usage: { cacheReadInputTokens: 900 }, + }), + ]); + + const transformed = executor.transformEventStreamToSSE(response, "kiro-model"); + const chunks = parseSSEJsonChunks(await transformed.text()); + const finish = chunks.find((chunk) => chunk.choices?.[0]?.finish_reason); + + assert.equal(finish.usage.cache_read_input_tokens, 900); + assert.equal(finish.usage.cache_creation_input_tokens, undefined); +}); + +// snake_case spellings appear on some Kiro frames; a cache count must not be +// dropped just because it is not the Bedrock camelCase spelling. +test("KiroExecutor accepts snake_case cache token spellings", async () => { + const executor = new KiroExecutor(); + const response = buildEventStreamResponse([ + buildEventFrame("assistantResponseEvent", { content: "ok" }), + buildEventFrame("metadataEvent", { + usage: { + prompt_tokens: 5, + completion_tokens: 1, + cache_read_input_tokens: 64, + cache_creation_input_tokens: 32, + }, + }), + ]); + + const transformed = executor.transformEventStreamToSSE(response, "kiro-model"); + const chunks = parseSSEJsonChunks(await transformed.text()); + const finish = chunks.find((chunk) => chunk.choices?.[0]?.finish_reason); + + assert.deepEqual(finish.usage, { + prompt_tokens: 5, + completion_tokens: 1, + total_tokens: 6, + cache_read_input_tokens: 64, + cache_creation_input_tokens: 32, + }); +}); + +// Live generateAssistantResponse sends NO token counts — only +// contextUsageEvent.contextUsagePercentage plus a meteringEvent credit figure. +// The synthesized usage is therefore an estimate, and the context budget the +// percentage applies to decides it. A fixed 200000 undercounts claude-sonnet-5 +// (1M window) by 5x, and those numbers feed usage_history and the API-key +// token-limit counters. +test("KiroExecutor scales the usage estimate by the model's own context window", async () => { + const executor = new KiroExecutor(); + const estimateFor = async (model) => { + const response = buildEventStreamResponse([ + buildEventFrame("assistantResponseEvent", { content: "x".repeat(400) }), + buildEventFrame("contextUsageEvent", { contextUsagePercentage: 10 }), + ]); + const chunks = parseSSEJsonChunks( + await executor.transformEventStreamToSSE(response, model).text() + ); + return chunks.find((chunk) => chunk.choices?.[0]?.finish_reason).usage; + }; + + // 10% of 200000, with the 100-token completion carved out of the total. + assert.deepEqual(await estimateFor("claude-sonnet-4.5"), { + prompt_tokens: 19900, + completion_tokens: 100, + total_tokens: 20000, + }); + + // 10% of 1000000 — a fixed 200000 budget would have reported 20000 here. + assert.deepEqual(await estimateFor("claude-sonnet-5"), { + prompt_tokens: 99900, + completion_tokens: 100, + total_tokens: 100000, + }); + + // 10% of 272000. + assert.equal((await estimateFor("gpt-5.6-sol")).total_tokens, 27200); + + // Registry models without their own contextLength inherit defaultContextLength, + // and an unknown id falls back to the same budget rather than reporting zero. + assert.equal((await estimateFor("glm-5")).total_tokens, 20000); + assert.equal((await estimateFor("not-a-kiro-model")).total_tokens, 20000); +}); + +// The percentage already covers the whole context, so adding a separately +// estimated completion on top would double-count it and inflate total_tokens +// past what Kiro reported. +test("KiroExecutor carves the completion estimate out of the reported total", async () => { + const executor = new KiroExecutor(); + const response = buildEventStreamResponse([ + buildEventFrame("assistantResponseEvent", { content: "y".repeat(800) }), + buildEventFrame("contextUsageEvent", { contextUsagePercentage: 5 }), + ]); + + const chunks = parseSSEJsonChunks( + await executor.transformEventStreamToSSE(response, "claude-sonnet-4.5").text() + ); + const usage = chunks.find((chunk) => chunk.choices?.[0]?.finish_reason).usage; + + // 5% of 200000 is 10000 and must stay the total, not become 10000 + 200. + assert.equal(usage.total_tokens, 10000); + assert.equal(usage.completion_tokens, 200); + assert.equal(usage.prompt_tokens, 9800); +}); + +// With no percentage there is no total to split, so the output estimate has to +// stand on its own instead of being silently dropped. +test("KiroExecutor still reports a completion estimate without a context percentage", async () => { + const executor = new KiroExecutor(); + const response = buildEventStreamResponse([ + buildEventFrame("assistantResponseEvent", { content: "z".repeat(400) }), + ]); + + const chunks = parseSSEJsonChunks( + await executor.transformEventStreamToSSE(response, "claude-sonnet-4.5").text() + ); + const usage = chunks.find((chunk) => chunk.choices?.[0]?.finish_reason).usage; + + assert.equal(usage.prompt_tokens, 0); + assert.equal(usage.completion_tokens, 100); + assert.equal(usage.total_tokens, 100); +}); + test("KiroExecutor.transformEventStreamToSSE surfaces native reasoning frames as reasoning_content", async () => { const executor = new KiroExecutor(); // Verified live wire format: Kiro streams adaptive-thinking reasoning as a diff --git a/tests/unit/kiro-available-models.test.ts b/tests/unit/kiro-available-models.test.ts index d1cf7ae0cb..0aa052050b 100644 --- a/tests/unit/kiro-available-models.test.ts +++ b/tests/unit/kiro-available-models.test.ts @@ -41,6 +41,82 @@ test("parseKiroModels reads CodeWhisperer ListAvailableModels shape", () => { assert.equal(models[0].owned_by, "kiro"); }); +test("parseKiroModels preserves live prompt-caching capability metadata", () => { + const [model] = parseKiroModels({ + models: [ + { + modelId: "claude-sonnet-4.5", + modelName: "Claude Sonnet 4.5", + promptCaching: { + supportsPromptCaching: true, + minimumTokensPerCacheCheckpoint: 1024, + maximumCacheCheckpointsPerRequest: 4, + }, + }, + ], + }); + + assert.deepEqual(model.promptCaching, { + supportsPromptCaching: true, + minimumTokensPerCacheCheckpoint: 1024, + maximumCacheCheckpointsPerRequest: 4, + }); +}); + +test("parseKiroModels keeps nonnumeric prompt-caching limits unknown", () => { + const [model] = parseKiroModels({ + models: [ + { + modelId: "claude-sonnet-4.5", + promptCaching: { + supportsPromptCaching: true, + minimumTokensPerCacheCheckpoint: null, + maximumCacheCheckpointsPerRequest: false, + }, + }, + ], + }); + + assert.deepEqual(model.promptCaching, { + supportsPromptCaching: true, + minimumTokensPerCacheCheckpoint: null, + maximumCacheCheckpointsPerRequest: null, + }); +}); + +test("fetchKiroAvailableModels carries upstream prompt-caching metadata to model variants", async () => { + const fetchImpl = (async () => + jsonResponse({ + models: [ + { + modelId: "claude-sonnet-4.5", + promptCaching: { + supportsPromptCaching: true, + minimumTokensPerCacheCheckpoint: 1024, + maximumCacheCheckpointsPerRequest: 4, + }, + }, + ], + })) as unknown as typeof fetch; + + const result = await fetchKiroAvailableModels({ + accessToken: "tok", + providerSpecificData: {}, + fetchImpl, + fallbackModels: FALLBACK, + }); + + assert.ok(result.models.length >= 1); + for (const model of result.models) { + assert.equal(model.upstreamModelId, "claude-sonnet-4.5"); + assert.deepEqual(model.promptCaching, { + supportsPromptCaching: true, + minimumTokensPerCacheCheckpoint: 1024, + maximumCacheCheckpointsPerRequest: 4, + }); + } +}); + test("resolveKiroRegion prefers stored region, then profileArn, else us-east-1", () => { assert.equal(resolveKiroRegion({ region: "eu-central-1" }), "eu-central-1"); assert.equal(