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:
Rad Wadud 2026-09-08 17:52:41 -05:00
parent 5b2b5420af
commit 990dea27d5
3 changed files with 151 additions and 12 deletions

View file

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

View file

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

View file

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