unity-mcp/MCPForUnity/Editor/Services/BridgeControlService.cs

using System;
using System.Threading.Tasks;
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
{
    /// <summary>
    /// Bridges the editor UI to the active transport (HTTP with WebSocket push, or stdio).
    /// </summary>
    public class BridgeControlService : IBridgeControlService
    {
        private readonly TransportManager _transportManager;
        private TransportMode _preferredMode = TransportMode.Http;

        public BridgeControlService()
        {
            _transportManager = MCPServiceLocator.TransportManager;
        }

        private TransportMode ResolvePreferredMode()
        {
            bool useHttp = EditorPrefs.GetBool(EditorPrefKeys.UseHttpTransport, true);
            _preferredMode = useHttp ? TransportMode.Http : TransportMode.Stdio;
            return _preferredMode;
        }

        private static BridgeVerificationResult BuildVerificationResult(TransportState state, TransportMode mode, bool pingSucceeded, string messageOverride = null, bool? handshakeOverride = null)
        {
            bool handshakeValid = handshakeOverride ?? (mode == TransportMode.Stdio ? state.IsConnected : true);
            string transportLabel = string.IsNullOrWhiteSpace(state.TransportName)
                ? mode.ToString().ToLowerInvariant()
                : state.TransportName;
            string detailSuffix = string.IsNullOrWhiteSpace(state.Details) ? string.Empty : $" [{state.Details}]";
            string message = messageOverride
                ?? state.Error
                ?? (state.IsConnected ? $"Transport '{transportLabel}' connected{detailSuffix}" : $"Transport '{transportLabel}' disconnected{detailSuffix}");

            return new BridgeVerificationResult
            {
                Success = pingSucceeded && handshakeValid,
                HandshakeValid = handshakeValid,
                PingSucceeded = pingSucceeded,
                Message = message
            };
        }

        public bool IsRunning => _transportManager.GetState().IsConnected;

        public int CurrentPort
        {
            get
            {
                var state = _transportManager.GetState();
                if (state.Port.HasValue)
                {
                    return state.Port.Value;
                }

                // Legacy fallback while the stdio bridge is still in play
                return StdioBridgeHost.GetCurrentPort();
            }
        }

        public bool IsAutoConnectMode => StdioBridgeHost.IsAutoConnectMode();
        public TransportMode? ActiveMode => _transportManager.ActiveMode;

        public async Task<bool> StartAsync()
        {
            var mode = ResolvePreferredMode();
            try
            {
                bool started = await _transportManager.StartAsync(mode);
                if (!started)
                {
                    McpLog.Warn($"Failed to start MCP transport: {mode}");
                }
                return started;
            }
            catch (Exception ex)
            {
                McpLog.Error($"Error starting MCP transport {mode}: {ex.Message}");
                return false;
            }
        }

        public async Task StopAsync()
        {
            try
            {
                await _transportManager.StopAsync();
            }
            catch (Exception ex)
            {
                McpLog.Warn($"Error stopping MCP transport: {ex.Message}");
            }
        }

        public async Task<BridgeVerificationResult> VerifyAsync()
        {
            var mode = _transportManager.ActiveMode ?? ResolvePreferredMode();
            bool pingSucceeded = await _transportManager.VerifyAsync();
            var state = _transportManager.GetState();
            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();

            if (mode == TransportMode.Stdio)
            {
                bool handshakeValid = state.IsConnected && port == CurrentPort;
                string message = handshakeValid
                    ? $"STDIO transport listening on port {CurrentPort}"
                    : $"STDIO transport port mismatch (expected {CurrentPort}, got {port})";
                return BuildVerificationResult(state, mode, pingSucceeded && handshakeValid, message, handshakeValid);
            }

            return BuildVerificationResult(state, mode, pingSucceeded);
        }

    }
}