From ea0115cbdbbf180062cc59e794371f62bf5875c3 Mon Sep 17 00:00:00 2001 From: Ingmar Krusch Date: Wed, 15 Jul 2026 21:58:06 +0200 Subject: [PATCH] fix(backend/bedrock): preserve system-prompt cache_control breakpoint (list form) (#2225) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description The LiteLLM backend flattened the Anthropic top-level `system` field to a joined string whenever it arrived as a **list of content blocks**, discarding each block's `cache_control`. LiteLLM's Bedrock Converse transform (`AmazonConverseConfig._transform_system_message`) only emits a `cachePoint` for content blocks that carry `cache_control`, never for a plain string. So on any `--backend bedrock` deployment the **system prefix was never cached**: every turn re-sent the full system prompt (typically 5k-25k tokens with Claude Code) at full input price. #1390 fixed the analogous case for `tool_result` blocks in `_convert_messages_for_litellm`, but the top-level `system` field handling in `send_message` / `stream_message` was out of scope there and still flattened. The cache hits observed on live Bedrock traffic came only from the tool-result / message-tail breakpoint, masking that the largest, most stable block was uncached. Closes # ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `headroom/backends/litellm.py`: factor the top-level `system` field conversion into a single `_system_field_to_message` helper. `str` stays string-content (unchanged behavior); a `list` maps to text blocks retaining each block's `cache_control`; non-dict entries coerce to a plain text block. Both call sites (`send_message` non-streaming, `stream_message` streaming) now call the helper, so they stay byte-identical. - `tests/test_bedrock_tool_result_cache_and_streaming_stats.py`: add `TestSystemFieldCacheControl` — list-with-`cache_control` retains it, plain-string is unchanged, list-without-`cache_control` produces list content with no marker, plus two end-to-end checks that drive the emitted message through `AmazonConverseConfig._transform_system_message` and assert a `cachePoint` is present for the cache_control case and absent otherwise. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ uv run pytest tests/test_bedrock_tool_result_cache_and_streaming_stats.py -q collected 13 items tests/test_bedrock_tool_result_cache_and_streaming_stats.py ............. [100%] 13 passed in 1.18s $ uv run ruff check headroom/backends/litellm.py tests/test_bedrock_tool_result_cache_and_streaming_stats.py All checks passed! ``` ## Real Behavior Proof - Environment: personal fork deployed as a real proxy (macOS launchd service, `headroom install apply`) with `--backend bedrock --mode cache`, fronting a live Claude Code session. Model `global.anthropic.claude-sonnet-5`, region eu-west-1. - Exact command / steps: ran a purpose-built probe that POSTs Anthropic-shape `/v1/messages` to the running proxy with a 7,692-token STABLE system prompt carrying a single `cache_control: {type: ephemeral}` breakpoint (and no other cache_control anywhere), a pinned `x-headroom-session-id`, across 5 sequential turns, reading the raw response `usage` each turn. - Observed result: **before the fix**, the response `usage` had no cache fields at all — `cache_creation_input_tokens` and `cache_read_input_tokens` both absent, nothing cached. **After the fix**, turn 1 shows `cache_creation_input_tokens=10164` (write) and turns 2-5 each show `cache_read_input_tokens=10164` (read) with `cache_creation=0` — write-once, then read the system prefix from Bedrock's cache on every subsequent turn. The proxy's `/stats` `prefix_cache` tracker registered all four later turns as hits (`hit_requests += 1` per turn, `bust_count = 0`). - Not tested: no change to the tool_result / message-tail breakpoint path (already handled by #1390 / #2144); this fix is scoped to the top-level `system` field only. The in-`messages` text-block flatten in `_convert_messages_for_litellm` is intentionally left untouched. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective - [x] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable ## Additional Notes - No linked issue number: found via independent investigation of a personal `--backend bedrock` deployment. - Companion to #2196 (`fix(proxy/bedrock): wire PrefixCacheTracker updates into Bedrock backend paths`) from the same investigation. #2196 wires the tracker; this fixes the system-prompt breakpoint that #1390 left flattened on the top-level `system` field. Co-authored-by: Claude Opus 4.8 (1M context) --- headroom/backends/litellm.py | 52 ++++++++++----- ...k_tool_result_cache_and_streaming_stats.py | 66 +++++++++++++++++++ 2 files changed, 101 insertions(+), 17 deletions(-) diff --git a/headroom/backends/litellm.py b/headroom/backends/litellm.py index 526a8d76e..36987b3e5 100644 --- a/headroom/backends/litellm.py +++ b/headroom/backends/litellm.py @@ -747,6 +747,39 @@ class LiteLLMBackend(Backend): return converted + def _system_field_to_message(self, system: Any) -> dict[str, Any]: + """Convert Anthropic's top-level `system` field to an OpenAI-style message. + + `system` can be a plain string or a list of content blocks, each of + which may carry its own `cache_control` breakpoint (Claude Code puts + the prompt-caching marker on the last system block). Flattening the + list to a joined string, as this code used to do, drops that + `cache_control` entirely: litellm's Bedrock Converse transformation + only emits a `cachePoint` when it sees content blocks with + `cache_control` on them, never for a plain string. That silently + broke prompt caching of the system prefix. #1390 covers the analogous + case for tool_result blocks in `_convert_messages_for_litellm` above; + this handles the top-level `system` field, which was out of scope + there. Preserve block structure and cache_control so the breakpoint + survives into the litellm call. + """ + if isinstance(system, str): + return {"role": "system", "content": system} + if isinstance(system, list): + blocks: list[dict[str, Any]] = [] + for s in system: + if isinstance(s, dict): + block: dict[str, Any] = {"type": "text", "text": s.get("text", "")} + if "cache_control" in s: + block["cache_control"] = s["cache_control"] + else: + block = {"type": "text", "text": str(s)} + blocks.append(block) + return {"role": "system", "content": blocks} + # Shouldn't happen in practice (None is filtered out via "system" in + # body), but stay defensive rather than raising. + return {"role": "system", "content": str(system)} + def _to_anthropic_response( self, litellm_response: Any, @@ -843,15 +876,7 @@ class LiteLLMBackend(Backend): # System prompt (Anthropic puts it in body, OpenAI in messages) if "system" in body: - system = body["system"] - if isinstance(system, str): - kwargs["messages"].insert(0, {"role": "system", "content": system}) - elif isinstance(system, list): - # Anthropic list format - system_text = " ".join( - s.get("text", "") if isinstance(s, dict) else str(s) for s in system - ) - kwargs["messages"].insert(0, {"role": "system", "content": system_text}) + kwargs["messages"].insert(0, self._system_field_to_message(body["system"])) # Provider-specific region config if self.region: @@ -956,14 +981,7 @@ class LiteLLMBackend(Backend): if "tool_choice" in body: kwargs["tool_choice"] = _convert_tool_choice(body["tool_choice"]) if "system" in body: - system = body["system"] - if isinstance(system, str): - kwargs["messages"].insert(0, {"role": "system", "content": system}) - elif isinstance(system, list): - system_text = " ".join( - s.get("text", "") if isinstance(s, dict) else str(s) for s in system - ) - kwargs["messages"].insert(0, {"role": "system", "content": system_text}) + kwargs["messages"].insert(0, self._system_field_to_message(body["system"])) # Provider-specific region config if self.region: diff --git a/tests/test_bedrock_tool_result_cache_and_streaming_stats.py b/tests/test_bedrock_tool_result_cache_and_streaming_stats.py index f29e0ccc7..59591f011 100644 --- a/tests/test_bedrock_tool_result_cache_and_streaming_stats.py +++ b/tests/test_bedrock_tool_result_cache_and_streaming_stats.py @@ -252,3 +252,69 @@ class TestStreamingCacheStatsCompletion: assert usage["input_tokens"] == 42 assert "cache_read_input_tokens" not in usage assert "cache_creation_input_tokens" not in usage + + +class TestSystemFieldCacheControl: + """The top-level Anthropic `system` field must not be flattened to a + plain string when it carries per-block cache_control, or Bedrock prompt + caching of the system prefix silently breaks (see module docstring, + #1390's uncovered case).""" + + def test_list_system_with_cache_control_preserved(self): + backend = _backend() + system = [ + {"type": "text", "text": "You are a helpful assistant."}, + {"type": "text", "text": "Long static prefix.", "cache_control": {"type": "ephemeral"}}, + ] + msg = backend._system_field_to_message(system) + assert msg["role"] == "system" + assert isinstance(msg["content"], list) + assert msg["content"][-1]["cache_control"] == {"type": "ephemeral"} + + def test_string_system_unaffected(self): + backend = _backend() + msg = backend._system_field_to_message("You are a helpful assistant.") + assert msg == {"role": "system", "content": "You are a helpful assistant."} + assert isinstance(msg["content"], str) + + def test_list_system_without_cache_control_has_no_cache_control_keys(self): + backend = _backend() + system = [ + {"type": "text", "text": "First block."}, + {"type": "text", "text": "Second block."}, + ] + msg = backend._system_field_to_message(system) + assert isinstance(msg["content"], list) + assert all("cache_control" not in block for block in msg["content"]) + + def test_bedrock_converse_transform_emits_cachepoint_for_list_with_cache_control(self): + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + + backend = _backend() + system = [ + {"type": "text", "text": "You are a helpful assistant."}, + {"type": "text", "text": "Long static prefix.", "cache_control": {"type": "ephemeral"}}, + ] + system_msg = backend._system_field_to_message(system) + messages = [system_msg, {"role": "user", "content": "hi"}] + + _, system_blocks = AmazonConverseConfig()._transform_system_message( + messages, model="global.anthropic.claude-sonnet-5" + ) + assert any("cachePoint" in block for block in system_blocks) + + def test_bedrock_converse_transform_omits_cachepoint_without_cache_control(self): + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + + backend = _backend() + system = [ + {"type": "text", "text": "First block."}, + {"type": "text", "text": "Second block."}, + ] + system_msg = backend._system_field_to_message(system) + messages = [system_msg, {"role": "user", "content": "hi"}] + + _, system_blocks = AmazonConverseConfig()._transform_system_message( + messages, model="global.anthropic.claude-sonnet-5" + ) + assert not any("cachePoint" in block for block in system_blocks)