diff --git a/src/adapters/kiro.ts b/src/adapters/kiro.ts index 9cb55e0625..9e43e891c0 100644 --- a/src/adapters/kiro.ts +++ b/src/adapters/kiro.ts @@ -32,6 +32,7 @@ import type { import type { ProviderAdapter } from "./base"; import type { AdapterFetchContext, AdapterRequest } from "./base"; import { extractKiroImages, normalizeKiroImages, type KiroImage } from "./kiro-images"; +import { sniffImageDimensions } from "./anthropic-image-guard"; import { fetchKiroWithRetry } from "./kiro-retry"; import { convertKiroToolContext } from "./kiro-tools"; import { neutralizeIdentity } from "./identity"; @@ -134,6 +135,52 @@ function messageLogText(msg: OcxMessage): string { }).filter(Boolean).join("\n"); } +function estimateKiroImageTokens(image: KiroImage): number { + const dimensions = sniffImageDimensions(image.source.bytes); + if (dimensions) { + return Math.max(256, Math.ceil(dimensions.width * dimensions.height / 750)); + } + const decodedBytes = Math.floor(image.source.bytes.length * 3 / 4); + return Math.max(256, Math.ceil(decodedBytes / 512)); +} + +function estimateKiroTokens(text: string, modelId?: string): number { + return estimateTokens(text, modelId ? `kiro/${modelId}` : "kiro"); +} + +function estimateKiroPayloadInputTokens(payload: Record, modelId: string): number { + const conversationState = (payload as { + conversationState?: { + history?: KiroHistoryEntry[]; + currentMessage?: KiroHistoryEntry; + }; + }).conversationState; + if (!conversationState) return 0; + + const parts: string[] = []; + let imageTokens = 0; + const entries = [ + ...(conversationState.history ?? []), + ...(conversationState.currentMessage ? [conversationState.currentMessage] : []), + ]; + for (const entry of entries) { + const user = entry.userInputMessage; + if (user) { + if (user.content) parts.push(user.content); + for (const image of user.images ?? []) imageTokens += estimateKiroImageTokens(image); + const context = user.userInputMessageContext; + if (context?.tools?.length) parts.push(serializeForUsage(context.tools)); + if (context?.toolResults?.length) parts.push(serializeForUsage(context.toolResults)); + } + const assistant = entry.assistantResponseMessage; + if (assistant) { + if (assistant.content) parts.push(assistant.content); + if (assistant.toolUses?.length) parts.push(serializeForUsage(assistant.toolUses)); + } + } + return estimateKiroTokens(parts.join("\n"), modelId) + imageTokens; +} + function shouldCountStablePromptOverhead(parsed: OcxParsedRequest): boolean { return !parsed.previousResponseId && !parsed.context.messages.some(m => m.role === "assistant"); } @@ -148,25 +195,21 @@ function estimateKiroInputTokens(parsed: OcxParsedRequest): number { if (parsed.context.tools?.length) parts.push(serializeForUsage(parsed.context.tools)); } - return estimateTokens(parts.join("\n"), parsed.modelId); + return estimateKiroTokens(parts.join("\n"), parsed.modelId); } function estimateKiroLogInputTokens(parsed: OcxParsedRequest): number { const parts = parsed.context.messages.map(messageLogText).filter(Boolean); if (parsed.context.systemPrompt?.length) parts.push(...parsed.context.systemPrompt); if (parsed.context.tools?.length) parts.push(serializeForUsage(parsed.context.tools)); - return Math.max(estimateKiroInputTokens(parsed), estimateTokens(parts.join("\n"), parsed.modelId)); + return Math.max(estimateKiroInputTokens(parsed), estimateKiroTokens(parts.join("\n"), parsed.modelId)); } -function configuredKiroContextWindow(provider: OcxProviderConfig, modelId: string | undefined): number | undefined { +function kiroUpstreamContextWindow(modelId: string | undefined): number | undefined { if (!modelId) return undefined; const normalizedModelId = normalizeKiroModelId(modelId); if (normalizedModelId === "auto") return undefined; - const window = - modelRecordValue(provider.modelContextWindows, modelId) - ?? modelRecordValue(provider.modelContextWindows, normalizedModelId) - ?? provider.contextWindow - ?? modelRecordValue(KIRO_MODEL_CONTEXT_WINDOWS, modelId) + const window = modelRecordValue(KIRO_MODEL_CONTEXT_WINDOWS, modelId) ?? modelRecordValue(KIRO_MODEL_CONTEXT_WINDOWS, normalizedModelId); return typeof window === "number" && Number.isFinite(window) && window > 0 ? window : undefined; } @@ -465,17 +508,26 @@ interface KiroAttemptResult { interface KiroFallbackAttempt { response: Response; inputTokens: number; + contextInputEstimate: number; nameMap: Map; conversationId: string; } +interface KiroContextWindowState { + value?: number; +} + type KiroFallbackFactory = ( conversationId: string | undefined, assistantText: string, sawReasoning: boolean, ) => Promise; -function mergeKiroUsage(first: OcxUsage | undefined, second: OcxUsage | undefined): OcxUsage | undefined { +function mergeKiroUsage( + first: OcxUsage | undefined, + second: OcxUsage | undefined, + preserveFirstContextGrowth = false, +): OcxUsage | undefined { if (!first) return second; if (!second) return first; const sumOptional = (key: keyof OcxUsage): number | undefined => { @@ -488,9 +540,23 @@ function mergeKiroUsage(first: OcxUsage | undefined, second: OcxUsage | undefine const totalTokens = typeof first.totalTokens === "number" && typeof second.totalTokens === "number" ? first.totalTokens + second.totalTokens : undefined; + const carriedContextTotal = preserveFirstContextGrowth && typeof first.contextTotalTokens === "number" + ? first.contextTotalTokens + second.outputTokens + : undefined; + const combinedOutputTokens = first.outputTokens + second.outputTokens; return { inputTokens: first.inputTokens + second.inputTokens, - outputTokens: first.outputTokens + second.outputTokens, + outputTokens: combinedOutputTokens, + ...(typeof first.contextTotalTokens === "number" || typeof second.contextTotalTokens === "number" + ? { + contextTotalTokens: Math.max( + first.contextTotalTokens ?? 0, + second.contextTotalTokens ?? 0, + carriedContextTotal ?? 0, + combinedOutputTokens, + ), + } + : {}), ...(totalTokens !== undefined ? { totalTokens } : {}), ...(sumOptional("cachedInputTokens") !== undefined ? { cachedInputTokens: sumOptional("cachedInputTokens") } : {}), ...(sumOptional("cacheReadInputTokens") !== undefined ? { cacheReadInputTokens: sumOptional("cacheReadInputTokens") } : {}), @@ -530,10 +596,11 @@ async function* parseKiroAttempt( mode: KiroCompletionMode, modelId: string | undefined, inputTokens: number, - contextWindow: number | undefined, + contextWindowState: KiroContextWindowState, nameMap: Map | undefined, conversationId: string | undefined, previousAssistantText?: string, + contextInputEstimate?: number, ): AsyncGenerator { const emptyResult = (): KiroAttemptResult => ({ assistantText: "", sawReasoning: false }); if (!response.body) { @@ -560,11 +627,28 @@ async function* parseKiroAttempt( const providerState = (): { kiro: { conversationId: string } } | undefined => returnedConversationId ? { kiro: { conversationId: returnedConversationId } } : undefined; - const usage = (): OcxUsage => authoritativeUsage ?? ({ + const contextUsageTotalFloor = (): number | undefined => { + if (contextUsagePercentage === undefined || !contextWindowState.value) return undefined; + const floor = Math.ceil(contextWindowState.value * Math.min(contextUsagePercentage, 100) / 100); + return Number.isFinite(floor) && floor > 0 ? floor : undefined; + }; + const usage = (): OcxUsage => { + const base = authoritativeUsage ?? { inputTokens, - outputTokens: estimateTokens(outputChars, modelId), + outputTokens: estimateKiroTokens(outputChars, modelId), estimated: true, - }); + }; + const estimatedContextTotal = contextInputEstimate !== undefined + ? contextInputEstimate + base.outputTokens + : undefined; + const authoritativeTurnTotal = base.inputTokens + base.outputTokens; + const contextTotal = Math.max( + estimatedContextTotal ?? 0, + contextUsageTotalFloor() ?? 0, + authoritativeTurnTotal, + ); + return contextTotal > 0 ? { ...base, contextTotalTokens: contextTotal } : base; + }; const classifiedTerminal = (failure: KiroErrorClassification): AdapterEvent => ({ type: "error", @@ -743,6 +827,9 @@ async function* parseKiroAttempt( if (isValidKiroConversationId(ev.conversationId)) returnedConversationId = ev.conversationId; break; case "content": + if (ev.modelId) { + contextWindowState.value = kiroUpstreamContextWindow(ev.modelId) ?? contextWindowState.value; + } if (open) { open = null; return { assistantText, sawReasoning, terminal: protocolTerminal(kiroTruncationErrorMessage("content arrived before tool stop")) }; @@ -843,7 +930,7 @@ async function* parseKiroAttempt( if (contextUsagePercentage !== undefined) { debugProviderDiagnostic("kiro", "context_usage", { contextUsagePercentage, - ...(contextWindow ? { configuredContextWindow: contextWindow } : {}), + ...(contextWindowState.value ? { upstreamContextWindow: contextWindowState.value } : {}), }); } debugProviderDiagnostic("kiro", "attempt_complete", { @@ -970,15 +1057,19 @@ export async function* parseKiroStream( conversationId?: string, completionMode: KiroCompletionMode = "disabled", fallbackFactory?: KiroFallbackFactory, + contextInputEstimate?: number, ): AsyncGenerator { + const contextWindowState: KiroContextWindowState = { value: contextWindow }; const first = parseKiroAttempt( response, completionMode, modelId, inputTokens, - contextWindow, + contextWindowState, nameMap, conversationId, + undefined, + contextInputEstimate, ); let firstNext = await first.next(); while (!firstNext.done) { @@ -1039,10 +1130,11 @@ export async function* parseKiroStream( "text_fallback", modelId, fallback.inputTokens, - contextWindow, + contextWindowState, fallback.nameMap, fallback.conversationId, firstResult.assistantText, + fallback.contextInputEstimate, ); let secondNext = await second.next(); while (!secondNext.done) { @@ -1054,7 +1146,8 @@ export async function* parseKiroStream( yield retryableKiroIncomplete( "empty_kiro_fallback", "Kiro's bounded completion retry ended without a terminal result", - mergeKiroUsage(firstResult.usage, secondResult.usage) ?? { inputTokens, outputTokens: 0, estimated: true }, + mergeKiroUsage(firstResult.usage, secondResult.usage, Boolean(firstResult.assistantText)) + ?? { inputTokens, outputTokens: 0, estimated: true }, secondResult.providerState ?? firstResult.providerState, ); return; @@ -1062,7 +1155,7 @@ export async function* parseKiroStream( if (secondResult.terminal.type === "done" || secondResult.terminal.type === "incomplete") { yield { ...secondResult.terminal, - usage: mergeKiroUsage(firstResult.usage, secondResult.terminal.usage), + usage: mergeKiroUsage(firstResult.usage, secondResult.terminal.usage, Boolean(firstResult.assistantText)), providerState: secondResult.terminal.providerState ?? firstResult.providerState, }; return; @@ -1070,7 +1163,7 @@ export async function* parseKiroStream( yield { ...secondResult.terminal, ...(secondResult.terminal.type === "error" - ? { usage: mergeKiroUsage(firstResult.usage, secondResult.terminal.usage) } + ? { usage: mergeKiroUsage(firstResult.usage, secondResult.terminal.usage, Boolean(firstResult.assistantText)) } : {}), }; } @@ -1080,6 +1173,7 @@ export function createKiroAdapter(provider: OcxProviderConfig): ProviderAdapter // Per-request closure (resolveAdapter builds a fresh adapter per request — server.ts:440 — so this // is race-free) carrying the heuristic input-token estimate from buildRequest into the stream. let inputTokens = 0; + let contextInputEstimate = 0; let modelId: string | undefined; let contextWindow: number | undefined; let toolNameMap: Map | undefined; @@ -1097,6 +1191,7 @@ export function createKiroAdapter(provider: OcxProviderConfig): ProviderAdapter conversationId: string; completionMode: KiroCompletionMode; inputTokens: number; + contextInputEstimate: number; }> => { if (typeof provider.apiKey !== "string" || provider.apiKey.trim() === "") { throw new Error("kiro token missing — run ocx login kiro"); @@ -1118,6 +1213,7 @@ export function createKiroAdapter(provider: OcxProviderConfig): ProviderAdapter if (profileArn) headers["x-amzn-kiro-profile-arn"] = profileArn; const built = buildKiroPayload(parsed, profileArn, forcedCompletionMode); await normalizeKiroImages(built.payload); + const contextInputEstimate = estimateKiroPayloadInputTokens(built.payload, parsed.modelId); const body = JSON.stringify(built.payload); debugProviderDiagnostic("kiro", "request", { region, @@ -1141,6 +1237,7 @@ export function createKiroAdapter(provider: OcxProviderConfig): ProviderAdapter conversationId: built.conversationId, completionMode: built.completionMode, inputTokens: estimateKiroInputTokens(parsed), + contextInputEstimate, }; }; @@ -1179,6 +1276,7 @@ export function createKiroAdapter(provider: OcxProviderConfig): ProviderAdapter return { response, inputTokens: retry.inputTokens, + contextInputEstimate: retry.contextInputEstimate, nameMap: retry.nameMap, conversationId: retry.conversationId, }; @@ -1189,8 +1287,9 @@ export function createKiroAdapter(provider: OcxProviderConfig): ProviderAdapter async buildRequest(parsed: OcxParsedRequest, incoming) { const built = await build(parsed); modelId = parsed.modelId; - contextWindow = configuredKiroContextWindow(provider, parsed.modelId); + contextWindow = kiroUpstreamContextWindow(parsed.modelId); inputTokens = built.inputTokens; + contextInputEstimate = built.contextInputEstimate; toolNameMap = built.nameMap; conversationId = built.conversationId; completionMode = built.completionMode; @@ -1209,6 +1308,7 @@ export function createKiroAdapter(provider: OcxProviderConfig): ProviderAdapter conversationId, completionMode, completionMode === "required" ? fallbackFactory : undefined, + contextInputEstimate, ); }, @@ -1239,6 +1339,7 @@ export function createKiroAdapter(provider: OcxProviderConfig): ProviderAdapter conversationId, completionMode, completionMode === "required" ? fallbackFactory : undefined, + contextInputEstimate, )) events.push(e); return events; }, diff --git a/src/bridge.ts b/src/bridge.ts index 09f11dd40e..6e7025f1e6 100644 --- a/src/bridge.ts +++ b/src/bridge.ts @@ -15,20 +15,29 @@ function sseEvent(name: string, data: Record): string { function responsesUsage(usage: OcxUsage | undefined): Record { if (!usage) return { input_tokens: 0, output_tokens: 0, total_tokens: 0 }; - // inputTokens is already inclusive of cache read/write (types.ts convention). - const inputTokens = usage.inputTokens; + // Stateful providers may report an absolute active-context checkpoint separately from their + // per-attempt usage. Split that checkpoint into input + output without adding output twice. + const inputTokens = usage.contextTotalTokens !== undefined + ? Math.max(0, usage.contextTotalTokens - usage.outputTokens) + : usage.inputTokens; const out: Record = { input_tokens: inputTokens, output_tokens: usage.outputTokens, - total_tokens: usageDisplayTotalTokens(usage) ?? inputTokens + usage.outputTokens, + total_tokens: usage.contextTotalTokens !== undefined + ? usage.contextTotalTokens + : usageDisplayTotalTokens(usage) ?? inputTokens + usage.outputTokens, }; const inputDetails: Record = {}; if (usage.cachedInputTokens !== undefined) { // cached_tokens carries cache READS only, matching OpenAI semantics. - inputDetails.cached_tokens = usage.cachedInputTokens; + inputDetails.cached_tokens = Math.min(usage.cachedInputTokens, inputTokens); } if (usage.cacheCreationInputTokens !== undefined) { - inputDetails.cache_write_tokens = usage.cacheCreationInputTokens; + const cacheRead = inputDetails.cached_tokens ?? 0; + inputDetails.cache_write_tokens = Math.min( + usage.cacheCreationInputTokens, + Math.max(0, inputTokens - cacheRead), + ); } if (Object.keys(inputDetails).length > 0) { out.input_tokens_details = inputDetails; diff --git a/src/types.ts b/src/types.ts index c5937600f8..9df4c114c3 100644 --- a/src/types.ts +++ b/src/types.ts @@ -308,6 +308,12 @@ export interface OcxUrlCitation { export interface OcxUsage { inputTokens: number; outputTokens: number; + /** + * Absolute active-context size after the response. Stateful providers can expose this separately + * from their per-attempt usage. Responses serialization derives the input side from + * `contextTotalTokens - outputTokens` so output is never added to an absolute checkpoint twice. + */ + contextTotalTokens?: number; totalTokens?: number; cachedInputTokens?: number; cacheReadInputTokens?: number; diff --git a/tests/bridge.test.ts b/tests/bridge.test.ts index 13f60ee414..00b419f187 100644 --- a/tests/bridge.test.ts +++ b/tests/bridge.test.ts @@ -129,6 +129,43 @@ describe("Responses bridge reasoning and usage parity", () => { }); }); + test("absolute context total drives Responses compaction without double-counting output", async () => { + const frames = await collectSse(bridgeToResponsesSSE(replay([ + { + type: "done", + usage: { + inputTokens: 58, + contextTotalTokens: 226_000, + outputTokens: 12, + estimated: true, + }, + }, + ]), "kiro/claude-opus-5")); + + const completed = frames.find(f => f.event === "response.completed")?.data.response as Record; + expect(completed.usage).toEqual({ + input_tokens: 225_988, + output_tokens: 12, + total_tokens: 226_000, + }); + }); + + test("consecutive context checkpoints remain absolute instead of accumulating in the bridge", async () => { + const totals: number[] = []; + for (const [contextTotalTokens, outputTokens] of [[10_000, 42], [10_300, 20]] as const) { + const frames = await collectSse(bridgeToResponsesSSE(replay([{ + type: "done", + usage: { inputTokens: 1, contextTotalTokens, outputTokens, estimated: true }, + }]), "kiro/claude-opus-5")); + const completed = frames.find(f => f.event === "response.completed")?.data.response as Record; + const usage = completed.usage as Record; + expect(usage.input_tokens).toBe(contextTotalTokens - outputTokens); + expect(usage.total_tokens).toBe(contextTotalTokens); + totals.push(usage.total_tokens); + } + expect(totals).toEqual([10_000, 10_300]); + }); + test("Anthropic cache read and write tokens pass through Responses usage without re-adding", async () => { const frames = await collectSse(bridgeToResponsesSSE(replay([ { @@ -153,6 +190,28 @@ describe("Responses bridge reasoning and usage parity", () => { }); }); + test("absolute context projection keeps cache details within derived input", async () => { + const frames = await collectSse(bridgeToResponsesSSE(replay([{ + type: "done", + usage: { + inputTokens: 200, + outputTokens: 10, + contextTotalTokens: 100, + cachedInputTokens: 150, + cacheReadInputTokens: 150, + cacheCreationInputTokens: 50, + }, + }]), "kiro/claude-opus-5")); + + const completed = frames.find(f => f.event === "response.completed")?.data.response as Record; + expect(completed.usage).toMatchObject({ + input_tokens: 90, + output_tokens: 10, + total_tokens: 100, + input_tokens_details: { cached_tokens: 90, cache_write_tokens: 0 }, + }); + }); + test("adapter heartbeat is non-visual in streaming and non-streaming responses", async () => { const events: AdapterEvent[] = [ { type: "heartbeat" }, diff --git a/tests/kiro-stream.test.ts b/tests/kiro-stream.test.ts index 5dcc624526..c2e9107cbc 100644 --- a/tests/kiro-stream.test.ts +++ b/tests/kiro-stream.test.ts @@ -312,6 +312,46 @@ describe("kiro adapter — parseStream", () => { }); }); + test("bounded fallback uses its rebuilt context estimate for the final absolute checkpoint", async () => { + const firstText = "p".repeat(7000); + const finalText = "f".repeat(3500); + globalThis.fetch = (async () => new Response(streamOf(eventFrame({ content: finalText })))) as typeof fetch; + const adapter = createKiroAdapter(provider); + const request = await adapter.buildRequest(parsedWith([{ role: "user", content: "do it" }], [bashTool])); + const initialContextEstimate = request.usageLog?.inputTokens ?? 0; + + const events = await collectAdapterEvents(adapter.parseStream(new Response(streamOf( + eventFrame({ content: firstText }), + )))); + const done = events.at(-1); + expect(done?.type).toBe("done"); + const usage = done?.type === "done" ? done.usage : undefined; + expect(usage?.outputTokens).toBe(estimateTokens(firstText, "claude-sonnet-4.5") + estimateTokens(finalText, "claude-sonnet-4.5")); + expect(usage?.contextTotalTokens).toBeGreaterThan( + initialContextEstimate + Math.max( + estimateTokens(firstText, "claude-sonnet-4.5"), + estimateTokens(finalText, "claude-sonnet-4.5"), + ), + ); + }); + + test("bounded fallback preserves definite growth after an upstream context checkpoint", async () => { + const finalText = "f".repeat(3500); + const finalOutputTokens = estimateTokens(finalText, "claude-sonnet-4.5"); + globalThis.fetch = (async () => new Response(streamOf(eventFrame({ content: finalText })))) as typeof fetch; + const adapter = createKiroAdapter(provider); + await adapter.buildRequest(parsedWith([{ role: "user", content: "do it" }], [bashTool])); + + const events = await collectAdapterEvents(adapter.parseStream(new Response(streamOf( + eventFrame({ content: "I am checking." }), + eventFrame({ contextUsagePercentage: 25 }), + )))); + const done = events.at(-1); + + expect(done?.type).toBe("done"); + if (done?.type === "done") expect(done.usage?.contextTotalTokens).toBe(50_000 + finalOutputTokens); + }); + test("keeps a private-completion fallback after reasoning-only output as the final answer", async () => { globalThis.fetch = (async () => new Response(streamOf(...completionFrames("Done.")))) as typeof fetch; const adapter = createKiroAdapter(provider); @@ -348,6 +388,24 @@ describe("kiro adapter — parseStream", () => { expect(events.at(-1)).toMatchObject({ type: "done", endTurn: true }); }); + test("reasoning-only fallback keeps absolute context above combined output", async () => { + const reasoning = "r".repeat(14_000); + const finalText = "f".repeat(14_000); + globalThis.fetch = (async () => new Response(streamOf(eventFrame({ content: finalText })))) as typeof fetch; + const adapter = createKiroAdapter(provider); + await adapter.buildRequest(parsedWith([{ role: "user", content: "solve" }], [bashTool])); + + const events = await collectAdapterEvents(adapter.parseStream(new Response(streamOf( + eventFrame({ content: `${reasoning}` }), + )))); + const done = events.at(-1); + + expect(done?.type).toBe("done"); + if (done?.type === "done") { + expect(done.usage?.contextTotalTokens).toBeGreaterThanOrEqual(done.usage?.outputTokens ?? 0); + } + }); + test("normal Responses cancellation aborts the adapter-owned fallback without another replay", async () => { const abort = new AbortController(); let fetches = 0; @@ -777,6 +835,7 @@ describe("kiro adapter — parseStream", () => { ); expect(done).toEqual({ inputTokens: 15, + contextTotalTokens: 204, cachedInputTokens: 3, cacheReadInputTokens: 3, cacheCreationInputTokens: 2, @@ -785,6 +844,26 @@ describe("kiro adapter — parseStream", () => { }); }); + test("authoritative turn usage floors a smaller payload context estimate", async () => { + const adapter = createKiroAdapter(provider); + await adapter.buildRequest(parsedWith([{ role: "user", content: "hi" }])); + const done = await doneUsage( + adapter, + eventFrame({ content: "answer" }), + eventFrame({ + tokenUsage: { + uncachedInputTokens: 500, + outputTokens: 4, + totalTokens: 504, + }, + }, "metadataEvent"), + ); + + expect(done.inputTokens).toBe(500); + expect(done.outputTokens).toBe(4); + expect(done.contextTotalTokens).toBe(504); + }); + test("invalid provider token usage is rejected instead of replacing estimates", async () => { const adapter = createKiroAdapter(provider); await adapter.buildRequest(parsedWith([{ role: "user", content: "hi" }])); @@ -823,7 +902,7 @@ describe("kiro adapter — parseStream", () => { expect((events[0] as { message: string }).message).toContain("Compact or reduce the history"); }); - test("Kiro contextUsagePercentage remains diagnostic and does not override totals", async () => { + test("Kiro contextUsagePercentage drives context pressure without overriding turn totals", async () => { const adapter = createKiroAdapter(provider); await adapter.buildRequest(parsedWith([{ role: "user", content: "x".repeat(700) }])); const done = await doneUsage( @@ -836,6 +915,15 @@ describe("kiro adapter — parseStream", () => { expect(done.outputTokens).toBe(100); expect(done.totalTokens).toBeUndefined(); expect(done.estimated).toBe(true); + expect(done.contextTotalTokens).toBe(50_000); + }); + + test("Kiro context percentage uses the native model window instead of a configured client cap", async () => { + const adapter = createKiroAdapter({ ...provider, contextWindow: 1_000_000 }); + await adapter.buildRequest(parsedWith([{ role: "user", content: "hi" }], undefined, "claude-sonnet-4.5")); + const done = await doneUsage(adapter, eventFrame({ content: "ok" }), eventFrame({ contextUsagePercentage: 25 })); + + expect(done.contextTotalTokens).toBe(50_000); }); test("Kiro auto ignores provider-level context window and falls back to heuristic totals", async () => { @@ -850,6 +938,29 @@ describe("kiro adapter — parseStream", () => { expect(done.inputTokens).toBe(200); expect(done.outputTokens).toBe(100); expect(done.totalTokens).toBeUndefined(); + expect(done.contextTotalTokens).toBe(300); + }); + + test("Kiro auto uses the concrete response model to decode context percentage", async () => { + const adapter = createKiroAdapter(provider); + await adapter.buildRequest(parsedWith([{ role: "user", content: "hi" }], undefined, "kiro-auto")); + const done = await doneUsage( + adapter, + eventFrame({ content: "ok", modelId: "claude-sonnet-4.5" }), + eventFrame({ contextUsagePercentage: 25 }), + ); + + expect(done.contextTotalTokens).toBe(50_000); + }); + + test("Kiro GPT routes use the Kiro token ratio without context percentage", async () => { + const adapter = createKiroAdapter(provider); + await adapter.buildRequest(parsedWith([{ role: "user", content: "x".repeat(3500) }], undefined, "gpt-5.6-sol")); + const done = await doneUsage(adapter, eventFrame({ content: "y".repeat(3500) })); + + expect(done.inputTokens).toBe(1000); + expect(done.outputTokens).toBe(1000); + expect(done.contextTotalTokens).toBe(2000); }); test("fresh payload includes history while usage counts only the current turn", async () => { @@ -875,6 +986,37 @@ describe("kiro adapter — parseStream", () => { expect(longBody.length).toBeGreaterThan(shortBody.length + 10_000); expect(longUsage.inputTokens).toBe(shortUsage.inputTokens); expect(longUsage.inputTokens).toBe(estimateTokens(latest, "claude-sonnet-4.5")); + expect(longUsage.contextTotalTokens).toBeGreaterThan(shortUsage.contextTotalTokens ?? 0); + }); + + test("context pressure follows the normalized Kiro payload while logs retain dropped reasoning", async () => { + const privateReasoning = "private-plan-".repeat(1000); + const adapter = createKiroAdapter(provider); + const request = await adapter.buildRequest(parsedWith([ + { role: "user", content: "old question" }, + { role: "assistant", content: [{ type: "thinking", thinking: privateReasoning }] }, + { role: "user", content: "latest question" }, + ])); + const usage = await doneUsage(adapter, eventFrame({ content: "ok" })); + + expect(request.body).not.toContain(privateReasoning); + expect(request.usageLog?.inputTokens).toBeGreaterThan((usage.contextTotalTokens ?? 0) + 1000); + expect(usage.contextTotalTokens).toBeLessThan(1000); + }); + + test("normalized images contribute conservative context tokens", async () => { + const onePixelPng = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="; + const adapter = createKiroAdapter(provider); + await adapter.buildRequest(parsedWith([{ + role: "user", + content: [ + { type: "text", text: "inspect" }, + { type: "image", imageUrl: `data:image/png;base64,${onePixelPng}` }, + ], + }])); + const usage = await doneUsage(adapter, eventFrame({ content: "ok" })); + + expect(usage.contextTotalTokens).toBeGreaterThanOrEqual(256 + usage.outputTokens); }); test("request log usage estimates the full Codex context while SSE usage stays current-turn", async () => { @@ -891,6 +1033,7 @@ describe("kiro adapter — parseStream", () => { expect(usage.inputTokens).toBe(estimateTokens(latest, "claude-sonnet-4.5")); expect(request.usageLog?.estimated).toBe(true); expect(request.usageLog?.inputTokens).toBeGreaterThan(usage.inputTokens + 4000); + expect(usage.contextTotalTokens).toBe((request.usageLog?.inputTokens ?? 0) + usage.outputTokens); }); test("resumed payload preserves the complete locally expanded history", async () => { diff --git a/tests/request-log.test.ts b/tests/request-log.test.ts index e7f86a3929..cf7689fc13 100644 --- a/tests/request-log.test.ts +++ b/tests/request-log.test.ts @@ -22,8 +22,14 @@ import { sealRequestAttemptIdentity, type RequestLogContext, } from "../src/server/request-log"; +import { bridgeToResponsesSSE } from "../src/bridge"; +import type { AdapterEvent } from "../src/types"; import type { PersistedUsageEntry } from "../src/usage/log"; +async function* replayAdapterEvents(events: AdapterEvent[]): AsyncGenerator { + for (const event of events) yield event; +} + function log(overrides: Partial): RequestLogEntry { return { requestId: "ocx-test", @@ -746,6 +752,36 @@ describe("request log metadata", () => { }); }); + test("deferred logging preserves a bridged Kiro absolute context checkpoint", async () => { + const entries: RequestLogEntry[] = []; + const body = bridgeToResponsesSSE(replayAdapterEvents([{ + type: "done", + usage: { + inputTokens: 58, + outputTokens: 100, + contextTotalTokens: 50_000, + estimated: true, + }, + }]), "kiro/claude-opus-5"); + const response = responseWithDeferredRequestLog( + new Response(body, { status: 200, headers: { "content-type": "text/event-stream" } }), + "ocx-test-kiro-context-checkpoint", + Date.now(), + { model: "kiro/claude-opus-5", provider: "kiro-p9d8524", usageLogInputTokens: 200 }, + entry => entries.push(entry), + ); + + const text = await response.text(); + expect(text).toContain('"input_tokens":49900'); + expect(text).toContain('"total_tokens":50000'); + expect(entries).toHaveLength(1); + expect(entries[0]).toMatchObject({ + usageStatus: "estimated", + totalTokens: 50_000, + usage: { inputTokens: 49_900, outputTokens: 100, totalTokens: 50_000, estimated: true }, + }); + }); + test("final logging shows numeric Kiro estimates even when SSE usage is absent", async () => { const entries: RequestLogEntry[] = []; const response = responseWithDeferredRequestLog(