diff --git a/agent/prompt_cache_boundary.py b/agent/prompt_cache_boundary.py index 8793306e9a7ce..b55ce55303ecc 100644 --- a/agent/prompt_cache_boundary.py +++ b/agent/prompt_cache_boundary.py @@ -19,6 +19,16 @@ The registry is process-local by design. A freshly fired webhook/cron invocation is always built and sent by the same process, which is the only window where the split pays off. Any miss (restart, eviction, historic message) falls back to the pre-existing whole-message policy. + +Split-shape lifetime: the split is applied only while the skill message is +one of the plan's marked endpoints (the last few cacheable messages). Once +later turns rotate it out of that window it ships as a single string block +again, which changes the block boundary once and re-ingests the prefix from +that message onward exactly one time in a long-lived session. Webhook/cron +invocations — the workload this exists for — send the skill turn as the +newest message every time, so they always hit the split shape; the one-time +re-ingest only affects long interactive sessions and nets out far below the +per-invocation full rewrite this removes. """ import threading @@ -32,10 +42,12 @@ _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 +# long-lived gateway process. Evict by total retained characters too (a +# conservative proxy for bytes: actual memory is 1–4x depending on the +# string's widest code point), always keeping the newest entry so a single +# oversized scaffold still gets a boundary instead of silently disabling +# the split. +_MAX_CHARS = 4 * 1024 * 1024 _lock = threading.Lock() _prefixes: "OrderedDict[str, None]" = OrderedDict() @@ -50,7 +62,7 @@ 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: + while len(_prefixes) > 1 and sum(map(len, _prefixes)) > _MAX_CHARS: _prefixes.popitem(last=False) @@ -65,17 +77,15 @@ def find_stable_prefix(content: str) -> Optional[str]: which would silently drop it back to whole-message caching. """ with _lock: - candidates = list(_prefixes) - best: Optional[str] = None - for prefix in candidates: - 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 + best: Optional[str] = None + for prefix in _prefixes: + 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: + # After the scan so the OrderedDict is never mutated mid-iteration. + _prefixes.move_to_end(best) + return best def clear_stable_prefixes() -> None: diff --git a/agent/skill_commands.py b/agent/skill_commands.py index 0809f95bed562..a14aebb1187ef 100644 --- a/agent/skill_commands.py +++ b/agent/skill_commands.py @@ -71,6 +71,23 @@ SKILL_SCAFFOLD_SQL_LIKE = _SKILL_INVOCATION_PREFIX + "%" SKILL_EXCERPT_JOINT = "\x1e" +def append_user_instruction(parts: list, instruction: str) -> str: + """Append the instruction line to ``parts``; return the stable prefix. + + Shared by every builder that ends a static skill scaffold with the + caller-supplied volatile instruction (single-skill invocations, cron job + prompts). The returned prefix ends exactly at the instruction marker, so + registering it with ``agent.prompt_cache_boundary`` lets the Anthropic + cache planner put a breakpoint on the scaffold instead of caching the + whole message as one atomic block (#81867). Keeping construction in one + place guarantees the registered prefix stays a byte-prefix of the built + message — the invariant the request-time split depends on. + """ + stable_prefix = "\n".join(parts) + "\n" + _SINGLE_SKILL_INSTRUCTION + parts.append(f"{_SINGLE_SKILL_INSTRUCTION}{instruction}") + return stable_prefix + + def extract_user_instruction_from_skill_message(content: Any) -> Optional[str]: """Recover the user's instruction from a slash-skill-expanded turn. @@ -370,15 +387,14 @@ def _build_skill_message( # one atomic block (#81867). The static instruction prose stays on # the stable side; the volatile instruction (webhook payload, ticket # IDs, timestamps) and any runtime note ride in the tail. - stable_prefix = "\n".join(parts) + "\n" + _SINGLE_SKILL_INSTRUCTION - parts.append(f"{_SINGLE_SKILL_INSTRUCTION}{user_instruction}") + stable_prefix = append_user_instruction(parts, user_instruction) if runtime_note: parts.append("") parts.append(f"[Runtime note: {runtime_note}]") message = "\n".join(parts) - if stable_prefix is not None and len(message) > len(stable_prefix): + if stable_prefix is not None and message.startswith(stable_prefix) and len(message) > len(stable_prefix): register_stable_prefix(stable_prefix) return message diff --git a/cron/scheduler.py b/cron/scheduler.py index 2812536e2bbed..71a54368d1ec7 100644 --- a/cron/scheduler.py +++ b/cron/scheduler.py @@ -2815,15 +2815,14 @@ def _build_job_prompt( stable_prefix = None if prompt: - from agent.skill_commands import _SINGLE_SKILL_INSTRUCTION + from agent.skill_commands import append_user_instruction parts.append("") # The skill blocks (and any skipped-skill notice) above are stable per # job config; the appended instruction carries the volatile per-run # data (cron hint + prompt + script output + run context). Declare # that boundary for the Anthropic cache planner (#81867). - stable_prefix = "\n".join(parts) + "\n" + _SINGLE_SKILL_INSTRUCTION - parts.append(f"{_SINGLE_SKILL_INSTRUCTION}{prompt}") + stable_prefix = append_user_instruction(parts, prompt) assembled = _scan_assembled_cron_prompt("\n".join(parts), job, has_skills=True) if stable_prefix and len(assembled) > len(stable_prefix) and assembled.startswith(stable_prefix): # Guarded because the injection scanner may sanitize (mutate) the diff --git a/tests/agent/test_prompt_cache_boundary.py b/tests/agent/test_prompt_cache_boundary.py index 0e62b24110258..0449f04ed40ba 100644 --- a/tests/agent/test_prompt_cache_boundary.py +++ b/tests/agent/test_prompt_cache_boundary.py @@ -64,6 +64,19 @@ def skills(tmp_path, monkeypatch): class TestRegistry: + def test_append_user_instruction_prefix_invariant(self): + """The shared builder helper must return a byte-prefix of the final + joined message — the invariant every registration site depends on.""" + from agent.skill_commands import append_user_instruction + + parts = ["scaffold line one", "", "skill body", ""] + instruction = "ticket=42 time=10:00" + stable_prefix = append_user_instruction(parts, instruction) + message = "\n".join(parts) + + assert message.startswith(stable_prefix) + assert message[len(stable_prefix):] == instruction + def test_requires_proper_prefix(self): register_stable_prefix("scaffold") assert find_stable_prefix("scaffold volatile") == "scaffold" @@ -105,12 +118,12 @@ class TestRegistry: 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): + def test_total_char_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) + monkeypatch.setattr(prompt_cache_boundary, "_MAX_CHARS", 100) register_stable_prefix("a" * 80) register_stable_prefix("b" * 80) @@ -121,7 +134,7 @@ class TestRegistry: def test_single_oversized_prefix_still_registers(self, monkeypatch): from agent import prompt_cache_boundary - monkeypatch.setattr(prompt_cache_boundary, "_MAX_BYTES", 10) + monkeypatch.setattr(prompt_cache_boundary, "_MAX_CHARS", 10) oversized = "x" * 500 register_stable_prefix(oversized)