coplaydev--unity-mcp/MCPForUnity/Editor/Services/StdioBridgeReloadHandler.cs

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
{
    /// <summary>
    /// Ensures the legacy stdio bridge resumes after domain reloads, mirroring the HTTP handler.
    /// </summary>
    [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");
            }
        }
    }
}