From aa40dbe7656a0cd58aa609cb8d6f24fb990c84e9 Mon Sep 17 00:00:00 2001 From: joaomarcos Date: Sun, 6 Sep 2026 14:45:14 -0300 Subject: [PATCH] fix(agent): skip the tools freeze once on a surface switch, not for the session MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The first cut gated the saved tool_names pin on "the surface drifted", which stays true for as long as the stored prompt names the old surface — i.e. until the next compaction. On the gateway path, where a fresh AIAgent is built per turn, that left the tools freeze off for every remaining turn, so a check_fn that flaps could reorder `tools[]` and break the tool cache block on its own. Gate it on the turn that actually ANNOUNCES the switch instead, and persist the fresh, toolset-correct names there. The next turn's row already holds this surface's tools, so the pin resumes immediately: skipped once, not disabled. --- agent/conversation_loop.py | 27 ++++++++++++++--------- tests/agent/test_system_prompt_restore.py | 22 ++++++++++++++++-- 2 files changed, 36 insertions(+), 13 deletions(-) diff --git a/agent/conversation_loop.py b/agent/conversation_loop.py index eca7e2eb67..ac02f88aa3 100644 --- a/agent/conversation_loop.py +++ b/agent/conversation_loop.py @@ -690,16 +690,21 @@ def _restore_or_build_system_prompt(agent, system_message, conversation_history) 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) + announced_switch = _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. + # NOT on the turn that announces 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, + # and is persisted so the very next turn pins THAT one — the freeze is skipped once, + # not disabled for the rest of the session. try: saved_tools = session_row.get("tool_names") if session_row else None - if saved_tools and not surface_changed: + if announced_switch: + from tools.mcp_tool_agent import persist_agent_tool_names + persist_agent_tool_names(agent) + elif saved_tools: from tools.mcp_tool_agent import restore_agent_tool_prefix restore_agent_tool_prefix(agent, json.loads(saved_tools)) except Exception: @@ -818,17 +823,17 @@ def _surface_already_announced(conversation_history, platform: str) -> bool: 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. + Returns whether the note was staged — the caller pairs that with re-persisting the tool + names, because this is the turn the surface's toolset actually changes under it. The note + 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. """ 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 + return False from agent.system_prompt import platform_surface_hint note = ( f"{_SURFACE_SWITCH_NOTE_PREFIX}{current} (the system prompt above was written for " diff --git a/tests/agent/test_system_prompt_restore.py b/tests/agent/test_system_prompt_restore.py index aecb954842..4600a9c863 100644 --- a/tests/agent/test_system_prompt_restore.py +++ b/tests/agent/test_system_prompt_restore.py @@ -116,7 +116,7 @@ class TestSurfaceSwitch: 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): + def test_tool_prefix_is_replaced_on_the_turn_that_announces_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 @@ -125,9 +125,27 @@ class TestSurfaceSwitch: """ from unittest.mock import patch - with patch("tools.mcp_tool_agent.restore_agent_tool_prefix") as pin: + with ( + patch("tools.mcp_tool_agent.restore_agent_tool_prefix") as pin, + patch("tools.mcp_tool_agent.persist_agent_tool_names") as persist, + ): self._restore(stored="desktop", current="tui", tool_names='["desktop_ui_tool"]') pin.assert_not_called() + persist.assert_called_once() + + def test_tool_prefix_is_pinned_again_on_the_next_turn(self): + # The freeze is skipped once, not disabled for the rest of the session: the announcing + # turn persisted this surface's tools, so the next turn pins those. + from unittest.mock import patch + + history = [ + {"role": "user", "content": f"hi\n\n{_SURFACE_SWITCH_NOTE_PREFIX}tui (...)]"}, + {"role": "assistant", "content": "hello"}, + ] + with patch("tools.mcp_tool_agent.restore_agent_tool_prefix") as pin: + self._restore(stored="desktop", current="tui", history=history, + tool_names='["read_file"]') + pin.assert_called_once() def test_tool_prefix_is_still_pinned_on_the_same_surface(self): from unittest.mock import patch