diff --git a/src/Aspire.Cli/Commands/AgentMcpCommand.cs b/src/Aspire.Cli/Commands/AgentMcpCommand.cs index ce9106bc126..36df5b2a4dd 100644 --- a/src/Aspire.Cli/Commands/AgentMcpCommand.cs +++ b/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 _knownTools; - private string? _selectedAppHostPath; - private Dictionary? _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 _logger; - private readonly IDocsIndexService _docsIndexService; /// /// 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()); _knownTools = new Dictionary { [KnownMcpTools.ListResources] = new ListResourcesTool(auxiliaryBackchannelMonitor, loggerFactory.CreateLogger()), @@ -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,13 +116,15 @@ protected override async Task 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; @@ -135,8 +132,6 @@ protected override async Task ExecuteAsync(ParseResult parseResult, Cancell private async ValueTask HandleListToolsAsync(RequestContext request, CancellationToken cancellationToken) { - _ = request; - _logger.LogDebug("MCP ListTools request received"); var tools = new List(); @@ -150,15 +145,14 @@ private async ValueTask HandleListToolsAsync(RequestContext new Tool + tools.AddRange(resourceToolMap.Select(x => new Tool { Name = x.Key, Description = x.Value.Tool.Description, @@ -213,17 +207,16 @@ private async ValueTask HandleCallToolAsync(RequestContext HandleCallToolAsync(RequestContext 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 RefreshResourceToolMapAsync(CancellationToken cancellationToken) - { - var refreshedMap = new Dictionary(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(); - var seenResources = new HashSet(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; - } - /// /// Gets the appropriate AppHost connection based on the selection logic. /// diff --git a/src/Aspire.Cli/Mcp/IMcpResourceToolRefreshService.cs b/src/Aspire.Cli/Mcp/IMcpResourceToolRefreshService.cs new file mode 100644 index 00000000000..b08357049e3 --- /dev/null +++ b/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; + +/// +/// Service responsible for refreshing resource-based MCP tools and sending tool list change notifications. +/// +internal interface IMcpResourceToolRefreshService +{ + /// + /// Attempts to get the current resource tool map if it is valid (not invalidated and AppHost hasn't changed). + /// + /// When this method returns true, contains the current resource tool map. + /// true if the tool map is valid and no refresh is needed; otherwise, false. + bool TryGetResourceToolMap(out IReadOnlyDictionary resourceToolMap); + + /// + /// Marks the resource tool map as needing a refresh. + /// + void InvalidateToolMap(); + + /// + /// Refreshes the resource tool map by discovering MCP tools from connected resources. + /// + /// The cancellation token. + /// The refreshed resource tool map. + Task> RefreshResourceToolMapAsync(CancellationToken cancellationToken); + + /// + /// Sends a tools list changed notification to connected MCP clients. + /// + /// The cancellation token. + Task SendToolsListChangedNotificationAsync(CancellationToken cancellationToken); + + /// + /// Sets the MCP server instance used for sending notifications. + /// + /// The MCP server, or null to clear. + void SetMcpServer(McpServer? server); +} + +/// +/// Represents an entry in the resource tool map. +/// +/// The name of the resource that exposes the tool. +/// The MCP tool definition. +internal sealed record ResourceToolEntry(string ResourceName, Tool Tool); diff --git a/src/Aspire.Cli/Mcp/KnownMcpTools.cs b/src/Aspire.Cli/Mcp/KnownMcpTools.cs index 78e65c1d01a..de9276b63c0 100644 --- a/src/Aspire.Cli/Mcp/KnownMcpTools.cs +++ b/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"; + /// + /// Gets all known MCP tool names. + /// + public static IReadOnlyList 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 diff --git a/src/Aspire.Cli/Mcp/McpResourceToolRefreshService.cs b/src/Aspire.Cli/Mcp/McpResourceToolRefreshService.cs new file mode 100644 index 00000000000..95ecf656214 --- /dev/null +++ b/src/Aspire.Cli/Mcp/McpResourceToolRefreshService.cs @@ -0,0 +1,126 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +using System.Diagnostics; +using Aspire.Cli.Backchannel; +using Microsoft.Extensions.Logging; +using ModelContextProtocol.Protocol; +using ModelContextProtocol.Server; + +namespace Aspire.Cli.Mcp; + +/// +/// Service responsible for refreshing resource-based MCP tools and sending tool list change notifications. +/// +internal sealed class McpResourceToolRefreshService : IMcpResourceToolRefreshService +{ + private readonly IAuxiliaryBackchannelMonitor _auxiliaryBackchannelMonitor; + private readonly ILogger _logger; + private readonly object _lock = new(); + private McpServer? _server; + private Dictionary _resourceToolMap = new(StringComparer.Ordinal); + private bool _invalidated = true; + private string? _selectedAppHostPath; + + public McpResourceToolRefreshService( + IAuxiliaryBackchannelMonitor auxiliaryBackchannelMonitor, + ILogger logger) + { + _auxiliaryBackchannelMonitor = auxiliaryBackchannelMonitor; + _logger = logger; + } + + /// + public bool TryGetResourceToolMap(out IReadOnlyDictionary resourceToolMap) + { + lock (_lock) + { + if (_invalidated || _selectedAppHostPath != _auxiliaryBackchannelMonitor.SelectedAppHostPath) + { + resourceToolMap = null!; + return false; + } + + resourceToolMap = _resourceToolMap; + return true; + } + } + + /// + public void InvalidateToolMap() + { + lock (_lock) + { + _invalidated = true; + } + } + + /// + public void SetMcpServer(McpServer? server) + { + _server = server; + } + + /// + public async Task SendToolsListChangedNotificationAsync(CancellationToken cancellationToken) + { + if (_server is { } server) + { + await server.SendNotificationAsync(NotificationMethods.ToolListChangedNotification, cancellationToken).ConfigureAwait(false); + } + } + + /// + public async Task> RefreshResourceToolMapAsync(CancellationToken cancellationToken) + { + _logger.LogDebug("Refreshing resource tool map."); + + var refreshedMap = new Dictionary(StringComparer.Ordinal); + + string? selectedAppHostPath = null; + try + { + var connection = await AppHostConnectionHelper.GetSelectedConnectionAsync(_auxiliaryBackchannelMonitor, _logger, cancellationToken).ConfigureAwait(false); + + if (connection is not null) + { + selectedAppHostPath = connection.AppHostInfo?.AppHostPath; + + var allResources = await connection.GetResourceSnapshotsAsync(cancellationToken).ConfigureAwait(false); + var resourcesWithTools = allResources.Where(r => r.McpServer is not null).ToList(); + + _logger.LogDebug("Resources with MCP tools received: {Count}", resourcesWithTools.Count); + + foreach (var resource in resourcesWithTools) + { + Debug.Assert(resource.McpServer is not null); + + foreach (var tool in resource.McpServer.Tools) + { + var exposedName = $"{resource.Name.Replace("-", "_")}_{tool.Name}"; + refreshedMap[exposedName] = new ResourceToolEntry(resource.Name, tool); + + _logger.LogDebug("{Tool}: {Description}", exposedName, tool.Description); + } + } + } + else + { + _logger.LogDebug("Unable to refresh resource tool map because there's no selected connection."); + } + } + 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."); + } + + lock (_lock) + { + _resourceToolMap = refreshedMap; + _selectedAppHostPath = selectedAppHostPath; + _invalidated = false; + return _resourceToolMap; + } + } +} diff --git a/src/Aspire.Cli/Mcp/Tools/RefreshToolsTool.cs b/src/Aspire.Cli/Mcp/Tools/RefreshToolsTool.cs index 9eee648bf78..502f00251a3 100644 --- a/src/Aspire.Cli/Mcp/Tools/RefreshToolsTool.cs +++ b/src/Aspire.Cli/Mcp/Tools/RefreshToolsTool.cs @@ -6,7 +6,7 @@ namespace Aspire.Cli.Mcp.Tools; -internal sealed class RefreshToolsTool(Func> refreshToolsAsync, Func sendToolsListChangedNotificationAsync) : CliMcpTool +internal sealed class RefreshToolsTool(IMcpResourceToolRefreshService refreshService) : CliMcpTool { public override string Name => KnownMcpTools.RefreshTools; @@ -21,12 +21,13 @@ public override async ValueTask CallToolAsync(CallToolContext co { _ = context; - var count = await refreshToolsAsync(cancellationToken).ConfigureAwait(false); - await sendToolsListChangedNotificationAsync(cancellationToken).ConfigureAwait(false); + var resourceToolMap = await refreshService.RefreshResourceToolMapAsync(cancellationToken).ConfigureAwait(false); + await refreshService.SendToolsListChangedNotificationAsync(cancellationToken).ConfigureAwait(false); + var totalToolCount = KnownMcpTools.All.Count + resourceToolMap.Count; return new CallToolResult { - Content = [new TextContentBlock { Text = $"Tools refreshed: {count} tools available" }] + Content = [new TextContentBlock { Text = $"Tools refreshed: {totalToolCount} tools available" }] }; } } diff --git a/tests/Aspire.Cli.Tests/Commands/AgentMcpCommandTests.cs b/tests/Aspire.Cli.Tests/Commands/AgentMcpCommandTests.cs index e0b3f071083..dd7cb3e3587 100644 --- a/tests/Aspire.Cli.Tests/Commands/AgentMcpCommandTests.cs +++ b/tests/Aspire.Cli.Tests/Commands/AgentMcpCommandTests.cs @@ -337,8 +337,8 @@ public async Task McpServer_CallTool_RefreshTools_ReturnsResult() var textContent = result.Content[0] as TextContentBlock; Assert.NotNull(textContent); - // Verify the exact text content with the correct tool count - var expectedToolCount = _agentMcpCommand.KnownTools.Count; + // Verify the text content indicates refresh success (resource tool count is 0 in this test, so total = known tools) + var expectedToolCount = KnownMcpTools.All.Count; Assert.Equal($"Tools refreshed: {expectedToolCount} tools available", textContent.Text); // Assert - Verify the ToolListChanged notification was received