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.
This commit is contained in:
mateo-berri 2026-09-14 18:19:51 -07:00
parent dad6e9be4a
commit 931bdb8c0b
2 changed files with 36 additions and 4 deletions

View file

@ -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]] = []

View file

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