From 990dea27d5c87c7c48dbc286c2efa3c6a610cf54 Mon Sep 17 00:00:00 2001 From: Rad Wadud <104943953+rad-p44@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:52:41 -0500 Subject: [PATCH 1/3] fix(headroom): protect cache_control-marked rows anywhere in history get_protected_indices() only protected system rows, the last user row, and the last assistant row. A message carrying its own Anthropic cache_control breakpoint further back in history (e.g. a large cached tool result from a few turns ago) was not protected, so the Headroom guardrail would send it to /v1/compress and rewrite it. The row came back byte-different but kept its cache_control marker, so the provider's prompt cache treated the next request as a miss on that prefix: a cache read silently became a cache write. This reproduces the production cache-hit-rate collapse reported in #39519 (~65-70% down to ~40-50% within 48h of enabling the guardrail). get_protected_indices() now also protects any message whose content -- directly on the message, or on any part of a list-of-parts content -- carries a cache_control marker, regardless of its position in history. Both compress() and the Headroom guardrail already share this function as their compression-eligibility policy, so both get the fix. Adds test coverage for cache_control on the message dict itself, on a content part, mid-history, and de-duplicated against already-protected indices. Updates the Headroom guardrail's PARTS_MESSAGES fixture, which previously relied on this exact gap for its all-text merge/flatten test coverage, to use a separate un-marked row (the cache_control-marked-row merge behavior is covered directly by compresr's own test, since a marked row no longer reaches that merge path through Headroom). Fixes #39519 --- litellm/compression/compress.py | 35 +++++++++- .../test_litellm/compression/test_compress.py | 60 ++++++++++++++++ .../guardrail_hooks/test_headroom.py | 68 ++++++++++++++++--- 3 files changed, 151 insertions(+), 12 deletions(-) diff --git a/litellm/compression/compress.py b/litellm/compression/compress.py index c646baf9d9e..62b05a4938f 100644 --- a/litellm/compression/compress.py +++ b/litellm/compression/compress.py @@ -205,21 +205,54 @@ def _extract_anthropic_tool_exchange_spans( return spans, None +def _message_has_cache_control(message: Mapping[str, object]) -> bool: + """True if ``message`` carries an Anthropic ``cache_control`` breakpoint. + + A breakpoint can sit directly on the message dict, or on any part of a + list-of-parts ``content`` (the shape Anthropic's own messages use). Either + placement pins the provider's KV-cache prefix to this row's exact bytes, so + either placement must protect the row the same way. + """ + if message.get("cache_control") is not None: + return True + content: Final = message.get("content") + if isinstance(content, list): + return any(isinstance(part, Mapping) and part.get("cache_control") is not None for part in content) + return False + + 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 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. Rewriting the row (even leaving the marker in place) + changes those bytes, so the next request misses the cache it thinks it is + reusing and silently pays a cache write instead of a cache read. This is + not limited to the last user/assistant row: a marker several turns back + (e.g. on a large cached tool result) needs the same protection. """ 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 + cache_control_indices: Final = tuple( + index for index, msg in enumerate(messages) if _message_has_cache_control(msg) + ) + seen: Final[set[int]] = set() + ordered: Final[list[int]] = [] + for index in system_indices + last_user + last_assistant + cache_control_indices: + if index not in seen: + seen.add(index) + ordered.append(index) + return tuple(ordered) def _combine_scores( diff --git a/tests/test_litellm/compression/test_compress.py b/tests/test_litellm/compression/test_compress.py index 6827c37dfd5..f9877ea2bc4 100644 --- a/tests/test_litellm/compression/test_compress.py +++ b/tests/test_litellm/compression/test_compress.py @@ -53,3 +53,63 @@ 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_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"}, + { + "role": "user", + "content": [ + {"type": "text", "text": "a large cached tool result"}, + ], + }, + {"role": "assistant", "content": "ack"}, + {"role": "user", "content": "live instruction"}, + ] + messages[2]["content"][0]["cache_control"] = {"type": "ephemeral"} + + # 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] + + +def test_cache_control_directly_on_message_is_protected(): + messages = [ + {"role": "user", "content": "old question", "cache_control": {"type": "ephemeral"}}, + {"role": "assistant", "content": "old answer"}, + {"role": "user", "content": "live instruction"}, + ] + + assert sorted(get_protected_indices(messages)) == [0, 1, 2] + + +def test_cache_control_protection_does_not_duplicate_already_protected_rows(): + # The last user row is already protected by role; marking it too must not + # produce a duplicate index. + messages = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "live", "cache_control": {"type": "ephemeral"}}, + ] + + protected = get_protected_indices(messages) + + assert sorted(protected) == [0, 1] + assert len(protected) == len(set(protected)) + + +def test_content_that_is_not_a_list_of_mappings_is_not_treated_as_cache_control(): + # Defensive: a plain string content, or a list of non-dict items, must not + # raise or be misread as carrying a breakpoint. + messages = [ + {"role": "assistant", "content": "plain string content"}, + {"role": "user", "content": ["not", "a", "dict", "list"]}, + {"role": "user", "content": "live instruction"}, + ] + + assert sorted(get_protected_indices(messages)) == [0, 2] 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..5cd42bd3f83 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -1795,14 +1795,18 @@ PARTS_MESSAGES = [ ], }, { + # No cache_control here on purpose: this row exercises the general + # multi-part flatten/merge mechanics (shared with compresr). A row + # carrying its own cache_control is a different, dedicated case -- + # see test_mid_history_cache_control_row_is_never_sent_for_compression + # (#39519): get_protected_indices withholds it from /v1/compress + # entirely rather than letting it be rewritten and re-merged, because + # rewriting the bytes under a live breakpoint busts the cache the + # marker is supposed to preserve. "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 +1895,17 @@ 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. + # Rewritten all-text row collapses to one part carrying the rewritten text. + # This fixture row carries no cache_control (see PARTS_MESSAGES): the + # last-declared-breakpoint-survives-the-merge behavior is a property of + # merge_rewritten_text_parts and is covered directly by compresr's + # test_all_text_row_merges_and_keeps_last_cache_control, since a + # cache_control-marked row never reaches this merge path through Headroom + # at all -- get_protected_indices withholds it before compression runs + # (see test_mid_history_cache_control_row_is_never_sent_for_compression). 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 +2530,45 @@ async def test_history_is_still_compressed(guardrail: HeadroomGuardrail): assert messages[3] == compressed_history[1] +# --------------------------------------------------------------------------- +# #39519: a mid-history row carrying its own Anthropic cache_control marker +# (e.g. a large tool result the client already cached several turns back) was +# still sent to /v1/compress and rewritten. It came back byte-different but +# kept its marker, so the next request's cache read silently became a cache +# write. get_protected_indices() now protects any cache_control-marked row, +# not just system/last-user/last-assistant, so it must never reach the wire. +# --------------------------------------------------------------------------- + +CACHE_MARKED_HISTORY_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": [{"type": "text", "text": "large cached file body " + "F" * 5000}], + "cache_control": {"type": "ephemeral"}, + }, + {"role": "assistant", "content": "Summarized the file for you."}, + {"role": "user", "content": "live instruction"}, +] + + +@pytest.mark.asyncio +async def test_mid_history_cache_control_row_is_never_sent_for_compression(guardrail: HeadroomGuardrail): + wire, result = await _wire_and_result(guardrail, CACHE_MARKED_HISTORY_MESSAGES) + + cached_row = CACHE_MARKED_HISTORY_MESSAGES[3] + assert cached_row not in wire + assert not any(row.get("tool_call_id") == "old_1" for row in wire) + # Byte-identical, marker intact -- the next request's cache read survives. + assert result["structured_messages"][3] == cached_row + + # --------------------------------------------------------------------------- # #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 36c1e5e17d1326f1a8f3dc7b25e86a69349d53e7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:34:31 -0700 Subject: [PATCH 2/3] refactor(compression): build the protected index set without mutation --- litellm/compression/compress.py | 33 ++++--------------- .../guardrail_hooks/test_headroom.py | 26 --------------- 2 files changed, 7 insertions(+), 52 deletions(-) diff --git a/litellm/compression/compress.py b/litellm/compression/compress.py index 62b05a4938f..c79e6aed57a 100644 --- a/litellm/compression/compress.py +++ b/litellm/compression/compress.py @@ -206,13 +206,6 @@ def _extract_anthropic_tool_exchange_spans( def _message_has_cache_control(message: Mapping[str, object]) -> bool: - """True if ``message`` carries an Anthropic ``cache_control`` breakpoint. - - A breakpoint can sit directly on the message dict, or on any part of a - list-of-parts ``content`` (the shape Anthropic's own messages use). Either - placement pins the provider's KV-cache prefix to this row's exact bytes, so - either placement must protect the row the same way. - """ if message.get("cache_control") is not None: return True content: Final = message.get("content") @@ -231,28 +224,16 @@ def get_protected_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int 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. Rewriting the row (even leaving the marker in place) - changes those bytes, so the next request misses the cache it thinks it is - reusing and silently pays a cache write instead of a cache read. This is - not limited to the last user/assistant row: a marker several turns back - (e.g. on a large cached tool result) needs the same protection. + 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. """ 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:] - cache_control_indices: Final = tuple( - index for index, msg in enumerate(messages) if _message_has_cache_control(msg) - ) - seen: Final[set[int]] = set() - ordered: Final[list[int]] = [] - for index in system_indices + last_user + last_assistant + cache_control_indices: - if index not in seen: - seen.add(index) - ordered.append(index) - return tuple(ordered) + 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)) def _combine_scores( 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 5cd42bd3f83..d4531398ba1 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -1795,14 +1795,6 @@ PARTS_MESSAGES = [ ], }, { - # No cache_control here on purpose: this row exercises the general - # multi-part flatten/merge mechanics (shared with compresr). A row - # carrying its own cache_control is a different, dedicated case -- - # see test_mid_history_cache_control_row_is_never_sent_for_compression - # (#39519): get_protected_indices withholds it from /v1/compress - # entirely rather than letting it be rewritten and re-merged, because - # rewriting the bytes under a live breakpoint busts the cache the - # marker is supposed to preserve. "role": "user", "content": [ {"type": "text", "text": "Earlier turn."}, @@ -1895,14 +1887,6 @@ 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 rewritten text. - # This fixture row carries no cache_control (see PARTS_MESSAGES): the - # last-declared-breakpoint-survives-the-merge behavior is a property of - # merge_rewritten_text_parts and is covered directly by compresr's - # test_all_text_row_merges_and_keeps_last_cache_control, since a - # cache_control-marked row never reaches this merge path through Headroom - # at all -- get_protected_indices withholds it before compression runs - # (see test_mid_history_cache_control_row_is_never_sent_for_compression). assert isinstance(history_content, list) assert len(history_content) == 1 assert history_content[0]["text"] == "compressed history. Retrieve more: hash=b573993006976af767214fac" @@ -2530,15 +2514,6 @@ async def test_history_is_still_compressed(guardrail: HeadroomGuardrail): assert messages[3] == compressed_history[1] -# --------------------------------------------------------------------------- -# #39519: a mid-history row carrying its own Anthropic cache_control marker -# (e.g. a large tool result the client already cached several turns back) was -# still sent to /v1/compress and rewritten. It came back byte-different but -# kept its marker, so the next request's cache read silently became a cache -# write. get_protected_indices() now protects any cache_control-marked row, -# not just system/last-user/last-assistant, so it must never reach the wire. -# --------------------------------------------------------------------------- - CACHE_MARKED_HISTORY_MESSAGES = [ {"role": "system", "content": "You are Claude Code. " + "S" * 5000}, {"role": "user", "content": "old question " + "Q" * 5000}, @@ -2565,7 +2540,6 @@ async def test_mid_history_cache_control_row_is_never_sent_for_compression(guard cached_row = CACHE_MARKED_HISTORY_MESSAGES[3] assert cached_row not in wire assert not any(row.get("tool_call_id") == "old_1" for row in wire) - # Byte-identical, marker intact -- the next request's cache read survives. assert result["structured_messages"][3] == cached_row From 931bdb8c0b50e825d249949b726bafd9f2825443 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 18:19:51 -0700 Subject: [PATCH 3/3] fix(compression): protect part-level cache_control rows in compress() too compress() scores text-only copies of the rows, so a content-part cache_control marker was gone by the time get_protected_indices ran and the pinned row could still be stubbed. Read protection from the original rows, which are index-aligned with the normalized copies, and add a regression test that fails without the change. --- litellm/compression/compress.py | 2 +- .../test_litellm/compression/test_compress.py | 38 +++++++++++++++++-- 2 files changed, 36 insertions(+), 4 deletions(-) diff --git a/litellm/compression/compress.py b/litellm/compression/compress.py index c79e6aed57a..b80f78a50c1 100644 --- a/litellm/compression/compress.py +++ b/litellm/compression/compress.py @@ -435,7 +435,7 @@ def compress( combined_scores = bm25_scores # Protected messages are never compressed - protected_indices: Final = get_protected_indices(normalized_messages) + protected_indices: Final = get_protected_indices(original_messages) kept_indices: set[int] = set(protected_indices) tool_exchange_spans: list[set[int]] = [] diff --git a/tests/test_litellm/compression/test_compress.py b/tests/test_litellm/compression/test_compress.py index f9877ea2bc4..6e908bcbdcd 100644 --- a/tests/test_litellm/compression/test_compress.py +++ b/tests/test_litellm/compression/test_compress.py @@ -6,7 +6,8 @@ never rewrite. It is consumed by compress() and by the Headroom guardrail, so the two agree on what "never compress this" means. """ -from litellm.compression.compress import get_protected_indices +from litellm.compression.compress import compress, get_protected_indices +from litellm.types.utils import CallTypes def test_protects_system_last_user_and_last_assistant(): @@ -66,13 +67,12 @@ def test_mid_history_cache_control_part_is_protected(): { "role": "user", "content": [ - {"type": "text", "text": "a large cached tool result"}, + {"type": "text", "text": "a large cached tool result", "cache_control": {"type": "ephemeral"}}, ], }, {"role": "assistant", "content": "ack"}, {"role": "user", "content": "live instruction"}, ] - messages[2]["content"][0]["cache_control"] = {"type": "ephemeral"} # index 3 = last assistant, index 4 = last user (both protected by role # regardless), index 2 = the cache_control-marked row itself. @@ -113,3 +113,35 @@ def test_content_that_is_not_a_list_of_mappings_is_not_treated_as_cache_control( ] assert sorted(get_protected_indices(messages)) == [0, 2] + + +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", + "content": [ + {"type": "text", "text": "cached tool result " * 2000, "cache_control": {"type": "ephemeral"}}, + ], + } + messages = [ + stale_log, + {"role": "assistant", "content": "old answer"}, + pinned, + {"role": "assistant", "content": "ack"}, + {"role": "user", "content": "live instruction"}, + ] + + result = compress( + messages, + model="gpt-4o", + call_type=CallTypes.anthropic_messages, + compression_trigger=1000, + compression_target=500, + ) + + assert len(result["messages"]) == len(messages) + assert result["messages"][2] == pinned + assert result["messages"][0] != stale_log + assert len(result["cache"]) >= 1