VolcanicAgenticCore/extensions/Volcanic.AgenticCore.Extensions.Mcp/McpServerConnection.cs

using System.Globalization;
using System.Text.Json;
using Microsoft.Extensions.Logging;
using ModelContextProtocol;
using ModelContextProtocol.Client;
using ModelContextProtocol.Protocol;
using Volcanic.AgenticCore.Abstractions.Extensions;
using Volcanic.AgenticCore.Abstractions.Processes;

namespace Volcanic.AgenticCore.Extensions.Mcp;

internal enum McpConnectionStatus { NotStarted, Running, Failed }

/// <summary>One server process for one session: started on the first call, stopped on session shutdown or when idle.</summary>
internal sealed class McpServerConnection : IAsyncDisposable
{
    private static readonly TimeSpan ExitGrace = TimeSpan.FromSeconds(2);
    private static readonly TimeSpan StderrGrace = TimeSpan.FromMilliseconds(150);
    private static readonly TimeSpan CatalogCheckCap = TimeSpan.FromSeconds(10);

    private readonly McpServerConfig _config;
    private readonly ILogger _logger;
    private readonly Func<string, string?> _getEnv;
    private readonly McpCatalog? _catalog;
    private readonly TimeProvider _time;
    private readonly SemaphoreSlim _gate = new(1, 1);
    private readonly Func<string, string> _redact;
    private readonly StderrTail _stderr;
    private readonly Lock _lock = new();
    private McpProcess? _process;
    private ITimer? _idleTimer;
    private long _timerGeneration;
    private int _inFlight;
    private int _disposed;
    private long _nextRequestId;

    /// <param name="redact">Applied to the server's stderr as it arrives and to logged errors; identity when null.</param>
    public McpServerConnection(string name, McpServerConfig config, ILogger logger, Func<string, string?> getEnv, McpCatalog? catalog = null, TimeProvider? time = null, Func<string, string>? redact = null)
    {
        Name = name;
        _redact = redact ?? (static text => text);
        _stderr = new StderrTail(redact);
        _config = config;
        _logger = logger;
        _getEnv = getEnv;
        _catalog = catalog;
        _time = time ?? TimeProvider.System;
    }

    public string Name { get; }
    public McpConnectionStatus Status { get; private set; }
    public string? LastError { get; private set; }
    public int StartCount { get; private set; }
    public IReadOnlyList<string> StderrTail => _stderr.Lines;
    /// <summary>True when the current process has ended but no call has noticed yet (the status still says running).</summary>
    public bool ProcessExited => Volatile.Read(ref _process)?.HasExited == true;

    private bool Disposed => Volatile.Read(ref _disposed) == 1;

    /// <summary>
    /// Calls <paramref name="mcpName"/> on the session's process, starting it first if needed. Throws
    /// <see cref="TimeoutException"/>, <see cref="OperationCanceledException"/>, <see cref="McpServerExitedException"/>
    /// or <see cref="McpServerStartException"/>; a JSON-RPC error from the server surfaces as the SDK's exception.
    /// </summary>
    public async Task<CallToolResult> CallAsync(string mcpName, JsonElement arguments, IExtensionContext context, CancellationToken cancellationToken)
    {
        var process = await EnsureStartedAsync(context, cancellationToken).ConfigureAwait(false);
        try
        {
            using var timeout = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
            timeout.CancelAfter(TimeSpan.FromSeconds(_config.CallTimeoutSeconds));
            var request = new CallToolRequestParams { Name = mcpName, Arguments = ToArguments(arguments) };
            var id = new RequestId($"volcanic-{Interlocked.Increment(ref _nextRequestId).ToString(CultureInfo.InvariantCulture)}");
            try
            {
                return await process.Client.SendRequestAsync<CallToolRequestParams, CallToolResult>(
                    RequestMethods.ToolsCall, request, McpJsonUtilities.DefaultOptions, id, timeout.Token).ConfigureAwait(false);
            }
            catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
            {
                await SendCancelledAsync(process, id, "aborted").ConfigureAwait(false);
                throw;
            }
            catch (OperationCanceledException) when (timeout.IsCancellationRequested)
            {
                await SendCancelledAsync(process, id, "timed out").ConfigureAwait(false);
                throw new TimeoutException($"{mcpName} timed out after {_config.CallTimeoutSeconds}s");
            }
            catch (Exception ex) when (ex is not McpProtocolException)
            {
                if (await process.WaitForExitAsync(ExitGrace).ConfigureAwait(false) is { } exit)
                    throw await OnExitedAsync(process, exit).ConfigureAwait(false);
                throw;
            }
        }
        finally
        {
            EndCall();
        }
    }

    /// <summary>Stops the process; the next call starts a new one. Idempotent.</summary>
    public async ValueTask StopAsync()
    {
        await _gate.WaitAsync().ConfigureAwait(false);
        try { await StopCoreAsync().ConfigureAwait(false); }
        finally { _gate.Release(); }
    }

    /// <summary>Stops the process and refuses later calls with <see cref="ObjectDisposedException"/>. Idempotent.</summary>
    public async ValueTask DisposeAsync()
    {
        if (Interlocked.Exchange(ref _disposed, 1) == 1) return;
        await StopAsync().ConfigureAwait(false);
    }

    /// <summary>Caller holds <see cref="_gate"/>. Resets a failed status too; <see cref="LastError"/> stays for <c>/mcp &lt;name&gt;</c>.</summary>
    private async ValueTask StopCoreAsync()
    {
        lock (_lock) { _idleTimer?.Dispose(); _idleTimer = null; _timerGeneration++; }
        if (Interlocked.Exchange(ref _process, null) is { } process) await process.DisposeAsync().ConfigureAwait(false);
        Status = McpConnectionStatus.NotStarted;
    }

    /// <summary>Returns a live process with the call counted as in flight (so the idle timer can't stop it underneath).</summary>
    private async Task<McpProcess> EnsureStartedAsync(IExtensionContext context, CancellationToken cancellationToken)
    {
        await _gate.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            ObjectDisposedException.ThrowIf(Disposed, this);
            if (_process is { } existing)
            {
                if (existing.HasExited) throw await OnExitedAsync(existing, await existing.Completion.ConfigureAwait(false)).ConfigureAwait(false);
                BeginCall();
                return existing;
            }

            ProcessSpawnContext spawn;
            try
            {
                spawn = await context.PrepareProcessAsync(McpLaunch.BuildRequest(Name, _config, context.Cwd, _getEnv), cancellationToken).ConfigureAwait(false);
            }
            catch (Exception ex) when (ex is not OperationCanceledException)
            {
                throw Fail(new McpServerStartException($"{Name}: process environment hook failed: {ex.Message}", ex));
            }

            _stderr.Clear();
            McpProcess process;
            try { process = await McpProcess.StartAsync(Name, _config, spawn, _stderr, _logger, cancellationToken).ConfigureAwait(false); }
            catch (McpServerStartException ex) { throw Fail(ex); }

            _process = process;
            StartCount++;
            Status = McpConnectionStatus.Running;
            LastError = null;
            _logger.LogInformation("MCP server {Name} started in {Cwd}", Name, spawn.Cwd);
            _ = Task.Run(() => WarnIfCatalogChangedAsync(process), CancellationToken.None); // only logs; never holds the gate
            BeginCall();
            return process;
        }
        finally
        {
            _gate.Release();
        }
    }

    private McpServerStartException Fail(McpServerStartException ex)
    {
        Status = McpConnectionStatus.Failed;
        LastError = ex.Message;
        _logger.LogWarning("{Message}", _redact(ex.Message));
        return ex;
    }

    private async Task<McpServerExitedException> OnExitedAsync(McpProcess process, ClientCompletionDetails? details)
    {
        Interlocked.CompareExchange(ref _process, null, process);
        await process.DisposeAsync().ConfigureAwait(false);
        await Task.Delay(StderrGrace).ConfigureAwait(false); // the last stderr lines race the exit
        var code = (details as StdioClientCompletionDetails)?.ExitCode?.ToString(CultureInfo.InvariantCulture) ?? "?";
        var tail = _stderr.Lines is { Count: > 0 } lines ? string.Join('\n', lines) : "no stderr output";
        var message = $"{Name} MCP server exited ({code}): {tail}. It restarts on the next call; browser state was lost.";
        Status = McpConnectionStatus.Failed;
        LastError = message;
        _logger.LogWarning("{Message}", message);
        return new McpServerExitedException(message);
    }

    private void BeginCall()
    {
        lock (_lock)
        {
            _inFlight++;
            _timerGeneration++;
            _idleTimer?.Dispose();
            _idleTimer = null;
        }
    }

    private void EndCall()
    {
        lock (_lock)
        {
            if (--_inFlight > 0 || _config.IdleTimeoutSeconds == 0 || Disposed) return;
            _idleTimer?.Dispose();
            var generation = ++_timerGeneration;
            _idleTimer = _time.CreateTimer(_ => _ = StopIfIdleAsync(generation), null, TimeSpan.FromSeconds(_config.IdleTimeoutSeconds), Timeout.InfiniteTimeSpan);
        }
    }

    /// <summary>A callback that fired before a later call reset the timer finds a newer generation and does nothing.</summary>
    private async Task StopIfIdleAsync(long generation)
    {
        try
        {
            await _gate.WaitAsync().ConfigureAwait(false);
            try
            {
                lock (_lock) { if (_inFlight > 0 || _timerGeneration != generation) return; }
                if (_process is null) return;
                _logger.LogInformation("MCP server {Name} was idle for {Seconds}s; stopping it", Name, _config.IdleTimeoutSeconds);
                await StopCoreAsync().ConfigureAwait(false);
            }
            finally
            {
                _gate.Release();
            }
        }
        catch (Exception ex) when (ex is not OutOfMemoryException)
        {
            _logger.LogWarning("MCP server {Name}: idle stop failed: {Message}", Name, ex.Message);
        }
    }

    private async Task WarnIfCatalogChangedAsync(McpProcess process)
    {
        if (_catalog is null) return;
        try
        {
            var cap = TimeSpan.FromSeconds(_config.StartupTimeoutSeconds);
            using var cts = new CancellationTokenSource(cap < CatalogCheckCap ? cap : CatalogCheckCap);
            var live = await process.Client.ListToolsAsync(cancellationToken: cts.Token).ConfigureAwait(false);
            var liveNames = live.Select(t => t.Name).Order(StringComparer.Ordinal);
            if (!liveNames.SequenceEqual(_catalog.Tools.Select(t => t.McpName).Order(StringComparer.Ordinal), StringComparer.Ordinal))
                _logger.LogWarning("MCP server {Name}: its tools changed since discovery; calls still go by name. Restart the host to pick up the new list.", Name);
        }
        catch (Exception ex) when (ex is not OutOfMemoryException)
        {
            _logger.LogDebug("MCP server {Name}: could not compare its tools with the catalogue: {Message}", Name, ex.Message);
        }
    }

    /// <summary>
    /// The SDK (2.2.0) doesn't send <c>notifications/cancelled</c> when a request's token fires (checked against the fake
    /// server), so the call carries its own id and the notification goes out here. Best effort: a dead server can't take it.
    /// </summary>
    private async Task SendCancelledAsync(McpProcess process, RequestId id, string reason)
    {
        try
        {
            using var cts = new CancellationTokenSource(ExitGrace);
            await process.Client.SendNotificationAsync(NotificationMethods.CancelledNotification,
                new CancelledNotificationParams { RequestId = id, Reason = reason }, McpJsonUtilities.DefaultOptions, cts.Token).ConfigureAwait(false);
        }
        catch (Exception ex) when (ex is not OutOfMemoryException)
        {
            _logger.LogDebug("MCP server {Name}: could not send notifications/cancelled: {Message}", Name, ex.Message);
        }
    }

    private static Dictionary<string, JsonElement> ToArguments(JsonElement arguments)
    {
        var result = new Dictionary<string, JsonElement>(StringComparer.Ordinal);
        if (arguments.ValueKind == JsonValueKind.Object)
            foreach (var property in arguments.EnumerateObject()) result[property.Name] = property.Value.Clone();
        return result;
    }
}