From 17cd543fab766632a864dc1a6f2faedba7ed07bd Mon Sep 17 00:00:00 2001 From: Marcus Sanatan Date: Thu, 27 Nov 2025 19:33:26 -0400 Subject: [PATCH] Fix stdio reloads (#402) * First pass at MCP client refactor * Restore original text instructions Well most of them, I modified a few * Move configurators to their own folder It's less clusterd * Remvoe override for Windsurf because we no longer need to use it * Add Antigravity configs Works like Windsurf, but it sucks ass * Add some docs for properties * Add comprehensive MCP client configurators documentation * Add missing imports (#7) * Handle Linux paths when unregistering CLI commands * Construct a JSON error in a much more secure fashion * Fix stdio auto-reconnect after domain reloads We mirror what we've done with the HTTP/websocket connection We also ensure the states from the stdio/HTTP connections are handled separately. Things now work as expected * Fix ActiveMode to return resolved transport mode instead of preferred mode The ActiveMode property now calls ResolvePreferredMode() to return the actual active transport mode rather than just the preferred mode setting. * Minor improvements for stdio bridge - Consolidated the !useHttp && isRunning checks into a single shouldResume flag. - Wrapped the fire-and-forget StopAsync in a continuation that logs faults (matching the HTTP handler pattern). - Wrapped StartAsync in a continuation that logs failures and only triggers the health check on success. * Refactor TransportManager to use switch expressions and improve error handling - Replace if-else chains with switch expressions for better readability and exhaustiveness checking - Add GetClient() helper method to centralize client retrieval logic - Wrap StopAsync in try-catch to log failures when stopping a failed transport - Use client.TransportName instead of mode.ToString() for consistent naming in error messages --- .../Editor/Constants/EditorPrefKeys.cs | 1 + .../Editor/Services/BridgeControlService.cs | 29 ++-- .../Services/HttpBridgeReloadHandler.cs | 16 ++- .../Services/StdioBridgeReloadHandler.cs | 104 ++++++++++++++ .../Services/StdioBridgeReloadHandler.cs.meta | 11 ++ .../Services/Transport/TransportManager.cs | 133 ++++++++++++------ .../Transport/Transports/StdioBridgeHost.cs | 50 ------- 7 files changed, 235 insertions(+), 109 deletions(-) create mode 100644 MCPForUnity/Editor/Services/StdioBridgeReloadHandler.cs create mode 100644 MCPForUnity/Editor/Services/StdioBridgeReloadHandler.cs.meta diff --git a/MCPForUnity/Editor/Constants/EditorPrefKeys.cs b/MCPForUnity/Editor/Constants/EditorPrefKeys.cs index ffaa31c..30dcd2b 100644 --- a/MCPForUnity/Editor/Constants/EditorPrefKeys.cs +++ b/MCPForUnity/Editor/Constants/EditorPrefKeys.cs @@ -11,6 +11,7 @@ namespace MCPForUnity.Editor.Constants internal const string ValidationLevel = "MCPForUnity.ValidationLevel"; internal const string UnitySocketPort = "MCPForUnity.UnitySocketPort"; internal const string ResumeHttpAfterReload = "MCPForUnity.ResumeHttpAfterReload"; + internal const string ResumeStdioAfterReload = "MCPForUnity.ResumeStdioAfterReload"; internal const string UvxPathOverride = "MCPForUnity.UvxPath"; internal const string ClaudeCliPathOverride = "MCPForUnity.ClaudeCliPath"; diff --git a/MCPForUnity/Editor/Services/BridgeControlService.cs b/MCPForUnity/Editor/Services/BridgeControlService.cs index c67efd1..4c18da0 100644 --- a/MCPForUnity/Editor/Services/BridgeControlService.cs +++ b/MCPForUnity/Editor/Services/BridgeControlService.cs @@ -49,13 +49,21 @@ namespace MCPForUnity.Editor.Services }; } - public bool IsRunning => _transportManager.GetState().IsConnected; + public bool IsRunning + { + get + { + var mode = ResolvePreferredMode(); + return _transportManager.IsRunning(mode); + } + } public int CurrentPort { get { - var state = _transportManager.GetState(); + var mode = ResolvePreferredMode(); + var state = _transportManager.GetState(mode); if (state.Port.HasValue) { return state.Port.Value; @@ -67,7 +75,7 @@ namespace MCPForUnity.Editor.Services } public bool IsAutoConnectMode => StdioBridgeHost.IsAutoConnectMode(); - public TransportMode? ActiveMode => _transportManager.ActiveMode; + public TransportMode? ActiveMode => ResolvePreferredMode(); public async Task StartAsync() { @@ -92,7 +100,8 @@ namespace MCPForUnity.Editor.Services { try { - await _transportManager.StopAsync(); + var mode = ResolvePreferredMode(); + await _transportManager.StopAsync(mode); } catch (Exception ex) { @@ -102,17 +111,17 @@ namespace MCPForUnity.Editor.Services public async Task VerifyAsync() { - var mode = _transportManager.ActiveMode ?? ResolvePreferredMode(); - bool pingSucceeded = await _transportManager.VerifyAsync(); - var state = _transportManager.GetState(); + var mode = ResolvePreferredMode(); + bool pingSucceeded = await _transportManager.VerifyAsync(mode); + var state = _transportManager.GetState(mode); return BuildVerificationResult(state, mode, pingSucceeded); } public BridgeVerificationResult Verify(int port) { - var mode = _transportManager.ActiveMode ?? ResolvePreferredMode(); - bool pingSucceeded = _transportManager.VerifyAsync().GetAwaiter().GetResult(); - var state = _transportManager.GetState(); + var mode = ResolvePreferredMode(); + bool pingSucceeded = _transportManager.VerifyAsync(mode).GetAwaiter().GetResult(); + var state = _transportManager.GetState(mode); if (mode == TransportMode.Stdio) { diff --git a/MCPForUnity/Editor/Services/HttpBridgeReloadHandler.cs b/MCPForUnity/Editor/Services/HttpBridgeReloadHandler.cs index 16b8bd8..0422a92 100644 --- a/MCPForUnity/Editor/Services/HttpBridgeReloadHandler.cs +++ b/MCPForUnity/Editor/Services/HttpBridgeReloadHandler.cs @@ -24,8 +24,8 @@ namespace MCPForUnity.Editor.Services { try { - var bridge = MCPServiceLocator.Bridge; - bool shouldResume = bridge.IsRunning && bridge.ActiveMode == TransportMode.Http; + var transport = MCPServiceLocator.TransportManager; + bool shouldResume = transport.IsRunning(TransportMode.Http); if (shouldResume) { @@ -36,9 +36,9 @@ namespace MCPForUnity.Editor.Services EditorPrefs.DeleteKey(EditorPrefKeys.ResumeHttpAfterReload); } - if (bridge.IsRunning) + if (shouldResume) { - var stopTask = bridge.StopAsync(); + var stopTask = transport.StopAsync(TransportMode.Http); stopTask.ContinueWith(t => { if (t.IsFaulted && t.Exception != null) @@ -59,7 +59,9 @@ namespace MCPForUnity.Editor.Services bool resume = false; try { - resume = EditorPrefs.GetBool(EditorPrefKeys.ResumeHttpAfterReload, false); + // Only resume HTTP if it is still the selected transport. + bool useHttp = EditorPrefs.GetBool(EditorPrefKeys.UseHttpTransport, true); + resume = useHttp && EditorPrefs.GetBool(EditorPrefKeys.ResumeHttpAfterReload, false); if (resume) { EditorPrefs.DeleteKey(EditorPrefKeys.ResumeHttpAfterReload); @@ -90,7 +92,7 @@ namespace MCPForUnity.Editor.Services { try { - var startTask = MCPServiceLocator.Bridge.StartAsync(); + var startTask = MCPServiceLocator.TransportManager.StartAsync(TransportMode.Http); startTask.ContinueWith(t => { if (t.IsFaulted) @@ -123,7 +125,7 @@ namespace MCPForUnity.Editor.Services { try { - bool started = await MCPServiceLocator.Bridge.StartAsync(); + bool started = await MCPServiceLocator.TransportManager.StartAsync(TransportMode.Http); if (!started) { McpLog.Warn("Failed to resume HTTP MCP bridge after domain reload"); diff --git a/MCPForUnity/Editor/Services/StdioBridgeReloadHandler.cs b/MCPForUnity/Editor/Services/StdioBridgeReloadHandler.cs new file mode 100644 index 0000000..3202035 --- /dev/null +++ b/MCPForUnity/Editor/Services/StdioBridgeReloadHandler.cs @@ -0,0 +1,104 @@ +using System; +using UnityEditor; +using MCPForUnity.Editor.Constants; +using MCPForUnity.Editor.Helpers; +using MCPForUnity.Editor.Services.Transport; +using MCPForUnity.Editor.Services.Transport.Transports; + +namespace MCPForUnity.Editor.Services +{ + /// + /// Ensures the legacy stdio bridge resumes after domain reloads, mirroring the HTTP handler. + /// + [InitializeOnLoad] + internal static class StdioBridgeReloadHandler + { + static StdioBridgeReloadHandler() + { + AssemblyReloadEvents.beforeAssemblyReload += OnBeforeAssemblyReload; + AssemblyReloadEvents.afterAssemblyReload += OnAfterAssemblyReload; + } + + private static void OnBeforeAssemblyReload() + { + try + { + // Only persist resume intent when stdio is the active transport and the bridge is running. + bool useHttp = EditorPrefs.GetBool(EditorPrefKeys.UseHttpTransport, true); + bool isRunning = MCPServiceLocator.TransportManager.IsRunning(TransportMode.Stdio); + bool shouldResume = !useHttp && isRunning; + + if (shouldResume) + { + EditorPrefs.SetBool(EditorPrefKeys.ResumeStdioAfterReload, true); + + // Stop only the stdio bridge; leave HTTP untouched if it is running concurrently. + var stopTask = MCPServiceLocator.TransportManager.StopAsync(TransportMode.Stdio); + stopTask.ContinueWith(t => + { + if (t.IsFaulted && t.Exception != null) + { + McpLog.Warn($"Error stopping stdio bridge before reload: {t.Exception.GetBaseException()?.Message}"); + } + }, System.Threading.Tasks.TaskScheduler.Default); + } + else + { + EditorPrefs.DeleteKey(EditorPrefKeys.ResumeStdioAfterReload); + } + } + catch (Exception ex) + { + McpLog.Warn($"Failed to persist stdio reload flag: {ex.Message}"); + } + } + + private static void OnAfterAssemblyReload() + { + bool resume = false; + try + { + resume = EditorPrefs.GetBool(EditorPrefKeys.ResumeStdioAfterReload, false); + bool useHttp = EditorPrefs.GetBool(EditorPrefKeys.UseHttpTransport, true); + resume = resume && !useHttp; + if (resume) + { + EditorPrefs.DeleteKey(EditorPrefKeys.ResumeStdioAfterReload); + } + } + catch (Exception ex) + { + McpLog.Warn($"Failed to read stdio reload flag: {ex.Message}"); + } + + if (!resume) + { + return; + } + + // Restart via TransportManager so state stays in sync; if it fails (port busy), rely on UI to retry. + TryStartBridgeImmediate(); + } + + private static void TryStartBridgeImmediate() + { + var startTask = MCPServiceLocator.TransportManager.StartAsync(TransportMode.Stdio); + startTask.ContinueWith(t => + { + if (t.IsFaulted) + { + var baseEx = t.Exception?.GetBaseException(); + McpLog.Warn($"Failed to resume stdio bridge after reload: {baseEx?.Message}"); + return; + } + if (!t.Result) + { + McpLog.Warn("Failed to resume stdio bridge after domain reload"); + return; + } + + MCPForUnity.Editor.Windows.MCPForUnityEditorWindow.RequestHealthVerification(); + }, System.Threading.Tasks.TaskScheduler.Default); + } + } +} diff --git a/MCPForUnity/Editor/Services/StdioBridgeReloadHandler.cs.meta b/MCPForUnity/Editor/Services/StdioBridgeReloadHandler.cs.meta new file mode 100644 index 0000000..d4e43fa --- /dev/null +++ b/MCPForUnity/Editor/Services/StdioBridgeReloadHandler.cs.meta @@ -0,0 +1,11 @@ +fileFormatVersion: 2 +guid: 6e603c72a87974cf5b495cd683165fbf +MonoImporter: + externalObjects: {} + serializedVersion: 2 + defaultReferences: [] + executionOrder: 0 + icon: {instanceID: 0} + userData: + assetBundleName: + assetBundleVariant: diff --git a/MCPForUnity/Editor/Services/Transport/TransportManager.cs b/MCPForUnity/Editor/Services/Transport/TransportManager.cs index 7a6afe9..d221ab8 100644 --- a/MCPForUnity/Editor/Services/Transport/TransportManager.cs +++ b/MCPForUnity/Editor/Services/Transport/TransportManager.cs @@ -10,8 +10,10 @@ namespace MCPForUnity.Editor.Services.Transport /// public class TransportManager { - private IMcpTransportClient _active; - private TransportMode? _activeMode; + private IMcpTransportClient _httpClient; + private IMcpTransportClient _stdioClient; + private TransportState _httpState = TransportState.Disconnected("http"); + private TransportState _stdioState = TransportState.Disconnected("stdio"); private Func _webSocketFactory; private Func _stdioFactory; @@ -22,8 +24,8 @@ namespace MCPForUnity.Editor.Services.Transport () => new StdioTransportClient()); } - public IMcpTransportClient ActiveTransport => _active; - public TransportMode? ActiveMode => _activeMode; + public IMcpTransportClient ActiveTransport => null; // Deprecated single-transport accessor + public TransportMode? ActiveMode => null; // Deprecated single-transport accessor public void Configure( Func webSocketFactory, @@ -33,68 +35,115 @@ namespace MCPForUnity.Editor.Services.Transport _stdioFactory = stdioFactory ?? throw new ArgumentNullException(nameof(stdioFactory)); } - public async Task StartAsync(TransportMode mode) + private IMcpTransportClient GetOrCreateClient(TransportMode mode) { - await StopAsync(); - - IMcpTransportClient next = mode switch + return mode switch { - TransportMode.Stdio => _stdioFactory(), - TransportMode.Http => _webSocketFactory(), - _ => throw new ArgumentOutOfRangeException(nameof(mode), mode, "Unsupported transport mode") - } ?? throw new InvalidOperationException($"Factory returned null for transport mode {mode}"); - - bool started = await next.StartAsync(); - if (!started) - { - await next.StopAsync(); - _active = null; - _activeMode = null; - return false; - } - - _active = next; - _activeMode = mode; - return true; + TransportMode.Http => _httpClient ??= _webSocketFactory(), + TransportMode.Stdio => _stdioClient ??= _stdioFactory(), + _ => throw new ArgumentOutOfRangeException(nameof(mode), mode, "Unsupported transport mode"), + }; } - public async Task StopAsync() + private IMcpTransportClient GetClient(TransportMode mode) { - if (_active != null) + return mode switch + { + TransportMode.Http => _httpClient, + TransportMode.Stdio => _stdioClient, + _ => throw new ArgumentOutOfRangeException(nameof(mode), mode, "Unsupported transport mode"), + }; + } + + public async Task StartAsync(TransportMode mode) + { + IMcpTransportClient client = GetOrCreateClient(mode); + + bool started = await client.StartAsync(); + if (!started) { try { - await _active.StopAsync(); + await client.StopAsync(); } catch (Exception ex) { - McpLog.Warn($"Error while stopping transport {_active.TransportName}: {ex.Message}"); - } - finally - { - _active = null; - _activeMode = null; + McpLog.Warn($"Error while stopping transport {client.TransportName}: {ex.Message}"); } + UpdateState(mode, TransportState.Disconnected(client.TransportName, "Failed to start")); + return false; + } + + UpdateState(mode, client.State ?? TransportState.Connected(client.TransportName)); + return true; + } + + public async Task StopAsync(TransportMode? mode = null) + { + async Task StopClient(IMcpTransportClient client, TransportMode clientMode) + { + if (client == null) return; + try { await client.StopAsync(); } + catch (Exception ex) { McpLog.Warn($"Error while stopping transport {client.TransportName}: {ex.Message}"); } + finally { UpdateState(clientMode, TransportState.Disconnected(client.TransportName)); } + } + + if (mode == null) + { + await StopClient(_httpClient, TransportMode.Http); + await StopClient(_stdioClient, TransportMode.Stdio); + return; + } + + if (mode == TransportMode.Http) + { + await StopClient(_httpClient, TransportMode.Http); + } + else + { + await StopClient(_stdioClient, TransportMode.Stdio); } } - public async Task VerifyAsync() + public async Task VerifyAsync(TransportMode mode) { - if (_active == null) + IMcpTransportClient client = GetClient(mode); + if (client == null) { return false; } - return await _active.VerifyAsync(); + + bool ok = await client.VerifyAsync(); + var state = client.State ?? TransportState.Disconnected(client.TransportName, "No state reported"); + UpdateState(mode, state); + return ok; } - public TransportState GetState() + public TransportState GetState(TransportMode mode) { - if (_active == null) + return mode switch { - return TransportState.Disconnected(_activeMode?.ToString()?.ToLowerInvariant() ?? "unknown", "Transport not started"); - } + TransportMode.Http => _httpState, + TransportMode.Stdio => _stdioState, + _ => throw new ArgumentOutOfRangeException(nameof(mode), mode, "Unsupported transport mode"), + }; + } - return _active.State ?? TransportState.Disconnected(_active.TransportName, "No state reported"); + public bool IsRunning(TransportMode mode) => GetState(mode).IsConnected; + + private void UpdateState(TransportMode mode, TransportState state) + { + switch (mode) + { + case TransportMode.Http: + _httpState = state; + break; + case TransportMode.Stdio: + _stdioState = state; + break; + default: + throw new ArgumentOutOfRangeException(nameof(mode), mode, "Unsupported transport mode"); + } } } diff --git a/MCPForUnity/Editor/Services/Transport/Transports/StdioBridgeHost.cs b/MCPForUnity/Editor/Services/Transport/Transports/StdioBridgeHost.cs index 5cc1585..ab127b4 100644 --- a/MCPForUnity/Editor/Services/Transport/Transports/StdioBridgeHost.cs +++ b/MCPForUnity/Editor/Services/Transport/Transports/StdioBridgeHost.cs @@ -57,7 +57,6 @@ namespace MCPForUnity.Editor.Services.Transport.Transports private static int mainThreadId; private static int currentUnityPort = 6400; private static bool isAutoConnectMode = false; - private static bool shouldRestartAfterReload = false; private const ulong MaxFrameBytes = 64UL * 1024 * 1024; private const int FrameIOTimeoutMs = 30000; @@ -162,8 +161,6 @@ namespace MCPForUnity.Editor.Services.Transport.Transports } } EditorApplication.quitting += Stop; - AssemblyReloadEvents.beforeAssemblyReload += OnBeforeAssemblyReload; - AssemblyReloadEvents.afterAssemblyReload += OnAfterAssemblyReload; EditorApplication.playModeStateChanged += _ => { if (ShouldAutoStartBridge()) @@ -406,10 +403,6 @@ namespace MCPForUnity.Editor.Services.Transport.Transports listenerTask = Task.Run(() => ListenerLoopAsync(cts.Token)); CommandRegistry.Initialize(); EditorApplication.update += ProcessCommands; - try { AssemblyReloadEvents.beforeAssemblyReload -= OnBeforeAssemblyReload; } catch { } - try { AssemblyReloadEvents.beforeAssemblyReload += OnBeforeAssemblyReload; } catch { } - try { AssemblyReloadEvents.afterAssemblyReload -= OnAfterAssemblyReload; } catch { } - try { AssemblyReloadEvents.afterAssemblyReload += OnAfterAssemblyReload; } catch { } try { EditorApplication.quitting -= Stop; } catch { } try { EditorApplication.quitting += Stop; } catch { } heartbeatSeq++; @@ -470,8 +463,6 @@ namespace MCPForUnity.Editor.Services.Transport.Transports } try { EditorApplication.update -= ProcessCommands; } catch { } - try { AssemblyReloadEvents.beforeAssemblyReload -= OnBeforeAssemblyReload; } catch { } - try { AssemblyReloadEvents.afterAssemblyReload -= OnAfterAssemblyReload; } catch { } try { EditorApplication.quitting -= Stop; } catch { } try @@ -1023,47 +1014,6 @@ namespace MCPForUnity.Editor.Services.Transport.Transports return false; } - private static void OnBeforeAssemblyReload() - { - if (isRunning) - { - shouldRestartAfterReload = true; - } - try { Stop(); } catch { } - } - - private static void OnAfterAssemblyReload() - { - WriteHeartbeat(false, "idle"); - LogBreadcrumb("Idle"); - bool shouldResume = ShouldAutoStartBridge() || shouldRestartAfterReload; - if (shouldRestartAfterReload) - { - shouldRestartAfterReload = false; - } - if (!shouldResume) - { - return; - } - - // If we're not compiling, try to bring the bridge up immediately to avoid depending on editor focus. - if (!IsCompiling()) - { - try - { - Start(); - return; // Successful immediate start; no need to schedule a delayed retry - } - catch (Exception ex) - { - // Fall through to delayed retry if immediate start fails - McpLog.Warn($"Immediate STDIO bridge restart after reload failed: {ex.Message}"); - } - } - - // Fallback path when compiling or if immediate start failed - ScheduleInitRetry(); - } private static void WriteHeartbeat(bool reloading, string reason = null) {