Playa/extensions/Volcanic.AgenticCore.Extensions.Mcp/Remote/RemoteMcpExtension.cs

using System.Text;
using Microsoft.Extensions.Logging;
using ModelContextProtocol.Protocol;
using Volcanic.AgenticCore.Abstractions.Extensions;
using Volcanic.AgenticCore.Abstractions.Extensions.Events;
using Volcanic.AgenticCore.Abstractions.Messages;

namespace Volcanic.AgenticCore.Extensions.Mcp.Remote;

/// <summary>
/// Remote MCP servers (Streamable HTTP or SSE) per chat, from the workspace's <c>.volcanic/mcp/servers.json</c>. Only the
/// owner's commands (<c>/mcp_add</c>, <c>/mcp_remove</c>, <c>/mcp_tools</c>, <c>/mcp_list</c>) change the servers; the
/// agent can't, because entries are sealed with a host secret. Tools appear as <c>&lt;server&gt;__&lt;tool&gt;</c>. A server
/// that can't be reached leaves the session without its tools, and the agent's system prompt says why. Instantiated
/// per session.
/// </summary>
public sealed class RemoteMcpExtension(RemoteMcpOptions options) : IExtension
{
    public const string ExtensionName = "mcp-remote";
    public const string ToolSeparator = "__";
    internal const string SectionKey = "mcp-remote";

    private readonly RemoteMcpOptions _options = options ?? throw new ArgumentNullException(nameof(options));
    private readonly TimeProvider _time = options.Time ?? TimeProvider.System;
    private readonly Lock _lock = new();
    private readonly SemaphoreSlim _sync = new(1, 1);
    private readonly Dictionary<string, ServerState> _servers = new(StringComparer.Ordinal);
    private readonly Dictionary<string, RemoteMcpTool> _tools = new(StringComparer.Ordinal);
    private IReadOnlyList<string> _fileProblems = [];
    private IExtensionApi? _api;
    private ILogger _logger = Microsoft.Extensions.Logging.Abstractions.NullLogger.Instance;
    private Func<string, string> _hostRedact = static text => text;
    private volatile RemoteServerCipher? _cipher;
    private volatile string? _cwd;

    public string Name => ExtensionName;

    internal RemoteMcpOptions Options => _options;

    public ValueTask ActivateAsync(IExtensionApi api, CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(api);
        _api = api;
        _logger = api.Logger;
        _hostRedact = McpExtension.RedactorOf(api);
        // Always registered: an unknown /mcp_add with a key in it would otherwise reach the model as a prompt.
        RemoteMcpCommands.Register(api, this);
        if (!_options.IsConfigured)
        {
            api.Logger.LogInformation("Remote MCP servers disabled: no encryption secret (set MCP_ENCRYPTION_SECRET)");
            return ValueTask.CompletedTask;
        }
        api.On<SessionStartEvent>(async (_, ctx, _) =>
        {
            Track(ctx);
            await SyncAsync(wait: false).ConfigureAwait(false); // connects in the background; the first run waits for it
        });
        api.On<BeforeAgentStartEvent>(async (e, ctx, _) =>
        {
            Track(ctx);
            await SyncAsync(wait: true).ConfigureAwait(false);
            if (PromptSection() is { } section) e.SystemPromptOptions.Sections[SectionKey] = section;
        });
        api.On<ToolResultEvent>((e, _, _) =>
        {
            for (var i = 0; i < e.Content.Count; i++)
                if (e.Content[i] is TextContent text && Redact(text.Text) is var redacted && redacted != text.Text)
                    e.Content[i] = new TextContent(redacted);
            return ValueTask.CompletedTask;
        });
        api.On<SessionShutdownEvent>(async (e, _, _) =>
        {
            if (e.Reason is not (SessionShutdownReason.Quit or SessionShutdownReason.Reload)) return; // /new keeps the connections
            List<ServerState> all;
            lock (_lock) { all = [.. _servers.Values]; _servers.Clear(); }
            await Task.WhenAll(all.Select(s => s.Connection.DisposeAsync().AsTask())).ConfigureAwait(false);
        });
        return ValueTask.CompletedTask;
    }

    // ---- State for the commands ----

    internal bool IsConfigured => _options.IsConfigured;

    internal RemoteServerCipher CipherFor(IExtensionContext ctx)
    {
        Track(ctx);
        return _cipher!;
    }

    internal sealed record ServerStatus(RemoteServerEntry Entry, string State, IReadOnlyList<string> ExposedTools, IReadOnlyList<Tool>? AllTools, string? ServerInfo);

    internal IReadOnlyList<ServerStatus> Statuses()
    {
        lock (_lock)
            return [.. _servers.Values.OrderBy(s => s.Entry.Name, StringComparer.Ordinal).Select(s => new ServerStatus(
                s.Entry,
                s.Connecting is not null && s.AllTools is null ? "connecting"
                    : s.Error is { } error ? $"unavailable: {error}"
                    : $"connected · {s.Exposed.Count} tools",
                [.. s.Exposed],
                s.AllTools,
                s.Connection.ServerInfoName is { } info ? $"{info} {s.Connection.ServerVersion}".Trim() : null))];
    }

    internal IReadOnlyList<string> FileProblems
    {
        get { lock (_lock) return _fileProblems; }
    }

    /// <summary>The host's secrets and every header value of this chat's servers become <c>[REDACTED]</c>.</summary>
    internal string Redact(string text)
    {
        List<string> secrets;
        lock (_lock) secrets = [.. _servers.Values.SelectMany(s => SecretsOf(s.Entry))];
        var result = text;
        foreach (var secret in secrets.Distinct(StringComparer.Ordinal).OrderByDescending(s => s.Length))
            result = result.Replace(secret, "[REDACTED]", StringComparison.Ordinal);
        return _hostRedact(result);
    }

    /// <summary>A header value, and its credential alone when it has a scheme ("Bearer abc…" → "abc…"). Short values are left alone.</summary>
    internal static IEnumerable<string> SecretsOf(RemoteServerEntry entry)
    {
        foreach (var value in entry.Headers.Values)
        {
            if (value.Length >= 6) yield return value;
            var space = value.IndexOf(' ', StringComparison.Ordinal);
            if (space > 0 && value[(space + 1)..].Trim() is { Length: >= 6 } credential) yield return credential;
        }
    }

    /// <summary>
    /// Brings the session in line with <c>servers.json</c>: drops removed servers, connects new or changed ones, retries
    /// failed ones whose delay has passed. With <paramref name="wait"/> it returns once every connect in flight has ended
    /// (each is bounded by <see cref="RemoteMcpOptions.ConnectTimeout"/>).
    /// </summary>
    internal async Task SyncAsync(bool wait)
    {
        if (_cipher is not { } cipher || _cwd is not { } cwd) return;
        var pending = new List<Task>();
        await _sync.WaitAsync().ConfigureAwait(false);
        try
        {
            var read = RemoteServersFile.Read(cwd, cipher);
            var problems = read.Error is { } error ? [error, .. read.Problems] : read.Problems;
            List<ServerState> removed = [];
            lock (_lock)
            {
                foreach (var problem in problems.Except(_fileProblems))
                    _logger.LogWarning("Remote MCP: {Problem}", problem);
                _fileProblems = problems;
                var wanted = read.Servers.ToDictionary(s => s.Name, StringComparer.Ordinal);
                foreach (var state in _servers.Values.ToList())
                {
                    if (wanted.TryGetValue(state.Entry.Name, out var entry) && entry.Digest() == state.Digest) continue;
                    removed.Add(state);
                    RemoveLocked(state);
                }
                foreach (var entry in read.Servers)
                {
                    if (!_servers.TryGetValue(entry.Name, out var state))
                    {
                        state = new ServerState(entry, new RemoteServerConnection(entry, _options));
                        _servers[entry.Name] = state;
                    }
                    if (state.Connecting is null && state.AllTools is null && (state.RetryAt is null || _time.GetUtcNow() >= state.RetryAt))
                        state.Connecting = Task.Run(() => ConnectAsync(state));
                    if (state.Connecting is { } connecting) pending.Add(connecting);
                }
            }
            foreach (var state in removed) _ = state.Connection.DisposeAsync().AsTask();
        }
        finally
        {
            _sync.Release();
        }
        if (wait) await Task.WhenAll(pending).ConfigureAwait(false);
    }

    private async Task ConnectAsync(ServerState state)
    {
        try
        {
            var tools = await state.Connection.ConnectAsync(CancellationToken.None).ConfigureAwait(false);
            lock (_lock)
            {
                if (!IsCurrent(state)) return;
                state.AllTools = tools;
                state.Error = null;
                state.Failures = 0;
                state.RetryAt = null;
                ApplyToolsLocked(state);
            }
            _logger.LogInformation("Remote MCP server {Name}: {Count} tools ({Server} {Version})", state.Entry.Name, state.Exposed.Count,
                state.Connection.ServerInfoName, state.Connection.ServerVersion);
        }
        catch (Exception ex) when (ex is RemoteServerException or ObjectDisposedException)
        {
            if (ex is ObjectDisposedException) return; // removed while connecting
            var error = Redact(ex.Message);
            lock (_lock)
            {
                if (!IsCurrent(state)) return;
                state.Failures++;
                state.Error = error;
                var delay = TimeSpan.FromTicks(Math.Min(_options.MaxRetryDelay.Ticks, _options.RetryDelay.Ticks * (1L << Math.Min(state.Failures - 1, 10))));
                state.RetryAt = _time.GetUtcNow() + delay;
            }
            _logger.LogWarning("Remote MCP server {Name} is unavailable: {Error}", state.Entry.Name, state.Error);
        }
        finally
        {
            lock (_lock) state.Connecting = null;
        }
    }

    private bool IsCurrent(ServerState state) => _servers.TryGetValue(state.Entry.Name, out var current) && ReferenceEquals(current, state);

    /// <summary>Registers (or re-binds) the allowed tools as <c>&lt;server&gt;__&lt;tool&gt;</c> and activates them.</summary>
    private void ApplyToolsLocked(ServerState state)
    {
        var api = _api!;
        var allowed = state.Entry.Tools is { } list ? list.ToHashSet(StringComparer.Ordinal) : null;
        var tools = state.AllTools!.Where(t => allowed is null || allowed.Contains(t.Name)).ToList();
        var exposed = ToolNaming.Expose([.. tools.Select(t => t.Name)], state.Entry.Name + ToolSeparator);
        var taken = api.GetAllTools().Where(t => t.Source != ExtensionName).Select(t => t.Name).ToHashSet(StringComparer.Ordinal);
        state.Exposed.Clear();
        for (var i = 0; i < tools.Count; i++)
        {
            var name = exposed[i];
            if (taken.Contains(name))
            {
                _logger.LogWarning("Remote MCP server {Server}: tool {Tool} clashes with an existing tool and is skipped", state.Entry.Name, name);
                continue;
            }
            if (_tools.TryGetValue(name, out var tool) && tool.Server is { } owner && owner != state.Entry.Name)
            {
                _logger.LogWarning("Remote MCP tool {Tool} is already offered by {Owner}; {Server}'s is skipped", name, owner, state.Entry.Name);
                continue;
            }
            if (tool is null)
            {
                tool = new RemoteMcpTool(name);
                tool.Bind(tools[i], state.Connection);
                api.RegisterTool(tool);
                _tools[name] = tool;
            }
            else tool.Bind(tools[i], state.Connection);
            state.Exposed.Add(name);
        }
        api.SetActiveTools(api.GetActiveTools().Union(state.Exposed, StringComparer.Ordinal));
    }

    private void RemoveLocked(ServerState state)
    {
        _servers.Remove(state.Entry.Name);
        if (state.Exposed.Count == 0) return;
        var gone = state.Exposed.ToHashSet(StringComparer.Ordinal);
        foreach (var name in gone) _tools[name].Unbind();
        _api!.SetActiveTools(_api.GetActiveTools().Where(n => !gone.Contains(n)));
    }

    /// <summary>What the agent is told: connected servers in one line each, and why the others are missing.</summary>
    internal string? PromptSection()
    {
        IReadOnlyList<ServerStatus> statuses = Statuses();
        var problems = FileProblems;
        if (statuses.Count == 0 && problems.Count == 0) return null;
        var sb = new StringBuilder("## Remote MCP servers\n");
        sb.Append("The owner manages this chat's MCP servers with /mcp_list, /mcp_add, /mcp_remove and /mcp_tools; you can't add or change them (never edit .volcanic/mcp/servers.json). ");
        sb.Append("A server's tools are named <server>__<tool>.\n");
        foreach (var s in statuses)
        {
            var url = RemoteServersFile.DisplayUrl(s.Entry.Url);
            sb.Append(s.State.StartsWith("unavailable", StringComparison.Ordinal)
                ? $"- {s.Entry.Name} ({url}) is UNAVAILABLE, its tools are missing in this conversation: {s.State["unavailable: ".Length..]}. It is retried automatically; tell the user if they need it.\n"
                : $"- {s.Entry.Name} ({url}): {s.State}\n");
        }
        foreach (var problem in problems) sb.Append($"- Not loaded: {problem}\n");
        return Redact(sb.ToString().TrimEnd());
    }

    private void Track(IExtensionContext ctx)
    {
        _cwd = ctx.Cwd;
        if (_cipher is null && _options.IsConfigured) _cipher = new RemoteServerCipher(_options.Secret!, ctx.SessionManager.Scope);
    }

    private sealed class ServerState(RemoteServerEntry entry, RemoteServerConnection connection)
    {
        public RemoteServerEntry Entry { get; } = entry;
        public string Digest { get; } = entry.Digest();
        public RemoteServerConnection Connection { get; } = connection;
        public Task? Connecting { get; set; }
        public IReadOnlyList<Tool>? AllTools { get; set; }
        public string? Error { get; set; }
        public int Failures { get; set; }
        public DateTimeOffset? RetryAt { get; set; }
        public List<string> Exposed { get; } = [];
    }
}