From bb9e92b189edc1367edc225a08cf8924657e82f8 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 22 Sep 2026 14:08:08 -0400 Subject: [PATCH 01/58] 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 b9c465f9f03..500874a1c93 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 ab058de62d4..3cdebd5e1a0 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net10.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net10.0/PublicAPI.Unshipped.txt @@ -1 +1,2 @@ #nullable enable +Microsoft.Agents.AI.Workflows.Declarative.Mcp.DefaultMcpToolHandler.InvokeToolInWorkflowSessionAsync(string! serverUrl, string? serverLabel, string! toolName, System.Collections.Generic.IDictionary? arguments, System.Collections.Generic.IDictionary? headers, string? connectionName, string! workflowSessionId, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net472/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net472/PublicAPI.Unshipped.txt index ab058de62d4..3cdebd5e1a0 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net472/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net472/PublicAPI.Unshipped.txt @@ -1 +1,2 @@ #nullable enable +Microsoft.Agents.AI.Workflows.Declarative.Mcp.DefaultMcpToolHandler.InvokeToolInWorkflowSessionAsync(string! serverUrl, string? serverLabel, string! toolName, System.Collections.Generic.IDictionary? arguments, System.Collections.Generic.IDictionary? headers, string? connectionName, string! workflowSessionId, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net8.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net8.0/PublicAPI.Unshipped.txt index ab058de62d4..3cdebd5e1a0 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net8.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net8.0/PublicAPI.Unshipped.txt @@ -1 +1,2 @@ #nullable enable +Microsoft.Agents.AI.Workflows.Declarative.Mcp.DefaultMcpToolHandler.InvokeToolInWorkflowSessionAsync(string! serverUrl, string? serverLabel, string! toolName, System.Collections.Generic.IDictionary? arguments, System.Collections.Generic.IDictionary? headers, string? connectionName, string! workflowSessionId, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net9.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net9.0/PublicAPI.Unshipped.txt index ab058de62d4..3cdebd5e1a0 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net9.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net9.0/PublicAPI.Unshipped.txt @@ -1 +1,2 @@ #nullable enable +Microsoft.Agents.AI.Workflows.Declarative.Mcp.DefaultMcpToolHandler.InvokeToolInWorkflowSessionAsync(string! serverUrl, string? serverLabel, string! toolName, System.Collections.Generic.IDictionary? arguments, System.Collections.Generic.IDictionary? headers, string? connectionName, string! workflowSessionId, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt index ab058de62d4..3cdebd5e1a0 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt @@ -1 +1,2 @@ #nullable enable +Microsoft.Agents.AI.Workflows.Declarative.Mcp.DefaultMcpToolHandler.InvokeToolInWorkflowSessionAsync(string! serverUrl, string? serverLabel, string! toolName, System.Collections.Generic.IDictionary? arguments, System.Collections.Generic.IDictionary? headers, string? connectionName, string! workflowSessionId, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/IMcpToolHandler.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/IMcpToolHandler.cs index 56b1c3deb4c..fbea21f35b4 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/IMcpToolHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/IMcpToolHandler.cs @@ -39,3 +39,34 @@ Task InvokeToolAsync( string? connectionName, CancellationToken cancellationToken = default); } + +/// +/// Defines the contract for MCP handlers that isolate stateful protocol sessions by workflow session. +/// +public interface IWorkflowScopedMcpToolHandler : IMcpToolHandler +{ + /// + /// Invokes an MCP tool within the framework-owned workflow session scope. + /// + /// The URL of the MCP server. + /// An optional label identifying the server connection. + /// The name of the tool to invoke. + /// Optional arguments to pass to the tool. + /// Optional headers to include in the request. + /// An optional connection name for managed connections. + /// The framework-owned identifier for the current workflow session. + /// A token to observe cancellation. + /// + /// A task representing the asynchronous operation. The result contains a + /// with the tool invocation output. + /// + Task InvokeToolInWorkflowSessionAsync( + string serverUrl, + string? serverLabel, + string toolName, + IDictionary? arguments, + IDictionary? headers, + string? connectionName, + string workflowSessionId, + CancellationToken cancellationToken = default); +} diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowContext.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowContext.cs index 6616aa5d00c..92ffbe56a16 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 46a5cae5fdd..53fedcdb0d4 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 aaff60b08b4..f8370d722eb 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 a36677538eb..bf09693fe40 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 a36677538eb..bf09693fe40 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 a36677538eb..bf09693fe40 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 a36677538eb..bf09693fe40 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 a36677538eb..bf09693fe40 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 ab94a9fa8aa..410316a0108 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 b8b35fffd64..68ff864085f 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/IWorkflowContext.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/IWorkflowContext.cs @@ -195,3 +195,18 @@ ValueTask ReadOrInitStateAsync(string key, Func initialStateFactory, Ca /// bool ConcurrentRunsEnabled { get; } } + +/// +/// Exposes the framework-owned identifier for a workflow session. +/// +/// +/// Runtime services can use this identifier to scope stateful resources to one +/// workflow session without accepting caller-provided cache partition keys. +/// +public interface IWorkflowSessionContext +{ + /// + /// Gets the framework-owned workflow session identifier. + /// + string SessionId { get; } +} diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunnerContext.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunnerContext.cs index 353c38f25c9..0db031c9a00 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 ab058de62d4..724a27a74a4 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net10.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net10.0/PublicAPI.Unshipped.txt @@ -1 +1,3 @@ #nullable enable +Microsoft.Agents.AI.Workflows.IWorkflowSessionContext +Microsoft.Agents.AI.Workflows.IWorkflowSessionContext.SessionId.get -> string! diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net472/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net472/PublicAPI.Unshipped.txt index ab058de62d4..724a27a74a4 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net472/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net472/PublicAPI.Unshipped.txt @@ -1 +1,3 @@ #nullable enable +Microsoft.Agents.AI.Workflows.IWorkflowSessionContext +Microsoft.Agents.AI.Workflows.IWorkflowSessionContext.SessionId.get -> string! diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net8.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net8.0/PublicAPI.Unshipped.txt index ab058de62d4..724a27a74a4 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net8.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net8.0/PublicAPI.Unshipped.txt @@ -1 +1,3 @@ #nullable enable +Microsoft.Agents.AI.Workflows.IWorkflowSessionContext +Microsoft.Agents.AI.Workflows.IWorkflowSessionContext.SessionId.get -> string! diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net9.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net9.0/PublicAPI.Unshipped.txt index ab058de62d4..724a27a74a4 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net9.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net9.0/PublicAPI.Unshipped.txt @@ -1 +1,3 @@ #nullable enable +Microsoft.Agents.AI.Workflows.IWorkflowSessionContext +Microsoft.Agents.AI.Workflows.IWorkflowSessionContext.SessionId.get -> string! diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt index ab058de62d4..724a27a74a4 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt @@ -1 +1,3 @@ #nullable enable +Microsoft.Agents.AI.Workflows.IWorkflowSessionContext +Microsoft.Agents.AI.Workflows.IWorkflowSessionContext.SessionId.get -> string! diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs index 3a481fcd0e6..37064cbbd92 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 5a5b318f055..21ddf07416e 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 bb34b6c9a26..ded8316083a 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 9638b435337..332db4cec00 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 7d3647d00bd..3edd91c63de 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 64eb59750fb..f0f4c77434b 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 0826302cf69..58cadb183dc 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 28b3a1a0d67..c0896c4c6f3 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/58] .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 500874a1c93..effa492f680 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 5d052c64d3d..631519288bb 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 92ffbe56a16..a60a7756f52 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 1126b43d8ff..2fcdea4ecce 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowExecutor.cs @@ -154,7 +154,8 @@ private async ValueTask AdvanceAsync(ChatMessage input, IWorkflowContext context // No state to restore if we're starting from the beginning. state.SetInitialized(); - DeclarativeWorkflowContext declarativeContext = new(context, state); + DeclarativeWorkflowContext declarativeContext = + await DeclarativeWorkflowContext.CreateAsync(context, state, cancellationToken).ConfigureAwait(false); // Conversation id resolution prefers state already persisted by a prior turn, // so multi-turn invocations reuse the same backend conversation rather than diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DelegateActionExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DelegateActionExecutor.cs index 716a3656dac..ade8dd1034b 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DelegateActionExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DelegateActionExecutor.cs @@ -61,7 +61,9 @@ public override async ValueTask HandleAsync(TMessage message, IWorkflowContext c { if (this._action is not null) { - await this._action.Invoke(new DeclarativeWorkflowContext(context, this._state), message, cancellationToken).ConfigureAwait(false); + DeclarativeWorkflowContext declarativeContext = + await DeclarativeWorkflowContext.CreateAsync(context, this._state, cancellationToken).ConfigureAwait(false); + await this._action.Invoke(declarativeContext, message, cancellationToken).ConfigureAwait(false); } if (this._emitResult) diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/ActionExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/ActionExecutor.cs index cf636effaf8..563ff86f147 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/ActionExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/ActionExecutor.cs @@ -76,7 +76,9 @@ public ValueTask ResetAsync() [SendsMessage(typeof(ActionExecutorResult))] public override async ValueTask HandleAsync(TMessage message, IWorkflowContext context, CancellationToken cancellationToken) { - object? result = await this.ExecuteAsync(new DeclarativeWorkflowContext(context, this._session.State), message, cancellationToken).ConfigureAwait(false); + DeclarativeWorkflowContext declarativeContext = + await DeclarativeWorkflowContext.CreateAsync(context, this._session.State, cancellationToken).ConfigureAwait(false); + object? result = await this.ExecuteAsync(declarativeContext, message, cancellationToken).ConfigureAwait(false); Debug.WriteLine($"RESULT #{this.Id} - {result ?? "(null)"}"); await context.SendResultMessageAsync(this.Id, result, cancellationToken).ConfigureAwait(false); diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/RootExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/RootExecutor.cs index 80f6e69b60d..bf6ab5871fe 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 f8370d722eb..aaff60b08b4 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 410316a0108..fbda33c132a 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/GroupChatManager.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/GroupChatManager.cs @@ -152,7 +152,7 @@ protected virtual ValueTask OnCheckpointRestoredAsync(IWorkflowContext context, internal async ValueTask CheckpointAsync(IWorkflowContext context, CancellationToken cancellationToken = default) { await context.QueueStateUpdateAsync(BaseStateKey, new GroupChatManagerState(this.IterationCount), cancellationToken: cancellationToken).ConfigureAwait(false); - await this.OnCheckpointingAsync(new PrefixingWorkflowContext(context, SubclassStateKeyPrefix), cancellationToken).ConfigureAwait(false); + await this.OnCheckpointingAsync(PrefixingWorkflowContext.Create(context, SubclassStateKeyPrefix), cancellationToken).ConfigureAwait(false); } // Root restore entry point invoked by the hosting GroupChatHost. Symmetric to CheckpointAsync. @@ -160,7 +160,7 @@ internal async ValueTask RestoreCheckpointAsync(IWorkflowContext context, Cancel { GroupChatManagerState? state = await context.ReadStateAsync(BaseStateKey, cancellationToken: cancellationToken).ConfigureAwait(false); this.IterationCount = state?.IterationCount ?? 0; - await this.OnCheckpointRestoredAsync(new PrefixingWorkflowContext(context, SubclassStateKeyPrefix), cancellationToken).ConfigureAwait(false); + await this.OnCheckpointRestoredAsync(PrefixingWorkflowContext.Create(context, SubclassStateKeyPrefix), cancellationToken).ConfigureAwait(false); } } @@ -169,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 37064cbbd92..e87714025c0 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 5193296db4d..eb627e24c78 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 f0f4c77434b..95280353c42 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 58cadb183dc..b43a4243a1e 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/58] 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 effa492f680..340b8c09a29 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 e87714025c0..7199cbef904 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 95280353c42..e4f6330e411 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 b43a4243a1e..b55f8086ee8 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/58] 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 3edd91c63de..fef61859393 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 e4f6330e411..cebf8a08a01 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 b55f8086ee8..48cabc94c19 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/58] 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 340b8c09a29..f80b8728aaa 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 70c6ada3aba..f17265f78c2 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 7199cbef904..14a027420ed 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 4199aa96f2c..1ba4cb17031 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 cebf8a08a01..820bd193231 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 48cabc94c19..877a759fdeb 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/58] 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 f80b8728aaa..9ddc283c7a2 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 f17265f78c2..47962fe5e5d 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 53fedcdb0d4..3b09518ca50 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 14a027420ed..8a8184f0a75 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 ded8316083a..be97b96aa82 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 820bd193231..4c3cec2529e 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 877a759fdeb..9e90d10dc5b 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/58] 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 9ddc283c7a2..bc31709a63c 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 8a8184f0a75..dbf2f9dc686 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 332db4cec00..0679795d516 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 5db554574f4..a612f9aef08 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 4c3cec2529e..d5d3f5ab6b5 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 9e90d10dc5b..494cb8e3938 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 c0896c4c6f3..3b450f4a340 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/58] 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 d5d3f5ab6b5..2658a5f8671 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 494cb8e3938..017d127ee47 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/58] 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 47962fe5e5d..f88035abd12 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/58] 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 3b09518ca50..ed74d30d39d 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/58] 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 2658a5f8671..757f1964aee 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 017d127ee47..751fe0587b6 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/58] 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 757f1964aee..5fb5444a394 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/58] 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 5fb5444a394..c2697596c23 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/58] 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 751fe0587b6..12f43875d3a 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/58] 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 2fee60b8c69..f335a1fbc85 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/58] 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 12f43875d3a..e8c8bbae9a0 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/58] 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 e8c8bbae9a0..ddc87f0ea04 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/58] 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 f335a1fbc85..772f6ea9b13 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 dbf2f9dc686..a57f9017693 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 c2697596c23..f34406f7350 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 ddc87f0ea04..ffee74c68ae 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/58] 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 a57f9017693..6b27a3dce27 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/58] 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 772f6ea9b13..f6b33e35574 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 6b27a3dce27..1e313753575 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/58] 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 f6b33e35574..976d2b99815 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/58] 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 11f7cf493cf..538e8bc3bbb 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 1e92d130de9..6b8a0f79968 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 f34406f7350..a0d84d8843f 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 ffee74c68ae..770801afd88 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/58] 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 770801afd88..8cdb26391d0 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/58] 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 69b6c1e9bcc..c0a2e075ae1 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunner.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunner.cs @@ -31,7 +31,14 @@ public static InProcessRunner CreateTopLevelRunner(Workflow workflow, ICheckpoin knownValidInputTypes: knownValidInputTypes); } - public static InProcessRunner CreateSubworkflowRunner(Workflow workflow, ICheckpointManager? checkpointManager, string? sessionId = null, object? existingOwnerSignoff = null, bool enableConcurrentRuns = false, IEnumerable? knownValidInputTypes = null) + public static InProcessRunner CreateSubworkflowRunner( + Workflow workflow, + ICheckpointManager? checkpointManager, + string? sessionId = null, + object? existingOwnerSignoff = null, + bool enableConcurrentRuns = false, + IEnumerable? knownValidInputTypes = null, + string? workflowSessionId = null) { return new InProcessRunner(workflow, checkpointManager, @@ -39,10 +46,19 @@ public static InProcessRunner CreateSubworkflowRunner(Workflow workflow, ICheckp existingOwnerSignoff: existingOwnerSignoff, enableConcurrentRuns: enableConcurrentRuns, knownValidInputTypes: knownValidInputTypes, - subworkflow: true); + subworkflow: true, + workflowSessionId: workflowSessionId); } - private InProcessRunner(Workflow workflow, ICheckpointManager? checkpointManager, string? sessionId = null, object? existingOwnerSignoff = null, bool subworkflow = false, bool enableConcurrentRuns = false, IEnumerable? knownValidInputTypes = null) + private InProcessRunner( + Workflow workflow, + ICheckpointManager? checkpointManager, + string? sessionId = null, + object? existingOwnerSignoff = null, + bool subworkflow = false, + bool enableConcurrentRuns = false, + IEnumerable? knownValidInputTypes = null, + string? workflowSessionId = null) { if (enableConcurrentRuns && !workflow.AllowConcurrent) { @@ -54,7 +70,16 @@ private InProcessRunner(Workflow workflow, ICheckpointManager? checkpointManager this.StartExecutorId = workflow.StartExecutorId; this.Workflow = Throw.IfNull(workflow); - this.RunContext = new InProcessRunnerContext(workflow, this.SessionId, checkpointingEnabled: checkpointManager != null, this.OutgoingEvents, this.StepTracer, existingOwnerSignoff, subworkflow, enableConcurrentRuns); + this.RunContext = new InProcessRunnerContext( + workflow, + this.SessionId, + checkpointingEnabled: checkpointManager != null, + this.OutgoingEvents, + this.StepTracer, + existingOwnerSignoff, + subworkflow, + enableConcurrentRuns, + workflowSessionId: workflowSessionId); this.CheckpointManager = checkpointManager; this._knownValidInputTypes = knownValidInputTypes != null @@ -65,6 +90,8 @@ private InProcessRunner(Workflow workflow, ICheckpointManager? checkpointManager /// public string SessionId { get; } + internal string WorkflowSessionId => this.RunContext.WorkflowSessionId; + /// public string StartExecutorId { get; } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunnerContext.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunnerContext.cs index 0db031c9a00..aa8bf398592 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunnerContext.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunnerContext.cs @@ -47,7 +47,8 @@ public InProcessRunnerContext( object? existingOwnershipSignoff = null, bool subworkflow = false, bool enableConcurrentRuns = false, - ILogger? logger = null) + ILogger? logger = null, + string? workflowSessionId = null) { if (enableConcurrentRuns) { @@ -62,6 +63,7 @@ public InProcessRunnerContext( this._workflow = workflow; this._sessionId = sessionId; + this.WorkflowSessionId = workflowSessionId ?? sessionId; this._edgeMap = new(this, this._workflow, stepTracer); this._outputFilter = new(workflow); @@ -72,6 +74,8 @@ public InProcessRunnerContext( } public WorkflowTelemetryContext TelemetryContext => this._workflow.TelemetryContext; + internal string WorkflowSessionId { get; } + public IExternalRequestSink RegisterPort(string executorId, RequestPort port) { if (!this._edgeMap.TryRegisterPort(this, executorId, port)) @@ -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 58e3a9e5239..a166d8b0e6b 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/WorkflowHostExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/WorkflowHostExecutor.cs @@ -20,6 +20,7 @@ internal class WorkflowHostExecutor : Executor, IAsyncDisposable private readonly Workflow _workflow; private readonly ProtocolDescriptor _workflowProtocol; private readonly object _ownershipToken; + private string? _workflowSessionId; private InProcessRunner? _activeRunner; private InMemoryCheckpointManager? _checkpointManager; @@ -56,6 +57,7 @@ protected override ProtocolBuilder ConfigureProtocol(ProtocolBuilder protocolBui private async ValueTask QueueExternalMessageAsync(PortableValue portableValue, IWorkflowContext context, CancellationToken cancellationToken) { + this.SetWorkflowSessionId(context); if (portableValue.Is(out ExternalResponse? response)) { response = this.CheckAndUnqualifyResponse(response); @@ -95,7 +97,8 @@ internal async ValueTask EnsureRunnerAsync() this._checkpointManager, this._sessionId, this._ownershipToken, - this.JoinContext.ConcurrentRunsEnabled); + this.JoinContext.ConcurrentRunsEnabled, + workflowSessionId: this._workflowSessionId); } return this._activeRunner; @@ -268,6 +271,7 @@ await context.QueueStateUpdateAsync(PendingResponsePortsStateKey, protected internal override async ValueTask OnCheckpointRestoredAsync(IWorkflowContext context, CancellationToken cancellationToken = default) { await base.OnCheckpointRestoredAsync(context, cancellationToken).ConfigureAwait(false); + this.SetWorkflowSessionId(context); InMemoryCheckpointManager manager = await context.ReadStateAsync(CheckpointManagerStateKey, cancellationToken: cancellationToken).ConfigureAwait(false) ?? new(); if (this._checkpointManager == manager) @@ -293,6 +297,14 @@ await context.ReadStateAsync>(PendingRespons await this.EnsureRunSendMessageAsync(resume: true, cancellationToken: cancellationToken).ConfigureAwait(false); } + private void SetWorkflowSessionId(IWorkflowContext context) + { + string parentSessionId = context is IWorkflowSessionContext sessionContext + ? sessionContext.SessionId + : this._sessionId; + this._workflowSessionId ??= SubworkflowBinding.CreateSubworkflowSessionId(parentSessionId, this.Id); + } + private async ValueTask ResetAsync() { if (this._run != null) diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/SubworkflowBinding.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/SubworkflowBinding.cs index 538e8bc3bbb..9519af12414 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 6b8a0f79968..89928c889c4 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 a0d84d8843f..a4f9b175bc6 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 8cdb26391d0..7661496ff27 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/58] 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 89928c889c4..bddef5e4f19 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/58] 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 976d2b99815..c209fa65ea9 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 a4f9b175bc6..685ae9369b3 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 7661496ff27..b655694d31a 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/58] 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 b655694d31a..a76215b1bf1 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/58] 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 3b450f4a340..9fa9b1eeb9a 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/58] 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 c209fa65ea9..53edfc30e54 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 b98bbd4f87a..fcdf522d483 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/58] 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 0679795d516..90eaa2cf36b 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 a9eb95473d9..73e89235c19 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 a612f9aef08..bb61fbf0eb5 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 fef61859393..5f13045d3be 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 685ae9369b3..2da5232b3b8 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 a76215b1bf1..5801f1347cc 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 9fa9b1eeb9a..7088025e8b1 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/58] 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 53edfc30e54..b82bfc6ec6b 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 fcdf522d483..6ea27b4314a 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 6d01c8366a4..9ca1e5928f7 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 db99b45426a..dac6b96ba72 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 b69838995d5..f2a839a8e35 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 bb61fbf0eb5..dd95d9fa7e8 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 2da5232b3b8..b59c9d8fe66 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 5801f1347cc..8df29de59f3 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 7088025e8b1..0066e0a7aa9 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/58] 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 8df29de59f3..c022d70806f 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/58] 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 6ea27b4314a..5390b822d64 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/58] 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 9ca1e5928f7..c9511dc0953 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 0ef70ca2f6e..4ef9fdd82b0 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/58] 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 b82bfc6ec6b..31d955d7103 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 5390b822d64..a9820c16dcf 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 8a7bb4e21a7..bfab8c96bfe 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 4ef9fdd82b0..605056bea23 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 0066e0a7aa9..94c1f3261a4 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/58] 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 013c7ad6e68..179e4eb64cf 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 5f13045d3be..7eb3e97c0d3 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 b59c9d8fe66..5a30a44bb36 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 c022d70806f..ce68f2eee65 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 94c1f3261a4..ac6bd924c25 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/58] 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 c9511dc0953..7b86cf15c51 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 dac6b96ba72..7f096c57427 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 bfab8c96bfe..301f5383eab 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 605056bea23..2faaa76b935 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 179e4eb64cf..8aadeee9bbe 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 dd95d9fa7e8..4aeacac41a7 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 ac6bd924c25..fe7650e6dc0 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/58] 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 31d955d7103..53cb8507218 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 a9820c16dcf..67390dd543b 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 301f5383eab..03e1e4f8d59 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 8aadeee9bbe..3bfc41fa1bd 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 9a6814af60d..bc56946f857 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 fe7650e6dc0..41c9cad89c6 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/58] 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 41c9cad89c6..64e9bce749a 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/58] 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 3bfc41fa1bd..b43cee25e07 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 4aeacac41a7..2a0d0c30f3e 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 5a30a44bb36..d17115e85d7 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 ce68f2eee65..ed754b69ea8 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 64e9bce749a..2233b000993 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/58] 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 d17115e85d7..e6f8f50f8d2 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/58] 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 e287fdb7b9b..7419362c53f 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)); From 839e4f0af358979d70dfc419f3037c634b3850e5 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Thu, 15 Oct 2026 12:27:12 -0400 Subject: [PATCH 43/58] fix: finish MCP cache cleanup and preserve entry checkpoint identity Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 48ae509d-294c-4bf5-864d-81baecbae204 --- .../agent_framework/_workflows/_workflow.py | 4 ++ python/packages/declarative/AGENTS.md | 5 +- .../_workflows/_declarative_builder.py | 9 +++- .../_workflows/_executors_control_flow.py | 7 --- .../_workflows/_mcp_handler.py | 14 ++++-- .../tests/test_default_mcp_tool_handler.py | 48 ++++++++++++++++++- .../tests/test_invoke_mcp_tool_executor.py | 36 ++++++++++++++ 7 files changed, 108 insertions(+), 15 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index ecf2bb7ad00..595f4c6fdc0 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -414,6 +414,7 @@ def __init__( # 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 + self._initial_state_initializer: Callable[[State], 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`` @@ -768,6 +769,9 @@ async def _execute_with_message_or_checkpoint( f"Expected input types: {start_executor.input_types}." ) + if self._initial_state_initializer is not None: + self._initial_state_initializer(self._runner.state) + await self._runner.context.send_message(entry_message) # Record the entry checkpoint capturing the seeded input before any # executor runs. The runner only checkpoints after each superstep. diff --git a/python/packages/declarative/AGENTS.md b/python/packages/declarative/AGENTS.md index 5b62a2c1fcc..2699cfb8565 100644 --- a/python/packages/declarative/AGENTS.md +++ b/python/packages/declarative/AGENTS.md @@ -29,9 +29,10 @@ 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, +Run scopes are selected before the initial entry checkpoint and preserved when +Entry replays. Sessionless WorkflowAgent turns select their own checkpointed scope, independent of explicit agent sessions, direct runs, and restored approval scopes. -Approval continuations keep their originating active scope until the next Entry. +Approval continuations keep their originating active scope until the next run with new input. 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/_declarative_builder.py b/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_builder.py index e6bcce76809..544916ea70f 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_builder.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_builder.py @@ -22,6 +22,8 @@ WorkflowBuilder, WorkflowEvent, ) +from agent_framework._workflows._const import WORKFLOW_AGENT_RUN_KEY +from agent_framework._workflows._state import State from ._declarative_base import ( ConditionResult, @@ -50,7 +52,7 @@ ) from ._executors_tools import TOOL_ACTION_EXECUTORS, InvokeFunctionToolExecutor from ._http_handler import HttpRequestHandler -from ._mcp_handler import MCPToolHandler +from ._mcp_handler import MCPToolHandler, activate_workflow_session_id logger = logging.getLogger(__name__) @@ -128,6 +130,10 @@ def _validate_mcp_response_batch(requests: Sequence[WorkflowEvent[Any]]) -> None ) +def _initialize_mcp_run_state(state: State) -> None: + activate_workflow_session_id(state, reset_unscoped=True, is_agent_run=state.get(WORKFLOW_AGENT_RUN_KEY) is True) + + class DeclarativeWorkflowBuilder: """Builds a Workflow graph from declarative YAML actions. @@ -262,6 +268,7 @@ def build(self) -> Workflow: workflow = builder.build() workflow._response_batch_validator = _validate_mcp_response_batch # pyright: ignore[reportPrivateUsage] + workflow._initial_state_initializer = _initialize_mcp_run_state # pyright: ignore[reportPrivateUsage] return workflow def _collect_explicit_ids(self, actions: list[dict[str, Any]] | None) -> None: 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 9826cf361be..51c36ffdd61 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,7 +21,6 @@ WorkflowContext, handler, ) -from agent_framework._workflows._const import WORKFLOW_AGENT_RUN_KEY from ._declarative_base import ( ActionComplete, @@ -414,13 +413,7 @@ async def handle_action( ctx: WorkflowContext[ActionComplete], ) -> None: """Simply pass through to continue the workflow.""" - 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, 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 2c0d0189460..e4f9569288e 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -813,14 +813,17 @@ 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)) + await self._await_cleanup_task(asyncio.create_task(self._release_entry_core(entry))) + + @staticmethod + async def _await_cleanup_task(cleanup: asyncio.Task[None]) -> None: cancelled = False - while not release.done(): + while not cleanup.done(): try: - await asyncio.shield(release) + await asyncio.shield(cleanup) except asyncio.CancelledError: cancelled = True - release.result() + cleanup.result() if cancelled: raise asyncio.CancelledError @@ -836,6 +839,9 @@ async def _release_entry_core(self, entry: _CacheEntry) -> None: await self._close_claimed_entry(entry) async def _close_claimed_entry(self, entry: _CacheEntry) -> None: + await self._await_cleanup_task(asyncio.create_task(self._close_claimed_entry_core(entry))) + + async def _close_claimed_entry_core(self, entry: _CacheEntry) -> None: cleanup_outcome: asyncio.Future[None] = asyncio.get_running_loop().create_future() try: await self._close_invocation_entry(entry, cleanup_outcome) 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 4340b34fed5..33d25c12279 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -17,7 +17,7 @@ import asyncio import json import sys -from typing import Any +from typing import Any, Literal from unittest.mock import AsyncMock, patch import httpx2 as httpx @@ -1089,6 +1089,52 @@ async def close(tool: FakeTool) -> None: cancel_first_close = False await asyncio.wait_for(handler.aclose(), timeout=1) + async def test_repeated_cancellation_during_eviction_finalization_completes_shutdown(self) -> None: + handler = DefaultMCPToolHandler(cache_max_size=1) + close_started = asyncio.Event() + release_close = asyncio.Event() + finalization_waiting = asyncio.Event() + original_close = FakeTool.close + + class ObservedLock(asyncio.Lock): + async def acquire(self) -> Literal[True]: + if self.locked(): + finalization_waiting.set() + return await super().acquire() + + handler._cache_lock = ObservedLock() + + async def gated_close(tool: FakeTool) -> None: + if tool.kwargs["url"] == "https://a/": + close_started.set() + await release_close.wait() + await original_close(tool) + + with _patch_tool(), patch.object(FakeTool, "close", gated_close): + await handler.invoke_tool(_invocation(server_url="https://a/")) + retired_entry = next(iter(handler._cache.values())) + invocation = asyncio.create_task(handler.invoke_tool(_invocation(server_url="https://b/"))) + await close_started.wait() + invocation.cancel() + await handler._cache_lock.acquire() + try: + release_close.set() + await asyncio.wait_for(finalization_waiting.wait(), timeout=1) + invocation.cancel() + await asyncio.sleep(0) + finally: + handler._cache_lock.release() + + await asyncio.wait({invocation}, timeout=1) + assert invocation.done() + with pytest.raises(asyncio.CancelledError): + invocation.result() + assert retired_entry.closed.is_set() + assert id(retired_entry) not in handler._retired + await asyncio.wait_for(handler.aclose(), timeout=1) + + assert all(tool.close_count == 1 for tool in FakeTool.instances) + @pytest.mark.asyncio async def test_repeated_cancellation_while_releasing_entry_does_not_leak_lease(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 a893767550f..d3a64f0669e 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,42 @@ 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("run_kind", ["direct", "sessionless_agent", "explicit_agent"]) + @pytest.mark.parametrize("streaming", [False, True]) + async def test_entry_checkpoint_replay_preserves_workflow_session_id(self, run_kind: str, streaming: 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) + session = AgentSession() if run_kind == "explicit_agent" else None + + if run_kind == "direct": + if streaming: + await workflow.run({}, stream=True, checkpoint_storage=storage).get_final_response() + else: + await workflow.run({}, checkpoint_storage=storage) + else: + agent = workflow.as_agent() + if streaming: + await agent.run( + "first turn", session=session, stream=True, checkpoint_storage=storage + ).get_final_response() + else: + await agent.run("first turn", session=session, checkpoint_storage=storage) + + [original] = handler.invocations + checkpoints = await storage.list_checkpoints(workflow_name=workflow.name) + [entry_checkpoint] = [checkpoint for checkpoint in checkpoints if checkpoint.iteration_count == 0] + workflow = factory.create_workflow_from_definition(definition) + await workflow.run(checkpoint_id=entry_checkpoint.checkpoint_id, checkpoint_storage=storage) + + assert len(handler.invocations) == 2 + assert original.workflow_session_id + assert handler.invocations[1].workflow_session_id == original.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]) From b4d4d1a3b6e4a7b4b8c905197f6218f685a2e49e Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Thu, 15 Oct 2026 13:54:25 -0400 Subject: [PATCH 44/58] fix: bound MCP client lifetimes and align workflow dependencies Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 48ae509d-294c-4bf5-864d-81baecbae204 --- .../DefaultMcpToolHandler.cs | 54 ++++--- .../DefaultMcpToolHandlerLifetimeTests.cs | 150 ++++++++++++++++-- python/packages/core/pyproject.toml | 4 +- python/packages/declarative/AGENTS.md | 8 + .../_workflows/_mcp_handler.py | 46 +++--- python/packages/declarative/pyproject.toml | 4 +- .../tests/test_default_mcp_tool_handler.py | 99 ++++++++++-- python/pyproject.toml | 4 +- python/uv.lock | 6 +- 9 files changed, 294 insertions(+), 81 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 53cb8507218..1bdef79957e 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs @@ -30,8 +30,10 @@ 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 in a bounded least-recently-used cache. Evicted sessions -/// are disposed after their active invocations finish. +/// connection name, and explicit headers in a bounded least-recently-used cache. The bound includes +/// connecting and closing sessions. At capacity, an idle session is disposed before its replacement +/// connects. If every slot is busy, new scopes fail with ; +/// existing scopes can still reuse their cached session. /// Non-cancellation cleanup failures are reported through warnings without replacing /// the invocation result or error. /// @@ -420,6 +422,28 @@ private async Task AcquireClientAsync( if (!this._clientCreations.TryGetValue(clientCacheKey, out clientCreation)) { + if (this._clients.Count + this._retiredClients.Count + this._clientCreations.Count >= this._clientCacheMaxSize) + { + LinkedListNode<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash)>? idleNode = + this._clientLru.First; + while (idleNode is not null && this._clients[idleNode.Value].ActiveInvocations != 0) + { + idleNode = idleNode.Next; + } + + if (idleNode is null) + { + throw new InvalidOperationException("MCP client capacity is exhausted; retry when an active invocation completes."); + } + + clientToDispose = this._clients[idleNode.Value]; + this._clients.Remove(idleNode.Value); + this._clientLru.Remove(idleNode); + clientToDispose.Evicted = true; + clientToDispose.DisposalClaimed = true; + this._retiredClients.Add(clientToDispose); + } + clientCreation = new(TaskCreationOptions.RunContinuationsAsynchronously); this._clientCreations[clientCacheKey] = clientCreation; clientCreationLifetime = new(TaskCreationOptions.RunContinuationsAsynchronously); @@ -456,6 +480,11 @@ private async Task AcquireClientAsync( Exception? clientCreationCleanupException = null; try { + if (clientToDispose is not null) + { + await this.DisposeCachedClientAsync(clientToDispose).ConfigureAwait(false); + } + await this._clientCreationSemaphore.WaitAsync(cancellationToken).ConfigureAwait(false); creationSemaphoreEntered = true; await this._clientLock.WaitAsync(cancellationToken).ConfigureAwait(false); @@ -516,22 +545,6 @@ await this.CompleteClientCreationLifetimeAsync( connection = null; 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; - } - } - result = newClient; clientCreation.TrySetResult(newClient); } @@ -553,11 +566,6 @@ await this.CompleteClientCreationLifetimeAsync( clientCreation.TrySetException(disposedException); throw disposedException; } - - if (clientToDispose is not null) - { - await this.DisposeCachedClientAsync(clientToDispose).ConfigureAwait(false); - } } catch (Exception 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 67390dd543b..04c6cb8eee7 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs @@ -24,6 +24,109 @@ namespace Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests; Justification = "Provider callbacks return completed tasks; each invocation and handler disposal is awaited before disposing caller-owned clients.")] public sealed class DefaultMcpToolHandlerLifetimeTests { + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task NoProvider_ActiveClientCapacity_RejectsNewScopesAndRecoversAsync(bool cancelActive) + { + // Arrange + ProtocolStub stub = new(); + using SemaphoreSlim started = new(0); + using SemaphoreSlim release = new(0); + int operations = 0; + stub.BeforeOperationAsync = async token => + { + if (Interlocked.Increment(ref operations) == 1) + { + started.Release(); + await release.WaitAsync(token); + } + }; + await using DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler, clientCacheMaxSize: 1); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + using CancellationTokenSource firstCancellation = CancellationTokenSource.CreateLinkedTokenSource(timeout.Token); + Task first = InvokeScopedAsync(handler, "workflow-a", "ping", firstCancellation.Token); + await started.WaitAsync(timeout.Token); + + try + { + // Act + await InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); + InvalidOperationException exception = await Assert.ThrowsAsync( + () => InvokeScopedAsync(handler, "workflow-b", "ping", timeout.Token)); + + // Assert + Assert.Contains("capacity", exception.Message); + Assert.Equal(1, stub.Initializations); + Assert.Equal(0, stub.Terminations); + } + finally + { + if (cancelActive) + { + firstCancellation.Cancel(); + } + + release.Release(); + try + { + await first; + } + catch (OperationCanceledException) when (cancelActive) + { + } + } + + await InvokeScopedAsync(handler, "workflow-b", "ping", timeout.Token); + Assert.Equal(2, stub.Initializations); + Assert.Equal(1, stub.Terminations); + } + + [Fact] + public async Task NoProvider_Replacement_WaitsForIdleClientCleanupBeforeConnectingAsync() + { + // Arrange + ProtocolStub stub = new(); + using SemaphoreSlim closeStarted = new(0); + using SemaphoreSlim releaseClose = new(0); + int disposals = 0; + await using DefaultMcpToolHandler handler = new( + null, stub.CreateMessageHandler, clientCacheMaxSize: 1, + clientConnectionDisposer: async connection => + { + if (Interlocked.Increment(ref disposals) == 1) + { + closeStarted.Release(); + await releaseClose.WaitAsync(); + } + + await connection.DisposeAsync(); + }); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + await InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); + Task replacement = InvokeScopedAsync(handler, "workflow-b", "ping", timeout.Token); + await closeStarted.WaitAsync(timeout.Token); + + try + { + // Act + Assert.Equal(1, stub.Initializations); + InvalidOperationException exception = await Assert.ThrowsAsync( + () => InvokeScopedAsync(handler, "workflow-c", "ping", timeout.Token)); + + // Assert + Assert.Contains("capacity", exception.Message); + } + finally + { + releaseClose.Release(); + await replacement; + } + + Assert.Equal(2, stub.Initializations); + Assert.Equal(1, stub.Terminations); + } + [Theory] [InlineData("ping")] [InlineData(DefaultMcpToolHandler.ListToolsToolName)] @@ -201,7 +304,7 @@ await handler.InvokeToolInWorkflowSessionAsync( } [Fact] - public async Task NoProvider_CacheEviction_DefersDisposalUntilActiveInvocationCompletesAsync() + public async Task NoProvider_CacheEviction_SkipsActiveClientUntilInvocationCompletesAsync() { // Arrange ProtocolStub stub = new(); @@ -216,22 +319,23 @@ public async Task NoProvider_CacheEviction_DefersDisposalUntilActiveInvocationCo await releaseFirst.WaitAsync(token); } }; - DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler, clientCacheMaxSize: 1); + DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler, clientCacheMaxSize: 2); using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); // Act Task first = InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); await firstStarted.WaitAsync(timeout.Token); await InvokeScopedAsync(handler, "workflow-b", "ping", timeout.Token); + await InvokeScopedAsync(handler, "workflow-c", "ping", timeout.Token); // Assert - Assert.Equal(0, stub.Terminations); + Assert.Equal(1, stub.Terminations); Task disposal = handler.DisposeAsync().AsTask(); Assert.False(disposal.IsCompleted); releaseFirst.Release(); await first; await disposal; - Assert.Equal(2, stub.Terminations); + Assert.Equal(3, stub.Terminations); } [Theory] @@ -239,8 +343,8 @@ public async Task NoProvider_CacheEviction_DefersDisposalUntilActiveInvocationCo [InlineData(false, true)] [InlineData(true, false)] [InlineData(true, true)] - public async Task NoProvider_EvictedHttpClientCleanupFailure_PreservesInvocationOutcomeAsync( - bool activeEviction, bool failOperation) + public async Task NoProvider_HttpClientCleanupFailure_PreservesInvocationOutcomeAsync( + bool activeDisposal, bool failOperation) { // Arrange ProtocolStub stub = new() { FailTransportDisposal = true }; @@ -249,15 +353,16 @@ public async Task NoProvider_EvictedHttpClientCleanupFailure_PreservesInvocation int operations = 0; stub.BeforeOperationAsync = async token => { - if (activeEviction && Interlocked.Increment(ref operations) == 1) + if (activeDisposal && Interlocked.Increment(ref operations) == 1) { firstStarted.Release(); await releaseFirst.WaitAsync(token); } }; - await using DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler, clientCacheMaxSize: 1); + DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler, clientCacheMaxSize: 1); using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); using CleanupTraceListener listener = new(); + Task? disposal = null; Trace.Listeners.Add(listener); try { @@ -265,11 +370,10 @@ public async Task NoProvider_EvictedHttpClientCleanupFailure_PreservesInvocation Task first = handler.InvokeToolInWorkflowSessionAsync( "https://first.example/api", null, "ping", null, null, null, "workflow-a", timeout.Token); Task outcome; - if (activeEviction) + if (activeDisposal) { await firstStarted.WaitAsync(timeout.Token); - await handler.InvokeToolInWorkflowSessionAsync( - "https://second.example/api", null, "ping", null, null, null, "workflow-b", timeout.Token); + disposal = handler.DisposeAsync().AsTask(); stub.FailOperation = failOperation; releaseFirst.Release(); outcome = first; @@ -294,12 +398,18 @@ await handler.InvokeToolInWorkflowSessionAsync( Assert.Equal("ok", Assert.IsType(Assert.Single(result.Outputs)).Text); } + if (disposal is not null) + { + await disposal; + } + Assert.Contains("Failed to dispose MCP HTTP client", listener.Output); Assert.Contains("transport cleanup failed", listener.Output); } finally { releaseFirst.Release(); + await (disposal ?? handler.DisposeAsync().AsTask()); Trace.Listeners.Remove(listener); } } @@ -354,18 +464,24 @@ public async Task NoProvider_ConcurrentWorkflowSessionCreations_AreBoundedByCach }; 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); + await initializationsStarted.WaitAsync(timeout.Token); + try + { + await Assert.ThrowsAsync( + () => InvokeScopedAsync(handler, "workflow-c", "ping", timeout.Token)); + } + finally + { + releaseInitializations.Release(2); + await Task.WhenAll(first, second); + } + Task third = InvokeScopedAsync(handler, "workflow-c", "ping", timeout.Token); - 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; diff --git a/python/packages/core/pyproject.toml b/python/packages/core/pyproject.toml index b5132f3b831..596e9ba82ff 100644 --- a/python/packages/core/pyproject.toml +++ b/python/packages/core/pyproject.toml @@ -4,7 +4,7 @@ description = "Microsoft Agent Framework for building AI Agents with Python. Thi authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}] readme = "README.md" requires-python = ">=3.10" -version = "1.22.0" +version = "1.22.1" license-files = ["LICENSE"] urls.homepage = "https://aka.ms/agent-framework" urls.source = "https://github.com/microsoft/agent-framework/tree/main/python" @@ -48,7 +48,7 @@ all = [ "agent-framework-chatkit>=1.0.0b261015,<2", "agent-framework-claude>=1.0.0b261015,<2", "agent-framework-copilotstudio>=1.0.0b261015,<2", - "agent-framework-declarative>=1.3.0,<2", + "agent-framework-declarative>=1.3.1,<2", "agent-framework-devui>=1.0.0b261015,<2", "agent-framework-durabletask>=1.0.0b260922,<2", "agent-framework-foundry>=1.15.0,<2", diff --git a/python/packages/declarative/AGENTS.md b/python/packages/declarative/AGENTS.md index 2699cfb8565..834249efb3d 100644 --- a/python/packages/declarative/AGENTS.md +++ b/python/packages/declarative/AGENTS.md @@ -2,6 +2,8 @@ YAML/JSON-based declarative agent and workflow definitions. +Requires Core 1.22.1 or later for workflow run initialization and continuation hooks. + ## Main Classes - **`AgentFactory`** - Creates agents from declarative definitions @@ -41,6 +43,12 @@ from an active invocation's context (including inherited child tasks and cleanup raises `RuntimeError` before changing handler state; shutdown must run outside that context or after invocation completion. +Without a provider, `cache_max_size` bounds live sessions including connecting and +closing clients, not only idle cache entries. At capacity, the least-recently-used +idle client is fully closed before its replacement connects. If every slot is +busy, new scopes receive an explicit capacity error; same-scope cache reuse +remains available. + Per-invocation sessions intentionally add connection overhead and lose server session continuity between invocations; shared session ownership requires an explicitly scoped custom `MCPToolHandler`. 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 e4f9569288e..49ab79df48e 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -246,9 +246,10 @@ 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. 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``. + connections. Active invocations lease their cache entry. Admission counts + cached clients, connections in progress, and clients still closing. At + capacity, an idle entry is closed before its replacement connects; if no + idle entry is available, a new scope receives a capacity error. 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 @@ -280,11 +281,12 @@ class DefaultMCPToolHandler: Args: 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 after any active invocation finishes. This also - limits simultaneous connection attempts. Defaults to ``32``. Does - not enable session caching when a provider is configured. + cache_max_size: Maximum number of live MCP clients in no-provider mode, + including connecting and closing clients. At capacity, the + least-recently-used idle entry is closed before connecting a new + client. If every slot is busy, new scopes receive a capacity error; + existing scopes can still reuse their cached client. Defaults to + ``32``. Does not limit provider-backed invocations. """ LIST_TOOLS_TOOL_NAME: ClassVar[str] = "tools/list" @@ -646,6 +648,7 @@ async def _get_or_create_entry(self, invocation: MCPToolInvocation) -> _CacheEnt # Phase 1: check the cache and either claim creation or wait for an # already in-flight creation. creating = False + evicted: _CacheEntry | None = None async with self._cache_lock: if self._closed: raise RuntimeError("DefaultMCPToolHandler is closed") @@ -656,6 +659,20 @@ async def _get_or_create_entry(self, invocation: MCPToolInvocation) -> _CacheEnt return existing inflight = self._inflight.get(key) if inflight is None: + if len(self._cache) + len(self._retired) + len(self._inflight) >= self._cache_max_size: + idle = next( + ((cached_key, entry) for cached_key, entry in self._cache.items() if entry.active_users == 0), + None, + ) + if idle is None: + raise RuntimeError( + "MCP client capacity is exhausted; retry when an active invocation completes" + ) + evicted_key, evicted = idle + del self._cache[evicted_key] + evicted.evicted = True + evicted.disposal_claimed = True + self._retired[id(evicted)] = evicted inflight = asyncio.get_running_loop().create_future() self._inflight[key] = inflight self._inflight_cleanup[key] = asyncio.get_running_loop().create_future() @@ -668,6 +685,8 @@ async def _get_or_create_entry(self, invocation: MCPToolInvocation) -> _CacheEnt # Phase 2: we own creation. Build the entry outside the lock. try: + if evicted is not None: + await self._close_claimed_entry(evicted) async with self._creation_semaphore: async with self._cache_lock: if self._closed: @@ -679,12 +698,11 @@ async def _get_or_create_entry(self, invocation: MCPToolInvocation) -> _CacheEnt raise exc.cleanup_exception from exc.creation_exception raise - # Phase 3: insert with LRU eviction; resolve the in-flight future. + # Phase 3: insert and resolve the in-flight future. # If ``aclose`` ran while we were connecting, ``_closed`` is now # True; don't insert into the cache (it has been drained), close # the just-built entry, and surface the closed-handler error to # all awaiters of the future. - evicted: _CacheEntry | None = None duplicate: _CacheEntry | None = None handler_closed = False try: @@ -705,12 +723,6 @@ async def _get_or_create_entry(self, invocation: MCPToolInvocation) -> _CacheEnt 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: @@ -736,8 +748,6 @@ async def _get_or_create_entry(self, invocation: MCPToolInvocation) -> _CacheEnt 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 diff --git a/python/packages/declarative/pyproject.toml b/python/packages/declarative/pyproject.toml index 05d9552dcf1..c7f4428dc8c 100644 --- a/python/packages/declarative/pyproject.toml +++ b/python/packages/declarative/pyproject.toml @@ -4,7 +4,7 @@ description = "Declarative specification support for Microsoft Agent Framework." authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}] readme = "README.md" requires-python = ">=3.10" -version = "1.3.0" +version = "1.3.1" license-files = ["LICENSE"] urls.homepage = "https://aka.ms/agent-framework" urls.source = "https://github.com/microsoft/agent-framework/tree/main/python" @@ -22,7 +22,7 @@ classifiers = [ "Typing :: Typed", ] dependencies = [ - "agent-framework-core>=1.20.0,<2", + "agent-framework-core>=1.22.1,<2", "httpx>=0.27,<0.29", "httpx2>=2.0.0,<3", "powerfx>=0.0.32,<0.0.35; python_version < '3.14'", 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 33d25c12279..d45f70b1ad3 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -960,6 +960,74 @@ async def test_header_provider_returns_captured_headers(self) -> None: class TestCache: + @pytest.mark.parametrize("cancel_active", [False, True]) + async def test_active_clients_reject_new_scopes_at_capacity(self, cancel_active: bool) -> None: + handler = DefaultMCPToolHandler(cache_max_size=1) + started = asyncio.Event() + release = asyncio.Event() + calls = 0 + + async def gated_call(_tool: FakeTool, tool_name: str, **arguments: Any) -> Any: + nonlocal calls + calls += 1 + if calls == 1: + started.set() + await release.wait() + return [Content.from_text("ok")] + + with _patch_tool(), patch.object(FakeTool, "call_tool", gated_call): + first = asyncio.create_task(handler.invoke_tool(_invocation(workflow_session_id="workflow-a"))) + await started.wait() + try: + reused = await handler.invoke_tool(_invocation(workflow_session_id="workflow-a")) + assert not reused.is_error + rejected = await handler.invoke_tool(_invocation(workflow_session_id="workflow-b")) + assert rejected.is_error + assert rejected.error_message is not None + assert "capacity" in rejected.error_message + assert len(FakeTool.instances) == 1 + assert FakeTool.instances[0].close_count == 0 + finally: + if cancel_active: + first.cancel() + release.set() + await asyncio.gather(first, return_exceptions=True) + + recovered = await handler.invoke_tool(_invocation(workflow_session_id="workflow-b")) + assert not recovered.is_error + await handler.aclose() + assert len(FakeTool.instances) == 2 + assert all(tool.close_count == 1 for tool in FakeTool.instances) + + async def test_replacement_waits_for_idle_client_cleanup_before_connecting(self) -> None: + handler = DefaultMCPToolHandler(cache_max_size=1) + close_started = asyncio.Event() + release_close = asyncio.Event() + original_close = FakeTool.close + + async def gated_close(tool: FakeTool) -> None: + if tool is FakeTool.instances[0]: + close_started.set() + await release_close.wait() + await original_close(tool) + + with _patch_tool(), patch.object(FakeTool, "close", gated_close): + await handler.invoke_tool(_invocation(workflow_session_id="workflow-a")) + replacement = asyncio.create_task(handler.invoke_tool(_invocation(workflow_session_id="workflow-b"))) + await close_started.wait() + try: + assert len(FakeTool.instances) == 1 + rejected = await handler.invoke_tool(_invocation(workflow_session_id="workflow-c")) + assert rejected.is_error + assert rejected.error_message is not None + assert "capacity" in rejected.error_message + finally: + release_close.set() + await replacement + await handler.aclose() + assert len(FakeTool.instances) == 2 + assert all(tool.close_count == 1 for tool in FakeTool.instances) + @pytest.mark.parametrize("explicit_session", [False, True]) async def test_sessionless_scope_reuses_its_own_cached_client(self, explicit_session: bool) -> None: state = State() @@ -1038,8 +1106,8 @@ async def test_lru_eviction_closes_old_entry(self) -> None: 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) + async def test_lru_eviction_skips_active_client_until_invocation_finishes(self) -> None: + handler = DefaultMCPToolHandler(cache_max_size=2) first_started = asyncio.Event() release_first = asyncio.Event() @@ -1053,9 +1121,12 @@ async def gated_call(tool: FakeTool, tool_name: str, **arguments: Any) -> Any: 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/")) + third = await handler.invoke_tool(_invocation(server_url="https://c/")) assert not second.is_error + assert not third.is_error assert FakeTool.instances[0].close_count == 0 + assert FakeTool.instances[1].close_count == 1 close_task = asyncio.create_task(handler.aclose()) await asyncio.sleep(0) assert not close_task.done() @@ -1067,6 +1138,7 @@ async def gated_call(tool: FakeTool, tool_name: str, **arguments: Any) -> Any: assert not first_result.is_error assert FakeTool.instances[0].close_count == 1 assert FakeTool.instances[1].close_count == 1 + assert FakeTool.instances[2].close_count == 1 @pytest.mark.asyncio async def test_lru_eviction_cleanup_cancellation_releases_new_entry(self) -> None: @@ -1163,7 +1235,7 @@ async def gated_call(_tool: FakeTool, _tool_name: str, **_arguments: Any) -> Any 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: + async def test_entry_creation_rejects_new_scopes_at_capacity(self) -> None: handler = DefaultMCPToolHandler(cache_max_size=2) connecting = 0 max_connecting = 0 @@ -1187,17 +1259,16 @@ async def gated_connect(tool: FakeTool) -> None: 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): - cancelled_result = await cancelled - assert cancelled_result is None - - release.set() - results = await asyncio.gather(first, second) + try: + rejected = await handler.invoke_tool(_invocation(workflow_session_id="workflow-c")) + assert rejected.is_error + assert rejected.error_message is not None + assert "capacity" in rejected.error_message + assert len(FakeTool.instances) == 2 + finally: + release.set() + results = await asyncio.gather(first, second) + await handler.aclose() assert all(not result.is_error for result in results) assert max_connecting == 2 diff --git a/python/pyproject.toml b/python/pyproject.toml index a318c267a89..b03331dad36 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -4,7 +4,7 @@ description = "Microsoft Agent Framework for building AI Agents with Python. Thi authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}] readme = "README.md" requires-python = ">=3.10" -version = "1.22.0" +version = "1.22.1" license-files = ["LICENSE"] urls.homepage = "https://aka.ms/agent-framework" urls.source = "https://github.com/microsoft/agent-framework/tree/main/python" @@ -24,7 +24,7 @@ classifiers = [ "Typing :: Typed", ] dependencies = [ - "agent-framework-core[all]==1.22.0", + "agent-framework-core[all]==1.22.1", ] [dependency-groups] diff --git a/python/uv.lock b/python/uv.lock index 6571372d261..012623e9eca 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -115,7 +115,7 @@ wheels = [ [[package]] name = "agent-framework" -version = "1.22.0" +version = "1.22.1" source = { virtual = "." } dependencies = [ { name = "agent-framework-core", extra = ["all"] }, @@ -441,7 +441,7 @@ requires-dist = [ [[package]] name = "agent-framework-core" -version = "1.22.0" +version = "1.22.1" source = { editable = "packages/core" } dependencies = [ { name = "msgspec" }, @@ -549,7 +549,7 @@ dev = [ [[package]] name = "agent-framework-declarative" -version = "1.3.0" +version = "1.3.1" source = { editable = "packages/declarative" } dependencies = [ { name = "agent-framework-core" }, From 31acf5bcf96f9143d5703769f455ea313ae23f3f Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Thu, 15 Oct 2026 14:09:09 -0400 Subject: [PATCH 45/58] fix: select built workspace wheels in meta install checks Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 48ae509d-294c-4bf5-864d-81baecbae204 --- .github/workflows/python-tests.yml | 17 ++- .../tests/test_default_mcp_tool_handler.py | 7 +- python/scripts/dependencies/README.md | 7 ++ .../tests/test_meta_package_install.py | 106 ++++++++++++++++++ 4 files changed, 134 insertions(+), 3 deletions(-) create mode 100644 python/scripts/dependencies/tests/test_meta_package_install.py diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml index cc08ec06fe1..2cbfcfe6daf 100644 --- a/.github/workflows/python-tests.yml +++ b/.github/workflows/python-tests.yml @@ -104,6 +104,17 @@ jobs: --no-build-logs \ --no-create-gitignore \ --out-dir "${{ github.workspace }}/wheelhouse" + - name: Constrain workspace packages to built wheels + shell: pwsh + run: | + $wheels = @(Get-ChildItem wheelhouse -Filter *.whl) + if ($wheels.Count -eq 0) { + throw "No built workspace wheels were found." + } + $wheels | Sort-Object Name | ForEach-Object { + $name = ($_.Name -split "-")[0] + "$name @ $([System.Uri]::new($_.FullName).AbsoluteUri)" + } | Set-Content wheelhouse-constraints.txt -Encoding utf8 - name: Resolve the end-user install timeout-minutes: 6 shell: pwsh @@ -118,6 +129,7 @@ jobs: --dry-run ` --ignore-installed ` --find-links wheelhouse ` + --constraint wheelhouse-constraints.txt ` --report "pip-resolution-${{ matrix.python-version }}.json" ` "agent-framework==$version" ` azure-identity @@ -138,5 +150,8 @@ jobs: run: | $python = ".meta-venv\Scripts\python.exe" $version = & $python -c "import tomllib; print(tomllib.load(open('python/pyproject.toml', 'rb'))['project']['version'])" - & $python -m pip install --find-links wheelhouse "agent-framework==$version" azure-identity + & $python -m pip install --find-links wheelhouse --constraint wheelhouse-constraints.txt "agent-framework==$version" azure-identity + if ($LASTEXITCODE -ne 0) { + exit $LASTEXITCODE + } & $python -c "import agent_framework; import azure.identity" 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 d45f70b1ad3..02a58acdae3 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -1023,8 +1023,11 @@ async def gated_close(tool: FakeTool) -> None: assert "capacity" in rejected.error_message finally: release_close.set() - await replacement - await handler.aclose() + try: + replacement_result = await replacement + finally: + await handler.aclose() + assert not replacement_result.is_error assert len(FakeTool.instances) == 2 assert all(tool.close_count == 1 for tool in FakeTool.instances) diff --git a/python/scripts/dependencies/README.md b/python/scripts/dependencies/README.md index d512c9b97e7..cb3a8684d41 100644 --- a/python/scripts/dependencies/README.md +++ b/python/scripts/dependencies/README.md @@ -8,6 +8,13 @@ This folder contains the Python workspace tooling for dependency maintenance: Run the commands below from the `python/` directory. +The meta-package install CI constrains workspace distributions to the wheels it +has just built. `--find-links` alone can prefer a registry wheel with the same +version but different dependency metadata. Both the resolution probe and actual +installation use the generated constraints; third-party dependencies still +resolve from the package index. `tests/test_meta_package_install.py` exercises +both paths against an isolated local index. + ## Files in this folder - `validate_dependency_bounds.py` diff --git a/python/scripts/dependencies/tests/test_meta_package_install.py b/python/scripts/dependencies/tests/test_meta_package_install.py new file mode 100644 index 00000000000..13b1a6ab5be --- /dev/null +++ b/python/scripts/dependencies/tests/test_meta_package_install.py @@ -0,0 +1,106 @@ +# Copyright (c) Microsoft. All rights reserved. + +import json +import re +import subprocess +import sys +from functools import partial +from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from threading import Thread +from zipfile import ZipFile + +import pytest +import yaml + + +def _write_wheel(directory: Path, name: str, version: str, dependencies: list[str]) -> Path: + directory.mkdir(parents=True, exist_ok=True) + normalized_name = name.replace("-", "_") + wheel = directory / f"{normalized_name}-{version}-py3-none-any.whl" + metadata = f"Metadata-Version: 2.1\nName: {name}\nVersion: {version}\n" + metadata += "".join(f"Requires-Dist: {dependency}\n" for dependency in dependencies) + dist_info = f"{normalized_name}-{version}.dist-info" + with ZipFile(wheel, "w") as archive: + archive.writestr(f"{dist_info}/METADATA", metadata) + archive.writestr( + f"{dist_info}/WHEEL", + "Wheel-Version: 1.0\nGenerator: test\nRoot-Is-Purelib: true\nTag: py3-none-any\n", + ) + archive.writestr(f"{dist_info}/RECORD", "") + return wheel + + +@pytest.mark.parametrize("step_name", ["Resolve the end-user install", "Install and import the meta package"]) +def test_meta_install_uses_built_wheels_instead_of_same_version_registry_metadata( + tmp_path: Path, step_name: str +) -> None: + workspace = tmp_path / "workspace with spaces" + wheelhouse = workspace / "wheelhouse" + _write_wheel(wheelhouse, "agent-framework", "1.0.1", ["agent-framework-core==1.0.1"]) + _write_wheel(wheelhouse, "agent-framework-core", "1.0.1", ["agent-framework-foundry-hosting==1.0.0b1", "mcp>=2,<3"]) + hosting = _write_wheel(wheelhouse, "agent-framework-foundry-hosting", "1.0.0b1", ["mcp>=2,<3"]) + registry = tmp_path / "registry" + registry_hosting = registry / "simple" / "agent-framework-foundry-hosting" + published = _write_wheel(registry_hosting, "agent-framework-foundry-hosting", "1.0.0b1", ["mcp<2"]) + (registry_hosting / "index.html").write_text(f'{published.name}', encoding="utf-8") + registry_mcp = registry / "simple" / "mcp" + mcp = _write_wheel(registry_mcp, "mcp", "2.2.0", []) + (registry_mcp / "index.html").write_text(f'{mcp.name}', encoding="utf-8") + + workflow_path = Path(__file__).resolve().parents[4] / ".github" / "workflows" / "python-tests.yml" + steps = yaml.safe_load(workflow_path.read_text(encoding="utf-8"))["jobs"]["python-meta-package-install"]["steps"] + for step in steps: + if step.get("name") == "Constrain workspace packages to built wheels": + subprocess.run( + ["pwsh", "-NoProfile", "-NonInteractive", "-Command", step["run"]], + cwd=workspace, + check=True, + capture_output=True, + text=True, + ) + + install_script = next(step["run"] for step in steps if step.get("name") == step_name) + selection_arguments = [ + value + for flag in ("--find-links", "--constraint") + for argument in re.findall(rf"{flag}\s+(\S+)", install_script) + for value in (flag, argument) + ] + report = workspace / "report.json" + server = ThreadingHTTPServer(("127.0.0.1", 0), partial(SimpleHTTPRequestHandler, directory=str(registry))) + server_thread = Thread(target=server.serve_forever, daemon=True) + server_thread.start() + try: + command = [ + sys.executable, + "-m", + "pip", + "--isolated", + "install", + "--disable-pip-version-check", + "--no-cache-dir", + "--ignore-installed", + "--index-url", + f"http://127.0.0.1:{server.server_port}/simple", + "--report", + str(report), + *selection_arguments, + "agent-framework==1.0.1", + ] + if step_name == "Resolve the end-user install": + command.append("--dry-run") + else: + command.extend(["--target", str(workspace / "installed")]) + result = subprocess.run(command, cwd=workspace, capture_output=True, text=True, timeout=30) + finally: + server.shutdown() + server.server_close() + server_thread.join() + + assert result.returncode == 0, result.stdout + result.stderr + installed = json.loads(report.read_text(encoding="utf-8"))["install"] + selected_hosting = next(item for item in installed if item["metadata"]["name"] == "agent-framework-foundry-hosting") + assert selected_hosting["download_info"]["url"] == hosting.as_uri() + selected_mcp = next(item for item in installed if item["metadata"]["name"] == "mcp") + assert selected_mcp["download_info"]["url"].startswith("http://127.0.0.1:") From ec94868e8c1611091e914163d1c4c4fe8777300c Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Thu, 15 Oct 2026 14:18:20 -0400 Subject: [PATCH 46/58] fix: preserve cancellation during MCP lifecycle cleanup Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 48ae509d-294c-4bf5-864d-81baecbae204 --- .../DefaultMcpToolHandler.cs | 10 +++- .../DefaultMcpToolHandlerLifetimeTests.cs | 47 +++++++++++++++++-- .../_workflows/_mcp_handler.py | 8 +--- .../tests/test_default_mcp_tool_handler.py | 45 ++++++++++++++++++ 4 files changed, 99 insertions(+), 11 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 1bdef79957e..cefa73fd174 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs @@ -482,7 +482,15 @@ private async Task AcquireClientAsync( { if (clientToDispose is not null) { - await this.DisposeCachedClientAsync(clientToDispose).ConfigureAwait(false); + try + { + await this.DisposeCachedClientAsync(clientToDispose).ConfigureAwait(false); + } + catch (Exception cleanupException) + { + clientCreationCleanupException = cleanupException; + throw; + } } await this._clientCreationSemaphore.WaitAsync(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 04c6cb8eee7..3a014979268 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 @@ -68,12 +68,13 @@ public async Task NoProvider_ActiveClientCapacity_RejectsNewScopesAndRecoversAsy } release.Release(); - try + if (cancelActive) { - await first; + await Assert.ThrowsAnyAsync(() => first); } - catch (OperationCanceledException) when (cancelActive) + else { + await first; } } @@ -127,6 +128,46 @@ public async Task NoProvider_Replacement_WaitsForIdleClientCleanupBeforeConnecti Assert.Equal(1, stub.Terminations); } + [Fact] + public async Task NoProvider_DisposalDuringIdleEviction_ReportsCleanupCancellationAsync() + { + // Arrange + OperationCanceledException cleanupException = new("session cleanup cancelled"); + ProtocolStub stub = new(); + using SemaphoreSlim closeStarted = new(0); + using SemaphoreSlim releaseClose = new(0); + DefaultMcpToolHandler handler = new( + null, stub.CreateMessageHandler, clientCacheMaxSize: 1, + clientConnectionDisposer: async connection => + { + closeStarted.Release(); + await releaseClose.WaitAsync(); + await connection.DisposeAsync(); + throw cleanupException; + }); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + await InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); + Task replacement = InvokeScopedAsync(handler, "workflow-b", "ping", timeout.Token); + await closeStarted.WaitAsync(timeout.Token); + + // Act + Task disposal = handler.DisposeAsync().AsTask(); + try + { + Assert.False(disposal.IsCompleted); + } + finally + { + releaseClose.Release(); + } + + // Assert + Assert.Same(cleanupException, await Assert.ThrowsAsync(() => replacement)); + Assert.Same(cleanupException, await Assert.ThrowsAsync(() => disposal.WaitAsync(timeout.Token))); + Assert.Equal(1, stub.Initializations); + Assert.Equal(1, stub.Terminations); + } + [Theory] [InlineData("ping")] [InlineData(DefaultMcpToolHandler.ListToolsToolName)] 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 49ab79df48e..52647c44c53 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -808,13 +808,7 @@ async def cleanup() -> None: inflight.set_exception(self._entry_creation_exception(creation_exception)) 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() + await self._await_cleanup_task(asyncio.create_task(cleanup())) @staticmethod def _entry_creation_exception(exc: BaseException) -> BaseException: 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 02a58acdae3..e9d5818c6be 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -1164,6 +1164,51 @@ async def close(tool: FakeTool) -> None: cancel_first_close = False await asyncio.wait_for(handler.aclose(), timeout=1) + @pytest.mark.parametrize("cancellation_count", [1, 2]) + async def test_cancellation_during_failed_creation_bookkeeping_propagates(self, cancellation_count: int) -> None: + handler = DefaultMCPToolHandler() + connect_started = asyncio.Event() + release_connect = asyncio.Event() + bookkeeping_waiting = asyncio.Event() + + class ObservedLock(asyncio.Lock): + async def acquire(self) -> Literal[True]: + if self.locked(): + bookkeeping_waiting.set() + return await super().acquire() + + handler._cache_lock = ObservedLock() + + async def failed_connect(_tool: FakeTool) -> None: + connect_started.set() + await release_connect.wait() + raise RuntimeError("connection failed") + + with _patch_tool(), patch.object(FakeTool, "connect", failed_connect): + invocation = asyncio.create_task(handler.invoke_tool(_invocation())) + try: + await asyncio.wait_for(connect_started.wait(), timeout=1) + await handler._cache_lock.acquire() + try: + release_connect.set() + await asyncio.wait_for(bookkeeping_waiting.wait(), timeout=1) + for _ in range(cancellation_count): + invocation.cancel() + await asyncio.sleep(0) + assert not invocation.done() + finally: + handler._cache_lock.release() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(invocation, timeout=1) + assert not handler._inflight + assert not handler._inflight_cleanup + finally: + release_connect.set() + await asyncio.wait_for(handler.aclose(), timeout=1) + + assert FakeTool.instances[0].close_count == 1 + async def test_repeated_cancellation_during_eviction_finalization_completes_shutdown(self) -> None: handler = DefaultMCPToolHandler(cache_max_size=1) close_started = asyncio.Event() From d6c95265583b0b1e581577278b99c00c3268b214 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Thu, 15 Oct 2026 14:29:14 -0400 Subject: [PATCH 47/58] fix: retain workflow identity on response continuations Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 48ae509d-294c-4bf5-864d-81baecbae204 --- .../DefaultMcpToolHandlerLifetimeTests.cs | 2 +- .../core/agent_framework/_workflows/_agent.py | 2 - python/packages/declarative/AGENTS.md | 2 + .../tests/test_invoke_mcp_tool_executor.py | 74 +++++++++++++++++++ 4 files changed, 77 insertions(+), 3 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 3a014979268..3ff71500d6c 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs @@ -163,7 +163,7 @@ public async Task NoProvider_DisposalDuringIdleEviction_ReportsCleanupCancellati // Assert Assert.Same(cleanupException, await Assert.ThrowsAsync(() => replacement)); - Assert.Same(cleanupException, await Assert.ThrowsAsync(() => disposal.WaitAsync(timeout.Token))); + Assert.Same(cleanupException, await Assert.ThrowsAsync(() => disposal)); Assert.Equal(1, stub.Initializations); Assert.Equal(1, stub.Terminations); } diff --git a/python/packages/core/agent_framework/_workflows/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index 0aee8aae69d..5181a867abb 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/python/packages/core/agent_framework/_workflows/_agent.py @@ -516,7 +516,6 @@ async def _run_core( function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, ) - self.workflow._set_agent_run_context(agent_session_id) # pyright: ignore[reportPrivateUsage] async for event in response_stream: yield event else: @@ -527,7 +526,6 @@ async def _run_core( function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, ) - 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: diff --git a/python/packages/declarative/AGENTS.md b/python/packages/declarative/AGENTS.md index 834249efb3d..f77552c858e 100644 --- a/python/packages/declarative/AGENTS.md +++ b/python/packages/declarative/AGENTS.md @@ -34,6 +34,8 @@ framework-marked WorkflowAgent turns retain the agent continuation scope. Run scopes are selected before the initial entry checkpoint and preserved when Entry replays. Sessionless WorkflowAgent turns select their own checkpointed scope, independent of explicit agent sessions, direct runs, and restored approval scopes. +Response-only continuations retain the originating agent identity even when supplied +a new `AgentSession` after checkpoint restoration. Approval continuations keep their originating active scope until the next run with new input. 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 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 d3a64f0669e..b0b9a97dd0c 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,80 @@ 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("pause_kind", ["Question", "RequestExternalInput"]) + @pytest.mark.parametrize("streaming", [False, True]) + async def test_response_only_agent_checkpoint_resume_preserves_workflow_session_id( + self, pause_kind: str, streaming: bool + ) -> None: + from agent_framework import InMemoryCheckpointStorage + + from agent_framework_declarative._workflows import ExternalInputResponse + + storage = InMemoryCheckpointStorage() + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + definition = { + "name": "mcp_continuation", + "actions": [ + {**_action(tool_name="before"), "id": "before"}, + { + "kind": pause_kind, + "id": "pause", + "question" if pause_kind == "Question" else "prompt": "Continue?", + "variable": "Local.answer", + }, + {**_action(tool_name="after"), "id": "after"}, + ], + } + workflow = factory.create_workflow_from_definition(definition) + agent = workflow.as_agent() + original_session = AgentSession() + if streaming: + response = await agent.run( + "first turn", session=original_session, stream=True, checkpoint_storage=storage + ).get_final_response() + else: + response = await agent.run("first turn", session=original_session, checkpoint_storage=storage) + + [request] = (await workflow._runner_context.get_pending_request_info_events()).values() + [function_call] = [ + content for message in response.messages for content in message.contents if content.type == "function_call" + ] + assert function_call.call_id == request.request_id + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + workflow = factory.create_workflow_from_definition(definition) + agent = workflow.as_agent() + resumed_session = AgentSession() + assert resumed_session.session_id != original_session.session_id + reply = Message( + role="user", + contents=[ + Content( + "function_result", + call_id=request.request_id, + result=ExternalInputResponse(user_input="continue"), + ) + ], + ) + if streaming: + await agent.run( + reply, + session=resumed_session, + stream=True, + checkpoint_id=checkpoint_id, + checkpoint_storage=storage, + ).get_final_response() + else: + await agent.run(reply, session=resumed_session, checkpoint_id=checkpoint_id, checkpoint_storage=storage) + + [original, resumed] = handler.invocations + assert original.tool_name == "before" + assert resumed.tool_name == "after" + assert original.workflow_session_id + assert resumed.workflow_session_id == original.workflow_session_id + assert not await workflow._runner_context.get_pending_request_info_events() + @pytest.mark.parametrize("run_kind", ["direct", "sessionless_agent", "explicit_agent"]) @pytest.mark.parametrize("streaming", [False, True]) async def test_entry_checkpoint_replay_preserves_workflow_session_id(self, run_kind: str, streaming: bool) -> None: From 30ff41ab00bfce2a67c1c4844ddb78d039d83131 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Thu, 15 Oct 2026 14:53:22 -0400 Subject: [PATCH 48/58] fix: retain originating workflow scope for external input Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 48ae509d-294c-4bf5-864d-81baecbae204 --- python/packages/declarative/AGENTS.md | 9 +- .../_workflows/_declarative_builder.py | 18 ++- .../_workflows/_executors_external_input.py | 16 ++ .../_workflows/_mcp_handler.py | 2 +- .../tests/test_invoke_mcp_tool_executor.py | 149 ++++++++++++++++++ 5 files changed, 186 insertions(+), 8 deletions(-) diff --git a/python/packages/declarative/AGENTS.md b/python/packages/declarative/AGENTS.md index f77552c858e..11c2373ea50 100644 --- a/python/packages/declarative/AGENTS.md +++ b/python/packages/declarative/AGENTS.md @@ -34,7 +34,7 @@ framework-marked WorkflowAgent turns retain the agent continuation scope. Run scopes are selected before the initial entry checkpoint and preserved when Entry replays. Sessionless WorkflowAgent turns select their own checkpointed scope, independent of explicit agent sessions, direct runs, and restored approval scopes. -Response-only continuations retain the originating agent identity even when supplied +Response-only continuations retain the originating workflow scope even when supplied a new `AgentSession` after checkpoint restoration. Approval continuations keep their originating active scope until the next run with new input. With a provider, every invocation (including `tools/list`) gets a fresh tool/session, @@ -81,6 +81,13 @@ arguments. Explicit caller tool arguments in options are preserved. - **`ExternalInputRequest`** / **`ExternalInputResponse`** - Human-in-the-loop support - **`AgentExternalInputRequest`** / **`AgentExternalInputResponse`** - Agent-level input requests +`Question` and `RequestExternalInput` requests retain their originating workflow +scope across fresh runs and checkpoint restores. Responses restore that scope +before updating state or continuing execution. Factory-built workflows reject +response batches spanning different scopes, including mixed MCP approvals and +external input. Legacy external-input requests without a valid originating scope +are rejected before consumption rather than inheriting the current scope. + ## Usage ```python 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 544916ea70f..6702288d72e 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_builder.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_builder.py @@ -43,7 +43,11 @@ IfConditionEvaluatorExecutor, JoinExecutor, ) -from ._executors_external_input import EXTERNAL_INPUT_EXECUTORS +from ._executors_external_input import ( + EXTERNAL_INPUT_EXECUTORS, + ExternalInputRequest, + get_external_input_workflow_session_id, +) from ._executors_http import HTTP_ACTION_EXECUTORS, HttpRequestActionExecutor from ._executors_mcp import ( MCP_ACTION_EXECUTORS, @@ -119,15 +123,17 @@ def _validate_mcp_response_batch(requests: Sequence[WorkflowEvent[Any]]) -> None: scopes: set[str | None] = set() for request in requests: - if isinstance(request.data, MCPToolApprovalRequest): + if isinstance(request.data, ExternalInputRequest): + scope = get_external_input_workflow_session_id(request.data) + elif 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) + else: + continue + scopes.add(scope) if len(scopes) > 1: - raise DeclarativeWorkflowError( - "MCP approval responses from different workflow sessions must be submitted in separate runs." - ) + raise DeclarativeWorkflowError("Responses from different workflow sessions must be submitted in separate runs.") def _initialize_mcp_run_state(state: State) -> None: diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_external_input.py b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_external_input.py index 44bd6aad8da..30208fd3b70 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_external_input.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_external_input.py @@ -21,6 +21,8 @@ ActionComplete, DeclarativeActionExecutor, ) +from ._errors import DeclarativeWorkflowError +from ._mcp_handler import get_or_create_workflow_session_id, restore_workflow_session_id def _get_prompt_text(action_def: dict[str, Any], primary_key: str, fallback_key: str) -> Any: @@ -67,12 +69,22 @@ class ExternalInputRequest: ``RequestExternalInput`` action (e.g. ``"approval"``) and it is propagated unchanged. metadata: Additional context (choices, output_property, timeout, etc.). + workflow_session_id: Framework-owned scope of the workflow that issued the request. """ request_id: str message: str request_type: str = "external" metadata: dict[str, Any] = field(default_factory=dict) # type: ignore + workflow_session_id: str | None = None + + +def get_external_input_workflow_session_id(request: ExternalInputRequest) -> str: + """Return the originating scope, rejecting requests with an unknown origin.""" + scope = request.workflow_session_id + if not isinstance(scope, str) or not scope: + raise DeclarativeWorkflowError("External input request is missing a valid originating workflow session.") + return scope @dataclass @@ -144,6 +156,7 @@ async def handle_action( "allow_free_text": allow_free_text, "default_value": default_value, }, + workflow_session_id=get_or_create_workflow_session_id(ctx.state), ), ExternalInputResponse, ) @@ -156,6 +169,7 @@ async def handle_response( ctx: WorkflowContext[ActionComplete], ) -> None: """Handle the user's response to the question.""" + restore_workflow_session_id(ctx.state, get_external_input_workflow_session_id(original_request)) state = self._get_state(ctx.state) output_property = original_request.metadata.get("output_property", "Local.answer") @@ -212,6 +226,7 @@ async def handle_action( message=str(evaluated_message), request_type=request_type, metadata=request_metadata, + workflow_session_id=get_or_create_workflow_session_id(ctx.state), ), ExternalInputResponse, ) @@ -224,6 +239,7 @@ async def handle_response( ctx: WorkflowContext[ActionComplete], ) -> None: """Handle the external input response.""" + restore_workflow_session_id(ctx.state, get_external_input_workflow_session_id(original_request)) state = self._get_state(ctx.state) output_property = original_request.metadata.get("output_property", "Local.externalInput") 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 52647c44c53..90d63d0b117 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -108,7 +108,7 @@ def activate_workflow_session_id(state: State, *, reset_unscoped: bool, is_agent 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. + # A response 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) 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 b0b9a97dd0c..3c53693f1b9 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -294,6 +294,155 @@ async def test_response_only_agent_checkpoint_resume_preserves_workflow_session_ assert resumed.workflow_session_id == original.workflow_session_id assert not await workflow._runner_context.get_pending_request_info_events() + @pytest.mark.parametrize("pause_kind", ["Question", "RequestExternalInput"]) + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("restore", [False, True]) + async def test_pending_external_input_restores_originating_workflow_session_id( + self, pause_kind: str, streaming: bool, restore: bool + ) -> None: + from agent_framework import InMemoryCheckpointStorage, WorkflowEvent + + from agent_framework_declarative._workflows import ExternalInputResponse + + storage = InMemoryCheckpointStorage() + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + definition = { + "name": "mcp_pending_input", + "actions": [ + {**_action(tool_name="before"), "id": "before"}, + { + "kind": pause_kind, + "id": "pause", + "question" if pause_kind == "Question" else "prompt": "Continue?", + "variable": "Local.answer", + }, + {**_action(tool_name="after"), "id": "after"}, + ], + } + workflow = factory.create_workflow_from_definition(definition) + requests: list[WorkflowEvent[Any]] = [] + for _ in range(2): + if streaming: + await workflow.run({}, stream=True, checkpoint_storage=storage).get_final_response() + else: + await workflow.run({}, checkpoint_storage=storage) + [request] = [ + event + for event in (await workflow._runner_context.get_pending_request_info_events()).values() + if event.request_id not in {previous.request_id for previous in requests} + ] + requests.append(request) + original_scopes = [invocation.workflow_session_id for invocation in handler.invocations] + assert all(original_scopes) + assert original_scopes[0] != original_scopes[1] + + 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 request, scope in zip(requests, original_scopes, strict=True): + responses = {request.request_id: ExternalInputResponse(user_input="continue")} + if streaming: + await workflow.run(responses=responses, stream=True).get_final_response() + else: + await workflow.run(responses=responses) + assert handler.invocations[-1].tool_name == "after" + assert handler.invocations[-1].workflow_session_id == scope + assert handler.call_count == 4 + assert not await workflow._runner_context.get_pending_request_info_events() + + @pytest.mark.parametrize("mixed_approval", [False, True]) + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("restore", [False, True]) + async def test_external_input_batches_reject_different_workflow_sessions( + self, mixed_approval: bool, streaming: bool, restore: bool + ) -> None: + from agent_framework import InMemoryCheckpointStorage + + from agent_framework_declarative._workflows import ExternalInputResponse, ToolApprovalResponse + + storage = InMemoryCheckpointStorage() + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + definition = { + "name": "mcp_input_batches", + "actions": [ + {**_action(tool_name="before"), "id": "before"}, + {"kind": "Question", "id": "pause", "question": "Continue?", "variable": "Local.answer"}, + {**_action(tool_name="after", require_approval=mixed_approval), "id": "after"}, + ], + } + workflow = factory.create_workflow_from_definition(definition) + await workflow.run({}, checkpoint_storage=storage) + [first] = (await workflow._runner_context.get_pending_request_info_events()).values() + await workflow.run({}, checkpoint_storage=storage) + [second] = [ + event + for event in (await workflow._runner_context.get_pending_request_info_events()).values() + if event.request_id != first.request_id + ] + if mixed_approval: + await workflow.run( + responses={second.request_id: ExternalInputResponse(user_input="continue")}, + checkpoint_storage=storage, + ) + [second] = [ + event + for event in (await workflow._runner_context.get_pending_request_info_events()).values() + if event.request_id != first.request_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) + + responses = { + first.request_id: ExternalInputResponse(user_input="continue"), + second.request_id: ( + ToolApprovalResponse(approved=True) if mixed_approval else ExternalInputResponse(user_input="continue") + ), + } + with pytest.raises(DeclarativeWorkflowError, match="different workflow sessions"): + if streaming: + await workflow.run(responses=responses, stream=True).get_final_response() + else: + await workflow.run(responses=responses) + assert handler.call_count == 2 + assert set(await workflow._runner_context.get_pending_request_info_events()) == set(responses) + + @pytest.mark.parametrize("pause_kind", ["Question", "RequestExternalInput"]) + async def test_external_input_without_originating_scope_is_not_consumed(self, pause_kind: str) -> None: + from agent_framework_declarative._workflows import ExternalInputResponse + + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + workflow = factory.create_workflow_from_definition({ + "name": "mcp_legacy_input", + "actions": [ + {**_action(tool_name="before"), "id": "before"}, + { + "kind": pause_kind, + "id": "pause", + "question" if pause_kind == "Question" else "prompt": "Continue?", + "variable": "Local.answer", + }, + {**_action(tool_name="after"), "id": "after"}, + ], + }) + await workflow.run({}) + [request] = (await workflow._runner_context.get_pending_request_info_events()).values() + request.data.workflow_session_id = None + + with pytest.raises(DeclarativeWorkflowError, match="originating workflow session"): + await workflow.run(responses={request.request_id: ExternalInputResponse(user_input="continue")}) + + assert handler.call_count == 1 + assert set(await workflow._runner_context.get_pending_request_info_events()) == {request.request_id} + @pytest.mark.parametrize("run_kind", ["direct", "sessionless_agent", "explicit_agent"]) @pytest.mark.parametrize("streaming", [False, True]) async def test_entry_checkpoint_replay_preserves_workflow_session_id(self, run_kind: str, streaming: bool) -> None: From 3b2f7b3ec96116f24ec881dd2cb6978b5716b537 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Thu, 15 Oct 2026 15:28:19 -0400 Subject: [PATCH 49/58] fix(python): preserve pending request scopes and recover legacy continuations Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 48ae509d-294c-4bf5-864d-81baecbae204 --- ...thon-declarative-request-scope-recovery.md | 67 ++++ python/packages/declarative/AGENTS.md | 23 +- python/packages/declarative/README.md | 35 ++ .../_workflows/_declarative_builder.py | 45 ++- .../_workflows/_executors_agents.py | 39 +- .../_workflows/_executors_external_input.py | 19 +- .../_workflows/_executors_mcp.py | 8 +- .../_workflows/_executors_tools.py | 9 + .../_workflows/_factory.py | 9 + .../_workflows/_mcp_handler.py | 15 + .../test_declarative_approval_binding.py | 43 ++- .../tests/test_function_tool_executor.py | 8 +- .../declarative/tests/test_graph_coverage.py | 2 + .../tests/test_invoke_mcp_tool_executor.py | 336 +++++++++++++++++- 14 files changed, 586 insertions(+), 72 deletions(-) create mode 100644 docs/decisions/0046-python-declarative-request-scope-recovery.md diff --git a/docs/decisions/0046-python-declarative-request-scope-recovery.md b/docs/decisions/0046-python-declarative-request-scope-recovery.md new file mode 100644 index 00000000000..7e59f804163 --- /dev/null +++ b/docs/decisions/0046-python-declarative-request-scope-recovery.md @@ -0,0 +1,67 @@ +--- +status: proposed +contact: baywet +date: 2026-10-15 +deciders: eavanvalkenburg, westey-m +--- + +# Preserve declarative request origins with opt-in legacy recovery + +## Context and Problem Statement + +A declarative workflow can retain requests from several fresh runs. Resuming a +question, tool approval, or agent external-input request must restore its +originating workflow scope rather than inherit the most recent run's MCP session. +Older checkpoints do not retain that origin, so the framework cannot infer which +session should be resumed. + +## Decision Drivers + +- All declarative pause types need the same origin and response-batch rules. +- Existing checkpoints need an explicit recovery path without guessing an origin. +- Repeated restoration must select the same migrated scope without sharing a newer run's session. + +## Considered Options + +- Use the current workflow scope for requests without an origin. This needs no + host changes and retains apparent continuity, but can attach an older request + to an unrelated fresh run. +- Reject all unknown origins and require restarting the workflow. This avoids + guessing, but prevents hosts from recovering otherwise usable pending requests. +- Automatically assign isolated scopes. This permits recovery without host + changes, but silently loses the previous MCP session's server-side state. +- Require explicit host opt-in to assign stable, isolated scopes. This makes the + continuity trade-off visible and permits recovery, but requires host changes + and separate resumes for legacy requests. + +## Decision Outcome + +Choose explicit host opt-in through +`WorkflowFactory(migrate_legacy_request_scopes=True)`, defaulting to `False`: + +- Questions/external input, function approvals, MCP approvals, and agent external + input retain their originating scope in the request. Response handlers restore + it before continuing. Agent requests also retain their own loop configuration. +- Factory-built workflows validate all submitted request origins before consuming + responses. Mixed-scope batches are rejected without changing request metadata. +- Without opt-in, a missing origin produces an actionable recovery error. +- With opt-in, a missing origin becomes + `legacy-request:`. The distinct namespace never selects + the active run's scope; preserving the request ID makes migration stable across + workers and repeated restoration. Legacy requests must be resumed separately. +- Existing valid origins are unchanged. Invalid stored origins are not migrated. + Migration does not change header verification or authorize an unverifiable + MCP header context. +- Newly emitted requests and subsequent checkpoints retain the migrated scope. + The old MCP session and its server-side history cannot be reconstructed. + Hosts requiring that continuity must restart rather than opt in. +- Legacy agent loop configuration remains available from its old checkpoint + slot; new requests own their configuration instead of replacing shared state. + +Checkpoint contents and framework request IDs remain trusted host-owned data. +Migration does not reconstruct state already overwritten in a legacy checkpoint. +This is a checkpoint-compatibility change with an explicit recovery path, not +transparent restoration of an old session. + +The proposed deciders are Python code owners; maintainer and +engineering-management approval is still required. diff --git a/python/packages/declarative/AGENTS.md b/python/packages/declarative/AGENTS.md index 11c2373ea50..4d9f8c5bc9d 100644 --- a/python/packages/declarative/AGENTS.md +++ b/python/packages/declarative/AGENTS.md @@ -81,12 +81,23 @@ arguments. Explicit caller tool arguments in options are preserved. - **`ExternalInputRequest`** / **`ExternalInputResponse`** - Human-in-the-loop support - **`AgentExternalInputRequest`** / **`AgentExternalInputResponse`** - Agent-level input requests -`Question` and `RequestExternalInput` requests retain their originating workflow -scope across fresh runs and checkpoint restores. Responses restore that scope -before updating state or continuing execution. Factory-built workflows reject -response batches spanning different scopes, including mixed MCP approvals and -external input. Legacy external-input requests without a valid originating scope -are rejected before consumption rather than inheriting the current scope. +`Question`, `RequestExternalInput`, function approvals, MCP approvals, and agent +external-input requests retain their originating workflow scope across fresh runs +and checkpoint restores. Responses restore that scope before updating state or +continuing execution. Agent external-input requests also retain their own +`ExternalLoopState` rather than overwriting another pending request's configuration. +Factory-built workflows reject response batches spanning different scopes across +all four request families before consuming any responses. + +Legacy requests without an originating scope are rejected by default with recovery +instructions. Hosts can opt in through +`WorkflowFactory(migrate_legacy_request_scopes=True)`: each submitted legacy request +receives a stable `legacy-request:` scope. Migration never +infers an origin from current state or recovers the old MCP session. Resume requests +separately, preserve their framework IDs, and checkpoint the migrated continuation. +Rejected mixed batches do not mutate request origins. Malformed stored origins +are rejected even with migration enabled. Legacy agent loop configuration is read +from the old checkpoint slot; newly emitted requests retain their own snapshot. ## Usage diff --git a/python/packages/declarative/README.md b/python/packages/declarative/README.md index 0f502ccd003..87be1f421be 100644 --- a/python/packages/declarative/README.md +++ b/python/packages/declarative/README.md @@ -91,3 +91,38 @@ This intentionally incurs connection and initialization overhead and does not retain server session state between provider-backed invocations. Applications that require shared session ownership must implement an explicitly scoped custom `MCPToolHandler`; there is no provider-backed session-cache opt-in. + +## Resuming pending workflow requests + +Questions, external-input requests, function approvals, MCP approvals, and agent +external-input requests retain the workflow scope in which they were issued. +Responses restore that scope even after another fresh run or a checkpoint restore. +Submit responses from different scopes in separate runs; mixed-scope batches are +rejected before any request is consumed. + +Older checkpoints may contain requests without an originating scope. By default, +resuming these requests raises an error without consuming them. To recover, the +host can explicitly opt in when recreating the workflow: + +```python +from agent_framework.declarative import WorkflowFactory + +factory = WorkflowFactory( + checkpoint_storage=storage, + mcp_tool_handler=mcp_handler, + migrate_legacy_request_scopes=True, +) +workflow = factory.create_workflow_from_yaml_path("workflow.yaml") +await workflow.run(checkpoint_id=checkpoint_id) +await workflow.run(responses={pending_request_id: response}) +``` + +Migration gives each legacy request a stable, separately isolated scope derived +from its framework request ID. It never reuses the latest run's scope and cannot +recover the old MCP session or its server-side history. Preserve request IDs, +resume legacy requests individually, and retain the resulting checkpoint. Newly +emitted requests carry the migrated scope, so subsequent restores no longer +require the opt-in for those requests. Any remaining unmigrated requests still do. +Invalid stored scopes are rejected even with migration enabled. Header verification +and any required MCP reapproval remain unchanged. Hosts that require the old +server-side session must restart the workflow instead of migrating it. 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 6702288d72e..9ed8200fd12 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_builder.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_builder.py @@ -15,6 +15,7 @@ import logging from collections.abc import Sequence +from functools import partial from typing import Any, cast from agent_framework import ( @@ -32,7 +33,7 @@ LoopIterationResult, ) from ._errors import DeclarativeWorkflowError -from ._executors_agents import AGENT_ACTION_EXECUTORS, InvokeAzureAgentExecutor +from ._executors_agents import AGENT_ACTION_EXECUTORS, AgentExternalInputRequest, InvokeAzureAgentExecutor from ._executors_basic import BASIC_ACTION_EXECUTORS from ._executors_control_flow import ( CONTROL_FLOW_EXECUTORS, @@ -46,7 +47,6 @@ from ._executors_external_input import ( EXTERNAL_INPUT_EXECUTORS, ExternalInputRequest, - get_external_input_workflow_session_id, ) from ._executors_http import HTTP_ACTION_EXECUTORS, HttpRequestActionExecutor from ._executors_mcp import ( @@ -54,9 +54,9 @@ InvokeMcpToolActionExecutor, MCPToolApprovalRequest, ) -from ._executors_tools import TOOL_ACTION_EXECUTORS, InvokeFunctionToolExecutor +from ._executors_tools import TOOL_ACTION_EXECUTORS, InvokeFunctionToolExecutor, ToolApprovalRequest from ._http_handler import HttpRequestHandler -from ._mcp_handler import MCPToolHandler, activate_workflow_session_id +from ._mcp_handler import MCPToolHandler, activate_workflow_session_id, get_request_workflow_session_id logger = logging.getLogger(__name__) @@ -120,20 +120,29 @@ } -def _validate_mcp_response_batch(requests: Sequence[WorkflowEvent[Any]]) -> None: - scopes: set[str | None] = set() +def _validate_response_batch( + requests: Sequence[WorkflowEvent[Any]], *, migrate_legacy_request_scopes: bool = False +) -> None: + scoped_requests: list[ + tuple[ExternalInputRequest | MCPToolApprovalRequest | ToolApprovalRequest | AgentExternalInputRequest, str] + ] = [] for request in requests: - if isinstance(request.data, ExternalInputRequest): - scope = get_external_input_workflow_session_id(request.data) - elif 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.") - else: + data = request.data + if not isinstance( + data, (ExternalInputRequest, MCPToolApprovalRequest, ToolApprovalRequest, AgentExternalInputRequest) + ): continue - scopes.add(scope) + scope = data.workflow_session_id + if scope is None and migrate_legacy_request_scopes: + if not isinstance(request.request_id, str) or not request.request_id: + raise DeclarativeWorkflowError("Cannot migrate a legacy request without a valid request ID.") + scope = f"legacy-request:{request.request_id}" + scoped_requests.append((data, get_request_workflow_session_id(scope))) + scopes = {scope for _, scope in scoped_requests} if len(scopes) > 1: raise DeclarativeWorkflowError("Responses from different workflow sessions must be submitted in separate runs.") + for data, scope in scoped_requests: + data.workflow_session_id = scope def _initialize_mcp_run_state(state: State) -> None: @@ -173,6 +182,7 @@ def __init__( http_request_handler: HttpRequestHandler | None = None, mcp_tool_handler: MCPToolHandler | None = None, env_config: DeclarativeEnvConfig | None = None, + migrate_legacy_request_scopes: bool = False, ): """Initialize the builder. @@ -195,6 +205,8 @@ def __init__( how the ``Env`` PowerFx symbol is populated for every executor built by this builder. Defaults to an empty configuration (``Env`` not exposed). + migrate_legacy_request_scopes: Opt in to assigning legacy pending requests + stable, isolated scopes instead of rejecting missing origins. """ self._yaml_def = yaml_definition self._workflow_id = workflow_id or yaml_definition.get("name", "declarative_workflow") @@ -210,6 +222,7 @@ def __init__( self._http_request_handler = http_request_handler self._mcp_tool_handler = mcp_tool_handler self._env_config: DeclarativeEnvConfig = env_config if env_config is not None else DeclarativeEnvConfig() + self._migrate_legacy_request_scopes = migrate_legacy_request_scopes # Resolve max_iterations: explicit arg > YAML maxTurns > core default resolved = max_iterations if max_iterations is not None else yaml_definition.get("maxTurns") if resolved is not None and (not isinstance(resolved, int) or resolved <= 0): @@ -273,7 +286,9 @@ def build(self) -> Workflow: executor.set_declarative_env_config(self._env_config) workflow = builder.build() - workflow._response_batch_validator = _validate_mcp_response_batch # pyright: ignore[reportPrivateUsage] + workflow._response_batch_validator = partial( # pyright: ignore[reportPrivateUsage] + _validate_response_batch, migrate_legacy_request_scopes=self._migrate_legacy_request_scopes + ) workflow._initial_state_initializer = _initialize_mcp_run_state # pyright: ignore[reportPrivateUsage] return workflow diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_agents.py b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_agents.py index 30de4b7eb00..c0d6cc4acc8 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_agents.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_agents.py @@ -12,7 +12,8 @@ - AutoSend behavior control """ -import contextlib +from __future__ import annotations + import json import logging import uuid @@ -34,6 +35,11 @@ DeclarativeActionExecutor, DeclarativeWorkflowState, ) +from ._mcp_handler import ( + get_or_create_workflow_session_id, + get_request_workflow_session_id, + restore_workflow_session_id, +) logger = logging.getLogger(__name__) @@ -432,6 +438,8 @@ async def on_request(request: AgentExternalInputRequest) -> ExternalInputRespons iteration: int = 0 messages: list[Message] = field(default_factory=lambda: cast(list[Message], [])) function_calls: list[Content] = field(default_factory=lambda: cast(list[Content], [])) + workflow_session_id: str | None = None + loop_state: ExternalLoopState | None = None @dataclass @@ -469,8 +477,8 @@ class AgentExternalInputResponse: class ExternalLoopState: """State saved for external loop resumption. - Stored in workflow state to allow the response_handler to - continue the loop with the same configuration. + Retained by each pending request so interleaved continuations + preserve their own configuration. """ agent_name: str @@ -1006,8 +1014,6 @@ async def handle_action( messages_path=messages_path, max_iterations=max_iterations, ) - ctx.state.set(EXTERNAL_LOOP_STATE_KEY, loop_state) - # Emit request for external input - workflow will yield here request = AgentExternalInputRequest( request_id=str(uuid.uuid4()), @@ -1016,6 +1022,8 @@ async def handle_action( iteration=iteration, messages=all_messages, function_calls=tool_calls, + workflow_session_id=get_or_create_workflow_session_id(ctx.state), + loop_state=loop_state, ) logger.info(f"InvokeAzureAgent: yielding for external input (iteration {iteration})") await ctx.request_info(request, AgentExternalInputResponse) @@ -1037,14 +1045,17 @@ async def handle_external_input_response( This is called when the workflow resumes after yielding for external input. It continues the agent invocation loop with the user's new input. """ + restore_workflow_session_id(ctx.state, get_request_workflow_session_id(original_request.workflow_session_id)) logger.debug( "handle_external_input_response: resuming with user_input='%s'", response.user_input[:100] if response.user_input else None, ) state = self._get_state(ctx.state) - # Retrieve saved loop state - loop_state: ExternalLoopState | None = ctx.state.get(EXTERNAL_LOOP_STATE_KEY) + # Older checkpoints store the loop configuration in a single shared slot. + loop_state = original_request.loop_state + if loop_state is None: + loop_state = ctx.state.get(EXTERNAL_LOOP_STATE_KEY) if loop_state is None: logger.error("InvokeAzureAgent: external loop state not found, cannot resume") await ctx.send_message(ActionComplete()) @@ -1081,9 +1092,7 @@ async def handle_external_input_response( ) if not should_continue: - # User input caused loop to exit - clean up and complete - with contextlib.suppress(KeyError): - ctx.state.delete(EXTERNAL_LOOP_STATE_KEY) + # User input caused loop to exit. await ctx.send_message(ActionComplete()) return @@ -1141,8 +1150,6 @@ async def handle_external_input_response( "InvokeAzureAgent: external loop exited due to condition=false " "(sending ActionComplete to continue workflow)" ) - with contextlib.suppress(KeyError): - ctx.state.delete(EXTERNAL_LOOP_STATE_KEY) await ctx.send_message(ActionComplete()) return @@ -1150,8 +1157,6 @@ async def handle_external_input_response( if iteration < max_iterations: # Update loop state for next iteration loop_state.iteration = iteration + 1 - ctx.state.set(EXTERNAL_LOOP_STATE_KEY, loop_state) - # Emit another request for external input request = AgentExternalInputRequest( request_id=str(uuid.uuid4()), @@ -1160,6 +1165,8 @@ async def handle_external_input_response( iteration=iteration, messages=all_messages, function_calls=tool_calls, + workflow_session_id=get_or_create_workflow_session_id(ctx.state), + loop_state=loop_state, ) logger.info(f"InvokeAzureAgent: yielding for external input (iteration {iteration})") await ctx.request_info(request, AgentExternalInputResponse) @@ -1167,10 +1174,6 @@ async def handle_external_input_response( logger.warning(f"InvokeAzureAgent: external loop exceeded max iterations ({max_iterations})") - # Loop complete - clean up and send completion - with contextlib.suppress(KeyError): - ctx.state.delete(EXTERNAL_LOOP_STATE_KEY) - await ctx.send_message(ActionComplete()) diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_external_input.py b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_external_input.py index 30208fd3b70..edf8b78eaa8 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_external_input.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_external_input.py @@ -21,8 +21,11 @@ ActionComplete, DeclarativeActionExecutor, ) -from ._errors import DeclarativeWorkflowError -from ._mcp_handler import get_or_create_workflow_session_id, restore_workflow_session_id +from ._mcp_handler import ( + get_or_create_workflow_session_id, + get_request_workflow_session_id, + restore_workflow_session_id, +) def _get_prompt_text(action_def: dict[str, Any], primary_key: str, fallback_key: str) -> Any: @@ -79,14 +82,6 @@ class ExternalInputRequest: workflow_session_id: str | None = None -def get_external_input_workflow_session_id(request: ExternalInputRequest) -> str: - """Return the originating scope, rejecting requests with an unknown origin.""" - scope = request.workflow_session_id - if not isinstance(scope, str) or not scope: - raise DeclarativeWorkflowError("External input request is missing a valid originating workflow session.") - return scope - - @dataclass class ExternalInputResponse: """Response to an ExternalInputRequest. @@ -169,7 +164,7 @@ async def handle_response( ctx: WorkflowContext[ActionComplete], ) -> None: """Handle the user's response to the question.""" - restore_workflow_session_id(ctx.state, get_external_input_workflow_session_id(original_request)) + restore_workflow_session_id(ctx.state, get_request_workflow_session_id(original_request.workflow_session_id)) state = self._get_state(ctx.state) output_property = original_request.metadata.get("output_property", "Local.answer") @@ -239,7 +234,7 @@ async def handle_response( ctx: WorkflowContext[ActionComplete], ) -> None: """Handle the external input response.""" - restore_workflow_session_id(ctx.state, get_external_input_workflow_session_id(original_request)) + restore_workflow_session_id(ctx.state, get_request_workflow_session_id(original_request.workflow_session_id)) state = self._get_state(ctx.state) output_property = original_request.metadata.get("output_property", "Local.externalInput") 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 cf8d8da63eb..53420b251ab 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, + get_request_workflow_session_id, restore_workflow_session_id, ) @@ -324,6 +325,8 @@ async def handle_approval_response( ctx: WorkflowContext[ActionComplete, str], ) -> None: """Resume the pinned operation only with the reviewed header context.""" + workflow_session_id = get_request_workflow_session_id(original_request.workflow_session_id) + restore_workflow_session_id(ctx.state, workflow_session_id) state = self._get_state(ctx.state) tool_name = original_request.tool_name @@ -335,11 +338,6 @@ 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", diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_tools.py b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_tools.py index 1d642c91c3a..7d56f8f874d 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_tools.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_tools.py @@ -34,6 +34,11 @@ DeclarativeWorkflowState, ) from ._executors_agents import TOOL_REGISTRY_KEY +from ._mcp_handler import ( + get_or_create_workflow_session_id, + get_request_workflow_session_id, + restore_workflow_session_id, +) logger = logging.getLogger(__name__) @@ -61,11 +66,13 @@ class ToolApprovalRequest: request_id: Unique identifier for this approval request. function_name: Evaluated function name to be invoked. arguments: Evaluated arguments to be passed to the function. + workflow_session_id: Framework-owned scope of the workflow that issued the request. """ request_id: str function_name: str arguments: dict[str, Any] + workflow_session_id: str | None = None @dataclass @@ -488,6 +495,7 @@ async def handle_action( request_id=request_id, function_name=function_name, arguments=arguments, + workflow_session_id=get_or_create_workflow_session_id(ctx.state), ) logger.info(f"{self.__class__.__name__}: requesting approval for '{function_name}'") await ctx.request_info(request, ToolApprovalResponse, request_id=request_id) @@ -524,6 +532,7 @@ async def handle_approval_response( Rejected calls store the rejection and complete without evaluating ``autoSend``. """ + restore_workflow_session_id(ctx.state, get_request_workflow_session_id(original_request.workflow_session_id)) state = self._get_state(ctx.state) function_name = original_request.function_name diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py b/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py index a8329a59734..b80062659b0 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py @@ -98,6 +98,7 @@ def __init__( mcp_tool_handler: MCPToolHandler | None = None, configuration: Mapping[str, str] | None = None, restrict_env_to_configuration: bool = True, + migrate_legacy_request_scopes: bool = False, ) -> None: """Initialize the workflow factory. @@ -141,6 +142,12 @@ def __init__( discovered in PowerFx expressions inside the workflow definition so unrelated environment variables never enter the PowerFx scope. + migrate_legacy_request_scopes: Opt in to resuming legacy pending requests + without an originating workflow scope. Each request receives a stable, + isolated scope derived from its framework request ID, never the active + run's scope. The old MCP session cannot be recovered; resume requests + separately and checkpoint the migrated continuation. Defaults to + ``False``, which rejects unknown origins before consuming responses. Examples: .. code-block:: python @@ -196,6 +203,7 @@ def __init__( self._mcp_tool_handler = mcp_tool_handler self._configuration: dict[str, str] = dict(configuration) if configuration else {} self._restrict_env_to_configuration = restrict_env_to_configuration + self._migrate_legacy_request_scopes = migrate_legacy_request_scopes def create_workflow_from_yaml_path( self, @@ -455,6 +463,7 @@ def _create_workflow( http_request_handler=self._http_request_handler, mcp_tool_handler=self._mcp_tool_handler, env_config=env_config, + migrate_legacy_request_scopes=self._migrate_legacy_request_scopes, ) workflow = graph_builder.build() except ValueError as e: 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 90d63d0b117..713f2d12a2c 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -41,6 +41,8 @@ import httpx2 as httpx +from ._errors import DeclarativeWorkflowError + if TYPE_CHECKING: from agent_framework import Content from agent_framework._workflows._state import State @@ -113,6 +115,19 @@ def restore_workflow_session_id(state: State, workflow_session_id: str) -> None: state.set(_WORKFLOW_SESSION_ID_KEY, workflow_session_id) +def get_request_workflow_session_id(workflow_session_id: str | None) -> str: + """Validate a pending request's scope without inferring it from the active run.""" + if workflow_session_id is None: + raise DeclarativeWorkflowError( + "Request is missing a valid originating workflow session. " + "Restore the checkpoint with WorkflowFactory(migrate_legacy_request_scopes=True) " + "to resume legacy requests in isolated scopes, or restart the workflow." + ) + if not isinstance(workflow_session_id, str) or not workflow_session_id: + raise DeclarativeWorkflowError("Invalid originating workflow session state.") + return 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_declarative_approval_binding.py b/python/packages/declarative/tests/test_declarative_approval_binding.py index 441f73891d0..a12ce501d1c 100644 --- a/python/packages/declarative/tests/test_declarative_approval_binding.py +++ b/python/packages/declarative/tests/test_declarative_approval_binding.py @@ -177,7 +177,9 @@ def my_tool(x: int) -> int: executor = InvokeFunctionToolExecutor(self._action(), tools={"my_tool": my_tool}) - request = ToolApprovalRequest(request_id="r-1", function_name="my_tool", arguments={"x": 1}) + request = ToolApprovalRequest( + request_id="r-1", function_name="my_tool", arguments={"x": 1}, workflow_session_id="workflow-a" + ) await executor.handle_approval_response(request, ToolApprovalResponse(approved=True), mock_context) assert call_log == [1] @@ -193,7 +195,9 @@ def my_tool(x: int) -> int: return x executor = InvokeFunctionToolExecutor(self._action(), tools={"my_tool": my_tool}) - request = ToolApprovalRequest(request_id="r-bool", function_name="my_tool", arguments={"x": 1}) + request = ToolApprovalRequest( + request_id="r-bool", function_name="my_tool", arguments={"x": 1}, workflow_session_id="workflow-a" + ) response = object.__new__(ToolApprovalResponse) response.approved = "true" # type: ignore[assignment] # ty: ignore[invalid-assignment] response.reason = None @@ -214,8 +218,12 @@ def my_tool(x: int) -> int: executor = InvokeFunctionToolExecutor(self._action(), tools={"my_tool": my_tool}) - request_a = ToolApprovalRequest(request_id="r-A", function_name="my_tool", arguments={"x": 1}) - request_b = ToolApprovalRequest(request_id="r-B", function_name="my_tool", arguments={"x": 999}) + request_a = ToolApprovalRequest( + request_id="r-A", function_name="my_tool", arguments={"x": 1}, workflow_session_id="workflow-a" + ) + request_b = ToolApprovalRequest( + request_id="r-B", function_name="my_tool", arguments={"x": 999}, workflow_session_id="workflow-b" + ) # Deliver response for B first, then for A. Each invocation must use its own payload. await executor.handle_approval_response(request_b, ToolApprovalResponse(approved=True), mock_context) @@ -238,7 +246,9 @@ def my_tool(x: int) -> int: # Poison the old key shape (no longer read by the executor). mock_state._data["_tool_approval_state_fn_action"] = {"function_name": "other", "arguments": {"x": 999}} - request = ToolApprovalRequest(request_id="r-3", function_name="my_tool", arguments={"x": 7}) + request = ToolApprovalRequest( + request_id="r-3", function_name="my_tool", arguments={"x": 7}, workflow_session_id="workflow-a" + ) await executor.handle_approval_response(request, ToolApprovalResponse(approved=True), mock_context) assert call_log == [7] @@ -258,7 +268,9 @@ def my_tool(x: int) -> int: # Pretend the executor that emitted the request is gone; a fresh one handles the response. fresh = InvokeFunctionToolExecutor(self._action(), tools={"my_tool": my_tool}) - request = ToolApprovalRequest(request_id="r-4", function_name="my_tool", arguments={"x": 42}) + request = ToolApprovalRequest( + request_id="r-4", function_name="my_tool", arguments={"x": 42}, workflow_session_id="workflow-a" + ) await fresh.handle_approval_response(request, ToolApprovalResponse(approved=True), mock_context) assert call_log == [42] @@ -275,7 +287,9 @@ def my_tool(x: int) -> int: executor = InvokeFunctionToolExecutor(self._action(), tools={"my_tool": my_tool}) - request = ToolApprovalRequest(request_id="r-5", function_name="my_tool", arguments={"x": 3}) + request = ToolApprovalRequest( + request_id="r-5", function_name="my_tool", arguments={"x": 3}, workflow_session_id="workflow-a" + ) await executor.handle_approval_response( request, ToolApprovalResponse(approved=False, reason="not authorized"), mock_context ) @@ -333,6 +347,7 @@ async def test_resume_uses_request_payload_fields(self, mock_state, mock_context server_label="prod", arguments={"q": "x"}, connection_name="conn-A", + workflow_session_id="workflow-a", ) await executor.handle_approval_response(request, ToolApprovalResponse(approved=True), mock_context) @@ -358,6 +373,7 @@ async def test_truthy_malformed_response_does_not_invoke(self, mock_state, mock_ server_label=None, arguments={"q": "x"}, connection_name=None, + workflow_session_id="workflow-a", ) response = object.__new__(ToolApprovalResponse) response.approved = 1 # type: ignore[assignment] # ty: ignore[invalid-assignment] @@ -380,6 +396,7 @@ async def test_concurrent_pending_mcp_approvals_do_not_swap(self, mock_state, mo server_label=None, arguments={"q": "alpha"}, connection_name="conn-A", + workflow_session_id="workflow-a", ) request_b = MCPToolApprovalRequest( request_id="r-B", @@ -388,6 +405,7 @@ async def test_concurrent_pending_mcp_approvals_do_not_swap(self, mock_state, mo server_label=None, arguments={"q": "beta"}, connection_name="conn-B", + workflow_session_id="workflow-b", ) await executor.handle_approval_response(request_b, ToolApprovalResponse(approved=True), mock_context) @@ -416,6 +434,7 @@ async def test_legacy_headers_require_fresh_approval(self, mock_state, mock_cont server_label=None, arguments={"q": "x"}, connection_name=None, + workflow_session_id="legacy-request:r-1", ) await executor.handle_approval_response(request, ToolApprovalResponse(approved=True), mock_context) @@ -534,6 +553,7 @@ async def test_mcp_resume_ignores_stale_state_at_old_approval_key(self, mock_sta server_label=None, arguments={"q": "real"}, connection_name=None, + workflow_session_id="workflow-a", ) await executor.handle_approval_response(request, ToolApprovalResponse(approved=True), mock_context) @@ -557,6 +577,7 @@ async def test_fresh_mcp_executor_resume_works(self, mock_state, mock_context) - server_label=None, arguments={"q": "fresh"}, connection_name=None, + workflow_session_id="workflow-a", ) await fresh.handle_approval_response(request, ToolApprovalResponse(approved=True), mock_context) @@ -625,6 +646,7 @@ async def test_resume_routes_output_to_pinned_conversation_not_mutated_state( arguments={"q": "x"}, connection_name=None, metadata={"conversation_id": "conv-original"}, + workflow_session_id="workflow-a", ) await executor.handle_approval_response(original_request, ToolApprovalResponse(approved=True), mock_context) @@ -632,8 +654,10 @@ async def test_resume_routes_output_to_pinned_conversation_not_mutated_state( assert state.get("System.conversations.conv-mutated.messages") == [] @pytest.mark.asyncio - async def test_resume_handles_legacy_request_without_new_fields(self, mock_state, mock_context) -> None: - """Resume tolerates payloads lacking ``connection_name`` / ``metadata`` (legacy pickle shape).""" + async def test_resume_handles_migrated_legacy_request_without_routing_fields( + self, mock_state, mock_context + ) -> None: + """Migrated legacy payloads may still lack ``connection_name`` / ``metadata``.""" @dataclass class _LegacyMCPApprovalRequest: @@ -643,6 +667,7 @@ class _LegacyMCPApprovalRequest: server_label: str | None arguments: dict[str, Any] header_names: list[str] + workflow_session_id: str = "legacy-request:r-1" _seed_state(mock_state) handler = _RecordingMcpHandler() diff --git a/python/packages/declarative/tests/test_function_tool_executor.py b/python/packages/declarative/tests/test_function_tool_executor.py index 3982831a9fc..211a3b698d2 100644 --- a/python/packages/declarative/tests/test_function_tool_executor.py +++ b/python/packages/declarative/tests/test_function_tool_executor.py @@ -763,11 +763,7 @@ def mock_state() -> MagicMock: mock_state._data = {} def mock_get(key: str, default: Any = None) -> Any: - if key not in mock_state._data: - if default is not None: - return default - raise KeyError(key) - return mock_state._data[key] + return mock_state._data.get(key, default) def mock_set(key: str, value: Any) -> None: mock_state._data[key] = value @@ -1205,6 +1201,7 @@ def my_tool(x: int) -> int: request_id="req-123", function_name="my_tool", arguments={"x": 7}, + workflow_session_id="workflow-a", ) response = ToolApprovalResponse(approved=True) @@ -1241,6 +1238,7 @@ def my_tool(x: int) -> int: request_id="req-456", function_name="my_tool", arguments={"x": 5}, + workflow_session_id="workflow-a", ) response = ToolApprovalResponse(approved=False, reason="Not authorized") diff --git a/python/packages/declarative/tests/test_graph_coverage.py b/python/packages/declarative/tests/test_graph_coverage.py index 5b150bf497f..f113d8ed2e0 100644 --- a/python/packages/declarative/tests/test_graph_coverage.py +++ b/python/packages/declarative/tests/test_graph_coverage.py @@ -2780,6 +2780,7 @@ async def test_handle_external_input_response_no_state(self, mock_context, mock_ agent_name="TestAgent", agent_response="Hello", iteration=1, + workflow_session_id="workflow-a", ) response = AgentExternalInputResponse(user_input="hi there") @@ -2828,6 +2829,7 @@ async def test_handle_external_input_response_agent_not_found(self, mock_context agent_name="NonExistentAgent", agent_response="Hello", iteration=1, + workflow_session_id="workflow-a", ) response = AgentExternalInputResponse(user_input="continue") 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 3c53693f1b9..45c2bd77bbf 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -9,7 +9,9 @@ """ import sys +from pathlib import Path from typing import Any +from unittest.mock import AsyncMock, MagicMock import httpx2 as httpx import pytest @@ -799,7 +801,7 @@ async def test_same_scope_approval_batch_is_supported(self, streaming: bool, app 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._declarative_builder import _validate_response_batch from agent_framework_declarative._workflows._executors_control_flow import JoinExecutor handler = StubMcpHandler(_ok()) @@ -809,7 +811,7 @@ async def test_same_scope_approval_batch_is_supported(self, streaming: bool, app 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 + workflow._response_batch_validator = _validate_response_batch paused = await workflow.run({}) approvals = paused.get_request_info_events() assert len(approvals) == 2 @@ -896,6 +898,335 @@ async def test_server_label_and_connection_name_forwarded(self) -> None: assert inv.connection_name == "azure-conn" +def _scope_pause_action(kind: str) -> dict[str, Any]: + actions: dict[str, dict[str, Any]] = { + "Question": {"kind": "Question", "question": "Continue?", "variable": "Local.answer"}, + "RequestExternalInput": { + "kind": "RequestExternalInput", + "prompt": "Continue?", + "variable": "Local.answer", + }, + "InvokeFunctionTool": {"kind": "InvokeFunctionTool", "functionName": "echo", "requireApproval": True}, + "InvokeAzureAgent": { + "kind": "InvokeAzureAgent", + "agent": "TestAgent", + "input": {"externalLoop": {"when": '=Local.userInput <> "stop"'}}, + }, + "InvokeMcpTool": _action(require_approval=True), + } + return actions[kind] + + +def _scope_pause_response(kind: str, user_input: str = "stop") -> Any: + from agent_framework_declarative._workflows import ( + AgentExternalInputResponse, + ExternalInputResponse, + ToolApprovalResponse, + ) + + if kind == "InvokeAzureAgent": + return AgentExternalInputResponse(user_input=user_input) + if kind in ("InvokeFunctionTool", "InvokeMcpTool"): + return ToolApprovalResponse(approved=True) + return ExternalInputResponse(user_input=user_input) + + +def _scope_checkpoint_storage(path: Path) -> Any: + from agent_framework import FileCheckpointStorage + + return FileCheckpointStorage( + path, + allowed_checkpoint_types=[ + "agent_framework_declarative._workflows._declarative_base:ActionComplete", + "agent_framework_declarative._workflows._declarative_base:ActionTrigger", + "agent_framework_declarative._workflows._executors_tools:ToolApprovalRequest", + "agent_framework_declarative._workflows._executors_tools:ToolApprovalResponse", + "agent_framework_declarative._workflows._executors_mcp:MCPToolApprovalRequest", + "agent_framework_declarative._workflows._executors_agents:AgentExternalInputRequest", + "agent_framework_declarative._workflows._executors_agents:AgentExternalInputResponse", + "agent_framework_declarative._workflows._executors_agents:ExternalLoopState", + "agent_framework_declarative._workflows._executors_external_input:ExternalInputRequest", + "agent_framework_declarative._workflows._executors_external_input:ExternalInputResponse", + ], + ) + + +class TestRequestScopeContinuations: + @pytest.mark.parametrize("pause_kind", ["InvokeFunctionTool", "InvokeAzureAgent"]) + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("restore", [False, True]) + async def test_remaining_request_types_restore_originating_scope( + self, tmp_path: Path, pause_kind: str, streaming: bool, restore: bool + ) -> None: + storage = _scope_checkpoint_storage(tmp_path) + handler = StubMcpHandler(_ok()) + function = MagicMock(return_value="done") + agent = MagicMock(run=AsyncMock(return_value="Need input")) + factory = WorkflowFactory(mcp_tool_handler=handler, agents={"TestAgent": agent}).register_tool("echo", function) + definition = { + "name": "request_continuations", + "actions": [ + {**_action(tool_name="before"), "id": "before"}, + {"kind": "SetValue", "path": "Local.userInput", "value": ""}, + {**_scope_pause_action(pause_kind), "id": "pause"}, + {**_action(tool_name="after"), "id": "after"}, + ], + } + workflow = factory.create_workflow_from_definition(definition) + requests: list[Any] = [] + for _ in range(2): + await workflow.run({}, checkpoint_storage=storage) + [request] = [ + event + for event in (await workflow._runner_context.get_pending_request_info_events()).values() + if event.request_id not in {previous.request_id for previous in requests} + ] + requests.append(request) + scopes = [invocation.workflow_session_id for invocation in handler.invocations] + assert scopes[0] and scopes[1] and scopes[0] != scopes[1] + 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) + assert set(await workflow._runner_context.get_pending_request_info_events()) == { + request.request_id for request in requests + } + + for request, scope in zip(requests, scopes, strict=True): + responses = {request.request_id: _scope_pause_response(pause_kind)} + if streaming: + await workflow.run(responses=responses, stream=True).get_final_response() + else: + await workflow.run(responses=responses) + assert handler.invocations[-1].tool_name == "after" + assert handler.invocations[-1].workflow_session_id == scope + assert handler.call_count == 4 + assert function.call_count == (2 if pause_kind == "InvokeFunctionTool" else 0) + assert not await workflow._runner_context.get_pending_request_info_events() + + @pytest.mark.parametrize("streaming", [False, True]) + async def test_repeated_agent_input_retains_its_own_loop_state(self, tmp_path: Path, streaming: bool) -> None: + from agent_framework_declarative._workflows._mcp_handler import get_or_create_workflow_session_id + + storage = _scope_checkpoint_storage(tmp_path) + handler = StubMcpHandler(_ok()) + agent = MagicMock(run=AsyncMock(return_value="Need input")) + factory = WorkflowFactory(mcp_tool_handler=handler, agents={"TestAgent": agent}) + definition = { + "name": "repeated_agent_input", + "actions": [ + {**_action(tool_name="before"), "id": "before"}, + {"kind": "SetValue", "path": "Local.userInput", "value": ""}, + {**_scope_pause_action("InvokeAzureAgent"), "id": "pause"}, + {**_action(tool_name="after"), "id": "after"}, + ], + } + workflow = factory.create_workflow_from_definition(definition) + await workflow.run({}) + [first] = (await workflow._runner_context.get_pending_request_info_events()).values() + await workflow.run({}) + [second] = [ + event + for event in (await workflow._runner_context.get_pending_request_info_events()).values() + if event.request_id != first.request_id + ] + first_scope, second_scope = [invocation.workflow_session_id for invocation in handler.invocations] + await workflow.run( + responses={first.request_id: _scope_pause_response("InvokeAzureAgent", "continue")}, + checkpoint_storage=storage, + ) + assert get_or_create_workflow_session_id(workflow._runner.state) == first_scope + [repeated] = [ + event + for event in (await workflow._runner_context.get_pending_request_info_events()).values() + if event.request_id != second.request_id + ] + 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 request, scope in ((second, second_scope), (repeated, first_scope)): + responses = {request.request_id: _scope_pause_response("InvokeAzureAgent")} + if streaming: + await workflow.run(responses=responses, stream=True).get_final_response() + else: + await workflow.run(responses=responses) + assert handler.invocations[-1].tool_name == "after" + assert handler.invocations[-1].workflow_session_id == scope + assert handler.call_count == 4 + assert agent.run.await_count == 3 + assert not await workflow._runner_context.get_pending_request_info_events() + + @pytest.mark.parametrize("pause_kind", ["InvokeFunctionTool", "InvokeAzureAgent"]) + @pytest.mark.parametrize("streaming", [False, True]) + async def test_mixed_request_families_reject_different_scopes(self, pause_kind: str, streaming: bool) -> None: + handler = StubMcpHandler(_ok()) + function = MagicMock(return_value="done") + agent = MagicMock(run=AsyncMock(return_value="Need input")) + factory = WorkflowFactory(mcp_tool_handler=handler, agents={"TestAgent": agent}).register_tool("echo", function) + workflow = factory.create_workflow_from_definition({ + "name": "mixed_request_scopes", + "actions": [ + {"kind": "Question", "question": "Continue?", "variable": "Local.answer"}, + {"kind": "SetValue", "path": "Local.userInput", "value": ""}, + {**_scope_pause_action(pause_kind), "id": "pause"}, + _action(), + ], + }) + await workflow.run({}) + [first] = (await workflow._runner_context.get_pending_request_info_events()).values() + await workflow.run({}) + [second] = [ + event + for event in (await workflow._runner_context.get_pending_request_info_events()).values() + if event.request_id != first.request_id + ] + await workflow.run(responses={second.request_id: _scope_pause_response("Question")}) + [second] = [ + event + for event in (await workflow._runner_context.get_pending_request_info_events()).values() + if event.request_id != first.request_id + ] + responses = { + first.request_id: _scope_pause_response("Question"), + second.request_id: _scope_pause_response(pause_kind), + } + with pytest.raises(DeclarativeWorkflowError, match="different workflow sessions"): + if streaming: + await workflow.run(responses=responses, stream=True).get_final_response() + else: + await workflow.run(responses=responses) + assert handler.call_count == 0 + function.assert_not_called() + assert agent.run.await_count == (1 if pause_kind == "InvokeAzureAgent" else 0) + assert set(await workflow._runner_context.get_pending_request_info_events()) == set(responses) + + async def test_legacy_mcp_origin_is_rejected_before_consumption(self) -> None: + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + workflow = factory.create_workflow_from_definition(_yaml(_action(require_approval=True))) + await workflow.run({}) + [legacy] = (await workflow._runner_context.get_pending_request_info_events()).values() + legacy.data.workflow_session_id = None + await workflow.run({}) + + with pytest.raises(DeclarativeWorkflowError, match="migrate_legacy_request_scopes"): + await workflow.run(responses={legacy.request_id: _scope_pause_response("InvokeMcpTool")}) + + assert handler.call_count == 0 + assert legacy.request_id in await workflow._runner_context.get_pending_request_info_events() + + @pytest.mark.parametrize( + "pause_kind", ["Question", "RequestExternalInput", "InvokeFunctionTool", "InvokeAzureAgent", "InvokeMcpTool"] + ) + async def test_opt_in_legacy_migration_survives_file_checkpoints(self, tmp_path: Path, pause_kind: str) -> None: + storage = _scope_checkpoint_storage(tmp_path) + handler = StubMcpHandler(_ok()) + agent = MagicMock(run=AsyncMock(return_value="Need input")) + function = MagicMock(return_value="done") + factory = WorkflowFactory(mcp_tool_handler=handler, agents={"TestAgent": agent}).register_tool("echo", function) + definition = { + "name": "legacy_request_migration", + "actions": [ + {**_action(tool_name="before"), "id": "before"}, + {"kind": "SetValue", "path": "Local.userInput", "value": ""}, + {**_scope_pause_action(pause_kind), "id": "pause"}, + {**_action(tool_name="after"), "id": "after"}, + {**_scope_pause_action("Question"), "id": "followup"}, + {**_action(tool_name="done"), "id": "done"}, + ], + } + workflow = factory.create_workflow_from_definition(definition) + await workflow.run({}, checkpoint_storage=storage) + await workflow.run({}, checkpoint_storage=storage) + original_scopes = {invocation.workflow_session_id for invocation in handler.invocations} + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + checkpoint = await storage.load(checkpoint_id) + assert len(checkpoint.pending_request_info_events) == 2 + for event in checkpoint.pending_request_info_events.values(): + event.data.workflow_session_id = None + if pause_kind == "InvokeAzureAgent": + from agent_framework_declarative._workflows._executors_agents import EXTERNAL_LOOP_STATE_KEY + + checkpoint.state[EXTERNAL_LOOP_STATE_KEY] = event.data.loop_state + del event.data.loop_state + await storage.save(checkpoint) + first_id, second_id = checkpoint.pending_request_info_events + + workflow = factory.create_workflow_from_definition(definition) + await workflow.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage) + with pytest.raises(DeclarativeWorkflowError, match="migrate_legacy_request_scopes"): + await workflow.run(responses={first_id: _scope_pause_response(pause_kind)}) + assert handler.call_count == 2 + assert set(await workflow._runner_context.get_pending_request_info_events()) == {first_id, second_id} + + migration_factory = WorkflowFactory( + mcp_tool_handler=handler, agents={"TestAgent": agent}, migrate_legacy_request_scopes=True + ).register_tool("echo", function) + workflow = migration_factory.create_workflow_from_definition(definition) + await workflow.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage) + with pytest.raises(DeclarativeWorkflowError, match="different workflow sessions"): + await workflow.run( + responses={ + first_id: _scope_pause_response(pause_kind), + second_id: _scope_pause_response(pause_kind), + } + ) + pending = await workflow._runner_context.get_pending_request_info_events() + assert all(event.data.workflow_session_id is None for event in pending.values()) + assert handler.call_count == 2 + function.assert_not_called() + + migrated_scopes: list[str] = [] + for _ in range(2): + workflow = migration_factory.create_workflow_from_definition(definition) + await workflow.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage) + before = handler.call_count + await workflow.run(responses={first_id: _scope_pause_response(pause_kind)}, checkpoint_storage=storage) + scope = handler.invocations[-1].workflow_session_id + assert scope and scope not in original_scopes + assert all(invocation.workflow_session_id == scope for invocation in handler.invocations[before:]) + migrated_scopes.append(scope) + assert migrated_scopes[0] == migrated_scopes[1] + + migrated_checkpoint_id = workflow.get_last_checkpoint_id() + assert migrated_checkpoint_id is not None + workflow = factory.create_workflow_from_definition(definition) + await workflow.run(checkpoint_id=migrated_checkpoint_id, checkpoint_storage=storage) + pending = await workflow._runner_context.get_pending_request_info_events() + [followup] = [event for event in pending.values() if event.request_id != second_id] + assert followup.data.workflow_session_id == migrated_scopes[0] + await workflow.run(responses={followup.request_id: _scope_pause_response("Question")}) + assert handler.invocations[-1].tool_name == "done" + assert handler.invocations[-1].workflow_session_id == migrated_scopes[0] + + migrated_checkpoint = await storage.load(migrated_checkpoint_id) + workflow = migration_factory.create_workflow_from_definition(definition) + await workflow.run(checkpoint_id=migrated_checkpoint.checkpoint_id, checkpoint_storage=storage) + await workflow.run(responses={second_id: _scope_pause_response(pause_kind)}) + assert handler.invocations[-1].tool_name == "after" + assert handler.invocations[-1].workflow_session_id not in original_scopes | set(migrated_scopes) + + @pytest.mark.parametrize("invalid_scope", ["", 42]) + @pytest.mark.parametrize("migrate", [False, True]) + async def test_migration_does_not_replace_invalid_scopes(self, invalid_scope: Any, migrate: bool) -> None: + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler, migrate_legacy_request_scopes=migrate) + workflow = factory.create_workflow_from_definition(_yaml(_action(require_approval=True))) + await workflow.run({}) + [request] = (await workflow._runner_context.get_pending_request_info_events()).values() + request.data.workflow_session_id = invalid_scope + + with pytest.raises(DeclarativeWorkflowError, match="Invalid originating workflow session"): + await workflow.run(responses={request.request_id: _scope_pause_response("InvokeMcpTool")}) + + assert handler.call_count == 0 + assert set(await workflow._runner_context.get_pending_request_info_events()) == {request.request_id} + + # ---------- Output handling ------------------------------------------------ @@ -1251,6 +1582,7 @@ async def test_approval_response_rejected_assigns_error(self, mock_state, mock_c server_url="https://mcp.example/api", server_label=None, arguments={}, + workflow_session_id="workflow-a", ), ToolApprovalResponse(approved=False, reason="not authorized"), mock_context, From 97ed49b4dc5dd3805beb92487a4dcd74584dc90b Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Thu, 15 Oct 2026 15:32:12 -0400 Subject: [PATCH 50/58] fix(python): use generator unpacking in continuation regression Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 48ae509d-294c-4bf5-864d-81baecbae204 --- .../packages/declarative/tests/test_invoke_mcp_tool_executor.py | 2 +- 1 file changed, 1 insertion(+), 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 45c2bd77bbf..5ced0c10803 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -1031,7 +1031,7 @@ async def test_repeated_agent_input_retains_its_own_loop_state(self, tmp_path: P for event in (await workflow._runner_context.get_pending_request_info_events()).values() if event.request_id != first.request_id ] - first_scope, second_scope = [invocation.workflow_session_id for invocation in handler.invocations] + first_scope, second_scope = (invocation.workflow_session_id for invocation in handler.invocations) await workflow.run( responses={first.request_id: _scope_pause_response("InvokeAzureAgent", "continue")}, checkpoint_storage=storage, From 16dc9b29e3e44e8e366957c02376c3a7e81b15d1 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Thu, 15 Oct 2026 16:08:03 -0400 Subject: [PATCH 51/58] fix(workflows): validate queued response origins before replay Apply the response batch policy before restoring stored or embedded continuations and include queued responses when validating new submissions. Run built-wheel selection regressions explicitly in an isolated Windows test environment. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 48ae509d-294c-4bf5-864d-81baecbae204 --- .github/workflows/python-tests.yml | 12 ++ ...thon-declarative-request-scope-recovery.md | 3 + .../agent_framework/_workflows/_runner.py | 13 +- .../_workflows/_runner_context.py | 19 +- .../agent_framework/_workflows/_workflow.py | 11 +- .../core/tests/workflow/test_checkpoint.py | 4 + python/packages/declarative/AGENTS.md | 3 + python/packages/declarative/README.md | 4 + .../tests/test_invoke_mcp_tool_executor.py | 198 ++++++++++++++++++ 9 files changed, 264 insertions(+), 3 deletions(-) diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml index 2cbfcfe6daf..4e0577e9306 100644 --- a/.github/workflows/python-tests.yml +++ b/.github/workflows/python-tests.yml @@ -86,6 +86,18 @@ jobs: run: | uv venv --python "${{ matrix.python-version }}" --seed .meta-venv & ".meta-venv\Scripts\python.exe" --version + - name: Test built-wheel selection + shell: pwsh + run: | + uv venv --python "${{ matrix.python-version }}" --seed .meta-test-venv + if ($LASTEXITCODE -ne 0) { + exit $LASTEXITCODE + } + uv pip install --python ".meta-test-venv\Scripts\python.exe" pytest pyyaml + if ($LASTEXITCODE -ne 0) { + exit $LASTEXITCODE + } + & ".meta-test-venv\Scripts\python.exe" -m pytest -o addopts= python/scripts/dependencies/tests/test_meta_package_install.py -q - name: Build release wheelhouse shell: bash run: | diff --git a/docs/decisions/0046-python-declarative-request-scope-recovery.md b/docs/decisions/0046-python-declarative-request-scope-recovery.md index 7e59f804163..b5aa029b392 100644 --- a/docs/decisions/0046-python-declarative-request-scope-recovery.md +++ b/docs/decisions/0046-python-declarative-request-scope-recovery.md @@ -44,6 +44,9 @@ Choose explicit host opt-in through it before continuing. Agent requests also retain their own loop configuration. - Factory-built workflows validate all submitted request origins before consuming responses. Mixed-scope batches are rejected without changing request metadata. +- Checkpoint restores validate queued response origins before restoring shared or + executor state, including embedded child checkpoints. Submitted and already + queued responses share the same batch validation before pending requests are consumed. - Without opt-in, a missing origin produces an actionable recovery error. - With opt-in, a missing origin becomes `legacy-request:`. The distinct namespace never selects diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index f41183cb1c2..a30f1e1c721 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -4,7 +4,7 @@ import logging import warnings from collections import defaultdict -from collections.abc import AsyncGenerator, Sequence +from collections.abc import AsyncGenerator, Callable, Sequence from typing import TYPE_CHECKING, Any, cast from ..exceptions import ( @@ -20,6 +20,7 @@ from ._runner_context import ( RunnerContext, WorkflowMessage, + _get_queued_response_info_events, # pyright: ignore[reportPrivateUsage] ) from ._state import State @@ -83,6 +84,7 @@ def __init__( # Checkpointing related attributes self._previous_checkpoint_id: CheckpointID | None = None + self._response_batch_validator: Callable[[Sequence[WorkflowEvent]], None] | None = None @property def context(self) -> RunnerContext: @@ -326,6 +328,8 @@ async def restore_from_checkpoint( "Please rebuild the original workflow before resuming." ) + self._validate_checkpoint_responses(checkpoint) + # Restore state. Clear first so import_state (which merges) does # not leak stale keys from a prior run on this Workflow instance. # This matters more now that Workflow.run() no longer wipes state @@ -389,6 +393,8 @@ async def restore_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None: ) try: + self._validate_checkpoint_responses(checkpoint) + # Clear first so import_state (which merges) does not leak stale keys from a # prior run on this Workflow instance. self._state.clear() @@ -401,6 +407,11 @@ async def restore_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None: logger.error(f"Failed to restore from checkpoint {checkpoint.checkpoint_id}: {e}") raise WorkflowCheckpointException(f"Failed to restore from checkpoint {checkpoint.checkpoint_id}") from e + def _validate_checkpoint_responses(self, checkpoint: WorkflowCheckpoint) -> None: + requests = _get_queued_response_info_events(checkpoint.messages) + if requests and self._response_batch_validator is not None: + self._response_batch_validator(requests) + async def _save_executor_states(self) -> None: """Populate executor state by calling checkpoint hooks on executors.""" for exec_id, executor in self._executors.items(): diff --git a/python/packages/core/agent_framework/_workflows/_runner_context.py b/python/packages/core/agent_framework/_workflows/_runner_context.py index edbbfea50c9..82b91558e73 100644 --- a/python/packages/core/agent_framework/_workflows/_runner_context.py +++ b/python/packages/core/agent_framework/_workflows/_runner_context.py @@ -5,7 +5,7 @@ import asyncio import logging import warnings -from collections.abc import Callable +from collections.abc import Callable, Mapping, Sequence from copy import copy from dataclasses import dataclass from enum import Enum @@ -96,6 +96,15 @@ def from_dict(data: dict[str, Any]) -> WorkflowMessage: ) +def _get_queued_response_info_events(messages: Mapping[str, Sequence[WorkflowMessage]]) -> list[WorkflowEvent[Any]]: + return [ + message.original_request_info_event + for batch in messages.values() + for message in batch + if message.type == MessageType.RESPONSE and message.original_request_info_event is not None + ] + + @runtime_checkable class RunnerContext(Protocol): """Protocol for the execution context used by the runner. @@ -128,6 +137,10 @@ async def has_messages(self) -> bool: """ ... + async def get_queued_response_info_events(self) -> list[WorkflowEvent[Any]]: + """Get original requests for queued responses without consuming messages.""" + ... + async def add_event(self, event: WorkflowEvent) -> None: """Add an event to the execution context. @@ -611,6 +624,10 @@ async def get_pending_request_info_events(self) -> dict[str, WorkflowEvent[Any]] """ return dict(self._pending_request_info_events) + async def get_queued_response_info_events(self) -> list[WorkflowEvent[Any]]: + """Get original requests for queued responses without consuming messages.""" + return _get_queued_response_info_events(self._messages) + async def cancel_request_info_events(self, request_ids: set[str]) -> dict[str, WorkflowEvent[Any]]: """Remove and return pending request_info events selected for cancellation.""" return { diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 595f4c6fdc0..6b791bdb0b6 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -414,6 +414,7 @@ def __init__( # 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 + self._runner._response_batch_validator = self._validate_response_batch # pyright: ignore[reportPrivateUsage] self._initial_state_initializer: Callable[[State], None] | None = None # Run-scoped pause checkpoint bookkeeping (owned by Workflow, not callers). @@ -1145,6 +1146,10 @@ async def _restore_and_send_responses( await self._runner.restore_from_checkpoint(checkpoint_id, checkpoint_storage) await self._send_responses_internal(responses) + def _validate_response_batch(self, requests: Sequence[WorkflowEvent]) -> None: + if self._response_batch_validator is not None: + self._response_batch_validator(requests) + async def _validate_responses_internal(self, responses: Mapping[str, Any]) -> dict[str, Any]: """Validate responses before a run can consume pending requests.""" from ._agent_executor import _validate_computer_tool_result # pyright: ignore[reportPrivateUsage] @@ -1164,7 +1169,11 @@ async def _validate_responses_internal(self, responses: Mapping[str, Any]) -> di _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]) + queued_requests = await self._runner.context.get_queued_response_info_events() + self._response_batch_validator([ + *queued_requests, + *(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/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index 6e13a414f8b..642dde75596 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -654,8 +654,12 @@ def _build_workflow() -> Any: # Resuming from the response-entry checkpoint replays the continuation and reproduces the output. replay_workflow = _build_workflow() + validated_requests: list[list[Any]] = [] + replay_workflow._response_batch_validator = lambda requests: validated_requests.append(list(requests)) replay_outputs = (await replay_workflow.run(checkpoint_id=response_entry.checkpoint_id)).get_outputs() assert replay_outputs == ["approve->yes"] + assert [[event.request_id for event in batch] for batch in validated_requests] == [[request_id]] + assert validated_requests[0][0].data == "approve" async def test_memory_checkpoint_storage_roundtrip_json_native_types(): diff --git a/python/packages/declarative/AGENTS.md b/python/packages/declarative/AGENTS.md index 4d9f8c5bc9d..dd1ea5e9c9f 100644 --- a/python/packages/declarative/AGENTS.md +++ b/python/packages/declarative/AGENTS.md @@ -95,6 +95,9 @@ instructions. Hosts can opt in through receives a stable `legacy-request:` scope. Migration never infers an origin from current state or recovers the old MCP session. Resume requests separately, preserve their framework IDs, and checkpoint the migrated continuation. +The same validation and migration run before restoring queued checkpoint responses, +including embedded child checkpoints. Newly submitted responses are validated +together with already queued responses before any pending request is consumed. Rejected mixed batches do not mutate request origins. Malformed stored origins are rejected even with migration enabled. Legacy agent loop configuration is read from the old checkpoint slot; newly emitted requests retain their own snapshot. diff --git a/python/packages/declarative/README.md b/python/packages/declarative/README.md index 87be1f421be..770341aac41 100644 --- a/python/packages/declarative/README.md +++ b/python/packages/declarative/README.md @@ -123,6 +123,10 @@ recover the old MCP session or its server-side history. Preserve request IDs, resume legacy requests individually, and retain the resulting checkpoint. Newly emitted requests carry the migrated scope, so subsequent restores no longer require the opt-in for those requests. Any remaining unmigrated requests still do. +The same policy applies when a checkpoint already contains queued responses: +restore validates and migrates their original requests before restoring workflow +state or replaying the continuation. Responses submitted alongside queued responses +must belong to the same scope; mixed batches leave pending requests unconsumed. Invalid stored scopes are rejected even with migration enabled. Header verification and any required MCP reapproval remain unchanged. Hosts that require the old server-side session must restart the workflow instead of migrating it. 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 5ced0c10803..85d6097a869 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -1226,6 +1226,204 @@ async def test_migration_does_not_replace_invalid_scopes(self, invalid_scope: An assert handler.call_count == 0 assert set(await workflow._runner_context.get_pending_request_info_events()) == {request.request_id} + @pytest.mark.parametrize( + "pause_kind", ["Question", "RequestExternalInput", "InvokeFunctionTool", "InvokeAzureAgent", "InvokeMcpTool"] + ) + @pytest.mark.parametrize( + ("restore_kind", "streaming"), + [("memory", False), ("memory", True), ("file", False), ("file", True), ("embedded", False)], + ) + async def test_legacy_queued_response_migration_precedes_replay( + self, tmp_path: Path, pause_kind: str, restore_kind: str, streaming: bool + ) -> None: + from agent_framework import InMemoryCheckpointStorage + from agent_framework._workflows._runner_context import MessageType + + storage = _scope_checkpoint_storage(tmp_path) if restore_kind == "file" else InMemoryCheckpointStorage() + handler = StubMcpHandler(_ok()) + agent = MagicMock(run=AsyncMock(return_value="Need input")) + function = MagicMock(return_value="done") + factory = WorkflowFactory( + checkpoint_storage=storage, mcp_tool_handler=handler, agents={"TestAgent": agent} + ).register_tool("echo", function) + definition = { + "name": "queued_legacy_response", + "actions": [ + {**_action(tool_name="before"), "id": "before"}, + {"kind": "SetValue", "path": "Local.userInput", "value": ""}, + {**_scope_pause_action(pause_kind), "id": "pause"}, + {**_action(tool_name="after"), "id": "after"}, + {**_scope_pause_action("Question"), "id": "followup"}, + ], + } + workflow = factory.create_workflow_from_definition(definition) + await workflow.run({}) + [request] = (await workflow._runner_context.get_pending_request_info_events()).values() + await workflow.run(responses={request.request_id: _scope_pause_response(pause_kind)}) + checkpoints = await storage.list_checkpoints(workflow_name=workflow.name) + [checkpoint] = [ + item + for item in checkpoints + if any(message.type == MessageType.RESPONSE for batch in item.messages.values() for message in batch) + ] + [queued_response] = [message for batch in checkpoint.messages.values() for message in batch] + original = queued_response.original_request_info_event + assert original is not None + assert original.request_id == request.request_id + assert not checkpoint.pending_request_info_events + original_scope = original.data.workflow_session_id + original.data.workflow_session_id = None + await storage.save(checkpoint) + before = handler.call_count + function.reset_mock() + + factory = WorkflowFactory( + checkpoint_storage=storage, + mcp_tool_handler=handler, + agents={"TestAgent": agent}, + migrate_legacy_request_scopes=True, + ).register_tool("echo", function) + workflow = factory.create_workflow_from_definition(definition) + if restore_kind == "embedded": + await workflow._runner.restore_checkpoint(await storage.load(checkpoint.checkpoint_id)) + _ = [event async for event in workflow._runner.run_until_convergence()] + elif streaming: + await workflow.run(checkpoint_id=checkpoint.checkpoint_id, stream=True).get_final_response() + else: + await workflow.run(checkpoint_id=checkpoint.checkpoint_id) + + replayed = handler.invocations[before:] + assert replayed + assert replayed[-1].tool_name == "after" + scope = replayed[-1].workflow_session_id + assert scope == f"legacy-request:{request.request_id}" + assert scope != original_scope + assert all(invocation.workflow_session_id == scope for invocation in replayed) + assert function.call_count == (1 if pause_kind == "InvokeFunctionTool" else 0) + [followup] = (await workflow._runner_context.get_pending_request_info_events()).values() + assert followup.data.workflow_session_id == scope + saved_id = workflow.get_last_checkpoint_id() + assert saved_id is not None + saved = await storage.load(saved_id) + assert saved.pending_request_info_events[followup.request_id].data.workflow_session_id == scope + + @pytest.mark.parametrize("restore_kind", ["file", "embedded"]) + @pytest.mark.parametrize("legacy", [False, True]) + async def test_queued_response_batch_is_validated_before_restoring_state( + self, tmp_path: Path, restore_kind: str, legacy: bool + ) -> None: + from agent_framework import InMemoryCheckpointStorage + from agent_framework.exceptions import WorkflowCheckpointException + + storage = _scope_checkpoint_storage(tmp_path) if restore_kind == "file" else InMemoryCheckpointStorage() + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory( + checkpoint_storage=storage, mcp_tool_handler=handler, migrate_legacy_request_scopes=True + ) + workflow = factory.create_workflow_from_definition(_yaml(_action(require_approval=True))) + await workflow.run({}) + await workflow.run({}) + requests = await workflow._runner_context.get_pending_request_info_events() + assert len(requests) == 2 + for request in requests.values(): + if legacy: + request.data.workflow_session_id = None + await workflow._runner_context.send_request_info_response( + request.request_id, _scope_pause_response("InvokeMcpTool") + ) + await workflow._runner.create_checkpoint_if_enabled() + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + checkpoint = await storage.load(checkpoint_id) + assert not checkpoint.pending_request_info_events + + workflow = factory.create_workflow_from_definition(_yaml(_action(require_approval=True))) + workflow._runner.state.set("preexisting", "keep") + workflow._runner.state.commit() + with pytest.raises(WorkflowCheckpointException) as exc: + if restore_kind == "embedded": + await workflow._runner.restore_checkpoint(checkpoint) + else: + await workflow.run(checkpoint_id=checkpoint_id) + + assert isinstance(exc.value.__cause__, DeclarativeWorkflowError) + assert "different workflow sessions" in str(exc.value.__cause__) + assert workflow._runner.state.get("preexisting") == "keep" + assert handler.call_count == 0 + if legacy: + for batch in checkpoint.messages.values(): + for message in batch: + original = message.original_request_info_event + assert original is not None + assert original.data.workflow_session_id is None + + @pytest.mark.parametrize("restore_kind", ["file", "embedded"]) + async def test_queued_and_new_responses_share_batch_validation(self, tmp_path: Path, restore_kind: str) -> None: + from agent_framework import InMemoryCheckpointStorage + + storage = _scope_checkpoint_storage(tmp_path) if restore_kind == "file" else InMemoryCheckpointStorage() + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(checkpoint_storage=storage, mcp_tool_handler=handler) + definition = _yaml(_action(require_approval=True)) + workflow = factory.create_workflow_from_definition(definition) + await workflow.run({}) + await workflow.run({}) + first_id, second_id = (await workflow._runner_context.get_pending_request_info_events()).keys() + await workflow._runner_context.send_request_info_response(first_id, _scope_pause_response("InvokeMcpTool")) + await workflow._runner.create_checkpoint_if_enabled() + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + + workflow = factory.create_workflow_from_definition(definition) + with pytest.raises(DeclarativeWorkflowError, match="different workflow sessions"): + if restore_kind == "embedded": + await workflow._runner.restore_checkpoint(await storage.load(checkpoint_id)) + await workflow.run(responses={second_id: _scope_pause_response("InvokeMcpTool")}) + else: + await workflow.run( + checkpoint_id=checkpoint_id, responses={second_id: _scope_pause_response("InvokeMcpTool")} + ) + + assert handler.call_count == 0 + assert set(await workflow._runner_context.get_pending_request_info_events()) == {second_id} + + @pytest.mark.parametrize("restore_kind", ["file", "embedded"]) + async def test_legacy_queued_response_requires_opt_in_before_restore( + self, tmp_path: Path, restore_kind: str + ) -> None: + from agent_framework import InMemoryCheckpointStorage + from agent_framework.exceptions import WorkflowCheckpointException + + storage = _scope_checkpoint_storage(tmp_path) if restore_kind == "file" else InMemoryCheckpointStorage() + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(checkpoint_storage=storage, mcp_tool_handler=handler) + definition = _yaml(_action(require_approval=True)) + workflow = factory.create_workflow_from_definition(definition) + await workflow.run({}) + [request] = (await workflow._runner_context.get_pending_request_info_events()).values() + request.data.workflow_session_id = None + await workflow._runner_context.send_request_info_response( + request.request_id, _scope_pause_response("InvokeMcpTool") + ) + await workflow._runner.create_checkpoint_if_enabled() + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + checkpoint = await storage.load(checkpoint_id) + + workflow = factory.create_workflow_from_definition(definition) + workflow._runner.state.set("preexisting", "keep") + workflow._runner.state.commit() + with pytest.raises(WorkflowCheckpointException) as exc: + if restore_kind == "embedded": + await workflow._runner.restore_checkpoint(checkpoint) + else: + await workflow.run(checkpoint_id=checkpoint_id) + + assert isinstance(exc.value.__cause__, DeclarativeWorkflowError) + assert "migrate_legacy_request_scopes" in str(exc.value.__cause__) + assert workflow._runner.state.get("preexisting") == "keep" + assert handler.call_count == 0 + # ---------- Output handling ------------------------------------------------ From 39f9174a8192441f32848d0e757530edff912441 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Thu, 15 Oct 2026 18:19:38 -0400 Subject: [PATCH 52/58] fix(workflows): preflight nested checkpoints before restoring state Validate each embedded workflow with its own policy on a detached candidate before replacing parent, child, or executor state. Share snapshot validation between preflight and restore, and retain migration metadata only after the complete tree passes validation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 48ae509d-294c-4bf5-864d-81baecbae204 --- ...thon-declarative-request-scope-recovery.md | 4 + python/packages/core/AGENTS.md | 3 + .../agent_framework/_workflows/_executor.py | 4 + .../agent_framework/_workflows/_runner.py | 52 +++--- .../_workflows/_workflow_executor.py | 6 + python/packages/declarative/AGENTS.md | 7 +- python/packages/declarative/README.md | 8 +- .../tests/test_invoke_mcp_tool_executor.py | 148 ++++++++++++++++++ 8 files changed, 203 insertions(+), 29 deletions(-) diff --git a/docs/decisions/0046-python-declarative-request-scope-recovery.md b/docs/decisions/0046-python-declarative-request-scope-recovery.md index b5aa029b392..3eb54cad62f 100644 --- a/docs/decisions/0046-python-declarative-request-scope-recovery.md +++ b/docs/decisions/0046-python-declarative-request-scope-recovery.md @@ -47,6 +47,10 @@ Choose explicit host opt-in through - Checkpoint restores validate queued response origins before restoring shared or executor state, including embedded child checkpoints. Submitted and already queued responses share the same batch validation before pending requests are consumed. +- Nested restoration preflights a detached copy of the complete embedded tree + with each child's own policy before applying live state or queues. Migration + metadata is confined to that candidate; a rejected child leaves earlier + siblings, parent executors, and the supplied checkpoint unchanged. - Without opt-in, a missing origin produces an actionable recovery error. - With opt-in, a missing origin becomes `legacy-request:`. The distinct namespace never selects diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index e1a9a90db2f..4fb0f649a33 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -324,6 +324,9 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID by that continuation, can atomically restore a supplied checkpoint before cancellation, and returns the resulting `WorkflowRunResult`. Cancelling a computer request also cancels the other pending requests in that agent's batch, retains already resolved sibling evidence in terminal output, and resets its session before new input. +- **Nested checkpoint restoration** - Stored and embedded restores preflight a detached checkpoint tree before + replacing live state or queues. `WorkflowExecutor` delegates its private validation hook to its child's runner, + so each child uses its own response-origin policy and migration metadata stays in the candidate snapshot. - **`WorkflowBuilder`** - Fluent API for building workflows, including explicit `output_from` / `intermediate_output_from` selection for caller-facing emissions. `output_from` is an allow-list for **Workflow Output**; unselected executor payloads are hidden unless diff --git a/python/packages/core/agent_framework/_workflows/_executor.py b/python/packages/core/agent_framework/_workflows/_executor.py index 7c1a6ffd524..25d3d86f5c7 100644 --- a/python/packages/core/agent_framework/_workflows/_executor.py +++ b/python/packages/core/agent_framework/_workflows/_executor.py @@ -533,6 +533,10 @@ async def _cancel_pending_request(self, request_id: str, ctx: WorkflowContext[An """Release executor-owned state for a cancelled workflow request.""" del request_id, ctx + def _validate_checkpoint_restore(self, state: dict[str, Any]) -> None: + """Validate a checkpoint snapshot without changing live executor state.""" + del state + async def on_checkpoint_save(self) -> dict[str, Any]: """Hook called when the workflow is being saved to a checkpoint. diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index a30f1e1c721..160642ab312 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. import asyncio +import copy import logging import warnings from collections import defaultdict @@ -321,14 +322,8 @@ async def restore_from_checkpoint( logger.error(f"Checkpoint {checkpoint_id} not found") raise WorkflowCheckpointException(f"Checkpoint {checkpoint_id} not found") - # Validate the loaded checkpoint against the workflow - if self._graph_signature_hash != checkpoint.graph_signature_hash: - raise WorkflowCheckpointException( - "Workflow graph has changed since the checkpoint was created. " - "Please rebuild the original workflow before resuming." - ) - - self._validate_checkpoint_responses(checkpoint) + checkpoint = copy.deepcopy(checkpoint) + self._validate_checkpoint_tree(checkpoint) # Restore state. Clear first so import_state (which merges) does # not leak stale keys from a prior run on this Workflow instance. @@ -386,14 +381,9 @@ async def restore_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None: WorkflowCheckpointException: If the checkpoint's graph signature does not match this runner's workflow, or if restoration otherwise fails. """ - if self._graph_signature_hash != checkpoint.graph_signature_hash: - raise WorkflowCheckpointException( - "Workflow graph has changed since the checkpoint was created. " - "Please rebuild the original workflow before resuming." - ) - try: - self._validate_checkpoint_responses(checkpoint) + checkpoint = copy.deepcopy(checkpoint) + self._validate_checkpoint_tree(checkpoint) # Clear first so import_state (which merges) does not leak stale keys from a # prior run on this Workflow instance. @@ -403,14 +393,24 @@ async def restore_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None: self._restore_edge_runner_states() await self._ctx.apply_checkpoint(checkpoint) self._mark_resumed(checkpoint) + except WorkflowCheckpointException: + raise except Exception as e: logger.error(f"Failed to restore from checkpoint {checkpoint.checkpoint_id}: {e}") raise WorkflowCheckpointException(f"Failed to restore from checkpoint {checkpoint.checkpoint_id}") from e - def _validate_checkpoint_responses(self, checkpoint: WorkflowCheckpoint) -> None: + def _validate_checkpoint_tree(self, checkpoint: WorkflowCheckpoint) -> None: + if self._graph_signature_hash != checkpoint.graph_signature_hash: + raise WorkflowCheckpointException( + "Workflow graph has changed since the checkpoint was created. " + "Please rebuild the original workflow before resuming." + ) requests = _get_queued_response_info_events(checkpoint.messages) if requests and self._response_batch_validator is not None: self._response_batch_validator(requests) + if EXECUTOR_STATE_KEY in checkpoint.state: + for _, executor, state in self._get_executor_states(checkpoint.state[EXECUTOR_STATE_KEY]): + executor._validate_checkpoint_restore(state) # pyright: ignore[reportPrivateUsage] async def _save_executor_states(self) -> None: """Populate executor state by calling checkpoint hooks on executors.""" @@ -431,14 +431,20 @@ async def _save_executor_states(self) -> None: async def _restore_executor_states(self) -> None: """Restore executor state by calling restore hooks on executors.""" - has_executor_states = self._state.has(EXECUTOR_STATE_KEY) - if not has_executor_states: + if not self._state.has(EXECUTOR_STATE_KEY): return - executor_states = self._state.get(EXECUTOR_STATE_KEY) + for executor_id, executor, state in self._get_executor_states(self._state.get(EXECUTOR_STATE_KEY)): + try: + await executor.on_checkpoint_restore(state) + except Exception as ex: # pragma: no cover - defensive + raise WorkflowCheckpointException(f"Executor {executor_id} on_checkpoint_restore failed") from ex + + def _get_executor_states(self, executor_states: Any) -> list[tuple[str, Executor, dict[str, Any]]]: if not isinstance(executor_states, dict): raise WorkflowCheckpointException("Executor states in shared state is not a dictionary. Unable to restore.") + snapshots: list[tuple[str, Executor, dict[str, Any]]] = [] for executor_id, state in executor_states.items(): # pyright: ignore[reportUnknownVariableType] if not isinstance(executor_id, str): raise WorkflowCheckpointException("Executor ID in executor states is not a string. Unable to restore.") @@ -450,12 +456,8 @@ async def _restore_executor_states(self) -> None: executor = self._executors.get(executor_id) if not executor: raise WorkflowCheckpointException(f"Executor {executor_id} not found during state restoration.") - - # Try the updated behavior only if backward compatibility did not restore - try: - await executor.on_checkpoint_restore(state) # pyright: ignore[reportUnknownArgumentType] - except Exception as ex: # pragma: no cover - defensive - raise WorkflowCheckpointException(f"Executor {executor_id} on_checkpoint_restore failed") from ex + snapshots.append((executor_id, executor, cast(dict[str, Any], state))) + return snapshots def _save_edge_runner_states(self) -> None: """Store edge runner delivery state, such as partially filled fan-in buffers, in state. diff --git a/python/packages/core/agent_framework/_workflows/_workflow_executor.py b/python/packages/core/agent_framework/_workflows/_workflow_executor.py index 7313e91f322..2651e5b41d3 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_executor.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_executor.py @@ -590,6 +590,12 @@ async def on_checkpoint_save(self) -> dict[str, Any]: "sub_workflow_checkpoint": await self.workflow._runner.build_checkpoint(), # pyright: ignore[reportPrivateUsage] } + @override + def _validate_checkpoint_restore(self, state: dict[str, Any]) -> None: + sub_workflow_checkpoint = state.get("sub_workflow_checkpoint") + if sub_workflow_checkpoint is not None: + self.workflow._runner._validate_checkpoint_tree(sub_workflow_checkpoint) # pyright: ignore[reportPrivateUsage] + @override async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: """Restore the WorkflowExecutor state from a checkpoint snapshot.""" diff --git a/python/packages/declarative/AGENTS.md b/python/packages/declarative/AGENTS.md index dd1ea5e9c9f..34a8ed0050d 100644 --- a/python/packages/declarative/AGENTS.md +++ b/python/packages/declarative/AGENTS.md @@ -96,8 +96,11 @@ receives a stable `legacy-request:` scope. Migration never infers an origin from current state or recovers the old MCP session. Resume requests separately, preserve their framework IDs, and checkpoint the migrated continuation. The same validation and migration run before restoring queued checkpoint responses, -including embedded child checkpoints. Newly submitted responses are validated -together with already queued responses before any pending request is consumed. +including embedded child checkpoints. The complete nested tree is validated on a +detached snapshot before any parent or child state, executor state, or queue is +replaced; a rejected descendant does not partially restore earlier siblings. +Newly submitted responses are validated together with already queued responses +before any pending request is consumed. Rejected mixed batches do not mutate request origins. Malformed stored origins are rejected even with migration enabled. Legacy agent loop configuration is read from the old checkpoint slot; newly emitted requests retain their own snapshot. diff --git a/python/packages/declarative/README.md b/python/packages/declarative/README.md index 770341aac41..428e92755eb 100644 --- a/python/packages/declarative/README.md +++ b/python/packages/declarative/README.md @@ -125,8 +125,12 @@ emitted requests carry the migrated scope, so subsequent restores no longer require the opt-in for those requests. Any remaining unmigrated requests still do. The same policy applies when a checkpoint already contains queued responses: restore validates and migrates their original requests before restoring workflow -state or replaying the continuation. Responses submitted alongside queued responses -must belong to the same scope; mixed batches leave pending requests unconsumed. +state or replaying the continuation. Nested workflows preflight the entire embedded +checkpoint tree using each child's own policy before replacing any live state, +executor state, or queues. Preflight rejection leaves the running workflows and supplied +checkpoint unchanged, including staged migration metadata. +Responses submitted alongside queued responses must belong to the same scope; +mixed batches leave pending requests unconsumed. Invalid stored scopes are rejected even with migration enabled. Header verification and any required MCP reapproval remain unchanged. Hosts that require the old server-side session must restart the workflow instead of migrating it. 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 85d6097a869..5daa2c4cb0a 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -1424,6 +1424,154 @@ async def test_legacy_queued_response_requires_opt_in_before_restore( assert workflow._runner.state.get("preexisting") == "keep" assert handler.call_count == 0 + @pytest.mark.parametrize("restore_kind", ["file", "embedded"]) + @pytest.mark.parametrize("depth", [1, 2]) + @pytest.mark.parametrize("mode", ["missing", "mixed", "valid", "pending"]) + async def test_nested_checkpoint_preflight_preserves_live_state( + self, tmp_path: Path, restore_kind: str, depth: int, mode: str + ) -> None: + from copy import deepcopy + + from agent_framework import ( + Executor, + Workflow, + WorkflowBuilder, + WorkflowCheckpoint, + WorkflowContext, + WorkflowExecutor, + handler, + ) + from agent_framework._workflows._checkpoint_encoding import encode_checkpoint_value + from agent_framework._workflows._runner_context import WorkflowMessage + from agent_framework.exceptions import WorkflowCheckpointException + + class RestoreRecorder(Executor): + def __init__(self, executor_id: str) -> None: + super().__init__(id=executor_id) + self.value = "snapshot" + self.restore_calls = 0 + + @handler + async def forward(self, message: str, ctx: WorkflowContext[str]) -> None: + await ctx.send_message(message) + + async def on_checkpoint_save(self) -> dict[str, Any]: + return {"value": self.value} + + async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: + self.restore_calls += 1 + self.value = state["value"] + + storage = _scope_checkpoint_storage(tmp_path) + mcp_handler = StubMcpHandler(_ok()) + leaves: list[Workflow] = [] + for index in range(2): + factory = WorkflowFactory( + mcp_tool_handler=mcp_handler, + migrate_legacy_request_scopes=index == 0 or mode in {"mixed", "valid"}, + ) + leaf = factory.create_workflow_from_definition(_yaml(_action(require_approval=True))) + await leaf.run({}) + if index == 1 and mode == "mixed": + await leaf.run({}) + requests = await leaf._runner_context.get_pending_request_info_events() + for request in requests.values(): + if mode != "mixed" or index == 0: + request.data.workflow_session_id = None + if index != 1 or mode != "pending": + await leaf._runner_context.send_request_info_response( + request.request_id, _scope_pause_response("InvokeMcpTool") + ) + leaves.append(leaf) + + workflows = list(leaves) + recorders: list[RestoreRecorder] = [] + root = leaves[1] + for level in range(depth): + recorder = RestoreRecorder(f"recorder_{level}") + builder = WorkflowBuilder(start_executor=recorder, checkpoint_storage=storage) + nested = WorkflowExecutor(root, id=f"nested_{level}") + if level == 0: + good = WorkflowExecutor(leaves[0], id="good") + builder.add_fan_out_edges(recorder, [good, nested]) + else: + builder.add_edge(recorder, nested) + root = builder.build() + workflows.append(root) + recorders.append(recorder) + for workflow in workflows: + workflow._runner.state.set("restore_marker", "snapshot") + workflow._runner.state.commit() + checkpoint = await root._runner.build_checkpoint() + checkpoint_before = encode_checkpoint_value(checkpoint.to_dict()) + await storage.save(checkpoint) + + for workflow in workflows: + workflow._runner.state.set("restore_marker", "live") + workflow._runner.state.commit() + workflow._runner.state.set("uncommitted", "keep") + await workflow._runner_context.send_message( + WorkflowMessage(data="live", source_id="live", target_id=next(iter(workflow.executors))) + ) + for recorder in recorders: + recorder.value = "live" + + async def snapshots() -> list[WorkflowCheckpoint]: + return [ + deepcopy( + await workflow._runner_context.build_checkpoint( + workflow.name, workflow.graph_signature_hash, workflow._runner.state, None, 0 + ) + ) + for workflow in workflows + ] + + before = await snapshots() + + async def restore() -> None: + if restore_kind == "embedded": + await root._runner.restore_checkpoint(checkpoint) + else: + await root._runner.restore_from_checkpoint(checkpoint.checkpoint_id) + + if mode in {"missing", "mixed"}: + with pytest.raises(WorkflowCheckpointException) as exc: + await restore() + cause: BaseException = exc.value + while cause.__cause__ is not None: + cause = cause.__cause__ + assert isinstance(cause, DeclarativeWorkflowError) + expected = "migrate_legacy_request_scopes" if mode == "missing" else "different workflow sessions" + assert expected in str(cause) + after = await snapshots() + for old, current in zip(before, after, strict=True): + assert encode_checkpoint_value(current.state) == encode_checkpoint_value(old.state) + assert encode_checkpoint_value(current.messages) == encode_checkpoint_value(old.messages) + assert encode_checkpoint_value(current.pending_request_info_events) == encode_checkpoint_value( + old.pending_request_info_events + ) + assert all(workflow._runner.state.get("uncommitted") == "keep" for workflow in workflows) + assert all(recorder.value == "live" and recorder.restore_calls == 0 for recorder in recorders) + else: + await restore() + assert all(workflow._runner.state.get("restore_marker") == "snapshot" for workflow in workflows) + assert all(recorder.value == "snapshot" and recorder.restore_calls == 1 for recorder in recorders) + assert mcp_handler.call_count == 0 + for index, leaf in enumerate(leaves): + queued = await leaf._runner_context.get_queued_response_info_events() + if mode == "pending" and index == 1: + [pending] = (await leaf._runner_context.get_pending_request_info_events()).values() + assert pending.data.workflow_session_id is None + assert not queued + continue + [request] = queued + scope = f"legacy-request:{request.request_id}" + assert request.data.workflow_session_id == scope + _ = [event async for event in leaf._runner.run_until_convergence()] + assert mcp_handler.invocations[-1].workflow_session_id == scope + assert encode_checkpoint_value(checkpoint.to_dict()) == checkpoint_before + assert mcp_handler.call_count == (2 if mode == "valid" else 1 if mode == "pending" else 0) + # ---------- Output handling ------------------------------------------------ From 27ec8f9946e69852f14928967cf37bd92b50b074 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Fri, 16 Oct 2026 09:15:12 -0400 Subject: [PATCH 53/58] fix(workflows): validate delegated responses before consuming state Complete failed MCP connection cleanup before releasing its cache slot, and preflight nested responses using child pending requests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: ee49ffae-d04c-4453-a2e2-446cf720e160 --- .../agent_framework/_workflows/_executor.py | 10 +- .../agent_framework/_workflows/_runner.py | 16 ++- .../agent_framework/_workflows/_workflow.py | 16 ++- .../_workflows/_workflow_executor.py | 13 +- .../_workflows/_mcp_handler.py | 10 +- .../tests/test_default_mcp_tool_handler.py | 52 +++++++ .../tests/test_invoke_mcp_tool_executor.py | 134 ++++++++++++++++++ 7 files changed, 231 insertions(+), 20 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_executor.py b/python/packages/core/agent_framework/_workflows/_executor.py index 25d3d86f5c7..56373f424fd 100644 --- a/python/packages/core/agent_framework/_workflows/_executor.py +++ b/python/packages/core/agent_framework/_workflows/_executor.py @@ -7,7 +7,7 @@ import inspect import logging import types -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Sequence from typing import Any, TypeVar, overload from ..observability import create_processing_span @@ -533,9 +533,13 @@ async def _cancel_pending_request(self, request_id: str, ctx: WorkflowContext[An """Release executor-owned state for a cancelled workflow request.""" del request_id, ctx - def _validate_checkpoint_restore(self, state: dict[str, Any]) -> None: + async def _validate_response_batch(self, requests: Sequence[WorkflowEvent]) -> None: + """Validate delegated responses before the parent consumes pending requests.""" + del requests + + def _validate_checkpoint_restore(self, state: dict[str, Any], requests: Sequence[WorkflowEvent] = ()) -> None: """Validate a checkpoint snapshot without changing live executor state.""" - del state + del state, requests async def on_checkpoint_save(self) -> dict[str, Any]: """Hook called when the workflow is being saved to a checkpoint. diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index 160642ab312..6fb2ea653ff 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -399,18 +399,28 @@ async def restore_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None: logger.error(f"Failed to restore from checkpoint {checkpoint.checkpoint_id}: {e}") raise WorkflowCheckpointException(f"Failed to restore from checkpoint {checkpoint.checkpoint_id}") from e - def _validate_checkpoint_tree(self, checkpoint: WorkflowCheckpoint) -> None: + def _validate_checkpoint_tree( + self, checkpoint: WorkflowCheckpoint, delegated_requests: Sequence[WorkflowEvent] = () + ) -> None: if self._graph_signature_hash != checkpoint.graph_signature_hash: raise WorkflowCheckpointException( "Workflow graph has changed since the checkpoint was created. " "Please rebuild the original workflow before resuming." ) - requests = _get_queued_response_info_events(checkpoint.messages) + requests = [ + *_get_queued_response_info_events(checkpoint.messages), + *( + checkpoint.pending_request_info_events[request.request_id] + for request in delegated_requests + if request.request_id in checkpoint.pending_request_info_events + ), + ] if requests and self._response_batch_validator is not None: self._response_batch_validator(requests) if EXECUTOR_STATE_KEY in checkpoint.state: for _, executor, state in self._get_executor_states(checkpoint.state[EXECUTOR_STATE_KEY]): - executor._validate_checkpoint_restore(state) # pyright: ignore[reportPrivateUsage] + delegated = [request for request in requests if request.source_executor_id == executor.id] + executor._validate_checkpoint_restore(state, delegated) # pyright: ignore[reportPrivateUsage] async def _save_executor_states(self) -> None: """Populate executor state by calling checkpoint hooks on executors.""" diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 6b791bdb0b6..eece664b463 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -1150,6 +1150,15 @@ def _validate_response_batch(self, requests: Sequence[WorkflowEvent]) -> None: if self._response_batch_validator is not None: self._response_batch_validator(requests) + async def _validate_response_requests(self, requests: Sequence[WorkflowEvent]) -> None: + queued_requests = await self._runner.context.get_queued_response_info_events() + batch = [*queued_requests, *requests] + self._validate_response_batch(batch) + for executor in self.executors.values(): + delegated = [request for request in batch if request.source_executor_id == executor.id] + if delegated: + await executor._validate_response_batch(delegated) # pyright: ignore[reportPrivateUsage] + async def _validate_responses_internal(self, responses: Mapping[str, Any]) -> dict[str, Any]: """Validate responses before a run can consume pending requests.""" from ._agent_executor import _validate_computer_tool_result # pyright: ignore[reportPrivateUsage] @@ -1168,12 +1177,7 @@ 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: - queued_requests = await self._runner.context.get_queued_response_info_events() - self._response_batch_validator([ - *queued_requests, - *(pending_requests[request_id] for request_id in coerced_responses), - ]) + await self._validate_response_requests([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/core/agent_framework/_workflows/_workflow_executor.py b/python/packages/core/agent_framework/_workflows/_workflow_executor.py index 2651e5b41d3..3c122afcfd2 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_executor.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_executor.py @@ -3,7 +3,7 @@ import logging import sys import types -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Any, cast @@ -591,10 +591,17 @@ async def on_checkpoint_save(self) -> dict[str, Any]: } @override - def _validate_checkpoint_restore(self, state: dict[str, Any]) -> None: + async def _validate_response_batch(self, requests: Sequence[WorkflowEvent]) -> None: + pending = await self.workflow._runner_context.get_pending_request_info_events() # pyright: ignore[reportPrivateUsage] + await self.workflow._validate_response_requests([ # pyright: ignore[reportPrivateUsage] + pending[request.request_id] for request in requests if request.request_id in pending + ]) + + @override + def _validate_checkpoint_restore(self, state: dict[str, Any], requests: Sequence[WorkflowEvent] = ()) -> None: sub_workflow_checkpoint = state.get("sub_workflow_checkpoint") if sub_workflow_checkpoint is not None: - self.workflow._runner._validate_checkpoint_tree(sub_workflow_checkpoint) # pyright: ignore[reportPrivateUsage] + self.workflow._runner._validate_checkpoint_tree(sub_workflow_checkpoint, requests) # pyright: ignore[reportPrivateUsage] @override async def on_checkpoint_restore(self, state: dict[str, Any]) -> 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 713f2d12a2c..1784a409ce4 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -906,13 +906,13 @@ def _header_provider(_kwargs: dict[str, Any]) -> dict[str, str]: else None ), ) + cleanup_outcome: asyncio.Future[None] = asyncio.get_running_loop().create_future() try: - if self._client_provider is not None: - await self._close_invocation_entry(failed_entry) - else: - await self._close_entry(failed_entry) + await self._close_invocation_entry(failed_entry, cleanup_outcome) except BaseException as cleanup_exception: - if self._client_provider is not None: + if self._client_provider is not None or ( + cleanup_outcome.done() and cleanup_outcome.exception() is None + ): raise raise _EntryCreationCleanupFailure(creation_exception, cleanup_exception) from cleanup_exception raise 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 e9d5818c6be..5ae6a13be37 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -1387,6 +1387,58 @@ 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_connect_cancellation_retains_slot_until_cleanup(self) -> None: + handler = DefaultMCPToolHandler(cache_max_size=1) + connect_started = asyncio.Event() + close_started = asyncio.Event() + release_close = asyncio.Event() + close_finished = asyncio.Event() + close_tasks: list[asyncio.Task[None]] = [] + + async def connect(_tool: FakeTool) -> None: + connect_started.set() + await asyncio.Event().wait() + + async def finish_close() -> None: + close_started.set() + await release_close.wait() + close_finished.set() + + async def close(_tool: FakeTool) -> None: + task = asyncio.create_task(finish_close()) + close_tasks.append(task) + await asyncio.shield(task) + + with _patch_tool(), patch.object(FakeTool, "connect", connect), patch.object(FakeTool, "close", close): + creator = asyncio.create_task(handler.invoke_tool(_invocation(workflow_session_id="first"))) + shutdown: asyncio.Task[None] | None = None + try: + await connect_started.wait() + creator.cancel() + await close_started.wait() + creator.cancel() + for _ in range(10): + await asyncio.sleep(0) + + assert not creator.done() + result = await handler.invoke_tool(_invocation(workflow_session_id="second")) + assert result.is_error + assert len(FakeTool.instances) == 1 + shutdown = asyncio.create_task(handler.aclose()) + await asyncio.sleep(0) + assert not shutdown.done() + assert not close_finished.is_set() + finally: + release_close.set() + with pytest.raises(asyncio.CancelledError): + await creator + await asyncio.gather(*close_tasks) + if shutdown is not None: + await shutdown + else: + await handler.aclose() + @pytest.mark.asyncio async def test_repeated_creator_cancellation_completes_inflight_failure(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 5daa2c4cb0a..e52a94fd954 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -1424,6 +1424,140 @@ async def test_legacy_queued_response_requires_opt_in_before_restore( assert workflow._runner.state.get("preexisting") == "keep" assert handler.call_count == 0 + @pytest.mark.parametrize("restore_kind", ["file", "embedded"]) + @pytest.mark.parametrize("depth", [1, 2]) + @pytest.mark.parametrize("mode", ["missing", "invalid", "migrate", "valid"]) + async def test_parent_queued_child_response_preflight_preserves_state( + self, tmp_path: Path, restore_kind: str, depth: int, mode: str + ) -> None: + from copy import deepcopy + + from agent_framework import WorkflowBuilder, WorkflowExecutor + from agent_framework._workflows._checkpoint_encoding import encode_checkpoint_value + from agent_framework._workflows._const import EXECUTOR_STATE_KEY + from agent_framework.exceptions import WorkflowCheckpointException + + storage = _scope_checkpoint_storage(tmp_path) + mcp_handler = StubMcpHandler(_ok()) + child = WorkflowFactory( + mcp_tool_handler=mcp_handler, migrate_legacy_request_scopes=mode == "migrate" + ).create_workflow_from_definition(_yaml(_action(require_approval=True))) + workflows = [child] + parent = child + for level in range(depth): + wrapper = WorkflowExecutor(parent, id=f"nested_{level}", propagate_request=True) + parent = WorkflowBuilder(start_executor=wrapper, checkpoint_storage=storage).build() + workflows.append(parent) + await parent.run({}) + [request] = (await parent._runner_context.get_pending_request_info_events()).values() + await parent._runner_context.send_request_info_response( + request.request_id, _scope_pause_response("InvokeMcpTool") + ) + checkpoint = deepcopy(await parent._runner.build_checkpoint()) + # The child's pending request, not the parent's copy, owns the response policy. + for batch in checkpoint.messages.values(): + for message in batch: + message.original_request_info_event = deepcopy(message.original_request_info_event) + child_checkpoint = checkpoint + for level in reversed(range(depth)): + child_checkpoint = child_checkpoint.state[EXECUTOR_STATE_KEY][f"nested_{level}"]["sub_workflow_checkpoint"] + original_scope = child_checkpoint.pending_request_info_events[request.request_id].data.workflow_session_id + if mode != "valid": + child_checkpoint.pending_request_info_events[request.request_id].data.workflow_session_id = ( + 42 if mode == "invalid" else None + ) + checkpoint_before = encode_checkpoint_value(checkpoint.to_dict()) + await storage.save(checkpoint) + + for workflow in workflows: + workflow._runner.state.set("restore_marker", "live") + workflow._runner.state.commit() + workflow._runner.state.set("uncommitted", "keep") + + async def snapshots() -> list[Any]: + result: list[Any] = [] + for workflow in workflows: + snapshot = await workflow._runner_context.build_checkpoint( + workflow.name, workflow.graph_signature_hash, workflow._runner.state, None, 0 + ) + result.append( + encode_checkpoint_value({ + "state": snapshot.state, + "messages": snapshot.messages, + "requests": snapshot.pending_request_info_events, + }) + ) + return result + + before = await snapshots() + + async def restore() -> None: + if restore_kind == "embedded": + await parent._runner.restore_checkpoint(checkpoint) + else: + await parent._runner.restore_from_checkpoint(checkpoint.checkpoint_id) + + if mode in {"missing", "invalid"}: + with pytest.raises(WorkflowCheckpointException): + await restore() + assert await snapshots() == before + for workflow in workflows: + assert workflow._runner.state.get("restore_marker") == "live" + assert workflow._runner.state.get("uncommitted") == "keep" + else: + await restore() + _ = [event async for event in parent._runner.run_until_convergence()] + assert mcp_handler.call_count == 1 + assert mcp_handler.invocations[0].workflow_session_id == ( + f"legacy-request:{request.request_id}" if mode == "migrate" else original_scope + ) + assert encode_checkpoint_value(checkpoint.to_dict()) == checkpoint_before + assert mcp_handler.call_count == (1 if mode in {"migrate", "valid"} else 0) + + @pytest.mark.parametrize("depth", [1, 2]) + @pytest.mark.parametrize("mode", ["missing", "invalid", "queued"]) + async def test_parent_child_response_rejection_preserves_pending_request(self, depth: int, mode: str) -> None: + from copy import deepcopy + + from agent_framework import WorkflowBuilder, WorkflowExecutor + + mcp_handler = StubMcpHandler(_ok()) + child = WorkflowFactory(mcp_tool_handler=mcp_handler).create_workflow_from_definition( + _yaml(_action(require_approval=True)) + ) + parent = child + for level in range(depth): + parent = WorkflowBuilder( + start_executor=WorkflowExecutor(parent, id=f"nested_{level}", propagate_request=True) + ).build() + await parent.run({}) + if mode == "queued": + await parent.run({}) + pending = await parent._runner_context.get_pending_request_info_events() + child_pending = await child._runner_context.get_pending_request_info_events() + request_ids = list(pending) + request_id = request_ids[-1] + original_scope = child_pending[request_id].data.workflow_session_id + if mode == "queued": + await child._runner_context.send_request_info_response( + request_ids[0], _scope_pause_response("InvokeMcpTool") + ) + else: + pending[request_id].data = deepcopy(pending[request_id].data) + child_pending[request_id].data.workflow_session_id = None if mode == "missing" else 42 + + with pytest.raises(DeclarativeWorkflowError): + await parent.run(responses={request_id: _scope_pause_response("InvokeMcpTool")}) + + assert set(await parent._runner_context.get_pending_request_info_events()) == set(request_ids) + assert request_id in await child._runner_context.get_pending_request_info_events() + assert mcp_handler.call_count == 0 + if mode != "queued": + child_pending[request_id].data.workflow_session_id = original_scope + await parent.run(responses={request_id: _scope_pause_response("InvokeMcpTool")}) + assert not await parent._runner_context.get_pending_request_info_events() + assert mcp_handler.call_count == 1 + @pytest.mark.parametrize("restore_kind", ["file", "embedded"]) @pytest.mark.parametrize("depth", [1, 2]) @pytest.mark.parametrize("mode", ["missing", "mixed", "valid", "pending"]) From 3e297469c6941272a04775faa262b266c38dc1d1 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Fri, 16 Oct 2026 09:46:55 -0400 Subject: [PATCH 54/58] fix(workflows): preflight delegated response migrations together Validate detached response batches throughout the workflow tree before applying migrations to live pending requests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: ee49ffae-d04c-4453-a2e2-446cf720e160 --- .../agent_framework/_workflows/_executor.py | 9 ++- .../agent_framework/_workflows/_workflow.py | 26 +++++++- .../_workflows/_workflow_executor.py | 12 +++- .../tests/test_invoke_mcp_tool_executor.py | 65 +++++++++++++++++++ 4 files changed, 104 insertions(+), 8 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_executor.py b/python/packages/core/agent_framework/_workflows/_executor.py index 56373f424fd..6a9947a72ba 100644 --- a/python/packages/core/agent_framework/_workflows/_executor.py +++ b/python/packages/core/agent_framework/_workflows/_executor.py @@ -31,6 +31,8 @@ logger = logging.getLogger(__name__) +_ResponseBatchValidation = tuple[Callable[[Sequence[WorkflowEvent]], None], Sequence[WorkflowEvent]] + # region Executor class Executor(RequestInfoMixin, DictConvertible): @@ -533,9 +535,12 @@ async def _cancel_pending_request(self, request_id: str, ctx: WorkflowContext[An """Release executor-owned state for a cancelled workflow request.""" del request_id, ctx - async def _validate_response_batch(self, requests: Sequence[WorkflowEvent]) -> None: - """Validate delegated responses before the parent consumes pending requests.""" + async def _get_response_validation_batches( + self, requests: Sequence[WorkflowEvent] + ) -> list[_ResponseBatchValidation]: + """Collect delegated response policies without changing pending requests.""" del requests + return [] def _validate_checkpoint_restore(self, state: dict[str, Any], requests: Sequence[WorkflowEvent] = ()) -> None: """Validate a checkpoint snapshot without changing live executor state.""" diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index eece664b463..85777174274 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -13,6 +13,7 @@ import weakref from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable, Collection, Mapping, Sequence from contextlib import aclosing +from copy import deepcopy from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Literal, cast, overload @@ -50,7 +51,10 @@ WorkflowRunState, _framework_event_origin, # type: ignore ) -from ._executor import Executor +from ._executor import ( + Executor, + _ResponseBatchValidation, # pyright: ignore[reportPrivateUsage] +) from ._model_utils import DictConvertible from ._runner import RunnerImpl from ._runner_context import RunnerContext, WorkflowMessage @@ -1151,13 +1155,29 @@ def _validate_response_batch(self, requests: Sequence[WorkflowEvent]) -> None: self._response_batch_validator(requests) async def _validate_response_requests(self, requests: Sequence[WorkflowEvent]) -> None: + validations = await self._get_response_validation_batches(requests) + # Preserve shared request data while preflighting the whole tree before live migrations. + detached_batches = deepcopy([batch for _, batch in validations]) + for (validator, _), batch in zip(validations, detached_batches, strict=True): + validator(batch) + for validator, batch in validations: + validator(batch) + + async def _get_response_validation_batches( + self, requests: Sequence[WorkflowEvent] + ) -> list[_ResponseBatchValidation]: queued_requests = await self._runner.context.get_queued_response_info_events() batch = [*queued_requests, *requests] - self._validate_response_batch(batch) + validations: list[_ResponseBatchValidation] = [] + if self._response_batch_validator is not None: + validations.append((self._validate_response_batch, batch)) for executor in self.executors.values(): delegated = [request for request in batch if request.source_executor_id == executor.id] if delegated: - await executor._validate_response_batch(delegated) # pyright: ignore[reportPrivateUsage] + validations.extend( + await executor._get_response_validation_batches(delegated) # pyright: ignore[reportPrivateUsage] + ) + return validations async def _validate_responses_internal(self, responses: Mapping[str, Any]) -> dict[str, Any]: """Validate responses before a run can consume pending requests.""" diff --git a/python/packages/core/agent_framework/_workflows/_workflow_executor.py b/python/packages/core/agent_framework/_workflows/_workflow_executor.py index 3c122afcfd2..780b5b74554 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_executor.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_executor.py @@ -24,7 +24,11 @@ WorkflowRunState, _framework_event_origin, # type: ignore[reportPrivateUsage] ) -from ._executor import Executor, handler +from ._executor import ( + Executor, + _ResponseBatchValidation, # pyright: ignore[reportPrivateUsage] + handler, +) from ._request_info_mixin import response_handler from ._runner_context import WorkflowMessage from ._typing_utils import is_instance_of @@ -591,9 +595,11 @@ async def on_checkpoint_save(self) -> dict[str, Any]: } @override - async def _validate_response_batch(self, requests: Sequence[WorkflowEvent]) -> None: + async def _get_response_validation_batches( + self, requests: Sequence[WorkflowEvent] + ) -> list[_ResponseBatchValidation]: pending = await self.workflow._runner_context.get_pending_request_info_events() # pyright: ignore[reportPrivateUsage] - await self.workflow._validate_response_requests([ # pyright: ignore[reportPrivateUsage] + return await self.workflow._get_response_validation_batches([ # pyright: ignore[reportPrivateUsage] pending[request.request_id] for request in requests if request.request_id in pending ]) 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 e52a94fd954..8044ec88f5a 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -1558,6 +1558,71 @@ async def test_parent_child_response_rejection_preserves_pending_request(self, d assert not await parent._runner_context.get_pending_request_info_events() assert mcp_handler.call_count == 1 + @pytest.mark.parametrize("depth", [1, 2]) + @pytest.mark.parametrize("queued", [False, True]) + @pytest.mark.parametrize("origin", [None, 42]) + async def test_sibling_response_rejection_preserves_unmigrated_origins( + self, depth: int, queued: bool, origin: Any + ) -> None: + from agent_framework import Executor, WorkflowBuilder, WorkflowContext, WorkflowExecutor, handler + + class Start(Executor): + @handler + async def forward(self, message: dict[str, Any], ctx: WorkflowContext[dict[str, Any]]) -> None: + await ctx.send_message(message) + + mcp_handler = StubMcpHandler(_ok()) + leaves = [ + WorkflowFactory( + mcp_tool_handler=mcp_handler, migrate_legacy_request_scopes=index == 0 + ).create_workflow_from_definition(_yaml(_action(require_approval=True))) + for index in range(2) + ] + wrappers: list[WorkflowExecutor] = [] + for index, leaf in enumerate(leaves): + child = leaf + for level in range(depth): + wrapper = WorkflowExecutor(child, id=f"child_{index}_{level}", propagate_request=True) + if level + 1 < depth: + child = WorkflowBuilder(start_executor=wrapper).build() + wrappers.append(wrapper) + start = Start(id="start") + parent = WorkflowBuilder(start_executor=start).add_fan_out_edges(start, wrappers).build() + await parent.run({}) + requests = [ + next(iter((await leaf._runner_context.get_pending_request_info_events()).values())) for leaf in leaves + ] + first, second = requests + second_scope = second.data.workflow_session_id + first.data.workflow_session_id = None + second.data.workflow_session_id = origin + if queued: + await parent._runner_context.send_request_info_response( + first.request_id, _scope_pause_response("InvokeMcpTool") + ) + pending_ids = set(await parent._runner_context.get_pending_request_info_events()) + responses = {request_id: _scope_pause_response("InvokeMcpTool") for request_id in pending_ids} + + with pytest.raises(DeclarativeWorkflowError): + await parent.run(responses=responses) + + assert first.data.workflow_session_id is None + assert second.data.workflow_session_id == origin + assert set(await parent._runner_context.get_pending_request_info_events()) == pending_ids + for leaf, request in zip(leaves, requests, strict=True): + assert set(await leaf._runner_context.get_pending_request_info_events()) == {request.request_id} + assert mcp_handler.call_count == 0 + + second.data.workflow_session_id = second_scope + await parent.run(responses=responses) + assert not await parent._runner_context.get_pending_request_info_events() + assert first.data.workflow_session_id == f"legacy-request:{first.request_id}" + assert {invocation.workflow_session_id for invocation in mcp_handler.invocations} == { + f"legacy-request:{first.request_id}", + second_scope, + } + assert mcp_handler.call_count == 2 + @pytest.mark.parametrize("restore_kind", ["file", "embedded"]) @pytest.mark.parametrize("depth", [1, 2]) @pytest.mark.parametrize("mode", ["missing", "mixed", "valid", "pending"]) From f35ed148a1b29efb0b0ade1c5e9317009d4b52cb Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Fri, 16 Oct 2026 09:57:52 -0400 Subject: [PATCH 55/58] fix(tests): initialize nested workflow wrapper before iteration Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: ee49ffae-d04c-4453-a2e2-446cf720e160 --- .../declarative/tests/test_invoke_mcp_tool_executor.py | 7 +++---- 1 file changed, 3 insertions(+), 4 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 8044ec88f5a..9f1fe99d90a 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -1580,11 +1580,10 @@ async def forward(self, message: dict[str, Any], ctx: WorkflowContext[dict[str, ] wrappers: list[WorkflowExecutor] = [] for index, leaf in enumerate(leaves): - child = leaf - for level in range(depth): + wrapper = WorkflowExecutor(leaf, id=f"child_{index}_0", propagate_request=True) + for level in range(1, depth): + child = WorkflowBuilder(start_executor=wrapper).build() wrapper = WorkflowExecutor(child, id=f"child_{index}_{level}", propagate_request=True) - if level + 1 < depth: - child = WorkflowBuilder(start_executor=wrapper).build() wrappers.append(wrapper) start = Start(id="start") parent = WorkflowBuilder(start_executor=start).add_fan_out_edges(start, wrappers).build() From 9e4d86f4378b09bada1241392e601071b1b2ed89 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Fri, 16 Oct 2026 10:08:33 -0400 Subject: [PATCH 56/58] fix(workflows): preflight queued wrapped child replies Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: ee49ffae-d04c-4453-a2e2-446cf720e160 --- .../agent_framework/_workflows/_runner.py | 24 ++- .../tests/test_invoke_mcp_tool_executor.py | 141 ++++++++++++++++++ 2 files changed, 163 insertions(+), 2 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index 6fb2ea653ff..7703c7f56f3 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -402,6 +402,8 @@ async def restore_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None: def _validate_checkpoint_tree( self, checkpoint: WorkflowCheckpoint, delegated_requests: Sequence[WorkflowEvent] = () ) -> None: + from ._workflow_executor import SubWorkflowResponseMessage + if self._graph_signature_hash != checkpoint.graph_signature_hash: raise WorkflowCheckpointException( "Workflow graph has changed since the checkpoint was created. " @@ -418,9 +420,27 @@ def _validate_checkpoint_tree( if requests and self._response_batch_validator is not None: self._response_batch_validator(requests) if EXECUTOR_STATE_KEY in checkpoint.state: + delegated_by_executor: defaultdict[str, list[WorkflowEvent]] = defaultdict(list) + for request in requests: + delegated_by_executor[request.source_executor_id].append(request) + # Wrapped replies retain child-local source IDs, so use the parent graph to route them. + for batch in checkpoint.messages.values(): + for message in batch: + if not isinstance(message.data, SubWorkflowResponseMessage): + continue + targets = { + edge.target_id + for runner in self._edge_runner_map.get(message.source_id, []) + for edge in runner._edge_group.edges # pyright: ignore[reportPrivateUsage] + if edge.source_id == message.source_id + and (message.target_id is None or edge.target_id == message.target_id) + } + for target_id in targets: + delegated_by_executor[target_id].append(message.data.source_event) for _, executor, state in self._get_executor_states(checkpoint.state[EXECUTOR_STATE_KEY]): - delegated = [request for request in requests if request.source_executor_id == executor.id] - executor._validate_checkpoint_restore(state, delegated) # pyright: ignore[reportPrivateUsage] + executor._validate_checkpoint_restore( # pyright: ignore[reportPrivateUsage] + state, delegated_by_executor[executor.id] + ) async def _save_executor_states(self) -> None: """Populate executor state by calling checkpoint hooks on executors.""" 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 9f1fe99d90a..6e8f81775a2 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -1514,6 +1514,147 @@ async def restore() -> None: assert encode_checkpoint_value(checkpoint.to_dict()) == checkpoint_before assert mcp_handler.call_count == (1 if mode in {"migrate", "valid"} else 0) + @pytest.mark.parametrize("restore_kind", ["file", "embedded"]) + @pytest.mark.parametrize("depth", [1, 2]) + @pytest.mark.parametrize("targeted", [False, True]) + @pytest.mark.parametrize("mode", ["missing", "invalid", "migrate", "valid"]) + async def test_wrapped_child_response_preflight_preserves_state( + self, tmp_path: Path, restore_kind: str, depth: int, targeted: bool, mode: str + ) -> None: + from copy import deepcopy + + from agent_framework import ( + Executor, + SubWorkflowRequestMessage, + WorkflowBuilder, + WorkflowContext, + WorkflowExecutor, + handler, + ) + from agent_framework._workflows._checkpoint_encoding import encode_checkpoint_value + from agent_framework._workflows._const import EXECUTOR_STATE_KEY + from agent_framework._workflows._runner_context import WorkflowMessage + from agent_framework.exceptions import WorkflowCheckpointException + + class RestoreRecorder(Executor): + def __init__(self) -> None: + super().__init__(id="coordinator") + self.value = "snapshot" + self.restore_calls = 0 + + @handler + async def forward(self, message: dict[str, Any], ctx: WorkflowContext[dict[str, Any]]) -> None: + await ctx.send_message(message) + + async def on_checkpoint_save(self) -> dict[str, Any]: + return {"value": self.value} + + async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: + self.restore_calls += 1 + self.value = state["value"] + + storage = _scope_checkpoint_storage(tmp_path) + mcp_handler = StubMcpHandler(_ok()) + leaves = [ + WorkflowFactory( + mcp_tool_handler=mcp_handler, migrate_legacy_request_scopes=index == 1 and mode == "migrate" + ).create_workflow_from_definition(_yaml(_action(require_approval=True))) + for index in range(2) + ] + workflows = list(leaves) + wrappers: list[WorkflowExecutor] = [] + for index, leaf in enumerate(leaves): + child = leaf + for level in range(1, depth): + child = WorkflowBuilder( + start_executor=WorkflowExecutor(child, id=f"nested_{index}_{level}", propagate_request=True) + ).build() + workflows.append(child) + wrappers.append(WorkflowExecutor(child, id=f"child_{index}")) + recorder = RestoreRecorder() + parent = ( + WorkflowBuilder(start_executor=recorder, checkpoint_storage=storage) + .add_fan_out_edges(recorder, wrappers) + .build() + ) + workflows.append(parent) + sibling_scope: str | None = None + for wrapper in wrappers: + await wrapper.workflow.run({}) + [request] = (await wrapper.workflow._runner_context.get_pending_request_info_events()).values() + if wrapper is wrappers[0]: + sibling_scope = request.data.workflow_session_id + reply = SubWorkflowRequestMessage(source_event=deepcopy(request), executor_id=wrapper.id).create_response( + _scope_pause_response("InvokeMcpTool") + ) + await parent._runner_context.send_message( + WorkflowMessage(data=reply, source_id=recorder.id, target_id=wrapper.id if targeted else None) + ) + for workflow in workflows: + workflow._runner.state.set("restore_marker", "snapshot") + workflow._runner.state.commit() + checkpoint = deepcopy(await parent._runner.build_checkpoint()) + child_checkpoint = checkpoint.state[EXECUTOR_STATE_KEY]["child_1"]["sub_workflow_checkpoint"] + for level in reversed(range(1, depth)): + child_checkpoint = child_checkpoint.state[EXECUTOR_STATE_KEY][f"nested_1_{level}"][ + "sub_workflow_checkpoint" + ] + [child_request] = child_checkpoint.pending_request_info_events.values() + original_scope = child_request.data.workflow_session_id + if mode != "valid": + child_request.data.workflow_session_id = 42 if mode == "invalid" else None + checkpoint_before = encode_checkpoint_value(checkpoint.to_dict()) + await storage.save(checkpoint) + + for workflow in workflows: + workflow._runner.state.set("restore_marker", "live") + workflow._runner.state.commit() + workflow._runner.state.set("uncommitted", "keep") + recorder.value = "live" + + async def snapshots() -> list[Any]: + result: list[Any] = [] + for workflow in workflows: + snapshot = await workflow._runner_context.build_checkpoint( + workflow.name, workflow.graph_signature_hash, workflow._runner.state, None, 0 + ) + result.append( + encode_checkpoint_value({ + "state": snapshot.state, + "messages": snapshot.messages, + "requests": snapshot.pending_request_info_events, + }) + ) + return result + + before = await snapshots() + + async def restore() -> None: + if restore_kind == "embedded": + await parent._runner.restore_checkpoint(checkpoint) + else: + await parent._runner.restore_from_checkpoint(checkpoint.checkpoint_id) + + if mode in {"missing", "invalid"}: + with pytest.raises(WorkflowCheckpointException): + await restore() + assert await snapshots() == before + assert recorder.value == "live" + assert recorder.restore_calls == 0 + assert all(workflow._runner.state.get("uncommitted") == "keep" for workflow in workflows) + else: + await restore() + _ = [event async for event in parent._runner.run_until_convergence()] + assert all(workflow._runner.state.get("restore_marker") == "snapshot" for workflow in workflows) + assert recorder.value == "snapshot" + assert recorder.restore_calls == 1 + assert {invocation.workflow_session_id for invocation in mcp_handler.invocations} == { + sibling_scope, + f"legacy-request:{child_request.request_id}" if mode == "migrate" else original_scope, + } + assert encode_checkpoint_value(checkpoint.to_dict()) == checkpoint_before + assert mcp_handler.call_count == (2 if mode in {"migrate", "valid"} else 0) + @pytest.mark.parametrize("depth", [1, 2]) @pytest.mark.parametrize("mode", ["missing", "invalid", "queued"]) async def test_parent_child_response_rejection_preserves_pending_request(self, depth: int, mode: str) -> None: From fde8d779e6d4926be215e8fb5556e0e6e401e212 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Fri, 16 Oct 2026 15:27:47 -0400 Subject: [PATCH 57/58] chore: remove packaging changes from workflow fix Restore package manifests, lockfile, and installation CI to main. Remove package-installation tests and the unreleased Core version claim. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 69fba54c-3d5c-4d7d-909f-8db1f7c467dd --- .github/workflows/python-tests.yml | 29 +---- python/packages/core/pyproject.toml | 4 +- python/packages/declarative/AGENTS.md | 2 - python/packages/declarative/pyproject.toml | 4 +- python/pyproject.toml | 4 +- python/scripts/dependencies/README.md | 7 -- .../tests/test_meta_package_install.py | 106 ------------------ python/uv.lock | 6 +- 8 files changed, 10 insertions(+), 152 deletions(-) delete mode 100644 python/scripts/dependencies/tests/test_meta_package_install.py diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml index 4e0577e9306..cc08ec06fe1 100644 --- a/.github/workflows/python-tests.yml +++ b/.github/workflows/python-tests.yml @@ -86,18 +86,6 @@ jobs: run: | uv venv --python "${{ matrix.python-version }}" --seed .meta-venv & ".meta-venv\Scripts\python.exe" --version - - name: Test built-wheel selection - shell: pwsh - run: | - uv venv --python "${{ matrix.python-version }}" --seed .meta-test-venv - if ($LASTEXITCODE -ne 0) { - exit $LASTEXITCODE - } - uv pip install --python ".meta-test-venv\Scripts\python.exe" pytest pyyaml - if ($LASTEXITCODE -ne 0) { - exit $LASTEXITCODE - } - & ".meta-test-venv\Scripts\python.exe" -m pytest -o addopts= python/scripts/dependencies/tests/test_meta_package_install.py -q - name: Build release wheelhouse shell: bash run: | @@ -116,17 +104,6 @@ jobs: --no-build-logs \ --no-create-gitignore \ --out-dir "${{ github.workspace }}/wheelhouse" - - name: Constrain workspace packages to built wheels - shell: pwsh - run: | - $wheels = @(Get-ChildItem wheelhouse -Filter *.whl) - if ($wheels.Count -eq 0) { - throw "No built workspace wheels were found." - } - $wheels | Sort-Object Name | ForEach-Object { - $name = ($_.Name -split "-")[0] - "$name @ $([System.Uri]::new($_.FullName).AbsoluteUri)" - } | Set-Content wheelhouse-constraints.txt -Encoding utf8 - name: Resolve the end-user install timeout-minutes: 6 shell: pwsh @@ -141,7 +118,6 @@ jobs: --dry-run ` --ignore-installed ` --find-links wheelhouse ` - --constraint wheelhouse-constraints.txt ` --report "pip-resolution-${{ matrix.python-version }}.json" ` "agent-framework==$version" ` azure-identity @@ -162,8 +138,5 @@ jobs: run: | $python = ".meta-venv\Scripts\python.exe" $version = & $python -c "import tomllib; print(tomllib.load(open('python/pyproject.toml', 'rb'))['project']['version'])" - & $python -m pip install --find-links wheelhouse --constraint wheelhouse-constraints.txt "agent-framework==$version" azure-identity - if ($LASTEXITCODE -ne 0) { - exit $LASTEXITCODE - } + & $python -m pip install --find-links wheelhouse "agent-framework==$version" azure-identity & $python -c "import agent_framework; import azure.identity" diff --git a/python/packages/core/pyproject.toml b/python/packages/core/pyproject.toml index 596e9ba82ff..b5132f3b831 100644 --- a/python/packages/core/pyproject.toml +++ b/python/packages/core/pyproject.toml @@ -4,7 +4,7 @@ description = "Microsoft Agent Framework for building AI Agents with Python. Thi authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}] readme = "README.md" requires-python = ">=3.10" -version = "1.22.1" +version = "1.22.0" license-files = ["LICENSE"] urls.homepage = "https://aka.ms/agent-framework" urls.source = "https://github.com/microsoft/agent-framework/tree/main/python" @@ -48,7 +48,7 @@ all = [ "agent-framework-chatkit>=1.0.0b261015,<2", "agent-framework-claude>=1.0.0b261015,<2", "agent-framework-copilotstudio>=1.0.0b261015,<2", - "agent-framework-declarative>=1.3.1,<2", + "agent-framework-declarative>=1.3.0,<2", "agent-framework-devui>=1.0.0b261015,<2", "agent-framework-durabletask>=1.0.0b260922,<2", "agent-framework-foundry>=1.15.0,<2", diff --git a/python/packages/declarative/AGENTS.md b/python/packages/declarative/AGENTS.md index 34a8ed0050d..d6c26bbee8c 100644 --- a/python/packages/declarative/AGENTS.md +++ b/python/packages/declarative/AGENTS.md @@ -2,8 +2,6 @@ YAML/JSON-based declarative agent and workflow definitions. -Requires Core 1.22.1 or later for workflow run initialization and continuation hooks. - ## Main Classes - **`AgentFactory`** - Creates agents from declarative definitions diff --git a/python/packages/declarative/pyproject.toml b/python/packages/declarative/pyproject.toml index c7f4428dc8c..05d9552dcf1 100644 --- a/python/packages/declarative/pyproject.toml +++ b/python/packages/declarative/pyproject.toml @@ -4,7 +4,7 @@ description = "Declarative specification support for Microsoft Agent Framework." authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}] readme = "README.md" requires-python = ">=3.10" -version = "1.3.1" +version = "1.3.0" license-files = ["LICENSE"] urls.homepage = "https://aka.ms/agent-framework" urls.source = "https://github.com/microsoft/agent-framework/tree/main/python" @@ -22,7 +22,7 @@ classifiers = [ "Typing :: Typed", ] dependencies = [ - "agent-framework-core>=1.22.1,<2", + "agent-framework-core>=1.20.0,<2", "httpx>=0.27,<0.29", "httpx2>=2.0.0,<3", "powerfx>=0.0.32,<0.0.35; python_version < '3.14'", diff --git a/python/pyproject.toml b/python/pyproject.toml index b03331dad36..a318c267a89 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -4,7 +4,7 @@ description = "Microsoft Agent Framework for building AI Agents with Python. Thi authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}] readme = "README.md" requires-python = ">=3.10" -version = "1.22.1" +version = "1.22.0" license-files = ["LICENSE"] urls.homepage = "https://aka.ms/agent-framework" urls.source = "https://github.com/microsoft/agent-framework/tree/main/python" @@ -24,7 +24,7 @@ classifiers = [ "Typing :: Typed", ] dependencies = [ - "agent-framework-core[all]==1.22.1", + "agent-framework-core[all]==1.22.0", ] [dependency-groups] diff --git a/python/scripts/dependencies/README.md b/python/scripts/dependencies/README.md index cb3a8684d41..d512c9b97e7 100644 --- a/python/scripts/dependencies/README.md +++ b/python/scripts/dependencies/README.md @@ -8,13 +8,6 @@ This folder contains the Python workspace tooling for dependency maintenance: Run the commands below from the `python/` directory. -The meta-package install CI constrains workspace distributions to the wheels it -has just built. `--find-links` alone can prefer a registry wheel with the same -version but different dependency metadata. Both the resolution probe and actual -installation use the generated constraints; third-party dependencies still -resolve from the package index. `tests/test_meta_package_install.py` exercises -both paths against an isolated local index. - ## Files in this folder - `validate_dependency_bounds.py` diff --git a/python/scripts/dependencies/tests/test_meta_package_install.py b/python/scripts/dependencies/tests/test_meta_package_install.py deleted file mode 100644 index 13b1a6ab5be..00000000000 --- a/python/scripts/dependencies/tests/test_meta_package_install.py +++ /dev/null @@ -1,106 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -import json -import re -import subprocess -import sys -from functools import partial -from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer -from pathlib import Path -from threading import Thread -from zipfile import ZipFile - -import pytest -import yaml - - -def _write_wheel(directory: Path, name: str, version: str, dependencies: list[str]) -> Path: - directory.mkdir(parents=True, exist_ok=True) - normalized_name = name.replace("-", "_") - wheel = directory / f"{normalized_name}-{version}-py3-none-any.whl" - metadata = f"Metadata-Version: 2.1\nName: {name}\nVersion: {version}\n" - metadata += "".join(f"Requires-Dist: {dependency}\n" for dependency in dependencies) - dist_info = f"{normalized_name}-{version}.dist-info" - with ZipFile(wheel, "w") as archive: - archive.writestr(f"{dist_info}/METADATA", metadata) - archive.writestr( - f"{dist_info}/WHEEL", - "Wheel-Version: 1.0\nGenerator: test\nRoot-Is-Purelib: true\nTag: py3-none-any\n", - ) - archive.writestr(f"{dist_info}/RECORD", "") - return wheel - - -@pytest.mark.parametrize("step_name", ["Resolve the end-user install", "Install and import the meta package"]) -def test_meta_install_uses_built_wheels_instead_of_same_version_registry_metadata( - tmp_path: Path, step_name: str -) -> None: - workspace = tmp_path / "workspace with spaces" - wheelhouse = workspace / "wheelhouse" - _write_wheel(wheelhouse, "agent-framework", "1.0.1", ["agent-framework-core==1.0.1"]) - _write_wheel(wheelhouse, "agent-framework-core", "1.0.1", ["agent-framework-foundry-hosting==1.0.0b1", "mcp>=2,<3"]) - hosting = _write_wheel(wheelhouse, "agent-framework-foundry-hosting", "1.0.0b1", ["mcp>=2,<3"]) - registry = tmp_path / "registry" - registry_hosting = registry / "simple" / "agent-framework-foundry-hosting" - published = _write_wheel(registry_hosting, "agent-framework-foundry-hosting", "1.0.0b1", ["mcp<2"]) - (registry_hosting / "index.html").write_text(f'{published.name}', encoding="utf-8") - registry_mcp = registry / "simple" / "mcp" - mcp = _write_wheel(registry_mcp, "mcp", "2.2.0", []) - (registry_mcp / "index.html").write_text(f'{mcp.name}', encoding="utf-8") - - workflow_path = Path(__file__).resolve().parents[4] / ".github" / "workflows" / "python-tests.yml" - steps = yaml.safe_load(workflow_path.read_text(encoding="utf-8"))["jobs"]["python-meta-package-install"]["steps"] - for step in steps: - if step.get("name") == "Constrain workspace packages to built wheels": - subprocess.run( - ["pwsh", "-NoProfile", "-NonInteractive", "-Command", step["run"]], - cwd=workspace, - check=True, - capture_output=True, - text=True, - ) - - install_script = next(step["run"] for step in steps if step.get("name") == step_name) - selection_arguments = [ - value - for flag in ("--find-links", "--constraint") - for argument in re.findall(rf"{flag}\s+(\S+)", install_script) - for value in (flag, argument) - ] - report = workspace / "report.json" - server = ThreadingHTTPServer(("127.0.0.1", 0), partial(SimpleHTTPRequestHandler, directory=str(registry))) - server_thread = Thread(target=server.serve_forever, daemon=True) - server_thread.start() - try: - command = [ - sys.executable, - "-m", - "pip", - "--isolated", - "install", - "--disable-pip-version-check", - "--no-cache-dir", - "--ignore-installed", - "--index-url", - f"http://127.0.0.1:{server.server_port}/simple", - "--report", - str(report), - *selection_arguments, - "agent-framework==1.0.1", - ] - if step_name == "Resolve the end-user install": - command.append("--dry-run") - else: - command.extend(["--target", str(workspace / "installed")]) - result = subprocess.run(command, cwd=workspace, capture_output=True, text=True, timeout=30) - finally: - server.shutdown() - server.server_close() - server_thread.join() - - assert result.returncode == 0, result.stdout + result.stderr - installed = json.loads(report.read_text(encoding="utf-8"))["install"] - selected_hosting = next(item for item in installed if item["metadata"]["name"] == "agent-framework-foundry-hosting") - assert selected_hosting["download_info"]["url"] == hosting.as_uri() - selected_mcp = next(item for item in installed if item["metadata"]["name"] == "mcp") - assert selected_mcp["download_info"]["url"].startswith("http://127.0.0.1:") diff --git a/python/uv.lock b/python/uv.lock index 012623e9eca..6571372d261 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -115,7 +115,7 @@ wheels = [ [[package]] name = "agent-framework" -version = "1.22.1" +version = "1.22.0" source = { virtual = "." } dependencies = [ { name = "agent-framework-core", extra = ["all"] }, @@ -441,7 +441,7 @@ requires-dist = [ [[package]] name = "agent-framework-core" -version = "1.22.1" +version = "1.22.0" source = { editable = "packages/core" } dependencies = [ { name = "msgspec" }, @@ -549,7 +549,7 @@ dev = [ [[package]] name = "agent-framework-declarative" -version = "1.3.1" +version = "1.3.0" source = { editable = "packages/declarative" } dependencies = [ { name = "agent-framework-core" }, From 2c52eeda221b12286fabccd44a200b7c437b4277 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Mon, 19 Oct 2026 15:20:42 -0400 Subject: [PATCH 58/58] chore: resets dotnet changes --- .../DefaultMcpToolHandler.cs | 694 ++--------------- .../PublicAPI/net10.0/PublicAPI.Unshipped.txt | 1 - .../PublicAPI/net472/PublicAPI.Unshipped.txt | 1 - .../PublicAPI/net8.0/PublicAPI.Unshipped.txt | 1 - .../PublicAPI/net9.0/PublicAPI.Unshipped.txt | 1 - .../netstandard2.0/PublicAPI.Unshipped.txt | 1 - .../IMcpToolHandler.cs | 31 - .../Interpreter/DeclarativeActionExecutor.cs | 4 +- .../Interpreter/DeclarativeWorkflowContext.cs | 26 +- .../DeclarativeWorkflowExecutor.cs | 3 +- .../Interpreter/DelegateActionExecutor.cs | 4 +- .../Kit/ActionExecutor.cs | 4 +- .../Kit/RootExecutor.cs | 3 +- .../ObjectModel/InvokeMcpToolExecutor.cs | 48 +- .../PublicAPI/net10.0/PublicAPI.Unshipped.txt | 2 - .../PublicAPI/net472/PublicAPI.Unshipped.txt | 2 - .../PublicAPI/net8.0/PublicAPI.Unshipped.txt | 2 - .../PublicAPI/net9.0/PublicAPI.Unshipped.txt | 2 - .../netstandard2.0/PublicAPI.Unshipped.txt | 2 - .../GroupChatManager.cs | 19 +- .../IWorkflowContext.cs | 15 - .../InProc/InProcessRunner.cs | 35 +- .../InProc/InProcessRunnerContext.cs | 10 +- .../PublicAPI/net10.0/PublicAPI.Unshipped.txt | 2 - .../PublicAPI/net472/PublicAPI.Unshipped.txt | 2 - .../PublicAPI/net8.0/PublicAPI.Unshipped.txt | 2 - .../PublicAPI/net9.0/PublicAPI.Unshipped.txt | 2 - .../netstandard2.0/PublicAPI.Unshipped.txt | 2 - .../Specialized/WorkflowHostExecutor.cs | 14 +- .../SubworkflowBinding.cs | 8 - .../DefaultMcpToolHandlerLifetimeTests.cs | 700 +----------------- .../DefaultMcpToolHandlerTests.cs | 37 +- .../Kit/IWorkflowContextExtensionsTests.cs | 22 +- .../Kit/RootExecutorTests.cs | 3 +- .../ObjectModel/InvokeMcpToolExecutorTest.cs | 139 +--- .../PowerFx/WorkflowFormulaStateTests.cs | 52 -- .../RepresentationTests.cs | 43 -- 37 files changed, 116 insertions(+), 1823 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 1f5664a36c8..9bd67e55eb4 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs @@ -6,7 +6,6 @@ using System.IO; using System.Linq; using System.Net.Http; -using System.Runtime.ExceptionServices; using System.Security.Cryptography; using System.Text; using System.Text.Json; @@ -29,17 +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, workflow invocations are cached by workflow session, server URL, label, -/// connection name, and explicit headers in a bounded least-recently-used cache. The bound includes -/// connecting and closing sessions. At capacity, an idle session is disposed before its replacement -/// connects. If every slot is busy, new scopes fail with ; -/// existing scopes can still reuse their cached session. +/// Without a provider, sessions are cached by 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 : IWorkflowScopedMcpToolHandler, IAsyncDisposable +public sealed class DefaultMcpToolHandler : IMcpToolHandler, IAsyncDisposable { - private const int DefaultClientCacheMaxSize = 32; private const string FilenameAdditionalPropertyName = "filename"; /// @@ -49,21 +43,12 @@ public sealed class DefaultMcpToolHandler : IWorkflowScopedMcpToolHandler, IAsyn public const string ListToolsToolName = "tools/list"; private static readonly JsonWriterOptions s_toolListJsonWriterOptions = new() { Indented = true }; - private static readonly OperationCanceledException s_clientCreationCancelledException = - new("The MCP client creator was cancelled."); private readonly Func>? _httpClientProvider; private readonly Func _httpMessageHandlerFactory; - private readonly Func _clientConnectionDisposer; - private readonly Dictionary<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash), CachedClient> _clients = []; - private readonly Dictionary<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash), TaskCompletionSource> _clientCreations = []; - private readonly HashSet> _clientCreationLifetimes = []; - private readonly LinkedList<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash)> _clientLru = []; - private readonly HashSet _retiredClients = []; - private readonly Dictionary _ownedHttpClients = []; + private readonly Dictionary<(string Url, string Label, string Connection, string HeadersHash), ClientConnection> _clients = []; + private readonly Dictionary _ownedHttpClients = []; private readonly SemaphoreSlim _clientLock = new(1, 1); - private readonly SemaphoreSlim _clientCreationSemaphore; - private readonly int _clientCacheMaxSize; private readonly AsyncLocal _providerInvocationContext = new(); private TaskCompletionSource? _providerInvocationsDrained; private int _activeProviderInvocations; @@ -98,20 +83,10 @@ public DefaultMcpToolHandler(Func>? internal DefaultMcpToolHandler( Func>? httpClientProvider, - Func httpMessageHandlerFactory, - int clientCacheMaxSize = DefaultClientCacheMaxSize, - Func? clientConnectionDisposer = null) + Func httpMessageHandlerFactory) { - if (clientCacheMaxSize <= 0) - { - throw new ArgumentOutOfRangeException(nameof(clientCacheMaxSize), "The MCP client cache size must be positive."); - } - this._httpClientProvider = httpClientProvider; this._httpMessageHandlerFactory = Throw.IfNull(httpMessageHandlerFactory); - this._clientConnectionDisposer = clientConnectionDisposer ?? (connection => connection.DisposeAsync()); - this._clientCacheMaxSize = clientCacheMaxSize; - this._clientCreationSemaphore = new(clientCacheMaxSize, clientCacheMaxSize); } /// @@ -123,26 +98,6 @@ 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)) { @@ -186,7 +141,7 @@ public async Task InvokeToolInWorkflowSessionAsync( try { ClientConnection invocationClient = await this.CreateClientAsync( - serverUrl.Trim(), serverLabel, headers, cancellationToken).ConfigureAwait(false); + serverUrl.Trim(), serverLabel, headers, httpClientCacheKey: null, cancellationToken).ConfigureAwait(false); await using (invocationClient.ConfigureAwait(false)) { return await InvokeClientAsync(invocationClient.Client, toolName, arguments, cancellationToken).ConfigureAwait(false); @@ -211,16 +166,8 @@ public async Task InvokeToolInWorkflowSessionAsync( } } - CachedClient cachedClient = await this.AcquireClientAsync( - serverUrl, serverLabel, headers, connectionName, workflowSessionId, cancellationToken).ConfigureAwait(false); - try - { - return await InvokeClientAsync(cachedClient.Connection.Client, toolName, arguments, cancellationToken).ConfigureAwait(false); - } - finally - { - await this.ReleaseClientAsync(cachedClient).ConfigureAwait(false); - } + McpClient client = await this.GetOrCreateClientAsync(serverUrl, serverLabel, headers, connectionName, cancellationToken).ConfigureAwait(false); + return await InvokeClientAsync(client, toolName, arguments, cancellationToken).ConfigureAwait(false); } private static async Task InvokeClientAsync( @@ -286,14 +233,12 @@ public async ValueTask DisposeAsync() } Task? providerInvocations; - List> clientCreationLifetimes; await this._clientLock.WaitAsync().ConfigureAwait(false); try { this.ThrowIfDisposing(); this._disposing = true; providerInvocations = this._providerInvocationsDrained?.Task; - clientCreationLifetimes = this._clientCreationLifetimes.Select(source => source.Task).ToList(); } finally { @@ -305,99 +250,32 @@ public async ValueTask DisposeAsync() await providerInvocations.ConfigureAwait(false); } - Exception? clientCreationException = null; - try - { - await Task.WhenAll(clientCreationLifetimes).ConfigureAwait(false); - } - catch (Exception exception) when (!IsFatalException(exception)) - { - clientCreationException = exception; - } - - List cachedClients; - List clientsToDispose = []; await this._clientLock.WaitAsync().ConfigureAwait(false); try { - cachedClients = [.. this._clients.Values, .. this._retiredClients]; - foreach (CachedClient client in cachedClients) + foreach (ClientConnection client in this._clients.Values) { - client.Evicted = true; - this._retiredClients.Add(client); - if (client.ActiveInvocations == 0 && !client.DisposalClaimed) - { - client.DisposalClaimed = true; - clientsToDispose.Add(client); - } + await client.DisposeAsync().ConfigureAwait(false); } this._clients.Clear(); - this._clientLru.Clear(); - } - finally - { - this._clientLock.Release(); - } - - Exception? clientCleanupException = null; - try - { - await DrainCleanupAsync( - clientsToDispose.Select(client => this.DisposeCachedClientAsync(client)), - cachedClients.Select(client => client.Disposed.Task)).ConfigureAwait(false); - } - catch (Exception exception) when (!IsFatalException(exception)) - { - clientCleanupException = exception; - } - finally - { - this._clientLock.Dispose(); - this._clientCreationSemaphore.Dispose(); - } - if (clientCreationException is not null) - { - ExceptionDispatchInfo.Capture(clientCreationException).Throw(); - } - - if (clientCleanupException is not null) - { - ExceptionDispatchInfo.Capture(clientCleanupException).Throw(); - } - } + // Dispose only HttpClients that the handler created (not user-provided ones) + foreach (HttpClient httpClient in this._ownedHttpClients.Values) + { + httpClient.Dispose(); + } - internal static async Task DrainCleanupAsync(IEnumerable cleanupTasks, IEnumerable completionTasks) - { - Exception? cleanupException = null; - try - { - await Task.WhenAll(cleanupTasks).ConfigureAwait(false); + this._ownedHttpClients.Clear(); } - catch (Exception exception) when (!IsFatalException(exception)) + finally { - cleanupException = exception; + this._clientLock.Release(); } - await Task.WhenAll(completionTasks).ConfigureAwait(false); - - if (cleanupException is not null) - { - ExceptionDispatchInfo.Capture(cleanupException).Throw(); - } + this._clientLock.Dispose(); } - 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() @@ -408,254 +286,28 @@ private void ThrowIfDisposing() } } - private async Task AcquireClientAsync( + private async Task GetOrCreateClientAsync( string serverUrl, string? serverLabel, IDictionary? headers, string? connectionName, - string workflowSessionId, CancellationToken cancellationToken) { string trimmedUrl = serverUrl.Trim(); - var clientCacheKey = BuildCacheKey(workflowSessionId, trimmedUrl, serverLabel, connectionName, headers); - CachedClient? clientToDispose = null; - TaskCompletionSource? clientCreation; - TaskCompletionSource? clientCreationLifetime = null; - bool ownsClientCreation = false; + var clientCacheKey = BuildCacheKey(trimmedUrl, serverLabel, connectionName, headers); await this._clientLock.WaitAsync(cancellationToken).ConfigureAwait(false); try { this.ThrowIfDisposing(); - if (this._clients.TryGetValue(clientCacheKey, out CachedClient? existingClient)) - { - existingClient.ActiveInvocations++; - this._clientLru.Remove(existingClient.LruNode); - this._clientLru.AddLast(existingClient.LruNode); - return existingClient; - } - - if (!this._clientCreations.TryGetValue(clientCacheKey, out clientCreation)) - { - if (this._clients.Count + this._retiredClients.Count + this._clientCreations.Count >= this._clientCacheMaxSize) - { - LinkedListNode<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash)>? idleNode = - this._clientLru.First; - while (idleNode is not null && this._clients[idleNode.Value].ActiveInvocations != 0) - { - idleNode = idleNode.Next; - } - - if (idleNode is null) - { - throw new InvalidOperationException("MCP client capacity is exhausted; retry when an active invocation completes."); - } - - clientToDispose = this._clients[idleNode.Value]; - this._clients.Remove(idleNode.Value); - this._clientLru.Remove(idleNode); - clientToDispose.Evicted = true; - clientToDispose.DisposalClaimed = true; - this._retiredClients.Add(clientToDispose); - } - - clientCreation = new(TaskCreationOptions.RunContinuationsAsynchronously); - this._clientCreations[clientCacheKey] = clientCreation; - clientCreationLifetime = new(TaskCreationOptions.RunContinuationsAsynchronously); - this._clientCreationLifetimes.Add(clientCreationLifetime); - ownsClientCreation = true; - } - } - finally - { - this._clientLock.Release(); - } - - if (!ownsClientCreation) - { - try - { - await WaitForClientCreationAsync(clientCreation.Task, cancellationToken).ConfigureAwait(false); - } - catch (OperationCanceledException exception) - when (ReferenceEquals(exception, s_clientCreationCancelledException)) - { - return await this.AcquireClientAsync( - serverUrl, serverLabel, headers, connectionName, workflowSessionId, cancellationToken).ConfigureAwait(false); - } - - return await this.AcquireClientAsync( - serverUrl, serverLabel, headers, connectionName, workflowSessionId, cancellationToken).ConfigureAwait(false); - } - - TaskCompletionSource ownedClientCreationLifetime = clientCreationLifetime ?? - throw new InvalidOperationException("Missing MCP client creation lifetime."); - ClientConnection? connection = null; - bool creationSemaphoreEntered = false; - Exception? clientCreationCleanupException = null; - try - { - if (clientToDispose is not null) - { - try - { - await this.DisposeCachedClientAsync(clientToDispose).ConfigureAwait(false); - } - catch (Exception cleanupException) - { - clientCreationCleanupException = cleanupException; - throw; - } - } - - await this._clientCreationSemaphore.WaitAsync(cancellationToken).ConfigureAwait(false); - creationSemaphoreEntered = true; - await this._clientLock.WaitAsync(cancellationToken).ConfigureAwait(false); - try - { - this.ThrowIfDisposing(); - } - finally - { - this._clientLock.Release(); - } - - connection = await this.CreateClientAsync( - trimmedUrl, serverLabel, headers, cancellationToken, - cleanupException => clientCreationCleanupException = cleanupException).ConfigureAwait(false); - } - catch (Exception exception) - { - Exception sharedCreationException = - exception is OperationCanceledException && cancellationToken.IsCancellationRequested - ? s_clientCreationCancelledException - : exception; - try - { - await this.CompleteClientCreationFailureAsync( - clientCacheKey, clientCreation, sharedCreationException).ConfigureAwait(false); - } - finally - { - if (creationSemaphoreEntered) - { - this._clientCreationSemaphore.Release(); - creationSemaphoreEntered = false; - } - - await this.CompleteClientCreationLifetimeAsync( - ownedClientCreationLifetime, clientCreationCleanupException).ConfigureAwait(false); - } - - throw; - } - - ObjectDisposedException? disposedException = null; - CachedClient? result = null; - await this._clientLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); - try - { - this._clientCreations.Remove(clientCacheKey); - if (this._disposing) - { - disposedException = new ObjectDisposedException(nameof(DefaultMcpToolHandler)); - } - else - { - LinkedListNode<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash)> node = - this._clientLru.AddLast(clientCacheKey); - CachedClient newClient = new(connection, node) { ActiveInvocations = 1 }; - connection = null; - this._clients[clientCacheKey] = newClient; - - result = newClient; - clientCreation.TrySetResult(newClient); - } - } - finally - { - this._clientLock.Release(); - } - - try - { - if (connection is not null) - { - await this._clientConnectionDisposer(connection).ConfigureAwait(false); - } - - if (disposedException is not null) - { - clientCreation.TrySetException(disposedException); - throw disposedException; - } - } - catch (Exception exception) - { - if (disposedException is not null && !ReferenceEquals(exception, disposedException)) - { - clientCreation.TrySetException(exception); - clientCreationCleanupException = exception; - } - - if (result is not null) - { - await this.ReleaseClientAsync(result).ConfigureAwait(false); - } - - throw; - } - finally - { - if (creationSemaphoreEntered) - { - this._clientCreationSemaphore.Release(); - } - - await this.CompleteClientCreationLifetimeAsync( - ownedClientCreationLifetime, clientCreationCleanupException).ConfigureAwait(false); - } - - return result ?? throw new InvalidOperationException("Failed to acquire MCP client."); - } - - private static async Task WaitForClientCreationAsync(Task clientCreation, CancellationToken cancellationToken) - { - if (!cancellationToken.CanBeCanceled || clientCreation.IsCompleted) - { - await clientCreation.ConfigureAwait(false); - return; - } - - TaskCompletionSource cancellation = new(TaskCreationOptions.RunContinuationsAsynchronously); - using CancellationTokenRegistration registration = cancellationToken.Register( - static state => ((TaskCompletionSource)state!).TrySetResult(true), - cancellation); - - if (await Task.WhenAny(clientCreation, cancellation.Task).ConfigureAwait(false) == cancellation.Task) - { - cancellationToken.ThrowIfCancellationRequested(); - } - - await clientCreation.ConfigureAwait(false); - } - - private async Task CompleteClientCreationFailureAsync( - (string WorkflowSession, string Url, string Label, string Connection, string HeadersHash) clientCacheKey, - TaskCompletionSource clientCreation, - Exception exception) - { - await this._clientLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); - try - { - if (this._clientCreations.TryGetValue(clientCacheKey, out TaskCompletionSource? existingCreation) && - ReferenceEquals(existingCreation, clientCreation)) + if (this._clients.TryGetValue(clientCacheKey, out ClientConnection? existingClient)) { - this._clientCreations.Remove(clientCacheKey); + return existingClient.Client; } - clientCreation.TrySetException(exception); - _ = clientCreation.Task.Exception; + ClientConnection newClient = await this.CreateClientAsync(trimmedUrl, serverLabel, headers, trimmedUrl, cancellationToken).ConfigureAwait(false); + this._clients[clientCacheKey] = newClient; + return newClient.Client; } finally { @@ -663,129 +315,60 @@ private async Task CompleteClientCreationFailureAsync( } } - private async Task CompleteClientCreationLifetimeAsync( - TaskCompletionSource clientCreationLifetime, - Exception? cleanupException = null) - { - await this._clientLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); - try - { - this._clientCreationLifetimes.Remove(clientCreationLifetime); - } - finally - { - this._clientLock.Release(); - } - - if (cleanupException is not null) - { - clientCreationLifetime.TrySetException(cleanupException); - } - else - { - clientCreationLifetime.TrySetResult(true); - } - } - - private async ValueTask ReleaseClientAsync(CachedClient client) - { - bool dispose = false; - await this._clientLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); - try - { - client.ActiveInvocations--; - if (client.ActiveInvocations == 0 && client.Evicted && !client.DisposalClaimed) - { - client.DisposalClaimed = true; - dispose = true; - } - } - finally - { - this._clientLock.Release(); - } - - if (dispose) - { - await this.DisposeCachedClientAsync(client).ConfigureAwait(false); - } - } - - private async Task DisposeCachedClientAsync(CachedClient client) - { - Exception? cleanupException = null; - try - { - await this._clientConnectionDisposer(client.Connection).ConfigureAwait(false); - } - catch (Exception exception) - { - cleanupException = exception; - throw; - } - finally - { - await this._clientLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); - try - { - this._retiredClients.Remove(client); - } - finally - { - this._clientLock.Release(); - } - - if (cleanupException is not null) - { - client.Disposed.TrySetException(cleanupException); - } - else - { - client.Disposed.TrySetResult(true); - } - } - } - /// - /// Builds the per-client cache key as a 5-tuple of - /// (workflowSessionId, trimmed serverUrl, serverLabel, connectionName, headers hash). - /// All five components participate so that separate workflow sessions and callers using - /// different labels/connections/headers receive distinct instances - /// even when targeting the same URL. + /// 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. /// - internal static (string WorkflowSession, string Url, string Label, string Connection, string HeadersHash) BuildCacheKey( - string workflowSessionId, + internal static (string Url, string Label, string Connection, string HeadersHash) BuildCacheKey( string trimmedUrl, string? serverLabel, string? connectionName, IDictionary? headers) => - (workflowSessionId, trimmedUrl, serverLabel ?? string.Empty, connectionName ?? string.Empty, ComputeHeadersHash(headers)); + (trimmedUrl, serverLabel ?? string.Empty, connectionName ?? string.Empty, ComputeHeadersHash(headers)); private async Task CreateClientAsync( string serverUrl, string? serverLabel, IDictionary? headers, - CancellationToken cancellationToken, - Action? reportCleanupFailure = null) + 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 && this._httpClientProvider is not null) - { - httpClient = this.CreatePinnedHttpClient(serverUrl); - ownsHttpClient = true; - } - else if (httpClient is null) + if (httpClient is null && + (httpClientCacheKey is null || !this._ownedHttpClients.TryGetValue(httpClientCacheKey, out httpClient))) { - ownedHttpClientLease = await this.AcquireOwnedHttpClientAsync(serverUrl, cancellationToken).ConfigureAwait(false); - httpClient = ownedHttpClientLease.Client; + // Pin credential headers to the configured server origin as defense-in-depth. Forcing + // StreamableHttp (below) already removes the primary vector (a server-advertised cross-origin + // SSE message endpoint), and AllowAutoRedirect=false blocks auto-redirects. This handler is the + // backstop: it guarantees the Authorization token and other credentials never leave the pinned + // origin even if a future change re-enables AutoDetect or redirects, or the SDK constructs a + // request to a new URI (AdditionalHeaders are re-stamped by the transport, so HttpClient's own + // redirect header-stripping does not cover them). + // Caveat for future maintainers: the handler strips credentials on every cross-origin request. + // That is safe today because this handler carries auth via static request headers, not the + // SDK's built-in OAuth flow. If OAuth support is ever added here, requests to the authorization + // server (a different origin than the MCP resource server) legitimately carry credentials, so + // those auth-server origins would need to be allow-listed to avoid breaking the token exchange. + OriginPinningHandler pinningHandler = new(new Uri(serverUrl)) { InnerHandler = this._httpMessageHandlerFactory() }; + httpClient = new HttpClient(pinningHandler); + if (httpClientCacheKey is null) + { + ownsHttpClient = true; + } + else + { + this._ownedHttpClients[httpClientCacheKey] = httpClient; + } } HttpClientTransportOptions transportOptions = new() @@ -800,108 +383,20 @@ private async Task CreateClientAsync( TransportMode = HttpTransportMode.StreamableHttp }; - HttpClientTransport? transport = null; + HttpClientTransport transport = new(transportOptions, httpClient, ownsHttpClient: ownsHttpClient); + try { - HttpClient resolvedHttpClient = httpClient ?? throw new InvalidOperationException("Failed to resolve MCP HTTP client."); - transport = new(transportOptions, resolvedHttpClient, ownsHttpClient: ownsHttpClient); McpClient client = await McpClient.CreateAsync(transport, cancellationToken: cancellationToken).ConfigureAwait(false); - ClientConnection connection = new(client, transport, ownedHttpClientLease); - ownedHttpClientLease = null; - return connection; + return new ClientConnection(client, transport); } catch { - try - { - try - { - if (transport is not null) - { - await DisposeResourceAsync(transport, "transport").ConfigureAwait(false); - } - } - finally - { - if (ownedHttpClientLease is not null) - { - await ownedHttpClientLease.DisposeAsync().ConfigureAwait(false); - } - } - } - catch (Exception cleanupException) - { - reportCleanupFailure?.Invoke(cleanupException); - throw; - } - + await DisposeResourceAsync(transport, "transport").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(); - } - - try - { - clientToDispose?.Dispose(); - } - catch (Exception exception) when (exception is not OperationCanceledException) - { - Trace.TraceWarning("Failed to dispose MCP HTTP client: {0}", exception); - } - } - - private HttpClient CreatePinnedHttpClient(string serverUrl) - { - // Pin credential headers to the configured server origin as defense-in-depth. Forcing - // StreamableHttp (below) already removes the primary vector (a server-advertised cross-origin - // SSE message endpoint), and AllowAutoRedirect=false blocks auto-redirects. This handler is the - // backstop: it guarantees the Authorization token and other credentials never leave the pinned - // origin even if a future change re-enables AutoDetect or redirects, or the SDK constructs a - // request to a new URI (AdditionalHeaders are re-stamped by the transport, so HttpClient's own - // redirect header-stripping does not cover them). - OriginPinningHandler pinningHandler = new(new Uri(serverUrl)) { InnerHandler = this._httpMessageHandlerFactory() }; - return new HttpClient(pinningHandler); - } - private static HttpMessageHandler CreateHttpMessageHandler() => new HttpClientHandler { @@ -914,56 +409,11 @@ private static HttpMessageHandler CreateHttpMessageHandler() => private sealed class ProviderInvocationContext(Task completion, ProviderInvocationContext? parent) { public Task Completion { get; } = completion; - public ProviderInvocationContext? Parent { get; } = parent; - } - - private sealed class CachedClient( - ClientConnection connection, - LinkedListNode<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash)> lruNode) - { - public ClientConnection Connection { get; } = connection; - - public LinkedListNode<(string WorkflowSession, string Url, string Label, string Connection, string HeadersHash)> LruNode { get; } = lruNode; - - public TaskCompletionSource Disposed { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously); - - public int ActiveInvocations { get; set; } - - public bool Evicted { get; set; } - - public bool DisposalClaimed { get; set; } - } - - private sealed class OwnedHttpClient(HttpClient client) - { - public HttpClient Client { get; } = client; - - public int ReferenceCount { get; set; } - } - - private sealed class OwnedHttpClientLease( - DefaultMcpToolHandler owner, - string serverUrl, - HttpClient client) : IAsyncDisposable - { - private bool _disposed; - - public HttpClient Client { get; } = client; - public async ValueTask DisposeAsync() - { - if (!this._disposed) - { - this._disposed = true; - await owner.ReleaseOwnedHttpClientAsync(serverUrl).ConfigureAwait(false); - } - } + public ProviderInvocationContext? Parent { get; } = parent; } - internal sealed class ClientConnection( - McpClient client, - IAsyncDisposable transport, - IAsyncDisposable? ownedHttpClientLease = null) : IAsyncDisposable + internal sealed class ClientConnection(McpClient client, IAsyncDisposable transport) : IAsyncDisposable { public McpClient Client { get; } = client; @@ -975,18 +425,8 @@ public async ValueTask DisposeAsync() } finally { - try - { - // McpClient owns the connected session, not the reusable transport factory. - await DisposeResourceAsync(transport, "transport").ConfigureAwait(false); - } - finally - { - if (ownedHttpClientLease is not null) - { - await ownedHttpClientLease.DisposeAsync().ConfigureAwait(false); - } - } + // McpClient owns the connected session, not the reusable transport factory. + await DisposeResourceAsync(transport, "transport").ConfigureAwait(false); } } } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net10.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/PublicAPI/net10.0/PublicAPI.Unshipped.txt index 3cdebd5e1a0..ab058de62d4 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,2 +1 @@ #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 3cdebd5e1a0..ab058de62d4 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,2 +1 @@ #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 3cdebd5e1a0..ab058de62d4 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,2 +1 @@ #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 3cdebd5e1a0..ab058de62d4 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,2 +1 @@ #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 3cdebd5e1a0..ab058de62d4 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,2 +1 @@ #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 fbea21f35b4..56b1c3deb4c 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/IMcpToolHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/IMcpToolHandler.cs @@ -39,34 +39,3 @@ Task InvokeToolAsync( string? connectionName, CancellationToken cancellationToken = default); } - -/// -/// Defines the contract for MCP handlers that isolate stateful protocol sessions by workflow session. -/// -public interface IWorkflowScopedMcpToolHandler : IMcpToolHandler -{ - /// - /// Invokes an MCP tool within the framework-owned workflow session scope. - /// - /// The URL of the MCP server. - /// An optional label identifying the server connection. - /// The name of the tool to invoke. - /// Optional arguments to pass to the tool. - /// Optional headers to include in the request. - /// An optional connection name for managed connections. - /// The framework-owned identifier for the current workflow session. - /// A token to observe cancellation. - /// - /// A task representing the asynchronous operation. The result contains a - /// with the tool invocation output. - /// - Task InvokeToolInWorkflowSessionAsync( - string serverUrl, - string? serverLabel, - string toolName, - IDictionary? arguments, - IDictionary? headers, - string? connectionName, - string workflowSessionId, - CancellationToken cancellationToken = default); -} diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeActionExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeActionExecutor.cs index 2c5ac9867cd..4902589514d 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeActionExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeActionExecutor.cs @@ -96,9 +96,7 @@ public override async ValueTask HandleAsync(ActionExecutorResult message, IWorkf try { - DeclarativeWorkflowContext declarativeContext = - await DeclarativeWorkflowContext.CreateAsync(context, this.State, cancellationToken).ConfigureAwait(false); - object? result = await this.ExecuteAsync(declarativeContext, cancellationToken).ConfigureAwait(false); + object? result = await this.ExecuteAsync(new DeclarativeWorkflowContext(context, this.State), 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 f88035abd12..183b2e785e7 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowContext.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowContext.cs @@ -15,10 +15,8 @@ namespace Microsoft.Agents.AI.Workflows.Declarative.Interpreter; -internal sealed class DeclarativeWorkflowContext : IWorkflowContext, IWorkflowSessionContext +internal sealed class DeclarativeWorkflowContext : IWorkflowContext { - internal const string WorkflowSessionIdStateKey = "__declarative_mcp_workflow_session_id"; - public static readonly FrozenSet ManagedScopes = [ VariableScopeNames.Local, @@ -26,36 +24,16 @@ internal sealed class DeclarativeWorkflowContext : IWorkflowContext, IWorkflowSe VariableScopeNames.Global, ]; - private DeclarativeWorkflowContext(IWorkflowContext source, WorkflowFormulaState state, string sessionId) + public DeclarativeWorkflowContext(IWorkflowContext source, WorkflowFormulaState state) { this.Source = source; this.State = state; - this.SessionId = sessionId; - } - - public static async ValueTask CreateAsync( - IWorkflowContext source, - WorkflowFormulaState state, - CancellationToken cancellationToken = default) - { - string generatedSessionId = Guid.NewGuid().ToString("N"); - string sessionId = source is IWorkflowSessionContext sessionContext - ? sessionContext.SessionId - : await source.ReadOrInitStateAsync( - WorkflowSessionIdStateKey, - () => generatedSessionId, - VariableScopeNames.System, - cancellationToken: cancellationToken).ConfigureAwait(false); - - return new(source, state, sessionId); } private IWorkflowContext Source { get; } public WorkflowFormulaState State { get; } public IReadOnlyDictionary? TraceContext => this.Source.TraceContext; - public string SessionId { get; } - /// public bool ConcurrentRunsEnabled => this.Source.ConcurrentRunsEnabled; diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowExecutor.cs index 2fcdea4ecce..1126b43d8ff 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowExecutor.cs @@ -154,8 +154,7 @@ private async ValueTask AdvanceAsync(ChatMessage input, IWorkflowContext context // No state to restore if we're starting from the beginning. state.SetInitialized(); - DeclarativeWorkflowContext declarativeContext = - await DeclarativeWorkflowContext.CreateAsync(context, state, cancellationToken).ConfigureAwait(false); + DeclarativeWorkflowContext declarativeContext = new(context, state); // 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 ade8dd1034b..716a3656dac 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DelegateActionExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DelegateActionExecutor.cs @@ -61,9 +61,7 @@ public override async ValueTask HandleAsync(TMessage message, IWorkflowContext c { if (this._action is not null) { - DeclarativeWorkflowContext declarativeContext = - await DeclarativeWorkflowContext.CreateAsync(context, this._state, cancellationToken).ConfigureAwait(false); - await this._action.Invoke(declarativeContext, message, cancellationToken).ConfigureAwait(false); + await this._action.Invoke(new DeclarativeWorkflowContext(context, this._state), 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 563ff86f147..cf636effaf8 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/ActionExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/ActionExecutor.cs @@ -76,9 +76,7 @@ public ValueTask ResetAsync() [SendsMessage(typeof(ActionExecutorResult))] public override async ValueTask HandleAsync(TMessage message, IWorkflowContext context, CancellationToken cancellationToken) { - DeclarativeWorkflowContext declarativeContext = - await DeclarativeWorkflowContext.CreateAsync(context, this._session.State, cancellationToken).ConfigureAwait(false); - object? result = await this.ExecuteAsync(declarativeContext, message, cancellationToken).ConfigureAwait(false); + object? result = await this.ExecuteAsync(new DeclarativeWorkflowContext(context, this._session.State), 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 9e008ed391b..3a18a06c948 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/RootExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/RootExecutor.cs @@ -63,8 +63,7 @@ public ValueTask ResetAsync() [SendsMessage(typeof(ActionExecutorResult))] public override async ValueTask HandleAsync(TInput message, IWorkflowContext context, CancellationToken cancellationToken) { - DeclarativeWorkflowContext declarativeContext = - await DeclarativeWorkflowContext.CreateAsync(context, this._state, cancellationToken).ConfigureAwait(false); + DeclarativeWorkflowContext declarativeContext = new(context, this._state); ChatMessage input = (this._inputTransform ?? DefaultInputTransform).Invoke(message); diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/InvokeMcpToolExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/InvokeMcpToolExecutor.cs index 5f76bb3621e..02dafb8ce2c 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/InvokeMcpToolExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/InvokeMcpToolExecutor.cs @@ -32,6 +32,7 @@ internal sealed class InvokeMcpToolExecutor( { private const string ApprovalSnapshotStateKey = nameof(_approvalSnapshots); private const string LegacyApprovalSnapshotStateKey = "_approvalSnapshot"; + /// /// Snapshots of evaluated parameters captured at approval-request time, keyed by /// per-invocation request id. Each pending approval lives here until the matching @@ -109,8 +110,7 @@ await this.RequestApprovalAsync( } // No approval required - invoke the tool directly - McpServerToolResultContent resultContent = await this.InvokeToolAsync( - context, + McpServerToolResultContent resultContent = await mcpToolHandler.InvokeToolAsync( serverUrl, serverLabel, toolName, @@ -193,8 +193,7 @@ public async ValueTask CaptureResponseAsync( string? conversationId = this.GetConversationId(); bool autoSend = snapshot.AutoSend ?? this.GetAutoSendValue(); - McpServerToolResultContent resultContent = await this.InvokeToolAsync( - context, + McpServerToolResultContent resultContent = await mcpToolHandler.InvokeToolAsync( snapshot.ServerUrl, snapshot.ServerLabel, snapshot.ToolName, @@ -207,47 +206,6 @@ public async ValueTask CaptureResponseAsync( await context.SendResultMessageAsync(this.Id, result: null, cancellationToken).ConfigureAwait(false); } - private async Task InvokeToolAsync( - IWorkflowContext context, - string serverUrl, - string? serverLabel, - string toolName, - IDictionary? arguments, - IDictionary? headers, - string? connectionName, - CancellationToken cancellationToken) - { - if (mcpToolHandler is IWorkflowScopedMcpToolHandler scopedHandler) - { - string generatedWorkflowSessionId = Guid.NewGuid().ToString("N"); - string workflowSessionId = context is IWorkflowSessionContext sessionContext - ? sessionContext.SessionId - : await context.ReadOrInitStateAsync( - DeclarativeWorkflowContext.WorkflowSessionIdStateKey, - () => generatedWorkflowSessionId, - VariableScopeNames.System, - cancellationToken).ConfigureAwait(false); - return await scopedHandler.InvokeToolInWorkflowSessionAsync( - serverUrl, - serverLabel, - toolName, - arguments, - headers, - connectionName, - workflowSessionId, - cancellationToken).ConfigureAwait(false); - } - - return await mcpToolHandler.InvokeToolAsync( - serverUrl, - serverLabel, - toolName, - arguments, - headers, - connectionName, - cancellationToken).ConfigureAwait(false); - } - /// /// Completes the MCP tool invocation by raising the completion event. /// diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net10.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net10.0/PublicAPI.Unshipped.txt index 5bee0f03310..8af74cb2369 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net10.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net10.0/PublicAPI.Unshipped.txt @@ -4,8 +4,6 @@ Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowedEnvi Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowProcessEnvironmentVariableFallback.get -> bool Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowProcessEnvironmentVariableFallback.init -> void Microsoft.Agents.AI.Workflows.Declarative.Events.ExternalInputResponse.RequestId.get -> string? -Microsoft.Agents.AI.Workflows.Declarative.IWorkflowScopedMcpToolHandler -Microsoft.Agents.AI.Workflows.Declarative.IWorkflowScopedMcpToolHandler.InvokeToolInWorkflowSessionAsync(string! serverUrl, string? serverLabel, string! toolName, System.Collections.Generic.IDictionary? arguments, System.Collections.Generic.IDictionary? headers, string? connectionName, string! workflowSessionId, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.ConvertValueWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, Microsoft.Agents.AI.Workflows.Declarative.Kit.VariableType! targetType, string! key, string? scopeName = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.ValueTask!> static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.EvaluateValueWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, string! expression, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.ValueTask!> static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.FormatTemplateWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, string! line) -> System.Threading.Tasks.ValueTask!> diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net472/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net472/PublicAPI.Unshipped.txt index 5bee0f03310..8af74cb2369 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net472/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net472/PublicAPI.Unshipped.txt @@ -4,8 +4,6 @@ Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowedEnvi Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowProcessEnvironmentVariableFallback.get -> bool Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowProcessEnvironmentVariableFallback.init -> void Microsoft.Agents.AI.Workflows.Declarative.Events.ExternalInputResponse.RequestId.get -> string? -Microsoft.Agents.AI.Workflows.Declarative.IWorkflowScopedMcpToolHandler -Microsoft.Agents.AI.Workflows.Declarative.IWorkflowScopedMcpToolHandler.InvokeToolInWorkflowSessionAsync(string! serverUrl, string? serverLabel, string! toolName, System.Collections.Generic.IDictionary? arguments, System.Collections.Generic.IDictionary? headers, string? connectionName, string! workflowSessionId, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.ConvertValueWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, Microsoft.Agents.AI.Workflows.Declarative.Kit.VariableType! targetType, string! key, string? scopeName = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.ValueTask!> static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.EvaluateValueWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, string! expression, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.ValueTask!> static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.FormatTemplateWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, string! line) -> System.Threading.Tasks.ValueTask!> diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net8.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net8.0/PublicAPI.Unshipped.txt index 5bee0f03310..8af74cb2369 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net8.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net8.0/PublicAPI.Unshipped.txt @@ -4,8 +4,6 @@ Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowedEnvi Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowProcessEnvironmentVariableFallback.get -> bool Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowProcessEnvironmentVariableFallback.init -> void Microsoft.Agents.AI.Workflows.Declarative.Events.ExternalInputResponse.RequestId.get -> string? -Microsoft.Agents.AI.Workflows.Declarative.IWorkflowScopedMcpToolHandler -Microsoft.Agents.AI.Workflows.Declarative.IWorkflowScopedMcpToolHandler.InvokeToolInWorkflowSessionAsync(string! serverUrl, string? serverLabel, string! toolName, System.Collections.Generic.IDictionary? arguments, System.Collections.Generic.IDictionary? headers, string? connectionName, string! workflowSessionId, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.ConvertValueWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, Microsoft.Agents.AI.Workflows.Declarative.Kit.VariableType! targetType, string! key, string? scopeName = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.ValueTask!> static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.EvaluateValueWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, string! expression, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.ValueTask!> static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.FormatTemplateWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, string! line) -> System.Threading.Tasks.ValueTask!> diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net9.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net9.0/PublicAPI.Unshipped.txt index 5bee0f03310..8af74cb2369 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net9.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/net9.0/PublicAPI.Unshipped.txt @@ -4,8 +4,6 @@ Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowedEnvi Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowProcessEnvironmentVariableFallback.get -> bool Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowProcessEnvironmentVariableFallback.init -> void Microsoft.Agents.AI.Workflows.Declarative.Events.ExternalInputResponse.RequestId.get -> string? -Microsoft.Agents.AI.Workflows.Declarative.IWorkflowScopedMcpToolHandler -Microsoft.Agents.AI.Workflows.Declarative.IWorkflowScopedMcpToolHandler.InvokeToolInWorkflowSessionAsync(string! serverUrl, string? serverLabel, string! toolName, System.Collections.Generic.IDictionary? arguments, System.Collections.Generic.IDictionary? headers, string? connectionName, string! workflowSessionId, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.ConvertValueWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, Microsoft.Agents.AI.Workflows.Declarative.Kit.VariableType! targetType, string! key, string? scopeName = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.ValueTask!> static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.EvaluateValueWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, string! expression, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.ValueTask!> static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.FormatTemplateWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, string! line) -> System.Threading.Tasks.ValueTask!> diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt index 5bee0f03310..8af74cb2369 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt @@ -4,8 +4,6 @@ Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowedEnvi Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowProcessEnvironmentVariableFallback.get -> bool Microsoft.Agents.AI.Workflows.Declarative.DeclarativeWorkflowOptions.AllowProcessEnvironmentVariableFallback.init -> void Microsoft.Agents.AI.Workflows.Declarative.Events.ExternalInputResponse.RequestId.get -> string? -Microsoft.Agents.AI.Workflows.Declarative.IWorkflowScopedMcpToolHandler -Microsoft.Agents.AI.Workflows.Declarative.IWorkflowScopedMcpToolHandler.InvokeToolInWorkflowSessionAsync(string! serverUrl, string? serverLabel, string! toolName, System.Collections.Generic.IDictionary? arguments, System.Collections.Generic.IDictionary? headers, string? connectionName, string! workflowSessionId, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.ConvertValueWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, Microsoft.Agents.AI.Workflows.Declarative.Kit.VariableType! targetType, string! key, string? scopeName = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.ValueTask!> static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.EvaluateValueWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, string! expression, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.ValueTask!> static Microsoft.Agents.AI.Workflows.Declarative.Kit.IWorkflowContextExtensions.FormatTemplateWithSensitivityAsync(this Microsoft.Agents.AI.Workflows.IWorkflowContext! context, string! line) -> System.Threading.Tasks.ValueTask!> diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/GroupChatManager.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/GroupChatManager.cs index fbda33c132a..ab94a9fa8aa 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(PrefixingWorkflowContext.Create(context, SubclassStateKeyPrefix), cancellationToken).ConfigureAwait(false); + await this.OnCheckpointingAsync(new PrefixingWorkflowContext(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(PrefixingWorkflowContext.Create(context, SubclassStateKeyPrefix), cancellationToken).ConfigureAwait(false); + await this.OnCheckpointRestoredAsync(new PrefixingWorkflowContext(context, SubclassStateKeyPrefix), cancellationToken).ConfigureAwait(false); } } @@ -169,16 +169,11 @@ 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 class PrefixingWorkflowContext(IWorkflowContext inner, string prefix) : IWorkflowContext +internal sealed 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; @@ -227,11 +222,3 @@ public async ValueTask QueueClearScopeAsync(string? scopeName = null, Cancellati private string Wrap(string key) => this._prefix + Throw.IfNullOrEmpty(key); } - -internal sealed class PrefixingWorkflowSessionContext( - IWorkflowContext inner, - string prefix, - IWorkflowSessionContext sessionContext) : PrefixingWorkflowContext(inner, prefix), IWorkflowSessionContext -{ - public string SessionId => sessionContext.SessionId; -} diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/IWorkflowContext.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/IWorkflowContext.cs index 68ff864085f..b8b35fffd64 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/IWorkflowContext.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/IWorkflowContext.cs @@ -195,18 +195,3 @@ ValueTask ReadOrInitStateAsync(string key, Func initialStateFactory, Ca /// bool ConcurrentRunsEnabled { get; } } - -/// -/// Exposes the framework-owned identifier for a workflow session. -/// -/// -/// Runtime services can use this identifier to scope stateful resources to one -/// workflow session without accepting caller-provided cache partition keys. -/// -public interface IWorkflowSessionContext -{ - /// - /// Gets the framework-owned workflow session identifier. - /// - string SessionId { get; } -} diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunner.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunner.cs index c0a2e075ae1..69b6c1e9bcc 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunner.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunner.cs @@ -31,14 +31,7 @@ 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, - string? workflowSessionId = null) + public static InProcessRunner CreateSubworkflowRunner(Workflow workflow, ICheckpointManager? checkpointManager, string? sessionId = null, object? existingOwnerSignoff = null, bool enableConcurrentRuns = false, IEnumerable? knownValidInputTypes = null) { return new InProcessRunner(workflow, checkpointManager, @@ -46,19 +39,10 @@ public static InProcessRunner CreateSubworkflowRunner( existingOwnerSignoff: existingOwnerSignoff, enableConcurrentRuns: enableConcurrentRuns, knownValidInputTypes: knownValidInputTypes, - subworkflow: true, - workflowSessionId: workflowSessionId); + subworkflow: true); } - private InProcessRunner( - Workflow workflow, - ICheckpointManager? checkpointManager, - string? sessionId = null, - object? existingOwnerSignoff = null, - bool subworkflow = false, - bool enableConcurrentRuns = false, - IEnumerable? knownValidInputTypes = null, - string? workflowSessionId = null) + private InProcessRunner(Workflow workflow, ICheckpointManager? checkpointManager, string? sessionId = null, object? existingOwnerSignoff = null, bool subworkflow = false, bool enableConcurrentRuns = false, IEnumerable? knownValidInputTypes = null) { if (enableConcurrentRuns && !workflow.AllowConcurrent) { @@ -70,16 +54,7 @@ private InProcessRunner( 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, - workflowSessionId: workflowSessionId); + this.RunContext = new InProcessRunnerContext(workflow, this.SessionId, checkpointingEnabled: checkpointManager != null, this.OutgoingEvents, this.StepTracer, existingOwnerSignoff, subworkflow, enableConcurrentRuns); this.CheckpointManager = checkpointManager; this._knownValidInputTypes = knownValidInputTypes != null @@ -90,8 +65,6 @@ private InProcessRunner( /// 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 aa8bf398592..353c38f25c9 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunnerContext.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunnerContext.cs @@ -47,8 +47,7 @@ public InProcessRunnerContext( object? existingOwnershipSignoff = null, bool subworkflow = false, bool enableConcurrentRuns = false, - ILogger? logger = null, - string? workflowSessionId = null) + ILogger? logger = null) { if (enableConcurrentRuns) { @@ -63,7 +62,6 @@ public InProcessRunnerContext( this._workflow = workflow; this._sessionId = sessionId; - this.WorkflowSessionId = workflowSessionId ?? sessionId; this._edgeMap = new(this, this._workflow, stepTracer); this._outputFilter = new(workflow); @@ -74,8 +72,6 @@ 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)) @@ -353,10 +349,8 @@ public IExternalRequestSink RegisterPort(RequestPort port) private sealed class BoundWorkflowContext( InProcessRunnerContext RunnerContext, string ExecutorId, - Dictionary? traceContext) : IWorkflowContext, IWorkflowSessionContext + Dictionary? traceContext) : IWorkflowContext { - public string SessionId => RunnerContext.WorkflowSessionId; - public ValueTask AddEventAsync(WorkflowEvent workflowEvent, CancellationToken cancellationToken = default) => RunnerContext.AddEventAsync(workflowEvent, cancellationToken); public ValueTask SendMessageAsync(object message, string? targetId = null, CancellationToken cancellationToken = default) diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net10.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI.Workflows/PublicAPI/net10.0/PublicAPI.Unshipped.txt index 724a27a74a4..ab058de62d4 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,3 +1 @@ #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 724a27a74a4..ab058de62d4 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,3 +1 @@ #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 724a27a74a4..ab058de62d4 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,3 +1 @@ #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 724a27a74a4..ab058de62d4 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,3 +1 @@ #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 724a27a74a4..ab058de62d4 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,3 +1 @@ #nullable enable -Microsoft.Agents.AI.Workflows.IWorkflowSessionContext -Microsoft.Agents.AI.Workflows.IWorkflowSessionContext.SessionId.get -> string! diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/WorkflowHostExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/WorkflowHostExecutor.cs index a166d8b0e6b..58e3a9e5239 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/WorkflowHostExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/WorkflowHostExecutor.cs @@ -20,7 +20,6 @@ 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; @@ -57,7 +56,6 @@ 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); @@ -97,8 +95,7 @@ internal async ValueTask EnsureRunnerAsync() this._checkpointManager, this._sessionId, this._ownershipToken, - this.JoinContext.ConcurrentRunsEnabled, - workflowSessionId: this._workflowSessionId); + this.JoinContext.ConcurrentRunsEnabled); } return this._activeRunner; @@ -271,7 +268,6 @@ 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) @@ -297,14 +293,6 @@ 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 9519af12414..11f7cf493cf 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/SubworkflowBinding.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/SubworkflowBinding.cs @@ -1,7 +1,6 @@ // 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; @@ -36,13 +35,6 @@ async ValueTask InitHostExecutorAsync(string sessionId) } } - internal static string CreateSubworkflowSessionId(string parentSessionId, string subworkflowId) => - string.Concat( - parentSessionId.Length.ToString(CultureInfo.InvariantCulture), - ":", - parentSessionId, - subworkflowId); - /// public override bool IsSharedInstance => false; diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs index 3ff71500d6c..3a481fcd0e6 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs @@ -24,150 +24,6 @@ namespace Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests; Justification = "Provider callbacks return completed tasks; each invocation and handler disposal is awaited before disposing caller-owned clients.")] public sealed class DefaultMcpToolHandlerLifetimeTests { - [Theory] - [InlineData(false)] - [InlineData(true)] - public async Task NoProvider_ActiveClientCapacity_RejectsNewScopesAndRecoversAsync(bool cancelActive) - { - // Arrange - ProtocolStub stub = new(); - using SemaphoreSlim started = new(0); - using SemaphoreSlim release = new(0); - int operations = 0; - stub.BeforeOperationAsync = async token => - { - if (Interlocked.Increment(ref operations) == 1) - { - started.Release(); - await release.WaitAsync(token); - } - }; - await using DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler, clientCacheMaxSize: 1); - using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); - using CancellationTokenSource firstCancellation = CancellationTokenSource.CreateLinkedTokenSource(timeout.Token); - Task first = InvokeScopedAsync(handler, "workflow-a", "ping", firstCancellation.Token); - await started.WaitAsync(timeout.Token); - - try - { - // Act - await InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); - InvalidOperationException exception = await Assert.ThrowsAsync( - () => InvokeScopedAsync(handler, "workflow-b", "ping", timeout.Token)); - - // Assert - Assert.Contains("capacity", exception.Message); - Assert.Equal(1, stub.Initializations); - Assert.Equal(0, stub.Terminations); - } - finally - { - if (cancelActive) - { - firstCancellation.Cancel(); - } - - release.Release(); - if (cancelActive) - { - await Assert.ThrowsAnyAsync(() => first); - } - else - { - await first; - } - } - - await InvokeScopedAsync(handler, "workflow-b", "ping", timeout.Token); - Assert.Equal(2, stub.Initializations); - Assert.Equal(1, stub.Terminations); - } - - [Fact] - public async Task NoProvider_Replacement_WaitsForIdleClientCleanupBeforeConnectingAsync() - { - // Arrange - ProtocolStub stub = new(); - using SemaphoreSlim closeStarted = new(0); - using SemaphoreSlim releaseClose = new(0); - int disposals = 0; - await using DefaultMcpToolHandler handler = new( - null, stub.CreateMessageHandler, clientCacheMaxSize: 1, - clientConnectionDisposer: async connection => - { - if (Interlocked.Increment(ref disposals) == 1) - { - closeStarted.Release(); - await releaseClose.WaitAsync(); - } - - await connection.DisposeAsync(); - }); - using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); - await InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); - Task replacement = InvokeScopedAsync(handler, "workflow-b", "ping", timeout.Token); - await closeStarted.WaitAsync(timeout.Token); - - try - { - // Act - Assert.Equal(1, stub.Initializations); - InvalidOperationException exception = await Assert.ThrowsAsync( - () => InvokeScopedAsync(handler, "workflow-c", "ping", timeout.Token)); - - // Assert - Assert.Contains("capacity", exception.Message); - } - finally - { - releaseClose.Release(); - await replacement; - } - - Assert.Equal(2, stub.Initializations); - Assert.Equal(1, stub.Terminations); - } - - [Fact] - public async Task NoProvider_DisposalDuringIdleEviction_ReportsCleanupCancellationAsync() - { - // Arrange - OperationCanceledException cleanupException = new("session cleanup cancelled"); - ProtocolStub stub = new(); - using SemaphoreSlim closeStarted = new(0); - using SemaphoreSlim releaseClose = new(0); - DefaultMcpToolHandler handler = new( - null, stub.CreateMessageHandler, clientCacheMaxSize: 1, - clientConnectionDisposer: async connection => - { - closeStarted.Release(); - await releaseClose.WaitAsync(); - await connection.DisposeAsync(); - throw cleanupException; - }); - using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); - await InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); - Task replacement = InvokeScopedAsync(handler, "workflow-b", "ping", timeout.Token); - await closeStarted.WaitAsync(timeout.Token); - - // Act - Task disposal = handler.DisposeAsync().AsTask(); - try - { - Assert.False(disposal.IsCompleted); - } - finally - { - releaseClose.Release(); - } - - // Assert - Assert.Same(cleanupException, await Assert.ThrowsAsync(() => replacement)); - Assert.Same(cleanupException, await Assert.ThrowsAsync(() => disposal)); - Assert.Equal(1, stub.Initializations); - Assert.Equal(1, stub.Terminations); - } - [Theory] [InlineData("ping")] [InlineData(DefaultMcpToolHandler.ListToolsToolName)] @@ -248,416 +104,6 @@ public async Task NoProvider_ReusesSessionUntilHandlerDisposalAsync() transport.Protected().Verify("Dispose", Times.AtLeastOnce(), ItExpr.Is(disposing => disposing)); } - [Fact] - public async Task NoProvider_SeparateWorkflowSessions_UseSeparateCachedSessionsAsync() - { - // Arrange - ProtocolStub stub = new(); - await using DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler); - using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); - - // Act - await InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); - await InvokeScopedAsync(handler, "workflow-b", "ping", timeout.Token); - await InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); - - // Assert - Assert.Equal(2, stub.Initializations); - Assert.Equal(0, stub.Terminations); - } - - [Fact] - public async Task DrainCleanupAsync_Cancellation_DrainsAllCleanupAndCompletionTasksAsync() - { - // Arrange - TaskCompletionSource firstDisposed = new(TaskCreationOptions.RunContinuationsAsynchronously); - TaskCompletionSource secondDisposed = new(TaskCreationOptions.RunContinuationsAsynchronously); - - async Task CancelledCleanupAsync() - { - try - { - await Task.Yield(); - throw new OperationCanceledException("session cleanup cancelled"); - } - finally - { - firstDisposed.TrySetResult(true); - } - } - - async Task SuccessfulCleanupAsync() - { - await Task.Yield(); - secondDisposed.TrySetResult(true); - } - - // Act - await Assert.ThrowsAnyAsync(() => DefaultMcpToolHandler.DrainCleanupAsync( - [CancelledCleanupAsync(), SuccessfulCleanupAsync()], - [firstDisposed.Task, secondDisposed.Task])); - - // Assert - Assert.True(await firstDisposed.Task); - Assert.True(await secondDisposed.Task); - } - - [Fact] - public async Task NoProvider_CacheEviction_DisposesLeastRecentlyUsedSessionAsync() - { - // Arrange - ProtocolStub stub = new(); - await using DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler, clientCacheMaxSize: 2); - using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); - - // Act - await InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); - await InvokeScopedAsync(handler, "workflow-b", "ping", timeout.Token); - await InvokeScopedAsync(handler, "workflow-c", "ping", timeout.Token); - - // Assert - Assert.Equal(3, stub.Initializations); - Assert.Equal(1, stub.Terminations); - } - - [Fact] - public async Task NoProvider_CacheEviction_DisposesUnusedOwnedHttpClientAsync() - { - // Arrange - ProtocolStub stub = new(); - await using DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler, clientCacheMaxSize: 1); - using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); - - // Act - await handler.InvokeToolInWorkflowSessionAsync( - "https://first.example/api", null, "ping", null, null, null, "workflow-a", timeout.Token); - await handler.InvokeToolInWorkflowSessionAsync( - "https://second.example/api", null, "ping", null, null, null, "workflow-b", timeout.Token); - - // Assert - Assert.Equal(2, stub.Initializations); - Assert.Equal(1, stub.Terminations); - Assert.Equal(2, stub.Handlers.Count); - stub.Handlers[0].Protected().Verify( - "Dispose", Times.AtLeastOnce(), ItExpr.Is(disposing => disposing)); - stub.Handlers[1].Protected().Verify( - "Dispose", Times.Never(), ItExpr.Is(disposing => disposing)); - } - - [Fact] - public async Task NoProvider_CacheEviction_SkipsActiveClientUntilInvocationCompletesAsync() - { - // Arrange - ProtocolStub stub = new(); - using SemaphoreSlim firstStarted = new(0); - using SemaphoreSlim releaseFirst = new(0); - int operations = 0; - stub.BeforeOperationAsync = async token => - { - if (Interlocked.Increment(ref operations) == 1) - { - firstStarted.Release(); - await releaseFirst.WaitAsync(token); - } - }; - DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler, clientCacheMaxSize: 2); - using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); - - // Act - Task first = InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); - await firstStarted.WaitAsync(timeout.Token); - await InvokeScopedAsync(handler, "workflow-b", "ping", timeout.Token); - await InvokeScopedAsync(handler, "workflow-c", "ping", timeout.Token); - - // Assert - Assert.Equal(1, stub.Terminations); - Task disposal = handler.DisposeAsync().AsTask(); - Assert.False(disposal.IsCompleted); - releaseFirst.Release(); - await first; - await disposal; - Assert.Equal(3, stub.Terminations); - } - - [Theory] - [InlineData(false, false)] - [InlineData(false, true)] - [InlineData(true, false)] - [InlineData(true, true)] - public async Task NoProvider_HttpClientCleanupFailure_PreservesInvocationOutcomeAsync( - bool activeDisposal, bool failOperation) - { - // Arrange - ProtocolStub stub = new() { FailTransportDisposal = true }; - using SemaphoreSlim firstStarted = new(0); - using SemaphoreSlim releaseFirst = new(0); - int operations = 0; - stub.BeforeOperationAsync = async token => - { - if (activeDisposal && Interlocked.Increment(ref operations) == 1) - { - firstStarted.Release(); - await releaseFirst.WaitAsync(token); - } - }; - DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler, clientCacheMaxSize: 1); - using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); - using CleanupTraceListener listener = new(); - Task? disposal = null; - Trace.Listeners.Add(listener); - try - { - // Act - Task first = handler.InvokeToolInWorkflowSessionAsync( - "https://first.example/api", null, "ping", null, null, null, "workflow-a", timeout.Token); - Task outcome; - if (activeDisposal) - { - await firstStarted.WaitAsync(timeout.Token); - disposal = handler.DisposeAsync().AsTask(); - stub.FailOperation = failOperation; - releaseFirst.Release(); - outcome = first; - } - else - { - await first; - stub.FailOperation = failOperation; - outcome = handler.InvokeToolInWorkflowSessionAsync( - "https://second.example/api", null, "ping", null, null, null, "workflow-b", timeout.Token); - } - - // Assert - if (failOperation) - { - await Assert.ThrowsAsync(() => outcome); - } - else - { - McpServerToolResultContent result = await outcome; - Assert.NotNull(result.Outputs); - Assert.Equal("ok", Assert.IsType(Assert.Single(result.Outputs)).Text); - } - - if (disposal is not null) - { - await disposal; - } - - Assert.Contains("Failed to dispose MCP HTTP client", listener.Output); - Assert.Contains("transport cleanup failed", listener.Output); - } - finally - { - releaseFirst.Release(); - await (disposal ?? handler.DisposeAsync().AsTask()); - Trace.Listeners.Remove(listener); - } - } - - [Fact] - public async Task NoProvider_DisposalDuringActiveInvocation_ReportsCleanupCancellationAsync() - { - // Arrange - OperationCanceledException cleanupException = new("session cleanup cancelled"); - ProtocolStub stub = new(); - using SemaphoreSlim invocationStarted = new(0); - using SemaphoreSlim releaseInvocation = new(0); - stub.BeforeOperationAsync = async token => - { - invocationStarted.Release(); - await releaseInvocation.WaitAsync(token); - }; - DefaultMcpToolHandler handler = new( - null, - stub.CreateMessageHandler, - clientConnectionDisposer: async connection => - { - await connection.DisposeAsync(); - throw cleanupException; - }); - using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); - - // Act - Task invocation = - InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); - await invocationStarted.WaitAsync(timeout.Token); - Task disposal = handler.DisposeAsync().AsTask(); - releaseInvocation.Release(); - - // Assert - Assert.Same(cleanupException, await Assert.ThrowsAsync(() => invocation)); - Assert.Same(cleanupException, await Assert.ThrowsAsync(() => disposal)); - Assert.Equal(1, stub.Terminations); - } - - [Fact] - public async Task NoProvider_ConcurrentWorkflowSessionCreations_AreBoundedByCacheSizeAsync() - { - // Arrange - ProtocolStub stub = new(); - using SemaphoreSlim initializationsStarted = new(0); - using SemaphoreSlim releaseInitializations = new(0); - stub.BeforeInitializationAsync = async token => - { - initializationsStarted.Release(); - await releaseInitializations.WaitAsync(token); - }; - await using DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler, clientCacheMaxSize: 2); - using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); - - // Act - Task first = InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); - await initializationsStarted.WaitAsync(timeout.Token); - Task second = InvokeScopedAsync(handler, "workflow-b", "ping", timeout.Token); - await initializationsStarted.WaitAsync(timeout.Token); - try - { - await Assert.ThrowsAsync( - () => InvokeScopedAsync(handler, "workflow-c", "ping", timeout.Token)); - } - finally - { - releaseInitializations.Release(2); - await Task.WhenAll(first, second); - } - - Task third = InvokeScopedAsync(handler, "workflow-c", "ping", timeout.Token); - await initializationsStarted.WaitAsync(timeout.Token); - releaseInitializations.Release(); - await third; - - // Assert - Assert.Equal(3, stub.Initializations); - Assert.Equal(1, stub.Terminations); - } - - [Fact] - public async Task NoProvider_DisposalDuringFailedCreation_PreservesCreationFailureAsync() - { - // Arrange - ProtocolStub stub = new() { FailInitialization = true }; - using SemaphoreSlim initializationStarted = new(0); - using SemaphoreSlim releaseInitialization = new(0); - stub.BeforeInitializationAsync = async token => - { - initializationStarted.Release(); - await releaseInitialization.WaitAsync(token); - }; - DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler); - using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); - - // Act - Task invocation = - InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); - await initializationStarted.WaitAsync(timeout.Token); - Task disposal = handler.DisposeAsync().AsTask(); - releaseInitialization.Release(); - - // Assert - await Assert.ThrowsAsync(() => invocation); - await disposal; - Assert.Single(stub.Handlers); - stub.Handlers[0].Protected().Verify( - "Dispose", Times.AtLeastOnce(), ItExpr.Is(disposing => disposing)); - } - - [Fact] - public async Task NoProvider_DisposalDuringFailedCreation_ReportsCleanupCancellationAsync() - { - // Arrange - OperationCanceledException cleanupException = new("transport cleanup cancelled"); - ProtocolStub stub = new() { FailInitialization = true, TransportDisposalException = cleanupException }; - using SemaphoreSlim initializationStarted = new(0); - using SemaphoreSlim releaseInitialization = new(0); - stub.BeforeInitializationAsync = async token => - { - initializationStarted.Release(); - await releaseInitialization.WaitAsync(token); - }; - DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler); - using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); - - // Act - Task invocation = - InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); - await initializationStarted.WaitAsync(timeout.Token); - Task disposal = handler.DisposeAsync().AsTask(); - releaseInitialization.Release(); - - // Assert - Assert.Same(cleanupException, await Assert.ThrowsAsync(() => invocation)); - Assert.Same(cleanupException, await Assert.ThrowsAsync(() => disposal)); - Assert.Single(stub.Handlers).Protected().Verify( - "Dispose", Times.AtLeastOnce(), ItExpr.Is(disposing => disposing)); - } - - [Fact] - public async Task NoProvider_DisposalDuringSuccessfulCreation_CleansUpOrphanAndCompletesDisposalAsync() - { - // Arrange - ProtocolStub stub = new(); - using SemaphoreSlim initializationStarted = new(0); - using SemaphoreSlim releaseInitialization = new(0); - stub.BeforeInitializationAsync = async token => - { - initializationStarted.Release(); - await releaseInitialization.WaitAsync(token); - }; - DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler); - using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); - - // Act - Task invocation = - InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); - await initializationStarted.WaitAsync(timeout.Token); - Task disposal = handler.DisposeAsync().AsTask(); - releaseInitialization.Release(); - - // Assert - await Assert.ThrowsAsync(() => invocation); - await disposal; - Assert.Equal(1, stub.Terminations); - Assert.Single(stub.Handlers).Protected().Verify( - "Dispose", Times.AtLeastOnce(), ItExpr.Is(disposing => disposing)); - } - - [Fact] - public async Task NoProvider_DisposalDuringSuccessfulCreation_ReportsOrphanCleanupCancellationAsync() - { - // Arrange - OperationCanceledException cleanupException = new("session cleanup cancelled"); - ProtocolStub stub = new(); - using SemaphoreSlim initializationStarted = new(0); - using SemaphoreSlim releaseInitialization = new(0); - stub.BeforeInitializationAsync = async token => - { - initializationStarted.Release(); - await releaseInitialization.WaitAsync(token); - }; - DefaultMcpToolHandler handler = new( - null, - stub.CreateMessageHandler, - clientConnectionDisposer: async connection => - { - await connection.DisposeAsync(); - throw cleanupException; - }); - using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); - - // Act - Task invocation = - InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); - await initializationStarted.WaitAsync(timeout.Token); - Task disposal = handler.DisposeAsync().AsTask(); - releaseInitialization.Release(); - - // Assert - Assert.Same(cleanupException, await Assert.ThrowsAsync(() => invocation)); - Assert.Same(cleanupException, await Assert.ThrowsAsync(() => disposal)); - Assert.Equal(1, stub.Terminations); - } - [Fact] public async Task NoProvider_DifferentConnectionNames_UseSeparateCachedSessionsAsync() { @@ -739,124 +185,6 @@ public async Task Provider_ConcurrentInvocations_DoNotCoalesceOrSerializeSession Assert.Equal(2, stub.Terminations); } - [Fact] - public async Task NoProvider_ConcurrentSameWorkflowInvocations_CoalesceSessionCreationAsync() - { - // Arrange - ProtocolStub stub = new(); - await using DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler); - using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); - - // Act - await Task.WhenAll( - InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token), - InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token)); - - // Assert - Assert.Equal(1, stub.Initializations); - } - - [Fact] - public async Task NoProvider_CancelledWaiter_DoesNotWaitForSharedCreationAsync() - { - // Arrange - ProtocolStub stub = new(); - using SemaphoreSlim initializationStarted = new(0); - using SemaphoreSlim releaseInitialization = new(0); - stub.BeforeInitializationAsync = async token => - { - initializationStarted.Release(); - await releaseInitialization.WaitAsync(token); - }; - await using DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler); - using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); - using CancellationTokenSource waiterCancellation = new(); - - // Act - Task creator = InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); - await initializationStarted.WaitAsync(timeout.Token); - Task waiter = InvokeScopedAsync(handler, "workflow-a", "ping", waiterCancellation.Token); - await Task.Yield(); - waiterCancellation.Cancel(); - Task completed = await Task.WhenAny(waiter, Task.Delay(TimeSpan.FromSeconds(1), timeout.Token)); - releaseInitialization.Release(); - await creator; - - // Assert - Assert.Same(waiter, completed); - await Assert.ThrowsAnyAsync(() => waiter); - Assert.Equal(1, stub.Initializations); - } - - [Fact] - public async Task NoProvider_CancelledCreator_DoesNotCancelSharedWaiterAsync() - { - // Arrange - ProtocolStub stub = new(); - using SemaphoreSlim initializationStarted = new(0); - int initializationAttempts = 0; - stub.BeforeInitializationAsync = async token => - { - initializationStarted.Release(); - if (Interlocked.Increment(ref initializationAttempts) == 1) - { - await Task.Delay(Timeout.Infinite, token); - } - }; - await using DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler); - using CancellationTokenSource creatorCancellation = new(); - using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); - - // Act - Task creator = - InvokeScopedAsync(handler, "workflow-a", "ping", creatorCancellation.Token); - await initializationStarted.WaitAsync(timeout.Token); - Task waiter = InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); - await Task.Yield(); - creatorCancellation.Cancel(); - await Assert.ThrowsAnyAsync(() => creator); - await initializationStarted.WaitAsync(timeout.Token); - McpServerToolResultContent result = await waiter; - - // Assert - Assert.NotNull(result); - Assert.Equal(2, initializationAttempts); - Assert.Equal(1, stub.Initializations); - } - - [Fact] - public async Task NoProvider_InitializationTimeout_IsSharedWithoutRetryAsync() - { - // Arrange - ProtocolStub stub = new(); - using SemaphoreSlim initializationStarted = new(0); - using SemaphoreSlim releaseInitialization = new(0); - OperationCanceledException timeoutException = new("initialization timed out"); - int initializationAttempts = 0; - stub.BeforeInitializationAsync = async _ => - { - Interlocked.Increment(ref initializationAttempts); - initializationStarted.Release(); - await releaseInitialization.WaitAsync(CancellationToken.None); - throw timeoutException; - }; - await using DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler); - using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); - - // Act - Task creator = InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); - await initializationStarted.WaitAsync(timeout.Token); - Task waiter = InvokeScopedAsync(handler, "workflow-a", "ping", timeout.Token); - await Task.Yield(); - releaseInitialization.Release(); - - // Assert - await Assert.ThrowsAnyAsync(() => creator); - await Assert.ThrowsAnyAsync(() => waiter); - Assert.Equal(1, initializationAttempts); - Assert.Equal(0, stub.Initializations); - } - [Fact] public async Task Provider_OperationFailure_DisposesSessionAndPreservesCallerClientAsync() { @@ -1160,15 +488,13 @@ public async Task Provider_CompletedNestedContext_StillGuardsActiveOuterInvocati Assert.Equal(2, stub.Terminations); } - [Theory] - [InlineData(false)] - [InlineData(true)] - public async Task InitializationAndCleanupFailure_PreservesInitializationErrorAsync(bool hasProvider) + [Fact] + public async Task Provider_InitializationAndCleanupFailure_PreservesInitializationErrorAsync() { // Arrange ProtocolStub stub = new() { FailInitialization = true, FailTransportDisposal = true }; await using DefaultMcpToolHandler handler = new( - hasProvider ? (_, _) => Task.FromResult(null) : null, stub.CreateMessageHandler); + (_, _) => Task.FromResult(null), stub.CreateMessageHandler); using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); using CleanupTraceListener listener = new(); Trace.Listeners.Add(listener); @@ -1178,9 +504,7 @@ public async Task InitializationAndCleanupFailure_PreservesInitializationErrorAs await Assert.ThrowsAsync(() => InvokeAsync(handler, "ping", timeout.Token)); // Assert - Assert.Contains( - hasProvider ? "Failed to dispose MCP transport" : "Failed to dispose MCP HTTP client", listener.Output); - Assert.Contains("transport cleanup failed", listener.Output); + Assert.Contains("Failed to dispose MCP transport", listener.Output); Assert.Single(stub.Handlers).Protected().Verify( "Dispose", Times.AtLeastOnce(), ItExpr.Is(disposing => disposing)); } @@ -1236,11 +560,6 @@ 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; @@ -1250,11 +569,9 @@ private sealed class ProtocolStub public int Terminations => this._terminations; public List> Handlers { get; } = []; public Func? BeforeOperationAsync { get; set; } - public Func? BeforeInitializationAsync { get; set; } public bool FailInitialization { get; set; } public bool FailOperation { get; set; } public bool FailTransportDisposal { get; set; } - public Exception? TransportDisposalException { get; set; } public HttpMessageHandler CreateMessageHandler() { @@ -1262,10 +579,10 @@ public HttpMessageHandler CreateMessageHandler() handler.Protected() .Setup>("SendAsync", ItExpr.IsAny(), ItExpr.IsAny()) .Returns(this.SendAsync); - if (this.FailTransportDisposal || this.TransportDisposalException is not null) + if (this.FailTransportDisposal) { handler.Protected().Setup("Dispose", ItExpr.Is(disposing => disposing)) - .Throws(this.TransportDisposalException ?? new InvalidOperationException("transport cleanup failed")); + .Throws(new InvalidOperationException("transport cleanup failed")); } this.Handlers.Add(handler); @@ -1312,11 +629,6 @@ private async Task SendAsync(HttpRequestMessage request, Ca string? sessionId = null; if (method == "initialize") { - if (this.BeforeInitializationAsync is not null) - { - await this.BeforeInitializationAsync(cancellationToken); - } - if (this.FailInitialization) { return EmptyResponse(HttpStatusCode.BadRequest, request); diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerTests.cs index b6eb32b667d..aa1fb0440aa 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerTests.cs @@ -471,7 +471,7 @@ public void ComputeHeadersHash_DifferentHeaders_ReturnsDifferentHash() // (InvokeToolAsync against a fake server) doesn't surface cache-hit behavior // without standing up a real MCP server — McpClient.CreateAsync fails before // _clients[key] = newClient runs, so nothing ever gets cached. - // Tuple equality on the returned 5-tuple verifies that the dimensions + // Tuple equality on the returned 4-tuple verifies that the dimensions // collectively discriminate cache entries. [Fact] @@ -481,32 +481,19 @@ public void BuildCacheKey_SameInputs_ReturnsEqualKeys() Dictionary headers = new() { ["Authorization"] = "Bearer token" }; // Act - var key1 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", "label", "conn", headers); - var key2 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", "label", "conn", headers); + var key1 = DefaultMcpToolHandler.BuildCacheKey("http://localhost/mcp", "label", "conn", headers); + var key2 = DefaultMcpToolHandler.BuildCacheKey("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("workflow", "http://localhost/mcp", "label", "connection-a", null); - var key2 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", "label", "connection-b", null); + var key1 = DefaultMcpToolHandler.BuildCacheKey("http://localhost/mcp", "label", "connection-a", null); + var key2 = DefaultMcpToolHandler.BuildCacheKey("http://localhost/mcp", "label", "connection-b", null); // Assert Assert.NotEqual(key2, key1); @@ -518,8 +505,8 @@ public void BuildCacheKey_DifferentConnectionName_ReturnsDifferentKeys() public void BuildCacheKey_DifferentServerLabel_ReturnsDifferentKeys() { // Act - var key1 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", "label-a", null, null); - var key2 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", "label-b", null, null); + var key1 = DefaultMcpToolHandler.BuildCacheKey("http://localhost/mcp", "label-a", null, null); + var key2 = DefaultMcpToolHandler.BuildCacheKey("http://localhost/mcp", "label-b", null, null); // Assert Assert.NotEqual(key2, key1); @@ -532,8 +519,8 @@ public void BuildCacheKey_CaseSensitiveUrlPath_ReturnsDifferentKeys() { // Arrange — RFC 3986: URL path is case-sensitive // Act - var key1 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/Tools", null, null, null); - var key2 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/tools", null, null, null); + var key1 = DefaultMcpToolHandler.BuildCacheKey("http://localhost/Tools", null, null, null); + var key2 = DefaultMcpToolHandler.BuildCacheKey("http://localhost/tools", null, null, null); // Assert Assert.NotEqual(key2, key1); @@ -547,8 +534,8 @@ public void BuildCacheKey_HeaderValuesCaseSensitive_ReturnsDifferentKeys() Dictionary headers2 = new() { ["Authorization"] = "Bearer abc" }; // Act - var key1 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", null, null, headers1); - var key2 = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", null, null, headers2); + var key1 = DefaultMcpToolHandler.BuildCacheKey("http://localhost/mcp", null, null, headers1); + var key2 = DefaultMcpToolHandler.BuildCacheKey("http://localhost/mcp", null, null, headers2); // Assert — header value case must propagate into the cache key Assert.NotEqual(key2, key1); @@ -559,7 +546,7 @@ public void BuildCacheKey_HeaderValuesCaseSensitive_ReturnsDifferentKeys() public void BuildCacheKey_NullLabelAndConnection_NormalizesToEmptyString() { // Act - var key = DefaultMcpToolHandler.BuildCacheKey("workflow", "http://localhost/mcp", null, null, null); + var key = DefaultMcpToolHandler.BuildCacheKey("http://localhost/mcp", null, null, null); // Assert — verifies null-safety contract callers rely on Assert.Equal(string.Empty, key.Label); diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/Kit/IWorkflowContextExtensionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/Kit/IWorkflowContextExtensionsTests.cs index 802f71cf48d..d26a1b6a84b 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/Kit/IWorkflowContextExtensionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/Kit/IWorkflowContextExtensionsTests.cs @@ -23,7 +23,7 @@ public async Task FormatTemplateAsync_WithSensitiveValue_ThrowsAsync() WorkflowFormulaState state = new(RecalcEngineFactory.Create()); state.Set("SOME_SECRET", FormulaValue.New("secret-value"), VariableScopeNames.Environment, SensitivityLevel.Sensitive); state.Bind(); - DeclarativeWorkflowContext context = await CreateContextAsync(state); + DeclarativeWorkflowContext context = new(new Mock().Object, state); // Act ValueTask FormatAsync() => context.FormatTemplateAsync("={Env.SOME_SECRET}"); @@ -40,7 +40,7 @@ public async Task FormatTemplateWithSensitivityAsync_WithSensitiveValue_ReturnsS WorkflowFormulaState state = new(RecalcEngineFactory.Create()); state.Set("SOME_SECRET", FormulaValue.New("secret-value"), VariableScopeNames.Environment, SensitivityLevel.Sensitive); state.Bind(); - DeclarativeWorkflowContext context = await CreateContextAsync(state); + DeclarativeWorkflowContext context = new(new Mock().Object, state); // Act EvaluationResult result = await context.FormatTemplateWithSensitivityAsync("={Env.SOME_SECRET}"); @@ -57,7 +57,7 @@ public async Task EvaluateValueAsync_WithSensitiveValue_ThrowsAsync() WorkflowFormulaState state = new(RecalcEngineFactory.Create()); state.Set(SystemScope.Names.LastMessageText, FormulaValue.New("secret-value"), VariableScopeNames.System, SensitivityLevel.Sensitive); state.Bind(); - DeclarativeWorkflowContext context = await CreateContextAsync(state); + DeclarativeWorkflowContext context = new(new Mock().Object, state); // Act ValueTask EvaluateAsync() => context.EvaluateValueAsync("System.LastMessageText"); @@ -74,7 +74,7 @@ public async Task EvaluateValueWithSensitivityAsync_WithSensitiveValue_ReturnsSe WorkflowFormulaState state = new(RecalcEngineFactory.Create()); state.Set(SystemScope.Names.LastMessageText, FormulaValue.New("secret-value"), VariableScopeNames.System, SensitivityLevel.Sensitive); state.Bind(); - DeclarativeWorkflowContext context = await CreateContextAsync(state); + DeclarativeWorkflowContext context = new(new Mock().Object, state); // Act EvaluationResult result = await context.EvaluateValueWithSensitivityAsync("System.LastMessageText"); @@ -91,7 +91,7 @@ public async Task QueueStateUpdateAsync_WithSensitivity_RebindsStateAsync() WorkflowFormulaState state = new(RecalcEngineFactory.Create()); state.Set("TestValue", FormulaValue.New("old-value")); state.Bind(); - DeclarativeWorkflowContext context = await CreateContextAsync(state); + DeclarativeWorkflowContext context = new(new Mock().Object, state); // Act await context.QueueStateUpdateAsync(PropertyPath.Create("Local.TestValue"), FormulaValue.New("new-value"), SensitivityLevel.Sensitive); @@ -113,8 +113,7 @@ public async Task ReadStateWithSensitivityAsync_QueuesSensitiveAssignmentAsync() source .Setup(c => c.ReadStateAsync(SystemScope.Names.LastMessageText, VariableScopeNames.System, default)) .Returns(new ValueTask("secret-value")); - source.As().SetupGet(c => c.SessionId).Returns("test-session"); - DeclarativeWorkflowContext context = await DeclarativeWorkflowContext.CreateAsync(source.Object, state); + DeclarativeWorkflowContext context = new(source.Object, state); // Act var evaluatedValue = await context.ReadStateWithSensitivityAsync(SystemScope.Names.LastMessageText, VariableScopeNames.System); @@ -178,7 +177,7 @@ public async Task GeneratedForeachPattern_WithSensitiveCollection_PreservesItemS VariableScopeNames.Environment, SensitivityLevel.Sensitive); state.Bind(); - DeclarativeWorkflowContext context = await CreateContextAsync(state); + DeclarativeWorkflowContext context = new(new Mock().Object, state); // Act EvaluationResult evaluatedValue = await context.EvaluateValueWithSensitivityAsync("Env.SensitiveItems"); @@ -214,11 +213,4 @@ public async Task ConvertValueWithSensitivityAsync_WithPlainContext_PreservesSen Assert.Equal(42M, result.Value); Assert.Equal(SensitivityLevel.Sensitive, result.Sensitivity); } - - private static async ValueTask CreateContextAsync(WorkflowFormulaState state) - { - Mock source = new(); - source.As().SetupGet(c => c.SessionId).Returns("test-session"); - return await DeclarativeWorkflowContext.CreateAsync(source.Object, state); - } } diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/Kit/RootExecutorTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/Kit/RootExecutorTests.cs index 4d156d8c115..4fedeb6d0f5 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/Kit/RootExecutorTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/Kit/RootExecutorTests.cs @@ -33,13 +33,12 @@ public async Task InitializeEnvironmentAsync_OnlyQueuesAllowedVariablesAsync() }; TestRootExecutor executor = new(options); Mock sourceContext = new(MockBehavior.Strict); - sourceContext.As().SetupGet(c => c.SessionId).Returns("test-session"); sourceContext.Setup(c => c.QueueStateUpdateAsync("ALLOWED", It.IsAny(), VariableScopeNames.Environment, It.IsAny())) .Returns(default(ValueTask)); sourceContext.Setup(c => c.QueueStateUpdateAsync("ALLOWED", SensitivityLevel.Sensitive, WorkflowFormulaState.GetSensitivityScopeName(VariableScopeNames.Environment), It.IsAny())) .Returns(default(ValueTask)); - DeclarativeWorkflowContext context = await DeclarativeWorkflowContext.CreateAsync(sourceContext.Object, executor.Session.State); + DeclarativeWorkflowContext context = new(sourceContext.Object, executor.Session.State); // Act await executor.InitializeAsync(context, "ALLOWED", "HIDDEN"); diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/ObjectModel/InvokeMcpToolExecutorTest.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/ObjectModel/InvokeMcpToolExecutorTest.cs index 5051f3a818c..b3cd7d8c8d1 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/ObjectModel/InvokeMcpToolExecutorTest.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/ObjectModel/InvokeMcpToolExecutorTest.cs @@ -129,86 +129,6 @@ public async Task InvokeMcpToolExecuteWithoutApprovalAsync() await this.ExecuteTestAsync(model); } - [Fact] - public async Task InvokeMcpToolSeparateWorkflowSessionsUseSeparateSessionScopesAsync() - { - // Arrange - this.State.InitializeSystem(); - RecordingScopedMcpToolHandler handler = new(); - InvokeMcpTool model = this.CreateModel( - displayName: nameof(InvokeMcpToolSeparateWorkflowSessionsUseSeparateSessionScopesAsync), - serverUrl: TestServerUrl, - toolName: TestToolName, - requireApproval: false); - MockAgentProvider agentProvider = new(); - InvokeMcpToolExecutor action = new(model, handler, agentProvider.Object, this.State); - Mock firstContext = CreateMockWorkflowContext(); - Mock secondContext = CreateMockWorkflowContext(); - firstContext.As().SetupGet(context => context.SessionId).Returns("workflow-a"); - secondContext.As().SetupGet(context => context.SessionId).Returns("workflow-b"); - - // Act - await action.HandleAsync(new ActionExecutorResult("first"), firstContext.Object, CancellationToken.None); - await action.HandleAsync(new ActionExecutorResult("second"), secondContext.Object, CancellationToken.None); - - // Assert - Assert.Equal(2, handler.WorkflowSessionIds.Count); - Assert.NotEqual(handler.WorkflowSessionIds[0], handler.WorkflowSessionIds[1]); - } - - [Fact] - public async Task InvokeMcpToolSameWorkflowSessionReusesSessionScopeAsync() - { - // Arrange - this.State.InitializeSystem(); - RecordingScopedMcpToolHandler handler = new(); - InvokeMcpTool model = this.CreateModel( - displayName: nameof(InvokeMcpToolSameWorkflowSessionReusesSessionScopeAsync), - serverUrl: TestServerUrl, - toolName: TestToolName, - requireApproval: false); - MockAgentProvider agentProvider = new(); - InvokeMcpToolExecutor action = new(model, handler, agentProvider.Object, this.State); - Mock context = CreateMockWorkflowContext(); - context.As().SetupGet(current => current.SessionId).Returns("workflow-a"); - - // Act - await action.HandleAsync(new ActionExecutorResult("first"), context.Object, CancellationToken.None); - await action.HandleAsync(new ActionExecutorResult("second"), context.Object, CancellationToken.None); - - // Assert - Assert.Equal(2, handler.WorkflowSessionIds.Count); - Assert.Equal(handler.WorkflowSessionIds[0], handler.WorkflowSessionIds[1]); - } - - [Fact] - public async Task InvokeMcpToolLegacyContextsPersistSeparateSessionScopesAsync() - { - // Arrange - this.State.InitializeSystem(); - RecordingScopedMcpToolHandler handler = new(); - InvokeMcpTool model = this.CreateModel( - displayName: nameof(InvokeMcpToolLegacyContextsPersistSeparateSessionScopesAsync), - serverUrl: TestServerUrl, - toolName: TestToolName, - requireApproval: false); - MockAgentProvider agentProvider = new(); - InvokeMcpToolExecutor action = new(model, handler, agentProvider.Object, this.State); - InvokeMcpToolExecutor reconstructedAction = new(model, handler, agentProvider.Object, this.State); - Mock firstContext = CreateMockWorkflowContextWithSessionState(); - Mock secondContext = CreateMockWorkflowContextWithSessionState(); - - // Act - await action.HandleAsync(new ActionExecutorResult("first"), firstContext.Object, CancellationToken.None); - await action.HandleAsync(new ActionExecutorResult("second"), secondContext.Object, CancellationToken.None); - await reconstructedAction.HandleAsync(new ActionExecutorResult("continued"), firstContext.Object, CancellationToken.None); - - // Assert - Assert.Equal(3, handler.WorkflowSessionIds.Count); - Assert.NotEqual(handler.WorkflowSessionIds[0], handler.WorkflowSessionIds[1]); - Assert.Equal(handler.WorkflowSessionIds[0], handler.WorkflowSessionIds[2]); - } - [Fact] public async Task InvokeMcpToolExecuteWithServerLabelAsync() { @@ -355,13 +275,12 @@ public async Task AutoSendUsesMcpApprovalSnapshotAsync(bool autoSend, bool resto nameof(AutoSendUsesMcpApprovalSnapshotAsync), TestServerUrl, TestToolName, requireApproval: true).ToBuilder(); builder.Output = new InvokeToolOutput.Builder { AutoSend = BoolExpression.Expression("Local.Control").ToBuilder() }; InvokeMcpTool model = AssignParent(builder); - RecordingScopedMcpToolHandler handler = new(); + MockMcpToolProvider handler = new(); MockAgentProvider provider = new(); - InvokeMcpToolExecutor action = new(model, handler, provider.Object, this.State); + InvokeMcpToolExecutor action = new(model, handler.Object, provider.Object, this.State); List requests = []; List events = []; Mock context = CreateMockWorkflowContextWithStateStore(requests); - context.As().SetupGet(current => current.SessionId).Returns("workflow-a"); context.Setup(c => c.AddEventAsync(It.IsAny(), It.IsAny())) .Callback((e, _) => events.Add(e)) .Returns(default(ValueTask)); @@ -369,7 +288,7 @@ public async Task AutoSendUsesMcpApprovalSnapshotAsync(bool autoSend, bool resto if (restore) { await InvokeProtectedMethodAsync(action, "OnCheckpointingAsync", context.Object, CancellationToken.None); - action = new InvokeMcpToolExecutor(model, handler, provider.Object, this.State); + action = new InvokeMcpToolExecutor(model, handler.Object, provider.Object, this.State); await InvokeProtectedMethodAsync(action, "OnCheckpointRestoredAsync", context.Object, CancellationToken.None); } this.State.Set("Control", FormulaValue.New(!autoSend)); @@ -380,7 +299,6 @@ public async Task AutoSendUsesMcpApprovalSnapshotAsync(bool autoSend, bool resto // Assert Assert.Equal(autoSend ? 1 : 0, events.OfType().Count()); - Assert.Equal("workflow-a", Assert.Single(handler.WorkflowSessionIds)); } [Fact] @@ -2606,28 +2524,6 @@ 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 @@ -2943,35 +2839,6 @@ public enum McpInvocationValueLocation #region Mock MCP Tool Provider - private sealed class RecordingScopedMcpToolHandler : IWorkflowScopedMcpToolHandler - { - public List WorkflowSessionIds { get; } = []; - - public Task InvokeToolAsync( - string serverUrl, - string? serverLabel, - string toolName, - IDictionary? arguments, - IDictionary? headers, - string? connectionName, - CancellationToken cancellationToken = default) => - throw new InvalidOperationException("The workflow-scoped overload must be used."); - - public Task InvokeToolInWorkflowSessionAsync( - string serverUrl, - string? serverLabel, - string toolName, - IDictionary? arguments, - IDictionary? headers, - string? connectionName, - string workflowSessionId, - CancellationToken cancellationToken = default) - { - this.WorkflowSessionIds.Add(workflowSessionId); - return Task.FromResult(new McpServerToolResultContent("mock-call-id") { Outputs = [] }); - } - } - /// /// Mock implementation of for unit testing purposes. /// diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/PowerFx/WorkflowFormulaStateTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/PowerFx/WorkflowFormulaStateTests.cs index 5d0ffd7660a..34db6f294eb 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/PowerFx/WorkflowFormulaStateTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/PowerFx/WorkflowFormulaStateTests.cs @@ -4,7 +4,6 @@ using System.Collections.Generic; using System.Threading; using System.Threading.Tasks; -using Microsoft.Agents.AI.Workflows.Declarative.Interpreter; using Microsoft.Agents.AI.Workflows.Declarative.PowerFx; using Microsoft.Agents.ObjectModel; using Microsoft.Agents.ObjectModel.Abstractions; @@ -165,34 +164,6 @@ public async Task RestoreAsync_KeepsCaseVariantVariablesIndependentAsync(bool re Assert.Equal(SensitivityLevel.None, this.State.GetSensitivity("ITEM")); } - [Fact] - public async Task DeclarativeContextFallbackSessionId_IsScopedToPersistedRunStateAsync() - { - // Arrange - Dictionary<(string? ScopeName, string Key), string> firstRunState = []; - Dictionary<(string? ScopeName, string Key), string> secondRunState = []; - IWorkflowContext firstContext = CreateContext(firstRunState); - IWorkflowContext restoredContext = CreateContext(firstRunState); - IWorkflowContext secondContext = CreateContext(secondRunState); - - // Act - DeclarativeWorkflowContext first = - await DeclarativeWorkflowContext.CreateAsync(firstContext, this.State); - DeclarativeWorkflowContext continued = - await DeclarativeWorkflowContext.CreateAsync(firstContext, this.State); - DeclarativeWorkflowContext restored = - await DeclarativeWorkflowContext.CreateAsync(restoredContext, this.State); - DeclarativeWorkflowContext second = - await DeclarativeWorkflowContext.CreateAsync(secondContext, this.State); - - // Assert - Assert.Equal(first.SessionId, continued.SessionId); - Assert.Equal(first.SessionId, restored.SessionId); - Assert.NotEqual(first.SessionId, second.SessionId); - Assert.True(firstRunState.ContainsKey((VariableScopeNames.System, "__declarative_mcp_workflow_session_id"))); - Assert.False(firstRunState.ContainsKey((null, "__declarative_mcp_workflow_session_id"))); - } - [Fact] public async Task RestoreAsync_RestoresPersistedSensitivityAsync() { @@ -211,27 +182,4 @@ public async Task RestoreAsync_RestoresPersistedSensitivityAsync() // Assert Assert.Equal(SensitivityLevel.Sensitive, this.State.GetSensitivity("secret")); } - - private static IWorkflowContext CreateContext(Dictionary<(string? ScopeName, string Key), string> state) - { - Mock context = new(); - context - .Setup(current => current.ReadOrInitStateAsync( - It.IsAny(), - It.IsAny>(), - It.IsAny(), - It.IsAny())) - .Returns((string key, Func factory, string? scopeName, CancellationToken cancellationToken) => - { - var scopedKey = (scopeName, key); - if (!state.TryGetValue(scopedKey, out string? value)) - { - value = factory(); - state[scopedKey] = value; - } - - return new ValueTask(value); - }); - return context.Object; - } } diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/RepresentationTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/RepresentationTests.cs index bddef5e4f19..1e92d130de9 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/RepresentationTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/RepresentationTests.cs @@ -8,7 +8,6 @@ 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; @@ -92,48 +91,6 @@ public async Task Test_SpecializedExecutor_InfosAsync() await RunExecutorBindingInfoMatchTestAsync(new RequestInfoExecutor(TestRequestPort)); } - [Fact] - public void SubworkflowSessionId_IsStableAndHierarchicallyScoped() - { - // Arrange - const string ParentSessionId = "parent-session"; - - // Act - string firstChild = SubworkflowBinding.CreateSubworkflowSessionId(ParentSessionId, "first"); - string sameChild = SubworkflowBinding.CreateSubworkflowSessionId(ParentSessionId, "first"); - string siblingChild = SubworkflowBinding.CreateSubworkflowSessionId(ParentSessionId, "second"); - string nestedChild = SubworkflowBinding.CreateSubworkflowSessionId(firstChild, "nested"); - - // Assert - Assert.Equal(firstChild, sameChild); - Assert.NotEqual(ParentSessionId, firstChild); - Assert.NotEqual(firstChild, siblingChild); - Assert.NotEqual(firstChild, nestedChild); - Assert.NotEqual(siblingChild, nestedChild); - } - - [Fact] - public async Task SubworkflowRunner_PreservesLegacyCheckpointSessionIdAsync() - { - // Arrange - const string LegacySessionId = "legacy-session"; - string workflowSessionId = SubworkflowBinding.CreateSubworkflowSessionId(LegacySessionId, "child"); - TestExecutor executor = new(); - Workflow workflow = new WorkflowBuilder(executor).Build(); - - // Act - InProcessRunner runner = InProcessRunner.CreateSubworkflowRunner( - workflow, - checkpointManager: new InMemoryCheckpointManager(), - sessionId: LegacySessionId, - workflowSessionId: workflowSessionId); - - // Assert - Assert.Equal(LegacySessionId, runner.SessionId); - Assert.Equal(workflowSessionId, runner.WorkflowSessionId); - await runner.RequestEndRunAsync(); - } - private static string Source(int id) => $"Source/{id}"; private static string Sink(int id) => $"Sink/{id}";