fix(agent): a surface switch must not re-prefill the whole request (#104414)

`_stored_prompt_matches_runtime` treated `Platform` as a runtime-identity field, so
answering a live session from another surface — desktop -> TUI, or a resume after a
dashboard restart whose chat is a PTY TUI child — declared the stored prompt stale and
rebuilt it. The system prompt is the first thing in the request, so changing any byte of
it moves the first divergent byte to the head of a 220K-token request and the entire
conversation behind it re-prefills: a session that was hitting 240000/240287 came back
at 1536/219861.

The guard was not wrong about correctness — a desktop-built prompt on a terminal session
advertises inline widgets and a MEDIA: channel the TUI does not have — but the surface is
advisory metadata about the renderer, not a cache domain. Model/provider and cwd drift
change what the prompt should SAY; the surface changes only one paragraph.

Reuse the stored bytes across a surface switch and correct the paragraph where it costs
nothing to cache: `_stage_surface_switch_note` stages a one-shot note carrying the CURRENT
surface's guidance on the same per-turn user-message channel the gateway's must-deliver
notes use. It lands after the cached prefix and is stamped into the byte-stable
`api_content` sidecar, so later turns replay it instead of re-prefilling, and the prompt
converges at the next compaction — a boundary that already breaks the cache.

The saved tool_names prefix is not pinned across a switch: the tool registry is
process-global, so `_merge_preserving_prefix` would carry a saved-but-unloaded tool
forward (under `coding_context: focus` desktop gets a desktop_ui toolset the TUI cannot
run). On the same surface the tools freeze is untouched.
This commit is contained in:
joaomarcos
2026-09-06 14:40:21 -03:00
committed by kshitij
parent 990473a79c
commit 80d6bda144
4 changed files with 242 additions and 15 deletions

View File

@@ -688,11 +688,18 @@ def _restore_or_build_system_prompt(agent, system_message, conversation_history)
# Continuing session — reuse the exact system prompt from the
# previous turn so the Anthropic cache prefix matches.
agent._cached_system_prompt = stored_prompt
# The reused bytes may describe the surface this conversation STARTED on; correct that
# at the tail of the request instead of rebuilding the prompt in front of it (#104414).
surface_changed = _stage_surface_switch_note(agent, stored_prompt, conversation_history)
# Same contract for tools[]: pin the array to the order this session already
# sent (tools freeze) instead of re-probing every check_fn on a fresh AIAgent.
# NOT across a surface switch: the saved names are the PREVIOUS surface's toolset
# (``coding_context: focus`` gives desktop a desktop_ui toolset the TUI has no use for)
# and the tool registry is process-global, so the merge would re-advertise tools this
# surface cannot run. The fresh, toolset-correct array wins there.
try:
saved_tools = session_row.get("tool_names") if session_row else None
if saved_tools:
if saved_tools and not surface_changed:
from tools.mcp_tool_agent import restore_agent_tool_prefix
restore_agent_tool_prefix(agent, json.loads(saved_tools))
except Exception:
@@ -753,6 +760,91 @@ def _restore_or_build_system_prompt(agent, system_message, conversation_history)
)
# A surface switch (desktop <-> TUI, a session resumed under a different host) changes only
# which interface renders the reply, but the surface guidance and the ``Platform:`` trailer are
# embedded in the persisted prompt. Rebuilding for it produced a system prompt that diverged
# from the cached one within its first blocks, so the ENTIRE request behind it re-prefilled —
# a 220K-token session came back at a 1% cache hit (#104414). The stored bytes are therefore
# kept and the CURRENT surface's guidance is delivered on the per-turn user-message channel
# instead: that lands after the cached prefix, is stamped into the byte-stable ``api_content``
# sidecar so later turns replay it unchanged, and the prompt itself converges at the next
# rebuild boundary (compaction).
_SURFACE_SWITCH_NOTE_PREFIX = "[System: This conversation is now being answered on a different interface: "
def _stored_prompt_platform(prompt: str) -> str:
"""The ``Platform:`` value the stored prompt was built with ("" when absent).
Last match wins, like the volatile-tier reads in ``_stored_prompt_matches_runtime``: the
trailer sits at the very end, so embedded project context cannot shadow it.
"""
prefix = "Platform:"
matches = [line[len(prefix):].strip() for line in prompt.splitlines() if line.startswith(prefix)]
return matches[-1] if matches else ""
def _transcript_row_texts(msg):
"""Every string the model actually saw for one transcript row: the wire sidecar first,
then the stored content (plain string, or the text parts of a multimodal list)."""
if not isinstance(msg, dict):
return
sidecar = msg.get("api_content")
if isinstance(sidecar, str):
yield sidecar
content = msg.get("content")
if isinstance(content, str):
yield content
elif isinstance(content, list):
for part in content:
if isinstance(part, dict) and isinstance(part.get("text"), str):
yield part["text"]
def _surface_already_announced(conversation_history, platform: str) -> bool:
"""True when the NEWEST surface note in the transcript already names ``platform``.
The note is stamped into the ``api_content`` sidecar, so every later turn replays it; the
gateway builds a fresh AIAgent per turn and would otherwise stack one copy per turn until
the next compaction. Only the newest note is consulted — an older one names the surface
the conversation has since left.
"""
for msg in reversed(conversation_history or []):
for text in _transcript_row_texts(msg):
if _SURFACE_SWITCH_NOTE_PREFIX in text:
return text.rsplit(_SURFACE_SWITCH_NOTE_PREFIX, 1)[1].startswith(f"{platform} (")
return False
def _stage_surface_switch_note(agent, stored_prompt: str, conversation_history) -> bool:
"""Stage the correction note when the reused prompt describes a different surface.
Returns whether the surface actually drifted — the caller also uses it to decide whether
the saved tool prefix may be pinned. The note itself is one-shot per turn (consumed by
``agent.turn_context.consume_surface_switch_note``) and is skipped when the transcript
already carries one for this surface, but that does not make the drift go away.
"""
current = str(getattr(agent, "platform", "") or "").strip()
stored = _stored_prompt_platform(stored_prompt)
if not current or not stored or stored == current:
return False
if _surface_already_announced(conversation_history, current):
return True
from agent.system_prompt import platform_surface_hint
note = (
f"{_SURFACE_SWITCH_NOTE_PREFIX}{current} (the system prompt above was written for "
f"{stored}). Its interface section no longer applies — use the guidance below for "
"formatting, file delivery and any interface-specific capability.]"
)
hint = platform_surface_hint(agent)
agent._surface_switch_note = f"{note}\n{hint}" if hint else note
logger.info(
"Session %s switched surface %s -> %s; keeping the stored system prompt and delivering "
"the new surface guidance as a turn note (prefix cache preserved).",
agent.session_id, stored, current,
)
return True
def _stored_prompt_matches_runtime(agent, prompt: str) -> bool:
"""Return False when the persisted runtime-identity lines are stale."""
@@ -780,8 +872,9 @@ def _stored_prompt_matches_runtime(agent, prompt: str) -> bool:
return candidate[len(prefix):].strip()
return ""
# Model/provider identity, then cwd drift, then runtime-surface drift (reusing a
# desktop-built prompt on a terminal session would inject the wrong runtime hints).
# Model/provider identity, then cwd drift. A cwd change is a real content change (context
# files, the workspace snapshot and the coding posture are all resolved from it), so it
# still rebuilds; the runtime surface does not — see the note above.
for label, attr in (("Model", "model"), ("Provider", "provider")):
stored = line_value(label)
current = str(getattr(agent, attr, "") or "").strip()
@@ -792,9 +885,10 @@ def _stored_prompt_matches_runtime(agent, prompt: str) -> bool:
stored_cwd = host_info_value("Current working directory")
if stored_cwd and stored_cwd != str(resolve_agent_cwd()):
return False
stored_platform = line_value("Platform")
current_platform = str(getattr(agent, "platform", "") or "").strip()
return not (stored_platform and current_platform and stored_platform != current_platform)
# Platform is deliberately NOT an identity field: a surface switch does not invalidate the
# stored bytes, it only makes their interface section out of date, and that is corrected by
# _stage_surface_switch_note without touching the cached prefix (#104414).
return True
# Named so _is_synthetic_compression_user_turn can recognize a crash-persisted nudge by

View File

@@ -457,6 +457,16 @@ def _timestamp_line(agent: Any) -> str:
return timestamp_line + "".join(f"\n{label}: {value}" for label, value in trailer if value)
def platform_surface_hint(agent: Any) -> str:
"""The rendering-surface guidance for ``agent.platform`` ("" when the surface has none).
Public because a session that changed surface mid-conversation keeps its stored prompt
(rebuilding it re-prefills the whole request) and delivers the CURRENT surface's guidance
through the per-turn user-message channel instead — see ``_SURFACE_SWITCH_NOTE_PREFIX``
in ``agent.conversation_loop`` (#104414)."""
return _platform_hint(agent) or ""
def _memory_parts(agent: Any) -> List[str]:
"""Built-in memory/USER.md blocks plus the external provider block (gated on
the same check ``inject_memory_provider_tools`` uses, so we never advertise
@@ -736,7 +746,7 @@ def format_tools_for_system_message(agent: Any) -> str:
__all__ = ["build_system_prompt_parts", "build_system_prompt", "invalidate_system_prompt",
"restore_plugin_prompt_sections", "format_tools_for_system_message"]
"platform_surface_hint", "restore_plugin_prompt_sections", "format_tools_for_system_message"]
# ---- BEGIN PLUGIN-COMPAT (revert-scheduled; see COMPAT_MANIFEST.md) ----

View File

@@ -124,6 +124,20 @@ def consume_gateway_turn_context_notes(agent: Any) -> str:
return notes if isinstance(notes, str) else ""
def consume_surface_switch_note(agent: Any) -> str:
"""Pop the one-shot surface-change note staged by the system-prompt restore.
Rides the same user-message channel as the gateway's must-deliver notes: a session that
moved between surfaces keeps its stored system prompt (rebuilding it re-prefills the whole
request) and gets the new surface's guidance here, behind the cached prefix (#104414).
"""
note = getattr(agent, "_surface_switch_note", "") or ""
if hasattr(agent, "_surface_switch_note"):
with suppress(Exception):
agent._surface_switch_note = ""
return note if isinstance(note, str) else ""
def append_notes_to_multimodal_content(content: Any, notes: str) -> bool:
"""Append must-deliver notes as a durable text part on a multimodal (list) user
message (the sidecar path returns ``None`` for non-string content)."""
@@ -663,11 +677,15 @@ def _collect_pre_llm_call_context(
def _merge_gateway_notes(
agent: Any, messages: List[Any], current_turn_user_idx: int, plugin_user_context: str
) -> str:
"""Gateway must-deliver notes ride the user-message injection channel (one-shot,
gateway-staged) so the ephemeral system prompt stays byte-stable. Multimodal (list)
content can't take the string sidecar — append a durable text part instead."""
_gateway_notes = consume_gateway_turn_context_notes(agent)
if not _gateway_notes:
"""Must-deliver per-turn notes ride the user-message injection channel (one-shot) so the
ephemeral system prompt stays byte-stable: the gateway's staged notes, then the
surface-switch correction. Multimodal (list) content can't take the string sidecar —
append a durable text part instead."""
_turn_notes = "\n\n".join(
part for part in (consume_gateway_turn_context_notes(agent),
consume_surface_switch_note(agent)) if part
)
if not _turn_notes:
return plugin_user_context
_gw_turn_content = (
messages[current_turn_user_idx].get("content")
@@ -676,10 +694,10 @@ def _merge_gateway_notes(
else None
)
if isinstance(_gw_turn_content, list):
append_notes_to_multimodal_content(_gw_turn_content, _gateway_notes)
append_notes_to_multimodal_content(_gw_turn_content, _turn_notes)
return plugin_user_context
return (
plugin_user_context + "\n\n" + _gateway_notes if plugin_user_context else _gateway_notes
plugin_user_context + "\n\n" + _turn_notes if plugin_user_context else _turn_notes
)

View File

@@ -20,7 +20,7 @@ from unittest.mock import MagicMock
import pytest
from agent.conversation_loop import _restore_or_build_system_prompt
from agent.conversation_loop import _SURFACE_SWITCH_NOTE_PREFIX, _restore_or_build_system_prompt
def _make_agent(session_db=None, prebuilt_prompt: str = "BUILT_PROMPT"):
@@ -40,6 +40,111 @@ def _make_agent(session_db=None, prebuilt_prompt: str = "BUILT_PROMPT"):
return agent
# ---------------------------------------------------------------------------
# Surface switch (#104414)
# ---------------------------------------------------------------------------
class TestSurfaceSwitch:
"""A desktop <-> TUI switch must not rebuild the system prompt.
The rebuild it used to trigger changed the first blocks of a 200K+ token request, so
the whole conversation re-prefilled at a ~1% cache hit. The stored bytes are now reused
and the new surface's guidance is delivered as a per-turn note behind the cached prefix.
"""
@staticmethod
def _stored(platform: str) -> str:
return (
"SYSTEM PROMPT BODY\n\nConversation started: Monday, January 05, 2026\n"
"Model: test-model\nProvider: openrouter\n"
f"Platform: {platform}"
)
def _restore(self, *, stored: str, current: str, history=None, tool_names=None):
db = MagicMock()
row = {"system_prompt": self._stored(stored)}
if tool_names is not None:
row["tool_names"] = tool_names
db.get_session.return_value = row
agent = _make_agent(session_db=db)
agent.platform = current
agent._platform_hint_overrides = None
agent._surface_switch_note = ""
agent._gateway_turn_context_notes = ""
_restore_or_build_system_prompt(
agent, None, history if history is not None else [{"role": "user", "content": "hi"}]
)
return agent
def test_switch_reuses_the_stored_prompt(self):
agent = self._restore(stored="desktop", current="tui")
assert agent._cached_system_prompt == self._stored("desktop")
agent._build_system_prompt.assert_not_called()
def test_switch_stages_the_new_surface_guidance(self):
agent = self._restore(stored="desktop", current="tui")
note = agent._surface_switch_note
assert note.startswith(f"{_SURFACE_SWITCH_NOTE_PREFIX}tui (")
# The correction carries the CURRENT surface's hint, so the model is not left
# following the desktop guidance still sitting in the reused prompt.
assert "terminal UI (TUI)" in note
def test_same_surface_stages_nothing(self):
assert self._restore(stored="cli", current="cli")._surface_switch_note == ""
def test_unknown_surface_stages_nothing(self):
assert self._restore(stored="", current="tui")._surface_switch_note == ""
def test_not_restaged_once_the_transcript_carries_it(self):
# The note is stamped into the byte-stable api_content sidecar, and the gateway
# builds a fresh AIAgent per turn — without the dedup every turn would add a copy.
history = [
{"role": "user", "content": "hi",
"api_content": f"hi\n\n{_SURFACE_SWITCH_NOTE_PREFIX}tui (...)]"},
{"role": "assistant", "content": "hello"},
]
agent = self._restore(stored="desktop", current="tui", history=history)
assert agent._surface_switch_note == ""
def test_a_further_switch_is_announced_again(self):
# Only the NEWEST note counts: it names the surface the conversation has left.
history = [
{"role": "user", "content": f"hi\n\n{_SURFACE_SWITCH_NOTE_PREFIX}tui (...)]"},
{"role": "assistant", "content": "hello"},
]
agent = self._restore(stored="desktop", current="cli", history=history)
assert agent._surface_switch_note.startswith(f"{_SURFACE_SWITCH_NOTE_PREFIX}cli (")
def test_tool_prefix_is_not_pinned_across_a_switch(self):
"""The saved names are the PREVIOUS surface's toolset.
The tool registry is process-global, so ``_merge_preserving_prefix`` would keep a
saved-but-not-loaded tool (e.g. the desktop_ui toolset a `coding_context: focus`
desktop session gets) and re-advertise it on a TUI session that cannot run it.
"""
from unittest.mock import patch
with patch("tools.mcp_tool_agent.restore_agent_tool_prefix") as pin:
self._restore(stored="desktop", current="tui", tool_names='["desktop_ui_tool"]')
pin.assert_not_called()
def test_tool_prefix_is_still_pinned_on_the_same_surface(self):
from unittest.mock import patch
with patch("tools.mcp_tool_agent.restore_agent_tool_prefix") as pin:
self._restore(stored="cli", current="cli", tool_names='["read_file"]')
pin.assert_called_once()
def test_note_rides_the_user_message_channel_once(self):
from agent.turn_context import _merge_gateway_notes, consume_surface_switch_note
agent = self._restore(stored="desktop", current="tui")
staged = agent._surface_switch_note
assert _merge_gateway_notes(agent, [{"role": "user", "content": "hi"}], 0, "") == staged
assert consume_surface_switch_note(agent) == ""
# ---------------------------------------------------------------------------
# Happy paths
# ---------------------------------------------------------------------------