diff --git a/scripts/test-layout/layout.json b/scripts/test-layout/layout.json index d9c536333c8..2ee6c63c2fc 100644 --- a/scripts/test-layout/layout.json +++ b/scripts/test-layout/layout.json @@ -1587,6 +1587,7 @@ "responses-translated-reset.test.ts": "responses", "responses-routed-web-search-fields.test.ts": "responses", "responses-self-named-namespace-scrub.test.ts": "responses", + "responses-combo-bridge-usage.test.ts": "responses", "responses-combo-exhausted-error.test.ts": "responses", "responses-combo-run-turn-classified-failure.test.ts": "responses", "responses-send-budget-counts.test.ts": "responses", @@ -1674,7 +1675,8 @@ "server-jev-combo-e2e.test.ts": "server", "combo-requested-effort.test.ts": "server", "server-kiro-completion-e2e.test.ts": "server", "server-kiro-device-login.test.ts": "server", "server-kiro-oauth-401-replay.test.ts": "server", - "server-kiro-refusal-e2e.test.ts": "server", "server-live-existing-call.test.ts": "server", + "server-kiro-refusal-e2e.test.ts": "server", "server-kiro-usage-log-e2e.test.ts": "server", + "server-live-existing-call.test.ts": "server", "server-live-frame-log.test.ts": "server", "server-live-realtime-fixtures.test.ts": "server", "server-live.test.ts": "server", diff --git a/src/adapters/kiro-events.ts b/src/adapters/kiro-events.ts index f49e73e2b5d..6f0db2ece29 100644 --- a/src/adapters/kiro-events.ts +++ b/src/adapters/kiro-events.ts @@ -109,6 +109,7 @@ function parseTokenUsage(eventType: string, value: unknown): OcxUsage | undefine totalTokens, ...(cacheRead !== undefined ? { cachedInputTokens: cacheRead, cacheReadInputTokens: cacheRead } : {}), ...(cacheWrite !== undefined ? { cacheCreationInputTokens: cacheWrite } : {}), + upstreamReported: true, }; } diff --git a/src/adapters/kiro/stream.ts b/src/adapters/kiro/stream.ts index dbd39e83764..fdeb027deef 100644 --- a/src/adapters/kiro/stream.ts +++ b/src/adapters/kiro/stream.ts @@ -187,6 +187,7 @@ function mergeKiroUsage( ...(sumOptional("providerCredits") !== undefined ? { providerCredits: sumOptional("providerCredits") } : {}), ...(sumOptional("reasoningOutputTokens") !== undefined ? { reasoningOutputTokens: sumOptional("reasoningOutputTokens") } : {}), ...(first.estimated || second.estimated ? { estimated: true } : {}), + ...(first.upstreamReported && second.upstreamReported ? { upstreamReported: true as const } : {}), }; } diff --git a/src/server/request-log.ts b/src/server/request-log.ts index 089e0a8a77a..9448cc0db8b 100644 --- a/src/server/request-log.ts +++ b/src/server/request-log.ts @@ -1623,7 +1623,10 @@ function finalizedUsage( const usageFallback = !finalUsage && estimate !== undefined ? { inputTokens: estimate, outputTokens: 0, estimated: true } : undefined; - const combinedInputTokens = finalUsage && estimate !== undefined + // The estimate is a floor under a count the adapter could only guess. A count the adapter read + // off the upstream's own usage frame stays exact: max() against a larger local estimate would + // replace it and mark it estimated. + const combinedInputTokens = finalUsage && !finalUsage.upstreamReported && estimate !== undefined ? Math.max(finalUsage.inputTokens, estimate) : undefined; const loggedUsage = finalUsage && combinedInputTokens !== undefined @@ -1905,6 +1908,7 @@ export function aggregateAttemptUsage( ...(reasoningOutputTokens !== undefined ? { reasoningOutputTokens } : {}), ...(providerCredits !== undefined ? { providerCredits } : {}), ...(status === "estimated" ? { estimated: true } : {}), + ...(usages.every(usage => usage.upstreamReported) ? { upstreamReported: true as const } : {}), }; return { usage: aggregate, status, totalTokens }; } diff --git a/src/server/responses/core-combo.ts b/src/server/responses/core-combo.ts index 64e8d7943dc..efb734a097b 100644 --- a/src/server/responses/core-combo.ts +++ b/src/server/responses/core-combo.ts @@ -931,6 +931,19 @@ export async function executeComboResponses( finishRequestAttempt(attempt, 499, Date.now() - started, childLog.usage); attemptRetained = true; }; + // The row is written from `logCtx`, which copies `childLog` only when the child commits. A + // bridge child binds its usage later, at the terminal frame, so the binding is forwarded here + // after that point. Left on `childLog`, the row would re-read the bridged wire instead of the + // adapter's usage, and that wire has no providerCredits and a context-derived input count. + // A binding that arrives before the commit waits for it, and a child that never commits + // forwards nothing: both flags belong to this target only. + let childCommitted = false; + let childUsageBound = false; + const forwardChildUsage = (): void => { + if (!childCommitted || !childUsageBound) return; + logCtx.usageFromBridge = true; + if (childLog.usage) logCtx.usage = childLog.usage; + }; const completedTarget = { provider: pick.target.provider, model: pick.target.model }; const writerGeneration = pick.writerGeneration; let consumedChildFailure: ConsumedComboFailure | undefined; @@ -1007,6 +1020,10 @@ export async function executeComboResponses( comboDispatchPermit: hopDecision?.allowed ? hopDecision.permit : undefined, comboReplaySnapshot, deferCodexResetDerivedCooldown, + onBridgeUsageBound: () => { + childUsageBound = true; + forwardChildUsage(); + }, // Attempt-relative TTFT is recorded HERE (not via childLog.firstOutputMs — a later // Object.assign(logCtx, childLog) would overwrite the request-relative value). onFirstOutput: () => { @@ -1093,6 +1110,8 @@ export async function executeComboResponses( options.onCodexAuthContextResolved?.(resolvedAuth); options.setTerminalOutcomeRecorder?.(terminalRecorder); callbackGate.commit(); + childCommitted = true; + forwardChildUsage(); return response; } diff --git a/src/server/responses/core-options.ts b/src/server/responses/core-options.ts index 3ee37545491..17a20911590 100644 --- a/src/server/responses/core-options.ts +++ b/src/server/responses/core-options.ts @@ -203,6 +203,13 @@ export interface HandleResponsesOptions { callerDirectAuth?: CallerDirectAuth | null; /** Internal recursion guard; callers outside this module must not set it. */ comboAttempt?: boolean; + /** + * Set by the combo loop for each child. A child logs into its own context, which the combo + * copies once when the child commits, and the bridge reports usage after that, at the terminal + * frame. The transport calls this when it binds that usage on the child's context, so the combo + * can carry the binding to the context that writes the row. + */ + onBridgeUsageBound?: () => void; /** Child-owned prepaid initial send; capacity remains charged until dispatch or unsent release. */ comboInitialSend?: { permit: SingleUseDispatchPermit; producerOwned?: boolean }; /** Exact externally booked combo hop, used for its child's spend preflight and send reports. */ diff --git a/src/server/responses/request-transport.ts b/src/server/responses/request-transport.ts index 397ffa2187d..e5c75f3184e 100644 --- a/src/server/responses/request-transport.ts +++ b/src/server/responses/request-transport.ts @@ -522,12 +522,11 @@ export async function prepareResponsesTransport( logCtx.usageFromBridge = true; if (usesApiKeyAccount(route.provider)) { logCtx.usage = logCtx.activeAttempt?.usage; - return; - } - if (usage) { + } else if (usage) { logCtx.usage = usage; if (logCtx.activeAttempt) logCtx.activeAttempt.usage = usage; } + options.onBridgeUsageBound?.(); }; const selectionIsCurrent = (binding: DispatchBinding | undefined): boolean => { if (!anthropicAdmissionIsCurrent()) return false; diff --git a/src/types/request.ts b/src/types/request.ts index edc20eaae0d..4b247feffec 100644 --- a/src/types/request.ts +++ b/src/types/request.ts @@ -562,6 +562,13 @@ export interface OcxUsage { cacheCreation1hInputTokens?: number; reasoningOutputTokens?: number; estimated?: boolean; + /** + * The adapter read these counts off the upstream's own usage frame. Only an adapter whose + * provider is estimated by default (usageForFinalLog) sets it, and a merge keeps it only when + * every leg had it. A figure taken off a bridged wire or guessed never carries it, so the + * provider default still marks those as estimated. + */ + upstreamReported?: true; /** * The raw upstream usage object for Responses-shaped upstreams (openai/codex#41980 parity): * codex-rs preserves the complete `response.usage` object through its own pipeline, so fields diff --git a/src/usage/log.ts b/src/usage/log.ts index e9e61d2da0c..13ecd2b38f8 100644 --- a/src/usage/log.ts +++ b/src/usage/log.ts @@ -536,6 +536,10 @@ export function usageTotalTokens(usage: OcxUsage | undefined): number | undefine * Providers whose adapters can only estimate usage (no authoritative per-turn frame). * Callers should pass the route ADAPTER when available; the name-prefix match is a * fallback for paths that only know the configured provider name (e.g. "cursor-mykey"). + * + * Kiro is on the list because most of what reaches the log is a guess: a figure taken off the + * bridged Responses wire derives its input from the context checkpoint. The adapter opts a + * measured figure out with `upstreamReported`, set where it reads Kiro's own tokenUsage frame. */ function isEstimatedUsageProvider(providerOrAdapter: string): boolean { return providerOrAdapter === "kiro" || providerOrAdapter.startsWith("kiro-") @@ -556,7 +560,9 @@ export function usageForFinalLog( ): OcxUsage | undefined { if (!usage) return undefined; if (locallyAnswered) return usage; - if (usage.estimated || isEstimatedUsageProvider(provider)) return { ...usage, estimated: true }; + if (usage.estimated || (isEstimatedUsageProvider(provider) && !usage.upstreamReported)) { + return { ...usage, estimated: true }; + } return usage; } diff --git a/structure/providers/kiro.md b/structure/providers/kiro.md index e0e26d4ff12..eae07ee3c1f 100644 --- a/structure/providers/kiro.md +++ b/structure/providers/kiro.md @@ -207,8 +207,14 @@ finite `contextUsagePercentage` (and a `tokenUsage` block) there as a fallback, from `metadataEvent` is legitimate rather than impossible. Both feed the same field, and any positive value overwrites an earlier one. -Spend arrives in `meteringEvent` as **credits, not tokens**. No captured response carried -`tokenUsage` on any event, which is why Kiro token usage stays estimated. The parser preserves +Spend arrives in `meteringEvent` as **credits, not tokens**. Kiro token usage stays estimated +unless a `metadataEvent` carries a `tokenUsage` block; no captured response did. That block is the +only frame that yields a measured count: the parser marks the usage `upstreamReported`, +`mergeKiroUsage` keeps the mark only when both legs of a two-leg turn have it, and +`usageForFinalLog` (`src/usage/log.ts`) then skips the provider-wide estimated mark while +`finalizedUsage` (`src/server/request-log.ts`) skips the local input floor. A turn without +`tokenUsage`, a leg without it, and a count read off the bridged Responses wire all keep the +estimated mark. The parser preserves `meteringEvent` unit/usage (`amount` is an alias) and optional `unitPlural`; credit readings populate `OcxUsage.providerCredits` independently of token metadata. The latest reading within a response is a snapshot; separate completion-fallback responses add their credits. Missing metering stays diff --git a/structure/transports/responses-spend.md b/structure/transports/responses-spend.md index 83b6912bca8..fd0cbc511aa 100644 --- a/structure/transports/responses-spend.md +++ b/structure/transports/responses-spend.md @@ -182,6 +182,15 @@ what those sends may cost, and it is the only bound here that survives a restart caller is `request-spend.ts`, installed on the execution budget at genuine ingress in `core.ts` and parked on the log context so `addFinalRequestLog` can settle it. Native Chat installs the same tracker before its independent physical-send ladder and charges it immediately before each dispatch, so that fast path cannot bypass root, identity or provider-pool ceilings. +The tracker settles with `logCtx.usage`, and so does the row. A combo copies its child's log +context into that one only when the child commits, while the bridge binds the adapter usage later, +at the terminal frame. `executeComboResponses` therefore passes each child `onBridgeUsageBound` +(`core-options.ts`), which `bindKeyUsageFromBridge` calls after it binds. After the child commits, +the combo sets `usageFromBridge` and `usage` on its own context; a binding that came earlier waits +for the commit, and a child that fails or never commits forwards nothing. Without that, the row +and the settlement would read the usage back off the bridged Responses wire, which carries no +`providerCredits` and derives its input count from the context checkpoint. + Unconfigured requests preserve the legacy counter observer: reservations increment the send counter, refunds decrement it, and named reports reconcile prepaid sends. Enforced requests use the seed and physical-start protocol below across Responses, native Messages and native Chat. diff --git a/tests/fixtures/test-layout-expected.json b/tests/fixtures/test-layout-expected.json index 28f9d9dd49f..5639600afb8 100644 --- a/tests/fixtures/test-layout-expected.json +++ b/tests/fixtures/test-layout-expected.json @@ -123,6 +123,7 @@ "codex-shim-overlay-readiness.test.ts": "codex-integration", "codex-shim-overlay-candidates.test.ts": "codex-integration", "droid-normalized-reasoning-defaults.test.ts": "clients", + "responses-combo-bridge-usage.test.ts": "responses", "responses-combo-exhausted-error.test.ts": "responses", "responses-combo-run-turn-classified-failure.test.ts": "responses", "cli-config-default-show.test.ts": "cli", @@ -1155,7 +1156,8 @@ "server-jev-combo-e2e.test.ts": "server", "server-kiro-completion-e2e.test.ts": "server", "combo-requested-effort.test.ts": "server", "server-kiro-device-login.test.ts": "server", "server-kiro-oauth-401-replay.test.ts": "server", - "server-kiro-refusal-e2e.test.ts": "server", "server-live-existing-call.test.ts": "server", + "server-kiro-refusal-e2e.test.ts": "server", "server-kiro-usage-log-e2e.test.ts": "server", + "server-live-existing-call.test.ts": "server", "server-live-frame-log.test.ts": "server", "server-live-realtime-fixtures.test.ts": "server", "server-live.test.ts": "server", "server-loopback-host-gate.test.ts": "server", "server-management-auth.test.ts": "server", diff --git a/tests/providers/kiro/kiro-stream.test.ts b/tests/providers/kiro/kiro-stream.test.ts index c09fac50430..2e2cc6850a4 100644 --- a/tests/providers/kiro/kiro-stream.test.ts +++ b/tests/providers/kiro/kiro-stream.test.ts @@ -1538,6 +1538,7 @@ describe("kiro adapter — parseStream", () => { cacheCreationInputTokens: 2, outputTokens: 4, totalTokens: 19, + upstreamReported: true, }); // The whole-payload checkpoint must exceed this turn's own total, which is the point of // reporting it separately. diff --git a/tests/responses/responses-combo-bridge-usage.test.ts b/tests/responses/responses-combo-bridge-usage.test.ts new file mode 100644 index 00000000000..a32e5af3b8b --- /dev/null +++ b/tests/responses/responses-combo-bridge-usage.test.ts @@ -0,0 +1,164 @@ +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { clearComboSelectionState, clearComboTargetCooldowns } from "../../src/combos"; +import { clearKeyCooldowns } from "../../src/providers/key-failover"; +import { createTranslatorBudget } from "../../src/lib/translator-budget"; +import { handleComboResponses } from "../../src/server/responses/core"; +import { executeComboResponses } from "../../src/server/responses/core-combo"; +import type { HandleResponsesOptions } from "../../src/server/responses/core-options"; +import type { RequestLogContext } from "../../src/server/request-log"; +import type { OcxConfig, OcxUsage } from "../../src/types"; +import { acquireOwnedSpendHome } from "../helpers/owned-spend-home"; + +/** + * A combo child logs into its own context and the combo copies it into the request context when + * the child commits. The bridge binds the adapter usage on the child, possibly before that commit + * and normally after it. Only the child that commits may hand its usage to the request context. + */ +let releaseSpendHome: (() => void) | undefined; +const reset = (): void => { + clearComboSelectionState(); + clearComboTargetCooldowns(); + clearKeyCooldowns(); +}; +beforeEach(() => { + reset(); + releaseSpendHome = acquireOwnedSpendHome(); +}); +afterEach(() => { + releaseSpendHome?.(); + releaseSpendHome = undefined; + reset(); +}); + +const provider = (name: string) => ({ + adapter: "openai-chat", + baseUrl: `https://${name}.example/v1`, + authMode: "key", + apiKey: `sk-${name}`, + models: [`model-${name}`], +}); +const config = { + defaultProvider: "primary", + providers: { primary: provider("primary"), fallback: provider("fallback") }, + combos: { + fan: { + strategy: "failover", + targets: [ + { provider: "primary", model: "model-primary" }, + { provider: "fallback", model: "model-fallback" }, + ], + }, + }, +} as unknown as OcxConfig; + +const usageOf = (credits: number): OcxUsage => ({ + inputTokens: 18, outputTokens: 2, providerCredits: credits, upstreamReported: true, +}); + +interface Child { + provider: string; + /** The request context the combo writes the row from. */ + parent: RequestLogContext; + log: RequestLogContext; + options: HandleResponsesOptions; +} + +/** Stands in for the bridge: binds usage on the child's own context, as bindKeyUsageFromBridge does. */ +function bind(child: Child, usage: OcxUsage): void { + child.log.usageFromBridge = true; + child.log.usage = usage; + if (child.log.activeAttempt) child.log.activeAttempt.usage = usage; + child.options.onBridgeUsageBound?.(); +} + +async function runCombo( + answer: (child: Child) => Response | Promise, + options: Partial = {}, +): Promise<{ response: Promise; logCtx: RequestLogContext; children: Child[] }> { + const children: Child[] = []; + const logCtx: RequestLogContext = { model: "", provider: "" }; + const req = new Request("http://localhost/v1/responses", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ model: "combo/fan", stream: false, input: "hello" }), + }); + const response = executeComboResponses(req, await req.clone().json(), "fan", config, logCtx, { + translatorBudget: createTranslatorBudget(), + ...options, + }, { + handleComboResponses, + handleResponses: async (_req, _config, childLog, childOptions) => { + const child: Child = { provider: childLog.provider, parent: logCtx, log: childLog, options: childOptions ?? {} }; + children.push(child); + return answer(child); + }, + }); + return { response, logCtx, children }; +} + +const ok = (): Response => Response.json({ status: "completed" }); +const refused = (): Response => Response.json( + { error: { message: "weekly usage limit reached", type: "rate_limit_error" } }, + { status: 429 }, +); + +describe("combo hands the committed child's bridge usage to the request context", () => { + test("usage bound after the commit reaches the request context", async () => { + const run = await runCombo(() => ok()); + expect((await run.response).status).toBe(200); + expect(run.logCtx.usageFromBridge).toBeUndefined(); + + const usage = usageOf(0.04); + bind(run.children[0]!, usage); + expect(run.logCtx.usageFromBridge).toBe(true); + expect(run.logCtx.usage).toBe(usage); + expect(run.logCtx.activeAttempt?.usage).toBe(usage); + }); + + test("usage bound before the commit is still handed over once the child commits", async () => { + const usage = usageOf(0.05); + const run = await runCombo(child => { + bind(child, usage); + // The commit has not happened yet, so nothing may reach the request context from here. + expect(child.parent.usage).toBeUndefined(); + return ok(); + }); + expect((await run.response).status).toBe(200); + expect(run.logCtx.usageFromBridge).toBe(true); + expect(run.logCtx.usage).toBe(usage); + }); + + test("a refused first target never lends its usage; only the committed target does", async () => { + const run = await runCombo(child => child.provider === "primary" ? refused() : ok()); + expect((await run.response).status).toBe(200); + expect(run.children.map(child => child.provider)).toEqual(["primary", "fallback"]); + const [primary, fallback] = run.children as [Child, Child]; + + // The refused target's stream can still report after the fallback has committed. + bind(primary, usageOf(9.99)); + expect(run.logCtx.usage).toBeUndefined(); + expect(run.logCtx.usageFromBridge).toBeUndefined(); + + const usage = usageOf(0.04); + bind(fallback, usage); + expect(run.logCtx.usage).toBe(usage); + expect(run.logCtx.usage?.providerCredits).toBe(0.04); + expect(run.logCtx.usageFromBridge).toBe(true); + }); + + test("a commit that throws forwards nothing for that child", async () => { + // The child reported its terminal before the commit, so the commit publishes it to the + // caller's callback. A callback that throws aborts the commit. + const run = await runCombo(child => { + child.options.onNativePassthroughTerminal?.("completed"); + return ok(); + }, { + onNativePassthroughTerminal: () => { throw new Error("terminal sink failed"); }, + }); + await expect(run.response).rejects.toThrow("terminal sink failed"); + + bind(run.children[0]!, usageOf(0.07)); + expect(run.logCtx.usage).toBeUndefined(); + expect(run.logCtx.usageFromBridge).toBeUndefined(); + }); +}); diff --git a/tests/server/server-kiro-usage-log-e2e.test.ts b/tests/server/server-kiro-usage-log-e2e.test.ts new file mode 100644 index 00000000000..4dc8a1eee1c --- /dev/null +++ b/tests/server/server-kiro-usage-log-e2e.test.ts @@ -0,0 +1,283 @@ +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { mkdtempSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { KIRO_COMPLETION_TOOL_NAME } from "../../src/adapters/kiro-constants"; +import { saveConfig } from "../../src/config"; +import { encodeMessage } from "../../src/lib/eventstream-decoder"; +import { saveCredential } from "../../src/oauth/store"; +import { startServer } from "../../src/server"; +import { clearRequestLogsForTests, getRequestLogEntries, type RequestLogEntry } from "../../src/server/request-log"; +import { readUsageEntries, resetUsageReadCacheForTests } from "../../src/usage/log"; +import type { OcxConfig } from "../../src/types"; +import { installIsolatedCodexHome, type IsolatedCodexHome } from "../helpers/isolated-codex-home"; +import { removeTreeWithRetry } from "../helpers/remove-tree"; + +// What Kiro itself reports for one turn, and what the usage ledger must record for it whichever +// client wire asked: reported counts, credits, no estimate. The prompt below is long enough that +// the adapter's own input estimate is far above these counts, so a row that fell back to the +// estimate cannot pass for a reported one. +const CREDITS = 0.04582331509121062; +const TOKENS = { uncachedInputTokens: 1000, cacheReadInputTokens: 400, cacheWriteInputTokens: 100, outputTokens: 50, totalTokens: 1550 }; +const REPORTED = { + inputTokens: 1500, + outputTokens: 50, + cachedInputTokens: 400, + cacheReadInputTokens: 400, + cacheCreationInputTokens: 100, + providerCredits: CREDITS, +}; +const LONG_PROMPT = "Summarise the repository layout. ".repeat(3000); + +const enc = new TextEncoder(); +const CHAT_ENDPOINT = "https://runtime.us-east-1.kiro.dev/"; +const originalFetch = globalThis.fetch; + +let testDir = ""; +let previousOpenCodexHome: string | undefined; +let previousRegion: string | undefined; +let isolatedCodexHome: IsolatedCodexHome | null = null; + +beforeEach(() => { + previousOpenCodexHome = process.env.OPENCODEX_HOME; + previousRegion = process.env.KIRO_REGION; + isolatedCodexHome = installIsolatedCodexHome("ocx-kiro-usage-codex-"); + testDir = mkdtempSync(join(tmpdir(), "ocx-kiro-usage-")); + process.env.OPENCODEX_HOME = testDir; + process.env.KIRO_REGION = "us-east-1"; + clearRequestLogsForTests(); +}); + +afterEach(() => { + globalThis.fetch = originalFetch; + if (previousOpenCodexHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = previousOpenCodexHome; + if (previousRegion === undefined) delete process.env.KIRO_REGION; + else process.env.KIRO_REGION = previousRegion; + isolatedCodexHome?.restore(); + isolatedCodexHome = null; + removeTreeWithRetry(testDir); +}); + +function eventFrame(eventType: string, payload: Record): Uint8Array { + return encodeMessage({ ":message-type": "event", ":event-type": eventType }, enc.encode(JSON.stringify(payload))); +} + +function kiroTurn(withTokenUsage: boolean): Uint8Array[] { + return [ + eventFrame("assistantResponseEvent", { content: "The layout is flat." }), + ...kiroMeters(withTokenUsage), + ]; +} + +function kiroMeters(withTokenUsage: boolean): Uint8Array[] { + return [ + eventFrame("meteringEvent", { unit: "credit", unitPlural: "credits", usage: CREDITS }), + ...(withTokenUsage ? [eventFrame("metadataEvent", { tokenUsage: TOKENS })] : []), + ]; +} + +/** With tools declared, Kiro answers in two legs: progress prose, then the private completion call. */ +function kiroTwoLegTurn(secondLegTokenUsage: boolean): Uint8Array[][] { + const input = JSON.stringify({ answer: "The layout is flat." }); + const toolUse = (extra: Record) => + eventFrame("toolUseEvent", { name: KIRO_COMPLETION_TOOL_NAME, toolUseId: "completion-1", ...extra }); + return [ + [eventFrame("assistantResponseEvent", { content: "Checking the layout." }), ...kiroMeters(true)], + [toolUse({}), toolUse({ input }), toolUse({ stop: true }), ...kiroMeters(secondLegTokenUsage)], + ]; +} + +function streamOf(frames: Uint8Array[]): ReadableStream { + let index = 0; + return new ReadableStream({ + pull(controller) { + if (index < frames.length) controller.enqueue(frames[index++]); + else controller.close(); + }, + }); +} + +type Route = "direct" | "combo"; + +async function startKiro(route: Route, turns: Uint8Array[][]): Promise> { + await saveCredential("kiro", { + access: "synthetic-access", + refresh: "synthetic-refresh", + expires: Date.now() + 3_600_000, + accountId: "kiro-usage-account", + source: "oauth", + }); + saveConfig({ + port: 0, + hostname: "127.0.0.1", + defaultProvider: "kiro", + providers: { kiro: { adapter: "kiro", baseUrl: CHAT_ENDPOINT.slice(0, -1), authMode: "oauth", models: ["claude-sonnet-4.5"] } }, + ...(route === "combo" + ? { combos: { claude: { strategy: "failover", targets: [{ provider: "kiro", model: "claude-sonnet-4.5" }] } } } + : {}), + } as unknown as OcxConfig); + globalThis.fetch = (async (input, init) => { + const url = input instanceof Request ? input.url : String(input); + if (url === CHAT_ENDPOINT) { + return new Response(streamOf(turns.shift() ?? []), { + headers: { "content-type": "application/vnd.amazon.eventstream" }, + }); + } + return originalFetch(input, init); + }) as typeof fetch; + return startServer(0); +} + +const MODEL: Record = { direct: "kiro/claude-sonnet-4.5", combo: "combo/claude" }; + +interface Entry { + label: string; + path: string; + body: (model: string) => Record; +} + +const ENTRIES: Entry[] = [ + { + label: "/v1/responses stream", + path: "/v1/responses", + body: model => ({ model, input: LONG_PROMPT, stream: true }), + }, + { + label: "/v1/messages stream", + path: "/v1/messages", + body: model => ({ model, max_tokens: 256, stream: true, messages: [{ role: "user", content: LONG_PROMPT }] }), + }, + { + label: "/v1/messages non-stream", + path: "/v1/messages", + body: model => ({ model, max_tokens: 256, stream: false, messages: [{ role: "user", content: LONG_PROMPT }] }), + }, + { + label: "/v1/chat/completions non-stream", + path: "/v1/chat/completions", + body: model => ({ model, stream: false, messages: [{ role: "user", content: LONG_PROMPT }] }), + }, +]; + +async function ask(proxy: ReturnType, entry: Entry, route: Route): Promise { + const response = await originalFetch(new URL(entry.path, proxy.url), { + method: "POST", + headers: { "content-type": "application/json", "anthropic-version": "2023-06-01" }, + body: JSON.stringify(entry.body(MODEL[route])), + }); + expect(response.status).toBe(200); + return response.text(); +} + +async function finalRow(): Promise<{ memory: RequestLogEntry; persisted: ReturnType[number] }> { + for (let attempt = 0; attempt < 100 && getRequestLogEntries().length === 0; attempt++) { + await new Promise(resolve => setTimeout(resolve, 20)); + } + const memory = getRequestLogEntries().at(-1); + if (!memory) throw new Error("no request log row was written"); + resetUsageReadCacheForTests(); + const persisted = readUsageEntries().find(row => row.requestId === memory.requestId); + if (!persisted) throw new Error("the row never reached usage.jsonl"); + return { memory, persisted }; +} + +describe("Kiro usage reaches the ledger as Kiro reported it, on every inbound wire", () => { + for (const route of ["direct", "combo"] as const) { + for (const entry of ENTRIES) { + test(`${entry.label} via ${route} records reported counts and credits`, async () => { + const proxy = await startKiro(route, [kiroTurn(true)]); + try { + await ask(proxy, entry, route); + const { memory, persisted } = await finalRow(); + for (const row of [memory, persisted]) { + expect(row.usageStatus).toBe("reported"); + expect(row.usage).toMatchObject(REPORTED); + expect(row.usage?.estimated).toBeUndefined(); + expect(row.totalTokens).toBe(1550); + } + } finally { + await proxy.stop(true); + } + }); + + test(`${entry.label} via ${route} stays estimated when Kiro sent no token usage`, async () => { + const proxy = await startKiro(route, [kiroTurn(false)]); + try { + await ask(proxy, entry, route); + const { memory, persisted } = await finalRow(); + for (const row of [memory, persisted]) { + expect(row.usageStatus).toBe("estimated"); + expect(row.usage?.estimated).toBe(true); + // The credits are metered separately from tokens and survive the missing token frame. + expect(row.usage?.providerCredits).toBe(CREDITS); + // The local estimate is a floor here, so it is the count that was written down. + expect(row.usage?.inputTokens).toBeGreaterThan(REPORTED.inputTokens); + } + } finally { + await proxy.stop(true); + } + }); + } + } + + test("the Claude wire still carries Kiro's output and cache counts", async () => { + const proxy = await startKiro("combo", [kiroTurn(true)]); + try { + const messages = ENTRIES.find(entry => entry.label === "/v1/messages non-stream")!; + const message = JSON.parse(await ask(proxy, messages, "combo")) as { usage: Record }; + expect(message.usage).toMatchObject({ + output_tokens: 50, + cache_read_input_tokens: 400, + cache_creation_input_tokens: 100, + }); + } finally { + await proxy.stop(true); + } + }); + + // Claude Code always declares tools, so its text answers take the two-leg path. + for (const route of ["direct", "combo"] as const) { + test(`a two-leg /v1/messages turn via ${route} sums both legs and stays reported`, async () => { + const proxy = await startKiro(route, kiroTwoLegTurn(true)); + try { + const tools = [{ name: "bash", description: "Run a command", input_schema: { type: "object" } }]; + await ask(proxy, { label: "", path: "/v1/messages", body: model => ({ + model, max_tokens: 256, stream: true, tools, messages: [{ role: "user", content: LONG_PROMPT }], + }) }, route); + const { memory, persisted } = await finalRow(); + for (const row of [memory, persisted]) { + expect(row.usageStatus).toBe("reported"); + expect(row.usage).toMatchObject({ + inputTokens: 2 * REPORTED.inputTokens, + outputTokens: 2 * REPORTED.outputTokens, + cacheReadInputTokens: 2 * REPORTED.cacheReadInputTokens, + cacheCreationInputTokens: 2 * REPORTED.cacheCreationInputTokens, + providerCredits: 2 * CREDITS, + }); + expect(row.usage?.estimated).toBeUndefined(); + } + } finally { + await proxy.stop(true); + } + }); + + test(`a two-leg turn via ${route} is estimated when one leg lacks token usage`, async () => { + const proxy = await startKiro(route, kiroTwoLegTurn(false)); + try { + const tools = [{ name: "bash", description: "Run a command", input_schema: { type: "object" } }]; + await ask(proxy, { label: "", path: "/v1/messages", body: model => ({ + model, max_tokens: 256, stream: true, tools, messages: [{ role: "user", content: LONG_PROMPT }], + }) }, route); + const { memory, persisted } = await finalRow(); + for (const row of [memory, persisted]) { + expect(row.usageStatus).toBe("estimated"); + expect(row.usage?.estimated).toBe(true); + expect(row.usage?.providerCredits).toBe(2 * CREDITS); + } + } finally { + await proxy.stop(true); + } + }); + } +}); diff --git a/tests/usage/request-log-estimate-cap.test.ts b/tests/usage/request-log-estimate-cap.test.ts index c8e89c222cc..80faa5ad3cd 100644 --- a/tests/usage/request-log-estimate-cap.test.ts +++ b/tests/usage/request-log-estimate-cap.test.ts @@ -1,5 +1,6 @@ import { describe, expect, test } from "bun:test"; import { + aggregateAttemptUsage, beginRequestAttempt, finishRequestAttempt, noteAttemptSend, @@ -141,6 +142,39 @@ describe("request log token-estimate context-window cap (codex-router PR #140)", expect(attempt.usage?.inputTokens).toBe(500_000); }); + test("a count the adapter read off the upstream usage frame is not floored and not estimated", () => { + const attempt = beginRequestAttempt(1, "kiro", "deepseek-3.2", "kiro"); + noteAttemptSend(attempt, 100_000); + finishRequestAttempt(attempt, 200, 10, { inputTokens: 50_000, outputTokens: 5, upstreamReported: true }); + expect(attempt.usage?.inputTokens).toBe(50_000); + expect(attempt.usage?.estimated).toBeUndefined(); + expect(attempt.usageStatus).toBe("reported"); + + // Without the mark the same figure is still a floor-merged estimate. + const unmarked = beginRequestAttempt(1, "kiro", "deepseek-3.2", "kiro"); + noteAttemptSend(unmarked, 100_000); + finishRequestAttempt(unmarked, 200, 10, { inputTokens: 50_000, outputTokens: 5 }); + expect(unmarked.usage?.inputTokens).toBe(100_000); + expect(unmarked.usageStatus).toBe("estimated"); + }); + + test("an aggregate carries the upstream-reported mark only when every attempt has usage with it", () => { + const marked = () => { + const attempt = beginRequestAttempt(1, "kiro", "deepseek-3.2", "kiro"); + finishRequestAttempt(attempt, 200, 10, { inputTokens: 10, outputTokens: 2, upstreamReported: true }); + return attempt; + }; + const unmarked = beginRequestAttempt(2, "kiro", "deepseek-3.2", "kiro"); + finishRequestAttempt(unmarked, 200, 10, { inputTokens: 10, outputTokens: 2 }); + + expect(aggregateAttemptUsage([marked(), marked()]).usage?.upstreamReported).toBe(true); + expect(aggregateAttemptUsage([marked(), unmarked]).usage?.upstreamReported).toBeUndefined(); + // No usage at all is not "every attempt reported": nothing is aggregated, so nothing is marked. + expect(aggregateAttemptUsage([])).toEqual({ status: "unreported" }); + const silent = beginRequestAttempt(1, "kiro", "deepseek-3.2", "kiro"); + expect(aggregateAttemptUsage([silent]).usage).toBeUndefined(); + }); + test("kiro auto has no fixed window and is never guessed", () => { const attempt = beginRequestAttempt(1, "kiro", "kiro-auto", "kiro"); noteAttemptSend(attempt, 500_000);