Procházet zdrojové kódy

fix(ai): normalize Bedrock cache usage (#38427)

Co-authored-by: Aiden Cline <63023139+rekram1-node@users.noreply.github.com>
opencode-agent[bot] před 2 týdny
rodič
revize
6e8aefcfa0

+ 10 - 9
packages/ai/src/protocols/bedrock-converse.ts

@@ -436,21 +436,22 @@ const mapFinishReason = (reason: string): FinishReason => {
   return "unknown"
 }
 
-// AWS Bedrock Converse reports `inputTokens` (inclusive total) with
-// `cacheReadInputTokens` and `cacheWriteInputTokens` as subsets. Pass
-// the total through and derive the non-cached breakdown. Bedrock does
-// not break reasoning out of `outputTokens` for any current model.
+// AWS reports inputTokens separately from cache reads and writes.
+// Bedrock does not break reasoning out of outputTokens for current models.
 const mapUsage = (usage: BedrockUsageSchema | undefined): Usage | undefined => {
   if (!usage) return undefined
-  const cacheTotal = (usage.cacheReadInputTokens ?? 0) + (usage.cacheWriteInputTokens ?? 0)
-  const nonCached = ProviderShared.subtractTokens(usage.inputTokens, cacheTotal)
+  const inputTokens = ProviderShared.sumTokens(
+    usage.inputTokens,
+    usage.cacheReadInputTokens,
+    usage.cacheWriteInputTokens,
+  )
   return new Usage({
-    inputTokens: usage.inputTokens,
+    inputTokens,
     outputTokens: usage.outputTokens,
-    nonCachedInputTokens: nonCached,
+    nonCachedInputTokens: usage.inputTokens,
     cacheReadInputTokens: usage.cacheReadInputTokens,
     cacheWriteInputTokens: usage.cacheWriteInputTokens,
-    totalTokens: ProviderShared.totalTokens(usage.inputTokens, usage.outputTokens, usage.totalTokens),
+    totalTokens: ProviderShared.totalTokens(inputTokens, usage.outputTokens, usage.totalTokens),
     providerMetadata: { bedrock: usage },
   })
 }

+ 4 - 3
packages/ai/src/schema/events.ts

@@ -34,11 +34,12 @@ import { ProviderFailureClassification } from "./errors"
  *
  * **Semantics by provider**:
  *
- * - OpenAI Chat / Responses / Gemini / Bedrock: provider reports inclusive
+ * - OpenAI Chat / Responses / Gemini: provider reports inclusive
  *   `inputTokens` and an inclusive `outputTokens`; mapper subtracts to
  *   derive the breakdown.
- * - Anthropic: provider reports the breakdown natively (`input_tokens` is
- *   non-cached only); mapper sums to derive the inclusive `inputTokens`.
+ * - Anthropic and Bedrock report the input breakdown natively: Anthropic's
+ *   `input_tokens` and Bedrock's `inputTokens` are non-cached only. Their
+ *   mappers sum the breakdown to derive the inclusive `inputTokens`.
  *   Anthropic does *not* break extended-thinking out of `output_tokens`, so
  *   `reasoningTokens` is `undefined` and `outputTokens` carries the
  *   combined total — a documented limitation of the Anthropic API.

Rozdílová data souboru nebyla zobrazena, protože soubor je příliš velký
+ 21 - 0
packages/ai/test/fixtures/recordings/bedrock-converse-cache/writes-then-reads-cachepoint-on-identical-second-call.json


+ 14 - 8
packages/ai/test/provider/bedrock-converse-cache.recorded.test.ts

@@ -13,12 +13,8 @@ const RECORDING_REGION = process.env.BEDROCK_RECORDING_REGION ?? "us-east-1"
 // call wouldn't deterministically prove cache mapping works. Override with
 // BEDROCK_CACHE_MODEL_ID if your account has access elsewhere.
 const model = AmazonBedrock.configure({
-  credentials: {
-    region: RECORDING_REGION,
-    accessKeyId: process.env.AWS_ACCESS_KEY_ID ?? "fixture",
-    secretAccessKey: process.env.AWS_SECRET_ACCESS_KEY ?? "fixture",
-    sessionToken: process.env.AWS_SESSION_TOKEN,
-  },
+  apiKey: process.env.AWS_BEARER_TOKEN_BEDROCK ?? "fixture",
+  region: RECORDING_REGION,
 }).model(process.env.BEDROCK_CACHE_MODEL_ID ?? "us.anthropic.claude-haiku-4-5-20251001-v1:0")
 
 const cacheRequest = LLM.request({
@@ -36,7 +32,7 @@ const recorded = recordedTests({
   prefix: "bedrock-converse-cache",
   provider: "amazon-bedrock",
   protocol: "bedrock-converse",
-  requires: ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"],
+  requires: ["AWS_BEARER_TOKEN_BEDROCK"],
   // Two identical requests in one cassette — replay walks the cassette in
   // recording order so the second call replays the cached-hit interaction.
 })
@@ -45,10 +41,20 @@ describe("Bedrock Converse cache recorded", () => {
   recorded.effect.with("writes then reads cachePoint on identical second call", { tags: ["cache"] }, () =>
     Effect.gen(function* () {
       const first = yield* LLMClient.generate(cacheRequest)
-      expect(first.usage?.cacheReadInputTokens ?? 0).toBeGreaterThanOrEqual(0)
+      expect(first.usage?.cacheWriteInputTokens ?? 0).toBeGreaterThan(0)
+      expect(first.usage?.inputTokens).toBe(
+        (first.usage?.nonCachedInputTokens ?? 0) +
+          (first.usage?.cacheReadInputTokens ?? 0) +
+          (first.usage?.cacheWriteInputTokens ?? 0),
+      )
 
       const second = yield* LLMClient.generate(cacheRequest)
       expect(second.usage?.cacheReadInputTokens ?? 0).toBeGreaterThan(0)
+      expect(second.usage?.inputTokens).toBe(
+        (second.usage?.nonCachedInputTokens ?? 0) +
+          (second.usage?.cacheReadInputTokens ?? 0) +
+          (second.usage?.cacheWriteInputTokens ?? 0),
+      )
     }),
   )
 })

+ 33 - 0
packages/ai/test/provider/bedrock-converse.test.ts

@@ -269,6 +269,39 @@ describe("Bedrock Converse route", () => {
     }),
   )
 
+  it.effect("adds cache reads and writes to Bedrock input usage", () =>
+    Effect.gen(function* () {
+      const body = eventStreamBody(
+        ["messageStart", { role: "assistant" }],
+        ["contentBlockDelta", { contentBlockIndex: 0, delta: { text: "Hello" } }],
+        ["contentBlockStop", { contentBlockIndex: 0 }],
+        ["messageStop", { stopReason: "end_turn" }],
+        [
+          "metadata",
+          {
+            usage: {
+              inputTokens: 5,
+              outputTokens: 2,
+              totalTokens: 12,
+              cacheReadInputTokens: 3,
+              cacheWriteInputTokens: 2,
+            },
+          },
+        ],
+      )
+      const response = yield* LLMClient.generate(baseRequest).pipe(Effect.provide(fixedBytes(body)))
+
+      expect(response.usage).toMatchObject({
+        inputTokens: 10,
+        nonCachedInputTokens: 5,
+        cacheReadInputTokens: 3,
+        cacheWriteInputTokens: 2,
+        outputTokens: 2,
+        totalTokens: 12,
+      })
+    }),
+  )
+
   it.effect("assembles streamed tool call input", () =>
     Effect.gen(function* () {
       const body = eventStreamBody(

Některé soubory nejsou zobrazeny, neboť je v těchto rozdílových datech změněno mnoho souborů