Clean MCP sessions and deduplicate tool notifications

Author: liuzqk

Server/src/services/tools/__init__.py

@@ -167,7 +167,9 @@ async def sync_tool_visibility_from_unity(
             len(enabled_tools), len(tools),
         )
 
-        PluginHub._sync_server_tool_visibility(enabled_tools)
+        tool_list_changed = PluginHub._tool_list_changed(enabled_tools)
+        if tool_list_changed:
+            PluginHub._sync_server_tool_visibility(enabled_tools)
 
         # Register custom (non-built-in) tools via CustomToolService.
         # The extended get_tool_states response includes is_built_in,
@@ -235,7 +237,7 @@ async def sync_tool_visibility_from_unity(
                 "Update MCPForUnity to enable custom tool sync in stdio mode."
             )
 
-        if notify:
+        if notify and tool_list_changed:
             await PluginHub._notify_mcp_tool_list_changed()
 
         # Build summary

Server/src/transport/plugin_hub.py

@@ -3,6 +3,8 @@
 from __future__ import annotations
 
 import asyncio
+import hashlib
+import json
 import logging
 import os
 import time
@@ -62,23 +64,39 @@ def _read_bounded_wait_env(name: str, default_s: float, max_s: float) -> float:
 _session_tracking_installed = False
 
 
+def _track_mcp_session(session: Any) -> None:
+    _active_mcp_sessions.add(session)
+
+
+def _untrack_mcp_session(session: Any) -> None:
+    _active_mcp_sessions.discard(session)
+
+
 def _install_session_tracking() -> None:
     """Patch *MiddlewareServerSession* to track active MCP client sessions."""
     global _session_tracking_installed
     if _session_tracking_installed:
         return
-    _session_tracking_installed = True
 
     from fastmcp.server.low_level import MiddlewareServerSession
 
     _original_aenter = MiddlewareServerSession.__aenter__
+    _original_aexit = MiddlewareServerSession.__aexit__
 
     async def _tracking_aenter(self):  # type: ignore[override]
         result = await _original_aenter(self)
-        _active_mcp_sessions.add(self)
+        _track_mcp_session(self)
         return result
 
+    async def _tracking_aexit(self, exc_type, exc_value, traceback):  # type: ignore[override]
+        try:
+            return await _original_aexit(self, exc_type, exc_value, traceback)
+        finally:
+            _untrack_mcp_session(self)
+
     MiddlewareServerSession.__aenter__ = _tracking_aenter  # type: ignore[assignment]
+    MiddlewareServerSession.__aexit__ = _tracking_aexit  # type: ignore[assignment]
+    _session_tracking_installed = True
 
 
 class PluginDisconnectedError(RuntimeError):
@@ -147,6 +165,7 @@ class PluginHub(WebSocketEndpoint):
     _last_pong: ClassVar[dict[str, float]] = {}
     # session_id -> ping task
     _ping_tasks: ClassVar[dict[str, asyncio.Task]] = {}
+    _published_tool_fingerprint: ClassVar[str | None] = None
 
     @classmethod
     def configure(
@@ -160,6 +179,7 @@ def configure(
         cls._loop = loop or asyncio.get_running_loop()
         # Ensure coordination primitives are bound to the configured loop
         cls._lock = asyncio.Lock()
+        cls._published_tool_fingerprint = None
         # Start tracking MCP client sessions for tool-change notifications
         if mcp is not None:
             _install_session_tracking()
@@ -530,13 +550,16 @@ async def _handle_register_tools(self, websocket: WebSocket, payload: RegisterTo
         logger.info(
             f"Registered {len(payload.tools)} tools for session {session_id}")
 
-        # Sync server-level FastMCP visibility so new MCP client sessions
-        # (e.g. new Claude Code conversations) see the correct tool set.
-        self._sync_server_tool_visibility(payload.tools)
+        if cls._tool_list_changed(payload.tools):
+            # Sync server-level FastMCP visibility so new MCP client sessions
+            # (e.g. new Claude Code conversations) see the correct tool set.
+            self._sync_server_tool_visibility(payload.tools)
 
-        # Notify any already-connected MCP clients (e.g. CC over stdio) that
-        # the tool list has changed so they re-fetch.
-        await cls._notify_mcp_tool_list_changed()
+            # Notify any already-connected MCP clients that the published tool
+            # schema changed so they re-fetch it.
+            await cls._notify_mcp_tool_list_changed()
+        else:
+            logger.debug("Unity tool schema unchanged; skipping tools/list_changed")
 
         try:
             from services.custom_tool_service import CustomToolService
@@ -555,6 +578,33 @@ async def _handle_register_tools(self, websocket: WebSocket, payload: RegisterTo
                 exc_info=exc,
             )
 
+    @classmethod
+    def _tool_list_changed(cls, registered_tools: list) -> bool:
+        serialized: list[str] = []
+        for tool in registered_tools:
+            if isinstance(tool, dict):
+                payload = tool
+            elif hasattr(tool, "model_dump"):
+                try:
+                    payload = tool.model_dump(mode="json")
+                except TypeError:
+                    payload = tool.model_dump()
+            else:
+                payload = vars(tool)
+
+            serialized.append(
+                json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str)
+            )
+
+        digest = hashlib.sha256(
+            "\n".join(sorted(serialized)).encode("utf-8")
+        ).hexdigest()
+        if digest == cls._published_tool_fingerprint:
+            return False
+
+        cls._published_tool_fingerprint = digest
+        return True
+
     @classmethod
     def _sync_server_tool_visibility(cls, registered_tools: list) -> None:
         """Sync FastMCP server-level tool group visibility to match Unity's state.
@@ -641,17 +691,20 @@ async def _notify_mcp_tool_list_changed(cls) -> None:
         sessions = list(_active_mcp_sessions)
         if not sessions:
             return
+        notified = 0
         for session in sessions:
             try:
                 await session.send_tool_list_changed()
+                notified += 1
             except Exception:
+                _untrack_mcp_session(session)
                 logger.debug(
-                    "Failed to notify MCP session of tool list change",
+                    "Failed to notify MCP session of tool list change; removed stale session",
                     exc_info=True,
                 )
         logger.info(
             "Sent tools/list_changed notification to %d MCP session(s)",
-            len(sessions),
+            notified,
         )
 
     async def _handle_command_result(self, payload: CommandResultMessage) -> None:

Server/tests/test_plugin_hub_session_lifecycle.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+import pytest
+
+from transport import plugin_hub
+from transport.plugin_hub import PluginHub
+
+
+class _FakeSession:
+    def __init__(self, fail: bool = False) -> None:
+        self.fail = fail
+        self.notifications = 0
+
+    async def send_tool_list_changed(self) -> None:
+        if self.fail:
+            raise RuntimeError("closed")
+        self.notifications += 1
+
+
+@pytest.fixture(autouse=True)
+def _reset_session_tracking() -> None:
+    plugin_hub._active_mcp_sessions.clear()
+    PluginHub._published_tool_fingerprint = None
+    yield
+    plugin_hub._active_mcp_sessions.clear()
+    PluginHub._published_tool_fingerprint = None
+
+
+def test_session_tracking_removes_exited_session() -> None:
+    session = _FakeSession()
+
+    plugin_hub._track_mcp_session(session)
+    assert session in plugin_hub._active_mcp_sessions
+
+    plugin_hub._untrack_mcp_session(session)
+    assert session not in plugin_hub._active_mcp_sessions
+
+
+def test_twenty_connect_disconnect_cycles_return_to_baseline() -> None:
+    sessions = [_FakeSession() for _ in range(20)]
+
+    for session in sessions:
+        plugin_hub._track_mcp_session(session)
+        plugin_hub._untrack_mcp_session(session)
+
+    assert list(plugin_hub._active_mcp_sessions) == []
+
+
+@pytest.mark.asyncio
+async def test_notification_prunes_closed_sessions() -> None:
+    active = _FakeSession()
+    closed = _FakeSession(fail=True)
+    plugin_hub._track_mcp_session(active)
+    plugin_hub._track_mcp_session(closed)
+
+    await PluginHub._notify_mcp_tool_list_changed()
+
+    assert active.notifications == 1
+    assert active in plugin_hub._active_mcp_sessions
+    assert closed not in plugin_hub._active_mcp_sessions
+
+
+@pytest.mark.asyncio
+async def test_unchanged_tools_do_not_rebroadcast() -> None:
+    sessions = [_FakeSession(), _FakeSession()]
+    for session in sessions:
+        plugin_hub._track_mcp_session(session)
+
+    tools = [{"name": "read_console", "description": "Console"}]
+    assert PluginHub._tool_list_changed(tools) is True
+    await PluginHub._notify_mcp_tool_list_changed()
+    assert PluginHub._tool_list_changed(tools) is False
+
+    assert [session.notifications for session in sessions] == [1, 1]
+
+
+def test_tool_fingerprint_deduplicates_reordered_payload() -> None:
+    first = [
+        {"name": "manage_scene", "description": "Scene"},
+        {"name": "read_console", "description": "Console"},
+    ]
+    reordered = list(reversed(first))
+    changed_schema = [
+        {"name": "manage_scene", "description": "Scene changed"},
+        {"name": "read_console", "description": "Console"},
+    ]
+
+    assert PluginHub._tool_list_changed(first) is True
+    assert PluginHub._tool_list_changed(reordered) is False
+    assert PluginHub._tool_list_changed(changed_schema) is True