From 5afedb7edb4af123b508faf43b48c4e66b30788f Mon Sep 17 00:00:00 2001 From: Rockford Lhotka Date: Mon, 16 Feb 2026 22:56:40 -0600 Subject: [PATCH] Add connection retry logic for dead downstream servers When a downstream MCP server crashes, the aggregator now automatically disconnects the stale connection and retries once with a fresh connection instead of permanently caching the dead client. Co-Authored-By: Claude Opus 4.6 --- .../Services/ConnectionManager.cs | 34 +++++ src/McpAggregator.Core/Services/ToolIndex.cs | 5 +- src/McpAggregator.Core/Tools/AdminTools.cs | 5 +- .../Tools/ToolProxyHandler.cs | 5 +- .../Controllers/AdminController.cs | 9 +- .../Services/ExecuteWithRetryTests.cs | 140 ++++++++++++++++++ 6 files changed, 187 insertions(+), 11 deletions(-) create mode 100644 test/McpAggregator.Core.Tests/Services/ExecuteWithRetryTests.cs diff --git a/src/McpAggregator.Core/Services/ConnectionManager.cs b/src/McpAggregator.Core/Services/ConnectionManager.cs index 9cfe6d9..a4d3a0e 100644 --- a/src/McpAggregator.Core/Services/ConnectionManager.cs +++ b/src/McpAggregator.Core/Services/ConnectionManager.cs @@ -96,6 +96,40 @@ public async Task DisconnectAsync(string serverName) } } + public async Task ExecuteWithRetryAsync( + string serverName, + Func> operation, + CancellationToken ct = default) + { + var client = await GetClientAsync(serverName, ct); + try + { + return await operation(client, ct); + } + catch (Exception ex) when (ShouldRetry(ex)) + { + _logger.LogWarning(ex, "Connection to '{Server}' appears broken, reconnecting", serverName); + await DisconnectAsync(serverName); + client = await GetClientAsync(serverName, ct); + return await operation(client, ct); + } + } + + private static bool ShouldRetry(Exception ex) + { + if (ex is OperationCanceledException or AggregatorException) + return false; + + return ex is System.IO.IOException + or System.Net.Http.HttpRequestException + or System.Net.Sockets.SocketException + or ObjectDisposedException + || ex.InnerException is System.IO.IOException + or System.Net.Http.HttpRequestException + or System.Net.Sockets.SocketException + or ObjectDisposedException; + } + public async Task CleanupIdleConnectionsAsync(CancellationToken ct = default) { var cutoff = DateTimeOffset.UtcNow - _options.ConnectionIdleTimeout; diff --git a/src/McpAggregator.Core/Services/ToolIndex.cs b/src/McpAggregator.Core/Services/ToolIndex.cs index a95b780..745a9e0 100644 --- a/src/McpAggregator.Core/Services/ToolIndex.cs +++ b/src/McpAggregator.Core/Services/ToolIndex.cs @@ -4,6 +4,7 @@ using McpAggregator.Core.Models; using Microsoft.Extensions.Logging; using Microsoft.Extensions.Options; +using ModelContextProtocol.Client; namespace McpAggregator.Core.Services; @@ -123,8 +124,8 @@ public async Task> GetToolsForServerAsync(string serverName, Ca return cached.Tools; } - var client = await _connectionManager.GetClientAsync(serverName, ct); - var mcpTools = await client.ListToolsAsync(cancellationToken: ct); + var mcpTools = await _connectionManager.ExecuteWithRetryAsync>(serverName, + async (client, token) => await client.ListToolsAsync(cancellationToken: token), ct); var tools = mcpTools.Select(t => new ToolDetail { diff --git a/src/McpAggregator.Core/Tools/AdminTools.cs b/src/McpAggregator.Core/Tools/AdminTools.cs index 619c079..dfaaf5c 100644 --- a/src/McpAggregator.Core/Tools/AdminTools.cs +++ b/src/McpAggregator.Core/Tools/AdminTools.cs @@ -2,6 +2,7 @@ using System.Text.Json; using McpAggregator.Core.Models; using McpAggregator.Core.Services; +using ModelContextProtocol.Client; using ModelContextProtocol.Server; namespace McpAggregator.Core.Tools; @@ -136,8 +137,8 @@ public static async Task UpdateSkill( try { - var client = await connectionManager.GetClientAsync(server.Name, ct); - var mcpTools = await client.ListToolsAsync(cancellationToken: ct); + var mcpTools = await connectionManager.ExecuteWithRetryAsync>(server.Name, + async (client, token) => await client.ListToolsAsync(cancellationToken: token), ct); var toolSummaries = mcpTools.Select(t => new ToolSummary { diff --git a/src/McpAggregator.Core/Tools/ToolProxyHandler.cs b/src/McpAggregator.Core/Tools/ToolProxyHandler.cs index 83d0716..0169c03 100644 --- a/src/McpAggregator.Core/Tools/ToolProxyHandler.cs +++ b/src/McpAggregator.Core/Tools/ToolProxyHandler.cs @@ -30,8 +30,6 @@ public async Task InvokeAsync( string? argumentsJson, CancellationToken ct = default) { - var client = await _connectionManager.GetClientAsync(serverName, ct); - IReadOnlyDictionary? args = null; if (!string.IsNullOrWhiteSpace(argumentsJson)) { @@ -45,7 +43,8 @@ public async Task InvokeAsync( try { - var result = await client.CallToolAsync(toolName, args, cancellationToken: cts.Token); + var result = await _connectionManager.ExecuteWithRetryAsync(serverName, + async (client, token) => await client.CallToolAsync(toolName, args, cancellationToken: token), cts.Token); var textContent = result.Content .OfType() diff --git a/src/McpAggregator.HttpServer/Controllers/AdminController.cs b/src/McpAggregator.HttpServer/Controllers/AdminController.cs index 23e10f2..0fb5376 100644 --- a/src/McpAggregator.HttpServer/Controllers/AdminController.cs +++ b/src/McpAggregator.HttpServer/Controllers/AdminController.cs @@ -1,6 +1,7 @@ using McpAggregator.Core.Models; using McpAggregator.Core.Services; using Microsoft.AspNetCore.Mvc; +using ModelContextProtocol.Client; namespace McpAggregator.HttpServer.Controllers; @@ -44,8 +45,8 @@ public async Task RegisterServer([FromBody] RegisterServerRequest { try { - var client = await _connectionManager.GetClientAsync(server.Name, ct); - var mcpTools = await client.ListToolsAsync(cancellationToken: ct); + var mcpTools = await _connectionManager.ExecuteWithRetryAsync>(server.Name, + async (client, token) => await client.ListToolsAsync(cancellationToken: token), ct); var toolSummaries = mcpTools.Select(t => new ToolSummary { @@ -80,8 +81,8 @@ public async Task RegenerateSummary(string name, CancellationToke try { - var client = await _connectionManager.GetClientAsync(server.Name, ct); - var mcpTools = await client.ListToolsAsync(cancellationToken: ct); + var mcpTools = await _connectionManager.ExecuteWithRetryAsync>(server.Name, + async (client, token) => await client.ListToolsAsync(cancellationToken: token), ct); var toolSummaries = mcpTools.Select(t => new ToolSummary { diff --git a/test/McpAggregator.Core.Tests/Services/ExecuteWithRetryTests.cs b/test/McpAggregator.Core.Tests/Services/ExecuteWithRetryTests.cs new file mode 100644 index 0000000..5c0d990 --- /dev/null +++ b/test/McpAggregator.Core.Tests/Services/ExecuteWithRetryTests.cs @@ -0,0 +1,140 @@ +using System.IO; +using System.Net.Sockets; +using McpAggregator.Core.Configuration; +using McpAggregator.Core.Exceptions; +using McpAggregator.Core.Models; +using McpAggregator.Core.Services; +using McpAggregator.Core.Storage; +using McpAggregator.Core.Tests.Helpers; +using Microsoft.Extensions.Logging.Abstractions; +using Rocks; + +namespace McpAggregator.Core.Tests.Services; + +[TestClass] +public class ExecuteWithRetryTests +{ + private static (ConnectionManager Manager, ServerRegistry Registry) CreateManager( + params RegisteredServer[] servers) + { + var expectations = new IRegistryPersistenceCreateExpectations(); + expectations.Setups.LoadAsync(Arg.Any()) + .ReturnValue(Task.FromResult(new RegistryData { Servers = [.. servers] })); + expectations.Setups.SaveAsync(Arg.Any(), Arg.Any()) + .ReturnValue(Task.CompletedTask); + var persistence = expectations.Instance(); + + var registry = new ServerRegistry( + persistence, + TestHelpers.OptionsOf(new AggregatorOptions()), + TestHelpers.NullLoggerOf()); + + var manager = new ConnectionManager( + registry, + TestHelpers.OptionsOf(new AggregatorOptions()), + NullLoggerFactory.Instance, + TestHelpers.NullLoggerOf()); + + return (manager, registry); + } + + [TestMethod] + public async Task ExecuteWithRetryAsync_SucceedsOnFirstTry_DoesNotRetry() + { + // We can't easily mock McpClient (concrete class), so we test ShouldRetry + // indirectly by verifying non-retryable exceptions propagate. + var (manager, registry) = CreateManager(TestHelpers.StdioServer("srv")); + await registry.EnsureLoadedAsync(); + + // GetClientAsync will fail with ServerUnavailableException (can't actually connect in tests) + // but that's wrapped by ConnectAsync — verify it propagates without retry + await Assert.ThrowsExceptionAsync( + () => manager.ExecuteWithRetryAsync("srv", + (client, ct) => Task.FromResult("ok"))); + } + + [TestMethod] + public void ShouldRetry_ReturnsFalse_ForOperationCanceledException() + { + Assert.IsFalse(InvokeShouldRetry(new OperationCanceledException())); + } + + [TestMethod] + public void ShouldRetry_ReturnsFalse_ForAggregatorException() + { + Assert.IsFalse(InvokeShouldRetry(new AggregatorException("test"))); + } + + [TestMethod] + public void ShouldRetry_ReturnsFalse_ForServerNotFoundException() + { + Assert.IsFalse(InvokeShouldRetry(new ServerNotFoundException("srv"))); + } + + [TestMethod] + public void ShouldRetry_ReturnsFalse_ForToolExecutionException() + { + Assert.IsFalse(InvokeShouldRetry(new ToolExecutionException("srv", "tool", "error"))); + } + + [TestMethod] + public void ShouldRetry_ReturnsTrue_ForIOException() + { + Assert.IsTrue(InvokeShouldRetry(new IOException("pipe broken"))); + } + + [TestMethod] + public void ShouldRetry_ReturnsTrue_ForSocketException() + { + Assert.IsTrue(InvokeShouldRetry(new SocketException())); + } + + [TestMethod] + public void ShouldRetry_ReturnsTrue_ForObjectDisposedException() + { + Assert.IsTrue(InvokeShouldRetry(new ObjectDisposedException("client"))); + } + + [TestMethod] + public void ShouldRetry_ReturnsTrue_ForHttpRequestException() + { + Assert.IsTrue(InvokeShouldRetry(new HttpRequestException("connection refused"))); + } + + [TestMethod] + public void ShouldRetry_ReturnsTrue_ForWrappedIOException() + { + var ex = new InvalidOperationException("outer", new IOException("inner")); + Assert.IsTrue(InvokeShouldRetry(ex)); + } + + [TestMethod] + public void ShouldRetry_ReturnsTrue_ForWrappedSocketException() + { + var ex = new InvalidOperationException("outer", new SocketException()); + Assert.IsTrue(InvokeShouldRetry(ex)); + } + + [TestMethod] + public void ShouldRetry_ReturnsFalse_ForGenericException() + { + Assert.IsFalse(InvokeShouldRetry(new InvalidOperationException("generic"))); + } + + [TestMethod] + public void ShouldRetry_ReturnsFalse_ForArgumentException() + { + Assert.IsFalse(InvokeShouldRetry(new ArgumentException("bad arg"))); + } + + /// + /// Invoke the private static ShouldRetry method via reflection. + /// + private static bool InvokeShouldRetry(Exception ex) + { + var method = typeof(ConnectionManager).GetMethod("ShouldRetry", + System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Static) + ?? throw new InvalidOperationException("ShouldRetry method not found"); + return (bool)method.Invoke(null, [ex])!; + } +}