From c513072213ae0e577c9ef249b784cd32fdc0b822 Mon Sep 17 00:00:00 2001 From: whatevertogo Date: Wed, 11 Feb 2026 21:41:05 +0800 Subject: [PATCH 01/10] fix: preserve tool enabled/disabled state across Unity restarts Fixes #709 - Remove logic in ToolDiscoveryService.EnsurePreferenceInitialized that forced built-in tools to be re-enabled on every Unity launch - Add ReregisterToolsAsync() to IMcpTransportClient for dynamic tool list updates when toggles change in the UI - Implement tool reregistration in WebSocketTransportClient - Add ReregisterToolsAsync() to McpToolsSection to trigger updates when tool toggles are flipped - Add unit tests for ToolDiscoveryService preference persistence Co-Authored-By: Claude Opus 4.6 --- .../Editor/Services/ToolDiscoveryService.cs | 10 -- .../Services/Transport/IMcpTransportClient.cs | 1 + .../Services/Transport/TransportManager.cs | 24 ++-- .../Transports/StdioTransportClient.cs | 7 ++ .../Transports/WebSocketTransportClient.cs | 23 ++++ .../Components/Tools/McpToolsSection.cs | 29 +++++ .../Services/ToolDiscoveryServiceTests.cs | 115 ++++++++++++++++++ .../ToolDiscoveryServiceTests.cs.meta | 11 ++ .../UnityMCPTests/Assets/Tests/Editor.meta | 8 ++ .../Assets/Tests/Editor/State.meta | 8 ++ 10 files changed, 216 insertions(+), 20 deletions(-) create mode 100644 TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/ToolDiscoveryServiceTests.cs create mode 100644 TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/ToolDiscoveryServiceTests.cs.meta create mode 100644 TestProjects/UnityMCPTests/Assets/Tests/Editor.meta create mode 100644 TestProjects/UnityMCPTests/Assets/Tests/Editor/State.meta diff --git a/MCPForUnity/Editor/Services/ToolDiscoveryService.cs b/MCPForUnity/Editor/Services/ToolDiscoveryService.cs index b5b86c0a2..ac7cc1455 100644 --- a/MCPForUnity/Editor/Services/ToolDiscoveryService.cs +++ b/MCPForUnity/Editor/Services/ToolDiscoveryService.cs @@ -226,16 +226,6 @@ private void EnsurePreferenceInitialized(ToolMetadata metadata) { bool defaultValue = metadata.AutoRegister || metadata.IsBuiltIn; EditorPrefs.SetBool(key, defaultValue); - return; - } - - if (metadata.IsBuiltIn && !metadata.AutoRegister) - { - bool currentValue = EditorPrefs.GetBool(key, metadata.AutoRegister); - if (currentValue == metadata.AutoRegister) - { - EditorPrefs.SetBool(key, true); - } } } diff --git a/MCPForUnity/Editor/Services/Transport/IMcpTransportClient.cs b/MCPForUnity/Editor/Services/Transport/IMcpTransportClient.cs index 3d8584fd9..44f1d3b8c 100644 --- a/MCPForUnity/Editor/Services/Transport/IMcpTransportClient.cs +++ b/MCPForUnity/Editor/Services/Transport/IMcpTransportClient.cs @@ -14,5 +14,6 @@ public interface IMcpTransportClient Task StartAsync(); Task StopAsync(); Task VerifyAsync(); + Task ReregisterToolsAsync(); } } diff --git a/MCPForUnity/Editor/Services/Transport/TransportManager.cs b/MCPForUnity/Editor/Services/Transport/TransportManager.cs index 1204e7014..b544273b4 100644 --- a/MCPForUnity/Editor/Services/Transport/TransportManager.cs +++ b/MCPForUnity/Editor/Services/Transport/TransportManager.cs @@ -42,16 +42,6 @@ private IMcpTransportClient GetOrCreateClient(TransportMode mode) }; } - private IMcpTransportClient GetClient(TransportMode mode) - { - return mode switch - { - TransportMode.Http => _httpClient, - TransportMode.Stdio => _stdioClient, - _ => throw new ArgumentOutOfRangeException(nameof(mode), mode, "Unsupported transport mode"), - }; - } - public async Task StartAsync(TransportMode mode) { IMcpTransportClient client = GetOrCreateClient(mode); @@ -128,6 +118,20 @@ public TransportState GetState(TransportMode mode) public bool IsRunning(TransportMode mode) => GetState(mode).IsConnected; + /// + /// Gets the active transport client for the specified mode. + /// Returns null if the client hasn't been created yet. + /// + public IMcpTransportClient GetClient(TransportMode mode) + { + return mode switch + { + TransportMode.Http => _httpClient, + TransportMode.Stdio => _stdioClient, + _ => throw new ArgumentOutOfRangeException(nameof(mode), mode, "Unsupported transport mode"), + }; + } + private void UpdateState(TransportMode mode, TransportState state) { switch (mode) diff --git a/MCPForUnity/Editor/Services/Transport/Transports/StdioTransportClient.cs b/MCPForUnity/Editor/Services/Transport/Transports/StdioTransportClient.cs index ea3ed1a22..9dadc4336 100644 --- a/MCPForUnity/Editor/Services/Transport/Transports/StdioTransportClient.cs +++ b/MCPForUnity/Editor/Services/Transport/Transports/StdioTransportClient.cs @@ -46,5 +46,12 @@ public Task VerifyAsync() return Task.FromResult(running); } + public Task ReregisterToolsAsync() + { + // Stdio transport doesn't support dynamic tool reregistration + // Tools are registered at server startup + return Task.CompletedTask; + } + } } diff --git a/MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs b/MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs index b94c0836f..6d2f76838 100644 --- a/MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs +++ b/MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs @@ -506,6 +506,29 @@ private async Task SendRegisterToolsAsync(CancellationToken token) McpLog.Info($"[WebSocket] Sent {tools.Count} tools registration", false); } + public async Task ReregisterToolsAsync() + { + if (!IsConnected || _lifecycleCts == null) + { + McpLog.Warn("[WebSocket] Cannot reregister tools: not connected"); + return; + } + + try + { + await SendRegisterToolsAsync(_lifecycleCts.Token).ConfigureAwait(false); + McpLog.Info("[WebSocket] Tool reregistration completed", false); + } + catch (System.OperationCanceledException) + { + McpLog.Warn("[WebSocket] Tool reregistration cancelled"); + } + catch (System.Exception ex) + { + McpLog.Error($"[WebSocket] Tool reregistration failed: {ex.Message}"); + } + } + private async Task HandleExecuteAsync(JObject payload, CancellationToken token) { string commandId = payload.Value("id"); diff --git a/MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs b/MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs index dc5a3eabf..8644a7ad6 100644 --- a/MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs +++ b/MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs @@ -1,9 +1,11 @@ using System; using System.Collections.Generic; using System.Linq; +using System.Threading; using MCPForUnity.Editor.Constants; using MCPForUnity.Editor.Helpers; using MCPForUnity.Editor.Services; +using MCPForUnity.Editor.Services.Transport; using MCPForUnity.Editor.Tools; using UnityEditor; using UnityEngine.UIElements; @@ -231,6 +233,30 @@ private void HandleToggleChange(ToolMetadata tool, bool enabled, bool updateSumm { UpdateSummary(); } + + // Trigger tool reregistration with connected MCP server + ReregisterToolsAsync(); + } + + private void ReregisterToolsAsync() + { + // Fire and forget - don't block UI + ThreadPool.QueueUserWorkItem(_ => + { + try + { + var transportManager = MCPServiceLocator.TransportManager; + var client = transportManager.GetClient(TransportMode.Http); + if (client != null && client.IsConnected) + { + client.ReregisterToolsAsync().Wait(); + } + } + catch (Exception ex) + { + McpLog.Warn($"Failed to reregister tools: {ex.Message}"); + } + }); } private void SetAllToolsState(bool enabled) @@ -253,6 +279,9 @@ private void SetAllToolsState(bool enabled) } UpdateSummary(); + + // Trigger tool reregistration after bulk change + ReregisterToolsAsync(); } private void UpdateSummary() diff --git a/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/ToolDiscoveryServiceTests.cs b/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/ToolDiscoveryServiceTests.cs new file mode 100644 index 000000000..3ae8b7e35 --- /dev/null +++ b/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/ToolDiscoveryServiceTests.cs @@ -0,0 +1,115 @@ +using NUnit.Framework; +using MCPForUnity.Editor.Constants; +using MCPForUnity.Editor.Services; +using UnityEditor; + +namespace MCPForUnity.Editor.Tests.EditMode.Services +{ + [TestFixture] + public class ToolDiscoveryServiceTests + { + private const string TestToolName = "test_tool_for_testing"; + + [SetUp] + public void SetUp() + { + // Clean up any test preferences + string testKey = EditorPrefKeys.ToolEnabledPrefix + TestToolName; + if (EditorPrefs.HasKey(testKey)) + { + EditorPrefs.DeleteKey(testKey); + } + } + + [TearDown] + public void TearDown() + { + // Clean up test preferences after each test + string testKey = EditorPrefKeys.ToolEnabledPrefix + TestToolName; + if (EditorPrefs.HasKey(testKey)) + { + EditorPrefs.DeleteKey(testKey); + } + } + + [Test] + public void SetToolEnabled_WritesToEditorPrefs() + { + // Arrange + var service = new ToolDiscoveryService(); + + // Act + service.SetToolEnabled(TestToolName, false); + + // Assert + string key = EditorPrefKeys.ToolEnabledPrefix + TestToolName; + Assert.IsTrue(EditorPrefs.HasKey(key), "Preference key should exist after SetToolEnabled"); + Assert.IsFalse(EditorPrefs.GetBool(key, true), "Preference should be set to false"); + } + + [Test] + public void IsToolEnabled_ReturnsTrue_WhenNoPreferenceExistsAndToolIsBuiltIn() + { + // Arrange - Ensure no preference exists + string key = EditorPrefKeys.ToolEnabledPrefix + TestToolName; + if (EditorPrefs.HasKey(key)) + { + EditorPrefs.DeleteKey(key); + } + + var service = new ToolDiscoveryService(); + + // Act - For a non-existent tool, IsToolEnabled should return false + // (since metadata.AutoRegister defaults to false for non-existent tools) + bool result = service.IsToolEnabled(TestToolName); + + // Assert - Non-existent tools return false (no metadata found) + Assert.IsFalse(result, "Non-existent tool should return false"); + } + + [Test] + public void IsToolEnabled_ReturnsStoredValue_WhenPreferenceExists() + { + // Arrange + string key = EditorPrefKeys.ToolEnabledPrefix + TestToolName; + EditorPrefs.SetBool(key, false); // Store false value + var service = new ToolDiscoveryService(); + + // Act + bool result = service.IsToolEnabled(TestToolName); + + // Assert + Assert.IsFalse(result, "Should return the stored preference value (false)"); + } + + [Test] + public void IsToolEnabled_ReturnsTrue_WhenPreferenceSetToTrue() + { + // Arrange + string key = EditorPrefKeys.ToolEnabledPrefix + TestToolName; + EditorPrefs.SetBool(key, true); + var service = new ToolDiscoveryService(); + + // Act + bool result = service.IsToolEnabled(TestToolName); + + // Assert + Assert.IsTrue(result, "Should return the stored preference value (true)"); + } + + [Test] + public void ToolToggle_PersistsAcrossServiceInstances() + { + // Arrange + var service1 = new ToolDiscoveryService(); + service1.SetToolEnabled(TestToolName, false); + + // Act - Create a new service instance + var service2 = new ToolDiscoveryService(); + bool result = service2.IsToolEnabled(TestToolName); + + // Assert - The disabled state should persist + Assert.IsFalse(result, "Tool state should persist across service instances"); + } + } +} diff --git a/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/ToolDiscoveryServiceTests.cs.meta b/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/ToolDiscoveryServiceTests.cs.meta new file mode 100644 index 000000000..cea36ebb1 --- /dev/null +++ b/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/ToolDiscoveryServiceTests.cs.meta @@ -0,0 +1,11 @@ +fileFormatVersion: 2 +guid: 7a8b9c0d1e2f3a4b5c6d7e8f9a0b1c2d +MonoImporter: + externalObjects: {} + serializedVersion: 2 + defaultReferences: [] + executionOrder: 0 + icon: {instanceID: 0} + userData: + assetBundleName: + assetBundleVariant: diff --git a/TestProjects/UnityMCPTests/Assets/Tests/Editor.meta b/TestProjects/UnityMCPTests/Assets/Tests/Editor.meta new file mode 100644 index 000000000..b6e169d0d --- /dev/null +++ b/TestProjects/UnityMCPTests/Assets/Tests/Editor.meta @@ -0,0 +1,8 @@ +fileFormatVersion: 2 +guid: ea478f98f983e4c4daafd9a07cf03e3b +folderAsset: yes +DefaultImporter: + externalObjects: {} + userData: + assetBundleName: + assetBundleVariant: diff --git a/TestProjects/UnityMCPTests/Assets/Tests/Editor/State.meta b/TestProjects/UnityMCPTests/Assets/Tests/Editor/State.meta new file mode 100644 index 000000000..a4a8007ae --- /dev/null +++ b/TestProjects/UnityMCPTests/Assets/Tests/Editor/State.meta @@ -0,0 +1,8 @@ +fileFormatVersion: 2 +guid: 778e8246ab09497419985acaee595f6c +folderAsset: yes +DefaultImporter: + externalObjects: {} + userData: + assetBundleName: + assetBundleVariant: From 09c61e10befde8aafc8a82e9cb4b35819dda13e0 Mon Sep 17 00:00:00 2001 From: whatevertogo Date: Wed, 11 Feb 2026 22:04:38 +0800 Subject: [PATCH 02/10] fix: avoid duplicate reregistration and harden tool preference tests --- .../Components/Tools/McpToolsSection.cs | 32 ++++++++++---- .../Services/ToolDiscoveryServiceTests.cs | 42 ++++++++++++++++++- 2 files changed, 66 insertions(+), 8 deletions(-) diff --git a/MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs b/MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs index 8644a7ad6..687fbf2eb 100644 --- a/MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs +++ b/MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs @@ -225,7 +225,11 @@ private VisualElement CreateToolRow(ToolMetadata tool) return row; } - private void HandleToggleChange(ToolMetadata tool, bool enabled, bool updateSummary = true) + private void HandleToggleChange( + ToolMetadata tool, + bool enabled, + bool updateSummary = true, + bool reregisterTools = true) { MCPServiceLocator.ToolDiscovery.SetToolEnabled(tool.Name, enabled); @@ -234,8 +238,11 @@ private void HandleToggleChange(ToolMetadata tool, bool enabled, bool updateSumm UpdateSummary(); } - // Trigger tool reregistration with connected MCP server - ReregisterToolsAsync(); + if (reregisterTools) + { + // Trigger tool reregistration with connected MCP server + ReregisterToolsAsync(); + } } private void ReregisterToolsAsync() @@ -261,11 +268,18 @@ private void ReregisterToolsAsync() private void SetAllToolsState(bool enabled) { + bool hasChanges = false; + foreach (var tool in allTools) { if (!toolToggleMap.TryGetValue(tool.Name, out var toggle)) { - MCPServiceLocator.ToolDiscovery.SetToolEnabled(tool.Name, enabled); + bool currentEnabled = MCPServiceLocator.ToolDiscovery.IsToolEnabled(tool.Name); + if (currentEnabled != enabled) + { + MCPServiceLocator.ToolDiscovery.SetToolEnabled(tool.Name, enabled); + hasChanges = true; + } continue; } @@ -275,13 +289,17 @@ private void SetAllToolsState(bool enabled) } toggle.SetValueWithoutNotify(enabled); - HandleToggleChange(tool, enabled, updateSummary: false); + HandleToggleChange(tool, enabled, updateSummary: false, reregisterTools: false); + hasChanges = true; } UpdateSummary(); - // Trigger tool reregistration after bulk change - ReregisterToolsAsync(); + if (hasChanges) + { + // Trigger a single reregistration after bulk change + ReregisterToolsAsync(); + } } private void UpdateSummary() diff --git a/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/ToolDiscoveryServiceTests.cs b/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/ToolDiscoveryServiceTests.cs index 3ae8b7e35..7e267c225 100644 --- a/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/ToolDiscoveryServiceTests.cs +++ b/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/ToolDiscoveryServiceTests.cs @@ -1,3 +1,4 @@ +using System.Linq; using NUnit.Framework; using MCPForUnity.Editor.Constants; using MCPForUnity.Editor.Services; @@ -48,7 +49,7 @@ public void SetToolEnabled_WritesToEditorPrefs() } [Test] - public void IsToolEnabled_ReturnsTrue_WhenNoPreferenceExistsAndToolIsBuiltIn() + public void IsToolEnabled_ReturnsFalse_WhenToolDoesNotExist() { // Arrange - Ensure no preference exists string key = EditorPrefKeys.ToolEnabledPrefix + TestToolName; @@ -111,5 +112,44 @@ public void ToolToggle_PersistsAcrossServiceInstances() // Assert - The disabled state should persist Assert.IsFalse(result, "Tool state should persist across service instances"); } + + [Test] + public void DiscoverAllTools_DoesNotOverrideStoredFalse_ForBuiltInAutoRegisterFalseTool() + { + // Arrange + var service = new ToolDiscoveryService(); + var builtInTool = service.DiscoverAllTools() + .FirstOrDefault(tool => tool.IsBuiltIn && !tool.AutoRegister); + + Assert.IsNotNull(builtInTool, "Expected at least one built-in tool with AutoRegister=false."); + + string key = EditorPrefKeys.ToolEnabledPrefix + builtInTool.Name; + bool hadOriginalKey = EditorPrefs.HasKey(key); + bool originalValue = hadOriginalKey && EditorPrefs.GetBool(key, true); + + try + { + EditorPrefs.SetBool(key, false); + service.InvalidateCache(); + + // Act + service.DiscoverAllTools(); + bool enabled = service.IsToolEnabled(builtInTool.Name); + + // Assert + Assert.IsFalse(enabled, $"Built-in tool '{builtInTool.Name}' should remain disabled when preference is false."); + } + finally + { + if (hadOriginalKey) + { + EditorPrefs.SetBool(key, originalValue); + } + else + { + EditorPrefs.DeleteKey(key); + } + } + } } } From 10754e581906db1f8048c2cf039e9709d72de8a4 Mon Sep 17 00:00:00 2001 From: whatevertogo Date: Wed, 11 Feb 2026 22:40:26 +0800 Subject: [PATCH 03/10] feat: implement tool filtering based on Unity session state in on_list_tools() --- .../transport/unity_instance_middleware.py | 142 ++++++++++++++++++ .../tests/test_transport_characterization.py | 74 +++++++++ 2 files changed, 216 insertions(+) diff --git a/Server/src/transport/unity_instance_middleware.py b/Server/src/transport/unity_instance_middleware.py index 7adbe31ab..d172a304d 100644 --- a/Server/src/transport/unity_instance_middleware.py +++ b/Server/src/transport/unity_instance_middleware.py @@ -55,6 +55,43 @@ def __init__(self): super().__init__() self._active_by_key: dict[str, str] = {} self._lock = RLock() + self._unity_managed_tool_names = { + "batch_execute", + "execute_menu_item", + "find_gameobjects", + "get_test_job", + "manage_asset", + "manage_components", + "manage_editor", + "manage_gameobject", + "manage_material", + "manage_prefabs", + "manage_scene", + "manage_script", + "manage_scriptable_object", + "manage_shader", + "manage_texture", + "manage_vfx", + "read_console", + "refresh_unity", + "run_tests", + } + self._tool_alias_to_unity_target = { + # Server-side script helpers route to Unity's manage_script command. + "apply_text_edits": "manage_script", + "create_script": "manage_script", + "delete_script": "manage_script", + "find_in_file": "manage_script", + "get_sha": "manage_script", + "script_apply_edits": "manage_script", + "validate_script": "manage_script", + } + self._server_only_tool_names = { + "debug_request_context", + "execute_custom_tool", + "manage_script_capabilities", + "set_active_instance", + } def get_session_key(self, ctx) -> str: """ @@ -260,3 +297,108 @@ async def on_read_resource(self, context: MiddlewareContext, call_next): """Inject active Unity instance into resource context if available.""" await self._inject_unity_instance(context) return await call_next(context) + + async def on_list_tools(self, context: MiddlewareContext, call_next): + """Filter MCP tool listing to the Unity-enabled set when session data is available.""" + await self._inject_unity_instance(context) + tools = await call_next(context) + + if not self._should_filter_tool_listing(): + return tools + + enabled_tool_names = await self._resolve_enabled_tool_names_for_context(context) + if not enabled_tool_names: + return tools + + filtered = [] + for tool in tools: + tool_name = getattr(tool, "name", None) + if self._is_tool_visible(tool_name, enabled_tool_names): + filtered.append(tool) + + return filtered + + def _should_filter_tool_listing(self) -> bool: + transport = (config.transport_mode or "stdio").lower() + return transport == "http" and PluginHub.is_configured() + + async def _resolve_enabled_tool_names_for_context( + self, + context: MiddlewareContext, + ) -> set[str] | None: + ctx = context.fastmcp_context + user_id = ctx.get_state("user_id") if config.http_remote_hosted else None + active_instance = ctx.get_state("unity_instance") + + project_hashes = self._resolve_candidate_project_hashes(active_instance) + if not project_hashes: + try: + sessions_data = await PluginHub.get_sessions(user_id=user_id) + sessions = sessions_data.sessions if sessions_data else {} + except Exception: + return None + + if not sessions: + return None + + if len(sessions) == 1: + only_session = next(iter(sessions.values())) + only_hash = getattr(only_session, "hash", None) + if only_hash: + project_hashes = [only_hash] + else: + # Multiple sessions without explicit selection: use a union so we don't + # hide tools that are valid in at least one visible Unity instance. + project_hashes = [ + session.hash + for session in sessions.values() + if getattr(session, "hash", None) + ] + + if not project_hashes: + return None + + enabled_tool_names: set[str] = set() + for project_hash in project_hashes: + try: + registered_tools = await PluginHub.get_tools_for_project(project_hash) + except Exception: + continue + + for tool in registered_tools: + tool_name = getattr(tool, "name", None) + if isinstance(tool_name, str) and tool_name: + enabled_tool_names.add(tool_name) + + return enabled_tool_names or None + + @staticmethod + def _resolve_candidate_project_hashes(active_instance: str | None) -> list[str]: + if not active_instance: + return [] + + if "@" in active_instance: + _, _, suffix = active_instance.rpartition("@") + return [suffix] if suffix else [] + + return [active_instance] + + def _is_tool_visible(self, tool_name: str | None, enabled_tool_names: set[str]) -> bool: + if not isinstance(tool_name, str) or not tool_name: + return True + + if tool_name in self._server_only_tool_names: + return True + + if tool_name in enabled_tool_names: + return True + + unity_target = self._tool_alias_to_unity_target.get(tool_name) + if unity_target: + return unity_target in enabled_tool_names + + # Keep unknown tools visible for forward compatibility. + if tool_name not in self._unity_managed_tool_names: + return True + + return False diff --git a/Server/tests/test_transport_characterization.py b/Server/tests/test_transport_characterization.py index a0b617934..f1b698bcb 100644 --- a/Server/tests/test_transport_characterization.py +++ b/Server/tests/test_transport_characterization.py @@ -18,6 +18,7 @@ from unittest.mock import AsyncMock, Mock, MagicMock, patch, call from datetime import datetime, timezone import uuid +from types import SimpleNamespace from transport.unity_instance_middleware import UnityInstanceMiddleware, get_unity_instance_middleware, set_unity_instance_middleware from transport.plugin_registry import PluginRegistry, PluginSession @@ -31,6 +32,7 @@ SessionDetails, ) from models.models import ToolDefinitionModel +from core.config import config # ============================================================================ @@ -266,6 +268,78 @@ async def mock_call_next(_ctx): if len(c[0]) > 0 and c[0][0] == "unity_instance"] assert len(calls) == 0 + @pytest.mark.asyncio + async def test_list_tools_filters_disabled_unity_tools_and_aliases(self, mock_context, monkeypatch): + """ + Current behavior: in HTTP mode with a connected Unity session, on_list_tools() + uses PluginHub-registered tool names to hide disabled Unity tools while keeping + server-only tools visible. Aliases like create_script follow manage_script state. + """ + middleware = UnityInstanceMiddleware() + middleware_ctx = Mock() + middleware_ctx.fastmcp_context = mock_context + + mock_context.set_state("unity_instance", "Project@abc123") + monkeypatch.setattr(config, "transport_mode", "http") + + available_tools = [ + SimpleNamespace(name="manage_scene"), + SimpleNamespace(name="manage_script"), + SimpleNamespace(name="set_active_instance"), + SimpleNamespace(name="manage_asset"), + SimpleNamespace(name="create_script"), + ] + + async def call_next(_ctx): + return available_tools + + with patch.object(middleware, "_inject_unity_instance", new=AsyncMock()): + with patch("transport.unity_instance_middleware.PluginHub.is_configured", return_value=True): + with patch("transport.unity_instance_middleware.PluginHub.get_tools_for_project", new_callable=AsyncMock) as mock_get_tools: + mock_get_tools.return_value = [ + SimpleNamespace(name="manage_scene"), + SimpleNamespace(name="manage_script"), + ] + + filtered = await middleware.on_list_tools(middleware_ctx, call_next) + + names = [tool.name for tool in filtered] + assert "manage_scene" in names + assert "create_script" in names + assert "set_active_instance" in names + assert "manage_asset" not in names + + @pytest.mark.asyncio + async def test_list_tools_skips_filter_when_no_enabled_set_available(self, mock_context, monkeypatch): + """ + Current behavior: if Unity has not yet registered tool definitions for a + session, on_list_tools() leaves the FastMCP list unchanged. + """ + middleware = UnityInstanceMiddleware() + middleware_ctx = Mock() + middleware_ctx.fastmcp_context = mock_context + + mock_context.set_state("unity_instance", "Project@abc123") + monkeypatch.setattr(config, "transport_mode", "http") + + original_tools = [ + SimpleNamespace(name="manage_scene"), + SimpleNamespace(name="manage_asset"), + SimpleNamespace(name="set_active_instance"), + ] + + async def call_next(_ctx): + return original_tools + + with patch.object(middleware, "_inject_unity_instance", new=AsyncMock()): + with patch("transport.unity_instance_middleware.PluginHub.is_configured", return_value=True): + with patch("transport.unity_instance_middleware.PluginHub.get_tools_for_project", new_callable=AsyncMock) as mock_get_tools: + mock_get_tools.return_value = [] + + filtered = await middleware.on_list_tools(middleware_ctx, call_next) + + assert [tool.name for tool in filtered] == [tool.name for tool in original_tools] + # ============================================================================ # AUTO-SELECT INSTANCE TESTS From fb17aeae1721ebb152e4b123b8d0a8e304502f0d Mon Sep 17 00:00:00 2001 From: whatevertogo Date: Wed, 11 Feb 2026 23:29:36 +0800 Subject: [PATCH 04/10] chore: delete State.meta file --- TestProjects/UnityMCPTests/Assets/Tests/Editor/State.meta | 8 -------- 1 file changed, 8 deletions(-) delete mode 100644 TestProjects/UnityMCPTests/Assets/Tests/Editor/State.meta diff --git a/TestProjects/UnityMCPTests/Assets/Tests/Editor/State.meta b/TestProjects/UnityMCPTests/Assets/Tests/Editor/State.meta deleted file mode 100644 index a4a8007ae..000000000 --- a/TestProjects/UnityMCPTests/Assets/Tests/Editor/State.meta +++ /dev/null @@ -1,8 +0,0 @@ -fileFormatVersion: 2 -guid: 778e8246ab09497419985acaee595f6c -folderAsset: yes -DefaultImporter: - externalObjects: {} - userData: - assetBundleName: - assetBundleVariant: From 3b624c1f4379edd7c9ec884cec7cc7f8bc7ae422 Mon Sep 17 00:00:00 2001 From: whatevertogo Date: Wed, 11 Feb 2026 23:57:53 +0800 Subject: [PATCH 05/10] fix: update tool filtering logic to handle user-specific tool lookups and improve error handling --- .../Components/Tools/McpToolsSection.cs | 22 +++--- Server/src/transport/plugin_hub.py | 10 ++- .../transport/unity_instance_middleware.py | 28 +++++-- .../tests/test_transport_characterization.py | 74 ++++++++++++++++++- .../ToolDiscoveryServiceTests.cs.meta | 2 +- 5 files changed, 114 insertions(+), 22 deletions(-) diff --git a/MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs b/MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs index 687fbf2eb..8890ef0f9 100644 --- a/MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs +++ b/MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs @@ -1,7 +1,7 @@ using System; using System.Collections.Generic; using System.Linq; -using System.Threading; +using System.Threading.Tasks; using MCPForUnity.Editor.Constants; using MCPForUnity.Editor.Helpers; using MCPForUnity.Editor.Services; @@ -247,21 +247,23 @@ private void HandleToggleChange( private void ReregisterToolsAsync() { - // Fire and forget - don't block UI - ThreadPool.QueueUserWorkItem(_ => + // Fire and forget - don't block UI thread + var transportManager = MCPServiceLocator.TransportManager; + var client = transportManager.GetClient(TransportMode.Http); + if (client == null || !client.IsConnected) + { + return; + } + + _ = Task.Run(async () => { try { - var transportManager = MCPServiceLocator.TransportManager; - var client = transportManager.GetClient(TransportMode.Http); - if (client != null && client.IsConnected) - { - client.ReregisterToolsAsync().Wait(); - } + await client.ReregisterToolsAsync().ConfigureAwait(false); } catch (Exception ex) { - McpLog.Warn($"Failed to reregister tools: {ex.Message}"); + McpLog.Warn($"Failed to reregister tools: {ex}"); } }); } diff --git a/Server/src/transport/plugin_hub.py b/Server/src/transport/plugin_hub.py index 1058c4136..be29d1fdb 100644 --- a/Server/src/transport/plugin_hub.py +++ b/Server/src/transport/plugin_hub.py @@ -318,12 +318,16 @@ async def get_sessions(cls, user_id: str | None = None) -> SessionList: ) @classmethod - async def get_tools_for_project(cls, project_hash: str) -> list[Any]: - """Retrieve tools registered for a active project hash.""" + async def get_tools_for_project( + cls, + project_hash: str, + user_id: str | None = None, + ) -> list[Any]: + """Retrieve tools registered for an active project hash.""" if cls._registry is None: return [] - session_id = await cls._registry.get_session_id_by_hash(project_hash) + session_id = await cls._registry.get_session_id_by_hash(project_hash, user_id=user_id) if not session_id: return [] diff --git a/Server/src/transport/unity_instance_middleware.py b/Server/src/transport/unity_instance_middleware.py index d172a304d..7c7119ee1 100644 --- a/Server/src/transport/unity_instance_middleware.py +++ b/Server/src/transport/unity_instance_middleware.py @@ -307,7 +307,7 @@ async def on_list_tools(self, context: MiddlewareContext, call_next): return tools enabled_tool_names = await self._resolve_enabled_tool_names_for_context(context) - if not enabled_tool_names: + if enabled_tool_names is None: return tools filtered = [] @@ -335,7 +335,13 @@ async def _resolve_enabled_tool_names_for_context( try: sessions_data = await PluginHub.get_sessions(user_id=user_id) sessions = sessions_data.sessions if sessions_data else {} - except Exception: + except Exception as exc: + logger.debug( + "Failed to fetch sessions for tool filtering (user_id=%s, %s)", + user_id, + type(exc).__name__, + exc_info=True, + ) return None if not sessions: @@ -359,10 +365,19 @@ async def _resolve_enabled_tool_names_for_context( return None enabled_tool_names: set[str] = set() + resolved_any_project = False for project_hash in project_hashes: try: - registered_tools = await PluginHub.get_tools_for_project(project_hash) - except Exception: + registered_tools = await PluginHub.get_tools_for_project(project_hash, user_id=user_id) + resolved_any_project = True + except Exception as exc: + logger.debug( + "Failed to fetch tools for project hash %s (user_id=%s, %s)", + project_hash, + user_id, + type(exc).__name__, + exc_info=True, + ) continue for tool in registered_tools: @@ -370,7 +385,10 @@ async def _resolve_enabled_tool_names_for_context( if isinstance(tool_name, str) and tool_name: enabled_tool_names.add(tool_name) - return enabled_tool_names or None + if not resolved_any_project: + return None + + return enabled_tool_names @staticmethod def _resolve_candidate_project_hashes(active_instance: str | None) -> list[str]: diff --git a/Server/tests/test_transport_characterization.py b/Server/tests/test_transport_characterization.py index f1b698bcb..dc98189fe 100644 --- a/Server/tests/test_transport_characterization.py +++ b/Server/tests/test_transport_characterization.py @@ -310,10 +310,11 @@ async def call_next(_ctx): assert "manage_asset" not in names @pytest.mark.asyncio - async def test_list_tools_skips_filter_when_no_enabled_set_available(self, mock_context, monkeypatch): + async def test_list_tools_filters_when_enabled_set_is_empty(self, mock_context, monkeypatch): """ - Current behavior: if Unity has not yet registered tool definitions for a - session, on_list_tools() leaves the FastMCP list unchanged. + Current behavior: when a Unity session is resolved but it has zero enabled + tools registered, Unity-managed tools are filtered out while server-only + and unknown tools remain visible. """ middleware = UnityInstanceMiddleware() middleware_ctx = Mock() @@ -325,7 +326,9 @@ async def test_list_tools_skips_filter_when_no_enabled_set_available(self, mock_ original_tools = [ SimpleNamespace(name="manage_scene"), SimpleNamespace(name="manage_asset"), + SimpleNamespace(name="create_script"), SimpleNamespace(name="set_active_instance"), + SimpleNamespace(name="custom_server_tool"), ] async def call_next(_ctx): @@ -338,8 +341,73 @@ async def call_next(_ctx): filtered = await middleware.on_list_tools(middleware_ctx, call_next) + names = [tool.name for tool in filtered] + assert "set_active_instance" in names + assert "custom_server_tool" in names + assert "manage_scene" not in names + assert "manage_asset" not in names + assert "create_script" not in names + + @pytest.mark.asyncio + async def test_list_tools_skips_filter_when_enabled_set_lookup_fails(self, mock_context, monkeypatch): + """ + Current behavior: if enabled-tool lookup fails unexpectedly, on_list_tools() + leaves the FastMCP list unchanged to avoid hiding tools due to transient + PluginHub failures. + """ + middleware = UnityInstanceMiddleware() + middleware_ctx = Mock() + middleware_ctx.fastmcp_context = mock_context + + mock_context.set_state("unity_instance", "Project@abc123") + monkeypatch.setattr(config, "transport_mode", "http") + + original_tools = [ + SimpleNamespace(name="manage_scene"), + SimpleNamespace(name="manage_asset"), + SimpleNamespace(name="set_active_instance"), + ] + + async def call_next(_ctx): + return original_tools + + with patch.object(middleware, "_inject_unity_instance", new=AsyncMock()): + with patch("transport.unity_instance_middleware.PluginHub.is_configured", return_value=True): + with patch("transport.unity_instance_middleware.PluginHub.get_tools_for_project", new_callable=AsyncMock) as mock_get_tools: + mock_get_tools.side_effect = RuntimeError("hub unavailable") + + filtered = await middleware.on_list_tools(middleware_ctx, call_next) + assert [tool.name for tool in filtered] == [tool.name for tool in original_tools] + @pytest.mark.asyncio + async def test_list_tools_uses_user_scoped_tool_lookup_in_hosted_mode(self, mock_context, monkeypatch): + """ + Current behavior: in remote-hosted HTTP mode, tool filtering fetches + Unity-registered tools scoped to the current user. + """ + middleware = UnityInstanceMiddleware() + middleware_ctx = Mock() + middleware_ctx.fastmcp_context = mock_context + + mock_context.set_state("unity_instance", "Project@abc123") + mock_context.set_state("user_id", "user-123") + monkeypatch.setattr(config, "transport_mode", "http") + monkeypatch.setattr(config, "http_remote_hosted", True) + + async def call_next(_ctx): + return [SimpleNamespace(name="manage_scene")] + + with patch.object(middleware, "_inject_unity_instance", new=AsyncMock()): + with patch("transport.unity_instance_middleware.PluginHub.is_configured", return_value=True): + with patch("transport.unity_instance_middleware.PluginHub.get_tools_for_project", new_callable=AsyncMock) as mock_get_tools: + mock_get_tools.return_value = [SimpleNamespace(name="manage_scene")] + + filtered = await middleware.on_list_tools(middleware_ctx, call_next) + + assert [tool.name for tool in filtered] == ["manage_scene"] + mock_get_tools.assert_awaited_once_with("abc123", user_id="user-123") + # ============================================================================ # AUTO-SELECT INSTANCE TESTS diff --git a/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/ToolDiscoveryServiceTests.cs.meta b/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/ToolDiscoveryServiceTests.cs.meta index cea36ebb1..ce33028b5 100644 --- a/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/ToolDiscoveryServiceTests.cs.meta +++ b/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/ToolDiscoveryServiceTests.cs.meta @@ -1,5 +1,5 @@ fileFormatVersion: 2 -guid: 7a8b9c0d1e2f3a4b5c6d7e8f9a0b1c2d +guid: 5f3b1c9b5dc24a52ae7053af3e2135ab MonoImporter: externalObjects: {} serializedVersion: 2 From c67e9f2826e4a81386bed389f74a15837edc2fec Mon Sep 17 00:00:00 2001 From: whatevertogo Date: Thu, 12 Feb 2026 00:25:26 +0800 Subject: [PATCH 06/10] fix: add user_id param to execute_custom_tool for remote-hosted mode - Pass user_id to execute_tool for user-scoped tool lookups - Update get_tool_definition to support user-scoped queries Co-Authored-By: Claude Opus 4.6 --- Server/src/services/custom_tool_service.py | 56 ++++++++++++++++--- Server/src/services/resources/custom_tools.py | 4 +- .../src/services/tools/execute_custom_tool.py | 10 +++- Server/src/transport/plugin_hub.py | 9 ++- 4 files changed, 68 insertions(+), 11 deletions(-) diff --git a/Server/src/services/custom_tool_service.py b/Server/src/services/custom_tool_service.py index 9dabaad8a..91349750c 100644 --- a/Server/src/services/custom_tool_service.py +++ b/Server/src/services/custom_tool_service.py @@ -10,6 +10,7 @@ from starlette.requests import Request from starlette.responses import JSONResponse +from core.config import config from models.models import MCPResponse, ToolDefinitionModel, ToolParameterModel from core.logging_decorator import log_execution from core.telemetry_decorator import telemetry_tool @@ -28,6 +29,23 @@ _MAX_POLL_SECONDS = 600 +def get_user_id_from_context(ctx: Context) -> str | None: + """Read user_id from request-scoped context in remote-hosted mode.""" + if not config.http_remote_hosted: + return None + + get_state = getattr(ctx, "get_state", None) + if not callable(get_state): + return None + + try: + user_id = get_state("user_id") + except Exception: + return None + + return user_id if isinstance(user_id, str) and user_id else None + + class RegisterToolsPayload(BaseModel): project_id: str project_hash: str | None = None @@ -84,16 +102,25 @@ async def register_tools(request: Request) -> JSONResponse: return JSONResponse(response.model_dump()) # --- Public API for MCP tools --------------------------------------- - async def list_registered_tools(self, project_id: str) -> list[ToolDefinitionModel]: + async def list_registered_tools( + self, + project_id: str, + user_id: str | None = None, + ) -> list[ToolDefinitionModel]: legacy = list(self._project_tools.get(project_id, {}).values()) - hub_tools = await PluginHub.get_tools_for_project(project_id) + hub_tools = await PluginHub.get_tools_for_project(project_id, user_id=user_id) return legacy + hub_tools - async def get_tool_definition(self, project_id: str, tool_name: str) -> ToolDefinitionModel | None: + async def get_tool_definition( + self, + project_id: str, + tool_name: str, + user_id: str | None = None, + ) -> ToolDefinitionModel | None: tool = self._project_tools.get(project_id, {}).get(tool_name) if tool: return tool - return await PluginHub.get_tool_definition(project_id, tool_name) + return await PluginHub.get_tool_definition(project_id, tool_name, user_id=user_id) async def execute_tool( self, @@ -101,13 +128,14 @@ async def execute_tool( tool_name: str, unity_instance: str | None, params: dict[str, object] | None = None, + user_id: str | None = None, ) -> MCPResponse: params = params or {} logger.info( f"Executing tool '{tool_name}' for project '{project_id}' (instance={unity_instance}) with params: {params}" ) - definition = await self.get_tool_definition(project_id, tool_name) + definition = await self.get_tool_definition(project_id, tool_name, user_id=user_id) if definition is None: return MCPResponse( success=False, @@ -119,6 +147,7 @@ async def execute_tool( unity_instance, tool_name, params, + user_id=user_id, ) if not definition.requires_polling: @@ -132,6 +161,7 @@ async def execute_tool( params, response, definition.poll_action or "status", + user_id=user_id, ) logger.info(f"Tool '{tool_name}' polled response: {result}") return result @@ -156,6 +186,7 @@ async def _poll_until_complete( initial_params: dict[str, object], initial_response, poll_action: str, + user_id: str | None = None, ) -> MCPResponse: poll_params = dict(initial_params) poll_params["action"] = poll_action or "status" @@ -180,7 +211,11 @@ async def _poll_until_complete( try: response = await send_with_unity_instance( - async_send_command_with_retry, unity_instance, tool_name, poll_params + async_send_command_with_retry, + unity_instance, + tool_name, + poll_params, + user_id=user_id, ) except Exception as exc: # pragma: no cover - network/domain reload variability logger.debug(f"Polling {tool_name} failed, will retry: {exc}") @@ -347,8 +382,15 @@ async def _handler(ctx: Context, **kwargs) -> MCPResponse: ) params = {k: v for k, v in kwargs.items() if v is not None} + user_id = get_user_id_from_context(ctx) service = CustomToolService.get_instance() - return await service.execute_tool(project_id, definition.name, unity_instance, params) + return await service.execute_tool( + project_id, + definition.name, + unity_instance, + params, + user_id=user_id, + ) _handler.__name__ = f"custom_tool_{definition.name}" _handler.__doc__ = definition.description or "" diff --git a/Server/src/services/resources/custom_tools.py b/Server/src/services/resources/custom_tools.py index a28ac2cbb..e0fe9c6b1 100644 --- a/Server/src/services/resources/custom_tools.py +++ b/Server/src/services/resources/custom_tools.py @@ -4,6 +4,7 @@ from models import MCPResponse from services.custom_tool_service import ( CustomToolService, + get_user_id_from_context, resolve_project_id_for_unity_instance, ToolDefinitionModel, ) @@ -42,7 +43,8 @@ async def get_custom_tools(ctx: Context) -> CustomToolsResourceResponse | MCPRes ) service = CustomToolService.get_instance() - tools = await service.list_registered_tools(project_id) + user_id = get_user_id_from_context(ctx) + tools = await service.list_registered_tools(project_id, user_id=user_id) data = CustomToolsData( project_id=project_id, diff --git a/Server/src/services/tools/execute_custom_tool.py b/Server/src/services/tools/execute_custom_tool.py index 1aef2adc4..8e66d4f19 100644 --- a/Server/src/services/tools/execute_custom_tool.py +++ b/Server/src/services/tools/execute_custom_tool.py @@ -4,6 +4,7 @@ from services.custom_tool_service import ( CustomToolService, + get_user_id_from_context, resolve_project_id_for_unity_instance, ) from services.registry import mcp_for_unity_tool @@ -40,4 +41,11 @@ async def execute_custom_tool(ctx: Context, tool_name: str, parameters: dict | N ) service = CustomToolService.get_instance() - return await service.execute_tool(project_id, tool_name, unity_instance, parameters) + user_id = get_user_id_from_context(ctx) + return await service.execute_tool( + project_id, + tool_name, + unity_instance, + parameters, + user_id=user_id, + ) diff --git a/Server/src/transport/plugin_hub.py b/Server/src/transport/plugin_hub.py index be29d1fdb..fa8aa3a98 100644 --- a/Server/src/transport/plugin_hub.py +++ b/Server/src/transport/plugin_hub.py @@ -338,12 +338,17 @@ async def get_tools_for_project( return list(session.tools.values()) @classmethod - async def get_tool_definition(cls, project_hash: str, tool_name: str) -> Any | None: + async def get_tool_definition( + cls, + project_hash: str, + tool_name: str, + user_id: str | None = None, + ) -> Any | None: """Retrieve a specific tool definition for an active project hash.""" if cls._registry is None: return None - session_id = await cls._registry.get_session_id_by_hash(project_hash) + session_id = await cls._registry.get_session_id_by_hash(project_hash, user_id=user_id) if not session_id: return None From fc81d9e8ac47000b42029dca19d7817d5767d658 Mon Sep 17 00:00:00 2001 From: whatevertogo Date: Thu, 12 Feb 2026 00:29:34 +0800 Subject: [PATCH 07/10] fix: preserve tool state and filter tools in client listings (fixes #709) Comprehensive fix for Issue #709: - Unity tools stay disabled across restarts (EditorPrefs persistence) - Client tool lists filter out disabled Unity tools (HTTP mode only) - Script helper tools follow their target tool's enabled state - Added graceful error fallback when tool lookup fails Co-Authored-By: Claude Opus 4.6 --- .../test_custom_tool_service_user_scope.py | 106 ++++++++++++++++++ .../tests/test_transport_characterization.py | 86 ++++++++++++++ 2 files changed, 192 insertions(+) create mode 100644 Server/tests/test_custom_tool_service_user_scope.py diff --git a/Server/tests/test_custom_tool_service_user_scope.py b/Server/tests/test_custom_tool_service_user_scope.py new file mode 100644 index 000000000..a291199c5 --- /dev/null +++ b/Server/tests/test_custom_tool_service_user_scope.py @@ -0,0 +1,106 @@ +from unittest.mock import AsyncMock, Mock, patch + +import pytest + +from core.config import config +from models.models import MCPResponse, ToolDefinitionModel +from services.custom_tool_service import CustomToolService +from services.resources.custom_tools import get_custom_tools +from services.tools.execute_custom_tool import execute_custom_tool + + +class _DummyMcp: + def custom_route(self, _path, methods=None): # noqa: ARG002 + def _decorator(fn): + return fn + + return _decorator + + +@pytest.mark.asyncio +async def test_list_registered_tools_threads_user_id_to_plugin_hub(): + service = CustomToolService(_DummyMcp()) + + with patch("services.custom_tool_service.PluginHub.get_tools_for_project", new_callable=AsyncMock) as mock_get: + mock_get.return_value = [] + await service.list_registered_tools("project-hash", user_id="user-1") + + mock_get.assert_awaited_once_with("project-hash", user_id="user-1") + + +@pytest.mark.asyncio +async def test_get_tool_definition_threads_user_id_to_plugin_hub(): + service = CustomToolService(_DummyMcp()) + + with patch("services.custom_tool_service.PluginHub.get_tool_definition", new_callable=AsyncMock) as mock_get: + mock_get.return_value = None + await service.get_tool_definition("project-hash", "my_tool", user_id="user-1") + + mock_get.assert_awaited_once_with("project-hash", "my_tool", user_id="user-1") + + +@pytest.mark.asyncio +async def test_execute_tool_threads_user_id_to_definition_lookup_and_transport(): + service = CustomToolService(_DummyMcp()) + definition = ToolDefinitionModel(name="my_tool", description="My tool", requires_polling=False) + + with patch.object(service, "get_tool_definition", new_callable=AsyncMock) as mock_get_definition: + with patch("services.custom_tool_service.send_with_unity_instance", new_callable=AsyncMock) as mock_send: + mock_get_definition.return_value = definition + mock_send.return_value = {"success": True, "message": "ok"} + + await service.execute_tool( + "project-hash", + "my_tool", + "Project@project-hash", + {"x": 1}, + user_id="user-1", + ) + + mock_get_definition.assert_awaited_once_with("project-hash", "my_tool", user_id="user-1") + mock_send.assert_awaited_once() + assert mock_send.call_args.kwargs["user_id"] == "user-1" + + +@pytest.mark.asyncio +async def test_execute_custom_tool_threads_user_id_from_context(monkeypatch): + monkeypatch.setattr(config, "http_remote_hosted", True) + + ctx = Mock() + state = {"unity_instance": "Project@project-hash", "user_id": "user-1"} + ctx.get_state = Mock(side_effect=lambda key: state.get(key)) + + service = Mock() + service.execute_tool = AsyncMock(return_value=MCPResponse(success=True, message="ok")) + + with patch("services.tools.execute_custom_tool.resolve_project_id_for_unity_instance", return_value="project-hash"): + with patch("services.tools.execute_custom_tool.CustomToolService.get_instance", return_value=service): + await execute_custom_tool(ctx, "my_tool", {}) + + service.execute_tool.assert_awaited_once_with( + "project-hash", + "my_tool", + "Project@project-hash", + {}, + user_id="user-1", + ) + + +@pytest.mark.asyncio +async def test_custom_tools_resource_threads_user_id_from_context(monkeypatch): + monkeypatch.setattr(config, "http_remote_hosted", True) + + ctx = Mock() + state = {"unity_instance": "Project@project-hash", "user_id": "user-1"} + ctx.get_state = Mock(side_effect=lambda key: state.get(key)) + + service = Mock() + service.list_registered_tools = AsyncMock( + return_value=[ToolDefinitionModel(name="my_tool", description="My tool")] + ) + + with patch("services.resources.custom_tools.resolve_project_id_for_unity_instance", return_value="project-hash"): + with patch("services.resources.custom_tools.CustomToolService.get_instance", return_value=service): + await get_custom_tools(ctx) + + service.list_registered_tools.assert_awaited_once_with("project-hash", user_id="user-1") diff --git a/Server/tests/test_transport_characterization.py b/Server/tests/test_transport_characterization.py index dc98189fe..b67b711ab 100644 --- a/Server/tests/test_transport_characterization.py +++ b/Server/tests/test_transport_characterization.py @@ -1007,6 +1007,92 @@ def test_plugin_hub_not_configured_sends_command_fails(self): with pytest.raises(RuntimeError, match="not configured"): asyncio.run(PluginHub.send_command("sess-id", "ping", {})) + @pytest.mark.asyncio + async def test_plugin_hub_get_tools_for_project_honors_user_scope( + self, + configured_plugin_hub, + plugin_registry, + monkeypatch, + ): + """ + Current behavior: in remote-hosted mode, get_tools_for_project() + resolves by (user_id, project_hash) so users with the same hash do not + see each other's tool registrations. + """ + monkeypatch.setattr(config, "http_remote_hosted", True) + + await plugin_registry.register( + session_id="sess-user-a", + project_name="Project", + project_hash="shared-hash", + unity_version="2022.3", + user_id="user-a", + ) + await plugin_registry.register( + session_id="sess-user-b", + project_name="Project", + project_hash="shared-hash", + unity_version="2022.3", + user_id="user-b", + ) + await plugin_registry.register_tools_for_session( + "sess-user-a", + [ToolDefinitionModel(name="tool_a", description="Tool A")], + ) + await plugin_registry.register_tools_for_session( + "sess-user-b", + [ToolDefinitionModel(name="tool_b", description="Tool B")], + ) + + tools_for_a = await PluginHub.get_tools_for_project("shared-hash", user_id="user-a") + tools_for_b = await PluginHub.get_tools_for_project("shared-hash", user_id="user-b") + + assert [tool.name for tool in tools_for_a] == ["tool_a"] + assert [tool.name for tool in tools_for_b] == ["tool_b"] + + @pytest.mark.asyncio + async def test_plugin_hub_get_tool_definition_honors_user_scope( + self, + configured_plugin_hub, + plugin_registry, + monkeypatch, + ): + """ + Current behavior: in remote-hosted mode, get_tool_definition() is + user-scoped for shared project hashes. + """ + monkeypatch.setattr(config, "http_remote_hosted", True) + + await plugin_registry.register( + session_id="sess-user-a", + project_name="Project", + project_hash="shared-hash", + unity_version="2022.3", + user_id="user-a", + ) + await plugin_registry.register( + session_id="sess-user-b", + project_name="Project", + project_hash="shared-hash", + unity_version="2022.3", + user_id="user-b", + ) + await plugin_registry.register_tools_for_session( + "sess-user-a", + [ToolDefinitionModel(name="tool_a", description="Tool A")], + ) + await plugin_registry.register_tools_for_session( + "sess-user-b", + [ToolDefinitionModel(name="tool_b", description="Tool B")], + ) + + tool_for_a = await PluginHub.get_tool_definition("shared-hash", "tool_a", user_id="user-a") + tool_for_b = await PluginHub.get_tool_definition("shared-hash", "tool_a", user_id="user-b") + + assert tool_for_a is not None + assert tool_for_a.name == "tool_a" + assert tool_for_b is None + # ============================================================================ # GLOBAL MIDDLEWARE SINGLETON TESTS From 1a9b08351fadc9825b55e0961c8735adac3707ca Mon Sep 17 00:00:00 2001 From: whatevertogo Date: Thu, 12 Feb 2026 01:37:26 +0800 Subject: [PATCH 08/10] fix: add unity_target parameter to mcp_for_unity_tool for improved tool visibility control --- Server/src/services/registry/tool_registry.py | 25 +- .../services/tools/debug_request_context.py | 1 + .../src/services/tools/execute_custom_tool.py | 1 + Server/src/services/tools/find_in_file.py | 1 + Server/src/services/tools/manage_script.py | 6 + .../src/services/tools/script_apply_edits.py | 1 + .../src/services/tools/set_active_instance.py | 1 + .../transport/unity_instance_middleware.py | 161 ++++++---- Server/tests/test_tool_registry_metadata.py | 64 ++++ .../tests/test_transport_characterization.py | 276 ++++++++++++++++-- 10 files changed, 462 insertions(+), 75 deletions(-) create mode 100644 Server/tests/test_tool_registry_metadata.py diff --git a/Server/src/services/registry/tool_registry.py b/Server/src/services/registry/tool_registry.py index 08526a5a5..732cde2c3 100644 --- a/Server/src/services/registry/tool_registry.py +++ b/Server/src/services/registry/tool_registry.py @@ -10,6 +10,7 @@ def mcp_for_unity_tool( name: str | None = None, description: str | None = None, + unity_target: str | None = "self", **kwargs ) -> Callable: """ @@ -20,6 +21,10 @@ def mcp_for_unity_tool( Args: name: Tool name (defaults to function name) description: Tool description + unity_target: Visibility target used by middleware filtering. + - "self" (default): tool follows its own enabled state. + - None: server-only tool, always visible in tool listing. + - "": alias tool that follows another Unity tool state. **kwargs: Additional arguments passed to @mcp.tool() Example: @@ -29,11 +34,29 @@ async def my_custom_tool(ctx: Context, ...): """ def decorator(func: Callable) -> Callable: tool_name = name if name is not None else func.__name__ + # Safety guard: unity_target is internal metadata and must never leak into mcp.tool kwargs. + tool_kwargs = dict(kwargs) # Create a copy to avoid side effects + if "unity_target" in tool_kwargs: + del tool_kwargs["unity_target"] + + if unity_target is None: + normalized_unity_target: str | None = None + elif isinstance(unity_target, str) and unity_target.strip(): + normalized_unity_target = ( + tool_name if unity_target == "self" else unity_target.strip() + ) + else: + raise ValueError( + f"Invalid unity_target for tool '{tool_name}': {unity_target!r}. " + "Expected None or a non-empty string." + ) + _tool_registry.append({ 'func': func, 'name': tool_name, 'description': description, - 'kwargs': kwargs + 'unity_target': normalized_unity_target, + 'kwargs': tool_kwargs, }) return func diff --git a/Server/src/services/tools/debug_request_context.py b/Server/src/services/tools/debug_request_context.py index d3539e2e2..6ce2913fa 100644 --- a/Server/src/services/tools/debug_request_context.py +++ b/Server/src/services/tools/debug_request_context.py @@ -13,6 +13,7 @@ @mcp_for_unity_tool( + unity_target=None, description="Return the current FastMCP request context details (client_id, session_id, and meta dump).", annotations=ToolAnnotations( title="Debug Request Context", diff --git a/Server/src/services/tools/execute_custom_tool.py b/Server/src/services/tools/execute_custom_tool.py index 8e66d4f19..6152255de 100644 --- a/Server/src/services/tools/execute_custom_tool.py +++ b/Server/src/services/tools/execute_custom_tool.py @@ -13,6 +13,7 @@ @mcp_for_unity_tool( name="execute_custom_tool", + unity_target=None, description="Execute a project-scoped custom tool registered by Unity.", annotations=ToolAnnotations( title="Execute Custom Tool", diff --git a/Server/src/services/tools/find_in_file.py b/Server/src/services/tools/find_in_file.py index 56372d2b9..7a33d95b9 100644 --- a/Server/src/services/tools/find_in_file.py +++ b/Server/src/services/tools/find_in_file.py @@ -66,6 +66,7 @@ def _split_uri(uri: str) -> tuple[str, str]: @mcp_for_unity_tool( + unity_target="manage_script", description="Searches a file with a regex pattern and returns line numbers and excerpts.", annotations=ToolAnnotations( title="Find in File", diff --git a/Server/src/services/tools/manage_script.py b/Server/src/services/tools/manage_script.py index dc9ab5fe9..b5df345a8 100644 --- a/Server/src/services/tools/manage_script.py +++ b/Server/src/services/tools/manage_script.py @@ -65,6 +65,7 @@ def _split_uri(uri: str) -> tuple[str, str]: @mcp_for_unity_tool( + unity_target="manage_script", description=( """Apply small text edits to a C# script identified by URI. IMPORTANT: This tool replaces EXACT character positions. Always verify content at target lines/columns BEFORE editing! @@ -375,6 +376,7 @@ async def _flip_async(): @mcp_for_unity_tool( + unity_target="manage_script", description="Create a new C# script at the given project path.", annotations=ToolAnnotations( title="Create Script", @@ -426,6 +428,7 @@ async def create_script( @mcp_for_unity_tool( + unity_target="manage_script", description="Delete a C# script by URI or Assets-relative path.", annotations=ToolAnnotations( title="Delete Script", @@ -454,6 +457,7 @@ async def delete_script( @mcp_for_unity_tool( + unity_target="manage_script", description="Validate a C# script and return diagnostics.", annotations=ToolAnnotations( title="Validate Script", @@ -575,6 +579,7 @@ async def manage_script( @mcp_for_unity_tool( + unity_target=None, description=( """Get manage_script capabilities (supported ops, limits, and guards). Returns: @@ -613,6 +618,7 @@ async def manage_script_capabilities(ctx: Context) -> dict[str, Any]: @mcp_for_unity_tool( + unity_target="manage_script", description="Get SHA256 and basic metadata for a Unity C# script without returning file contents. Requires uri (script path under Assets/ or mcpforunity://path/Assets/... or file://...).", annotations=ToolAnnotations( title="Get SHA", diff --git a/Server/src/services/tools/script_apply_edits.py b/Server/src/services/tools/script_apply_edits.py index 2ada06137..cc5772e94 100644 --- a/Server/src/services/tools/script_apply_edits.py +++ b/Server/src/services/tools/script_apply_edits.py @@ -312,6 +312,7 @@ def _err(code: str, message: str, *, expected: dict[str, Any] | None = None, rew @mcp_for_unity_tool( name="script_apply_edits", + unity_target="manage_script", description=( """Structured C# edits (methods/classes) with safer boundaries - prefer this over raw text. Best practices: diff --git a/Server/src/services/tools/set_active_instance.py b/Server/src/services/tools/set_active_instance.py index 16d642075..30582867f 100644 --- a/Server/src/services/tools/set_active_instance.py +++ b/Server/src/services/tools/set_active_instance.py @@ -12,6 +12,7 @@ @mcp_for_unity_tool( + unity_target=None, description="Set the active Unity instance for this client/session. Accepts Name@hash or hash.", annotations=ToolAnnotations( title="Set Active Instance", diff --git a/Server/src/transport/unity_instance_middleware.py b/Server/src/transport/unity_instance_middleware.py index 7c7119ee1..196b7dfc3 100644 --- a/Server/src/transport/unity_instance_middleware.py +++ b/Server/src/transport/unity_instance_middleware.py @@ -6,10 +6,12 @@ """ from threading import RLock import logging +import time from fastmcp.server.middleware import Middleware, MiddlewareContext from core.config import config +from services.registry import get_registered_tools from transport.plugin_hub import PluginHub logger = logging.getLogger("mcp-for-unity-server") @@ -55,43 +57,14 @@ def __init__(self): super().__init__() self._active_by_key: dict[str, str] = {} self._lock = RLock() - self._unity_managed_tool_names = { - "batch_execute", - "execute_menu_item", - "find_gameobjects", - "get_test_job", - "manage_asset", - "manage_components", - "manage_editor", - "manage_gameobject", - "manage_material", - "manage_prefabs", - "manage_scene", - "manage_script", - "manage_scriptable_object", - "manage_shader", - "manage_texture", - "manage_vfx", - "read_console", - "refresh_unity", - "run_tests", - } - self._tool_alias_to_unity_target = { - # Server-side script helpers route to Unity's manage_script command. - "apply_text_edits": "manage_script", - "create_script": "manage_script", - "delete_script": "manage_script", - "find_in_file": "manage_script", - "get_sha": "manage_script", - "script_apply_edits": "manage_script", - "validate_script": "manage_script", - } - self._server_only_tool_names = { - "debug_request_context", - "execute_custom_tool", - "manage_script_capabilities", - "set_active_instance", - } + self._metadata_lock = RLock() + self._unity_managed_tool_names: set[str] = set() + self._tool_alias_to_unity_target: dict[str, str] = {} + self._server_only_tool_names: set[str] = set() + self._tool_visibility_signature: tuple[tuple[str, str], ...] = () + self._last_tool_visibility_refresh = 0.0 + self._tool_visibility_refresh_interval_seconds = 0.5 + self._has_logged_empty_registry_warning = False def get_session_key(self, ctx) -> str: """ @@ -306,6 +279,7 @@ async def on_list_tools(self, context: MiddlewareContext, call_next): if not self._should_filter_tool_listing(): return tools + self._refresh_tool_visibility_metadata_from_registry() enabled_tool_names = await self._resolve_enabled_tool_names_for_context(context) if enabled_tool_names is None: return tools @@ -329,21 +303,31 @@ async def _resolve_enabled_tool_names_for_context( ctx = context.fastmcp_context user_id = ctx.get_state("user_id") if config.http_remote_hosted else None active_instance = ctx.get_state("unity_instance") - project_hashes = self._resolve_candidate_project_hashes(active_instance) - if not project_hashes: - try: - sessions_data = await PluginHub.get_sessions(user_id=user_id) - sessions = sessions_data.sessions if sessions_data else {} - except Exception as exc: - logger.debug( - "Failed to fetch sessions for tool filtering (user_id=%s, %s)", - user_id, - type(exc).__name__, - exc_info=True, - ) - return None + try: + sessions_data = await PluginHub.get_sessions(user_id=user_id) + sessions = sessions_data.sessions if sessions_data else {} + except Exception as exc: + logger.debug( + "Failed to fetch sessions for tool filtering (user_id=%s, %s)", + user_id, + type(exc).__name__, + exc_info=True, + ) + return None + + session_hashes = { + getattr(session, "hash", None) + for session in sessions.values() + if getattr(session, "hash", None) + } + if project_hashes: + active_hash = project_hashes[0] + # Stale active_instance should not hide all Unity-managed tools. + if active_hash not in session_hashes: + return None + else: if not sessions: return None @@ -355,11 +339,7 @@ async def _resolve_enabled_tool_names_for_context( else: # Multiple sessions without explicit selection: use a union so we don't # hide tools that are valid in at least one visible Unity instance. - project_hashes = [ - session.hash - for session in sessions.values() - if getattr(session, "hash", None) - ] + project_hashes = [hash_value for hash_value in session_hashes if hash_value] if not project_hashes: return None @@ -390,6 +370,77 @@ async def _resolve_enabled_tool_names_for_context( return enabled_tool_names + def _refresh_tool_visibility_metadata_from_registry(self) -> None: + now = time.monotonic() + if now - self._last_tool_visibility_refresh < self._tool_visibility_refresh_interval_seconds: + return + + with self._metadata_lock: + now = time.monotonic() + if now - self._last_tool_visibility_refresh < self._tool_visibility_refresh_interval_seconds: + return + + try: + registry_tools = get_registered_tools() + except Exception: + logger.warning( + "Failed to refresh tool visibility metadata from registry; keeping previous metadata.", + exc_info=True, + ) + self._last_tool_visibility_refresh = now + return + + if not registry_tools and not self._has_logged_empty_registry_warning: + logger.warning( + "Tool registry is empty during tool-list filtering; treating tools as unknown/visible." + ) + self._has_logged_empty_registry_warning = True + elif registry_tools: + self._has_logged_empty_registry_warning = False + + unity_managed_tool_names: set[str] = set() + tool_alias_to_unity_target: dict[str, str] = {} + server_only_tool_names: set[str] = set() + signature_entries: list[tuple[str, str]] = [] + + for tool_info in registry_tools: + tool_name = tool_info.get("name") + if not isinstance(tool_name, str) or not tool_name: + continue + + unity_target = tool_info.get("unity_target", tool_name) + if unity_target is None: + server_only_tool_names.add(tool_name) + signature_entries.append((tool_name, "")) + continue + + if not isinstance(unity_target, str) or not unity_target: + logger.debug( + "Skipping tool visibility metadata with invalid unity_target: %s", + tool_info, + ) + continue + + if unity_target == tool_name: + unity_managed_tool_names.add(tool_name) + signature_entries.append((tool_name, unity_target)) + continue + + tool_alias_to_unity_target[tool_name] = unity_target + unity_managed_tool_names.add(unity_target) + signature_entries.append((tool_name, unity_target)) + + signature = tuple(sorted(signature_entries, key=lambda item: item[0])) + if signature == self._tool_visibility_signature: + self._last_tool_visibility_refresh = now + return + + self._unity_managed_tool_names = unity_managed_tool_names + self._tool_alias_to_unity_target = tool_alias_to_unity_target + self._server_only_tool_names = server_only_tool_names + self._tool_visibility_signature = signature + self._last_tool_visibility_refresh = now + @staticmethod def _resolve_candidate_project_hashes(active_instance: str | None) -> list[str]: if not active_instance: diff --git a/Server/tests/test_tool_registry_metadata.py b/Server/tests/test_tool_registry_metadata.py new file mode 100644 index 000000000..2d7329da1 --- /dev/null +++ b/Server/tests/test_tool_registry_metadata.py @@ -0,0 +1,64 @@ +import pytest + +from services.registry import get_registered_tools, mcp_for_unity_tool +import services.registry.tool_registry as tool_registry_module + + +@pytest.fixture(autouse=True) +def restore_tool_registry_state(): + original_registry = list(tool_registry_module._tool_registry) + try: + yield + finally: + tool_registry_module._tool_registry[:] = original_registry + + +def test_tool_registry_defaults_unity_target_to_tool_name(): + @mcp_for_unity_tool() + def _default_target_tool(): + return None + + registered_tools = get_registered_tools() + tool_info = next(item for item in registered_tools if item["name"] == "_default_target_tool") + assert tool_info["unity_target"] == "_default_target_tool" + + +def test_tool_registry_supports_server_only_and_alias_targets(): + @mcp_for_unity_tool(unity_target=None) + def _server_only_tool(): + return None + + @mcp_for_unity_tool(unity_target="manage_script") + def _manage_script_alias_tool(): + return None + + registered_tools = get_registered_tools() + server_only = next(item for item in registered_tools if item["name"] == "_server_only_tool") + alias_tool = next(item for item in registered_tools if item["name"] == "_manage_script_alias_tool") + + assert server_only["unity_target"] is None + assert alias_tool["unity_target"] == "manage_script" + + +def test_tool_registry_does_not_leak_unity_target_into_tool_kwargs(): + @mcp_for_unity_tool(unity_target="manage_script", annotations={"title": "x"}) + def _non_leaking_target_tool(): + return None + + registered_tools = get_registered_tools() + tool_info = next(item for item in registered_tools if item["name"] == "_non_leaking_target_tool") + assert tool_info["unity_target"] == "manage_script" + assert "unity_target" not in tool_info["kwargs"] + assert tool_info["kwargs"]["annotations"] == {"title": "x"} + + +def test_tool_registry_rejects_invalid_unity_target_values(): + with pytest.raises(ValueError, match="Invalid unity_target"): + @mcp_for_unity_tool(unity_target="") + def _invalid_empty_target_tool(): + return None + + with pytest.raises(ValueError, match="Invalid unity_target"): + @mcp_for_unity_tool(unity_target=123) # type: ignore[arg-type] + def _invalid_non_string_target_tool(): + return None diff --git a/Server/tests/test_transport_characterization.py b/Server/tests/test_transport_characterization.py index b67b711ab..c9684172f 100644 --- a/Server/tests/test_transport_characterization.py +++ b/Server/tests/test_transport_characterization.py @@ -39,6 +39,19 @@ # FIXTURES # ============================================================================ + +def _tool_registry_for_visibility_tests() -> list[dict]: + return [ + {"name": "manage_scene", "unity_target": "manage_scene"}, + {"name": "manage_script", "unity_target": "manage_script"}, + {"name": "manage_asset", "unity_target": "manage_asset"}, + {"name": "create_script", "unity_target": "manage_script"}, + {"name": "find_in_file", "unity_target": "manage_script"}, + {"name": "script_apply_edits", "unity_target": "manage_script"}, + {"name": "set_active_instance", "unity_target": None}, + {"name": "execute_custom_tool", "unity_target": None}, + ] + @pytest.fixture def mock_context(): """Create a mock FastMCP context.""" @@ -295,13 +308,25 @@ async def call_next(_ctx): with patch.object(middleware, "_inject_unity_instance", new=AsyncMock()): with patch("transport.unity_instance_middleware.PluginHub.is_configured", return_value=True): - with patch("transport.unity_instance_middleware.PluginHub.get_tools_for_project", new_callable=AsyncMock) as mock_get_tools: - mock_get_tools.return_value = [ - SimpleNamespace(name="manage_scene"), - SimpleNamespace(name="manage_script"), - ] - - filtered = await middleware.on_list_tools(middleware_ctx, call_next) + with patch("transport.unity_instance_middleware.get_registered_tools", return_value=_tool_registry_for_visibility_tests()): + with patch("transport.unity_instance_middleware.PluginHub.get_sessions", new_callable=AsyncMock) as mock_get_sessions: + with patch("transport.unity_instance_middleware.PluginHub.get_tools_for_project", new_callable=AsyncMock) as mock_get_tools: + mock_get_sessions.return_value = SessionList( + sessions={ + "session-1": SessionDetails( + project="Project", + hash="abc123", + unity_version="2022.3", + connected_at="2025-01-26T00:00:00Z", + ) + } + ) + mock_get_tools.return_value = [ + SimpleNamespace(name="manage_scene"), + SimpleNamespace(name="manage_script"), + ] + + filtered = await middleware.on_list_tools(middleware_ctx, call_next) names = [tool.name for tool in filtered] assert "manage_scene" in names @@ -336,10 +361,22 @@ async def call_next(_ctx): with patch.object(middleware, "_inject_unity_instance", new=AsyncMock()): with patch("transport.unity_instance_middleware.PluginHub.is_configured", return_value=True): - with patch("transport.unity_instance_middleware.PluginHub.get_tools_for_project", new_callable=AsyncMock) as mock_get_tools: - mock_get_tools.return_value = [] - - filtered = await middleware.on_list_tools(middleware_ctx, call_next) + with patch("transport.unity_instance_middleware.get_registered_tools", return_value=_tool_registry_for_visibility_tests()): + with patch("transport.unity_instance_middleware.PluginHub.get_sessions", new_callable=AsyncMock) as mock_get_sessions: + with patch("transport.unity_instance_middleware.PluginHub.get_tools_for_project", new_callable=AsyncMock) as mock_get_tools: + mock_get_sessions.return_value = SessionList( + sessions={ + "session-1": SessionDetails( + project="Project", + hash="abc123", + unity_version="2022.3", + connected_at="2025-01-26T00:00:00Z", + ) + } + ) + mock_get_tools.return_value = [] + + filtered = await middleware.on_list_tools(middleware_ctx, call_next) names = [tool.name for tool in filtered] assert "set_active_instance" in names @@ -373,10 +410,22 @@ async def call_next(_ctx): with patch.object(middleware, "_inject_unity_instance", new=AsyncMock()): with patch("transport.unity_instance_middleware.PluginHub.is_configured", return_value=True): - with patch("transport.unity_instance_middleware.PluginHub.get_tools_for_project", new_callable=AsyncMock) as mock_get_tools: - mock_get_tools.side_effect = RuntimeError("hub unavailable") - - filtered = await middleware.on_list_tools(middleware_ctx, call_next) + with patch("transport.unity_instance_middleware.get_registered_tools", return_value=_tool_registry_for_visibility_tests()): + with patch("transport.unity_instance_middleware.PluginHub.get_sessions", new_callable=AsyncMock) as mock_get_sessions: + with patch("transport.unity_instance_middleware.PluginHub.get_tools_for_project", new_callable=AsyncMock) as mock_get_tools: + mock_get_sessions.return_value = SessionList( + sessions={ + "session-1": SessionDetails( + project="Project", + hash="abc123", + unity_version="2022.3", + connected_at="2025-01-26T00:00:00Z", + ) + } + ) + mock_get_tools.side_effect = RuntimeError("hub unavailable") + + filtered = await middleware.on_list_tools(middleware_ctx, call_next) assert [tool.name for tool in filtered] == [tool.name for tool in original_tools] @@ -400,14 +449,203 @@ async def call_next(_ctx): with patch.object(middleware, "_inject_unity_instance", new=AsyncMock()): with patch("transport.unity_instance_middleware.PluginHub.is_configured", return_value=True): - with patch("transport.unity_instance_middleware.PluginHub.get_tools_for_project", new_callable=AsyncMock) as mock_get_tools: - mock_get_tools.return_value = [SimpleNamespace(name="manage_scene")] - - filtered = await middleware.on_list_tools(middleware_ctx, call_next) + with patch("transport.unity_instance_middleware.get_registered_tools", return_value=_tool_registry_for_visibility_tests()): + with patch("transport.unity_instance_middleware.PluginHub.get_sessions", new_callable=AsyncMock) as mock_get_sessions: + with patch("transport.unity_instance_middleware.PluginHub.get_tools_for_project", new_callable=AsyncMock) as mock_get_tools: + mock_get_sessions.return_value = SessionList( + sessions={ + "session-1": SessionDetails( + project="Project", + hash="abc123", + unity_version="2022.3", + connected_at="2025-01-26T00:00:00Z", + ) + } + ) + mock_get_tools.return_value = [SimpleNamespace(name="manage_scene")] + + filtered = await middleware.on_list_tools(middleware_ctx, call_next) assert [tool.name for tool in filtered] == ["manage_scene"] mock_get_tools.assert_awaited_once_with("abc123", user_id="user-123") + @pytest.mark.asyncio + async def test_list_tools_skips_filter_when_active_instance_hash_is_stale(self, mock_context, monkeypatch): + middleware = UnityInstanceMiddleware() + middleware_ctx = Mock() + middleware_ctx.fastmcp_context = mock_context + + mock_context.set_state("unity_instance", "Project@stale-hash") + monkeypatch.setattr(config, "transport_mode", "http") + + original_tools = [ + SimpleNamespace(name="manage_scene"), + SimpleNamespace(name="manage_asset"), + SimpleNamespace(name="set_active_instance"), + ] + + async def call_next(_ctx): + return original_tools + + with patch.object(middleware, "_inject_unity_instance", new=AsyncMock()): + with patch("transport.unity_instance_middleware.PluginHub.is_configured", return_value=True): + with patch("transport.unity_instance_middleware.get_registered_tools", return_value=_tool_registry_for_visibility_tests()): + with patch("transport.unity_instance_middleware.PluginHub.get_sessions", new_callable=AsyncMock) as mock_get_sessions: + with patch("transport.unity_instance_middleware.PluginHub.get_tools_for_project", new_callable=AsyncMock) as mock_get_tools: + mock_get_sessions.return_value = SessionList( + sessions={ + "session-1": SessionDetails( + project="Project", + hash="abc123", + unity_version="2022.3", + connected_at="2025-01-26T00:00:00Z", + ) + } + ) + + filtered = await middleware.on_list_tools(middleware_ctx, call_next) + + assert [tool.name for tool in filtered] == [tool.name for tool in original_tools] + mock_get_tools.assert_not_called() + + @pytest.mark.asyncio + async def test_list_tools_hides_alias_when_target_tool_is_disabled(self, mock_context, monkeypatch): + middleware = UnityInstanceMiddleware() + middleware_ctx = Mock() + middleware_ctx.fastmcp_context = mock_context + + mock_context.set_state("unity_instance", "Project@abc123") + monkeypatch.setattr(config, "transport_mode", "http") + + original_tools = [ + SimpleNamespace(name="manage_scene"), + SimpleNamespace(name="manage_script"), + SimpleNamespace(name="create_script"), + SimpleNamespace(name="set_active_instance"), + ] + + async def call_next(_ctx): + return original_tools + + with patch.object(middleware, "_inject_unity_instance", new=AsyncMock()): + with patch("transport.unity_instance_middleware.PluginHub.is_configured", return_value=True): + with patch("transport.unity_instance_middleware.get_registered_tools", return_value=_tool_registry_for_visibility_tests()): + with patch("transport.unity_instance_middleware.PluginHub.get_sessions", new_callable=AsyncMock) as mock_get_sessions: + with patch("transport.unity_instance_middleware.PluginHub.get_tools_for_project", new_callable=AsyncMock) as mock_get_tools: + mock_get_sessions.return_value = SessionList( + sessions={ + "session-1": SessionDetails( + project="Project", + hash="abc123", + unity_version="2022.3", + connected_at="2025-01-26T00:00:00Z", + ) + } + ) + # manage_script is disabled; alias create_script should also be hidden. + mock_get_tools.return_value = [SimpleNamespace(name="manage_scene")] + + filtered = await middleware.on_list_tools(middleware_ctx, call_next) + + names = [tool.name for tool in filtered] + assert "manage_scene" in names + assert "set_active_instance" in names + assert "manage_script" not in names + assert "create_script" not in names + + @pytest.mark.asyncio + async def test_list_tools_keeps_all_visible_when_tool_registry_is_empty(self, mock_context, monkeypatch): + middleware = UnityInstanceMiddleware() + middleware_ctx = Mock() + middleware_ctx.fastmcp_context = mock_context + + mock_context.set_state("unity_instance", "Project@abc123") + monkeypatch.setattr(config, "transport_mode", "http") + + original_tools = [ + SimpleNamespace(name="manage_scene"), + SimpleNamespace(name="set_active_instance"), + SimpleNamespace(name="execute_custom_tool"), + ] + + async def call_next(_ctx): + return original_tools + + with patch.object(middleware, "_inject_unity_instance", new=AsyncMock()): + with patch("transport.unity_instance_middleware.PluginHub.is_configured", return_value=True): + with patch("transport.unity_instance_middleware.get_registered_tools", return_value=[]): + with patch("transport.unity_instance_middleware.PluginHub.get_sessions", new_callable=AsyncMock) as mock_get_sessions: + with patch("transport.unity_instance_middleware.PluginHub.get_tools_for_project", new_callable=AsyncMock) as mock_get_tools: + mock_get_sessions.return_value = SessionList( + sessions={ + "session-1": SessionDetails( + project="Project", + hash="abc123", + unity_version="2022.3", + connected_at="2025-01-26T00:00:00Z", + ) + } + ) + mock_get_tools.return_value = [] + + filtered = await middleware.on_list_tools(middleware_ctx, call_next) + + assert [tool.name for tool in filtered] == [tool.name for tool in original_tools] + + @pytest.mark.asyncio + async def test_list_tools_uses_union_of_enabled_tools_across_multiple_sessions(self, mock_context, monkeypatch): + middleware = UnityInstanceMiddleware() + middleware_ctx = Mock() + middleware_ctx.fastmcp_context = mock_context + + monkeypatch.setattr(config, "transport_mode", "http") + + original_tools = [ + SimpleNamespace(name="manage_scene"), + SimpleNamespace(name="manage_asset"), + SimpleNamespace(name="manage_script"), + ] + + async def call_next(_ctx): + return original_tools + + async def get_tools_side_effect(project_hash, user_id=None): # noqa: ARG001 + if project_hash == "hash-a": + return [SimpleNamespace(name="manage_scene")] + if project_hash == "hash-b": + return [SimpleNamespace(name="manage_asset")] + return [] + + with patch.object(middleware, "_inject_unity_instance", new=AsyncMock()): + with patch("transport.unity_instance_middleware.PluginHub.is_configured", return_value=True): + with patch("transport.unity_instance_middleware.get_registered_tools", return_value=_tool_registry_for_visibility_tests()): + with patch("transport.unity_instance_middleware.PluginHub.get_sessions", new_callable=AsyncMock) as mock_get_sessions: + with patch("transport.unity_instance_middleware.PluginHub.get_tools_for_project", new_callable=AsyncMock) as mock_get_tools: + mock_get_sessions.return_value = SessionList( + sessions={ + "session-a": SessionDetails( + project="ProjectA", + hash="hash-a", + unity_version="2022.3", + connected_at="2025-01-26T00:00:00Z", + ), + "session-b": SessionDetails( + project="ProjectB", + hash="hash-b", + unity_version="2022.3", + connected_at="2025-01-26T00:00:00Z", + ), + } + ) + mock_get_tools.side_effect = get_tools_side_effect + + filtered = await middleware.on_list_tools(middleware_ctx, call_next) + + names = [tool.name for tool in filtered] + assert "manage_scene" in names + assert "manage_asset" in names + assert "manage_script" not in names + # ============================================================================ # AUTO-SELECT INSTANCE TESTS From 07fee1bfc41c28550f206a1f954404cdfc8a3027 Mon Sep 17 00:00:00 2001 From: whatevertogo Date: Thu, 12 Feb 2026 10:35:11 +0800 Subject: [PATCH 09/10] fix: update tool filtering logic to defer filtering when no tools are registered --- .../transport/unity_instance_middleware.py | 5 +- .../tests/test_transport_characterization.py | 67 +++++++++++++++++-- 2 files changed, 67 insertions(+), 5 deletions(-) diff --git a/Server/src/transport/unity_instance_middleware.py b/Server/src/transport/unity_instance_middleware.py index 196b7dfc3..41b4e8baf 100644 --- a/Server/src/transport/unity_instance_middleware.py +++ b/Server/src/transport/unity_instance_middleware.py @@ -349,7 +349,10 @@ async def _resolve_enabled_tool_names_for_context( for project_hash in project_hashes: try: registered_tools = await PluginHub.get_tools_for_project(project_hash, user_id=user_id) - resolved_any_project = True + # Only mark as resolved if tools are actually registered. + # An empty list means register_tools hasn't been sent yet. + if registered_tools: + resolved_any_project = True except Exception as exc: logger.debug( "Failed to fetch tools for project hash %s (user_id=%s, %s)", diff --git a/Server/tests/test_transport_characterization.py b/Server/tests/test_transport_characterization.py index c9684172f..3b68677ac 100644 --- a/Server/tests/test_transport_characterization.py +++ b/Server/tests/test_transport_characterization.py @@ -335,11 +335,12 @@ async def call_next(_ctx): assert "manage_asset" not in names @pytest.mark.asyncio - async def test_list_tools_filters_when_enabled_set_is_empty(self, mock_context, monkeypatch): + async def test_list_tools_skips_filter_when_no_tools_registered_yet(self, mock_context, monkeypatch): """ - Current behavior: when a Unity session is resolved but it has zero enabled - tools registered, Unity-managed tools are filtered out while server-only - and unknown tools remain visible. + When a Unity session is connected but register_tools has not been sent yet + (empty registered_tools), defer filtering to avoid hiding tools that may + be valid once register_tools arrives. This prevents clients that cache + early list_tools responses from getting persistently incomplete tool lists. """ middleware = UnityInstanceMiddleware() middleware_ctx = Mock() @@ -374,10 +375,68 @@ async def call_next(_ctx): ) } ) + # Simulate register_tools not yet sent mock_get_tools.return_value = [] filtered = await middleware.on_list_tools(middleware_ctx, call_next) + names = [tool.name for tool in filtered] + # All tools should be visible when register_tools hasn't been sent yet + assert "manage_scene" in names + assert "manage_asset" in names + assert "create_script" in names + assert "set_active_instance" in names + assert "custom_server_tool" in names + + @pytest.mark.asyncio + async def test_list_tools_filters_when_all_tools_disabled(self, mock_context, monkeypatch): + """ + When register_tools has been sent with an empty tool list (all tools disabled), + Unity-managed tools are filtered out while server-only tools remain visible. + This differs from the "no tools registered yet" case where we defer filtering. + """ + middleware = UnityInstanceMiddleware() + middleware_ctx = Mock() + middleware_ctx.fastmcp_context = mock_context + + mock_context.set_state("unity_instance", "Project@abc123") + monkeypatch.setattr(config, "transport_mode", "http") + + original_tools = [ + SimpleNamespace(name="manage_scene"), + SimpleNamespace(name="manage_asset"), + SimpleNamespace(name="create_script"), + SimpleNamespace(name="set_active_instance"), + SimpleNamespace(name="custom_server_tool"), + ] + + async def call_next(_ctx): + return original_tools + + # Simulate a registered tool that indicates all tools are disabled + disabled_tool = SimpleNamespace(name="_marker_tool_indicates_registration_sent") + registered_tools = [disabled_tool] + + with patch.object(middleware, "_inject_unity_instance", new=AsyncMock()): + with patch("transport.unity_instance_middleware.PluginHub.is_configured", return_value=True): + with patch("transport.unity_instance_middleware.get_registered_tools", return_value=_tool_registry_for_visibility_tests()): + with patch("transport.unity_instance_middleware.PluginHub.get_sessions", new_callable=AsyncMock) as mock_get_sessions: + with patch("transport.unity_instance_middleware.PluginHub.get_tools_for_project", new_callable=AsyncMock) as mock_get_tools: + mock_get_sessions.return_value = SessionList( + sessions={ + "session-1": SessionDetails( + project="Project", + hash="abc123", + unity_version="2022.3", + connected_at="2025-01-26T00:00:00Z", + ) + } + ) + # register_tools has been sent (non-empty list), but all tools disabled + mock_get_tools.return_value = registered_tools + + filtered = await middleware.on_list_tools(middleware_ctx, call_next) + names = [tool.name for tool in filtered] assert "set_active_instance" in names assert "custom_server_tool" in names From 6347c976a6c6cf6abcb4f13c1d109dcd6ae05e74 Mon Sep 17 00:00:00 2001 From: David Sarno Date: Thu, 12 Feb 2026 14:28:10 -0800 Subject: [PATCH 10/10] fix: add missing using UnityEditor.UIElements for IntegerField IntegerField lives in UnityEditor.UIElements, not UnityEngine.UIElements. Without this import the project fails to compile. --- MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs | 1 + 1 file changed, 1 insertion(+) diff --git a/MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs b/MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs index 8890ef0f9..4cf09aa26 100644 --- a/MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs +++ b/MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs @@ -8,6 +8,7 @@ using MCPForUnity.Editor.Services.Transport; using MCPForUnity.Editor.Tools; using UnityEditor; +using UnityEditor.UIElements; using UnityEngine.UIElements; namespace MCPForUnity.Editor.Windows.Components.Tools