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

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

@ -1,368 +1,363 @@
using System; using System;
using System.Collections.Generic; using System.Collections.Generic;
using System.Linq; using System.Linq;
using System.Threading.Tasks; using System.Threading.Tasks;
using MCPForUnity.Editor.Helpers; using MCPForUnity.Editor.Helpers;
using MCPForUnity.Editor.Services; using MCPForUnity.Editor.Services;
using MCPForUnity.Editor.Windows.Components.ClientConfig; using MCPForUnity.Editor.Windows.Components.ClientConfig;
using MCPForUnity.Editor.Windows.Components.Connection; using MCPForUnity.Editor.Windows.Components.Connection;
using MCPForUnity.Editor.Windows.Components.Settings; using MCPForUnity.Editor.Windows.Components.Settings;
using UnityEditor; using UnityEditor;
using UnityEditor.UIElements; using UnityEditor.UIElements;
using UnityEngine; using UnityEngine;
using UnityEngine.UIElements; using UnityEngine.UIElements;
using MCPForUnity.Editor.Constants; using MCPForUnity.Editor.Constants;
using MCPForUnity.Editor.Helpers; using MCPForUnity.Editor.Windows.Components.Tools;
using MCPForUnity.Editor.Services;
using MCPForUnity.Editor.Windows.Components.Settings; namespace MCPForUnity.Editor.Windows
using MCPForUnity.Editor.Windows.Components.Connection; {
using MCPForUnity.Editor.Windows.Components.ClientConfig; public class MCPForUnityEditorWindow : EditorWindow
using MCPForUnity.Editor.Windows.Components.Tools; {
// Section controllers
namespace MCPForUnity.Editor.Windows private McpSettingsSection settingsSection;
{ private McpConnectionSection connectionSection;
public class MCPForUnityEditorWindow : EditorWindow private McpClientConfigSection clientConfigSection;
{ private McpToolsSection toolsSection;
// Section controllers
private McpSettingsSection settingsSection; private ToolbarToggle settingsTabToggle;
private McpConnectionSection connectionSection; private ToolbarToggle toolsTabToggle;
private McpClientConfigSection clientConfigSection; private VisualElement settingsPanel;
private McpToolsSection toolsSection; private VisualElement toolsPanel;
private ToolbarToggle settingsTabToggle; private static readonly HashSet<MCPForUnityEditorWindow> OpenWindows = new();
private ToolbarToggle toolsTabToggle; private bool guiCreated = false;
private VisualElement settingsPanel; private double lastRefreshTime = 0;
private VisualElement toolsPanel; private const double RefreshDebounceSeconds = 0.5;
private static readonly HashSet<MCPForUnityEditorWindow> OpenWindows = new(); private enum ActivePanel
private bool guiCreated = false; {
private double lastRefreshTime = 0; Settings,
private const double RefreshDebounceSeconds = 0.5; Tools
}
private enum ActivePanel
{ internal static void CloseAllWindows()
Settings, {
Tools var windows = OpenWindows.Where(window => window != null).ToArray();
} foreach (var window in windows)
{
internal static void CloseAllWindows() window.Close();
{ }
var windows = OpenWindows.Where(window => window != null).ToArray(); }
foreach (var window in windows)
{ public static void ShowWindow()
window.Close(); {
} var window = GetWindow<MCPForUnityEditorWindow>("MCP For Unity");
} window.minSize = new Vector2(500, 600);
}
public static void ShowWindow()
{ // Helper to check and manage open windows from other classes
var window = GetWindow<MCPForUnityEditorWindow>("MCP For Unity"); public static bool HasAnyOpenWindow()
window.minSize = new Vector2(500, 600); {
} return OpenWindows.Count > 0;
}
// Helper to check and manage open windows from other classes
public static bool HasAnyOpenWindow() public static void CloseAllOpenWindows()
{ {
return OpenWindows.Count > 0; if (OpenWindows.Count == 0)
} return;
public static void CloseAllOpenWindows() // Copy to array to avoid modifying the collection while iterating
{ var arr = new MCPForUnityEditorWindow[OpenWindows.Count];
if (OpenWindows.Count == 0) OpenWindows.CopyTo(arr);
return; foreach (var window in arr)
{
// Copy to array to avoid modifying the collection while iterating try
var arr = new MCPForUnityEditorWindow[OpenWindows.Count]; {
OpenWindows.CopyTo(arr); window?.Close();
foreach (var window in arr) }
{ catch (Exception ex)
try {
{ McpLog.Warn($"Error closing MCP window: {ex.Message}");
window?.Close(); }
} }
catch (Exception ex) }
{
McpLog.Warn($"Error closing MCP window: {ex.Message}"); public void CreateGUI()
} {
} // Guard against repeated CreateGUI calls (e.g., domain reloads)
} if (guiCreated)
return;
public void CreateGUI()
{ string basePath = AssetPathUtility.GetMcpPackageRootPath();
// Guard against repeated CreateGUI calls (e.g., domain reloads)
if (guiCreated) // Load main window UXML
return; var visualTree = AssetDatabase.LoadAssetAtPath<VisualTreeAsset>(
$"{basePath}/Editor/Windows/MCPForUnityEditorWindow.uxml"
string basePath = AssetPathUtility.GetMcpPackageRootPath(); );
// Load main window UXML if (visualTree == null)
var visualTree = AssetDatabase.LoadAssetAtPath<VisualTreeAsset>( {
$"{basePath}/Editor/Windows/MCPForUnityEditorWindow.uxml" McpLog.Error(
); $"Failed to load UXML at: {basePath}/Editor/Windows/MCPForUnityEditorWindow.uxml"
);
if (visualTree == null) return;
{ }
McpLog.Error(
$"Failed to load UXML at: {basePath}/Editor/Windows/MCPForUnityEditorWindow.uxml" visualTree.CloneTree(rootVisualElement);
);
return; // Load main window USS
} var mainStyleSheet = AssetDatabase.LoadAssetAtPath<StyleSheet>(
$"{basePath}/Editor/Windows/MCPForUnityEditorWindow.uss"
visualTree.CloneTree(rootVisualElement); );
if (mainStyleSheet != null)
// Load main window USS {
var mainStyleSheet = AssetDatabase.LoadAssetAtPath<StyleSheet>( rootVisualElement.styleSheets.Add(mainStyleSheet);
$"{basePath}/Editor/Windows/MCPForUnityEditorWindow.uss" }
);
if (mainStyleSheet != null) // Load common USS
{ var commonStyleSheet = AssetDatabase.LoadAssetAtPath<StyleSheet>(
rootVisualElement.styleSheets.Add(mainStyleSheet); $"{basePath}/Editor/Windows/Components/Common.uss"
} );
if (commonStyleSheet != null)
// Load common USS {
var commonStyleSheet = AssetDatabase.LoadAssetAtPath<StyleSheet>( rootVisualElement.styleSheets.Add(commonStyleSheet);
$"{basePath}/Editor/Windows/Components/Common.uss" }
);
if (commonStyleSheet != null) settingsPanel = rootVisualElement.Q<VisualElement>("settings-panel");
{ toolsPanel = rootVisualElement.Q<VisualElement>("tools-panel");
rootVisualElement.styleSheets.Add(commonStyleSheet); var settingsContainer = rootVisualElement.Q<VisualElement>("settings-container");
} var toolsContainer = rootVisualElement.Q<VisualElement>("tools-container");
settingsPanel = rootVisualElement.Q<VisualElement>("settings-panel"); if (settingsPanel == null || toolsPanel == null)
toolsPanel = rootVisualElement.Q<VisualElement>("tools-panel"); {
var settingsContainer = rootVisualElement.Q<VisualElement>("settings-container"); McpLog.Error("Failed to find tab panels in UXML");
var toolsContainer = rootVisualElement.Q<VisualElement>("tools-container"); return;
}
if (settingsPanel == null || toolsPanel == null)
{ if (settingsContainer == null)
McpLog.Error("Failed to find tab panels in UXML"); {
return; McpLog.Error("Failed to find settings-container in UXML");
} return;
}
if (settingsContainer == null)
{ if (toolsContainer == null)
McpLog.Error("Failed to find settings-container in UXML"); {
return; McpLog.Error("Failed to find tools-container in UXML");
} return;
}
if (toolsContainer == null)
{ SetupTabs();
McpLog.Error("Failed to find tools-container in UXML");
return; // Load and initialize Settings section
} var settingsTree = AssetDatabase.LoadAssetAtPath<VisualTreeAsset>(
$"{basePath}/Editor/Windows/Components/Settings/McpSettingsSection.uxml"
SetupTabs(); );
if (settingsTree != null)
// Load and initialize Settings section {
var settingsTree = AssetDatabase.LoadAssetAtPath<VisualTreeAsset>( var settingsRoot = settingsTree.Instantiate();
$"{basePath}/Editor/Windows/Components/Settings/McpSettingsSection.uxml" settingsContainer.Add(settingsRoot);
); settingsSection = new McpSettingsSection(settingsRoot);
if (settingsTree != null) settingsSection.OnGitUrlChanged += () =>
{ clientConfigSection?.UpdateManualConfiguration();
var settingsRoot = settingsTree.Instantiate(); settingsSection.OnHttpServerCommandUpdateRequested += () =>
settingsContainer.Add(settingsRoot); connectionSection?.UpdateHttpServerCommandDisplay();
settingsSection = new McpSettingsSection(settingsRoot); }
settingsSection.OnGitUrlChanged += () =>
clientConfigSection?.UpdateManualConfiguration(); // Load and initialize Connection section
settingsSection.OnHttpServerCommandUpdateRequested += () => var connectionTree = AssetDatabase.LoadAssetAtPath<VisualTreeAsset>(
connectionSection?.UpdateHttpServerCommandDisplay(); $"{basePath}/Editor/Windows/Components/Connection/McpConnectionSection.uxml"
} );
if (connectionTree != null)
// Load and initialize Connection section {
var connectionTree = AssetDatabase.LoadAssetAtPath<VisualTreeAsset>( var connectionRoot = connectionTree.Instantiate();
$"{basePath}/Editor/Windows/Components/Connection/McpConnectionSection.uxml" settingsContainer.Add(connectionRoot);
); connectionSection = new McpConnectionSection(connectionRoot);
if (connectionTree != null) connectionSection.OnManualConfigUpdateRequested += () =>
{ clientConfigSection?.UpdateManualConfiguration();
var connectionRoot = connectionTree.Instantiate(); }
settingsContainer.Add(connectionRoot);
connectionSection = new McpConnectionSection(connectionRoot); // Load and initialize Client Configuration section
connectionSection.OnManualConfigUpdateRequested += () => var clientConfigTree = AssetDatabase.LoadAssetAtPath<VisualTreeAsset>(
clientConfigSection?.UpdateManualConfiguration(); $"{basePath}/Editor/Windows/Components/ClientConfig/McpClientConfigSection.uxml"
} );
if (clientConfigTree != null)
// Load and initialize Client Configuration section {
var clientConfigTree = AssetDatabase.LoadAssetAtPath<VisualTreeAsset>( var clientConfigRoot = clientConfigTree.Instantiate();
$"{basePath}/Editor/Windows/Components/ClientConfig/McpClientConfigSection.uxml" settingsContainer.Add(clientConfigRoot);
); clientConfigSection = new McpClientConfigSection(clientConfigRoot);
if (clientConfigTree != null) }
{
var clientConfigRoot = clientConfigTree.Instantiate(); // Load and initialize Tools section
settingsContainer.Add(clientConfigRoot); var toolsTree = AssetDatabase.LoadAssetAtPath<VisualTreeAsset>(
clientConfigSection = new McpClientConfigSection(clientConfigRoot); $"{basePath}/Editor/Windows/Components/Tools/McpToolsSection.uxml"
} );
if (toolsTree != null)
// Load and initialize Tools section {
var toolsTree = AssetDatabase.LoadAssetAtPath<VisualTreeAsset>( var toolsRoot = toolsTree.Instantiate();
$"{basePath}/Editor/Windows/Components/Tools/McpToolsSection.uxml" toolsContainer.Add(toolsRoot);
); toolsSection = new McpToolsSection(toolsRoot);
if (toolsTree != null) toolsSection.Refresh();
{ }
var toolsRoot = toolsTree.Instantiate(); else
toolsContainer.Add(toolsRoot); {
toolsSection = new McpToolsSection(toolsRoot); McpLog.Warn("Failed to load tools section UXML. Tool configuration will be unavailable.");
toolsSection.Refresh(); }
} guiCreated = true;
else
{ // Initial updates
McpLog.Warn("Failed to load tools section UXML. Tool configuration will be unavailable."); RefreshAllData();
} }
guiCreated = true;
private void OnEnable()
// Initial updates {
RefreshAllData(); EditorApplication.update += OnEditorUpdate;
} OpenWindows.Add(this);
}
private void OnEnable()
{ private void OnDisable()
EditorApplication.update += OnEditorUpdate; {
OpenWindows.Add(this); EditorApplication.update -= OnEditorUpdate;
} OpenWindows.Remove(this);
guiCreated = false;
private void OnDisable() }
{
EditorApplication.update -= OnEditorUpdate; private void OnFocus()
OpenWindows.Remove(this); {
guiCreated = false; // Only refresh data if UI is built
} if (rootVisualElement == null || rootVisualElement.childCount == 0)
return;
private void OnFocus()
{ RefreshAllData();
// Only refresh data if UI is built }
if (rootVisualElement == null || rootVisualElement.childCount == 0)
return; private void OnEditorUpdate()
{
RefreshAllData(); if (rootVisualElement == null || rootVisualElement.childCount == 0)
} return;
private void OnEditorUpdate() connectionSection?.UpdateConnectionStatus();
{ }
if (rootVisualElement == null || rootVisualElement.childCount == 0)
return; private void RefreshAllData()
{
connectionSection?.UpdateConnectionStatus(); // Debounce rapid successive calls (e.g., from OnFocus being called multiple times)
} double currentTime = EditorApplication.timeSinceStartup;
if (currentTime - lastRefreshTime < RefreshDebounceSeconds)
private void RefreshAllData() {
{ return;
// Debounce rapid successive calls (e.g., from OnFocus being called multiple times) }
double currentTime = EditorApplication.timeSinceStartup; lastRefreshTime = currentTime;
if (currentTime - lastRefreshTime < RefreshDebounceSeconds)
{ connectionSection?.UpdateConnectionStatus();
return;
} if (MCPServiceLocator.Bridge.IsRunning)
lastRefreshTime = currentTime; {
_ = connectionSection?.VerifyBridgeConnectionAsync();
connectionSection?.UpdateConnectionStatus(); }
if (MCPServiceLocator.Bridge.IsRunning) settingsSection?.UpdatePathOverrides();
{ clientConfigSection?.RefreshSelectedClient();
_ = connectionSection?.VerifyBridgeConnectionAsync(); }
}
private void SetupTabs()
settingsSection?.UpdatePathOverrides(); {
clientConfigSection?.RefreshSelectedClient(); settingsTabToggle = rootVisualElement.Q<ToolbarToggle>("settings-tab");
} toolsTabToggle = rootVisualElement.Q<ToolbarToggle>("tools-tab");
private void SetupTabs() settingsPanel?.RemoveFromClassList("hidden");
{ toolsPanel?.RemoveFromClassList("hidden");
settingsTabToggle = rootVisualElement.Q<ToolbarToggle>("settings-tab");
toolsTabToggle = rootVisualElement.Q<ToolbarToggle>("tools-tab"); if (settingsTabToggle != null)
{
settingsPanel?.RemoveFromClassList("hidden"); settingsTabToggle.RegisterValueChangedCallback(evt =>
toolsPanel?.RemoveFromClassList("hidden"); {
if (!evt.newValue)
if (settingsTabToggle != null) {
{ if (toolsTabToggle != null && !toolsTabToggle.value)
settingsTabToggle.RegisterValueChangedCallback(evt => {
{ settingsTabToggle.SetValueWithoutNotify(true);
if (!evt.newValue) }
{ return;
if (toolsTabToggle != null && !toolsTabToggle.value) }
{
settingsTabToggle.SetValueWithoutNotify(true); SwitchPanel(ActivePanel.Settings);
} });
return; }
}
if (toolsTabToggle != null)
SwitchPanel(ActivePanel.Settings); {
}); toolsTabToggle.RegisterValueChangedCallback(evt =>
} {
if (!evt.newValue)
if (toolsTabToggle != null) {
{ if (settingsTabToggle != null && !settingsTabToggle.value)
toolsTabToggle.RegisterValueChangedCallback(evt => {
{ toolsTabToggle.SetValueWithoutNotify(true);
if (!evt.newValue) }
{ return;
if (settingsTabToggle != null && !settingsTabToggle.value) }
{
toolsTabToggle.SetValueWithoutNotify(true); SwitchPanel(ActivePanel.Tools);
} });
return; }
}
var savedPanel = EditorPrefs.GetString(EditorPrefKeys.EditorWindowActivePanel, ActivePanel.Settings.ToString());
SwitchPanel(ActivePanel.Tools); if (!Enum.TryParse(savedPanel, out ActivePanel initialPanel))
}); {
} initialPanel = ActivePanel.Settings;
}
var savedPanel = EditorPrefs.GetString(EditorPrefKeys.EditorWindowActivePanel, ActivePanel.Settings.ToString());
if (!Enum.TryParse(savedPanel, out ActivePanel initialPanel)) SwitchPanel(initialPanel);
{ }
initialPanel = ActivePanel.Settings;
} private void SwitchPanel(ActivePanel panel)
{
SwitchPanel(initialPanel); bool showSettings = panel == ActivePanel.Settings;
}
if (settingsPanel != null)
private void SwitchPanel(ActivePanel panel) {
{ settingsPanel.style.display = showSettings ? DisplayStyle.Flex : DisplayStyle.None;
bool showSettings = panel == ActivePanel.Settings; }
if (settingsPanel != null) if (toolsPanel != null)
{ {
settingsPanel.style.display = showSettings ? DisplayStyle.Flex : DisplayStyle.None; toolsPanel.style.display = showSettings ? DisplayStyle.None : DisplayStyle.Flex;
} }
if (toolsPanel != null) settingsTabToggle?.SetValueWithoutNotify(showSettings);
{ toolsTabToggle?.SetValueWithoutNotify(!showSettings);
toolsPanel.style.display = showSettings ? DisplayStyle.None : DisplayStyle.Flex;
} EditorPrefs.SetString(EditorPrefKeys.EditorWindowActivePanel, panel.ToString());
}
settingsTabToggle?.SetValueWithoutNotify(showSettings);
toolsTabToggle?.SetValueWithoutNotify(!showSettings); internal static void RequestHealthVerification()
{
EditorPrefs.SetString(EditorPrefKeys.EditorWindowActivePanel, panel.ToString()); foreach (var window in OpenWindows)
} {
window?.ScheduleHealthCheck();
internal static void RequestHealthVerification() }
{ }
foreach (var window in OpenWindows)
{ private void ScheduleHealthCheck()
window?.ScheduleHealthCheck(); {
} EditorApplication.delayCall += async () =>
} {
// Ensure window and components are still valid before execution
private void ScheduleHealthCheck() if (this == null || connectionSection == null)
{ {
EditorApplication.delayCall += async () => return;
{ }
// Ensure window and components are still valid before execution
if (this == null || connectionSection == null) try
{ {
return; await connectionSection.VerifyBridgeConnectionAsync();
} }
catch (Exception ex)
try {
{ // Log but don't crash if verification fails during cleanup
await connectionSection.VerifyBridgeConnectionAsync(); McpLog.Warn($"Health check verification failed: {ex.Message}");
} }
catch (Exception ex) };
{ }
// Log but don't crash if verification fails during cleanup }
McpLog.Warn($"Health check verification failed: {ex.Message}"); }
}
};
}
}
}