mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-17 23:51:30 +00:00
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
This commit is contained in:
parent
5b2b5420af
commit
990dea27d5
3 changed files with 151 additions and 12 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue