From e25239570c841163bc6038d6b8758edb27b68628 Mon Sep 17 00:00:00 2001
From: Tony Rogers <tony@tonyrogers.me>
Date: Tue, 1 Sep 2026 17:27:11 -0400
Subject: [PATCH] fix(bedrock): fold cache tokens into non-streaming converse
 usage
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The non-streaming Converse conversion copied inputTokens raw, dropping
cacheReadInputTokens / cacheWriteInputTokens even though the internal
TokenUsage parses them — AWS defines inputTokens as exclusive of cached
tokens, so prompt usage was understated whenever caching was active.
Fold them in, matching the streaming path.
---
 .../src/types/assistant_content.rs            | 40 ++++++++++++++++++-
 1 file changed, 39 insertions(+), 1 deletion(-)

diff --git a/rig-integrations/rig-bedrock/src/types/assistant_content.rs b/rig-integrations/rig-bedrock/src/types/assistant_content.rs
index d4a2e8e8cd..9521ec21b9 100644
--- a/rig-integrations/rig-bedrock/src/types/assistant_content.rs
+++ b/rig-integrations/rig-bedrock/src/types/assistant_content.rs
@@ -42,11 +42,16 @@ impl TryFrom<AwsConverseOutput> for completion::CompletionResponse<AwsConverseOu
             )),
         }?;
 
+        // Folds cache read/write tokens into `input_tokens`, matching the
+        // streaming path — Bedrock reports `inputTokens` exclusive of cached
+        // tokens, while `totalTokens` already includes them.
         let usage = value
             .0
             .usage()
             .map(|usage| completion::Usage {
-                input_tokens: usage.input_tokens as u64,
+                input_tokens: usage.input_tokens as u64
+                    + usage.cache_read_input_tokens.unwrap_or_default() as u64
+                    + usage.cache_write_input_tokens.unwrap_or_default() as u64,
                 output_tokens: usage.output_tokens as u64,
                 total_tokens: usage.total_tokens as u64,
             })
@@ -191,6 +196,39 @@ mod tests {
         );
     }
 
+    #[test]
+    fn aws_converse_output_folds_cache_tokens_into_input() {
+        let message = aws_bedrock::Message::builder()
+            .role(aws_bedrock::ConversationRole::Assistant)
+            .content(aws_bedrock::ContentBlock::Text("txt".into()))
+            .build()
+            .unwrap();
+        let converse_output =
+            aws_sdk_bedrockruntime::operation::converse::ConverseOutput::builder()
+                .output(aws_bedrock::ConverseOutput::Message(message))
+                .stop_reason(aws_bedrock::StopReason::EndTurn)
+                .build()
+                .unwrap();
+        let mut internal: InternalConverseOutput = converse_output.try_into().unwrap();
+        // Figures shaped like a Converse response with caching active:
+        // inputTokens excludes cached tokens, totalTokens includes them.
+        internal.usage = Some(crate::types::converse_output::TokenUsage {
+            input_tokens: 100,
+            output_tokens: 50,
+            total_tokens: 20_150,
+            cache_read_input_tokens: Some(18_000),
+            cache_write_input_tokens: Some(2_000),
+        });
+
+        let completion: completion::CompletionResponse<AwsConverseOutput> =
+            AwsConverseOutput(internal).try_into().unwrap();
+        assert_eq!(
+            completion.usage.input_tokens, 20_100,
+            "cache tokens fold into input"
+        );
+        assert_eq!(completion.usage.total_tokens, 20_150);
+    }
+
     #[test]
     fn aws_content_block_to_assistant_content() {
         let content_block = aws_bedrock::ContentBlock::Text("text".into());