From 5f6364baf484eb2f6979293069af397ec2135725 Mon Sep 17 00:00:00 2001 From: Hannes Rudolph Date: Wed, 11 Feb 2026 18:30:28 -0700 Subject: [PATCH] fix: Bedrock inputTokens now emits total tokens per ApiStreamUsageChunk contract Stop overriding usage.inputTokens with noCacheTokens before passing to normalizeProviderUsage. The normalizer already extracts non-cached tokens from inputTokenDetails.noCacheTokens via the anthropic base profile. Also fix calculateCost to receive inputTokensTotal (it derives uncached internally), preventing a double-subtraction of cache tokens. Adds dedicated Bedrock usage-metrics tests covering cache-detail scenarios. --- .../__tests__/bedrock-usage-metrics.spec.ts | 187 ++++++++++++++++++ src/api/providers/__tests__/bedrock.spec.ts | 7 +- src/api/providers/bedrock.ts | 19 +- 3 files changed, 197 insertions(+), 16 deletions(-) create mode 100644 src/api/providers/__tests__/bedrock-usage-metrics.spec.ts diff --git a/src/api/providers/__tests__/bedrock-usage-metrics.spec.ts b/src/api/providers/__tests__/bedrock-usage-metrics.spec.ts new file mode 100644 index 0000000000..0e22e8571e --- /dev/null +++ b/src/api/providers/__tests__/bedrock-usage-metrics.spec.ts @@ -0,0 +1,187 @@ +// npx vitest run src/api/providers/__tests__/bedrock-usage-metrics.spec.ts + +vi.mock("@roo-code/telemetry", () => ({ + TelemetryService: { + instance: { + captureException: vi.fn(), + }, + }, +})) + +vi.mock("@aws-sdk/credential-providers", () => ({ + fromIni: vi.fn().mockReturnValue({ + accessKeyId: "profile-access-key", + secretAccessKey: "profile-secret-key", + }), +})) + +const { mockStreamText, mockGenerateText } = vi.hoisted(() => ({ + mockStreamText: vi.fn(), + mockGenerateText: vi.fn(), +})) + +vi.mock("ai", async (importOriginal) => { + const actual = await importOriginal() + return { + ...actual, + streamText: mockStreamText, + generateText: mockGenerateText, + } +}) + +vi.mock("@ai-sdk/amazon-bedrock", () => ({ + createAmazonBedrock: vi.fn(() => vi.fn(() => ({ modelId: "test", provider: "bedrock" }))), +})) + +import { AwsBedrockHandler } from "../bedrock" + +describe("AwsBedrockHandler usage metrics", () => { + beforeEach(() => { + vi.clearAllMocks() + }) + + function createHandler() { + return new AwsBedrockHandler({ + apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0", + awsAccessKey: "test-access-key", + awsSecretKey: "test-secret-key", + awsRegion: "us-east-1", + }) + } + + async function collectChunks(stream: AsyncGenerator) { + const chunks: any[] = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + return chunks + } + + it("emits inputTokens as total (not noCacheTokens) when cache details are present", async () => { + async function* mockFullStream() { + yield { type: "text-delta", text: "Hello" } + } + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream(), + usage: Promise.resolve({ + inputTokens: 13_071, + outputTokens: 93, + inputTokenDetails: { + noCacheTokens: 10, + cacheWriteTokens: 489, + cacheReadTokens: 12_572, + }, + }), + providerMetadata: Promise.resolve({}), + }) + + const handler = createHandler() + const chunks = await collectChunks( + handler.createMessage("You are a helpful assistant.", [{ role: "user", content: "Hello" }]), + ) + + const usageChunk = chunks.find((c) => c.type === "usage") + expect(usageChunk).toBeDefined() + + // inputTokens must be the total (13,071), not noCacheTokens (10). + // This aligns with the ApiStreamUsageChunk contract: + // "inputTokens: Total input tokens (cached + non-cached)." + expect(usageChunk.inputTokens).toBe(13_071) + expect(usageChunk.nonCachedInputTokens).toBe(10) + expect(usageChunk.outputTokens).toBe(93) + expect(usageChunk.cacheWriteTokens).toBe(489) + expect(usageChunk.cacheReadTokens).toBe(12_572) + }) + + it("emits inputTokens as total when cache metrics come from providerMetadata.bedrock.usage", async () => { + async function* mockFullStream() { + yield { type: "text-delta", text: "Hello" } + } + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream(), + usage: Promise.resolve({ + inputTokens: 500, + outputTokens: 50, + }), + providerMetadata: Promise.resolve({ + bedrock: { + usage: { + cacheReadInputTokens: 300, + cacheWriteInputTokens: 100, + }, + }, + }), + }) + + const handler = createHandler() + const chunks = await collectChunks( + handler.createMessage("You are a helpful assistant.", [{ role: "user", content: "Hello" }]), + ) + + const usageChunk = chunks.find((c) => c.type === "usage") + expect(usageChunk).toBeDefined() + expect(usageChunk.inputTokens).toBe(500) + expect(usageChunk.outputTokens).toBe(50) + expect(usageChunk.cacheReadTokens).toBe(300) + expect(usageChunk.cacheWriteTokens).toBe(100) + // Non-cached should be derived: 500 - 300 - 100 = 100 + expect(usageChunk.nonCachedInputTokens).toBe(100) + }) + + it("handles basic usage without cache details", async () => { + async function* mockFullStream() { + yield { type: "text-delta", text: "Hello" } + } + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream(), + usage: Promise.resolve({ + inputTokens: 100, + outputTokens: 50, + }), + providerMetadata: Promise.resolve({}), + }) + + const handler = createHandler() + const chunks = await collectChunks( + handler.createMessage("You are a helpful assistant.", [{ role: "user", content: "Hello" }]), + ) + + const usageChunk = chunks.find((c) => c.type === "usage") + expect(usageChunk).toBeDefined() + expect(usageChunk.inputTokens).toBe(100) + expect(usageChunk.outputTokens).toBe(50) + }) + + it("calculates cost correctly with cache tokens using total inputTokens", async () => { + async function* mockFullStream() { + yield { type: "text-delta", text: "Hello" } + } + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream(), + usage: Promise.resolve({ + inputTokens: 13_071, + outputTokens: 93, + inputTokenDetails: { + noCacheTokens: 10, + cacheWriteTokens: 489, + cacheReadTokens: 12_572, + }, + }), + providerMetadata: Promise.resolve({}), + }) + + const handler = createHandler() + const chunks = await collectChunks( + handler.createMessage("You are a helpful assistant.", [{ role: "user", content: "Hello" }]), + ) + + const usageChunk = chunks.find((c) => c.type === "usage") + expect(usageChunk).toBeDefined() + // totalCost should be > 0 since we have token usage + expect(usageChunk.totalCost).toBeGreaterThan(0) + }) +}) diff --git a/src/api/providers/__tests__/bedrock.spec.ts b/src/api/providers/__tests__/bedrock.spec.ts index 16c45e9457..aafec75985 100644 --- a/src/api/providers/__tests__/bedrock.spec.ts +++ b/src/api/providers/__tests__/bedrock.spec.ts @@ -489,7 +489,7 @@ describe("AwsBedrockHandler", () => { }) } - it("uses non-cached input tokens from AI SDK v6 usage details", async () => { + it("emits total inputTokens (not noCacheTokens) from AI SDK v6 usage details", async () => { async function* mockFullStream() { yield { type: "text-delta", text: "Hello" } } @@ -520,9 +520,12 @@ describe("AwsBedrockHandler", () => { } const usageChunk = chunks.find((chunk) => chunk.type === "usage") + // inputTokens must be the total (13,071) per ApiStreamUsageChunk contract, + // with non-cached count in a separate field. expect(usageChunk).toMatchObject({ type: "usage", - inputTokens: 10, + inputTokens: 13_071, + nonCachedInputTokens: 10, outputTokens: 93, cacheWriteTokens: 489, cacheReadTokens: 12_572, diff --git a/src/api/providers/bedrock.ts b/src/api/providers/bedrock.ts index fe2390cbbf..1f75ba17e2 100644 --- a/src/api/providers/bedrock.ts +++ b/src/api/providers/bedrock.ts @@ -349,25 +349,16 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH info: ModelInfo, providerMetadata?: Record>, ): ApiStreamUsageChunk { - const inputTokens = usage.inputTokenDetails?.noCacheTokens ?? usage.inputTokens ?? 0 + // Keep inputTokens as the total (not noCacheTokens) so normalizeProviderUsage + // correctly populates inputTokensTotal. The normalizer reads non-cached tokens + // from inputTokenDetails.noCacheTokens via the anthropic base profile. + const inputTokens = usage.inputTokens ?? 0 const outputTokens = usage.outputTokens ?? 0 // The AI SDK exposes reasoningTokens as a top-level field on usage, and also // under outputTokenDetails.reasoningTokens — there is no .details property. const reasoningTokens = usage.reasoningTokens ?? usage.outputTokenDetails?.reasoningTokens ?? 0 - // Extract cache metrics primarily from usage (AI SDK standard locations), - // falling back to providerMetadata.bedrock.usage for provider-specific fields. - const bedrockUsage = providerMetadata?.bedrock?.usage as - | { cacheReadInputTokens?: number; cacheWriteInputTokens?: number } - | undefined - const cacheReadTokens = - usage.inputTokenDetails?.cacheReadTokens ?? - usage.cachedInputTokens ?? - bedrockUsage?.cacheReadInputTokens ?? - 0 - const cacheWriteTokens = usage.inputTokenDetails?.cacheWriteTokens ?? bedrockUsage?.cacheWriteInputTokens ?? 0 - // For prompt routers, the AI SDK surfaces the invoked model ID in // providerMetadata.bedrock.trace.promptRouter.invokedModelId. // When present, look up that model's pricing info for accurate cost calculation. @@ -410,7 +401,7 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH return { ...normalized.chunk, totalCost: this.calculateCost({ - inputTokens: normalized.canonical.inputTokensNonCached ?? inputTokens, + inputTokens: normalized.canonical.inputTokensTotal, outputTokens: normalized.canonical.outputTokens, cacheWriteTokens: normalized.canonical.cacheWriteTokens, cacheReadTokens: normalized.canonical.cacheReadTokens,