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 1d5bc1cc799f..2c827cf2bfa0 100644 --- a/agent/conversation_loop.py +++ b/agent/conversation_loop.py @@ -72,7 +72,10 @@ from agent.model_metadata import ( save_context_length, ) from agent.process_bootstrap import _install_safe_stdio -from agent.prompt_caching import apply_anthropic_cache_control +from agent.prompt_caching import ( + apply_anthropic_cache_control, + strip_anthropic_cache_control, +) from agent.retry_utils import ( adaptive_rate_limit_backoff, is_zai_coding_overload_error, @@ -444,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" @@ -836,6 +822,104 @@ def _sync_failover_system_message(agent, api_messages, active_system_prompt): return sp +def _ensure_cached_system_prompt_static(agent, system_message=None) -> None: + """Rebuild ``_cached_system_prompt_static`` when caching becomes active. + + Sessions restored under a cache-off primary skip the static-prefix rebuild + (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). + + 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( + messages: List[Dict[str, Any]], + guidance: Any, +) -> List[Dict[str, Any]]: + """Remove MoA reference guidance previously attached by ``_attach_reference_guidance``. + + Thin wrapper over :func:`agent.moa_loop.peel_reference_guidance` (kept + adjacent to the attach so the forward/inverse shapes evolve together). + Lazy import mirrors the module's other moa_loop touchpoints. + """ + from agent.moa_loop import peel_reference_guidance + + return peel_reference_guidance(messages, guidance) + + +def _redecorate_prompt_cache_for_provider( + agent, + api_messages: List[Dict[str, Any]], + *, + system_message=None, + moa_prepared: Optional[Dict[str, Any]] = None, +) -> tuple[List[Dict[str, Any]], Optional[Dict[str, Any]]]: + """Strip and re-apply cache_control for the *current* provider policy. + + Decoration runs once per call block before the retry loop for the primary + provider. ``try_activate_fallback`` refreshes ``_use_prompt_caching`` / + ``_use_native_cache_layout`` but the nine failover ``continue`` paths reused + the old ``api_messages`` (#72626). Mirror ``_reapply_reasoning_echo_for_provider`` + by reshaping at the top of each retry attempt. + + The source list is the mutated in-flight request (image shrink / ASCII / + reasoning_details recoveries already applied) — never a pristine + pre-decoration snapshot. MoA guidance is peeled, the base is redecorated, + then ``rebase_prepared_request`` re-attaches guidance outside the cached + span. + """ + messages: List[Dict[str, Any]] = [ + dict(m) if isinstance(m, dict) else m for m in (api_messages or []) + ] + prepared = moa_prepared + guidance = prepared.get("guidance") if isinstance(prepared, dict) else None + if guidance: + messages = _peel_moa_guidance(messages, guidance) + + strip_anthropic_cache_control(messages) + + # Direct attribute access matches the call-block decoration site — the + # flags are unconditionally initialized on AIAgent, and a getattr + # default here would mask a real init bug as silent cache-off. + if agent._use_prompt_caching: + _ensure_cached_system_prompt_static(agent, system_message=system_message) + static = getattr(agent, "_cached_system_prompt_static", None) + messages = apply_anthropic_cache_control( + messages, + cache_ttl=agent._cache_ttl, + native_anthropic=agent._use_native_cache_layout, + static_system_prefix=static if isinstance(static, str) else None, + ) + + if ( + prepared is not None + and getattr(agent, "provider", None) == "moa" + ): + # No `and guidance` here: guidance=None is a real prepared shape + # (all-references-failed / silent degraded policy builds the + # prepared request without attaching guidance), and the MoA facade + # sends prepared["messages"] — not api_kwargs["messages"] — so the + # rebase must refresh the prepared object even when there is no + # guidance to re-attach. rebase_prepared_request handles falsy + # guidance by copying the messages and skipping the attach. + completions = getattr(getattr(agent.client, "chat", None), "completions", None) + rebase = getattr(completions, "rebase_prepared_request", None) + if callable(rebase): + prepared = rebase(prepared, messages) + messages = prepared["messages"] + + return messages, prepared + + def _apply_context_engine_selection( agent: Any, api_messages: List[Dict[str, Any]], @@ -1909,6 +1993,18 @@ def run_conversation( # unless the active provider needs it) so the fallback request # isn't sent with stale, primary-shaped reasoning fields. agent._reapply_reasoning_echo_for_provider(api_messages) + # Same story for prompt-cache decoration (#72626): try_activate_ + # fallback refreshes the policy flags, but the decorated list + # still carries the primary's breakpoints (or none). Strip and + # re-render for the current provider before building kwargs. + api_messages, _moa_prepared_request = ( + _redecorate_prompt_cache_for_provider( + agent, + api_messages, + system_message=system_message, + moa_prepared=_moa_prepared_request, + ) + ) api_kwargs = agent._build_api_kwargs(api_messages) if agent._force_ascii_payload: _sanitize_structure_non_ascii(api_kwargs) diff --git a/agent/moa_loop.py b/agent/moa_loop.py index ed0f41a817aa..173816c8cbd1 100644 --- a/agent/moa_loop.py +++ b/agent/moa_loop.py @@ -1337,6 +1337,63 @@ def _attach_reference_guidance(agg_messages: list[dict[str, Any]], guidance: str agg_messages.append({"role": "user", "content": guidance}) +def peel_reference_guidance( + messages: list[dict[str, Any]], + guidance: Any, +) -> list[dict[str, Any]]: + """Remove reference guidance previously attached by ``_attach_reference_guidance``. + + Exact inverse of the three attach shapes above (string merge, trailing + text part, appended user message) — kept adjacent so the two evolve + together; a drifting separator or shape would make the peel silently + no-op and let a cache breakpoint land on the turn-varying guidance + block (the bug class #72626 fixes). + + Used by the failover redecoration chokepoint: redecoration must run on + the base transcript so the last cache breakpoint does not land on the + guidance; callers then rebase via ``rebase_prepared_request``. + + Returns a new list (input list and its messages are not mutated). + """ + if not guidance or not messages: + return messages + guidance_text = str(guidance) + last = messages[-1] + if not isinstance(last, dict) or last.get("role") != "user": + return messages + content = last.get("content") + if content == guidance_text: + # Attach shape (c): guidance was appended as its own user message. + return list(messages[:-1]) + suffix = "\n\n" + guidance_text + if isinstance(content, str) and content.endswith(suffix): + # Attach shape (a): merged into a trailing string user turn. + peeled = dict(last) + peeled["content"] = content[: -len(suffix)] + return [*messages[:-1], peeled] + if isinstance(content, list) and content: + last_part = content[-1] + if isinstance(last_part, dict) and last_part.get("type", "text") == "text": + text = last_part.get("text") or "" + if text == suffix or text == guidance_text: + # Attach shape (b): guidance rode as its own trailing part. + peeled = dict(last) + peeled["content"] = list(content[:-1]) + if not peeled["content"]: + # The guidance part was the only content — mirror the + # string shape (c) and drop the whole message rather + # than leaving an empty-content user turn behind. + return list(messages[:-1]) + return [*messages[:-1], peeled] + if text.endswith(suffix): + new_part = dict(last_part) + new_part["text"] = text[: -len(suffix)] + peeled = dict(last) + peeled["content"] = [*content[:-1], new_part] + return [*messages[:-1], peeled] + return messages + + class MoAChatCompletions: """OpenAI-chat-compatible facade where the aggregator is the acting model.""" diff --git a/agent/prompt_caching.py b/agent/prompt_caching.py index 09b51e61d018..2e606cd377c3 100644 --- a/agent/prompt_caching.py +++ b/agent/prompt_caching.py @@ -119,6 +119,59 @@ def _apply_system_cache_markers( return 1 +def strip_anthropic_cache_control( + api_messages: List[Dict[str, Any]], +) -> List[Dict[str, Any]]: + """Remove ``cache_control`` markers and undo decoration-produced list shapes. + + Used before re-applying decoration after a mid-turn provider failover so + the mutated, undecorated shape (image shrink / ASCII cleanup / etc.) is + preserved while markers match the *new* provider's cache policy (#72626). + + Flattening back to a plain string is restricted to the exact shapes + :func:`apply_anthropic_cache_control` produces from string content — + a single ``{"type": "text"}`` part, or the two-part ``[static, volatile]`` + system split — so the ``""``-join is provably byte-exact. Organic + multi-part text (merged user turns, imported transcripts) and parts + carrying extra keys (``citations`` etc.) keep their structure; only + per-part markers are removed. Marker removal is copy-on-write on the + part dicts: content parts may alias the persistent conversation history + (the per-call copy is shallow), and stripping must never rewrite the + stored transcript. + + Mutates the top-level message dicts of ``api_messages`` in place and + returns the same list. + """ + for msg in api_messages: + if not isinstance(msg, dict): + continue + msg.pop("cache_control", None) + content = msg.get("content") + if not isinstance(content, list): + continue + if any(isinstance(part, dict) and "cache_control" in part for part in content): + content = [ + {k: v for k, v in part.items() if k != "cache_control"} + if isinstance(part, dict) and "cache_control" in part + else part + for part in content + ] + msg["content"] = content + decoration_shape = content and all( + isinstance(part, dict) + and part.get("type", "text") == "text" + and isinstance(part.get("text"), str) + and set(part.keys()) <= {"type", "text"} + for part in content + ) and ( + len(content) == 1 + or (msg.get("role") == "system" and len(content) == 2) + ) + if decoration_shape: + msg["content"] = "".join(part["text"] for part in content) + return api_messages + + def apply_anthropic_cache_control( api_messages: List[Dict[str, Any]], cache_ttl: str = "5m", 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/apps/desktop/src/app/chat/right-rail/preview.tsx b/apps/desktop/src/app/chat/right-rail/preview.tsx index 219a28aacfc8..9d9a6b3950e5 100644 --- a/apps/desktop/src/app/chat/right-rail/preview.tsx +++ b/apps/desktop/src/app/chat/right-rail/preview.tsx @@ -10,7 +10,7 @@ import { ContextMenuSeparator, ContextMenuTrigger } from '@/components/ui/context-menu' -import { PANE_TAB_STRIP_LINE, PaneTab, PaneTabLabel } from '@/components/ui/pane-tab' +import { PaneTab, PaneTabLabel } from '@/components/ui/pane-tab' import { Tip } from '@/components/ui/tooltip' import { translateNow, useI18n } from '@/i18n' import { formatCombo } from '@/lib/keybinds/combo' @@ -131,12 +131,7 @@ export function ChatPreviewRail({ onRestartServer, setTitlebarToolGroup }: ChatP // titlebar-height so it opens below the band. 0px elsewhere → unchanged. style={{ paddingTop: 'var(--right-rail-top-inset, 0px)' }} > -