Skip to content

Commit 0c8d2aa

Browse files
dsarnoJordonh18
andauthored
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>
1 parent fd44ab3 commit 0c8d2aa

File tree

2 files changed

+401
-369
lines changed

2 files changed

+401
-369
lines changed

MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs

Lines changed: 38 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,11 @@
88
using System.Threading.Tasks;
99
using MCPForUnity.Editor.Config;
1010
using MCPForUnity.Editor.Helpers;
11+
using MCPForUnity.Editor.Services;
1112
using MCPForUnity.Editor.Services.Transport;
1213
using Newtonsoft.Json;
1314
using Newtonsoft.Json.Linq;
15+
using UnityEditor;
1416
using UnityEngine;
1517

1618
namespace MCPForUnity.Editor.Services.Transport.Transports
@@ -65,6 +67,39 @@ public WebSocketTransportClient(IToolDiscoveryService toolDiscoveryService = nul
6567
public string TransportName => TransportDisplayName;
6668
public TransportState State => _state;
6769

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+
68103
public async Task<bool> StartAsync()
69104
{
70105
// Capture identity values on the main thread before any async context switching
@@ -421,7 +456,9 @@ private async Task SendRegisterToolsAsync(CancellationToken token)
421456
{
422457
if (_toolDiscoveryService == null) return;
423458

424-
var tools = _toolDiscoveryService.GetEnabledTools();
459+
token.ThrowIfCancellationRequested();
460+
var tools = await GetEnabledToolsOnMainThreadAsync(token).ConfigureAwait(false);
461+
token.ThrowIfCancellationRequested();
425462
McpLog.Info($"[WebSocket] Preparing to register {tools.Count} tool(s) with the bridge.");
426463
var toolsArray = new JArray();
427464

0 commit comments

Comments
 (0)