diff --git a/tests/agent/test_failover_identity.py b/tests/agent/test_failover_identity.py index b9466137ab02..884fcbb44ca0 100644 --- a/tests/agent/test_failover_identity.py +++ b/tests/agent/test_failover_identity.py @@ -11,7 +11,10 @@ gpt-5.4-mini after a Codex usage-limit 429. from types import SimpleNamespace from agent.chat_completion_helpers import rewrite_prompt_model_identity -from agent.conversation_loop import _sync_failover_system_message +from agent.conversation_loop import ( + _redecorate_prompt_cache_for_provider, + _sync_failover_system_message, +) from agent.prompt_caching import apply_anthropic_cache_control @@ -34,6 +37,40 @@ def _agent(prompt=_PROMPT, ephemeral=None): ) +def _count_cache_markers(messages): + count = 0 + for msg in messages: + if "cache_control" in msg: + count += 1 + content = msg.get("content") + if isinstance(content, list): + for part in content: + if isinstance(part, dict) and "cache_control" in part: + count += 1 + return count + + +def _cache_agent( + *, + use_caching, + native=False, + prompt=_PROMPT, + static=None, + cache_ttl="5m", + provider="openai", +): + return SimpleNamespace( + _cached_system_prompt=prompt, + _cached_system_prompt_static=static, + ephemeral_system_prompt=None, + _use_prompt_caching=use_caching, + _use_native_cache_layout=native, + _cache_ttl=cache_ttl, + provider=provider, + client=None, + ) + + class TestRewritePromptModelIdentity: def test_swaps_identity_lines_to_fallback_runtime(self): agent = _agent() @@ -189,3 +226,153 @@ class TestSyncFailoverPreservesCacheDecoration: ] _sync_failover_system_message(agent, api_messages, _PROMPT) assert api_messages[0]["content"] == agent._cached_system_prompt + + +class TestRedecoratePromptCacheOnPolicyChange: + """#72626: failover must re-render breakpoints for the *new* cache policy. + + ``TestSyncFailoverPreservesCacheDecoration`` only covers same-policy + identity sync (native→native). These cases cover the policy-change class. + """ + + _STATIC = "You are a helpful assistant.\n\nStable brief.\n" + + def test_cache_off_to_cache_on_adds_breakpoints(self): + prompt = self._STATIC + "Model: gpt-5.4-mini\nProvider: openai" + # Primary never decorated (cache-off). + undecorated = [ + {"role": "system", "content": prompt}, + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi"}, + {"role": "user", "content": "again"}, + ] + assert _count_cache_markers(undecorated) == 0 + + agent = _cache_agent( + use_caching=True, + native=True, + prompt=prompt, + static=self._STATIC, + provider="anthropic", + ) + decorated, _ = _redecorate_prompt_cache_for_provider(agent, undecorated) + assert _count_cache_markers(decorated) >= 2 + assert isinstance(decorated[0]["content"], list) + assert decorated[0]["content"][0]["text"] == self._STATIC + + def test_cache_on_to_cache_off_strips_breakpoints(self): + prompt = self._STATIC + "Model: claude\nProvider: anthropic" + decorated = apply_anthropic_cache_control( + [ + {"role": "system", "content": prompt}, + {"role": "user", "content": "hello"}, + ], + native_anthropic=True, + static_system_prefix=self._STATIC, + ) + assert _count_cache_markers(decorated) >= 2 + + agent = _cache_agent( + use_caching=False, + native=False, + prompt=prompt, + static=self._STATIC, + provider="custom", + ) + stripped, _ = _redecorate_prompt_cache_for_provider(agent, decorated) + assert _count_cache_markers(stripped) == 0 + assert stripped[0]["content"] == prompt + + def test_native_to_envelope_relocates_breakpoints_to_carriers(self): + prompt = "Model: claude\nProvider: anthropic" + # Native layout can mark empty assistant turns; envelope cannot. + native = apply_anthropic_cache_control( + [ + {"role": "system", "content": prompt}, + {"role": "user", "content": "do it"}, + {"role": "assistant", "content": "", "tool_calls": [{"id": "1"}]}, + {"role": "tool", "content": "ok", "tool_call_id": "1"}, + {"role": "user", "content": "thanks"}, + ], + native_anthropic=True, + ) + agent = _cache_agent( + use_caching=True, + native=False, + prompt=prompt, + provider="openrouter", + ) + envelope, _ = _redecorate_prompt_cache_for_provider(agent, native) + # Empty assistant must not carry a wasted top-level marker on envelope. + empty_assistant = next( + m for m in envelope if m.get("role") == "assistant" and not m.get("content") + ) + assert "cache_control" not in empty_assistant + assert _count_cache_markers(envelope) >= 1 + + def test_preserves_mutated_tool_result_bytes(self): + # Redecoration must use the in-flight mutated list, not a pristine + # pre-decoration snapshot — image-shrink / ASCII recoveries live here. + prompt = "sys" + messages = [ + {"role": "system", "content": prompt}, + {"role": "user", "content": "run"}, + {"role": "tool", "content": "SHRUNK_IMAGE_PAYLOAD", "tool_call_id": "t1"}, + ] + agent = _cache_agent(use_caching=True, native=True, prompt=prompt) + out, _ = _redecorate_prompt_cache_for_provider(agent, messages) + tool = next(m for m in out if m.get("role") == "tool") + text = ( + tool["content"] + if isinstance(tool["content"], str) + else tool["content"][0]["text"] + ) + assert text == "SHRUNK_IMAGE_PAYLOAD" + + def test_moa_guidance_stays_outside_last_breakpoint(self): + prompt = "sys" + guidance = ( + "[Mixture of Agents context — use this as private guidance for the " + "normal Hermes agent loop.]\nAggregator: agg\n\nadvice" + ) + base = [ + {"role": "system", "content": prompt}, + {"role": "user", "content": "task"}, + {"role": "assistant", "content": "ok"}, + {"role": "user", "content": "task\n\n" + guidance}, + ] + decorated = apply_anthropic_cache_control(base, native_anthropic=True) + + class _Completions: + def rebase_prepared_request(self, prepared, messages): + from agent.moa_loop import _attach_reference_guidance + + out = [dict(m) for m in messages] + _attach_reference_guidance(out, prepared["guidance"]) + return {**prepared, "messages": out} + + agent = _cache_agent(use_caching=True, native=True, prompt=prompt, provider="moa") + agent.client = SimpleNamespace(chat=SimpleNamespace(completions=_Completions())) + prepared = {"guidance": guidance, "messages": decorated} + out, new_prepared = _redecorate_prompt_cache_for_provider( + agent, decorated, moa_prepared=prepared + ) + assert new_prepared is not None + # Guidance must be present, but the cache marker on the last user + # turn's content parts must terminate *before* the guidance text. + last = out[-1] + assert last.get("role") == "user" + content = last["content"] + if isinstance(content, list): + marked = [p for p in content if isinstance(p, dict) and "cache_control" in p] + guidance_parts = [ + p for p in content + if isinstance(p, dict) and guidance in (p.get("text") or "") + ] + assert marked, "base transcript should still carry breakpoints" + assert guidance_parts, "guidance must be re-attached" + # Marked part should not contain the guidance body. + assert all(guidance not in (p.get("text") or "") for p in marked) + else: + assert guidance in content + diff --git a/tests/agent/test_prompt_caching.py b/tests/agent/test_prompt_caching.py index ba21e4c91951..039a6626b7b0 100644 --- a/tests/agent/test_prompt_caching.py +++ b/tests/agent/test_prompt_caching.py @@ -5,6 +5,7 @@ from agent.prompt_caching import ( _apply_cache_marker, _can_carry_marker, apply_anthropic_cache_control, + strip_anthropic_cache_control, ) @@ -317,7 +318,10 @@ class TestNormalizationOrdering: from agent import conversation_loop src = inspect.getsource(conversation_loop) - mark = src.index("apply_anthropic_cache_control(\n") + # Anchor on the call-block decoration (before the retry loop), not the + # mid-failover redecoration helper which also calls apply_*. + anchor = src.index("Runs LAST, after every message mutation above") + mark = src.index("apply_anthropic_cache_control(\n", anchor) for earlier in ( 'am["content"].strip()', # whitespace normalization "_sanitize_api_messages(api_messages)", # orphan sweep @@ -327,3 +331,65 @@ class TestNormalizationOrdering: assert src.index(earlier) < mark, ( f"{earlier!r} must run before cache breakpoints are injected" ) + + +class TestStripAnthropicCacheControl: + """strip must undo decoration so failover can re-render for a new policy.""" + + def test_removes_top_level_and_part_markers(self): + messages = apply_anthropic_cache_control( + [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "yo"}, + ], + native_anthropic=True, + ) + assert any( + "cache_control" in (m if isinstance(m.get("content"), str) else {}) + or ( + isinstance(m.get("content"), list) + and any( + isinstance(p, dict) and "cache_control" in p for p in m["content"] + ) + ) + or "cache_control" in m + for m in messages + ) + strip_anthropic_cache_control(messages) + for msg in messages: + assert "cache_control" not in msg + content = msg.get("content") + if isinstance(content, list): + for part in content: + if isinstance(part, dict): + assert "cache_control" not in part + + def test_flattens_system_static_volatile_back_to_string(self): + static = "You are helpful.\n" + full = static + "Model: claude\nProvider: anthropic" + messages = apply_anthropic_cache_control( + [{"role": "system", "content": full}, {"role": "user", "content": "hi"}], + native_anthropic=True, + static_system_prefix=static, + ) + assert isinstance(messages[0]["content"], list) + strip_anthropic_cache_control(messages) + assert messages[0]["content"] == full + + def test_preserves_multimodal_part_structure(self): + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "see", "cache_control": {"type": "ephemeral"}}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,xx"}}, + ], + } + ] + strip_anthropic_cache_control(messages) + content = messages[0]["content"] + assert isinstance(content, list) and len(content) == 2 + assert content[0] == {"type": "text", "text": "see"} + assert content[1]["type"] == "image_url" +