Merge pull request #41161 from BerriAI/litellm_headroom_protect_cached_prefix

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

litellm/compression/compress.py

@@ -214,26 +214,33 @@ def _message_has_cache_control(message: Mapping[str, object]) -> bool:
     return False
 
 
+def _cached_prefix_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int, ...]:
+    last_breakpoint: Final = max(
+        (index for index, msg in enumerate(messages) if _message_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
-    - Any message carrying an Anthropic cache_control breakpoint
+    - 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. A cache_control
-    breakpoint pins the provider's prompt-cache prefix to that row's exact
-    bytes, so rewriting a marked row anywhere in history turns the next
-    request's cache read into a cache write.
+    breakpoint pins the provider's prompt-cache prefix to the exact bytes of every
+    row up to it, so rewriting any row inside that prefix 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:]
     assistant_indices: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "assistant")
-    cache_control_indices: Final = tuple(index for index, msg in enumerate(messages) if _message_has_cache_control(msg))
-    return tuple(dict.fromkeys(system_indices + last_user + assistant_indices[-1:] + cache_control_indices))
+    return tuple(dict.fromkeys(system_indices + last_user + assistant_indices[-1:] + _cached_prefix_indices(messages)))
 
 
 def _combine_scores(

tests/test_litellm/compression/test_compress.py

@@ -56,11 +56,87 @@ def test_no_user_or_assistant_rows():
     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
+
+
 def test_mid_history_cache_control_part_is_protected():
-    # A large cached tool result from a few turns back, not the last user or
-    # last assistant row -- exactly the row a provider prompt-cache pins to
-    # exact bytes. Rewriting it (even leaving the marker on) changes those
-    # bytes and turns the next request's cache read into a cache write.
     messages = [
         {"role": "user", "content": "old question"},
         {"role": "assistant", "content": "old answer"},
@@ -74,9 +150,7 @@ def test_mid_history_cache_control_part_is_protected():
         {"role": "user", "content": "live instruction"},
     ]
 
-    # index 3 = last assistant, index 4 = last user (both protected by role
-    # regardless), index 2 = the cache_control-marked row itself.
-    assert sorted(get_protected_indices(messages)) == [2, 3, 4]
+    assert sorted(get_protected_indices(messages)) == [0, 1, 2, 3, 4]
 
 
 def test_cache_control_directly_on_message_is_protected():
@@ -116,8 +190,6 @@ def test_content_that_is_not_a_list_of_mappings_is_not_treated_as_cache_control(
 
 
 def test_compress_keeps_part_level_cache_control_row_verbatim():
-    # compress() scores text-only copies of the rows, where a part-level marker
-    # is gone; protection has to read the original rows or the pinned row is stubbed.
     stale_log = {"role": "user", "content": [{"type": "text", "text": "stale log line " * 2000}]}
     pinned = {
         "role": "user",
@@ -126,9 +198,9 @@ def test_compress_keeps_part_level_cache_control_row_verbatim():
         ],
     }
     messages = [
-        stale_log,
-        {"role": "assistant", "content": "old answer"},
         pinned,
+        {"role": "assistant", "content": "old answer"},
+        stale_log,
         {"role": "assistant", "content": "ack"},
         {"role": "user", "content": "live instruction"},
     ]
@@ -142,6 +214,6 @@ def test_compress_keeps_part_level_cache_control_row_verbatim():
     )
 
     assert len(result["messages"]) == len(messages)
-    assert result["messages"][2] == pinned
-    assert result["messages"][0] != stale_log
+    assert result["messages"][0] == pinned
+    assert result["messages"][2] != stale_log
     assert len(result["cache"]) >= 1

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

@@ -900,7 +900,7 @@ async def test_service_declared_ccr_hashes_drive_injection_and_validation(guardr
     )
 
     assert has_headroom_retrieve_tool(result.get("tools") or [])
-    (issued, _expiry), = guardrail._issued_hashes_by_call_id.values()
+    ((issued, _expiry),) = guardrail._issued_hashes_by_call_id.values()
     assert issued == frozenset({"98ca69107318", "b573993006976af767214fac"})
 
 
@@ -953,7 +953,6 @@ async def test_anthropic_assistant_history_never_reaches_compression_service(gua
     assert result["messages"][1]["content"] == [{"type": "text", "text": table}]
 
 
-
 def test_has_headroom_retrieve_tool_recognizes_anthropic_native_shape():
     """By the time an Anthropic Messages API response reaches the agentic-loop
     gate, the OpenAI-shaped tool this guardrail injects (type: "function")
@@ -2342,9 +2341,7 @@ async def test_streaming_responses_resolves_ccr_retrieval_end_to_end(
     )
     assert streamed_text == final_answer
     assert not any("function_call" in str(getattr(event, "type", "")) for event in events)
-    assert not any(
-        getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events
-    )
+    assert not any(getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events)
     mock_get.assert_called_once()
     assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0])
 
@@ -2399,9 +2396,7 @@ def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end(
         getattr(event, "delta", "") for event in events if getattr(event, "type", None) == "response.output_text.delta"
     )
     assert streamed_text == final_answer
-    assert not any(
-        getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events
-    )
+    assert not any(getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events)
     mock_get.assert_called_once()
     assert len(upstream.calls) == 2
     assert not json.loads(upstream.calls[1].request.content).get("stream")
@@ -2514,6 +2509,38 @@ 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]
+
+
 CACHE_MARKED_HISTORY_MESSAGES = [
     {"role": "system", "content": "You are Claude Code. " + "S" * 5000},
     {"role": "user", "content": "old question " + "Q" * 5000},
@@ -2529,6 +2556,7 @@ CACHE_MARKED_HISTORY_MESSAGES = [
         "cache_control": {"type": "ephemeral"},
     },
     {"role": "assistant", "content": "Summarized the file for you."},
+    {"role": "tool", "tool_call_id": "tail", "content": "volatile tail output " + "T" * 5000},
     {"role": "user", "content": "live instruction"},
 ]