From fa7e2778195953900039aa81a6d91732e7b46fe5 Mon Sep 17 00:00:00 2001 From: Mansur Besleney Date: Fri, 9 Jan 2026 14:37:55 +0300 Subject: [PATCH] Refactor MCP tool handling and fix telemetry typos Extracted AbpMcpServerTool to its own file and improved tool name validation in McpHttpClientService. Enhanced error handling and logging, made cache validity configurable, and fixed typos in ActivityNameConsts for AbpCli command telemetry constants. --- .../Cli/Commands/Services/AbpMcpServerTool.cs | 43 +++++++++++++ .../Commands/Services/McpHttpClientService.cs | 61 ++++++++++++------- .../Cli/Commands/Services/McpServerService.cs | 43 +++---------- .../Commands/Services/McpToolsCacheService.cs | 7 ++- .../Telemetry/Constants/ActivityNameConsts.cs | 22 +++---- 5 files changed, 104 insertions(+), 72 deletions(-) create mode 100644 framework/src/Volo.Abp.Cli.Core/Volo/Abp/Cli/Commands/Services/AbpMcpServerTool.cs diff --git a/framework/src/Volo.Abp.Cli.Core/Volo/Abp/Cli/Commands/Services/AbpMcpServerTool.cs b/framework/src/Volo.Abp.Cli.Core/Volo/Abp/Cli/Commands/Services/AbpMcpServerTool.cs new file mode 100644 index 0000000000..71f0f206f6 --- /dev/null +++ b/framework/src/Volo.Abp.Cli.Core/Volo/Abp/Cli/Commands/Services/AbpMcpServerTool.cs @@ -0,0 +1,43 @@ +using System; +using System.Collections.Generic; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using ModelContextProtocol.Protocol; +using ModelContextProtocol.Server; + +namespace Volo.Abp.Cli.Commands.Services; + +internal class AbpMcpServerTool : McpServerTool +{ + private readonly string _name; + private readonly string _description; + private readonly JsonElement _inputSchema; + private readonly Func, CancellationToken, ValueTask> _handler; + + public AbpMcpServerTool( + string name, + string description, + JsonElement inputSchema, + Func, CancellationToken, ValueTask> handler) + { + _name = name; + _description = description; + _inputSchema = inputSchema; + _handler = handler; + } + + public override Tool ProtocolTool => new Tool + { + Name = _name, + Description = _description, + InputSchema = _inputSchema + }; + + public override IReadOnlyList Metadata => Array.Empty(); + + public override ValueTask InvokeAsync(RequestContext context, CancellationToken cancellationToken) + { + return _handler(context, cancellationToken); + } +} diff --git a/framework/src/Volo.Abp.Cli.Core/Volo/Abp/Cli/Commands/Services/McpHttpClientService.cs b/framework/src/Volo.Abp.Cli.Core/Volo/Abp/Cli/Commands/Services/McpHttpClientService.cs index ed848e5326..3aae328bd7 100644 --- a/framework/src/Volo.Abp.Cli.Core/Volo/Abp/Cli/Commands/Services/McpHttpClientService.cs +++ b/framework/src/Volo.Abp.Cli.Core/Volo/Abp/Cli/Commands/Services/McpHttpClientService.cs @@ -1,5 +1,6 @@ using System; using System.Collections.Generic; +using System.Linq; using System.Net; using System.Net.Http; using System.Text; @@ -31,7 +32,8 @@ public class McpHttpClientService : ITransientDependency private readonly ILogger _logger; private readonly IMcpLogger _mcpLogger; private readonly MemoryService _memoryService; - private string _cachedServerUrl; + private readonly Lazy> _cachedServerUrlLazy; + private List _validToolNames; public McpHttpClientService( CliHttpClientFactory httpClientFactory, @@ -43,39 +45,43 @@ public class McpHttpClientService : ITransientDependency _logger = logger; _mcpLogger = mcpLogger; _memoryService = memoryService; + _cachedServerUrlLazy = new Lazy>(GetMcpServerUrlInternalAsync); } private async Task GetMcpServerUrlAsync() { - // Return cached URL if already resolved - if (_cachedServerUrl != null) - { - return _cachedServerUrl; - } + return await _cachedServerUrlLazy.Value; + } + private async Task GetMcpServerUrlInternalAsync() + { // 1. Check environment variable (highest priority) var envUrl = Environment.GetEnvironmentVariable(CliConsts.McpServerUrlEnvironmentVariable); if (!string.IsNullOrWhiteSpace(envUrl)) { - _cachedServerUrl = envUrl.TrimEnd('/'); - return _cachedServerUrl; + return envUrl.TrimEnd('/'); } // 2. Check persisted setting var persistedUrl = await _memoryService.GetAsync(CliConsts.MemoryKeys.McpServerUrl); if (!string.IsNullOrWhiteSpace(persistedUrl)) { - _cachedServerUrl = persistedUrl.TrimEnd('/'); - return _cachedServerUrl; + return persistedUrl.TrimEnd('/'); } // 3. Return default - _cachedServerUrl = CliConsts.DefaultMcpServerUrl; - return _cachedServerUrl; + return CliConsts.DefaultMcpServerUrl; } public async Task CallToolAsync(string toolName, JsonElement arguments) { + // Validate toolName against whitelist to prevent malicious input + if (_validToolNames != null && !_validToolNames.Contains(toolName)) + { + _mcpLogger.Warning(LogSource, $"Attempted to call unknown tool: {toolName}"); + return CreateErrorResponse($"Unknown tool: {toolName}"); + } + var baseUrl = await GetMcpServerUrlAsync(); var url = $"{baseUrl}/tools/call"; @@ -141,9 +147,14 @@ public class McpHttpClientService : ITransientDependency }, JsonSerializerOptionsWeb); } - private Exception CreateToolDefinitionException(string userMessage) + private CliUsageException CreateToolDefinitionException(string userMessage) { - return new Exception($"Failed to fetch tool definitions: {userMessage}"); + return new CliUsageException($"Failed to fetch tool definitions: {userMessage}"); + } + + private CliUsageException CreateToolDefinitionException(string userMessage, Exception innerException) + { + return new CliUsageException($"Failed to fetch tool definitions: {userMessage}", innerException); } private string GetSanitizedHttpErrorMessage(HttpStatusCode statusCode) @@ -201,33 +212,37 @@ public class McpHttpClientService : ITransientDependency // The API returns { tools: [...] } format var result = JsonSerializer.Deserialize(responseContent, JsonSerializerOptionsWeb); - - return result?.Tools ?? new List(); + var tools = result?.Tools ?? new List(); + + // Cache tool names for validation + _validToolNames = tools.Select(t => t.Name).ToList(); + + return tools; } catch (HttpRequestException ex) { - throw CreateHttpException(ex, "Network error fetching tool definitions"); + throw CreateHttpExceptionWithInner(ex, "Network error fetching tool definitions"); } catch (TaskCanceledException ex) { - throw CreateHttpException(ex, "Timeout fetching tool definitions"); + throw CreateHttpExceptionWithInner(ex, "Timeout fetching tool definitions"); } catch (JsonException ex) { - throw CreateHttpException(ex, "JSON parsing error"); + throw CreateHttpExceptionWithInner(ex, "JSON parsing error"); } - catch (Exception ex) when (ex.Message.StartsWith("Failed to fetch tool definitions:")) + catch (CliUsageException) { // Already sanitized, rethrow as-is throw; } catch (Exception ex) { - throw CreateHttpException(ex, "Unexpected error fetching tool definitions"); + throw CreateHttpExceptionWithInner(ex, "Unexpected error fetching tool definitions"); } } - private Exception CreateHttpException(Exception ex, string context) + private CliUsageException CreateHttpExceptionWithInner(Exception ex, string context) { _mcpLogger.Error(LogSource, context, ex); @@ -239,7 +254,7 @@ public class McpHttpClientService : ITransientDependency _ => "An unexpected error occurred. Please try again later." }; - return CreateToolDefinitionException(userMessage); + return CreateToolDefinitionException(userMessage, ex); } private class McpToolsResponse diff --git a/framework/src/Volo.Abp.Cli.Core/Volo/Abp/Cli/Commands/Services/McpServerService.cs b/framework/src/Volo.Abp.Cli.Core/Volo/Abp/Cli/Commands/Services/McpServerService.cs index 0c8c4ba21d..30370b8bc6 100644 --- a/framework/src/Volo.Abp.Cli.Core/Volo/Abp/Cli/Commands/Services/McpServerService.cs +++ b/framework/src/Volo.Abp.Cli.Core/Volo/Abp/Cli/Commands/Services/McpServerService.cs @@ -15,6 +15,7 @@ namespace Volo.Abp.Cli.Commands.Services; public class McpServerService : ITransientDependency { private const string LogSource = nameof(McpServerService); + private const int MaxLogResponseLength = 500; private static class ToolErrorMessages { @@ -173,7 +174,12 @@ public class McpServerService : ITransientDependency catch (Exception ex) { _mcpLogger.Error(LogSource, $"Failed to deserialize response as CallToolResult: {ex.Message}"); - _mcpLogger.Debug(LogSource, $"Response was: {resultJson.Substring(0, Math.Min(500, resultJson.Length))}"); + + var logResponse = resultJson.Length <= MaxLogResponseLength + ? resultJson + : resultJson.Substring(0, MaxLogResponseLength); + _mcpLogger.Debug(LogSource, $"Response was: {logResponse}"); + return null; } } @@ -189,39 +195,4 @@ public class McpServerService : ITransientDependency _mcpLogger.Debug(LogSource, $"Tool '{toolName}' executed successfully"); } } - - private class AbpMcpServerTool : McpServerTool - { - private readonly string _name; - private readonly string _description; - private readonly JsonElement _inputSchema; - private readonly Func, CancellationToken, ValueTask> _handler; - - public AbpMcpServerTool( - string name, - string description, - JsonElement inputSchema, - Func, CancellationToken, ValueTask> handler) - { - _name = name; - _description = description; - _inputSchema = inputSchema; - _handler = handler; - } - - public override Tool ProtocolTool => new Tool - { - Name = _name, - Description = _description, - InputSchema = _inputSchema - }; - - public override IReadOnlyList Metadata => Array.Empty(); - - public override ValueTask InvokeAsync(RequestContext context, CancellationToken cancellationToken) - { - return _handler(context, cancellationToken); - } - } - } diff --git a/framework/src/Volo.Abp.Cli.Core/Volo/Abp/Cli/Commands/Services/McpToolsCacheService.cs b/framework/src/Volo.Abp.Cli.Core/Volo/Abp/Cli/Commands/Services/McpToolsCacheService.cs index ec5669d2e8..741e30a840 100644 --- a/framework/src/Volo.Abp.Cli.Core/Volo/Abp/Cli/Commands/Services/McpToolsCacheService.cs +++ b/framework/src/Volo.Abp.Cli.Core/Volo/Abp/Cli/Commands/Services/McpToolsCacheService.cs @@ -16,6 +16,7 @@ namespace Volo.Abp.Cli.Commands.Services; public class McpToolsCacheService : ITransientDependency { private const string LogSource = nameof(McpToolsCacheService); + private const int CacheValidityHours = 24; private readonly McpHttpClientService _mcpHttpClient; private readonly MemoryService _memoryService; @@ -110,8 +111,8 @@ public class McpToolsCacheService : ITransientDependency if (DateTime.TryParse(lastFetchTimeString, CultureInfo.InvariantCulture, DateTimeStyles.None, out var lastFetchTime)) { - // Check if less than 24 hours old - if (DateTime.Now.Subtract(lastFetchTime).TotalHours < 24) + // Check if less than configured hours old + if (DateTime.Now.Subtract(lastFetchTime).TotalHours < CacheValidityHours) { return true; } @@ -167,6 +168,8 @@ public class McpToolsCacheService : ITransientDependency PropertyNamingPolicy = JsonNamingPolicy.CamelCase }); + // Using synchronous File.WriteAllText is acceptable here since cache writes are not on the critical path + // and we need to support multiple target frameworks File.WriteAllText(CliPaths.McpToolsCache, json); // Set restrictive file permissions (user read/write only) diff --git a/framework/src/Volo.Abp.Core/Volo/Abp/Internal/Telemetry/Constants/ActivityNameConsts.cs b/framework/src/Volo.Abp.Core/Volo/Abp/Internal/Telemetry/Constants/ActivityNameConsts.cs index aabbf142e3..fa6d41a9a7 100644 --- a/framework/src/Volo.Abp.Core/Volo/Abp/Internal/Telemetry/Constants/ActivityNameConsts.cs +++ b/framework/src/Volo.Abp.Core/Volo/Abp/Internal/Telemetry/Constants/ActivityNameConsts.cs @@ -58,17 +58,17 @@ public static class ActivityNameConsts public const string AbpStudioSuiteOpen = "AbpStudio.Suite.Open"; public const string AbpStudioGlobalSecretsManage = "AbpStudio.GlobalSecrets.Manage"; public const string AbpStudioGlobalMetadataManage = "AbpStudio.GlobalMetadata.Manage"; - public const string AbpCliCommandsNewSolution = "AbpCli.Comands.NewSolution"; - public const string AbpCliCommandsNewModule = "AbpCli.Comands.NewModule"; - public const string AbpCliCommandsNewPackage = "AbpCli.Comands.NewPackage"; - public const string AbpCliCommandsUpdate = "AbpCli.Comands.Update"; - public const string AbpCliCommandsClean = "AbpCli.Comands.Clean"; - public const string AbpCliCommandsAddPackage = "AbpCli.Comands.AddPackage"; - public const string AbpCliCommandsAddPackageRef = "AbpCli.Comands.AddPackageRef"; - public const string AbpCliCommandsInstallModule = "AbpCli.Comands.InstallModule"; - public const string AbpCliCommandsInstallLocalModule = "AbpCli.Comands.InstallLocalModule"; - public const string AbpCliCommandsListModules = "AbpCli.Comands.ListModules"; - public const string AbpCliCommandsMcp = "AbpCli.Comands.Mcp"; + public const string AbpCliCommandsNewSolution = "AbpCli.Commands.NewSolution"; + public const string AbpCliCommandsNewModule = "AbpCli.Commands.NewModule"; + public const string AbpCliCommandsNewPackage = "AbpCli.Commands.NewPackage"; + public const string AbpCliCommandsUpdate = "AbpCli.Commands.Update"; + public const string AbpCliCommandsClean = "AbpCli.Commands.Clean"; + public const string AbpCliCommandsAddPackage = "AbpCli.Commands.AddPackage"; + public const string AbpCliCommandsAddPackageRef = "AbpCli.Commands.AddPackageRef"; + public const string AbpCliCommandsInstallModule = "AbpCli.Commands.InstallModule"; + public const string AbpCliCommandsInstallLocalModule = "AbpCli.Commands.InstallLocalModule"; + public const string AbpCliCommandsListModules = "AbpCli.Commands.ListModules"; + public const string AbpCliCommandsMcp = "AbpCli.Commands.Mcp"; public const string AbpCliRun = "AbpCli.Run"; public const string AbpCliExit = "AbpCli.Exit"; public const string ApplicationRun = "Application.Run";