diff --git a/Server/src/services/tools/__init__.py b/Server/src/services/tools/__init__.py index b587663d2..33c168aff 100644 --- a/Server/src/services/tools/__init__.py +++ b/Server/src/services/tools/__init__.py @@ -235,8 +235,10 @@ async def sync_tool_visibility_from_unity( "Update MCPForUnity to enable custom tool sync in stdio mode." ) - if notify: - await PluginHub._notify_mcp_tool_list_changed() + await PluginHub._record_and_notify_tool_list_change( + enabled_tools, + notify=notify, + ) # Build summary from services.registry import get_group_tool_names diff --git a/Server/src/transport/plugin_hub.py b/Server/src/transport/plugin_hub.py index 45aeff4a3..bc312637f 100644 --- a/Server/src/transport/plugin_hub.py +++ b/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 @@ -10,6 +12,7 @@ import weakref from typing import TYPE_CHECKING, Any, ClassVar +import anyio from starlette.endpoints import WebSocketEndpoint from starlette.websockets import WebSocket, WebSocketState @@ -60,25 +63,57 @@ def _read_bounded_wait_env(name: str, default_s: float, max_s: float) -> float: # session so we can send ``tools/list_changed`` notifications later. _active_mcp_sessions: weakref.WeakSet = weakref.WeakSet() _session_tracking_installed = False +_SESSION_TRACKING_STATE_ATTRIBUTE = "_mcpforunity_active_sessions" + + +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 + global _active_mcp_sessions, _session_tracking_installed from fastmcp.server.low_level import MiddlewareServerSession + existing_sessions = getattr( + MiddlewareServerSession, + _SESSION_TRACKING_STATE_ATTRIBUTE, + None, + ) + if isinstance(existing_sessions, weakref.WeakSet): + _active_mcp_sessions = existing_sessions + _session_tracking_installed = True + return + if _session_tracking_installed: + return + _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] + setattr( + MiddlewareServerSession, + _SESSION_TRACKING_STATE_ATTRIBUTE, + _active_mcp_sessions, + ) + _session_tracking_installed = True class PluginDisconnectedError(RuntimeError): @@ -147,6 +182,9 @@ 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 + _pending_tool_list_notifications: ClassVar[weakref.WeakSet] = weakref.WeakSet() + _TOOL_LIST_NOTIFY_TIMEOUT_SECONDS = 1.0 @classmethod def configure( @@ -160,6 +198,8 @@ 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 + cls._pending_tool_list_notifications.clear() # Start tracking MCP client sessions for tool-change notifications if mcp is not None: _install_session_tracking() @@ -534,10 +574,6 @@ async def _handle_register_tools(self, websocket: WebSocket, payload: RegisterTo # (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() - try: from services.custom_tool_service import CustomToolService @@ -555,8 +591,104 @@ async def _handle_register_tools(self, websocket: WebSocket, payload: RegisterTo exc_info=exc, ) + await cls._record_and_notify_tool_list_change(payload.tools) + + @classmethod + def _tool_fingerprint(cls, tools: list) -> str: + serialized: list[str] = [] + for tool in tools: + if hasattr(tool, "to_mcp_tool"): + tool = tool.to_mcp_tool() + + 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, + ) + ) + + return hashlib.sha256( + "\n".join(sorted(serialized)).encode("utf-8") + ).hexdigest() + + @classmethod + async def _record_published_tool_list(cls, registered_tools: list) -> bool: + """Record a stable fingerprint after the publication transaction succeeds.""" + published_tools = registered_tools + if cls._mcp is not None: + try: + published_tools = list(await cls._mcp.list_tools()) + except Exception: + logger.warning( + "Failed to inspect the published FastMCP tool list; " + "leaving its fingerprint unchanged for retry", + exc_info=True, + ) + return False + + try: + digest = cls._tool_fingerprint(published_tools) + except Exception: + logger.warning( + "Failed to fingerprint the published FastMCP tool list; " + "leaving its fingerprint unchanged for retry", + exc_info=True, + ) + return False + if digest == cls._published_tool_fingerprint: + return False + + cls._published_tool_fingerprint = digest + return True + + @classmethod + async def _record_and_notify_tool_list_change( + cls, + registered_tools: list, + notify: bool = True, + ) -> bool: + changed = await cls._record_published_tool_list(registered_tools) + if not notify: + return changed + + if changed: + cls._pending_tool_list_notifications.update( + list(_active_mcp_sessions) + ) + + if cls._pending_tool_list_notifications: + # Retry only sessions that have not yet observed the latest + # published schema, avoiding duplicate notifications to peers that + # already succeeded. + await cls._notify_mcp_tool_list_changed( + list(cls._pending_tool_list_notifications) + ) + elif changed: + logger.debug( + "Published tool schema changed but no MCP sessions are active; " + "skipping tools/list_changed" + ) + else: + logger.debug( + "Published tool schema unchanged; skipping tools/list_changed" + ) + + return changed + @classmethod - def _sync_server_tool_visibility(cls, registered_tools: list) -> None: + def _sync_server_tool_visibility(cls, registered_tools: list) -> bool: """Sync FastMCP server-level tool group visibility to match Unity's state. When Unity sends ``register_tools``, some groups may have been toggled @@ -573,7 +705,7 @@ def _sync_server_tool_visibility(cls, registered_tools: list) -> None: """ mcp = cls._mcp if mcp is None: - return + return True try: from services.registry import get_group_tool_names, TOOL_GROUPS @@ -621,14 +753,19 @@ def _sync_server_tool_visibility(cls, registered_tools: list) -> None: len(mcp._transforms), cls._unity_transform_start or 0, ) + return True except Exception: logger.debug( "Failed to sync server-level tool visibility", exc_info=True, ) + return False @classmethod - async def _notify_mcp_tool_list_changed(cls) -> None: + async def _notify_mcp_tool_list_changed( + cls, + sessions: list[Any] | None = None, + ) -> None: """Send ``tools/list_changed`` to every connected MCP client session. After server-level tool visibility is updated (e.g. when Unity reports @@ -638,20 +775,46 @@ async def _notify_mcp_tool_list_changed(cls) -> None: transforms but do **not** push notifications to already-connected sessions — we do that here. """ - sessions = list(_active_mcp_sessions) + sessions = list(_active_mcp_sessions) if sessions is None else sessions if not sessions: return - for session in sessions: + + async def notify_session(session: Any) -> bool: try: - await session.send_tool_list_changed() + await asyncio.wait_for( + session.send_tool_list_changed(), + timeout=cls._TOOL_LIST_NOTIFY_TIMEOUT_SECONDS, + ) + cls._pending_tool_list_notifications.discard(session) + return True + except (anyio.BrokenResourceError, anyio.ClosedResourceError): + _untrack_mcp_session(session) + cls._pending_tool_list_notifications.discard(session) + logger.debug( + "Failed to notify MCP session of tool list change; removed stale session", + exc_info=True, + ) + except asyncio.TimeoutError: + cls._pending_tool_list_notifications.add(session) + logger.debug( + "Timed out notifying MCP session of tool list change; " + "keeping session for the next publication", + exc_info=True, + ) except Exception: + cls._pending_tool_list_notifications.add(session) logger.debug( - "Failed to notify MCP session of tool list change", + "Failed to notify MCP session of tool list change; " + "keeping session because closure was not confirmed", exc_info=True, ) + return False + + notified = sum(await asyncio.gather(*(notify_session(session) for session in sessions))) logger.info( - "Sent tools/list_changed notification to %d MCP session(s)", - len(sessions), + "Sent tools/list_changed notification to %d MCP session(s); %d active", + notified, + len(_active_mcp_sessions), ) async def _handle_command_result(self, payload: CommandResultMessage) -> None: diff --git a/Server/tests/integration/test_stdio_custom_tool_sync.py b/Server/tests/integration/test_stdio_custom_tool_sync.py index d7dd35765..267ad5d60 100644 --- a/Server/tests/integration/test_stdio_custom_tool_sync.py +++ b/Server/tests/integration/test_stdio_custom_tool_sync.py @@ -130,6 +130,32 @@ async def test_sync_skips_builtin_tools(): mock_get_instance.assert_not_called() +@pytest.mark.asyncio +async def test_sync_records_published_tools_when_visibility_update_fails(): + """Pending notifications should retry even if visibility sync is transiently unavailable.""" + response = _make_unity_response([BUILTIN_TOOL]) + + with patch( + "transport.legacy.unity_connection.async_send_command_with_retry", + new_callable=AsyncMock, + return_value=response, + ), patch( + "transport.plugin_hub.PluginHub._sync_server_tool_visibility", + return_value=False, + ), patch( + "transport.plugin_hub.PluginHub._record_and_notify_tool_list_change", + new_callable=AsyncMock, + ) as mock_record_and_notify, patch( + "services.custom_tool_service.CustomToolService.get_instance", + ) as mock_get_instance: + from services.tools import sync_tool_visibility_from_unity + result = await sync_tool_visibility_from_unity(notify=True) + + assert result["synced"] is True + mock_get_instance.assert_not_called() + mock_record_and_notify.assert_awaited_once() + + @pytest.mark.asyncio async def test_sync_skips_when_no_extended_metadata(): """When Unity returns old-format data (no is_built_in), skip custom tool registration.""" diff --git a/Server/tests/test_plugin_hub_session_lifecycle.py b/Server/tests/test_plugin_hub_session_lifecycle.py new file mode 100644 index 000000000..12ed6adae --- /dev/null +++ b/Server/tests/test_plugin_hub_session_lifecycle.py @@ -0,0 +1,220 @@ +from __future__ import annotations + +import asyncio +import sys +import types +import weakref + +import anyio +import pytest + +from transport import plugin_hub +from transport.plugin_hub import PluginHub + + +class _FakeSession: + def __init__( + self, + failure: Exception | None = None, + delay: float = 0, + ) -> None: + self.failure = failure + self.delay = delay + self.notifications = 0 + + async def send_tool_list_changed(self) -> None: + if self.delay: + await asyncio.sleep(self.delay) + if self.failure: + raise self.failure + self.notifications += 1 + + +@pytest.fixture(autouse=True) +def _reset_session_tracking() -> None: + plugin_hub._active_mcp_sessions.clear() + PluginHub._published_tool_fingerprint = None + PluginHub._pending_tool_list_notifications.clear() + yield + plugin_hub._active_mcp_sessions.clear() + PluginHub._published_tool_fingerprint = None + PluginHub._pending_tool_list_notifications.clear() + + +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(failure=anyio.ClosedResourceError()) + 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_notification_keeps_session_after_unconfirmed_failure() -> None: + transient = _FakeSession(failure=RuntimeError("temporary")) + plugin_hub._track_mcp_session(transient) + + await PluginHub._notify_mcp_tool_list_changed() + + assert transient in plugin_hub._active_mcp_sessions + assert transient in PluginHub._pending_tool_list_notifications + + +@pytest.mark.asyncio +async def test_notification_timeout_is_bounded_without_pruning_live_session( + monkeypatch: pytest.MonkeyPatch, +) -> None: + slow = _FakeSession(delay=0.1) + plugin_hub._track_mcp_session(slow) + monkeypatch.setattr(PluginHub, "_TOOL_LIST_NOTIFY_TIMEOUT_SECONDS", 0.01) + + await PluginHub._notify_mcp_tool_list_changed() + + assert slow.notifications == 0 + assert slow in plugin_hub._active_mcp_sessions + assert slow in PluginHub._pending_tool_list_notifications + + +@pytest.mark.asyncio +async def test_patched_session_context_is_tracked_symmetrically( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _MiddlewareServerSession: + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc_value, traceback): + return None + + low_level = types.ModuleType("fastmcp.server.low_level") + low_level.MiddlewareServerSession = _MiddlewareServerSession + monkeypatch.setitem(sys.modules, "fastmcp.server.low_level", low_level) + monkeypatch.setattr(plugin_hub, "_session_tracking_installed", False) + + plugin_hub._install_session_tracking() + session = _MiddlewareServerSession() + await session.__aenter__() + preserved_sessions = plugin_hub._active_mcp_sessions + + assert session in preserved_sessions + + monkeypatch.setattr(plugin_hub, "_active_mcp_sessions", weakref.WeakSet()) + monkeypatch.setattr(plugin_hub, "_session_tracking_installed", False) + plugin_hub._install_session_tracking() + + assert plugin_hub._active_mcp_sessions is preserved_sessions + await session.__aexit__(None, None, None) + assert session 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 await PluginHub._record_and_notify_tool_list_change(tools) is True + assert [session.notifications for session in sessions] == [1, 1] + + assert await PluginHub._record_and_notify_tool_list_change(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_fingerprint(first) == PluginHub._tool_fingerprint(reordered) + assert PluginHub._tool_fingerprint(first) != PluginHub._tool_fingerprint(changed_schema) + + +@pytest.mark.asyncio +async def test_fingerprint_is_not_committed_when_server_inspection_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _FailingMcp: + async def list_tools(self) -> list: + raise RuntimeError("inspection failed") + + monkeypatch.setattr(PluginHub, "_mcp", _FailingMcp()) + + assert await PluginHub._record_published_tool_list([{"name": "one"}]) is False + assert PluginHub._published_tool_fingerprint is None + + +@pytest.mark.asyncio +async def test_actual_server_tool_registration_changes_publication_fingerprint( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _FakeMcp: + def __init__(self) -> None: + self.tools: list[dict] = [] + + async def list_tools(self) -> list[dict]: + return self.tools + + mcp = _FakeMcp() + monkeypatch.setattr(PluginHub, "_mcp", mcp) + desired = [{"name": "custom"}] + + assert await PluginHub._record_published_tool_list(desired) is True + mcp.tools.append({"name": "custom"}) + assert await PluginHub._record_published_tool_list(desired) is True + + +@pytest.mark.asyncio +async def test_failed_notification_retries_only_pending_session() -> None: + delivered = _FakeSession() + transient = _FakeSession(failure=RuntimeError("temporary")) + plugin_hub._track_mcp_session(delivered) + plugin_hub._track_mcp_session(transient) + tools = [{"name": "read_console", "description": "Console"}] + + assert await PluginHub._record_and_notify_tool_list_change(tools) is True + + assert delivered.notifications == 1 + assert transient.notifications == 0 + assert delivered not in PluginHub._pending_tool_list_notifications + assert transient in PluginHub._pending_tool_list_notifications + + transient.failure = None + assert await PluginHub._record_and_notify_tool_list_change(tools) is False + + assert delivered.notifications == 1 + assert transient.notifications == 1 + assert list(PluginHub._pending_tool_list_notifications) == []