refactor(cache): one scaffold-split builder and one prior-decoration strip pass

Simplify pass over the skill-scaffold split fix:

- _scaffold_parts builds the [scaffold, tail] parts for both the marked split
  (_apply_cache_marker) and the unmarked one (_split_unmarked_scaffolds). The
  stripper only recognises the unmarked split while both stay byte-identical,
  so they now come from one place.
- _strip_prior_decoration owns the "strip every list-content message before
  re-planning" gate. build_prompt_cache_plan no longer strips and then hands
  the result to apply_anthropic_cache_control, which stripped it again.
- strip_anthropic_cache_control checks the registry only for plain-text
  two-part user messages, on the text it already joins for the flatten. The
  old check also ran for non-text parts, where a lookup miss compared
  None == None.
- apply_anthropic_cache_control docstring: drop the duplicated "Idempotent"
  paragraph and the claim that only marked messages pay the copy cost (false
  since the gate widened).
- Tests: cover the direct native tool-cache layout, which plans without
  apply_anthropic_cache_control (red with its _split_unmarked_scaffolds call
  removed). The re-plan test no longer depends on the marker window size.
This commit is contained in:
kshitijk4poor
2026-10-06 13:50:02 +05:30
committed by kshitij
parent a0e066d614
commit 90194c64ff
2 changed files with 61 additions and 51 deletions

View File

@@ -55,21 +55,19 @@ def _apply_cache_marker(msg: dict, cache_marker: dict, native_anthropic: bool =
if not (role in ("tool", "assistant") and not native_anthropic):
msg["cache_control"] = cache_marker
elif isinstance(content, str):
stable_prefix = _scaffold_prefix(msg)
if stable_prefix is not None:
# Builder-declared boundary: the scaffold carries the breakpoint and the volatile
# tail rides unmarked. Request-local only — the stored message stays a string.
msg["content"] = [_text_part(stable_prefix, cache_marker), _text_part(content[len(stable_prefix):])]
else:
msg["content"] = [_text_part(content, cache_marker)]
# Builder-declared boundary: the scaffold carries the breakpoint and the volatile tail
# rides unmarked. Request-local only — the stored message stays a string.
msg["content"] = _scaffold_parts(msg, cache_marker) or [_text_part(content, cache_marker)]
elif isinstance(content, list) and content and isinstance(content[-1], dict):
content[-1]["cache_control"] = cache_marker
def _scaffold_prefix(msg: dict) -> str | None:
"""The registered stable scaffold of a string user message (``find_stable_prefix`` keeps the tail non-blank)."""
def _scaffold_parts(msg: dict, cache_marker: dict | None = None) -> list | None:
"""``[scaffold, tail]`` parts of a string user message with a registered scaffold, the
marker (if any) on the scaffold; ``find_stable_prefix`` keeps the tail non-blank."""
content = msg.get("content")
return find_stable_prefix(content) if msg.get("role") == "user" and isinstance(content, str) else None
prefix = find_stable_prefix(content) if msg.get("role") == "user" and isinstance(content, str) else None
return None if prefix is None else [_text_part(prefix, cache_marker), _text_part(content[len(prefix):])]
def _split_unmarked_scaffolds(messages: List[Dict[str, Any]]) -> None:
@@ -82,9 +80,8 @@ def _split_unmarked_scaffolds(messages: List[Dict[str, Any]]) -> None:
registered and the failover stripper can still recognise its split.
"""
for i, msg in enumerate(messages):
prefix = _scaffold_prefix(msg) if isinstance(msg, dict) else None
if prefix is not None:
messages[i] = {**msg, "content": [_text_part(prefix), _text_part(msg["content"][len(prefix):])]}
if isinstance(msg, dict) and (parts := _scaffold_parts(msg)):
messages[i] = {**msg, "content": parts}
def _can_carry_marker(msg: dict, native_anthropic: bool, tool_part_markers: bool = True) -> bool:
@@ -208,14 +205,10 @@ def strip_anthropic_cache_control(api_messages: List[Dict[str, Any]]) -> List[Di
if not isinstance(content, list):
continue
role = msg.get("role")
# The skill split is the only decoration marking the FIRST part of a user message,
# so the shape alone identifies it even after the prefix registry evicted the entry.
two_parts = role == "user" and len(content) == 2 and all(isinstance(p, dict) for p in content)
skill_split_shape = two_parts and "cache_control" not in content[1] and (
"cache_control" in content[0]
# Unmarked split (outside the marker window): only the registry can tell it from a
# genuine two-part user message.
or find_stable_prefix("".join(str(p.get("text", "")) for p in content)) == content[0].get("text"))
user_pair = role == "user" and len(content) == 2 and all(isinstance(p, dict) for p in content)
# The skill split is the only decoration marking the FIRST part of a user message, so the
# marked shape identifies it even after the prefix registry evicted the entry.
first_part_marked = user_pair and "cache_control" in content[0] and "cache_control" not in content[1]
if _has_part_marker(content):
content = msg["content"] = [
{k: v for k, v in part.items() if k != "cache_control"}
@@ -225,11 +218,31 @@ def strip_anthropic_cache_control(api_messages: List[Dict[str, Any]]) -> List[Di
isinstance(part, dict) and part.get("type", "text") == "text"
and isinstance(part.get("text"), str) and set(part.keys()) <= {"type", "text"} for part in content
)
if plain_text_parts and (len(content) == 1 or (role == "system" and len(content) == 2) or skill_split_shape):
msg["content"] = "".join(part["text"] for part in content)
if not plain_text_parts:
continue
joined = "".join(part["text"] for part in content)
# An UNMARKED split (outside the marker window) only the registry tells from a genuine
# two-part user message.
skill_split = first_part_marked or (user_pair and find_stable_prefix(joined) == content[0]["text"])
if len(content) == 1 or (role == "system" and len(content) == 2) or skill_split:
msg["content"] = joined
return api_messages
def _strip_prior_decoration(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""Re-canonicalise planner output before re-planning (in place on the list, copy-on-write per message).
Every list-content message is stripped, not just marked ones: an UNMARKED skill split (outside
the previous request's marker window) left as two parts would get the new breakpoint on its
volatile tail instead of the scaffold. A shallow copy suffices because the strip is
copy-on-write on content parts.
"""
for i, msg in enumerate(messages):
if isinstance(msg, dict) and ("cache_control" in msg or isinstance(msg.get("content"), list)):
messages[i] = strip_anthropic_cache_control([dict(msg)])[0]
return messages
def strip_anthropic_tool_cache_control(tools: List[Dict[str, Any]] | None) -> List[Dict[str, Any]]:
"""Return copied tools without request-local Anthropic cache markers."""
cleaned = copy.deepcopy(tools or [])
@@ -289,20 +302,15 @@ def build_prompt_cache_plan(
) -> PromptCachePlan:
"""Build copy-on-write cache sections for one resolved request destination
(``tool_part_markers=False`` keeps markers off role:tool parts on LiteLLM-style routes)."""
messages = list(api_messages or [])
for i, msg in enumerate(messages):
if isinstance(msg, dict) and (
"cache_control" in msg or isinstance(msg.get("content"), list)
):
messages[i] = strip_anthropic_cache_control([dict(msg)])[0]
planned_tools = strip_anthropic_tool_cache_control(tools)
if not direct_native_tool_cache or not planned_tools:
# apply_anthropic_cache_control re-canonicalises prior decoration itself.
planned_messages = apply_anthropic_cache_control(
messages, cache_ttl=cache_ttl, native_anthropic=native_anthropic,
list(api_messages or []), cache_ttl=cache_ttl, native_anthropic=native_anthropic,
static_system_prefix=static_system_prefix, tool_part_markers=tool_part_markers)
return PromptCachePlan(messages=planned_messages, tools=planned_tools)
messages = _strip_prior_decoration(list(api_messages or []))
marker = _build_marker(cache_ttl)
if messages and isinstance(messages[0], dict) and messages[0].get("role") == "system":
# Tool-cache layout: only the static prefix carries a system-side marker; the
@@ -327,28 +335,17 @@ def apply_anthropic_cache_control(
With a matching ``static_system_prefix`` the prefix and full system prompt each get a
marker and the remaining two go to the latest cacheable non-system messages; otherwise
the legacy system-and-3 layout applies. Idempotent: pre-existing markers are stripped from
a per-message copy first. Returns a shallow list copy with deep copies of modified messages.
Idempotent: pre-existing ``cache_control`` markers are stripped from a per-message copy before new ones
are placed, so calling this twice (or handing it messages a prior call already marked) can never
accumulate past 4 markers. Only messages that already carry a marker pay the copy cost — a shallow
top-level copy suffices because :func:`strip_anthropic_cache_control` is copy-on-write on content parts
— and the rest of the copy-on-write contract is unchanged (#90971).
the legacy system-and-3 layout applies. Idempotent: prior decoration is stripped from a
per-message copy first, so a second call (or messages a prior call already planned) can never
accumulate past 4 markers (#90971). Returns a shallow list copy with deep copies of modified
messages.
"""
if not api_messages:
return api_messages
messages = list(api_messages)
messages = _strip_prior_decoration(list(api_messages))
marker = _build_marker(cache_ttl)
for i, msg in enumerate(messages):
# Same gate as build_prompt_cache_plan: an UNMARKED skill split (outside the marker
# window on the previous request) carries no marker, and left as two parts the new
# breakpoint would land on its volatile tail instead of the scaffold.
if isinstance(msg, dict) and ("cache_control" in msg or isinstance(msg.get("content"), list)):
messages[i] = strip_anthropic_cache_control([dict(msg)])[0]
breakpoints_used = 0
if messages[0].get("role") == "system":
messages[0] = copy.deepcopy(messages[0])

View File

@@ -23,7 +23,6 @@ from agent.prompt_cache_boundary import (
register_stable_prefix,
)
from agent.prompt_caching import (
_has_part_marker,
apply_anthropic_cache_control,
build_prompt_cache_plan,
strip_anthropic_cache_control,
@@ -225,13 +224,27 @@ class TestRequestLocalSplit:
original += [{"role": role, "content": f"turn {i}"} for i in range(4) for role in ("assistant", "user")]
unmarked = apply_anthropic_cache_control(copy.deepcopy(original))
assert not _has_part_marker(unmarked[0]["content"])
assert all("cache_control" not in part for part in unmarked[0]["content"])
again = apply_anthropic_cache_control(copy.deepcopy(unmarked[:3]))
assert again == apply_anthropic_cache_control(copy.deepcopy(original[:3]))
again = apply_anthropic_cache_control(copy.deepcopy(unmarked[:1]))
assert again == apply_anthropic_cache_control(copy.deepcopy(original[:1]))
assert again[0]["content"][0]["cache_control"] == MARKER
assert "cache_control" not in again[0]["content"][1]
def test_direct_tool_cache_plan_keeps_the_split_outside_the_marker_window(self):
"""The native tool-cache layout plans without apply_anthropic_cache_control, so it needs
the same unmarked split or its third request re-writes the scaffold."""
scaffold = "stable scaffold\n\n" + _SINGLE_SKILL_INSTRUCTION
register_stable_prefix(scaffold)
history = [{"role": "user", "content": scaffold + "ticket=one"}]
history += [{"role": role, "content": f"turn {i}"} for i in range(4) for role in ("assistant", "user")]
tools = [{"type": "function", "function": {"name": "terminal", "parameters": {}}}]
plan = build_prompt_cache_plan(history, tools, native_anthropic=True, direct_native_tool_cache=True)
assert plan.messages[0]["content"] == [{"type": "text", "text": scaffold}, {"type": "text", "text": "ticket=one"}]
assert isinstance(history[0]["content"], str)
def test_strip_flattens_even_after_the_prefix_was_evicted(self):
"""Mid-turn failover re-decorates a request built many messages ago
(#72626). If flattening depended on the registry still holding the