From aa294302fe501b5c2bf6a196288530c3384f6b26 Mon Sep 17 00:00:00 2001 From: Lucien <164857476+dakjdakd@users.noreply.github.com> Date: Fri, 2 Oct 2026 14:56:19 +0000 Subject: [PATCH] fix(bedrock): preserve the service total when folding cache input tokens (#12586) Co-authored-by: Douwe Osinga --- .github/workflows/mcp-conformance.yml | 1 - Justfile | 3 ++- .../src/conversation/token_usage.rs | 2 +- crates/goose/src/providers/formats/bedrock.rs | 17 +++++++++++------ 4 files changed, 14 insertions(+), 9 deletions(-) diff --git a/.github/workflows/mcp-conformance.yml b/.github/workflows/mcp-conformance.yml index 52fa00496f..4109dbe83b 100644 --- a/.github/workflows/mcp-conformance.yml +++ b/.github/workflows/mcp-conformance.yml @@ -45,7 +45,6 @@ jobs: path: | target/debug/goose target/debug/mcp_conformance_driver - retention-days: 1 if-no-files-found: error conformance: diff --git a/Justfile b/Justfile index 4b234d6cb3..e52de147a6 100644 --- a/Justfile +++ b/Justfile @@ -489,7 +489,8 @@ mcp-conformance version="2025-11-25" suite="all" conformance_version="0.2.0-alph if [ -n "{{baseline}}" ]; then baseline_args=(--expected-failures "{{baseline}}") fi - GOOSE_DISABLE_KEYRING=1 npx -y @modelcontextprotocol/conformance@{{conformance_version}} client --command "target/debug/mcp_conformance_driver" --spec-version "{{version}}" --suite "{{suite}}" ${baseline_args[@]+"${baseline_args[@]}"} + export GOOSE_BIN="$PWD/target/debug/goose" + GOOSE_DISABLE_KEYRING=1 npx -y @modelcontextprotocol/conformance@{{conformance_version}} client --command "$PWD/target/debug/mcp_conformance_driver" --spec-version "{{version}}" --suite "{{suite}}" ${baseline_args[@]+"${baseline_args[@]}"} build-test-tools: cargo build -p goose-test diff --git a/crates/goose-provider-types/src/conversation/token_usage.rs b/crates/goose-provider-types/src/conversation/token_usage.rs index 2b7f3abb16..a5d2dd0fba 100644 --- a/crates/goose-provider-types/src/conversation/token_usage.rs +++ b/crates/goose-provider-types/src/conversation/token_usage.rs @@ -174,7 +174,7 @@ impl Usage { } /// For providers whose reported `input_tokens`/`total_tokens` exclude - /// cache tokens (e.g. Anthropic, Bedrock): folds the cache breakdown in. + /// cache tokens (e.g. Anthropic): folds the cache breakdown in. pub fn from_cache_exclusive_input( input_tokens: Option, output_tokens: Option, diff --git a/crates/goose/src/providers/formats/bedrock.rs b/crates/goose/src/providers/formats/bedrock.rs index 584a37d3b0..611ecf1b3e 100644 --- a/crates/goose/src/providers/formats/bedrock.rs +++ b/crates/goose/src/providers/formats/bedrock.rs @@ -612,13 +612,18 @@ pub fn from_bedrock_role(role: &bedrock::ConversationRole) -> Result { } pub fn from_bedrock_usage(usage: &bedrock::TokenUsage) -> Usage { - Usage::from_cache_exclusive_input( - Some(usage.input_tokens), + let cache_read = usage.cache_read_input_tokens; + let cache_write = usage.cache_write_input_tokens; + let input_tokens = usage + .input_tokens + .saturating_add(cache_read.unwrap_or(0)) + .saturating_add(cache_write.unwrap_or(0)); + Usage::new( + Some(input_tokens), Some(usage.output_tokens), Some(usage.total_tokens), - usage.cache_read_input_tokens, - usage.cache_write_input_tokens, ) + .with_cache_tokens(cache_read, cache_write) } pub fn from_bedrock_json(document: &Document) -> Result { @@ -1106,11 +1111,11 @@ mod tests { } #[test] - fn test_from_bedrock_usage_folds_cache_tokens_into_input() { + fn test_from_bedrock_usage_includes_cache_in_input_and_preserves_total() { let usage = bedrock::TokenUsage::builder() .input_tokens(7) .output_tokens(50) - .total_tokens(57) + .total_tokens(6057) .cache_read_input_tokens(5000) .cache_write_input_tokens(1000) .build()