969 lines
39 KiB
C#
969 lines
39 KiB
C#
using System;
|
|
using System.Collections.Generic;
|
|
using System.Collections.Concurrent;
|
|
using System.IO;
|
|
using System.Linq;
|
|
using System.Net;
|
|
using System.Net.Sockets;
|
|
using System.Text;
|
|
using System.Threading;
|
|
using System.Threading.Tasks;
|
|
|
|
namespace io.github.hatayama.uLoopMCP
|
|
{
|
|
// Related classes:
|
|
// - McpServerController: Manages the lifecycle of this server.
|
|
// - UnityCommandExecutor: Executes commands received from clients.
|
|
// - JsonRpcProcessor: Handles JSON-RPC 2.0 message processing.
|
|
/// <summary>
|
|
/// Represents a connected client
|
|
/// </summary>
|
|
public class ConnectedClient
|
|
{
|
|
public readonly string Endpoint;
|
|
public readonly string ClientName;
|
|
public readonly DateTime ConnectedAt;
|
|
public readonly int Port;
|
|
public readonly NetworkStream Stream;
|
|
|
|
public ConnectedClient(string endpoint, NetworkStream stream, int port, string clientName = McpConstants.UNKNOWN_CLIENT_NAME)
|
|
{
|
|
Endpoint = endpoint;
|
|
Stream = stream; // Allow null stream for UI display purposes
|
|
ClientName = clientName;
|
|
ConnectedAt = DateTime.Now;
|
|
Port = port;
|
|
}
|
|
|
|
// Private constructor for WithClientName to preserve ConnectedAt
|
|
private ConnectedClient(string endpoint, NetworkStream stream, int port, string clientName, DateTime connectedAt)
|
|
{
|
|
Endpoint = endpoint;
|
|
Stream = stream; // Allow null stream for UI display purposes
|
|
ClientName = clientName;
|
|
ConnectedAt = connectedAt;
|
|
Port = port;
|
|
}
|
|
|
|
public ConnectedClient WithClientName(string clientName)
|
|
{
|
|
return new ConnectedClient(Endpoint, Stream, Port, clientName, ConnectedAt);
|
|
}
|
|
|
|
}
|
|
|
|
/// <summary>
|
|
/// Unity MCP Bridge TCP/IP Server.
|
|
/// Accepts connections from the TypeScript MCP Server and handles JSON-RPC 2.0 communication.
|
|
/// </summary>
|
|
public class McpBridgeServer : IDisposable
|
|
{
|
|
// Note: Domain reload progress is now tracked via McpSessionManager
|
|
|
|
// Events for server lifecycle notifications
|
|
public static event System.Action OnServerStopping;
|
|
public static event System.Action OnServerStarted;
|
|
|
|
// Events for individual tool management
|
|
public static event System.Action<ConnectedClient> OnToolConnected;
|
|
public static event System.Action<string> OnToolDisconnected;
|
|
public static event System.Action OnAllToolsCleared;
|
|
|
|
// Fired from thread pool when ServerLoopAsync exits while _isRunning is still true.
|
|
// Subscribers must marshal to main thread before accessing Unity APIs.
|
|
public static event System.Action OnServerLoopExited;
|
|
|
|
// HResult error codes for normal disconnection detection
|
|
private static readonly HashSet<int> NormalDisconnectionHResults = new()
|
|
{
|
|
unchecked((int)0x800703E3), // ERROR_OPERATION_ABORTED
|
|
unchecked((int)0x80070040), // ERROR_NETNAME_DELETED
|
|
unchecked((int)0x80072745), // ERROR_CONNECTION_ABORTED
|
|
unchecked((int)0x80072746) // ERROR_CONNECTION_RESET
|
|
};
|
|
|
|
private TcpListener _tcpListener;
|
|
private CancellationTokenSource _cancellationTokenSource;
|
|
private Task _serverTask;
|
|
private readonly ConcurrentDictionary<int, Task> _clientTasks = new();
|
|
// Read from thread pool (ServerLoopAsync), written from main thread (StopServer)
|
|
private volatile bool _isRunning = false;
|
|
|
|
// Guard against concurrent cleanup from ServerLoopAsync finally + external disposal
|
|
private int _unexpectedExitCleanupStarted = 0;
|
|
|
|
// Client management for broadcasting notifications
|
|
private readonly ConcurrentDictionary<string, ConnectedClient> _connectedClients = new();
|
|
|
|
/// <summary>
|
|
/// Whether the server is running.
|
|
/// </summary>
|
|
public bool IsRunning => _isRunning;
|
|
|
|
/// <summary>
|
|
/// The server's port number.
|
|
/// </summary>
|
|
public int Port { get; private set; } = McpEditorSettings.GetCustomPort();
|
|
|
|
/// <summary>
|
|
/// Event on client connection.
|
|
/// </summary>
|
|
public event Action<string> OnClientConnected;
|
|
|
|
/// <summary>
|
|
/// Event on client disconnection.
|
|
/// </summary>
|
|
public event Action<string> OnClientDisconnected;
|
|
|
|
/// <summary>
|
|
/// Event on error.
|
|
/// </summary>
|
|
public event Action<string> OnError;
|
|
|
|
/// <summary>
|
|
/// Generate unique client key using Endpoint
|
|
/// </summary>
|
|
private string GenerateClientKey(string endpoint)
|
|
{
|
|
// Use endpoint as unique identifier
|
|
return endpoint;
|
|
}
|
|
|
|
/// <summary>
|
|
/// Get list of connected clients sorted by name
|
|
/// </summary>
|
|
public IReadOnlyCollection<ConnectedClient> GetConnectedClients()
|
|
{
|
|
return _connectedClients.Values.OrderBy(client => client.ClientName).ToArray();
|
|
}
|
|
|
|
/// <summary>
|
|
/// Update client name for a connected client
|
|
/// </summary>
|
|
public void UpdateClientName(string clientEndpoint, string clientName)
|
|
{
|
|
// Find client by endpoint (backward compatibility)
|
|
ConnectedClient targetClient = _connectedClients.Values
|
|
.FirstOrDefault(c => c.Endpoint == clientEndpoint);
|
|
|
|
if (targetClient != null)
|
|
{
|
|
string clientKey = GenerateClientKey(targetClient.Endpoint);
|
|
ConnectedClient updatedClient = targetClient.WithClientName(clientName);
|
|
bool updateResult = _connectedClients.TryUpdate(clientKey, updatedClient, targetClient);
|
|
|
|
// Clear reconnecting flags when client name is successfully set (client is now fully connected)
|
|
if (updateResult && clientName != McpConstants.UNKNOWN_CLIENT_NAME)
|
|
{
|
|
McpServerController.ClearReconnectingFlag();
|
|
|
|
// Notify tool connected
|
|
OnToolConnected?.Invoke(updatedClient);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// <summary>
|
|
/// Checks if the specified port is in use.
|
|
/// Delegates to NetworkUtility for consistent port checking behavior.
|
|
/// </summary>
|
|
/// <param name="port">The port number to check.</param>
|
|
/// <returns>True if the port is in use.</returns>
|
|
public static bool IsPortInUse(int port)
|
|
{
|
|
return NetworkUtility.IsPortInUse(port);
|
|
}
|
|
|
|
/// <summary>
|
|
/// Starts the server.
|
|
/// </summary>
|
|
/// <param name="port">
|
|
/// The port number to bind to. Use -1 to fall back to the saved custom port
|
|
/// from <see cref="McpEditorSettings.GetCustomPort"/>. Defaults to -1.
|
|
/// </param>
|
|
public void StartServer(int port = -1, bool clearServerStartingLockWhenReady = true)
|
|
{
|
|
if (_isRunning)
|
|
{
|
|
return;
|
|
}
|
|
|
|
Port = port == -1 ? McpEditorSettings.GetCustomPort() : port;
|
|
_cancellationTokenSource = new CancellationTokenSource();
|
|
_unexpectedExitCleanupStarted = 0;
|
|
|
|
try
|
|
{
|
|
_tcpListener = new TcpListener(IPAddress.Loopback, Port);
|
|
_tcpListener.Start();
|
|
_isRunning = true;
|
|
|
|
TcpListener tcpListener = _tcpListener;
|
|
CancellationToken cancellationToken = _cancellationTokenSource.Token;
|
|
_serverTask = Task.Run(() => ServerLoopAsync(tcpListener, cancellationToken));
|
|
|
|
// Safety net: log if the server task faults unexpectedly.
|
|
// Primary detection is in ServerLoopAsync's finally block; this catches unhandled exceptions in Task.Run itself.
|
|
_serverTask.ContinueWith(task =>
|
|
{
|
|
if (task.IsFaulted)
|
|
{
|
|
VibeLogger.LogError(
|
|
"server_task_faulted",
|
|
$"Server task faulted unexpectedly: {task.Exception?.GetBaseException().Message}",
|
|
new { exceptionType = task.Exception?.GetBaseException().GetType().Name }
|
|
);
|
|
}
|
|
}, TaskScheduler.Default);
|
|
|
|
// Server is now ready to accept connections - clean up compilation/reload locks.
|
|
CompilationLockService.DeleteLockFile();
|
|
DomainReloadDetectionService.DeleteLockFile();
|
|
if (clearServerStartingLockWhenReady)
|
|
{
|
|
ServerStartingLockService.DeleteLockFile();
|
|
}
|
|
|
|
// Notify that server has started
|
|
OnServerStarted?.Invoke();
|
|
|
|
}
|
|
catch (SocketException ex) when (ex.SocketErrorCode == SocketError.AddressAlreadyInUse)
|
|
{
|
|
_isRunning = false;
|
|
string errorMessage = $"Port {Port} is already in use. Please choose a different port.";
|
|
OnError?.Invoke(errorMessage);
|
|
throw new InvalidOperationException(errorMessage, ex);
|
|
}
|
|
catch (Exception ex)
|
|
{
|
|
_isRunning = false;
|
|
string errorMessage = $"Failed to start MCP Server: {ex.Message}";
|
|
OnError?.Invoke(errorMessage);
|
|
throw;
|
|
}
|
|
}
|
|
|
|
/// <summary>
|
|
/// Stops the server.
|
|
/// </summary>
|
|
public void StopServer()
|
|
{
|
|
StopServerCore(ServerStopMode.Normal);
|
|
}
|
|
|
|
/// <summary>
|
|
/// Domain reload can suspend editor-thread continuations before the server task observes shutdown.
|
|
/// </summary>
|
|
public void StopServerBeforeDomainReload()
|
|
{
|
|
StopServerCore(ServerStopMode.BeforeDomainReload);
|
|
}
|
|
|
|
private void StopServerCore(ServerStopMode mode)
|
|
{
|
|
bool wasRunning = _isRunning;
|
|
CancellationTokenSource cancellationTokenSource = Interlocked.Exchange(ref _cancellationTokenSource, null);
|
|
TcpListener tcpListener = Interlocked.Exchange(ref _tcpListener, null);
|
|
Task serverTask = mode == ServerStopMode.BeforeDomainReload
|
|
? _serverTask
|
|
: Interlocked.Exchange(ref _serverTask, null);
|
|
|
|
if (!wasRunning &&
|
|
cancellationTokenSource == null &&
|
|
tcpListener == null &&
|
|
serverTask == null)
|
|
{
|
|
return;
|
|
}
|
|
|
|
if (wasRunning && mode != ServerStopMode.UnexpectedLoopExit)
|
|
{
|
|
ServerShutdownReason shutdownReason = McpEditorSettings.GetIsDomainReloadInProgress()
|
|
? ServerShutdownReason.DomainReload
|
|
: ServerShutdownReason.EditorQuit;
|
|
|
|
SendShutdownNotification(shutdownReason);
|
|
OnServerStopping?.Invoke();
|
|
}
|
|
|
|
_isRunning = false;
|
|
|
|
DisconnectAllClientsCore(notifyLifecycleEvents: mode != ServerStopMode.UnexpectedLoopExit);
|
|
|
|
cancellationTokenSource?.Cancel();
|
|
tcpListener?.Stop();
|
|
|
|
if (mode == ServerStopMode.BeforeDomainReload)
|
|
{
|
|
return;
|
|
}
|
|
|
|
if (mode == ServerStopMode.UnexpectedLoopExit)
|
|
{
|
|
cancellationTokenSource?.Dispose();
|
|
return;
|
|
}
|
|
|
|
serverTask?.Wait(TimeSpan.FromSeconds(McpServerConfig.SHUTDOWN_TIMEOUT_SECONDS));
|
|
DisposeCancellationTokenSourceAfterClientTasks(cancellationTokenSource);
|
|
}
|
|
|
|
private void DisposeCancellationTokenSourceAfterClientTasks(
|
|
CancellationTokenSource cancellationTokenSource)
|
|
{
|
|
Task[] clientTasks = _clientTasks.Values.ToArray();
|
|
if (clientTasks.Length == 0)
|
|
{
|
|
cancellationTokenSource?.Dispose();
|
|
return;
|
|
}
|
|
|
|
if (MainThreadSwitcher.IsMainThread)
|
|
{
|
|
_ = ObserveClientTasksThenDisposeCancellationTokenSourceAsync(
|
|
clientTasks,
|
|
cancellationTokenSource);
|
|
return;
|
|
}
|
|
|
|
WaitForClientTasksToComplete(clientTasks);
|
|
cancellationTokenSource?.Dispose();
|
|
}
|
|
|
|
private static void WaitForClientTasksToComplete(Task[] clientTasks)
|
|
{
|
|
Task completionTask = CreateObservedClientTaskCompletionTask(clientTasks);
|
|
bool completed = completionTask.Wait(TimeSpan.FromSeconds(McpServerConfig.SHUTDOWN_TIMEOUT_SECONDS));
|
|
if (!completed)
|
|
{
|
|
LogClientTaskShutdownTimeout(clientTasks.Length);
|
|
}
|
|
}
|
|
|
|
private static async Task ObserveClientTasksThenDisposeCancellationTokenSourceAsync(
|
|
Task[] clientTasks,
|
|
CancellationTokenSource cancellationTokenSource)
|
|
{
|
|
try
|
|
{
|
|
Task completionTask = CreateObservedClientTaskCompletionTask(clientTasks);
|
|
Task timeoutTask = Task.Delay(TimeSpan.FromSeconds(McpServerConfig.SHUTDOWN_TIMEOUT_SECONDS));
|
|
Task finishedTask = await Task.WhenAny(completionTask, timeoutTask);
|
|
if (!ReferenceEquals(finishedTask, completionTask))
|
|
{
|
|
LogClientTaskShutdownTimeout(clientTasks.Length);
|
|
}
|
|
}
|
|
finally
|
|
{
|
|
cancellationTokenSource?.Dispose();
|
|
}
|
|
}
|
|
|
|
private static Task CreateObservedClientTaskCompletionTask(Task[] clientTasks)
|
|
{
|
|
return Task.WhenAll(clientTasks).ContinueWith(task =>
|
|
{
|
|
if (task.IsFaulted)
|
|
{
|
|
_ = task.Exception;
|
|
}
|
|
}, CancellationToken.None, TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default);
|
|
}
|
|
|
|
private static void LogClientTaskShutdownTimeout(int activeClientTasks)
|
|
{
|
|
VibeLogger.LogWarning(
|
|
"client_tasks_shutdown_timeout",
|
|
"Timed out waiting for client handlers to finish during normal server shutdown.",
|
|
new { activeClientTasks }
|
|
);
|
|
}
|
|
|
|
/// <summary>
|
|
/// Explicitly disconnect all connected clients
|
|
/// This ensures TypeScript clients receive proper close events
|
|
/// </summary>
|
|
private void DisconnectAllClients()
|
|
{
|
|
DisconnectAllClientsCore(notifyLifecycleEvents: true);
|
|
}
|
|
|
|
/// <summary>
|
|
/// OnAllToolsCleared subscribers (UI components) require the main thread.
|
|
/// Background-thread callers (CleanupAfterUnexpectedLoopExit) must suppress
|
|
/// lifecycle events to avoid cross-thread Unity API violations.
|
|
/// </summary>
|
|
private void DisconnectAllClientsCore(bool notifyLifecycleEvents)
|
|
{
|
|
if (_connectedClients.IsEmpty)
|
|
{
|
|
return;
|
|
}
|
|
|
|
List<string> clientsToRemove = new List<string>();
|
|
|
|
foreach (KeyValuePair<string, ConnectedClient> client in _connectedClients)
|
|
{
|
|
try
|
|
{
|
|
// Close the NetworkStream to send proper close event to TypeScript client
|
|
if (client.Value.Stream != null && client.Value.Stream.CanWrite)
|
|
{
|
|
client.Value.Stream.Close();
|
|
}
|
|
clientsToRemove.Add(client.Key);
|
|
}
|
|
catch (Exception ex)
|
|
{
|
|
OnError?.Invoke($"Error disconnecting client {client.Key}: {ex.Message}");
|
|
clientsToRemove.Add(client.Key); // Remove even if disconnect failed
|
|
}
|
|
}
|
|
|
|
// Remove all clients from the connected clients list
|
|
foreach (string clientKey in clientsToRemove)
|
|
{
|
|
_connectedClients.TryRemove(clientKey, out _);
|
|
}
|
|
|
|
// Lifecycle events require main thread — only fire when explicitly requested
|
|
if (notifyLifecycleEvents && !McpEditorSettings.GetIsDomainReloadInProgress())
|
|
{
|
|
OnAllToolsCleared?.Invoke();
|
|
}
|
|
}
|
|
|
|
/// <summary>
|
|
/// StopServer() guards on _isRunning==true, but by the time this runs _isRunning may already
|
|
/// be false or the normal shutdown path may race with the finally block.
|
|
/// A separate cleanup path that skips the _isRunning guard is needed.
|
|
/// Lifecycle events are deferred to OnServerLoopExited → EditorApplication.delayCall
|
|
/// because this runs on the thread pool where Unity APIs are unsafe.
|
|
/// </summary>
|
|
private void CleanupAfterUnexpectedLoopExit()
|
|
{
|
|
if (Interlocked.Exchange(ref _unexpectedExitCleanupStarted, 1) != 0)
|
|
{
|
|
return;
|
|
}
|
|
|
|
StopServerCore(ServerStopMode.UnexpectedLoopExit);
|
|
}
|
|
|
|
/// <summary>
|
|
/// The server's main loop.
|
|
/// </summary>
|
|
private async Task ServerLoopAsync(TcpListener listener, CancellationToken cancellationToken)
|
|
{
|
|
System.Diagnostics.Debug.Assert(listener != null, "TcpListener must be captured before starting the server loop.");
|
|
try
|
|
{
|
|
while (!cancellationToken.IsCancellationRequested && _isRunning)
|
|
{
|
|
try
|
|
{
|
|
TcpClient client = await AcceptTcpClientAsync(listener, cancellationToken);
|
|
if (client != null)
|
|
{
|
|
string clientEndpoint = client.Client.RemoteEndPoint?.ToString() ?? McpServerConfig.UNKNOWN_CLIENT_ENDPOINT;
|
|
OnClientConnected?.Invoke(clientEndpoint);
|
|
|
|
Task clientTask = Task.Run(() => HandleClientAsync(client, cancellationToken));
|
|
TrackClientTask(clientTask);
|
|
}
|
|
}
|
|
catch (ObjectDisposedException)
|
|
{
|
|
// Expected when StopServer() disposes TcpListener while accept is pending.
|
|
// If _isRunning is still true here, this is an unexpected disposal — finally block handles state cleanup.
|
|
if (_isRunning)
|
|
{
|
|
VibeLogger.LogWarning(
|
|
"server_loop_disposed_while_running",
|
|
"TcpListener disposed while server was still marked as running. Exiting loop."
|
|
);
|
|
}
|
|
break;
|
|
}
|
|
catch (ThreadAbortException ex)
|
|
{
|
|
// Log and re-throw ThreadAbortException
|
|
if (!McpEditorSettings.GetIsDomainReloadInProgress())
|
|
{
|
|
OnError?.Invoke($"Unexpected thread abort: {ex.Message}");
|
|
}
|
|
throw;
|
|
}
|
|
catch (Exception ex)
|
|
{
|
|
if (!cancellationToken.IsCancellationRequested)
|
|
{
|
|
string errorMessage = $"Server loop error: {ex.Message}";
|
|
OnError?.Invoke(errorMessage);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
finally
|
|
{
|
|
bool wasUnexpectedExit = ShouldTreatLoopExitAsUnexpected(_isRunning, cancellationToken);
|
|
if (wasUnexpectedExit)
|
|
{
|
|
VibeLogger.LogWarning(
|
|
"server_loop_unexpected_exit",
|
|
"ServerLoopAsync exited while _isRunning was still true. Cleaning up and triggering recovery.",
|
|
new { cancellationRequested = cancellationToken.IsCancellationRequested }
|
|
);
|
|
|
|
CleanupAfterUnexpectedLoopExit();
|
|
OnServerLoopExited?.Invoke();
|
|
}
|
|
}
|
|
}
|
|
|
|
private static bool ShouldTreatLoopExitAsUnexpected(
|
|
bool isRunning,
|
|
CancellationToken cancellationToken)
|
|
{
|
|
// A canceled loop belongs to an intentional shutdown and must not recover a newer server.
|
|
return isRunning && !cancellationToken.IsCancellationRequested;
|
|
}
|
|
|
|
/// <summary>
|
|
/// Asynchronously accepts a client from the TcpListener.
|
|
/// </summary>
|
|
private async Task<TcpClient> AcceptTcpClientAsync(TcpListener listener, CancellationToken cancellationToken)
|
|
{
|
|
try
|
|
{
|
|
return await AcceptTcpClientAsyncCore(listener, cancellationToken);
|
|
}
|
|
catch (ThreadAbortException ex)
|
|
{
|
|
// Log and re-throw ThreadAbortException
|
|
if (!McpEditorSettings.GetIsDomainReloadInProgress())
|
|
{
|
|
OnError?.Invoke($"Unexpected thread abort: {ex.Message}");
|
|
}
|
|
throw;
|
|
}
|
|
catch (OperationCanceledException)
|
|
{
|
|
return null;
|
|
}
|
|
}
|
|
|
|
internal static Task<TcpClient> AcceptTcpClientAsyncForTests(
|
|
TcpListener listener,
|
|
CancellationToken cancellationToken)
|
|
{
|
|
return AcceptTcpClientAsyncCore(listener, cancellationToken);
|
|
}
|
|
|
|
internal int GetActiveClientTaskCountForTests()
|
|
{
|
|
return _clientTasks.Count;
|
|
}
|
|
|
|
internal static bool ShouldTreatLoopExitAsUnexpectedForTests(
|
|
bool isRunning,
|
|
CancellationToken cancellationToken)
|
|
{
|
|
return ShouldTreatLoopExitAsUnexpected(isRunning, cancellationToken);
|
|
}
|
|
|
|
internal void TrackClientTaskForTests(Task clientTask)
|
|
{
|
|
TrackClientTask(clientTask);
|
|
}
|
|
|
|
private static async Task<TcpClient> AcceptTcpClientAsyncCore(
|
|
TcpListener listener,
|
|
CancellationToken cancellationToken)
|
|
{
|
|
System.Diagnostics.Debug.Assert(listener != null, "TcpListener must be available while accepting clients.");
|
|
if (listener == null || cancellationToken.IsCancellationRequested)
|
|
{
|
|
return null;
|
|
}
|
|
|
|
using CancellationTokenRegistration cancellationRegistration =
|
|
cancellationToken.Register(() => listener.Stop());
|
|
|
|
try
|
|
{
|
|
return await listener.AcceptTcpClientAsync();
|
|
}
|
|
catch (ObjectDisposedException) when (cancellationToken.IsCancellationRequested)
|
|
{
|
|
return null;
|
|
}
|
|
catch (SocketException) when (cancellationToken.IsCancellationRequested)
|
|
{
|
|
return null;
|
|
}
|
|
catch (InvalidOperationException) when (cancellationToken.IsCancellationRequested)
|
|
{
|
|
return null;
|
|
}
|
|
}
|
|
|
|
private void TrackClientTask(Task clientTask)
|
|
{
|
|
System.Diagnostics.Debug.Assert(clientTask != null, "Client handler task must be created before tracking.");
|
|
if (clientTask == null)
|
|
{
|
|
return;
|
|
}
|
|
|
|
bool added = _clientTasks.TryAdd(clientTask.Id, clientTask);
|
|
System.Diagnostics.Debug.Assert(added, "Client handler task id should be unique while tracking.");
|
|
|
|
clientTask.ContinueWith(completedTask =>
|
|
{
|
|
_clientTasks.TryRemove(completedTask.Id, out _);
|
|
if (completedTask.IsFaulted)
|
|
{
|
|
Exception exception = completedTask.Exception?.GetBaseException();
|
|
OnError?.Invoke($"Client handler task faulted: {exception?.Message}");
|
|
}
|
|
}, CancellationToken.None, TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default);
|
|
}
|
|
|
|
/// <summary>
|
|
/// Handles communication with the client using Content-Length framing.
|
|
/// </summary>
|
|
private async Task HandleClientAsync(TcpClient client, CancellationToken cancellationToken)
|
|
{
|
|
string clientEndpoint = client.Client.RemoteEndPoint?.ToString() ?? McpServerConfig.UNKNOWN_CLIENT_ENDPOINT;
|
|
|
|
// Initialize new components for Content-Length framing
|
|
DynamicBufferManager bufferManager = null;
|
|
MessageReassembler messageReassembler = null;
|
|
|
|
try
|
|
{
|
|
using (client)
|
|
using (NetworkStream stream = client.GetStream())
|
|
{
|
|
|
|
// Check for existing connection from same endpoint and close it
|
|
string clientKey = GenerateClientKey(clientEndpoint);
|
|
if (_connectedClients.TryGetValue(clientKey, out ConnectedClient existingClient))
|
|
{
|
|
existingClient.Stream?.Close();
|
|
_connectedClients.TryRemove(clientKey, out _);
|
|
|
|
// Notify tool disconnected
|
|
OnToolDisconnected?.Invoke(existingClient.ClientName);
|
|
}
|
|
|
|
// Add new client to connected clients for notification broadcasting
|
|
int clientPort = (client.Client.RemoteEndPoint as System.Net.IPEndPoint)?.Port ?? 0;
|
|
ConnectedClient connectedClient = new ConnectedClient(clientEndpoint, stream, clientPort);
|
|
_connectedClients.TryAdd(clientKey, connectedClient);
|
|
|
|
// Initialize new framing components
|
|
bufferManager = new DynamicBufferManager();
|
|
messageReassembler = new MessageReassembler(bufferManager);
|
|
|
|
// Start with initial buffer size
|
|
byte[] buffer = bufferManager.GetBuffer(BufferConfig.INITIAL_BUFFER_SIZE);
|
|
|
|
while (!cancellationToken.IsCancellationRequested && client.Connected)
|
|
{
|
|
int bytesRead = await stream.ReadAsync(buffer, 0, buffer.Length, cancellationToken);
|
|
|
|
if (bytesRead == 0)
|
|
{
|
|
break; // Client disconnected.
|
|
}
|
|
|
|
// Add received data to message reassembler
|
|
messageReassembler.AddData(buffer, bytesRead);
|
|
|
|
// Extract any complete messages
|
|
string[] completeJsonMessages = messageReassembler.ExtractCompleteMessages();
|
|
|
|
foreach (string requestJson in completeJsonMessages)
|
|
{
|
|
if (string.IsNullOrWhiteSpace(requestJson)) continue;
|
|
|
|
// JSON-RPC processing and response sending with client context
|
|
string responseJson = await JsonRpcProcessor.ProcessRequest(requestJson, clientEndpoint);
|
|
|
|
// Only send response if it's not null (notifications return null)
|
|
if (!string.IsNullOrEmpty(responseJson))
|
|
{
|
|
// Check stream and client state before attempting write
|
|
if (!stream.CanWrite || !client.Connected || cancellationToken.IsCancellationRequested)
|
|
{
|
|
return; // Skip the write operation
|
|
}
|
|
|
|
// Send response with Content-Length framing
|
|
string framedResponse = CreateContentLengthFrame(responseJson);
|
|
byte[] responseData = Encoding.UTF8.GetBytes(framedResponse);
|
|
|
|
await stream.WriteAsync(responseData, 0, responseData.Length, cancellationToken);
|
|
}
|
|
}
|
|
|
|
// Validate reassembler state and clear if needed
|
|
if (!messageReassembler.ValidateState())
|
|
{
|
|
}
|
|
}
|
|
}
|
|
}
|
|
catch (ThreadAbortException)
|
|
{
|
|
// Treat as normal behavior if a domain reload is in progress.
|
|
// No need to log thread aborts during domain reload
|
|
}
|
|
catch (OperationCanceledException)
|
|
{
|
|
// Normal cancellation during server shutdown or domain reload
|
|
// No logging needed as this is expected behavior during Unity Editor operations
|
|
}
|
|
catch (IOException ex)
|
|
{
|
|
// I/O errors are usually normal disconnections - only log as info instead of warning
|
|
if (IsNormalDisconnectionException(ex))
|
|
{
|
|
// Log normal disconnections as info level
|
|
}
|
|
}
|
|
catch (Exception ex)
|
|
{
|
|
OnError?.Invoke(ex.Message);
|
|
}
|
|
finally
|
|
{
|
|
// Dispose of framing components
|
|
try
|
|
{
|
|
messageReassembler?.Dispose();
|
|
bufferManager?.Dispose();
|
|
}
|
|
catch (Exception ex)
|
|
{
|
|
OnError?.Invoke($"Error during client disposal: {ex.Message}");
|
|
}
|
|
|
|
// Remove client from connected clients list
|
|
// Find client by endpoint to get the correct key
|
|
ConnectedClient clientToRemove = _connectedClients.Values
|
|
.FirstOrDefault(c => c.Endpoint == clientEndpoint);
|
|
|
|
if (clientToRemove != null)
|
|
{
|
|
string clientKey = GenerateClientKey(clientToRemove.Endpoint);
|
|
_connectedClients.TryRemove(clientKey, out _);
|
|
|
|
// Notify tool disconnected
|
|
OnToolDisconnected?.Invoke(clientToRemove.ClientName);
|
|
}
|
|
|
|
client.Close();
|
|
OnClientDisconnected?.Invoke(clientEndpoint);
|
|
}
|
|
}
|
|
|
|
/// <summary>
|
|
/// Sends a pre-formatted JSON-RPC notification to all connected clients using Content-Length framing.
|
|
/// </summary>
|
|
/// <param name="notificationJson">The complete JSON-RPC notification string</param>
|
|
public void SendNotificationToClients(string notificationJson)
|
|
{
|
|
if (_connectedClients.IsEmpty)
|
|
{
|
|
return;
|
|
}
|
|
|
|
// Frame the notification with Content-Length header
|
|
string framedNotification = CreateContentLengthFrame(notificationJson);
|
|
byte[] notificationData = Encoding.UTF8.GetBytes(framedNotification);
|
|
|
|
SendNotificationDataAsync(notificationData).Forget();
|
|
}
|
|
|
|
/// <summary>
|
|
/// Sends a shutdown notification to all connected clients.
|
|
/// This should be called before disconnecting clients so TypeScript side can
|
|
/// differentiate between domain reload (temporary) and editor quit (permanent).
|
|
/// </summary>
|
|
/// <param name="reason">The reason for server shutdown</param>
|
|
public void SendShutdownNotification(ServerShutdownReason reason)
|
|
{
|
|
if (_connectedClients.IsEmpty)
|
|
{
|
|
return;
|
|
}
|
|
|
|
// Create JSON-RPC 2.0 notification with shutdown reason
|
|
string notificationJson = $"{{\"jsonrpc\":\"2.0\",\"method\":\"notifications/server/shutdown\",\"params\":{{\"reason\":\"{reason}\"}}}}";
|
|
|
|
// Frame the notification with Content-Length header
|
|
string framedNotification = CreateContentLengthFrame(notificationJson);
|
|
byte[] notificationData = Encoding.UTF8.GetBytes(framedNotification);
|
|
|
|
// Send synchronously to ensure it's sent before connection closes
|
|
SendNotificationDataSync(notificationData);
|
|
}
|
|
|
|
/// <summary>
|
|
/// Send notification data to all connected clients synchronously.
|
|
/// Used for shutdown notifications to ensure delivery before connection closes.
|
|
/// </summary>
|
|
private void SendNotificationDataSync(byte[] notificationData)
|
|
{
|
|
foreach (KeyValuePair<string, ConnectedClient> client in _connectedClients)
|
|
{
|
|
try
|
|
{
|
|
if (client.Value.Stream?.CanWrite == true)
|
|
{
|
|
client.Value.Stream.Write(notificationData, 0, notificationData.Length);
|
|
client.Value.Stream.Flush();
|
|
}
|
|
}
|
|
catch (Exception)
|
|
{
|
|
// Ignore errors during shutdown notification - client may already be disconnected
|
|
}
|
|
}
|
|
}
|
|
|
|
/// <summary>
|
|
/// Send notification data to all connected clients
|
|
/// </summary>
|
|
private async Task SendNotificationDataAsync(byte[] notificationData)
|
|
{
|
|
List<string> clientsToRemove = new List<string>();
|
|
|
|
foreach (KeyValuePair<string, ConnectedClient> client in _connectedClients)
|
|
{
|
|
try
|
|
{
|
|
if (client.Value.Stream?.CanWrite == true)
|
|
{
|
|
await client.Value.Stream.WriteAsync(notificationData, 0, notificationData.Length);
|
|
}
|
|
else
|
|
{
|
|
clientsToRemove.Add(client.Key);
|
|
}
|
|
}
|
|
catch (Exception ex)
|
|
{
|
|
OnError?.Invoke($"Error writing notification to client {client.Key}: {ex.Message}");
|
|
clientsToRemove.Add(client.Key);
|
|
}
|
|
}
|
|
|
|
// Remove disconnected clients
|
|
foreach (string clientKey in clientsToRemove)
|
|
{
|
|
if (_connectedClients.TryRemove(clientKey, out ConnectedClient removedClient))
|
|
{
|
|
// Notify tool disconnected
|
|
OnToolDisconnected?.Invoke(removedClient.ClientName);
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
/// <summary>
|
|
/// Creates a Content-Length framed message for JSON-RPC 2.0 communication.
|
|
/// </summary>
|
|
/// <param name="jsonContent">The JSON content to frame</param>
|
|
/// <returns>The framed message with Content-Length header</returns>
|
|
private string CreateContentLengthFrame(string jsonContent)
|
|
{
|
|
if (string.IsNullOrEmpty(jsonContent))
|
|
{
|
|
return string.Empty;
|
|
}
|
|
|
|
// Calculate content length in bytes (UTF-8 encoding)
|
|
int contentLength = Encoding.UTF8.GetByteCount(jsonContent);
|
|
|
|
// Create the framed message: Content-Length: <n>\r\n\r\n<json_content>
|
|
return $"Content-Length: {contentLength}\r\n\r\n{jsonContent}";
|
|
}
|
|
|
|
/// <summary>
|
|
/// Determines if the given exception represents a normal client disconnection.
|
|
/// </summary>
|
|
/// <param name="ex">The exception to evaluate</param>
|
|
/// <returns>True if the exception represents a normal disconnection, false otherwise</returns>
|
|
private static bool IsNormalDisconnectionException(Exception ex)
|
|
{
|
|
switch (ex)
|
|
{
|
|
case SocketException sockEx:
|
|
return sockEx.SocketErrorCode is SocketError.ConnectionReset or
|
|
SocketError.ConnectionAborted or
|
|
SocketError.OperationAborted or
|
|
SocketError.Shutdown or
|
|
SocketError.NotConnected;
|
|
|
|
case ObjectDisposedException:
|
|
return true;
|
|
|
|
case IOException ioEx when ioEx.InnerException is SocketException innerSockEx:
|
|
return innerSockEx.SocketErrorCode is SocketError.ConnectionReset or
|
|
SocketError.ConnectionAborted or
|
|
SocketError.OperationAborted or
|
|
SocketError.Shutdown or
|
|
SocketError.NotConnected;
|
|
|
|
case IOException ioEx:
|
|
// Check HResult codes for common disconnection scenarios
|
|
return NormalDisconnectionHResults.Contains(ioEx.HResult) ||
|
|
IsNormalDisconnectionByInnerException(ioEx);
|
|
|
|
default:
|
|
return false;
|
|
}
|
|
}
|
|
|
|
/// <summary>
|
|
/// Recursively checks inner exceptions for normal disconnection scenarios
|
|
/// </summary>
|
|
/// <param name="ex">The exception to check</param>
|
|
/// <returns>True if any inner exception indicates a normal disconnection</returns>
|
|
private static bool IsNormalDisconnectionByInnerException(Exception ex)
|
|
{
|
|
Exception innerEx = ex.InnerException;
|
|
while (innerEx != null)
|
|
{
|
|
if (IsNormalDisconnectionException(innerEx))
|
|
{
|
|
return true;
|
|
}
|
|
innerEx = innerEx.InnerException;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
/// <summary>
|
|
/// Releases resources.
|
|
/// </summary>
|
|
public void Dispose()
|
|
{
|
|
StopServer();
|
|
}
|
|
|
|
private enum ServerStopMode
|
|
{
|
|
Normal,
|
|
BeforeDomainReload,
|
|
UnexpectedLoopExit
|
|
}
|
|
}
|
|
}
|