diff --git a/MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs b/MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs index 0648193e7..61983c356 100644 --- a/MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs +++ b/MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs @@ -8,9 +8,11 @@ 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,26 @@ public WebSocketTransportClient(IToolDiscoveryService toolDiscoveryService = nul public string TransportName => TransportDisplayName; public TransportState State => _state; + private Task> GetEnabledToolsOnMainThreadAsync() + { + var tcs = new TaskCompletionSource>(TaskCreationOptions.RunContinuationsAsynchronously); + + EditorApplication.delayCall += () => + { + try + { + var tools = _toolDiscoveryService?.GetEnabledTools() ?? new List(); + tcs.TrySetResult(tools); + } + catch (Exception ex) + { + tcs.TrySetException(ex); + } + }; + + return tcs.Task; + } + public async Task StartAsync() { // Capture identity values on the main thread before any async context switching @@ -421,7 +443,9 @@ private async Task SendRegisterToolsAsync(CancellationToken token) { if (_toolDiscoveryService == null) return; - var tools = _toolDiscoveryService.GetEnabledTools(); + token.ThrowIfCancellationRequested(); + var tools = await GetEnabledToolsOnMainThreadAsync().ConfigureAwait(false); + token.ThrowIfCancellationRequested(); McpLog.Info($"[WebSocket] Preparing to register {tools.Count} tool(s) with the bridge."); var toolsArray = new JArray();