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.
This commit is contained in:
dsad 2026-07-26 19:10:05 +00:00 committed by Teknium
parent c2ee5039ee
commit 3c4220cd9c
2 changed files with 124 additions and 1 deletions

View file

@ -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

View file

@ -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