From 3c4220cd9cfc7a7e55852453d16ebf7039c2504e Mon Sep 17 00:00:00 2001 From: dsad <65560494+necoweb3@users.noreply.github.com> Date: Sun, 26 Jul 2026 19:10:05 +0000 Subject: [PATCH] fix(agent): keep system cache breakpoints across provider failover `apply_anthropic_cache_control` runs once per call block, before the retry loop, and splits the system prompt into `[static prefix, volatile tail]` text blocks carrying the cache_control breakpoints. A failover fires *inside* that retry loop, and `_sync_failover_system_message` assigns a bare string over `api_messages[0]["content"]` to refresh the `Model:`/`Provider:` identity lines. That drops the block list and both breakpoints. `convert_messages_to_anthropic` only emits system cache blocks for list content (`isinstance(content, list)`), so the retried request ships `system` as one plain string: zero breakpoints, nothing written to cache, and the whole system prompt re-billed at full write price. The next call misses too, so a failover costs two full-price prompts instead of one write + one read. Measured with the real functions, same history, native Anthropic layout: normal turn system=BLOCK LIST(2) system breakpoints=2 after failover system=plain string system breakpoints=0 `_sync_failover_system_message`'s own docstring notes this fires on "every gateway turn, since fallback re-activates per message while the primary is down" -- so a gateway running on a degraded primary pays it on every message. Fix: rewrite the decorated blocks in place instead of flattening them. `rewrite_prompt_model_identity` only touches the LAST `Model:`/`Provider:` lines, and those live in the volatile tail, so the static prefix stays byte-identical and its cache entry keeps matching -- the failover retry keeps both breakpoints AND the warm prefix. Shapes we cannot safely patch fall back to the existing plain-string assignment. This is the same class the retry loop already guards against eight lines below the gap, for reasoning fields: # api_messages is built once, before this retry loop, while the primary # provider is active. [...] so the fallback request isn't sent with stale, # primary-shaped reasoning fields. agent._reapply_reasoning_echo_for_provider(api_messages) There was no equivalent for cache decoration. --- agent/conversation_loop.py | 38 +++++++++++- tests/agent/test_failover_identity.py | 87 +++++++++++++++++++++++++++ 2 files changed, 124 insertions(+), 1 deletion(-) diff --git a/agent/conversation_loop.py b/agent/conversation_loop.py index a73f08fb24cf..1f72873edc10 100644 --- a/agent/conversation_loop.py +++ b/agent/conversation_loop.py @@ -774,6 +774,41 @@ def _compression_deferred_result( } +def _rewrite_system_content_blocks(system_message: dict, effective: str) -> bool: + """Rewrite a cache-decorated system message in place, keeping its blocks. + + ``apply_anthropic_cache_control`` runs once per call block, *before* the + retry loop, and splits the system prompt into ``[static prefix, volatile + tail]`` text blocks carrying the cache_control breakpoints. Assigning a bare + string over that list drops both breakpoints, so the failover retry ships + the whole system prompt uncached and re-bills it in full. + + ``rewrite_prompt_model_identity`` only touches the LAST ``Model:`` / + ``Provider:`` lines, and those live in the volatile tail — so the static + prefix stays byte-identical and its cache entry keeps matching. Returns + False when the shape is not one we can safely patch, so the caller falls + back to the plain-string assignment. + """ + content = system_message.get("content") + if not isinstance(content, list) or not content: + return False + if not all( + isinstance(part, dict) and part.get("type") == "text" for part in content + ): + return False + if len(content) == 1: + content[0]["text"] = effective + return True + if len(content) == 2: + head = content[0].get("text") or "" + if head and effective.startswith(head): + tail = effective[len(head):] + if tail: + content[1]["text"] = tail + return True + return False + + def _sync_failover_system_message(agent, api_messages, active_system_prompt): """Refresh the in-flight system message after a provider failover. @@ -796,7 +831,8 @@ def _sync_failover_system_message(agent, api_messages, active_system_prompt): effective = sp if agent.ephemeral_system_prompt: effective = (effective + "\n\n" + agent.ephemeral_system_prompt).strip() - api_messages[0]["content"] = effective + if not _rewrite_system_content_blocks(api_messages[0], effective): + api_messages[0]["content"] = effective return sp diff --git a/tests/agent/test_failover_identity.py b/tests/agent/test_failover_identity.py index 1937da6b6439..b9466137ab02 100644 --- a/tests/agent/test_failover_identity.py +++ b/tests/agent/test_failover_identity.py @@ -12,6 +12,7 @@ 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.prompt_caching import apply_anthropic_cache_control _PROMPT = ( @@ -102,3 +103,89 @@ class TestSyncFailoverSystemMessage: assert api_messages == [{"role": "user", "content": "hi"}] # Still returns the cached prompt for subsequent call-block rebuilds. assert result == agent._cached_system_prompt + + +class TestSyncFailoverPreservesCacheDecoration: + """The sync must not flatten a cache-decorated system message. + + ``apply_anthropic_cache_control`` runs once per call block, before the retry + loop, splitting the system prompt into ``[static prefix, volatile tail]`` + blocks that carry the cache_control breakpoints. A failover fires *inside* + that retry loop, so overwriting the list with a bare string drops both + breakpoints and the retried request re-bills the whole system prompt. + """ + + _STATIC = "You are a helpful assistant.\n\nStable brief.\n" + + def _decorated(self, prompt): + messages = [ + {"role": "system", "content": prompt}, + {"role": "user", "content": "what model are you?"}, + ] + return apply_anthropic_cache_control( + messages, + cache_ttl=None, + native_anthropic=True, + static_system_prefix=self._STATIC, + ) + + def test_keeps_breakpoints_and_static_prefix(self): + prompt = self._STATIC + "Model: gpt-5.4-mini\nProvider: openai-codex" + agent = _agent(prompt=prompt) + api_messages = self._decorated(prompt) + assert isinstance(api_messages[0]["content"], list) + + rewrite_prompt_model_identity(agent, "gemma4:e2b-mlx", "custom") + _sync_failover_system_message(agent, api_messages, prompt) + + content = api_messages[0]["content"] + assert isinstance(content, list), "cache decoration was flattened to a string" + assert len(content) == 2 + assert all(part.get("cache_control") for part in content), ( + "the failover retry would ship zero system cache breakpoints" + ) + # The static prefix must stay byte-identical or its cache entry misses. + assert content[0]["text"] == self._STATIC + # The identity refresh still lands, in the volatile tail. + assert "Model: gemma4:e2b-mlx" in content[1]["text"] + assert "Provider: custom" in content[1]["text"] + + def test_keeps_single_block_shape(self): + prompt = "Model: gpt-5.4-mini\nProvider: openai-codex" + agent = _agent(prompt=prompt) + # No static prefix match -> the single-block fallback layout. + api_messages = apply_anthropic_cache_control( + [{"role": "system", "content": prompt}], + cache_ttl=None, + native_anthropic=True, + static_system_prefix=None, + ) + assert isinstance(api_messages[0]["content"], list) + assert len(api_messages[0]["content"]) == 1 + + rewrite_prompt_model_identity(agent, "gemma4:e2b-mlx", "custom") + _sync_failover_system_message(agent, api_messages, prompt) + + content = api_messages[0]["content"] + assert isinstance(content, list) and len(content) == 1 + assert content[0].get("cache_control") + assert "Model: gemma4:e2b-mlx" in content[0]["text"] + + def test_ephemeral_prompt_lands_in_the_volatile_tail(self): + prompt = self._STATIC + "Model: gpt-5.4-mini\nProvider: openai-codex" + agent = _agent(prompt=prompt, ephemeral="Stay terse.") + api_messages = self._decorated(prompt) + + _sync_failover_system_message(agent, api_messages, prompt) + + content = api_messages[0]["content"] + assert content[0]["text"] == self._STATIC + assert content[1]["text"].endswith("Stay terse.") + + def test_unknown_block_shape_falls_back_to_string(self): + agent = _agent() + api_messages = [ + {"role": "system", "content": [{"type": "image", "source": {}}]}, + ] + _sync_failover_system_message(agent, api_messages, _PROMPT) + assert api_messages[0]["content"] == agent._cached_system_prompt