Fix/websocket queue to main thread #443
Conversation
WalkthroughA new async helper schedules enabled-tool retrieval on the Unity main thread via EditorApplication.delayCall and replaces the prior synchronous collection in SendRegisterToolsAsync; a separate editor window file had import/format cleanup only. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes
Possibly related PRs
Suggested reviewers
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: 0
🧹 Nitpick comments (1)
MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs (1)
70-99: Well-implemented main thread marshalling.The use of
EditorApplication.delayCallwithTaskCompletionSourceproperly handles Unity's main-thread requirement for tool discovery. The cancellation handling, race condition checks (IsCompleted), and exception handling are all correctly implemented.Consider moving
registration.Dispose()to afinallyblock to guarantee cleanup even ifGetEnabledTools()throws:EditorApplication.delayCall += () => { try { - // Clean up the registration once we're running - registration.Dispose(); - if (tcs.Task.IsCompleted) { return; } var tools = _toolDiscoveryService?.GetEnabledTools() ?? new List<ToolMetadata>(); tcs.TrySetResult(tools); } catch (Exception ex) { tcs.TrySetException(ex); + } + finally + { + registration.Dispose(); } };
📜 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
🧬 Code graph analysis (1)
MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs (2)
MCPForUnity/Editor/Services/IToolDiscoveryService.cs (4)
List(44-44)List(54-54)ToolMetadata(8-22)ToolMetadata(49-49)MCPForUnity/Editor/Services/ToolDiscoveryService.cs (5)
List(19-60)List(72-77)List(159-193)ToolMetadata(62-70)ToolMetadata(107-157)
🔇 Additional comments (2)
MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs (2)
457-459: LGTM! Proper async tool retrieval with cancellation checks.The replacement of synchronous tool retrieval with the new async method correctly ensures tool discovery runs on Unity's main thread. The cancellation checks before and after the operation provide appropriate safeguards.
15-15: LGTM!Required for
EditorApplication.delayCallused in the new marshalling method.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs (1)
70-101: Main-thread marshalling viadelayCallis solid; consider unsubscribing on cancellationThis TCS +
EditorApplication.delayCallpattern cleanly movesGetEnabledToolsonto the main thread, and the cancellation registration avoids hangs during shutdown. One optional refinement: also unsubscribe thedelayCallcallback when the token is cancelled before the callback runs, so you don’t keep a stale closure hanging off the editor event.An example refactor keeping behavior but explicitly unsubscribing:
- private Task<List<ToolMetadata>> GetEnabledToolsOnMainThreadAsync(CancellationToken token) - { - var tcs = new TaskCompletionSource<List<ToolMetadata>>(TaskCreationOptions.RunContinuationsAsynchronously); - - // Register cancellation to break the deadlock if StopAsync is called while waiting for main thread - var registration = token.Register(() => tcs.TrySetCanceled()); - - EditorApplication.delayCall += () => - { - try - { - if (tcs.Task.IsCompleted) - { - return; - } - - var tools = _toolDiscoveryService?.GetEnabledTools() ?? new List<ToolMetadata>(); - tcs.TrySetResult(tools); - } - catch (Exception ex) - { - tcs.TrySetException(ex); - } - finally - { - // Ensure registration is disposed even if discovery throws - registration.Dispose(); - } - }; - - return tcs.Task; - } + private Task<List<ToolMetadata>> GetEnabledToolsOnMainThreadAsync(CancellationToken token) + { + var tcs = new TaskCompletionSource<List<ToolMetadata>>(TaskCreationOptions.RunContinuationsAsynchronously); + + EditorApplication.CallbackFunction callback = null; + CancellationTokenRegistration registration = default; + + callback = () => + { + try + { + if (tcs.Task.IsCompleted) + { + return; + } + + var tools = _toolDiscoveryService?.GetEnabledTools() ?? new List<ToolMetadata>(); + tcs.TrySetResult(tools); + } + catch (Exception ex) + { + tcs.TrySetException(ex); + } + finally + { + // Ensure registration is disposed and the callback is unsubscribed even if discovery throws + registration.Dispose(); + EditorApplication.delayCall -= callback; + } + }; + + EditorApplication.delayCall += callback; + + // Register cancellation to break the deadlock if StopAsync is called while waiting for main thread + registration = token.Register(() => + { + EditorApplication.delayCall -= callback; + tcs.TrySetCanceled(); + }); + + return tcs.Task; + }This keeps the main-thread behavior and cancellation guarantees, while also cleaning up the event subscription eagerly on cancel.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs(3 hunks)MCPForUnity/Editor/Windows/MCPForUnityEditorWindow.cs(1 hunks)
✅ Files skipped from review due to trivial changes (1)
- MCPForUnity/Editor/Windows/MCPForUnityEditorWindow.cs
🧰 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)
11-16: Editor-only using directives look appropriateBringing in
MCPForUnity.Editor.ServicesandUnityEditorhere matches the use ofIToolDiscoveryServiceandEditorApplication.delayCallin this editor-only transport class; no issues from my side.Please just confirm this file is compiled into an editor-only assembly (path suggests it is) so that
UnityEditorreferences never end up in player builds.
455-463: Cancellation-aware tool registration flow looks goodThe added
ThrowIfCancellationRequestedcalls before and afterGetEnabledToolsOnMainThreadAsyncensure you don’t proceed with tool registration once the connection is being torn down, and theConfigureAwait(false)keeps the heavy JSON work off the main thread. No further changes suggested.
2 participants
Summary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings.