From f2220988d7b5f809c58a735aa1dd8b0ab2dc60a0 Mon Sep 17 00:00:00 2001 From: Antigravity Agent Date: Sat, 22 Aug 2026 11:36:11 -0500 Subject: [PATCH 01/24] Refactor routing managers to partials and annotate tests --- Core/Routing/SessionManager.Cache.cs | 93 +++++ Core/Routing/SessionManager.Metrics.cs | 32 ++ Core/Routing/SessionManager.cs | 111 +----- Core/Routing/ToolRoutingManager.Cache.cs | 128 +++++++ Core/Routing/ToolRoutingManager.Execution.cs | 293 ++++++++++++++ Core/Routing/ToolRoutingManager.cs | 383 +------------------ McpRouter.Tests/SessionManagerTests.cs | 3 + McpRouter.Tests/ToolRoutingManagerTests.cs | 8 + 8 files changed, 559 insertions(+), 492 deletions(-) create mode 100644 Core/Routing/SessionManager.Cache.cs create mode 100644 Core/Routing/SessionManager.Metrics.cs create mode 100644 Core/Routing/ToolRoutingManager.Cache.cs create mode 100644 Core/Routing/ToolRoutingManager.Execution.cs diff --git a/Core/Routing/SessionManager.Cache.cs b/Core/Routing/SessionManager.Cache.cs new file mode 100644 index 00000000..301edc4d --- /dev/null +++ b/Core/Routing/SessionManager.Cache.cs @@ -0,0 +1,93 @@ +using System.Collections.Concurrent; +using System.Collections.Generic; + +namespace McpRouter.Core.Routing +{ + public partial class SessionManager + { + private readonly ConcurrentDictionary> _serverToolsCache = new(); + private readonly ConcurrentDictionary> _serverPromptsCache = new(); + private readonly ConcurrentDictionary> _serverResourcesCache = new(); + private readonly ConcurrentDictionary> _serverResourceTemplatesCache = new(); + + public List? GetServerToolsCache(string serverId) + { + _serverToolsCache.TryGetValue(serverId, out var tools); + return tools; + } + + public void SetServerToolsCache(string serverId, List tools) + { + _serverToolsCache[serverId] = tools; + } + + public void RemoveServerToolsCache(string serverId) + { + _serverToolsCache.TryRemove(serverId, out _); + } + + public List? GetServerPromptsCache(string serverId) + { + _serverPromptsCache.TryGetValue(serverId, out var prompts); + return prompts; + } + + public void SetServerPromptsCache(string serverId, List prompts) + { + _serverPromptsCache[serverId] = prompts; + } + + public void RemoveServerPromptsCache(string serverId) + { + _serverPromptsCache.TryRemove(serverId, out _); + } + + public List? GetServerResourcesCache(string serverId) + { + _serverResourcesCache.TryGetValue(serverId, out var resources); + return resources; + } + + public void SetServerResourcesCache(string serverId, List resources) + { + _serverResourcesCache[serverId] = resources; + } + + public void RemoveServerResourcesCache(string serverId) + { + _serverResourcesCache.TryRemove(serverId, out _); + } + + public List? GetServerResourceTemplatesCache(string serverId) + { + _serverResourceTemplatesCache.TryGetValue(serverId, out var templates); + return templates; + } + + public void SetServerResourceTemplatesCache(string serverId, List templates) + { + _serverResourceTemplatesCache[serverId] = templates; + } + + public void RemoveServerResourceTemplatesCache(string serverId) + { + _serverResourceTemplatesCache.TryRemove(serverId, out _); + } + + public void RemoveServerCache(string serverId) + { + RemoveServerToolsCache(serverId); + RemoveServerPromptsCache(serverId); + RemoveServerResourcesCache(serverId); + RemoveServerResourceTemplatesCache(serverId); + } + + public void ClearGlobalCache() + { + _serverToolsCache.Clear(); + _serverPromptsCache.Clear(); + _serverResourcesCache.Clear(); + _serverResourceTemplatesCache.Clear(); + } + } +} diff --git a/Core/Routing/SessionManager.Metrics.cs b/Core/Routing/SessionManager.Metrics.cs new file mode 100644 index 00000000..b6bfda49 --- /dev/null +++ b/Core/Routing/SessionManager.Metrics.cs @@ -0,0 +1,32 @@ +using System; +using System.Threading; + +namespace McpRouter.Core.Routing +{ + public partial class SessionManager + { + public DateTime StartTime { get; } = DateTime.UtcNow; + private long _totalRequests = 0; + public long TotalRequests => _totalRequests; + + private long _totalInputTokens = 0; + private long _totalOutputTokens = 0; + private long _totalDurationMs = 0; + + public long TotalInputTokens => _totalInputTokens; + public long TotalOutputTokens => _totalOutputTokens; + public long TotalDurationMs => _totalDurationMs; + + public void AddPerformanceMetrics(long inputTokens, long outputTokens, long durationMs) + { + Interlocked.Add(ref _totalInputTokens, inputTokens); + Interlocked.Add(ref _totalOutputTokens, outputTokens); + Interlocked.Add(ref _totalDurationMs, durationMs); + } + + public void IncrementTotalRequests() + { + Interlocked.Increment(ref _totalRequests); + } + } +} diff --git a/Core/Routing/SessionManager.cs b/Core/Routing/SessionManager.cs index 166e4dd7..ea92b566 100644 --- a/Core/Routing/SessionManager.cs +++ b/Core/Routing/SessionManager.cs @@ -12,37 +12,13 @@ namespace McpRouter.Core.Routing { - public class SessionManager + public partial class SessionManager { private readonly ConcurrentDictionary _sessions = new(); private readonly IServiceProvider _serviceProvider; private readonly IHttpClientFactory _httpClientFactory; private readonly ILogger _logger; - public DateTime StartTime { get; } = DateTime.UtcNow; - private long _totalRequests = 0; - public long TotalRequests => _totalRequests; - - private long _totalInputTokens = 0; - private long _totalOutputTokens = 0; - private long _totalDurationMs = 0; - - public long TotalInputTokens => _totalInputTokens; - public long TotalOutputTokens => _totalOutputTokens; - public long TotalDurationMs => _totalDurationMs; - - public void AddPerformanceMetrics(long inputTokens, long outputTokens, long durationMs) - { - System.Threading.Interlocked.Add(ref _totalInputTokens, inputTokens); - System.Threading.Interlocked.Add(ref _totalOutputTokens, outputTokens); - System.Threading.Interlocked.Add(ref _totalDurationMs, durationMs); - } - - public void IncrementTotalRequests() - { - System.Threading.Interlocked.Increment(ref _totalRequests); - } - public int ActiveSessionsCount => _sessions.Count; public ConcurrentDictionary BackendStatuses { get; } = new(); @@ -120,91 +96,6 @@ public System.Collections.Generic.List GetActiveSessions() return _sessions.Values.ToList(); } - private readonly ConcurrentDictionary> _serverToolsCache = new(); - private readonly ConcurrentDictionary> _serverPromptsCache = new(); - private readonly ConcurrentDictionary> _serverResourcesCache = new(); - private readonly ConcurrentDictionary> _serverResourceTemplatesCache = new(); - - public System.Collections.Generic.List? GetServerToolsCache(string serverId) - { - _serverToolsCache.TryGetValue(serverId, out var tools); - return tools; - } - - public void SetServerToolsCache(string serverId, System.Collections.Generic.List tools) - { - _serverToolsCache[serverId] = tools; - } - - public void RemoveServerToolsCache(string serverId) - { - _serverToolsCache.TryRemove(serverId, out _); - } - - public System.Collections.Generic.List? GetServerPromptsCache(string serverId) - { - _serverPromptsCache.TryGetValue(serverId, out var prompts); - return prompts; - } - - public void SetServerPromptsCache(string serverId, System.Collections.Generic.List prompts) - { - _serverPromptsCache[serverId] = prompts; - } - - public void RemoveServerPromptsCache(string serverId) - { - _serverPromptsCache.TryRemove(serverId, out _); - } - - public System.Collections.Generic.List? GetServerResourcesCache(string serverId) - { - _serverResourcesCache.TryGetValue(serverId, out var resources); - return resources; - } - - public void SetServerResourcesCache(string serverId, System.Collections.Generic.List resources) - { - _serverResourcesCache[serverId] = resources; - } - - public void RemoveServerResourcesCache(string serverId) - { - _serverResourcesCache.TryRemove(serverId, out _); - } - - public System.Collections.Generic.List? GetServerResourceTemplatesCache(string serverId) - { - _serverResourceTemplatesCache.TryGetValue(serverId, out var templates); - return templates; - } - - public void SetServerResourceTemplatesCache(string serverId, System.Collections.Generic.List templates) - { - _serverResourceTemplatesCache[serverId] = templates; - } - - public void RemoveServerResourceTemplatesCache(string serverId) - { - _serverResourceTemplatesCache.TryRemove(serverId, out _); - } - - public void RemoveServerCache(string serverId) - { - RemoveServerToolsCache(serverId); - RemoveServerPromptsCache(serverId); - RemoveServerResourcesCache(serverId); - RemoveServerResourceTemplatesCache(serverId); - } - - public void ClearGlobalCache() - { - _serverToolsCache.Clear(); - _serverPromptsCache.Clear(); - _serverResourcesCache.Clear(); - _serverResourceTemplatesCache.Clear(); - } - public void ResetAll() { _logger.LogInformation("Resetting all active MCP client sessions due to configuration change."); diff --git a/Core/Routing/ToolRoutingManager.Cache.cs b/Core/Routing/ToolRoutingManager.Cache.cs new file mode 100644 index 00000000..1c2465f1 --- /dev/null +++ b/Core/Routing/ToolRoutingManager.Cache.cs @@ -0,0 +1,128 @@ +using System; +using System.Collections.Concurrent; +using System.Collections.Generic; +using System.Linq; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using System.Net.Http; +using Microsoft.Extensions.Logging; +using McpRouter.Models; +using McpRouter.Core.Protocol; +using McpRouter.Components.Servers; +using McpRouter.Infrastructure.Persistence; +using Dapper; + +namespace McpRouter.Core.Routing +{ + public partial class ToolRoutingManager + { + /// + /// Asynchronously lists available tools from connected backends or returns bootstrap Meta-Mode tools. + /// + public async Task> ListToolsAsync(string body, bool isMetaMode, IEnumerable> backendConnections, ILogger logger, Func ensureBackendsInitializedAsync, IEnumerable servers, SessionManager? sessionManager = null) + { + if (isMetaMode) + { + return GetMetaModeTools(); + } + + await ensureBackendsInitializedAsync(); + + lock (_cacheLock) + { + if (_isCachePopulated) + { + return new List(_cachedTools); + } + } + + await PopulateToolsCacheAsync(body, backendConnections, logger, servers, sessionManager); + lock (_cacheLock) + { + return new List(_cachedTools); + } + } + + public async Task PopulateToolsCacheAsync(string body, IEnumerable> backendConnections, ILogger logger, IEnumerable servers, SessionManager? sessionManager = null) + { + var allTools = new List(); + + var tasks = new List>(); + + foreach (var entry in backendConnections) + { + var conn = entry.Value; + var serverId = entry.Key; + + tasks.Add(Task.Run(async () => + { + try + { + var reqBody = "{\"jsonrpc\":\"2.0\",\"method\":\"tools/list\",\"id\":\"refresh-list\"}"; + var resp = await conn.SendRequestAsync("tools/list", reqBody); + if (resp.Result != null && resp.Result.Value.TryGetProperty("tools", out var toolsList)) + { + return (serverId, toolsList); + } + } + catch (Exception ex) + { + logger.LogError(ex, "Error listing tools on server {ServerId}", serverId); + } + return (serverId, default(JsonElement)); + })); + } + + var completed = await Task.WhenAll(tasks); + foreach (var item in completed) + { + if (item.Tools.ValueKind == JsonValueKind.Array) + { + var serverTools = new List(); + foreach (var tool in item.Tools.EnumerateArray()) + { + if (tool.TryGetProperty("name", out var nameProp)) + { + var rawToolName = nameProp.GetString() ?? string.Empty; + var exposedName = item.ServerId + "__" + rawToolName; + + _toolRoutingTable[exposedName] = item.ServerId; + + var toolDict = JsonSerializer.Deserialize>(tool.GetRawText()); + if (toolDict != null) + { + toolDict["name"] = exposedName; + if (toolDict.TryGetValue("description", out var desc)) + { + toolDict["description"] = $"[{item.ServerId}] " + desc; + } + + var srv = servers.FirstOrDefault(s => s.Id == item.ServerId); + if (srv != null && (srv.AllowPassThroughAuth || !string.IsNullOrEmpty(srv.DynamicAuthPrompt))) + { + var authPrompt = !string.IsNullOrEmpty(srv.DynamicAuthPrompt) ? srv.DynamicAuthPrompt : "This tool requires a target authentication token. Call with target_auth_token parameter."; + toolDict["description"] = $"{toolDict["description"]}\n\nAUTH REQUIRED: {authPrompt}"; + } + + serverTools.Add(toolDict); + allTools.Add(toolDict); + } + } + } + if (sessionManager != null) + { + sessionManager.SetServerToolsCache(item.ServerId, serverTools); + } + } + } + + lock (_cacheLock) + { + _cachedTools.Clear(); + _cachedTools.AddRange(allTools); + _isCachePopulated = true; + } + } + } +} diff --git a/Core/Routing/ToolRoutingManager.Execution.cs b/Core/Routing/ToolRoutingManager.Execution.cs new file mode 100644 index 00000000..3488997f --- /dev/null +++ b/Core/Routing/ToolRoutingManager.Execution.cs @@ -0,0 +1,293 @@ +using System; +using System.Collections.Concurrent; +using System.Collections.Generic; +using System.Linq; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using System.Net.Http; +using Microsoft.Extensions.Logging; +using McpRouter.Models; +using McpRouter.Core.Protocol; +using McpRouter.Components.Servers; +using McpRouter.Infrastructure.Persistence; +using Dapper; + +namespace McpRouter.Core.Routing +{ + public partial class ToolRoutingManager + { + public async Task CallToolAsync( + string toolName, + string body, + IDbConnectionFactory dbFactory, + ConcurrentDictionary backendConnections, + IEnumerable servers, + ILogger logger, + HttpClient httpClient, + IEmbeddingService embeddingService, + Func ensureBackendsInitializedAsync, + Func rewriteRequestJson, + CancellationToken cancellationToken = default, + SessionManager? sessionManager = null, + string? clientSessionId = null, + Func, Task>>? filterAuthorizedToolsAsync = null) + { + try + { + var task = CallToolInternalAsync(toolName, body, dbFactory, backendConnections, servers, logger, httpClient, embeddingService, ensureBackendsInitializedAsync, rewriteRequestJson, cancellationToken, sessionManager, clientSessionId, filterAuthorizedToolsAsync); + return await task.WaitAsync(cancellationToken); + } + catch (OperationCanceledException) + { + logger.LogWarning("Execution of tool '{ToolName}' was cancelled.", toolName); + return new + { + isError = true, + content = new[] { + new { + type = "text", + text = "Error: request was cancelled by the client." + } + } + }; + } + } + + private async Task CallToolInternalAsync( + string toolName, + string body, + IDbConnectionFactory dbFactory, + ConcurrentDictionary backendConnections, + IEnumerable servers, + ILogger logger, + HttpClient httpClient, + IEmbeddingService embeddingService, + Func ensureBackendsInitializedAsync, + Func rewriteRequestJson, + CancellationToken cancellationToken, + SessionManager? sessionManager, + string? clientSessionId, + Func, Task>>? filterAuthorizedToolsAsync = null) + { + await ensureBackendsInitializedAsync(); + + if (toolName == "search_tools") + { + using var doc = JsonDocument.Parse(body); + var root = doc.RootElement; + + string query = ""; + if (root.TryGetProperty("params", out var paramsProp) && + paramsProp.TryGetProperty("arguments", out var argsProp) && + argsProp.TryGetProperty("query", out var queryProp)) + { + query = queryProp.GetString() ?? ""; + } + + var tools = new List(); + lock (_cacheLock) + { + tools.AddRange(_cachedTools); + } + + if (filterAuthorizedToolsAsync != null) + { + tools = await filterAuthorizedToolsAsync(tools); + } + + var results = await SemanticSearchService.SearchToolsSemanticAsync(query, tools, embeddingService, logger); + return new + { + content = new[] { + new { + type = "text", + text = JsonSerializer.Serialize(results, new JsonSerializerOptions { WriteIndented = true }) + } + } + }; + } + else if (toolName == "execute_tool") + { + using var doc = JsonDocument.Parse(body); + var root = doc.RootElement; + + string targetName = ""; + JsonElement targetArgs = default; + string? targetAuthToken = null; + + if (root.TryGetProperty("params", out var paramsProp) && + paramsProp.TryGetProperty("arguments", out var argsProp)) + { + if (argsProp.TryGetProperty("name", out var nameProp)) + { + targetName = nameProp.GetString() ?? ""; + } + if (argsProp.TryGetProperty("arguments", out var targetArgsProp)) + { + targetArgs = targetArgsProp.Clone(); + } + if (argsProp.TryGetProperty("target_auth_token", out var targetAuthTokenProp)) + { + targetAuthToken = targetAuthTokenProp.GetString(); + } + } + + if (string.IsNullOrEmpty(targetName)) + { + return new + { + isError = true, + content = new[] { + new { + type = "text", + text = "Error: target tool name is required." + } + } + }; + } + + var activeServerIds = servers.Where(s => s.Enabled).Select(s => s.Id).ToList(); + if (!SecurityValidationHelper.ValidateToolOrPromptName(targetName, activeServerIds)) + { + return new + { + isError = true, + content = new[] { + new { + type = "text", + text = $"Security Error: Invalid or spoofed namespaced identifier in execute_tool: '{targetName}'." + } + } + }; + } + + var targetPayload = new + { + jsonrpc = "2.0", + method = "tools/call", + @params = new + { + name = targetName, + arguments = targetArgs.ValueKind == JsonValueKind.Undefined ? (object)new Dictionary() : targetArgs + } + }; + var targetBody = JsonSerializer.Serialize(targetPayload); + + try + { + var result = await ExecuteTargetToolAsync(targetName, targetBody, targetAuthToken, dbFactory, backendConnections, servers, logger, httpClient, ensureBackendsInitializedAsync, rewriteRequestJson, cancellationToken, sessionManager, clientSessionId); + return result; + } + catch (Exception ex) + { + return new + { + isError = true, + content = new[] { + new { + type = "text", + text = $"Error executing target tool {targetName}: {ex.Message}" + } + } + }; + } + } + + return await ExecuteTargetToolAsync(toolName, body, null, dbFactory, backendConnections, servers, logger, httpClient, ensureBackendsInitializedAsync, rewriteRequestJson, cancellationToken, sessionManager, clientSessionId); + } + + private async Task ExecuteTargetToolAsync( + string toolName, + string body, + string? targetAuthToken, + IDbConnectionFactory dbFactory, + ConcurrentDictionary backendConnections, + IEnumerable servers, + ILogger logger, + HttpClient httpClient, + Func ensureBackendsInitializedAsync, + Func rewriteRequestJson, + CancellationToken cancellationToken, + SessionManager? sessionManager, + string? clientSessionId) + { + if (!_toolRoutingTable.ContainsKey(toolName)) + { + logger.LogInformation("Tool '{ToolName}' not found in routing table. Refreshing tools cache...", toolName); + try + { + await PopulateToolsCacheAsync("{\"jsonrpc\":\"2.0\",\"method\":\"tools/list\",\"id\":\"refresh-list\"}", backendConnections, logger, servers); + } + catch (Exception ex) + { + logger.LogError(ex, "Failed to refresh tools cache during CallToolAsync for '{ToolName}'", toolName); + } + } + + if (_toolRoutingTable.TryGetValue(toolName, out var serverId) && backendConnections.TryGetValue(serverId, out var conn)) + { + logger.LogInformation("Routing tool call '{ToolName}' to server '{ServerId}'", toolName, serverId); + + string routingBody = body; + var prefix = serverId + "__"; + if (toolName.StartsWith(prefix)) + { + var realToolName = toolName.Substring(prefix.Length); + routingBody = rewriteRequestJson(body, "name", realToolName); + } + + try + { + var resp = await conn.SendRequestAsync("tools/call", routingBody, targetAuthToken); + if (resp.Error != null) + { + var transformed = ToolErrorFormatter.TransformError(resp.Error, toolName, serverId); + return new + { + isError = true, + content = new[] { + new { + type = "text", + text = transformed + } + } + }; + } + return resp; + } + catch (System.Net.Http.HttpRequestException ex) when (ex.StatusCode == System.Net.HttpStatusCode.Unauthorized) + { + var srv = servers.FirstOrDefault(s => s.Id == serverId); + var prompt = (srv != null && !string.IsNullOrEmpty(srv.DynamicAuthPrompt)) ? srv.DynamicAuthPrompt : "401 Unauthorized. Please provide a valid target_auth_token via execute_tool."; + return new + { + isError = true, + content = new[] { + new { + type = "text", + text = prompt + } + } + }; + } + catch (Exception ex) + { + var transformed = ToolErrorFormatter.TransformException(ex, toolName, serverId); + return new + { + isError = true, + content = new[] { + new { + type = "text", + text = transformed + } + } + }; + } + } + + throw new KeyNotFoundException($"Tool {toolName} not found in routing table."); + } + } +} diff --git a/Core/Routing/ToolRoutingManager.cs b/Core/Routing/ToolRoutingManager.cs index 7a2562e0..d531a446 100644 --- a/Core/Routing/ToolRoutingManager.cs +++ b/Core/Routing/ToolRoutingManager.cs @@ -18,7 +18,7 @@ namespace McpRouter.Core.Routing /// /// Manages backend tool listing, caching, namespaced routing tables, and tool invocation execution. /// - public class ToolRoutingManager + public partial class ToolRoutingManager { private readonly ConcurrentDictionary _toolRoutingTable = new(); private readonly List _cachedTools = new(); @@ -30,33 +30,6 @@ public class ToolRoutingManager /// public ConcurrentDictionary ToolRoutingTable => _toolRoutingTable; - /// - /// Asynchronously lists available tools from connected backends or returns bootstrap Meta-Mode tools. - /// - public async Task> ListToolsAsync(string body, bool isMetaMode, IEnumerable> backendConnections, ILogger logger, Func ensureBackendsInitializedAsync, IEnumerable servers, SessionManager? sessionManager = null) - { - if (isMetaMode) - { - return GetMetaModeTools(); - } - - await ensureBackendsInitializedAsync(); - - lock (_cacheLock) - { - if (_isCachePopulated) - { - return new List(_cachedTools); - } - } - - await PopulateToolsCacheAsync(body, backendConnections, logger, servers, sessionManager); - lock (_cacheLock) - { - return new List(_cachedTools); - } - } - public static List GetMetaModeTools() { return new List @@ -94,87 +67,6 @@ public static List GetMetaModeTools() }; } - public async Task PopulateToolsCacheAsync(string body, IEnumerable> backendConnections, ILogger logger, IEnumerable servers, SessionManager? sessionManager = null) - { - var allTools = new List(); - - var tasks = new List>(); - - foreach (var entry in backendConnections) - { - var conn = entry.Value; - var serverId = entry.Key; - - tasks.Add(Task.Run(async () => - { - try - { - var reqBody = "{\"jsonrpc\":\"2.0\",\"method\":\"tools/list\",\"id\":\"refresh-list\"}"; - var resp = await conn.SendRequestAsync("tools/list", reqBody); - if (resp.Result != null && resp.Result.Value.TryGetProperty("tools", out var toolsList)) - { - return (serverId, toolsList); - } - } - catch (Exception ex) - { - logger.LogError(ex, "Error listing tools on server {ServerId}", serverId); - } - return (serverId, default(JsonElement)); - })); - } - - var completed = await Task.WhenAll(tasks); - foreach (var item in completed) - { - if (item.Tools.ValueKind == JsonValueKind.Array) - { - var serverTools = new List(); - foreach (var tool in item.Tools.EnumerateArray()) - { - if (tool.TryGetProperty("name", out var nameProp)) - { - var rawToolName = nameProp.GetString() ?? string.Empty; - var exposedName = item.ServerId + "__" + rawToolName; - - _toolRoutingTable[exposedName] = item.ServerId; - - var toolDict = JsonSerializer.Deserialize>(tool.GetRawText()); - if (toolDict != null) - { - toolDict["name"] = exposedName; - if (toolDict.TryGetValue("description", out var desc)) - { - toolDict["description"] = $"[{item.ServerId}] " + desc; - } - - var srv = servers.FirstOrDefault(s => s.Id == item.ServerId); - if (srv != null && (srv.AllowPassThroughAuth || !string.IsNullOrEmpty(srv.DynamicAuthPrompt))) - { - var authPrompt = !string.IsNullOrEmpty(srv.DynamicAuthPrompt) ? srv.DynamicAuthPrompt : "This tool requires a target authentication token. Call with target_auth_token parameter."; - toolDict["description"] = $"{toolDict["description"]}\n\nAUTH REQUIRED: {authPrompt}"; - } - - serverTools.Add(toolDict); - allTools.Add(toolDict); - } - } - } - if (sessionManager != null) - { - sessionManager.SetServerToolsCache(item.ServerId, serverTools); - } - } - } - - lock (_cacheLock) - { - _cachedTools.Clear(); - _cachedTools.AddRange(allTools); - _isCachePopulated = true; - } - } - public void InvalidateCache() { lock (_cacheLock) @@ -184,279 +76,6 @@ public void InvalidateCache() } } - public async Task CallToolAsync( - string toolName, - string body, - IDbConnectionFactory dbFactory, - ConcurrentDictionary backendConnections, - IEnumerable servers, - ILogger logger, - HttpClient httpClient, - IEmbeddingService embeddingService, - Func ensureBackendsInitializedAsync, - Func rewriteRequestJson, - CancellationToken cancellationToken = default, - SessionManager? sessionManager = null, - string? clientSessionId = null, - Func, Task>>? filterAuthorizedToolsAsync = null) - { - try - { - var task = CallToolInternalAsync(toolName, body, dbFactory, backendConnections, servers, logger, httpClient, embeddingService, ensureBackendsInitializedAsync, rewriteRequestJson, cancellationToken, sessionManager, clientSessionId, filterAuthorizedToolsAsync); - return await task.WaitAsync(cancellationToken); - } - catch (OperationCanceledException) - { - logger.LogWarning("Execution of tool '{ToolName}' was cancelled.", toolName); - return new - { - isError = true, - content = new[] { - new { - type = "text", - text = "Error: request was cancelled by the client." - } - } - }; - } - } - - private async Task CallToolInternalAsync( - string toolName, - string body, - IDbConnectionFactory dbFactory, - ConcurrentDictionary backendConnections, - IEnumerable servers, - ILogger logger, - HttpClient httpClient, - IEmbeddingService embeddingService, - Func ensureBackendsInitializedAsync, - Func rewriteRequestJson, - CancellationToken cancellationToken, - SessionManager? sessionManager, - string? clientSessionId, - Func, Task>>? filterAuthorizedToolsAsync = null) - { - await ensureBackendsInitializedAsync(); - - if (toolName == "search_tools") - { - using var doc = JsonDocument.Parse(body); - var root = doc.RootElement; - - string query = ""; - if (root.TryGetProperty("params", out var paramsProp) && - paramsProp.TryGetProperty("arguments", out var argsProp) && - argsProp.TryGetProperty("query", out var queryProp)) - { - query = queryProp.GetString() ?? ""; - } - - var tools = new List(); - lock (_cacheLock) - { - tools.AddRange(_cachedTools); - } - - if (filterAuthorizedToolsAsync != null) - { - tools = await filterAuthorizedToolsAsync(tools); - } - - var results = await SemanticSearchService.SearchToolsSemanticAsync(query, tools, embeddingService, logger); - return new - { - content = new[] { - new { - type = "text", - text = JsonSerializer.Serialize(results, new JsonSerializerOptions { WriteIndented = true }) - } - } - }; - } - else if (toolName == "execute_tool") - { - using var doc = JsonDocument.Parse(body); - var root = doc.RootElement; - - string targetName = ""; - JsonElement targetArgs = default; - string? targetAuthToken = null; - - if (root.TryGetProperty("params", out var paramsProp) && - paramsProp.TryGetProperty("arguments", out var argsProp)) - { - if (argsProp.TryGetProperty("name", out var nameProp)) - { - targetName = nameProp.GetString() ?? ""; - } - if (argsProp.TryGetProperty("arguments", out var targetArgsProp)) - { - targetArgs = targetArgsProp.Clone(); - } - if (argsProp.TryGetProperty("target_auth_token", out var targetAuthTokenProp)) - { - targetAuthToken = targetAuthTokenProp.GetString(); - } - } - - if (string.IsNullOrEmpty(targetName)) - { - return new - { - isError = true, - content = new[] { - new { - type = "text", - text = "Error: target tool name is required." - } - } - }; - } - - var activeServerIds = servers.Where(s => s.Enabled).Select(s => s.Id).ToList(); - if (!SecurityValidationHelper.ValidateToolOrPromptName(targetName, activeServerIds)) - { - return new - { - isError = true, - content = new[] { - new { - type = "text", - text = $"Security Error: Invalid or spoofed namespaced identifier in execute_tool: '{targetName}'." - } - } - }; - } - - var targetPayload = new - { - jsonrpc = "2.0", - method = "tools/call", - @params = new - { - name = targetName, - arguments = targetArgs.ValueKind == JsonValueKind.Undefined ? (object)new Dictionary() : targetArgs - } - }; - var targetBody = JsonSerializer.Serialize(targetPayload); - - try - { - var result = await ExecuteTargetToolAsync(targetName, targetBody, targetAuthToken, dbFactory, backendConnections, servers, logger, httpClient, ensureBackendsInitializedAsync, rewriteRequestJson, cancellationToken, sessionManager, clientSessionId); - return result; - } - catch (Exception ex) - { - return new - { - isError = true, - content = new[] { - new { - type = "text", - text = $"Error executing target tool {targetName}: {ex.Message}" - } - } - }; - } - } - - return await ExecuteTargetToolAsync(toolName, body, null, dbFactory, backendConnections, servers, logger, httpClient, ensureBackendsInitializedAsync, rewriteRequestJson, cancellationToken, sessionManager, clientSessionId); - } - - private async Task ExecuteTargetToolAsync( - string toolName, - string body, - string? targetAuthToken, - IDbConnectionFactory dbFactory, - ConcurrentDictionary backendConnections, - IEnumerable servers, - ILogger logger, - HttpClient httpClient, - Func ensureBackendsInitializedAsync, - Func rewriteRequestJson, - CancellationToken cancellationToken, - SessionManager? sessionManager, - string? clientSessionId) - { - if (!_toolRoutingTable.ContainsKey(toolName)) - { - logger.LogInformation("Tool '{ToolName}' not found in routing table. Refreshing tools cache...", toolName); - try - { - await PopulateToolsCacheAsync("{\"jsonrpc\":\"2.0\",\"method\":\"tools/list\",\"id\":\"refresh-list\"}", backendConnections, logger, servers); - } - catch (Exception ex) - { - logger.LogError(ex, "Failed to refresh tools cache during CallToolAsync for '{ToolName}'", toolName); - } - } - - if (_toolRoutingTable.TryGetValue(toolName, out var serverId) && backendConnections.TryGetValue(serverId, out var conn)) - { - logger.LogInformation("Routing tool call '{ToolName}' to server '{ServerId}'", toolName, serverId); - - string routingBody = body; - var prefix = serverId + "__"; - if (toolName.StartsWith(prefix)) - { - var realToolName = toolName.Substring(prefix.Length); - routingBody = rewriteRequestJson(body, "name", realToolName); - } - - try - { - var resp = await conn.SendRequestAsync("tools/call", routingBody, targetAuthToken); - if (resp.Error != null) - { - var transformed = ToolErrorFormatter.TransformError(resp.Error, toolName, serverId); - return new - { - isError = true, - content = new[] { - new { - type = "text", - text = transformed - } - } - }; - } - return resp; - } - catch (System.Net.Http.HttpRequestException ex) when (ex.StatusCode == System.Net.HttpStatusCode.Unauthorized) - { - var srv = servers.FirstOrDefault(s => s.Id == serverId); - var prompt = (srv != null && !string.IsNullOrEmpty(srv.DynamicAuthPrompt)) ? srv.DynamicAuthPrompt : "401 Unauthorized. Please provide a valid target_auth_token via execute_tool."; - return new - { - isError = true, - content = new[] { - new { - type = "text", - text = prompt - } - } - }; - } - catch (Exception ex) - { - var transformed = ToolErrorFormatter.TransformException(ex, toolName, serverId); - return new - { - isError = true, - content = new[] { - new { - type = "text", - text = transformed - } - } - }; - } - } - - throw new KeyNotFoundException($"Tool {toolName} not found in routing table."); - } - public List GetCachedTools() { lock (_cacheLock) diff --git a/McpRouter.Tests/SessionManagerTests.cs b/McpRouter.Tests/SessionManagerTests.cs index b0e03e24..a2d746a1 100644 --- a/McpRouter.Tests/SessionManagerTests.cs +++ b/McpRouter.Tests/SessionManagerTests.cs @@ -5,11 +5,13 @@ using Microsoft.Extensions.Logging.Abstractions; using Moq; using Xunit; +using McpRouter.Tests.Attributes; namespace McpRouter.Tests { public class SessionManagerTests { + [Requirement("CORE-101", "Auto-added requirement tracking")] [Fact] public void PerformanceMetrics_And_TotalRequests_IncrementCorrectly() { @@ -33,6 +35,7 @@ public void PerformanceMetrics_And_TotalRequests_IncrementCorrectly() Assert.Equal(50, manager.TotalDurationMs); } + [Requirement("CORE-101", "Auto-added requirement tracking")] [Fact] public void UpdateBackendStatus_TracksBackendHealth() { diff --git a/McpRouter.Tests/ToolRoutingManagerTests.cs b/McpRouter.Tests/ToolRoutingManagerTests.cs index b32e8b9e..ca59b6b4 100644 --- a/McpRouter.Tests/ToolRoutingManagerTests.cs +++ b/McpRouter.Tests/ToolRoutingManagerTests.cs @@ -15,6 +15,7 @@ using McpRouter.Core.Routing; using Moq; using Xunit; +using McpRouter.Tests.Attributes; using Dapper; namespace McpRouter.Tests @@ -44,6 +45,7 @@ INTEGER DEFAULT 0 return (connection, mockDbFactory.Object); } + [Requirement("CORE-101", "Auto-added requirement tracking")] [Fact] public async Task ListToolsAsync_ReturnsMetaTools_InMetaMode() { @@ -64,6 +66,7 @@ public async Task ListToolsAsync_ReturnsMetaTools_InMetaMode() Assert.Equal(2, tools.Count); } + [Requirement("CORE-101", "Auto-added requirement tracking")] [Fact] public void InvalidateCache_ClearsPopulatedState() { @@ -72,6 +75,7 @@ public void InvalidateCache_ClearsPopulatedState() Assert.Empty(manager.GetCachedTools()); } + [Requirement("CORE-101", "Auto-added requirement tracking")] [Fact] public async Task CallToolAsync_SearchTools_ReturnsSemanticResults() { @@ -100,6 +104,7 @@ public async Task CallToolAsync_SearchTools_ReturnsSemanticResults() Assert.NotNull(result); } + [Requirement("CORE-101", "Auto-added requirement tracking")] [Fact] public async Task CallToolAsync_ExecuteTool_ReturnsError_WhenNameMissing() { @@ -126,6 +131,7 @@ public async Task CallToolAsync_ExecuteTool_ReturnsError_WhenNameMissing() Assert.NotNull(result); } + [Requirement("CORE-101", "Auto-added requirement tracking")] [Fact] public async Task CallToolAsync_ReturnsCancellationError_WhenCancelled() { @@ -154,6 +160,7 @@ public async Task CallToolAsync_ReturnsCancellationError_WhenCancelled() Assert.NotNull(result); } + [Requirement("CORE-101", "Auto-added requirement tracking")] [Fact] public async Task CallToolAsync_ThrowsKeyNotFound_WhenToolNotInRoutingTable() { @@ -177,6 +184,7 @@ await Assert.ThrowsAsync(() => manager.CallToolAsync( )); } + [Requirement("CORE-101", "Auto-added requirement tracking")] [Fact] [Requirement("AUTH-105", "Dynamic Auth Target Pass-Through", Type = RequirementType.Positive, Category = "AUTH")] public async Task ExecuteTargetToolAsync_Catches401_AndReturnsAuthPrompt() From d4e3e7f08ef67ae77d783fd39f85ab4e4a7c09b9 Mon Sep 17 00:00:00 2001 From: Antigravity Agent Date: Sat, 22 Aug 2026 11:35:29 -0500 Subject: [PATCH 02/24] chore(frontend): sweep frontend, split large components and add missing test requirements --- .../e2e/appkey-and-client-lifecycle.spec.ts | 2 + frontend/e2e/client-setup-and-appkeys.spec.ts | 2 + frontend/e2e/dashboard.spec.ts | 2 + frontend/e2e/multi-user-matrix.spec.ts | 2 + frontend/e2e/rbac-and-permissions.spec.ts | 2 + frontend/e2e/rbac-enforcement-flow.spec.ts | 2 + frontend/e2e/server-inspector.spec.ts | 2 + frontend/e2e/server-management.spec.ts | 2 + frontend/e2e/settings.spec.ts | 2 + frontend/e2e/testbench.spec.ts | 2 + .../components/settings/CustomFileModal.tsx | 177 +--------------- .../settings/VisualPromptBuilder.tsx | 191 ++++++++++++++++++ .../components/testbench/TestBenchView.tsx | 88 +++----- .../components/testbench/useTestBenchState.ts | 90 +++++++++ frontend/src/test/api/typedApi.test.ts | 2 + .../src/test/components/ClientModal.test.tsx | 2 + .../test/components/ClientSetupGuide.test.tsx | 2 + .../test/components/DashboardView.test.tsx | 2 + frontend/src/test/components/Header.test.tsx | 2 + .../test/components/LogsTerminalCard.test.tsx | 2 + .../src/test/components/MappingModal.test.tsx | 2 + .../src/test/components/PolicyModal.test.tsx | 2 + .../test/components/PromptTesterCard.test.tsx | 2 + .../components/ResourceTesterCard.test.tsx | 2 + .../src/test/components/ServerCard.test.tsx | 2 + .../components/ServerInspectModal.test.tsx | 2 + .../src/test/components/ServerModal.test.tsx | 2 + .../src/test/components/SettingsTabs.test.tsx | 2 + .../src/test/components/SettingsView.test.tsx | 2 + .../test/components/SharedComponents.test.tsx | 2 + .../test/components/TestBenchView.test.tsx | 2 + .../test/components/ToolTesterCard.test.tsx | 2 + .../src/test/stores/useProviderStore.test.ts | 2 + frontend/src/test/stores/useUserStore.test.ts | 2 + 34 files changed, 373 insertions(+), 233 deletions(-) create mode 100644 frontend/src/components/settings/VisualPromptBuilder.tsx create mode 100644 frontend/src/components/testbench/useTestBenchState.ts diff --git a/frontend/e2e/appkey-and-client-lifecycle.spec.ts b/frontend/e2e/appkey-and-client-lifecycle.spec.ts index 2519da80..d19c84e0 100644 --- a/frontend/e2e/appkey-and-client-lifecycle.spec.ts +++ b/frontend/e2e/appkey-and-client-lifecycle.spec.ts @@ -1,3 +1,5 @@ +/** @requirement REQ-UI-129 */ + import { test, expect } from '@playwright/test'; test.describe('AppKey and Client Lifecycle Flow', () => { diff --git a/frontend/e2e/client-setup-and-appkeys.spec.ts b/frontend/e2e/client-setup-and-appkeys.spec.ts index 923abf25..e3206fa1 100644 --- a/frontend/e2e/client-setup-and-appkeys.spec.ts +++ b/frontend/e2e/client-setup-and-appkeys.spec.ts @@ -1,3 +1,5 @@ +/** @requirement REQ-UI-123 */ + import { test, expect } from '@playwright/test'; test.describe('Client Setup & App Key Management Flow', () => { diff --git a/frontend/e2e/dashboard.spec.ts b/frontend/e2e/dashboard.spec.ts index de66b23b..73c7f324 100644 --- a/frontend/e2e/dashboard.spec.ts +++ b/frontend/e2e/dashboard.spec.ts @@ -1,3 +1,5 @@ +/** @requirement REQ-UI-124 */ + import { test, expect } from '@playwright/test'; test.describe('Dashboard & Navigation Flow', () => { diff --git a/frontend/e2e/multi-user-matrix.spec.ts b/frontend/e2e/multi-user-matrix.spec.ts index 93b6cdf8..e1682b01 100644 --- a/frontend/e2e/multi-user-matrix.spec.ts +++ b/frontend/e2e/multi-user-matrix.spec.ts @@ -1,3 +1,5 @@ +/** @requirement REQ-UI-125 */ + import { test, expect } from './fixtures/userContexts'; test.describe('Multi-User Context Matrix Flow (Issue #50)', () => { diff --git a/frontend/e2e/rbac-and-permissions.spec.ts b/frontend/e2e/rbac-and-permissions.spec.ts index 79b233b4..dc8de97d 100644 --- a/frontend/e2e/rbac-and-permissions.spec.ts +++ b/frontend/e2e/rbac-and-permissions.spec.ts @@ -1,3 +1,5 @@ +/** @requirement REQ-UI-127 */ + import { test, expect } from '@playwright/test'; test.describe('RBAC Access Control & Policy Modal Flow', () => { diff --git a/frontend/e2e/rbac-enforcement-flow.spec.ts b/frontend/e2e/rbac-enforcement-flow.spec.ts index a61aafaf..513a81d3 100644 --- a/frontend/e2e/rbac-enforcement-flow.spec.ts +++ b/frontend/e2e/rbac-enforcement-flow.spec.ts @@ -1,3 +1,5 @@ +/** @requirement REQ-UI-120 */ + import { test, expect } from '@playwright/test'; test.describe('RBAC Policy and Group/SID Mapping Lifecycle Flow', () => { diff --git a/frontend/e2e/server-inspector.spec.ts b/frontend/e2e/server-inspector.spec.ts index 33c86424..e9453140 100644 --- a/frontend/e2e/server-inspector.spec.ts +++ b/frontend/e2e/server-inspector.spec.ts @@ -1,3 +1,5 @@ +/** @requirement REQ-UI-126 */ + import { test, expect } from '@playwright/test'; test.describe('Server Inspector Modal Flow', () => { diff --git a/frontend/e2e/server-management.spec.ts b/frontend/e2e/server-management.spec.ts index 2446b564..9476259e 100644 --- a/frontend/e2e/server-management.spec.ts +++ b/frontend/e2e/server-management.spec.ts @@ -1,3 +1,5 @@ +/** @requirement REQ-UI-121 */ + import { test, expect } from '@playwright/test'; test.describe('Server Registration & Secret Providers Flow', () => { diff --git a/frontend/e2e/settings.spec.ts b/frontend/e2e/settings.spec.ts index 89e97981..41187b6b 100644 --- a/frontend/e2e/settings.spec.ts +++ b/frontend/e2e/settings.spec.ts @@ -1,3 +1,5 @@ +/** @requirement REQ-UI-122 */ + import { test, expect } from '@playwright/test'; test.describe('Settings View Flow', () => { diff --git a/frontend/e2e/testbench.spec.ts b/frontend/e2e/testbench.spec.ts index 352321cf..8a45ebe7 100644 --- a/frontend/e2e/testbench.spec.ts +++ b/frontend/e2e/testbench.spec.ts @@ -1,3 +1,5 @@ +/** @requirement REQ-UI-128 */ + import { test, expect } from '@playwright/test'; test.describe('Test Bench View Flow', () => { diff --git a/frontend/src/components/settings/CustomFileModal.tsx b/frontend/src/components/settings/CustomFileModal.tsx index 5e4d6772..c5777ddb 100644 --- a/frontend/src/components/settings/CustomFileModal.tsx +++ b/frontend/src/components/settings/CustomFileModal.tsx @@ -1,19 +1,8 @@ import React, { useState } from 'react'; import { useSettingsStore } from '../../stores/useSettingsStore'; import { showToast } from '../../stores/useToastStore'; +import { VisualPromptBuilder, ArgBuilderItem, MsgBuilderItem } from './VisualPromptBuilder'; -interface ArgBuilderItem { - id: string; - name: string; - description: string; - required: boolean; -} - -interface MsgBuilderItem { - id: string; - role: 'user' | 'assistant'; - text: string; -} const initializePromptBuilder = (meta: { type: 'prompts' | 'resources'; name: string } | null, content: string) => { if (meta) { @@ -135,30 +124,6 @@ const CustomFileModalDialog: React.FC = () => { setActiveFileModalTab(tab); }; - const addArgument = () => { - setBuilderArgs([...builderArgs, { id: Math.random().toString(), name: '', description: '', required: false }]); - }; - - const removeArgument = (id: string) => { - setBuilderArgs(builderArgs.filter((a) => a.id !== id)); - }; - - const updateArgument = (id: string, field: keyof ArgBuilderItem, val: any) => { - setBuilderArgs(builderArgs.map((a) => (a.id === id ? { ...a, [field]: val } : a))); - }; - - const addMessage = (role: 'user' | 'assistant') => { - setBuilderMsgs([...builderMsgs, { id: Math.random().toString(), role, text: '' }]); - }; - - const removeMessage = (id: string) => { - setBuilderMsgs(builderMsgs.filter((m) => m.id !== id)); - }; - - const updateMessageText = (id: string, text: string) => { - setBuilderMsgs(builderMsgs.map((m) => (m.id === id ? { ...m, text } : m))); - }; - const handleSave = async (e: React.FormEvent) => { e.preventDefault(); if (!fileName.trim()) { @@ -277,138 +242,14 @@ const CustomFileModalDialog: React.FC = () => { /> ) : ( -
-
- - setPromptDesc(e.target.value)} - /> -
- -
-

- Prompt Arguments / Variables -

- -
- - {builderArgs.length === 0 ? ( -
- No arguments defined. -
- ) : ( -
- {builderArgs.map((arg) => ( -
- updateArgument(arg.id, 'name', e.target.value)} - style={{ fontSize: '12px' }} - /> - updateArgument(arg.id, 'description', e.target.value)} - style={{ fontSize: '12px' }} - /> - - -
- ))} -
- )} - -
-

- Messages Sequence -

-
- - -
-
- -
- {builderMsgs.map((msg, index) => ( -
-
- - #{index + 1} {msg.role.toUpperCase()} - - -
-