fix: preserve tool toggle state and filter tools in client listings (rebased #723)#735
fix: preserve tool toggle state and filter tools in client listings (rebased #723)#735dsarno merged 11 commits intoCoplayDev:betafrom
Conversation
Fixes CoplayDev#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 <noreply@anthropic.com>
… and improve error handling
- 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 <noreply@anthropic.com>
…oplayDev#709) Comprehensive fix for Issue CoplayDev#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 <noreply@anthropic.com>
…ol visibility control
IntegerField lives in UnityEditor.UIElements, not UnityEngine.UIElements. Without this import the project fails to compile.
…hub.com/whatevertogo/unity-mcp into fix/723-rebased # Conflicts: # MCPForUnity/Editor/Services/Transport/TransportManager.cs
Reviewer's GuideImplements persistent Unity tool enable/disable state with server-side visibility filtering and user-scoped tool registration, plus client support for re-registering tools over HTTP/WebSocket, while tightening tool registry metadata and adding comprehensive tests. Sequence diagram for Unity-aware MCP tool listing with visibility filteringsequenceDiagram
actor User
participant UnityEditor as UnityEditor
participant HttpClient as HttpMcpClient
participant Server as FastMCP_Server
participant Middleware as UnityInstanceMiddleware
participant PluginHub as PluginHub
participant Registry as ToolRegistry
User->>UnityEditor: Open tools panel / request tools
UnityEditor->>HttpClient: list_tools()
HttpClient->>Server: HTTP list_tools request
Server->>Middleware: on_list_tools(context, call_next)
activate Middleware
Middleware->>Middleware: _inject_unity_instance(context)
Middleware->>Server: call_next(context)
Server-->>Middleware: tools (all_tools)
Middleware->>Middleware: _should_filter_tool_listing()
alt transport is http and PluginHub.is_configured()
Middleware->>Middleware: _refresh_tool_visibility_metadata_from_registry()
activate Middleware
Middleware->>Registry: get_registered_tools()
Registry-->>Middleware: registry_tools
Middleware->>Middleware: update _unity_managed_tool_names
Middleware->>Middleware: update _tool_alias_to_unity_target
Middleware->>Middleware: update _server_only_tool_names
deactivate Middleware
Middleware->>Middleware: _resolve_enabled_tool_names_for_context(context)
activate Middleware
Middleware->>PluginHub: get_sessions(user_id)
PluginHub-->>Middleware: SessionList
Middleware->>Middleware: _resolve_candidate_project_hashes(active_instance)
loop for each project_hash
Middleware->>PluginHub: get_tools_for_project(project_hash, user_id)
PluginHub->>Registry: get_session_id_by_hash(project_hash, user_id)
Registry-->>PluginHub: session_id
PluginHub-->>Middleware: tools_for_project
Middleware->>Middleware: accumulate enabled_tool_names
end
Middleware-->>Middleware: enabled_tool_names
deactivate Middleware
loop for tool in tools
Middleware->>Middleware: _is_tool_visible(tool.name, enabled_tool_names)
alt visible
Middleware->>Middleware: add tool to filtered list
end
end
Middleware-->>Server: filtered tools
else no filtering
Middleware-->>Server: tools (unmodified)
end
deactivate Middleware
Server-->>HttpClient: tool list (filtered)
HttpClient-->>UnityEditor: tool list (filtered)
UnityEditor-->>User: Render enabled tools
Sequence diagram for Unity tool toggle and server re-registrationsequenceDiagram
actor User
participant ToolsUI as McpToolsSection
participant ToolDiscovery as ToolDiscoveryService
participant TransportMgr as TransportManager
participant HttpClient as HttpTransportClient
participant WsClient as WebSocketTransportClient
participant Server as FastMCP_Server
User->>ToolsUI: Toggle tool checkbox
ToolsUI->>ToolsUI: HandleToggleChange(tool, enabled, updateSummary, reregisterTools)
ToolsUI->>ToolDiscovery: SetToolEnabled(tool.Name, enabled)
ToolDiscovery-->>ToolsUI: preference updated
alt updateSummary
ToolsUI->>ToolsUI: UpdateSummary()
end
alt reregisterTools
ToolsUI->>ToolsUI: ReregisterToolsAsync()
ToolsUI->>TransportMgr: GetClient(TransportMode.Http)
TransportMgr-->>ToolsUI: HttpClient or null
alt client is null or !IsConnected
ToolsUI-->>User: Local state updated, no server re-register
else client connected
ToolsUI->>HttpClient: ReregisterToolsAsync()
activate HttpClient
HttpClient->>WsClient: (HTTP path) or internal transport to server
WsClient->>WsClient: SendRegisterToolsAsync(token)
WsClient->>Server: register_tools payload
Server-->>WsClient: register_tools response
WsClient-->>HttpClient: reregistration complete
deactivate HttpClient
HttpClient-->>ToolsUI: ReregisterToolsAsync completed
end
end
User->>ToolsUI: Click Enable All / Disable All
ToolsUI->>ToolsUI: SetAllToolsState(enabled)
ToolsUI->>ToolDiscovery: IsToolEnabled(tool.Name) for each tool
loop tools
alt state differs
ToolsUI->>ToolDiscovery: SetToolEnabled(tool.Name, enabled)
ToolsUI->>ToolsUI: HandleToggleChange(tool, enabled, updateSummary=false, reregisterTools=false)
end
end
ToolsUI->>ToolsUI: UpdateSummary()
alt any changes
ToolsUI->>ToolsUI: ReregisterToolsAsync() (single bulk call)
end
Class diagram for transport clients and tool re-registrationclassDiagram
class IMcpTransportClient {
<<interface>>
+bool IsConnected
+Task~bool~ StartAsync()
+Task StopAsync()
+Task~bool~ VerifyAsync()
+Task ReregisterToolsAsync()
}
class WebSocketTransportClient {
-CancellationTokenSource _lifecycleCts
+bool IsConnected
+Task~bool~ StartAsync()
+Task StopAsync()
+Task~bool~ VerifyAsync()
+Task ReregisterToolsAsync()
-Task SendRegisterToolsAsync(CancellationToken token)
}
class StdioTransportClient {
+bool IsConnected
+Task~bool~ StartAsync()
+Task StopAsync()
+Task~bool~ VerifyAsync()
+Task ReregisterToolsAsync()
}
class TransportManager {
-IMcpTransportClient _httpClient
-IMcpTransportClient _stdioClient
+Task~bool~ StartAsync(TransportMode mode)
+Task StopAsync(TransportMode mode)
+Task~bool~ VerifyAsync(TransportMode mode)
+void ForceStop(TransportMode mode)
+IMcpTransportClient GetClient(TransportMode mode)
-IMcpTransportClient GetOrCreateClient(TransportMode mode)
-void UpdateState(TransportMode mode, TransportState state)
}
class McpToolsSection {
-IEnumerable~ToolMetadata~ allTools
-Dictionary~string, Toggle~ toolToggleMap
-void HandleToggleChange(ToolMetadata tool, bool enabled, bool updateSummary, bool reregisterTools)
-void ReregisterToolsAsync()
-void SetAllToolsState(bool enabled)
-void UpdateSummary()
}
IMcpTransportClient <|.. WebSocketTransportClient
IMcpTransportClient <|.. StdioTransportClient
TransportManager --> IMcpTransportClient : uses
McpToolsSection --> TransportManager : uses
McpToolsSection ..> IMcpTransportClient : calls ReregisterToolsAsync()
class ToolMetadata {
+string Name
+bool AutoRegister
+bool IsBuiltIn
}
class Toggle {
+void SetValueWithoutNotify(bool value)
}
McpToolsSection --> ToolMetadata
McpToolsSection --> Toggle
Class diagram for tool registry decorator and visibility metadataclassDiagram
class ToolRegistryModule {
+list~dict~ _tool_registry
+Callable mcp_for_unity_tool(str name, str description, str unity_target, **kwargs)
}
class mcp_for_unity_tool_decorator {
-str tool_name
-str_or_none normalized_unity_target
-dict tool_kwargs
+Callable __call__(Callable func)
}
ToolRegistryModule o--> mcp_for_unity_tool_decorator
class UnityInstanceMiddleware {
-dict~str, str~ _active_by_key
-RLock _lock
-RLock _metadata_lock
-set~str~ _unity_managed_tool_names
-dict~str, str~ _tool_alias_to_unity_target
-set~str~ _server_only_tool_names
-tuple~tuple~str, str~~ _tool_visibility_signature
-float _last_tool_visibility_refresh
-float _tool_visibility_refresh_interval_seconds
-bool _has_logged_empty_registry_warning
+async on_list_tools(MiddlewareContext context, Callable call_next)
-bool _should_filter_tool_listing()
+async set_on_request_handlers()
-async set_on_list_tools_handler()
-async _resolve_enabled_tool_names_for_context(MiddlewareContext context)
-void _refresh_tool_visibility_metadata_from_registry()
-static list~str~ _resolve_candidate_project_hashes(str active_instance)
-bool _is_tool_visible(str tool_name, set~str~ enabled_tool_names)
}
class PluginHub {
+async classmethod get_sessions(str user_id)
+async classmethod get_tools_for_project(str project_hash, str user_id)
+async classmethod get_tool_definition(str project_hash, str tool_name, str user_id)
}
class ToolInfoDict {
+str name
+str_or_none unity_target
+dict kwargs
}
UnityInstanceMiddleware --> ToolRegistryModule : calls get_registered_tools()
UnityInstanceMiddleware --> PluginHub : queries sessions and tools
ToolRegistryModule --> ToolInfoDict : stores metadata
Class diagram for CustomToolService and user-scoped tool executionclassDiagram
class CustomToolService {
-dict~str, dict~str, ToolDefinitionModel~~ _project_tools
+static CustomToolService get_instance()
+async JSONResponse register_tools(Request request)
+async list_registered_tools(str project_id, str user_id)
+async get_tool_definition(str project_id, str tool_name, str user_id)
+async execute_tool(str project_id, str tool_name, str unity_instance, dict params, str user_id) MCPResponse
-async MCPResponse _poll_until_complete(str project_id, str tool_name, str unity_instance, dict initial_params, MCPResponse initial_response, str poll_action, str user_id)
}
class ToolDefinitionModel {
+str name
+bool requires_polling
+str_or_none poll_action
+str description
}
class Context {
+object get_state(str key)
}
class CustomToolsResource {
+async CustomToolsResourceResponse get_custom_tools(Context ctx)
}
class ExecuteCustomToolHandler {
+async MCPResponse execute_custom_tool(Context ctx, str tool_name, dict parameters)
}
class HelperFunctions {
+str_or_none get_user_id_from_context(Context ctx)
}
class PluginHub {
+async classmethod get_tools_for_project(str project_hash, str user_id)
+async classmethod get_tool_definition(str project_hash, str tool_name, str user_id)
}
CustomToolService --> ToolDefinitionModel
CustomToolsResource --> CustomToolService : list_registered_tools(project_id, user_id)
ExecuteCustomToolHandler --> CustomToolService : execute_tool(project_id, tool_name, unity_instance, parameters, user_id)
CustomToolService --> PluginHub : get_tools_for_project, get_tool_definition
HelperFunctions --> Context
CustomToolsResource --> HelperFunctions : get_user_id_from_context(ctx)
ExecuteCustomToolHandler --> HelperFunctions : get_user_id_from_context(ctx)
File-Level Changes
Assessment against linked issues
Possibly linked issues
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
|
Caution Review failedThe pull request is closed. 📝 WalkthroughWalkthroughThis PR implements dynamic tool reregistration and user-scoped visibility filtering to fix disabled tools reappearing and remaining visible to clients. The system removes forced-enable logic, adds transport-layer reregistration when tools are toggled, introduces registry-based metadata for tool targeting, propagates user identities through service methods, and adds server-side filtering of tools based on enabled state and user context. Changes
Sequence Diagram(s)sequenceDiagram
participant User as User (UI)
participant UI as McpToolsSection
participant Transport as WebSocket Transport
participant Server as MCP Server
User->>UI: Toggle tool enabled/disabled
UI->>UI: Update tool state in batch
UI->>UI: Detect changes occurred
UI->>Transport: ReregisterToolsAsync()
Transport->>Transport: Check if connected
alt Connected
Transport->>Server: SendRegisterToolsAsync()
Server->>Server: Update registered tools
Server-->>Transport: Success
Transport-->>UI: Task completed
else Not Connected
Transport-->>UI: Warning logged, Task completed
end
sequenceDiagram
participant Client as MCP Client
participant Middleware as UnityInstanceMiddleware
participant PluginHub as PluginHub
participant Registry as Tool Registry
Client->>Middleware: Request list_tools
Middleware->>Middleware: Inject Unity instance data
alt Should filter (HTTP mode)
Middleware->>Registry: Refresh visibility metadata
Registry-->>Middleware: Tool visibility sets/mappings
Middleware->>PluginHub: Get sessions and projects
PluginHub-->>Middleware: Session/project data
Middleware->>PluginHub: Get tools_for_project(user_id)
PluginHub-->>Middleware: User-scoped tools
Middleware->>Middleware: Filter by visibility and enabled state
end
Middleware-->>Client: Filtered tool list
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Poem
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Hey - I've found 2 issues
Prompt for AI Agents
Please address the comments from this code review:
## Individual Comments
### Comment 1
<location> `Server/src/transport/unity_instance_middleware.py:396-400` </location>
<code_context>
+ 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
</code_context>
<issue_to_address>
**suggestion (bug_risk):** Empty registry warning message does not match actual behavior when previous metadata exists
When `registry_tools` is empty you log that tools are treated as unknown/visible, but the existing `_unity_managed_tool_names`, `_tool_alias_to_unity_target`, and `_server_only_tool_names` remain unchanged. If the registry was previously populated, tools still follow the old metadata. Please either update the log message to state that existing metadata is retained, or clear these metadata structures when the registry is empty so behavior matches the message.
</issue_to_address>
### Comment 2
<location> `Server/src/services/custom_tool_service.py:32-41` </location>
<code_context>
_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
+
+
</code_context>
<issue_to_address>
**suggestion (bug_risk):** Broad exception handling around ctx.get_state may hide programming/configuration issues
This helper currently catches all `Exception` from `ctx.get_state("user_id")` and returns `None`, which can mask real context/configuration bugs (e.g., attribute/type errors). Please narrow the `except` to the specific failure modes you expect from `get_state`, or at least log the exception (e.g., at debug level), so genuine integration issues remain visible while still tolerating missing state in normal operation.
Suggested implementation:
```python
import logging
from starlette.requests import Request
```
```python
_MAX_POLL_SECONDS = 600
logger = logging.getLogger(__name__)
```
```python
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 KeyError:
# Expected case: user_id state is not present
return None
except Exception as exc:
# Log unexpected integration/configuration issues while remaining tolerant
logger.debug("Failed to read 'user_id' from context state", exc_info=exc)
return None
return user_id if isinstance(user_id, str) and user_id else None
```
If `ctx.get_state` can raise more specific, known exceptions (e.g. a custom context error), you should replace or augment the generic `Exception` handler with those specific types so that truly unexpected errors can still propagate.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
| 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 |
There was a problem hiding this comment.
suggestion (bug_risk): Empty registry warning message does not match actual behavior when previous metadata exists
When registry_tools is empty you log that tools are treated as unknown/visible, but the existing _unity_managed_tool_names, _tool_alias_to_unity_target, and _server_only_tool_names remain unchanged. If the registry was previously populated, tools still follow the old metadata. Please either update the log message to state that existing metadata is retained, or clear these metadata structures when the registry is empty so behavior matches the message.
| 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: |
There was a problem hiding this comment.
suggestion (bug_risk): Broad exception handling around ctx.get_state may hide programming/configuration issues
This helper currently catches all Exception from ctx.get_state("user_id") and returns None, which can mask real context/configuration bugs (e.g., attribute/type errors). Please narrow the except to the specific failure modes you expect from get_state, or at least log the exception (e.g., at debug level), so genuine integration issues remain visible while still tolerating missing state in normal operation.
Suggested implementation:
import logging
from starlette.requests import Request_MAX_POLL_SECONDS = 600
logger = logging.getLogger(__name__)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 KeyError:
# Expected case: user_id state is not present
return None
except Exception as exc:
# Log unexpected integration/configuration issues while remaining tolerant
logger.debug("Failed to read 'user_id' from context state", exc_info=exc)
return None
return user_id if isinstance(user_id, str) and user_id else NoneIf ctx.get_state can raise more specific, known exceptions (e.g. a custom context error), you should replace or augment the generic Exception handler with those specific types so that truly unexpected errors can still propagate.
Rebased version of #723 by @whatevertogo to resolve merge conflicts after #688 landed on beta.
Original PR: #723 | Original author: @whatevertogo
Fixes #709 - Disabled tools now persist across Unity Editor restarts and are correctly filtered from client tool listings.
Conflict resolution: kept both ForceStop() from #688 and public GetClient() from #723 in TransportManager.cs.
Bug fix included: added missing using UnityEditor.UIElements for IntegerField (pushed directly to author branch).
Summary by Sourcery
Align Unity tool enablement between client and server, ensuring disabled tools persist across sessions and are correctly filtered from tool listings in HTTP/hosted modes.
Bug Fixes:
Enhancements:
Tests:
Summary by CodeRabbit
New Features
Bug Fixes