From 96cbf3c5c41f1856e87b1385e3aee1645607be3a Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Mon, 19 Oct 2026 15:19:40 -0400 Subject: [PATCH 1/3] fix: mcp workflow isolation --- .../DefaultMcpToolHandler.cs | 694 +++++++++++++++-- .../PublicAPI/net10.0/PublicAPI.Unshipped.txt | 1 + .../PublicAPI/net472/PublicAPI.Unshipped.txt | 1 + .../PublicAPI/net8.0/PublicAPI.Unshipped.txt | 1 + .../PublicAPI/net9.0/PublicAPI.Unshipped.txt | 1 + .../netstandard2.0/PublicAPI.Unshipped.txt | 1 + .../IMcpToolHandler.cs | 31 + .../Interpreter/DeclarativeActionExecutor.cs | 4 +- .../Interpreter/DeclarativeWorkflowContext.cs | 26 +- .../DeclarativeWorkflowExecutor.cs | 3 +- .../Interpreter/DelegateActionExecutor.cs | 4 +- .../Kit/ActionExecutor.cs | 4 +- .../Kit/RootExecutor.cs | 3 +- .../ObjectModel/InvokeMcpToolExecutor.cs | 48 +- .../PublicAPI/net10.0/PublicAPI.Unshipped.txt | 2 + .../PublicAPI/net472/PublicAPI.Unshipped.txt | 2 + .../PublicAPI/net8.0/PublicAPI.Unshipped.txt | 2 + .../PublicAPI/net9.0/PublicAPI.Unshipped.txt | 2 + .../netstandard2.0/PublicAPI.Unshipped.txt | 2 + .../GroupChatManager.cs | 19 +- .../IWorkflowContext.cs | 15 + .../InProc/InProcessRunner.cs | 35 +- .../InProc/InProcessRunnerContext.cs | 10 +- .../PublicAPI/net10.0/PublicAPI.Unshipped.txt | 2 + .../PublicAPI/net472/PublicAPI.Unshipped.txt | 2 + .../PublicAPI/net8.0/PublicAPI.Unshipped.txt | 2 + .../PublicAPI/net9.0/PublicAPI.Unshipped.txt | 2 + .../netstandard2.0/PublicAPI.Unshipped.txt | 2 + .../Specialized/WorkflowHostExecutor.cs | 14 +- .../SubworkflowBinding.cs | 8 + .../DefaultMcpToolHandlerLifetimeTests.cs | 700 +++++++++++++++++- .../DefaultMcpToolHandlerTests.cs | 37 +- .../Kit/IWorkflowContextExtensionsTests.cs | 22 +- .../Kit/RootExecutorTests.cs | 3 +- .../ObjectModel/InvokeMcpToolExecutorTest.cs | 139 +++- .../PowerFx/WorkflowFormulaStateTests.cs | 52 ++ .../RepresentationTests.cs | 43 ++ 37 files changed, 1823 insertions(+), 116 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs index 9bd67e55eb4..1f5664a36c8 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs @@ -6,6 +6,7 @@ using System.IO; using System.Linq; using System.Net.Http; +using System.Runtime.ExceptionServices; using System.Security.Cryptography; using System.Text; using System.Text.Json; @@ -28,12 +29,17 @@ namespace Microsoft.Agents.AI.Workflows.Declarative.Mcp; /// a pre-configured for each server. /// Provider-backed invocations create and dispose a separate MCP session for every call, including /// tools/list, because provider authentication is not represented in the session cache key. -/// Without a provider, sessions are cached by server URL, label, connection name, and explicit headers. +/// Without a provider, workflow invocations are cached by workflow session, server URL, label, +/// connection name, and explicit headers in a bounded least-recently-used cache. The bound includes +/// connecting and closing sessions. At capacity, an idle session is disposed before its replacement +/// connects. If every slot is busy, new scopes fail with ; +/// existing scopes can still reuse their cached session. /// Non-cancellation cleanup failures are reported through warnings without replacing /// the invocation result or error. /// -public sealed class DefaultMcpToolHandler : IMcpToolHandler, IAsyncDisposable +public sealed class DefaultMcpToolHandler : IWorkflowScopedMcpToolHandler, IAsyncDisposable { + private const int DefaultClientCacheMaxSize = 32; private const string FilenameAdditionalPropertyName = "filename"; /// @@ -43,12 +49,21 @@ public sealed class DefaultMcpToolHandler : IMcpToolHandler, IAsyncDisposable public const string ListToolsToolName = "tools/list"; private static readonly JsonWriterOptions s_toolListJsonWriterOptions = new() { Indented = true }; + private static readonly OperationCanceledException s_clientCreationCancelledException = + new("The MCP client creator was cancelled."); private readonly Func>? _httpClientProvider; private readonly Func _httpMessageHandlerFactory; - private readonly Dictionary<(string Url, string Label, string Connection, string HeadersHash), ClientConnection> _clients = []; - private readonly Dictionary _ownedHttpClients = []; + private readonly Func _clientConnectionDisposer; + private readonly Dictionary<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash), CachedClient> _clients = []; + private readonly Dictionary<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash), TaskCompletionSource> _clientCreations = []; + private readonly HashSet> _clientCreationLifetimes = []; + private readonly LinkedList<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash)> _clientLru = []; + private readonly HashSet _retiredClients = []; + private readonly Dictionary _ownedHttpClients = []; private readonly SemaphoreSlim _clientLock = new(1, 1); + private readonly SemaphoreSlim _clientCreationSemaphore; + private readonly int _clientCacheMaxSize; private readonly AsyncLocal _providerInvocationContext = new(); private TaskCompletionSource? _providerInvocationsDrained; private int _activeProviderInvocations; @@ -83,10 +98,20 @@ public DefaultMcpToolHandler(Func>? internal DefaultMcpToolHandler( Func>? httpClientProvider, - Func httpMessageHandlerFactory) + Func httpMessageHandlerFactory, + int clientCacheMaxSize = DefaultClientCacheMaxSize, + Func? clientConnectionDisposer = null) { + if (clientCacheMaxSize <= 0) + { + throw new ArgumentOutOfRangeException(nameof(clientCacheMaxSize), "The MCP client cache size must be positive."); + } + this._httpClientProvider = httpClientProvider; this._httpMessageHandlerFactory = Throw.IfNull(httpMessageHandlerFactory); + this._clientConnectionDisposer = clientConnectionDisposer ?? (connection => connection.DisposeAsync()); + this._clientCacheMaxSize = clientCacheMaxSize; + this._clientCreationSemaphore = new(clientCacheMaxSize, clientCacheMaxSize); } /// @@ -98,6 +123,26 @@ public async Task InvokeToolAsync( IDictionary? headers, string? connectionName, CancellationToken cancellationToken = default) + => await this.InvokeToolInWorkflowSessionAsync( + serverUrl, + serverLabel, + toolName, + arguments, + headers, + connectionName, + workflowSessionId: string.Empty, + cancellationToken).ConfigureAwait(false); + + /// + public async Task InvokeToolInWorkflowSessionAsync( + string serverUrl, + string? serverLabel, + string toolName, + IDictionary? arguments, + IDictionary? headers, + string? connectionName, + string workflowSessionId, + CancellationToken cancellationToken = default) { if (IsListToolsToolName(toolName)) { @@ -141,7 +186,7 @@ public async Task InvokeToolAsync( try { ClientConnection invocationClient = await this.CreateClientAsync( - serverUrl.Trim(), serverLabel, headers, httpClientCacheKey: null, cancellationToken).ConfigureAwait(false); + serverUrl.Trim(), serverLabel, headers, cancellationToken).ConfigureAwait(false); await using (invocationClient.ConfigureAwait(false)) { return await InvokeClientAsync(invocationClient.Client, toolName, arguments, cancellationToken).ConfigureAwait(false); @@ -166,8 +211,16 @@ public async Task InvokeToolAsync( } } - McpClient client = await this.GetOrCreateClientAsync(serverUrl, serverLabel, headers, connectionName, cancellationToken).ConfigureAwait(false); - return await InvokeClientAsync(client, toolName, arguments, cancellationToken).ConfigureAwait(false); + CachedClient cachedClient = await this.AcquireClientAsync( + serverUrl, serverLabel, headers, connectionName, workflowSessionId, cancellationToken).ConfigureAwait(false); + try + { + return await InvokeClientAsync(cachedClient.Connection.Client, toolName, arguments, cancellationToken).ConfigureAwait(false); + } + finally + { + await this.ReleaseClientAsync(cachedClient).ConfigureAwait(false); + } } private static async Task InvokeClientAsync( @@ -233,12 +286,14 @@ public async ValueTask DisposeAsync() } Task? providerInvocations; + List> clientCreationLifetimes; await this._clientLock.WaitAsync().ConfigureAwait(false); try { this.ThrowIfDisposing(); this._disposing = true; providerInvocations = this._providerInvocationsDrained?.Task; + clientCreationLifetimes = this._clientCreationLifetimes.Select(source => source.Task).ToList(); } finally { @@ -250,32 +305,99 @@ public async ValueTask DisposeAsync() await providerInvocations.ConfigureAwait(false); } + Exception? clientCreationException = null; + try + { + await Task.WhenAll(clientCreationLifetimes).ConfigureAwait(false); + } + catch (Exception exception) when (!IsFatalException(exception)) + { + clientCreationException = exception; + } + + List cachedClients; + List clientsToDispose = []; await this._clientLock.WaitAsync().ConfigureAwait(false); try { - foreach (ClientConnection client in this._clients.Values) + cachedClients = [.. this._clients.Values, .. this._retiredClients]; + foreach (CachedClient client in cachedClients) { - await client.DisposeAsync().ConfigureAwait(false); + client.Evicted = true; + this._retiredClients.Add(client); + if (client.ActiveInvocations == 0 && !client.DisposalClaimed) + { + client.DisposalClaimed = true; + clientsToDispose.Add(client); + } } this._clients.Clear(); - - // Dispose only HttpClients that the handler created (not user-provided ones) - foreach (HttpClient httpClient in this._ownedHttpClients.Values) - { - httpClient.Dispose(); - } - - this._ownedHttpClients.Clear(); + this._clientLru.Clear(); } finally { this._clientLock.Release(); } - this._clientLock.Dispose(); + Exception? clientCleanupException = null; + try + { + await DrainCleanupAsync( + clientsToDispose.Select(client => this.DisposeCachedClientAsync(client)), + cachedClients.Select(client => client.Disposed.Task)).ConfigureAwait(false); + } + catch (Exception exception) when (!IsFatalException(exception)) + { + clientCleanupException = exception; + } + finally + { + this._clientLock.Dispose(); + this._clientCreationSemaphore.Dispose(); + } + + if (clientCreationException is not null) + { + ExceptionDispatchInfo.Capture(clientCreationException).Throw(); + } + + if (clientCleanupException is not null) + { + ExceptionDispatchInfo.Capture(clientCleanupException).Throw(); + } } + internal static async Task DrainCleanupAsync(IEnumerable cleanupTasks, IEnumerable completionTasks) + { + Exception? cleanupException = null; + try + { + await Task.WhenAll(cleanupTasks).ConfigureAwait(false); + } + catch (Exception exception) when (!IsFatalException(exception)) + { + cleanupException = exception; + } + + await Task.WhenAll(completionTasks).ConfigureAwait(false); + + if (cleanupException is not null) + { + ExceptionDispatchInfo.Capture(cleanupException).Throw(); + } + } + + private static bool IsFatalException(Exception exception) => + exception is OutOfMemoryException and not InsufficientMemoryException + or StackOverflowException + or AccessViolationException + or AppDomainUnloadedException + or BadImageFormatException + or CannotUnloadAppDomainException + or InvalidProgramException + or ThreadAbortException; + [System.Diagnostics.CodeAnalysis.SuppressMessage("Maintainability", "CA1513:Use ObjectDisposedException throw helper", Justification = "The helper is not available on .NET Framework or .NET Standard 2.0.")] private void ThrowIfDisposing() @@ -286,28 +408,254 @@ private void ThrowIfDisposing() } } - private async Task GetOrCreateClientAsync( + private async Task AcquireClientAsync( string serverUrl, string? serverLabel, IDictionary? headers, string? connectionName, + string workflowSessionId, CancellationToken cancellationToken) { string trimmedUrl = serverUrl.Trim(); - var clientCacheKey = BuildCacheKey(trimmedUrl, serverLabel, connectionName, headers); + var clientCacheKey = BuildCacheKey(workflowSessionId, trimmedUrl, serverLabel, connectionName, headers); + CachedClient? clientToDispose = null; + TaskCompletionSource? clientCreation; + TaskCompletionSource? clientCreationLifetime = null; + bool ownsClientCreation = false; await this._clientLock.WaitAsync(cancellationToken).ConfigureAwait(false); try { this.ThrowIfDisposing(); - if (this._clients.TryGetValue(clientCacheKey, out ClientConnection? existingClient)) + if (this._clients.TryGetValue(clientCacheKey, out CachedClient? existingClient)) + { + existingClient.ActiveInvocations++; + this._clientLru.Remove(existingClient.LruNode); + this._clientLru.AddLast(existingClient.LruNode); + return existingClient; + } + + if (!this._clientCreations.TryGetValue(clientCacheKey, out clientCreation)) + { + if (this._clients.Count + this._retiredClients.Count + this._clientCreations.Count >= this._clientCacheMaxSize) + { + LinkedListNode<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash)>? idleNode = + this._clientLru.First; + while (idleNode is not null && this._clients[idleNode.Value].ActiveInvocations != 0) + { + idleNode = idleNode.Next; + } + + if (idleNode is null) + { + throw new InvalidOperationException("MCP client capacity is exhausted; retry when an active invocation completes."); + } + + clientToDispose = this._clients[idleNode.Value]; + this._clients.Remove(idleNode.Value); + this._clientLru.Remove(idleNode); + clientToDispose.Evicted = true; + clientToDispose.DisposalClaimed = true; + this._retiredClients.Add(clientToDispose); + } + + clientCreation = new(TaskCreationOptions.RunContinuationsAsynchronously); + this._clientCreations[clientCacheKey] = clientCreation; + clientCreationLifetime = new(TaskCreationOptions.RunContinuationsAsynchronously); + this._clientCreationLifetimes.Add(clientCreationLifetime); + ownsClientCreation = true; + } + } + finally + { + this._clientLock.Release(); + } + + if (!ownsClientCreation) + { + try + { + await WaitForClientCreationAsync(clientCreation.Task, cancellationToken).ConfigureAwait(false); + } + catch (OperationCanceledException exception) + when (ReferenceEquals(exception, s_clientCreationCancelledException)) + { + return await this.AcquireClientAsync( + serverUrl, serverLabel, headers, connectionName, workflowSessionId, cancellationToken).ConfigureAwait(false); + } + + return await this.AcquireClientAsync( + serverUrl, serverLabel, headers, connectionName, workflowSessionId, cancellationToken).ConfigureAwait(false); + } + + TaskCompletionSource ownedClientCreationLifetime = clientCreationLifetime ?? + throw new InvalidOperationException("Missing MCP client creation lifetime."); + ClientConnection? connection = null; + bool creationSemaphoreEntered = false; + Exception? clientCreationCleanupException = null; + try + { + if (clientToDispose is not null) + { + try + { + await this.DisposeCachedClientAsync(clientToDispose).ConfigureAwait(false); + } + catch (Exception cleanupException) + { + clientCreationCleanupException = cleanupException; + throw; + } + } + + await this._clientCreationSemaphore.WaitAsync(cancellationToken).ConfigureAwait(false); + creationSemaphoreEntered = true; + await this._clientLock.WaitAsync(cancellationToken).ConfigureAwait(false); + try + { + this.ThrowIfDisposing(); + } + finally + { + this._clientLock.Release(); + } + + connection = await this.CreateClientAsync( + trimmedUrl, serverLabel, headers, cancellationToken, + cleanupException => clientCreationCleanupException = cleanupException).ConfigureAwait(false); + } + catch (Exception exception) + { + Exception sharedCreationException = + exception is OperationCanceledException && cancellationToken.IsCancellationRequested + ? s_clientCreationCancelledException + : exception; + try + { + await this.CompleteClientCreationFailureAsync( + clientCacheKey, clientCreation, sharedCreationException).ConfigureAwait(false); + } + finally + { + if (creationSemaphoreEntered) + { + this._clientCreationSemaphore.Release(); + creationSemaphoreEntered = false; + } + + await this.CompleteClientCreationLifetimeAsync( + ownedClientCreationLifetime, clientCreationCleanupException).ConfigureAwait(false); + } + + throw; + } + + ObjectDisposedException? disposedException = null; + CachedClient? result = null; + await this._clientLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); + try + { + this._clientCreations.Remove(clientCacheKey); + if (this._disposing) + { + disposedException = new ObjectDisposedException(nameof(DefaultMcpToolHandler)); + } + else + { + LinkedListNode<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash)> node = + this._clientLru.AddLast(clientCacheKey); + CachedClient newClient = new(connection, node) { ActiveInvocations = 1 }; + connection = null; + this._clients[clientCacheKey] = newClient; + + result = newClient; + clientCreation.TrySetResult(newClient); + } + } + finally + { + this._clientLock.Release(); + } + + try + { + if (connection is not null) + { + await this._clientConnectionDisposer(connection).ConfigureAwait(false); + } + + if (disposedException is not null) + { + clientCreation.TrySetException(disposedException); + throw disposedException; + } + } + catch (Exception exception) + { + if (disposedException is not null && !ReferenceEquals(exception, disposedException)) + { + clientCreation.TrySetException(exception); + clientCreationCleanupException = exception; + } + + if (result is not null) + { + await this.ReleaseClientAsync(result).ConfigureAwait(false); + } + + throw; + } + finally + { + if (creationSemaphoreEntered) + { + this._clientCreationSemaphore.Release(); + } + + await this.CompleteClientCreationLifetimeAsync( + ownedClientCreationLifetime, clientCreationCleanupException).ConfigureAwait(false); + } + + return result ?? throw new InvalidOperationException("Failed to acquire MCP client."); + } + + private static async Task WaitForClientCreationAsync(Task clientCreation, CancellationToken cancellationToken) + { + if (!cancellationToken.CanBeCanceled || clientCreation.IsCompleted) + { + await clientCreation.ConfigureAwait(false); + return; + } + + TaskCompletionSource cancellation = new(TaskCreationOptions.RunContinuationsAsynchronously); + using CancellationTokenRegistration registration = cancellationToken.Register( + static state => ((TaskCompletionSource)state!).TrySetResult(true), + cancellation); + + if (await Task.WhenAny(clientCreation, cancellation.Task).ConfigureAwait(false) == cancellation.Task) + { + cancellationToken.ThrowIfCancellationRequested(); + } + + await clientCreation.ConfigureAwait(false); + } + + private async Task CompleteClientCreationFailureAsync( + (string WorkflowSession, string Url, string Label, string Connection, string HeadersHash) clientCacheKey, + TaskCompletionSource clientCreation, + Exception exception) + { + await this._clientLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); + try + { + if (this._clientCreations.TryGetValue(clientCacheKey, out TaskCompletionSource? existingCreation) && + ReferenceEquals(existingCreation, clientCreation)) { - return existingClient.Client; + this._clientCreations.Remove(clientCacheKey); } - ClientConnection newClient = await this.CreateClientAsync(trimmedUrl, serverLabel, headers, trimmedUrl, cancellationToken).ConfigureAwait(false); - this._clients[clientCacheKey] = newClient; - return newClient.Client; + clientCreation.TrySetException(exception); + _ = clientCreation.Task.Exception; } finally { @@ -315,60 +663,129 @@ private async Task GetOrCreateClientAsync( } } + private async Task CompleteClientCreationLifetimeAsync( + TaskCompletionSource clientCreationLifetime, + Exception? cleanupException = null) + { + await this._clientLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); + try + { + this._clientCreationLifetimes.Remove(clientCreationLifetime); + } + finally + { + this._clientLock.Release(); + } + + if (cleanupException is not null) + { + clientCreationLifetime.TrySetException(cleanupException); + } + else + { + clientCreationLifetime.TrySetResult(true); + } + } + + private async ValueTask ReleaseClientAsync(CachedClient client) + { + bool dispose = false; + await this._clientLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); + try + { + client.ActiveInvocations--; + if (client.ActiveInvocations == 0 && client.Evicted && !client.DisposalClaimed) + { + client.DisposalClaimed = true; + dispose = true; + } + } + finally + { + this._clientLock.Release(); + } + + if (dispose) + { + await this.DisposeCachedClientAsync(client).ConfigureAwait(false); + } + } + + private async Task DisposeCachedClientAsync(CachedClient client) + { + Exception? cleanupException = null; + try + { + await this._clientConnectionDisposer(client.Connection).ConfigureAwait(false); + } + catch (Exception exception) + { + cleanupException = exception; + throw; + } + finally + { + await this._clientLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); + try + { + this._retiredClients.Remove(client); + } + finally + { + this._clientLock.Release(); + } + + if (cleanupException is not null) + { + client.Disposed.TrySetException(cleanupException); + } + else + { + client.Disposed.TrySetResult(true); + } + } + } + /// - /// Builds the per-client cache key as a 4-tuple of - /// (trimmed serverUrl, serverLabel, connectionName, headers hash). All four components - /// participate so that callers using different labels/connections/headers receive - /// distinct instances even when targeting the same URL. + /// Builds the per-client cache key as a 5-tuple of + /// (workflowSessionId, trimmed serverUrl, serverLabel, connectionName, headers hash). + /// All five components participate so that separate workflow sessions and callers using + /// different labels/connections/headers receive distinct instances + /// even when targeting the same URL. /// - internal static (string Url, string Label, string Connection, string HeadersHash) BuildCacheKey( + internal static (string WorkflowSession, string Url, string Label, string Connection, string HeadersHash) BuildCacheKey( + string workflowSessionId, string trimmedUrl, string? serverLabel, string? connectionName, IDictionary? headers) => - (trimmedUrl, serverLabel ?? string.Empty, connectionName ?? string.Empty, ComputeHeadersHash(headers)); + (workflowSessionId, trimmedUrl, serverLabel ?? string.Empty, connectionName ?? string.Empty, ComputeHeadersHash(headers)); private async Task CreateClientAsync( string serverUrl, string? serverLabel, IDictionary? headers, - string? httpClientCacheKey, - CancellationToken cancellationToken) + CancellationToken cancellationToken, + Action? reportCleanupFailure = null) { - // Only the no-provider path shares handler-owned HTTP clients. HttpClient? httpClient = null; bool ownsHttpClient = false; + OwnedHttpClientLease? ownedHttpClientLease = null; if (this._httpClientProvider is not null) { httpClient = await this._httpClientProvider(serverUrl, cancellationToken).ConfigureAwait(false); } - if (httpClient is null && - (httpClientCacheKey is null || !this._ownedHttpClients.TryGetValue(httpClientCacheKey, out httpClient))) + if (httpClient is null && this._httpClientProvider is not null) { - // Pin credential headers to the configured server origin as defense-in-depth. Forcing - // StreamableHttp (below) already removes the primary vector (a server-advertised cross-origin - // SSE message endpoint), and AllowAutoRedirect=false blocks auto-redirects. This handler is the - // backstop: it guarantees the Authorization token and other credentials never leave the pinned - // origin even if a future change re-enables AutoDetect or redirects, or the SDK constructs a - // request to a new URI (AdditionalHeaders are re-stamped by the transport, so HttpClient's own - // redirect header-stripping does not cover them). - // Caveat for future maintainers: the handler strips credentials on every cross-origin request. - // That is safe today because this handler carries auth via static request headers, not the - // SDK's built-in OAuth flow. If OAuth support is ever added here, requests to the authorization - // server (a different origin than the MCP resource server) legitimately carry credentials, so - // those auth-server origins would need to be allow-listed to avoid breaking the token exchange. - OriginPinningHandler pinningHandler = new(new Uri(serverUrl)) { InnerHandler = this._httpMessageHandlerFactory() }; - httpClient = new HttpClient(pinningHandler); - if (httpClientCacheKey is null) - { - ownsHttpClient = true; - } - else - { - this._ownedHttpClients[httpClientCacheKey] = httpClient; - } + httpClient = this.CreatePinnedHttpClient(serverUrl); + ownsHttpClient = true; + } + else if (httpClient is null) + { + ownedHttpClientLease = await this.AcquireOwnedHttpClientAsync(serverUrl, cancellationToken).ConfigureAwait(false); + httpClient = ownedHttpClientLease.Client; } HttpClientTransportOptions transportOptions = new() @@ -383,20 +800,108 @@ private async Task CreateClientAsync( TransportMode = HttpTransportMode.StreamableHttp }; - HttpClientTransport transport = new(transportOptions, httpClient, ownsHttpClient: ownsHttpClient); - + HttpClientTransport? transport = null; try { + HttpClient resolvedHttpClient = httpClient ?? throw new InvalidOperationException("Failed to resolve MCP HTTP client."); + transport = new(transportOptions, resolvedHttpClient, ownsHttpClient: ownsHttpClient); McpClient client = await McpClient.CreateAsync(transport, cancellationToken: cancellationToken).ConfigureAwait(false); - return new ClientConnection(client, transport); + ClientConnection connection = new(client, transport, ownedHttpClientLease); + ownedHttpClientLease = null; + return connection; } catch { - await DisposeResourceAsync(transport, "transport").ConfigureAwait(false); + try + { + try + { + if (transport is not null) + { + await DisposeResourceAsync(transport, "transport").ConfigureAwait(false); + } + } + finally + { + if (ownedHttpClientLease is not null) + { + await ownedHttpClientLease.DisposeAsync().ConfigureAwait(false); + } + } + } + catch (Exception cleanupException) + { + reportCleanupFailure?.Invoke(cleanupException); + throw; + } + throw; } } + private async Task AcquireOwnedHttpClientAsync( + string serverUrl, + CancellationToken cancellationToken) + { + await this._clientLock.WaitAsync(cancellationToken).ConfigureAwait(false); + try + { + if (!this._ownedHttpClients.TryGetValue(serverUrl, out OwnedHttpClient? entry)) + { + entry = new(this.CreatePinnedHttpClient(serverUrl)); + this._ownedHttpClients[serverUrl] = entry; + } + + entry.ReferenceCount++; + return new(this, serverUrl, entry.Client); + } + finally + { + this._clientLock.Release(); + } + } + + private async ValueTask ReleaseOwnedHttpClientAsync(string serverUrl) + { + HttpClient? clientToDispose = null; + await this._clientLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); + try + { + OwnedHttpClient entry = this._ownedHttpClients[serverUrl]; + if (--entry.ReferenceCount == 0) + { + this._ownedHttpClients.Remove(serverUrl); + clientToDispose = entry.Client; + } + } + finally + { + this._clientLock.Release(); + } + + try + { + clientToDispose?.Dispose(); + } + catch (Exception exception) when (exception is not OperationCanceledException) + { + Trace.TraceWarning("Failed to dispose MCP HTTP client: {0}", exception); + } + } + + private HttpClient CreatePinnedHttpClient(string serverUrl) + { + // Pin credential headers to the configured server origin as defense-in-depth. Forcing + // StreamableHttp (below) already removes the primary vector (a server-advertised cross-origin + // SSE message endpoint), and AllowAutoRedirect=false blocks auto-redirects. This handler is the + // backstop: it guarantees the Authorization token and other credentials never leave the pinned + // origin even if a future change re-enables AutoDetect or redirects, or the SDK constructs a + // request to a new URI (AdditionalHeaders are re-stamped by the transport, so HttpClient's own + // redirect header-stripping does not cover them). + OriginPinningHandler pinningHandler = new(new Uri(serverUrl)) { InnerHandler = this._httpMessageHandlerFactory() }; + return new HttpClient(pinningHandler); + } + private static HttpMessageHandler CreateHttpMessageHandler() => new HttpClientHandler { @@ -409,11 +914,56 @@ private static HttpMessageHandler CreateHttpMessageHandler() => private sealed class ProviderInvocationContext(Task completion, ProviderInvocationContext? parent) { public Task Completion { get; } = completion; - public ProviderInvocationContext? Parent { get; } = parent; } - internal sealed class ClientConnection(McpClient client, IAsyncDisposable transport) : IAsyncDisposable + private sealed class CachedClient( + ClientConnection connection, + LinkedListNode<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash)> lruNode) + { + public ClientConnection Connection { get; } = connection; + + public LinkedListNode<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash)> LruNode { get; } = lruNode; + + public TaskCompletionSource Disposed { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously); + + public int ActiveInvocations { get; set; } + + public bool Evicted { get; set; } + + public bool DisposalClaimed { get; set; } + } + + private sealed class OwnedHttpClient(HttpClient client) + { + public HttpClient Client { get; } = client; + + public int ReferenceCount { get; set; } + } + + private sealed class OwnedHttpClientLease( + DefaultMcpToolHandler owner, + string serverUrl, + HttpClient client) : IAsyncDisposable + { + private bool _disposed; + + public HttpClient Client { get; } = client; + + public async ValueTask DisposeAsync() + { + if (!this._disposed) + { + this._disposed = true; + await owner.ReleaseOwnedHttpClientAsync(serverUrl).ConfigureAwait(false); + } + } + } + + internal sealed class ClientConnection( + McpClient client, + IAsyncDisposable transport, + IAsyncDisposable? ownedHttpClientLease = null) : IAsyncDisposable { public McpClient Client { get; } = client; @@ -425,8 +975,18 @@ public async ValueTask DisposeAsync() } finally { - // McpClient owns the connected session, not the reusable transport factory. - await DisposeResourceAsync(transport, "transport").ConfigureAwait(false); + try + { + // McpClient owns the connected session, not the reusable transport factory. + await DisposeResourceAsync(transport, "transport").ConfigureAwait(false); + } + finally + { + if (ownedHttpClientLease is not null) + { + await ownedHttpClientLease.DisposeAsync().ConfigureAwait(false); + } + } } } } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net10.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net10.0/PublicAPI.Unshipped.txt index ab058de62d4..3cdebd5e1a0 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net10.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net10.0/PublicAPI.Unshipped.txt @@ -1 +1,2 @@ #nullable enable +Microsoft.Agents.AI.Workflows.Declarative.Mcp.DefaultMcpToolHandler.InvokeToolInWorkflowSessionAsync(string! serverUrl, string? serverLabel, string! toolName, System.Collections.Generic.IDictionary? arguments, System.Collections.Generic.IDictionary? headers, string? connectionName, string! workflowSessionId, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net472/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net472/PublicAPI.Unshipped.txt index ab058de62d4..3cdebd5e1a0 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net472/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net472/PublicAPI.Unshipped.txt @@ -1 +1,2 @@ #nullable enable +Microsoft.Agents.AI.Workflows.Declarative.Mcp.DefaultMcpToolHandler.InvokeToolInWorkflowSessionAsync(string! serverUrl, string? serverLabel, string! toolName, System.Collections.Generic.IDictionary? arguments, System.Collections.Generic.IDictionary? headers, string? connectionName, string! workflowSessionId, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net8.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net8.0/PublicAPI.Unshipped.txt index ab058de62d4..3cdebd5e1a0 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net8.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net8.0/PublicAPI.Unshipped.txt @@ -1 +1,2 @@ #nullable enable +Microsoft.Agents.AI.Workflows.Declarative.Mcp.DefaultMcpToolHandler.InvokeToolInWorkflowSessionAsync(string! serverUrl, string? serverLabel, string! toolName, System.Collections.Generic.IDictionary? arguments, System.Collections.Generic.IDictionary? headers, string? connectionName, string! workflowSessionId, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net9.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net9.0/PublicAPI.Unshipped.txt index ab058de62d4..3cdebd5e1a0 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net9.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net9.0/PublicAPI.Unshipped.txt @@ -1 +1,2 @@ #nullable enable +Microsoft.Agents.AI.Workflows.Declarative.Mcp.DefaultMcpToolHandler.InvokeToolInWorkflowSessionAsync(string! serverUrl, string? serverLabel, string! toolName, System.Collections.Generic.IDictionary? arguments, System.Collections.Generic.IDictionary? headers, string? connectionName, string! workflowSessionId, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt index ab058de62d4..3cdebd5e1a0 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt @@ -1 +1,2 @@ #nullable enable +Microsoft.Agents.AI.Workflows.Declarative.Mcp.DefaultMcpToolHandler.InvokeToolInWorkflowSessionAsync(string! serverUrl, string? serverLabel, string! toolName, System.Collections.Generic.IDictionary? arguments, System.Collections.Generic.IDictionary? headers, string? connectionName, string! workflowSessionId, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/IMcpToolHandler.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/IMcpToolHandler.cs index 56b1c3deb4c..fbea21f35b4 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/IMcpToolHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/IMcpToolHandler.cs @@ -39,3 +39,34 @@ Task InvokeToolAsync( string? connectionName, CancellationToken cancellationToken = default); } + +/// +/// Defines the contract for MCP handlers that isolate stateful protocol sessions by workflow session. +/// +public interface IWorkflowScopedMcpToolHandler : IMcpToolHandler +{ + /// + /// Invokes an MCP tool within the framework-owned workflow session scope. + /// + /// The URL of the MCP server. + /// An optional label identifying the server connection. + /// The name of the tool to invoke. + /// Optional arguments to pass to the tool. + /// Optional headers to include in the request. + /// An optional connection name for managed connections. + /// The framework-owned identifier for the current workflow session. + /// A token to observe cancellation. + /// + /// A task representing the asynchronous operation. The result contains a + /// with the tool invocation output. + /// + Task InvokeToolInWorkflowSessionAsync( + string serverUrl, + string? serverLabel, + string toolName, + IDictionary? arguments, + IDictionary? headers, + string? connectionName, + string workflowSessionId, + CancellationToken cancellationToken = default); +} diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeActionExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeActionExecutor.cs index 4902589514d..2c5ac9867cd 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeActionExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeActionExecutor.cs @@ -96,7 +96,9 @@ public override async ValueTask HandleAsync(ActionExecutorResult message, IWorkf try { - object? result = await this.ExecuteAsync(new DeclarativeWorkflowContext(context, this.State), cancellationToken).ConfigureAwait(false); + DeclarativeWorkflowContext declarativeContext = + await DeclarativeWorkflowContext.CreateAsync(context, this.State, cancellationToken).ConfigureAwait(false); + object? result = await this.ExecuteAsync(declarativeContext, cancellationToken).ConfigureAwait(false); Debug.WriteLine($"RESULT #{this.Id} - {result ?? "(null)"}"); if (this.EmitResultEvent) diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowContext.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowContext.cs index 183b2e785e7..f88035abd12 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowContext.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowContext.cs @@ -15,8 +15,10 @@ namespace Microsoft.Agents.AI.Workflows.Declarative.Interpreter; -internal sealed class DeclarativeWorkflowContext : IWorkflowContext +internal sealed class DeclarativeWorkflowContext : IWorkflowContext, IWorkflowSessionContext { + internal const string WorkflowSessionIdStateKey = "__declarative_mcp_workflow_session_id"; + public static readonly FrozenSet ManagedScopes = [ VariableScopeNames.Local, @@ -24,16 +26,36 @@ internal sealed class DeclarativeWorkflowContext : IWorkflowContext VariableScopeNames.Global, ]; - public DeclarativeWorkflowContext(IWorkflowContext source, WorkflowFormulaState state) + private DeclarativeWorkflowContext(IWorkflowContext source, WorkflowFormulaState state, string sessionId) { this.Source = source; this.State = state; + this.SessionId = sessionId; + } + + public static async ValueTask CreateAsync( + IWorkflowContext source, + WorkflowFormulaState state, + CancellationToken cancellationToken = default) + { + string generatedSessionId = Guid.NewGuid().ToString("N"); + string sessionId = source is IWorkflowSessionContext sessionContext + ? sessionContext.SessionId + : await source.ReadOrInitStateAsync( + WorkflowSessionIdStateKey, + () => generatedSessionId, + VariableScopeNames.System, + cancellationToken: cancellationToken).ConfigureAwait(false); + + return new(source, state, sessionId); } private IWorkflowContext Source { get; } public WorkflowFormulaState State { get; } public IReadOnlyDictionary? TraceContext => this.Source.TraceContext; + public string SessionId { get; } + /// public bool ConcurrentRunsEnabled => this.Source.ConcurrentRunsEnabled; diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowExecutor.cs index 1126b43d8ff..2fcdea4ecce 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowExecutor.cs @@ -154,7 +154,8 @@ private async ValueTask AdvanceAsync(ChatMessage input, IWorkflowContext context // No state to restore if we're starting from the beginning. state.SetInitialized(); - DeclarativeWorkflowContext declarativeContext = new(context, state); + DeclarativeWorkflowContext declarativeContext = + await DeclarativeWorkflowContext.CreateAsync(context, state, cancellationToken).ConfigureAwait(false); // Conversation id resolution prefers state already persisted by a prior turn, // so multi-turn invocations reuse the same backend conversation rather than diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DelegateActionExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DelegateActionExecutor.cs index 716a3656dac..ade8dd1034b 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DelegateActionExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DelegateActionExecutor.cs @@ -61,7 +61,9 @@ public override async ValueTask HandleAsync(TMessage message, IWorkflowContext c { if (this._action is not null) { - await this._action.Invoke(new DeclarativeWorkflowContext(context, this._state), message, cancellationToken).ConfigureAwait(false); + DeclarativeWorkflowContext declarativeContext = + await DeclarativeWorkflowContext.CreateAsync(context, this._state, cancellationToken).ConfigureAwait(false); + await this._action.Invoke(declarativeContext, message, cancellationToken).ConfigureAwait(false); } if (this._emitResult) diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/ActionExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/ActionExecutor.cs index cf636effaf8..563ff86f147 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/ActionExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/ActionExecutor.cs @@ -76,7 +76,9 @@ public ValueTask ResetAsync() [SendsMessage(typeof(ActionExecutorResult))] public override async ValueTask HandleAsync(TMessage message, IWorkflowContext context, CancellationToken cancellationToken) { - object? result = await this.ExecuteAsync(new DeclarativeWorkflowContext(context, this._session.State), message, cancellationToken).ConfigureAwait(false); + DeclarativeWorkflowContext declarativeContext = + await DeclarativeWorkflowContext.CreateAsync(context, this._session.State, cancellationToken).ConfigureAwait(false); + object? result = await this.ExecuteAsync(declarativeContext, message, cancellationToken).ConfigureAwait(false); Debug.WriteLine($"RESULT #{this.Id} - {result ?? "(null)"}"); await context.SendResultMessageAsync(this.Id, result, cancellationToken).ConfigureAwait(false); diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/RootExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/RootExecutor.cs index 3a18a06c948..9e008ed391b 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/RootExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/RootExecutor.cs @@ -63,7 +63,8 @@ public ValueTask ResetAsync() [SendsMessage(typeof(ActionExecutorResult))] public override async ValueTask HandleAsync(TInput message, IWorkflowContext context, CancellationToken cancellationToken) { - DeclarativeWorkflowContext declarativeContext = new(context, this._state); + DeclarativeWorkflowContext declarativeContext = + await DeclarativeWorkflowContext.CreateAsync(context, this._state, cancellationToken).ConfigureAwait(false); ChatMessage input = (this._inputTransform ?? DefaultInputTransform).Invoke(message); diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/InvokeMcpToolExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/InvokeMcpToolExecutor.cs index 02dafb8ce2c..5f76bb3621e 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/InvokeMcpToolExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/InvokeMcpToolExecutor.cs @@ -32,7 +32,6 @@ internal sealed class InvokeMcpToolExecutor( { private const string ApprovalSnapshotStateKey = nameof(_approvalSnapshots); private const string LegacyApprovalSnapshotStateKey = "_approvalSnapshot"; - /// /// Snapshots of evaluated parameters captured at approval-request time, keyed by /// per-invocation request id. Each pending approval lives here until the matching @@ -110,7 +109,8 @@ await this.RequestApprovalAsync( } // No approval required - invoke the tool directly - McpServerToolResultContent resultContent = await mcpToolHandler.InvokeToolAsync( + McpServerToolResultContent resultContent = await this.InvokeToolAsync( + context, serverUrl, serverLabel, toolName, @@ -193,7 +193,8 @@ public async ValueTask CaptureResponseAsync( string? conversationId = this.GetConversationId(); bool autoSend = snapshot.AutoSend ?? this.GetAutoSendValue(); - McpServerToolResultContent resultContent = await mcpToolHandler.InvokeToolAsync( + McpServerToolResultContent resultContent = await this.InvokeToolAsync( + context, snapshot.ServerUrl, snapshot.ServerLabel, snapshot.ToolName, @@ -206,6 +207,47 @@ public async ValueTask CaptureResponseAsync( await context.SendResultMessageAsync(this.Id, result: null, cancellationToken).ConfigureAwait(false); } + private async Task InvokeToolAsync( + IWorkflowContext context, + string serverUrl, + string? serverLabel, + string toolName, + IDictionary? arguments, + IDictionary? headers, + string? connectionName, + CancellationToken cancellationToken) + { + if (mcpToolHandler is IWorkflowScopedMcpToolHandler scopedHandler) + { + string generatedWorkflowSessionId = Guid.NewGuid().ToString("N"); + string workflowSessionId = context is IWorkflowSessionContext sessionContext + ? sessionContext.SessionId + : await context.ReadOrInitStateAsync( + DeclarativeWorkflowContext.WorkflowSessionIdStateKey, + () => generatedWorkflowSessionId, + VariableScopeNames.System, + cancellationToken).ConfigureAwait(false); + return await scopedHandler.InvokeToolInWorkflowSessionAsync( + serverUrl, + serverLabel, + toolName, + arguments, + headers, + connectionName, + workflowSessionId, + cancellationToken).ConfigureAwait(false); + } + + return await mcpToolHandler.InvokeToolAsync( + serverUrl, + serverLabel, + toolName, + arguments, + headers, + connectionName, + cancellationToken).ConfigureAwait(false); + } + /// /// Completes the MCP tool invocation by raising the completion event. /// diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net10.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net10.0/PublicAPI.Unshipped.txt index 8af74cb2369..5bee0f03310 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net10.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net10.0/PublicAPI.Unshipped.txt @@ -4,6 +4,8 @@ Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowedEnvi Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowProcessEnvironmentVariableFallback.get -> bool Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowProcessEnvironmentVariableFallback.init -> void Microsoft.Agents.AI.Workflows.Declarative.Events.ExternalInputResponse.RequestId.get -> string? +Microsoft.Agents.AI.Workflows.Declarative.IWorkflowScopedMcpToolHandler +Microsoft.Agents.AI.Workflows.Declarative.IWorkflowScopedMcpToolHandler.InvokeToolInWorkflowSessionAsync(string! serverUrl, string? serverLabel, string! toolName, System.Collections.Generic.IDictionary? arguments, System.Collections.Generic.IDictionary? headers, string? connectionName, string! workflowSessionId, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.ConvertValueWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, Microsoft.Agents.AI.Workflows.Declarative.Kit.VariableType! targetType, string! key, string? scopeName = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.ValueTask!> static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.EvaluateValueWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, string! expression, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.ValueTask!> static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.FormatTemplateWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, string! line) -> System.Threading.Tasks.ValueTask!> diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net472/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net472/PublicAPI.Unshipped.txt index 8af74cb2369..5bee0f03310 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net472/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net472/PublicAPI.Unshipped.txt @@ -4,6 +4,8 @@ Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowedEnvi Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowProcessEnvironmentVariableFallback.get -> bool Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowProcessEnvironmentVariableFallback.init -> void Microsoft.Agents.AI.Workflows.Declarative.Events.ExternalInputResponse.RequestId.get -> string? +Microsoft.Agents.AI.Workflows.Declarative.IWorkflowScopedMcpToolHandler +Microsoft.Agents.AI.Workflows.Declarative.IWorkflowScopedMcpToolHandler.InvokeToolInWorkflowSessionAsync(string! serverUrl, string? serverLabel, string! toolName, System.Collections.Generic.IDictionary? arguments, System.Collections.Generic.IDictionary? headers, string? connectionName, string! workflowSessionId, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.ConvertValueWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, Microsoft.Agents.AI.Workflows.Declarative.Kit.VariableType! targetType, string! key, string? scopeName = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.ValueTask!> static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.EvaluateValueWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, string! expression, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.ValueTask!> static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.FormatTemplateWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, string! line) -> System.Threading.Tasks.ValueTask!> diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net8.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net8.0/PublicAPI.Unshipped.txt index 8af74cb2369..5bee0f03310 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net8.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net8.0/PublicAPI.Unshipped.txt @@ -4,6 +4,8 @@ Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowedEnvi Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowProcessEnvironmentVariableFallback.get -> bool Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowProcessEnvironmentVariableFallback.init -> void Microsoft.Agents.AI.Workflows.Declarative.Events.ExternalInputResponse.RequestId.get -> string? +Microsoft.Agents.AI.Workflows.Declarative.IWorkflowScopedMcpToolHandler +Microsoft.Agents.AI.Workflows.Declarative.IWorkflowScopedMcpToolHandler.InvokeToolInWorkflowSessionAsync(string! serverUrl, string? serverLabel, string! toolName, System.Collections.Generic.IDictionary? arguments, System.Collections.Generic.IDictionary? headers, string? connectionName, string! workflowSessionId, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.ConvertValueWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, Microsoft.Agents.AI.Workflows.Declarative.Kit.VariableType! targetType, string! key, string? scopeName = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.ValueTask!> static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.EvaluateValueWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, string! expression, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.ValueTask!> static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.FormatTemplateWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, string! line) -> System.Threading.Tasks.ValueTask!> diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net9.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net9.0/PublicAPI.Unshipped.txt index 8af74cb2369..5bee0f03310 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net9.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net9.0/PublicAPI.Unshipped.txt @@ -4,6 +4,8 @@ Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowedEnvi Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowProcessEnvironmentVariableFallback.get -> bool Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowProcessEnvironmentVariableFallback.init -> void Microsoft.Agents.AI.Workflows.Declarative.Events.ExternalInputResponse.RequestId.get -> string? +Microsoft.Agents.AI.Workflows.Declarative.IWorkflowScopedMcpToolHandler +Microsoft.Agents.AI.Workflows.Declarative.IWorkflowScopedMcpToolHandler.InvokeToolInWorkflowSessionAsync(string! serverUrl, string? serverLabel, string! toolName, System.Collections.Generic.IDictionary? arguments, System.Collections.Generic.IDictionary? headers, string? connectionName, string! workflowSessionId, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.ConvertValueWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, Microsoft.Agents.AI.Workflows.Declarative.Kit.VariableType! targetType, string! key, string? scopeName = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.ValueTask!> static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.EvaluateValueWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, string! expression, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.ValueTask!> static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.FormatTemplateWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, string! line) -> System.Threading.Tasks.ValueTask!> diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt index 8af74cb2369..5bee0f03310 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt @@ -4,6 +4,8 @@ Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowedEnvi Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowProcessEnvironmentVariableFallback.get -> bool Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowProcessEnvironmentVariableFallback.init -> void Microsoft.Agents.AI.Workflows.Declarative.Events.ExternalInputResponse.RequestId.get -> string? +Microsoft.Agents.AI.Workflows.Declarative.IWorkflowScopedMcpToolHandler +Microsoft.Agents.AI.Workflows.Declarative.IWorkflowScopedMcpToolHandler.InvokeToolInWorkflowSessionAsync(string! serverUrl, string? serverLabel, string! toolName, System.Collections.Generic.IDictionary? arguments, System.Collections.Generic.IDictionary? headers, string? connectionName, string! workflowSessionId, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.ConvertValueWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, Microsoft.Agents.AI.Workflows.Declarative.Kit.VariableType! targetType, string! key, string? scopeName = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.ValueTask!> static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.EvaluateValueWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, string! expression, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.ValueTask!> static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.FormatTemplateWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, string! line) -> System.Threading.Tasks.ValueTask!> diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/GroupChatManager.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/GroupChatManager.cs index ab94a9fa8aa..fbda33c132a 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/GroupChatManager.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/GroupChatManager.cs @@ -152,7 +152,7 @@ protected virtual ValueTask OnCheckpointRestoredAsync(IWorkflowContext context, internal async ValueTask CheckpointAsync(IWorkflowContext context, CancellationToken cancellationToken = default) { await context.QueueStateUpdateAsync(BaseStateKey, new GroupChatManagerState(this.IterationCount), cancellationToken: cancellationToken).ConfigureAwait(false); - await this.OnCheckpointingAsync(new PrefixingWorkflowContext(context, SubclassStateKeyPrefix), cancellationToken).ConfigureAwait(false); + await this.OnCheckpointingAsync(PrefixingWorkflowContext.Create(context, SubclassStateKeyPrefix), cancellationToken).ConfigureAwait(false); } // Root restore entry point invoked by the hosting GroupChatHost. Symmetric to CheckpointAsync. @@ -160,7 +160,7 @@ internal async ValueTask RestoreCheckpointAsync(IWorkflowContext context, Cancel { GroupChatManagerState? state = await context.ReadStateAsync(BaseStateKey, cancellationToken: cancellationToken).ConfigureAwait(false); this.IterationCount = state?.IterationCount ?? 0; - await this.OnCheckpointRestoredAsync(new PrefixingWorkflowContext(context, SubclassStateKeyPrefix), cancellationToken).ConfigureAwait(false); + await this.OnCheckpointRestoredAsync(PrefixingWorkflowContext.Create(context, SubclassStateKeyPrefix), cancellationToken).ConfigureAwait(false); } } @@ -169,11 +169,16 @@ internal sealed record GroupChatManagerState(int IterationCount); // IWorkflowContext decorator that prepends a fixed prefix to every state key passed through it. // All non-state members (events, message sending, output yielding, halt requests, trace context, // and runtime characteristics) delegate directly to the wrapped context. -internal sealed class PrefixingWorkflowContext(IWorkflowContext inner, string prefix) : IWorkflowContext +internal class PrefixingWorkflowContext(IWorkflowContext inner, string prefix) : IWorkflowContext { private readonly IWorkflowContext _inner = Throw.IfNull(inner); private readonly string _prefix = Throw.IfNullOrEmpty(prefix); + public static IWorkflowContext Create(IWorkflowContext inner, string prefix) => + inner is IWorkflowSessionContext sessionContext + ? new PrefixingWorkflowSessionContext(inner, prefix, sessionContext) + : new PrefixingWorkflowContext(inner, prefix); + public IReadOnlyDictionary? TraceContext => this._inner.TraceContext; public bool ConcurrentRunsEnabled => this._inner.ConcurrentRunsEnabled; @@ -222,3 +227,11 @@ public async ValueTask QueueClearScopeAsync(string? scopeName = null, Cancellati private string Wrap(string key) => this._prefix + Throw.IfNullOrEmpty(key); } + +internal sealed class PrefixingWorkflowSessionContext( + IWorkflowContext inner, + string prefix, + IWorkflowSessionContext sessionContext) : PrefixingWorkflowContext(inner, prefix), IWorkflowSessionContext +{ + public string SessionId => sessionContext.SessionId; +} diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/IWorkflowContext.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/IWorkflowContext.cs index b8b35fffd64..68ff864085f 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/IWorkflowContext.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/IWorkflowContext.cs @@ -195,3 +195,18 @@ ValueTask ReadOrInitStateAsync(string key, Func initialStateFactory, Ca /// bool ConcurrentRunsEnabled { get; } } + +/// +/// Exposes the framework-owned identifier for a workflow session. +/// +/// +/// Runtime services can use this identifier to scope stateful resources to one +/// workflow session without accepting caller-provided cache partition keys. +/// +public interface IWorkflowSessionContext +{ + /// + /// Gets the framework-owned workflow session identifier. + /// + string SessionId { get; } +} diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunner.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunner.cs index 69b6c1e9bcc..c0a2e075ae1 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunner.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunner.cs @@ -31,7 +31,14 @@ public static InProcessRunner CreateTopLevelRunner(Workflow workflow, ICheckpoin knownValidInputTypes: knownValidInputTypes); } - public static InProcessRunner CreateSubworkflowRunner(Workflow workflow, ICheckpointManager? checkpointManager, string? sessionId = null, object? existingOwnerSignoff = null, bool enableConcurrentRuns = false, IEnumerable? knownValidInputTypes = null) + public static InProcessRunner CreateSubworkflowRunner( + Workflow workflow, + ICheckpointManager? checkpointManager, + string? sessionId = null, + object? existingOwnerSignoff = null, + bool enableConcurrentRuns = false, + IEnumerable? knownValidInputTypes = null, + string? workflowSessionId = null) { return new InProcessRunner(workflow, checkpointManager, @@ -39,10 +46,19 @@ public static InProcessRunner CreateSubworkflowRunner(Workflow workflow, ICheckp existingOwnerSignoff: existingOwnerSignoff, enableConcurrentRuns: enableConcurrentRuns, knownValidInputTypes: knownValidInputTypes, - subworkflow: true); + subworkflow: true, + workflowSessionId: workflowSessionId); } - private InProcessRunner(Workflow workflow, ICheckpointManager? checkpointManager, string? sessionId = null, object? existingOwnerSignoff = null, bool subworkflow = false, bool enableConcurrentRuns = false, IEnumerable? knownValidInputTypes = null) + private InProcessRunner( + Workflow workflow, + ICheckpointManager? checkpointManager, + string? sessionId = null, + object? existingOwnerSignoff = null, + bool subworkflow = false, + bool enableConcurrentRuns = false, + IEnumerable? knownValidInputTypes = null, + string? workflowSessionId = null) { if (enableConcurrentRuns && !workflow.AllowConcurrent) { @@ -54,7 +70,16 @@ private InProcessRunner(Workflow workflow, ICheckpointManager? checkpointManager this.StartExecutorId = workflow.StartExecutorId; this.Workflow = Throw.IfNull(workflow); - this.RunContext = new InProcessRunnerContext(workflow, this.SessionId, checkpointingEnabled: checkpointManager != null, this.OutgoingEvents, this.StepTracer, existingOwnerSignoff, subworkflow, enableConcurrentRuns); + this.RunContext = new InProcessRunnerContext( + workflow, + this.SessionId, + checkpointingEnabled: checkpointManager != null, + this.OutgoingEvents, + this.StepTracer, + existingOwnerSignoff, + subworkflow, + enableConcurrentRuns, + workflowSessionId: workflowSessionId); this.CheckpointManager = checkpointManager; this._knownValidInputTypes = knownValidInputTypes != null @@ -65,6 +90,8 @@ private InProcessRunner(Workflow workflow, ICheckpointManager? checkpointManager /// public string SessionId { get; } + internal string WorkflowSessionId => this.RunContext.WorkflowSessionId; + /// public string StartExecutorId { get; } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunnerContext.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunnerContext.cs index 353c38f25c9..aa8bf398592 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunnerContext.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunnerContext.cs @@ -47,7 +47,8 @@ public InProcessRunnerContext( object? existingOwnershipSignoff = null, bool subworkflow = false, bool enableConcurrentRuns = false, - ILogger? logger = null) + ILogger? logger = null, + string? workflowSessionId = null) { if (enableConcurrentRuns) { @@ -62,6 +63,7 @@ public InProcessRunnerContext( this._workflow = workflow; this._sessionId = sessionId; + this.WorkflowSessionId = workflowSessionId ?? sessionId; this._edgeMap = new(this, this._workflow, stepTracer); this._outputFilter = new(workflow); @@ -72,6 +74,8 @@ public InProcessRunnerContext( } public WorkflowTelemetryContext TelemetryContext => this._workflow.TelemetryContext; + internal string WorkflowSessionId { get; } + public IExternalRequestSink RegisterPort(string executorId, RequestPort port) { if (!this._edgeMap.TryRegisterPort(this, executorId, port)) @@ -349,8 +353,10 @@ public IExternalRequestSink RegisterPort(RequestPort port) private sealed class BoundWorkflowContext( InProcessRunnerContext RunnerContext, string ExecutorId, - Dictionary? traceContext) : IWorkflowContext + Dictionary? traceContext) : IWorkflowContext, IWorkflowSessionContext { + public string SessionId => RunnerContext.WorkflowSessionId; + public ValueTask AddEventAsync(WorkflowEvent workflowEvent, CancellationToken cancellationToken = default) => RunnerContext.AddEventAsync(workflowEvent, cancellationToken); public ValueTask SendMessageAsync(object message, string? targetId = null, CancellationToken cancellationToken = default) diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net10.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net10.0/PublicAPI.Unshipped.txt index ab058de62d4..724a27a74a4 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net10.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net10.0/PublicAPI.Unshipped.txt @@ -1 +1,3 @@ #nullable enable +Microsoft.Agents.AI.Workflows.IWorkflowSessionContext +Microsoft.Agents.AI.Workflows.IWorkflowSessionContext.SessionId.get -> string! diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net472/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net472/PublicAPI.Unshipped.txt index ab058de62d4..724a27a74a4 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net472/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net472/PublicAPI.Unshipped.txt @@ -1 +1,3 @@ #nullable enable +Microsoft.Agents.AI.Workflows.IWorkflowSessionContext +Microsoft.Agents.AI.Workflows.IWorkflowSessionContext.SessionId.get -> string! diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net8.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net8.0/PublicAPI.Unshipped.txt index ab058de62d4..724a27a74a4 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net8.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net8.0/PublicAPI.Unshipped.txt @@ -1 +1,3 @@ #nullable enable +Microsoft.Agents.AI.Workflows.IWorkflowSessionContext +Microsoft.Agents.AI.Workflows.IWorkflowSessionContext.SessionId.get -> string! diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net9.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net9.0/PublicAPI.Unshipped.txt index ab058de62d4..724a27a74a4 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net9.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net9.0/PublicAPI.Unshipped.txt @@ -1 +1,3 @@ #nullable enable +Microsoft.Agents.AI.Workflows.IWorkflowSessionContext +Microsoft.Agents.AI.Workflows.IWorkflowSessionContext.SessionId.get -> string! diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt index ab058de62d4..724a27a74a4 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt @@ -1 +1,3 @@ #nullable enable +Microsoft.Agents.AI.Workflows.IWorkflowSessionContext +Microsoft.Agents.AI.Workflows.IWorkflowSessionContext.SessionId.get -> string! diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/WorkflowHostExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/WorkflowHostExecutor.cs index 58e3a9e5239..a166d8b0e6b 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/WorkflowHostExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/WorkflowHostExecutor.cs @@ -20,6 +20,7 @@ internal class WorkflowHostExecutor : Executor, IAsyncDisposable private readonly Workflow _workflow; private readonly ProtocolDescriptor _workflowProtocol; private readonly object _ownershipToken; + private string? _workflowSessionId; private InProcessRunner? _activeRunner; private InMemoryCheckpointManager? _checkpointManager; @@ -56,6 +57,7 @@ protected override ProtocolBuilder ConfigureProtocol(ProtocolBuilder protocolBui private async ValueTask QueueExternalMessageAsync(PortableValue portableValue, IWorkflowContext context, CancellationToken cancellationToken) { + this.SetWorkflowSessionId(context); if (portableValue.Is(out ExternalResponse? response)) { response = this.CheckAndUnqualifyResponse(response); @@ -95,7 +97,8 @@ internal async ValueTask EnsureRunnerAsync() this._checkpointManager, this._sessionId, this._ownershipToken, - this.JoinContext.ConcurrentRunsEnabled); + this.JoinContext.ConcurrentRunsEnabled, + workflowSessionId: this._workflowSessionId); } return this._activeRunner; @@ -268,6 +271,7 @@ await context.QueueStateUpdateAsync(PendingResponsePortsStateKey, protected internal override async ValueTask OnCheckpointRestoredAsync(IWorkflowContext context, CancellationToken cancellationToken = default) { await base.OnCheckpointRestoredAsync(context, cancellationToken).ConfigureAwait(false); + this.SetWorkflowSessionId(context); InMemoryCheckpointManager manager = await context.ReadStateAsync(CheckpointManagerStateKey, cancellationToken: cancellationToken).ConfigureAwait(false) ?? new(); if (this._checkpointManager == manager) @@ -293,6 +297,14 @@ await context.ReadStateAsync>(PendingRespons await this.EnsureRunSendMessageAsync(resume: true, cancellationToken: cancellationToken).ConfigureAwait(false); } + private void SetWorkflowSessionId(IWorkflowContext context) + { + string parentSessionId = context is IWorkflowSessionContext sessionContext + ? sessionContext.SessionId + : this._sessionId; + this._workflowSessionId ??= SubworkflowBinding.CreateSubworkflowSessionId(parentSessionId, this.Id); + } + private async ValueTask ResetAsync() { if (this._run != null) diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/SubworkflowBinding.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/SubworkflowBinding.cs index 11f7cf493cf..9519af12414 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/SubworkflowBinding.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/SubworkflowBinding.cs @@ -1,6 +1,7 @@ // Copyright (c) Microsoft. All rights reserved. using System; +using System.Globalization; using System.Threading.Tasks; using Microsoft.Agents.AI.Workflows.Specialized; using Microsoft.Shared.Diagnostics; @@ -35,6 +36,13 @@ async ValueTask InitHostExecutorAsync(string sessionId) } } + internal static string CreateSubworkflowSessionId(string parentSessionId, string subworkflowId) => + string.Concat( + parentSessionId.Length.ToString(CultureInfo.InvariantCulture), + ":", + parentSessionId, + subworkflowId); + /// public override bool IsSharedInstance => false; diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs index 3a481fcd0e6..3ff71500d6c 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs @@ -24,6 +24,150 @@ namespace Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests; Justification = "Provider callbacks return completed tasks; each invocation and handler disposal is awaited before disposing caller-owned clients.")] public sealed class DefaultMcpToolHandlerLifetimeTests { + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task NoProvider_ActiveClientCapacity_RejectsNewScopesAndRecoversAsync(bool cancelActive) + { + // Arrange + ProtocolStub stub = new(); + using SemaphoreSlim started = new(0); + using SemaphoreSlim release = new(0); + int operations = 0; + stub.BeforeOperationAsync = async token => + { + if (Interlocked.Increment(ref operations) == 1) + { + started.Release(); + await release.WaitAsync(token); + } + }; + await using DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler, clientCacheMaxSize: 1); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + using CancellationTokenSource firstCancellation = CancellationTokenSource.CreateLinkedTokenSource(timeout.Token); + Task first = InvokeScopedAsync(handler, "workflow-a", "ping", firstCancellation.Token); + await started.WaitAsync(timeout.Token); + + try + { + // Act + await InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); + InvalidOperationException exception = await Assert.ThrowsAsync( + () => InvokeScopedAsync(handler, "workflow-b", "ping", timeout.Token)); + + // Assert + Assert.Contains("capacity", exception.Message); + Assert.Equal(1, stub.Initializations); + Assert.Equal(0, stub.Terminations); + } + finally + { + if (cancelActive) + { + firstCancellation.Cancel(); + } + + release.Release(); + if (cancelActive) + { + await Assert.ThrowsAnyAsync(() => first); + } + else + { + await first; + } + } + + await InvokeScopedAsync(handler, "workflow-b", "ping", timeout.Token); + Assert.Equal(2, stub.Initializations); + Assert.Equal(1, stub.Terminations); + } + + [Fact] + public async Task NoProvider_Replacement_WaitsForIdleClientCleanupBeforeConnectingAsync() + { + // Arrange + ProtocolStub stub = new(); + using SemaphoreSlim closeStarted = new(0); + using SemaphoreSlim releaseClose = new(0); + int disposals = 0; + await using DefaultMcpToolHandler handler = new( + null, stub.CreateMessageHandler, clientCacheMaxSize: 1, + clientConnectionDisposer: async connection => + { + if (Interlocked.Increment(ref disposals) == 1) + { + closeStarted.Release(); + await releaseClose.WaitAsync(); + } + + await connection.DisposeAsync(); + }); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + await InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); + Task replacement = InvokeScopedAsync(handler, "workflow-b", "ping", timeout.Token); + await closeStarted.WaitAsync(timeout.Token); + + try + { + // Act + Assert.Equal(1, stub.Initializations); + InvalidOperationException exception = await Assert.ThrowsAsync( + () => InvokeScopedAsync(handler, "workflow-c", "ping", timeout.Token)); + + // Assert + Assert.Contains("capacity", exception.Message); + } + finally + { + releaseClose.Release(); + await replacement; + } + + Assert.Equal(2, stub.Initializations); + Assert.Equal(1, stub.Terminations); + } + + [Fact] + public async Task NoProvider_DisposalDuringIdleEviction_ReportsCleanupCancellationAsync() + { + // Arrange + OperationCanceledException cleanupException = new("session cleanup cancelled"); + ProtocolStub stub = new(); + using SemaphoreSlim closeStarted = new(0); + using SemaphoreSlim releaseClose = new(0); + DefaultMcpToolHandler handler = new( + null, stub.CreateMessageHandler, clientCacheMaxSize: 1, + clientConnectionDisposer: async connection => + { + closeStarted.Release(); + await releaseClose.WaitAsync(); + await connection.DisposeAsync(); + throw cleanupException; + }); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + await InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); + Task replacement = InvokeScopedAsync(handler, "workflow-b", "ping", timeout.Token); + await closeStarted.WaitAsync(timeout.Token); + + // Act + Task disposal = handler.DisposeAsync().AsTask(); + try + { + Assert.False(disposal.IsCompleted); + } + finally + { + releaseClose.Release(); + } + + // Assert + Assert.Same(cleanupException, await Assert.ThrowsAsync(() => replacement)); + Assert.Same(cleanupException, await Assert.ThrowsAsync(() => disposal)); + Assert.Equal(1, stub.Initializations); + Assert.Equal(1, stub.Terminations); + } + [Theory] [InlineData("ping")] [InlineData(DefaultMcpToolHandler.ListToolsToolName)] @@ -104,6 +248,416 @@ public async Task NoProvider_ReusesSessionUntilHandlerDisposalAsync() transport.Protected().Verify("Dispose", Times.AtLeastOnce(), ItExpr.Is(disposing => disposing)); } + [Fact] + public async Task NoProvider_SeparateWorkflowSessions_UseSeparateCachedSessionsAsync() + { + // Arrange + ProtocolStub stub = new(); + await using DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + + // Act + await InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); + await InvokeScopedAsync(handler, "workflow-b", "ping", timeout.Token); + await InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); + + // Assert + Assert.Equal(2, stub.Initializations); + Assert.Equal(0, stub.Terminations); + } + + [Fact] + public async Task DrainCleanupAsync_Cancellation_DrainsAllCleanupAndCompletionTasksAsync() + { + // Arrange + TaskCompletionSource firstDisposed = new(TaskCreationOptions.RunContinuationsAsynchronously); + TaskCompletionSource secondDisposed = new(TaskCreationOptions.RunContinuationsAsynchronously); + + async Task CancelledCleanupAsync() + { + try + { + await Task.Yield(); + throw new OperationCanceledException("session cleanup cancelled"); + } + finally + { + firstDisposed.TrySetResult(true); + } + } + + async Task SuccessfulCleanupAsync() + { + await Task.Yield(); + secondDisposed.TrySetResult(true); + } + + // Act + await Assert.ThrowsAnyAsync(() => DefaultMcpToolHandler.DrainCleanupAsync( + [CancelledCleanupAsync(), SuccessfulCleanupAsync()], + [firstDisposed.Task, secondDisposed.Task])); + + // Assert + Assert.True(await firstDisposed.Task); + Assert.True(await secondDisposed.Task); + } + + [Fact] + public async Task NoProvider_CacheEviction_DisposesLeastRecentlyUsedSessionAsync() + { + // Arrange + ProtocolStub stub = new(); + await using DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler, clientCacheMaxSize: 2); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + + // Act + await InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); + await InvokeScopedAsync(handler, "workflow-b", "ping", timeout.Token); + await InvokeScopedAsync(handler, "workflow-c", "ping", timeout.Token); + + // Assert + Assert.Equal(3, stub.Initializations); + Assert.Equal(1, stub.Terminations); + } + + [Fact] + public async Task NoProvider_CacheEviction_DisposesUnusedOwnedHttpClientAsync() + { + // Arrange + ProtocolStub stub = new(); + await using DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler, clientCacheMaxSize: 1); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + + // Act + await handler.InvokeToolInWorkflowSessionAsync( + "https://first.example/api", null, "ping", null, null, null, "workflow-a", timeout.Token); + await handler.InvokeToolInWorkflowSessionAsync( + "https://second.example/api", null, "ping", null, null, null, "workflow-b", timeout.Token); + + // Assert + Assert.Equal(2, stub.Initializations); + Assert.Equal(1, stub.Terminations); + Assert.Equal(2, stub.Handlers.Count); + stub.Handlers[0].Protected().Verify( + "Dispose", Times.AtLeastOnce(), ItExpr.Is(disposing => disposing)); + stub.Handlers[1].Protected().Verify( + "Dispose", Times.Never(), ItExpr.Is(disposing => disposing)); + } + + [Fact] + public async Task NoProvider_CacheEviction_SkipsActiveClientUntilInvocationCompletesAsync() + { + // Arrange + ProtocolStub stub = new(); + using SemaphoreSlim firstStarted = new(0); + using SemaphoreSlim releaseFirst = new(0); + int operations = 0; + stub.BeforeOperationAsync = async token => + { + if (Interlocked.Increment(ref operations) == 1) + { + firstStarted.Release(); + await releaseFirst.WaitAsync(token); + } + }; + DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler, clientCacheMaxSize: 2); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + + // Act + Task first = InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); + await firstStarted.WaitAsync(timeout.Token); + await InvokeScopedAsync(handler, "workflow-b", "ping", timeout.Token); + await InvokeScopedAsync(handler, "workflow-c", "ping", timeout.Token); + + // Assert + Assert.Equal(1, stub.Terminations); + Task disposal = handler.DisposeAsync().AsTask(); + Assert.False(disposal.IsCompleted); + releaseFirst.Release(); + await first; + await disposal; + Assert.Equal(3, stub.Terminations); + } + + [Theory] + [InlineData(false, false)] + [InlineData(false, true)] + [InlineData(true, false)] + [InlineData(true, true)] + public async Task NoProvider_HttpClientCleanupFailure_PreservesInvocationOutcomeAsync( + bool activeDisposal, bool failOperation) + { + // Arrange + ProtocolStub stub = new() { FailTransportDisposal = true }; + using SemaphoreSlim firstStarted = new(0); + using SemaphoreSlim releaseFirst = new(0); + int operations = 0; + stub.BeforeOperationAsync = async token => + { + if (activeDisposal && Interlocked.Increment(ref operations) == 1) + { + firstStarted.Release(); + await releaseFirst.WaitAsync(token); + } + }; + DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler, clientCacheMaxSize: 1); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + using CleanupTraceListener listener = new(); + Task? disposal = null; + Trace.Listeners.Add(listener); + try + { + // Act + Task first = handler.InvokeToolInWorkflowSessionAsync( + "https://first.example/api", null, "ping", null, null, null, "workflow-a", timeout.Token); + Task outcome; + if (activeDisposal) + { + await firstStarted.WaitAsync(timeout.Token); + disposal = handler.DisposeAsync().AsTask(); + stub.FailOperation = failOperation; + releaseFirst.Release(); + outcome = first; + } + else + { + await first; + stub.FailOperation = failOperation; + outcome = handler.InvokeToolInWorkflowSessionAsync( + "https://second.example/api", null, "ping", null, null, null, "workflow-b", timeout.Token); + } + + // Assert + if (failOperation) + { + await Assert.ThrowsAsync(() => outcome); + } + else + { + McpServerToolResultContent result = await outcome; + Assert.NotNull(result.Outputs); + Assert.Equal("ok", Assert.IsType(Assert.Single(result.Outputs)).Text); + } + + if (disposal is not null) + { + await disposal; + } + + Assert.Contains("Failed to dispose MCP HTTP client", listener.Output); + Assert.Contains("transport cleanup failed", listener.Output); + } + finally + { + releaseFirst.Release(); + await (disposal ?? handler.DisposeAsync().AsTask()); + Trace.Listeners.Remove(listener); + } + } + + [Fact] + public async Task NoProvider_DisposalDuringActiveInvocation_ReportsCleanupCancellationAsync() + { + // Arrange + OperationCanceledException cleanupException = new("session cleanup cancelled"); + ProtocolStub stub = new(); + using SemaphoreSlim invocationStarted = new(0); + using SemaphoreSlim releaseInvocation = new(0); + stub.BeforeOperationAsync = async token => + { + invocationStarted.Release(); + await releaseInvocation.WaitAsync(token); + }; + DefaultMcpToolHandler handler = new( + null, + stub.CreateMessageHandler, + clientConnectionDisposer: async connection => + { + await connection.DisposeAsync(); + throw cleanupException; + }); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + + // Act + Task invocation = + InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); + await invocationStarted.WaitAsync(timeout.Token); + Task disposal = handler.DisposeAsync().AsTask(); + releaseInvocation.Release(); + + // Assert + Assert.Same(cleanupException, await Assert.ThrowsAsync(() => invocation)); + Assert.Same(cleanupException, await Assert.ThrowsAsync(() => disposal)); + Assert.Equal(1, stub.Terminations); + } + + [Fact] + public async Task NoProvider_ConcurrentWorkflowSessionCreations_AreBoundedByCacheSizeAsync() + { + // Arrange + ProtocolStub stub = new(); + using SemaphoreSlim initializationsStarted = new(0); + using SemaphoreSlim releaseInitializations = new(0); + stub.BeforeInitializationAsync = async token => + { + initializationsStarted.Release(); + await releaseInitializations.WaitAsync(token); + }; + await using DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler, clientCacheMaxSize: 2); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + + // Act + Task first = InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); + await initializationsStarted.WaitAsync(timeout.Token); + Task second = InvokeScopedAsync(handler, "workflow-b", "ping", timeout.Token); + await initializationsStarted.WaitAsync(timeout.Token); + try + { + await Assert.ThrowsAsync( + () => InvokeScopedAsync(handler, "workflow-c", "ping", timeout.Token)); + } + finally + { + releaseInitializations.Release(2); + await Task.WhenAll(first, second); + } + + Task third = InvokeScopedAsync(handler, "workflow-c", "ping", timeout.Token); + await initializationsStarted.WaitAsync(timeout.Token); + releaseInitializations.Release(); + await third; + + // Assert + Assert.Equal(3, stub.Initializations); + Assert.Equal(1, stub.Terminations); + } + + [Fact] + public async Task NoProvider_DisposalDuringFailedCreation_PreservesCreationFailureAsync() + { + // Arrange + ProtocolStub stub = new() { FailInitialization = true }; + using SemaphoreSlim initializationStarted = new(0); + using SemaphoreSlim releaseInitialization = new(0); + stub.BeforeInitializationAsync = async token => + { + initializationStarted.Release(); + await releaseInitialization.WaitAsync(token); + }; + DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + + // Act + Task invocation = + InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); + await initializationStarted.WaitAsync(timeout.Token); + Task disposal = handler.DisposeAsync().AsTask(); + releaseInitialization.Release(); + + // Assert + await Assert.ThrowsAsync(() => invocation); + await disposal; + Assert.Single(stub.Handlers); + stub.Handlers[0].Protected().Verify( + "Dispose", Times.AtLeastOnce(), ItExpr.Is(disposing => disposing)); + } + + [Fact] + public async Task NoProvider_DisposalDuringFailedCreation_ReportsCleanupCancellationAsync() + { + // Arrange + OperationCanceledException cleanupException = new("transport cleanup cancelled"); + ProtocolStub stub = new() { FailInitialization = true, TransportDisposalException = cleanupException }; + using SemaphoreSlim initializationStarted = new(0); + using SemaphoreSlim releaseInitialization = new(0); + stub.BeforeInitializationAsync = async token => + { + initializationStarted.Release(); + await releaseInitialization.WaitAsync(token); + }; + DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + + // Act + Task invocation = + InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); + await initializationStarted.WaitAsync(timeout.Token); + Task disposal = handler.DisposeAsync().AsTask(); + releaseInitialization.Release(); + + // Assert + Assert.Same(cleanupException, await Assert.ThrowsAsync(() => invocation)); + Assert.Same(cleanupException, await Assert.ThrowsAsync(() => disposal)); + Assert.Single(stub.Handlers).Protected().Verify( + "Dispose", Times.AtLeastOnce(), ItExpr.Is(disposing => disposing)); + } + + [Fact] + public async Task NoProvider_DisposalDuringSuccessfulCreation_CleansUpOrphanAndCompletesDisposalAsync() + { + // Arrange + ProtocolStub stub = new(); + using SemaphoreSlim initializationStarted = new(0); + using SemaphoreSlim releaseInitialization = new(0); + stub.BeforeInitializationAsync = async token => + { + initializationStarted.Release(); + await releaseInitialization.WaitAsync(token); + }; + DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + + // Act + Task invocation = + InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); + await initializationStarted.WaitAsync(timeout.Token); + Task disposal = handler.DisposeAsync().AsTask(); + releaseInitialization.Release(); + + // Assert + await Assert.ThrowsAsync(() => invocation); + await disposal; + Assert.Equal(1, stub.Terminations); + Assert.Single(stub.Handlers).Protected().Verify( + "Dispose", Times.AtLeastOnce(), ItExpr.Is(disposing => disposing)); + } + + [Fact] + public async Task NoProvider_DisposalDuringSuccessfulCreation_ReportsOrphanCleanupCancellationAsync() + { + // Arrange + OperationCanceledException cleanupException = new("session cleanup cancelled"); + ProtocolStub stub = new(); + using SemaphoreSlim initializationStarted = new(0); + using SemaphoreSlim releaseInitialization = new(0); + stub.BeforeInitializationAsync = async token => + { + initializationStarted.Release(); + await releaseInitialization.WaitAsync(token); + }; + DefaultMcpToolHandler handler = new( + null, + stub.CreateMessageHandler, + clientConnectionDisposer: async connection => + { + await connection.DisposeAsync(); + throw cleanupException; + }); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + + // Act + Task invocation = + InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); + await initializationStarted.WaitAsync(timeout.Token); + Task disposal = handler.DisposeAsync().AsTask(); + releaseInitialization.Release(); + + // Assert + Assert.Same(cleanupException, await Assert.ThrowsAsync(() => invocation)); + Assert.Same(cleanupException, await Assert.ThrowsAsync(() => disposal)); + Assert.Equal(1, stub.Terminations); + } + [Fact] public async Task NoProvider_DifferentConnectionNames_UseSeparateCachedSessionsAsync() { @@ -185,6 +739,124 @@ public async Task Provider_ConcurrentInvocations_DoNotCoalesceOrSerializeSession Assert.Equal(2, stub.Terminations); } + [Fact] + public async Task NoProvider_ConcurrentSameWorkflowInvocations_CoalesceSessionCreationAsync() + { + // Arrange + ProtocolStub stub = new(); + await using DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + + // Act + await Task.WhenAll( + InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token), + InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token)); + + // Assert + Assert.Equal(1, stub.Initializations); + } + + [Fact] + public async Task NoProvider_CancelledWaiter_DoesNotWaitForSharedCreationAsync() + { + // Arrange + ProtocolStub stub = new(); + using SemaphoreSlim initializationStarted = new(0); + using SemaphoreSlim releaseInitialization = new(0); + stub.BeforeInitializationAsync = async token => + { + initializationStarted.Release(); + await releaseInitialization.WaitAsync(token); + }; + await using DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + using CancellationTokenSource waiterCancellation = new(); + + // Act + Task creator = InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); + await initializationStarted.WaitAsync(timeout.Token); + Task waiter = InvokeScopedAsync(handler, "workflow-a", "ping", waiterCancellation.Token); + await Task.Yield(); + waiterCancellation.Cancel(); + Task completed = await Task.WhenAny(waiter, Task.Delay(TimeSpan.FromSeconds(1), timeout.Token)); + releaseInitialization.Release(); + await creator; + + // Assert + Assert.Same(waiter, completed); + await Assert.ThrowsAnyAsync(() => waiter); + Assert.Equal(1, stub.Initializations); + } + + [Fact] + public async Task NoProvider_CancelledCreator_DoesNotCancelSharedWaiterAsync() + { + // Arrange + ProtocolStub stub = new(); + using SemaphoreSlim initializationStarted = new(0); + int initializationAttempts = 0; + stub.BeforeInitializationAsync = async token => + { + initializationStarted.Release(); + if (Interlocked.Increment(ref initializationAttempts) == 1) + { + await Task.Delay(Timeout.Infinite, token); + } + }; + await using DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler); + using CancellationTokenSource creatorCancellation = new(); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + + // Act + Task creator = + InvokeScopedAsync(handler, "workflow-a", "ping", creatorCancellation.Token); + await initializationStarted.WaitAsync(timeout.Token); + Task waiter = InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); + await Task.Yield(); + creatorCancellation.Cancel(); + await Assert.ThrowsAnyAsync(() => creator); + await initializationStarted.WaitAsync(timeout.Token); + McpServerToolResultContent result = await waiter; + + // Assert + Assert.NotNull(result); + Assert.Equal(2, initializationAttempts); + Assert.Equal(1, stub.Initializations); + } + + [Fact] + public async Task NoProvider_InitializationTimeout_IsSharedWithoutRetryAsync() + { + // Arrange + ProtocolStub stub = new(); + using SemaphoreSlim initializationStarted = new(0); + using SemaphoreSlim releaseInitialization = new(0); + OperationCanceledException timeoutException = new("initialization timed out"); + int initializationAttempts = 0; + stub.BeforeInitializationAsync = async _ => + { + Interlocked.Increment(ref initializationAttempts); + initializationStarted.Release(); + await releaseInitialization.WaitAsync(CancellationToken.None); + throw timeoutException; + }; + await using DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + + // Act + Task creator = InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); + await initializationStarted.WaitAsync(timeout.Token); + Task waiter = InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); + await Task.Yield(); + releaseInitialization.Release(); + + // Assert + await Assert.ThrowsAnyAsync(() => creator); + await Assert.ThrowsAnyAsync(() => waiter); + Assert.Equal(1, initializationAttempts); + Assert.Equal(0, stub.Initializations); + } + [Fact] public async Task Provider_OperationFailure_DisposesSessionAndPreservesCallerClientAsync() { @@ -488,13 +1160,15 @@ public async Task Provider_CompletedNestedContext_StillGuardsActiveOuterInvocati Assert.Equal(2, stub.Terminations); } - [Fact] - public async Task Provider_InitializationAndCleanupFailure_PreservesInitializationErrorAsync() + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task InitializationAndCleanupFailure_PreservesInitializationErrorAsync(bool hasProvider) { // Arrange ProtocolStub stub = new() { FailInitialization = true, FailTransportDisposal = true }; await using DefaultMcpToolHandler handler = new( - (_, _) => Task.FromResult(null), stub.CreateMessageHandler); + hasProvider ? (_, _) => Task.FromResult(null) : null, stub.CreateMessageHandler); using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); using CleanupTraceListener listener = new(); Trace.Listeners.Add(listener); @@ -504,7 +1178,9 @@ public async Task Provider_InitializationAndCleanupFailure_PreservesInitializati await Assert.ThrowsAsync(() => InvokeAsync(handler, "ping", timeout.Token)); // Assert - Assert.Contains("Failed to dispose MCP transport", listener.Output); + Assert.Contains( + hasProvider ? "Failed to dispose MCP transport" : "Failed to dispose MCP HTTP client", listener.Output); + Assert.Contains("transport cleanup failed", listener.Output); Assert.Single(stub.Handlers).Protected().Verify( "Dispose", Times.AtLeastOnce(), ItExpr.Is(disposing => disposing)); } @@ -560,6 +1236,11 @@ private static Task InvokeAsync( DefaultMcpToolHandler handler, string toolName, CancellationToken cancellationToken) => handler.InvokeToolAsync("https://mcp.example/api", null, toolName, null, null, null, cancellationToken); + private static Task InvokeScopedAsync( + DefaultMcpToolHandler handler, string workflowSessionId, string toolName, CancellationToken cancellationToken) => + handler.InvokeToolInWorkflowSessionAsync( + "https://mcp.example/api", null, toolName, null, null, null, workflowSessionId, cancellationToken); + private sealed class ProtocolStub { private int _initializations; @@ -569,9 +1250,11 @@ private sealed class ProtocolStub public int Terminations => this._terminations; public List> Handlers { get; } = []; public Func? BeforeOperationAsync { get; set; } + public Func? BeforeInitializationAsync { get; set; } public bool FailInitialization { get; set; } public bool FailOperation { get; set; } public bool FailTransportDisposal { get; set; } + public Exception? TransportDisposalException { get; set; } public HttpMessageHandler CreateMessageHandler() { @@ -579,10 +1262,10 @@ public HttpMessageHandler CreateMessageHandler() handler.Protected() .Setup>("SendAsync", ItExpr.IsAny(), ItExpr.IsAny()) .Returns(this.SendAsync); - if (this.FailTransportDisposal) + if (this.FailTransportDisposal || this.TransportDisposalException is not null) { handler.Protected().Setup("Dispose", ItExpr.Is(disposing => disposing)) - .Throws(new InvalidOperationException("transport cleanup failed")); + .Throws(this.TransportDisposalException ?? new InvalidOperationException("transport cleanup failed")); } this.Handlers.Add(handler); @@ -629,6 +1312,11 @@ private async Task SendAsync(HttpRequestMessage request, Ca string? sessionId = null; if (method == "initialize") { + if (this.BeforeInitializationAsync is not null) + { + await this.BeforeInitializationAsync(cancellationToken); + } + if (this.FailInitialization) { return EmptyResponse(HttpStatusCode.BadRequest, request); diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerTests.cs index aa1fb0440aa..b6eb32b667d 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerTests.cs @@ -471,7 +471,7 @@ public void ComputeHeadersHash_DifferentHeaders_ReturnsDifferentHash() // (InvokeToolAsync against a fake server) doesn't surface cache-hit behavior // without standing up a real MCP server — McpClient.CreateAsync fails before // _clients[key] = newClient runs, so nothing ever gets cached. - // Tuple equality on the returned 4-tuple verifies that the dimensions + // Tuple equality on the returned 5-tuple verifies that the dimensions // collectively discriminate cache entries. [Fact] @@ -481,19 +481,32 @@ public void BuildCacheKey_SameInputs_ReturnsEqualKeys() Dictionary headers = new() { ["Authorization"] = "Bearer token" }; // Act - var key1 = DefaultMcpToolHandler.BuildCacheKey("http://localhost/mcp", "label", "conn", headers); - var key2 = DefaultMcpToolHandler.BuildCacheKey("http://localhost/mcp", "label", "conn", headers); + var key1 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", "label", "conn", headers); + var key2 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", "label", "conn", headers); // Assert Assert.Equal(key2, key1); } + [Fact] + public void BuildCacheKey_DifferentWorkflowSession_ReturnsDifferentKeys() + { + // Act + var key1 = DefaultMcpToolHandler.BuildCacheKey("workflow-a", "http://localhost/mcp", "label", "conn", null); + var key2 = DefaultMcpToolHandler.BuildCacheKey("workflow-b", "http://localhost/mcp", "label", "conn", null); + + // Assert + Assert.NotEqual(key2, key1); + Assert.Equal("workflow-a", key1.WorkflowSession); + Assert.Equal("workflow-b", key2.WorkflowSession); + } + [Fact] public void BuildCacheKey_DifferentConnectionName_ReturnsDifferentKeys() { // Act - var key1 = DefaultMcpToolHandler.BuildCacheKey("http://localhost/mcp", "label", "connection-a", null); - var key2 = DefaultMcpToolHandler.BuildCacheKey("http://localhost/mcp", "label", "connection-b", null); + var key1 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", "label", "connection-a", null); + var key2 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", "label", "connection-b", null); // Assert Assert.NotEqual(key2, key1); @@ -505,8 +518,8 @@ public void BuildCacheKey_DifferentConnectionName_ReturnsDifferentKeys() public void BuildCacheKey_DifferentServerLabel_ReturnsDifferentKeys() { // Act - var key1 = DefaultMcpToolHandler.BuildCacheKey("http://localhost/mcp", "label-a", null, null); - var key2 = DefaultMcpToolHandler.BuildCacheKey("http://localhost/mcp", "label-b", null, null); + var key1 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", "label-a", null, null); + var key2 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", "label-b", null, null); // Assert Assert.NotEqual(key2, key1); @@ -519,8 +532,8 @@ public void BuildCacheKey_CaseSensitiveUrlPath_ReturnsDifferentKeys() { // Arrange — RFC 3986: URL path is case-sensitive // Act - var key1 = DefaultMcpToolHandler.BuildCacheKey("http://localhost/Tools", null, null, null); - var key2 = DefaultMcpToolHandler.BuildCacheKey("http://localhost/tools", null, null, null); + var key1 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/Tools", null, null, null); + var key2 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/tools", null, null, null); // Assert Assert.NotEqual(key2, key1); @@ -534,8 +547,8 @@ public void BuildCacheKey_HeaderValuesCaseSensitive_ReturnsDifferentKeys() Dictionary headers2 = new() { ["Authorization"] = "Bearer abc" }; // Act - var key1 = DefaultMcpToolHandler.BuildCacheKey("http://localhost/mcp", null, null, headers1); - var key2 = DefaultMcpToolHandler.BuildCacheKey("http://localhost/mcp", null, null, headers2); + var key1 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", null, null, headers1); + var key2 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", null, null, headers2); // Assert — header value case must propagate into the cache key Assert.NotEqual(key2, key1); @@ -546,7 +559,7 @@ public void BuildCacheKey_HeaderValuesCaseSensitive_ReturnsDifferentKeys() public void BuildCacheKey_NullLabelAndConnection_NormalizesToEmptyString() { // Act - var key = DefaultMcpToolHandler.BuildCacheKey("http://localhost/mcp", null, null, null); + var key = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", null, null, null); // Assert — verifies null-safety contract callers rely on Assert.Equal(string.Empty, key.Label); diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/Kit/IWorkflowContextExtensionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/Kit/IWorkflowContextExtensionsTests.cs index d26a1b6a84b..802f71cf48d 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/Kit/IWorkflowContextExtensionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/Kit/IWorkflowContextExtensionsTests.cs @@ -23,7 +23,7 @@ public async Task FormatTemplateAsync_WithSensitiveValue_ThrowsAsync() WorkflowFormulaState state = new(RecalcEngineFactory.Create()); state.Set("SOME_SECRET", FormulaValue.New("secret-value"), VariableScopeNames.Environment, SensitivityLevel.Sensitive); state.Bind(); - DeclarativeWorkflowContext context = new(new Mock().Object, state); + DeclarativeWorkflowContext context = await CreateContextAsync(state); // Act ValueTask FormatAsync() => context.FormatTemplateAsync("={Env.SOME_SECRET}"); @@ -40,7 +40,7 @@ public async Task FormatTemplateWithSensitivityAsync_WithSensitiveValue_ReturnsS WorkflowFormulaState state = new(RecalcEngineFactory.Create()); state.Set("SOME_SECRET", FormulaValue.New("secret-value"), VariableScopeNames.Environment, SensitivityLevel.Sensitive); state.Bind(); - DeclarativeWorkflowContext context = new(new Mock().Object, state); + DeclarativeWorkflowContext context = await CreateContextAsync(state); // Act EvaluationResult result = await context.FormatTemplateWithSensitivityAsync("={Env.SOME_SECRET}"); @@ -57,7 +57,7 @@ public async Task EvaluateValueAsync_WithSensitiveValue_ThrowsAsync() WorkflowFormulaState state = new(RecalcEngineFactory.Create()); state.Set(SystemScope.Names.LastMessageText, FormulaValue.New("secret-value"), VariableScopeNames.System, SensitivityLevel.Sensitive); state.Bind(); - DeclarativeWorkflowContext context = new(new Mock().Object, state); + DeclarativeWorkflowContext context = await CreateContextAsync(state); // Act ValueTask EvaluateAsync() => context.EvaluateValueAsync("System.LastMessageText"); @@ -74,7 +74,7 @@ public async Task EvaluateValueWithSensitivityAsync_WithSensitiveValue_ReturnsSe WorkflowFormulaState state = new(RecalcEngineFactory.Create()); state.Set(SystemScope.Names.LastMessageText, FormulaValue.New("secret-value"), VariableScopeNames.System, SensitivityLevel.Sensitive); state.Bind(); - DeclarativeWorkflowContext context = new(new Mock().Object, state); + DeclarativeWorkflowContext context = await CreateContextAsync(state); // Act EvaluationResult result = await context.EvaluateValueWithSensitivityAsync("System.LastMessageText"); @@ -91,7 +91,7 @@ public async Task QueueStateUpdateAsync_WithSensitivity_RebindsStateAsync() WorkflowFormulaState state = new(RecalcEngineFactory.Create()); state.Set("TestValue", FormulaValue.New("old-value")); state.Bind(); - DeclarativeWorkflowContext context = new(new Mock().Object, state); + DeclarativeWorkflowContext context = await CreateContextAsync(state); // Act await context.QueueStateUpdateAsync(PropertyPath.Create("Local.TestValue"), FormulaValue.New("new-value"), SensitivityLevel.Sensitive); @@ -113,7 +113,8 @@ public async Task ReadStateWithSensitivityAsync_QueuesSensitiveAssignmentAsync() source .Setup(c => c.ReadStateAsync(SystemScope.Names.LastMessageText, VariableScopeNames.System, default)) .Returns(new ValueTask("secret-value")); - DeclarativeWorkflowContext context = new(source.Object, state); + source.As().SetupGet(c => c.SessionId).Returns("test-session"); + DeclarativeWorkflowContext context = await DeclarativeWorkflowContext.CreateAsync(source.Object, state); // Act var evaluatedValue = await context.ReadStateWithSensitivityAsync(SystemScope.Names.LastMessageText, VariableScopeNames.System); @@ -177,7 +178,7 @@ public async Task GeneratedForeachPattern_WithSensitiveCollection_PreservesItemS VariableScopeNames.Environment, SensitivityLevel.Sensitive); state.Bind(); - DeclarativeWorkflowContext context = new(new Mock().Object, state); + DeclarativeWorkflowContext context = await CreateContextAsync(state); // Act EvaluationResult evaluatedValue = await context.EvaluateValueWithSensitivityAsync("Env.SensitiveItems"); @@ -213,4 +214,11 @@ public async Task ConvertValueWithSensitivityAsync_WithPlainContext_PreservesSen Assert.Equal(42M, result.Value); Assert.Equal(SensitivityLevel.Sensitive, result.Sensitivity); } + + private static async ValueTask CreateContextAsync(WorkflowFormulaState state) + { + Mock source = new(); + source.As().SetupGet(c => c.SessionId).Returns("test-session"); + return await DeclarativeWorkflowContext.CreateAsync(source.Object, state); + } } diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/Kit/RootExecutorTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/Kit/RootExecutorTests.cs index 4fedeb6d0f5..4d156d8c115 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/Kit/RootExecutorTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/Kit/RootExecutorTests.cs @@ -33,12 +33,13 @@ public async Task InitializeEnvironmentAsync_OnlyQueuesAllowedVariablesAsync() }; TestRootExecutor executor = new(options); Mock sourceContext = new(MockBehavior.Strict); + sourceContext.As().SetupGet(c => c.SessionId).Returns("test-session"); sourceContext.Setup(c => c.QueueStateUpdateAsync("ALLOWED", It.IsAny(), VariableScopeNames.Environment, It.IsAny())) .Returns(default(ValueTask)); sourceContext.Setup(c => c.QueueStateUpdateAsync("ALLOWED", SensitivityLevel.Sensitive, WorkflowFormulaState.GetSensitivityScopeName(VariableScopeNames.Environment), It.IsAny())) .Returns(default(ValueTask)); - DeclarativeWorkflowContext context = new(sourceContext.Object, executor.Session.State); + DeclarativeWorkflowContext context = await DeclarativeWorkflowContext.CreateAsync(sourceContext.Object, executor.Session.State); // Act await executor.InitializeAsync(context, "ALLOWED", "HIDDEN"); diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/ObjectModel/InvokeMcpToolExecutorTest.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/ObjectModel/InvokeMcpToolExecutorTest.cs index b3cd7d8c8d1..5051f3a818c 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/ObjectModel/InvokeMcpToolExecutorTest.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/ObjectModel/InvokeMcpToolExecutorTest.cs @@ -129,6 +129,86 @@ public async Task InvokeMcpToolExecuteWithoutApprovalAsync() await this.ExecuteTestAsync(model); } + [Fact] + public async Task InvokeMcpToolSeparateWorkflowSessionsUseSeparateSessionScopesAsync() + { + // Arrange + this.State.InitializeSystem(); + RecordingScopedMcpToolHandler handler = new(); + InvokeMcpTool model = this.CreateModel( + displayName: nameof(InvokeMcpToolSeparateWorkflowSessionsUseSeparateSessionScopesAsync), + serverUrl: TestServerUrl, + toolName: TestToolName, + requireApproval: false); + MockAgentProvider agentProvider = new(); + InvokeMcpToolExecutor action = new(model, handler, agentProvider.Object, this.State); + Mock firstContext = CreateMockWorkflowContext(); + Mock secondContext = CreateMockWorkflowContext(); + firstContext.As().SetupGet(context => context.SessionId).Returns("workflow-a"); + secondContext.As().SetupGet(context => context.SessionId).Returns("workflow-b"); + + // Act + await action.HandleAsync(new ActionExecutorResult("first"), firstContext.Object, CancellationToken.None); + await action.HandleAsync(new ActionExecutorResult("second"), secondContext.Object, CancellationToken.None); + + // Assert + Assert.Equal(2, handler.WorkflowSessionIds.Count); + Assert.NotEqual(handler.WorkflowSessionIds[0], handler.WorkflowSessionIds[1]); + } + + [Fact] + public async Task InvokeMcpToolSameWorkflowSessionReusesSessionScopeAsync() + { + // Arrange + this.State.InitializeSystem(); + RecordingScopedMcpToolHandler handler = new(); + InvokeMcpTool model = this.CreateModel( + displayName: nameof(InvokeMcpToolSameWorkflowSessionReusesSessionScopeAsync), + serverUrl: TestServerUrl, + toolName: TestToolName, + requireApproval: false); + MockAgentProvider agentProvider = new(); + InvokeMcpToolExecutor action = new(model, handler, agentProvider.Object, this.State); + Mock context = CreateMockWorkflowContext(); + context.As().SetupGet(current => current.SessionId).Returns("workflow-a"); + + // Act + await action.HandleAsync(new ActionExecutorResult("first"), context.Object, CancellationToken.None); + await action.HandleAsync(new ActionExecutorResult("second"), context.Object, CancellationToken.None); + + // Assert + Assert.Equal(2, handler.WorkflowSessionIds.Count); + Assert.Equal(handler.WorkflowSessionIds[0], handler.WorkflowSessionIds[1]); + } + + [Fact] + public async Task InvokeMcpToolLegacyContextsPersistSeparateSessionScopesAsync() + { + // Arrange + this.State.InitializeSystem(); + RecordingScopedMcpToolHandler handler = new(); + InvokeMcpTool model = this.CreateModel( + displayName: nameof(InvokeMcpToolLegacyContextsPersistSeparateSessionScopesAsync), + serverUrl: TestServerUrl, + toolName: TestToolName, + requireApproval: false); + MockAgentProvider agentProvider = new(); + InvokeMcpToolExecutor action = new(model, handler, agentProvider.Object, this.State); + InvokeMcpToolExecutor reconstructedAction = new(model, handler, agentProvider.Object, this.State); + Mock firstContext = CreateMockWorkflowContextWithSessionState(); + Mock secondContext = CreateMockWorkflowContextWithSessionState(); + + // Act + await action.HandleAsync(new ActionExecutorResult("first"), firstContext.Object, CancellationToken.None); + await action.HandleAsync(new ActionExecutorResult("second"), secondContext.Object, CancellationToken.None); + await reconstructedAction.HandleAsync(new ActionExecutorResult("continued"), firstContext.Object, CancellationToken.None); + + // Assert + Assert.Equal(3, handler.WorkflowSessionIds.Count); + Assert.NotEqual(handler.WorkflowSessionIds[0], handler.WorkflowSessionIds[1]); + Assert.Equal(handler.WorkflowSessionIds[0], handler.WorkflowSessionIds[2]); + } + [Fact] public async Task InvokeMcpToolExecuteWithServerLabelAsync() { @@ -275,12 +355,13 @@ public async Task AutoSendUsesMcpApprovalSnapshotAsync(bool autoSend, bool resto nameof(AutoSendUsesMcpApprovalSnapshotAsync), TestServerUrl, TestToolName, requireApproval: true).ToBuilder(); builder.Output = new InvokeToolOutput.Builder { AutoSend = BoolExpression.Expression("Local.Control").ToBuilder() }; InvokeMcpTool model = AssignParent(builder); - MockMcpToolProvider handler = new(); + RecordingScopedMcpToolHandler handler = new(); MockAgentProvider provider = new(); - InvokeMcpToolExecutor action = new(model, handler.Object, provider.Object, this.State); + InvokeMcpToolExecutor action = new(model, handler, provider.Object, this.State); List requests = []; List events = []; Mock context = CreateMockWorkflowContextWithStateStore(requests); + context.As().SetupGet(current => current.SessionId).Returns("workflow-a"); context.Setup(c => c.AddEventAsync(It.IsAny(), It.IsAny())) .Callback((e, _) => events.Add(e)) .Returns(default(ValueTask)); @@ -288,7 +369,7 @@ public async Task AutoSendUsesMcpApprovalSnapshotAsync(bool autoSend, bool resto if (restore) { await InvokeProtectedMethodAsync(action, "OnCheckpointingAsync", context.Object, CancellationToken.None); - action = new InvokeMcpToolExecutor(model, handler.Object, provider.Object, this.State); + action = new InvokeMcpToolExecutor(model, handler, provider.Object, this.State); await InvokeProtectedMethodAsync(action, "OnCheckpointRestoredAsync", context.Object, CancellationToken.None); } this.State.Set("Control", FormulaValue.New(!autoSend)); @@ -299,6 +380,7 @@ public async Task AutoSendUsesMcpApprovalSnapshotAsync(bool autoSend, bool resto // Assert Assert.Equal(autoSend ? 1 : 0, events.OfType().Count()); + Assert.Equal("workflow-a", Assert.Single(handler.WorkflowSessionIds)); } [Fact] @@ -2524,6 +2606,28 @@ private static Mock CreateMockWorkflowContext(List CreateMockWorkflowContextWithSessionState() + { + Mock context = CreateMockWorkflowContext(); + Dictionary<(string? Scope, string Key), string> state = []; + context.Setup(current => current.ReadOrInitStateAsync( + It.IsAny(), + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Returns((string key, Func initialStateFactory, string? scopeName, CancellationToken _) => + { + if (!state.TryGetValue((scopeName, key), out string? value)) + { + value = initialStateFactory(); + state[(scopeName, key)] = value; + } + + return new ValueTask(value); + }); + return context; + } + /// /// Creates a mock workflow context that actually stores state values (for checkpoint/restore tests). /// Optionally accepts an externally-owned state store so callers can drive multi-step @@ -2839,6 +2943,35 @@ public enum McpInvocationValueLocation #region Mock MCP Tool Provider + private sealed class RecordingScopedMcpToolHandler : IWorkflowScopedMcpToolHandler + { + public List WorkflowSessionIds { get; } = []; + + public Task InvokeToolAsync( + string serverUrl, + string? serverLabel, + string toolName, + IDictionary? arguments, + IDictionary? headers, + string? connectionName, + CancellationToken cancellationToken = default) => + throw new InvalidOperationException("The workflow-scoped overload must be used."); + + public Task InvokeToolInWorkflowSessionAsync( + string serverUrl, + string? serverLabel, + string toolName, + IDictionary? arguments, + IDictionary? headers, + string? connectionName, + string workflowSessionId, + CancellationToken cancellationToken = default) + { + this.WorkflowSessionIds.Add(workflowSessionId); + return Task.FromResult(new McpServerToolResultContent("mock-call-id") { Outputs = [] }); + } + } + /// /// Mock implementation of for unit testing purposes. /// diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/PowerFx/WorkflowFormulaStateTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/PowerFx/WorkflowFormulaStateTests.cs index 34db6f294eb..5d0ffd7660a 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/PowerFx/WorkflowFormulaStateTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/PowerFx/WorkflowFormulaStateTests.cs @@ -4,6 +4,7 @@ using System.Collections.Generic; using System.Threading; using System.Threading.Tasks; +using Microsoft.Agents.AI.Workflows.Declarative.Interpreter; using Microsoft.Agents.AI.Workflows.Declarative.PowerFx; using Microsoft.Agents.ObjectModel; using Microsoft.Agents.ObjectModel.Abstractions; @@ -164,6 +165,34 @@ public async Task RestoreAsync_KeepsCaseVariantVariablesIndependentAsync(bool re Assert.Equal(SensitivityLevel.None, this.State.GetSensitivity("ITEM")); } + [Fact] + public async Task DeclarativeContextFallbackSessionId_IsScopedToPersistedRunStateAsync() + { + // Arrange + Dictionary<(string? ScopeName, string Key), string> firstRunState = []; + Dictionary<(string? ScopeName, string Key), string> secondRunState = []; + IWorkflowContext firstContext = CreateContext(firstRunState); + IWorkflowContext restoredContext = CreateContext(firstRunState); + IWorkflowContext secondContext = CreateContext(secondRunState); + + // Act + DeclarativeWorkflowContext first = + await DeclarativeWorkflowContext.CreateAsync(firstContext, this.State); + DeclarativeWorkflowContext continued = + await DeclarativeWorkflowContext.CreateAsync(firstContext, this.State); + DeclarativeWorkflowContext restored = + await DeclarativeWorkflowContext.CreateAsync(restoredContext, this.State); + DeclarativeWorkflowContext second = + await DeclarativeWorkflowContext.CreateAsync(secondContext, this.State); + + // Assert + Assert.Equal(first.SessionId, continued.SessionId); + Assert.Equal(first.SessionId, restored.SessionId); + Assert.NotEqual(first.SessionId, second.SessionId); + Assert.True(firstRunState.ContainsKey((VariableScopeNames.System, "__declarative_mcp_workflow_session_id"))); + Assert.False(firstRunState.ContainsKey((null, "__declarative_mcp_workflow_session_id"))); + } + [Fact] public async Task RestoreAsync_RestoresPersistedSensitivityAsync() { @@ -182,4 +211,27 @@ public async Task RestoreAsync_RestoresPersistedSensitivityAsync() // Assert Assert.Equal(SensitivityLevel.Sensitive, this.State.GetSensitivity("secret")); } + + private static IWorkflowContext CreateContext(Dictionary<(string? ScopeName, string Key), string> state) + { + Mock context = new(); + context + .Setup(current => current.ReadOrInitStateAsync( + It.IsAny(), + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Returns((string key, Func factory, string? scopeName, CancellationToken cancellationToken) => + { + var scopedKey = (scopeName, key); + if (!state.TryGetValue(scopedKey, out string? value)) + { + value = factory(); + state[scopedKey] = value; + } + + return new ValueTask(value); + }); + return context.Object; + } } diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/RepresentationTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/RepresentationTests.cs index 1e92d130de9..bddef5e4f19 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/RepresentationTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/RepresentationTests.cs @@ -8,6 +8,7 @@ using System.Threading; using System.Threading.Tasks; using Microsoft.Agents.AI.Workflows.Checkpointing; +using Microsoft.Agents.AI.Workflows.InProc; using Microsoft.Agents.AI.Workflows.Sample; using Microsoft.Agents.AI.Workflows.Specialized; using Microsoft.Extensions.AI; @@ -91,6 +92,48 @@ public async Task Test_SpecializedExecutor_InfosAsync() await RunExecutorBindingInfoMatchTestAsync(new RequestInfoExecutor(TestRequestPort)); } + [Fact] + public void SubworkflowSessionId_IsStableAndHierarchicallyScoped() + { + // Arrange + const string ParentSessionId = "parent-session"; + + // Act + string firstChild = SubworkflowBinding.CreateSubworkflowSessionId(ParentSessionId, "first"); + string sameChild = SubworkflowBinding.CreateSubworkflowSessionId(ParentSessionId, "first"); + string siblingChild = SubworkflowBinding.CreateSubworkflowSessionId(ParentSessionId, "second"); + string nestedChild = SubworkflowBinding.CreateSubworkflowSessionId(firstChild, "nested"); + + // Assert + Assert.Equal(firstChild, sameChild); + Assert.NotEqual(ParentSessionId, firstChild); + Assert.NotEqual(firstChild, siblingChild); + Assert.NotEqual(firstChild, nestedChild); + Assert.NotEqual(siblingChild, nestedChild); + } + + [Fact] + public async Task SubworkflowRunner_PreservesLegacyCheckpointSessionIdAsync() + { + // Arrange + const string LegacySessionId = "legacy-session"; + string workflowSessionId = SubworkflowBinding.CreateSubworkflowSessionId(LegacySessionId, "child"); + TestExecutor executor = new(); + Workflow workflow = new WorkflowBuilder(executor).Build(); + + // Act + InProcessRunner runner = InProcessRunner.CreateSubworkflowRunner( + workflow, + checkpointManager: new InMemoryCheckpointManager(), + sessionId: LegacySessionId, + workflowSessionId: workflowSessionId); + + // Assert + Assert.Equal(LegacySessionId, runner.SessionId); + Assert.Equal(workflowSessionId, runner.WorkflowSessionId); + await runner.RequestEndRunAsync(); + } + private static string Source(int id) => $"Source/{id}"; private static string Sink(int id) => $"Sink/{id}"; From 32ce3efe1a59a9902764e4e8c583e59655e72b4f Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 20 Oct 2026 10:48:06 -0400 Subject: [PATCH 2/3] refactor: separate MCP client caching and lifetime management Extract workflow-scoped caching, connection creation, provider lifetime coordination, and resource cleanup from DefaultMcpToolHandler without changing session behavior. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 94e932cb-8d65-4169-8c8b-3b5ceff9d09d --- .../DefaultMcpToolHandler.cs | 876 +----------------- .../McpClientCache.cs | 569 ++++++++++++ .../McpClientConnection.cs | 55 ++ .../McpClientFactory.cs | 218 +++++ .../McpClientLifetime.cs | 132 +++ .../DefaultMcpToolHandlerLifetimeTests.cs | 4 +- .../DefaultMcpToolHandlerTests.cs | 46 +- 7 files changed, 1016 insertions(+), 884 deletions(-) create mode 100644 dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/McpClientCache.cs create mode 100644 dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/McpClientConnection.cs create mode 100644 dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/McpClientFactory.cs create mode 100644 dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/McpClientLifetime.cs diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs index 1f5664a36c8..3cc8e8d01f2 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs @@ -6,8 +6,6 @@ using System.IO; using System.Linq; using System.Net.Http; -using System.Runtime.ExceptionServices; -using System.Security.Cryptography; using System.Text; using System.Text.Json; using System.Threading; @@ -49,25 +47,9 @@ public sealed class DefaultMcpToolHandler : IWorkflowScopedMcpToolHandler, IAsyn public const string ListToolsToolName = "tools/list"; private static readonly JsonWriterOptions s_toolListJsonWriterOptions = new() { Indented = true }; - private static readonly OperationCanceledException s_clientCreationCancelledException = - new("The MCP client creator was cancelled."); - private readonly Func>? _httpClientProvider; - private readonly Func _httpMessageHandlerFactory; - private readonly Func _clientConnectionDisposer; - private readonly Dictionary<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash), CachedClient> _clients = []; - private readonly Dictionary<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash), TaskCompletionSource> _clientCreations = []; - private readonly HashSet> _clientCreationLifetimes = []; - private readonly LinkedList<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash)> _clientLru = []; - private readonly HashSet _retiredClients = []; - private readonly Dictionary _ownedHttpClients = []; - private readonly SemaphoreSlim _clientLock = new(1, 1); - private readonly SemaphoreSlim _clientCreationSemaphore; - private readonly int _clientCacheMaxSize; - private readonly AsyncLocal _providerInvocationContext = new(); - private TaskCompletionSource? _providerInvocationsDrained; - private int _activeProviderInvocations; - private bool _disposing; + private readonly McpClientFactory _clientFactory; + private readonly McpClientCache _clientCache; /// /// Initializes a new instance of the class. @@ -92,7 +74,7 @@ public sealed class DefaultMcpToolHandler : IWorkflowScopedMcpToolHandler, IAsyn /// /// public DefaultMcpToolHandler(Func>? httpClientProvider = null) - : this(httpClientProvider, CreateHttpMessageHandler) + : this(httpClientProvider, McpClientFactory.CreateHttpMessageHandler) { } @@ -100,18 +82,11 @@ internal DefaultMcpToolHandler( Func>? httpClientProvider, Func httpMessageHandlerFactory, int clientCacheMaxSize = DefaultClientCacheMaxSize, - Func? clientConnectionDisposer = null) + Func? clientConnectionDisposer = null) { - if (clientCacheMaxSize <= 0) - { - throw new ArgumentOutOfRangeException(nameof(clientCacheMaxSize), "The MCP client cache size must be positive."); - } - - this._httpClientProvider = httpClientProvider; - this._httpMessageHandlerFactory = Throw.IfNull(httpMessageHandlerFactory); - this._clientConnectionDisposer = clientConnectionDisposer ?? (connection => connection.DisposeAsync()); - this._clientCacheMaxSize = clientCacheMaxSize; - this._clientCreationSemaphore = new(clientCacheMaxSize, clientCacheMaxSize); + McpClientLifetime lifetime = new(); + this._clientFactory = new(httpClientProvider, httpMessageHandlerFactory, lifetime); + this._clientCache = new(this._clientFactory, lifetime, clientCacheMaxSize, clientConnectionDisposer); } /// @@ -164,63 +139,18 @@ public async Task InvokeToolInWorkflowSessionAsync( } } - if (this._httpClientProvider is not null) + if (this._clientFactory.HasHttpClientProvider) { - TaskCompletionSource invocationCompleted = new(TaskCreationOptions.RunContinuationsAsynchronously); - await this._clientLock.WaitAsync(cancellationToken).ConfigureAwait(false); - try - { - this.ThrowIfDisposing(); - if (this._activeProviderInvocations++ == 0) - { - this._providerInvocationsDrained = new(TaskCreationOptions.RunContinuationsAsynchronously); - } - } - finally - { - this._clientLock.Release(); - } - - ProviderInvocationContext? previousInvocation = this._providerInvocationContext.Value; - this._providerInvocationContext.Value = new(invocationCompleted.Task, previousInvocation); - try - { - ClientConnection invocationClient = await this.CreateClientAsync( - serverUrl.Trim(), serverLabel, headers, cancellationToken).ConfigureAwait(false); - await using (invocationClient.ConfigureAwait(false)) - { - return await InvokeClientAsync(invocationClient.Client, toolName, arguments, cancellationToken).ConfigureAwait(false); - } - } - finally - { - await this._clientLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); - try - { - if (--this._activeProviderInvocations == 0) - { - this._providerInvocationsDrained!.SetResult(true); - } - } - finally - { - invocationCompleted.SetResult(true); - this._providerInvocationContext.Value = previousInvocation; - this._clientLock.Release(); - } - } + return await this._clientFactory.InvokeProviderAsync( + serverUrl, serverLabel, headers, + client => InvokeClientAsync(client, toolName, arguments, cancellationToken), + cancellationToken).ConfigureAwait(false); } - CachedClient cachedClient = await this.AcquireClientAsync( - serverUrl, serverLabel, headers, connectionName, workflowSessionId, cancellationToken).ConfigureAwait(false); - try - { - return await InvokeClientAsync(cachedClient.Connection.Client, toolName, arguments, cancellationToken).ConfigureAwait(false); - } - finally - { - await this.ReleaseClientAsync(cachedClient).ConfigureAwait(false); - } + return await this._clientCache.InvokeAsync( + serverUrl, serverLabel, headers, connectionName, workflowSessionId, + client => InvokeClientAsync(client, toolName, arguments, cancellationToken), + cancellationToken).ConfigureAwait(false); } private static async Task InvokeClientAsync( @@ -275,779 +205,7 @@ internal static McpServerToolResultContent CreateListToolsResultContent(IEnumera /// Disposal is requested from an active provider-backed invocation, including a child execution context. /// Dispose the handler from its owning scope instead. /// - public async ValueTask DisposeAsync() - { - for (ProviderInvocationContext? context = this._providerInvocationContext.Value; context is not null; context = context.Parent) - { - if (!context.Completion.IsCompleted) - { - throw new InvalidOperationException("Cannot dispose the MCP handler from an active provider-backed invocation."); - } - } - - Task? providerInvocations; - List> clientCreationLifetimes; - await this._clientLock.WaitAsync().ConfigureAwait(false); - try - { - this.ThrowIfDisposing(); - this._disposing = true; - providerInvocations = this._providerInvocationsDrained?.Task; - clientCreationLifetimes = this._clientCreationLifetimes.Select(source => source.Task).ToList(); - } - finally - { - this._clientLock.Release(); - } - - if (providerInvocations is not null) - { - await providerInvocations.ConfigureAwait(false); - } - - Exception? clientCreationException = null; - try - { - await Task.WhenAll(clientCreationLifetimes).ConfigureAwait(false); - } - catch (Exception exception) when (!IsFatalException(exception)) - { - clientCreationException = exception; - } - - List cachedClients; - List clientsToDispose = []; - await this._clientLock.WaitAsync().ConfigureAwait(false); - try - { - cachedClients = [.. this._clients.Values, .. this._retiredClients]; - foreach (CachedClient client in cachedClients) - { - client.Evicted = true; - this._retiredClients.Add(client); - if (client.ActiveInvocations == 0 && !client.DisposalClaimed) - { - client.DisposalClaimed = true; - clientsToDispose.Add(client); - } - } - - this._clients.Clear(); - this._clientLru.Clear(); - } - finally - { - this._clientLock.Release(); - } - - Exception? clientCleanupException = null; - try - { - await DrainCleanupAsync( - clientsToDispose.Select(client => this.DisposeCachedClientAsync(client)), - cachedClients.Select(client => client.Disposed.Task)).ConfigureAwait(false); - } - catch (Exception exception) when (!IsFatalException(exception)) - { - clientCleanupException = exception; - } - finally - { - this._clientLock.Dispose(); - this._clientCreationSemaphore.Dispose(); - } - - if (clientCreationException is not null) - { - ExceptionDispatchInfo.Capture(clientCreationException).Throw(); - } - - if (clientCleanupException is not null) - { - ExceptionDispatchInfo.Capture(clientCleanupException).Throw(); - } - } - - internal static async Task DrainCleanupAsync(IEnumerable cleanupTasks, IEnumerable completionTasks) - { - Exception? cleanupException = null; - try - { - await Task.WhenAll(cleanupTasks).ConfigureAwait(false); - } - catch (Exception exception) when (!IsFatalException(exception)) - { - cleanupException = exception; - } - - await Task.WhenAll(completionTasks).ConfigureAwait(false); - - if (cleanupException is not null) - { - ExceptionDispatchInfo.Capture(cleanupException).Throw(); - } - } - - private static bool IsFatalException(Exception exception) => - exception is OutOfMemoryException and not InsufficientMemoryException - or StackOverflowException - or AccessViolationException - or AppDomainUnloadedException - or BadImageFormatException - or CannotUnloadAppDomainException - or InvalidProgramException - or ThreadAbortException; - - [System.Diagnostics.CodeAnalysis.SuppressMessage("Maintainability", "CA1513:Use ObjectDisposedException throw helper", - Justification = "The helper is not available on .NET Framework or .NET Standard 2.0.")] - private void ThrowIfDisposing() - { - if (this._disposing) - { - throw new ObjectDisposedException(nameof(DefaultMcpToolHandler)); - } - } - - private async Task AcquireClientAsync( - string serverUrl, - string? serverLabel, - IDictionary? headers, - string? connectionName, - string workflowSessionId, - CancellationToken cancellationToken) - { - string trimmedUrl = serverUrl.Trim(); - var clientCacheKey = BuildCacheKey(workflowSessionId, trimmedUrl, serverLabel, connectionName, headers); - CachedClient? clientToDispose = null; - TaskCompletionSource? clientCreation; - TaskCompletionSource? clientCreationLifetime = null; - bool ownsClientCreation = false; - - await this._clientLock.WaitAsync(cancellationToken).ConfigureAwait(false); - try - { - this.ThrowIfDisposing(); - if (this._clients.TryGetValue(clientCacheKey, out CachedClient? existingClient)) - { - existingClient.ActiveInvocations++; - this._clientLru.Remove(existingClient.LruNode); - this._clientLru.AddLast(existingClient.LruNode); - return existingClient; - } - - if (!this._clientCreations.TryGetValue(clientCacheKey, out clientCreation)) - { - if (this._clients.Count + this._retiredClients.Count + this._clientCreations.Count >= this._clientCacheMaxSize) - { - LinkedListNode<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash)>? idleNode = - this._clientLru.First; - while (idleNode is not null && this._clients[idleNode.Value].ActiveInvocations != 0) - { - idleNode = idleNode.Next; - } - - if (idleNode is null) - { - throw new InvalidOperationException("MCP client capacity is exhausted; retry when an active invocation completes."); - } - - clientToDispose = this._clients[idleNode.Value]; - this._clients.Remove(idleNode.Value); - this._clientLru.Remove(idleNode); - clientToDispose.Evicted = true; - clientToDispose.DisposalClaimed = true; - this._retiredClients.Add(clientToDispose); - } - - clientCreation = new(TaskCreationOptions.RunContinuationsAsynchronously); - this._clientCreations[clientCacheKey] = clientCreation; - clientCreationLifetime = new(TaskCreationOptions.RunContinuationsAsynchronously); - this._clientCreationLifetimes.Add(clientCreationLifetime); - ownsClientCreation = true; - } - } - finally - { - this._clientLock.Release(); - } - - if (!ownsClientCreation) - { - try - { - await WaitForClientCreationAsync(clientCreation.Task, cancellationToken).ConfigureAwait(false); - } - catch (OperationCanceledException exception) - when (ReferenceEquals(exception, s_clientCreationCancelledException)) - { - return await this.AcquireClientAsync( - serverUrl, serverLabel, headers, connectionName, workflowSessionId, cancellationToken).ConfigureAwait(false); - } - - return await this.AcquireClientAsync( - serverUrl, serverLabel, headers, connectionName, workflowSessionId, cancellationToken).ConfigureAwait(false); - } - - TaskCompletionSource ownedClientCreationLifetime = clientCreationLifetime ?? - throw new InvalidOperationException("Missing MCP client creation lifetime."); - ClientConnection? connection = null; - bool creationSemaphoreEntered = false; - Exception? clientCreationCleanupException = null; - try - { - if (clientToDispose is not null) - { - try - { - await this.DisposeCachedClientAsync(clientToDispose).ConfigureAwait(false); - } - catch (Exception cleanupException) - { - clientCreationCleanupException = cleanupException; - throw; - } - } - - await this._clientCreationSemaphore.WaitAsync(cancellationToken).ConfigureAwait(false); - creationSemaphoreEntered = true; - await this._clientLock.WaitAsync(cancellationToken).ConfigureAwait(false); - try - { - this.ThrowIfDisposing(); - } - finally - { - this._clientLock.Release(); - } - - connection = await this.CreateClientAsync( - trimmedUrl, serverLabel, headers, cancellationToken, - cleanupException => clientCreationCleanupException = cleanupException).ConfigureAwait(false); - } - catch (Exception exception) - { - Exception sharedCreationException = - exception is OperationCanceledException && cancellationToken.IsCancellationRequested - ? s_clientCreationCancelledException - : exception; - try - { - await this.CompleteClientCreationFailureAsync( - clientCacheKey, clientCreation, sharedCreationException).ConfigureAwait(false); - } - finally - { - if (creationSemaphoreEntered) - { - this._clientCreationSemaphore.Release(); - creationSemaphoreEntered = false; - } - - await this.CompleteClientCreationLifetimeAsync( - ownedClientCreationLifetime, clientCreationCleanupException).ConfigureAwait(false); - } - - throw; - } - - ObjectDisposedException? disposedException = null; - CachedClient? result = null; - await this._clientLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); - try - { - this._clientCreations.Remove(clientCacheKey); - if (this._disposing) - { - disposedException = new ObjectDisposedException(nameof(DefaultMcpToolHandler)); - } - else - { - LinkedListNode<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash)> node = - this._clientLru.AddLast(clientCacheKey); - CachedClient newClient = new(connection, node) { ActiveInvocations = 1 }; - connection = null; - this._clients[clientCacheKey] = newClient; - - result = newClient; - clientCreation.TrySetResult(newClient); - } - } - finally - { - this._clientLock.Release(); - } - - try - { - if (connection is not null) - { - await this._clientConnectionDisposer(connection).ConfigureAwait(false); - } - - if (disposedException is not null) - { - clientCreation.TrySetException(disposedException); - throw disposedException; - } - } - catch (Exception exception) - { - if (disposedException is not null && !ReferenceEquals(exception, disposedException)) - { - clientCreation.TrySetException(exception); - clientCreationCleanupException = exception; - } - - if (result is not null) - { - await this.ReleaseClientAsync(result).ConfigureAwait(false); - } - - throw; - } - finally - { - if (creationSemaphoreEntered) - { - this._clientCreationSemaphore.Release(); - } - - await this.CompleteClientCreationLifetimeAsync( - ownedClientCreationLifetime, clientCreationCleanupException).ConfigureAwait(false); - } - - return result ?? throw new InvalidOperationException("Failed to acquire MCP client."); - } - - private static async Task WaitForClientCreationAsync(Task clientCreation, CancellationToken cancellationToken) - { - if (!cancellationToken.CanBeCanceled || clientCreation.IsCompleted) - { - await clientCreation.ConfigureAwait(false); - return; - } - - TaskCompletionSource cancellation = new(TaskCreationOptions.RunContinuationsAsynchronously); - using CancellationTokenRegistration registration = cancellationToken.Register( - static state => ((TaskCompletionSource)state!).TrySetResult(true), - cancellation); - - if (await Task.WhenAny(clientCreation, cancellation.Task).ConfigureAwait(false) == cancellation.Task) - { - cancellationToken.ThrowIfCancellationRequested(); - } - - await clientCreation.ConfigureAwait(false); - } - - private async Task CompleteClientCreationFailureAsync( - (string WorkflowSession, string Url, string Label, string Connection, string HeadersHash) clientCacheKey, - TaskCompletionSource clientCreation, - Exception exception) - { - await this._clientLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); - try - { - if (this._clientCreations.TryGetValue(clientCacheKey, out TaskCompletionSource? existingCreation) && - ReferenceEquals(existingCreation, clientCreation)) - { - this._clientCreations.Remove(clientCacheKey); - } - - clientCreation.TrySetException(exception); - _ = clientCreation.Task.Exception; - } - finally - { - this._clientLock.Release(); - } - } - - private async Task CompleteClientCreationLifetimeAsync( - TaskCompletionSource clientCreationLifetime, - Exception? cleanupException = null) - { - await this._clientLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); - try - { - this._clientCreationLifetimes.Remove(clientCreationLifetime); - } - finally - { - this._clientLock.Release(); - } - - if (cleanupException is not null) - { - clientCreationLifetime.TrySetException(cleanupException); - } - else - { - clientCreationLifetime.TrySetResult(true); - } - } - - private async ValueTask ReleaseClientAsync(CachedClient client) - { - bool dispose = false; - await this._clientLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); - try - { - client.ActiveInvocations--; - if (client.ActiveInvocations == 0 && client.Evicted && !client.DisposalClaimed) - { - client.DisposalClaimed = true; - dispose = true; - } - } - finally - { - this._clientLock.Release(); - } - - if (dispose) - { - await this.DisposeCachedClientAsync(client).ConfigureAwait(false); - } - } - - private async Task DisposeCachedClientAsync(CachedClient client) - { - Exception? cleanupException = null; - try - { - await this._clientConnectionDisposer(client.Connection).ConfigureAwait(false); - } - catch (Exception exception) - { - cleanupException = exception; - throw; - } - finally - { - await this._clientLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); - try - { - this._retiredClients.Remove(client); - } - finally - { - this._clientLock.Release(); - } - - if (cleanupException is not null) - { - client.Disposed.TrySetException(cleanupException); - } - else - { - client.Disposed.TrySetResult(true); - } - } - } - - /// - /// Builds the per-client cache key as a 5-tuple of - /// (workflowSessionId, trimmed serverUrl, serverLabel, connectionName, headers hash). - /// All five components participate so that separate workflow sessions and callers using - /// different labels/connections/headers receive distinct instances - /// even when targeting the same URL. - /// - internal static (string WorkflowSession, string Url, string Label, string Connection, string HeadersHash) BuildCacheKey( - string workflowSessionId, - string trimmedUrl, - string? serverLabel, - string? connectionName, - IDictionary? headers) => - (workflowSessionId, trimmedUrl, serverLabel ?? string.Empty, connectionName ?? string.Empty, ComputeHeadersHash(headers)); - - private async Task CreateClientAsync( - string serverUrl, - string? serverLabel, - IDictionary? headers, - CancellationToken cancellationToken, - Action? reportCleanupFailure = null) - { - HttpClient? httpClient = null; - bool ownsHttpClient = false; - OwnedHttpClientLease? ownedHttpClientLease = null; - - if (this._httpClientProvider is not null) - { - httpClient = await this._httpClientProvider(serverUrl, cancellationToken).ConfigureAwait(false); - } - - if (httpClient is null && this._httpClientProvider is not null) - { - httpClient = this.CreatePinnedHttpClient(serverUrl); - ownsHttpClient = true; - } - else if (httpClient is null) - { - ownedHttpClientLease = await this.AcquireOwnedHttpClientAsync(serverUrl, cancellationToken).ConfigureAwait(false); - httpClient = ownedHttpClientLease.Client; - } - - HttpClientTransportOptions transportOptions = new() - { - Endpoint = new Uri(serverUrl), - Name = serverLabel ?? "McpClient", - AdditionalHeaders = headers, - // Use Streamable HTTP rather than AutoDetect so the client does not negotiate down to the - // legacy HTTP+SSE transport, which trusts a server-advertised message endpoint. That - // server-controlled endpoint is the primary vector for redirecting the Authorization token - // to a different origin; Streamable HTTP keeps every request on the configured origin. - TransportMode = HttpTransportMode.StreamableHttp - }; - - HttpClientTransport? transport = null; - try - { - HttpClient resolvedHttpClient = httpClient ?? throw new InvalidOperationException("Failed to resolve MCP HTTP client."); - transport = new(transportOptions, resolvedHttpClient, ownsHttpClient: ownsHttpClient); - McpClient client = await McpClient.CreateAsync(transport, cancellationToken: cancellationToken).ConfigureAwait(false); - ClientConnection connection = new(client, transport, ownedHttpClientLease); - ownedHttpClientLease = null; - return connection; - } - catch - { - try - { - try - { - if (transport is not null) - { - await DisposeResourceAsync(transport, "transport").ConfigureAwait(false); - } - } - finally - { - if (ownedHttpClientLease is not null) - { - await ownedHttpClientLease.DisposeAsync().ConfigureAwait(false); - } - } - } - catch (Exception cleanupException) - { - reportCleanupFailure?.Invoke(cleanupException); - throw; - } - - throw; - } - } - - private async Task AcquireOwnedHttpClientAsync( - string serverUrl, - CancellationToken cancellationToken) - { - await this._clientLock.WaitAsync(cancellationToken).ConfigureAwait(false); - try - { - if (!this._ownedHttpClients.TryGetValue(serverUrl, out OwnedHttpClient? entry)) - { - entry = new(this.CreatePinnedHttpClient(serverUrl)); - this._ownedHttpClients[serverUrl] = entry; - } - - entry.ReferenceCount++; - return new(this, serverUrl, entry.Client); - } - finally - { - this._clientLock.Release(); - } - } - - private async ValueTask ReleaseOwnedHttpClientAsync(string serverUrl) - { - HttpClient? clientToDispose = null; - await this._clientLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); - try - { - OwnedHttpClient entry = this._ownedHttpClients[serverUrl]; - if (--entry.ReferenceCount == 0) - { - this._ownedHttpClients.Remove(serverUrl); - clientToDispose = entry.Client; - } - } - finally - { - this._clientLock.Release(); - } - - try - { - clientToDispose?.Dispose(); - } - catch (Exception exception) when (exception is not OperationCanceledException) - { - Trace.TraceWarning("Failed to dispose MCP HTTP client: {0}", exception); - } - } - - private HttpClient CreatePinnedHttpClient(string serverUrl) - { - // Pin credential headers to the configured server origin as defense-in-depth. Forcing - // StreamableHttp (below) already removes the primary vector (a server-advertised cross-origin - // SSE message endpoint), and AllowAutoRedirect=false blocks auto-redirects. This handler is the - // backstop: it guarantees the Authorization token and other credentials never leave the pinned - // origin even if a future change re-enables AutoDetect or redirects, or the SDK constructs a - // request to a new URI (AdditionalHeaders are re-stamped by the transport, so HttpClient's own - // redirect header-stripping does not cover them). - OriginPinningHandler pinningHandler = new(new Uri(serverUrl)) { InnerHandler = this._httpMessageHandlerFactory() }; - return new HttpClient(pinningHandler); - } - - private static HttpMessageHandler CreateHttpMessageHandler() => - new HttpClientHandler - { - // Keep cookies out of pooled transport state and prevent credential-bearing redirects. - UseCookies = false, - AllowAutoRedirect = false, - CheckCertificateRevocationList = true - }; - - private sealed class ProviderInvocationContext(Task completion, ProviderInvocationContext? parent) - { - public Task Completion { get; } = completion; - public ProviderInvocationContext? Parent { get; } = parent; - } - - private sealed class CachedClient( - ClientConnection connection, - LinkedListNode<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash)> lruNode) - { - public ClientConnection Connection { get; } = connection; - - public LinkedListNode<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash)> LruNode { get; } = lruNode; - - public TaskCompletionSource Disposed { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously); - - public int ActiveInvocations { get; set; } - - public bool Evicted { get; set; } - - public bool DisposalClaimed { get; set; } - } - - private sealed class OwnedHttpClient(HttpClient client) - { - public HttpClient Client { get; } = client; - - public int ReferenceCount { get; set; } - } - - private sealed class OwnedHttpClientLease( - DefaultMcpToolHandler owner, - string serverUrl, - HttpClient client) : IAsyncDisposable - { - private bool _disposed; - - public HttpClient Client { get; } = client; - - public async ValueTask DisposeAsync() - { - if (!this._disposed) - { - this._disposed = true; - await owner.ReleaseOwnedHttpClientAsync(serverUrl).ConfigureAwait(false); - } - } - } - - internal sealed class ClientConnection( - McpClient client, - IAsyncDisposable transport, - IAsyncDisposable? ownedHttpClientLease = null) : IAsyncDisposable - { - public McpClient Client { get; } = client; - - public async ValueTask DisposeAsync() - { - try - { - await DisposeResourceAsync(this.Client, "session").ConfigureAwait(false); - } - finally - { - try - { - // McpClient owns the connected session, not the reusable transport factory. - await DisposeResourceAsync(transport, "transport").ConfigureAwait(false); - } - finally - { - if (ownedHttpClientLease is not null) - { - await ownedHttpClientLease.DisposeAsync().ConfigureAwait(false); - } - } - } - } - } - - private static async ValueTask DisposeResourceAsync(T resource, string resourceName) - where T : IAsyncDisposable - { - try - { - await resource.DisposeAsync().ConfigureAwait(false); - } - catch (Exception exception) when (exception is not OperationCanceledException) - { - Trace.TraceWarning("Failed to dispose MCP {0}: {1}", resourceName, exception); - } - } - - /// - /// Computes a deterministic, order-independent hash of the header set. - /// Header names are lower-cased for case-insensitive matching (RFC 7230 §3.2). - /// Header values remain case-sensitive (RFC 7235 — credentials are case-sensitive). - /// -#pragma warning disable CA1308 // RFC 7230 §3.2 requires lower-cased header names for case-insensitive comparison; CA1308's uppercase preference does not apply here - internal static string ComputeHeadersHash(IDictionary? headers) - { - if (headers is null || headers.Count == 0) - { - return string.Empty; - } - - // Sort by lower-cased key for deterministic ordering, preserving value case. - SortedDictionary sorted = new(StringComparer.Ordinal); - foreach (KeyValuePair header in headers) - { - sorted[header.Key.ToLowerInvariant()] = header.Value; - } - - StringBuilder payload = new(); - foreach (KeyValuePair kvp in sorted) - { - payload.Append(kvp.Key).Append(':').Append(kvp.Value).Append('\n'); - } - - byte[] inputBytes = Encoding.UTF8.GetBytes(payload.ToString()); -#if NET5_0_OR_GREATER - byte[] hashBytes = SHA256.HashData(inputBytes); -#else - using SHA256 sha256 = SHA256.Create(); - byte[] hashBytes = sha256.ComputeHash(inputBytes); -#endif - - // Convert to hex string (compatible with net472/netstandard2.0) - StringBuilder hex = new(hashBytes.Length * 2); - foreach (byte b in hashBytes) - { - hex.Append(b.ToString("X2", System.Globalization.CultureInfo.InvariantCulture)); - } - - return hex.ToString(); - } -#pragma warning restore CA1308 + public ValueTask DisposeAsync() => this._clientCache.DisposeAsync(); private static void ThrowIfListToolsArgumentsSpecified(IDictionary? arguments) { diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/McpClientCache.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/McpClientCache.cs new file mode 100644 index 00000000000..42545f2fa1a --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/McpClientCache.cs @@ -0,0 +1,569 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Runtime.ExceptionServices; +using System.Security.Cryptography; +using System.Text; +using System.Threading; +using System.Threading.Tasks; +using ModelContextProtocol.Client; + +namespace Microsoft.Agents.AI.Workflows.Declarative.Mcp; + +/// +/// Manages bounded workflow-scoped MCP sessions, concurrent creation, and active client leases. +/// +internal sealed class McpClientCache : IAsyncDisposable +{ + private static readonly OperationCanceledException s_clientCreationCancelledException = + new("The MCP client creator was cancelled."); + + private readonly McpClientFactory _clientFactory; + private readonly McpClientLifetime _lifetime; + private readonly Func _clientConnectionDisposer; + private readonly Dictionary<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash), CachedClient> _clients = []; + private readonly Dictionary<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash), TaskCompletionSource> _clientCreations = []; + private readonly HashSet> _clientCreationLifetimes = []; + private readonly LinkedList<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash)> _clientLru = []; + private readonly HashSet _retiredClients = []; + private readonly SemaphoreSlim _clientCreationSemaphore; + private readonly int _clientCacheMaxSize; + + public McpClientCache( + McpClientFactory clientFactory, + McpClientLifetime lifetime, + int clientCacheMaxSize, + Func? clientConnectionDisposer) + { + if (clientCacheMaxSize <= 0) + { + throw new ArgumentOutOfRangeException(nameof(clientCacheMaxSize), "The MCP client cache size must be positive."); + } + + this._clientFactory = clientFactory; + this._lifetime = lifetime; + this._clientConnectionDisposer = clientConnectionDisposer ?? (connection => connection.DisposeAsync()); + this._clientCacheMaxSize = clientCacheMaxSize; + this._clientCreationSemaphore = new(clientCacheMaxSize, clientCacheMaxSize); + } + + public async Task InvokeAsync( + string serverUrl, + string? serverLabel, + IDictionary? headers, + string? connectionName, + string workflowSessionId, + Func> invocation, + CancellationToken cancellationToken) + { + CachedClient cachedClient = await this.AcquireClientAsync( + serverUrl, serverLabel, headers, connectionName, workflowSessionId, cancellationToken).ConfigureAwait(false); + try + { + return await invocation(cachedClient.Connection.Client).ConfigureAwait(false); + } + finally + { + await this.ReleaseClientAsync(cachedClient).ConfigureAwait(false); + } + } + + public async ValueTask DisposeAsync() + { + this._lifetime.ThrowIfProviderInvocationActive(); + + Task providerInvocations; + List> clientCreationLifetimes; + await this._lifetime.SyncLock.WaitAsync().ConfigureAwait(false); + try + { + providerInvocations = this._lifetime.BeginDisposalAsync(); + clientCreationLifetimes = this._clientCreationLifetimes.Select(source => source.Task).ToList(); + } + finally + { + this._lifetime.SyncLock.Release(); + } + + await providerInvocations.ConfigureAwait(false); + + Exception? clientCreationException = null; + try + { + await Task.WhenAll(clientCreationLifetimes).ConfigureAwait(false); + } + catch (Exception exception) when (!McpClientLifetime.IsFatalException(exception)) + { + clientCreationException = exception; + } + + List cachedClients; + List clientsToDispose = []; + await this._lifetime.SyncLock.WaitAsync().ConfigureAwait(false); + try + { + cachedClients = [.. this._clients.Values, .. this._retiredClients]; + foreach (CachedClient client in cachedClients) + { + client.Evicted = true; + this._retiredClients.Add(client); + if (client.ActiveInvocations == 0 && !client.DisposalClaimed) + { + client.DisposalClaimed = true; + clientsToDispose.Add(client); + } + } + + this._clients.Clear(); + this._clientLru.Clear(); + } + finally + { + this._lifetime.SyncLock.Release(); + } + + Exception? clientCleanupException = null; + try + { + await McpClientLifetime.DrainCleanupAsync( + clientsToDispose.Select(client => this.DisposeCachedClientAsync(client)), + cachedClients.Select(client => client.Disposed.Task)).ConfigureAwait(false); + } + catch (Exception exception) when (!McpClientLifetime.IsFatalException(exception)) + { + clientCleanupException = exception; + } + finally + { + this._lifetime.Dispose(); + this._clientCreationSemaphore.Dispose(); + } + + if (clientCreationException is not null) + { + ExceptionDispatchInfo.Capture(clientCreationException).Throw(); + } + + if (clientCleanupException is not null) + { + ExceptionDispatchInfo.Capture(clientCleanupException).Throw(); + } + } + + private async Task AcquireClientAsync( + string serverUrl, + string? serverLabel, + IDictionary? headers, + string? connectionName, + string workflowSessionId, + CancellationToken cancellationToken) + { + string trimmedUrl = serverUrl.Trim(); + var clientCacheKey = BuildCacheKey(workflowSessionId, trimmedUrl, serverLabel, connectionName, headers); + CachedClient? clientToDispose = null; + TaskCompletionSource? clientCreation; + TaskCompletionSource? clientCreationLifetime = null; + bool ownsClientCreation = false; + + await this._lifetime.SyncLock.WaitAsync(cancellationToken).ConfigureAwait(false); + try + { + this._lifetime.ThrowIfDisposing(); + if (this._clients.TryGetValue(clientCacheKey, out CachedClient? existingClient)) + { + existingClient.ActiveInvocations++; + this._clientLru.Remove(existingClient.LruNode); + this._clientLru.AddLast(existingClient.LruNode); + return existingClient; + } + + if (!this._clientCreations.TryGetValue(clientCacheKey, out clientCreation)) + { + if (this._clients.Count + this._retiredClients.Count + this._clientCreations.Count >= this._clientCacheMaxSize) + { + LinkedListNode<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash)>? idleNode = + this._clientLru.First; + while (idleNode is not null && this._clients[idleNode.Value].ActiveInvocations != 0) + { + idleNode = idleNode.Next; + } + + if (idleNode is null) + { + throw new InvalidOperationException("MCP client capacity is exhausted; retry when an active invocation completes."); + } + + clientToDispose = this._clients[idleNode.Value]; + this._clients.Remove(idleNode.Value); + this._clientLru.Remove(idleNode); + clientToDispose.Evicted = true; + clientToDispose.DisposalClaimed = true; + this._retiredClients.Add(clientToDispose); + } + + clientCreation = new(TaskCreationOptions.RunContinuationsAsynchronously); + this._clientCreations[clientCacheKey] = clientCreation; + clientCreationLifetime = new(TaskCreationOptions.RunContinuationsAsynchronously); + this._clientCreationLifetimes.Add(clientCreationLifetime); + ownsClientCreation = true; + } + } + finally + { + this._lifetime.SyncLock.Release(); + } + + if (!ownsClientCreation) + { + try + { + await WaitForClientCreationAsync(clientCreation.Task, cancellationToken).ConfigureAwait(false); + } + catch (OperationCanceledException exception) + when (ReferenceEquals(exception, s_clientCreationCancelledException)) + { + return await this.AcquireClientAsync( + serverUrl, serverLabel, headers, connectionName, workflowSessionId, cancellationToken).ConfigureAwait(false); + } + + return await this.AcquireClientAsync( + serverUrl, serverLabel, headers, connectionName, workflowSessionId, cancellationToken).ConfigureAwait(false); + } + + TaskCompletionSource ownedClientCreationLifetime = clientCreationLifetime ?? + throw new InvalidOperationException("Missing MCP client creation lifetime."); + McpClientConnection? connection = null; + bool creationSemaphoreEntered = false; + Exception? clientCreationCleanupException = null; + try + { + if (clientToDispose is not null) + { + try + { + await this.DisposeCachedClientAsync(clientToDispose).ConfigureAwait(false); + } + catch (Exception cleanupException) + { + clientCreationCleanupException = cleanupException; + throw; + } + } + + await this._clientCreationSemaphore.WaitAsync(cancellationToken).ConfigureAwait(false); + creationSemaphoreEntered = true; + await this._lifetime.SyncLock.WaitAsync(cancellationToken).ConfigureAwait(false); + try + { + this._lifetime.ThrowIfDisposing(); + } + finally + { + this._lifetime.SyncLock.Release(); + } + + connection = await this._clientFactory.CreateClientAsync( + trimmedUrl, serverLabel, headers, cancellationToken, + cleanupException => clientCreationCleanupException = cleanupException).ConfigureAwait(false); + } + catch (Exception exception) + { + Exception sharedCreationException = + exception is OperationCanceledException && cancellationToken.IsCancellationRequested + ? s_clientCreationCancelledException + : exception; + try + { + await this.CompleteClientCreationFailureAsync( + clientCacheKey, clientCreation, sharedCreationException).ConfigureAwait(false); + } + finally + { + if (creationSemaphoreEntered) + { + this._clientCreationSemaphore.Release(); + creationSemaphoreEntered = false; + } + + await this.CompleteClientCreationLifetimeAsync( + ownedClientCreationLifetime, clientCreationCleanupException).ConfigureAwait(false); + } + + throw; + } + + ObjectDisposedException? disposedException = null; + CachedClient? result = null; + await this._lifetime.SyncLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); + try + { + this._clientCreations.Remove(clientCacheKey); + if (this._lifetime.IsDisposing) + { + disposedException = new ObjectDisposedException(nameof(DefaultMcpToolHandler)); + } + else + { + LinkedListNode<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash)> node = + this._clientLru.AddLast(clientCacheKey); + CachedClient newClient = new(connection, node) { ActiveInvocations = 1 }; + connection = null; + this._clients[clientCacheKey] = newClient; + + result = newClient; + clientCreation.TrySetResult(newClient); + } + } + finally + { + this._lifetime.SyncLock.Release(); + } + + try + { + if (connection is not null) + { + await this._clientConnectionDisposer(connection).ConfigureAwait(false); + } + + if (disposedException is not null) + { + clientCreation.TrySetException(disposedException); + throw disposedException; + } + } + catch (Exception exception) + { + if (disposedException is not null && !ReferenceEquals(exception, disposedException)) + { + clientCreation.TrySetException(exception); + clientCreationCleanupException = exception; + } + + if (result is not null) + { + await this.ReleaseClientAsync(result).ConfigureAwait(false); + } + + throw; + } + finally + { + if (creationSemaphoreEntered) + { + this._clientCreationSemaphore.Release(); + } + + await this.CompleteClientCreationLifetimeAsync( + ownedClientCreationLifetime, clientCreationCleanupException).ConfigureAwait(false); + } + + return result ?? throw new InvalidOperationException("Failed to acquire MCP client."); + } + + private static async Task WaitForClientCreationAsync(Task clientCreation, CancellationToken cancellationToken) + { + if (!cancellationToken.CanBeCanceled || clientCreation.IsCompleted) + { + await clientCreation.ConfigureAwait(false); + return; + } + + TaskCompletionSource cancellation = new(TaskCreationOptions.RunContinuationsAsynchronously); + using CancellationTokenRegistration registration = cancellationToken.Register( + static state => ((TaskCompletionSource)state!).TrySetResult(true), + cancellation); + + if (await Task.WhenAny(clientCreation, cancellation.Task).ConfigureAwait(false) == cancellation.Task) + { + cancellationToken.ThrowIfCancellationRequested(); + } + + await clientCreation.ConfigureAwait(false); + } + + private async Task CompleteClientCreationFailureAsync( + (string WorkflowSession, string Url, string Label, string Connection, string HeadersHash) clientCacheKey, + TaskCompletionSource clientCreation, + Exception exception) + { + await this._lifetime.SyncLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); + try + { + if (this._clientCreations.TryGetValue(clientCacheKey, out TaskCompletionSource? existingCreation) && + ReferenceEquals(existingCreation, clientCreation)) + { + this._clientCreations.Remove(clientCacheKey); + } + + clientCreation.TrySetException(exception); + _ = clientCreation.Task.Exception; + } + finally + { + this._lifetime.SyncLock.Release(); + } + } + + private async Task CompleteClientCreationLifetimeAsync( + TaskCompletionSource clientCreationLifetime, + Exception? cleanupException = null) + { + await this._lifetime.SyncLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); + try + { + this._clientCreationLifetimes.Remove(clientCreationLifetime); + } + finally + { + this._lifetime.SyncLock.Release(); + } + + if (cleanupException is not null) + { + clientCreationLifetime.TrySetException(cleanupException); + } + else + { + clientCreationLifetime.TrySetResult(true); + } + } + + private async ValueTask ReleaseClientAsync(CachedClient client) + { + bool dispose = false; + await this._lifetime.SyncLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); + try + { + client.ActiveInvocations--; + if (client.ActiveInvocations == 0 && client.Evicted && !client.DisposalClaimed) + { + client.DisposalClaimed = true; + dispose = true; + } + } + finally + { + this._lifetime.SyncLock.Release(); + } + + if (dispose) + { + await this.DisposeCachedClientAsync(client).ConfigureAwait(false); + } + } + + private async Task DisposeCachedClientAsync(CachedClient client) + { + Exception? cleanupException = null; + try + { + await this._clientConnectionDisposer(client.Connection).ConfigureAwait(false); + } + catch (Exception exception) + { + cleanupException = exception; + throw; + } + finally + { + await this._lifetime.SyncLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); + try + { + this._retiredClients.Remove(client); + } + finally + { + this._lifetime.SyncLock.Release(); + } + + if (cleanupException is not null) + { + client.Disposed.TrySetException(cleanupException); + } + else + { + client.Disposed.TrySetResult(true); + } + } + } + + /// + /// Builds the per-client cache key as a 5-tuple of + /// (workflowSessionId, trimmed serverUrl, serverLabel, connectionName, headers hash). + /// All five components participate so that separate workflow sessions and callers using + /// different labels/connections/headers receive distinct instances + /// even when targeting the same URL. + /// + internal static (string WorkflowSession, string Url, string Label, string Connection, string HeadersHash) BuildCacheKey( + string workflowSessionId, + string trimmedUrl, + string? serverLabel, + string? connectionName, + IDictionary? headers) => + (workflowSessionId, trimmedUrl, serverLabel ?? string.Empty, connectionName ?? string.Empty, ComputeHeadersHash(headers)); + + /// + /// Computes a deterministic, order-independent hash of the header set. + /// Header names are lower-cased for case-insensitive matching (RFC 7230 §3.2). + /// Header values remain case-sensitive (RFC 7235 — credentials are case-sensitive). + /// +#pragma warning disable CA1308 // RFC 7230 §3.2 requires lower-cased header names for case-insensitive comparison; CA1308's uppercase preference does not apply here + internal static string ComputeHeadersHash(IDictionary? headers) + { + if (headers is null || headers.Count == 0) + { + return string.Empty; + } + + // Sort by lower-cased key for deterministic ordering, preserving value case. + SortedDictionary sorted = new(StringComparer.Ordinal); + foreach (KeyValuePair header in headers) + { + sorted[header.Key.ToLowerInvariant()] = header.Value; + } + + StringBuilder payload = new(); + foreach (KeyValuePair kvp in sorted) + { + payload.Append(kvp.Key).Append(':').Append(kvp.Value).Append('\n'); + } + + byte[] inputBytes = Encoding.UTF8.GetBytes(payload.ToString()); +#if NET5_0_OR_GREATER + byte[] hashBytes = SHA256.HashData(inputBytes); +#else + using SHA256 sha256 = SHA256.Create(); + byte[] hashBytes = sha256.ComputeHash(inputBytes); +#endif + + // Convert to hex string (compatible with net472/netstandard2.0) + StringBuilder hex = new(hashBytes.Length * 2); + foreach (byte b in hashBytes) + { + hex.Append(b.ToString("X2", System.Globalization.CultureInfo.InvariantCulture)); + } + + return hex.ToString(); + } +#pragma warning restore CA1308 + + private sealed class CachedClient( + McpClientConnection connection, + LinkedListNode<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash)> lruNode) + { + public McpClientConnection Connection { get; } = connection; + + public LinkedListNode<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash)> LruNode { get; } = lruNode; + + public TaskCompletionSource Disposed { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously); + + public int ActiveInvocations { get; set; } + + public bool Evicted { get; set; } + + public bool DisposalClaimed { get; set; } + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/McpClientConnection.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/McpClientConnection.cs new file mode 100644 index 00000000000..b9bd10c99b7 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/McpClientConnection.cs @@ -0,0 +1,55 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Diagnostics; +using System.Threading.Tasks; +using ModelContextProtocol.Client; + +namespace Microsoft.Agents.AI.Workflows.Declarative.Mcp; + +/// +/// Owns an MCP session, its transport, and any lease on a pooled HTTP client. +/// +internal sealed class McpClientConnection( + McpClient client, + IAsyncDisposable transport, + IAsyncDisposable? ownedHttpClientLease = null) : IAsyncDisposable +{ + public McpClient Client { get; } = client; + + public async ValueTask DisposeAsync() + { + try + { + await DisposeResourceAsync(this.Client, "session").ConfigureAwait(false); + } + finally + { + try + { + // McpClient owns the connected session, not the reusable transport factory. + await DisposeResourceAsync(transport, "transport").ConfigureAwait(false); + } + finally + { + if (ownedHttpClientLease is not null) + { + await ownedHttpClientLease.DisposeAsync().ConfigureAwait(false); + } + } + } + } + + internal static async ValueTask DisposeResourceAsync(T resource, string resourceName) + where T : IAsyncDisposable + { + try + { + await resource.DisposeAsync().ConfigureAwait(false); + } + catch (Exception exception) when (exception is not OperationCanceledException) + { + Trace.TraceWarning("Failed to dispose MCP {0}: {1}", resourceName, exception); + } + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/McpClientFactory.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/McpClientFactory.cs new file mode 100644 index 00000000000..4420fa2c690 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/McpClientFactory.cs @@ -0,0 +1,218 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Diagnostics; +using System.Net.Http; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Shared.Diagnostics; +using ModelContextProtocol.Client; + +namespace Microsoft.Agents.AI.Workflows.Declarative.Mcp; + +/// +/// Creates MCP connections and manages ownership of their HTTP clients. +/// +internal sealed class McpClientFactory( + Func>? httpClientProvider, + Func httpMessageHandlerFactory, + McpClientLifetime lifetime) +{ + private readonly Func _httpMessageHandlerFactory = Throw.IfNull(httpMessageHandlerFactory); + private readonly Dictionary _ownedHttpClients = []; + + public bool HasHttpClientProvider => httpClientProvider is not null; + + public Task InvokeProviderAsync( + string serverUrl, + string? serverLabel, + IDictionary? headers, + Func> invocation, + CancellationToken cancellationToken) + => lifetime.InvokeProviderAsync(async () => + { + McpClientConnection connection = await this.CreateClientAsync( + serverUrl.Trim(), serverLabel, headers, cancellationToken).ConfigureAwait(false); + await using (connection.ConfigureAwait(false)) + { + return await invocation(connection.Client).ConfigureAwait(false); + } + }, cancellationToken); + + public async Task CreateClientAsync( + string serverUrl, + string? serverLabel, + IDictionary? headers, + CancellationToken cancellationToken, + Action? reportCleanupFailure = null) + { + HttpClient? httpClient = null; + bool ownsHttpClient = false; + OwnedHttpClientLease? ownedHttpClientLease = null; + + if (httpClientProvider is not null) + { + httpClient = await httpClientProvider(serverUrl, cancellationToken).ConfigureAwait(false); + } + + if (httpClient is null && httpClientProvider is not null) + { + httpClient = this.CreatePinnedHttpClient(serverUrl); + ownsHttpClient = true; + } + else if (httpClient is null) + { + ownedHttpClientLease = await this.AcquireOwnedHttpClientAsync(serverUrl, cancellationToken).ConfigureAwait(false); + httpClient = ownedHttpClientLease.Client; + } + + HttpClientTransportOptions transportOptions = new() + { + Endpoint = new Uri(serverUrl), + Name = serverLabel ?? "McpClient", + AdditionalHeaders = headers, + // Use Streamable HTTP rather than AutoDetect so the client does not negotiate down to the + // legacy HTTP+SSE transport, which trusts a server-advertised message endpoint. That + // server-controlled endpoint is the primary vector for redirecting the Authorization token + // to a different origin; Streamable HTTP keeps every request on the configured origin. + TransportMode = HttpTransportMode.StreamableHttp + }; + + HttpClientTransport? transport = null; + try + { + HttpClient resolvedHttpClient = httpClient ?? throw new InvalidOperationException("Failed to resolve MCP HTTP client."); + transport = new(transportOptions, resolvedHttpClient, ownsHttpClient: ownsHttpClient); + McpClient client = await McpClient.CreateAsync(transport, cancellationToken: cancellationToken).ConfigureAwait(false); + McpClientConnection connection = new(client, transport, ownedHttpClientLease); + ownedHttpClientLease = null; + return connection; + } + catch + { + try + { + try + { + if (transport is not null) + { + await McpClientConnection.DisposeResourceAsync(transport, "transport").ConfigureAwait(false); + } + } + finally + { + if (ownedHttpClientLease is not null) + { + await ownedHttpClientLease.DisposeAsync().ConfigureAwait(false); + } + } + } + catch (Exception cleanupException) + { + reportCleanupFailure?.Invoke(cleanupException); + throw; + } + + throw; + } + } + + private async Task AcquireOwnedHttpClientAsync( + string serverUrl, + CancellationToken cancellationToken) + { + await lifetime.SyncLock.WaitAsync(cancellationToken).ConfigureAwait(false); + try + { + if (!this._ownedHttpClients.TryGetValue(serverUrl, out OwnedHttpClient? entry)) + { + entry = new(this.CreatePinnedHttpClient(serverUrl)); + this._ownedHttpClients[serverUrl] = entry; + } + + entry.ReferenceCount++; + return new(this, serverUrl, entry.Client); + } + finally + { + lifetime.SyncLock.Release(); + } + } + + private async ValueTask ReleaseOwnedHttpClientAsync(string serverUrl) + { + HttpClient? clientToDispose = null; + await lifetime.SyncLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); + try + { + OwnedHttpClient entry = this._ownedHttpClients[serverUrl]; + if (--entry.ReferenceCount == 0) + { + this._ownedHttpClients.Remove(serverUrl); + clientToDispose = entry.Client; + } + } + finally + { + lifetime.SyncLock.Release(); + } + + try + { + clientToDispose?.Dispose(); + } + catch (Exception exception) when (exception is not OperationCanceledException) + { + Trace.TraceWarning("Failed to dispose MCP HTTP client: {0}", exception); + } + } + + private HttpClient CreatePinnedHttpClient(string serverUrl) + { + // Pin credential headers to the configured server origin as defense-in-depth. Forcing + // StreamableHttp already removes the primary vector (a server-advertised cross-origin + // SSE message endpoint), and AllowAutoRedirect=false blocks auto-redirects. This handler is the + // backstop: it guarantees the Authorization token and other credentials never leave the pinned + // origin even if a future change re-enables AutoDetect or redirects, or the SDK constructs a + // request to a new URI (AdditionalHeaders are re-stamped by the transport, so HttpClient's own + // redirect header-stripping does not cover them). + OriginPinningHandler pinningHandler = new(new Uri(serverUrl)) { InnerHandler = this._httpMessageHandlerFactory() }; + return new HttpClient(pinningHandler); + } + + internal static HttpMessageHandler CreateHttpMessageHandler() => + new HttpClientHandler + { + // Keep cookies out of pooled transport state and prevent credential-bearing redirects. + UseCookies = false, + AllowAutoRedirect = false, + CheckCertificateRevocationList = true + }; + + private sealed class OwnedHttpClient(HttpClient client) + { + public HttpClient Client { get; } = client; + + public int ReferenceCount { get; set; } + } + + private sealed class OwnedHttpClientLease( + McpClientFactory owner, + string serverUrl, + HttpClient client) : IAsyncDisposable + { + private bool _disposed; + + public HttpClient Client { get; } = client; + + public async ValueTask DisposeAsync() + { + if (!this._disposed) + { + this._disposed = true; + await owner.ReleaseOwnedHttpClientAsync(serverUrl).ConfigureAwait(false); + } + } + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/McpClientLifetime.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/McpClientLifetime.cs new file mode 100644 index 00000000000..c7168a27cd4 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/McpClientLifetime.cs @@ -0,0 +1,132 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Runtime.ExceptionServices; +using System.Threading; +using System.Threading.Tasks; + +namespace Microsoft.Agents.AI.Workflows.Declarative.Mcp; + +/// +/// Coordinates invocation admission and shutdown across cached and provider-backed clients. +/// +internal sealed class McpClientLifetime : IDisposable +{ + private readonly AsyncLocal _providerInvocationContext = new(); + private TaskCompletionSource? _providerInvocationsDrained; + private int _activeProviderInvocations; + + public SemaphoreSlim SyncLock { get; } = new(1, 1); + + public bool IsDisposing { get; private set; } + + public async Task InvokeProviderAsync(Func> invocation, CancellationToken cancellationToken) + { + TaskCompletionSource invocationCompleted = new(TaskCreationOptions.RunContinuationsAsynchronously); + await this.SyncLock.WaitAsync(cancellationToken).ConfigureAwait(false); + try + { + this.ThrowIfDisposing(); + if (this._activeProviderInvocations++ == 0) + { + this._providerInvocationsDrained = new(TaskCreationOptions.RunContinuationsAsynchronously); + } + } + finally + { + this.SyncLock.Release(); + } + + ProviderInvocationContext? previousInvocation = this._providerInvocationContext.Value; + this._providerInvocationContext.Value = new(invocationCompleted.Task, previousInvocation); + try + { + return await invocation().ConfigureAwait(false); + } + finally + { + await this.SyncLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); + try + { + if (--this._activeProviderInvocations == 0) + { + this._providerInvocationsDrained!.SetResult(true); + } + } + finally + { + invocationCompleted.SetResult(true); + this._providerInvocationContext.Value = previousInvocation; + this.SyncLock.Release(); + } + } + } + + public void ThrowIfProviderInvocationActive() + { + for (ProviderInvocationContext? context = this._providerInvocationContext.Value; context is not null; context = context.Parent) + { + if (!context.Completion.IsCompleted) + { + throw new InvalidOperationException("Cannot dispose the MCP handler from an active provider-backed invocation."); + } + } + } + + // Called with SyncLock held so shutdown and invocation admission are atomic. + public Task BeginDisposalAsync() + { + this.ThrowIfDisposing(); + this.IsDisposing = true; + return this._providerInvocationsDrained?.Task ?? Task.CompletedTask; + } + + public void Dispose() => this.SyncLock.Dispose(); + + [System.Diagnostics.CodeAnalysis.SuppressMessage("Maintainability", "CA1513:Use ObjectDisposedException throw helper", + Justification = "The helper is not available on .NET Framework or .NET Standard 2.0.")] + public void ThrowIfDisposing() + { + if (this.IsDisposing) + { + throw new ObjectDisposedException(nameof(DefaultMcpToolHandler)); + } + } + + internal static async Task DrainCleanupAsync(IEnumerable cleanupTasks, IEnumerable completionTasks) + { + Exception? cleanupException = null; + try + { + await Task.WhenAll(cleanupTasks).ConfigureAwait(false); + } + catch (Exception exception) when (!IsFatalException(exception)) + { + cleanupException = exception; + } + + await Task.WhenAll(completionTasks).ConfigureAwait(false); + + if (cleanupException is not null) + { + ExceptionDispatchInfo.Capture(cleanupException).Throw(); + } + } + + internal static bool IsFatalException(Exception exception) => + exception is OutOfMemoryException and not InsufficientMemoryException + or StackOverflowException + or AccessViolationException + or AppDomainUnloadedException + or BadImageFormatException + or CannotUnloadAppDomainException + or InvalidProgramException + or ThreadAbortException; + + private sealed class ProviderInvocationContext(Task completion, ProviderInvocationContext? parent) + { + public Task Completion { get; } = completion; + public ProviderInvocationContext? Parent { get; } = parent; + } +} diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs index 3ff71500d6c..b33593b3caf 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs @@ -293,7 +293,7 @@ async Task SuccessfulCleanupAsync() } // Act - await Assert.ThrowsAnyAsync(() => DefaultMcpToolHandler.DrainCleanupAsync( + await Assert.ThrowsAnyAsync(() => McpClientLifetime.DrainCleanupAsync( [CancelledCleanupAsync(), SuccessfulCleanupAsync()], [firstDisposed.Task, secondDisposed.Task])); @@ -1204,7 +1204,7 @@ public async Task ClientConnection_SessionCleanupFailure_StillCleansUpTransportA Mock transport = new(); transport.Setup(t => t.DisposeAsync()).Returns(new ValueTask(Task.FromException( new InvalidOperationException("transport cleanup failed")))); - DefaultMcpToolHandler.ClientConnection connection = new(client.Object, transport.Object); + McpClientConnection connection = new(client.Object, transport.Object); using CleanupTraceListener listener = new(); Trace.Listeners.Add(listener); try diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerTests.cs index b6eb32b667d..a363501c87e 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerTests.cs @@ -379,7 +379,7 @@ await handler.InvokeToolAsync( public void ComputeHeadersHash_WithNullHeaders_ReturnsEmptyString() { // Act - string result = DefaultMcpToolHandler.ComputeHeadersHash(null); + string result = McpClientCache.ComputeHeadersHash(null); // Assert Assert.Equal(string.Empty, result); @@ -389,7 +389,7 @@ public void ComputeHeadersHash_WithNullHeaders_ReturnsEmptyString() public void ComputeHeadersHash_WithEmptyHeaders_ReturnsEmptyString() { // Act - string result = DefaultMcpToolHandler.ComputeHeadersHash(new Dictionary()); + string result = McpClientCache.ComputeHeadersHash(new Dictionary()); // Assert Assert.Equal(string.Empty, result); @@ -411,8 +411,8 @@ public void ComputeHeadersHash_SameHeadersDifferentOrder_ReturnsSameHash() }; // Act - string hash1 = DefaultMcpToolHandler.ComputeHeadersHash(headers1); - string hash2 = DefaultMcpToolHandler.ComputeHeadersHash(headers2); + string hash1 = McpClientCache.ComputeHeadersHash(headers1); + string hash2 = McpClientCache.ComputeHeadersHash(headers2); // Assert Assert.Equal(hash2, hash1); @@ -426,8 +426,8 @@ public void ComputeHeadersHash_SameKeysDifferentCaseKeys_ReturnsSameHash() Dictionary headers2 = new() { ["authorization"] = "Bearer token" }; // Act - string hash1 = DefaultMcpToolHandler.ComputeHeadersHash(headers1); - string hash2 = DefaultMcpToolHandler.ComputeHeadersHash(headers2); + string hash1 = McpClientCache.ComputeHeadersHash(headers1); + string hash2 = McpClientCache.ComputeHeadersHash(headers2); // Assert Assert.Equal(hash2, hash1); @@ -441,8 +441,8 @@ public void ComputeHeadersHash_SameKeysDifferentCaseValues_ReturnsDifferentHash( Dictionary headers2 = new() { ["Authorization"] = "Bearer abc" }; // Act - string hash1 = DefaultMcpToolHandler.ComputeHeadersHash(headers1); - string hash2 = DefaultMcpToolHandler.ComputeHeadersHash(headers2); + string hash1 = McpClientCache.ComputeHeadersHash(headers1); + string hash2 = McpClientCache.ComputeHeadersHash(headers2); // Assert Assert.NotEqual(hash2, hash1); @@ -456,8 +456,8 @@ public void ComputeHeadersHash_DifferentHeaders_ReturnsDifferentHash() Dictionary headers2 = new() { ["Authorization"] = "Bearer token2" }; // Act - string hash1 = DefaultMcpToolHandler.ComputeHeadersHash(headers1); - string hash2 = DefaultMcpToolHandler.ComputeHeadersHash(headers2); + string hash1 = McpClientCache.ComputeHeadersHash(headers1); + string hash2 = McpClientCache.ComputeHeadersHash(headers2); // Assert Assert.NotEqual(hash2, hash1); @@ -481,8 +481,8 @@ public void BuildCacheKey_SameInputs_ReturnsEqualKeys() Dictionary headers = new() { ["Authorization"] = "Bearer token" }; // Act - var key1 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", "label", "conn", headers); - var key2 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", "label", "conn", headers); + var key1 = McpClientCache.BuildCacheKey("workflow", "http://localhost/mcp", "label", "conn", headers); + var key2 = McpClientCache.BuildCacheKey("workflow", "http://localhost/mcp", "label", "conn", headers); // Assert Assert.Equal(key2, key1); @@ -492,8 +492,8 @@ public void BuildCacheKey_SameInputs_ReturnsEqualKeys() public void BuildCacheKey_DifferentWorkflowSession_ReturnsDifferentKeys() { // Act - var key1 = DefaultMcpToolHandler.BuildCacheKey("workflow-a", "http://localhost/mcp", "label", "conn", null); - var key2 = DefaultMcpToolHandler.BuildCacheKey("workflow-b", "http://localhost/mcp", "label", "conn", null); + var key1 = McpClientCache.BuildCacheKey("workflow-a", "http://localhost/mcp", "label", "conn", null); + var key2 = McpClientCache.BuildCacheKey("workflow-b", "http://localhost/mcp", "label", "conn", null); // Assert Assert.NotEqual(key2, key1); @@ -505,8 +505,8 @@ public void BuildCacheKey_DifferentWorkflowSession_ReturnsDifferentKeys() public void BuildCacheKey_DifferentConnectionName_ReturnsDifferentKeys() { // Act - var key1 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", "label", "connection-a", null); - var key2 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", "label", "connection-b", null); + var key1 = McpClientCache.BuildCacheKey("workflow", "http://localhost/mcp", "label", "connection-a", null); + var key2 = McpClientCache.BuildCacheKey("workflow", "http://localhost/mcp", "label", "connection-b", null); // Assert Assert.NotEqual(key2, key1); @@ -518,8 +518,8 @@ public void BuildCacheKey_DifferentConnectionName_ReturnsDifferentKeys() public void BuildCacheKey_DifferentServerLabel_ReturnsDifferentKeys() { // Act - var key1 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", "label-a", null, null); - var key2 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", "label-b", null, null); + var key1 = McpClientCache.BuildCacheKey("workflow", "http://localhost/mcp", "label-a", null, null); + var key2 = McpClientCache.BuildCacheKey("workflow", "http://localhost/mcp", "label-b", null, null); // Assert Assert.NotEqual(key2, key1); @@ -532,8 +532,8 @@ public void BuildCacheKey_CaseSensitiveUrlPath_ReturnsDifferentKeys() { // Arrange — RFC 3986: URL path is case-sensitive // Act - var key1 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/Tools", null, null, null); - var key2 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/tools", null, null, null); + var key1 = McpClientCache.BuildCacheKey("workflow", "http://localhost/Tools", null, null, null); + var key2 = McpClientCache.BuildCacheKey("workflow", "http://localhost/tools", null, null, null); // Assert Assert.NotEqual(key2, key1); @@ -547,8 +547,8 @@ public void BuildCacheKey_HeaderValuesCaseSensitive_ReturnsDifferentKeys() Dictionary headers2 = new() { ["Authorization"] = "Bearer abc" }; // Act - var key1 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", null, null, headers1); - var key2 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", null, null, headers2); + var key1 = McpClientCache.BuildCacheKey("workflow", "http://localhost/mcp", null, null, headers1); + var key2 = McpClientCache.BuildCacheKey("workflow", "http://localhost/mcp", null, null, headers2); // Assert — header value case must propagate into the cache key Assert.NotEqual(key2, key1); @@ -559,7 +559,7 @@ public void BuildCacheKey_HeaderValuesCaseSensitive_ReturnsDifferentKeys() public void BuildCacheKey_NullLabelAndConnection_NormalizesToEmptyString() { // Act - var key = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", null, null, null); + var key = McpClientCache.BuildCacheKey("workflow", "http://localhost/mcp", null, null, null); // Assert — verifies null-safety contract callers rely on Assert.Equal(string.Empty, key.Label); From 2c895157078e024a040525829961b918d727862d Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 20 Oct 2026 15:19:33 -0400 Subject: [PATCH 3/3] docs: propose workflow-owned MCP sessions Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 94e932cb-8d65-4169-8c8b-3b5ceff9d09d --- ...0047-dotnet-workflow-owned-mcp-sessions.md | 505 ++++++++++++++++++ 1 file changed, 505 insertions(+) create mode 100644 docs/decisions/0047-dotnet-workflow-owned-mcp-sessions.md diff --git a/docs/decisions/0047-dotnet-workflow-owned-mcp-sessions.md b/docs/decisions/0047-dotnet-workflow-owned-mcp-sessions.md new file mode 100644 index 00000000000..7f16a1635e8 --- /dev/null +++ b/docs/decisions/0047-dotnet-workflow-owned-mcp-sessions.md @@ -0,0 +1,505 @@ +--- +status: proposed +contact: baywet +date: 2026-10-20 +deciders: [baywet, westey-m] +informed: [rhurey] +--- + +# .NET workflow-owned MCP sessions and factory-managed HTTP transport + +## Context and Problem Statement + +An initialized MCP `McpClient` represents a protocol connection that can have +associated server-side session state. It is not interchangeable with an arbitrary +idle client for the same server. A workflow must use its own connection when +continuing work, including after waiting for approval or another external event. + +The implementation in [#9545](https://github.com/microsoft/agent-framework/pull/9545) +combines several resource-management policies: workflow-specific MCP session +lookup, bounded least-recently-used eviction, active-call tracking, concurrent +initialization, shutdown draining, and reference-counted sharing of handler-owned +`HttpClient` instances. Extracting these into separate classes improves code +organization but does not remove the policy coupling. + +There are two different concerns: + +| Concern | Required isolation and ownership | +|---|---| +| MCP protocol session | Belongs to a logical workflow execution/session and its connection configuration | +| HTTP transport | Can reuse handlers and connection pools when transport configuration permits | + +The earlier handler already retained initialized MCP connections. Workflow +partitioning changes which connection is selected; it does not require a custom +HTTP-client pool. In particular, no active tool call does not imply that a +workflow has finished. Evicting an idle but live MCP session can change the +behavior of its next call. + +This ADR proposes a different ownership model for the .NET implementation. It is +**not implemented by adding this document**, does not declare the current PR's +behavior changed, and does not change the Python implementation. + +## Decision Drivers + +- Make workflow isolation a property of resource ownership rather than a + handler-wide cache key alone. +- Preserve MCP session continuity throughout a live workflow session. +- Keep active calls safe during completion, cancellation, and disposal. +- Respect ownership and configuration of user-supplied HTTP clients. +- Delegate reusable HTTP handler management to standard .NET infrastructure. +- Support console applications, worker services, and web applications. +- Avoid introducing an ASP.NET requirement or a DI container per MCP connection. +- Preserve existing transport protections and authentication boundaries. +- Make abandonment, checkpoint restore, resource limits, and migration behavior + explicit rather than hiding them behind eviction. + +## Considered Options + +1. Retain the handler-wide bounded MCP cache and reference-counted HTTP sharing. +2. Own MCP connections in a workflow session, with dedicated HTTP-client wrappers + and optional factory-managed handler reuse. +3. Own MCP connections in a workflow session, but create a dedicated HTTP handler + for every internally created client. +4. Share one application-wide `HttpClient` while retaining workflow-owned MCP + sessions. +5. Create a new MCP session for every tool invocation. + +## Decision Outcome + +Proposed chosen option: **own MCP connections in a workflow session, with +dedicated HTTP-client wrappers and optional factory-managed handler reuse**. + +The workflow execution/session owns its live connections. The tool handler +coordinates protocol operations but does not independently retain all workflows' +connections in a global LRU cache. The application owns its HTTP factory and any +HTTP clients it supplies. Exact public API names and runtime integration remain +subject to design review before implementation. + +### Workflow ownership boundary + +"Workflow instance" means the **logical execution/session**, not the reusable +`Workflow` graph, an executor, an individual tool call, a superstep, or the +duration of one `RunAsync`/resume operation. + +The current object model includes `Workflow`, disposable `Run` and +`StreamingRun` handles, and an internal `WorkflowSession` used by workflow-backed +agents. These are not automatically equivalent ownership boundaries. In +particular, agent execution can create successive run handles from checkpoints. +Implementation must connect both direct workflow execution and workflow-backed +agents to a session-level resource owner; attaching resources only to a temporary +run handle would not satisfy this ADR. + +Conceptual ownership: + +```text +Application + +-- IHttpClientFactory / pooled HTTP handler infrastructure + +-- caller-owned HttpClients, if supplied + +-- logical workflow session A + | +-- MCP connection for configuration X + | +-- McpClient and transport + | +-- owned HttpClient wrapper OR borrowed caller HttpClient + +-- logical workflow session B + +-- distinct MCP connection for configuration X + +-- distinct McpClient and transport + +-- owned HttpClient wrapper OR borrowed caller HttpClient +``` + +Sessions A and B must not share their MCP protocol session even when their +HTTP-client wrappers share a factory-managed handler and connection pool. +Different caller-provided HTTP clients may also legitimately share a transport +pool; this does not authorize sharing MCP sessions. +This proposal does not change existing graph/executor concurrency restrictions. + +Within one workflow owner, connection lookup still distinguishes server endpoint, +connection configuration, and authentication identity. Existing distinctions +such as connection name, label, and explicit headers must not be collapsed +without a separate compatibility decision. Changed credentials/configuration +must not silently select an incompatible existing MCP session. + +Session ownership is a runtime capability, not just an arbitrary string supplied +by an action. Trusted workflow identity remains useful for continuations, +diagnostics, and child scopes, but a matching string alone must not grant access +to another owner's live resources. Sibling subworkflows receive separate child +owners; the parent coordinates their shutdown before releasing its resources. +Checkpoint identifiers must retain their existing meaning. + +### HTTP ownership and factory integration + +| Source of HTTP client | Use and disposal contract | +|---|---| +| User-supplied client | Borrow as supplied; do not mutate shared client configuration or dispose it | +| Application-provided factory | Obtain a dedicated wrapper for the MCP connection; dispose the wrapper with that connection; let the factory manage pooled handlers | +| No supplied client or factory | Preserve a non-DI path that creates and owns a dedicated client with documented defaults | + +An MCP transport must not dispose a borrowed HTTP client. Owned wrappers must be +disposed exactly once, including when initialization fails. MCP and transport +cleanup must release owned HTTP resources even when earlier cleanup fails. +Supplying an already initialized `McpClient`, if exposed by a future API, likewise +requires an explicit ownership contract; accepting it must not implicitly allow +cross-workflow sharing. + +`IHttpClientFactory` belongs to `Microsoft.Extensions.Http` and does not require +ASP.NET. The application can use the Generic Host or a `ServiceCollection`. +The library must not create a service provider/factory for each MCP connection. +Retaining the non-DI path avoids making hosting infrastructure mandatory merely +to construct a handler. + +Use a bounded set of stable factory configuration names, not names derived from +workflow IDs, credentials, or arbitrary server URLs. The factory handles reuse of +the underlying pipelines. There is no additional handler-owned HTTP-client +dictionary, eviction policy, or reference-counted lease layer. + +Factory reuse is not exclusive checkout/return: multiple HTTP clients can use +the same underlying handler concurrently. Distinct transport policies, such as +proxies or client certificates, can still require separate configurations. + +### Transport policy and long-lived connections + +Internally created clients retain explicit defaults: cookies disabled, automatic +redirects disabled, certificate revocation checking, Streamable HTTP, and +origin-bound credential handling. Factory defaults alone are not a substitute +for this configuration. + +The current `OriginPinningHandler` stores one origin. It cannot be placed in a +single pooled pipeline for unrelated origins. To preserve this behavior while +sharing underlying handlers, use a connection-specific outer policy handler over +a factory-managed pipeline, following the documented `IHttpMessageHandlerFactory` +pattern, or an equivalently validated request-specific policy. The policy must +run after credential injection in outbound processing. Workflow identity, +credentials, and origin state must not be stored in shared mutable handlers. + +An outer policy handler and `HttpClient` wrapper may still be constructed per +connection. The objective is reuse of the expensive underlying HTTP handlers +and connection pools, not eliminating every `new HttpClient(...)` expression. +Caller-supplied clients remain under the caller's transport-policy contract; +the library must not silently replace their handler pipelines. + +Factory-created clients normally are short-lived. An MCP connection can instead +retain its wrapper for a long time. Factory handler expiry does not replace the +handler under that wrapper and must not terminate its active requests. On +supported runtimes, explicitly configure `SocketsHttpHandler.PooledConnectionLifetime` +to refresh network connections independently of the MCP session. This does not +renew or transfer server-side MCP state, and does not interrupt an active stream +merely because its connection age has elapsed. + +The package targets .NET 8, 9, and 10, .NET Standard 2.0, and .NET Framework 4.7.2. +Modern handler APIs are not uniformly available across these targets. Compatible +registration, DNS-refresh behavior, and supported factory-package versions must +be established per target; this proposal does not silently drop older targets. + +### Lifecycle and protection of active invocations + +An owner progresses through: + +```text +Open -> Closing -> Disposed +``` + +Waiting for external input is still `Open`, not a transition to `Closing`. + +1. **Acquire:** atomically confirm that the owner is open and register the work. + Resolve or initialize its connection. Concurrent first use of the same + configuration must not publish duplicate connections. +2. **Invoke:** keep the connection alive for the entire asynchronous operation. + Release the operation registration in `finally`, including on failure. +3. **Close:** reject new work and coordinate cancellation, where requested. + Await initialization and operations already admitted. Closing must also + account for a connection whose initialization completes during shutdown. +4. **Dispose:** after draining, release MCP clients, transports, and owned HTTP + wrappers. Do not dispose borrowed clients or the application's factory. + +Structured execution can supply the operation tracking if the runtime guarantees +that all tasks are joined before disposal. Otherwise retain a small per-owner +acquire/drain mechanism. Workflow ownership does not remove the need for +initialization synchronization, cleanup of unpublished resources, or safe +shutdown ordering. + +Cancellation is a request, not evidence that an operation has stopped. Cancelling +one waiter must not dispose a connection still used by another operation. +Session closure cancels shared initialization only according to an explicit +owner policy. Disposal requested from inside an admitted operation must not wait +on that same operation and deadlock; defer draining to the owner or reject that +reentrant request explicitly. + +Completion, abandonment, and process shutdown must have observable release paths. +A recoverable execution failure or cancellation of one run does not automatically +mean that the logical session has been abandoned; the owner must decide whether +to retain it for retry or close it. +A host shutdown deadline may stop waiting, but must report incomplete cleanup; +it must not pretend that a still-running operation drained successfully. +Cleanup errors must follow the repository's error-reporting conventions rather +than disappear. + +### Waiting, checkpoints, and capacity + +Returning control to the caller, yielding output, or writing a checkpoint does +not by itself release a live session. Approval continuations must return to the +original owner, even when they arrive through a new execution call. + +Live `McpClient`, transport, and HTTP objects are not checkpoint data. Restoring +in another process, or after explicitly releasing the original owner, requires +new connections. Retaining a workflow ID or MCP session identifier is not proof +that a server session can be reattached. Any durable reattachment requires +separate SDK/server support and an explicit contract; this ADR does not promise it. +Forking a checkpoint into an independent execution requires a distinct owner. + +There is no LRU eviction of still-open workflow sessions, including sessions +waiting for external input. Resource limits belong to explicit admission and +session-retention policy. Hosts may configure quotas, backpressure, or documented +expiration/abandonment, with appropriate tenant/owner boundaries. Exhaustion must +be observable; it must not silently replace a live MCP session with a new one. +This replaces the current fixed handler-wide limit rather than asserting that +resource use can be unlimited. + +### Authentication-provider compatibility + +The existing `httpClientProvider` callback runs for every invocation and creates +a fresh MCP session for every call, even when it returns the same HTTP client or +`null`. The callback may select different credentials on successive invocations. + +Do not reinterpret that callback as a once-per-session factory or begin retaining +its MCP connections silently. Preserve its existing semantics through a legacy +adapter, and offer an explicit opt-in connection-scoped contract for the new +ownership model. A session-aware application must define credential identity and +refresh behavior, for example through an application-owned authentication +handler. Obtaining a new token must not depend implicitly on recreating the MCP +session before every tool call. + +This compatibility mode intentionally retains per-call MCP sessions; it does +not gain state continuity merely because HTTP factory pooling is introduced. + +### Examples + +#### Non-web application registration + +The following uses existing .NET APIs to demonstrate that no ASP.NET application +or running host is required. It shows registration and ownership only, not the +complete origin-protection and handler configuration described above. + +```csharp +using System.Net.Http; +using Microsoft.Extensions.DependencyInjection; + +ServiceCollection services = new(); +services.AddHttpClient("WorkflowMcp"); + +using ServiceProvider applicationServices = services.BuildServiceProvider(); +IHttpClientFactory httpClients = + applicationServices.GetRequiredService(); + +using HttpClient connectionHttpClient = httpClients.CreateClient("WorkflowMcp"); +``` + +The connection, not an individual tool call, would hold `connectionHttpClient`. +The application service provider must outlive all sessions using its factory. +Two wrappers created under the same configuration can share an underlying +handler. The factory remains responsible for that handler after wrapper disposal. + +For a long-lived wrapper on modern .NET, the documented configuration pattern is: + +```csharp +services.AddHttpClient("WorkflowMcp") + .UseSocketsHttpHandler((handler, _) => + { + handler.UseCookies = false; + handler.AllowAutoRedirect = false; + handler.SslOptions.CertificateRevocationCheckMode = + System.Security.Cryptography.X509Certificates.X509RevocationMode.Online; + handler.PooledConnectionLifetime = TimeSpan.FromMinutes(5); + }) + .SetHandlerLifetime(Timeout.InfiniteTimeSpan); +``` + +Five minutes is illustrative, not a required default. Connection-specific origin +policy is still required. This example is not a cross-target implementation recipe. + +#### Logical workflow ownership across approval + +The following is **proposed pseudocode**. These session/resource APIs do not +currently exist under the names shown and are not a public API commitment. + +```csharp +await using var session = + await runtime.OpenWorkflowSessionAsync(workflow, httpClients); + +var pendingApproval = await session.RunUntilWaitingAsync(input); +var approvalResponse = await application.GetApprovalAsync(pendingApproval); +await session.ResumeAsync(approvalResponse); +await session.CompleteAsync(); +``` + +The owner remains alive throughout the approval wait and resume. `CompleteAsync` +closes admission and drains work; leaving the scope is an explicit release or +abandonment boundary. Disposing a request-local HTTP/DI scope while waiting must +not accidentally dispose this longer-lived workflow owner. Web applications +must retain it independently of an individual request, or explicitly choose a +checkpoint/connection-recreation policy instead. + +Two sessions created from the same workflow graph each get distinct MCP clients. +A tool invocation in either session cannot dispose the other's HTTP wrapper. +No extra workflow ID is required in YAML; ownership comes from the runtime. + +#### Shutdown racing with an active tool invocation + +```text +Call A: admitted -> initialize/get MCP connection -> await tools/call -> release +Owner: request close -> reject Call B -> await A -> dispose +HTTP factory: shared handler remains usable by other sessions +``` + +An approval wait with no active tool calls does not initiate the closing sequence. +A cancelled caller releases its registration only after its operation unwinds. +HTTP handler rotation neither ends the MCP session nor completes the tool call. + +### Consequences + +- Good, because MCP session identity and ownership follow the workflow object + model, while HTTP pooling follows application infrastructure. +- Good, because workflow completion provides a natural resource-release boundary + without evicting another still-live workflow's state. +- Good, because the handler-wide LRU, retired-entry policy, and HTTP-client + reference counts are no longer the primary lifetime mechanism. +- Neutral, because per-owner connection lookup and concurrency coordination still + exist; this is not a claim that all bookkeeping can be removed. +- Bad, because runtime ownership must span direct runs, workflow-backed agents, + approval continuations, and child workflows rather than just the MCP package. +- Bad, because hosts must handle abandoned/waiting sessions explicitly, and + retaining them can consume more resources than the current bounded cache. +- Bad, because factory integration adds package/configuration considerations, + especially for older targets and connection-specific outer policy handlers. + +### User experience and potential breaking changes + +Adding this ADR is documentation-only. The following impacts describe a future +implementation, not behavior already delivered by this PR. + +| Area | Proposed experience and compatibility impact | +|---|---| +| Workflow authoring | No proposed YAML change or manual workflow-ID argument; runtime ownership remains transparent to workflow definitions | +| Application setup | Optional factory registration/injection for handler reuse; do not require ASP.NET or remove non-DI construction | +| Session cleanup | A logical owner must be retained and explicitly completed/disposed; applications that treated a shared handler as the sole owner may require changes | +| Approval and resume | Same-owner resumes preserve the MCP session between calls rather than depending on idle-cache retention | +| Capacity | Remove implicit live-session LRU replacement and the fixed global 32-entry behavior; configure admission/retention explicitly instead | +| Provider callbacks | Preserve per-invocation legacy behavior; adopting a connection-scoped provider is explicit and can change token-refresh timing | +| Custom handlers/contexts | Existing `IMcpToolHandler` implementations must not acquire new mandatory members silently; use adapters or an additive scoped contract | +| Constructor APIs | Prefer additive options/registration; a new overload must not make existing calls such as `new DefaultMcpToolHandler(null)` ambiguous | +| Checkpoint restoration | Serialized workflow state does not serialize a live MCP connection; no new cross-process continuity guarantee | +| Handler disposal | A shared handler must not unexpectedly destroy workflow-owned sessions; changing its existing drain/disposal contract requires explicit migration | +| Dependencies and targets | Factory integration can introduce `Microsoft.Extensions.Http` dependencies; preserve target support and document registration/version requirements | +| Diagnostics | Factory logging is additional observable behavior; configure redaction and do not log credentials, authentication partitions, or sensitive URLs | + +Removing existing constructors or interfaces, requiring a new session-owner +argument everywhere, changing callback frequency, or changing handler-disposal +semantics can be source, binary, or behavioral breaking changes. An adapter that +uses a short-lived owner per legacy invocation cannot claim state continuity. +Any compatibility layer retaining sessions must define its explicit lifetime, +not conceal another global cache. + +Released public API changes require the repository's API-baseline and package +compatibility process. Any intentional breaking change needs release notes and +the repository's breaking-change designation when implementation is proposed; +the documentation-only addition is not itself an API break. + +## Validation + +Before implementation is accepted, add focused regression coverage for: + +- Distinct sessions from the same workflow graph cannot share MCP state, while + successive calls and approval resumes in one live owner reuse its connection. +- Concurrent first use initializes one connection per configuration; changed + connection/authentication identity does not select an incompatible connection. +- Closing races with acquisition, initialization, and active calls without early + disposal, leaked unpublished connections, duplicate cleanup, or deadlock. +- Waiter cancellation does not terminate another invocation's work; initialization + failure and cleanup failure remain observable and permit appropriate retry. +- Child workflow resources are isolated and drained in the required owner order. +- Owned HTTP wrappers are disposed once; borrowed clients and application factories + are not disposed; wrapper disposal does not disrupt other factory clients. +- Shared handlers contain no per-workflow or per-caller mutable authentication + state, and origin and redirect behavior remain consistent with the current + transport contract. +- Waiting sessions remain live under unrelated workflow activity; admission + limits do not evict them or silently recreate their sessions. +- Same-process resume, explicit abandonment, shutdown, and checkpoint restore + follow their documented continuity and release contracts. +- Existing provider callbacks and custom handlers retain their documented behavior. + +Validate HTTP configuration and public API compatibility on every supported +target, and exercise console and hosted applications. Compare initialization +count, resource retention, shutdown behavior, and handler/connection reuse before +claiming a performance improvement. + +## Pros and Cons of the Options + +### 1. Handler-wide bounded cache and reference-counted HTTP sharing + +- Good, because it preserves the current entry points and bounds retained entries. +- Neutral, because synchronization and invocation protection already exist. +- Bad, because idle eviction can end a live workflow's stateful MCP session. +- Bad, because global capacity couples unrelated workflow owners and transport + sharing requires another ownership mechanism. + +### 2. Workflow-owned MCP connections with optional HTTP factory reuse + +- Good, because the ownership boundary matches protocol state and deterministic + release, while the factory manages reusable HTTP infrastructure. +- Neutral, because each session still needs a small connection registry and + synchronization; factory integration can remain optional. +- Bad, because explicit runtime lifecycle integration and migration are required, + and long-lived HTTP wrappers need deliberate connection-refresh configuration. + +### 3. Workflow-owned MCP connections with dedicated HTTP handlers + +- Good, because ownership and origin policy are straightforward without DI. +- Neutral, because this remains a valid non-DI fallback. +- Bad, because every connection can create another network connection pool, + increasing connection establishment and socket pressure at high session counts. + +### 4. One application-wide HttpClient + +- Good, because compatible requests can share one wrapper and its connection pools. +- Neutral, because HTTP supports concurrent requests to multiple origins. +- Bad, because it does not solve MCP ownership, and one common pipeline cannot + carry different per-origin pinning, proxy, or client-certificate configurations + without additional request-aware design. + +### 5. A new MCP session for every tool invocation + +- Good, because connection lifetime is limited to one operation. +- Neutral, because the existing provider-backed path deliberately uses this model. +- Bad, because it adds initialization per call and cannot preserve MCP session + state between calls; HTTP handler reuse does not restore that state. + +## More Information + +### Public documentation + +- [IHttpClientFactory with .NET](https://learn.microsoft.com/en-us/dotnet/core/extensions/httpclient-factory): + registration through `AddHttpClient`, including non-web Generic Host usage. +- [HTTP client lifetime management](https://learn.microsoft.com/en-us/dotnet/core/extensions/httpclient-factory#httpclient-lifetime-management): + new wrappers, pooled handlers, default two-minute handler lifetime, safe wrapper + disposal, and the expectation of short-lived factory clients. +- [Using the factory with SocketsHttpHandler](https://learn.microsoft.com/en-us/dotnet/core/extensions/httpclient-factory#using-ihttpclientfactory-together-with-socketshttphandler): + explicit `PooledConnectionLifetime` for long-lived connections. +- [Message handler scopes](https://learn.microsoft.com/en-us/dotnet/core/extensions/httpclient-factory#message-handler-scopes-in-ihttpclientfactory): + handler scopes are separate from application scopes; the documented outer + handler pattern uses `IHttpMessageHandlerFactory` for the pooled pipeline. +- [Named clients](https://learn.microsoft.com/en-us/dotnet/core/extensions/httpclient-factory#named-clients): + avoid an unbounded number of client names. + +### Implementation scope and remaining design work + +This proposal targets the .NET declarative MCP integration and the workflow +runtime ownership hooks it needs. It does not redesign Python, change protocol +result conversion, or introduce durable MCP reconnection. + +Before implementation, deciders must approve the concrete owner/API shape, +legacy-handler disposal adapter, ownership across run-handle replacement, +credential-refresh contract, capacity/abandonment configuration, and cross-target +HTTP configuration. Confirm EM/architect involvement through the listed review +participants before acceptance. Request `westey-m`'s review and inform `rhurey`; +PR approval is the decision record, not an assertion made by this document. + +Implement and validate the new lifecycle before removing existing guards. +HTTP factory integration alone must not be treated as fixing MCP session lifetime.