Fix: #709 - Tool state persistence and filtering#722
Fix: #709 - Tool state persistence and filtering#722whatevertogo wants to merge 3 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>
Reviewer's GuideAdds Unity-aware tool filtering on the server and wiring on the Unity client to persist tool enablement and re‑register tools over HTTP WebSocket, plus tests for both transport filtering and editor tool discovery persistence. Sequence diagram for Unity tool toggle and HTTP WebSocket tool reregistrationsequenceDiagram
actor UnityUser
participant McpToolsSection
participant ToolDiscoveryService
participant TransportManager
participant IMcpTransportClient
participant WebSocketTransportClient
participant MCPServer
UnityUser->>McpToolsSection: Toggle tool checkbox
McpToolsSection->>ToolDiscoveryService: SetToolEnabled(toolName, enabled)
McpToolsSection->>McpToolsSection: HandleToggleChange(tool, enabled, updateSummary, reregisterTools)
alt updateSummary
McpToolsSection->>McpToolsSection: UpdateSummary()
end
alt reregisterTools
McpToolsSection->>McpToolsSection: ReregisterToolsAsync() (ThreadPool)
McpToolsSection->>TransportManager: GetClient(TransportMode_Http)
TransportManager-->>McpToolsSection: IMcpTransportClient or null
alt client is null or not connected
McpToolsSection-->>McpToolsSection: Skip reregistration
else client is WebSocketTransportClient and connected
McpToolsSection->>IMcpTransportClient: ReregisterToolsAsync()
activate WebSocketTransportClient
WebSocketTransportClient->>WebSocketTransportClient: SendRegisterToolsAsync(_lifecycleCts_Token)
WebSocketTransportClient->>MCPServer: Send tools registration
MCPServer-->>WebSocketTransportClient: Ack registration
deactivate WebSocketTransportClient
end
end
Sequence diagram for server-side Unity-aware tool list filteringsequenceDiagram
participant UnityClient
participant UnityInstanceMiddleware
participant PluginHub
participant FastMcpContext
participant MCPServerBackend
UnityClient->>UnityInstanceMiddleware: list_tools request
UnityInstanceMiddleware->>UnityInstanceMiddleware: on_list_tools(context, call_next)
UnityInstanceMiddleware->>UnityInstanceMiddleware: _inject_unity_instance(context)
UnityInstanceMiddleware->>MCPServerBackend: call_next(context)
MCPServerBackend-->>UnityInstanceMiddleware: tools
UnityInstanceMiddleware->>UnityInstanceMiddleware: _should_filter_tool_listing()
alt filtering disabled
UnityInstanceMiddleware-->>UnityClient: tools
else filtering enabled
UnityInstanceMiddleware->>FastMcpContext: get_state(user_id)
UnityInstanceMiddleware->>FastMcpContext: get_state(unity_instance)
UnityInstanceMiddleware->>UnityInstanceMiddleware: _resolve_candidate_project_hashes(active_instance)
alt project_hashes empty
UnityInstanceMiddleware->>PluginHub: get_sessions(user_id)
PluginHub-->>UnityInstanceMiddleware: sessions
alt no sessions
UnityInstanceMiddleware-->>UnityClient: tools
else one or more sessions
UnityInstanceMiddleware->>UnityInstanceMiddleware: derive project_hashes from sessions
end
end
UnityInstanceMiddleware->>UnityInstanceMiddleware: enabled_tool_names = {}
loop project_hash in project_hashes
UnityInstanceMiddleware->>PluginHub: get_tools_for_project(project_hash)
PluginHub-->>UnityInstanceMiddleware: registered_tools
UnityInstanceMiddleware->>UnityInstanceMiddleware: collect tool.name into enabled_tool_names
end
UnityInstanceMiddleware->>UnityInstanceMiddleware: filtered = []
loop tool in tools
UnityInstanceMiddleware->>UnityInstanceMiddleware: _is_tool_visible(tool.name, enabled_tool_names)
alt visible
UnityInstanceMiddleware->>UnityInstanceMiddleware: add tool to filtered
end
end
UnityInstanceMiddleware-->>UnityClient: filtered
end
Updated class diagram for Unity transport clients, tools UI, and server middlewareclassDiagram
class IMcpTransportClient {
<<interface>>
+bool IsConnected
+Task StartAsync()
+Task StopAsync()
+Task~bool~ VerifyAsync()
+Task ReregisterToolsAsync()
}
class WebSocketTransportClient {
+bool IsConnected
-CancellationTokenSource _lifecycleCts
+Task StartAsync()
+Task StopAsync()
+Task~bool~ VerifyAsync()
+Task ReregisterToolsAsync()
-Task SendRegisterToolsAsync(CancellationToken token)
}
class StdioTransportClient {
+bool IsConnected
+Task StartAsync()
+Task StopAsync()
+Task~bool~ VerifyAsync()
+Task ReregisterToolsAsync()
}
class TransportManager {
-IMcpTransportClient _httpClient
-IMcpTransportClient _stdioClient
+IMcpTransportClient GetOrCreateClient(TransportMode mode)
+IMcpTransportClient GetClient(TransportMode mode)
+Task~bool~ StartAsync(TransportMode mode)
+Task StopAsync(TransportMode mode)
+Task~bool~ VerifyAsync(TransportMode mode)
+TransportState GetState(TransportMode mode)
+bool IsRunning(TransportMode mode)
}
class McpToolsSection {
-Dictionary~string, Toggle~ toolToggleMap
-List~ToolMetadata~ allTools
-VisualElement CreateToolRow(ToolMetadata tool)
-void HandleToggleChange(ToolMetadata tool, bool enabled, bool updateSummary, bool reregisterTools)
-void ReregisterToolsAsync()
-void SetAllToolsState(bool enabled)
-void UpdateSummary()
}
class ToolDiscoveryService {
+void SetToolEnabled(string toolName, bool enabled)
+bool IsToolEnabled(string toolName)
-void EnsurePreferenceInitialized(ToolMetadata metadata)
}
class UnityInstanceMiddleware {
-dict~string, string~ _active_by_key
-set~string~ _unity_managed_tool_names
-dict~string, string~ _tool_alias_to_unity_target
-set~string~ _server_only_tool_names
+async Task on_list_tools(MiddlewareContext context, Func~MiddlewareContext, Task~ call_next)
-bool _should_filter_tool_listing()
-async Task~set~string~~ _resolve_enabled_tool_names_for_context(MiddlewareContext context)
-static List~string~ _resolve_candidate_project_hashes(string active_instance)
-bool _is_tool_visible(string tool_name, set~string~ enabled_tool_names)
}
class MCPServiceLocator {
<<static>>
+ToolDiscoveryService ToolDiscovery
+TransportManager TransportManager
}
IMcpTransportClient <|.. WebSocketTransportClient
IMcpTransportClient <|.. StdioTransportClient
TransportManager --> IMcpTransportClient : uses
McpToolsSection --> ToolDiscoveryService : uses
McpToolsSection --> TransportManager : uses
McpToolsSection --> MCPServiceLocator : accesses
MCPServiceLocator --> ToolDiscoveryService : provides
MCPServiceLocator --> TransportManager : provides
UnityInstanceMiddleware --> PluginHub : uses
UnityInstanceMiddleware --> MiddlewareContext : uses
UnityInstanceMiddleware --> FastMcpContext : uses
File-Level Changes
Assessment against linked issues
Possibly linked issues
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
📝 WalkthroughWalkthroughThis PR resolves issue Changes
Sequence Diagram(s)sequenceDiagram
participant UI as McpToolsSection
participant TM as TransportManager
participant TC as TransportClient
participant Server
UI->>UI: HandleToggleChange
UI->>UI: SetAllToolsState with reregisterTools=false
loop For each tool
UI->>UI: Toggle without reregistration
end
UI->>TM: ReregisterToolsAsync()
TM->>TC: GetClient() returns HTTP client
TM->>TC: ReregisterToolsAsync()
TC->>Server: SendRegisterToolsAsync()
Server->>Server: on_list_tools filters by enabled state
Server-->>TC: Filtered tool list
TC-->>TM: Success logged
TM-->>UI: Complete
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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, and left some high level feedback:
- The hard-coded tool name sets and alias mappings in
UnityInstanceMiddleware(_unity_managed_tool_names, _tool_alias_to_unity_target, _server_only_tool_names) may become a maintenance hotspot; consider centralizing these definitions or deriving them from a single Unity-side source to reduce the risk of drift. - In
McpToolsSection.ReregisterToolsAsync, usingThreadPool.QueueUserWorkItemwithclient.ReregisterToolsAsync().Wait()introduces a sync-over-async pattern on a thread-pool thread; consider usingTask.Runwithawait client.ReregisterToolsAsync()to keep it fully async and avoid potential blocking or deadlock issues.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The hard-coded tool name sets and alias mappings in `UnityInstanceMiddleware` (_unity_managed_tool_names, _tool_alias_to_unity_target, _server_only_tool_names) may become a maintenance hotspot; consider centralizing these definitions or deriving them from a single Unity-side source to reduce the risk of drift.
- In `McpToolsSection.ReregisterToolsAsync`, using `ThreadPool.QueueUserWorkItem` with `client.ReregisterToolsAsync().Wait()` introduces a sync-over-async pattern on a thread-pool thread; consider using `Task.Run` with `await client.ReregisterToolsAsync()` to keep it fully async and avoid potential blocking or deadlock issues.
## Individual Comments
### Comment 1
<location> `Server/src/transport/unity_instance_middleware.py:352-353` </location>
<code_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)
</code_context>
<issue_to_address>
**issue (bug_risk):** Guard against missing `hash` attribute when building `project_hashes` for multiple sessions.
In this comprehension, `session.hash` is accessed before the `if` condition is evaluated, so a session without a `hash` attribute will still raise `AttributeError`. Use `getattr` for both access and filtering, e.g.:
```python
project_hashes = [
h
for session in sessions.values()
if (h := getattr(session, "hash", None))
]
```
or an equivalent version without the walrus operator.
</issue_to_address>
### Comment 2
<location> `MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs:509-530` </location>
<code_context>
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}");
</code_context>
<issue_to_address>
**suggestion (bug_risk):** Handle potential races where `_lifecycleCts` is disposed between the initial check and use.
In `ReregisterToolsAsync`, there’s a race between the `_lifecycleCts != null` check and `SendRegisterToolsAsync(_lifecycleCts.Token)`: if the CTS is disposed in between (e.g., during shutdown), you may get `ObjectDisposedException` instead of `OperationCanceledException`.
To harden this:
- capture `_lifecycleCts` into a local and use that consistently, and
- consider handling `ObjectDisposedException` the same as `OperationCanceledException` (treat as benign cancellation rather than an error) to avoid noisy shutdown logs.
```suggestion
public async Task ReregisterToolsAsync()
{
// Capture CTS locally to avoid races with concurrent disposal/change
var lifecycleCts = _lifecycleCts;
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.ObjectDisposedException)
{
// Treat disposal during shutdown the same as cancellation to avoid noisy logs
McpLog.Warn("[WebSocket] Tool reregistration cancelled");
}
catch (System.Exception ex)
{
McpLog.Error($"[WebSocket] Tool reregistration failed: {ex.Message}");
}
}
```
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
| project_hashes = [ | ||
| session.hash |
There was a problem hiding this comment.
issue (bug_risk): Guard against missing hash attribute when building project_hashes for multiple sessions.
In this comprehension, session.hash is accessed before the if condition is evaluated, so a session without a hash attribute will still raise AttributeError. Use getattr for both access and filtering, e.g.:
project_hashes = [
h
for session in sessions.values()
if (h := getattr(session, "hash", None))
]or an equivalent version without the walrus operator.
| 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}"); | ||
| } | ||
| } |
There was a problem hiding this comment.
suggestion (bug_risk): Handle potential races where _lifecycleCts is disposed between the initial check and use.
In ReregisterToolsAsync, there’s a race between the _lifecycleCts != null check and SendRegisterToolsAsync(_lifecycleCts.Token): if the CTS is disposed in between (e.g., during shutdown), you may get ObjectDisposedException instead of OperationCanceledException.
To harden this:
- capture
_lifecycleCtsinto a local and use that consistently, and - consider handling
ObjectDisposedExceptionthe same asOperationCanceledException(treat as benign cancellation rather than an error) to avoid noisy shutdown logs.
| 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}"); | |
| } | |
| } | |
| public async Task ReregisterToolsAsync() | |
| { | |
| // Capture CTS locally to avoid races with concurrent disposal/change | |
| var lifecycleCts = _lifecycleCts; | |
| 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.ObjectDisposedException) | |
| { | |
| // Treat disposal during shutdown the same as cancellation to avoid noisy logs | |
| McpLog.Warn("[WebSocket] Tool reregistration cancelled"); | |
| } | |
| catch (System.Exception ex) | |
| { | |
| McpLog.Error($"[WebSocket] Tool reregistration failed: {ex.Message}"); | |
| } | |
| } |
There was a problem hiding this comment.
Pull request overview
Fixes issue #709 by persisting tool enabled/disabled state across Unity restarts and ensuring MCP clients only see Unity tools that are enabled in the Unity editor (while preserving server-only tools and selected aliases).
Changes:
- Stop overwriting stored tool enabled/disabled preferences during tool discovery (EditorPrefs persistence fix).
- Add server-side tool-list filtering in
UnityInstanceMiddleware.on_list_tools()based on PluginHub-registered tool names, with alias handling. - Add Unity edit-mode tests and server characterization tests to cover persistence and tool-list filtering behavior.
Reviewed changes
Copilot reviewed 12 out of 12 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
MCPForUnity/Editor/Services/ToolDiscoveryService.cs |
Fixes preference initialization so stored disabled states aren’t overwritten on discovery/restart. |
MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs |
Triggers tool re-registration with the connected server when tool toggles change (including bulk enable/disable). |
MCPForUnity/Editor/Services/Transport/IMcpTransportClient.cs |
Extends transport interface with ReregisterToolsAsync(). |
MCPForUnity/Editor/Services/Transport/TransportManager.cs |
Exposes GetClient() publicly so UI can request the active client instance. |
MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs |
Implements tool re-registration by resending tool registration over the WS connection. |
MCPForUnity/Editor/Services/Transport/Transports/StdioTransportClient.cs |
Implements no-op re-registration (stdio doesn’t support dynamic tool re-registration). |
Server/src/transport/unity_instance_middleware.py |
Filters tool listing to Unity-enabled set (plus server-only + alias mapping) when in HTTP mode with PluginHub configured. |
Server/tests/test_transport_characterization.py |
Adds characterization tests verifying filtering behavior and alias/server-only exceptions. |
TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/ToolDiscoveryServiceTests.cs |
Adds edit-mode tests ensuring tool state persistence and no override of stored disabled values. |
TestProjects/UnityMCPTests/Assets/Tests/Editor.meta / .../State.meta / ...ToolDiscoveryServiceTests.cs.meta |
Adds Unity asset metadata for the new test folders/files. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| client.ReregisterToolsAsync().Wait(); | ||
| } | ||
| } | ||
| catch (Exception ex) | ||
| { | ||
| McpLog.Warn($"Failed to reregister tools: {ex.Message}"); |
There was a problem hiding this comment.
ReregisterToolsAsync() blocks on an async call using .Wait(). This can wrap exceptions in AggregateException (often resulting in an unhelpful One or more errors occurred message) and unnecessarily blocks a ThreadPool thread. Prefer an async flow (e.g., Task.Run(async () => await client.ReregisterToolsAsync())) or at least use GetAwaiter().GetResult() and log the inner exception details.
| client.ReregisterToolsAsync().Wait(); | |
| } | |
| } | |
| catch (Exception ex) | |
| { | |
| McpLog.Warn($"Failed to reregister tools: {ex.Message}"); | |
| // Use GetAwaiter().GetResult() to avoid AggregateException wrapping | |
| client.ReregisterToolsAsync().GetAwaiter().GetResult(); | |
| } | |
| } | |
| catch (Exception ex) | |
| { | |
| // Log full exception details, including any inner exceptions | |
| if (ex is AggregateException aggregateException && aggregateException.InnerException != null) | |
| { | |
| McpLog.Warn($"Failed to reregister tools: {aggregateException.InnerException}"); | |
| } | |
| else | |
| { | |
| McpLog.Warn($"Failed to reregister tools: {ex}"); | |
| } |
| try: | ||
| sessions_data = await PluginHub.get_sessions(user_id=user_id) | ||
| sessions = sessions_data.sessions if sessions_data else {} | ||
| except Exception: |
There was a problem hiding this comment.
The except Exception: return None around PluginHub.get_sessions(...) swallows all errors without logging. This can make it hard to understand why tool filtering is skipped, and can hide real bugs. Consider logging (debug/warn) and re-raising SystemExit/KeyboardInterrupt (consistent with the rest of this file).
| except Exception: | |
| except (KeyboardInterrupt, SystemExit): | |
| # Preserve process-control semantics. | |
| raise | |
| except Exception as exc: | |
| logger.warning( | |
| "Failed to fetch sessions from PluginHub for user %r; " | |
| "skipping tool filtering: %s", | |
| user_id, | |
| exc, | |
| ) |
| for project_hash in project_hashes: | ||
| try: | ||
| registered_tools = await PluginHub.get_tools_for_project(project_hash) | ||
| except Exception: |
There was a problem hiding this comment.
The except Exception: continue around PluginHub.get_tools_for_project(...) silently drops failures and can yield an empty enabled-tool set, causing the middleware to return the unfiltered tool list. Consider logging the exception (at least debug) so failures don't silently negate filtering behavior.
| except Exception: | |
| except Exception: | |
| logger.debug( | |
| "Failed to fetch tools for project %s in UnityInstanceMiddleware; skipping this project.", | |
| project_hash, | |
| exc_info=True, | |
| ) |
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Fix all issues with AI agents
In `@MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs`:
- Around line 248-267: ReregisterToolsAsync currently hardcodes
TransportMode.Http when calling MCPServiceLocator.TransportManager.GetClient,
which ignores the user's configured transport; change it to use the active
transport mode from MCPServiceLocator.Bridge.ActiveMode when obtaining the
client (keep the existing null and IsConnected checks and the try/catch around
client.ReregisterToolsAsync().Wait()). Update references in the
ReregisterToolsAsync method to call GetClient with
MCPServiceLocator.Bridge.ActiveMode so reregistration respects the user's
preferred transport.
In `@Server/src/transport/unity_instance_middleware.py`:
- Around line 362-371: The current loop over project_hashes swallows all
exceptions from PluginHub.get_tools_for_project by using bare except and
continue; change the except to capture the exception as e (e.g., except
Exception as e) and log the failure (including project_hash and exception
details/traceback) before continuing so failures are discoverable — update the
except block around PluginHub.get_tools_for_project to call the module or
instance logger (e.g., logger.exception or logger.error with exc_info=True)
referencing project_hash and the caught exception, then continue to the next
hash.
- Around line 335-339: The try/except around PluginHub.get_sessions currently
swallows all exceptions; change it to "except Exception as e:" and log the
exception at debug (or debug/exception) level before returning None so
programming errors are visible; specifically update the block around
PluginHub.get_sessions/user_id to capture the exception into a variable (e) and
call the module's logger (e.g., logger.debug or logger.exception) with a clear
message including e and any contextual info (user_id, sessions_data) prior to
returning None.
🧹 Nitpick comments (5)
Server/src/transport/unity_instance_middleware.py (1)
58-94: Hardcoded tool name sets will drift as tools are added/removed.The
_unity_managed_tool_names,_tool_alias_to_unity_target, and_server_only_tool_namesare static copies of the current tool inventory. New tools added in Unity won't appear here, and removed tools will linger. The fail-open approach in_is_tool_visible(line 401) mitigates this for new tools — they'll remain visible even when disabled, but won't be filterable until this list is updated.Consider whether these sets could be derived dynamically (e.g., from the
PluginHubregistry or a shared manifest), or at least add a comment noting the maintenance requirement.MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs (1)
509-530:ReregisterToolsAsyncimplementation is correct and well-guarded.The pre-checks, delegation to
SendRegisterToolsAsync, and structured exception handling are all appropriate. The broadcatch (Exception)at line 526 provides a safety net for edge cases like the_lifecycleCtsbeing disposed concurrently.Minor style nit: Lines 522 and 526 use fully-qualified
System.OperationCanceledExceptionandSystem.Exception, but the file already hasusing System;at line 1, so theSystem.prefix is unnecessary. The rest of the file uses the unqualified forms (e.g., lines 326, 336, 558, 566).Optional cleanup
- catch (System.OperationCanceledException) + catch (OperationCanceledException) { McpLog.Warn("[WebSocket] Tool reregistration cancelled"); } - catch (System.Exception ex) + catch (Exception ex) { McpLog.Error($"[WebSocket] Tool reregistration failed: {ex.Message}"); }MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs (1)
251-259: ConsiderTask.RunoverThreadPool.QueueUserWorkItem+.Wait().
.Wait()on aTaskinside a thread pool callback works here but risksAggregateExceptionwrapping.Task.Run(async () => await ...)is more idiomatic for fire-and-forget async and avoids synchronous blocking:Suggested change
- ThreadPool.QueueUserWorkItem(_ => + _ = 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(); } } catch (Exception ex) { McpLog.Warn($"Failed to reregister tools: {ex.Message}"); } });TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/ToolDiscoveryServiceTests.cs (2)
116-153: Test relies on the test project containing a built-in tool withAutoRegister=false.Line 124 asserts that at least one such tool exists. If the test project's tool configuration ever changes (e.g., all built-in tools become
AutoRegister=true), this test will fail with the message but no obvious link back to the cause. Consider adding a brief comment documenting which tool(s) this currently matches, to ease future debugging.
51-59: Redundant cleanup duplicatesSetUplogic.Lines 55-59 manually clean up the preference key, but
SetUp(lines 17-22) already does this before each test. The duplicate is harmless but unnecessary.
| 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}"); | ||
| } | ||
| }); | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check if BridgeService exposes preferred mode publicly and how Stdio transport implements ReregisterToolsAsync
echo "=== BridgeControlService public API ==="
rg -n "PreferredMode|_preferredMode|ResolvePreferredMode" --type=cs -C2
echo ""
echo "=== StdioTransportClient.ReregisterToolsAsync ==="
ast-grep --pattern $'class StdioTransportClient {
$$$
ReregisterToolsAsync() {
$$$
}
$$$
}'
echo ""
echo "=== All ReregisterToolsAsync implementations ==="
rg -n "ReregisterToolsAsync" --type=cs -C3Repository: CoplayDev/unity-mcp
Length of output: 9942
🏁 Script executed:
#!/bin/bash
echo "=== MCPServiceLocator API ==="
rg -n "class MCPServiceLocator|public static.*BridgeService|BridgeService.*get" --type=cs -A2 -B2
echo ""
echo "=== TransportManager.GetClient implementation ==="
ast-grep --pattern $'GetClient($MODE) {
$$$
}'
echo ""
echo "=== How GetClient is used in other places ==="
rg -n "GetClient\(" --type=cs -B2 -A2 | head -80Repository: CoplayDev/unity-mcp
Length of output: 2246
🏁 Script executed:
#!/bin/bash
echo "=== MCPServiceLocator full implementation ==="
cat -n MCPForUnity/Editor/Services/MCPServiceLocator.cs
echo ""
echo "=== TransportManager.GetClient full implementation ==="
rg -n "public IMcpTransportClient GetClient" --type=cs -A10Repository: CoplayDev/unity-mcp
Length of output: 7093
🏁 Script executed:
#!/bin/bash
echo "=== IBridgeControlService interface ==="
rg -n "interface IBridgeControlService" --type=cs -A20
echo ""
echo "=== TransportMode enum ==="
rg -n "enum TransportMode" --type=cs -A10
echo ""
echo "=== Checking if WebSocketTransportClient is actually used ==="
rg -n "WebSocketTransportClient|WebSocket" --type=cs | head -20Repository: CoplayDev/unity-mcp
Length of output: 5587
Use the active transport mode instead of hardcoding Http.
ReregisterToolsAsync always targets TransportMode.Http (line 256), but this ignores the user's configured transport preference. If the user has configured Stdio as the active transport (via BridgeControlService.ResolvePreferredMode), the Http client may not be initialized, causing GetClient to return null and reregistration to silently fail.
Use MCPServiceLocator.Bridge.ActiveMode to respect the user's configuration:
Proposed fix
private void ReregisterToolsAsync()
{
// Fire and forget - don't block UI
ThreadPool.QueueUserWorkItem(_ =>
{
try
{
var transportManager = MCPServiceLocator.TransportManager;
- var client = transportManager.GetClient(TransportMode.Http);
+ var mode = MCPServiceLocator.Bridge.ActiveMode;
+ if (!mode.HasValue)
+ return;
+ var client = transportManager.GetClient(mode.Value);
if (client != null && client.IsConnected)
{
client.ReregisterToolsAsync().Wait();
}
}
catch (Exception ex)
{
McpLog.Warn($"Failed to reregister tools: {ex.Message}");
}
});
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| 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 ReregisterToolsAsync() | |
| { | |
| // Fire and forget - don't block UI | |
| ThreadPool.QueueUserWorkItem(_ => | |
| { | |
| try | |
| { | |
| var transportManager = MCPServiceLocator.TransportManager; | |
| var mode = MCPServiceLocator.Bridge.ActiveMode; | |
| if (!mode.HasValue) | |
| return; | |
| var client = transportManager.GetClient(mode.Value); | |
| if (client != null && client.IsConnected) | |
| { | |
| client.ReregisterToolsAsync().Wait(); | |
| } | |
| } | |
| catch (Exception ex) | |
| { | |
| McpLog.Warn($"Failed to reregister tools: {ex.Message}"); | |
| } | |
| }); | |
| } |
🤖 Prompt for AI Agents
In `@MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs` around lines
248 - 267, ReregisterToolsAsync currently hardcodes TransportMode.Http when
calling MCPServiceLocator.TransportManager.GetClient, which ignores the user's
configured transport; change it to use the active transport mode from
MCPServiceLocator.Bridge.ActiveMode when obtaining the client (keep the existing
null and IsConnected checks and the try/catch around
client.ReregisterToolsAsync().Wait()). Update references in the
ReregisterToolsAsync method to call GetClient with
MCPServiceLocator.Bridge.ActiveMode so reregistration respects the user's
preferred transport.
| try: | ||
| sessions_data = await PluginHub.get_sessions(user_id=user_id) | ||
| sessions = sessions_data.sessions if sessions_data else {} | ||
| except Exception: | ||
| return None |
There was a problem hiding this comment.
Silent exception swallowing hides bugs.
The bare except Exception: return None at line 338 will silently swallow programming errors (e.g., TypeError, KeyError from unexpected data shapes). At minimum, log the exception at debug level so it's diagnosable.
Proposed fix
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: %s",
+ exc,
+ exc_info=True,
+ )
return None🧰 Tools
🪛 Ruff (0.15.0)
[warning] 338-338: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
In `@Server/src/transport/unity_instance_middleware.py` around lines 335 - 339,
The try/except around PluginHub.get_sessions currently swallows all exceptions;
change it to "except Exception as e:" and log the exception at debug (or
debug/exception) level before returning None so programming errors are visible;
specifically update the block around PluginHub.get_sessions/user_id to capture
the exception into a variable (e) and call the module's logger (e.g.,
logger.debug or logger.exception) with a clear message including e and any
contextual info (user_id, sessions_data) prior to returning None.
| 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) |
There was a problem hiding this comment.
try/except/continue silently discards errors — add logging.
Ruff S112 flags this pattern. If get_tools_for_project fails for one hash due to a bug (not a transient network error), this silently skips it, making issues very hard to diagnose.
Proposed fix
for project_hash in project_hashes:
try:
registered_tools = await PluginHub.get_tools_for_project(project_hash)
- except Exception:
+ except Exception as exc:
+ logger.debug(
+ "Failed to get tools for project hash %s: %s",
+ project_hash,
+ exc,
+ exc_info=True,
+ )
continue🧰 Tools
🪛 Ruff (0.15.0)
[error] 365-366: try-except-continue detected, consider logging the exception
(S112)
[warning] 365-365: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
In `@Server/src/transport/unity_instance_middleware.py` around lines 362 - 371,
The current loop over project_hashes swallows all exceptions from
PluginHub.get_tools_for_project by using bare except and continue; change the
except to capture the exception as e (e.g., except Exception as e) and log the
failure (including project_hash and exception details/traceback) before
continuing so failures are discoverable — update the except block around
PluginHub.get_tools_for_project to call the module or instance logger (e.g.,
logger.exception or logger.error with exc_info=True) referencing project_hash
and the caught exception, then continue to the next hash.
|
Closing to recreate with proper template |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 10754e5819
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| enabled_tool_names: set[str] = set() | ||
| for project_hash in project_hashes: | ||
| try: | ||
| registered_tools = await PluginHub.get_tools_for_project(project_hash) |
There was a problem hiding this comment.
Scope tool lookup to the current user in hosted mode
In remote-hosted deployments, this lookup ignores user_id even though _resolve_enabled_tool_names_for_context already derived it, so filtering uses PluginHub.get_tools_for_project(project_hash) with only a hash. PluginRegistry stores hosted sessions under (user_id, project_hash), which means hash-only resolution can return no session and the enabled set stays empty; on_list_tools then falls back to returning the full unfiltered tool list for that user.
Useful? React with 👍 / 👎.
| if isinstance(tool_name, str) and tool_name: | ||
| enabled_tool_names.add(tool_name) | ||
|
|
||
| return enabled_tool_names or None |
There was a problem hiding this comment.
Keep filtering when Unity reports zero enabled tools
This converts an empty enabled-tool set into None, and on_list_tools interprets that as 'no data, skip filtering'. If a Unity instance intentionally disables all Unity-managed tools and re-registers an empty tool list, the middleware will incorrectly return the full unfiltered list instead of hiding disabled tools, so tool listing no longer matches Unity’s actual enabled state.
Useful? React with 👍 / 👎.
Fixes #709. Unity restart no longer resets tool states. Client tool lists now respect Unity's enabled tools.
Summary by Sourcery
Align Unity client tool enablement state with server-side tool registration and ensure tool preferences persist across sessions and service instances.
New Features:
Bug Fixes:
Enhancements:
Tests:
Summary by CodeRabbit