diff --git a/agent/prompt_cache_boundary.py b/agent/prompt_cache_boundary.py new file mode 100644 index 0000000000000..49df158aed694 --- /dev/null +++ b/agent/prompt_cache_boundary.py @@ -0,0 +1,73 @@ +"""Builder-declared stable prefixes for Anthropic prompt caching (#81867). + +Skill, webhook, and cron builders concatenate a large static scaffold +(activation note + expanded skill body) with a small volatile invocation +tail (ticket payload, timestamps, run context) into one user-message +string. Only the builder knows the exact byte where the volatile tail +begins, so it registers the stable prefix here at construction time; the +cache planner consults the registry to place a cache breakpoint at that +boundary instead of caching the whole message as one atomic block. + +This deliberately avoids re-parsing scaffold marker strings out of the +message at request time: markers can legitimately appear inside skill +bodies or inside event payloads (e.g. a helpdesk ticket quoting an agent +transcript), and any delimiter-search heuristic then either shrinks the +cached prefix or — worse — silently absorbs volatile bytes into it, +reintroducing the per-invocation cache miss this exists to fix. + +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. +""" + +import threading +from collections import OrderedDict +from typing import Optional + +# A couple dozen distinct active scaffolds (webhook routes x skills x cron +# jobs) is generous for one gateway process; beyond that, oldest entries +# fall back to whole-message caching rather than growing unboundedly. +_MAX_ENTRIES = 32 + +_lock = threading.Lock() +_prefixes: "OrderedDict[str, None]" = OrderedDict() + + +def register_stable_prefix(prefix: str) -> None: + """Record ``prefix`` as the stable scaffold of a just-built message.""" + if not prefix: + return + with _lock: + _prefixes[prefix] = None + _prefixes.move_to_end(prefix) + while len(_prefixes) > _MAX_ENTRIES: + _prefixes.popitem(last=False) + + +def find_stable_prefix(content: str) -> Optional[str]: + """Longest registered prefix that is a *proper* prefix of ``content``. + + Proper (``len(content) > len(prefix)``) so the split never produces an + empty volatile text block, which Anthropic rejects on the wire. + """ + 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 + 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: + _prefixes.clear() diff --git a/agent/prompt_caching.py b/agent/prompt_caching.py index 2dd14c9f1fc2c..2914e94c87c57 100644 --- a/agent/prompt_caching.py +++ b/agent/prompt_caching.py @@ -14,6 +14,8 @@ 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 + @dataclass(frozen=True) class PromptCachePlan: @@ -58,6 +60,23 @@ def _apply_cache_marker(msg: dict, cache_marker: dict, native_anthropic: bool = return if isinstance(content, str): + if role == "user": + stable_prefix = find_stable_prefix(content) + if stable_prefix is not None: + # Builder-declared boundary (#81867): the scaffold carries the + # breakpoint, the volatile invocation tail rides unmarked so a + # changed ticket ID or timestamp no longer invalidates the + # whole skill body. Request-local only — the canonical session + # message stays a plain string. + msg["content"] = [ + { + "type": "text", + "text": stable_prefix, + "cache_control": cache_marker, + }, + {"type": "text", "text": content[len(stable_prefix):]}, + ] + return msg["content"] = [ {"type": "text", "text": content, "cache_control": cache_marker} ] @@ -211,6 +230,14 @@ 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"]) + ) ) if decoration_shape: msg["content"] = "".join(part["text"] for part in content) diff --git a/agent/skill_commands.py b/agent/skill_commands.py index eb0d646fecfb8..0809f95bed562 100644 --- a/agent/skill_commands.py +++ b/agent/skill_commands.py @@ -12,6 +12,7 @@ from pathlib import Path from typing import Any, Dict, Optional from hermes_constants import display_hermes_home +from agent.prompt_cache_boundary import register_stable_prefix from agent.skill_preprocessing import ( expand_inline_shell as _expand_inline_shell, load_skills_config as _load_skills_config, @@ -360,15 +361,26 @@ def _build_skill_message( f"(e.g. `node {skill_dir}/scripts/foo.js`)." ) + stable_prefix = None if user_instruction: parts.append("") - parts.append(f"The user has provided the following instruction alongside the skill invocation: {user_instruction}") + # Everything before the caller-supplied instruction is a stable + # scaffold; declare the exact boundary so the Anthropic cache planner + # can put a breakpoint on it instead of caching the whole message as + # 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}") if runtime_note: parts.append("") parts.append(f"[Runtime note: {runtime_note}]") - return "\n".join(parts) + message = "\n".join(parts) + if stable_prefix is not None and len(message) > len(stable_prefix): + register_stable_prefix(stable_prefix) + return message def scan_skill_commands() -> Dict[str, Dict[str, Any]]: diff --git a/cron/scheduler.py b/cron/scheduler.py index c0b56f7671b3e..2812536e2bbed 100644 --- a/cron/scheduler.py +++ b/cron/scheduler.py @@ -2813,9 +2813,26 @@ def _build_job_prompt( ) parts.insert(0, notice) + stable_prefix = None if prompt: - parts.extend(["", f"The user has provided the following instruction alongside the skill invocation: {prompt}"]) - return _scan_assembled_cron_prompt("\n".join(parts), job, has_skills=True) + from agent.skill_commands import _SINGLE_SKILL_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}") + 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 + # assembled bytes; a mismatch simply falls back to whole-message + # caching. + from agent.prompt_cache_boundary import register_stable_prefix + + register_stable_prefix(stable_prefix) + return assembled def _scan_assembled_cron_prompt( diff --git a/tests/agent/test_prompt_cache_boundary.py b/tests/agent/test_prompt_cache_boundary.py new file mode 100644 index 0000000000000..b786f13bdb9e3 --- /dev/null +++ b/tests/agent/test_prompt_cache_boundary.py @@ -0,0 +1,309 @@ +"""Builder-declared stable-prefix cache boundaries (#81867). + +Webhook/cron skill invocations concatenate a large static scaffold with a +small volatile tail into one user string; without a declared boundary the +whole message is cached as one atomic block and a changed ticket ID or +timestamp forces a full cache rewrite. These tests cover the registry, the +request-local split, the failover round-trip, and — via the real builders — +that the boundary comes from construction, not from re-parsing marker +strings (so payloads that quote the marker cannot poison the stable prefix). +""" + +import copy + +import pytest +from unittest.mock import patch + +import agent.skill_bundles as skill_bundles +import agent.skill_commands as skill_commands +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 ( + apply_anthropic_cache_control, + build_prompt_cache_plan, + strip_anthropic_cache_control, +) +from agent.skill_commands import _SINGLE_SKILL_INSTRUCTION + +MARKER = {"type": "ephemeral"} + +SKILL_BODY = "Inspect the report carefully and preserve the stable instructions." + + +@pytest.fixture(autouse=True) +def _isolated_registry(): + clear_stable_prefixes() + yield + clear_stable_prefixes() + + +def _write_skill(skills_dir, name, body=SKILL_BODY): + skill_dir = skills_dir / name + skill_dir.mkdir(parents=True, exist_ok=True) + (skill_dir / "SKILL.md").write_text( + f"---\nname: {name}\ndescription: Description for {name}\n---\n\n# {name}\n\n{body}\n" + ) + return skill_dir + + +@pytest.fixture() +def skills(tmp_path, monkeypatch): + skills_dir = tmp_path / "skills" + _write_skill(skills_dir, "triage") + monkeypatch.setattr(skills_tool, "SKILLS_DIR", skills_dir) + monkeypatch.setattr(skill_commands, "_skill_commands", {}) + monkeypatch.setattr(skill_commands, "_skill_commands_platform", None) + monkeypatch.setattr(skill_bundles, "_bundles_cache", {}) + monkeypatch.setattr(skill_bundles, "_bundles_cache_mtime", None) + skill_commands.scan_skill_commands() + return skills_dir + + +class TestRegistry: + def test_requires_proper_prefix(self): + register_stable_prefix("scaffold") + assert find_stable_prefix("scaffold volatile") == "scaffold" + # Exact match would leave an empty volatile block — never split. + assert find_stable_prefix("scaffold") is None + assert find_stable_prefix("other") is None + + def test_longest_registered_prefix_wins(self): + register_stable_prefix("scaffold") + register_stable_prefix("scaffold extended") + assert find_stable_prefix("scaffold extended tail") == "scaffold extended" + + def test_lru_eviction_falls_back_safely(self): + from agent import prompt_cache_boundary + + for index in range(prompt_cache_boundary._MAX_ENTRIES + 1): + register_stable_prefix(f"scaffold-{index} ") + assert find_stable_prefix("scaffold-0 volatile") is None + assert find_stable_prefix("scaffold-1 volatile") == "scaffold-1 " + + def test_empty_prefix_never_registered(self): + register_stable_prefix("") + assert not is_registered_stable_prefix("") + + +class TestRequestLocalSplit: + def test_volatile_tail_does_not_change_marked_prefix(self): + scaffold = "stable skill scaffold\n\n" + _SINGLE_SKILL_INSTRUCTION + register_stable_prefix(scaffold) + + first = apply_anthropic_cache_control( + [{"role": "user", "content": scaffold + "ticket=one time=10:00"}] + )[0]["content"] + second = apply_anthropic_cache_control( + [{"role": "user", "content": scaffold + "ticket=two time=10:01"}] + )[0]["content"] + + assert len(first) == len(second) == 2 + assert first[0] == second[0] + assert first[0] == {"type": "text", "text": scaffold, "cache_control": MARKER} + assert "cache_control" not in first[1] + assert first[1]["text"] == "ticket=one time=10:00" + assert second[1]["text"] == "ticket=two time=10:01" + + def test_unregistered_message_keeps_whole_block_layout(self): + content = "stable-looking scaffold\n\n" + _SINGLE_SKILL_INSTRUCTION + "tail" + marked = apply_anthropic_cache_control([{"role": "user", "content": content}]) + assert marked[0]["content"] == [ + {"type": "text", "text": content, "cache_control": MARKER} + ] + + def test_assistant_string_matching_a_prefix_is_not_split(self): + register_stable_prefix("scaffold ") + marked = apply_anthropic_cache_control( + [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "scaffold reply"}, + ] + ) + assert marked[1]["content"] == [ + {"type": "text", "text": "scaffold reply", "cache_control": MARKER} + ] + + def test_canonical_message_stays_a_string_in_plan(self): + scaffold = "stable scaffold " + register_stable_prefix(scaffold) + original = scaffold + "volatile" + messages = [{"role": "user", "content": original}] + + plan = build_prompt_cache_plan(messages, [], native_anthropic=True) + + assert messages == [{"role": "user", "content": original}] + assert isinstance(plan.messages[0]["content"], list) + assert plan.messages[0]["content"][0]["cache_control"] == MARKER + assert "cache_control" not in plan.messages[0]["content"][1] + + def test_strip_reconstructs_exact_string_and_redecorates_identically(self): + 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)) + first_wire = copy.deepcopy(marked) + stripped = strip_anthropic_cache_control(marked) + + assert stripped == original + assert apply_anthropic_cache_control(copy.deepcopy(stripped)) == first_wire + + def test_strip_leaves_organic_two_part_user_content_structured(self): + organic = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "part one"}, + {"type": "text", "text": "part two"}, + ], + } + ] + stripped = strip_anthropic_cache_control(copy.deepcopy(organic)) + assert stripped == organic + + +class TestRealBuilders: + def test_two_invocations_share_the_marked_scaffold(self, skills): + first = skill_commands.build_skill_invocation_message( + "/triage", user_instruction="ticket=one time=10:00" + ) + second = skill_commands.build_skill_invocation_message( + "/triage", user_instruction="ticket=two time=10:01" + ) + + first_blocks = apply_anthropic_cache_control( + [{"role": "user", "content": first}], native_anthropic=True + )[0]["content"] + second_blocks = apply_anthropic_cache_control( + [{"role": "user", "content": second}], native_anthropic=True + )[0]["content"] + + assert len(first_blocks) == len(second_blocks) == 2 + assert first_blocks[0] == second_blocks[0] + assert first_blocks[0]["cache_control"] == MARKER + assert SKILL_BODY in first_blocks[0]["text"] + assert "ticket=one" not in first_blocks[0]["text"] + assert "cache_control" not in first_blocks[1] + assert first_blocks[1]["text"] == "ticket=one time=10:00" + + def test_payload_quoting_the_marker_cannot_poison_the_stable_prefix(self, skills): + """The differentiator vs marker-search heuristics: a ticket that quotes + the instruction marker (e.g. a pasted agent transcript) must stay + entirely in the volatile tail, or two invocations stop sharing the + cached prefix.""" + hostile = ( + "ticket=one\n\n" + _SINGLE_SKILL_INSTRUCTION + "quoted transcript line" + ) + benign = "ticket=two" + + hostile_blocks = apply_anthropic_cache_control( + [ + { + "role": "user", + "content": skill_commands.build_skill_invocation_message( + "/triage", user_instruction=hostile + ), + } + ] + )[0]["content"] + benign_blocks = apply_anthropic_cache_control( + [ + { + "role": "user", + "content": skill_commands.build_skill_invocation_message( + "/triage", user_instruction=benign + ), + } + ] + )[0]["content"] + + assert hostile_blocks[0] == benign_blocks[0] + assert hostile_blocks[1]["text"] == hostile + assert benign_blocks[1]["text"] == benign + + def test_bare_invocation_keeps_whole_block_layout(self, skills): + message = skill_commands.build_skill_invocation_message("/triage") + blocks = apply_anthropic_cache_control( + [{"role": "user", "content": message}] + )[0]["content"] + assert blocks == [{"type": "text", "text": message, "cache_control": MARKER}] + + def test_builder_round_trip_survives_failover_strip(self, skills): + original = [ + { + "role": "user", + "content": skill_commands.build_skill_invocation_message( + "/triage", user_instruction="ticket=one" + ), + } + ] + marked = apply_anthropic_cache_control(copy.deepcopy(original)) + assert strip_anthropic_cache_control(marked) == original + + +class TestCronBuilder: + def _prompts(self, jobs): + from cron.scheduler import _build_job_prompt + + with patch( + "agent.skill_bundles.resolve_bundle_command_key", return_value=None + ), patch( + "tools.skills_tool.skill_view", side_effect=self._skill_view + ), patch( + "tools.skill_usage.bump_use" + ): + return [_build_job_prompt(job) for job in jobs] + + @staticmethod + def _skill_view(name: str) -> str: + import json + + if name == "missing": + return json.dumps({"success": False, "error": "missing"}) + return json.dumps({"success": True, "content": f"Stable content for {name}."}) + + def test_cron_runs_share_the_marked_scaffold(self): + common = {"id": "job-cache", "name": "cache boundary", "skills": ["alpha", "beta"]} + first, second = self._prompts( + [ + {**common, "prompt": "ticket=one time=10:00"}, + {**common, "prompt": "ticket=two time=10:01"}, + ] + ) + + first_blocks = apply_anthropic_cache_control( + [{"role": "user", "content": first}] + )[0]["content"] + second_blocks = apply_anthropic_cache_control( + [{"role": "user", "content": second}] + )[0]["content"] + + assert isinstance(first, str) and isinstance(second, str) + assert len(first_blocks) == len(second_blocks) == 2 + assert first_blocks[0] == second_blocks[0] + assert first_blocks[0]["cache_control"] == MARKER + assert "Stable content for alpha." in first_blocks[0]["text"] + assert "ticket=one" not in first_blocks[0]["text"] + assert first_blocks[1]["text"].endswith("ticket=one time=10:00") + + def test_missing_skill_notice_stays_in_the_stable_prefix(self): + common = {"id": "job-skip", "name": "skip notice", "skills": ["missing", "alpha"]} + first, second = self._prompts( + [{**common, "prompt": "ticket=one"}, {**common, "prompt": "ticket=two"}] + ) + + first_blocks = apply_anthropic_cache_control( + [{"role": "user", "content": first}] + )[0]["content"] + second_blocks = apply_anthropic_cache_control( + [{"role": "user", "content": second}] + )[0]["content"] + + assert first_blocks[0] == second_blocks[0] + assert "could not be found" in first_blocks[0]["text"] + assert first_blocks[1]["text"].endswith("ticket=one")