From bb9e92b189edc1367edc225a08cf8924657e82f8 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 22 Sep 2026 14:08:08 -0400 Subject: [PATCH 01/42] fix: isolate MCP sessions by workflow session Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7ed91283-cb21-49a7-98f4-b1c088d0b6b6 --- .../DefaultMcpToolHandler.cs | 47 ++++++++--- .../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/DeclarativeWorkflowContext.cs | 7 +- .../ObjectModel/InvokeMcpToolExecutor.cs | 43 +++++++++- .../PowerFx/WorkflowFormulaState.cs | 2 + .../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 | 4 +- .../IWorkflowContext.cs | 15 ++++ .../InProc/InProcessRunnerContext.cs | 4 +- .../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 + .../DefaultMcpToolHandlerLifetimeTests.cs | 40 +++++++++ .../DefaultMcpToolHandlerTests.cs | 37 ++++++--- .../ObjectModel/InvokeMcpToolExecutorTest.cs | 82 +++++++++++++++++++ python/packages/declarative/AGENTS.md | 3 + .../_workflows/_executors_mcp.py | 9 +- .../_workflows/_mcp_handler.py | 51 +++++++++--- .../tests/test_default_mcp_tool_handler.py | 71 +++++++++++----- .../tests/test_invoke_mcp_tool_executor.py | 29 +++++++ 31 files changed, 439 insertions(+), 61 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 b9c465f9f0..500874a1c9 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs @@ -28,11 +28,12 @@ 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. /// 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 string FilenameAdditionalPropertyName = "filename"; @@ -46,7 +47,7 @@ public sealed class DefaultMcpToolHandler : IMcpToolHandler, IAsyncDisposable private readonly Func>? _httpClientProvider; private readonly Func _httpMessageHandlerFactory; - private readonly Dictionary<(string Url, string Label, string Connection, string HeadersHash), ClientConnection> _clients = []; + private readonly Dictionary<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash), ClientConnection> _clients = []; private readonly Dictionary _ownedHttpClients = []; private readonly SemaphoreSlim _clientLock = new(1, 1); private readonly AsyncLocal _providerInvocationContext = new(); @@ -98,6 +99,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)) { @@ -151,7 +172,8 @@ public async Task InvokeToolAsync( } } - McpClient client = await this.GetOrCreateClientAsync(serverUrl, serverLabel, headers, connectionName, cancellationToken).ConfigureAwait(false); + McpClient client = await this.GetOrCreateClientAsync( + serverUrl, serverLabel, headers, connectionName, workflowSessionId, cancellationToken).ConfigureAwait(false); return await InvokeClientAsync(client, toolName, arguments, cancellationToken).ConfigureAwait(false); } @@ -276,10 +298,11 @@ private async Task GetOrCreateClientAsync( 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); await this._clientLock.WaitAsync(cancellationToken).ConfigureAwait(false); try @@ -301,17 +324,19 @@ private async Task GetOrCreateClientAsync( } /// - /// 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, 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 ab058de62d..3cdebd5e1a 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 ab058de62d..3cdebd5e1a 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 ab058de62d..3cdebd5e1a 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 ab058de62d..3cdebd5e1a 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 ab058de62d..3cdebd5e1a 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 56b1c3deb4..fbea21f35b 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/DeclarativeWorkflowContext.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowContext.cs index 6616aa5d00..92ffbe56a1 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowContext.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowContext.cs @@ -15,7 +15,7 @@ namespace Microsoft.Agents.AI.Workflows.Declarative.Interpreter; -internal sealed class DeclarativeWorkflowContext : IWorkflowContext +internal sealed class DeclarativeWorkflowContext : IWorkflowContext, IWorkflowSessionContext { public static readonly FrozenSet ManagedScopes = [ @@ -28,12 +28,17 @@ public DeclarativeWorkflowContext(IWorkflowContext source, WorkflowFormulaState { this.Source = source; this.State = state; + this.SessionId = source is IWorkflowSessionContext sessionContext + ? sessionContext.SessionId + : state.FallbackWorkflowSessionId; } 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/ObjectModel/InvokeMcpToolExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/InvokeMcpToolExecutor.cs index 46a5cae5fd..53fedcdb0d 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/InvokeMcpToolExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/InvokeMcpToolExecutor.cs @@ -31,6 +31,7 @@ internal sealed class InvokeMcpToolExecutor( { private const string ApprovalSnapshotStateKey = nameof(_approvalSnapshots); private const string LegacyApprovalSnapshotStateKey = "_approvalSnapshot"; + private readonly string _fallbackWorkflowSessionId = Guid.NewGuid().ToString("N"); /// /// Snapshots of evaluated parameters captured at approval-request time, keyed by @@ -115,7 +116,8 @@ public static bool RequiresNothing(object? message) => } // No approval required - invoke the tool directly - McpServerToolResultContent resultContent = await mcpToolHandler.InvokeToolAsync( + McpServerToolResultContent resultContent = await this.InvokeToolAsync( + context, serverUrl, serverLabel, toolName, @@ -172,7 +174,8 @@ public async ValueTask CaptureResponseAsync( Dictionary? headers = this.GetHeaders(); - McpServerToolResultContent resultContent = await mcpToolHandler.InvokeToolAsync( + McpServerToolResultContent resultContent = await this.InvokeToolAsync( + context, snapshot.ServerUrl, snapshot.ServerLabel, snapshot.ToolName, @@ -184,6 +187,42 @@ public async ValueTask CaptureResponseAsync( await this.ProcessResultAsync(context, resultContent, cancellationToken).ConfigureAwait(false); } + private Task InvokeToolAsync( + IWorkflowContext context, + string serverUrl, + string? serverLabel, + string toolName, + IDictionary? arguments, + IDictionary? headers, + string? connectionName, + CancellationToken cancellationToken) + { + if (mcpToolHandler is IWorkflowScopedMcpToolHandler scopedHandler) + { + string workflowSessionId = context is IWorkflowSessionContext sessionContext + ? sessionContext.SessionId + : this._fallbackWorkflowSessionId; + return scopedHandler.InvokeToolInWorkflowSessionAsync( + serverUrl, + serverLabel, + toolName, + arguments, + headers, + connectionName, + workflowSessionId, + cancellationToken); + } + + return mcpToolHandler.InvokeToolAsync( + serverUrl, + serverLabel, + toolName, + arguments, + headers, + connectionName, + cancellationToken); + } + /// /// Completes the MCP tool invocation by raising the completion event. /// diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PowerFx/WorkflowFormulaState.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PowerFx/WorkflowFormulaState.cs index aaff60b08b..f8370d722e 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PowerFx/WorkflowFormulaState.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PowerFx/WorkflowFormulaState.cs @@ -37,6 +37,8 @@ internal sealed class WorkflowFormulaState public WorkflowExpressionEngine Evaluator { get; } + public string FallbackWorkflowSessionId { get; } = System.Guid.NewGuid().ToString("N"); + public WorkflowFormulaState(RecalcEngine engine) { this._scopes = VariableScopeNames.AllScopes.ToDictionary(scopeName => GetScopeName(scopeName), _ => new WorkflowScope()); 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 a36677538e..bf09693fe4 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 @@ -1,2 +1,4 @@ #nullable enable 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! 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 a36677538e..bf09693fe4 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 @@ -1,2 +1,4 @@ #nullable enable 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! 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 a36677538e..bf09693fe4 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 @@ -1,2 +1,4 @@ #nullable enable 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! 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 a36677538e..bf09693fe4 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 @@ -1,2 +1,4 @@ #nullable enable 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! 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 a36677538e..bf09693fe4 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 @@ -1,2 +1,4 @@ #nullable enable 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! diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/GroupChatManager.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/GroupChatManager.cs index ab94a9fa8a..410316a010 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/GroupChatManager.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/GroupChatManager.cs @@ -169,7 +169,7 @@ 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 sealed class PrefixingWorkflowContext(IWorkflowContext inner, string prefix) : IWorkflowContext, IWorkflowSessionContext { private readonly IWorkflowContext _inner = Throw.IfNull(inner); private readonly string _prefix = Throw.IfNullOrEmpty(prefix); @@ -178,6 +178,8 @@ internal sealed class PrefixingWorkflowContext(IWorkflowContext inner, string pr public bool ConcurrentRunsEnabled => this._inner.ConcurrentRunsEnabled; + public string SessionId => ((IWorkflowSessionContext)this._inner).SessionId; + public ValueTask AddEventAsync(WorkflowEvent workflowEvent, CancellationToken cancellationToken = default) => this._inner.AddEventAsync(workflowEvent, cancellationToken); diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/IWorkflowContext.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/IWorkflowContext.cs index b8b35fffd6..68ff864085 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/InProcessRunnerContext.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunnerContext.cs index 353c38f25c..0db031c9a0 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunnerContext.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunnerContext.cs @@ -349,8 +349,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._sessionId; + 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 ab058de62d..724a27a74a 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 ab058de62d..724a27a74a 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 ab058de62d..724a27a74a 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 ab058de62d..724a27a74a 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 ab058de62d..724a27a74a 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/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs index 3a481fcd0e..37064cbbd9 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 @@ -104,6 +104,24 @@ 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 NoProvider_DifferentConnectionNames_UseSeparateCachedSessionsAsync() { @@ -185,6 +203,23 @@ 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 Provider_OperationFailure_DisposesSessionAndPreservesCallerClientAsync() { @@ -560,6 +595,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; 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 5a5b318f05..21ddf07416 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 @@ -418,7 +418,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] @@ -428,19 +428,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); @@ -452,8 +465,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); @@ -466,8 +479,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); @@ -481,8 +494,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); @@ -493,7 +506,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/ObjectModel/InvokeMcpToolExecutorTest.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/ObjectModel/InvokeMcpToolExecutorTest.cs index bb34b6c9a2..ded8316083 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 @@ -1,5 +1,6 @@ // Copyright (c) Microsoft. All rights reserved. +using System; using System.Collections.Concurrent; using System.Collections.Generic; using System.Linq; @@ -127,6 +128,58 @@ 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 InvokeMcpToolExecuteWithServerLabelAsync() { @@ -1778,6 +1831,35 @@ private InvokeMcpTool CreateModelWithVariableServerUrl(string displayName, strin #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/python/packages/declarative/AGENTS.md b/python/packages/declarative/AGENTS.md index 9638b43533..332db4cec0 100644 --- a/python/packages/declarative/AGENTS.md +++ b/python/packages/declarative/AGENTS.md @@ -15,6 +15,9 @@ YAML/JSON-based declarative agent and workflow definitions. ## MCP Handler Lifetimes `DefaultMCPToolHandler` caches/coalesces sessions only without a `client_provider`. +Cache identity includes a framework-owned workflow session ID in addition to +endpoint, label, connection, and headers, so separate workflow instances do not +share a stateful MCP protocol session while continuations of one workflow do. With a provider, every invocation (including `tools/list`) gets a fresh tool/session, even if the provider returns `None` or a shared HTTP client. Invocation cleanup closes the session and any internally owned fallback client, never caller-owned HTTP clients. diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_mcp.py b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_mcp.py index 7d3647d00b..3edd91c63d 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_mcp.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_mcp.py @@ -48,7 +48,12 @@ DeclarativeWorkflowState, ) from ._executors_tools import ToolApprovalResponse -from ._mcp_handler import MCPToolHandler, MCPToolInvocation, MCPToolResult +from ._mcp_handler import ( + MCPToolHandler, + MCPToolInvocation, + MCPToolResult, + _get_or_create_workflow_session_id, +) __all__ = [ "MCP_ACTION_EXECUTORS", @@ -236,6 +241,7 @@ async def handle_action( arguments=arguments, headers=headers, connection_name=connection_name, + workflow_session_id=_get_or_create_workflow_session_id(ctx.state), ) if require_approval: await self._request_approval( @@ -342,6 +348,7 @@ async def handle_approval_response( arguments=original_request.arguments, headers=self._evaluate_headers(state, self._action_def.get("headers")), connection_name=getattr(original_request, "connection_name", None), + workflow_session_id=_get_or_create_workflow_session_id(ctx.state), ) if invocation.headers or original_request.header_names: binding = getattr(original_request, "header_binding", None) diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py index 64eb59750f..f0f4c77434 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -30,6 +30,7 @@ import hashlib import json import logging +import uuid from collections import OrderedDict from collections.abc import Awaitable, Callable from contextvars import ContextVar, Token @@ -40,6 +41,7 @@ if TYPE_CHECKING: from agent_framework import Content + from agent_framework._workflows._state import State __all__ = [ "ClientProvider", @@ -52,6 +54,17 @@ logger = logging.getLogger(__name__) _DEFAULT_CACHE_MAX_SIZE = 32 +_WORKFLOW_SESSION_ID_KEY = "_declarative_mcp_workflow_session_id" + + +def _get_or_create_workflow_session_id(state: State) -> str: + workflow_session_id = state.get(_WORKFLOW_SESSION_ID_KEY) + if workflow_session_id is None: + workflow_session_id = uuid.uuid4().hex + state.set(_WORKFLOW_SESSION_ID_KEY, workflow_session_id) + if not isinstance(workflow_session_id, str) or not workflow_session_id: + raise ValueError("Invalid MCP workflow session state.") + return workflow_session_id @dataclass @@ -73,6 +86,9 @@ class MCPToolInvocation: - ``connection_name``: Optional Foundry connection name forwarded for handlers that resolve auth/credentials by connection. The default handler does not consume this field. + - ``workflow_session_id``: Framework-owned identifier for the current + workflow session. The default handler uses it to prevent separate + workflows from sharing one stateful MCP protocol session. """ server_url: str @@ -81,6 +97,7 @@ class MCPToolInvocation: arguments: dict[str, Any] = field(default_factory=dict) # type: ignore[reportUnknownVariableType] headers: dict[str, str] = field(default_factory=dict) # type: ignore[reportUnknownVariableType] connection_name: str | None = None + workflow_session_id: str | None = None def _empty_outputs() -> list[Any]: @@ -160,12 +177,13 @@ class DefaultMCPToolHandler: Without a ``client_provider``, caches one :class:`agent_framework.MCPStreamableHTTPTool` instance per - ``(server_url, server_label, connection_name, headers_hash)`` in a bounded - LRU. The cache prevents re-establishing an MCP session for every - invocation while ensuring different header sets (auth tokens) cannot - share a session — matches the .NET design intent while bounding - cardinality. ``server_label`` and ``connection_name`` also participate - in the key to distinguish logical connections. + ``(workflow_session_id, server_url, server_label, connection_name, + headers_hash)`` in a bounded LRU. The cache prevents re-establishing an MCP + session for every invocation while ensuring separate workflow sessions and + different header sets (auth tokens) cannot share a session — matches the + .NET design intent while bounding cardinality. ``server_label`` and + ``connection_name`` also participate in the key to distinguish logical + connections. Header *names* are lower-cased inside the hash payload only — the headers passed on the wire keep the caller's original casing — so two YAML actions that spell ``Authorization`` differently still share a @@ -228,14 +246,14 @@ def __init__( raise ValueError(f"cache_max_size must be positive, got {cache_max_size}") self._client_provider = client_provider self._cache_max_size = cache_max_size - self._cache: OrderedDict[tuple[str, str, str, str], _CacheEntry] = OrderedDict() + self._cache: OrderedDict[tuple[str, str, str, str, str], _CacheEntry] = OrderedDict() # Outer lock guards the cache + in-flight-future map only — never # held across network I/O. self._cache_lock = asyncio.Lock() # Per-key in-flight futures: while one task is connecting, other # tasks awaiting the same key will await the same future and share # the resulting cache entry. - self._inflight: dict[tuple[str, str, str, str], asyncio.Future[_CacheEntry]] = {} + self._inflight: dict[tuple[str, str, str, str, str], asyncio.Future[_CacheEntry]] = {} # Completion signals only: provider-backed calls never share entries. self._active_invocations: set[asyncio.Future[None]] = set() # Keep ancestry so a completed nested call cannot hide an active parent @@ -483,6 +501,7 @@ async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None: async def _get_or_create_entry(self, invocation: MCPToolInvocation) -> _CacheEntry: """Look up (or create) the cached MCP client for this invocation.""" key = self._cache_key( + invocation.workflow_session_id, invocation.server_url, invocation.server_label, invocation.connection_name, @@ -648,16 +667,18 @@ async def _close_entry(self, entry: _CacheEntry) -> None: @staticmethod def _cache_key( + workflow_session_id: str | None, server_url: str, server_label: str | None, connection_name: str | None, headers: dict[str, str] | None, - ) -> tuple[str, str, str, str]: + ) -> tuple[str, str, str, str, str]: """Build an order-independent cache key for the invocation identity. Used only without a ``client_provider``. The key includes - ``server_label`` and ``connection_name`` to distinguish logical - connections. + ``workflow_session_id`` to isolate stateful MCP protocol sessions + between workflows, plus ``server_label`` and ``connection_name`` to + distinguish logical connections. Header *names* are lower-cased inside the hash payload only so that ``Authorization`` and ``authorization`` map to the same @@ -669,4 +690,10 @@ def _cache_key( normalized = sorted((k.lower(), v) for k, v in headers.items()) payload = json.dumps(normalized, ensure_ascii=False) headers_hash = hashlib.sha256(payload.encode("utf-8")).hexdigest() - return (server_url, server_label or "", connection_name or "", headers_hash) + return ( + workflow_session_id or "", + server_url, + server_label or "", + connection_name or "", + headers_hash, + ) diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index 0826302cf6..58cadb183d 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -23,11 +23,13 @@ import httpx import pytest from agent_framework import Content +from agent_framework._workflows._state import State from agent_framework.exceptions import ToolExecutionException from agent_framework_declarative._workflows._mcp_handler import ( DefaultMCPToolHandler, MCPToolInvocation, + _get_or_create_workflow_session_id, ) pytestmark = pytest.mark.skipif( @@ -827,6 +829,22 @@ def test_invalid_cache_size_raises(self) -> None: DefaultMCPToolHandler(cache_max_size=-3) +class TestWorkflowSessionId: + def test_separate_workflow_states_get_separate_ids(self) -> None: + first = _get_or_create_workflow_session_id(State()) + second = _get_or_create_workflow_session_id(State()) + + assert first != second + + def test_same_workflow_state_reuses_id(self) -> None: + state = State() + + first = _get_or_create_workflow_session_id(state) + second = _get_or_create_workflow_session_id(state) + + assert first == second + + # ---------- Tool kwargs ---------------------------------------------------- @@ -879,12 +897,20 @@ class TestCache: async def test_same_url_and_headers_hit_cache(self) -> None: handler = DefaultMCPToolHandler() with _patch_tool(): - await handler.invoke_tool(_invocation(headers={"X": "1"})) - await handler.invoke_tool(_invocation(headers={"X": "1"})) + await handler.invoke_tool(_invocation(headers={"X": "1"}, workflow_session_id="workflow-a")) + await handler.invoke_tool(_invocation(headers={"X": "1"}, workflow_session_id="workflow-a")) # One tool created, connect called once. assert len(FakeTool.instances) == 1 assert FakeTool.instances[0].connect_count == 1 + @pytest.mark.asyncio + async def test_separate_workflow_sessions_use_separate_entries(self) -> None: + handler = DefaultMCPToolHandler() + with _patch_tool(): + await handler.invoke_tool(_invocation(workflow_session_id="workflow-a")) + await handler.invoke_tool(_invocation(workflow_session_id="workflow-b")) + assert len(FakeTool.instances) == 2 + @pytest.mark.asyncio async def test_different_headers_create_separate_entries(self) -> None: handler = DefaultMCPToolHandler() @@ -950,10 +976,10 @@ async def slow_connect(self: FakeTool) -> None: with _patch_tool(), patch.object(FakeTool, "connect", slow_connect): results = await asyncio.gather( - handler.invoke_tool(_invocation(headers={"X": "1"})), - handler.invoke_tool(_invocation(headers={"X": "1"})), - handler.invoke_tool(_invocation(headers={"X": "1"})), - handler.invoke_tool(_invocation(headers={"X": "1"})), + handler.invoke_tool(_invocation(headers={"X": "1"}, workflow_session_id="workflow-a")), + handler.invoke_tool(_invocation(headers={"X": "1"}, workflow_session_id="workflow-a")), + handler.invoke_tool(_invocation(headers={"X": "1"}, workflow_session_id="workflow-a")), + handler.invoke_tool(_invocation(headers={"X": "1"}, workflow_session_id="workflow-a")), ) assert all(not r.is_error for r in results) # Only one tool was created and connected, despite 4 concurrent calls. @@ -1235,38 +1261,43 @@ def boom(**_a: Any) -> Any: class TestCacheKey: def test_key_order_independent(self) -> None: - k1 = DefaultMCPToolHandler._cache_key("https://x/", None, None, {"A": "1", "B": "2"}) - k2 = DefaultMCPToolHandler._cache_key("https://x/", None, None, {"B": "2", "A": "1"}) + k1 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", None, None, {"A": "1", "B": "2"}) + k2 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", None, None, {"B": "2", "A": "1"}) assert k1 == k2 def test_key_distinguishes_values(self) -> None: - k1 = DefaultMCPToolHandler._cache_key("https://x/", None, None, {"A": "1"}) - k2 = DefaultMCPToolHandler._cache_key("https://x/", None, None, {"A": "2"}) + k1 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", None, None, {"A": "1"}) + k2 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", None, None, {"A": "2"}) assert k1 != k2 def test_empty_headers_use_fixed_hash(self) -> None: - k1 = DefaultMCPToolHandler._cache_key("https://x/", None, None, None) - k2 = DefaultMCPToolHandler._cache_key("https://x/", None, None, {}) + k1 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", None, None, None) + k2 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", None, None, {}) assert k1 == k2 + def test_key_distinguishes_workflow_session(self) -> None: + k1 = DefaultMCPToolHandler._cache_key("workflow-a", "https://x/", None, None, None) + k2 = DefaultMCPToolHandler._cache_key("workflow-b", "https://x/", None, None, None) + assert k1 != k2 + def test_key_distinguishes_connection_name(self) -> None: - k1 = DefaultMCPToolHandler._cache_key("https://x/", None, "conn-A", None) - k2 = DefaultMCPToolHandler._cache_key("https://x/", None, "conn-B", None) + k1 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", None, "conn-A", None) + k2 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", None, "conn-B", None) assert k1 != k2 def test_key_distinguishes_server_label(self) -> None: - k1 = DefaultMCPToolHandler._cache_key("https://x/", "Lbl-A", None, None) - k2 = DefaultMCPToolHandler._cache_key("https://x/", "Lbl-B", None, None) + k1 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", "Lbl-A", None, None) + k2 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", "Lbl-B", None, None) assert k1 != k2 def test_key_collapses_header_name_case(self) -> None: - k1 = DefaultMCPToolHandler._cache_key("https://x/", None, None, {"Authorization": "tk"}) - k2 = DefaultMCPToolHandler._cache_key("https://x/", None, None, {"authorization": "tk"}) + k1 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", None, None, {"Authorization": "tk"}) + k2 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", None, None, {"authorization": "tk"}) assert k1 == k2 def test_key_keeps_header_value_case(self) -> None: - k1 = DefaultMCPToolHandler._cache_key("https://x/", None, None, {"X": "Bearer-A"}) - k2 = DefaultMCPToolHandler._cache_key("https://x/", None, None, {"X": "bearer-a"}) + k1 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", None, None, {"X": "Bearer-A"}) + k2 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", None, None, {"X": "bearer-a"}) assert k1 != k2 diff --git a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py index 28b3a1a0d6..c0896c4c6f 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -164,6 +164,35 @@ async def test_basic_invocation_forwards_required_fields(self) -> None: assert inv.headers == {} assert inv.arguments == {} assert inv.connection_name is None + assert inv.workflow_session_id + + @pytest.mark.asyncio + async def test_separate_workflows_receive_separate_session_ids(self) -> None: + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + first = factory.create_workflow_from_definition(_yaml(_action())) + second = factory.create_workflow_from_definition(_yaml(_action())) + + await first.run({}) + await second.run({}) + + assert len(handler.invocations) == 2 + assert handler.invocations[0].workflow_session_id + assert handler.invocations[1].workflow_session_id + assert handler.invocations[0].workflow_session_id != handler.invocations[1].workflow_session_id + + @pytest.mark.asyncio + async def test_same_workflow_continuation_reuses_session_id(self) -> None: + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + workflow = factory.create_workflow_from_definition(_yaml(_action())) + + await workflow.run({}) + await workflow.run({}) + + assert len(handler.invocations) == 2 + assert handler.invocations[0].workflow_session_id + assert handler.invocations[0].workflow_session_id == handler.invocations[1].workflow_session_id @pytest.mark.asyncio async def test_arguments_evaluated_and_preserves_none(self) -> None: From 5232c71a8ae40da6d62f37a631c6cef9747f5297 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 22 Sep 2026 15:30:48 -0400 Subject: [PATCH 02/42] .NET/Python: Address MCP session review feedback Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../DefaultMcpToolHandler.cs | 169 ++++++++++++++++-- .../Interpreter/DeclarativeActionExecutor.cs | 4 +- .../Interpreter/DeclarativeWorkflowContext.cs | 22 ++- .../DeclarativeWorkflowExecutor.cs | 3 +- .../Interpreter/DelegateActionExecutor.cs | 4 +- .../Kit/ActionExecutor.cs | 4 +- .../Kit/RootExecutor.cs | 3 +- .../PowerFx/WorkflowFormulaState.cs | 2 - .../GroupChatManager.cs | 21 ++- .../DefaultMcpToolHandlerLifetimeTests.cs | 52 ++++++ .../PowerFx/WorkflowFormulaStateTests.cs | 54 ++++++ .../_workflows/_mcp_handler.py | 71 +++++++- .../tests/test_default_mcp_tool_handler.py | 71 ++++++++ 13 files changed, 443 insertions(+), 37 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 500874a1c9..effa492f68 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs @@ -29,12 +29,14 @@ namespace Microsoft.Agents.AI.Workflows.Declarative.Mcp; /// 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, workflow invocations are cached by workflow session, server URL, label, -/// connection name, and explicit headers. +/// connection name, and explicit headers in a bounded least-recently-used cache. Evicted sessions +/// are disposed after their active invocations finish. /// Non-cancellation cleanup failures are reported through warnings without replacing /// the invocation result or error. /// public sealed class DefaultMcpToolHandler : IWorkflowScopedMcpToolHandler, IAsyncDisposable { + private const int DefaultClientCacheMaxSize = 32; private const string FilenameAdditionalPropertyName = "filename"; /// @@ -47,9 +49,12 @@ public sealed class DefaultMcpToolHandler : IWorkflowScopedMcpToolHandler, IAsyn private readonly Func>? _httpClientProvider; private readonly Func _httpMessageHandlerFactory; - private readonly Dictionary<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash), ClientConnection> _clients = []; + private readonly Dictionary<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash), CachedClient> _clients = []; + 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 int _clientCacheMaxSize; private readonly AsyncLocal _providerInvocationContext = new(); private TaskCompletionSource? _providerInvocationsDrained; private int _activeProviderInvocations; @@ -84,10 +89,17 @@ public DefaultMcpToolHandler(Func>? internal DefaultMcpToolHandler( Func>? httpClientProvider, - Func httpMessageHandlerFactory) + Func httpMessageHandlerFactory, + int clientCacheMaxSize = DefaultClientCacheMaxSize) { + 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._clientCacheMaxSize = clientCacheMaxSize; } /// @@ -172,9 +184,16 @@ public async Task InvokeToolInWorkflowSessionAsync( } } - McpClient client = await this.GetOrCreateClientAsync( + CachedClient cachedClient = await this.AcquireClientAsync( serverUrl, serverLabel, headers, connectionName, workflowSessionId, cancellationToken).ConfigureAwait(false); - return await InvokeClientAsync(client, toolName, arguments, 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( @@ -257,15 +276,41 @@ public async ValueTask DisposeAsync() await providerInvocations.ConfigureAwait(false); } + 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(); + this._clientLru.Clear(); + } + finally + { + this._clientLock.Release(); + } + + foreach (CachedClient client in clientsToDispose) + { + await this.DisposeCachedClientAsync(client).ConfigureAwait(false); + } + + await Task.WhenAll(cachedClients.Select(client => client.Disposed.Task)).ConfigureAwait(false); + + await this._clientLock.WaitAsync().ConfigureAwait(false); + try + { // Dispose only HttpClients that the handler created (not user-provided ones) foreach (HttpClient httpClient in this._ownedHttpClients.Values) @@ -293,7 +338,7 @@ private void ThrowIfDisposing() } } - private async Task GetOrCreateClientAsync( + private async Task AcquireClientAsync( string serverUrl, string? serverLabel, IDictionary? headers, @@ -303,24 +348,105 @@ private async Task GetOrCreateClientAsync( { string trimmedUrl = serverUrl.Trim(); var clientCacheKey = BuildCacheKey(workflowSessionId, trimmedUrl, serverLabel, connectionName, headers); + CachedClient? clientToDispose = null; + CachedClient result; 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)) { - return existingClient.Client; + existingClient.ActiveInvocations++; + this._clientLru.Remove(existingClient.LruNode); + this._clientLru.AddLast(existingClient.LruNode); + result = existingClient; } + else + { + ClientConnection connection = await this.CreateClientAsync( + trimmedUrl, serverLabel, headers, trimmedUrl, cancellationToken).ConfigureAwait(false); + LinkedListNode<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash)> node = + this._clientLru.AddLast(clientCacheKey); + CachedClient newClient = new(connection, node) { ActiveInvocations = 1 }; + this._clients[clientCacheKey] = newClient; + + if (this._clients.Count > this._clientCacheMaxSize) + { + LinkedListNode<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash)> evictedNode = + this._clientLru.First!; + this._clientLru.RemoveFirst(); + CachedClient evictedClient = this._clients[evictedNode.Value]; + this._clients.Remove(evictedNode.Value); + evictedClient.Evicted = true; + this._retiredClients.Add(evictedClient); + if (evictedClient.ActiveInvocations == 0) + { + evictedClient.DisposalClaimed = true; + clientToDispose = evictedClient; + } + } - ClientConnection newClient = await this.CreateClientAsync(trimmedUrl, serverLabel, headers, trimmedUrl, cancellationToken).ConfigureAwait(false); - this._clients[clientCacheKey] = newClient; - return newClient.Client; + result = newClient; + } } finally { this._clientLock.Release(); } + + if (clientToDispose is not null) + { + await this.DisposeCachedClientAsync(clientToDispose).ConfigureAwait(false); + } + + return result; + } + + 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) + { + try + { + await client.Connection.DisposeAsync().ConfigureAwait(false); + } + finally + { + await this._clientLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); + try + { + this._retiredClients.Remove(client); + } + finally + { + this._clientLock.Release(); + } + + client.Disposed.TrySetResult(true); + } } /// @@ -418,6 +544,23 @@ private sealed class ProviderInvocationContext(Task completion, ProviderInvocati 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; } + } + internal sealed class ClientConnection(McpClient client, IAsyncDisposable transport) : IAsyncDisposable { public McpClient Client { get; } = client; 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 5d052c64d3..631519288b 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeActionExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeActionExecutor.cs @@ -89,7 +89,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 92ffbe56a1..a60a7756f5 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowContext.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowContext.cs @@ -17,6 +17,8 @@ namespace Microsoft.Agents.AI.Workflows.Declarative.Interpreter; internal sealed class DeclarativeWorkflowContext : IWorkflowContext, IWorkflowSessionContext { + private const string WorkflowSessionIdStateKey = "__declarative_mcp_workflow_session_id"; + public static readonly FrozenSet ManagedScopes = [ VariableScopeNames.Local, @@ -24,13 +26,27 @@ internal sealed class DeclarativeWorkflowContext : IWorkflowContext, IWorkflowSe VariableScopeNames.Global, ]; - public DeclarativeWorkflowContext(IWorkflowContext source, WorkflowFormulaState state) + private DeclarativeWorkflowContext(IWorkflowContext source, WorkflowFormulaState state, string sessionId) { this.Source = source; this.State = state; - this.SessionId = source is IWorkflowSessionContext sessionContext + this.SessionId = sessionId; + } + + public static async ValueTask CreateAsync( + IWorkflowContext source, + WorkflowFormulaState state, + CancellationToken cancellationToken = default) + { + string sessionId = source is IWorkflowSessionContext sessionContext ? sessionContext.SessionId - : state.FallbackWorkflowSessionId; + : await source.ReadOrInitStateAsync( + WorkflowSessionIdStateKey, + static () => Guid.NewGuid().ToString("N"), + scopeName: null, + cancellationToken: cancellationToken).ConfigureAwait(false); + + return new(source, state, sessionId); } private IWorkflowContext Source { get; } 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 1126b43d8f..2fcdea4ecc 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 716a3656da..ade8dd1034 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 cf636effaf..563ff86f14 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 80f6e69b60..bf6ab5871f 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/RootExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/RootExecutor.cs @@ -57,7 +57,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/PowerFx/WorkflowFormulaState.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PowerFx/WorkflowFormulaState.cs index f8370d722e..aaff60b08b 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PowerFx/WorkflowFormulaState.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PowerFx/WorkflowFormulaState.cs @@ -37,8 +37,6 @@ internal sealed class WorkflowFormulaState public WorkflowExpressionEngine Evaluator { get; } - public string FallbackWorkflowSessionId { get; } = System.Guid.NewGuid().ToString("N"); - public WorkflowFormulaState(RecalcEngine engine) { this._scopes = VariableScopeNames.AllScopes.ToDictionary(scopeName => GetScopeName(scopeName), _ => new WorkflowScope()); diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/GroupChatManager.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/GroupChatManager.cs index 410316a010..fbda33c132 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,17 +169,20 @@ 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, IWorkflowSessionContext +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; - public string SessionId => ((IWorkflowSessionContext)this._inner).SessionId; - public ValueTask AddEventAsync(WorkflowEvent workflowEvent, CancellationToken cancellationToken = default) => this._inner.AddEventAsync(workflowEvent, cancellationToken); @@ -224,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/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs index 37064cbbd9..e87714025c 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 @@ -122,6 +122,58 @@ public async Task NoProvider_SeparateWorkflowSessions_UseSeparateCachedSessionsA Assert.Equal(0, stub.Terminations); } + [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_DefersDisposalUntilActiveInvocationCompletesAsync() + { + // 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: 1); + 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); + + // Assert + Assert.Equal(0, stub.Terminations); + Task disposal = handler.DisposeAsync().AsTask(); + Assert.False(disposal.IsCompleted); + releaseFirst.Release(); + await first; + await disposal; + Assert.Equal(2, stub.Terminations); + } + [Fact] public async Task NoProvider_DifferentConnectionNames_UseSeparateCachedSessionsAsync() { 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 5193296db4..eb627e24c7 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 @@ -1,8 +1,14 @@ // Copyright (c) Microsoft. All rights reserved. +using System; +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.PowerFx.Types; +using Moq; namespace Microsoft.Agents.AI.Workflows.Declarative.UnitTests.PowerFx; @@ -81,4 +87,52 @@ public void SetOverwritesExistingValue() FormulaValue result = this.State.Get("key1"); Assert.Equal(newValue, result); } + + [Fact] + public async Task DeclarativeContextFallbackSessionId_IsScopedToPersistedRunStateAsync() + { + // Arrange + Dictionary firstRunState = []; + Dictionary 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); + } + + private static IWorkflowContext CreateContext(Dictionary 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) => + { + if (!state.TryGetValue(key, out string? value)) + { + value = factory(); + state[key] = value; + } + + return new ValueTask(value); + }); + return context.Object; + } } diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py index f0f4c77434..95280353c4 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -170,6 +170,10 @@ class _CacheEntry: tool: Any # MCPStreamableHTTPTool — typed Any to avoid import at module load owned_httpx_client: httpx.AsyncClient | None + active_users: int = 0 + evicted: bool = False + disposal_claimed: bool = False + closed: asyncio.Event = field(default_factory=asyncio.Event) class DefaultMCPToolHandler: @@ -183,7 +187,9 @@ class DefaultMCPToolHandler: different header sets (auth tokens) cannot share a session — matches the .NET design intent while bounding cardinality. ``server_label`` and ``connection_name`` also participate in the key to distinguish logical - connections. + connections. Active invocations lease their cache entry, so eviction waits + for the final user before closing it. Concurrent entry creation is also + capped at ``cache_max_size``. Header *names* are lower-cased inside the hash payload only — the headers passed on the wire keep the caller's original casing — so two YAML actions that spell ``Authorization`` differently still share a @@ -217,8 +223,9 @@ class DefaultMCPToolHandler: client_provider: Optional per-invocation ``httpx.AsyncClient`` provider. cache_max_size: Maximum number of cached MCP clients in no-provider mode. When exceeded, the least-recently-used entry is evicted and its - owned client closed. Defaults to ``32``. Does not enable session - caching when a provider is configured. + owned client closed after any active invocation finishes. This also + limits simultaneous connection attempts. Defaults to ``32``. Does + not enable session caching when a provider is configured. """ LIST_TOOLS_TOOL_NAME: ClassVar[str] = "tools/list" @@ -247,9 +254,11 @@ def __init__( self._client_provider = client_provider self._cache_max_size = cache_max_size self._cache: OrderedDict[tuple[str, str, str, str, str], _CacheEntry] = OrderedDict() + self._retired: dict[int, _CacheEntry] = {} # Outer lock guards the cache + in-flight-future map only — never # held across network I/O. self._cache_lock = asyncio.Lock() + self._creation_semaphore = asyncio.Semaphore(cache_max_size) # Per-key in-flight futures: while one task is connecting, other # tasks awaiting the same key will await the same future and share # the resulting cache entry. @@ -321,7 +330,9 @@ async def invoke_tool(self, invocation: MCPToolInvocation) -> MCPToolResult: return await self._invoke_entry(entry, invocation) finally: - if completion is not None: + if self._client_provider is None and entry is not None: + await self._release_entry(entry) + elif completion is not None: try: if entry is not None: await self._close_invocation_entry(entry) @@ -466,7 +477,16 @@ async def aclose(self) -> None: return self._closed = True entries = list(self._cache.values()) + entry_ids = {id(entry) for entry in entries} + entries.extend(entry for entry in self._retired.values() if id(entry) not in entry_ids) self._cache.clear() + entries_to_close: list[_CacheEntry] = [] + for entry in entries: + entry.evicted = True + self._retired[id(entry)] = entry + if entry.active_users == 0 and not entry.disposal_claimed: + entry.disposal_claimed = True + entries_to_close.append(entry) inflight_futures = list(self._inflight.values()) active_invocations = list(self._active_invocations) @@ -485,8 +505,10 @@ async def aclose(self) -> None: logger.debug("DefaultMCPToolHandler: in-flight future raised during aclose", exc_info=True) continue + for entry in entries_to_close: + await self._close_claimed_entry(entry) for entry in entries: - await self._close_entry(entry) + await entry.closed.wait() async def __aenter__(self) -> DefaultMCPToolHandler: return self @@ -517,6 +539,7 @@ async def _get_or_create_entry(self, invocation: MCPToolInvocation) -> _CacheEnt existing = self._cache.get(key) if existing is not None: self._cache.move_to_end(key) + existing.active_users += 1 return existing inflight = self._inflight.get(key) if inflight is None: @@ -525,11 +548,16 @@ async def _get_or_create_entry(self, invocation: MCPToolInvocation) -> _CacheEnt creating = True if not creating: - return await inflight + await inflight + return await self._get_or_create_entry(invocation) # Phase 2: we own creation. Build the entry outside the lock. try: - entry = await self._create_entry(invocation) + async with self._creation_semaphore: + async with self._cache_lock: + if self._closed: + raise RuntimeError("DefaultMCPToolHandler is closed") + entry = await self._create_entry(invocation) except BaseException as exc: async with self._cache_lock: self._inflight.pop(key, None) @@ -561,11 +589,17 @@ async def _get_or_create_entry(self, invocation: MCPToolInvocation) -> _CacheEnt self._cache.move_to_end(key) duplicate = entry entry = existing + entry.active_users += 1 else: + entry.active_users = 1 self._cache[key] = entry self._cache.move_to_end(key) if len(self._cache) > self._cache_max_size: _evicted_key, evicted = self._cache.popitem(last=False) + evicted.evicted = True + self._retired[id(evicted)] = evicted + if evicted.active_users == 0: + evicted.disposal_claimed = True if not inflight.done(): inflight.set_result(entry) @@ -582,10 +616,29 @@ async def _get_or_create_entry(self, invocation: MCPToolInvocation) -> _CacheEnt raise err if duplicate is not None: await self._close_entry(duplicate) - if evicted is not None: - await self._close_entry(evicted) + if evicted is not None and evicted.disposal_claimed: + await self._close_claimed_entry(evicted) return entry + async def _release_entry(self, entry: _CacheEntry) -> None: + close_entry = False + async with self._cache_lock: + entry.active_users -= 1 + if entry.active_users == 0 and entry.evicted and not entry.disposal_claimed: + entry.disposal_claimed = True + close_entry = True + + if close_entry: + await self._close_claimed_entry(entry) + + async def _close_claimed_entry(self, entry: _CacheEntry) -> None: + try: + await self._close_invocation_entry(entry) + finally: + async with self._cache_lock: + self._retired.pop(id(entry), None) + entry.closed.set() + async def _create_entry(self, invocation: MCPToolInvocation) -> _CacheEntry: """Construct (and connect) a fresh MCP client for ``invocation``.""" from agent_framework import MCPStreamableHTTPTool diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index 58cadb183d..b43a4243a1 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -943,6 +943,77 @@ async def test_lru_eviction_closes_old_entry(self) -> None: assert FakeTool.instances[1].close_count == 0 assert FakeTool.instances[2].close_count == 0 + @pytest.mark.asyncio + async def test_lru_eviction_defers_close_until_active_invocation_finishes(self) -> None: + handler = DefaultMCPToolHandler(cache_max_size=1) + first_started = asyncio.Event() + release_first = asyncio.Event() + + async def gated_call(tool: FakeTool, tool_name: str, **arguments: Any) -> Any: + if tool.kwargs["url"] == "https://a/": + first_started.set() + await release_first.wait() + return [Content.from_text("ok")] + + with _patch_tool(), patch.object(FakeTool, "call_tool", gated_call): + first = asyncio.create_task(handler.invoke_tool(_invocation(server_url="https://a/"))) + await first_started.wait() + second = await handler.invoke_tool(_invocation(server_url="https://b/")) + + assert not second.is_error + assert FakeTool.instances[0].close_count == 0 + close_task = asyncio.create_task(handler.aclose()) + await asyncio.sleep(0) + assert not close_task.done() + + release_first.set() + first_result = await first + await close_task + + assert not first_result.is_error + assert FakeTool.instances[0].close_count == 1 + assert FakeTool.instances[1].close_count == 1 + + @pytest.mark.asyncio + async def test_entry_creation_is_bounded_and_cancelled_waiter_is_cleaned_up(self) -> None: + handler = DefaultMCPToolHandler(cache_max_size=2) + connecting = 0 + max_connecting = 0 + capacity_reached = asyncio.Event() + release = asyncio.Event() + original_connect = FakeTool.connect + + async def gated_connect(tool: FakeTool) -> None: + nonlocal connecting, max_connecting + connecting += 1 + max_connecting = max(max_connecting, connecting) + if connecting == 2: + capacity_reached.set() + try: + await release.wait() + await original_connect(tool) + finally: + connecting -= 1 + + with _patch_tool(), patch.object(FakeTool, "connect", gated_connect): + first = asyncio.create_task(handler.invoke_tool(_invocation(workflow_session_id="workflow-a"))) + second = asyncio.create_task(handler.invoke_tool(_invocation(workflow_session_id="workflow-b"))) + await capacity_reached.wait() + cancelled = asyncio.create_task(handler.invoke_tool(_invocation(workflow_session_id="workflow-c"))) + await asyncio.sleep(0) + + assert len(FakeTool.instances) == 2 + cancelled.cancel() + with pytest.raises(asyncio.CancelledError): + await cancelled + + release.set() + results = await asyncio.gather(first, second) + + assert all(not result.is_error for result in results) + assert max_connecting == 2 + assert not handler._inflight + @pytest.mark.asyncio async def test_repeated_use_keeps_lru_alive(self) -> None: handler = DefaultMCPToolHandler(cache_max_size=2) From 777056f8558a8ea36164d7a4ebfaeaf845984eec Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Wed, 23 Sep 2026 09:11:31 -0400 Subject: [PATCH 03/42] Address MCP review feedback Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 89548cb9-9583-47c1-a00f-675aebf810f8 --- .../DefaultMcpToolHandler.cs | 152 +++++++++++++++--- .../DefaultMcpToolHandlerLifetimeTests.cs | 35 ++++ .../_workflows/_mcp_handler.py | 2 +- .../tests/test_default_mcp_tool_handler.py | 38 ++++- 4 files changed, 203 insertions(+), 24 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 effa492f68..340b8c09a2 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs @@ -50,6 +50,7 @@ public sealed class DefaultMcpToolHandler : IWorkflowScopedMcpToolHandler, IAsyn private readonly Func>? _httpClientProvider; private readonly Func _httpMessageHandlerFactory; 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 LinkedList<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash)> _clientLru = []; private readonly HashSet _retiredClients = []; private readonly Dictionary _ownedHttpClients = []; @@ -259,12 +260,14 @@ public async ValueTask DisposeAsync() } Task? providerInvocations; + List> clientCreations; await this._clientLock.WaitAsync().ConfigureAwait(false); try { this.ThrowIfDisposing(); this._disposing = true; providerInvocations = this._providerInvocationsDrained?.Task; + clientCreations = this._clientCreations.Values.Select(source => source.Task).ToList(); } finally { @@ -276,6 +279,22 @@ public async ValueTask DisposeAsync() await providerInvocations.ConfigureAwait(false); } + foreach (Task clientCreation in clientCreations) + { + try + { + await clientCreation.ConfigureAwait(false); + } + catch (OperationCanceledException) + { + // The invocation that owned this creation was cancelled; disposal only needs its cleanup to complete. + } + catch (Exception exception) + { + Trace.TraceWarning("MCP client creation failed during disposal: {0}", exception); + } + } + List cachedClients; List clientsToDispose = []; await this._clientLock.WaitAsync().ConfigureAwait(false); @@ -349,7 +368,8 @@ private async Task AcquireClientAsync( string trimmedUrl = serverUrl.Trim(); var clientCacheKey = BuildCacheKey(workflowSessionId, trimmedUrl, serverLabel, connectionName, headers); CachedClient? clientToDispose = null; - CachedClient result; + TaskCompletionSource? clientCreation; + bool ownsClientCreation = false; await this._clientLock.WaitAsync(cancellationToken).ConfigureAwait(false); try @@ -360,15 +380,57 @@ private async Task AcquireClientAsync( existingClient.ActiveInvocations++; this._clientLru.Remove(existingClient.LruNode); this._clientLru.AddLast(existingClient.LruNode); - result = existingClient; + return existingClient; + } + + if (!this._clientCreations.TryGetValue(clientCacheKey, out clientCreation)) + { + clientCreation = new(TaskCreationOptions.RunContinuationsAsynchronously); + this._clientCreations[clientCacheKey] = clientCreation; + ownsClientCreation = true; + } + } + finally + { + this._clientLock.Release(); + } + + if (!ownsClientCreation) + { + await clientCreation.Task.ConfigureAwait(false); + return await this.AcquireClientAsync( + serverUrl, serverLabel, headers, connectionName, workflowSessionId, cancellationToken).ConfigureAwait(false); + } + + ClientConnection? connection = null; + try + { + connection = await this.CreateClientAsync( + trimmedUrl, serverLabel, headers, trimmedUrl, cancellationToken).ConfigureAwait(false); + } + catch (Exception exception) + { + await this.CompleteClientCreationFailureAsync(clientCacheKey, clientCreation, exception).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)); + clientCreation.TrySetException(disposedException); } else { - ClientConnection connection = await this.CreateClientAsync( - trimmedUrl, serverLabel, headers, trimmedUrl, cancellationToken).ConfigureAwait(false); 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; if (this._clients.Count > this._clientCacheMaxSize) @@ -388,6 +450,7 @@ private async Task AcquireClientAsync( } result = newClient; + clientCreation.TrySetResult(newClient); } } finally @@ -395,12 +458,44 @@ private async Task AcquireClientAsync( this._clientLock.Release(); } + if (connection is not null) + { + await connection.DisposeAsync().ConfigureAwait(false); + } + if (clientToDispose is not null) { await this.DisposeCachedClientAsync(clientToDispose).ConfigureAwait(false); } - return result; + if (disposedException is not null) + { + throw disposedException; + } + + return result ?? throw new InvalidOperationException("Failed to acquire MCP client."); + } + + 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); + } + finally + { + this._clientLock.Release(); + } } private async ValueTask ReleaseClientAsync(CachedClient client) @@ -480,25 +575,26 @@ private async Task CreateClientAsync( 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 && httpClientCacheKey is 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). - OriginPinningHandler pinningHandler = new(new Uri(serverUrl)) { InnerHandler = this._httpMessageHandlerFactory() }; - httpClient = new HttpClient(pinningHandler); - if (httpClientCacheKey is null) + httpClient = this.CreatePinnedHttpClient(serverUrl); + ownsHttpClient = true; + } + else if (httpClient is null && httpClientCacheKey is not null) + { + string cacheKey = httpClientCacheKey; + await this._clientLock.WaitAsync(cancellationToken).ConfigureAwait(false); + try { - ownsHttpClient = true; + if (!this._ownedHttpClients.TryGetValue(cacheKey, out httpClient)) + { + httpClient = this.CreatePinnedHttpClient(serverUrl); + this._ownedHttpClients[cacheKey] = httpClient; + } } - else + finally { - this._ownedHttpClients[httpClientCacheKey] = httpClient; + this._clientLock.Release(); } } @@ -514,7 +610,8 @@ private async Task CreateClientAsync( TransportMode = HttpTransportMode.StreamableHttp }; - HttpClientTransport transport = new(transportOptions, httpClient, ownsHttpClient: ownsHttpClient); + HttpClient resolvedHttpClient = httpClient ?? throw new InvalidOperationException("Failed to resolve MCP HTTP client."); + HttpClientTransport transport = new(transportOptions, resolvedHttpClient, ownsHttpClient: ownsHttpClient); try { @@ -528,6 +625,19 @@ private async Task CreateClientAsync( } } + 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 { 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 e87714025c..7199cbef90 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 @@ -174,6 +174,35 @@ public async Task NoProvider_CacheEviction_DefersDisposalUntilActiveInvocationCo Assert.Equal(2, stub.Terminations); } + [Fact] + public async Task NoProvider_ConcurrentWorkflowSessionCreations_DoNotSerializeHandshakeAsync() + { + // 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)); + using CancellationTokenSource concurrencyTimeout = new(TimeSpan.FromSeconds(2)); + + // 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(concurrencyTimeout.Token); + releaseInitializations.Release(2); + await Task.WhenAll(first, second); + + // Assert + Assert.Equal(2, stub.Initializations); + Assert.Equal(0, stub.Terminations); + } + [Fact] public async Task NoProvider_DifferentConnectionNames_UseSeparateCachedSessionsAsync() { @@ -661,6 +690,7 @@ 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; } @@ -721,6 +751,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/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py index 95280353c4..e4f6330e41 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -548,7 +548,7 @@ async def _get_or_create_entry(self, invocation: MCPToolInvocation) -> _CacheEnt creating = True if not creating: - await inflight + _ = await asyncio.shield(inflight) return await self._get_or_create_entry(invocation) # Phase 2: we own creation. Build the entry outside the lock. diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index b43a4243a1..b55f8086ee 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -968,9 +968,10 @@ async def gated_call(tool: FakeTool, tool_name: str, **arguments: Any) -> Any: release_first.set() first_result = await first - await close_task + close_result = await close_task assert not first_result.is_error + assert close_result is None assert FakeTool.instances[0].close_count == 1 assert FakeTool.instances[1].close_count == 1 @@ -1005,7 +1006,8 @@ async def gated_connect(tool: FakeTool) -> None: assert len(FakeTool.instances) == 2 cancelled.cancel() with pytest.raises(asyncio.CancelledError): - await cancelled + cancelled_result = await cancelled + assert cancelled_result is None release.set() results = await asyncio.gather(first, second) @@ -1014,6 +1016,38 @@ async def gated_connect(tool: FakeTool) -> None: assert max_connecting == 2 assert not handler._inflight + @pytest.mark.asyncio + async def test_cancelled_waiter_does_not_cancel_shared_inflight_creation(self) -> None: + handler = DefaultMCPToolHandler() + connect_started = asyncio.Event() + release_connect = asyncio.Event() + original_connect = FakeTool.connect + + async def gated_connect(tool: FakeTool) -> None: + connect_started.set() + await release_connect.wait() + await original_connect(tool) + + with _patch_tool(), patch.object(FakeTool, "connect", gated_connect): + creator = asyncio.create_task(handler.invoke_tool(_invocation(workflow_session_id="workflow-a"))) + await connect_started.wait() + waiter = asyncio.create_task(handler.invoke_tool(_invocation(workflow_session_id="workflow-a"))) + await asyncio.sleep(0) + + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + waiter_result = await waiter + assert waiter_result is None + + release_connect.set() + creator_result = await creator + follow_up_result = await handler.invoke_tool(_invocation(workflow_session_id="workflow-a")) + + assert not creator_result.is_error + assert not follow_up_result.is_error + assert len(FakeTool.instances) == 1 + assert FakeTool.instances[0].connect_count == 1 + @pytest.mark.asyncio async def test_repeated_use_keeps_lru_alive(self) -> None: handler = DefaultMCPToolHandler(cache_max_size=2) From 218291558e2236e41c9b8244f9d4a9647bd65b05 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Wed, 23 Sep 2026 10:27:48 -0400 Subject: [PATCH 04/42] chore: linting --- .../_workflows/_executors_mcp.py | 6 +++--- .../_workflows/_mcp_handler.py | 2 +- .../declarative/tests/test_default_mcp_tool_handler.py | 10 +++++----- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_mcp.py b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_mcp.py index 3edd91c63d..fef6185939 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_mcp.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_mcp.py @@ -52,7 +52,7 @@ MCPToolHandler, MCPToolInvocation, MCPToolResult, - _get_or_create_workflow_session_id, + get_or_create_workflow_session_id, ) __all__ = [ @@ -241,7 +241,7 @@ async def handle_action( arguments=arguments, headers=headers, connection_name=connection_name, - workflow_session_id=_get_or_create_workflow_session_id(ctx.state), + workflow_session_id=get_or_create_workflow_session_id(ctx.state), ) if require_approval: await self._request_approval( @@ -348,7 +348,7 @@ async def handle_approval_response( arguments=original_request.arguments, headers=self._evaluate_headers(state, self._action_def.get("headers")), connection_name=getattr(original_request, "connection_name", None), - workflow_session_id=_get_or_create_workflow_session_id(ctx.state), + workflow_session_id=get_or_create_workflow_session_id(ctx.state), ) if invocation.headers or original_request.header_names: binding = getattr(original_request, "header_binding", None) diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py index e4f6330e41..cebf8a08a0 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -57,7 +57,7 @@ _WORKFLOW_SESSION_ID_KEY = "_declarative_mcp_workflow_session_id" -def _get_or_create_workflow_session_id(state: State) -> str: +def get_or_create_workflow_session_id(state: State) -> str: workflow_session_id = state.get(_WORKFLOW_SESSION_ID_KEY) if workflow_session_id is None: workflow_session_id = uuid.uuid4().hex diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index b55f8086ee..48cabc94c1 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -29,7 +29,7 @@ from agent_framework_declarative._workflows._mcp_handler import ( DefaultMCPToolHandler, MCPToolInvocation, - _get_or_create_workflow_session_id, + get_or_create_workflow_session_id, ) pytestmark = pytest.mark.skipif( @@ -831,16 +831,16 @@ def test_invalid_cache_size_raises(self) -> None: class TestWorkflowSessionId: def test_separate_workflow_states_get_separate_ids(self) -> None: - first = _get_or_create_workflow_session_id(State()) - second = _get_or_create_workflow_session_id(State()) + first = get_or_create_workflow_session_id(State()) + second = get_or_create_workflow_session_id(State()) assert first != second def test_same_workflow_state_reuses_id(self) -> None: state = State() - first = _get_or_create_workflow_session_id(state) - second = _get_or_create_workflow_session_id(state) + first = get_or_create_workflow_session_id(state) + second = get_or_create_workflow_session_id(state) assert first == second From f1dde52dc303521306caf6a385ef35d7ea6ccb84 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 23 Sep 2026 17:38:19 +0000 Subject: [PATCH 05/42] Address MCP review lifecycle comments Co-authored-by: baywet <7905502+baywet@users.noreply.github.com> --- .../DefaultMcpToolHandler.cs | 93 +++++++++++++++++-- .../Interpreter/DeclarativeWorkflowContext.cs | 2 +- .../DefaultMcpToolHandlerLifetimeTests.cs | 44 ++++++++- .../PowerFx/WorkflowFormulaStateTests.cs | 13 ++- .../_workflows/_mcp_handler.py | 12 ++- .../tests/test_default_mcp_tool_handler.py | 21 +++++ 6 files changed, 162 insertions(+), 23 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 340b8c09a2..f80b8728aa 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs @@ -55,6 +55,7 @@ public sealed class DefaultMcpToolHandler : IWorkflowScopedMcpToolHandler, IAsyn 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; @@ -101,6 +102,7 @@ internal DefaultMcpToolHandler( this._httpClientProvider = httpClientProvider; this._httpMessageHandlerFactory = Throw.IfNull(httpMessageHandlerFactory); this._clientCacheMaxSize = clientCacheMaxSize; + this._clientCreationSemaphore = new(clientCacheMaxSize, clientCacheMaxSize); } /// @@ -345,6 +347,7 @@ public async ValueTask DisposeAsync() } this._clientLock.Dispose(); + this._clientCreationSemaphore.Dispose(); } [System.Diagnostics.CodeAnalysis.SuppressMessage("Maintainability", "CA1513:Use ObjectDisposedException throw helper", @@ -397,20 +400,45 @@ private async Task AcquireClientAsync( if (!ownsClientCreation) { - await clientCreation.Task.ConfigureAwait(false); + await WaitForClientCreationAsync(clientCreation.Task, cancellationToken).ConfigureAwait(false); return await this.AcquireClientAsync( serverUrl, serverLabel, headers, connectionName, workflowSessionId, cancellationToken).ConfigureAwait(false); } ClientConnection? connection = null; + bool creationSemaphoreEntered = false; try { + 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, trimmedUrl, cancellationToken).ConfigureAwait(false); } catch (Exception exception) { - await this.CompleteClientCreationFailureAsync(clientCacheKey, clientCreation, exception).ConfigureAwait(false); + try + { + await this.CompleteClientCreationFailureAsync(clientCacheKey, clientCreation, exception).ConfigureAwait(false); + } + finally + { + if (creationSemaphoreEntered) + { + this._clientCreationSemaphore.Release(); + creationSemaphoreEntered = false; + } + } + throw; } @@ -423,7 +451,6 @@ private async Task AcquireClientAsync( if (this._disposing) { disposedException = new ObjectDisposedException(nameof(DefaultMcpToolHandler)); - clientCreation.TrySetException(disposedException); } else { @@ -458,22 +485,68 @@ private async Task AcquireClientAsync( this._clientLock.Release(); } - if (connection is not null) + try + { + if (connection is not null) + { + await connection.DisposeAsync().ConfigureAwait(false); + } + + if (disposedException is not null) + { + clientCreation.TrySetException(disposedException); + throw disposedException; + } + + if (clientToDispose is not null) + { + await this.DisposeCachedClientAsync(clientToDispose).ConfigureAwait(false); + } + } + catch (Exception exception) + { + if (disposedException is not null) + { + clientCreation.TrySetException(exception); + } + + if (result is not null) + { + await this.ReleaseClientAsync(result).ConfigureAwait(false); + } + + throw; + } + finally { - await connection.DisposeAsync().ConfigureAwait(false); + if (creationSemaphoreEntered) + { + this._clientCreationSemaphore.Release(); + } } - if (clientToDispose is not null) + 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 this.DisposeCachedClientAsync(clientToDispose).ConfigureAwait(false); + await clientCreation.ConfigureAwait(false); + return; } - if (disposedException is not null) + 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) { - throw disposedException; + cancellationToken.ThrowIfCancellationRequested(); } - return result ?? throw new InvalidOperationException("Failed to acquire MCP client."); + await clientCreation.ConfigureAwait(false); } private async Task CompleteClientCreationFailureAsync( 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 70c6ada3ab..f17265f78c 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowContext.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowContext.cs @@ -43,7 +43,7 @@ public static async ValueTask CreateAsync( : await source.ReadOrInitStateAsync( WorkflowSessionIdStateKey, static () => Guid.NewGuid().ToString("N"), - scopeName: null, + VariableScopeNames.System, cancellationToken: cancellationToken).ConfigureAwait(false); return new(source, state, sessionId); 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 7199cbef90..14a027420e 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 @@ -175,7 +175,7 @@ public async Task NoProvider_CacheEviction_DefersDisposalUntilActiveInvocationCo } [Fact] - public async Task NoProvider_ConcurrentWorkflowSessionCreations_DoNotSerializeHandshakeAsync() + public async Task NoProvider_ConcurrentWorkflowSessionCreations_AreBoundedByCacheSizeAsync() { // Arrange ProtocolStub stub = new(); @@ -195,12 +195,18 @@ public async Task NoProvider_ConcurrentWorkflowSessionCreations_DoNotSerializeHa await initializationsStarted.WaitAsync(timeout.Token); Task second = InvokeScopedAsync(handler, "workflow-b", "ping", timeout.Token); await initializationsStarted.WaitAsync(concurrencyTimeout.Token); + Task third = InvokeScopedAsync(handler, "workflow-c", "ping", timeout.Token); + using CancellationTokenSource gateTimeout = new(TimeSpan.FromMilliseconds(250)); + await Assert.ThrowsAnyAsync(() => initializationsStarted.WaitAsync(gateTimeout.Token)); releaseInitializations.Release(2); await Task.WhenAll(first, second); + await initializationsStarted.WaitAsync(timeout.Token); + releaseInitializations.Release(); + await third; // Assert - Assert.Equal(2, stub.Initializations); - Assert.Equal(0, stub.Terminations); + Assert.Equal(3, stub.Initializations); + Assert.Equal(1, stub.Terminations); } [Fact] @@ -301,6 +307,38 @@ await Task.WhenAll( 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 Provider_OperationFailure_DisposesSessionAndPreservesCallerClientAsync() { 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 4199aa96f2..1ba4cb1703 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 @@ -92,8 +92,8 @@ public void SetOverwritesExistingValue() public async Task DeclarativeContextFallbackSessionId_IsScopedToPersistedRunStateAsync() { // Arrange - Dictionary firstRunState = []; - Dictionary secondRunState = []; + 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); @@ -112,6 +112,8 @@ public async Task DeclarativeContextFallbackSessionId_IsScopedToPersistedRunStat 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] @@ -133,7 +135,7 @@ public async Task RestoreAsync_RestoresPersistedSensitivityAsync() Assert.Equal(SensitivityLevel.Sensitive, this.State.GetSensitivity("secret")); } - private static IWorkflowContext CreateContext(Dictionary state) + private static IWorkflowContext CreateContext(Dictionary<(string? ScopeName, string Key), string> state) { Mock context = new(); context @@ -144,10 +146,11 @@ private static IWorkflowContext CreateContext(Dictionary state) It.IsAny())) .Returns((string key, Func factory, string? scopeName, CancellationToken cancellationToken) => { - if (!state.TryGetValue(key, out string? value)) + var scopedKey = (scopeName, key); + if (!state.TryGetValue(scopedKey, out string? value)) { value = factory(); - state[key] = value; + state[scopedKey] = value; } return new ValueTask(value); diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py index cebf8a08a0..820bd19323 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -614,10 +614,14 @@ async def _get_or_create_entry(self, invocation: MCPToolInvocation) -> _CacheEnt inflight.set_exception(err) inflight.exception() raise err - if duplicate is not None: - await self._close_entry(duplicate) - if evicted is not None and evicted.disposal_claimed: - await self._close_claimed_entry(evicted) + try: + if duplicate is not None: + await self._close_entry(duplicate) + if evicted is not None and evicted.disposal_claimed: + await self._close_claimed_entry(evicted) + except BaseException: + await self._release_entry(entry) + raise return entry async def _release_entry(self, entry: _CacheEntry) -> None: diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index 48cabc94c1..877a759fde 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -975,6 +975,27 @@ async def gated_call(tool: FakeTool, tool_name: str, **arguments: Any) -> Any: assert FakeTool.instances[0].close_count == 1 assert FakeTool.instances[1].close_count == 1 + @pytest.mark.asyncio + async def test_lru_eviction_cleanup_cancellation_releases_new_entry(self) -> None: + handler = DefaultMCPToolHandler(cache_max_size=1) + cancel_first_close = True + original_close = FakeTool.close + + async def close(tool: FakeTool) -> None: + if cancel_first_close and tool.kwargs["url"] == "https://a/": + tool.close_count += 1 + raise asyncio.CancelledError + await original_close(tool) + + with _patch_tool(), patch.object(FakeTool, "close", close): + await handler.invoke_tool(_invocation(server_url="https://a/")) + with pytest.raises(asyncio.CancelledError): + await handler.invoke_tool(_invocation(server_url="https://b/")) + + assert all(entry.active_users == 0 for entry in handler._cache.values()) + cancel_first_close = False + await asyncio.wait_for(handler.aclose(), timeout=1) + @pytest.mark.asyncio async def test_entry_creation_is_bounded_and_cancelled_waiter_is_cleaned_up(self) -> None: handler = DefaultMCPToolHandler(cache_max_size=2) From 3293b653d43fda44738ee37097aceddc89956c93 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 23 Sep 2026 18:17:34 +0000 Subject: [PATCH 06/42] Address MCP session review feedback Co-authored-by: baywet <7905502+baywet@users.noreply.github.com> --- .../DefaultMcpToolHandler.cs | 160 +++++++++++++----- .../Interpreter/DeclarativeWorkflowContext.cs | 2 +- .../ObjectModel/InvokeMcpToolExecutor.cs | 18 +- .../DefaultMcpToolHandlerLifetimeTests.cs | 24 +++ .../ObjectModel/InvokeMcpToolExecutorTest.cs | 50 ++++++ .../_workflows/_mcp_handler.py | 10 +- .../tests/test_default_mcp_tool_handler.py | 41 +++++ 7 files changed, 246 insertions(+), 59 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 f80b8728aa..9ddc283c7a 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs @@ -53,7 +53,7 @@ public sealed class DefaultMcpToolHandler : IWorkflowScopedMcpToolHandler, IAsyn private readonly Dictionary<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash), TaskCompletionSource> _clientCreations = []; private readonly LinkedList<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash)> _clientLru = []; private readonly HashSet _retiredClients = []; - private readonly Dictionary _ownedHttpClients = []; + private readonly Dictionary _ownedHttpClients = []; private readonly SemaphoreSlim _clientLock = new(1, 1); private readonly SemaphoreSlim _clientCreationSemaphore; private readonly int _clientCacheMaxSize; @@ -162,7 +162,7 @@ public async Task InvokeToolInWorkflowSessionAsync( 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); @@ -329,23 +329,6 @@ public async ValueTask DisposeAsync() await Task.WhenAll(cachedClients.Select(client => client.Disposed.Task)).ConfigureAwait(false); - await this._clientLock.WaitAsync().ConfigureAwait(false); - try - { - - // Dispose only HttpClients that the handler created (not user-provided ones) - foreach (HttpClient httpClient in this._ownedHttpClients.Values) - { - httpClient.Dispose(); - } - - this._ownedHttpClients.Clear(); - } - finally - { - this._clientLock.Release(); - } - this._clientLock.Dispose(); this._clientCreationSemaphore.Dispose(); } @@ -422,7 +405,7 @@ private async Task AcquireClientAsync( } connection = await this.CreateClientAsync( - trimmedUrl, serverLabel, headers, trimmedUrl, cancellationToken).ConfigureAwait(false); + trimmedUrl, serverLabel, headers, cancellationToken).ConfigureAwait(false); } catch (Exception exception) { @@ -636,39 +619,26 @@ private async Task CreateClientAsync( string serverUrl, string? serverLabel, IDictionary? headers, - string? httpClientCacheKey, CancellationToken cancellationToken) { - // 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) + if (httpClient is null && this._httpClientProvider is not null) { httpClient = this.CreatePinnedHttpClient(serverUrl); ownsHttpClient = true; } - else if (httpClient is null && httpClientCacheKey is not null) + else if (httpClient is null) { - string cacheKey = httpClientCacheKey; - await this._clientLock.WaitAsync(cancellationToken).ConfigureAwait(false); - try - { - if (!this._ownedHttpClients.TryGetValue(cacheKey, out httpClient)) - { - httpClient = this.CreatePinnedHttpClient(serverUrl); - this._ownedHttpClients[cacheKey] = httpClient; - } - } - finally - { - this._clientLock.Release(); - } + ownedHttpClientLease = await this.AcquireOwnedHttpClientAsync(serverUrl, cancellationToken).ConfigureAwait(false); + httpClient = ownedHttpClientLease.Client; } HttpClientTransportOptions transportOptions = new() @@ -683,21 +653,80 @@ private async Task CreateClientAsync( TransportMode = HttpTransportMode.StreamableHttp }; - HttpClient resolvedHttpClient = httpClient ?? throw new InvalidOperationException("Failed to resolve MCP HTTP client."); - HttpClientTransport transport = new(transportOptions, resolvedHttpClient, 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 + { + if (transport is not null) + { + await DisposeResourceAsync(transport, "transport").ConfigureAwait(false); + } + } + finally + { + if (ownedHttpClientLease is not null) + { + await ownedHttpClientLease.DisposeAsync().ConfigureAwait(false); + } + } + 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(); + } + + clientToDispose?.Dispose(); + } + private HttpClient CreatePinnedHttpClient(string serverUrl) { // Pin credential headers to the configured server origin as defense-in-depth. Forcing @@ -744,7 +773,36 @@ private sealed class CachedClient( public bool DisposalClaimed { get; set; } } - internal sealed class ClientConnection(McpClient client, IAsyncDisposable transport) : IAsyncDisposable + 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; @@ -756,8 +814,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/Interpreter/DeclarativeWorkflowContext.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowContext.cs index f17265f78c..47962fe5e5 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowContext.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowContext.cs @@ -17,7 +17,7 @@ namespace Microsoft.Agents.AI.Workflows.Declarative.Interpreter; internal sealed class DeclarativeWorkflowContext : IWorkflowContext, IWorkflowSessionContext { - private const string WorkflowSessionIdStateKey = "__declarative_mcp_workflow_session_id"; + internal const string WorkflowSessionIdStateKey = "__declarative_mcp_workflow_session_id"; public static readonly FrozenSet ManagedScopes = [ 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 53fedcdb0d..3b09518ca5 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/InvokeMcpToolExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/InvokeMcpToolExecutor.cs @@ -31,8 +31,6 @@ internal sealed class InvokeMcpToolExecutor( { private const string ApprovalSnapshotStateKey = nameof(_approvalSnapshots); private const string LegacyApprovalSnapshotStateKey = "_approvalSnapshot"; - private readonly string _fallbackWorkflowSessionId = Guid.NewGuid().ToString("N"); - /// /// Snapshots of evaluated parameters captured at approval-request time, keyed by /// per-invocation request id. Each pending approval lives here until the matching @@ -187,7 +185,7 @@ public async ValueTask CaptureResponseAsync( await this.ProcessResultAsync(context, resultContent, cancellationToken).ConfigureAwait(false); } - private Task InvokeToolAsync( + private async Task InvokeToolAsync( IWorkflowContext context, string serverUrl, string? serverLabel, @@ -201,8 +199,12 @@ private Task InvokeToolAsync( { string workflowSessionId = context is IWorkflowSessionContext sessionContext ? sessionContext.SessionId - : this._fallbackWorkflowSessionId; - return scopedHandler.InvokeToolInWorkflowSessionAsync( + : await context.ReadOrInitStateAsync( + DeclarativeWorkflowContext.WorkflowSessionIdStateKey, + static () => Guid.NewGuid().ToString("N"), + VariableScopeNames.System, + cancellationToken).ConfigureAwait(false); + return await scopedHandler.InvokeToolInWorkflowSessionAsync( serverUrl, serverLabel, toolName, @@ -210,17 +212,17 @@ private Task InvokeToolAsync( headers, connectionName, workflowSessionId, - cancellationToken); + cancellationToken).ConfigureAwait(false); } - return mcpToolHandler.InvokeToolAsync( + return await mcpToolHandler.InvokeToolAsync( serverUrl, serverLabel, toolName, arguments, headers, connectionName, - cancellationToken); + cancellationToken).ConfigureAwait(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 14a027420e..8a8184f0a7 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 @@ -140,6 +140,30 @@ public async Task NoProvider_CacheEviction_DisposesLeastRecentlyUsedSessionAsync 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_DefersDisposalUntilActiveInvocationCompletesAsync() { 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 ded8316083..be97b96aa8 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 @@ -180,6 +180,34 @@ public async Task InvokeMcpToolSameWorkflowSessionReusesSessionScopeAsync() 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() { @@ -1603,6 +1631,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 diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py index 820bd19323..4c3cec2529 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -608,11 +608,13 @@ async def _get_or_create_entry(self, invocation: MCPToolInvocation) -> _CacheEnt # error so the caller (and any other awaiters) surface a # consistent "handler is closed" failure rather than receiving # an entry we are about to close behind their back. - await self._close_entry(entry) err = RuntimeError("DefaultMCPToolHandler is closed") - if not inflight.done(): - inflight.set_exception(err) - inflight.exception() + try: + await self._close_invocation_entry(entry) + finally: + if not inflight.done(): + inflight.set_exception(err) + inflight.exception() raise err try: if duplicate is not None: diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index 877a759fde..9e90d10dc5 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -1263,6 +1263,47 @@ async def gated_connect(self: FakeTool) -> None: assert result.is_error is True assert "closed" in (result.error_message or "").lower() + @pytest.mark.asyncio + async def test_cancelled_creator_does_not_block_aclose(self) -> None: + handler = DefaultMCPToolHandler() + connect_started = asyncio.Event() + release_connect = asyncio.Event() + close_started = asyncio.Event() + release_close = asyncio.Event() + original_connect = FakeTool.connect + original_close_entry = handler._close_entry + + async def gated_connect(self: FakeTool) -> None: + connect_started.set() + await release_connect.wait() + await original_connect(self) + + async def gated_close_entry(entry: Any) -> None: + close_started.set() + await release_close.wait() + await original_close_entry(entry) + + with ( + _patch_tool(), + patch.object(FakeTool, "connect", gated_connect), + patch.object(handler, "_close_entry", gated_close_entry), + ): + invoke_task = asyncio.create_task(handler.invoke_tool(_invocation(headers={"X": "1"}))) + await connect_started.wait() + close_task = asyncio.create_task(handler.aclose()) + await asyncio.sleep(0) + release_connect.set() + await close_started.wait() + invoke_task.cancel() + release_close.set() + + with pytest.raises(asyncio.CancelledError): + await invoke_task + await asyncio.wait_for(close_task, timeout=1) + + assert FakeTool.instances[0].close_count == 1 + assert not handler._inflight + # ---------- Result normalisation ------------------------------------------ From 2ce112b0346fe77f8f99b7b885877ed077fd057c Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 23 Sep 2026 19:23:23 +0000 Subject: [PATCH 07/42] Fix MCP creation cancellation and run isolation Co-authored-by: baywet <7905502+baywet@users.noreply.github.com> --- .../DefaultMcpToolHandler.cs | 11 ++- .../DefaultMcpToolHandlerLifetimeTests.cs | 36 +++++++++ python/packages/declarative/AGENTS.md | 4 +- .../_workflows/_executors_control_flow.py | 4 + .../_workflows/_mcp_handler.py | 81 +++++++++++++------ .../tests/test_default_mcp_tool_handler.py | 28 +++++++ .../tests/test_invoke_mcp_tool_executor.py | 21 ++++- 7 files changed, 156 insertions(+), 29 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 9ddc283c7a..bc31709a63 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs @@ -383,7 +383,16 @@ private async Task AcquireClientAsync( if (!ownsClientCreation) { - await WaitForClientCreationAsync(clientCreation.Task, cancellationToken).ConfigureAwait(false); + try + { + await WaitForClientCreationAsync(clientCreation.Task, cancellationToken).ConfigureAwait(false); + } + catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested) + { + return await this.AcquireClientAsync( + serverUrl, serverLabel, headers, connectionName, workflowSessionId, cancellationToken).ConfigureAwait(false); + } + return await this.AcquireClientAsync( serverUrl, serverLabel, headers, connectionName, workflowSessionId, cancellationToken).ConfigureAwait(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 8a8184f0a7..dbf2f9dc68 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 @@ -363,6 +363,42 @@ public async Task NoProvider_CancelledWaiter_DoesNotWaitForSharedCreationAsync() 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 Provider_OperationFailure_DisposesSessionAndPreservesCallerClientAsync() { diff --git a/python/packages/declarative/AGENTS.md b/python/packages/declarative/AGENTS.md index 332db4cec0..0679795d51 100644 --- a/python/packages/declarative/AGENTS.md +++ b/python/packages/declarative/AGENTS.md @@ -16,8 +16,8 @@ YAML/JSON-based declarative agent and workflow definitions. `DefaultMCPToolHandler` caches/coalesces sessions only without a `client_provider`. Cache identity includes a framework-owned workflow session ID in addition to -endpoint, label, connection, and headers, so separate workflow instances do not -share a stateful MCP protocol session while continuations of one workflow do. +endpoint, label, connection, and headers, so separate fresh runs do not share a +stateful MCP protocol session while continuations and checkpoint restores do. With a provider, every invocation (including `tools/list`) gets a fresh tool/session, even if the provider returns `None` or a shared HTTP client. Invocation cleanup closes the session and any internally owned fallback client, never caller-owned HTTP clients. diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_control_flow.py b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_control_flow.py index 5db554574f..a612f9aef0 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_control_flow.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_control_flow.py @@ -410,6 +410,10 @@ async def handle_action( ) -> None: """Simply pass through to continue the workflow.""" await self._ensure_state_initialized(ctx, trigger) + if self._action_def.get("kind") == "Entry": + from ._mcp_handler import reset_workflow_session_id + + reset_workflow_session_id(ctx.state) await ctx.send_message(ActionComplete()) diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py index 4c3cec2529..d5d3f5ab6b 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -67,6 +67,10 @@ def get_or_create_workflow_session_id(state: State) -> str: return workflow_session_id +def reset_workflow_session_id(state: State) -> None: + state.set(_WORKFLOW_SESSION_ID_KEY, uuid.uuid4().hex) + + @dataclass class MCPToolInvocation: """Description of an MCP tool call to be dispatched by a :class:`MCPToolHandler`. @@ -577,31 +581,35 @@ async def _get_or_create_entry(self, invocation: MCPToolInvocation) -> _CacheEnt evicted: _CacheEntry | None = None duplicate: _CacheEntry | None = None handler_closed = False - async with self._cache_lock: - self._inflight.pop(key, None) - if self._closed: - handler_closed = True - else: - existing = self._cache.get(key) - if existing is not None: - # Another writer beat us; prefer the existing entry and - # discard ours after the lock is released. - self._cache.move_to_end(key) - duplicate = entry - entry = existing - entry.active_users += 1 + try: + async with self._cache_lock: + self._inflight.pop(key, None) + if self._closed: + handler_closed = True else: - entry.active_users = 1 - self._cache[key] = entry - self._cache.move_to_end(key) - if len(self._cache) > self._cache_max_size: - _evicted_key, evicted = self._cache.popitem(last=False) - evicted.evicted = True - self._retired[id(evicted)] = evicted - if evicted.active_users == 0: - evicted.disposal_claimed = True - if not inflight.done(): - inflight.set_result(entry) + existing = self._cache.get(key) + if existing is not None: + # Another writer beat us; prefer the existing entry and + # discard ours after the lock is released. + self._cache.move_to_end(key) + duplicate = entry + entry = existing + entry.active_users += 1 + else: + entry.active_users = 1 + self._cache[key] = entry + self._cache.move_to_end(key) + if len(self._cache) > self._cache_max_size: + _evicted_key, evicted = self._cache.popitem(last=False) + evicted.evicted = True + self._retired[id(evicted)] = evicted + if evicted.active_users == 0: + evicted.disposal_claimed = True + if not inflight.done(): + inflight.set_result(entry) + except BaseException as exc: + await self._abort_entry_creation(key, inflight, entry, exc) + raise if handler_closed: # Close our orphaned entry; resolve the future with a clear @@ -626,6 +634,31 @@ async def _get_or_create_entry(self, invocation: MCPToolInvocation) -> _CacheEnt raise return entry + async def _abort_entry_creation( + self, + key: tuple[str, str, str, str, str], + inflight: asyncio.Future[_CacheEntry], + entry: _CacheEntry, + exc: BaseException, + ) -> None: + async def cleanup() -> None: + try: + await self._close_entry(entry) + finally: + async with self._cache_lock: + self._inflight.pop(key, None) + if not inflight.done(): + inflight.set_exception(exc) + inflight.exception() + + cleanup_task = asyncio.create_task(cleanup()) + while not cleanup_task.done(): + try: + await asyncio.shield(cleanup_task) + except asyncio.CancelledError: + continue + cleanup_task.result() + async def _release_entry(self, entry: _CacheEntry) -> None: close_entry = False async with self._cache_lock: diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index 9e90d10dc5..494cb8e393 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -1304,6 +1304,34 @@ async def gated_close_entry(entry: Any) -> None: assert FakeTool.instances[0].close_count == 1 assert not handler._inflight + @pytest.mark.asyncio + async def test_cancellation_while_waiting_for_phase_three_lock_cleans_up(self) -> None: + handler = DefaultMCPToolHandler() + entry_created = asyncio.Event() + release_creation = asyncio.Event() + original_create_entry = handler._create_entry + + async def gated_create_entry(invocation: MCPToolInvocation) -> Any: + entry = await original_create_entry(invocation) + entry_created.set() + await release_creation.wait() + return entry + + with _patch_tool(), patch.object(handler, "_create_entry", gated_create_entry): + invocation = asyncio.create_task(handler.invoke_tool(_invocation(headers={"X": "1"}))) + await entry_created.wait() + await handler._cache_lock.acquire() + release_creation.set() + await asyncio.sleep(0) + invocation.cancel() + handler._cache_lock.release() + + with pytest.raises(asyncio.CancelledError): + await invocation + + assert FakeTool.instances[0].close_count == 1 + assert not handler._inflight + # ---------- Result normalisation ------------------------------------------ diff --git a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py index c0896c4c6f..3b450f4a34 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -182,7 +182,7 @@ async def test_separate_workflows_receive_separate_session_ids(self) -> None: assert handler.invocations[0].workflow_session_id != handler.invocations[1].workflow_session_id @pytest.mark.asyncio - async def test_same_workflow_continuation_reuses_session_id(self) -> None: + async def test_fresh_runs_on_same_workflow_receive_separate_session_ids(self) -> None: handler = StubMcpHandler(_ok()) factory = WorkflowFactory(mcp_tool_handler=handler) workflow = factory.create_workflow_from_definition(_yaml(_action())) @@ -192,7 +192,24 @@ async def test_same_workflow_continuation_reuses_session_id(self) -> None: assert len(handler.invocations) == 2 assert handler.invocations[0].workflow_session_id - assert handler.invocations[0].workflow_session_id == handler.invocations[1].workflow_session_id + assert handler.invocations[0].workflow_session_id != handler.invocations[1].workflow_session_id + + @pytest.mark.asyncio + async def test_continuation_reuses_workflow_session_id(self) -> None: + from agent_framework_declarative._workflows import ToolApprovalResponse + from agent_framework_declarative._workflows._mcp_handler import get_or_create_workflow_session_id + + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + workflow = factory.create_workflow_from_definition(_yaml(_action(require_approval=True))) + + paused = await workflow.run({}) + [approval] = paused.get_request_info_events() + workflow_session_id = get_or_create_workflow_session_id(workflow._runner.state) # pyright: ignore[reportPrivateUsage] + await workflow.run(responses={approval.request_id: ToolApprovalResponse(approved=True)}) + + assert handler.last_invocation is not None + assert handler.last_invocation.workflow_session_id == workflow_session_id @pytest.mark.asyncio async def test_arguments_evaluated_and_preserves_none(self) -> None: From bda47d1892552567d174cc9023c94ff9351c3381 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Mon, 28 Sep 2026 10:23:35 -0400 Subject: [PATCH 08/42] fix: retry MCP waiter after creator cancellation Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 166199ae-915e-404b-837a-722c92fb28d4 --- .../_workflows/_mcp_handler.py | 7 +++- .../tests/test_default_mcp_tool_handler.py | 37 ++++++++++++++++++- 2 files changed, 41 insertions(+), 3 deletions(-) diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py index d5d3f5ab6b..2658a5f867 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -552,7 +552,12 @@ async def _get_or_create_entry(self, invocation: MCPToolInvocation) -> _CacheEnt creating = True if not creating: - _ = await asyncio.shield(inflight) + try: + _ = await asyncio.shield(inflight) + except asyncio.CancelledError: + current_task = asyncio.current_task() + if current_task is not None and current_task.cancelling(): + raise return await self._get_or_create_entry(invocation) # Phase 2: we own creation. Build the entry outside the lock. diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index 494cb8e393..017d127ee4 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -1069,6 +1069,39 @@ async def gated_connect(tool: FakeTool) -> None: assert len(FakeTool.instances) == 1 assert FakeTool.instances[0].connect_count == 1 + @pytest.mark.asyncio + async def test_cancelled_creator_does_not_cancel_shared_inflight_waiter(self) -> None: + handler = DefaultMCPToolHandler() + first_connect_started = asyncio.Event() + release_first_connect = asyncio.Event() + original_connect = FakeTool.connect + + async def gated_connect(tool: FakeTool) -> None: + if len(FakeTool.instances) == 1: + first_connect_started.set() + await release_first_connect.wait() + await original_connect(tool) + + with _patch_tool(), patch.object(FakeTool, "connect", gated_connect): + creator = asyncio.create_task(handler.invoke_tool(_invocation(workflow_session_id="workflow-a"))) + await first_connect_started.wait() + waiter = asyncio.create_task(handler.invoke_tool(_invocation(workflow_session_id="workflow-a"))) + await asyncio.sleep(0) + + creator.cancel() + release_first_connect.set() + with pytest.raises(asyncio.CancelledError): + _ = await creator + + waiter_result = await waiter + follow_up_result = await handler.invoke_tool(_invocation(workflow_session_id="workflow-a")) + + assert not waiter_result.is_error + assert not follow_up_result.is_error + assert len(FakeTool.instances) == 2 + assert FakeTool.instances[0].close_count == 1 + assert FakeTool.instances[1].connect_count == 1 + @pytest.mark.asyncio async def test_repeated_use_keeps_lru_alive(self) -> None: handler = DefaultMCPToolHandler(cache_max_size=2) @@ -1298,7 +1331,7 @@ async def gated_close_entry(entry: Any) -> None: release_close.set() with pytest.raises(asyncio.CancelledError): - await invoke_task + _ = await invoke_task await asyncio.wait_for(close_task, timeout=1) assert FakeTool.instances[0].close_count == 1 @@ -1327,7 +1360,7 @@ async def gated_create_entry(invocation: MCPToolInvocation) -> Any: handler._cache_lock.release() with pytest.raises(asyncio.CancelledError): - await invocation + _ = await invocation assert FakeTool.instances[0].close_count == 1 assert not handler._inflight From f5c133f9fb2c20f91c7e3304d14fe2a5353dda51 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Mon, 28 Sep 2026 12:11:27 -0400 Subject: [PATCH 09/42] chore: Use generated session ID for workflow session context Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- .../Interpreter/DeclarativeWorkflowContext.cs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 47962fe5e5..f88035abd1 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowContext.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowContext.cs @@ -38,11 +38,12 @@ public static async ValueTask CreateAsync( WorkflowFormulaState state, CancellationToken cancellationToken = default) { + string generatedSessionId = Guid.NewGuid().ToString("N"); string sessionId = source is IWorkflowSessionContext sessionContext ? sessionContext.SessionId : await source.ReadOrInitStateAsync( WorkflowSessionIdStateKey, - static () => Guid.NewGuid().ToString("N"), + () => generatedSessionId, VariableScopeNames.System, cancellationToken: cancellationToken).ConfigureAwait(false); From 5f21a68eec337c2fd11b94f4925aef79df3c6c46 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Mon, 28 Sep 2026 12:11:42 -0400 Subject: [PATCH 10/42] chore: Use generated workflow session ID for context Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- .../ObjectModel/InvokeMcpToolExecutor.cs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 3b09518ca5..ed74d30d39 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/InvokeMcpToolExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/InvokeMcpToolExecutor.cs @@ -197,11 +197,12 @@ private async Task InvokeToolAsync( { if (mcpToolHandler is IWorkflowScopedMcpToolHandler scopedHandler) { + string generatedWorkflowSessionId = Guid.NewGuid().ToString("N"); string workflowSessionId = context is IWorkflowSessionContext sessionContext ? sessionContext.SessionId : await context.ReadOrInitStateAsync( DeclarativeWorkflowContext.WorkflowSessionIdStateKey, - static () => Guid.NewGuid().ToString("N"), + () => generatedWorkflowSessionId, VariableScopeNames.System, cancellationToken).ConfigureAwait(false); return await scopedHandler.InvokeToolInWorkflowSessionAsync( From 97c636dd713cd9458352d19f53cbb867a743378f Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Mon, 28 Sep 2026 15:41:49 -0400 Subject: [PATCH 11/42] fix: make MCP shutdown cancellation-safe Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 166199ae-915e-404b-837a-722c92fb28d4 --- .../_workflows/_mcp_handler.py | 50 ++++++++++++------- .../tests/test_default_mcp_tool_handler.py | 29 +++++++++++ 2 files changed, 60 insertions(+), 19 deletions(-) diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py index 2658a5f867..757f1964ae 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -277,6 +277,7 @@ def __init__( # Set by ``aclose`` to prevent post-close cache insertions and to # reject new ``invoke_tool`` calls. Once set, never cleared. self._closed = False + self._shutdown_task: asyncio.Task[None] | None = None async def invoke_tool(self, invocation: MCPToolInvocation) -> MCPToolResult: """Invoke ``invocation.tool_name`` on an MCP client for the server. @@ -461,7 +462,7 @@ async def aclose(self) -> None: are rejected after connecting and cleaned up by their invocation. Concurrent shutdown calls wait for those invocations as well. - Idempotent. In no-provider mode, a second call returns immediately. + Idempotent. Concurrent and subsequent calls await the same shutdown task. Drains any in-flight ``_create_entry`` tasks before returning so their resources are cleaned up; the in-flight tasks see ``self._closed`` in phase 3 of :meth:`_get_or_create_entry`, close their own entry, and resolve @@ -477,26 +478,37 @@ async def aclose(self) -> None: "DefaultMCPToolHandler.aclose() cannot be called from an active provider-backed invocation" ) async with self._cache_lock: - if self._closed and self._client_provider is None: - return - self._closed = True - entries = list(self._cache.values()) - entry_ids = {id(entry) for entry in entries} - entries.extend(entry for entry in self._retired.values() if id(entry) not in entry_ids) - self._cache.clear() - entries_to_close: list[_CacheEntry] = [] - for entry in entries: - entry.evicted = True - self._retired[id(entry)] = entry - if entry.active_users == 0 and not entry.disposal_claimed: - entry.disposal_claimed = True - entries_to_close.append(entry) - inflight_futures = list(self._inflight.values()) - active_invocations = list(self._active_invocations) + if self._shutdown_task is None: + self._closed = True + entries = list(self._cache.values()) + entry_ids = {id(entry) for entry in entries} + entries.extend(entry for entry in self._retired.values() if id(entry) not in entry_ids) + self._cache.clear() + entries_to_close: list[_CacheEntry] = [] + for entry in entries: + entry.evicted = True + self._retired[id(entry)] = entry + if entry.active_users == 0 and not entry.disposal_claimed: + entry.disposal_claimed = True + entries_to_close.append(entry) + inflight_futures = list(self._inflight.values()) + active_invocations = list(self._active_invocations) + self._shutdown_task = asyncio.create_task( + self._drain_shutdown(entries, entries_to_close, inflight_futures, active_invocations) + ) + shutdown_task = self._shutdown_task + + await asyncio.shield(shutdown_task) + async def _drain_shutdown( + self, + entries: list[_CacheEntry], + entries_to_close: list[_CacheEntry], + inflight_futures: list[asyncio.Future[_CacheEntry]], + active_invocations: list[asyncio.Future[None]], + ) -> None: for completion in active_invocations: - # Cancelling shutdown must not cancel an invocation's completion signal. - await asyncio.shield(completion) + await completion # Wait for in-flight creations to finish their self-cleanup. Each # in-flight task self-closes its entry under the closed-flag branch diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index 017d127ee4..751fe0587b 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -1246,6 +1246,35 @@ async def test_aclose_is_idempotent(self) -> None: await handler.aclose() assert FakeTool.instances[0].close_count == 1 + @pytest.mark.asyncio + async def test_cancelled_aclose_continues_draining_cached_entries(self) -> None: + handler = DefaultMCPToolHandler() + first_close_started = asyncio.Event() + release_first_close = asyncio.Event() + original_close = FakeTool.close + + async def gated_close(tool: FakeTool) -> None: + if tool.kwargs["url"] == "https://a/": + first_close_started.set() + await release_first_close.wait() + await original_close(tool) + + with _patch_tool(), patch.object(FakeTool, "close", gated_close): + await handler.invoke_tool(_invocation(server_url="https://a/", headers={"X": "1"})) + await handler.invoke_tool(_invocation(server_url="https://b/", headers={"X": "1"})) + + shutdown = asyncio.create_task(handler.aclose()) + await first_close_started.wait() + shutdown.cancel() + with pytest.raises(asyncio.CancelledError): + _ = await shutdown + + release_first_close.set() + await asyncio.wait_for(handler.aclose(), timeout=1) + + assert [tool.close_count for tool in FakeTool.instances] == [1, 1] + assert all(tool._httpx_client is not None and tool._httpx_client.is_closed for tool in FakeTool.instances) + @pytest.mark.asyncio async def test_invoke_after_close_returns_error_result(self) -> None: """Post-close ``invoke_tool`` surfaces a tool error rather than crashing.""" From bfcd02e46cbd230580dd661697ee881a08866f3a Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 29 Sep 2026 08:35:57 -0400 Subject: [PATCH 12/42] chore: Potential fix for pull request finding 'Statement has no effect' Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com> --- .../_workflows/_mcp_handler.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py index 757f1964ae..5fb5444a39 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -507,8 +507,16 @@ async def _drain_shutdown( inflight_futures: list[asyncio.Future[_CacheEntry]], active_invocations: list[asyncio.Future[None]], ) -> None: - for completion in active_invocations: - await completion + if active_invocations: + invocation_results = await asyncio.gather(*active_invocations, return_exceptions=True) + for result in invocation_results: + if isinstance(result, BaseException): + if isinstance(result, (KeyboardInterrupt, SystemExit, GeneratorExit)): + raise result + logger.debug( + "DefaultMCPToolHandler: active invocation raised during aclose", + exc_info=(type(result), result, result.__traceback__), + ) # Wait for in-flight creations to finish their self-cleanup. Each # in-flight task self-closes its entry under the closed-flag branch From d8d8c61b7b1eb3b2b45c71b28c37c76de4816e9e Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 29 Sep 2026 08:42:39 -0400 Subject: [PATCH 13/42] fix: support MCP cancellation on Python 3.10 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 936efea1-d581-4dea-813f-1edcfd2ea3bf --- .../_workflows/_mcp_handler.py | 21 ++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py index 5fb5444a39..c2697596c2 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -33,6 +33,7 @@ import uuid from collections import OrderedDict from collections.abc import Awaitable, Callable +from contextlib import suppress from contextvars import ContextVar, Token from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, ClassVar, Protocol, cast, runtime_checkable @@ -180,6 +181,10 @@ class _CacheEntry: closed: asyncio.Event = field(default_factory=asyncio.Event) +class _EntryCreationCancelled(Exception): + """Signal waiters to retry after the task creating their entry was cancelled.""" + + class DefaultMCPToolHandler: """Default :class:`MCPToolHandler` backed by :class:`agent_framework.MCPStreamableHTTPTool`. @@ -572,12 +577,8 @@ async def _get_or_create_entry(self, invocation: MCPToolInvocation) -> _CacheEnt creating = True if not creating: - try: + with suppress(_EntryCreationCancelled): _ = await asyncio.shield(inflight) - except asyncio.CancelledError: - current_task = asyncio.current_task() - if current_task is not None and current_task.cancelling(): - raise return await self._get_or_create_entry(invocation) # Phase 2: we own creation. Build the entry outside the lock. @@ -591,7 +592,7 @@ async def _get_or_create_entry(self, invocation: MCPToolInvocation) -> _CacheEnt async with self._cache_lock: self._inflight.pop(key, None) if not inflight.done(): - inflight.set_exception(exc if isinstance(exc, BaseException) else RuntimeError(str(exc))) + inflight.set_exception(self._entry_creation_exception(exc)) # Mark the exception retrieved to suppress noisy "Future exception # was never retrieved" warnings when there are no other awaiters # (other awaiters still see the exception through their ``await``). @@ -673,7 +674,7 @@ async def cleanup() -> None: async with self._cache_lock: self._inflight.pop(key, None) if not inflight.done(): - inflight.set_exception(exc) + inflight.set_exception(self._entry_creation_exception(exc)) inflight.exception() cleanup_task = asyncio.create_task(cleanup()) @@ -684,6 +685,12 @@ async def cleanup() -> None: continue cleanup_task.result() + @staticmethod + def _entry_creation_exception(exc: BaseException) -> BaseException: + if isinstance(exc, asyncio.CancelledError): + return _EntryCreationCancelled() + return exc + async def _release_entry(self, entry: _CacheEntry) -> None: close_entry = False async with self._cache_lock: From 46abaeeea07c6cd3bc1b95592807da7f12ee3ff1 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 29 Sep 2026 09:01:05 -0400 Subject: [PATCH 14/42] test: avoid assigning None from MCP shutdown Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 936efea1-d581-4dea-813f-1edcfd2ea3bf --- .../declarative/tests/test_default_mcp_tool_handler.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index 751fe0587b..12f43875d3 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -968,10 +968,9 @@ async def gated_call(tool: FakeTool, tool_name: str, **arguments: Any) -> Any: release_first.set() first_result = await first - close_result = await close_task + await close_task assert not first_result.is_error - assert close_result is None assert FakeTool.instances[0].close_count == 1 assert FakeTool.instances[1].close_count == 1 @@ -1267,7 +1266,7 @@ async def gated_close(tool: FakeTool) -> None: await first_close_started.wait() shutdown.cancel() with pytest.raises(asyncio.CancelledError): - _ = await shutdown + await shutdown release_first_close.set() await asyncio.wait_for(handler.aclose(), timeout=1) From b470205bda929e01991b5ed462eba4260cbed42f Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 29 Sep 2026 09:03:50 -0400 Subject: [PATCH 15/42] fix: observe failed MCP client creation tasks Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 936efea1-d581-4dea-813f-1edcfd2ea3bf --- .../DefaultMcpToolHandler.cs | 1 + 1 file changed, 1 insertion(+) 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 2fee60b8c6..f335a1fbc8 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs @@ -556,6 +556,7 @@ private async Task CompleteClientCreationFailureAsync( } clientCreation.TrySetException(exception); + _ = clientCreation.Task.Exception; } finally { From f47001bc7b18cd3b42c4416bd24ee7199281dcc8 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 29 Sep 2026 09:19:51 -0400 Subject: [PATCH 16/42] chore: Potential fix for pull request finding 'Statement has no effect' Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com> --- .../packages/declarative/tests/test_default_mcp_tool_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index 12f43875d3..e8c8bbae9a 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -1266,7 +1266,7 @@ async def gated_close(tool: FakeTool) -> None: await first_close_started.wait() shutdown.cancel() with pytest.raises(asyncio.CancelledError): - await shutdown + await asyncio.gather(shutdown) release_first_close.set() await asyncio.wait_for(handler.aclose(), timeout=1) From 80988a00b07cf23fbea1624c40be250b9dc27e54 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 29 Sep 2026 09:22:38 -0400 Subject: [PATCH 17/42] test: await MCP shutdown through gather Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 936efea1-d581-4dea-813f-1edcfd2ea3bf --- .../packages/declarative/tests/test_default_mcp_tool_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index e8c8bbae9a..ddc87f0ea0 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -968,7 +968,7 @@ async def gated_call(tool: FakeTool, tool_name: str, **arguments: Any) -> Any: release_first.set() first_result = await first - await close_task + await asyncio.gather(close_task) assert not first_result.is_error assert FakeTool.instances[0].close_count == 1 From 9548e5bf9e9862cac6cad8a6f7c3b4db11976408 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 29 Sep 2026 10:16:55 -0400 Subject: [PATCH 18/42] Fix MCP shutdown cancellation draining Ensure cleanup-originated cancellation is surfaced only after all claimed MCP sessions and completion signals have drained in both implementations. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: aafca87b-cf2e-4c72-b5cb-ffcad6e1b1e0 --- .../DefaultMcpToolHandler.cs | 33 ++++++++++++++--- .../DefaultMcpToolHandlerLifetimeTests.cs | 36 +++++++++++++++++++ .../_workflows/_mcp_handler.py | 9 +++-- .../tests/test_default_mcp_tool_handler.py | 23 ++++++++++++ 4 files changed, 94 insertions(+), 7 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 f335a1fbc8..772f6ea9b1 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; @@ -322,15 +323,37 @@ public async ValueTask DisposeAsync() this._clientLock.Release(); } - foreach (CachedClient client in clientsToDispose) + try { - await this.DisposeCachedClientAsync(client).ConfigureAwait(false); + await DrainCleanupAsync( + clientsToDispose.Select(client => this.DisposeCachedClientAsync(client)), + cachedClients.Select(client => client.Disposed.Task)).ConfigureAwait(false); + } + finally + { + this._clientLock.Dispose(); + this._clientCreationSemaphore.Dispose(); + } + } + + internal static async Task DrainCleanupAsync(IEnumerable cleanupTasks, IEnumerable completionTasks) + { + Exception? cleanupException = null; + try + { + await Task.WhenAll(cleanupTasks).ConfigureAwait(false); + } + catch (Exception exception) + { + cleanupException = exception; } - await Task.WhenAll(cachedClients.Select(client => client.Disposed.Task)).ConfigureAwait(false); + await Task.WhenAll(completionTasks).ConfigureAwait(false); - this._clientLock.Dispose(); - this._clientCreationSemaphore.Dispose(); + if (cleanupException is not null) + { + ExceptionDispatchInfo.Capture(cleanupException).Throw(); + } } [System.Diagnostics.CodeAnalysis.SuppressMessage("Maintainability", "CA1513:Use ObjectDisposedException throw helper", 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 dbf2f9dc68..a57f901769 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 @@ -122,6 +122,42 @@ public async Task NoProvider_SeparateWorkflowSessions_UseSeparateCachedSessionsA 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(firstDisposed.Task.IsCompletedSuccessfully); + Assert.True(secondDisposed.Task.IsCompletedSuccessfully); + } + [Fact] public async Task NoProvider_CacheEviction_DisposesLeastRecentlyUsedSessionAsync() { diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py index c2697596c2..f34406f735 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -534,10 +534,15 @@ async def _drain_shutdown( logger.debug("DefaultMCPToolHandler: in-flight future raised during aclose", exc_info=True) continue - for entry in entries_to_close: - await self._close_claimed_entry(entry) + close_results = await asyncio.gather( + *(self._close_claimed_entry(entry) for entry in entries_to_close), + return_exceptions=True, + ) for entry in entries: await entry.closed.wait() + for result in close_results: + if isinstance(result, BaseException): + raise result async def __aenter__(self) -> DefaultMCPToolHandler: return self diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index ddc87f0ea0..ffee74c68a 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -1274,6 +1274,29 @@ async def gated_close(tool: FakeTool) -> None: assert [tool.close_count for tool in FakeTool.instances] == [1, 1] assert all(tool._httpx_client is not None and tool._httpx_client.is_closed for tool in FakeTool.instances) + @pytest.mark.asyncio + async def test_cleanup_cancellation_continues_draining_cached_entries(self) -> None: + handler = DefaultMCPToolHandler() + original_close = FakeTool.close + + async def cancel_first_close(tool: FakeTool) -> None: + if tool.kwargs["url"] == "https://a/": + tool.close_count += 1 + raise asyncio.CancelledError + await original_close(tool) + + with _patch_tool(), patch.object(FakeTool, "close", cancel_first_close): + await handler.invoke_tool(_invocation(server_url="https://a/", headers={"X": "1"})) + await handler.invoke_tool(_invocation(server_url="https://b/", headers={"X": "1"})) + + with pytest.raises(asyncio.CancelledError): + await handler.aclose() + with pytest.raises(asyncio.CancelledError): + await handler.aclose() + + assert [tool.close_count for tool in FakeTool.instances] == [1, 1] + assert all(tool._httpx_client is not None and tool._httpx_client.is_closed for tool in FakeTool.instances) + @pytest.mark.asyncio async def test_invoke_after_close_returns_error_result(self) -> None: """Post-close ``invoke_tool`` surfaces a tool error rather than crashing.""" From af94e4fdec99a6526cc2b113648f0a4178e71859 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 29 Sep 2026 10:29:57 -0400 Subject: [PATCH 19/42] Fix net472 MCP lifetime test build Use awaited task results instead of Task.IsCompletedSuccessfully so the regression test compiles on .NET Framework. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: aafca87b-cf2e-4c72-b5cb-ffcad6e1b1e0 --- .../DefaultMcpToolHandlerLifetimeTests.cs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 a57f901769..6b27a3dce2 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 @@ -154,8 +154,8 @@ await Assert.ThrowsAnyAsync(() => DefaultMcpToolHand [firstDisposed.Task, secondDisposed.Task])); // Assert - Assert.True(firstDisposed.Task.IsCompletedSuccessfully); - Assert.True(secondDisposed.Task.IsCompletedSuccessfully); + Assert.True(await firstDisposed.Task); + Assert.True(await secondDisposed.Task); } [Fact] From dd617dc66da1b1597dc230f42c7e1fa68d7c5b8e Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 29 Sep 2026 10:36:37 -0400 Subject: [PATCH 20/42] Wait for MCP client creators during disposal Track each cache client creator through its final semaphore release so handler disposal cannot tear down synchronization resources while creation is still exiting. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: aafca87b-cf2e-4c72-b5cb-ffcad6e1b1e0 --- .../DefaultMcpToolHandler.cs | 44 ++++++++++++------- .../DefaultMcpToolHandlerLifetimeTests.cs | 30 +++++++++++++ 2 files changed, 59 insertions(+), 15 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 772f6ea9b1..f6b33e3557 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs @@ -52,6 +52,7 @@ public sealed class DefaultMcpToolHandler : IWorkflowScopedMcpToolHandler, IAsyn private readonly Func _httpMessageHandlerFactory; 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 = []; @@ -263,14 +264,14 @@ public async ValueTask DisposeAsync() } Task? providerInvocations; - List> clientCreations; + List> clientCreationLifetimes; await this._clientLock.WaitAsync().ConfigureAwait(false); try { this.ThrowIfDisposing(); this._disposing = true; providerInvocations = this._providerInvocationsDrained?.Task; - clientCreations = this._clientCreations.Values.Select(source => source.Task).ToList(); + clientCreationLifetimes = this._clientCreationLifetimes.Select(source => source.Task).ToList(); } finally { @@ -282,20 +283,9 @@ public async ValueTask DisposeAsync() await providerInvocations.ConfigureAwait(false); } - foreach (Task clientCreation in clientCreations) + foreach (Task clientCreationLifetime in clientCreationLifetimes) { - try - { - await clientCreation.ConfigureAwait(false); - } - catch (OperationCanceledException) - { - // The invocation that owned this creation was cancelled; disposal only needs its cleanup to complete. - } - catch (Exception exception) - { - Trace.TraceWarning("MCP client creation failed during disposal: {0}", exception); - } + await clientCreationLifetime.ConfigureAwait(false); } List cachedClients; @@ -378,6 +368,7 @@ private async Task AcquireClientAsync( 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); @@ -396,6 +387,8 @@ private async Task AcquireClientAsync( { clientCreation = new(TaskCreationOptions.RunContinuationsAsynchronously); this._clientCreations[clientCacheKey] = clientCreation; + clientCreationLifetime = new(TaskCreationOptions.RunContinuationsAsynchronously); + this._clientCreationLifetimes.Add(clientCreationLifetime); ownsClientCreation = true; } } @@ -420,6 +413,8 @@ private async Task 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; try @@ -452,6 +447,8 @@ private async Task AcquireClientAsync( this._clientCreationSemaphore.Release(); creationSemaphoreEntered = false; } + + await this.CompleteClientCreationLifetimeAsync(ownedClientCreationLifetime).ConfigureAwait(false); } throw; @@ -538,6 +535,8 @@ private async Task AcquireClientAsync( { this._clientCreationSemaphore.Release(); } + + await this.CompleteClientCreationLifetimeAsync(ownedClientCreationLifetime).ConfigureAwait(false); } return result ?? throw new InvalidOperationException("Failed to acquire MCP client."); @@ -587,6 +586,21 @@ private async Task CompleteClientCreationFailureAsync( } } + private async Task CompleteClientCreationLifetimeAsync(TaskCompletionSource clientCreationLifetime) + { + await this._clientLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); + try + { + this._clientCreationLifetimes.Remove(clientCreationLifetime); + } + finally + { + this._clientLock.Release(); + } + + clientCreationLifetime.TrySetResult(true); + } + private async ValueTask ReleaseClientAsync(CachedClient client) { bool dispose = 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 6b27a3dce2..1e31375357 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 @@ -269,6 +269,36 @@ public async Task NoProvider_ConcurrentWorkflowSessionCreations_AreBoundedByCach 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_DifferentConnectionNames_UseSeparateCachedSessionsAsync() { From f33b4c6251bfeb58bfdd81994e4af803a850a894 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 29 Sep 2026 10:44:02 -0400 Subject: [PATCH 21/42] chore: Potential fix for pull request finding 'Generic catch clause' Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com> --- .../DefaultMcpToolHandler.cs | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) 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 f6b33e3557..976d2b9981 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs @@ -333,7 +333,7 @@ internal static async Task DrainCleanupAsync(IEnumerable cleanupTasks, IEn { await Task.WhenAll(cleanupTasks).ConfigureAwait(false); } - catch (Exception exception) + catch (Exception exception) when (!IsFatalException(exception)) { cleanupException = exception; } @@ -346,6 +346,16 @@ internal static async Task DrainCleanupAsync(IEnumerable cleanupTasks, IEn } } + 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() From 18720f0570b5c204aad3d9b0e0870c253ce1adc2 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 29 Sep 2026 12:57:34 -0400 Subject: [PATCH 22/42] Harden nested MCP session lifetimes Give subworkflows stable hierarchical cache scopes and ensure Python cached-entry leases release fully despite repeated caller cancellation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: aafca87b-cf2e-4c72-b5cb-ffcad6e1b1e0 --- .../SubworkflowBinding.cs | 16 ++++++++++- .../RepresentationTests.cs | 20 ++++++++++++++ .../_workflows/_mcp_handler.py | 12 +++++++++ .../tests/test_default_mcp_tool_handler.py | 27 +++++++++++++++++++ 4 files changed, 74 insertions(+), 1 deletion(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/SubworkflowBinding.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/SubworkflowBinding.cs index 11f7cf493c..538e8bc3bb 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; @@ -31,10 +32,23 @@ async ValueTask InitHostExecutorAsync(string sessionId) { ProtocolDescriptor workflowProtocol = await workflow.DescribeProtocolAsync().ConfigureAwait(false); - return new WorkflowHostExecutor(id, workflow, workflowProtocol, sessionId, ownershipToken, options); + return new WorkflowHostExecutor( + id, + workflow, + workflowProtocol, + CreateSubworkflowSessionId(sessionId, id), + ownershipToken, + options); } } + 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.UnitTests/RepresentationTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/RepresentationTests.cs index 1e92d130de..6b8a0f7996 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/RepresentationTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/RepresentationTests.cs @@ -91,6 +91,26 @@ 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); + } + private static string Source(int id) => $"Source/{id}"; private static string Sink(int id) => $"Sink/{id}"; diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py index f34406f735..a0d84d8843 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -697,6 +697,18 @@ def _entry_creation_exception(exc: BaseException) -> BaseException: return exc async def _release_entry(self, entry: _CacheEntry) -> None: + release = asyncio.create_task(self._release_entry_core(entry)) + cancelled = False + while not release.done(): + try: + await asyncio.shield(release) + except asyncio.CancelledError: + cancelled = True + release.result() + if cancelled: + raise asyncio.CancelledError + + async def _release_entry_core(self, entry: _CacheEntry) -> None: close_entry = False async with self._cache_lock: entry.active_users -= 1 diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index ffee74c68a..770801afd8 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -995,6 +995,33 @@ async def close(tool: FakeTool) -> None: cancel_first_close = False await asyncio.wait_for(handler.aclose(), timeout=1) + @pytest.mark.asyncio + async def test_repeated_cancellation_while_releasing_entry_does_not_leak_lease(self) -> None: + handler = DefaultMCPToolHandler() + invocation_started = asyncio.Event() + release_invocation = asyncio.Event() + + async def gated_call(_tool: FakeTool, _tool_name: str, **_arguments: Any) -> Any: + invocation_started.set() + await release_invocation.wait() + return [Content.from_text("ok")] + + with _patch_tool(), patch.object(FakeTool, "call_tool", gated_call): + invocation = asyncio.create_task(handler.invoke_tool(_invocation(headers={"X": "1"}))) + await invocation_started.wait() + await handler._cache_lock.acquire() + invocation.cancel() + await asyncio.sleep(0) + invocation.cancel() + handler._cache_lock.release() + + with pytest.raises(asyncio.CancelledError): + await invocation + await asyncio.wait_for(handler.aclose(), timeout=1) + + assert FakeTool.instances[0].close_count == 1 + assert all(entry.active_users == 0 for entry in handler._cache.values()) + @pytest.mark.asyncio async def test_entry_creation_is_bounded_and_cancelled_waiter_is_cleaned_up(self) -> None: handler = DefaultMCPToolHandler(cache_max_size=2) From bdee27b576fc66934588d29cc7710ebeb1afc550 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 29 Sep 2026 13:43:41 -0400 Subject: [PATCH 23/42] Clarify cancelled MCP task coordination Await the cancelled invocation through asyncio.gather so the test's coordination effect is explicit to code-quality analysis. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: aafca87b-cf2e-4c72-b5cb-ffcad6e1b1e0 --- .../packages/declarative/tests/test_default_mcp_tool_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index 770801afd8..8cdb26391d 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -1016,7 +1016,7 @@ async def gated_call(_tool: FakeTool, _tool_name: str, **_arguments: Any) -> Any handler._cache_lock.release() with pytest.raises(asyncio.CancelledError): - await invocation + await asyncio.gather(invocation) await asyncio.wait_for(handler.aclose(), timeout=1) assert FakeTool.instances[0].close_count == 1 From 1283c999d76b0c886fc9a774847c9191fbe1b2b4 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 29 Sep 2026 14:10:54 -0400 Subject: [PATCH 24/42] Harden MCP restore and cancellation handling Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: aafca87b-cf2e-4c72-b5cb-ffcad6e1b1e0 --- .../InProc/InProcessRunner.cs | 35 ++++++++- .../InProc/InProcessRunnerContext.cs | 8 +- .../Specialized/WorkflowHostExecutor.cs | 14 +++- .../SubworkflowBinding.cs | 8 +- .../RepresentationTests.cs | 23 ++++++ .../_workflows/_mcp_handler.py | 37 +++++++-- .../tests/test_default_mcp_tool_handler.py | 76 +++++++++++++++++++ 7 files changed, 179 insertions(+), 22 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunner.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunner.cs index 69b6c1e9bc..c0a2e075ae 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 0db031c9a0..aa8bf39859 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)) @@ -351,7 +355,7 @@ private sealed class BoundWorkflowContext( string ExecutorId, Dictionary? traceContext) : IWorkflowContext, IWorkflowSessionContext { - public string SessionId => RunnerContext._sessionId; + public string SessionId => RunnerContext.WorkflowSessionId; public ValueTask AddEventAsync(WorkflowEvent workflowEvent, CancellationToken cancellationToken = default) => RunnerContext.AddEventAsync(workflowEvent, cancellationToken); diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/WorkflowHostExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/WorkflowHostExecutor.cs index 58e3a9e523..a166d8b0e6 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 538e8bc3bb..9519af1241 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/SubworkflowBinding.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/SubworkflowBinding.cs @@ -32,13 +32,7 @@ async ValueTask InitHostExecutorAsync(string sessionId) { ProtocolDescriptor workflowProtocol = await workflow.DescribeProtocolAsync().ConfigureAwait(false); - return new WorkflowHostExecutor( - id, - workflow, - workflowProtocol, - CreateSubworkflowSessionId(sessionId, id), - ownershipToken, - options); + return new WorkflowHostExecutor(id, workflow, workflowProtocol, sessionId, ownershipToken, options); } } diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/RepresentationTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/RepresentationTests.cs index 6b8a0f7996..89928c889c 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; @@ -111,6 +112,28 @@ public void SubworkflowSessionId_IsStableAndHierarchicallyScoped() Assert.NotEqual(siblingChild, nestedChild); } + [Fact] + public async Task SubworkflowRunner_PreservesLegacyCheckpointSessionId() + { + // 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}"; diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py index a0d84d8843..a4f9b175bc 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -179,6 +179,7 @@ class _CacheEntry: evicted: bool = False disposal_claimed: bool = False closed: asyncio.Event = field(default_factory=asyncio.Event) + close_exception: BaseException | None = None class _EntryCreationCancelled(Exception): @@ -540,6 +541,9 @@ async def _drain_shutdown( ) for entry in entries: await entry.closed.wait() + for entry in entries: + if entry.close_exception is not None: + raise entry.close_exception for result in close_results: if isinstance(result, BaseException): raise result @@ -594,14 +598,7 @@ async def _get_or_create_entry(self, invocation: MCPToolInvocation) -> _CacheEnt raise RuntimeError("DefaultMCPToolHandler is closed") entry = await self._create_entry(invocation) except BaseException as exc: - async with self._cache_lock: - self._inflight.pop(key, None) - if not inflight.done(): - inflight.set_exception(self._entry_creation_exception(exc)) - # Mark the exception retrieved to suppress noisy "Future exception - # was never retrieved" warnings when there are no other awaiters - # (other awaiters still see the exception through their ``await``). - inflight.exception() + await self._complete_inflight_failure(key, inflight, exc) raise # Phase 3: insert with LRU eviction; resolve the in-flight future. @@ -690,6 +687,27 @@ async def cleanup() -> None: continue cleanup_task.result() + async def _complete_inflight_failure( + self, + key: tuple[str, str, str, str, str], + inflight: asyncio.Future[_CacheEntry], + exc: BaseException, + ) -> None: + async def cleanup() -> None: + async with self._cache_lock: + self._inflight.pop(key, None) + if not inflight.done(): + inflight.set_exception(self._entry_creation_exception(exc)) + inflight.exception() + + cleanup_task = asyncio.create_task(cleanup()) + while not cleanup_task.done(): + try: + await asyncio.shield(cleanup_task) + except asyncio.CancelledError: + continue + cleanup_task.result() + @staticmethod def _entry_creation_exception(exc: BaseException) -> BaseException: if isinstance(exc, asyncio.CancelledError): @@ -722,6 +740,9 @@ async def _release_entry_core(self, entry: _CacheEntry) -> None: async def _close_claimed_entry(self, entry: _CacheEntry) -> None: try: await self._close_invocation_entry(entry) + except BaseException as exc: + entry.close_exception = exc + raise finally: async with self._cache_lock: self._retired.pop(id(entry), None) diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index 8cdb26391d..7661496ff2 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -1128,6 +1128,47 @@ async def gated_connect(tool: FakeTool) -> None: assert FakeTool.instances[0].close_count == 1 assert FakeTool.instances[1].connect_count == 1 + @pytest.mark.asyncio + async def test_repeated_creator_cancellation_completes_inflight_failure(self) -> None: + handler = DefaultMCPToolHandler() + + class GatedCreationSemaphore: + def __init__(self) -> None: + self.entered = asyncio.Event() + self.release = asyncio.Event() + + async def __aenter__(self) -> None: + self.entered.set() + await self.release.wait() + + async def __aexit__(self, *_args: Any) -> None: + pass + + creation_gate: Any = GatedCreationSemaphore() + handler._creation_semaphore = creation_gate + + with _patch_tool(): + creator = asyncio.create_task(handler.invoke_tool(_invocation(workflow_session_id="workflow-a"))) + await creation_gate.entered.wait() + assert handler._inflight + waiter = asyncio.create_task(handler.invoke_tool(_invocation(workflow_session_id="workflow-a"))) + await asyncio.sleep(0) + + await handler._cache_lock.acquire() + creator.cancel() + await asyncio.sleep(0) + creator.cancel() + handler._cache_lock.release() + creation_gate.release.set() + + with pytest.raises(asyncio.CancelledError): + await asyncio.gather(creator) + waiter_result = await asyncio.wait_for(waiter, timeout=1) + await asyncio.wait_for(handler.aclose(), timeout=1) + + assert not waiter_result.is_error + assert not handler._inflight + @pytest.mark.asyncio async def test_repeated_use_keeps_lru_alive(self) -> None: handler = DefaultMCPToolHandler(cache_max_size=2) @@ -1324,6 +1365,41 @@ async def cancel_first_close(tool: FakeTool) -> None: assert [tool.close_count for tool in FakeTool.instances] == [1, 1] assert all(tool._httpx_client is not None and tool._httpx_client.is_closed for tool in FakeTool.instances) + @pytest.mark.asyncio + async def test_active_entry_cleanup_cancellation_is_reported_by_aclose(self) -> None: + handler = DefaultMCPToolHandler() + invocation_started = asyncio.Event() + release_invocation = asyncio.Event() + + async def gated_call(_tool: FakeTool, _tool_name: str, **_arguments: Any) -> Any: + invocation_started.set() + await release_invocation.wait() + return [Content.from_text("ok")] + + async def cancelled_close(tool: FakeTool) -> None: + tool.close_count += 1 + raise asyncio.CancelledError + + with ( + _patch_tool(), + patch.object(FakeTool, "call_tool", gated_call), + patch.object(FakeTool, "close", cancelled_close), + ): + invocation = asyncio.create_task(handler.invoke_tool(_invocation(headers={"X": "1"}))) + await invocation_started.wait() + shutdown = asyncio.create_task(handler.aclose()) + await asyncio.sleep(0) + assert handler._closed + release_invocation.set() + + with pytest.raises(asyncio.CancelledError): + await asyncio.gather(invocation) + with pytest.raises(asyncio.CancelledError): + await asyncio.gather(shutdown) + + assert FakeTool.instances[0].close_count == 1 + assert not handler._retired + @pytest.mark.asyncio async def test_invoke_after_close_returns_error_result(self) -> None: """Post-close ``invoke_tool`` surfaces a tool error rather than crashing.""" From 447f8b688bfcab7e9de80e438e3a5cf5400a0bbe Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 29 Sep 2026 14:25:52 -0400 Subject: [PATCH 25/42] Fix async workflow test naming Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: aafca87b-cf2e-4c72-b5cb-ffcad6e1b1e0 --- .../RepresentationTests.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/RepresentationTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/RepresentationTests.cs index 89928c889c..bddef5e4f1 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/RepresentationTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/RepresentationTests.cs @@ -113,7 +113,7 @@ public void SubworkflowSessionId_IsStableAndHierarchicallyScoped() } [Fact] - public async Task SubworkflowRunner_PreservesLegacyCheckpointSessionId() + public async Task SubworkflowRunner_PreservesLegacyCheckpointSessionIdAsync() { // Arrange const string LegacySessionId = "legacy-session"; From a5adae566efa248a14aa9f1b22b478e6967301a2 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Mon, 5 Oct 2026 12:01:49 -0400 Subject: [PATCH 26/42] Surface MCP cleanup failures during shutdown Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 05de6561-b8b7-4bea-81b0-36684365896c --- .../DefaultMcpToolHandler.cs | 66 ++++++++++++++--- .../DefaultMcpToolHandlerLifetimeTests.cs | 72 +++++++++++++++++++ .../_workflows/_mcp_handler.py | 47 ++++++++++-- .../tests/test_default_mcp_tool_handler.py | 35 +++++++++ 4 files changed, 207 insertions(+), 13 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 976d2b9981..c209fa65ea 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs @@ -50,6 +50,7 @@ public sealed class DefaultMcpToolHandler : IWorkflowScopedMcpToolHandler, IAsyn 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 = []; @@ -94,7 +95,8 @@ public DefaultMcpToolHandler(Func>? internal DefaultMcpToolHandler( Func>? httpClientProvider, Func httpMessageHandlerFactory, - int clientCacheMaxSize = DefaultClientCacheMaxSize) + int clientCacheMaxSize = DefaultClientCacheMaxSize, + Func? clientConnectionDisposer = null) { if (clientCacheMaxSize <= 0) { @@ -103,6 +105,7 @@ internal DefaultMcpToolHandler( this._httpClientProvider = httpClientProvider; this._httpMessageHandlerFactory = Throw.IfNull(httpMessageHandlerFactory); + this._clientConnectionDisposer = clientConnectionDisposer ?? (connection => connection.DisposeAsync()); this._clientCacheMaxSize = clientCacheMaxSize; this._clientCreationSemaphore = new(clientCacheMaxSize, clientCacheMaxSize); } @@ -283,9 +286,14 @@ public async ValueTask DisposeAsync() await providerInvocations.ConfigureAwait(false); } - foreach (Task clientCreationLifetime in clientCreationLifetimes) + Exception? clientCreationException = null; + try + { + await Task.WhenAll(clientCreationLifetimes).ConfigureAwait(false); + } + catch (Exception exception) when (!IsFatalException(exception)) { - await clientCreationLifetime.ConfigureAwait(false); + clientCreationException = exception; } List cachedClients; @@ -313,17 +321,32 @@ public async ValueTask DisposeAsync() 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) @@ -465,6 +488,7 @@ private async Task AcquireClientAsync( } ObjectDisposedException? disposedException = null; + Exception? clientCreationCleanupException = null; CachedClient? result = null; await this._clientLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); try @@ -511,7 +535,7 @@ private async Task AcquireClientAsync( { if (connection is not null) { - await connection.DisposeAsync().ConfigureAwait(false); + await this._clientConnectionDisposer(connection).ConfigureAwait(false); } if (disposedException is not null) @@ -530,6 +554,7 @@ private async Task AcquireClientAsync( if (disposedException is not null) { clientCreation.TrySetException(exception); + clientCreationCleanupException = exception; } if (result is not null) @@ -546,7 +571,8 @@ private async Task AcquireClientAsync( this._clientCreationSemaphore.Release(); } - await this.CompleteClientCreationLifetimeAsync(ownedClientCreationLifetime).ConfigureAwait(false); + await this.CompleteClientCreationLifetimeAsync( + ownedClientCreationLifetime, clientCreationCleanupException).ConfigureAwait(false); } return result ?? throw new InvalidOperationException("Failed to acquire MCP client."); @@ -596,7 +622,9 @@ private async Task CompleteClientCreationFailureAsync( } } - private async Task CompleteClientCreationLifetimeAsync(TaskCompletionSource clientCreationLifetime) + private async Task CompleteClientCreationLifetimeAsync( + TaskCompletionSource clientCreationLifetime, + Exception? cleanupException = null) { await this._clientLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); try @@ -608,7 +636,14 @@ private async Task CompleteClientCreationLifetimeAsync(TaskCompletionSource + { + 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() { @@ -299,6 +335,42 @@ public async Task NoProvider_DisposalDuringFailedCreation_PreservesCreationFailu "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() { diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py index a4f9b175bc..685ae9369b 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -273,6 +273,7 @@ def __init__( # tasks awaiting the same key will await the same future and share # the resulting cache entry. self._inflight: dict[tuple[str, str, str, str, str], asyncio.Future[_CacheEntry]] = {} + self._inflight_cleanup: dict[tuple[str, str, str, str, str], asyncio.Future[None]] = {} # Completion signals only: provider-backed calls never share entries. self._active_invocations: set[asyncio.Future[None]] = set() # Keep ancestry so a completed nested call cannot hide an active parent @@ -498,9 +499,16 @@ async def aclose(self) -> None: entry.disposal_claimed = True entries_to_close.append(entry) inflight_futures = list(self._inflight.values()) + inflight_cleanup_futures = list(self._inflight_cleanup.values()) active_invocations = list(self._active_invocations) self._shutdown_task = asyncio.create_task( - self._drain_shutdown(entries, entries_to_close, inflight_futures, active_invocations) + self._drain_shutdown( + entries, + entries_to_close, + inflight_futures, + inflight_cleanup_futures, + active_invocations, + ) ) shutdown_task = self._shutdown_task @@ -511,6 +519,7 @@ async def _drain_shutdown( entries: list[_CacheEntry], entries_to_close: list[_CacheEntry], inflight_futures: list[asyncio.Future[_CacheEntry]], + inflight_cleanup_futures: list[asyncio.Future[None]], active_invocations: list[asyncio.Future[None]], ) -> None: if active_invocations: @@ -534,6 +543,7 @@ async def _drain_shutdown( except BaseException: logger.debug("DefaultMCPToolHandler: in-flight future raised during aclose", exc_info=True) continue + cleanup_results = await asyncio.gather(*inflight_cleanup_futures, return_exceptions=True) close_results = await asyncio.gather( *(self._close_claimed_entry(entry) for entry in entries_to_close), @@ -544,6 +554,9 @@ async def _drain_shutdown( for entry in entries: if entry.close_exception is not None: raise entry.close_exception + for result in cleanup_results: + if isinstance(result, BaseException): + raise result for result in close_results: if isinstance(result, BaseException): raise result @@ -583,6 +596,7 @@ async def _get_or_create_entry(self, invocation: MCPToolInvocation) -> _CacheEnt if inflight is None: inflight = asyncio.get_running_loop().create_future() self._inflight[key] = inflight + self._inflight_cleanup[key] = asyncio.get_running_loop().create_future() creating = True if not creating: @@ -645,13 +659,16 @@ async def _get_or_create_entry(self, invocation: MCPToolInvocation) -> _CacheEnt # consistent "handler is closed" failure rather than receiving # an entry we are about to close behind their back. err = RuntimeError("DefaultMCPToolHandler is closed") + cleanup = self._inflight_cleanup.pop(key) try: - await self._close_invocation_entry(entry) + await self._close_invocation_entry(entry, cleanup) finally: if not inflight.done(): inflight.set_exception(err) inflight.exception() raise err + cleanup = self._inflight_cleanup.pop(key) + cleanup.set_result(None) try: if duplicate is not None: await self._close_entry(duplicate) @@ -670,8 +687,15 @@ async def _abort_entry_creation( exc: BaseException, ) -> None: async def cleanup() -> None: + cleanup_outcome = self._inflight_cleanup.pop(key) try: await self._close_entry(entry) + except BaseException as cleanup_exc: + cleanup_outcome.set_exception(cleanup_exc) + cleanup_outcome.exception() + raise + else: + cleanup_outcome.set_result(None) finally: async with self._cache_lock: self._inflight.pop(key, None) @@ -696,6 +720,8 @@ async def _complete_inflight_failure( async def cleanup() -> None: async with self._cache_lock: self._inflight.pop(key, None) + cleanup_outcome = self._inflight_cleanup.pop(key) + cleanup_outcome.set_result(None) if not inflight.done(): inflight.set_exception(self._entry_creation_exception(exc)) inflight.exception() @@ -798,7 +824,11 @@ def _header_provider(_kwargs: dict[str, Any]) -> dict[str, str]: owned_client = cast("httpx.AsyncClient | None", getattr(tool, "_httpx_client", None)) return _CacheEntry(tool=tool, owned_httpx_client=owned_client) - async def _close_invocation_entry(self, entry: _CacheEntry) -> None: + async def _close_invocation_entry( + self, + entry: _CacheEntry, + cleanup_outcome: asyncio.Future[None] | None = None, + ) -> None: """Finish invocation cleanup even if the caller is cancelled again.""" # MCPStreamableHTTPTool dispatches connect/close to its lifecycle owner, # keeping the SDK's cancel-scope entry and exit on that same task. @@ -809,7 +839,16 @@ async def _close_invocation_entry(self, entry: _CacheEntry) -> None: await asyncio.shield(cleanup) except asyncio.CancelledError: cancelled = True - cleanup.result() # Propagate cancellation/errors from cleanup itself. + try: + cleanup.result() # Propagate cancellation/errors from cleanup itself. + except BaseException as exc: + if cleanup_outcome is not None and not cleanup_outcome.done(): + cleanup_outcome.set_exception(exc) + cleanup_outcome.exception() + raise + else: + if cleanup_outcome is not None and not cleanup_outcome.done(): + cleanup_outcome.set_result(None) if cancelled: raise asyncio.CancelledError diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index 7661496ff2..b655694d31 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -1450,6 +1450,41 @@ async def gated_connect(self: FakeTool) -> None: assert result.is_error is True assert "closed" in (result.error_message or "").lower() + async def test_aclose_reports_inflight_creation_cleanup_cancellation(self) -> None: + handler = DefaultMCPToolHandler() + connect_started = asyncio.Event() + release_connect = asyncio.Event() + original_connect = FakeTool.connect + + async def gated_connect(self: FakeTool) -> None: + connect_started.set() + await release_connect.wait() + await original_connect(self) + + async def cancelled_close(tool: FakeTool) -> None: + tool.close_count += 1 + raise asyncio.CancelledError + + with ( + _patch_tool(), + patch.object(FakeTool, "connect", gated_connect), + patch.object(FakeTool, "close", cancelled_close), + ): + invocation = asyncio.create_task(handler.invoke_tool(_invocation(headers={"X": "1"}))) + await connect_started.wait() + shutdown = asyncio.create_task(handler.aclose()) + await asyncio.sleep(0) + release_connect.set() + + with pytest.raises(asyncio.CancelledError): + await invocation + with pytest.raises(asyncio.CancelledError): + await shutdown + + assert FakeTool.instances[0].close_count == 1 + assert not handler._inflight + assert not handler._inflight_cleanup + @pytest.mark.asyncio async def test_cancelled_creator_does_not_block_aclose(self) -> None: handler = DefaultMCPToolHandler() From 2ac83cf9e6d502efbe63f21bf404323c756cf877 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Mon, 5 Oct 2026 12:27:44 -0400 Subject: [PATCH 27/42] Fix cancellation assertions in MCP shutdown test Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 05de6561-b8b7-4bea-81b0-36684365896c --- .../declarative/tests/test_default_mcp_tool_handler.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index b655694d31..a76215b1bf 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -1477,9 +1477,10 @@ async def cancelled_close(tool: FakeTool) -> None: release_connect.set() with pytest.raises(asyncio.CancelledError): - await invocation + _ = await invocation + await asyncio.wait({shutdown}) with pytest.raises(asyncio.CancelledError): - await shutdown + shutdown.result() assert FakeTool.instances[0].close_count == 1 assert not handler._inflight From d1eab3f55c55dbc25508cdc4dabc133572ef36a4 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Mon, 5 Oct 2026 12:58:35 -0400 Subject: [PATCH 28/42] Make MCP approval checkpoint test deterministic Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 05de6561-b8b7-4bea-81b0-36684365896c --- .../declarative/tests/test_invoke_mcp_tool_executor.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py index 3b450f4a34..9fa9b1eeb9 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -500,7 +500,14 @@ def build(): assert "context-second" not in repr(checkpoints) assert "context-first" not in repr(approval.data) assert "context-second" not in repr(approval.data) - checkpoint = max(checkpoints, key=lambda item: item.timestamp) + # Wall-clock timestamps can tie on Windows. Select the checkpoint after + # the context change, where only the original approval remains pending. + checkpoint = next( + item + for item in checkpoints + if approval.request_id in item.pending_request_info_events + and change_request.request_id not in item.pending_request_info_events + ) assert checkpoint.state[_HEADER_BINDING_KEY] not in repr(approval.data) if restore: workflow = build() From abcd9c7a3ae3595ddf3e83759acdf4e75fc0b4b6 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Mon, 5 Oct 2026 13:14:44 -0400 Subject: [PATCH 29/42] Exclude expected MCP disposal rejection Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 05de6561-b8b7-4bea-81b0-36684365896c --- .../DefaultMcpToolHandler.cs | 2 +- .../DefaultMcpToolHandlerLifetimeTests.cs | 30 +++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) 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 c209fa65ea..53edfc30e5 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs @@ -551,7 +551,7 @@ private async Task AcquireClientAsync( } catch (Exception exception) { - if (disposedException is not null) + if (disposedException is not null && !ReferenceEquals(exception, disposedException)) { clientCreation.TrySetException(exception); clientCreationCleanupException = exception; 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 b98bbd4f87..fcdf522d48 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 @@ -335,6 +335,36 @@ public async Task NoProvider_DisposalDuringFailedCreation_PreservesCreationFailu "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() { From e4058ba0b4bedf25ffa4ed920f87464b0943744a Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Mon, 5 Oct 2026 15:01:24 -0400 Subject: [PATCH 30/42] Fix declarative MCP session continuity Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 2d4a3a89-ec77-4ce1-86dc-1e7db84cf4b8 --- python/packages/declarative/AGENTS.md | 2 + .../_workflows/_declarative_base.py | 11 ++++- .../_workflows/_executors_control_flow.py | 3 +- .../_workflows/_executors_mcp.py | 12 ++++- .../_workflows/_mcp_handler.py | 4 +- .../tests/test_default_mcp_tool_handler.py | 47 +++++++++++++++++++ .../tests/test_invoke_mcp_tool_executor.py | 44 +++++++++++++++++ 7 files changed, 119 insertions(+), 4 deletions(-) diff --git a/python/packages/declarative/AGENTS.md b/python/packages/declarative/AGENTS.md index 0679795d51..90eaa2cf36 100644 --- a/python/packages/declarative/AGENTS.md +++ b/python/packages/declarative/AGENTS.md @@ -37,6 +37,8 @@ workflow-local HMAC key held separately in trusted host checkpoint state. Only the opaque binding and header names enter the approval payload; raw headers are not checkpointed. Changed or unverifiable headers produce a replacement request for the same pinned operation, with a fresh request ID and no dispatch. +Pending requests also retain their originating workflow session ID across later +fresh runs and checkpoint restores. Fresh executors verify unchanged approvals using the checkpointed key; legacy requests or missing verification state require reapproval for non-empty headers. Custom handlers remain responsible for identity changes diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py b/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py index a9eb95473d..73e89235c1 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py @@ -1149,6 +1149,15 @@ def _get_state(self, state: State) -> DeclarativeWorkflowState: """Get the declarative workflow state wrapper.""" return DeclarativeWorkflowState(state, env_config=self._declarative_env_config) + @staticmethod + def _is_agent_continuation(state: DeclarativeWorkflowState, trigger: Any) -> bool: + """Return whether an agent-facing message list continues initialized state.""" + return ( + isinstance(trigger, list) + and all(isinstance(message, Message) for message in trigger) # pyright: ignore[reportUnknownVariableType] + and state.is_initialized() + ) + async def _ensure_state_initialized( self, ctx: WorkflowContext[Any, Any], @@ -1201,7 +1210,7 @@ async def _ensure_state_initialized( # reuse this workflow instance across unrelated conversations. # State left on the instance is not evidence of session ownership. # Explicit checkpoint/HIL restores use the runner's resume path. - is_continuation = not isinstance(trigger, Message) and state.is_initialized() + is_continuation = self._is_agent_continuation(state, trigger) # Locate the trailing user message in the trigger. last_user_index = -1 diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_control_flow.py b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_control_flow.py index a612f9aef0..bb61fbf0eb 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_control_flow.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_control_flow.py @@ -409,8 +409,9 @@ async def handle_action( ctx: WorkflowContext[ActionComplete], ) -> None: """Simply pass through to continue the workflow.""" + is_agent_continuation = self._is_agent_continuation(self._get_state(ctx.state), trigger) await self._ensure_state_initialized(ctx, trigger) - if self._action_def.get("kind") == "Entry": + if self._action_def.get("kind") == "Entry" and not is_agent_continuation: from ._mcp_handler import reset_workflow_session_id reset_workflow_session_id(ctx.state) diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_mcp.py b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_mcp.py index fef6185939..5f13045d3b 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_mcp.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_mcp.py @@ -86,6 +86,8 @@ class MCPToolApprovalRequest: (e.g. ``conversation_id``) for use by the resume handler. header_binding: Opaque binding of the reviewed headers. The verification key is retained separately in trusted workflow state. + workflow_session_id: Framework-owned MCP session identifier from the + run that originated the approval request. """ request_id: str @@ -97,6 +99,7 @@ class MCPToolApprovalRequest: connection_name: str | None = None metadata: dict[str, Any] = field(default_factory=lambda: {}) header_binding: str | None = None + workflow_session_id: str | None = None # --------------------------------------------------------------------------- @@ -307,6 +310,7 @@ async def _request_approval( header_binding=( self._bind_headers(ctx, request_id, invocation.headers, create_key=True) if invocation.headers else None ), + workflow_session_id=invocation.workflow_session_id, ) logger.info("%s: requesting approval for MCP tool '%s'", self.__class__.__name__, invocation.tool_name) await ctx.request_info(request, ToolApprovalResponse, request_id=request_id) @@ -341,6 +345,12 @@ async def handle_approval_response( await ctx.send_message(ActionComplete()) return + workflow_session_id = getattr(original_request, "workflow_session_id", None) + if workflow_session_id is None: + workflow_session_id = get_or_create_workflow_session_id(ctx.state) + elif not isinstance(workflow_session_id, str) or not workflow_session_id: + raise ValueError("Invalid MCP approval workflow session state.") + invocation = MCPToolInvocation( server_url=original_request.server_url, tool_name=tool_name, @@ -348,7 +358,7 @@ async def handle_approval_response( arguments=original_request.arguments, headers=self._evaluate_headers(state, self._action_def.get("headers")), connection_name=getattr(original_request, "connection_name", None), - workflow_session_id=get_or_create_workflow_session_id(ctx.state), + workflow_session_id=workflow_session_id, ) if invocation.headers or original_request.header_names: binding = getattr(original_request, "header_binding", None) diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py index 685ae9369b..2da5232b3b 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -687,7 +687,8 @@ async def _abort_entry_creation( exc: BaseException, ) -> None: async def cleanup() -> None: - cleanup_outcome = self._inflight_cleanup.pop(key) + async with self._cache_lock: + cleanup_outcome = self._inflight_cleanup[key] try: await self._close_entry(entry) except BaseException as cleanup_exc: @@ -699,6 +700,7 @@ async def cleanup() -> None: finally: async with self._cache_lock: self._inflight.pop(key, None) + self._inflight_cleanup.pop(key, None) if not inflight.done(): inflight.set_exception(self._entry_creation_exception(exc)) inflight.exception() diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index a76215b1bf..5801f1347c 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -1555,6 +1555,53 @@ async def gated_create_entry(invocation: MCPToolInvocation) -> Any: assert FakeTool.instances[0].close_count == 1 assert not handler._inflight + @pytest.mark.asyncio + async def test_aclose_reports_cancelled_abort_cleanup(self) -> None: + handler = DefaultMCPToolHandler() + entry_created = asyncio.Event() + release_creation = asyncio.Event() + close_started = asyncio.Event() + release_close = asyncio.Event() + original_create_entry = handler._create_entry + + async def gated_create_entry(invocation: MCPToolInvocation) -> Any: + entry = await original_create_entry(invocation) + entry_created.set() + await release_creation.wait() + return entry + + async def cancelled_close_entry(_entry: Any) -> None: + close_started.set() + await release_close.wait() + raise asyncio.CancelledError + + with ( + _patch_tool(), + patch.object(handler, "_create_entry", gated_create_entry), + patch.object(handler, "_close_entry", cancelled_close_entry), + ): + invocation = asyncio.create_task(handler.invoke_tool(_invocation(headers={"X": "1"}))) + await entry_created.wait() + await handler._cache_lock.acquire() + release_creation.set() + await asyncio.sleep(0) + invocation.cancel() + handler._cache_lock.release() + await close_started.wait() + + shutdown = asyncio.create_task(handler.aclose()) + await asyncio.sleep(0) + assert not shutdown.done() + release_close.set() + + with pytest.raises(asyncio.CancelledError): + _ = await invocation + with pytest.raises(asyncio.CancelledError): + _ = await shutdown + + assert not handler._inflight + assert not handler._inflight_cleanup + # ---------- Result normalisation ------------------------------------------ diff --git a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py index 9fa9b1eeb9..7088025e8b 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -194,6 +194,20 @@ async def test_fresh_runs_on_same_workflow_receive_separate_session_ids(self) -> assert handler.invocations[0].workflow_session_id assert handler.invocations[0].workflow_session_id != handler.invocations[1].workflow_session_id + @pytest.mark.asyncio + async def test_as_agent_continuation_reuses_workflow_session_id(self) -> None: + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + workflow = factory.create_workflow_from_definition(_yaml(_action())) + agent = workflow.as_agent(name="mcp-agent") + + await agent.run("first turn") + await agent.run("second turn") + + assert len(handler.invocations) == 2 + assert handler.invocations[0].workflow_session_id + assert handler.invocations[0].workflow_session_id == handler.invocations[1].workflow_session_id + @pytest.mark.asyncio async def test_continuation_reuses_workflow_session_id(self) -> None: from agent_framework_declarative._workflows import ToolApprovalResponse @@ -211,6 +225,32 @@ async def test_continuation_reuses_workflow_session_id(self) -> None: assert handler.last_invocation is not None assert handler.last_invocation.workflow_session_id == workflow_session_id + @pytest.mark.asyncio + async def test_pending_approval_uses_originating_workflow_session_id(self) -> None: + from agent_framework_declarative._workflows import ToolApprovalResponse + + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + workflow = factory.create_workflow_from_definition(_yaml(_action(require_approval=True))) + + first = await workflow.run({}) + [first_approval] = first.get_request_info_events() + first_session_id = first_approval.data.workflow_session_id + + second = await workflow.run({}) + [second_approval] = [ + event for event in second.get_request_info_events() if event.request_id != first_approval.request_id + ] + + assert first_session_id + assert second_approval.data.workflow_session_id + assert first_session_id != second_approval.data.workflow_session_id + + await workflow.run(responses={first_approval.request_id: ToolApprovalResponse(approved=True)}) + + assert handler.last_invocation is not None + assert handler.last_invocation.workflow_session_id == first_session_id + @pytest.mark.asyncio async def test_arguments_evaluated_and_preserves_none(self) -> None: handler = StubMcpHandler(_ok()) @@ -491,6 +531,8 @@ def build(): requests = paused.get_request_info_events() approval = next(event for event in requests if isinstance(event.data, MCPToolApprovalRequest)) change_request = next(event for event in requests if event is not approval) + approval_session_id = approval.data.workflow_session_id + assert approval_session_id selector = "second" if change_context else "first" await workflow.run(responses={change_request.request_id: ExternalInputResponse(user_input=selector)}) assert handler.call_count == 0 @@ -519,6 +561,7 @@ def build(): assert not resumed.get_request_info_events() assert handler.last_invocation is not None assert handler.last_invocation.headers == {"Authorization": "context-first"} + assert handler.last_invocation.workflow_session_id == approval_session_id return assert handler.call_count == 0 [replacement] = resumed.get_request_info_events() @@ -531,6 +574,7 @@ def build(): assert handler.last_invocation is not None assert handler.last_invocation.headers == {"Authorization": "context-second"} assert handler.last_invocation.arguments == {"q": "reviewed"} + assert handler.last_invocation.workflow_session_id == approval_session_id @pytest.mark.asyncio async def test_approval_required_emits_request_and_yields(self, mock_state, mock_context) -> None: # type: ignore[no-untyped-def] From c5ccc2d647fa583956f90c34dff994b45e0f43b0 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 6 Oct 2026 10:18:13 -0400 Subject: [PATCH 31/42] Fix MCP session concurrency review findings Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: b95469e1-3f0c-4a0e-a762-38ae9b9edde5 --- .../DefaultMcpToolHandler.cs | 13 +++- .../DefaultMcpToolHandlerLifetimeTests.cs | 33 ++++++++++ .../core/agent_framework/_workflows/_agent.py | 7 +++ .../core/agent_framework/_workflows/_const.py | 3 + .../agent_framework/_workflows/_workflow.py | 4 ++ .../_workflows/_executors_control_flow.py | 6 +- .../_workflows/_mcp_handler.py | 62 ++++++++++++++++--- .../tests/test_default_mcp_tool_handler.py | 38 +++++++++++- .../tests/test_invoke_mcp_tool_executor.py | 23 ++++++- 9 files changed, 171 insertions(+), 18 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 53edfc30e5..b82bfc6ec6 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs @@ -47,6 +47,8 @@ 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; @@ -436,7 +438,8 @@ private async Task AcquireClientAsync( { await WaitForClientCreationAsync(clientCreation.Task, cancellationToken).ConfigureAwait(false); } - catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested) + catch (OperationCanceledException exception) + when (ReferenceEquals(exception, s_clientCreationCancelledException)) { return await this.AcquireClientAsync( serverUrl, serverLabel, headers, connectionName, workflowSessionId, cancellationToken).ConfigureAwait(false); @@ -469,9 +472,14 @@ private async Task AcquireClientAsync( } catch (Exception exception) { + Exception sharedCreationException = + exception is OperationCanceledException && cancellationToken.IsCancellationRequested + ? s_clientCreationCancelledException + : exception; try { - await this.CompleteClientCreationFailureAsync(clientCacheKey, clientCreation, exception).ConfigureAwait(false); + await this.CompleteClientCreationFailureAsync( + clientCacheKey, clientCreation, sharedCreationException).ConfigureAwait(false); } finally { @@ -857,7 +865,6 @@ private static HttpMessageHandler CreateHttpMessageHandler() => 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 fcdf522d48..6ea27b4314 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 @@ -567,6 +567,39 @@ public async Task NoProvider_CancelledCreator_DoesNotCancelSharedWaiterAsync() 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 + Assert.Same(timeoutException, await Assert.ThrowsAsync(() => creator)); + Assert.Same(timeoutException, await Assert.ThrowsAsync(() => waiter)); + Assert.Equal(1, initializationAttempts); + Assert.Equal(0, stub.Initializations); + } + [Fact] public async Task Provider_OperationFailure_DisposesSessionAndPreservesCallerClientAsync() { diff --git a/python/packages/core/agent_framework/_workflows/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index 6d01c8366a..9ca1e5928f 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/python/packages/core/agent_framework/_workflows/_agent.py @@ -330,6 +330,7 @@ async def _run_impl( output_events: list[WorkflowEvent[Any]] = [] async for event in self._run_core( session_messages, + session.session_id if session is not None else None, checkpoint_id, checkpoint_storage, streaming=False, @@ -416,6 +417,7 @@ async def _run_stream_impl( all_updates: list[AgentResponseUpdate] = [] async for event in self._run_core( session_messages, + session.session_id if session is not None else None, checkpoint_id, checkpoint_storage, streaming=True, @@ -437,6 +439,7 @@ async def _run_stream_impl( async def _run_core( self, input_messages: Sequence[Message], + agent_session_id: str | None, checkpoint_id: str | None, checkpoint_storage: CheckpointStorage | None, streaming: bool, @@ -451,6 +454,7 @@ async def _run_core( Args: input_messages: Normalized input messages to process. + agent_session_id: Framework-local identity of the supplied agent session. checkpoint_id: ID of checkpoint to restore from. checkpoint_storage: Runtime checkpoint storage. streaming: Whether to use streaming workflow methods. @@ -461,6 +465,8 @@ async def _run_core( Yields: WorkflowEvent objects from the workflow execution. """ + self.workflow._set_agent_session_id(agent_session_id) # pyright: ignore[reportPrivateUsage] + # Restore the workflow state if a checkpoint is provided if checkpoint_id is not None: if checkpoint_storage is None: @@ -484,6 +490,7 @@ async def _run_core( if not input_messages: logger.info("No input messages provided; the workflow has been restored to the checkpoint state.") return + self.workflow._set_agent_session_id(agent_session_id) # pyright: ignore[reportPrivateUsage] final_state = self._workflow.status logger.debug(f"Workflow state: {final_state}") diff --git a/python/packages/core/agent_framework/_workflows/_const.py b/python/packages/core/agent_framework/_workflows/_const.py index db99b45426..dac6b96ba7 100644 --- a/python/packages/core/agent_framework/_workflows/_const.py +++ b/python/packages/core/agent_framework/_workflows/_const.py @@ -17,6 +17,9 @@ # to pass kwargs from workflow.run() through to agent.run() and @tool functions. WORKFLOW_RUN_KWARGS_KEY = "_workflow_run_kwargs" +# State key for the AgentSession identity supplied through WorkflowAgent. +WORKFLOW_AGENT_SESSION_ID_KEY = "_workflow_agent_session_id" + # Key used to store collision-free, executor-aware workflow invocation kwargs. # WORKFLOW_RUN_KWARGS_KEY remains in the legacy dict format for compatibility with # older first-party packages and checkpoints. diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index b69838995d..f2a839a8e3 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -34,6 +34,7 @@ RAW_FUNCTION_INVOCATION_KWARGS_KEY, RESOLVED_WORKFLOW_RUN_KWARGS_KEY, ROUTED_WORKFLOW_RUN_KWARGS_KEY, + WORKFLOW_AGENT_SESSION_ID_KEY, WORKFLOW_RUN_KWARGS_KEY, ) from ._edge import ( @@ -428,6 +429,9 @@ def status(self) -> WorkflowRunState: """ return self._status + def _set_agent_session_id(self, session_id: str | None) -> None: + self._runner.state.set(WORKFLOW_AGENT_SESSION_ID_KEY, session_id) + def to_dict(self) -> dict[str, Any]: """Serialize the workflow definition into a JSON-ready dictionary.""" data: dict[str, Any] = { diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_control_flow.py b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_control_flow.py index bb61fbf0eb..dd95d9fa7e 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_control_flow.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_control_flow.py @@ -411,10 +411,10 @@ async def handle_action( """Simply pass through to continue the workflow.""" is_agent_continuation = self._is_agent_continuation(self._get_state(ctx.state), trigger) await self._ensure_state_initialized(ctx, trigger) - if self._action_def.get("kind") == "Entry" and not is_agent_continuation: - from ._mcp_handler import reset_workflow_session_id + if self._action_def.get("kind") == "Entry": + from ._mcp_handler import activate_workflow_session_id - reset_workflow_session_id(ctx.state) + activate_workflow_session_id(ctx.state, reset_unscoped=not is_agent_continuation) await ctx.send_message(ActionComplete()) diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py index 2da5232b3b..b59c9d8fe6 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -56,9 +56,29 @@ _DEFAULT_CACHE_MAX_SIZE = 32 _WORKFLOW_SESSION_ID_KEY = "_declarative_mcp_workflow_session_id" +_AGENT_WORKFLOW_SESSION_NAMESPACE_KEY = "_declarative_mcp_agent_workflow_session_namespace" +_WORKFLOW_AGENT_SESSION_ID_KEY = "_workflow_agent_session_id" def get_or_create_workflow_session_id(state: State) -> str: + agent_session_id = state.get(_WORKFLOW_AGENT_SESSION_ID_KEY) + if agent_session_id is not None: + if not isinstance(agent_session_id, str) or not agent_session_id: + raise ValueError("Invalid agent session state.") + namespace_value = state.get(_AGENT_WORKFLOW_SESSION_NAMESPACE_KEY) + if namespace_value is None: + namespace_value = uuid.uuid4().hex + state.set(_AGENT_WORKFLOW_SESSION_NAMESPACE_KEY, namespace_value) + if not isinstance(namespace_value, str): + raise ValueError("Invalid MCP agent workflow session state.") + try: + namespace = uuid.UUID(namespace_value) + except ValueError as exc: + raise ValueError("Invalid MCP agent workflow session state.") from exc + workflow_session_id = uuid.uuid5(namespace, agent_session_id).hex + state.set(_WORKFLOW_SESSION_ID_KEY, workflow_session_id) + return workflow_session_id + workflow_session_id = state.get(_WORKFLOW_SESSION_ID_KEY) if workflow_session_id is None: workflow_session_id = uuid.uuid4().hex @@ -68,8 +88,11 @@ def get_or_create_workflow_session_id(state: State) -> str: return workflow_session_id -def reset_workflow_session_id(state: State) -> None: - state.set(_WORKFLOW_SESSION_ID_KEY, uuid.uuid4().hex) +def activate_workflow_session_id(state: State, *, reset_unscoped: bool) -> None: + if state.get(_WORKFLOW_AGENT_SESSION_ID_KEY) is not None: + get_or_create_workflow_session_id(state) + elif reset_unscoped: + state.set(_WORKFLOW_SESSION_ID_KEY, uuid.uuid4().hex) @dataclass @@ -186,6 +209,14 @@ class _EntryCreationCancelled(Exception): """Signal waiters to retry after the task creating their entry was cancelled.""" +class _EntryCreationCleanupFailure(BaseException): + """Keep failed-handshake cleanup separate from the connection failure.""" + + def __init__(self, creation_exception: BaseException, cleanup_exception: BaseException) -> None: + self.creation_exception = creation_exception + self.cleanup_exception = cleanup_exception + + class DefaultMCPToolHandler: """Default :class:`MCPToolHandler` backed by :class:`agent_framework.MCPStreamableHTTPTool`. @@ -613,6 +644,8 @@ async def _get_or_create_entry(self, invocation: MCPToolInvocation) -> _CacheEnt entry = await self._create_entry(invocation) except BaseException as exc: await self._complete_inflight_failure(key, inflight, exc) + if isinstance(exc, _EntryCreationCleanupFailure): + raise exc.cleanup_exception from exc.creation_exception raise # Phase 3: insert with LRU eviction; resolve the in-flight future. @@ -723,9 +756,15 @@ async def cleanup() -> None: async with self._cache_lock: self._inflight.pop(key, None) cleanup_outcome = self._inflight_cleanup.pop(key) - cleanup_outcome.set_result(None) + creation_exception = exc + if isinstance(exc, _EntryCreationCleanupFailure): + creation_exception = exc.creation_exception + cleanup_outcome.set_exception(exc.cleanup_exception) + cleanup_outcome.exception() + else: + cleanup_outcome.set_result(None) if not inflight.done(): - inflight.set_exception(self._entry_creation_exception(exc)) + inflight.set_exception(self._entry_creation_exception(creation_exception)) inflight.exception() cleanup_task = asyncio.create_task(cleanup()) @@ -801,7 +840,7 @@ def _header_provider(_kwargs: dict[str, Any]) -> dict[str, str]: ) try: await tool.connect() - except BaseException: + except BaseException as creation_exception: failed_entry = _CacheEntry( tool=tool, owned_httpx_client=( @@ -810,10 +849,15 @@ def _header_provider(_kwargs: dict[str, Any]) -> dict[str, str]: else None ), ) - if self._client_provider is not None: - await self._close_invocation_entry(failed_entry) - else: - await self._close_entry(failed_entry) + try: + if self._client_provider is not None: + await self._close_invocation_entry(failed_entry) + else: + await self._close_entry(failed_entry) + except BaseException as cleanup_exception: + if self._client_provider is not None: + raise + raise _EntryCreationCleanupFailure(creation_exception, cleanup_exception) from cleanup_exception raise # ``MCPStreamableHTTPTool.get_mcp_client`` lazily creates an diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index 5801f1347c..8df29de59f 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -1486,6 +1486,40 @@ async def cancelled_close(tool: FakeTool) -> None: assert not handler._inflight assert not handler._inflight_cleanup + async def test_aclose_reports_failed_handshake_cleanup_cancellation(self) -> None: + handler = DefaultMCPToolHandler() + connect_started = asyncio.Event() + release_connect = asyncio.Event() + + async def failed_connect(_tool: FakeTool) -> None: + connect_started.set() + await release_connect.wait() + raise RuntimeError("handshake failed") + + async def cancelled_close(tool: FakeTool) -> None: + tool.close_count += 1 + raise asyncio.CancelledError + + with ( + _patch_tool(), + patch.object(FakeTool, "connect", failed_connect), + patch.object(FakeTool, "close", cancelled_close), + ): + invocation = asyncio.create_task(handler.invoke_tool(_invocation(headers={"X": "1"}))) + await connect_started.wait() + shutdown = asyncio.create_task(handler.aclose()) + await asyncio.sleep(0) + release_connect.set() + + with pytest.raises(asyncio.CancelledError): + await invocation + with pytest.raises(asyncio.CancelledError): + await shutdown + + assert FakeTool.instances[0].close_count == 1 + assert not handler._inflight + assert not handler._inflight_cleanup + @pytest.mark.asyncio async def test_cancelled_creator_does_not_block_aclose(self) -> None: handler = DefaultMCPToolHandler() @@ -1595,9 +1629,9 @@ async def cancelled_close_entry(_entry: Any) -> None: release_close.set() with pytest.raises(asyncio.CancelledError): - _ = await invocation + await invocation with pytest.raises(asyncio.CancelledError): - _ = await shutdown + await shutdown assert not handler._inflight assert not handler._inflight_cleanup diff --git a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py index 7088025e8b..0066e0a7aa 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -26,7 +26,7 @@ reason="PowerFx engine not available (requires dotnet runtime)", ) -from agent_framework import Content, Message # noqa: E402 +from agent_framework import AgentSession, Content, Message # noqa: E402 from agent_framework.exceptions import ToolExecutionException # noqa: E402 from agent_framework_declarative._workflows import ( # noqa: E402 @@ -208,6 +208,27 @@ async def test_as_agent_continuation_reuses_workflow_session_id(self) -> None: assert handler.invocations[0].workflow_session_id assert handler.invocations[0].workflow_session_id == handler.invocations[1].workflow_session_id + @pytest.mark.asyncio + async def test_as_agent_sessions_isolate_and_reuse_workflow_session_ids(self) -> None: + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + workflow = factory.create_workflow_from_definition(_yaml(_action())) + agent = workflow.as_agent(name="mcp-agent") + session_a = AgentSession() + session_b = AgentSession() + + await agent.run("session A first turn", session=session_a) + await agent.run("session B first turn", session=session_b) + await agent.run("session A second turn", session=session_a) + + assert len(handler.invocations) == 3 + session_a_id = handler.invocations[0].workflow_session_id + session_b_id = handler.invocations[1].workflow_session_id + assert session_a_id + assert session_b_id + assert session_a_id != session_b_id + assert handler.invocations[2].workflow_session_id == session_a_id + @pytest.mark.asyncio async def test_continuation_reuses_workflow_session_id(self) -> None: from agent_framework_declarative._workflows import ToolApprovalResponse From 9acfbe483cf6441a68b1d263a5e0dc36c1f3d8e9 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 6 Oct 2026 10:33:57 -0400 Subject: [PATCH 32/42] Stabilize MCP cancellation assertions Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: b95469e1-3f0c-4a0e-a762-38ae9b9edde5 --- .../declarative/tests/test_default_mcp_tool_handler.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index 8df29de59f..c022d70806 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -1512,9 +1512,10 @@ async def cancelled_close(tool: FakeTool) -> None: release_connect.set() with pytest.raises(asyncio.CancelledError): - await invocation + _ = await invocation + await asyncio.wait({shutdown}) with pytest.raises(asyncio.CancelledError): - await shutdown + shutdown.result() assert FakeTool.instances[0].close_count == 1 assert not handler._inflight @@ -1629,9 +1630,10 @@ async def cancelled_close_entry(_entry: Any) -> None: release_close.set() with pytest.raises(asyncio.CancelledError): - await invocation + _ = await invocation + await asyncio.wait({shutdown}) with pytest.raises(asyncio.CancelledError): - await shutdown + shutdown.result() assert not handler._inflight assert not handler._inflight_cleanup From 29d114dd7af34eb26bd70ad9017e6182b12f04cc Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 6 Oct 2026 10:43:57 -0400 Subject: [PATCH 33/42] Fix cross-framework cancellation assertion Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: b95469e1-3f0c-4a0e-a762-38ae9b9edde5 --- .../DefaultMcpToolHandlerLifetimeTests.cs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 6ea27b4314..5390b822d6 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 @@ -594,8 +594,8 @@ public async Task NoProvider_InitializationTimeout_IsSharedWithoutRetryAsync() releaseInitialization.Release(); // Assert - Assert.Same(timeoutException, await Assert.ThrowsAsync(() => creator)); - Assert.Same(timeoutException, await Assert.ThrowsAsync(() => waiter)); + await Assert.ThrowsAnyAsync(() => creator); + await Assert.ThrowsAnyAsync(() => waiter); Assert.Equal(1, initializationAttempts); Assert.Equal(0, stub.Initializations); } From e227494747bb3e584a6aafaa5b136016433c48fd Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Wed, 7 Oct 2026 08:33:31 -0400 Subject: [PATCH 34/42] Fix workflow agent session identity on rejected concurrent runs Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: a12ce7d2-57d1-4cf3-9b43-8b14e5ef4de0 --- .../core/agent_framework/_workflows/_agent.py | 37 ++++++++++------ .../tests/workflow/test_workflow_agent.py | 44 +++++++++++++++++++ 2 files changed, 67 insertions(+), 14 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index 9ca1e5928f..c9511dc095 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/python/packages/core/agent_framework/_workflows/_agent.py @@ -465,8 +465,6 @@ async def _run_core( Yields: WorkflowEvent objects from the workflow execution. """ - self.workflow._set_agent_session_id(agent_session_id) # pyright: ignore[reportPrivateUsage] - # Restore the workflow state if a checkpoint is provided if checkpoint_id is not None: if checkpoint_storage is None: @@ -474,23 +472,26 @@ async def _run_core( logger.debug(f"Restoring workflow from checkpoint {checkpoint_id}") # Restore the workflow from checkpoint if streaming: - async for _ in self.workflow.run( + restore_stream = self.workflow.run( stream=True, checkpoint_id=checkpoint_id, checkpoint_storage=checkpoint_storage, tools=tools, - ): + ) + self.workflow._set_agent_session_id(agent_session_id) # pyright: ignore[reportPrivateUsage] + async for _ in restore_stream: pass else: - _ = await self.workflow.run( + restore_result = self.workflow.run( checkpoint_id=checkpoint_id, checkpoint_storage=checkpoint_storage, tools=tools, ) + self.workflow._set_agent_session_id(agent_session_id) # pyright: ignore[reportPrivateUsage] + _ = await restore_result if not input_messages: logger.info("No input messages provided; the workflow has been restored to the checkpoint state.") return - self.workflow._set_agent_session_id(agent_session_id) # pyright: ignore[reportPrivateUsage] final_state = self._workflow.status logger.debug(f"Workflow state: {final_state}") @@ -505,43 +506,51 @@ async def _run_core( pending_requests = await self.workflow._runner_context.get_pending_request_info_events() # pyright: ignore[reportPrivateUsage] function_responses = self._extract_function_responses(input_messages, pending_requests) if streaming: - async for event in self.workflow.run( + response_stream = self.workflow.run( responses=function_responses, stream=True, checkpoint_storage=checkpoint_storage, tools=tools, function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, - ): + ) + self.workflow._set_agent_session_id(agent_session_id) # pyright: ignore[reportPrivateUsage] + async for event in response_stream: yield event else: - for event in await self.workflow.run( + response_result = self.workflow.run( responses=function_responses, checkpoint_storage=checkpoint_storage, tools=tools, function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, - ): + ) + self.workflow._set_agent_session_id(agent_session_id) # pyright: ignore[reportPrivateUsage] + for event in await response_result: yield event elif final_state == WorkflowRunState.IDLE: if streaming: - async for event in self.workflow.run( + response_stream = self.workflow.run( message=input_messages, stream=True, checkpoint_storage=checkpoint_storage, tools=tools, function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, - ): + ) + self.workflow._set_agent_session_id(agent_session_id) # pyright: ignore[reportPrivateUsage] + async for event in response_stream: yield event else: - for event in await self.workflow.run( + response_result = self.workflow.run( message=input_messages, checkpoint_storage=checkpoint_storage, tools=tools, function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, - ): + ) + self.workflow._set_agent_session_id(agent_session_id) # pyright: ignore[reportPrivateUsage] + for event in await response_result: yield event else: raise AgentException(f"The underlying workflow is in an invalid state to restart: {final_state}.") diff --git a/python/packages/core/tests/workflow/test_workflow_agent.py b/python/packages/core/tests/workflow/test_workflow_agent.py index 0ef70ca2f6..4ef9fdd82b 100644 --- a/python/packages/core/tests/workflow/test_workflow_agent.py +++ b/python/packages/core/tests/workflow/test_workflow_agent.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. +import asyncio import uuid from collections.abc import Awaitable, Sequence from dataclasses import dataclass @@ -29,10 +30,12 @@ WorkflowBuilder, WorkflowContext, WorkflowEvent, + WorkflowException, executor, handler, response_handler, ) +from agent_framework._workflows._const import WORKFLOW_AGENT_SESSION_ID_KEY from agent_framework._workflows._typing_utils import deserialize_type @@ -193,6 +196,47 @@ async def handle_message( class TestWorkflowAgent: """Test cases for WorkflowAgent end-to-end functionality.""" + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("rejected_streaming", [False, True]) + async def test_rejected_concurrent_run_preserves_agent_session_id( + self, streaming: bool, rejected_streaming: bool, monkeypatch: pytest.MonkeyPatch + ) -> None: + started = asyncio.Event() + release = asyncio.Event() + workflow = WorkflowBuilder(start_executor=SimpleExecutor(id="start", response_text="Accepted")).build() + agent = workflow.as_agent() + session = AgentSession() + rejected_session = AgentSession() + original_has_messages = workflow._runner.context.has_messages + + async def gated_has_messages() -> bool: + started.set() + await release.wait() + return await original_has_messages() + + monkeypatch.setattr(workflow._runner.context, "has_messages", gated_has_messages) + + async def run_agent(agent_session: AgentSession, stream: bool) -> AgentResponse: + if stream: + return await agent.run("Hello", session=agent_session, stream=True).get_final_response() + return await agent.run("Hello", session=agent_session) + + task = asyncio.create_task(run_agent(session, streaming)) + try: + await started.wait() + assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == session.session_id + with pytest.raises(WorkflowException, match="Workflow is already running"): + await run_agent(rejected_session, rejected_streaming) + assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == session.session_id + finally: + release.set() + result = await task + + assert result.text == "Accepted: Hello" + assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == session.session_id + await run_agent(rejected_session, rejected_streaming) + assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == rejected_session.session_id + async def test_end_to_end_basic_workflow(self): """Test basic end-to-end workflow execution with 2 executors emitting AgentResponse.""" # Create workflow with two executors From 0b293d3e497aa987cd2c9454a2b8f365d94dd861 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Wed, 7 Oct 2026 10:19:15 -0400 Subject: [PATCH 35/42] Fix fresh workflow session isolation and MCP cleanup cancellation Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: a12ce7d2-57d1-4cf3-9b43-8b14e5ef4de0 --- .../DefaultMcpToolHandler.cs | 33 ++++++++++----- .../DefaultMcpToolHandlerLifetimeTests.cs | 35 +++++++++++++++- .../agent_framework/_workflows/_workflow.py | 3 ++ .../tests/workflow/test_workflow_agent.py | 41 +++++++++++++++++++ .../tests/test_invoke_mcp_tool_executor.py | 15 +++++++ 5 files changed, 114 insertions(+), 13 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 b82bfc6ec6..31d955d710 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs @@ -453,6 +453,7 @@ private async Task AcquireClientAsync( throw new InvalidOperationException("Missing MCP client creation lifetime."); ClientConnection? connection = null; bool creationSemaphoreEntered = false; + Exception? clientCreationCleanupException = null; try { await this._clientCreationSemaphore.WaitAsync(cancellationToken).ConfigureAwait(false); @@ -468,7 +469,8 @@ private async Task AcquireClientAsync( } connection = await this.CreateClientAsync( - trimmedUrl, serverLabel, headers, cancellationToken).ConfigureAwait(false); + trimmedUrl, serverLabel, headers, cancellationToken, + cleanupException => clientCreationCleanupException = cleanupException).ConfigureAwait(false); } catch (Exception exception) { @@ -489,14 +491,14 @@ await this.CompleteClientCreationFailureAsync( creationSemaphoreEntered = false; } - await this.CompleteClientCreationLifetimeAsync(ownedClientCreationLifetime).ConfigureAwait(false); + await this.CompleteClientCreationLifetimeAsync( + ownedClientCreationLifetime, clientCreationCleanupException).ConfigureAwait(false); } throw; } ObjectDisposedException? disposedException = null; - Exception? clientCreationCleanupException = null; CachedClient? result = null; await this._clientLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); try @@ -732,7 +734,8 @@ private async Task CreateClientAsync( string serverUrl, string? serverLabel, IDictionary? headers, - CancellationToken cancellationToken) + CancellationToken cancellationToken, + Action? reportCleanupFailure = null) { HttpClient? httpClient = null; bool ownsHttpClient = false; @@ -780,18 +783,26 @@ private async Task CreateClientAsync( { try { - if (transport is not null) + try { - await DisposeResourceAsync(transport, "transport").ConfigureAwait(false); + if (transport is not null) + { + await DisposeResourceAsync(transport, "transport").ConfigureAwait(false); + } } - } - finally - { - if (ownedHttpClientLease is not null) + finally { - await ownedHttpClientLease.DisposeAsync().ConfigureAwait(false); + if (ownedHttpClientLease is not null) + { + await ownedHttpClientLease.DisposeAsync().ConfigureAwait(false); + } } } + catch (Exception cleanupException) + { + reportCleanupFailure?.Invoke(cleanupException); + throw; + } throw; } 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 5390b822d6..a9820c16dc 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 @@ -335,6 +335,36 @@ public async Task NoProvider_DisposalDuringFailedCreation_PreservesCreationFailu "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() { @@ -993,6 +1023,7 @@ private sealed class ProtocolStub public bool FailInitialization { get; set; } public bool FailOperation { get; set; } public bool FailTransportDisposal { get; set; } + public Exception? TransportDisposalException { get; set; } public HttpMessageHandler CreateMessageHandler() { @@ -1000,10 +1031,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); diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 8a7bb4e21a..bfab8c96bf 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -868,6 +868,9 @@ def run( "Workflow is already running; concurrent runs are not allowed on the same instance." ) + if message is not None: + self._set_agent_session_id(None) + # No run is active, so any runtime checkpoint storage override still set on the # context is stale - left over from a prior run whose stream was dropped before # its async-generator finalizer ran. Clear it so this run starts clean and does diff --git a/python/packages/core/tests/workflow/test_workflow_agent.py b/python/packages/core/tests/workflow/test_workflow_agent.py index 4ef9fdd82b..605056bea2 100644 --- a/python/packages/core/tests/workflow/test_workflow_agent.py +++ b/python/packages/core/tests/workflow/test_workflow_agent.py @@ -19,6 +19,7 @@ Executor, FinishReason, HistoryProvider, + InMemoryCheckpointStorage, InMemoryHistoryProvider, Message, ResponseStream, @@ -196,6 +197,46 @@ async def handle_message( class TestWorkflowAgent: """Test cases for WorkflowAgent end-to-end functionality.""" + @pytest.mark.parametrize("streaming", [False, True]) + async def test_direct_fresh_runs_clear_agent_session_id(self, streaming: bool) -> None: + workflow = WorkflowBuilder(start_executor=SimpleExecutor(id="start", response_text="Accepted")).build() + agent = workflow.as_agent() + session = AgentSession() + await agent.run("Agent turn", session=session) + assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == session.session_id + + for text in ("First direct run", "Second direct run"): + messages = [Message(role="user", contents=[Content.from_text(text)])] + if streaming: + await workflow.run(messages, stream=True).get_final_response() + else: + await workflow.run(messages) + assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) is None + + @pytest.mark.parametrize("streaming", [False, True]) + async def test_direct_continuations_preserve_agent_session_id(self, streaming: bool) -> None: + workflow = WorkflowBuilder(start_executor=RequestingExecutor(id="request")).build() + storage = InMemoryCheckpointStorage() + session = AgentSession() + await workflow.as_agent().run("Agent turn", session=session, checkpoint_storage=storage) + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + pending_requests = await workflow._runner_context.get_pending_request_info_events() + assert len(pending_requests) == 1 + responses = {request_id: "Answer" for request_id in pending_requests} + + if streaming: + await workflow.run(responses=responses, stream=True).get_final_response() + assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == session.session_id + await workflow.run( + checkpoint_id=checkpoint_id, checkpoint_storage=storage, stream=True + ).get_final_response() + else: + await workflow.run(responses=responses) + assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == session.session_id + await workflow.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage) + assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == session.session_id + @pytest.mark.parametrize("streaming", [False, True]) @pytest.mark.parametrize("rejected_streaming", [False, True]) async def test_rejected_concurrent_run_preserves_agent_session_id( diff --git a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py index 0066e0a7aa..94c1f3261a 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -229,6 +229,21 @@ async def test_as_agent_sessions_isolate_and_reuse_workflow_session_ids(self) -> assert session_a_id != session_b_id assert handler.invocations[2].workflow_session_id == session_a_id + async def test_direct_fresh_runs_after_agent_turn_receive_separate_session_ids(self) -> None: + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + workflow = factory.create_workflow_from_definition(_yaml(_action())) + agent = workflow.as_agent(name="mcp-agent") + + await agent.run("agent turn", session=AgentSession()) + await workflow.run({}) + await workflow.run({}) + + session_ids = [invocation.workflow_session_id for invocation in handler.invocations] + assert len(session_ids) == 3 + assert all(session_ids) + assert len(set(session_ids)) == 3 + @pytest.mark.asyncio async def test_continuation_reuses_workflow_session_id(self) -> None: from agent_framework_declarative._workflows import ToolApprovalResponse From c8ef63f0d2608684e27cd10444b5f3e00606b90c Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Wed, 7 Oct 2026 10:58:40 -0400 Subject: [PATCH 36/42] Preserve originating MCP scope through approval continuations Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: a12ce7d2-57d1-4cf3-9b43-8b14e5ef4de0 --- python/packages/declarative/AGENTS.md | 3 +- .../_workflows/_executors_mcp.py | 12 ++-- .../_workflows/_mcp_handler.py | 8 +++ .../tests/test_default_mcp_tool_handler.py | 22 ++++++++ .../tests/test_invoke_mcp_tool_executor.py | 56 +++++++++++++++---- 5 files changed, 83 insertions(+), 18 deletions(-) diff --git a/python/packages/declarative/AGENTS.md b/python/packages/declarative/AGENTS.md index 013c7ad6e6..179e4eb64c 100644 --- a/python/packages/declarative/AGENTS.md +++ b/python/packages/declarative/AGENTS.md @@ -40,7 +40,8 @@ Only the opaque binding and header names enter the approval payload; raw headers are not checkpointed. Changed or unverifiable headers produce a replacement request for the same pinned operation, with a fresh request ID and no dispatch. Pending requests also retain their originating workflow session ID across later -fresh runs and checkpoint restores. +fresh runs and checkpoint restores. Resuming an MCP approval restores that scope +for downstream actions instead of inheriting a later run's MCP session. Fresh executors verify unchanged approvals using the checkpointed key; legacy requests or missing verification state require reapproval for non-empty headers. Custom handlers remain responsible for identity changes diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_mcp.py b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_mcp.py index 5f13045d3b..7eb3e97c0d 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_mcp.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_mcp.py @@ -53,6 +53,7 @@ MCPToolInvocation, MCPToolResult, get_or_create_workflow_session_id, + restore_workflow_session_id, ) __all__ = [ @@ -334,6 +335,11 @@ async def handle_approval_response( output_messages_path = _get_output_path(self._action_def, "messages") output_result_path = _get_output_path(self._action_def, "result") + workflow_session_id = getattr(original_request, "workflow_session_id", None) + if workflow_session_id is None: + workflow_session_id = get_or_create_workflow_session_id(ctx.state) + restore_workflow_session_id(ctx.state, workflow_session_id) + if response.approved is not True: logger.info( "%s: MCP tool '%s' rejected: %s", @@ -345,12 +351,6 @@ async def handle_approval_response( await ctx.send_message(ActionComplete()) return - workflow_session_id = getattr(original_request, "workflow_session_id", None) - if workflow_session_id is None: - workflow_session_id = get_or_create_workflow_session_id(ctx.state) - elif not isinstance(workflow_session_id, str) or not workflow_session_id: - raise ValueError("Invalid MCP approval workflow session state.") - invocation = MCPToolInvocation( server_url=original_request.server_url, tool_name=tool_name, diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py index b59c9d8fe6..5a30a44bb3 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -95,6 +95,14 @@ def activate_workflow_session_id(state: State, *, reset_unscoped: bool) -> None: state.set(_WORKFLOW_SESSION_ID_KEY, uuid.uuid4().hex) +def restore_workflow_session_id(state: State, workflow_session_id: str) -> None: + if not isinstance(workflow_session_id, str) or not workflow_session_id: + raise ValueError("Invalid MCP approval workflow session state.") + # An approval resumes its originating scope, not the latest agent turn's scope. + state.set(_WORKFLOW_AGENT_SESSION_ID_KEY, None) + state.set(_WORKFLOW_SESSION_ID_KEY, workflow_session_id) + + @dataclass class MCPToolInvocation: """Description of an MCP tool call to be dispatched by a :class:`MCPToolHandler`. diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index c022d70806..ce68f2eee6 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -30,6 +30,7 @@ DefaultMCPToolHandler, MCPToolInvocation, get_or_create_workflow_session_id, + restore_workflow_session_id, ) pytestmark = pytest.mark.skipif( @@ -830,6 +831,27 @@ def test_invalid_cache_size_raises(self) -> None: class TestWorkflowSessionId: + def test_restored_approval_scope_overrides_current_agent_session(self) -> None: + state = State() + state.set("_workflow_agent_session_id", "session-b") + agent_scope = get_or_create_workflow_session_id(state) + + restore_workflow_session_id(state, "approval-scope-a") + state.commit() + + assert get_or_create_workflow_session_id(state) == "approval-scope-a" + assert get_or_create_workflow_session_id(state) != agent_scope + + @pytest.mark.parametrize("workflow_session_id", ["", None, 42]) + def test_invalid_restored_approval_scope_preserves_current_scope(self, workflow_session_id: Any) -> None: + state = State() + current_scope = get_or_create_workflow_session_id(state) + + with pytest.raises(ValueError, match="Invalid MCP approval workflow session state"): + restore_workflow_session_id(state, workflow_session_id) + + assert get_or_create_workflow_session_id(state) == current_scope + def test_separate_workflow_states_get_separate_ids(self) -> None: first = get_or_create_workflow_session_id(State()) second = get_or_create_workflow_session_id(State()) diff --git a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py index 94c1f3261a..ac6bd924c2 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -261,31 +261,65 @@ async def test_continuation_reuses_workflow_session_id(self) -> None: assert handler.last_invocation is not None assert handler.last_invocation.workflow_session_id == workflow_session_id - @pytest.mark.asyncio - async def test_pending_approval_uses_originating_workflow_session_id(self) -> None: + @pytest.mark.parametrize("restore", [False, True]) + @pytest.mark.parametrize("agent_sessions", [False, True]) + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("approved", [False, True]) + async def test_pending_approval_uses_originating_workflow_session_id( + self, restore: bool, agent_sessions: bool, streaming: bool, approved: bool + ) -> None: + from agent_framework import InMemoryCheckpointStorage + from agent_framework_declarative._workflows import ToolApprovalResponse + storage = InMemoryCheckpointStorage() handler = StubMcpHandler(_ok()) factory = WorkflowFactory(mcp_tool_handler=handler) - workflow = factory.create_workflow_from_definition(_yaml(_action(require_approval=True))) + downstream = _action(tool_name="downstream") + downstream["id"] = "downstream" + definition = {"name": "mcp_test", "actions": [_action(require_approval=True), downstream]} + workflow = factory.create_workflow_from_definition(definition) + agent = workflow.as_agent(name="mcp-agent") - first = await workflow.run({}) - [first_approval] = first.get_request_info_events() + if agent_sessions: + await agent.run("first turn", session=AgentSession(), checkpoint_storage=storage) + else: + await workflow.run({}, checkpoint_storage=storage) + [first_approval] = (await workflow._runner_context.get_pending_request_info_events()).values() first_session_id = first_approval.data.workflow_session_id - second = await workflow.run({}) + await workflow.run({}, checkpoint_storage=storage) [second_approval] = [ - event for event in second.get_request_info_events() if event.request_id != first_approval.request_id + event + for event in (await workflow._runner_context.get_pending_request_info_events()).values() + if event.request_id != first_approval.request_id ] assert first_session_id assert second_approval.data.workflow_session_id assert first_session_id != second_approval.data.workflow_session_id - await workflow.run(responses={first_approval.request_id: ToolApprovalResponse(approved=True)}) - - assert handler.last_invocation is not None - assert handler.last_invocation.workflow_session_id == first_session_id + if restore: + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + workflow = factory.create_workflow_from_definition(definition) + await workflow.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage) + + for approval in (first_approval, second_approval): + responses = {approval.request_id: ToolApprovalResponse(approved=approved)} + if streaming: + await workflow.run(responses=responses, stream=True).get_final_response() + else: + await workflow.run(responses=responses) + + expected_tools = ["search", "downstream"] if approved else ["downstream"] + assert [invocation.tool_name for invocation in handler.invocations] == expected_tools * 2 + expected_sessions = [ + session_id + for session_id in (first_session_id, second_approval.data.workflow_session_id) + for _ in expected_tools + ] + assert [invocation.workflow_session_id for invocation in handler.invocations] == expected_sessions @pytest.mark.asyncio async def test_arguments_evaluated_and_preserves_none(self) -> None: From b10b5c41d97e33b6adef82f646c2e0612631fafd Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Wed, 7 Oct 2026 13:00:12 -0400 Subject: [PATCH 37/42] Distinguish agent turns from direct workflow MCP sessions Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: a12ce7d2-57d1-4cf3-9b43-8b14e5ef4de0 --- .../core/agent_framework/_workflows/_agent.py | 12 ++--- .../core/agent_framework/_workflows/_const.py | 3 ++ .../agent_framework/_workflows/_workflow.py | 7 ++- .../tests/workflow/test_workflow_agent.py | 45 +++++++++++++++---- python/packages/declarative/AGENTS.md | 2 + .../_workflows/_executors_control_flow.py | 5 ++- .../tests/test_invoke_mcp_tool_executor.py | 43 +++++++++++++----- 7 files changed, 87 insertions(+), 30 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index c9511dc095..7b86cf15c5 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/python/packages/core/agent_framework/_workflows/_agent.py @@ -478,7 +478,7 @@ async def _run_core( checkpoint_storage=checkpoint_storage, tools=tools, ) - self.workflow._set_agent_session_id(agent_session_id) # pyright: ignore[reportPrivateUsage] + self.workflow._set_agent_run_context(agent_session_id) # pyright: ignore[reportPrivateUsage] async for _ in restore_stream: pass else: @@ -487,7 +487,7 @@ async def _run_core( checkpoint_storage=checkpoint_storage, tools=tools, ) - self.workflow._set_agent_session_id(agent_session_id) # pyright: ignore[reportPrivateUsage] + self.workflow._set_agent_run_context(agent_session_id) # pyright: ignore[reportPrivateUsage] _ = await restore_result if not input_messages: logger.info("No input messages provided; the workflow has been restored to the checkpoint state.") @@ -514,7 +514,7 @@ async def _run_core( function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, ) - self.workflow._set_agent_session_id(agent_session_id) # pyright: ignore[reportPrivateUsage] + self.workflow._set_agent_run_context(agent_session_id) # pyright: ignore[reportPrivateUsage] async for event in response_stream: yield event else: @@ -525,7 +525,7 @@ async def _run_core( function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, ) - self.workflow._set_agent_session_id(agent_session_id) # pyright: ignore[reportPrivateUsage] + self.workflow._set_agent_run_context(agent_session_id) # pyright: ignore[reportPrivateUsage] for event in await response_result: yield event elif final_state == WorkflowRunState.IDLE: @@ -538,7 +538,7 @@ async def _run_core( function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, ) - self.workflow._set_agent_session_id(agent_session_id) # pyright: ignore[reportPrivateUsage] + self.workflow._set_agent_run_context(agent_session_id) # pyright: ignore[reportPrivateUsage] async for event in response_stream: yield event else: @@ -549,7 +549,7 @@ async def _run_core( function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, ) - self.workflow._set_agent_session_id(agent_session_id) # pyright: ignore[reportPrivateUsage] + self.workflow._set_agent_run_context(agent_session_id) # pyright: ignore[reportPrivateUsage] for event in await response_result: yield event else: diff --git a/python/packages/core/agent_framework/_workflows/_const.py b/python/packages/core/agent_framework/_workflows/_const.py index dac6b96ba7..7f096c5742 100644 --- a/python/packages/core/agent_framework/_workflows/_const.py +++ b/python/packages/core/agent_framework/_workflows/_const.py @@ -20,6 +20,9 @@ # State key for the AgentSession identity supplied through WorkflowAgent. WORKFLOW_AGENT_SESSION_ID_KEY = "_workflow_agent_session_id" +# State key identifying WorkflowAgent turns, including turns without an AgentSession. +WORKFLOW_AGENT_RUN_KEY = "_workflow_agent_run" + # Key used to store collision-free, executor-aware workflow invocation kwargs. # WORKFLOW_RUN_KWARGS_KEY remains in the legacy dict format for compatibility with # older first-party packages and checkpoints. diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index bfab8c96bf..301f5383ea 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -35,6 +35,7 @@ RAW_FUNCTION_INVOCATION_KWARGS_KEY, RESOLVED_WORKFLOW_RUN_KWARGS_KEY, ROUTED_WORKFLOW_RUN_KWARGS_KEY, + WORKFLOW_AGENT_RUN_KEY, WORKFLOW_AGENT_SESSION_ID_KEY, WORKFLOW_RUN_KWARGS_KEY, ) @@ -430,7 +431,8 @@ def status(self) -> WorkflowRunState: """ return self._status - def _set_agent_session_id(self, session_id: str | None) -> None: + def _set_agent_run_context(self, session_id: str | None) -> None: + self._runner.state.set(WORKFLOW_AGENT_RUN_KEY, True) self._runner.state.set(WORKFLOW_AGENT_SESSION_ID_KEY, session_id) def to_dict(self) -> dict[str, Any]: @@ -869,7 +871,8 @@ def run( ) if message is not None: - self._set_agent_session_id(None) + self._runner.state.set(WORKFLOW_AGENT_RUN_KEY, False) + self._runner.state.set(WORKFLOW_AGENT_SESSION_ID_KEY, None) # No run is active, so any runtime checkpoint storage override still set on the # context is stale - left over from a prior run whose stream was dropped before diff --git a/python/packages/core/tests/workflow/test_workflow_agent.py b/python/packages/core/tests/workflow/test_workflow_agent.py index 605056bea2..2faaa76b93 100644 --- a/python/packages/core/tests/workflow/test_workflow_agent.py +++ b/python/packages/core/tests/workflow/test_workflow_agent.py @@ -36,7 +36,7 @@ handler, response_handler, ) -from agent_framework._workflows._const import WORKFLOW_AGENT_SESSION_ID_KEY +from agent_framework._workflows._const import WORKFLOW_AGENT_RUN_KEY, WORKFLOW_AGENT_SESSION_ID_KEY from agent_framework._workflows._typing_utils import deserialize_type @@ -198,12 +198,16 @@ class TestWorkflowAgent: """Test cases for WorkflowAgent end-to-end functionality.""" @pytest.mark.parametrize("streaming", [False, True]) - async def test_direct_fresh_runs_clear_agent_session_id(self, streaming: bool) -> None: + @pytest.mark.parametrize("explicit_session", [False, True]) + async def test_direct_fresh_runs_clear_agent_session_id(self, streaming: bool, explicit_session: bool) -> None: workflow = WorkflowBuilder(start_executor=SimpleExecutor(id="start", response_text="Accepted")).build() agent = workflow.as_agent() - session = AgentSession() + session = AgentSession() if explicit_session else None await agent.run("Agent turn", session=session) - assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == session.session_id + assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == ( + session.session_id if session is not None else None + ) + assert workflow._runner.state.get(WORKFLOW_AGENT_RUN_KEY) is True for text in ("First direct run", "Second direct run"): messages = [Message(role="user", contents=[Content.from_text(text)])] @@ -212,6 +216,7 @@ async def test_direct_fresh_runs_clear_agent_session_id(self, streaming: bool) - else: await workflow.run(messages) assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) is None + assert workflow._runner.state.get(WORKFLOW_AGENT_RUN_KEY) is False @pytest.mark.parametrize("streaming", [False, True]) async def test_direct_continuations_preserve_agent_session_id(self, streaming: bool) -> None: @@ -228,19 +233,23 @@ async def test_direct_continuations_preserve_agent_session_id(self, streaming: b if streaming: await workflow.run(responses=responses, stream=True).get_final_response() assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == session.session_id + assert workflow._runner.state.get(WORKFLOW_AGENT_RUN_KEY) is True await workflow.run( checkpoint_id=checkpoint_id, checkpoint_storage=storage, stream=True ).get_final_response() else: await workflow.run(responses=responses) assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == session.session_id + assert workflow._runner.state.get(WORKFLOW_AGENT_RUN_KEY) is True await workflow.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage) assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == session.session_id + assert workflow._runner.state.get(WORKFLOW_AGENT_RUN_KEY) is True @pytest.mark.parametrize("streaming", [False, True]) @pytest.mark.parametrize("rejected_streaming", [False, True]) + @pytest.mark.parametrize("first_agent", [False, True]) async def test_rejected_concurrent_run_preserves_agent_session_id( - self, streaming: bool, rejected_streaming: bool, monkeypatch: pytest.MonkeyPatch + self, streaming: bool, rejected_streaming: bool, first_agent: bool, monkeypatch: pytest.MonkeyPatch ) -> None: started = asyncio.Event() release = asyncio.Event() @@ -262,19 +271,37 @@ async def run_agent(agent_session: AgentSession, stream: bool) -> AgentResponse: return await agent.run("Hello", session=agent_session, stream=True).get_final_response() return await agent.run("Hello", session=agent_session) - task = asyncio.create_task(run_agent(session, streaming)) + async def run_first() -> AgentResponse: + if first_agent: + return await run_agent(session, streaming) + messages = [Message(role="user", contents=[Content.from_text("Hello")])] + if streaming: + result = await workflow.run(messages, stream=True).get_final_response() + else: + result = await workflow.run(messages) + [output] = result.get_outputs() + assert isinstance(output, AgentResponse) + return output + + expected_session_id = session.session_id if first_agent else None + task = asyncio.create_task(run_first()) try: await started.wait() - assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == session.session_id + assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == expected_session_id + assert workflow._runner.state.get(WORKFLOW_AGENT_RUN_KEY) is first_agent with pytest.raises(WorkflowException, match="Workflow is already running"): await run_agent(rejected_session, rejected_streaming) - assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == session.session_id + with pytest.raises(WorkflowException, match="Workflow is already running"): + await workflow.run([Message(role="user", contents=[Content.from_text("Rejected")])]) + assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == expected_session_id + assert workflow._runner.state.get(WORKFLOW_AGENT_RUN_KEY) is first_agent finally: release.set() result = await task assert result.text == "Accepted: Hello" - assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == session.session_id + assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == expected_session_id + assert workflow._runner.state.get(WORKFLOW_AGENT_RUN_KEY) is first_agent await run_agent(rejected_session, rejected_streaming) assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == rejected_session.session_id diff --git a/python/packages/declarative/AGENTS.md b/python/packages/declarative/AGENTS.md index 179e4eb64c..8aadeee9bb 100644 --- a/python/packages/declarative/AGENTS.md +++ b/python/packages/declarative/AGENTS.md @@ -20,6 +20,8 @@ Both state classes validate object attribute names but leave dictionary keys unc Cache identity includes a framework-owned workflow session ID in addition to endpoint, label, connection, and headers, so separate fresh runs do not share a stateful MCP protocol session while continuations and checkpoint restores do. +Direct fresh runs reset this scope even for message-list inputs; only +framework-marked WorkflowAgent turns retain the agent continuation scope. With a provider, every invocation (including `tools/list`) gets a fresh tool/session, even if the provider returns `None` or a shared HTTP client. Invocation cleanup closes the session and any internally owned fallback client, never caller-owned HTTP clients. diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_control_flow.py b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_control_flow.py index dd95d9fa7e..4aeacac41a 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_control_flow.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_control_flow.py @@ -21,6 +21,7 @@ WorkflowContext, handler, ) +from agent_framework._workflows._const import WORKFLOW_AGENT_RUN_KEY from ._declarative_base import ( ActionComplete, @@ -409,7 +410,9 @@ async def handle_action( ctx: WorkflowContext[ActionComplete], ) -> None: """Simply pass through to continue the workflow.""" - is_agent_continuation = self._is_agent_continuation(self._get_state(ctx.state), trigger) + is_agent_continuation = ctx.state.get(WORKFLOW_AGENT_RUN_KEY) is True and self._is_agent_continuation( + self._get_state(ctx.state), trigger + ) await self._ensure_state_initialized(ctx, trigger) if self._action_def.get("kind") == "Entry": from ._mcp_handler import activate_workflow_session_id diff --git a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py index ac6bd924c2..fe7650e6dc 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -181,28 +181,38 @@ async def test_separate_workflows_receive_separate_session_ids(self) -> None: assert handler.invocations[1].workflow_session_id assert handler.invocations[0].workflow_session_id != handler.invocations[1].workflow_session_id - @pytest.mark.asyncio - async def test_fresh_runs_on_same_workflow_receive_separate_session_ids(self) -> None: + @pytest.mark.parametrize("message_list", [False, True]) + @pytest.mark.parametrize("streaming", [False, True]) + async def test_fresh_runs_on_same_workflow_receive_separate_session_ids( + self, message_list: bool, streaming: bool + ) -> None: handler = StubMcpHandler(_ok()) factory = WorkflowFactory(mcp_tool_handler=handler) workflow = factory.create_workflow_from_definition(_yaml(_action())) - await workflow.run({}) - await workflow.run({}) + for text in ("first run", "second run"): + inputs = [Message(role="user", contents=[Content.from_text(text)])] if message_list else {} + if streaming: + await workflow.run(inputs, stream=True).get_final_response() + else: + await workflow.run(inputs) assert len(handler.invocations) == 2 assert handler.invocations[0].workflow_session_id assert handler.invocations[0].workflow_session_id != handler.invocations[1].workflow_session_id - @pytest.mark.asyncio - async def test_as_agent_continuation_reuses_workflow_session_id(self) -> None: + @pytest.mark.parametrize("streaming", [False, True]) + async def test_as_agent_continuation_reuses_workflow_session_id(self, streaming: bool) -> None: handler = StubMcpHandler(_ok()) factory = WorkflowFactory(mcp_tool_handler=handler) workflow = factory.create_workflow_from_definition(_yaml(_action())) agent = workflow.as_agent(name="mcp-agent") - await agent.run("first turn") - await agent.run("second turn") + for text in ("first turn", "second turn"): + if streaming: + await agent.run(text, stream=True).get_final_response() + else: + await agent.run(text) assert len(handler.invocations) == 2 assert handler.invocations[0].workflow_session_id @@ -229,15 +239,24 @@ async def test_as_agent_sessions_isolate_and_reuse_workflow_session_ids(self) -> assert session_a_id != session_b_id assert handler.invocations[2].workflow_session_id == session_a_id - async def test_direct_fresh_runs_after_agent_turn_receive_separate_session_ids(self) -> None: + @pytest.mark.parametrize("message_list", [False, True]) + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("explicit_session", [False, True]) + async def test_direct_fresh_runs_after_agent_turn_receive_separate_session_ids( + self, message_list: bool, streaming: bool, explicit_session: bool + ) -> None: handler = StubMcpHandler(_ok()) factory = WorkflowFactory(mcp_tool_handler=handler) workflow = factory.create_workflow_from_definition(_yaml(_action())) agent = workflow.as_agent(name="mcp-agent") - await agent.run("agent turn", session=AgentSession()) - await workflow.run({}) - await workflow.run({}) + await agent.run("agent turn", session=AgentSession() if explicit_session else None) + for text in ("first direct run", "second direct run"): + inputs = [Message(role="user", contents=[Content.from_text(text)])] if message_list else {} + if streaming: + await workflow.run(inputs, stream=True).get_final_response() + else: + await workflow.run(inputs) session_ids = [invocation.workflow_session_id for invocation in handler.invocations] assert len(session_ids) == 3 From 2cfd09da77bbc3b423e7ddd486d1926e63dc2e6c Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Wed, 7 Oct 2026 14:54:51 -0400 Subject: [PATCH 38/42] Reject mixed MCP approval batches and preserve cleanup outcomes Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: a12ce7d2-57d1-4cf3-9b43-8b14e5ef4de0 --- .../DefaultMcpToolHandler.cs | 9 +- .../DefaultMcpToolHandlerLifetimeTests.cs | 82 ++++++++++++++++++- .../agent_framework/_workflows/_workflow.py | 3 + python/packages/declarative/AGENTS.md | 2 + .../_workflows/_declarative_builder.py | 26 +++++- .../tests/test_invoke_mcp_tool_executor.py | 69 ++++++++++++++-- 6 files changed, 177 insertions(+), 14 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 31d955d710..53cb850721 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs @@ -848,7 +848,14 @@ private async ValueTask ReleaseOwnedHttpClientAsync(string serverUrl) this._clientLock.Release(); } - clientToDispose?.Dispose(); + 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) 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 a9820c16dc..67390dd543 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 @@ -234,6 +234,76 @@ public async Task NoProvider_CacheEviction_DefersDisposalUntilActiveInvocationCo Assert.Equal(2, stub.Terminations); } + [Theory] + [InlineData(false, false)] + [InlineData(false, true)] + [InlineData(true, false)] + [InlineData(true, true)] + public async Task NoProvider_EvictedHttpClientCleanupFailure_PreservesInvocationOutcomeAsync( + bool activeEviction, 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 (activeEviction && Interlocked.Increment(ref operations) == 1) + { + firstStarted.Release(); + await releaseFirst.WaitAsync(token); + } + }; + await using DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler, clientCacheMaxSize: 1); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + using CleanupTraceListener listener = new(); + 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 (activeEviction) + { + await firstStarted.WaitAsync(timeout.Token); + await handler.InvokeToolInWorkflowSessionAsync( + "https://second.example/api", null, "ping", null, null, null, "workflow-b", timeout.Token); + 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); + } + + Assert.Contains("Failed to dispose MCP HTTP client", listener.Output); + Assert.Contains("transport cleanup failed", listener.Output); + } + finally + { + releaseFirst.Release(); + Trace.Listeners.Remove(listener); + } + } + [Fact] public async Task NoProvider_DisposalDuringActiveInvocation_ReportsCleanupCancellationAsync() { @@ -933,13 +1003,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); @@ -949,7 +1021,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)); } diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 301f5383ea..03e1e4f8d5 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -413,6 +413,7 @@ def __init__( # ever iterating, the weakref dereferences to ``None`` once Python collects it, # so a subsequent ``run()`` is allowed. self._active_run: weakref.ref[ResponseStream[WorkflowEvent, WorkflowRunResult]] | None = None + self._response_batch_validator: Callable[[Sequence[WorkflowEvent]], None] | None = None # Run-scoped pause checkpoint bookkeeping (owned by Workflow, not callers). # Captured at the start of each ``_run_core`` so ``resolve_pause_checkpoint_id`` @@ -1151,6 +1152,8 @@ async def _validate_responses_internal(self, responses: Mapping[str, Any]) -> di if isinstance(pending_request.data, Content) and pending_request.data.type == "computer_tool_call": _validate_computer_tool_result(pending_request.data, response) coerced_responses[request_id] = response + if self._response_batch_validator is not None: + self._response_batch_validator([pending_requests[request_id] for request_id in coerced_responses]) return coerced_responses async def _send_responses_internal(self, responses: Mapping[str, Any]) -> None: diff --git a/python/packages/declarative/AGENTS.md b/python/packages/declarative/AGENTS.md index 8aadeee9bb..3bfc41fa1b 100644 --- a/python/packages/declarative/AGENTS.md +++ b/python/packages/declarative/AGENTS.md @@ -44,6 +44,8 @@ request for the same pinned operation, with a fresh request ID and no dispatch. Pending requests also retain their originating workflow session ID across later fresh runs and checkpoint restores. Resuming an MCP approval restores that scope for downstream actions instead of inheriting a later run's MCP session. +Factory-built workflows reject response batches containing MCP approvals from +different originating scopes before dispatch; resume each scope in a separate run. Fresh executors verify unchanged approvals using the checkpointed key; legacy requests or missing verification state require reapproval for non-empty headers. Custom handlers remain responsible for identity changes diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_builder.py b/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_builder.py index 9a6814af60..bc56946f85 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_builder.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_builder.py @@ -14,11 +14,13 @@ from __future__ import annotations import logging +from collections.abc import Sequence from typing import Any, cast from agent_framework import ( Workflow, WorkflowBuilder, + WorkflowEvent, ) from ._declarative_base import ( @@ -41,7 +43,11 @@ ) from ._executors_external_input import EXTERNAL_INPUT_EXECUTORS from ._executors_http import HTTP_ACTION_EXECUTORS, HttpRequestActionExecutor -from ._executors_mcp import MCP_ACTION_EXECUTORS, InvokeMcpToolActionExecutor +from ._executors_mcp import ( + MCP_ACTION_EXECUTORS, + InvokeMcpToolActionExecutor, + MCPToolApprovalRequest, +) from ._executors_tools import TOOL_ACTION_EXECUTORS, InvokeFunctionToolExecutor from ._http_handler import HttpRequestHandler from ._mcp_handler import MCPToolHandler @@ -108,6 +114,20 @@ } +def _validate_mcp_response_batch(requests: Sequence[WorkflowEvent[Any]]) -> None: + scopes: set[str | None] = set() + for request in requests: + if isinstance(request.data, MCPToolApprovalRequest): + scope = getattr(request.data, "workflow_session_id", None) + if scope is not None and (not isinstance(scope, str) or not scope): + raise ValueError("Invalid MCP approval workflow session state.") + scopes.add(scope) + if len(scopes) > 1: + raise DeclarativeWorkflowError( + "MCP approval responses from different workflow sessions must be submitted in separate runs." + ) + + class DeclarativeWorkflowBuilder: """Builds a Workflow graph from declarative YAML actions. @@ -236,7 +256,9 @@ def build(self) -> Workflow: if isinstance(executor, DeclarativeActionExecutor): executor.set_declarative_env_config(self._env_config) - return builder.build() + workflow = builder.build() + workflow._response_batch_validator = _validate_mcp_response_batch # pyright: ignore[reportPrivateUsage] + return workflow def _validate_workflow(self, actions: list[dict[str, Any]]) -> None: """Validate the workflow definition before building. diff --git a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py index fe7650e6dc..41c9cad89c 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -284,8 +284,9 @@ async def test_continuation_reuses_workflow_session_id(self) -> None: @pytest.mark.parametrize("agent_sessions", [False, True]) @pytest.mark.parametrize("streaming", [False, True]) @pytest.mark.parametrize("approved", [False, True]) + @pytest.mark.parametrize("second_approved", [False, True]) async def test_pending_approval_uses_originating_workflow_session_id( - self, restore: bool, agent_sessions: bool, streaming: bool, approved: bool + self, restore: bool, agent_sessions: bool, streaming: bool, approved: bool, second_approved: bool ) -> None: from agent_framework import InMemoryCheckpointStorage @@ -318,28 +319,82 @@ async def test_pending_approval_uses_originating_workflow_session_id( assert second_approval.data.workflow_session_id assert first_session_id != second_approval.data.workflow_session_id + checkpoint_id = None if restore: checkpoint_id = workflow.get_last_checkpoint_id() assert checkpoint_id is not None workflow = factory.create_workflow_from_definition(definition) await workflow.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage) + decisions = {first_approval.request_id: approved, second_approval.request_id: second_approved} + batch = {request_id: ToolApprovalResponse(approved=decision) for request_id, decision in decisions.items()} + with pytest.raises(DeclarativeWorkflowError, match="different workflow sessions"): + if streaming: + await workflow.run( + responses=batch, + stream=True, + checkpoint_id=checkpoint_id, + checkpoint_storage=storage, + ).get_final_response() + else: + await workflow.run(responses=batch, checkpoint_id=checkpoint_id, checkpoint_storage=storage) + assert handler.call_count == 0 + assert set(await workflow._runner_context.get_pending_request_info_events()) == set(batch) + for approval in (first_approval, second_approval): - responses = {approval.request_id: ToolApprovalResponse(approved=approved)} + responses = {approval.request_id: ToolApprovalResponse(approved=decisions[approval.request_id])} if streaming: await workflow.run(responses=responses, stream=True).get_final_response() else: await workflow.run(responses=responses) - expected_tools = ["search", "downstream"] if approved else ["downstream"] - assert [invocation.tool_name for invocation in handler.invocations] == expected_tools * 2 + expected_tools = [ + tool + for decision in (approved, second_approved) + for tool in (["search", "downstream"] if decision else ["downstream"]) + ] + assert [invocation.tool_name for invocation in handler.invocations] == expected_tools expected_sessions = [ - session_id - for session_id in (first_session_id, second_approval.data.workflow_session_id) - for _ in expected_tools + approval.data.workflow_session_id + for approval in (first_approval, second_approval) + for _ in range(2 if decisions[approval.request_id] else 1) ] assert [invocation.workflow_session_id for invocation in handler.invocations] == expected_sessions + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("approved", [False, True]) + async def test_same_scope_approval_batch_is_supported(self, streaming: bool, approved: bool) -> None: + from agent_framework import WorkflowBuilder + + from agent_framework_declarative._workflows import InvokeMcpToolActionExecutor, ToolApprovalResponse + from agent_framework_declarative._workflows._declarative_builder import _validate_mcp_response_batch + from agent_framework_declarative._workflows._executors_control_flow import JoinExecutor + + handler = StubMcpHandler(_ok()) + entry = JoinExecutor({"kind": "Entry"}, id="entry") + actions = [ + InvokeMcpToolActionExecutor({**_action(require_approval=True), "id": name}, mcp_tool_handler=handler) + for name in ("first", "second") + ] + workflow = WorkflowBuilder(start_executor=entry).add_fan_out_edges(entry, actions).build() + workflow._response_batch_validator = _validate_mcp_response_batch + paused = await workflow.run({}) + approvals = paused.get_request_info_events() + assert len(approvals) == 2 + scopes = {approval.data.workflow_session_id for approval in approvals} + assert len(scopes) == 1 + assert all(scopes) + + responses = {approval.request_id: ToolApprovalResponse(approved=approved) for approval in approvals} + if streaming: + await workflow.run(responses=responses, stream=True).get_final_response() + else: + await workflow.run(responses=responses) + + assert handler.call_count == (2 if approved else 0) + assert all(invocation.workflow_session_id in scopes for invocation in handler.invocations) + assert not await workflow._runner_context.get_pending_request_info_events() + @pytest.mark.asyncio async def test_arguments_evaluated_and_preserves_none(self) -> None: handler = StubMcpHandler(_ok()) From 8204ea0728ee794a46cf9dfd816efd96c0ab5da4 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Wed, 7 Oct 2026 15:09:08 -0400 Subject: [PATCH 39/42] Annotate MCP workflow regression inputs for Zuban Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: a12ce7d2-57d1-4cf3-9b43-8b14e5ef4de0 --- .../declarative/tests/test_invoke_mcp_tool_executor.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py index 41c9cad89c..64e9bce749 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -191,7 +191,9 @@ async def test_fresh_runs_on_same_workflow_receive_separate_session_ids( workflow = factory.create_workflow_from_definition(_yaml(_action())) for text in ("first run", "second run"): - inputs = [Message(role="user", contents=[Content.from_text(text)])] if message_list else {} + inputs: list[Message] | dict[str, Any] = ( + [Message(role="user", contents=[Content.from_text(text)])] if message_list else {} + ) if streaming: await workflow.run(inputs, stream=True).get_final_response() else: @@ -252,7 +254,9 @@ async def test_direct_fresh_runs_after_agent_turn_receive_separate_session_ids( await agent.run("agent turn", session=AgentSession() if explicit_session else None) for text in ("first direct run", "second direct run"): - inputs = [Message(role="user", contents=[Content.from_text(text)])] if message_list else {} + inputs: list[Message] | dict[str, Any] = ( + [Message(role="user", contents=[Content.from_text(text)])] if message_list else {} + ) if streaming: await workflow.run(inputs, stream=True).get_final_response() else: From c3cb5b8972f4a5533353141fdebdc1c48c83433c Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Wed, 7 Oct 2026 15:39:44 -0400 Subject: [PATCH 40/42] Isolate sessionless workflow agent MCP scopes Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: a12ce7d2-57d1-4cf3-9b43-8b14e5ef4de0 --- python/packages/declarative/AGENTS.md | 3 + .../_workflows/_executors_control_flow.py | 7 +- .../_workflows/_mcp_handler.py | 11 +- .../tests/test_default_mcp_tool_handler.py | 69 +++++++ .../tests/test_invoke_mcp_tool_executor.py | 169 ++++++++++++++++++ 5 files changed, 254 insertions(+), 5 deletions(-) diff --git a/python/packages/declarative/AGENTS.md b/python/packages/declarative/AGENTS.md index 3bfc41fa1b..b43cee25e0 100644 --- a/python/packages/declarative/AGENTS.md +++ b/python/packages/declarative/AGENTS.md @@ -22,6 +22,9 @@ endpoint, label, connection, and headers, so separate fresh runs do not share a stateful MCP protocol session while continuations and checkpoint restores do. Direct fresh runs reset this scope even for message-list inputs; only framework-marked WorkflowAgent turns retain the agent continuation scope. +Sessionless WorkflowAgent turns select their own checkpointed scope at Entry, +independent of explicit agent sessions, direct runs, and restored approval scopes. +Approval continuations keep their originating active scope until the next Entry. With a provider, every invocation (including `tools/list`) gets a fresh tool/session, even if the provider returns `None` or a shared HTTP client. Invocation cleanup closes the session and any internally owned fallback client, never caller-owned HTTP clients. diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_control_flow.py b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_control_flow.py index 4aeacac41a..2a0d0c30f3 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_control_flow.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_control_flow.py @@ -410,14 +410,13 @@ async def handle_action( ctx: WorkflowContext[ActionComplete], ) -> None: """Simply pass through to continue the workflow.""" - is_agent_continuation = ctx.state.get(WORKFLOW_AGENT_RUN_KEY) is True and self._is_agent_continuation( - self._get_state(ctx.state), trigger - ) + is_agent_run = ctx.state.get(WORKFLOW_AGENT_RUN_KEY) is True + is_agent_continuation = is_agent_run and self._is_agent_continuation(self._get_state(ctx.state), trigger) await self._ensure_state_initialized(ctx, trigger) if self._action_def.get("kind") == "Entry": from ._mcp_handler import activate_workflow_session_id - activate_workflow_session_id(ctx.state, reset_unscoped=not is_agent_continuation) + activate_workflow_session_id(ctx.state, reset_unscoped=not is_agent_continuation, is_agent_run=is_agent_run) await ctx.send_message(ActionComplete()) diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py index 5a30a44bb3..d17115e85d 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -57,6 +57,7 @@ _DEFAULT_CACHE_MAX_SIZE = 32 _WORKFLOW_SESSION_ID_KEY = "_declarative_mcp_workflow_session_id" _AGENT_WORKFLOW_SESSION_NAMESPACE_KEY = "_declarative_mcp_agent_workflow_session_namespace" +_SESSIONLESS_AGENT_WORKFLOW_SESSION_ID_KEY = "_declarative_mcp_sessionless_agent_workflow_session_id" _WORKFLOW_AGENT_SESSION_ID_KEY = "_workflow_agent_session_id" @@ -88,9 +89,17 @@ def get_or_create_workflow_session_id(state: State) -> str: return workflow_session_id -def activate_workflow_session_id(state: State, *, reset_unscoped: bool) -> None: +def activate_workflow_session_id(state: State, *, reset_unscoped: bool, is_agent_run: bool = False) -> None: if state.get(_WORKFLOW_AGENT_SESSION_ID_KEY) is not None: get_or_create_workflow_session_id(state) + elif is_agent_run: + workflow_session_id = state.get(_SESSIONLESS_AGENT_WORKFLOW_SESSION_ID_KEY) + if workflow_session_id is None: + workflow_session_id = uuid.uuid4().hex + state.set(_SESSIONLESS_AGENT_WORKFLOW_SESSION_ID_KEY, workflow_session_id) + if not isinstance(workflow_session_id, str) or not workflow_session_id: + raise ValueError("Invalid MCP sessionless agent workflow session state.") + state.set(_WORKFLOW_SESSION_ID_KEY, workflow_session_id) elif reset_unscoped: state.set(_WORKFLOW_SESSION_ID_KEY, uuid.uuid4().hex) diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index ce68f2eee6..ed754b69ea 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -29,6 +29,7 @@ from agent_framework_declarative._workflows._mcp_handler import ( DefaultMCPToolHandler, MCPToolInvocation, + activate_workflow_session_id, get_or_create_workflow_session_id, restore_workflow_session_id, ) @@ -831,6 +832,47 @@ def test_invalid_cache_size_raises(self) -> None: class TestWorkflowSessionId: + def test_sessionless_activation_does_not_adopt_direct_scope(self) -> None: + state = State() + direct_scope = get_or_create_workflow_session_id(state) + + activate_workflow_session_id(state, reset_unscoped=False, is_agent_run=True) + sessionless_scope = get_or_create_workflow_session_id(state) + assert sessionless_scope != direct_scope + + activate_workflow_session_id(state, reset_unscoped=True) + assert get_or_create_workflow_session_id(state) not in (direct_scope, sessionless_scope) + state.commit() + + activate_workflow_session_id(state, reset_unscoped=False, is_agent_run=True) + assert get_or_create_workflow_session_id(state) == sessionless_scope + + def test_approval_restoration_preserves_sessionless_scope(self) -> None: + state = State() + activate_workflow_session_id(state, reset_unscoped=True, is_agent_run=True) + sessionless_scope = get_or_create_workflow_session_id(state) + state.set("_workflow_agent_session_id", "session-a") + explicit_scope = get_or_create_workflow_session_id(state) + + restore_workflow_session_id(state, explicit_scope) + state.commit() + assert get_or_create_workflow_session_id(state) == explicit_scope + assert explicit_scope != sessionless_scope + + activate_workflow_session_id(state, reset_unscoped=False, is_agent_run=True) + assert get_or_create_workflow_session_id(state) == sessionless_scope + + @pytest.mark.parametrize("scope", ["", 42, False]) + def test_invalid_sessionless_scope_preserves_current_scope(self, scope: Any) -> None: + state = State() + current_scope = get_or_create_workflow_session_id(state) + state.set("_declarative_mcp_sessionless_agent_workflow_session_id", scope) + + with pytest.raises(ValueError, match="Invalid MCP sessionless agent workflow session state"): + activate_workflow_session_id(state, reset_unscoped=True, is_agent_run=True) + + assert get_or_create_workflow_session_id(state) == current_scope + def test_restored_approval_scope_overrides_current_agent_session(self) -> None: state = State() state.set("_workflow_agent_session_id", "session-b") @@ -915,6 +957,33 @@ async def test_header_provider_returns_captured_headers(self) -> None: class TestCache: + @pytest.mark.parametrize("explicit_session", [False, True]) + async def test_sessionless_scope_reuses_its_own_cached_client(self, explicit_session: bool) -> None: + state = State() + with _patch_tool(): + async with DefaultMCPToolHandler() as handler: + activate_workflow_session_id(state, reset_unscoped=True, is_agent_run=True) + await handler.invoke_tool(_invocation(workflow_session_id=get_or_create_workflow_session_id(state))) + FakeTool.instances[0].call_handler = lambda **_: [Content.from_text("sessionless client")] + + if explicit_session: + state.set("_workflow_agent_session_id", "session-a") + activate_workflow_session_id(state, reset_unscoped=True, is_agent_run=explicit_session) + await handler.invoke_tool(_invocation(workflow_session_id=get_or_create_workflow_session_id(state))) + + state.set("_workflow_agent_session_id", None) + activate_workflow_session_id(state, reset_unscoped=False, is_agent_run=True) + result = await handler.invoke_tool( + _invocation(workflow_session_id=get_or_create_workflow_session_id(state)) + ) + + assert not result.is_error + assert len(result.outputs) == 1 + assert result.outputs[0].text == "sessionless client" + assert len(FakeTool.instances) == 2 + assert all(tool.connect_count == 1 for tool in FakeTool.instances) + assert all(tool.close_count == 1 for tool in FakeTool.instances) + @pytest.mark.asyncio async def test_same_url_and_headers_hit_cache(self) -> None: handler = DefaultMCPToolHandler() diff --git a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py index 64e9bce749..2233b00099 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -220,6 +220,175 @@ async def test_as_agent_continuation_reuses_workflow_session_id(self, streaming: assert handler.invocations[0].workflow_session_id assert handler.invocations[0].workflow_session_id == handler.invocations[1].workflow_session_id + @pytest.mark.parametrize("intervening_run", ["explicit_session", "direct_dict", "direct_message_list"]) + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("restore", [False, True]) + async def test_sessionless_agent_scope_survives_intervening_runs( + self, intervening_run: str, streaming: bool, restore: bool + ) -> None: + from agent_framework import InMemoryCheckpointStorage + + storage = InMemoryCheckpointStorage() + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + definition = _yaml(_action()) + workflow = factory.create_workflow_from_definition(definition) + agent = workflow.as_agent(name="mcp-agent") + + async def agent_turn(text: str, session: AgentSession | None = None) -> None: + if streaming: + await agent.run(text, session=session, stream=True, checkpoint_storage=storage).get_final_response() + else: + await agent.run(text, session=session, checkpoint_storage=storage) + + await agent_turn("first sessionless turn") + if intervening_run == "explicit_session": + await agent_turn("explicit session turn", AgentSession()) + else: + inputs: list[Message] | dict[str, Any] = ( + [Message(role="user", contents=[Content.from_text("direct run")])] + if intervening_run == "direct_message_list" + else {} + ) + if streaming: + await workflow.run(inputs, stream=True, checkpoint_storage=storage).get_final_response() + else: + await workflow.run(inputs, checkpoint_storage=storage) + + if restore: + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + workflow = factory.create_workflow_from_definition(definition) + await workflow.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage) + agent = workflow.as_agent(name="mcp-agent") + + await agent_turn("second sessionless turn") + + session_ids = [invocation.workflow_session_id for invocation in handler.invocations] + assert len(session_ids) == 3 + assert all(session_ids) + assert session_ids[0] != session_ids[1] + assert session_ids[0] == session_ids[2] + + @pytest.mark.parametrize("message_list", [False, True]) + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("restore", [False, True]) + async def test_sessionless_agent_does_not_inherit_direct_scope( + self, message_list: bool, streaming: bool, restore: bool + ) -> None: + from agent_framework import InMemoryCheckpointStorage + + storage = InMemoryCheckpointStorage() + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + definition = _yaml(_action()) + workflow = factory.create_workflow_from_definition(definition) + inputs: list[Message] | dict[str, Any] = ( + [Message(role="user", contents=[Content.from_text("direct run")])] if message_list else {} + ) + if streaming: + await workflow.run(inputs, stream=True, checkpoint_storage=storage).get_final_response() + else: + await workflow.run(inputs, checkpoint_storage=storage) + + if restore: + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + workflow = factory.create_workflow_from_definition(definition) + await workflow.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage) + + agent = workflow.as_agent(name="mcp-agent") + for text in ("first sessionless turn", "second sessionless turn"): + if streaming: + await agent.run(text, stream=True).get_final_response() + else: + await agent.run(text) + + session_ids = [invocation.workflow_session_id for invocation in handler.invocations] + assert len(session_ids) == 3 + assert all(session_ids) + assert session_ids[0] != session_ids[1] + assert session_ids[1] == session_ids[2] + + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("approved", [False, True]) + @pytest.mark.parametrize("restore", [False, True]) + async def test_sessionless_scope_survives_approval_origin_restoration( + self, streaming: bool, approved: bool, restore: bool + ) -> None: + from agent_framework import InMemoryCheckpointStorage + + from agent_framework_declarative._workflows import ToolApprovalResponse + + storage = InMemoryCheckpointStorage() + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + downstream = {**_action(tool_name="downstream"), "id": "downstream"} + definition = _yaml(_action(require_approval=True)) + definition["actions"].append(downstream) + workflow = factory.create_workflow_from_definition(definition) + agent = workflow.as_agent(name="mcp-agent") + + async def agent_turn(text: str, session: AgentSession | None = None) -> None: + if streaming: + await agent.run(text, session=session, stream=True, checkpoint_storage=storage).get_final_response() + else: + await agent.run(text, session=session, checkpoint_storage=storage) + + await agent_turn("first sessionless turn") + [sessionless_approval] = (await workflow._runner_context.get_pending_request_info_events()).values() + responses = {sessionless_approval.request_id: ToolApprovalResponse(approved=approved)} + if streaming: + await workflow.run(responses=responses, stream=True).get_final_response() + else: + await workflow.run(responses=responses) + assert [invocation.workflow_session_id for invocation in handler.invocations] == [ + sessionless_approval.data.workflow_session_id + ] * (2 if approved else 1) + + await agent_turn("explicit session turn", AgentSession()) + [explicit_approval] = (await workflow._runner_context.get_pending_request_info_events()).values() + assert sessionless_approval.data.workflow_session_id + assert explicit_approval.data.workflow_session_id + assert sessionless_approval.data.workflow_session_id != explicit_approval.data.workflow_session_id + + if restore: + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + workflow = factory.create_workflow_from_definition(definition) + await workflow.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage) + agent = workflow.as_agent(name="mcp-agent") + + responses = {explicit_approval.request_id: ToolApprovalResponse(approved=approved)} + if streaming: + await workflow.run(responses=responses, stream=True).get_final_response() + else: + await workflow.run(responses=responses) + assert [invocation.workflow_session_id for invocation in handler.invocations] == [ + scope + for scope in (sessionless_approval.data.workflow_session_id, explicit_approval.data.workflow_session_id) + for _ in range(2 if approved else 1) + ] + + await agent_turn("second sessionless turn") + [new_approval] = (await workflow._runner_context.get_pending_request_info_events()).values() + assert new_approval.data.workflow_session_id == sessionless_approval.data.workflow_session_id + responses = {new_approval.request_id: ToolApprovalResponse(approved=approved)} + if streaming: + await workflow.run(responses=responses, stream=True).get_final_response() + else: + await workflow.run(responses=responses) + assert [invocation.workflow_session_id for invocation in handler.invocations] == [ + scope + for scope in ( + sessionless_approval.data.workflow_session_id, + explicit_approval.data.workflow_session_id, + sessionless_approval.data.workflow_session_id, + ) + for _ in range(2 if approved else 1) + ] + assert not await workflow._runner_context.get_pending_request_info_events() + @pytest.mark.asyncio async def test_as_agent_sessions_isolate_and_reuse_workflow_session_ids(self) -> None: handler = StubMcpHandler(_ok()) From ab67a017947af462d3833030179e4fd65a0b29d0 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Thu, 8 Oct 2026 15:11:23 -0400 Subject: [PATCH 41/42] chore: Update _close_claimed_entry to manage cleanup outcome Modify _close_claimed_entry to handle cleanup outcome. Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- .../agent_framework_declarative/_workflows/_mcp_handler.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py index d17115e85d..e6f8f50f8d 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -822,10 +822,11 @@ async def _release_entry_core(self, entry: _CacheEntry) -> None: await self._close_claimed_entry(entry) async def _close_claimed_entry(self, entry: _CacheEntry) -> None: + cleanup_outcome: asyncio.Future[None] = asyncio.get_running_loop().create_future() try: - await self._close_invocation_entry(entry) + await self._close_invocation_entry(entry, cleanup_outcome) except BaseException as exc: - entry.close_exception = exc + entry.close_exception = cleanup_outcome.exception() if cleanup_outcome.done() else exc raise finally: async with self._cache_lock: From c8df5159b570fa24ecaa3536c2fc95df7aa25b2f Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Thu, 8 Oct 2026 15:29:43 -0400 Subject: [PATCH 42/42] Fix background-agent stale-waiter test thread-pool starvation Use a dedicated thread for blocking lock coordination and remove the unnecessary wait for continuation scheduling. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: f1abc0c2-cf28-4547-9bac-d682b4357e03 --- .../BackgroundAgents/BackgroundAgentsProviderTests.cs | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/BackgroundAgents/BackgroundAgentsProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/BackgroundAgents/BackgroundAgentsProviderTests.cs index e287fdb7b9..7419362c53 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/BackgroundAgents/BackgroundAgentsProviderTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/BackgroundAgents/BackgroundAgentsProviderTests.cs @@ -452,7 +452,6 @@ public async Task WaitForFirstCompletion_StaleWaiterLeavesContinuationRunningAsy // Arrange var firstResponse = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); var secondResponse = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); - var continuationStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); int runCount = 0; var agent = CreateMockAgentWithCallback("Research", () => { @@ -461,7 +460,6 @@ public async Task WaitForFirstCompletion_StaleWaiterLeavesContinuationRunningAsy return firstResponse.Task; } - continuationStarted.SetResult(true); return secondResponse.Task; }); var (tools, provider, session) = await CreateToolsWithSessionAsync(agent); @@ -489,8 +487,9 @@ await startBackgroundTask.InvokeAsync(new AIFunctionArguments Assert.False(wait.IsCompleted); // Act — hold the lock until the continuation is registered so the old waiter observes the new run. + // Use a dedicated thread so blocking under the lock does not starve the background runs. ValueTask continuation = default; - await Task.Run(() => + await Task.Factory.StartNew(() => { lock (runtimeState.SyncRoot) { @@ -506,10 +505,9 @@ await Task.Run(() => continuationRun = runtimeState.InFlightTasks[1]; continuationCancellation = runtimeState.TaskCancellations[1]; Assert.NotSame(firstRun, continuationRun); - Assert.True(SpinWait.SpinUntil(() => continuationStarted.Task.IsCompleted, TimeSpan.FromSeconds(5))); Assert.False(wait.IsCompleted); } - }); + }, CancellationToken.None, TaskCreationOptions.LongRunning, TaskScheduler.Default); // Assert — the stale waiter preserves the continuation's metadata and runtime resources. Assert.Equal("Task 1 continued with new input.", GetStringResult(await continuation));