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] 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