Bedrock Converse の input_tokens がキャッシュトークンを二重減算していたバグを修正

AWS Bedrock の Chat Converse 実装で、input_tokens がキャッシュトークンを二度引くことで実使用トークン数を過小評価していた問題を解消しました。これによりコスト計算や利用統計が正確に取得できるようになります。

背景

問題発覚の経緯として、Converse::Chat#input_tokens が cacheReadInputTokens と cacheWriteInputTokens を inputTokens から減算していた点が指摘されました。AWS の公式ドキュメントによれば、inputTokens はすでにキャッシュされていないトークン数だけを示し、キャッシュの読み書きは別項目として報告されます。

実際の影響は、キャッシュが有効になるマルチターン会話で input_tokens がゼロになる、あるいは負の値になるケースが頻発し、費用計算やメトリクスが大幅に過小評価されていた点です。特に LangSmith など外部トラッキングツールで負のコストが報告される問題が顕在化していました。

修正の必要性は、同一モデルを扱う Anthropic プロバイダーでは正しい実装が既に存在し、Bedrock でも同様の振る舞いに統一すべきという設計方針に合致します。

技術的な変更

input_tokens メソッドのロジックを lib/ruby_llm/protocols/converse/chat.rb でシンプルにし、usage['inputTokens'] をそのまま返すように変更しました。これによりキャッシュバケットの二重減算が排除されます。

@@
-          input_tokens = usage['inputTokens']
-          return unless input_tokens
-
-          [input_tokens.to_i - usage['cacheReadInputTokens'].to_i - usage['cacheWriteInputTokens'].to_i, 0].max
+          # AWS Bedrock reports inputTokens as already non-cached; cacheReadInputTokens and
+          # cacheWriteInputTokens are separate buckets, not folded into inputTokens. Subtracting
+          # them (as inclusive providers require) understates input and floors to zero on cache
+          # hits. See https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html
+          usage['inputTokens']

テストケースの更新は spec/ruby_llm/protocols/converse/chat_spec.rb に反映され、inputTokens がそのまま露出されること、キャッシュトークンが別属性として取得できること、そしてキャッシュが入力を上回っても input_tokens が 0 になることがないことを検証しています。

@@
-      'inputTokens' => 100,
+      'inputTokens' => 50,
@@
-      expect(message.input_tokens).to eq(60)
+      expect(message.input_tokens).to eq(50)
@@
+    it 'does not subtract cache buckets or floor to zero when the cached prefix exceeds fresh input' do
+      response_body = {
+        'modelId' => 'anthropic.claude-sonnet-4-5-20250929-v1:0',
+        'output' => {
+          'message' => {
+            'content' => [{ 'text' => 'Hi!' }]
+          }
+        },
+        'usage' => {
+          'inputTokens' => 3,
+          'outputTokens' => 5,
+          'cacheReadInputTokens' => 7714,
+          'cacheWriteInputTokens' => 327
+        }
+      }
+
+      response = instance_double(Faraday::Response, body: response_body)
+      message = described_class.parse_completion_response(response)
+
+      expect(message.input_tokens).to eq(3)
+      expect(message.cached_tokens).to eq(7714)
+      expect(message.cache_creation_tokens).to eq(327)
+    end

設計判断

既存実装の拡張方針として、Converse プロトコルは input_tokens をそのまま受け入れ、キャッシュ情報は別プロパティ (cached_tokens, cache_creation_tokens) に保持する設計に統一しました。これにより、プロバイダー間でのトークン意味合いの不整合が解消されます。

代替案の検討では、キャッシュ情報を合算した新たなフィールドを導入する案もありましたが、外部ツール（LangSmith など）が期待する個別のメトリクスを保護するため、現行の分離設計を維持しました。結果として、後方互換性を損なうことなく正確なコスト計算が可能になっています。

影響範囲は Converse プロトコルを利用する全ての Bedrock 統合に及びますが、input_tokens を boolean あるいは他形式で扱っているコードは存在しないため、実運用上の破壊的変更はありません。

まとめ

今回の修正は、AWS Bedrock のトークンレポート仕様に忠実に従い、input_tokens の二重減算を排除したことが核心です。正確な入力トークン数が取得できるようになったことで、料金算出や使用量モニタリングが信頼できるようになり、他プロバイダーとの一貫性も確保できました。