fix(headroom): protect the cached prefix through the last cache_control breakpoint

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/compression/compress.py

@@ -205,21 +205,41 @@ def _extract_anthropic_tool_exchange_spans(
     return spans, None
 
 
+def _has_cache_control(message: Mapping[str, object]) -> bool:
+    if message.get("cache_control") is not None:
+        return True
+    content: Final = message.get("content")
+    return isinstance(content, list) and any(
+        isinstance(part, Mapping) and part.get("cache_control") is not None for part in content
+    )
+
+
+def _cached_prefix_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int, ...]:
+    last_breakpoint: Final = max(
+        (index for index, msg in enumerate(messages) if _has_cache_control(msg)),
+        default=-1,
+    )
+    return tuple(range(last_breakpoint + 1))
+
+
 def get_protected_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int, ...]:
     """
     Return indices of messages that must never be compressed:
     - All system messages
     - The last user message
     - The last assistant message
+    - Every message up to and including the last one carrying an Anthropic cache_control breakpoint
 
     The last user message is what the model is being asked to act on right now,
     so compressing it replaces the live instruction with a marker. Compression
     guardrails share this policy; see the Headroom guardrail.
+    The provider caches the exact bytes of that prefix, so rewriting any row inside
+    it turns the next request's cache read into a cache write.
     """
     system_indices: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "system")
     last_user: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "user")[-1:]
-    last_assistant = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "assistant")[-1:]
-    return system_indices + last_user + last_assistant
+    last_assistant: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "assistant")[-1:]
+    return tuple(dict.fromkeys(system_indices + last_user + last_assistant + _cached_prefix_indices(messages)))
 
 
 def _combine_scores(

tests/test_litellm/compression/test_compress.py

@@ -53,3 +53,87 @@ def test_every_system_row_is_protected():
 def test_no_user_or_assistant_rows():
     assert sorted(get_protected_indices([{"role": "system", "content": "sys"}])) == [0]
     assert get_protected_indices([]) == ()
+
+
+def test_rows_before_last_cache_control_breakpoint_are_protected():
+    messages = [
+        {"role": "system", "content": "sys"},
+        {"role": "user", "content": "old question"},
+        {
+            "role": "assistant",
+            "content": "old answer",
+            "tool_calls": [
+                {"id": "t1", "type": "function", "function": {"name": "Read", "arguments": "{}"}}
+            ],
+        },
+        {"role": "tool", "tool_call_id": "t1", "content": "large file body"},
+        {
+            "role": "user",
+            "content": [{"type": "text", "text": "cached turn", "cache_control": {"type": "ephemeral"}}],
+        },
+        {
+            "role": "assistant",
+            "content": "ack",
+            "tool_calls": [
+                {"id": "t2", "type": "function", "function": {"name": "Bash", "arguments": "{}"}}
+            ],
+        },
+        {"role": "tool", "tool_call_id": "t2", "content": "later tool output"},
+        {"role": "user", "content": "live instruction"},
+    ]
+
+    protected = sorted(get_protected_indices(messages))
+
+    assert protected == [0, 1, 2, 3, 4, 5, 7]
+    assert 6 not in protected
+
+
+def test_cache_control_directly_on_message_protects_prefix():
+    messages = [
+        {"role": "system", "content": "sys"},
+        {"role": "tool", "tool_call_id": "before", "content": "large file body"},
+        {"role": "user", "content": "old question"},
+        {
+            "role": "tool",
+            "tool_call_id": "marked",
+            "content": "cached tool",
+            "cache_control": {"type": "ephemeral"},
+        },
+        {"role": "tool", "tool_call_id": "after", "content": "later tool output"},
+        {"role": "assistant", "content": "ack"},
+        {"role": "user", "content": "live instruction"},
+    ]
+
+    protected = sorted(get_protected_indices(messages))
+
+    assert 1 in protected
+    assert 3 in protected
+    assert 4 not in protected
+
+
+def test_no_cache_control_leaves_history_compressible():
+    messages = [
+        {"role": "system", "content": "sys"},
+        {"role": "user", "content": "old question"},
+        {"role": "assistant", "content": "old answer"},
+        {"role": "tool", "tool_call_id": "t1", "content": "large file body"},
+        {"role": "user", "content": "live instruction"},
+    ]
+
+    assert sorted(get_protected_indices(messages)) == [0, 2, 4]
+
+
+def test_non_mapping_content_parts_are_not_cache_control():
+    messages = [
+        {"role": "system", "content": "sys"},
+        {"role": "user", "content": ["not", "a", "dict"]},
+        {"role": "assistant", "content": "old answer"},
+        {"role": "tool", "tool_call_id": "t1", "content": "plain string"},
+        {"role": "user", "content": "live instruction"},
+    ]
+
+    protected = sorted(get_protected_indices(messages))
+
+    assert protected == [0, 2, 4]
+    assert 1 not in protected
+    assert 3 not in protected

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py

@@ -1797,12 +1797,8 @@ PARTS_MESSAGES = [
     {
         "role": "user",
         "content": [
-            {"type": "text", "text": "Earlier turn.", "cache_control": {"type": "ephemeral"}},
-            {
-                "type": "text",
-                "text": "Second block. " + "B" * 5000,
-                "cache_control": {"type": "ephemeral", "ttl": "1h"},
-            },
+            {"type": "text", "text": "Earlier turn."},
+            {"type": "text", "text": "Second block. " + "B" * 5000},
         ],
     },
     {
@@ -1891,14 +1887,9 @@ async def test_apply_guardrail_restores_rewritten_all_text_row(
 
     messages = result["structured_messages"]
     history_content = messages[1]["content"]
-    # Rewritten all-text row collapses to one part carrying the LAST declared
-    # breakpoint: an Anthropic breakpoint caches the prefix ending at its
-    # part, so after the merge the last one (and its TTL) still describes the
-    # row.
     assert isinstance(history_content, list)
     assert len(history_content) == 1
     assert history_content[0]["text"] == "compressed history. Retrieve more: hash=b573993006976af767214fac"
-    assert history_content[0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"}
     # Mixed row passes through byte-identical.
     assert messages[2]["content"] == PARTS_MESSAGES[2]["content"]
     # The service-declared hash still drives retrieve-tool injection on a restored row.
@@ -2523,6 +2514,42 @@ async def test_history_is_still_compressed(guardrail: HeadroomGuardrail):
     assert messages[3] == compressed_history[1]
 
 
+CACHED_PREFIX_MESSAGES = [
+    {"role": "system", "content": "You are Claude Code. " + "S" * 5000},
+    {"role": "user", "content": "old question " + "Q" * 5000},
+    {
+        "role": "assistant",
+        "content": "Reading the file now.",
+        "tool_calls": [
+            {"id": "old_1", "type": "function", "function": {"name": "Read", "arguments": "{}"}}
+        ],
+    },
+    {"role": "tool", "tool_call_id": "old_1", "content": "large file body " + "F" * 5000},
+    {
+        "role": "user",
+        "content": [{"type": "text", "text": "cached turn", "cache_control": {"type": "ephemeral"}}],
+    },
+    {
+        "role": "assistant",
+        "content": "Listing now.",
+        "tool_calls": [
+            {"id": "new_1", "type": "function", "function": {"name": "Bash", "arguments": "{}"}}
+        ],
+    },
+    {"role": "tool", "tool_call_id": "new_1", "content": "volatile tail output " + "T" * 5000},
+    {"role": "assistant", "content": "Finished listing."},
+    {"role": "user", "content": "live instruction"},
+]
+
+
+@pytest.mark.asyncio
+async def test_rows_before_last_cache_control_breakpoint_are_never_sent(guardrail: HeadroomGuardrail):
+    wire, result = await _wire_and_result(guardrail, CACHED_PREFIX_MESSAGES)
+
+    assert [row.get("tool_call_id") for row in wire] == ["new_1"]
+    assert result["structured_messages"][:5] == CACHED_PREFIX_MESSAGES[:5]
+
+
 # ---------------------------------------------------------------------------
 # #38558: a client that runs its own tool loop (e.g. Claude Code via the MCP
 # gateway) executes headroom_retrieve and echoes the recovered original content