From b2b026dc224dc80d99d71e4ad982467d221e5d49 Mon Sep 17 00:00:00 2001 From: jango <91889514+jangomango76@users.noreply.github.com> Date: Fri, 4 Sep 2026 13:53:10 -0700 Subject: [PATCH] fix(bot-mode): retain message_agent across tool rebuilds (#102864) --- tests/tools/test_refresh_agent_mcp_tools.py | 113 ++++++++++++++++++++ tools/mcp_tool_agent.py | 21 +++- 2 files changed, 133 insertions(+), 1 deletion(-) diff --git a/tests/tools/test_refresh_agent_mcp_tools.py b/tests/tools/test_refresh_agent_mcp_tools.py index f577c31c00..fa61f69ea1 100644 --- a/tests/tools/test_refresh_agent_mcp_tools.py +++ b/tests/tools/test_refresh_agent_mcp_tools.py @@ -11,6 +11,8 @@ freezing any particular tool list. import threading import types +import pytest + from tools import mcp_tool from tools import mcp_tool_agent as _mcp_agent @@ -338,3 +340,114 @@ def test_reprobe_tool_availability_drops_cached_check_fn_verdicts(monkeypatch): assert registry_mod._check_fn_cached(probe) is True assert ("sentinel",) not in model_tools._tool_defs_cache + + +# --------------------------------------------------------------------------- +# Bot Mode dynamic capability: every snapshot rebuild re-runs its auth gate +# --------------------------------------------------------------------------- + + +class _BotModeDB: + def __init__(self, home, title): + self.db_path = str(home / "state.db") + self._title = title + + def get_session_title(self, _session_id): + return self._title + + +@pytest.fixture +def managed_bot_home(tmp_path): + home = tmp_path / ".hermes" + profile = home / "profiles" / "researcher" + profile.mkdir(parents=True) + (profile / "profile.yaml").write_text( + "ui_meta:\n hermes-bots:\n shape: cloud\n", + encoding="utf-8", + ) + return home + + +def _bot_mode_agent(home, *, title="Bot Chat"): + agent = _agent(["read_file"]) + agent._session_db = _BotModeDB(home, title) + agent.session_id = "session-1" + agent._session_title_hint = None + agent._bot_mode_protocol = True + return agent + + +def _message_agent_schema_count(agent): + return sum( + t.get("function", {}).get("name") == "message_agent" + for t in agent.tools + if isinstance(t, dict) + ) + + +def _assert_tool_snapshot_coherent(agent): + names = {t["function"]["name"] for t in agent.tools} + assert agent.valid_tool_names == names + + +@pytest.mark.parametrize("rebuild", ["compaction", "reload", "between_turns", "resume"]) +def test_authorized_message_agent_survives_every_snapshot_rebuild( + managed_bot_home, monkeypatch, rebuild +): + """Compaction, live refreshes and eviction/resume all preserve the guarded tool.""" + from tools.bot_mode_dm import ensure_message_agent_tool + from tools import registry as registry_mod + + agent = _bot_mode_agent(managed_bot_home) + _serve(monkeypatch, [_tool("read_file")]) + entry = types.SimpleNamespace(name="read_file", schema=_tool("read_file")["function"]) + monkeypatch.setattr(registry_mod.registry, "get_all_entries", lambda: [entry], raising=False) + monkeypatch.setattr( + registry_mod.registry, + "get_entry", + lambda name, **_kw: entry if name == "read_file" else None, + raising=False, + ) + + assert ensure_message_agent_tool(agent) is True + + def rebuild_snapshot(): + if rebuild == "resume": + agent.tools = [_tool("read_file")] + agent.valid_tool_names = {"read_file"} + _mcp_agent.restore_agent_tool_prefix(agent, ["read_file", "message_agent"]) + else: + _mcp_agent.refresh_agent_mcp_tools( + agent, + content_aware=rebuild == "compaction", + preserve_prefix=rebuild == "between_turns", + ) + + for _ in range(2): + rebuild_snapshot() + assert _message_agent_schema_count(agent) == 1 + assert "message_agent" in agent.valid_tool_names + _assert_tool_snapshot_coherent(agent) + + +@pytest.mark.parametrize( + ("title", "managed"), + [("Ordinary chat", True), ("Bot Chat", False)], +) +def test_snapshot_rebuild_never_grants_message_agent_to_unauthorized_sessions( + tmp_path, managed_bot_home, monkeypatch, title, managed +): + """Ordinary and unmanaged chats remain fail-closed across repeated rebuilds.""" + home = managed_bot_home if managed else tmp_path / "unmanaged" + home.mkdir(exist_ok=True) + agent = _bot_mode_agent(home, title=title) + # Even a stale/leaked dynamic capability is scrubbed unless the live gate re-authorizes it. + agent.tools.append(_tool("message_agent")) + agent.valid_tool_names.add("message_agent") + _serve(monkeypatch, [_tool("read_file")]) + + for _ in range(2): + _mcp_agent.refresh_agent_mcp_tools(agent, content_aware=True) + assert _message_agent_schema_count(agent) == 0 + assert "message_agent" not in agent.valid_tool_names + _assert_tool_snapshot_coherent(agent) diff --git a/tools/mcp_tool_agent.py b/tools/mcp_tool_agent.py index c7f924c4dc..d707add554 100644 --- a/tools/mcp_tool_agent.py +++ b/tools/mcp_tool_agent.py @@ -2,6 +2,7 @@ AIAgent's tools/tool names, preserving the cached tools[] prefix across rebuilds, and re-injecting post-build tools.""" +import copy import logging import json import threading @@ -88,7 +89,8 @@ def refresh_agent_mcp_tools( rebuilt. Shared by the TUI RPC, gateway reload, late-binding thread and between-turns refresh: respects the toolset filter, diffs by tool NAME (a count compare misses an equal-size swap), re-injects the memory-provider / context-engine tools ``agent_init`` - appends after ``get_tool_definitions``, publishes ``(tools, valid_tool_names)`` together. + appends after ``get_tool_definitions`` plus guarded session capabilities, and publishes + ``(tools, valid_tool_names)`` together. ``preserve_prefix``: for rebuilds inside a live conversation the tool array is a cached request prefix and any moved byte re-prefills the whole history — existing tools keep their @@ -105,6 +107,7 @@ def refresh_agent_mcp_tools( new_names = {_def_name(t) for t in new_defs} # Post-build families re-appended on LOCALS only; live attributes untouched until publish. staged_engine_names = _reinject_post_build_tools(agent, new_defs, new_names) + _reinject_authorized_dynamic_tools(agent, new_defs, new_names) # Registry membership is read OUTSIDE ``_agent_tools_lock``: taking ``registry._lock`` # under the tools lock would be the first nesting of the two. prefix_registered: Optional[set] = None @@ -164,6 +167,7 @@ def restore_agent_tool_prefix(agent, saved_names: list) -> bool: saved_defs = [d for d in map(_saved_def, saved_names) if d is not None] registered_names = {entry.name for entry in registry.get_all_entries()} merged, merged_names = _merge_preserving_prefix(saved_defs, fresh_defs, registered_names) + _reinject_authorized_dynamic_tools(agent, merged, merged_names) with _agent_tools_lock: if merged == fresh_defs: return False @@ -192,6 +196,21 @@ def _merge_preserving_prefix(current_defs: list, new_defs: list, registered_name return merged, {_def_name(t) for t in merged} +def _reinject_authorized_dynamic_tools(agent, tools_list: list, name_set: set) -> None: + """Re-authorize session capabilities against a staged snapshot before publication.""" + try: + from tools.bot_mode_dm import MESSAGE_AGENT_TOOL_NAME, ensure_message_agent_tool + + tools_list[:] = [entry for entry in tools_list if _def_name(entry) != MESSAGE_AGENT_TOOL_NAME] + name_set.discard(MESSAGE_AGENT_TOOL_NAME) + staged_agent = copy.copy(agent) + staged_agent.tools = tools_list + staged_agent.valid_tool_names = name_set + ensure_message_agent_tool(staged_agent) + except Exception: + logger.debug("Dynamic tool re-injection skipped", exc_info=True) + + def _reinject_post_build_tools(agent, tools_list: list, name_set: set) -> set: """Append memory-provider and context-engine tools onto the caller's staged ``tools_list`` / ``name_set`` (never the live agent attributes), mirroring ``agent_init``'s post-build