diff --git a/open-sse/executors/kiro.ts b/open-sse/executors/kiro.ts index 65c4e6186a..815d3a6348 100644 --- a/open-sse/executors/kiro.ts +++ b/open-sse/executors/kiro.ts @@ -61,7 +61,7 @@ type KiroStreamState = { contextUsagePercentage?: number; hasContextUsage?: boolean; hasMeteringEvent?: boolean; - usage?: UsageSummary; + usage?: Partial; hasReasoningContent?: boolean; reasoningChunkCount?: number; // Inline-thinking splitter state (populated only when thinkingExpected=true). @@ -185,8 +185,7 @@ function resolveKiroMaxInputTokens(model: string): number { * inflate `total_tokens`. */ function ensureKiroUsage(state: KiroStreamState, model: string) { - if (state.usage) return; - + if (state.usage?.total_tokens !== undefined) return; const estimatedOutputTokens = state.totalContentLength && state.totalContentLength > 0 ? Math.max(1, Math.floor(state.totalContentLength / 4)) @@ -198,11 +197,11 @@ function ensureKiroUsage(state: KiroStreamState, model: string) { : 0; 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 = { + ...state.usage, prompt_tokens: 0, completion_tokens: estimatedOutputTokens, total_tokens: estimatedOutputTokens, @@ -213,6 +212,7 @@ function ensureKiroUsage(state: KiroStreamState, model: string) { const promptTokens = Math.max(0, estimatedTotalTokens - estimatedOutputTokens); state.usage = { + ...state.usage, prompt_tokens: promptTokens, completion_tokens: estimatedOutputTokens, total_tokens: promptTokens + estimatedOutputTokens, diff --git a/tests/unit/executor-kiro.test.ts b/tests/unit/executor-kiro.test.ts index 38856c07be..86cd2b713b 100644 --- a/tests/unit/executor-kiro.test.ts +++ b/tests/unit/executor-kiro.test.ts @@ -339,8 +339,12 @@ test("KiroExecutor keeps cache tokens that arrive without input/output totals", 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); + assert.deepEqual(finish.usage, { + prompt_tokens: 19999, + completion_tokens: 1, + total_tokens: 20000, + cache_read_input_tokens: 900, + }); }); // snake_case spellings appear on some Kiro frames; a cache count must not be