From aadb4efb291b90fdf0f385ba35744d0549552dc6 Mon Sep 17 00:00:00 2001 From: "mars.yu" Date: Fri, 28 Aug 2026 15:59:12 +0800 Subject: [PATCH 1/2] Pool MCP connections per scope, keyed by the identity they carry GetMcpClientAsync built a fresh transport and a fresh McpClient on every call, and Dispose did nothing, so a turn that listed a server's tools and then called three of them opened four connections and closed none of them. Pooling a connection means reusing whatever headers IMcpClientHeaderProvider answered with, and that is an identity. The pool is therefore an instance field of this class, which is registered per DI scope -- one HTTP request, one crontab run, one queued message -- so everything sharing a pool is already the same caller, and one user's connection cannot be handed to another. That is structural rather than a rule someone has to remember. The pool key also folds in a SHA-256 of the headers a connection opens with, so the guarantee survives this class later being registered with a longer lifetime: two credentials land on two entries even inside one pool. It is a hash of secrets, so it is never logged, and a test pins that down along with the three identities OneBrainMcpHeaderProvider can answer with never sharing an entry. Headers are now resolved once and handed to both the key and the transport. Resolving separately for each let the two disagree, and the key is the thing keeping one caller's connection away from another. Entries hold Lazy> so concurrent callers wanting the same server open one connection between them rather than one each. A failed connection is removed rather than cached, and McpToolExecutor now drops the pooled client when a call fails: keeping a dead one fails every remaining call in the scope, while discarding a live one costs a single reconnect. McpClient only implements IAsyncDisposable, so the manager implements both disposal interfaces. Async scopes get DisposeAsync; scopes created with CreateScope tear down synchronously and get a bounded wait instead, because a wedged transport must not hang the unit of work that is trying to finish. Co-Authored-By: Claude Opus 5 --- .../BotSharp.Core/BotSharp.Core.csproj | 4 + .../MCP/Managers/McpClientManager.cs | 222 +++++++++++++++++- .../Routing/Executor/MCPToolExecutor.cs | 9 + .../Mcp/McpClientPoolKeyTests.cs | 126 ++++++++++ 4 files changed, 352 insertions(+), 9 deletions(-) create mode 100644 tests/BotSharp.Core.UnitTests/Mcp/McpClientPoolKeyTests.cs diff --git a/src/Infrastructure/BotSharp.Core/BotSharp.Core.csproj b/src/Infrastructure/BotSharp.Core/BotSharp.Core.csproj index 744b2465c..8c3c32f09 100644 --- a/src/Infrastructure/BotSharp.Core/BotSharp.Core.csproj +++ b/src/Infrastructure/BotSharp.Core/BotSharp.Core.csproj @@ -301,4 +301,8 @@ + + + + diff --git a/src/Infrastructure/BotSharp.Core/MCP/Managers/McpClientManager.cs b/src/Infrastructure/BotSharp.Core/MCP/Managers/McpClientManager.cs index 8e21d74ef..0060f38c1 100644 --- a/src/Infrastructure/BotSharp.Core/MCP/Managers/McpClientManager.cs +++ b/src/Infrastructure/BotSharp.Core/MCP/Managers/McpClientManager.cs @@ -1,12 +1,47 @@ +using System.Collections.Concurrent; +using System.Security.Cryptography; using BotSharp.Core.MCP.Settings; using ModelContextProtocol.Client; namespace BotSharp.Core.MCP.Managers; -public class McpClientManager : IDisposable +/// +/// Hands out MCP clients, pooled for the lifetime of the DI scope that resolved this manager. +/// +/// +/// +/// WHY THE POOL IS SCOPED, AND HAS TO STAY SCOPED. A connection carries whatever +/// answered with, which is how a host calls a server as +/// the signed-in user instead of with one shared credential. Reusing a connection therefore +/// reuses an identity. This class is registered per scope — one HTTP request, one crontab run, +/// one queued message — so everything sharing a pool is already the same caller, and one user's +/// connection cannot be handed to another. That is structural, not a rule someone has to +/// remember. +/// +/// +/// The pool key carries a fingerprint of the headers a connection opens with, so the guarantee +/// survives this class later being registered with a longer lifetime: two credentials land on +/// two entries even inside one pool. The fingerprint is a hash of secrets and is never logged. +/// +/// +/// Before pooling, every tool call opened its own connection and none of them were ever closed: +/// this method built a fresh transport per call and did nothing. A turn +/// that listed tools and then called three of them opened four connections and leaked all four. +/// +/// +public class McpClientManager : IDisposable, IAsyncDisposable { + private const string KeySeparator = "|"; + + /// + /// How long a synchronous scope teardown waits for connections to close before giving up. + /// + private static readonly TimeSpan SyncCloseTimeout = TimeSpan.FromSeconds(5); + private readonly IServiceProvider _services; private readonly ILogger _logger; + private readonly ConcurrentDictionary _pool = new(); + private volatile bool _disposed; public McpClientManager( IServiceProvider services, @@ -16,17 +51,95 @@ public McpClientManager( _logger = logger; } + /// + /// The client for , opening one if this scope has not already. + /// Answers null rather than throwing when the server is unknown, disabled or unreachable, + /// which is the contract every caller is already written against. + /// public async Task GetMcpClientAsync(string serverId) { + if (_disposed) + { + return null; + } + + McpServerConfigModel config; + Dictionary? headers; + string key; + try { var settings = _services.GetRequiredService(); - var config = settings.McpServerConfigs.Where(x => x.Id == serverId).FirstOrDefault(); - if (config == null || !config.Enabled) + var found = settings.McpServerConfigs?.FirstOrDefault(x => x.Id == serverId); + if (found == null || !found.Enabled) { return null; } + config = found; + + // Resolved once, here, and handed to the transport below. Resolving separately for + // the key and for the connection would let the two disagree, and the key is the + // thing keeping one caller's connection away from another. + headers = ResolveHeaders(config); + key = BuildPoolKey(serverId, headers); + } + catch (Exception ex) + { + _logger.LogWarning(ex, $"Error when loading mcp client {serverId}"); + return null; + } + + // Lazy rather than a bare Task, so parallel tool calls that all want this server open + // one connection between them instead of one each. + var entry = _pool.GetOrAdd(key, _ => new PooledClient(serverId, new Lazy>( + () => CreateClientAsync(config, headers), + LazyThreadSafetyMode.ExecutionAndPublication))); + + McpClient? client = null; + try + { + client = await entry.Client.Value; + } + catch (Exception ex) + { + _logger.LogWarning(ex, $"Error when loading mcp client {serverId}"); + } + + if (client == null) + { + // A failure must not stay cached, or every later call in this scope gets it back. + // Removed by value, so a retry that already replaced the entry survives. + _pool.TryRemove(new KeyValuePair(key, entry)); + } + + return client; + } + + /// + /// Drops and closes this scope's connection to so the next call + /// opens a fresh one. Call it when a request over that connection failed at the transport + /// level: keeping a dead client fails every remaining call in the turn, while discarding a + /// live one costs a single reconnect, and that asymmetry says always discard. + /// + public async Task InvalidateAsync(string serverId) + { + var prefix = serverId + KeySeparator; + foreach (var key in _pool.Keys.Where(x => x.StartsWith(prefix, StringComparison.Ordinal)).ToList()) + { + if (_pool.TryRemove(key, out var entry)) + { + await CloseAsync(entry); + } + } + } + + private async Task CreateClientAsync(McpServerConfigModel config, Dictionary? headers) + { + try + { + var settings = _services.GetRequiredService(); + IClientTransport? transport = null; if (config.HttpConfig != null) { @@ -34,7 +147,7 @@ public McpClientManager( { Name = config.Name, Endpoint = new Uri(config.HttpConfig.EndPoint), - AdditionalHeaders = ResolveHeaders(config.Id, config.HttpConfig.AdditionalHeaders), + AdditionalHeaders = headers, ConnectionTimeout = config.HttpConfig.ConnectionTimeout }); } @@ -44,7 +157,7 @@ public McpClientManager( { Name = config.Name, Endpoint = new Uri(config.SseConfig.EndPoint), - AdditionalHeaders = ResolveHeaders(config.Id, config.SseConfig.AdditionalHeaders), + AdditionalHeaders = headers, ConnectionTimeout = config.SseConfig.ConnectionTimeout }); } @@ -69,7 +182,7 @@ public McpClientManager( } catch (Exception ex) { - _logger.LogWarning(ex, $"Error when loading mcp client {serverId}"); + _logger.LogWarning(ex, $"Error when loading mcp client {config.Id}"); return null; } } @@ -81,16 +194,107 @@ public McpClientManager( /// /// No provider is registered by default, and a provider is free to answer with what it was /// given, so a host without one — or with one that does not recognise this server — gets the - /// configured headers back untouched. + /// configured headers back untouched. A stdio server opens with no connection headers at + /// all, so the provider is not consulted for one. /// - private Dictionary? ResolveHeaders(string serverId, Dictionary? configured) + private Dictionary? ResolveHeaders(McpServerConfigModel config) { + if (config.HttpConfig == null && config.SseConfig == null) + { + return null; + } + + var configured = config.HttpConfig?.AdditionalHeaders ?? config.SseConfig?.AdditionalHeaders; var provider = _services.GetService(); - return provider == null ? configured : provider.GetHeaders(serverId, configured); + return provider == null ? configured : provider.GetHeaders(config.Id, configured); + } + + /// + /// The server id plus a fingerprint of the headers the connection will carry, so an entry is + /// shared only between calls that authenticate identically. Hashed rather than kept, because + /// those headers hold credentials; the result is treated as a secret and never logged. + /// + internal static string BuildPoolKey(string serverId, Dictionary? headers) + { + if (headers == null || headers.Count == 0) + { + return serverId + KeySeparator; + } + + var canonical = new StringBuilder(); + foreach (var pair in headers.OrderBy(x => x.Key, StringComparer.Ordinal)) + { + canonical.Append(pair.Key).Append(' ').Append(pair.Value).Append('\n'); + } + + var hash = SHA256.HashData(Encoding.UTF8.GetBytes(canonical.ToString())); + return serverId + KeySeparator + Convert.ToHexString(hash); + } + + private async Task CloseAsync(PooledClient entry) + { + try + { + var client = await entry.Client.Value; + if (client != null) + { + await client.DisposeAsync(); + } + } + catch (Exception ex) + { + _logger.LogWarning(ex, $"Error when closing mcp client {entry.ServerId}"); + } + } + + public async ValueTask DisposeAsync() + { + if (_disposed) + { + return; + } + + _disposed = true; + + foreach (var pair in _pool.ToArray()) + { + if (_pool.TryRemove(pair.Key, out var entry)) + { + await CloseAsync(entry); + } + } + + GC.SuppressFinalize(this); } + /// + /// Scopes made with CreateScope — crontab runs, queue consumers — tear down synchronously, + /// and only offers DisposeAsync, so this waits for it. The wait is + /// bounded: a wedged transport must not hang the unit of work that is trying to finish. An + /// async scope, an ASP.NET Core request among them, calls instead + /// and never comes through here. + /// public void Dispose() { + if (_disposed) + { + return; + } + + try + { + if (!DisposeAsync().AsTask().Wait(SyncCloseTimeout)) + { + _logger.LogWarning("Timed out closing pooled MCP clients; leaving them to the transport."); + } + } + catch (Exception ex) + { + _logger.LogWarning(ex, "Error when closing pooled MCP clients."); + } + GC.SuppressFinalize(this); } + + private sealed record PooledClient(string ServerId, Lazy> Client); } diff --git a/src/Infrastructure/BotSharp.Core/Routing/Executor/MCPToolExecutor.cs b/src/Infrastructure/BotSharp.Core/Routing/Executor/MCPToolExecutor.cs index 08be74599..d189bb4b1 100644 --- a/src/Infrastructure/BotSharp.Core/Routing/Executor/MCPToolExecutor.cs +++ b/src/Infrastructure/BotSharp.Core/Routing/Executor/MCPToolExecutor.cs @@ -51,6 +51,15 @@ public async Task ExecuteAsync(RoleDialogModel message) } catch (Exception ex) { + // The connection is pooled for the rest of this scope, so a transport-level failure + // here would poison every later call in the turn. Drop it and let the next call + // reconnect; discarding a connection that was actually fine costs one reconnect. + var clientManager = _services.GetService(); + if (clientManager != null) + { + await clientManager.InvalidateAsync(_mcpServerId); + } + message.Content = $"Error when calling tool {_functionName} of MCP server {_mcpServerId}. {ex.Message}"; return false; } diff --git a/tests/BotSharp.Core.UnitTests/Mcp/McpClientPoolKeyTests.cs b/tests/BotSharp.Core.UnitTests/Mcp/McpClientPoolKeyTests.cs new file mode 100644 index 000000000..a15d4daf7 --- /dev/null +++ b/tests/BotSharp.Core.UnitTests/Mcp/McpClientPoolKeyTests.cs @@ -0,0 +1,126 @@ +using BotSharp.Core.MCP.Managers; +using Xunit; + +namespace BotSharp.Core.UnitTests.Mcp; + +/// +/// Pins down the pool key that decides which callers may share an MCP connection. +/// +/// The manager is registered per DI scope, so in practice a pool only ever holds one caller's +/// connections and identities cannot mix. The key is the second line of defence: it folds in the +/// headers the connection opens with, so the guarantee still holds if someone later registers the +/// manager with a longer lifetime. These tests exist so that property cannot be quietly lost -- +/// getting it wrong hands one user's connection, and therefore one user's credential, to another. +/// +public class McpClientPoolKeyTests +{ + private const string ServerId = "sumo-logic"; + + [Fact] + public void SameHeaders_ShareOneEntry() + { + var a = McpClientManager.BuildPoolKey(ServerId, new Dictionary + { + ["Authorization"] = "Bearer alice-token", + ["X-Tenant"] = "lessen" + }); + + // Same pairs, different insertion order: the key is canonicalised, so these are one entry. + var b = McpClientManager.BuildPoolKey(ServerId, new Dictionary + { + ["X-Tenant"] = "lessen", + ["Authorization"] = "Bearer alice-token" + }); + + Assert.Equal(a, b); + } + + [Fact] + public void DifferentCredential_NeverSharesAnEntry() + { + var alice = McpClientManager.BuildPoolKey(ServerId, new Dictionary + { + ["Authorization"] = "Bearer alice-token" + }); + + var bob = McpClientManager.BuildPoolKey(ServerId, new Dictionary + { + ["Authorization"] = "Bearer bob-token" + }); + + Assert.NotEqual(alice, bob); + } + + /// + /// The three identities OneBrainMcpHeaderProvider can answer with -- the caller's own token, + /// the credential configured for the server, and X-API-KEY minted from a user id -- are + /// different callers, not interchangeable ways of naming one. None may share a connection. + /// + [Fact] + public void DifferentIdentityKinds_NeverShareAnEntry() + { + var callerToken = McpClientManager.BuildPoolKey(ServerId, new Dictionary + { + ["Authorization"] = "Bearer caller-token" + }); + + var configuredCredential = McpClientManager.BuildPoolKey(ServerId, new Dictionary + { + ["Authorization"] = "Bearer service-credential" + }); + + var mintedApiKey = McpClientManager.BuildPoolKey(ServerId, new Dictionary + { + ["X-API-KEY"] = "mesh-key-42" + }); + + Assert.Equal(3, new HashSet { callerToken, configuredCredential, mintedApiKey }.Count); + } + + [Fact] + public void SameHeaders_DifferentServers_DoNotShareAnEntry() + { + var headers = new Dictionary { ["Authorization"] = "Bearer alice-token" }; + + Assert.NotEqual( + McpClientManager.BuildPoolKey("sumo-logic", headers), + McpClientManager.BuildPoolKey("meshstage", headers)); + } + + /// + /// A stdio server is not consulted for headers, and an http server may simply have none + /// configured. Both land on one stable entry per server, distinct from any authenticated one. + /// + [Fact] + public void NoHeaders_IsStable_AndDistinctFromAuthenticated() + { + var fromNull = McpClientManager.BuildPoolKey(ServerId, null); + var fromEmpty = McpClientManager.BuildPoolKey(ServerId, new Dictionary()); + var authenticated = McpClientManager.BuildPoolKey(ServerId, new Dictionary + { + ["Authorization"] = "Bearer alice-token" + }); + + Assert.Equal(fromNull, fromEmpty); + Assert.NotEqual(fromNull, authenticated); + } + + /// + /// The key is derived from credentials, so it must not carry one. Keys reach logs and dumps by + /// accident far more easily than the headers themselves do. + /// + [Fact] + public void Key_DoesNotCarryTheCredential() + { + const string secret = "Bearer alice-super-secret-token"; + + var key = McpClientManager.BuildPoolKey(ServerId, new Dictionary + { + ["Authorization"] = secret + }); + + Assert.DoesNotContain("alice-super-secret-token", key); + Assert.DoesNotContain(secret, key); + Assert.StartsWith(ServerId + "|", key); + } +} From 628cdb06a6d87eeeca98d8421c0c2995334e4d81 Mon Sep 17 00:00:00 2001 From: "mars.yu" Date: Fri, 28 Aug 2026 16:41:58 +0800 Subject: [PATCH 2/2] Share the HTTP connection under MCP clients, not the clients themselves Pooling MCP clients, as the previous commit did, shares more than a socket. A client is a session: CreateAsync performs the initialize handshake, the server answers with a session id, and subscriptions and long-running tool tasks (ListTasksAsync, GetTaskResultAsync) live on it. Two callers on one session would see each other's tasks, and no per-request header can undo that, because it is server-side state rather than an authorization question. With IMcpClientHeaderProvider opening connections as the signed-in user, sharing a session would mean sharing an identity as well. So sessions are not shared at all now: every GetMcpClientAsync call opens its own and the caller owns it. The three call sites hold it in an await using, which closes the session on the server instead of leaving it to time out -- the leak the empty Dispose used to cause, and the reason the pool existed. What is shared instead is the layer that carries no identity. The HttpClient comes from IHttpClientFactory, named per server, so connections to one server reuse a pooled HttpMessageHandler. CreateClient hands back a fresh HttpClient each time, so one caller's headers are never seen by another. Building the transport with its own HttpClient, as this did before, gave every connection a private handler and therefore a private socket pool -- the usual way to exhaust sockets and to keep talking to an address DNS has already moved. AddBotSharpMCP now calls AddHttpClient so the factory it depends on is present. The call is idempotent, and a host that already registered one is unaffected. Timeout is left at the factory default. No configured tool is expected to run for 100 seconds, but that cap is one the SDK's own client may not have had, so a comment records the symptom and the one-line fix should a server keep a GET open for the length of its session. Co-Authored-By: Claude Opus 5 --- .../BotSharp.Core/BotSharp.Core.csproj | 4 - .../MCP/BotSharpMCPExtensions.cs | 5 + .../MCP/Hooks/MCPToolAgentHook.cs | 2 +- .../MCP/Managers/McpClientManager.cs | 264 ++++-------------- .../BotSharp.Core/MCP/Services/McpService.cs | 2 +- .../Routing/Executor/MCPToolExecutor.cs | 14 +- .../Mcp/McpClientPoolKeyTests.cs | 126 --------- 7 files changed, 69 insertions(+), 348 deletions(-) delete mode 100644 tests/BotSharp.Core.UnitTests/Mcp/McpClientPoolKeyTests.cs diff --git a/src/Infrastructure/BotSharp.Core/BotSharp.Core.csproj b/src/Infrastructure/BotSharp.Core/BotSharp.Core.csproj index 8c3c32f09..744b2465c 100644 --- a/src/Infrastructure/BotSharp.Core/BotSharp.Core.csproj +++ b/src/Infrastructure/BotSharp.Core/BotSharp.Core.csproj @@ -301,8 +301,4 @@ - - - - diff --git a/src/Infrastructure/BotSharp.Core/MCP/BotSharpMCPExtensions.cs b/src/Infrastructure/BotSharp.Core/MCP/BotSharpMCPExtensions.cs index 8eeee7b35..8da28baae 100644 --- a/src/Infrastructure/BotSharp.Core/MCP/BotSharpMCPExtensions.cs +++ b/src/Infrastructure/BotSharp.Core/MCP/BotSharpMCPExtensions.cs @@ -22,6 +22,11 @@ public static IServiceCollection AddBotSharpMCP(this IServiceCollection services if (settings != null && settings.Enabled && !settings.McpServerConfigs.IsNullOrEmpty()) { + // McpClientManager opens every connection over a client from this factory, so that + // connections to one server share a pooled handler instead of each building its own. + // Idempotent, and a host that already called it is unaffected. + services.AddHttpClient(); + services.AddScoped(); services.AddScoped(); } diff --git a/src/Infrastructure/BotSharp.Core/MCP/Hooks/MCPToolAgentHook.cs b/src/Infrastructure/BotSharp.Core/MCP/Hooks/MCPToolAgentHook.cs index ede3dd111..66313e683 100644 --- a/src/Infrastructure/BotSharp.Core/MCP/Hooks/MCPToolAgentHook.cs +++ b/src/Infrastructure/BotSharp.Core/MCP/Hooks/MCPToolAgentHook.cs @@ -50,7 +50,7 @@ private async Task> GetMcpContent(Agent agent) var mcps = agent.McpTools?.Where(x => !x.Disabled) ?? []; foreach (var item in mcps) { - var mcpClient = await mcpClientManager.GetMcpClientAsync(item.ServerId); + await using var mcpClient = await mcpClientManager.GetMcpClientAsync(item.ServerId); if (mcpClient == null) continue; var tools = await mcpClient.ListToolsAsync(); diff --git a/src/Infrastructure/BotSharp.Core/MCP/Managers/McpClientManager.cs b/src/Infrastructure/BotSharp.Core/MCP/Managers/McpClientManager.cs index 0060f38c1..e4249f3d3 100644 --- a/src/Infrastructure/BotSharp.Core/MCP/Managers/McpClientManager.cs +++ b/src/Infrastructure/BotSharp.Core/MCP/Managers/McpClientManager.cs @@ -1,47 +1,41 @@ -using System.Collections.Concurrent; -using System.Security.Cryptography; using BotSharp.Core.MCP.Settings; using ModelContextProtocol.Client; +using System.Net.Http; namespace BotSharp.Core.MCP.Managers; /// -/// Hands out MCP clients, pooled for the lifetime of the DI scope that resolved this manager. +/// Opens MCP clients. Each call returns a client of its own, which the caller owns and must +/// dispose; what is shared between callers is the HTTP connection underneath it. /// /// /// -/// WHY THE POOL IS SCOPED, AND HAS TO STAY SCOPED. A connection carries whatever -/// answered with, which is how a host calls a server as -/// the signed-in user instead of with one shared credential. Reusing a connection therefore -/// reuses an identity. This class is registered per scope — one HTTP request, one crontab run, -/// one queued message — so everything sharing a pool is already the same caller, and one user's -/// connection cannot be handed to another. That is structural, not a rule someone has to -/// remember. +/// WHY NOTHING ABOVE THE SOCKET IS SHARED. An MCP client is a session: CreateAsync performs the +/// initialize handshake, the server answers with a session id, and subscriptions and long-running +/// tool tasks (ListTasksAsync, GetTaskResultAsync) live on that session. Handing one session to +/// two callers would show one of them the other's tasks, and no per-request header can undo that +/// because it is server-side state rather than an authorization question. Since +/// lets a host open a connection as the signed-in user, +/// sharing a session would also mean sharing an identity. So sessions are never shared. /// /// -/// The pool key carries a fingerprint of the headers a connection opens with, so the guarantee -/// survives this class later being registered with a longer lifetime: two credentials land on -/// two entries even inside one pool. The fingerprint is a hash of secrets and is never logged. +/// WHAT IS SHARED. The HttpClient comes from IHttpClientFactory, named per server, so every +/// connection to one server reuses a pooled HttpMessageHandler -- the same TCP and TLS the +/// factory would give any other caller. That layer carries no identity: the credential lives in +/// the transport's headers, and CreateClient hands back a fresh HttpClient each time, so headers +/// set for one caller are never seen by another. This is what makes a per-call session cheap: +/// the handshake runs over an already-warm connection. /// /// -/// Before pooling, every tool call opened its own connection and none of them were ever closed: -/// this method built a fresh transport per call and did nothing. A turn -/// that listed tools and then called three of them opened four connections and leaked all four. +/// Building the transport with its own HttpClient, as this did before, gave every MCP connection +/// a private handler and therefore a private socket pool -- the usual way to exhaust sockets and +/// to keep talking to an address DNS has already moved. /// /// -public class McpClientManager : IDisposable, IAsyncDisposable +public class McpClientManager { - private const string KeySeparator = "|"; - - /// - /// How long a synchronous scope teardown waits for connections to close before giving up. - /// - private static readonly TimeSpan SyncCloseTimeout = TimeSpan.FromSeconds(5); - private readonly IServiceProvider _services; private readonly ILogger _logger; - private readonly ConcurrentDictionary _pool = new(); - private volatile bool _disposed; public McpClientManager( IServiceProvider services, @@ -52,112 +46,39 @@ public McpClientManager( } /// - /// The client for , opening one if this scope has not already. - /// Answers null rather than throwing when the server is unknown, disabled or unreachable, - /// which is the contract every caller is already written against. + /// Opens a client for . The caller owns it and must dispose it + /// -- an undisposed client leaves its session open on the server until the server times it out. + /// Answers null rather than throwing when the server is unknown, disabled or unreachable. /// public async Task GetMcpClientAsync(string serverId) { - if (_disposed) - { - return null; - } - - McpServerConfigModel config; - Dictionary? headers; - string key; - try { var settings = _services.GetRequiredService(); - var found = settings.McpServerConfigs?.FirstOrDefault(x => x.Id == serverId); - if (found == null || !found.Enabled) + var config = settings.McpServerConfigs?.FirstOrDefault(x => x.Id == serverId); + if (config == null || !config.Enabled) { return null; } - config = found; - - // Resolved once, here, and handed to the transport below. Resolving separately for - // the key and for the connection would let the two disagree, and the key is the - // thing keeping one caller's connection away from another. - headers = ResolveHeaders(config); - key = BuildPoolKey(serverId, headers); - } - catch (Exception ex) - { - _logger.LogWarning(ex, $"Error when loading mcp client {serverId}"); - return null; - } - - // Lazy rather than a bare Task, so parallel tool calls that all want this server open - // one connection between them instead of one each. - var entry = _pool.GetOrAdd(key, _ => new PooledClient(serverId, new Lazy>( - () => CreateClientAsync(config, headers), - LazyThreadSafetyMode.ExecutionAndPublication))); - - McpClient? client = null; - try - { - client = await entry.Client.Value; - } - catch (Exception ex) - { - _logger.LogWarning(ex, $"Error when loading mcp client {serverId}"); - } - - if (client == null) - { - // A failure must not stay cached, or every later call in this scope gets it back. - // Removed by value, so a retry that already replaced the entry survives. - _pool.TryRemove(new KeyValuePair(key, entry)); - } - - return client; - } - - /// - /// Drops and closes this scope's connection to so the next call - /// opens a fresh one. Call it when a request over that connection failed at the transport - /// level: keeping a dead client fails every remaining call in the turn, while discarding a - /// live one costs a single reconnect, and that asymmetry says always discard. - /// - public async Task InvalidateAsync(string serverId) - { - var prefix = serverId + KeySeparator; - foreach (var key in _pool.Keys.Where(x => x.StartsWith(prefix, StringComparison.Ordinal)).ToList()) - { - if (_pool.TryRemove(key, out var entry)) - { - await CloseAsync(entry); - } - } - } - - private async Task CreateClientAsync(McpServerConfigModel config, Dictionary? headers) - { - try - { - var settings = _services.GetRequiredService(); - IClientTransport? transport = null; if (config.HttpConfig != null) { - transport = new HttpClientTransport(new HttpClientTransportOptions + transport = CreateHttpTransport(config, new HttpClientTransportOptions { Name = config.Name, Endpoint = new Uri(config.HttpConfig.EndPoint), - AdditionalHeaders = headers, + AdditionalHeaders = ResolveHeaders(config.Id, config.HttpConfig.AdditionalHeaders), ConnectionTimeout = config.HttpConfig.ConnectionTimeout }); } else if (config.SseConfig != null) { - transport = new HttpClientTransport(new HttpClientTransportOptions + transport = CreateHttpTransport(config, new HttpClientTransportOptions { Name = config.Name, Endpoint = new Uri(config.SseConfig.EndPoint), - AdditionalHeaders = headers, + AdditionalHeaders = ResolveHeaders(config.Id, config.SseConfig.AdditionalHeaders), ConnectionTimeout = config.SseConfig.ConnectionTimeout }); } @@ -182,119 +103,50 @@ public async Task InvalidateAsync(string serverId) } catch (Exception ex) { - _logger.LogWarning(ex, $"Error when loading mcp client {config.Id}"); + _logger.LogWarning(ex, $"Error when loading mcp client {serverId}"); return null; } } /// - /// The headers to open a connection with: the ones from configuration, unless the host has - /// registered an that wants to adjust them. + /// A transport over an HttpClient from the factory, named for this server so its handler -- + /// and therefore its connection pool -- is reused by every later connection to the same + /// server. The instance itself is fresh per call, which is what keeps one caller's headers + /// out of another's request. /// - /// - /// No provider is registered by default, and a provider is free to answer with what it was - /// given, so a host without one — or with one that does not recognise this server — gets the - /// configured headers back untouched. A stdio server opens with no connection headers at - /// all, so the provider is not consulted for one. - /// - private Dictionary? ResolveHeaders(McpServerConfigModel config) + private HttpClientTransport CreateHttpTransport(McpServerConfigModel config, HttpClientTransportOptions options) { - if (config.HttpConfig == null && config.SseConfig == null) - { - return null; - } + var factory = _services.GetRequiredService(); + var http = factory.CreateClient(HttpClientName(config.Id)); - var configured = config.HttpConfig?.AdditionalHeaders ?? config.SseConfig?.AdditionalHeaders; - var provider = _services.GetService(); - return provider == null ? configured : provider.GetHeaders(config.Id, configured); + // Timeout is left at the factory default (100s) deliberately: no configured tool is + // expected to run that long. Note this is a cap the SDK's own HttpClient may not have + // had, so it arrived with this change -- a server whose transport keeps a GET open for + // the session (SSE, or streamable HTTP with a standalone listening stream) would be cut + // off at 100s no matter how quick its tools are. The symptom is a tool call failing with + // a canceled request; the fix is Timeout.InfiniteTimeSpan here. + + return new HttpClientTransport(options, http, loggerFactory: null, ownsHttpClient: true); } /// - /// The server id plus a fingerprint of the headers the connection will carry, so an entry is - /// shared only between calls that authenticate identically. Hashed rather than kept, because - /// those headers hold credentials; the result is treated as a secret and never logged. + /// One handler pool per server, so a slow or unhealthy server cannot occupy the connections + /// of the others. /// - internal static string BuildPoolKey(string serverId, Dictionary? headers) - { - if (headers == null || headers.Count == 0) - { - return serverId + KeySeparator; - } - - var canonical = new StringBuilder(); - foreach (var pair in headers.OrderBy(x => x.Key, StringComparer.Ordinal)) - { - canonical.Append(pair.Key).Append(' ').Append(pair.Value).Append('\n'); - } - - var hash = SHA256.HashData(Encoding.UTF8.GetBytes(canonical.ToString())); - return serverId + KeySeparator + Convert.ToHexString(hash); - } - - private async Task CloseAsync(PooledClient entry) - { - try - { - var client = await entry.Client.Value; - if (client != null) - { - await client.DisposeAsync(); - } - } - catch (Exception ex) - { - _logger.LogWarning(ex, $"Error when closing mcp client {entry.ServerId}"); - } - } - - public async ValueTask DisposeAsync() - { - if (_disposed) - { - return; - } - - _disposed = true; - - foreach (var pair in _pool.ToArray()) - { - if (_pool.TryRemove(pair.Key, out var entry)) - { - await CloseAsync(entry); - } - } - - GC.SuppressFinalize(this); - } + private static string HttpClientName(string serverId) => $"mcp:{serverId}"; /// - /// Scopes made with CreateScope — crontab runs, queue consumers — tear down synchronously, - /// and only offers DisposeAsync, so this waits for it. The wait is - /// bounded: a wedged transport must not hang the unit of work that is trying to finish. An - /// async scope, an ASP.NET Core request among them, calls instead - /// and never comes through here. + /// The headers to open a connection with: the ones from configuration, unless the host has + /// registered an that wants to adjust them. /// - public void Dispose() + /// + /// No provider is registered by default, and a provider is free to answer with what it was + /// given, so a host without one -- or with one that does not recognise this server -- gets the + /// configured headers back untouched. + /// + private Dictionary? ResolveHeaders(string serverId, Dictionary? configured) { - if (_disposed) - { - return; - } - - try - { - if (!DisposeAsync().AsTask().Wait(SyncCloseTimeout)) - { - _logger.LogWarning("Timed out closing pooled MCP clients; leaving them to the transport."); - } - } - catch (Exception ex) - { - _logger.LogWarning(ex, "Error when closing pooled MCP clients."); - } - - GC.SuppressFinalize(this); + var provider = _services.GetService(); + return provider == null ? configured : provider.GetHeaders(serverId, configured); } - - private sealed record PooledClient(string ServerId, Lazy> Client); } diff --git a/src/Infrastructure/BotSharp.Core/MCP/Services/McpService.cs b/src/Infrastructure/BotSharp.Core/MCP/Services/McpService.cs index 7dd20daf6..5b7d1c33c 100644 --- a/src/Infrastructure/BotSharp.Core/MCP/Services/McpService.cs +++ b/src/Infrastructure/BotSharp.Core/MCP/Services/McpService.cs @@ -28,7 +28,7 @@ public async Task> GetServerConfigsAsync() foreach (var config in configs) { - var client = await clientManager.GetMcpClientAsync(config.Id); + await using var client = await clientManager.GetMcpClientAsync(config.Id); if (client == null) continue; var tools = await client.ListToolsAsync(); diff --git a/src/Infrastructure/BotSharp.Core/Routing/Executor/MCPToolExecutor.cs b/src/Infrastructure/BotSharp.Core/Routing/Executor/MCPToolExecutor.cs index d189bb4b1..06d345ff1 100644 --- a/src/Infrastructure/BotSharp.Core/Routing/Executor/MCPToolExecutor.cs +++ b/src/Infrastructure/BotSharp.Core/Routing/Executor/MCPToolExecutor.cs @@ -27,7 +27,10 @@ public async Task ExecuteAsync(RoleDialogModel message) Dictionary argDict = JsonToDictionary(message.FunctionArgs); var clientManager = _services.GetRequiredService(); - var client = await clientManager.GetMcpClientAsync(_mcpServerId); + + // The client is a session of its own, so this call owns it. Disposing closes the + // session on the server; the connection underneath it stays in the factory's pool. + await using var client = await clientManager.GetMcpClientAsync(_mcpServerId); if (client == null) { @@ -51,15 +54,6 @@ public async Task ExecuteAsync(RoleDialogModel message) } catch (Exception ex) { - // The connection is pooled for the rest of this scope, so a transport-level failure - // here would poison every later call in the turn. Drop it and let the next call - // reconnect; discarding a connection that was actually fine costs one reconnect. - var clientManager = _services.GetService(); - if (clientManager != null) - { - await clientManager.InvalidateAsync(_mcpServerId); - } - message.Content = $"Error when calling tool {_functionName} of MCP server {_mcpServerId}. {ex.Message}"; return false; } diff --git a/tests/BotSharp.Core.UnitTests/Mcp/McpClientPoolKeyTests.cs b/tests/BotSharp.Core.UnitTests/Mcp/McpClientPoolKeyTests.cs deleted file mode 100644 index a15d4daf7..000000000 --- a/tests/BotSharp.Core.UnitTests/Mcp/McpClientPoolKeyTests.cs +++ /dev/null @@ -1,126 +0,0 @@ -using BotSharp.Core.MCP.Managers; -using Xunit; - -namespace BotSharp.Core.UnitTests.Mcp; - -/// -/// Pins down the pool key that decides which callers may share an MCP connection. -/// -/// The manager is registered per DI scope, so in practice a pool only ever holds one caller's -/// connections and identities cannot mix. The key is the second line of defence: it folds in the -/// headers the connection opens with, so the guarantee still holds if someone later registers the -/// manager with a longer lifetime. These tests exist so that property cannot be quietly lost -- -/// getting it wrong hands one user's connection, and therefore one user's credential, to another. -/// -public class McpClientPoolKeyTests -{ - private const string ServerId = "sumo-logic"; - - [Fact] - public void SameHeaders_ShareOneEntry() - { - var a = McpClientManager.BuildPoolKey(ServerId, new Dictionary - { - ["Authorization"] = "Bearer alice-token", - ["X-Tenant"] = "lessen" - }); - - // Same pairs, different insertion order: the key is canonicalised, so these are one entry. - var b = McpClientManager.BuildPoolKey(ServerId, new Dictionary - { - ["X-Tenant"] = "lessen", - ["Authorization"] = "Bearer alice-token" - }); - - Assert.Equal(a, b); - } - - [Fact] - public void DifferentCredential_NeverSharesAnEntry() - { - var alice = McpClientManager.BuildPoolKey(ServerId, new Dictionary - { - ["Authorization"] = "Bearer alice-token" - }); - - var bob = McpClientManager.BuildPoolKey(ServerId, new Dictionary - { - ["Authorization"] = "Bearer bob-token" - }); - - Assert.NotEqual(alice, bob); - } - - /// - /// The three identities OneBrainMcpHeaderProvider can answer with -- the caller's own token, - /// the credential configured for the server, and X-API-KEY minted from a user id -- are - /// different callers, not interchangeable ways of naming one. None may share a connection. - /// - [Fact] - public void DifferentIdentityKinds_NeverShareAnEntry() - { - var callerToken = McpClientManager.BuildPoolKey(ServerId, new Dictionary - { - ["Authorization"] = "Bearer caller-token" - }); - - var configuredCredential = McpClientManager.BuildPoolKey(ServerId, new Dictionary - { - ["Authorization"] = "Bearer service-credential" - }); - - var mintedApiKey = McpClientManager.BuildPoolKey(ServerId, new Dictionary - { - ["X-API-KEY"] = "mesh-key-42" - }); - - Assert.Equal(3, new HashSet { callerToken, configuredCredential, mintedApiKey }.Count); - } - - [Fact] - public void SameHeaders_DifferentServers_DoNotShareAnEntry() - { - var headers = new Dictionary { ["Authorization"] = "Bearer alice-token" }; - - Assert.NotEqual( - McpClientManager.BuildPoolKey("sumo-logic", headers), - McpClientManager.BuildPoolKey("meshstage", headers)); - } - - /// - /// A stdio server is not consulted for headers, and an http server may simply have none - /// configured. Both land on one stable entry per server, distinct from any authenticated one. - /// - [Fact] - public void NoHeaders_IsStable_AndDistinctFromAuthenticated() - { - var fromNull = McpClientManager.BuildPoolKey(ServerId, null); - var fromEmpty = McpClientManager.BuildPoolKey(ServerId, new Dictionary()); - var authenticated = McpClientManager.BuildPoolKey(ServerId, new Dictionary - { - ["Authorization"] = "Bearer alice-token" - }); - - Assert.Equal(fromNull, fromEmpty); - Assert.NotEqual(fromNull, authenticated); - } - - /// - /// The key is derived from credentials, so it must not carry one. Keys reach logs and dumps by - /// accident far more easily than the headers themselves do. - /// - [Fact] - public void Key_DoesNotCarryTheCredential() - { - const string secret = "Bearer alice-super-secret-token"; - - var key = McpClientManager.BuildPoolKey(ServerId, new Dictionary - { - ["Authorization"] = secret - }); - - Assert.DoesNotContain("alice-super-secret-token", key); - Assert.DoesNotContain(secret, key); - Assert.StartsWith(ServerId + "|", key); - } -}