VolcanicAgenticCore/extensions/Volcanic.AgenticCore.Extensions.Mcp/Remote/RemoteServerConnection.cs

using System.Globalization;
using System.Text.Json;
using ModelContextProtocol;
using ModelContextProtocol.Client;
using ModelContextProtocol.Protocol;

namespace Volcanic.AgenticCore.Extensions.Mcp.Remote;

/// <summary>
/// One remote server for one chat session: an SDK client over HTTP (Streamable HTTP, or SSE). Connected by
/// <see cref="ConnectAsync"/>; a call that fails in the transport drops the client and the next call connects again.
/// </summary>
internal sealed class RemoteServerConnection : IAsyncDisposable
{
    private static readonly HttpClient SharedHttp = new(new SocketsHttpHandler { PooledConnectionLifetime = TimeSpan.FromMinutes(5) })
    {
        Timeout = Timeout.InfiniteTimeSpan, // Streamable HTTP keeps a GET stream open; every request has its own deadline
    };

    private readonly RemoteMcpOptions _options;
    private readonly SemaphoreSlim _gate = new(1, 1);
    private McpClient? _client;
    private HttpClient? _ownHttp;
    private int _disposed;
    private long _nextRequestId;

    public RemoteServerConnection(RemoteServerEntry entry, RemoteMcpOptions options)
    {
        Entry = entry;
        _options = options;
    }

    public RemoteServerEntry Entry { get; }
    public string Name => Entry.Name;
    public string? ServerInfoName { get; private set; }
    public string? ServerVersion { get; private set; }

    /// <summary>Connects (if needed) and lists the tools. Throws <see cref="RemoteServerException"/> with a readable message.</summary>
    public async Task<IReadOnlyList<Tool>> ConnectAsync(CancellationToken cancellationToken)
    {
        await _gate.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            ObjectDisposedException.ThrowIf(Volatile.Read(ref _disposed) == 1, this);
            using var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
            cts.CancelAfter(_options.ConnectTimeout);
            try
            {
                var client = await EnsureClientAsync(cts.Token).ConfigureAwait(false);
                var tools = await client.ListToolsAsync(cancellationToken: cts.Token).ConfigureAwait(false);
                return [.. tools.Select(t => t.ProtocolTool)];
            }
            catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested)
            {
                await DropClientAsync().ConfigureAwait(false);
                throw new RemoteServerException($"{Name}: no answer from {RemoteServersFile.DisplayUrl(Entry.Url)} within {_options.ConnectTimeout.TotalSeconds:0}s");
            }
            catch (Exception ex) when (ex is not OperationCanceledException and not ObjectDisposedException)
            {
                await DropClientAsync().ConfigureAwait(false);
                throw new RemoteServerException($"{Name}: cannot connect to {RemoteServersFile.DisplayUrl(Entry.Url)}: {Describe(ex)}", ex);
            }
        }
        finally
        {
            _gate.Release();
        }
    }

    /// <summary>
    /// Calls a tool, connecting first if the last call dropped the client. Throws <see cref="TimeoutException"/>,
    /// <see cref="OperationCanceledException"/>, <see cref="RemoteServerException"/>, or the SDK's protocol exception.
    /// </summary>
    public async Task<CallToolResult> CallAsync(string mcpName, JsonElement arguments, CancellationToken cancellationToken)
    {
        McpClient client;
        await _gate.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            ObjectDisposedException.ThrowIf(Volatile.Read(ref _disposed) == 1, this);
            using var connect = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
            connect.CancelAfter(_options.ConnectTimeout);
            try
            {
                client = await EnsureClientAsync(connect.Token).ConfigureAwait(false);
            }
            catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested)
            {
                throw new RemoteServerException($"{Name} MCP server did not answer within {_options.ConnectTimeout.TotalSeconds:0}s; try again later.");
            }
            catch (Exception ex) when (ex is not OperationCanceledException)
            {
                throw new RemoteServerException($"{Name} MCP server is unreachable: {Describe(ex)}", ex);
            }
        }
        finally
        {
            _gate.Release();
        }

        using var timeout = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
        timeout.CancelAfter(_options.CallTimeout);
        var id = new RequestId($"volcanic-{Interlocked.Increment(ref _nextRequestId).ToString(CultureInfo.InvariantCulture)}");
        var request = new CallToolRequestParams { Name = mcpName, Arguments = ToArguments(arguments) };
        try
        {
            return await client.SendRequestAsync<CallToolRequestParams, CallToolResult>(
                RequestMethods.ToolsCall, request, McpJsonUtilities.DefaultOptions, id, timeout.Token).ConfigureAwait(false);
        }
        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
        {
            await SendCancelledAsync(client, id, "aborted").ConfigureAwait(false);
            throw;
        }
        catch (OperationCanceledException) when (timeout.IsCancellationRequested)
        {
            await SendCancelledAsync(client, id, "timed out").ConfigureAwait(false);
            throw new TimeoutException($"{mcpName} timed out after {_options.CallTimeout.TotalSeconds:0}s");
        }
        catch (Exception ex) when (ex is not McpProtocolException)
        {
            await DropClientAsync(client).ConfigureAwait(false);
            throw new RemoteServerException($"{Name} MCP server failed: {Describe(ex)}. The next call reconnects; the server may or may not have run this one.", ex);
        }
    }

    public async ValueTask DisposeAsync()
    {
        if (Interlocked.Exchange(ref _disposed, 1) == 1) return;
        await _gate.WaitAsync().ConfigureAwait(false);
        try { await DropClientAsync().ConfigureAwait(false); }
        finally { _gate.Release(); }
    }

    /// <summary>Caller holds the gate.</summary>
    private async Task<McpClient> EnsureClientAsync(CancellationToken cancellationToken)
    {
        if (_client is { } existing) return existing;
        var http = SharedHttp;
        if (_options.HttpHandlerFactory is { } factory) http = _ownHttp ??= new HttpClient(factory()) { Timeout = Timeout.InfiniteTimeSpan };
        var transport = new HttpClientTransport(new HttpClientTransportOptions
        {
            Name = Name,
            Endpoint = new Uri(Entry.Url),
            TransportMode = Entry.Transport switch
            {
                RemoteTransport.Http => HttpTransportMode.StreamableHttp,
                RemoteTransport.Sse => HttpTransportMode.Sse,
                _ => HttpTransportMode.AutoDetect,
            },
            AdditionalHeaders = new Dictionary<string, string>(Entry.Headers, StringComparer.OrdinalIgnoreCase),
            ConnectionTimeout = _options.ConnectTimeout,
        }, http, loggerFactory: null, ownsHttpClient: false);
        var client = await McpClient.CreateAsync(
            transport,
            new McpClientOptions { ClientInfo = new Implementation { Name = "volcanic-agentic-core", Version = "1.0.0" }, InitializationTimeout = _options.ConnectTimeout },
            loggerFactory: null,
            cancellationToken).ConfigureAwait(false);
        ServerInfoName = client.ServerInfo?.Name;
        ServerVersion = client.ServerInfo?.Version;
        _client = client;
        return client;
    }

    /// <summary>Drops the current client (or only <paramref name="expected"/>, if it is still the current one).</summary>
    private async Task DropClientAsync(McpClient? expected = null)
    {
        var client = expected is null ? Interlocked.Exchange(ref _client, null) : Interlocked.CompareExchange(ref _client, null, expected) == expected ? expected : null;
        if (client is null) return;
        try { await client.DisposeAsync().ConfigureAwait(false); }
        catch (Exception) { /* a broken transport may throw on dispose */ }
    }

    private static async Task SendCancelledAsync(McpClient client, RequestId id, string reason)
    {
        try
        {
            using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(2));
            await client.SendNotificationAsync(NotificationMethods.CancelledNotification,
                new CancelledNotificationParams { RequestId = id, Reason = reason }, McpJsonUtilities.DefaultOptions, cts.Token).ConfigureAwait(false);
        }
        catch (Exception ex) when (ex is not OutOfMemoryException) { /* best effort */ }
    }

    /// <summary>The innermost useful message: "Connection refused", "401 (Unauthorized)" and the like.</summary>
    internal static string Describe(Exception ex)
    {
        var message = ex.Message;
        for (var inner = ex.InnerException; inner is not null; inner = inner.InnerException)
            if (inner is HttpRequestException or System.Net.Sockets.SocketException) message = $"{message} ({inner.Message})";
        return 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;
    }
}

/// <summary>A remote server could not be reached or broke off; the message is meant for the agent and the owner.</summary>
internal sealed class RemoteServerException(string message, Exception? inner = null) : Exception(message, inner);