Fix WebSocket tool registration to run on main thread #442
Conversation
WalkthroughThe change introduces thread-safe tool access in WebSocketTransportClient by adding a new async helper method that marshals tool discovery operations to Unity's main thread via EditorApplication.delayCall, replacing the previous synchronous call in SendRegisterToolsAsync. Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~15 minutes
Possibly related issues
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs(3 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-09-03T16:00:55.839Z
Learnt from: dsarno
Repo: CoplayDev/unity-mcp PR: 0
File: :0-0
Timestamp: 2025-09-03T16:00:55.839Z
Learning: ComponentResolver in UnityMcpBridge/Editor/Tools/ManageGameObject.cs is a nested static class within ManageGameObject, not a sibling type. The `using static MCPForUnity.Editor.Tools.ManageGameObject;` import is required to access ComponentResolver methods directly without the outer class qualifier.
Applied to files:
MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs
🔇 Additional comments (2)
MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs (2)
15-15: LGTM!The
UnityEditornamespace is correctly added forEditorApplication.delayCallused in the new helper method.
446-448: Good addition of cancellation checks and async marshaling.The changes correctly fix the main-thread violation by awaiting the new helper method, and the cancellation checks before and after the async call improve responsiveness during connection teardown.
| private Task<List<ToolMetadata>> GetEnabledToolsOnMainThreadAsync() | ||
| { | ||
| var tcs = new TaskCompletionSource<List<ToolMetadata>>(TaskCreationOptions.RunContinuationsAsynchronously); | ||
|
|
||
| EditorApplication.delayCall += () => | ||
| { | ||
| try | ||
| { | ||
| var tools = _toolDiscoveryService?.GetEnabledTools() ?? new List<ToolMetadata>(); | ||
| tcs.TrySetResult(tools); | ||
| } | ||
| catch (Exception ex) | ||
| { | ||
| tcs.TrySetException(ex); | ||
| } | ||
| }; | ||
|
|
||
| return tcs.Task; | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Consider adding cancellation/timeout safeguards to prevent hanging.
The implementation correctly marshals tool discovery to Unity's main thread. However, if EditorApplication.delayCall callback never executes (e.g., during Unity shutdown or domain reload), the TaskCompletionSource would hang indefinitely, potentially blocking Dispose() which awaits the receive loop.
Consider registering a cancellation callback to handle scenarios where Unity stops executing delayCall callbacks:
-private Task<List<ToolMetadata>> GetEnabledToolsOnMainThreadAsync()
+private Task<List<ToolMetadata>> GetEnabledToolsOnMainThreadAsync(CancellationToken cancellationToken = default)
{
var tcs = new TaskCompletionSource<List<ToolMetadata>>(TaskCreationOptions.RunContinuationsAsynchronously);
+
+ // Register cancellation to prevent indefinite hanging
+ if (cancellationToken.CanBeCanceled)
+ {
+ cancellationToken.Register(() => tcs.TrySetCanceled(cancellationToken));
+ }
EditorApplication.delayCall += () =>
{
try
{
var tools = _toolDiscoveryService?.GetEnabledTools() ?? new List<ToolMetadata>();
tcs.TrySetResult(tools);
}
catch (Exception ex)
{
tcs.TrySetException(ex);
}
};
return tcs.Task;
}Then update the call site:
token.ThrowIfCancellationRequested();
-var tools = await GetEnabledToolsOnMainThreadAsync().ConfigureAwait(false);
+var tools = await GetEnabledToolsOnMainThreadAsync(token).ConfigureAwait(false);
token.ThrowIfCancellationRequested();Committable suggestion skipped: line range outside the PR's diff.
|
Thanks @Jordonh18 ! I added a guardrail and merged this as part of #443. |
|
ahh perfect 😊 |
3 participants
Issue
Unity logged
[WebSocket] Unexpected receive error: HasKey can only be called from the main threadduring WebSocket registration. The background receive loop invokedToolDiscoveryService.GetEnabledTools(), which touchesEditorPrefs.HasKey, a Unity Editor API that must run on the main thread.Root Cause
Tool discovery was executed on a background thread in the WebSocket registration flow, violating Unity’s main-thread requirement for Editor APIs.
Fix
GetEnabledToolsOnMainThreadAsync()to marshal tool discovery onto the Unity main thread viaEditorApplication.delayCall.SendRegisterToolsAsyncto await this main-thread helper with cancellation checks.Files Touched
MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.csNotes
Summary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings.