fix(bot-mode): retain message_agent across tool rebuilds (#102864)

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)

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