From 9c59feee7cca3de0cb9727e9463cc5a3f66bf27b Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 15 Sep 2026 01:07:41 +0000 Subject: [PATCH 1/2] 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 | 24 +++++- .../test_litellm/compression/test_compress.py | 84 +++++++++++++++++++ .../guardrail_hooks/test_headroom.py | 49 ++++++++--- 3 files changed, 144 insertions(+), 13 deletions(-) diff --git a/litellm/compression/compress.py b/litellm/compression/compress.py index c646baf9d9e..99410a533f9 100644 --- a/litellm/compression/compress.py +++ b/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( diff --git a/tests/test_litellm/compression/test_compress.py b/tests/test_litellm/compression/test_compress.py index 6827c37dfd5..57992ffcc55 100644 --- a/tests/test_litellm/compression/test_compress.py +++ b/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 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index d8eeb8d2b8a..315c85936f3 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/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 From e390dfbb64160e3aa6a32a47ab597e8002a7d437 Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 15 Sep 2026 01:13:17 +0000 Subject: [PATCH 2/2] fix(headroom): format protected index assignment Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/compression/compress.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/compression/compress.py b/litellm/compression/compress.py index 99410a533f9..0af9618382f 100644 --- a/litellm/compression/compress.py +++ b/litellm/compression/compress.py @@ -238,7 +238,9 @@ def get_protected_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int """ 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: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "assistant")[-1:] + 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)))