Refactor resource MCP refresh

src/Aspire.Cli/Commands/AgentMcpCommand.cs

@@ -2,7 +2,6 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.CommandLine;
-using System.Diagnostics.CodeAnalysis;
 using System.Globalization;
 using System.Text.Json;
 using Aspire.Cli.Backchannel;
@@ -32,15 +31,12 @@ namespace Aspire.Cli.Commands;
 internal sealed class AgentMcpCommand : BaseCommand
 {
     private readonly Dictionary<string, CliMcpTool> _knownTools;
-    private string? _selectedAppHostPath;
-    private Dictionary<string, (string ResourceName, Tool Tool)>? _resourceToolMap;
+    private readonly IMcpResourceToolRefreshService _resourceToolRefreshService;
     private McpServer? _server;
     private readonly IAuxiliaryBackchannelMonitor _auxiliaryBackchannelMonitor;
-    private readonly CliExecutionContext _executionContext;
     private readonly IMcpTransportFactory _transportFactory;
     private readonly ILoggerFactory _loggerFactory;
     private readonly ILogger<AgentMcpCommand> _logger;
-    private readonly IDocsIndexService _docsIndexService;
 
     /// <summary>
     /// Gets the dictionary of known MCP tools. Exposed for testing purposes.
@@ -64,11 +60,10 @@ public AgentMcpCommand(
         : base("mcp", AgentCommandStrings.McpCommand_Description, features, updateNotifier, executionContext, interactionService, telemetry)
     {
         _auxiliaryBackchannelMonitor = auxiliaryBackchannelMonitor;
-        _executionContext = executionContext;
         _transportFactory = transportFactory;
         _loggerFactory = loggerFactory;
         _logger = logger;
-        _docsIndexService = docsIndexService;
+        _resourceToolRefreshService = new McpResourceToolRefreshService(auxiliaryBackchannelMonitor, loggerFactory.CreateLogger<McpResourceToolRefreshService>());
         _knownTools = new Dictionary<string, CliMcpTool>
         {
             [KnownMcpTools.ListResources] = new ListResourcesTool(auxiliaryBackchannelMonitor, loggerFactory.CreateLogger<ListResourcesTool>()),
@@ -81,7 +76,7 @@ public AgentMcpCommand(
             [KnownMcpTools.ListAppHosts] = new ListAppHostsTool(auxiliaryBackchannelMonitor, executionContext),
             [KnownMcpTools.ListIntegrations] = new ListIntegrationsTool(packagingService, executionContext, auxiliaryBackchannelMonitor),
             [KnownMcpTools.Doctor] = new DoctorTool(environmentChecker),
-            [KnownMcpTools.RefreshTools] = new RefreshToolsTool(RefreshResourceToolMapAsync, SendToolsListChangedNotificationAsync),
+            [KnownMcpTools.RefreshTools] = new RefreshToolsTool(_resourceToolRefreshService),
             [KnownMcpTools.ListDocs] = new ListDocsTool(docsIndexService),
             [KnownMcpTools.SearchDocs] = new SearchDocsTool(docsSearchService, docsIndexService),
             [KnownMcpTools.GetDoc] = new GetDocTool(docsIndexService)
@@ -121,22 +116,22 @@ protected override async Task<int> ExecuteAsync(ParseResult parseResult, Cancell
         var transport = _transportFactory.CreateTransport();
         await using var server = McpServer.Create(transport, options, _loggerFactory);
 
-        // Keep a reference to the server for sending notifications
+        // Configure the refresh service with the server
+        _resourceToolRefreshService.SetMcpServer(server);
         _server = server;
 
         // Starts the MCP server, it's blocking until cancellation is requested
         await server.RunAsync(cancellationToken);
 
         // Clear the server reference on exit
+        _resourceToolRefreshService.SetMcpServer(null);
         _server = null;
 
         return ExitCodeConstants.Success;
     }
 
     private async ValueTask<ListToolsResult> HandleListToolsAsync(RequestContext<ListToolsRequestParams> request, CancellationToken cancellationToken)
     {
-        _ = request;
-
         _logger.LogDebug("MCP ListTools request received");
 
         var tools = new List<Tool>();
@@ -150,15 +145,14 @@ private async ValueTask<ListToolsResult> HandleListToolsAsync(RequestContext<Lis
 
         try
         {
-            // Detect if the tools list should be refreshed due to AppHost selection change
-            if (_resourceToolMap is null || _selectedAppHostPath != _auxiliaryBackchannelMonitor.SelectedAppHostPath)
+            // Refresh resource tools if needed (e.g., AppHost selection changed or invalidated)
+            if (!_resourceToolRefreshService.TryGetResourceToolMap(out var resourceToolMap))
             {
-                await RefreshResourceToolMapAsync(cancellationToken);
-                await SendToolsListChangedNotificationAsync(cancellationToken).ConfigureAwait(false);
-                _selectedAppHostPath = _auxiliaryBackchannelMonitor.SelectedAppHostPath;
+                resourceToolMap = await _resourceToolRefreshService.RefreshResourceToolMapAsync(cancellationToken);
+                await _resourceToolRefreshService.SendToolsListChangedNotificationAsync(cancellationToken).ConfigureAwait(false);
             }
 
-            tools.AddRange(_resourceToolMap.Select(x => new Tool
+            tools.AddRange(resourceToolMap.Select(x => new Tool
             {
                 Name = x.Key,
                 Description = x.Value.Tool.Description,
@@ -213,17 +207,16 @@ private async ValueTask<CallToolResult> HandleCallToolAsync(RequestContext<CallT
 
         var toolsRefreshed = false;
 
-        // Detect if the tools list should be refreshed due to AppHost selection change
-        if (_resourceToolMap is null || _selectedAppHostPath != _auxiliaryBackchannelMonitor.SelectedAppHostPath)
+        // Refresh resource tools if needed (e.g., AppHost selection changed or invalidated)
+        if (!_resourceToolRefreshService.TryGetResourceToolMap(out var resourceToolMap))
         {
-            await RefreshResourceToolMapAsync(cancellationToken);
-            _selectedAppHostPath = _auxiliaryBackchannelMonitor.SelectedAppHostPath;
+            resourceToolMap = await _resourceToolRefreshService.RefreshResourceToolMapAsync(cancellationToken);
+            await _resourceToolRefreshService.SendToolsListChangedNotificationAsync(cancellationToken).ConfigureAwait(false);
             toolsRefreshed = true;
-            await SendToolsListChangedNotificationAsync(cancellationToken).ConfigureAwait(false);
         }
 
         // Resource MCP tools are invoked via the AppHost backchannel (AppHost proxies to the resource MCP endpoint).
-        if (_resourceToolMap.TryGetValue(toolName, out var resourceAndTool))
+        if (resourceToolMap.TryGetValue(toolName, out var resourceAndTool))
         {
             var connection = await GetSelectedConnectionAsync(cancellationToken).ConfigureAwait(false);
             if (connection == null)
@@ -255,7 +248,7 @@ private async ValueTask<CallToolResult> HandleCallToolAsync(RequestContext<CallT
         // If we haven't refreshed yet, try refreshing once more in case the resource list changed
         if (!toolsRefreshed)
         {
-            _resourceToolMap = null;
+            _resourceToolRefreshService.InvalidateToolMap();
             return await HandleCallToolAsync(request, cancellationToken).ConfigureAwait(false);
         }
 
@@ -331,81 +324,6 @@ private async ValueTask<CallToolResult> CallDashboardToolAsync(
         }
     }
 
-    private Task SendToolsListChangedNotificationAsync(CancellationToken cancellationToken)
-    {
-        var server = _server;
-        if (server is null)
-        {
-            throw new InvalidOperationException("MCP server is not running.");
-        }
-
-        return server.SendNotificationAsync(NotificationMethods.ToolListChangedNotification, cancellationToken);
-    }
-
-    [MemberNotNull(nameof(_resourceToolMap))]
-    private async Task<int> RefreshResourceToolMapAsync(CancellationToken cancellationToken)
-    {
-        var refreshedMap = new Dictionary<string, (string, Tool)>(StringComparer.Ordinal);
-
-        try
-        {
-            var connection = await GetSelectedConnectionAsync(cancellationToken).ConfigureAwait(false);
-
-            if (connection is not null)
-            {
-                // Collect initial snapshots from the stream
-                // The stream yields initial snapshots for all resources first
-                var resourcesWithTools = new List<ResourceSnapshot>();
-                var seenResources = new HashSet<string>(StringComparer.Ordinal);
-
-                await foreach (var snapshot in connection.WatchResourceSnapshotsAsync(cancellationToken).ConfigureAwait(false))
-                {
-                    // Stop after we've seen all resources once (initial batch)
-                    if (!seenResources.Add(snapshot.Name))
-                    {
-                        break;
-                    }
-
-                    if (snapshot.McpServer is not null)
-                    {
-                        resourcesWithTools.Add(snapshot);
-                    }
-                }
-
-                _logger.LogDebug("Resources with MCP tools received: {Count}", resourcesWithTools.Count);
-
-                foreach (var resource in resourcesWithTools)
-                {
-                    if (resource.McpServer is null)
-                    {
-                        continue;
-                    }
-
-                    foreach (var tool in resource.McpServer.Tools)
-                    {
-                        var exposedName = $"{resource.Name.Replace("-", "_")}_{tool.Name}";
-                        refreshedMap[exposedName] = (resource.Name, tool);
-
-                        _logger.LogDebug("{Tool}: {Description}", exposedName, tool.Description);
-                    }
-                }
-            }
-
-        }
-        catch (Exception ex)
-        {
-            // Don't fail refresh_tools if resource discovery fails; still emit notification.
-            _logger.LogDebug(ex, "Failed to refresh resource MCP tool routing map");
-        }
-        finally
-        {
-            // Ensure _resourceToolMap is always non-null when exiting, even if connection is null or an exception occurs.
-            _resourceToolMap = refreshedMap;
-        }
-
-        return _resourceToolMap.Count + KnownTools.Count;
-    }
-
     /// <summary>
     /// Gets the appropriate AppHost connection based on the selection logic.
     /// </summary>

src/Aspire.Cli/Mcp/IMcpResourceToolRefreshService.cs

@@ -0,0 +1,51 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using ModelContextProtocol.Protocol;
+using ModelContextProtocol.Server;
+
+namespace Aspire.Cli.Mcp;
+
+/// <summary>
+/// Service responsible for refreshing resource-based MCP tools and sending tool list change notifications.
+/// </summary>
+internal interface IMcpResourceToolRefreshService
+{
+    /// <summary>
+    /// Attempts to get the current resource tool map if it is valid (not invalidated and AppHost hasn't changed).
+    /// </summary>
+    /// <param name="resourceToolMap">When this method returns <c>true</c>, contains the current resource tool map.</param>
+    /// <returns><c>true</c> if the tool map is valid and no refresh is needed; otherwise, <c>false</c>.</returns>
+    bool TryGetResourceToolMap(out IReadOnlyDictionary<string, ResourceToolEntry> resourceToolMap);
+
+    /// <summary>
+    /// Marks the resource tool map as needing a refresh.
+    /// </summary>
+    void InvalidateToolMap();
+
+    /// <summary>
+    /// Refreshes the resource tool map by discovering MCP tools from connected resources.
+    /// </summary>
+    /// <param name="cancellationToken">The cancellation token.</param>
+    /// <returns>The refreshed resource tool map.</returns>
+    Task<IReadOnlyDictionary<string, ResourceToolEntry>> RefreshResourceToolMapAsync(CancellationToken cancellationToken);
+
+    /// <summary>
+    /// Sends a tools list changed notification to connected MCP clients.
+    /// </summary>
+    /// <param name="cancellationToken">The cancellation token.</param>
+    Task SendToolsListChangedNotificationAsync(CancellationToken cancellationToken);
+
+    /// <summary>
+    /// Sets the MCP server instance used for sending notifications.
+    /// </summary>
+    /// <param name="server">The MCP server, or null to clear.</param>
+    void SetMcpServer(McpServer? server);
+}
+
+/// <summary>
+/// Represents an entry in the resource tool map.
+/// </summary>
+/// <param name="ResourceName">The name of the resource that exposes the tool.</param>
+/// <param name="Tool">The MCP tool definition.</param>
+internal sealed record ResourceToolEntry(string ResourceName, Tool Tool);

src/Aspire.Cli/Mcp/KnownMcpTools.cs

@@ -24,6 +24,27 @@ internal static class KnownMcpTools
     internal const string SearchDocs = "search_docs";
     internal const string GetDoc = "get_doc";
 
+    /// <summary>
+    /// Gets all known MCP tool names.
+    /// </summary>
+    public static IReadOnlyList<string> All { get; } =
+    [
+        ListResources,
+        ListConsoleLogs,
+        ExecuteResourceCommand,
+        ListStructuredLogs,
+        ListTraces,
+        ListTraceStructuredLogs,
+        SelectAppHost,
+        ListAppHosts,
+        ListIntegrations,
+        Doctor,
+        RefreshTools,
+        ListDocs,
+        SearchDocs,
+        GetDoc
+    ];
+
     public static bool IsLocalTool(string toolName) => toolName is
         SelectAppHost or
         ListAppHosts or