refactor(compression): build the protected index set without mutation

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(

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