using System; using System.Threading; using System.Threading.Tasks; using MCPForUnity.Editor.Constants; using MCPForUnity.Editor.Helpers; using MCPForUnity.Editor.Services.Transport; using MCPForUnity.Editor.Services.Transport.Transports; using MCPForUnity.Editor.Windows; using UnityEditor; namespace MCPForUnity.Editor.Services { /// /// Ensures the legacy stdio bridge resumes after domain reloads, mirroring the HTTP handler. /// [InitializeOnLoad] internal static class StdioBridgeReloadHandler { private static readonly TimeSpan[] ResumeRetrySchedule = { TimeSpan.Zero, TimeSpan.FromSeconds(1), TimeSpan.FromSeconds(3), TimeSpan.FromSeconds(5), TimeSpan.FromSeconds(10), TimeSpan.FromSeconds(30) }; private static CancellationTokenSource _retryCts; static StdioBridgeReloadHandler() { AssemblyReloadEvents.beforeAssemblyReload += OnBeforeAssemblyReload; AssemblyReloadEvents.afterAssemblyReload += OnAfterAssemblyReload; EditorApplication.quitting += CancelRetries; } private static void CancelRetries() { try { _retryCts?.Cancel(); } catch { } } private static void OnBeforeAssemblyReload() { // Cancel any in-flight retry loop before the next reload. CancelRetries(); try { // Only persist resume intent when stdio is the active transport and the bridge is running. bool useHttp = EditorConfigurationCache.Instance.UseHttpTransport; // Check both TransportManager AND StdioBridgeHost directly, because CI starts via StdioBridgeHost // bypassing TransportManager state. bool tmRunning = MCPServiceLocator.TransportManager.IsRunning(TransportMode.Stdio); bool hostRunning = StdioBridgeHost.IsRunning; bool isRunning = tmRunning || hostRunning; bool shouldResume = !useHttp && isRunning; if (shouldResume) { EditorPrefs.SetBool(EditorPrefKeys.ResumeStdioAfterReload, true); } else { EditorPrefs.DeleteKey(EditorPrefKeys.ResumeStdioAfterReload); } if (isRunning) { // Stop only stdio before reload. This is centralized here so resume-flag updates // and teardown cannot race each other via separate beforeAssemblyReload handlers. var stopTask = MCPServiceLocator.TransportManager.StopAsync(TransportMode.Stdio); try { stopTask.Wait(500); } catch { } // Legacy safety: stdio may have been started outside TransportManager state. try { StdioBridgeHost.Stop(); } catch { } } if (shouldResume) { // Write reloading status so clients don't think we vanished. StdioBridgeHost.WriteHeartbeat(true, "reloading"); } } catch (Exception ex) { McpLog.Warn($"Failed to persist stdio reload flag: {ex.Message}"); } } private static void OnAfterAssemblyReload() { bool resume = false; try { bool resumeFlag = EditorPrefs.GetBool(EditorPrefKeys.ResumeStdioAfterReload, false); bool useHttp = EditorConfigurationCache.Instance.UseHttpTransport; resume = resumeFlag && !useHttp; // If we're not going to resume, clear the flag immediately to avoid stuck "Resuming..." state if (!resume) { EditorPrefs.DeleteKey(EditorPrefKeys.ResumeStdioAfterReload); } } catch (Exception ex) { McpLog.Warn($"Failed to read stdio reload flag: {ex.Message}"); } if (!resume) { return; } // If the editor is not compiling, attempt an immediate restart without relying on editor focus. bool isCompiling = EditorApplication.isCompiling; try { var pipeline = Type.GetType("UnityEditor.Compilation.CompilationPipeline, UnityEditor"); var prop = pipeline?.GetProperty("isCompiling", System.Reflection.BindingFlags.Public | System.Reflection.BindingFlags.Static); if (prop != null) isCompiling |= (bool)prop.GetValue(null); } catch { } if (!isCompiling) { _ = ResumeStdioWithRetriesAsync(); return; } // Fallback when compiling: schedule on the editor loop EditorApplication.delayCall += () => { _ = ResumeStdioWithRetriesAsync(); }; } private static async Task ResumeStdioWithRetriesAsync() { // Cancel any previous retry loop and create a fresh token. CancelRetries(); var cts = _retryCts = new CancellationTokenSource(); var token = cts.Token; Exception lastException = null; for (int i = 0; i < ResumeRetrySchedule.Length; i++) { if (token.IsCancellationRequested) return; int attempt = i + 1; McpLog.Debug($"[Stdio Reload] Resume attempt {attempt}/{ResumeRetrySchedule.Length}"); TimeSpan delay = ResumeRetrySchedule[i]; if (delay > TimeSpan.Zero) { McpLog.Debug($"[Stdio Reload] Waiting {delay.TotalSeconds:0.#}s before resume attempt {attempt}"); try { await Task.Delay(delay, token); } catch (OperationCanceledException) { return; } } // Abort retries if the user switched transports while we were waiting. if (EditorConfigurationCache.Instance.UseHttpTransport) { try { EditorPrefs.DeleteKey(EditorPrefKeys.ResumeStdioAfterReload); } catch { } return; } try { bool started = await MCPServiceLocator.TransportManager.StartAsync(TransportMode.Stdio); if (started) { McpLog.Debug($"[Stdio Reload] Resume succeeded on attempt {attempt}"); try { EditorPrefs.DeleteKey(EditorPrefKeys.ResumeStdioAfterReload); } catch { } MCPForUnityEditorWindow.RequestHealthVerification(); return; } var state = MCPServiceLocator.TransportManager.GetState(TransportMode.Stdio); string reason = string.IsNullOrWhiteSpace(state?.Error) ? "no error detail" : state.Error; McpLog.Debug($"[Stdio Reload] Resume attempt {attempt} failed: {reason}"); } catch (Exception ex) { lastException = ex; McpLog.Debug($"[Stdio Reload] Resume attempt {attempt} threw: {ex.Message}"); } } try { EditorPrefs.DeleteKey(EditorPrefKeys.ResumeStdioAfterReload); } catch { } // Clear the stale "reloading" heartbeat so clients stop seeing reloading=true. // The bridge isn't running, so clients will get connection-refused (recoverable) // instead of hanging on a zombie socket or being rejected by the preflight check. try { StdioBridgeHost.WriteHeartbeat(false, "stopped"); } catch { } if (lastException != null) { McpLog.Warn($"Failed to resume stdio bridge after domain reload: {lastException.Message}"); } else { McpLog.Warn("Failed to resume stdio bridge after domain reload"); } } } }