From 0fc52b055f08a211b42508b427f68677e5a596db Mon Sep 17 00:00:00 2001 From: joaomarcos Date: Fri, 21 Aug 2026 12:04:10 +0530 Subject: [PATCH] fix(prompt_caching): make apply_anthropic_cache_control idempotent on pre-decorated input apply_anthropic_cache_control never stripped pre-existing cache_control markers before placing new ones, so calling it twice (or handing it messages a prior call already marked) accumulated markers past Anthropic's 4-breakpoint limit and produced HTTP 400 'cache_control can only be specified up to 4 times'. Strip any pre-existing markers from per-message copies before marking, mirroring the strip-then-mark pattern build_prompt_cache_plan already uses. Only messages that already carry a marker pay the copy cost; the copy-on-write contract (caller-owned messages are never mutated) is preserved. Repeated calls now converge to byte-identical output. Salvaged from #90972 by @JoaoMarcos44 (net diff of the PR's commit stack, intermediate reverts collapsed). Related: #90971 --- agent/prompt_caching.py | 17 +++ .../test_prompt_caching_idempotency.py | 121 ++++++++++++++++++ 2 files changed, 138 insertions(+) create mode 100644 tests/run_agent/test_prompt_caching_idempotency.py diff --git a/agent/prompt_caching.py b/agent/prompt_caching.py index 8643ad533e..433da9ae66 100644 --- a/agent/prompt_caching.py +++ b/agent/prompt_caching.py @@ -445,6 +445,12 @@ def apply_anthropic_cache_control( non-system messages. Without that prefix, the legacy system-and-3 layout is retained. + Idempotent: pre-existing ``cache_control`` markers are stripped from a + deep 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 deep-copy + cost; the rest of the copy-on-write contract is unchanged (#90971). + Returns: Shallow copy of message list with selective deep copies of modified messages. """ @@ -454,6 +460,17 @@ def apply_anthropic_cache_control( messages = list(api_messages) marker = _build_marker(cache_ttl) + for i, msg in enumerate(messages): + if not isinstance(msg, dict): + continue + content = msg.get("content") + has_marker = "cache_control" in msg or ( + isinstance(content, list) + and any(isinstance(part, dict) and "cache_control" in part for part in content) + ) + if has_marker: + messages[i] = strip_anthropic_cache_control([copy.deepcopy(msg)])[0] + breakpoints_used = 0 if messages[0].get("role") == "system": diff --git a/tests/run_agent/test_prompt_caching_idempotency.py b/tests/run_agent/test_prompt_caching_idempotency.py new file mode 100644 index 0000000000..b04f954a2b --- /dev/null +++ b/tests/run_agent/test_prompt_caching_idempotency.py @@ -0,0 +1,121 @@ +"""Tests for Anthropic prompt caching idempotency and breakpoint bounds (#90971).""" + +import copy + +from agent.prompt_caching import ( + apply_anthropic_cache_control, + build_prompt_cache_plan, + _count_cache_markers, + _can_carry_marker, +) + + +class TestPromptCachingIdempotency: + def test_apply_anthropic_cache_control_empty_messages(self): + """Empty messages list is a safe no-op.""" + assert apply_anthropic_cache_control([]) == [] + + def test_apply_anthropic_cache_control_never_exceeds_four_markers(self): + """Realistic conversations never push breakpoints_used past what + _apply_system_cache_markers can return ({0, 1, 2}), so the marker + total always stays within the 4-breakpoint API limit. + """ + messages = [{"role": "system", "content": "STATIC_PREFIX rest of the prompt"}] + for i in range(8): + messages.append({"role": "user", "content": f"Hello {i}"}) + messages.append({"role": "assistant", "content": f"Hi {i}"}) + + result = apply_anthropic_cache_control(messages, static_system_prefix="STATIC_PREFIX") + assert _count_cache_markers(result, []) <= 4 + + def test_apply_anthropic_cache_control_is_idempotent(self): + """Calling apply_anthropic_cache_control repeatedly on its own output + (no intervening strip_anthropic_cache_control) must converge to the + exact same marker placement, not merely stay under budget: a test + that only checks `<= 4` would still pass if a later round moved the + breakpoints somewhere else, or dropped every marker. Before the + idempotency fix, a second call on already-marked messages pushed the + total to 5, reproducing the `cache_control can only be specified up + to 4 times` HTTP 400 (#90971). + """ + messages = [{"role": "system", "content": "STATIC_PREFIX rest of the prompt"}] + for i in range(8): + messages.append({"role": "user", "content": f"Hello {i}"}) + messages.append({"role": "assistant", "content": f"Hi {i}"}) + + round1 = apply_anthropic_cache_control(messages, static_system_prefix="STATIC_PREFIX") + round2 = apply_anthropic_cache_control(round1, static_system_prefix="STATIC_PREFIX") + round3 = apply_anthropic_cache_control(round2, static_system_prefix="STATIC_PREFIX") + + assert round1 == round2 == round3 + assert _count_cache_markers(round1, []) <= 4 + + def test_apply_anthropic_cache_control_does_not_mutate_caller_messages(self): + """A caller's live message list must never be mutated in place, even + when it already carries stale cache_control markers (e.g. replayed + history). The function's contract is copy-on-write. + """ + caller_history = [ + {"role": "user", "content": f"u{i}", "cache_control": {"type": "ephemeral"}} + for i in range(5) + ] + snapshot = copy.deepcopy(caller_history) + + apply_anthropic_cache_control(caller_history) + + assert caller_history == snapshot + + def test_build_prompt_cache_plan_dynamic_tool_accounting(self): + """build_prompt_cache_plan never exceeds 4 markers with tool-cache layout.""" + tools = [ + {"type": "function", "function": {"name": "tool_a"}}, + {"type": "function", "function": {"name": "tool_b"}}, + {"type": "function", "function": {"name": "tool_c"}}, + ] + messages = [ + {"role": "system", "content": "PREFIX_STATIC System prompt"}, + {"role": "user", "content": "Run tool"}, + {"role": "assistant", "content": "Calling", "tool_calls": [{"name": "tool_a"}]}, + {"role": "tool", "content": "output", "tool_name": "tool_a"}, + {"role": "assistant", "content": "Done!"}, + ] + + plan = build_prompt_cache_plan( + messages, + tools, + native_anthropic=True, + direct_native_tool_cache=True, + static_system_prefix="PREFIX_STATIC", + ) + + assert plan.marker_count <= 4 + # Exactly 1 tool marker on the last tool + assert "cache_control" in plan.tools[-1] + assert "cache_control" not in plan.tools[0] + assert "cache_control" not in plan.tools[1] + + def test_build_prompt_cache_plan_direct_tool_cache_with_no_tools(self): + """When direct_native_tool_cache=True but tools is empty, falls back safely.""" + messages = [ + {"role": "system", "content": "System prompt"}, + {"role": "user", "content": "Question"}, + {"role": "assistant", "content": "Answer"}, + ] + plan = build_prompt_cache_plan( + messages, + [], + native_anthropic=True, + direct_native_tool_cache=True, + ) + assert plan.marker_count <= 4 + assert len(plan.tools) == 0 + + def test_can_carry_marker_envelope_vs_native(self): + """_can_carry_marker properly filters empty turns on non-native layouts.""" + empty_assistant = {"role": "assistant", "content": None} + assert _can_carry_marker(empty_assistant, native_anthropic=False) is False + assert _can_carry_marker(empty_assistant, native_anthropic=True) is True + + normal_user = {"role": "user", "content": "Hello"} + assert _can_carry_marker(normal_user, native_anthropic=False) is True + assert _can_carry_marker(normal_user, native_anthropic=True) is True