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

using Microsoft.Extensions.Logging;
using ModelContextProtocol.Client;
using Volcanic.AgenticCore.Abstractions.Processes;

namespace Volcanic.AgenticCore.Extensions.Mcp;

/// <summary>
/// Process-wide cache of <c>tools/list</c> results, keyed by server name and a digest of its unexpanded configuration.
/// Concurrent requests share one discovery. Success is kept for the process lifetime, failure for <see cref="FailureCooldown"/>.
/// </summary>
public sealed class McpCatalogCache
{
    /// <summary>How long a failed discovery is remembered before the next request tries again.</summary>
    public static readonly TimeSpan FailureCooldown = TimeSpan.FromSeconds(60);

    /// <summary>Explicit process-wide state, like <c>ClaudeLoginCoordinator.Shared</c>.</summary>
    public static McpCatalogCache Shared { get; } = new(TimeProvider.System);

    private readonly TimeProvider _time;
    private readonly Lock _lock = new();
    private readonly Dictionary<string, Task<McpCatalog>> _tasks = new(StringComparer.Ordinal);
    private readonly Dictionary<string, DateTimeOffset> _failedAt = new(StringComparer.Ordinal);

    internal McpCatalogCache(TimeProvider time) => _time = time;

    /// <summary>
    /// The server's catalogue, discovering it if nothing is cached or a failure's cooldown has passed. The shared discovery
    /// runs detached from any caller's token; each caller waits with its own. Throws <see cref="McpServerStartException"/>.
    /// <paramref name="redact"/> is applied to the server's stderr before it is logged or put in an error message.
    /// </summary>
    internal Task<McpCatalog> GetAsync(string name, McpServerConfig config, IProcessSpawnPolicy processes, Func<string, string?> getEnv, ILogger logger, CancellationToken cancellationToken, Func<string, string>? redact = null)
    {
        var key = McpLaunch.CacheKey(name, config);
        Task<McpCatalog> task;
        lock (_lock)
        {
            _tasks.TryGetValue(key, out var existing);
            var expired = existing is { IsFaulted: true }
                && _failedAt.TryGetValue(key, out var failedAt)
                && _time.GetUtcNow() - failedAt >= FailureCooldown;
            if (existing is null || expired)
            {
                _failedAt.Remove(key);
                existing = Task.Run(() => DiscoverAsync(key, name, config, processes, getEnv, logger, redact), CancellationToken.None);
                _tasks[key] = existing;
            }
            task = existing;
        }
        return task.WaitAsync(cancellationToken);
    }

    private async Task<McpCatalog> DiscoverAsync(string key, string name, McpServerConfig config, IProcessSpawnPolicy processes, Func<string, string?> getEnv, ILogger logger, Func<string, string>? redact)
    {
        var dir = Path.Combine(Path.GetTempPath(), "volcanic-mcp-discovery", Guid.NewGuid().ToString("N"));
        Directory.CreateDirectory(dir);
        var stderr = new StderrTail(redact);
        McpProcess? process = null;
        try
        {
            ProcessSpawnContext spawn;
            try { spawn = processes.Prepare(McpLaunch.BuildRequest(name, config, dir, getEnv)); }
            catch (Exception ex) { throw new McpServerStartException($"{name}: process environment hook failed: {ex.Message}", ex); }

            process = await McpProcess.StartAsync(name, config, spawn, stderr, logger, CancellationToken.None).ConfigureAwait(false);
            using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(config.StartupTimeoutSeconds));
            IList<McpClientTool> tools;
            try { tools = await process.Client.ListToolsAsync(cancellationToken: cts.Token).ConfigureAwait(false); }
            catch (OperationCanceledException)
            {
                throw new McpServerStartException($"{name} MCP server did not list its tools within {config.StartupTimeoutSeconds}s{McpProcess.TailSuffix(stderr)}");
            }
            catch (Exception ex)
            {
                throw new McpServerStartException($"{name} MCP server failed to list its tools: {ex.Message}{McpProcess.TailSuffix(stderr)}", ex);
            }
            var catalog = McpCatalog.From(name, config, process.Client.ServerInfo, [.. tools.Select(t => t.ProtocolTool)]);
            logger.LogInformation("MCP server {Name}: {Count} tools ({Server} {Version})", name, catalog.Tools.Count, catalog.ServerInfoName, catalog.ServerVersion);
            return catalog;
        }
        catch (Exception ex)
        {
            lock (_lock) _failedAt[key] = _time.GetUtcNow(); // recorded before the task faults
            logger.LogWarning("MCP server {Name} is unavailable: {Message}", name, redact is null ? ex.Message : redact(ex.Message));
            if (ex is McpServerStartException) throw;
            throw new McpServerStartException($"{name} MCP server discovery failed: {ex.Message}", ex);
        }
        finally
        {
            // Detached: the SDK spends ShutdownTimeout on every dispose, and no caller needs to wait for it.
            _ = Task.Run(() => CleanUpAsync(name, process, dir, logger), CancellationToken.None);
        }
    }

    private static async Task CleanUpAsync(string name, McpProcess? process, string dir, ILogger logger)
    {
        try
        {
            if (process is not null) await process.DisposeAsync().ConfigureAwait(false);
        }
        catch (Exception ex)
        {
            logger.LogDebug(ex, "MCP server {Name}: stopping the discovery process failed", name);
        }

        // Windows can hold the cwd for a moment after the process exits.
        for (var attempt = 1; ; attempt++)
        {
            try
            {
                if (Directory.Exists(dir)) Directory.Delete(dir, recursive: true);
                return;
            }
            catch (Exception ex) when (ex is IOException or UnauthorizedAccessException)
            {
                if (attempt == 5)
                {
                    logger.LogDebug(ex, "MCP server {Name}: could not delete the discovery directory {Dir}", name, dir);
                    return;
                }
                await Task.Delay(TimeSpan.FromMilliseconds(200 * attempt)).ConfigureAwait(false);
            }
        }
    }
}