diff --git a/agent/conversation_compression.py b/agent/conversation_compression.py index 0f56f9a92c43..506181f2d390 100644 --- a/agent/conversation_compression.py +++ b/agent/conversation_compression.py @@ -2040,14 +2040,13 @@ def compress_context( # (same startswith gate as the restore path); otherwise the # request layer falls back to the legacy single-breakpoint # layout with the prompt bytes untouched. - try: - from agent.system_prompt import build_system_prompt_parts as _build_parts + from agent.system_prompt import reconstruct_static_prefix - _static = _build_parts(agent, system_message=system_message)["stable"] - if _static and cached_system_prompt.startswith(_static): - agent._cached_system_prompt_static = _static - except Exception: - pass + reconstruct_static_prefix( + agent, + system_message=system_message, + log_label="compression keep-prompt", + ) else: new_system_prompt = agent._build_system_prompt(system_message) agent._cached_system_prompt = new_system_prompt diff --git a/agent/conversation_loop.py b/agent/conversation_loop.py index 6c90d13842dd..8e50acaaf3f9 100644 --- a/agent/conversation_loop.py +++ b/agent/conversation_loop.py @@ -447,30 +447,13 @@ def _restore_or_build_system_prompt(agent, system_message, conversation_history) # first turn — flip-flopping the wire shape mid-conversation and # silently degrading to the legacy single-breakpoint layout. # - # Safety: the rebuilt stable tier is used ONLY when the restored - # prompt literally starts with it (checked here AND re-checked by - # ``_apply_system_cache_markers``'s ``startswith`` gate). If any - # stable-tier input changed since the prompt was persisted (skills - # edited, identity changed), the prefix mismatches, ``_static`` - # stays None, and the request falls back to the legacy layout with - # the restored prompt bytes untouched — never a rewritten prompt. - # - # Gated on ``_use_prompt_caching`` so non-Anthropic routes skip the - # rebuild entirely (the static prefix is only consumed by - # ``apply_anthropic_cache_control``). - if getattr(agent, "_use_prompt_caching", False): - try: - from agent.system_prompt import build_system_prompt_parts as _build_parts + # ``reconstruct_static_prefix`` gates on ``_use_prompt_caching`` (so + # non-Anthropic routes skip the rebuild), applies the startswith + # safety gate (stored prompt bytes are never rewritten), and + # fails open to the legacy cache layout. + from agent.system_prompt import reconstruct_static_prefix - _static = _build_parts(agent, system_message=system_message)["stable"] - if _static and stored_prompt.startswith(_static): - agent._cached_system_prompt_static = _static - except Exception: - # Fail-open: restore continues with the legacy cache layout. - logger.debug( - "static system-prefix reconstruction failed on restore", - exc_info=True, - ) + reconstruct_static_prefix(agent, system_message=system_message) return if stored_prompt: stored_state = "stale_runtime" @@ -846,29 +829,16 @@ def _ensure_cached_system_prompt_static(agent, system_message=None) -> None: (gated on ``_use_prompt_caching`` at restore time). A later failover to a cache-on provider would otherwise redecorate with ``static_system_prefix= None`` and silently fall back to the legacy system-plus-3 layout (#72626). - """ - if not getattr(agent, "_use_prompt_caching", False): - return - stored = getattr(agent, "_cached_system_prompt", None) - if not isinstance(stored, str) or not stored: - return - existing = getattr(agent, "_cached_system_prompt_static", None) - if isinstance(existing, str) and existing and stored.startswith(existing): - return - try: - from agent.system_prompt import build_system_prompt_parts as _build_parts - static = _build_parts(agent, system_message=system_message)["stable"] - if static and stored.startswith(static): - agent._cached_system_prompt_static = static - else: - agent._cached_system_prompt_static = None - except Exception: - logger.debug( - "static system-prefix reconstruction failed on failover redecoration", - exc_info=True, - ) - agent._cached_system_prompt_static = None + Thin wrapper over :func:`agent.system_prompt.reconstruct_static_prefix`, + which memoizes failed rebuilds so this stays cheap on the retry-loop hot + path (it runs at the top of every attempt). + """ + from agent.system_prompt import reconstruct_static_prefix + + reconstruct_static_prefix( + agent, system_message=system_message, log_label="failover redecoration" + ) def _peel_moa_guidance( diff --git a/agent/system_prompt.py b/agent/system_prompt.py index d92072d1ac1e..8b8832ca2807 100644 --- a/agent/system_prompt.py +++ b/agent/system_prompt.py @@ -26,6 +26,7 @@ Pure helpers that read the agent's state. AIAgent keeps thin forwarders. from __future__ import annotations import json +import logging import os from typing import Any, Dict, List, Optional @@ -51,6 +52,8 @@ from agent.runtime_cwd import resolve_context_cwd from hermes_constants import get_hermes_home from utils import is_truthy_value +logger = logging.getLogger(__name__) + def _ra(): """Lazy reference to the ``run_agent`` module. @@ -582,6 +585,60 @@ def invalidate_system_prompt(agent: Any) -> None: agent._memory_store.load_from_disk() +def reconstruct_static_prefix( + agent: Any, + system_message: Optional[str] = None, + *, + log_label: str = "restore", +) -> None: + """Reconstruct ``_cached_system_prompt_static`` for a stored prompt. + + The static prefix is not persisted (only the full prompt is), so any + path that adopts a stored/kept ``_cached_system_prompt`` — session + restore, the compression keep-prompt path, or a failover to a cache-on + provider mid-turn (#72626) — must rebuild the stable tier to regain the + two-block ``[static, volatile]`` system layout. + + Safety: the rebuilt stable tier is used ONLY when the stored prompt + literally starts with it (checked here AND re-checked by + ``_apply_system_cache_markers``'s ``startswith`` gate). If any + stable-tier input changed since the prompt was persisted (skills + edited, identity changed), the prefix mismatches, the static stays + None, and requests fall back to the legacy layout with the stored + prompt bytes untouched — never a rewritten prompt. + + A failed reconstruction is memoized per stored prompt + (``_static_rebuild_failed_for``): ``build_system_prompt_parts`` does + real file I/O (SOUL.md, context files, memory), and callers on the + retry-loop hot path must not re-run it every attempt when the inputs + haven't changed. A legitimately changed stored prompt retries once. + """ + if not getattr(agent, "_use_prompt_caching", False): + return + stored = getattr(agent, "_cached_system_prompt", None) + if not isinstance(stored, str) or not stored: + return + existing = getattr(agent, "_cached_system_prompt_static", None) + if isinstance(existing, str) and existing and stored.startswith(existing): + return + if getattr(agent, "_static_rebuild_failed_for", None) == stored: + return + try: + static = build_system_prompt_parts(agent, system_message=system_message)["stable"] + if static and stored.startswith(static): + agent._cached_system_prompt_static = static + agent._static_rebuild_failed_for = None + return + except Exception: + logger.debug( + "static system-prefix reconstruction failed on %s", + log_label, + exc_info=True, + ) + agent._cached_system_prompt_static = None + agent._static_rebuild_failed_for = stored + + def format_tools_for_system_message(agent: Any) -> str: """Format tool definitions for the system message in the trajectory format. diff --git a/tests/agent/test_system_prompt_restore.py b/tests/agent/test_system_prompt_restore.py index 72564325d93e..fc77427125a9 100644 --- a/tests/agent/test_system_prompt_restore.py +++ b/tests/agent/test_system_prompt_restore.py @@ -350,5 +350,77 @@ class TestStaticPrefixReconstructionOnRestore: assert agent._cached_system_prompt_static is None +class TestReconstructStaticPrefixMemoization: + """A failed static rebuild must not re-run the parts builder every call. + + ``reconstruct_static_prefix`` sits on the retry-loop hot path via the + failover redecoration chokepoint (#72626); ``build_system_prompt_parts`` + does real file I/O (SOUL.md, context files, memory), so a persistent + stable-tier mismatch must be checked once per stored prompt, not on + every attempt of every API call. + """ + + def _agent(self, stored): + agent = _make_agent() + agent._use_prompt_caching = True + agent._cached_system_prompt = stored + agent._cached_system_prompt_static = None + return agent + + def test_failed_rebuild_is_memoized_per_stored_prompt(self): + from unittest.mock import patch as _patch + + from agent.system_prompt import reconstruct_static_prefix + + stored = "STORED PROMPT\n\ntail" + agent = self._agent(stored) + with _patch( + "agent.system_prompt.build_system_prompt_parts", + return_value={"stable": "MISMATCH", "context": "", "volatile": ""}, + ) as build: + reconstruct_static_prefix(agent) + reconstruct_static_prefix(agent) + reconstruct_static_prefix(agent) + assert build.call_count == 1 + assert agent._cached_system_prompt_static is None + + def test_changed_stored_prompt_retries_once(self): + from unittest.mock import patch as _patch + + from agent.system_prompt import reconstruct_static_prefix + + agent = self._agent("OLD STORED") + with _patch( + "agent.system_prompt.build_system_prompt_parts", + return_value={"stable": "MISMATCH", "context": "", "volatile": ""}, + ) as build: + reconstruct_static_prefix(agent) + # A new stored prompt (e.g. after compression) invalidates the + # failure memo and gets exactly one fresh attempt. + agent._cached_system_prompt = "NEW STORED" + reconstruct_static_prefix(agent) + reconstruct_static_prefix(agent) + assert build.call_count == 2 + + def test_success_clears_failure_memo_and_early_returns(self): + from unittest.mock import patch as _patch + + from agent.system_prompt import reconstruct_static_prefix + + stable = "STATIC HEAD" + stored = stable + "\n\nvolatile" + agent = self._agent(stored) + with _patch( + "agent.system_prompt.build_system_prompt_parts", + return_value={"stable": stable, "context": "", "volatile": ""}, + ) as build: + reconstruct_static_prefix(agent) + reconstruct_static_prefix(agent) + # Second call early-returns on the already-valid static prefix. + assert build.call_count == 1 + assert agent._cached_system_prompt_static == stable + assert getattr(agent, "_static_rebuild_failed_for", None) is None + + if __name__ == "__main__": pytest.main([__file__, "-v"])