From 4c5be0c295d7d4d0cc04d780edb86b7f3af1bb5b Mon Sep 17 00:00:00 2001 From: joaomarcos Date: Sat, 8 Aug 2026 20:12:07 -0300 Subject: [PATCH] perf(cache): harden the stable-prefix boundary against eviction and memory growth Follow-up review of the builder-declared cache boundary (#81867) found three ways the split could silently stop paying off, or keep paying more than it should, on a long-lived gateway process. Flattening no longer consults the registry. `strip_anthropic_cache_control` matched the decorated split by looking the first block up in the prefix registry, so a mid-turn failover that re-decorates a request built many messages earlier (#72626) would fail to flatten once _MAX_ENTRIES newer scaffolds had been registered in between, and would hand the next provider the two-part shape instead of the canonical string. The split is now matched by its shape: a marker on the *first* part of a user message is something no other decoration produces (list content otherwise gets its marker on the last part, and the two-part [static, volatile] split is role-gated to system), so the ""-join stays provably byte-exact without any process state. This drops `is_registered_stable_prefix` and one lock acquisition per stripped message. Lookups now refresh LRU position. A scaffold fired every minute by cron could be evicted by a burst of one-off skill invocations while still being the hottest prefix in the process, silently reverting it to whole-message caching. Registration now also evicts by total retained bytes (4 MiB). Entries hold whole expanded skill bodies, so a 32-entry cap alone does not bound memory. The newest entry is always kept, so a single oversized scaffold still gets a boundary instead of disabling the split. Tests: eviction-then-failover round-trip, LRU refresh on hit, byte-cap eviction, and oversized-single-entry survival. Co-Authored-By: Claude Fable 5 --- agent/prompt_cache_boundary.py | 23 ++++++--- agent/prompt_caching.py | 33 ++++++++---- tests/agent/test_prompt_cache_boundary.py | 62 ++++++++++++++++++++++- 3 files changed, 99 insertions(+), 19 deletions(-) diff --git a/agent/prompt_cache_boundary.py b/agent/prompt_cache_boundary.py index 49df158aed694..8793306e9a7ce 100644 --- a/agent/prompt_cache_boundary.py +++ b/agent/prompt_cache_boundary.py @@ -30,6 +30,13 @@ from typing import Optional # fall back to whole-message caching rather than growing unboundedly. _MAX_ENTRIES = 32 +# Entries hold whole expanded skill bodies, so an entry count alone does not +# bound memory — a handful of large skills can retain tens of MB in a +# long-lived gateway process. Evict by total retained bytes too, always +# keeping the newest entry so a single oversized scaffold still gets a +# boundary instead of silently disabling the split. +_MAX_BYTES = 4 * 1024 * 1024 + _lock = threading.Lock() _prefixes: "OrderedDict[str, None]" = OrderedDict() @@ -43,6 +50,8 @@ def register_stable_prefix(prefix: str) -> None: _prefixes.move_to_end(prefix) while len(_prefixes) > _MAX_ENTRIES: _prefixes.popitem(last=False) + while len(_prefixes) > 1 and sum(map(len, _prefixes)) > _MAX_BYTES: + _prefixes.popitem(last=False) def find_stable_prefix(content: str) -> Optional[str]: @@ -50,6 +59,10 @@ def find_stable_prefix(content: str) -> Optional[str]: Proper (``len(content) > len(prefix)``) so the split never produces an empty volatile text block, which Anthropic rejects on the wire. + + A hit refreshes the entry's LRU position: a scaffold fired every minute + by cron must not be evicted by a burst of one-off skill invocations, + which would silently drop it back to whole-message caching. """ with _lock: candidates = list(_prefixes) @@ -58,15 +71,13 @@ def find_stable_prefix(content: str) -> Optional[str]: if len(content) > len(prefix) and content.startswith(prefix): if best is None or len(prefix) > len(best): best = prefix + if best is not None: + with _lock: + if best in _prefixes: + _prefixes.move_to_end(best) return best -def is_registered_stable_prefix(text: str) -> bool: - """Exact-match check used when flattening a decorated split back.""" - with _lock: - return text in _prefixes - - def clear_stable_prefixes() -> None: """Test isolation helper.""" with _lock: diff --git a/agent/prompt_caching.py b/agent/prompt_caching.py index 2914e94c87c57..a9488345c6959 100644 --- a/agent/prompt_caching.py +++ b/agent/prompt_caching.py @@ -14,7 +14,7 @@ import copy from dataclasses import dataclass from typing import Any, Dict, List -from agent.prompt_cache_boundary import find_stable_prefix, is_registered_stable_prefix +from agent.prompt_cache_boundary import find_stable_prefix @dataclass(frozen=True) @@ -193,8 +193,11 @@ def strip_anthropic_cache_control( Flattening back to a plain string is restricted to the exact shapes :func:`apply_anthropic_cache_control` produces from string content — - a single ``{"type": "text"}`` part, or the two-part ``[static, volatile]`` - system split — so the ``""``-join is provably byte-exact. Organic + a single ``{"type": "text"}`` part, the two-part ``[static, volatile]`` + system split, or the two-part builder-declared skill split (recognised + by its marker-on-the-first-part shape, so flattening never depends on + the prefix registry still holding the entry) — so the ``""``-join is + provably byte-exact. Organic multi-part text (merged user turns, imported transcripts) and parts carrying extra keys (``citations`` etc.) keep their structure; only per-part markers are removed. Marker removal is copy-on-write on the @@ -213,6 +216,21 @@ def strip_anthropic_cache_control( content = msg.get("content") if not isinstance(content, list): continue + # Two-part skill-invocation split (#81867). The builder-declared + # boundary is the only decoration that marks the *first* part of a + # user message: list content otherwise receives its marker on the + # last part, and the two-part [static, volatile] split is role-gated + # to system. So the shape alone identifies it, and flattening stays + # correct even when the prefix registry has since evicted the entry + # (failover re-decorates a request built many messages ago, #72626). + skill_split_shape = ( + msg.get("role") == "user" + and len(content) == 2 + and isinstance(content[0], dict) + and isinstance(content[1], dict) + and "cache_control" in content[0] + and "cache_control" not in content[1] + ) if any(isinstance(part, dict) and "cache_control" in part for part in content): content = [ {k: v for k, v in part.items() if k != "cache_control"} @@ -230,14 +248,7 @@ def strip_anthropic_cache_control( ) and ( len(content) == 1 or (msg.get("role") == "system" and len(content) == 2) - or ( - # Two-part skill-invocation split: the first part is byte-for- - # byte a builder-registered scaffold, so the ""-join provably - # reconstructs the original string. - msg.get("role") == "user" - and len(content) == 2 - and is_registered_stable_prefix(content[0]["text"]) - ) + or skill_split_shape ) if decoration_shape: msg["content"] = "".join(part["text"] for part in content) diff --git a/tests/agent/test_prompt_cache_boundary.py b/tests/agent/test_prompt_cache_boundary.py index b786f13bdb9e3..0e62b24110258 100644 --- a/tests/agent/test_prompt_cache_boundary.py +++ b/tests/agent/test_prompt_cache_boundary.py @@ -20,7 +20,6 @@ import tools.skills_tool as skills_tool from agent.prompt_cache_boundary import ( clear_stable_prefixes, find_stable_prefix, - is_registered_stable_prefix, register_stable_prefix, ) from agent.prompt_caching import ( @@ -87,7 +86,47 @@ class TestRegistry: def test_empty_prefix_never_registered(self): register_stable_prefix("") - assert not is_registered_stable_prefix("") + assert find_stable_prefix("anything") is None + + def test_lookup_refreshes_the_entry_lru_position(self): + """A cron scaffold fired every minute must survive a burst of one-off + invocations; without the refresh it silently drops back to + whole-message caching while still being the hottest prefix.""" + from agent import prompt_cache_boundary + + register_stable_prefix("hot-scaffold ") + for index in range(prompt_cache_boundary._MAX_ENTRIES - 1): + register_stable_prefix(f"cold-{index} ") + + assert find_stable_prefix("hot-scaffold volatile") == "hot-scaffold " + + register_stable_prefix("newcomer ") + + assert find_stable_prefix("hot-scaffold volatile") == "hot-scaffold " + assert find_stable_prefix("cold-0 volatile") is None + + def test_total_byte_cap_evicts_oldest_and_keeps_newest(self, monkeypatch): + """Entries retain whole skill bodies, so the entry count alone does + not bound memory.""" + from agent import prompt_cache_boundary + + monkeypatch.setattr(prompt_cache_boundary, "_MAX_BYTES", 100) + + register_stable_prefix("a" * 80) + register_stable_prefix("b" * 80) + + assert find_stable_prefix("a" * 80 + "tail") is None + assert find_stable_prefix("b" * 80 + "tail") == "b" * 80 + + def test_single_oversized_prefix_still_registers(self, monkeypatch): + from agent import prompt_cache_boundary + + monkeypatch.setattr(prompt_cache_boundary, "_MAX_BYTES", 10) + oversized = "x" * 500 + + register_stable_prefix(oversized) + + assert find_stable_prefix(oversized + "tail") == oversized class TestRequestLocalSplit: @@ -153,6 +192,25 @@ class TestRequestLocalSplit: assert stripped == original assert apply_anthropic_cache_control(copy.deepcopy(stripped)) == first_wire + 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 + entry, a busy gateway that registered _MAX_ENTRIES newer scaffolds in + between would ship the split shape to the next provider instead of + the canonical string.""" + from agent import prompt_cache_boundary + + scaffold = "stable scaffold\n\n" + _SINGLE_SKILL_INSTRUCTION + register_stable_prefix(scaffold) + original = [{"role": "user", "content": scaffold + "ticket=one"}] + marked = apply_anthropic_cache_control(copy.deepcopy(original)) + + for index in range(prompt_cache_boundary._MAX_ENTRIES + 1): + register_stable_prefix(f"unrelated-scaffold-{index} ") + assert find_stable_prefix(scaffold + "ticket=one") is None + + assert strip_anthropic_cache_control(marked) == original + def test_strip_leaves_organic_two_part_user_content_structured(self): organic = [ {