fix(bedrock): guard cache-detail split against partial/unrecognized ttl entries

Address review feedback on #36762:
- Only use the parsed 5m/1h split when it fully accounts for
  cacheWriteInputTokens; an unrecognized ttl or missing entry now falls
  back to the aggregate (previous behavior) instead of silently
  understating cost.
- Mark TypedDict fields ReadOnly (AWS response data, never constructed
  by us) to satisfy the repo's type-discipline lint gate.
- Trim comments and add Final to locals per repo style.

Co-Authored-By: pi (Claude/GPT via @earendil-works/pi-coding-agent) <noreply@earendil.works>
This commit is contained in:
Daniel Vainshtein 2026-08-13 14:06:22 +03:00
parent 97290b4e0e
commit 42a2b5f057
3 changed files with 42 additions and 21 deletions

View file

@ -1773,17 +1773,17 @@ class AmazonConverseConfig(BaseConfig):
@staticmethod
def _parse_cache_details(usage: ConverseTokenUsageBlock) -> "CacheCreationTokenDetails | None":
"""
Split Converse's aggregate cacheWriteInputTokens into the 5m/1h TTL
breakdown from `cacheDetails`, so cost calc can bill each tier
correctly instead of defaulting the whole write to the 5m rate.
https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_CacheDetail.html
"""
cache_details = usage.get("cacheDetails")
"""https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_CacheDetail.html"""
cache_details: Final = usage.get("cacheDetails")
if not cache_details:
return None
tokens_5m = sum(d["inputTokens"] for d in cache_details if d.get("ttl") == "5m")
tokens_1h = sum(d["inputTokens"] for d in cache_details if d.get("ttl") == "1h")
tokens_5m: Final = sum(d["inputTokens"] for d in cache_details if d.get("ttl") == "5m")
tokens_1h: Final = sum(d["inputTokens"] for d in cache_details if d.get("ttl") == "1h")
# An unrecognized ttl or a partial breakdown would silently understate
# the cache-write cost, so only use the split when it fully accounts
# for the aggregate; otherwise fall back to the aggregate-only (5m) cost.
if tokens_5m + tokens_1h != usage.get("cacheWriteInputTokens", 0):
return None
return CacheCreationTokenDetails(
ephemeral_5m_input_tokens=tokens_5m,
ephemeral_1h_input_tokens=tokens_1h,

View file

@ -2,7 +2,7 @@ import json
from enum import Enum
from typing import TYPE_CHECKING, Any, Final, Literal
from typing_extensions import Required, TypedDict, override
from typing_extensions import ReadOnly, Required, TypedDict, override
from .openai import ChatCompletionToolCallChunk
@ -217,21 +217,21 @@ class ConverseResponseOutputBlock(TypedDict):
class CacheDetailBlock(TypedDict):
"""Per-TTL cache-write breakdown. https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_CacheDetail.html"""
"""Per-TTL cache-write breakdown, read-only AWS response data. https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_CacheDetail.html"""
inputTokens: int
ttl: Literal["5m", "1h"]
inputTokens: ReadOnly[int]
ttl: ReadOnly[Literal["5m", "1h"]]
class ConverseTokenUsageBlock(TypedDict, total=False):
inputTokens: Required[int]
outputTokens: Required[int]
totalTokens: Required[int]
cacheReadInputTokenCount: int
cacheReadInputTokens: int
cacheWriteInputTokenCount: int
cacheWriteInputTokens: int
cacheDetails: list[CacheDetailBlock]
inputTokens: Required[ReadOnly[int]]
outputTokens: Required[ReadOnly[int]]
totalTokens: Required[ReadOnly[int]]
cacheReadInputTokenCount: ReadOnly[int]
cacheReadInputTokens: ReadOnly[int]
cacheWriteInputTokenCount: ReadOnly[int]
cacheWriteInputTokens: ReadOnly[int]
cacheDetails: ReadOnly[list[CacheDetailBlock]] # mutable-ok: AWS response array, never mutated after parsing
class ServiceTierBlock(TypedDict):

View file

@ -75,6 +75,27 @@ def test_transform_usage_with_cache_details():
assert details.ephemeral_5m_input_tokens == 288
def test_transform_usage_with_mismatched_cache_details_falls_back():
"""An unrecognized ttl or partial breakdown must not silently understate
cache-write cost, so the split is only used when it fully accounts for
cacheWriteInputTokens."""
usage = ConverseTokenUsageBlock(
**{
"inputTokens": 76,
"outputTokens": 259,
"totalTokens": 335,
"cacheWriteInputTokens": 362,
"cacheDetails": [{"inputTokens": 74, "ttl": "1h"}], # missing the 5m entry
}
)
config = AmazonConverseConfig()
openai_usage = config._transform_usage(usage)
assert (
getattr(openai_usage.prompt_tokens_details, "cache_creation_token_details", None)
is None
)
def test_transform_usage_without_cache_details_stays_none():
"""No cacheDetails in the response (older models/regions) should leave
cache_creation_token_details unset, same as before this field existed."""