Playa/extensions/Volcanic.AgenticCore.Extensions.Mcp/McpExtension.cs

using Microsoft.Extensions.Logging;
using Volcanic.AgenticCore.Abstractions.Extensions;
using Volcanic.AgenticCore.Abstractions.Extensions.Events;

namespace Volcanic.AgenticCore.Extensions.Mcp;

/// <summary>
/// MCP client: registers the tools of the configured stdio MCP servers. The catalogue is discovered once per host process;
/// each session gets its own server process, started on the first call.
/// </summary>
public sealed class McpExtension : IExtension
{
    public const string CommandName = "mcp";

    private readonly McpOptions _options;
    private readonly McpCatalogCache _cache;
    private readonly Func<string, string?> _getEnv;
    private readonly TimeProvider _time;
    private readonly List<McpServerState> _servers = [];

    /// <summary>For DI: the host binds <see cref="McpOptions"/>, applies the environment and validates it.</summary>
    public McpExtension(McpOptions options)
        : this(options, McpCatalogCache.Shared, Environment.GetEnvironmentVariable, TimeProvider.System) { }

    /// <summary>For a dropped-in assembly: only the <c>PLAYWRIGHT_*</c> variables and the user settings.</summary>
    public McpExtension()
        : this(FromEnvironment(Environment.GetEnvironmentVariable), McpCatalogCache.Shared, Environment.GetEnvironmentVariable, TimeProvider.System) { }

    internal McpExtension(McpOptions options, McpCatalogCache cache, Func<string, string?> getEnv, TimeProvider time)
    {
        ArgumentNullException.ThrowIfNull(options);
        _options = options;
        _cache = cache;
        _getEnv = getEnv;
        _time = time;
    }

    public string Name => "mcp";

    internal McpOptions Options => _options;

    public async ValueTask ActivateAsync(IExtensionApi api, CancellationToken cancellationToken)
    {
        var logger = api.Logger;
        var redact = RedactorOf(api);
        var entries = McpServerSet.Resolve(_options, McpUserSettings.Load(McpUserSettings.AgentDir(_getEnv), logger), logger);
        var states = await Task.WhenAll(entries.Select(e => PrepareAsync(e, api, redact, cancellationToken))).ConfigureAwait(false);
        _servers.AddRange(states);
        RegisterTools(api);
        api.RegisterCommand(CommandName, McpCommand.Create(api, _servers, redact));
        api.On<BeforeAgentStartEvent>((e, _, _) =>
        {
            AddPromptSections(api, e.SystemPromptOptions);
            return ValueTask.CompletedTask;
        });
        api.On<SessionShutdownEvent>(async (e, _, _) =>
        {
            var final = e.Reason is SessionShutdownReason.Quit or SessionShutdownReason.Reload;
            // Concurrently: SDK 2.2 spends ShutdownTimeout on every live server it disposes.
            await Task.WhenAll(_servers
                .Where(s => s.Connection is not null)
                .Select(s => final ? s.Connection!.DisposeAsync().AsTask() : s.Connection!.StopAsync().AsTask())).ConfigureAwait(false);
        });
    }

    /// <summary>The host's redactor, or identity for an API without one (it then has no secrets to hide).</summary>
    internal static Func<string, string> RedactorOf(IExtensionApi api)
    {
        try
        {
            var redactor = api.Redactor;
            return redactor.Redact;
        }
        catch (NotSupportedException)
        {
            return static text => text;
        }
    }

    private async Task<McpServerState> PrepareAsync(McpServerEntry entry, IExtensionApi api, Func<string, string> redact, CancellationToken cancellationToken)
    {
        if (entry.Disabled) return new McpServerState(entry);
        if (entry.Problem is { } problem)
        {
            api.Logger.LogError("MCP server {Name} is misconfigured: {Problem}", entry.Name, problem);
            return new McpServerState(entry) { Error = problem };
        }
        try
        {
            var catalog = await _cache.GetAsync(entry.Name, entry.Config, api.Processes, _getEnv, api.Logger, cancellationToken, redact).ConfigureAwait(false);
            return new McpServerState(entry)
            {
                Catalog = catalog,
                Connection = new McpServerConnection(entry.Name, entry.Config, api.Logger, _getEnv, catalog, _time, redact),
            };
        }
        catch (Exception ex) when (ex is not OperationCanceledException)
        {
            return new McpServerState(entry) { Error = ex.Message }; // the cache has logged it
        }
    }

    private void RegisterTools(IExtensionApi api)
    {
        var builtIns = api.GetAllTools().Where(t => t.Source == "built-in").Select(t => t.Name).ToHashSet(StringComparer.Ordinal);
        var owners = new Dictionary<string, string>(StringComparer.Ordinal);
        foreach (var server in _servers)
        {
            if (server.Catalog is not { } catalog || server.Connection is not { } connection) continue;
            foreach (var tool in catalog.Tools)
            {
                if (builtIns.Contains(tool.ExposedName))
                {
                    api.Logger.LogWarning("MCP server {Server}: tool {Tool} would replace the built-in tool of the same name and is not registered; set a toolPrefix for this server.", server.Entry.Name, tool.ExposedName);
                    continue;
                }
                if (owners.TryGetValue(tool.ExposedName, out var owner))
                {
                    api.Logger.LogWarning("MCP tool {Tool} is offered by servers {Servers}; the first one wins.", tool.ExposedName, $"{owner} and {server.Entry.Name}");
                    continue;
                }
                owners[tool.ExposedName] = server.Entry.Name;
                api.RegisterTool(new McpTool(tool, connection, server.Entry.Config.Parallel));
                server.RegisteredTools.Add(tool.ExposedName);
            }
        }
    }

    private void AddPromptSections(IExtensionApi api, SystemPromptOptions options)
    {
        var active = api.GetActiveTools().ToHashSet(StringComparer.Ordinal);
        var hasSend = active.Contains("send"); // registered but inactive: the model can't call it
        foreach (var server in _servers)
        {
            if (!server.RegisteredTools.Any(active.Contains)) continue;
            if (server.Entry.PromptSection(hasSend) is { Length: > 0 } section) options.Sections[server.Entry.SectionKey] = section;
        }
    }

    private static McpOptions FromEnvironment(Func<string, string?> getEnv)
    {
        var options = new McpOptions();
        options.ApplyEnvironment(getEnv);
        return options;
    }
}

/// <summary>One server in one session: its catalogue (if discovery worked), its connection and the tools it got.</summary>
internal sealed class McpServerState(McpServerEntry entry)
{
    public McpServerEntry Entry { get; } = entry;
    public McpCatalog? Catalog { get; init; }
    public McpServerConnection? Connection { get; init; }
    public string? Error { get; init; }
    public List<string> RegisteredTools { get; } = [];
}