|
8 | 8 | using System.Threading.Tasks; |
9 | 9 | using MCPForUnity.Editor.Config; |
10 | 10 | using MCPForUnity.Editor.Helpers; |
| 11 | +using MCPForUnity.Editor.Services; |
11 | 12 | using MCPForUnity.Editor.Services.Transport; |
12 | 13 | using Newtonsoft.Json; |
13 | 14 | using Newtonsoft.Json.Linq; |
| 15 | +using UnityEditor; |
14 | 16 | using UnityEngine; |
15 | 17 |
|
16 | 18 | namespace MCPForUnity.Editor.Services.Transport.Transports |
@@ -65,6 +67,39 @@ public WebSocketTransportClient(IToolDiscoveryService toolDiscoveryService = nul |
65 | 67 | public string TransportName => TransportDisplayName; |
66 | 68 | public TransportState State => _state; |
67 | 69 |
|
| 70 | + private Task<List<ToolMetadata>> GetEnabledToolsOnMainThreadAsync(CancellationToken token) |
| 71 | + { |
| 72 | + var tcs = new TaskCompletionSource<List<ToolMetadata>>(TaskCreationOptions.RunContinuationsAsynchronously); |
| 73 | + |
| 74 | + // Register cancellation to break the deadlock if StopAsync is called while waiting for main thread |
| 75 | + var registration = token.Register(() => tcs.TrySetCanceled()); |
| 76 | + |
| 77 | + EditorApplication.delayCall += () => |
| 78 | + { |
| 79 | + try |
| 80 | + { |
| 81 | + if (tcs.Task.IsCompleted) |
| 82 | + { |
| 83 | + return; |
| 84 | + } |
| 85 | + |
| 86 | + var tools = _toolDiscoveryService?.GetEnabledTools() ?? new List<ToolMetadata>(); |
| 87 | + tcs.TrySetResult(tools); |
| 88 | + } |
| 89 | + catch (Exception ex) |
| 90 | + { |
| 91 | + tcs.TrySetException(ex); |
| 92 | + } |
| 93 | + finally |
| 94 | + { |
| 95 | + // Ensure registration is disposed even if discovery throws |
| 96 | + registration.Dispose(); |
| 97 | + } |
| 98 | + }; |
| 99 | + |
| 100 | + return tcs.Task; |
| 101 | + } |
| 102 | + |
68 | 103 | public async Task<bool> StartAsync() |
69 | 104 | { |
70 | 105 | // Capture identity values on the main thread before any async context switching |
@@ -421,7 +456,9 @@ private async Task SendRegisterToolsAsync(CancellationToken token) |
421 | 456 | { |
422 | 457 | if (_toolDiscoveryService == null) return; |
423 | 458 |
|
424 | | - var tools = _toolDiscoveryService.GetEnabledTools(); |
| 459 | + token.ThrowIfCancellationRequested(); |
| 460 | + var tools = await GetEnabledToolsOnMainThreadAsync(token).ConfigureAwait(false); |
| 461 | + token.ThrowIfCancellationRequested(); |
425 | 462 | McpLog.Info($"[WebSocket] Preparing to register {tools.Count} tool(s) with the bridge."); |
426 | 463 | var toolsArray = new JArray(); |
427 | 464 |
|
|
0 commit comments