PHXH
/
unity-mcp
51
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix/websocket queue to main thread (#443) 

* feat: Implement async method to retrieve enabled tools on the main thread

* fix: cancelable main-thread tool discovery

* chore: dispose cancellation registration and dedupe usings

---------

Co-authored-by: Jordon Harrison <Jordon.Harrison@outlook.com>
main
dsarno 2025-12-08 08:16:09 -08:00 committed by GitHub
parent fd44ab3b5d
commit 0c8d2aac42
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 401 additions and 369 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

39
MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs
View File

@ -8,9 +8,11 @@ using System.Threading;
using System.Threading.Tasks;
using MCPForUnity.Editor.Config;
using MCPForUnity.Editor.Helpers;
using MCPForUnity.Editor.Services;
using MCPForUnity.Editor.Services.Transport;
using Newtonsoft.Json;
using Newtonsoft.Json.Linq;
using UnityEditor;
using UnityEngine;
namespace MCPForUnity.Editor.Services.Transport.Transports
@ -65,6 +67,39 @@ namespace MCPForUnity.Editor.Services.Transport.Transports
public string TransportName => TransportDisplayName;
public TransportState State => _state;
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;
}
public async Task<bool> StartAsync()
{
// Capture identity values on the main thread before any async context switching
@ -421,7 +456,9 @@ namespace MCPForUnity.Editor.Services.Transport.Transports
{
if (_toolDiscoveryService == null) return;
var tools = _toolDiscoveryService.GetEnabledTools();
token.ThrowIfCancellationRequested();
var tools = await GetEnabledToolsOnMainThreadAsync(token).ConfigureAwait(false);
token.ThrowIfCancellationRequested();
McpLog.Info($"[WebSocket] Preparing to register {tools.Count} tool(s) with the bridge.");
var toolsArray = new JArray();

5
MCPForUnity/Editor/Windows/MCPForUnityEditorWindow.cs
View File

@ -12,11 +12,6 @@ using UnityEditor.UIElements;
using UnityEngine;
using UnityEngine.UIElements;
using MCPForUnity.Editor.Constants;
using MCPForUnity.Editor.Helpers;
using MCPForUnity.Editor.Services;
using MCPForUnity.Editor.Windows.Components.Settings;
using MCPForUnity.Editor.Windows.Components.Connection;
using MCPForUnity.Editor.Windows.Components.ClientConfig;
using MCPForUnity.Editor.Windows.Components.Tools;
namespace MCPForUnity.Editor.Windows