From bb9e92b189edc1367edc225a08cf8924657e82f8 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 22 Sep 2026 14:08:08 -0400 Subject: [PATCH 01/76] 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/76] .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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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/76] 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}"; From 626b3fb50a3623bc0361bdb26e23932773e3bb62 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 27 Oct 2026 10:43:50 -0400 Subject: [PATCH 59/76] Python: avoid caching unscoped MCP invocations Use per-invocation MCP sessions when the workflow scope is missing or empty, and reject empty scopes at the cache boundary. Cover isolation, real SDK lifecycle, cancellation, and shutdown while preserving scoped cache reuse. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 8e076c3e-5f45-4049-9cf5-f51ed5971ba8 --- python/packages/declarative/AGENTS.md | 7 +- python/packages/declarative/README.md | 10 +- .../_workflows/_mcp_handler.py | 35 +++-- .../tests/test_default_mcp_tool_handler.py | 132 ++++++++++++++++-- .../references/declarative.md | 2 + 5 files changed, 157 insertions(+), 29 deletions(-) diff --git a/python/packages/declarative/AGENTS.md b/python/packages/declarative/AGENTS.md index d6c26bbee8c..03fe1d1db89 100644 --- a/python/packages/declarative/AGENTS.md +++ b/python/packages/declarative/AGENTS.md @@ -23,7 +23,9 @@ series because response caching and default-header integration use its internals ## MCP Handler Lifetimes -`DefaultMCPToolHandler` caches/coalesces sessions only without a `client_provider`. +`DefaultMCPToolHandler` caches/coalesces sessions only without a `client_provider` +and with a non-empty `workflow_session_id`. Missing or empty IDs use the same +per-invocation lifecycle as provider-backed calls, without cache reuse or coalescing. 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. @@ -43,11 +45,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 +For cache-managed scopes, `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. +Provider-backed and unscoped per-invocation sessions do not use this cache bound. Per-invocation sessions intentionally add connection overhead and lose server session continuity between invocations; diff --git a/python/packages/declarative/README.md b/python/packages/declarative/README.md index 428e92755eb..35a0e66418e 100644 --- a/python/packages/declarative/README.md +++ b/python/packages/declarative/README.md @@ -72,7 +72,10 @@ on HTTPX internals. Compatibility tests cover HTTPX 0.27.0 and 0.28.1. ## MCP handler session lifetime -`DefaultMCPToolHandler()` reuses MCP sessions through a bounded LRU cache. +`DefaultMCPToolHandler()` reuses MCP sessions through a bounded LRU cache only +when the invocation has a non-empty `workflow_session_id`. Direct calls with a +missing or empty ID instead create and close a fresh session per invocation, +including `tools/list`; concurrent unscoped calls never share a connection. When configured with `client_provider`, it instead invokes the provider and creates a fresh MCP tool/session for every invocation, including `tools/list`. Sessions are never reused in this mode, even if the provider returns `None` @@ -81,14 +84,15 @@ or the same `httpx.AsyncClient` instance. Each invocation closes its tool/session on success, failure, or cancellation. Internally created fallback HTTP clients are also closed; caller-supplied HTTP clients remain caller-owned and are never closed by the handler. Calling -`aclose()` rejects new invocations and waits for active provider-backed +`aclose()` rejects new invocations and waits for active provider-backed or unscoped invocations to finish cleaning up. Calling it from an active invocation's context (including provider callbacks, inherited child tasks, and cleanup) raises `RuntimeError` before changing handler state, rather than waiting on itself. Close the handler outside that context or after the invocation completes. This intentionally incurs connection and initialization overhead and does not -retain server session state between provider-backed invocations. Applications +retain server session state between provider-backed or unscoped invocations. +`cache_max_size` bounds cache-managed clients, not these per-invocation sessions. Applications that require shared session ownership must implement an explicitly scoped custom `MCPToolHandler`; there is no provider-backed session-cache opt-in. 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 1784a409ce4..ef610d13837 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -149,7 +149,8 @@ class MCPToolInvocation: 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. + workflows from sharing one stateful MCP protocol session. Missing or + empty identifiers use a fresh session that is closed after each invocation. """ server_url: str @@ -253,7 +254,7 @@ def __init__(self, creation_exception: BaseException, cleanup_exception: BaseExc class DefaultMCPToolHandler: """Default :class:`MCPToolHandler` backed by :class:`agent_framework.MCPStreamableHTTPTool`. - Without a ``client_provider``, caches one + Without a ``client_provider`` and with a non-empty workflow session ID, caches one :class:`agent_framework.MCPStreamableHTTPTool` instance per ``(workflow_session_id, server_url, server_label, connection_name, headers_hash)`` in a bounded LRU. The cache prevents re-establishing an MCP @@ -269,6 +270,8 @@ class DefaultMCPToolHandler: headers passed on the wire keep the caller's original casing — so two YAML actions that spell ``Authorization`` differently still share a cache entry. + Invocations without a workflow session ID bypass the cache and close their + tool/session after each invocation, using the same lifecycle as provider-backed calls. Construction modes: @@ -296,12 +299,12 @@ class DefaultMCPToolHandler: Args: client_provider: Optional per-invocation ``httpx.AsyncClient`` provider. - cache_max_size: Maximum number of live MCP clients in no-provider mode, + cache_max_size: Maximum number of live cache-managed MCP clients, 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. + ``32``. Does not limit provider-backed or unscoped invocations. """ LIST_TOOLS_TOOL_NAME: ClassVar[str] = "tools/list" @@ -340,7 +343,7 @@ def __init__( # 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. + # Completion signals only: non-cacheable calls never share entries. self._active_invocations: set[asyncio.Future[None]] = set() # Keep ancestry so a completed nested call cannot hide an active parent # in the context inherited by child tasks, including cleanup tasks. @@ -376,9 +379,10 @@ async def invoke_tool(self, invocation: MCPToolInvocation) -> MCPToolResult: entry: _CacheEntry | None = None completion: asyncio.Future[None] | None = None context_token: Token[tuple[asyncio.Future[None], ...]] | None = None + use_cache = self._client_provider is None and bool(invocation.workflow_session_id) try: try: - if self._client_provider is None: + if use_cache: entry = await self._get_or_create_entry(invocation) else: async with self._cache_lock: @@ -408,7 +412,7 @@ async def invoke_tool(self, invocation: MCPToolInvocation) -> MCPToolResult: return await self._invoke_entry(entry, invocation) finally: - if self._client_provider is None and entry is not None: + if use_cache and entry is not None: await self._release_entry(entry) elif completion is not None: try: @@ -539,12 +543,12 @@ async def _invoke_list_tools(entry: _CacheEntry) -> MCPToolResult: return MCPToolResult(outputs=[Content.from_text(json.dumps(payload, indent=2, allow_nan=False))]) async def aclose(self) -> None: - """Close cached clients and wait for provider-backed invocations to clean up. + """Close cached clients and wait for non-cacheable invocations to clean up. Caller-supplied :class:`httpx.AsyncClient` instances (returned by the ``client_provider`` callback) are NOT closed. - Provider-backed calls already executing may finish; pending connections + Provider-backed and unscoped calls already executing may finish; pending connections are rejected after connecting and cleaned up by their invocation. Concurrent shutdown calls wait for those invocations as well. @@ -555,13 +559,14 @@ async def aclose(self) -> None: their future with ``RuntimeError("DefaultMCPToolHandler is closed")``. Raises: - RuntimeError: If called from an active provider-backed invocation's + RuntimeError: If called from an active provider-backed or unscoped invocation's context, including inherited child tasks and cleanup. Rejected before changing handler state to avoid waiting on itself. """ if any(not completion.done() for completion in self._invocation_context.get()): raise RuntimeError( "DefaultMCPToolHandler.aclose() cannot be called from an active provider-backed invocation" + " or unscoped invocation" ) async with self._cache_lock: if self._shutdown_task is None: @@ -910,8 +915,10 @@ def _header_provider(_kwargs: dict[str, Any]) -> dict[str, str]: try: await self._close_invocation_entry(failed_entry, cleanup_outcome) except BaseException as cleanup_exception: - if self._client_provider is not None or ( - cleanup_outcome.done() and cleanup_outcome.exception() is None + if ( + self._client_provider is not None + or not invocation.workflow_session_id + or (cleanup_outcome.done() and cleanup_outcome.exception() is None) ): raise raise _EntryCreationCleanupFailure(creation_exception, cleanup_exception) from cleanup_exception @@ -988,6 +995,8 @@ def _cache_key( that ``Authorization`` and ``authorization`` map to the same cache entry. Header values remain case-sensitive (per RFC 7235). """ + if not workflow_session_id: + raise ValueError("MCP session caching requires a non-empty workflow session ID") if not headers: headers_hash = "0" else: @@ -995,7 +1004,7 @@ def _cache_key( payload = json.dumps(normalized, ensure_ascii=False) headers_hash = hashlib.sha256(payload.encode("utf-8")).hexdigest() return ( - workflow_session_id or "", + workflow_session_id, server_url, server_label or "", connection_name or "", 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 5ae6a13be37..5e0fb9543bb 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -144,22 +144,123 @@ def _patch_tool() -> Any: def _invocation( - *, server_url: str = "https://mcp.example/api", tool_name: str = "search", **overrides: Any + *, + server_url: str = "https://mcp.example/api", + tool_name: str = "search", + workflow_session_id: str | None = "test-workflow", + **overrides: Any, ) -> MCPToolInvocation: return MCPToolInvocation( server_url=server_url, tool_name=tool_name, + workflow_session_id=workflow_session_id, **overrides, ) +class TestUnscopedInvocations: + @pytest.mark.parametrize("workflow_session_id", [None, ""]) + @pytest.mark.parametrize("tool_name", ["search", "tools/list"]) + async def test_sequential_calls_do_not_reuse_sessions( + self, workflow_session_id: str | None, tool_name: str + ) -> None: + with _patch_tool(): + async with DefaultMCPToolHandler() as handler: + for index in range(1, 3): + result = await handler.invoke_tool( + _invocation( + workflow_session_id=workflow_session_id, tool_name=tool_name, headers={"X-Test": "1"} + ) + ) + assert not result.is_error + assert len(FakeTool.instances) == index + tool = FakeTool.instances[-1] + assert tool.connect_count == 1 + assert tool.close_count == 1 + assert tool._httpx_client is not None + assert tool._httpx_client.is_closed + assert not handler._cache + assert not handler._inflight + assert all(tool.close_count == 1 for tool in FakeTool.instances) + + @pytest.mark.parametrize("workflow_session_id", [None, ""]) + async def test_concurrent_calls_do_not_coalesce_connections(self, workflow_session_id: str | None) -> None: + all_connecting = asyncio.Event() + release = asyncio.Event() + original_connect = FakeTool.connect + + async def gated_connect(tool: FakeTool) -> None: + await original_connect(tool) + if len(FakeTool.instances) == 2: + all_connecting.set() + await release.wait() + + with _patch_tool(), patch.object(FakeTool, "connect", gated_connect): + async with DefaultMCPToolHandler() as handler: + tasks = [ + asyncio.create_task( + handler.invoke_tool(_invocation(workflow_session_id=workflow_session_id, tool_name=name)) + ) + for name in ("search", "tools/list") + ] + try: + await asyncio.wait_for(all_connecting.wait(), timeout=5) + assert len({id(tool.session) for tool in FakeTool.instances}) == 2 + assert all(tool.close_count == 0 for tool in FakeTool.instances) + assert not handler._cache + assert not handler._inflight + finally: + release.set() + results = await asyncio.gather(*tasks) + assert all(not result.is_error for result in results) + assert all(tool.close_count == 1 for tool in FakeTool.instances) + + @pytest.mark.parametrize("workflow_session_id", [None, ""]) + @pytest.mark.parametrize("cancel_active", [False, True]) + async def test_shutdown_waits_for_invocation_cleanup( + self, workflow_session_id: str | None, cancel_active: bool + ) -> None: + started = asyncio.Event() + release = asyncio.Event() + + async def gated_call(_tool: FakeTool, tool_name: str, **arguments: Any) -> Any: + started.set() + await release.wait() + return [Content.from_text("ok")] + + with _patch_tool(), patch.object(FakeTool, "call_tool", gated_call): + handler = DefaultMCPToolHandler() + invocation = asyncio.create_task( + handler.invoke_tool(_invocation(workflow_session_id=workflow_session_id, headers={"X-Test": "1"})) + ) + await started.wait() + shutdown = asyncio.create_task(handler.aclose()) + try: + await asyncio.sleep(0) + assert not shutdown.done() + assert FakeTool.instances[0].close_count == 0 + if cancel_active: + invocation.cancel() + finally: + release.set() + results = await asyncio.gather(invocation, shutdown, return_exceptions=True) + if cancel_active: + assert isinstance(results[0], asyncio.CancelledError) + else: + assert not invocation.result().is_error + shutdown.result() + assert FakeTool.instances[0].close_count == 1 + assert FakeTool.instances[0]._httpx_client is not None + assert FakeTool.instances[0]._httpx_client.is_closed + + # ---------- Provider-backed invocation lifetimes -------------------------- class TestProviderLifetimes: - @pytest.mark.parametrize("returns_none", [False, True]) + @pytest.mark.parametrize("client_source", ["provider", "fallback", "unscoped"]) async def test_actual_sdk_lifecycle_and_hook_cleanup( - self, returns_none: bool, caplog: pytest.LogCaptureFixture + self, client_source: str, caplog: pytest.LogCaptureFixture ) -> None: import agent_framework from agent_framework import MCPStreamableHTTPTool @@ -226,17 +327,23 @@ def create_tool(**kwargs: Any) -> MCPStreamableHTTPTool: tools.append(tool) return tool - caller_client = None if returns_none else RecordingClient() + caller_client = RecordingClient() if client_source == "provider" else None provider = AsyncMock(return_value=caller_client) try: with ( patch.object(httpx, "AsyncClient", RecordingClient), patch.object(agent_framework, "MCPStreamableHTTPTool", side_effect=create_tool), ): - async with DefaultMCPToolHandler(client_provider=provider) as handler: + async with DefaultMCPToolHandler( + client_provider=None if client_source == "unscoped" else provider + ) as handler: for index, tool_name in enumerate(("ok", "tools/list", "ok", "tools/list"), start=1): result = await handler.invoke_tool( - _invocation(tool_name=tool_name, headers={"X-Test": "lifecycle"}) + _invocation( + tool_name=tool_name, + headers={"X-Test": "lifecycle"}, + workflow_session_id=None if client_source == "unscoped" else "test-workflow", + ) ) assert not result.is_error if tool_name == "ok": @@ -244,7 +351,7 @@ def create_tool(**kwargs: Any) -> MCPStreamableHTTPTool: else: assert result.outputs[0].text is not None assert json.loads(result.outputs[0].text)["tools"][0]["name"] == "ok" - assert provider.await_count == index + assert provider.await_count == (0 if client_source == "unscoped" else index) assert len(tools) == index assert len(initialized) == index assert not terminated @@ -256,7 +363,7 @@ def create_tool(**kwargs: Any) -> MCPStreamableHTTPTool: assert tool._lifecycle_owner_task is None assert tool._header_hook_client is None assert all(client.event_hooks["request"] == [caller_hook] for client in clients) - if returns_none: + if caller_client is None: assert len(clients) == index assert all(client.is_closed and client.close_count == 1 for client in clients) else: @@ -624,7 +731,8 @@ async def close(tool: FakeTool) -> None: assert tool._httpx_client.is_closed @pytest.mark.parametrize("connect_fails", [False, True]) - async def test_cleanup_cancellation_signals_completion(self, connect_fails: bool) -> None: + @pytest.mark.parametrize("unscoped", [False, True]) + async def test_cleanup_cancellation_signals_completion(self, connect_fails: bool, unscoped: bool) -> None: original_connect = FakeTool.connect completions: list[asyncio.Future[None]] = [] @@ -643,10 +751,12 @@ async def close(tool: FakeTool) -> None: patch.object(FakeTool, "connect", connect), patch.object(FakeTool, "close", close), ): - handler = DefaultMCPToolHandler(client_provider=AsyncMock(return_value=None)) + handler = DefaultMCPToolHandler(client_provider=None if unscoped else AsyncMock(return_value=None)) # Python 3.10 may drop the message when a cancelled task's result is retrieved again. with pytest.raises(asyncio.CancelledError): - await handler.invoke_tool(_invocation(headers={"X-Test": "1"})) + await handler.invoke_tool( + _invocation(headers={"X-Test": "1"}, workflow_session_id=None if unscoped else "test-workflow") + ) assert handler._invocation_context.get() == () assert not handler._active_invocations assert len(completions) == 1 diff --git a/skills/agent-framework-builder-python/references/declarative.md b/skills/agent-framework-builder-python/references/declarative.md index ce46be674cf..f3005a3f8c0 100644 --- a/skills/agent-framework-builder-python/references/declarative.md +++ b/skills/agent-framework-builder-python/references/declarative.md @@ -24,6 +24,8 @@ maturity. Check current package status and public docstrings before making compa - Allowlist external destinations and register only resources the definition may invoke. - Review the documented HTTP body limits and handler ownership rules; use a custom handler when host policy requires different authentication, retries, SSRF controls, or transport behavior. +- Direct `DefaultMCPToolHandler` calls without a non-empty `workflow_session_id` use a fresh session per invocation; + only explicitly scoped calls without a client provider reuse cached MCP sessions. ## Validation From cdedd6942d74ccb6f3707bf10be9628ed5097a86 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 27 Oct 2026 10:56:08 -0400 Subject: [PATCH 60/76] Python: preserve legacy runner context response handling Make queued-response inspection an optional capability and adapt legacy contexts through their existing non-persisting checkpoint snapshots. Preserve queued-response validation without draining messages or consuming pending requests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 8e076c3e-5f45-4049-9cf5-f51ed5971ba8 --- python/packages/core/AGENTS.md | 3 + .../_workflows/_runner_context.py | 13 ++-- .../agent_framework/_workflows/_workflow.py | 17 ++++- .../core/tests/workflow/test_workflow.py | 72 +++++++++++++++++++ 4 files changed, 99 insertions(+), 6 deletions(-) diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 59460d1ce7e..c4be44f275b 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -341,6 +341,9 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID - **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. +- **Runner context compatibility** - Queued-response inspection is an optional capability, not a required + `RunnerContext` member. Contexts without it use their non-persisting `build_checkpoint` snapshot so response + validation still includes queued replies without draining messages or consuming pending requests. - **`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/_runner_context.py b/python/packages/core/agent_framework/_workflows/_runner_context.py index 82b91558e73..4d1e40bb827 100644 --- a/python/packages/core/agent_framework/_workflows/_runner_context.py +++ b/python/packages/core/agent_framework/_workflows/_runner_context.py @@ -105,6 +105,15 @@ def _get_queued_response_info_events(messages: Mapping[str, Sequence[WorkflowMes ] +@runtime_checkable +class _SupportsQueuedResponseInfoEvents(Protocol): + """Optional fast path for inspecting queued responses without building a snapshot.""" + + async def get_queued_response_info_events(self) -> list[WorkflowEvent[Any]]: + """Get original requests for queued responses without consuming messages.""" + ... + + @runtime_checkable class RunnerContext(Protocol): """Protocol for the execution context used by the runner. @@ -137,10 +146,6 @@ 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. diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 85777174274..c6931a29b6e 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -57,7 +57,12 @@ ) from ._model_utils import DictConvertible from ._runner import RunnerImpl -from ._runner_context import RunnerContext, WorkflowMessage +from ._runner_context import ( + RunnerContext, + WorkflowMessage, + _get_queued_response_info_events, # pyright: ignore[reportPrivateUsage] + _SupportsQueuedResponseInfoEvents, # pyright: ignore[reportPrivateUsage] +) from ._state import State from ._typing_utils import is_instance_of, try_coerce_to_type from ._validation import ValidationTypeEnum, WorkflowValidationError @@ -1166,7 +1171,15 @@ async def _validate_response_requests(self, requests: Sequence[WorkflowEvent]) - async def _get_response_validation_batches( self, requests: Sequence[WorkflowEvent] ) -> list[_ResponseBatchValidation]: - queued_requests = await self._runner.context.get_queued_response_info_events() + context = self._runner.context + if isinstance(context, _SupportsQueuedResponseInfoEvents): + queued_requests = await context.get_queued_response_info_events() + else: + # Legacy contexts expose their queues through non-persisting checkpoint snapshots. + checkpoint = await context.build_checkpoint( + self.name, self.graph_signature_hash, self._runner.state, None, 0 + ) + queued_requests = _get_queued_response_info_events(checkpoint.messages) batch = [*queued_requests, *requests] validations: list[_ResponseBatchValidation] = [] if self._response_batch_validator is not None: diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index 76a285e6d93..70a18dcc640 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/python/packages/core/tests/workflow/test_workflow.py @@ -26,6 +26,7 @@ InProcRunnerContext, Message, ResponseStream, + RunnerContext, WorkflowBuilder, WorkflowCheckpoint, WorkflowCheckpointException, @@ -131,6 +132,77 @@ def test_coerce_request_info_response_rejects_mismatched_type() -> None: _coerce_request_info_response("not-an-int", int, "typed") +class LegacyRunnerContext: + """Delegate the original context contract without the new queue-inspection method.""" + + def __init__(self) -> None: + self.inner = InProcRunnerContext() + + def __getattr__(self, name: str) -> Any: + if name == "get_queued_response_info_events": + raise AttributeError(name) + return getattr(self.inner, name) + + +@pytest.mark.parametrize("streaming", [False, True]) +async def test_response_continuation_supports_legacy_runner_context(streaming: bool) -> None: + context = LegacyRunnerContext() + executor = MockExecutorRequestApproval(id="approver") + workflow = WorkflowBuilder(start_executor=executor).build() + workflow._runner._ctx = cast(RunnerContext, context) + result = await workflow.run(NumberMessage(data=7)) + requests = result.get_request_info_events() + assert len(requests) == 1 + request_id = requests[0].request_id + assert request_id is not None + + if streaming: + events = [event async for event in workflow.run(responses={request_id: ApprovalMessage(True)}, stream=True)] + assert [event.data for event in events if event.type == "output"] == [7] + else: + resumed = await workflow.run(responses={request_id: ApprovalMessage(True)}) + assert resumed.get_outputs() == [7] + assert await context.inner.get_pending_request_info_events() == {} + + +async def test_legacy_runner_context_validates_queued_responses_without_consuming_them() -> None: + class BatchApprover(Executor): + @handler + async def start(self, message: str, ctx: WorkflowContext) -> None: + await ctx.request_info("first", str, request_id="first") + await ctx.request_info("second", str, request_id="second") + + @response_handler + async def respond(self, original_request: str, response: str, ctx: WorkflowContext[Any, str]) -> None: + await ctx.yield_output(f"{original_request}:{response}") + + context = LegacyRunnerContext() + workflow = WorkflowBuilder(start_executor=BatchApprover(id="approver")).build() + workflow._runner._ctx = cast(RunnerContext, context) + await workflow.run("start") + await context.inner.send_request_info_response("first", "yes") + pending_before = await context.inner.get_pending_request_info_events() + validated: list[list[str | None]] = [] + + def validate(requests: Sequence[WorkflowEvent]) -> None: + validated.append([request.request_id for request in requests]) + if len(requests) > 1: + raise ValueError("mixed response batch") + + workflow._response_batch_validator = validate + with pytest.raises(ValueError, match="mixed response batch"): + await workflow.run(responses={"second": "yes"}) + + assert validated == [["first", "second"]] + assert await context.inner.get_pending_request_info_events() == pending_before + queued = await context.inner.get_queued_response_info_events() + assert [request.request_id for request in queued] == ["first"] + + workflow._response_batch_validator = None + resumed = await workflow.run(responses={"second": "yes"}) + assert sorted(resumed.get_outputs()) == ["first:yes", "second:yes"] + + async def test_public_stream_close_immediately_drains_live_executor_without_gc() -> None: started, cancelled, closed, release = (asyncio.Event() for _ in range(4)) cleanup: list[str] = [] From c7bfe991b0ae66aaa4a3d9931ceb70d17b410442 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 27 Oct 2026 11:52:57 -0400 Subject: [PATCH 61/76] fix(tests): narrow runner context before queue inspection Assert the concrete context used by the nested checkpoint test before accessing its optional queue-inspection capability. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 8e076c3e-5f45-4049-9cf5-f51ed5971ba8 --- .../declarative/tests/test_invoke_mcp_tool_executor.py | 5 ++++- 1 file changed, 4 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 6e8f81775a2..bf54344c08b 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -1773,6 +1773,7 @@ async def test_nested_checkpoint_preflight_preserves_live_state( from agent_framework import ( Executor, + InProcRunnerContext, Workflow, WorkflowBuilder, WorkflowCheckpoint, @@ -1897,7 +1898,9 @@ async def restore() -> None: 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() + context = leaf._runner_context + assert isinstance(context, InProcRunnerContext) + queued = await 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 From be2ef000e309f3c8f3a4b1824282c24256b2a1db Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 27 Oct 2026 12:05:03 -0400 Subject: [PATCH 62/76] fix(workflows): preserve agent context on rejected checkpoint restore Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 8e076c3e-5f45-4049-9cf5-f51ed5971ba8 --- .../core/agent_framework/_workflows/_agent.py | 2 - .../tests/workflow/test_workflow_agent.py | 63 +++++++++++++++++++ 2 files changed, 63 insertions(+), 2 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index 32ab247f636..a120fab3b0d 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/python/packages/core/agent_framework/_workflows/_agent.py @@ -537,7 +537,6 @@ async def _run_core( checkpoint_storage=checkpoint_storage, tools=tools, ) - self.workflow._set_agent_run_context(agent_session_id) # pyright: ignore[reportPrivateUsage] async for _ in restore_stream: pass else: @@ -546,7 +545,6 @@ async def _run_core( checkpoint_storage=checkpoint_storage, tools=tools, ) - 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.") diff --git a/python/packages/core/tests/workflow/test_workflow_agent.py b/python/packages/core/tests/workflow/test_workflow_agent.py index 585c79c09d4..1a6a8bf2834 100644 --- a/python/packages/core/tests/workflow/test_workflow_agent.py +++ b/python/packages/core/tests/workflow/test_workflow_agent.py @@ -29,6 +29,7 @@ UsageDetails, WorkflowAgent, WorkflowBuilder, + WorkflowCheckpointException, WorkflowContext, WorkflowEvent, WorkflowException, @@ -246,6 +247,68 @@ async def test_direct_continuations_preserve_agent_session_id(self, streaming: b 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("prior_agent_run", [False, True]) + async def test_rejected_checkpoint_preserves_agent_context(self, streaming: bool, prior_agent_run: bool) -> None: + workflow = WorkflowBuilder(start_executor=SimpleExecutor(id="start", response_text="Accepted")).build() + agent = workflow.as_agent() + storage = InMemoryCheckpointStorage() + await agent.run("Checkpoint turn", session=AgentSession(), checkpoint_storage=storage) + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + checkpoint = await storage.load(checkpoint_id) + assert checkpoint is not None + checkpoint.graph_signature_hash = "different-graph" + await storage.save(checkpoint) + + if prior_agent_run: + await agent.run("Live turn", session=AgentSession()) + else: + await workflow.run([Message(role="user", contents=[Content.from_text("Live turn")])]) + live_state = workflow._runner.state.export_state() + + with pytest.raises(WorkflowCheckpointException, match="graph has changed"): + if streaming: + await agent.run( + session=AgentSession(), + checkpoint_id=checkpoint_id, + checkpoint_storage=storage, + stream=True, + ).get_final_response() + else: + await agent.run(session=AgentSession(), checkpoint_id=checkpoint_id, checkpoint_storage=storage) + + assert workflow._runner.state.export_state() == live_state + + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("fresh_input", [False, True]) + async def test_checkpoint_restore_uses_checkpoint_agent_context(self, streaming: bool, fresh_input: bool) -> None: + workflow = WorkflowBuilder(start_executor=SimpleExecutor(id="start", response_text="Accepted")).build() + agent = workflow.as_agent() + storage = InMemoryCheckpointStorage() + checkpoint_session = AgentSession() + await agent.run("Checkpoint turn", session=checkpoint_session, checkpoint_storage=storage) + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + await agent.run("Live turn", session=AgentSession()) + caller_session = AgentSession() + message = "Fresh turn" if fresh_input else None + + if streaming: + await agent.run( + message, + session=caller_session, + checkpoint_id=checkpoint_id, + checkpoint_storage=storage, + stream=True, + ).get_final_response() + else: + await agent.run(message, session=caller_session, checkpoint_id=checkpoint_id, checkpoint_storage=storage) + + expected_session = caller_session if fresh_input else checkpoint_session + assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == expected_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]) From a9fe090bb1e8949c20a7fa4be96d9ab349f0825d Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 27 Oct 2026 14:46:40 -0400 Subject: [PATCH 63/76] fix(workflows): forward agent context to nested fresh runs Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 8e076c3e-5f45-4049-9cf5-f51ed5971ba8 --- python/packages/core/AGENTS.md | 2 + .../_workflows/_workflow_executor.py | 9 +++- .../tests/test_invoke_mcp_tool_executor.py | 52 +++++++++++++++++++ .../references/declarative.md | 2 + 4 files changed, 64 insertions(+), 1 deletion(-) diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index c4be44f275b..4c22ae2ba8a 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -338,6 +338,8 @@ 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 agent turns** - `WorkflowExecutor` forwards the parent agent-run context after a fresh child run is + accepted, before executing it. Response-only continuations and checkpoint restores keep their saved context. - **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. diff --git a/python/packages/core/agent_framework/_workflows/_workflow_executor.py b/python/packages/core/agent_framework/_workflows/_workflow_executor.py index 780b5b74554..9bdd8b34fc0 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_executor.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_executor.py @@ -1,5 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. +from __future__ import annotations + import logging import sys import types @@ -16,6 +18,8 @@ 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, ) from ._edge_runner import gather_cancelling_siblings_on_error @@ -485,12 +489,15 @@ async def process_workflow(self, input_data: object, ctx: WorkflowContext[Any, A fi_kwargs, ci_kwargs = self._get_child_invocation_kwargs(ctx) # Run the sub-workflow and collect all events, passing parent kwargs - result = await self.workflow.run( + child_run = self.workflow.run( input_data, tools=ctx.get_runtime_tools(), function_invocation_kwargs=fi_kwargs, client_kwargs=ci_kwargs, ) + if ctx.get_state(WORKFLOW_AGENT_RUN_KEY) is True: + self.workflow._set_agent_run_context(ctx.get_state(WORKFLOW_AGENT_SESSION_ID_KEY)) # pyright: ignore[reportPrivateUsage] + result = await child_run logger.debug(f"WorkflowExecutor {self.id} sub-workflow {self.workflow.id} completed with {len(result)} events") 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 bf54344c08b..c47687bc406 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -671,6 +671,58 @@ 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 + @pytest.mark.parametrize("message_list", [False, True]) + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("explicit_session", [False, True]) + @pytest.mark.parametrize("depth", [1, 2]) + async def test_nested_agent_turns_reuse_isolated_workflow_session_ids( + self, streaming: bool, explicit_session: bool, depth: int, message_list: bool + ) -> None: + from agent_framework import Executor, WorkflowBuilder, WorkflowContext, WorkflowExecutor, handler + + class Start(Executor): + @handler + async def forward(self, message: list[Message], ctx: WorkflowContext[list[Message]]) -> None: + await ctx.send_message(message) + + mcp_handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=mcp_handler) + wrappers: list[WorkflowExecutor] = [] + for index in range(2): + child = factory.create_workflow_from_definition(_yaml(_action(tool_name=f"tool_{index}"))) + wrapper = WorkflowExecutor(child, id=f"child_{index}_0") + for level in range(1, depth): + child = WorkflowBuilder(start_executor=wrapper).build() + wrapper = WorkflowExecutor(child, id=f"child_{index}_{level}") + wrappers.append(wrapper) + start = Start(id="start") + parent = WorkflowBuilder(start_executor=start).add_fan_out_edges(start, wrappers).build() + agent = parent.as_agent() + session_a = AgentSession() if explicit_session else None + message = [Message(role="user", contents=[Content.from_text("Agent turn")])] if message_list else "Agent turn" + for session in (session_a, AgentSession(), session_a): + if streaming: + await agent.run(message, session=session, stream=True).get_final_response() + else: + await agent.run(message, session=session) + for _ in range(2): + inputs = [Message(role="user", contents=[Content.from_text("Direct turn")])] + if streaming: + await parent.run(inputs, stream=True).get_final_response() + else: + await parent.run(inputs) + + scopes_by_tool = [ + [invocation.workflow_session_id for invocation in mcp_handler.invocations if invocation.tool_name == tool] + for tool in ("tool_0", "tool_1") + ] + for scopes in scopes_by_tool: + assert len(scopes) == 5 + assert all(scopes) + assert scopes[2] == scopes[0] + assert len({scopes[0], scopes[1], scopes[3], scopes[4]}) == 4 + assert set(scopes_by_tool[0]).isdisjoint(scopes_by_tool[1]) + @pytest.mark.parametrize("message_list", [False, True]) @pytest.mark.parametrize("streaming", [False, True]) @pytest.mark.parametrize("explicit_session", [False, True]) diff --git a/skills/agent-framework-builder-python/references/declarative.md b/skills/agent-framework-builder-python/references/declarative.md index f3005a3f8c0..3966738af49 100644 --- a/skills/agent-framework-builder-python/references/declarative.md +++ b/skills/agent-framework-builder-python/references/declarative.md @@ -26,6 +26,8 @@ maturity. Check current package status and public docstrings before making compa different authentication, retries, SSRF controls, or transport behavior. - Direct `DefaultMCPToolHandler` calls without a non-empty `workflow_session_id` use a fresh session per invocation; only explicitly scoped calls without a client provider reuse cached MCP sessions. +- Declarative children wrapped in `WorkflowExecutor` retain MCP session continuity across parent agent turns. + Each binding must own a unique child workflow instance; its MCP scope remains isolated from sibling bindings. ## Validation From 0d767ce528bf36a4e646d4dc2418698c28282197 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 27 Oct 2026 14:54:21 -0400 Subject: [PATCH 64/76] fix(workflows): normalize postponed workflow annotation Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 8e076c3e-5f45-4049-9cf5-f51ed5971ba8 --- .../core/agent_framework/_workflows/_workflow_executor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/packages/core/agent_framework/_workflows/_workflow_executor.py b/python/packages/core/agent_framework/_workflows/_workflow_executor.py index 9bdd8b34fc0..730e5469b0f 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_executor.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_executor.py @@ -301,7 +301,7 @@ async def handle_subworkflow_request( def __init__( self, - workflow: "Workflow", + workflow: Workflow, id: str, allow_direct_output: bool = False, propagate_request: bool = False, From 7b957d19f0e888053155346c88e8c7446d84423c Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 27 Oct 2026 15:27:04 -0400 Subject: [PATCH 65/76] fix(workflows): preserve child scopes on parent entry replay Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 8e076c3e-5f45-4049-9cf5-f51ed5971ba8 --- ...thon-declarative-request-scope-recovery.md | 20 ++++++ python/packages/core/AGENTS.md | 7 +- .../core/agent_framework/_workflows/_const.py | 5 ++ .../agent_framework/_workflows/_workflow.py | 39 +++++++++++ .../_workflows/_workflow_executor.py | 5 +- python/packages/declarative/AGENTS.md | 3 + .../_workflows/_mcp_handler.py | 12 +++- .../tests/test_invoke_mcp_tool_executor.py | 64 +++++++++++++++++++ .../references/declarative.md | 1 + 9 files changed, 148 insertions(+), 8 deletions(-) diff --git a/docs/decisions/0046-python-declarative-request-scope-recovery.md b/docs/decisions/0046-python-declarative-request-scope-recovery.md index 3eb54cad62f..7dbfb61c973 100644 --- a/docs/decisions/0046-python-declarative-request-scope-recovery.md +++ b/docs/decisions/0046-python-declarative-request-scope-recovery.md @@ -70,5 +70,25 @@ 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. +## Nested entry checkpoint replay + +Selecting a random child scope only when the child first executes is insufficient: +the parent's entry checkpoint predates that selection. Rebuilding the graph and +replaying that checkpoint would select a different child scope. Pre-seeding every +descendant's mutable state before the parent snapshot would instead couple parent +checkpoint creation to child initialization. + +Select a parent scope before its entry checkpoint and derive each child scope from +that scope and the binding ID. Agent turns select a stable scope for their session; +direct fresh runs select new scopes. Declarative response handlers restore the +originating parent scope before further nested invocation. Child scope selection +occurs only after the child accepts a fresh run, not before checkpoint validation +or on response-only continuations. + +Older parent checkpoints without this scope cannot reconstruct a child's original +scope if that child had not started. Such parents receive a new isolated scope +with a warning; subsequent checkpoints retain it. Existing request-origin +validation and explicit legacy-request migration rules are unchanged. + The proposed deciders are Python code owners; maintainer and engineering-management approval is still required. diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 4c22ae2ba8a..ef13d5402cb 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -338,8 +338,11 @@ 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 agent turns** - `WorkflowExecutor` forwards the parent agent-run context after a fresh child run is - accepted, before executing it. Response-only continuations and checkpoint restores keep their saved context. +- **Nested workflow scopes** - Fresh runs select a checkpoint-backed scope before the entry snapshot. + `WorkflowExecutor` derives a child scope from the parent scope and binding ID after the child run is accepted. + Agent turns retain their session scope; direct fresh runs select a new scope. Entry replay therefore preserves + child scopes even when child execution had not started at the snapshot. Response-only continuations restore + their originating scope before further nested invocation. - **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. diff --git a/python/packages/core/agent_framework/_workflows/_const.py b/python/packages/core/agent_framework/_workflows/_const.py index 7f096c57427..5c7e5c7731a 100644 --- a/python/packages/core/agent_framework/_workflows/_const.py +++ b/python/packages/core/agent_framework/_workflows/_const.py @@ -23,6 +23,11 @@ # State key identifying WorkflowAgent turns, including turns without an AgentSession. WORKFLOW_AGENT_RUN_KEY = "_workflow_agent_run" +# Checkpoint-backed workflow scope and the binding-specific scope inherited by a child. +WORKFLOW_SESSION_ID_KEY = "_workflow_session_id" +WORKFLOW_PARENT_SESSION_ID_KEY = "_workflow_parent_session_id" +WORKFLOW_AGENT_SESSION_NAMESPACE_KEY = "_workflow_agent_session_namespace" + # 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 c6931a29b6e..0adb8fc88dc 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -38,7 +38,10 @@ ROUTED_WORKFLOW_RUN_KWARGS_KEY, WORKFLOW_AGENT_RUN_KEY, WORKFLOW_AGENT_SESSION_ID_KEY, + WORKFLOW_AGENT_SESSION_NAMESPACE_KEY, + WORKFLOW_PARENT_SESSION_ID_KEY, WORKFLOW_RUN_KWARGS_KEY, + WORKFLOW_SESSION_ID_KEY, ) from ._edge import ( EdgeGroup, @@ -447,6 +450,40 @@ 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 _initialize_workflow_session_id(self) -> None: + state = self._runner.state + scope = state.get(WORKFLOW_PARENT_SESSION_ID_KEY) + if scope is None: + if state.get(WORKFLOW_AGENT_RUN_KEY) is True: + namespace = state.get(WORKFLOW_AGENT_SESSION_NAMESPACE_KEY) + if namespace is None: + namespace = uuid.uuid4().hex + state.set(WORKFLOW_AGENT_SESSION_NAMESPACE_KEY, namespace) + if not isinstance(namespace, str): + raise ValueError("Invalid workflow agent session namespace.") + scope = uuid.uuid5(uuid.UUID(namespace), json.dumps([state.get(WORKFLOW_AGENT_SESSION_ID_KEY)])).hex + else: + scope = uuid.uuid4().hex + if not isinstance(scope, str) or not scope: + raise ValueError("Invalid workflow session scope.") + state.set(WORKFLOW_SESSION_ID_KEY, scope) + + def _set_parent_run_context(self, parent_state: State, binding_id: str) -> None: + parent_scope = parent_state.get(WORKFLOW_SESSION_ID_KEY) + if parent_scope is None: + logger.warning("Parent workflow has no checkpointed session scope; assigning a new isolated scope.") + parent_scope = uuid.uuid4().hex + parent_state.set(WORKFLOW_SESSION_ID_KEY, parent_scope) + if not isinstance(parent_scope, str) or not parent_scope: + raise ValueError("Invalid parent workflow session scope.") + state = self._runner.state + state.set( + WORKFLOW_PARENT_SESSION_ID_KEY, + uuid.uuid5(uuid.NAMESPACE_URL, json.dumps([parent_scope, binding_id])).hex, + ) + state.set(WORKFLOW_AGENT_RUN_KEY, parent_state.get(WORKFLOW_AGENT_RUN_KEY) is True) + state.set(WORKFLOW_AGENT_SESSION_ID_KEY, parent_state.get(WORKFLOW_AGENT_SESSION_ID_KEY)) + def to_dict(self) -> dict[str, Any]: """Serialize the workflow definition into a JSON-ready dictionary.""" data: dict[str, Any] = { @@ -779,6 +816,7 @@ async def _execute_with_message_or_checkpoint( f"Expected input types: {start_executor.input_types}." ) + self._initialize_workflow_session_id() if self._initial_state_initializer is not None: self._initial_state_initializer(self._runner.state) @@ -895,6 +933,7 @@ def run( if message is not None: self._runner.state.set(WORKFLOW_AGENT_RUN_KEY, False) self._runner.state.set(WORKFLOW_AGENT_SESSION_ID_KEY, None) + self._runner.state.set(WORKFLOW_PARENT_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/agent_framework/_workflows/_workflow_executor.py b/python/packages/core/agent_framework/_workflows/_workflow_executor.py index 730e5469b0f..d4d02a92ff2 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_executor.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_executor.py @@ -18,8 +18,6 @@ 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, ) from ._edge_runner import gather_cancelling_siblings_on_error @@ -495,8 +493,7 @@ async def process_workflow(self, input_data: object, ctx: WorkflowContext[Any, A function_invocation_kwargs=fi_kwargs, client_kwargs=ci_kwargs, ) - if ctx.get_state(WORKFLOW_AGENT_RUN_KEY) is True: - self.workflow._set_agent_run_context(ctx.get_state(WORKFLOW_AGENT_SESSION_ID_KEY)) # pyright: ignore[reportPrivateUsage] + self.workflow._set_parent_run_context(ctx.state, self.id) # pyright: ignore[reportPrivateUsage] result = await child_run logger.debug(f"WorkflowExecutor {self.id} sub-workflow {self.workflow.id} completed with {len(result)} events") diff --git a/python/packages/declarative/AGENTS.md b/python/packages/declarative/AGENTS.md index 03fe1d1db89..a276644fc70 100644 --- a/python/packages/declarative/AGENTS.md +++ b/python/packages/declarative/AGENTS.md @@ -36,6 +36,9 @@ Entry replays. Sessionless WorkflowAgent turns select their own checkpointed sco independent of explicit agent sessions, direct runs, and restored approval scopes. Response-only continuations retain the originating workflow scope even when supplied a new `AgentSession` after checkpoint restoration. +Nested workflow bindings derive their scopes from the checkpointed parent scope and +binding ID, so replaying a parent entry checkpoint preserves child scopes without +sharing them across sibling bindings. 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/agent_framework_declarative/_workflows/_mcp_handler.py b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py index ef610d13837..f2438e77f18 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -40,6 +40,7 @@ from typing import TYPE_CHECKING, Any, ClassVar, Protocol, cast, runtime_checkable import httpx2 as httpx +from agent_framework._workflows._const import WORKFLOW_PARENT_SESSION_ID_KEY, WORKFLOW_SESSION_ID_KEY from ._errors import DeclarativeWorkflowError @@ -66,7 +67,7 @@ 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 agent_session_id is not None and state.get(WORKFLOW_PARENT_SESSION_ID_KEY) is 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) @@ -93,7 +94,12 @@ def get_or_create_workflow_session_id(state: State) -> str: 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: + parent_scope = state.get(WORKFLOW_PARENT_SESSION_ID_KEY) + if parent_scope is not None: + if not isinstance(parent_scope, str) or not parent_scope: + raise ValueError("Invalid parent workflow session scope.") + state.set(_WORKFLOW_SESSION_ID_KEY, parent_scope) + elif 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) @@ -105,6 +111,7 @@ def activate_workflow_session_id(state: State, *, reset_unscoped: bool, is_agent state.set(_WORKFLOW_SESSION_ID_KEY, workflow_session_id) elif reset_unscoped: state.set(_WORKFLOW_SESSION_ID_KEY, uuid.uuid4().hex) + state.set(WORKFLOW_SESSION_ID_KEY, get_or_create_workflow_session_id(state)) def restore_workflow_session_id(state: State, workflow_session_id: str) -> None: @@ -113,6 +120,7 @@ def restore_workflow_session_id(state: State, workflow_session_id: str) -> None: # 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) + state.set(WORKFLOW_SESSION_ID_KEY, workflow_session_id) def get_request_workflow_session_id(workflow_session_id: str | None) -> str: 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 c47687bc406..ecdd19e7cd7 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -671,6 +671,70 @@ 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 + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("run_kind", ["direct", "sessionless_agent", "explicit_agent"]) + @pytest.mark.parametrize("depth", [1, 2]) + async def test_parent_entry_replay_preserves_child_workflow_session_ids( + self, streaming: bool, run_kind: str, depth: int + ) -> None: + from agent_framework import ( + Executor, + InMemoryCheckpointStorage, + WorkflowBuilder, + WorkflowContext, + WorkflowExecutor, + ) + from agent_framework import handler as workflow_handler + + class Start(Executor): + @workflow_handler + async def forward(self, message: list[Message], ctx: WorkflowContext[list[Message]]) -> None: + await ctx.send_message(message) + + mcp_handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=mcp_handler) + storage = InMemoryCheckpointStorage() + + def build(): + wrappers = [] + for index in range(2): + child = factory.create_workflow_from_definition(_yaml(_action(tool_name=f"tool_{index}"))) + wrapper = WorkflowExecutor(child, id=f"child_{index}_0") + for level in range(1, depth): + child = WorkflowBuilder(start_executor=wrapper).build() + wrapper = WorkflowExecutor(child, id=f"child_{index}_{level}") + wrappers.append(wrapper) + start = Start(id="start") + return WorkflowBuilder(name="parent", start_executor=start).add_fan_out_edges(start, wrappers).build() + + parent = build() + messages = [Message(role="user", contents=[Content.from_text("First turn")])] + if run_kind == "direct": + await parent.run(messages, checkpoint_storage=storage) + else: + await parent.as_agent().run( + messages, + session=AgentSession() if run_kind == "explicit_agent" else None, + checkpoint_storage=storage, + ) + original_scopes = { + invocation.tool_name: invocation.workflow_session_id for invocation in mcp_handler.invocations + } + checkpoints = await storage.list_checkpoints(workflow_name=parent.name) + [entry] = [checkpoint for checkpoint in checkpoints if checkpoint.iteration_count == 0] + parent = build() + if streaming: + await parent.run( + checkpoint_id=entry.checkpoint_id, checkpoint_storage=storage, stream=True + ).get_final_response() + else: + await parent.run(checkpoint_id=entry.checkpoint_id, checkpoint_storage=storage) + + assert len(mcp_handler.invocations) == 4 + assert len(set(original_scopes.values())) == 2 + for invocation in mcp_handler.invocations[2:]: + assert invocation.workflow_session_id == original_scopes[invocation.tool_name] + @pytest.mark.parametrize("message_list", [False, True]) @pytest.mark.parametrize("streaming", [False, True]) @pytest.mark.parametrize("explicit_session", [False, True]) diff --git a/skills/agent-framework-builder-python/references/declarative.md b/skills/agent-framework-builder-python/references/declarative.md index 3966738af49..0a228dc71c2 100644 --- a/skills/agent-framework-builder-python/references/declarative.md +++ b/skills/agent-framework-builder-python/references/declarative.md @@ -28,6 +28,7 @@ maturity. Check current package status and public docstrings before making compa only explicitly scoped calls without a client provider reuse cached MCP sessions. - Declarative children wrapped in `WorkflowExecutor` retain MCP session continuity across parent agent turns. Each binding must own a unique child workflow instance; its MCP scope remains isolated from sibling bindings. + Replaying a parent entry checkpoint also retains child scopes when rebuilding the same workflow graph. ## Validation From 8639a7411478728ea02bab908bf04d52c8993f4c Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 27 Oct 2026 15:55:20 -0400 Subject: [PATCH 66/76] fix(workflows): restore parent scopes for child continuations Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 8e076c3e-5f45-4049-9cf5-f51ed5971ba8 --- ...thon-declarative-request-scope-recovery.md | 24 ++++++ python/packages/core/AGENTS.md | 5 +- .../core/agent_framework/_workflows/_const.py | 3 + .../agent_framework/_workflows/_runner.py | 22 +++++- .../agent_framework/_workflows/_workflow.py | 16 +++- .../_workflows/_workflow_executor.py | 74 +++++++++++++++++++ .../core/tests/workflow/test_sub_workflow.py | 59 +++++++++++++++ python/packages/declarative/AGENTS.md | 4 + .../tests/test_invoke_mcp_tool_executor.py | 68 +++++++++++++++++ .../references/declarative.md | 2 + 10 files changed, 273 insertions(+), 4 deletions(-) diff --git a/docs/decisions/0046-python-declarative-request-scope-recovery.md b/docs/decisions/0046-python-declarative-request-scope-recovery.md index 7dbfb61c973..f72ec99db5a 100644 --- a/docs/decisions/0046-python-declarative-request-scope-recovery.md +++ b/docs/decisions/0046-python-declarative-request-scope-recovery.md @@ -90,5 +90,29 @@ scope if that child had not started. Such parents receive a new isolated scope with a warning; subsequent checkpoints retain it. Existing request-origin validation and explicit legacy-request migration rules are unchanged. +## Nested response boundaries + +Restoring only a child's scope leaves its ancestors on a later fresh run's scope. +If runs A and B pause in the same child and B is answered before A, a downstream +binding could therefore select B's MCP session while processing A's continuation. + +Retain the originating parent run context for each forwarded child request in +framework-owned parent shared state. Unlike transient wrapper fields, this state +is checkpointed with its owning workflow, including at every embedded boundary. +After a child successfully resumes or cancels a request, restore the saved parent +context before forwarding outputs or capturing the origins of follow-up requests. +Release completed origins, preserve still-pending origins across new inputs, and +reject mixed-parent-scope response batches before consuming any requests, including +batches spanning sibling bindings. Cancellation batches obey the same parent-scope +rule because cancellation can also produce downstream outputs. +Checkpoint preflight applies the same validation +to queued propagated and wrapped replies before replacing live state. + +Older checkpoints without boundary origins cannot reconstruct those ancestors' +original scopes. A successful child continuation therefore receives a new isolated +parent scope with a warning rather than reusing the latest parent's session. +This does not bypass the child's declarative request-origin checks or explicit +legacy-request migration policy. + The proposed deciders are Python code owners; maintainer and engineering-management approval is still required. diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index ef13d5402cb..55da0745208 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -342,7 +342,10 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID `WorkflowExecutor` derives a child scope from the parent scope and binding ID after the child run is accepted. Agent turns retain their session scope; direct fresh runs select a new scope. Entry replay therefore preserves child scopes even when child execution had not started at the snapshot. Response-only continuations restore - their originating scope before further nested invocation. + their originating scope before further nested invocation. Each forwarded child request retains its parent + run context in checkpointed parent shared state. Responses and cancellations restore that context at every + boundary before forwarding outputs, then release the completed request's origin. Mixed-parent-scope response + batches are rejected before consuming requests, including across sibling bindings. - **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. diff --git a/python/packages/core/agent_framework/_workflows/_const.py b/python/packages/core/agent_framework/_workflows/_const.py index 5c7e5c7731a..e36f22392ed 100644 --- a/python/packages/core/agent_framework/_workflows/_const.py +++ b/python/packages/core/agent_framework/_workflows/_const.py @@ -1,5 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. +from __future__ import annotations + # Default maximum iterations for workflow execution. DEFAULT_MAX_ITERATIONS = 100 @@ -27,6 +29,7 @@ WORKFLOW_SESSION_ID_KEY = "_workflow_session_id" WORKFLOW_PARENT_SESSION_ID_KEY = "_workflow_parent_session_id" WORKFLOW_AGENT_SESSION_NAMESPACE_KEY = "_workflow_agent_session_namespace" +WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY = "_workflow_request_parent_contexts" # Key used to store collision-free, executor-aware workflow invocation kwargs. # WORKFLOW_RUN_KWARGS_KEY remains in the legacy dict format for compatibility with diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index 7703c7f56f3..63cbcb44533 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -1,5 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. +from __future__ import annotations + import asyncio import copy import logging @@ -108,7 +110,7 @@ def reset_iteration_count(self) -> None: """ self._iteration = 0 - async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]: + async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent]: """Run the workflow until no more messages are sent. The runner's checkpoint responsibility is intentionally narrow: it creates a @@ -402,7 +404,11 @@ 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 + from ._const import WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY + from ._workflow_executor import ( + SubWorkflowResponseMessage, + _validate_request_parent_contexts, # pyright: ignore[reportPrivateUsage] + ) if self._graph_signature_hash != checkpoint.graph_signature_hash: raise WorkflowCheckpointException( @@ -417,6 +423,18 @@ def _validate_checkpoint_tree( if request.request_id in checkpoint.pending_request_info_events ), ] + _validate_request_parent_contexts( + checkpoint.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {}), + [ + *requests, + *( + message.data.source_event + for batch in checkpoint.messages.values() + for message in batch + if isinstance(message.data, SubWorkflowResponseMessage) + ), + ], + ) if requests and self._response_batch_validator is not None: self._response_batch_validator(requests) if EXECUTOR_STATE_KEY in checkpoint.state: diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 0adb8fc88dc..8ae3544fb13 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -40,6 +40,7 @@ WORKFLOW_AGENT_SESSION_ID_KEY, WORKFLOW_AGENT_SESSION_NAMESPACE_KEY, WORKFLOW_PARENT_SESSION_ID_KEY, + WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, WORKFLOW_RUN_KWARGS_KEY, WORKFLOW_SESSION_ID_KEY, ) @@ -1210,6 +1211,8 @@ async def _validate_response_requests(self, requests: Sequence[WorkflowEvent]) - async def _get_response_validation_batches( self, requests: Sequence[WorkflowEvent] ) -> list[_ResponseBatchValidation]: + from ._workflow_executor import _validate_request_parent_contexts # pyright: ignore[reportPrivateUsage] + context = self._runner.context if isinstance(context, _SupportsQueuedResponseInfoEvents): queued_requests = await context.get_queued_response_info_events() @@ -1220,7 +1223,12 @@ async def _get_response_validation_batches( ) queued_requests = _get_queued_response_info_events(checkpoint.messages) batch = [*queued_requests, *requests] - validations: list[_ResponseBatchValidation] = [] + parent_contexts = self._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {}) + + def validate_parent_contexts(requests: Sequence[WorkflowEvent]) -> None: + _validate_request_parent_contexts(parent_contexts, requests) + + validations: list[_ResponseBatchValidation] = [(validate_parent_contexts, batch)] if self._response_batch_validator is not None: validations.append((self._validate_response_batch, batch)) for executor in self.executors.values(): @@ -1618,7 +1626,13 @@ async def cancel_pending_requests( raise ValueError("Pending workflow request IDs must be non-empty strings.") async def apply_cancellations() -> None: + from ._workflow_executor import _validate_request_parent_contexts # pyright: ignore[reportPrivateUsage] + pending = await self._runner.context.get_pending_request_info_events() + _validate_request_parent_contexts( + self._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {}), + [request for request_id, request in pending.items() if request_id in selected_ids], + ) computer_executor_ids = { event.source_executor_id for request_id, event in pending.items() diff --git a/python/packages/core/agent_framework/_workflows/_workflow_executor.py b/python/packages/core/agent_framework/_workflows/_workflow_executor.py index d4d02a92ff2..786cbeec086 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_executor.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_executor.py @@ -5,6 +5,7 @@ import logging import sys import types +import uuid from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Any, cast @@ -18,7 +19,12 @@ 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_PARENT_SESSION_ID_KEY, + WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, WORKFLOW_RUN_KWARGS_KEY, + WORKFLOW_SESSION_ID_KEY, ) from ._edge_runner import gather_cancelling_siblings_on_error from ._events import ( @@ -45,6 +51,45 @@ logger = logging.getLogger(__name__) +_PARENT_RUN_CONTEXT_KEYS = ( + WORKFLOW_SESSION_ID_KEY, + WORKFLOW_PARENT_SESSION_ID_KEY, + WORKFLOW_AGENT_RUN_KEY, + WORKFLOW_AGENT_SESSION_ID_KEY, +) + + +def _validate_request_parent_contexts( + contexts: Any, requests: Sequence[WorkflowEvent] = () +) -> dict[str, dict[str, Any]]: + """Preflight checkpointed parent origins before any child responses are consumed.""" + if not isinstance(contexts, dict): + raise ValueError("Invalid parent workflow request contexts.") + contexts = cast(dict[str, dict[str, Any]], contexts) + for request_id, value in contexts.items(): + if not isinstance(request_id, str) or not isinstance(value, dict): + raise ValueError("Invalid parent workflow request context.") + context = cast(dict[str, Any], value) + if set(context) != set(_PARENT_RUN_CONTEXT_KEYS): + raise ValueError("Invalid parent workflow request context.") + scope = context[WORKFLOW_SESSION_ID_KEY] + inherited_scope = context[WORKFLOW_PARENT_SESSION_ID_KEY] + agent_session = context[WORKFLOW_AGENT_SESSION_ID_KEY] + if ( + not isinstance(scope, str) + or not scope + or (inherited_scope is not None and (not isinstance(inherited_scope, str) or not inherited_scope)) + or not isinstance(context[WORKFLOW_AGENT_RUN_KEY], bool) + or (agent_session is not None and not isinstance(agent_session, str)) + ): + raise ValueError("Invalid parent workflow request context.") + scopes = { + contexts[request.request_id][WORKFLOW_SESSION_ID_KEY] for request in requests if request.request_id in contexts + } + if len(scopes) > 1: + raise ValueError("Responses span multiple parent workflow scopes; submit each scope in a separate run.") + return contexts + def _scope_invocation_kwargs_to_subworkflow( invocation_kwargs: WorkflowInvocationKwargs | Mapping[str, Any], @@ -292,6 +337,8 @@ async def handle_subworkflow_request( - Event processing is ordered - outputs are forwarded before requests - Responses are forwarded to the sub-workflow as they arrive; the sub-workflow tracks its own pending requests and resumes when they are answered + - Forwarded requests retain their parent run context in checkpointed parent shared state. + Responses and cancellations restore it before forwarding resumed outputs. - The WorkflowExecutor keeps no per-execution bookkeeping; the sub-workflow is the single source of truth for its pending requests. Starting a new execution while the sub-workflow still has pending requests logs a warning and is only safe when the wrapped workflow is stateless @@ -495,6 +542,7 @@ async def process_workflow(self, input_data: object, ctx: WorkflowContext[Any, A ) self.workflow._set_parent_run_context(ctx.state, self.id) # pyright: ignore[reportPrivateUsage] result = await child_run + self._remove_reemitted_pending_requests(result, pending_requests) logger.debug(f"WorkflowExecutor {self.id} sub-workflow {self.workflow.id} completed with {len(result)} events") @@ -583,6 +631,7 @@ async def _cancel_pending_request(self, request_id: str, ctx: WorkflowContext[An function_invocation_kwargs=fi_kwargs, client_kwargs=ci_kwargs, ) + self._restore_parent_run_context(request_id, ctx) self._remove_reemitted_pending_requests(result, pending_requests) await self._process_workflow_result(result, ctx) @@ -714,6 +763,11 @@ async def _forward_intermediate_output(output: Any) -> None: for event in request_info_events: request_id = event.request_id response_type = event.response_type + contexts = _validate_request_parent_contexts(ctx.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {})) + origin = {key: ctx.state.get(key) for key in _PARENT_RUN_CONTEXT_KEYS} + origin[WORKFLOW_AGENT_RUN_KEY] = origin[WORKFLOW_AGENT_RUN_KEY] is True + contexts.setdefault(request_id, origin) + ctx.state.set(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, contexts) if self._propagate_request: # In a workflow where the parent workflow does not handle the request, the request # should be propagated via the `request_info` mechanism to an external source. And @@ -779,5 +833,25 @@ async def _handle_response( responses={request_id: response}, tools=ctx.get_runtime_tools(), ) + self._restore_parent_run_context(request_id, ctx) self._remove_reemitted_pending_requests(result, pending_requests) await self._process_workflow_result(result, ctx) + + def _restore_parent_run_context(self, request_id: str, ctx: WorkflowContext[Any]) -> None: + contexts = _validate_request_parent_contexts(ctx.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {})) + origin = contexts.pop(request_id, None) + if origin is None: + logger.warning( + "WorkflowExecutor %s request %s has no checkpointed parent origin; assigning a new isolated scope.", + self.id, + request_id, + ) + origin = { + WORKFLOW_SESSION_ID_KEY: uuid.uuid4().hex, + WORKFLOW_PARENT_SESSION_ID_KEY: None, + WORKFLOW_AGENT_RUN_KEY: False, + WORKFLOW_AGENT_SESSION_ID_KEY: None, + } + for key in _PARENT_RUN_CONTEXT_KEYS: + ctx.state.set(key, origin[key]) + ctx.state.set(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, contexts) diff --git a/python/packages/core/tests/workflow/test_sub_workflow.py b/python/packages/core/tests/workflow/test_sub_workflow.py index cfd144b816b..2f937ce2560 100644 --- a/python/packages/core/tests/workflow/test_sub_workflow.py +++ b/python/packages/core/tests/workflow/test_sub_workflow.py @@ -220,6 +220,65 @@ async def test_basic_sub_workflow() -> None: assert parent.result.is_valid is True +@pytest.mark.parametrize("restore", [False, True]) +async def test_interleaved_sibling_responses_reject_mixed_parent_scopes(restore: bool) -> None: + from agent_framework._workflows._const import WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, WORKFLOW_SESSION_ID_KEY + + storage = InMemoryCheckpointStorage() + + def build() -> Workflow: + start = EmailFormatValidator() + children = [ + WorkflowExecutor( + WorkflowBuilder(start_executor=EmailDomainValidator()).build(), + id=name, + propagate_request=True, + allow_direct_output=True, + ) + for name in ("first", "second") + ] + return ( + WorkflowBuilder(start_executor=start, checkpoint_storage=storage).add_fan_out_edges(start, children).build() + ) + + workflow = build() + await workflow.run(EmailValidationRequest("a@first.example")) + requests_a = await workflow._runner_context.get_pending_request_info_events() + scope_a = workflow._runner.state.get(WORKFLOW_SESSION_ID_KEY) + await workflow.run(EmailValidationRequest("b@second.example")) + requests = await workflow._runner_context.get_pending_request_info_events() + requests_b = {key: request for key, request in requests.items() if key not in requests_a} + scope_b = workflow._runner.state.get(WORKFLOW_SESSION_ID_KEY) + if restore: + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + workflow = build() + await workflow.run(checkpoint_id=checkpoint_id) + + first_a = next(iter(requests_a.values())) + second_b = next( + request for request in requests_b.values() if request.source_executor_id != first_a.source_executor_id + ) + with pytest.raises(ValueError, match="multiple parent workflow scopes"): + await workflow.run(responses={first_a.request_id: True, second_b.request_id: True}) + with pytest.raises(ValueError, match="multiple parent workflow scopes"): + await workflow.cancel_pending_requests([first_a.request_id, second_b.request_id]) + assert set(await workflow._runner_context.get_pending_request_info_events()) == set(requests) + assert workflow._runner.state.get(WORKFLOW_SESSION_ID_KEY) == scope_b + + for batch, scope, email in ( + (requests_b, scope_b, "b@second.example"), + (requests_a, scope_a, "a@first.example"), + ): + result = await workflow.run(responses=dict.fromkeys(batch, True)) + outputs = result.get_outputs() + assert len(outputs) == 2 + assert {output.email for output in outputs} == {email} + assert workflow._runner.state.get(WORKFLOW_SESSION_ID_KEY) == scope + assert workflow._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY) == {} + assert not await workflow._runner_context.get_pending_request_info_events() + + async def test_propagated_sub_workflow_response_receives_runtime_tools() -> None: """Runtime tools supplied on parent resume reach the child response handler.""" captured_runtime_tools: list[Any] = [] diff --git a/python/packages/declarative/AGENTS.md b/python/packages/declarative/AGENTS.md index a276644fc70..c729fd02b39 100644 --- a/python/packages/declarative/AGENTS.md +++ b/python/packages/declarative/AGENTS.md @@ -39,6 +39,10 @@ a new `AgentSession` after checkpoint restoration. Nested workflow bindings derive their scopes from the checkpointed parent scope and binding ID, so replaying a parent entry checkpoint preserves child scopes without sharing them across sibling bindings. +Forwarded child requests also checkpoint the originating context at each parent +boundary. Interleaved continuations restore those ancestor scopes before forwarding +outputs to subsequent bindings; responses from different parent scopes must be +submitted separately. 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 ecdd19e7cd7..4195aa76500 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -735,6 +735,74 @@ def build(): for invocation in mcp_handler.invocations[2:]: assert invocation.workflow_session_id == original_scopes[invocation.tool_name] + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("depth", [1, 2]) + @pytest.mark.parametrize("restore", [False, True]) + async def test_resumed_child_restores_parent_scope_before_next_binding( + self, tmp_path: Path, streaming: bool, depth: int, restore: bool + ) -> None: + from agent_framework import Workflow, WorkflowBuilder, WorkflowExecutor + from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY + + mcp_handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=mcp_handler) + storage = _scope_checkpoint_storage(tmp_path) + + def build() -> Workflow: + child = factory.create_workflow_from_definition({ + "name": "first_child", + "actions": [ + _action(tool_name="before"), + _scope_pause_action("Question"), + {"kind": "SendActivity", "activity": {"text": "Completed"}}, + ], + }) + for level in range(1, depth): + nested = WorkflowExecutor(child, id=f"nested_{level}", propagate_request=True, allow_direct_output=True) + child = WorkflowBuilder(start_executor=nested).build() + first = WorkflowExecutor(child, id="first", propagate_request=True) + second = WorkflowExecutor( + factory.create_workflow_from_definition(_yaml(_action(tool_name="after"))), id="second" + ) + return WorkflowBuilder(start_executor=first, checkpoint_storage=storage).add_edge(first, second).build() + + parent = build() + await parent.run({}) + [request_a] = (await parent._runner_context.get_pending_request_info_events()).values() + scope_a = parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) + await parent.run({}) + requests = await parent._runner_context.get_pending_request_info_events() + [request_b] = [request for request in requests.values() if request.request_id != request_a.request_id] + scope_b = parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) + assert scope_a != scope_b + if restore: + checkpoint_id = parent.get_last_checkpoint_id() + assert checkpoint_id is not None + parent = build() + await parent.run(checkpoint_id=checkpoint_id) + + before_rejection = parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) + with pytest.raises(ValueError): + await parent.run(responses={request_a.request_id: 42}) + assert parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) == before_rejection + assert set(await parent._runner_context.get_pending_request_info_events()) == set(requests) + + for request, scope in ((request_b, scope_b), (request_a, scope_a)): + responses = {request.request_id: _scope_pause_response("Question")} + if streaming: + await parent.run(responses=responses, stream=True).get_final_response() + else: + await parent.run(responses=responses) + assert parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) == scope + + scopes = [ + invocation.workflow_session_id for invocation in mcp_handler.invocations if invocation.tool_name == "after" + ] + assert len(scopes) == 4 + assert scopes[0] != scopes[1] + assert scopes[2:] == [scopes[1], scopes[0]] + assert not await parent._runner_context.get_pending_request_info_events() + @pytest.mark.parametrize("message_list", [False, True]) @pytest.mark.parametrize("streaming", [False, True]) @pytest.mark.parametrize("explicit_session", [False, True]) diff --git a/skills/agent-framework-builder-python/references/declarative.md b/skills/agent-framework-builder-python/references/declarative.md index 0a228dc71c2..715532b4525 100644 --- a/skills/agent-framework-builder-python/references/declarative.md +++ b/skills/agent-framework-builder-python/references/declarative.md @@ -29,6 +29,8 @@ maturity. Check current package status and public docstrings before making compa - Declarative children wrapped in `WorkflowExecutor` retain MCP session continuity across parent agent turns. Each binding must own a unique child workflow instance; its MCP scope remains isolated from sibling bindings. Replaying a parent entry checkpoint also retains child scopes when rebuilding the same workflow graph. + Child responses restore the originating parent scope before invoking downstream bindings, including after + checkpoint restoration. Submit responses from different parent scopes in separate runs. ## Validation From e1b150d4c7355ad2781a3894f5d175f8b211e738 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 27 Oct 2026 16:30:19 -0400 Subject: [PATCH 67/76] fix(workflows): preserve core request origins during continuations Capture ordinary request origins and restore them before response handlers and cancellation. Validate response trees without copying arbitrary request payloads, and defer metadata migration until validation succeeds. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 8e076c3e-5f45-4049-9cf5-f51ed5971ba8 --- ...thon-declarative-request-scope-recovery.md | 18 ++++ python/packages/core/AGENTS.md | 14 ++- .../agent_framework/_workflows/_executor.py | 9 +- .../_workflows/_request_context.py | 89 +++++++++++++++++++ .../agent_framework/_workflows/_runner.py | 12 ++- .../agent_framework/_workflows/_workflow.py | 33 ++++--- .../_workflows/_workflow_context.py | 2 + .../_workflows/_workflow_executor.py | 76 ++-------------- .../core/tests/workflow/test_sub_workflow.py | 33 +++++++ python/packages/declarative/AGENTS.md | 4 + .../_workflows/_declarative_builder.py | 13 ++- .../tests/test_invoke_mcp_tool_executor.py | 72 +++++++++++++++ .../references/declarative.md | 1 + 13 files changed, 272 insertions(+), 104 deletions(-) create mode 100644 python/packages/core/agent_framework/_workflows/_request_context.py diff --git a/docs/decisions/0046-python-declarative-request-scope-recovery.md b/docs/decisions/0046-python-declarative-request-scope-recovery.md index f72ec99db5a..0d57536199d 100644 --- a/docs/decisions/0046-python-declarative-request-scope-recovery.md +++ b/docs/decisions/0046-python-declarative-request-scope-recovery.md @@ -114,5 +114,23 @@ parent scope with a warning rather than reusing the latest parent's session. This does not bypass the child's declarative request-origin checks or explicit legacy-request migration policy. +## Ordinary Core requests and live response preflight + +Ordinary `WorkflowContext.request_info()` requests can also pause before a child +binding. Capturing only forwarded child requests leaves these continuations on the +latest fresh run's scope. Capture their framework-owned run metadata in the same +checkpointed state, and restore it before Core invokes the response handler, +including for queued checkpoint replies. Core cancellation handlers restore the +same metadata. Propagated child responses use this path once; wrapped child +messages retain their boundary-specific capture and restoration. + +Deep-copying all live request payloads to preflight migrations imposes an unrelated +serialization requirement on non-checkpointed, in-process workflows. Instead, +response policies validate without mutation and return deferred updates to the +scope metadata they own. Apply these updates only after every policy in the tree +has validated. This preserves cross-child rejection atomicity without copying user +payloads or changing their identity. Stored checkpoint restoration still preflights +its detached candidate tree before replacing live state. + The proposed deciders are Python code owners; maintainer and engineering-management approval is still required. diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 55da0745208..e00ee3589fc 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -342,13 +342,19 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID `WorkflowExecutor` derives a child scope from the parent scope and binding ID after the child run is accepted. Agent turns retain their session scope; direct fresh runs select a new scope. Entry replay therefore preserves child scopes even when child execution had not started at the snapshot. Response-only continuations restore - their originating scope before further nested invocation. Each forwarded child request retains its parent - run context in checkpointed parent shared state. Responses and cancellations restore that context at every - boundary before forwarding outputs, then release the completed request's origin. Mixed-parent-scope response - batches are rejected before consuming requests, including across sibling bindings. + their originating scope before further nested invocation. Every `WorkflowContext.request_info()` request + retains its run context in checkpointed shared state, including ordinary Core requests before a child binding. + Core restores and releases that context before invoking response or cancellation handlers, including queued + checkpoint replies. Wrapped child replies use the same helpers at their message boundary; propagated replies + are restored only once. Mixed-parent-scope response batches are rejected before consuming requests, including + across sibling bindings. - **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. +- **Response policy staging** - Live response validation does not copy request payloads. Policies validate + without mutation and may return deferred metadata updates; all policies in the workflow tree must succeed + before those updates are applied. This preserves in-process non-copyable payloads and atomic legacy-scope + migration without imposing checkpoint serialization requirements on ordinary requests. - **Runner context compatibility** - Queued-response inspection is an optional capability, not a required `RunnerContext` member. Contexts without it use their non-persisting `build_checkpoint` snapshot so response validation still includes queued replies without draining messages or consuming pending requests. diff --git a/python/packages/core/agent_framework/_workflows/_executor.py b/python/packages/core/agent_framework/_workflows/_executor.py index 6a9947a72ba..37421e0dae7 100644 --- a/python/packages/core/agent_framework/_workflows/_executor.py +++ b/python/packages/core/agent_framework/_workflows/_executor.py @@ -1,5 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. +from __future__ import annotations + import asyncio import contextlib import copy @@ -17,6 +19,7 @@ _framework_event_origin, # type: ignore[reportPrivateUsage] ) from ._model_utils import DictConvertible +from ._request_context import restore_request_context from ._request_info_mixin import RequestInfoMixin from ._runner_context import MessageType, RunnerContext, WorkflowMessage from ._state import State @@ -31,7 +34,9 @@ logger = logging.getLogger(__name__) -_ResponseBatchValidation = tuple[Callable[[Sequence[WorkflowEvent]], None], Sequence[WorkflowEvent]] +# Metadata migrations are applied only after every validator in the tree succeeds. +_ResponseBatchValidator = Callable[[Sequence[WorkflowEvent]], Callable[[], None] | None] +_ResponseBatchValidation = tuple[_ResponseBatchValidator, Sequence[WorkflowEvent]] # region Executor @@ -311,6 +316,8 @@ async def execute( invoke_event = WorkflowEvent.executor_invoked(self.id, copy.deepcopy(message)) await context.add_event(invoke_event) try: + if context.request_id is not None: + restore_request_context(state, context.request_id) await handler(message, context) except Exception as exc: # Surface structured executor failure before propagating diff --git a/python/packages/core/agent_framework/_workflows/_request_context.py b/python/packages/core/agent_framework/_workflows/_request_context.py new file mode 100644 index 00000000000..3c9b26dbf51 --- /dev/null +++ b/python/packages/core/agent_framework/_workflows/_request_context.py @@ -0,0 +1,89 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +import logging +import uuid +from collections.abc import Sequence +from typing import Any, cast + +from ._const import ( + WORKFLOW_AGENT_RUN_KEY, + WORKFLOW_AGENT_SESSION_ID_KEY, + WORKFLOW_PARENT_SESSION_ID_KEY, + WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, + WORKFLOW_SESSION_ID_KEY, +) +from ._events import WorkflowEvent +from ._state import State + +logger = logging.getLogger(__name__) + +_RUN_CONTEXT_KEYS = ( + WORKFLOW_SESSION_ID_KEY, + WORKFLOW_PARENT_SESSION_ID_KEY, + WORKFLOW_AGENT_RUN_KEY, + WORKFLOW_AGENT_SESSION_ID_KEY, +) + + +def validate_request_contexts(contexts: Any, requests: Sequence[WorkflowEvent] = ()) -> dict[str, dict[str, Any]]: + """Preflight checkpointed request origins before any responses are consumed.""" + if not isinstance(contexts, dict): + raise ValueError("Invalid parent workflow request contexts.") + contexts = cast(dict[str, dict[str, Any]], contexts) + for request_id, value in contexts.items(): + if not isinstance(request_id, str) or not isinstance(value, dict): + raise ValueError("Invalid parent workflow request context.") + context = cast(dict[str, Any], value) + if set(context) != set(_RUN_CONTEXT_KEYS): + raise ValueError("Invalid parent workflow request context.") + scope = context[WORKFLOW_SESSION_ID_KEY] + inherited_scope = context[WORKFLOW_PARENT_SESSION_ID_KEY] + agent_session = context[WORKFLOW_AGENT_SESSION_ID_KEY] + if ( + not isinstance(scope, str) + or not scope + or (inherited_scope is not None and (not isinstance(inherited_scope, str) or not inherited_scope)) + or not isinstance(context[WORKFLOW_AGENT_RUN_KEY], bool) + or (agent_session is not None and not isinstance(agent_session, str)) + ): + raise ValueError("Invalid parent workflow request context.") + scopes = { + contexts[request.request_id][WORKFLOW_SESSION_ID_KEY] for request in requests if request.request_id in contexts + } + if len(scopes) > 1: + raise ValueError("Responses span multiple parent workflow scopes; submit each scope in a separate run.") + return contexts + + +def capture_request_context(state: State, request_id: str) -> None: + """Associate a request with framework-owned run metadata, without copying its payload.""" + contexts = validate_request_contexts(state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {})) + origin = {key: state.get(key) for key in _RUN_CONTEXT_KEYS} + if origin[WORKFLOW_SESSION_ID_KEY] is None: + logger.warning("Workflow request %s has no active run scope; assigning a new isolated scope.", request_id) + origin[WORKFLOW_SESSION_ID_KEY] = uuid.uuid4().hex + state.set(WORKFLOW_SESSION_ID_KEY, origin[WORKFLOW_SESSION_ID_KEY]) + origin[WORKFLOW_AGENT_RUN_KEY] = origin[WORKFLOW_AGENT_RUN_KEY] is True + contexts.setdefault(request_id, origin) + state.set(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, contexts) + + +def restore_request_context(state: State, request_id: str) -> None: + """Restore and release a request's run metadata before its continuation executes.""" + contexts = validate_request_contexts(state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {})) + origin = contexts.pop(request_id, None) + if origin is None: + logger.warning( + "Workflow request %s has no checkpointed run origin; assigning a new isolated scope.", request_id + ) + origin = { + WORKFLOW_SESSION_ID_KEY: uuid.uuid4().hex, + WORKFLOW_PARENT_SESSION_ID_KEY: None, + WORKFLOW_AGENT_RUN_KEY: False, + WORKFLOW_AGENT_SESSION_ID_KEY: None, + } + for key in _RUN_CONTEXT_KEYS: + state.set(key, origin[key]) + state.set(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, contexts) diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index 63cbcb44533..d70931c440d 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -20,6 +20,7 @@ from ._edge_runner import EdgeRunner, create_edge_runner, gather_cancelling_siblings_on_error from ._events import WorkflowEvent from ._executor import Executor +from ._request_context import validate_request_contexts from ._runner_context import ( RunnerContext, WorkflowMessage, @@ -405,10 +406,7 @@ def _validate_checkpoint_tree( self, checkpoint: WorkflowCheckpoint, delegated_requests: Sequence[WorkflowEvent] = () ) -> None: from ._const import WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY - from ._workflow_executor import ( - SubWorkflowResponseMessage, - _validate_request_parent_contexts, # pyright: ignore[reportPrivateUsage] - ) + from ._workflow_executor import SubWorkflowResponseMessage if self._graph_signature_hash != checkpoint.graph_signature_hash: raise WorkflowCheckpointException( @@ -423,7 +421,9 @@ def _validate_checkpoint_tree( if request.request_id in checkpoint.pending_request_info_events ), ] - _validate_request_parent_contexts( + if requests and self._response_batch_validator is not None: + self._response_batch_validator(requests) + validate_request_contexts( checkpoint.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {}), [ *requests, @@ -435,8 +435,6 @@ 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: diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 8ae3544fb13..faecb5e3c1a 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -13,7 +13,6 @@ 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 @@ -58,8 +57,10 @@ from ._executor import ( Executor, _ResponseBatchValidation, # pyright: ignore[reportPrivateUsage] + _ResponseBatchValidator, # pyright: ignore[reportPrivateUsage] ) from ._model_utils import DictConvertible +from ._request_context import restore_request_context, validate_request_contexts from ._runner import RunnerImpl from ._runner_context import ( RunnerContext, @@ -426,7 +427,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 + self._response_batch_validator: _ResponseBatchValidator | None = None self._runner._response_batch_validator = self._validate_response_batch # pyright: ignore[reportPrivateUsage] self._initial_state_initializer: Callable[[State], None] | None = None @@ -1197,22 +1198,20 @@ async def _restore_and_send_responses( def _validate_response_batch(self, requests: Sequence[WorkflowEvent]) -> None: if self._response_batch_validator is not None: - self._response_batch_validator(requests) + migration = self._response_batch_validator(requests) + if migration is not None: + migration() 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) + migrations = [validator(batch) for validator, batch in validations] + for migration in migrations: + if migration is not None: + migration() async def _get_response_validation_batches( self, requests: Sequence[WorkflowEvent] ) -> list[_ResponseBatchValidation]: - from ._workflow_executor import _validate_request_parent_contexts # pyright: ignore[reportPrivateUsage] - context = self._runner.context if isinstance(context, _SupportsQueuedResponseInfoEvents): queued_requests = await context.get_queued_response_info_events() @@ -1226,11 +1225,12 @@ async def _get_response_validation_batches( parent_contexts = self._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {}) def validate_parent_contexts(requests: Sequence[WorkflowEvent]) -> None: - _validate_request_parent_contexts(parent_contexts, requests) + validate_request_contexts(parent_contexts, requests) - validations: list[_ResponseBatchValidation] = [(validate_parent_contexts, batch)] + validations: list[_ResponseBatchValidation] = [] if self._response_batch_validator is not None: - validations.append((self._validate_response_batch, batch)) + validations.append((self._response_batch_validator, batch)) + validations.append((validate_parent_contexts, batch)) for executor in self.executors.values(): delegated = [request for request in batch if request.source_executor_id == executor.id] if delegated: @@ -1626,10 +1626,8 @@ async def cancel_pending_requests( raise ValueError("Pending workflow request IDs must be non-empty strings.") async def apply_cancellations() -> None: - from ._workflow_executor import _validate_request_parent_contexts # pyright: ignore[reportPrivateUsage] - pending = await self._runner.context.get_pending_request_info_events() - _validate_request_parent_contexts( + validate_request_contexts( self._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {}), [request for request_id, request in pending.items() if request_id in selected_ids], ) @@ -1658,6 +1656,7 @@ async def apply_cancellations() -> None: state=self._runner.state, runner_context=self._runner.context, ) + restore_request_context(self._runner.state, request_id) await executor._cancel_pending_request(request_id, context) # pyright: ignore[reportPrivateUsage] if checkpoint_storage is not None: diff --git a/python/packages/core/agent_framework/_workflows/_workflow_context.py b/python/packages/core/agent_framework/_workflows/_workflow_context.py index 0a887c20d9d..4d96f81d94f 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_context.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_context.py @@ -20,6 +20,7 @@ WorkflowEventSource, _framework_event_origin, # type: ignore ) +from ._request_context import capture_request_context from ._runner_context import RunnerContext, WorkflowMessage from ._state import State from ._typing_utils import contains_typevar @@ -433,6 +434,7 @@ async def request_info(self, request_data: object, response_type: type, *, reque response_type=response_type, ) await self._runner_context.add_request_info_event(request_info_event) + capture_request_context(self._state, request_info_event.request_id) def get_state(self, key: str, default: Any = None) -> Any: """Get a value from the workflow state.""" diff --git a/python/packages/core/agent_framework/_workflows/_workflow_executor.py b/python/packages/core/agent_framework/_workflows/_workflow_executor.py index 786cbeec086..da292a4d5db 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_executor.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_executor.py @@ -5,7 +5,6 @@ import logging import sys import types -import uuid from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Any, cast @@ -19,12 +18,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_PARENT_SESSION_ID_KEY, - WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, WORKFLOW_RUN_KWARGS_KEY, - WORKFLOW_SESSION_ID_KEY, ) from ._edge_runner import gather_cancelling_siblings_on_error from ._events import ( @@ -37,6 +31,7 @@ _ResponseBatchValidation, # pyright: ignore[reportPrivateUsage] handler, ) +from ._request_context import capture_request_context, restore_request_context from ._request_info_mixin import response_handler from ._runner_context import WorkflowMessage from ._typing_utils import is_instance_of @@ -51,45 +46,6 @@ logger = logging.getLogger(__name__) -_PARENT_RUN_CONTEXT_KEYS = ( - WORKFLOW_SESSION_ID_KEY, - WORKFLOW_PARENT_SESSION_ID_KEY, - WORKFLOW_AGENT_RUN_KEY, - WORKFLOW_AGENT_SESSION_ID_KEY, -) - - -def _validate_request_parent_contexts( - contexts: Any, requests: Sequence[WorkflowEvent] = () -) -> dict[str, dict[str, Any]]: - """Preflight checkpointed parent origins before any child responses are consumed.""" - if not isinstance(contexts, dict): - raise ValueError("Invalid parent workflow request contexts.") - contexts = cast(dict[str, dict[str, Any]], contexts) - for request_id, value in contexts.items(): - if not isinstance(request_id, str) or not isinstance(value, dict): - raise ValueError("Invalid parent workflow request context.") - context = cast(dict[str, Any], value) - if set(context) != set(_PARENT_RUN_CONTEXT_KEYS): - raise ValueError("Invalid parent workflow request context.") - scope = context[WORKFLOW_SESSION_ID_KEY] - inherited_scope = context[WORKFLOW_PARENT_SESSION_ID_KEY] - agent_session = context[WORKFLOW_AGENT_SESSION_ID_KEY] - if ( - not isinstance(scope, str) - or not scope - or (inherited_scope is not None and (not isinstance(inherited_scope, str) or not inherited_scope)) - or not isinstance(context[WORKFLOW_AGENT_RUN_KEY], bool) - or (agent_session is not None and not isinstance(agent_session, str)) - ): - raise ValueError("Invalid parent workflow request context.") - scopes = { - contexts[request.request_id][WORKFLOW_SESSION_ID_KEY] for request in requests if request.request_id in contexts - } - if len(scopes) > 1: - raise ValueError("Responses span multiple parent workflow scopes; submit each scope in a separate run.") - return contexts - def _scope_invocation_kwargs_to_subworkflow( invocation_kwargs: WorkflowInvocationKwargs | Mapping[str, Any], @@ -631,7 +587,6 @@ async def _cancel_pending_request(self, request_id: str, ctx: WorkflowContext[An function_invocation_kwargs=fi_kwargs, client_kwargs=ci_kwargs, ) - self._restore_parent_run_context(request_id, ctx) self._remove_reemitted_pending_requests(result, pending_requests) await self._process_workflow_result(result, ctx) @@ -763,11 +718,6 @@ async def _forward_intermediate_output(output: Any) -> None: for event in request_info_events: request_id = event.request_id response_type = event.response_type - contexts = _validate_request_parent_contexts(ctx.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {})) - origin = {key: ctx.state.get(key) for key in _PARENT_RUN_CONTEXT_KEYS} - origin[WORKFLOW_AGENT_RUN_KEY] = origin[WORKFLOW_AGENT_RUN_KEY] is True - contexts.setdefault(request_id, origin) - ctx.state.set(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, contexts) if self._propagate_request: # In a workflow where the parent workflow does not handle the request, the request # should be propagated via the `request_info` mechanism to an external source. And @@ -776,6 +726,7 @@ async def _forward_intermediate_output(output: Any) -> None: else: # In a workflow where the parent workflow has an executor that may intercept the # request and handle it directly, a message should be sent. + capture_request_context(ctx.state, request_id) await ctx.send_message(SubWorkflowRequestMessage(source_event=event, executor_id=self.id)) # Handle final state @@ -833,25 +784,8 @@ async def _handle_response( responses={request_id: response}, tools=ctx.get_runtime_tools(), ) - self._restore_parent_run_context(request_id, ctx) + # Propagated replies already restored their context before their response handler. + if ctx.request_id != request_id: + restore_request_context(ctx.state, request_id) self._remove_reemitted_pending_requests(result, pending_requests) await self._process_workflow_result(result, ctx) - - def _restore_parent_run_context(self, request_id: str, ctx: WorkflowContext[Any]) -> None: - contexts = _validate_request_parent_contexts(ctx.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {})) - origin = contexts.pop(request_id, None) - if origin is None: - logger.warning( - "WorkflowExecutor %s request %s has no checkpointed parent origin; assigning a new isolated scope.", - self.id, - request_id, - ) - origin = { - WORKFLOW_SESSION_ID_KEY: uuid.uuid4().hex, - WORKFLOW_PARENT_SESSION_ID_KEY: None, - WORKFLOW_AGENT_RUN_KEY: False, - WORKFLOW_AGENT_SESSION_ID_KEY: None, - } - for key in _PARENT_RUN_CONTEXT_KEYS: - ctx.state.set(key, origin[key]) - ctx.state.set(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, contexts) diff --git a/python/packages/core/tests/workflow/test_sub_workflow.py b/python/packages/core/tests/workflow/test_sub_workflow.py index 2f937ce2560..474992706ec 100644 --- a/python/packages/core/tests/workflow/test_sub_workflow.py +++ b/python/packages/core/tests/workflow/test_sub_workflow.py @@ -279,6 +279,39 @@ def build() -> Workflow: assert not await workflow._runner_context.get_pending_request_info_events() +@pytest.mark.parametrize("streaming", [False, True]) +async def test_response_validation_preserves_noncopyable_request_payload(streaming: bool) -> None: + from threading import Lock + + payload: dict[str, Any] = {"lock": Lock()} + + class Requester(Executor): + @handler + async def start(self, message: str, ctx: WorkflowContext) -> None: + await ctx.request_info(payload, bool) + + @response_handler + async def respond( + self, + original: dict[str, Any], + approved: bool, + ctx: WorkflowContext[Never, str], # type: ignore[valid-type] + ) -> None: + assert original is payload + assert original["lock"] is payload["lock"] + await ctx.yield_output("approved" if approved else "denied") + + workflow = WorkflowBuilder(start_executor=Requester(id="requester")).build() + await workflow.run("start") + [request_id] = (await workflow._runner_context.get_pending_request_info_events()).keys() + if streaming: + result = await workflow.run(responses={request_id: True}, stream=True).get_final_response() + else: + result = await workflow.run(responses={request_id: True}) + assert result.get_outputs() == ["approved"] + assert not await workflow._runner_context.get_pending_request_info_events() + + async def test_propagated_sub_workflow_response_receives_runtime_tools() -> None: """Runtime tools supplied on parent resume reach the child response handler.""" captured_runtime_tools: list[Any] = [] diff --git a/python/packages/declarative/AGENTS.md b/python/packages/declarative/AGENTS.md index c729fd02b39..d6dae130d1b 100644 --- a/python/packages/declarative/AGENTS.md +++ b/python/packages/declarative/AGENTS.md @@ -43,6 +43,8 @@ Forwarded child requests also checkpoint the originating context at each parent boundary. Interleaved continuations restore those ancestor scopes before forwarding outputs to subsequent bindings; responses from different parent scopes must be submitted separately. +Ordinary Core `ctx.request_info()` requests before a child binding retain and +restore their origins too, including when a queued response is replayed. 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 @@ -75,6 +77,8 @@ 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. +Response policies stage only the legacy scope-field updates until all nested +policies pass; live request payloads are not deep-copied for preflight. 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 9ed8200fd12..a298170f3b9 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_builder.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_builder.py @@ -14,7 +14,7 @@ from __future__ import annotations import logging -from collections.abc import Sequence +from collections.abc import Callable, Sequence from functools import partial from typing import Any, cast @@ -122,7 +122,7 @@ def _validate_response_batch( requests: Sequence[WorkflowEvent[Any]], *, migrate_legacy_request_scopes: bool = False -) -> None: +) -> Callable[[], None]: scoped_requests: list[ tuple[ExternalInputRequest | MCPToolApprovalRequest | ToolApprovalRequest | AgentExternalInputRequest, str] ] = [] @@ -141,8 +141,13 @@ def _validate_response_batch( 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 apply_migration() -> None: + for data, scope in scoped_requests: + if data.workflow_session_id is None: + data.workflow_session_id = scope + + return apply_migration def _initialize_mcp_run_state(state: State) -> None: 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 4195aa76500..1e4faef66a1 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -183,6 +183,78 @@ 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.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("restore_kind", ["memory", "file", "queued"]) + async def test_core_parent_request_restores_origin_before_child( + self, tmp_path: Path, streaming: bool, restore_kind: str + ) -> None: + from agent_framework import ( + Executor, + Workflow, + WorkflowBuilder, + WorkflowContext, + WorkflowExecutor, + response_handler, + ) + from agent_framework import handler as workflow_handler + from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY + + mcp_handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=mcp_handler) + storage = _scope_checkpoint_storage(tmp_path) + + class Requester(Executor): + @workflow_handler + async def start(self, message: str, ctx: WorkflowContext) -> None: + await ctx.request_info(message, str) + + @response_handler + async def respond(self, original: str, response: str, ctx: WorkflowContext[dict[str, Any]]) -> None: + await ctx.send_message({"original": original, "response": response}) + + def build() -> Workflow: + start = Requester(id="requester") + child = WorkflowExecutor(factory.create_workflow_from_definition(_yaml(_action())), id="child") + return WorkflowBuilder(start_executor=start, checkpoint_storage=storage).add_edge(start, child).build() + + parent = build() + await parent.run("A") + [request_a] = (await parent._runner_context.get_pending_request_info_events()).values() + scope_a = parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) + await parent.run("B") + requests = await parent._runner_context.get_pending_request_info_events() + [request_b] = [request for request in requests.values() if request.request_id != request_a.request_id] + if restore_kind == "file": + checkpoint_id = parent.get_last_checkpoint_id() + assert checkpoint_id is not None + parent = build() + await parent.run(checkpoint_id=checkpoint_id) + + with pytest.raises(ValueError, match="multiple parent workflow scopes"): + await parent.run(responses={request_a.request_id: "continue", request_b.request_id: "continue"}) + assert set(await parent._runner_context.get_pending_request_info_events()) == set(requests) + assert mcp_handler.call_count == 0 + await parent.run(responses={request_b.request_id: "continue"}) + if restore_kind == "queued": + await parent._runner_context.send_request_info_response(request_a.request_id, "continue") + await parent._runner.create_checkpoint_if_enabled() + checkpoint_id = parent.get_last_checkpoint_id() + assert checkpoint_id is not None + parent = build() + if streaming: + await parent.run(checkpoint_id=checkpoint_id, stream=True).get_final_response() + else: + await parent.run(checkpoint_id=checkpoint_id) + elif streaming: + await parent.run(responses={request_a.request_id: "continue"}, stream=True).get_final_response() + else: + await parent.run(responses={request_a.request_id: "continue"}) + + assert parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) == scope_a + assert len(mcp_handler.invocations) == 2 + assert mcp_handler.invocations[0].workflow_session_id != mcp_handler.invocations[1].workflow_session_id + assert not await parent._runner_context.get_pending_request_info_events() + @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( diff --git a/skills/agent-framework-builder-python/references/declarative.md b/skills/agent-framework-builder-python/references/declarative.md index 715532b4525..0f938d9f00e 100644 --- a/skills/agent-framework-builder-python/references/declarative.md +++ b/skills/agent-framework-builder-python/references/declarative.md @@ -31,6 +31,7 @@ maturity. Check current package status and public docstrings before making compa Replaying a parent entry checkpoint also retains child scopes when rebuilding the same workflow graph. Child responses restore the originating parent scope before invoking downstream bindings, including after checkpoint restoration. Submit responses from different parent scopes in separate runs. + The same origin recovery applies to ordinary Core `ctx.request_info()` pauses before invoking a child. ## Validation From 16b16558019eba94216a02de29af6fb57d7512ba Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 27 Oct 2026 17:59:26 -0400 Subject: [PATCH 68/76] fix(workflows): preflight wrapped replies and queued cancellations Validate graph-routed wrapped replies before response consumption and each live dispatch batch. Include queued replies in cancellation preflight while preserving optional runner-context inspection and deferred metadata updates. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 8e076c3e-5f45-4049-9cf5-f51ed5971ba8 --- ...thon-declarative-request-scope-recovery.md | 9 ++ python/packages/core/AGENTS.md | 6 +- .../agent_framework/_workflows/_runner.py | 62 ++++++--- .../_workflows/_runner_context.py | 12 +- .../agent_framework/_workflows/_workflow.py | 33 +++-- .../core/tests/workflow/test_checkpoint.py | 5 +- .../core/tests/workflow/test_sub_workflow.py | 127 ++++++++++++++++++ .../core/tests/workflow/test_workflow.py | 71 +++++++++- .../references/declarative.md | 2 + 9 files changed, 279 insertions(+), 48 deletions(-) diff --git a/docs/decisions/0046-python-declarative-request-scope-recovery.md b/docs/decisions/0046-python-declarative-request-scope-recovery.md index 0d57536199d..854d53552a6 100644 --- a/docs/decisions/0046-python-declarative-request-scope-recovery.md +++ b/docs/decisions/0046-python-declarative-request-scope-recovery.md @@ -124,6 +124,15 @@ including for queued checkpoint replies. Core cancellation handlers restore the same metadata. Propagated child responses use this path once; wrapped child messages retain their boundary-specific capture and restoration. +Preflight includes both propagated responses and wrapped child replies already +queued for delivery. Route wrapped replies through the parent graph to validate +their child's own policy without treating child scopes as parent-policy scopes. +Repeat validation before each live message batch is drained, so a rejected batch +retains its queued replies. Cancellation validates queued replies together with +the complete selected cancellation set before consuming requests or invoking +cancellation handlers. Legacy runner contexts use a non-persisting snapshot when +optional queued-message inspection is unavailable. + Deep-copying all live request payloads to preflight migrations imposes an unrelated serialization requirement on non-checkpointed, in-process workflows. Instead, response policies validate without mutation and return deferred updates to the diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index e00ee3589fc..4d69ac2e4bc 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -347,7 +347,9 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID Core restores and releases that context before invoking response or cancellation handlers, including queued checkpoint replies. Wrapped child replies use the same helpers at their message boundary; propagated replies are restored only once. Mixed-parent-scope response batches are rejected before consuming requests, including - across sibling bindings. + across sibling bindings. Validation includes queued wrapped replies and runs before each live message + batch is drained, routing wrapped replies to their child's own policy. Cancellation preflight combines + queued replies with the complete selected cancellation set before consuming any requests. - **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. @@ -355,7 +357,7 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID without mutation and may return deferred metadata updates; all policies in the workflow tree must succeed before those updates are applied. This preserves in-process non-copyable payloads and atomic legacy-scope migration without imposing checkpoint serialization requirements on ordinary requests. -- **Runner context compatibility** - Queued-response inspection is an optional capability, not a required +- **Runner context compatibility** - Queued-message inspection is an optional capability, not a required `RunnerContext` member. Contexts without it use their non-persisting `build_checkpoint` snapshot so response validation still includes queued replies without draining messages or consuming pending requests. - **`WorkflowBuilder`** - Fluent API for building workflows, including explicit diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index d70931c440d..61648f9fc4d 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -7,7 +7,7 @@ import logging import warnings from collections import defaultdict -from collections.abc import AsyncGenerator, Callable, Sequence +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence from typing import TYPE_CHECKING, Any, cast from ..exceptions import ( @@ -25,6 +25,7 @@ RunnerContext, WorkflowMessage, _get_queued_response_info_events, # pyright: ignore[reportPrivateUsage] + _SupportsQueuedMessages, # pyright: ignore[reportPrivateUsage] ) from ._state import State @@ -89,6 +90,7 @@ def __init__( # Checkpointing related attributes self._previous_checkpoint_id: CheckpointID | None = None self._response_batch_validator: Callable[[Sequence[WorkflowEvent]], None] | None = None + self._queued_response_validator: Callable[[], Awaitable[None]] | None = None @property def context(self) -> RunnerContext: @@ -235,6 +237,8 @@ async def _deliver_messages_for_edge_runner(edge_runner: EdgeRunner) -> None: tasks = [_deliver_messages_for_edge_runner(edge_runner) for edge_runner in associated_edge_runners] await gather_cancelling_siblings_on_error(*tasks) + if self._queued_response_validator is not None: + await self._queued_response_validator() message_batches = await self._ctx.drain_messages() await gather_cancelling_siblings_on_error( *( @@ -243,6 +247,37 @@ async def _deliver_messages_for_edge_runner(edge_runner: EdgeRunner) -> None: ) ) + async def _get_queued_messages(self) -> Mapping[str, Sequence[WorkflowMessage]]: + if isinstance(self._ctx, _SupportsQueuedMessages): + return await self._ctx.get_queued_messages() + # Legacy contexts expose their queues through non-persisting checkpoint snapshots. + checkpoint = await self._ctx.build_checkpoint( + self._workflow_name, self._graph_signature_hash, self._state, None, self._iteration + ) + return checkpoint.messages + + def _get_wrapped_response_requests( + self, messages: Mapping[str, Sequence[WorkflowMessage]] + ) -> dict[str, list[WorkflowEvent]]: + from ._workflow_executor import SubWorkflowResponseMessage + + requests: defaultdict[str, list[WorkflowEvent]] = defaultdict(list) + for batch in messages.values(): + for message in batch: + if not isinstance(message.data, SubWorkflowResponseMessage): + continue + # Wrapped replies retain child-local source IDs; route through the parent graph. + 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: + requests[target_id].append(message.data.source_event) + return dict(requests) + async def _prepare_checkpoint_state(self) -> None: """Persist executor snapshots into committed shared state. @@ -406,7 +441,6 @@ def _validate_checkpoint_tree( self, checkpoint: WorkflowCheckpoint, delegated_requests: Sequence[WorkflowEvent] = () ) -> None: from ._const import WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY - from ._workflow_executor import SubWorkflowResponseMessage if self._graph_signature_hash != checkpoint.graph_signature_hash: raise WorkflowCheckpointException( @@ -423,36 +457,20 @@ def _validate_checkpoint_tree( ] if requests and self._response_batch_validator is not None: self._response_batch_validator(requests) + wrapped_requests = self._get_wrapped_response_requests(checkpoint.messages) validate_request_contexts( checkpoint.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {}), [ *requests, - *( - message.data.source_event - for batch in checkpoint.messages.values() - for message in batch - if isinstance(message.data, SubWorkflowResponseMessage) - ), + *(request for batch in wrapped_requests.values() for request in batch), ], ) 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 target_id, batch in wrapped_requests.items(): + delegated_by_executor[target_id].extend(batch) for _, executor, state in self._get_executor_states(checkpoint.state[EXECUTOR_STATE_KEY]): executor._validate_checkpoint_restore( # pyright: ignore[reportPrivateUsage] state, delegated_by_executor[executor.id] diff --git a/python/packages/core/agent_framework/_workflows/_runner_context.py b/python/packages/core/agent_framework/_workflows/_runner_context.py index 4d1e40bb827..71bc526d77c 100644 --- a/python/packages/core/agent_framework/_workflows/_runner_context.py +++ b/python/packages/core/agent_framework/_workflows/_runner_context.py @@ -106,11 +106,11 @@ def _get_queued_response_info_events(messages: Mapping[str, Sequence[WorkflowMes @runtime_checkable -class _SupportsQueuedResponseInfoEvents(Protocol): - """Optional fast path for inspecting queued responses without building a snapshot.""" +class _SupportsQueuedMessages(Protocol): + """Optional fast path for inspecting queued messages without building a snapshot.""" - async def get_queued_response_info_events(self) -> list[WorkflowEvent[Any]]: - """Get original requests for queued responses without consuming messages.""" + async def get_queued_messages(self) -> Mapping[str, Sequence[WorkflowMessage]]: + """Inspect queued messages without consuming them or copying their payloads.""" ... @@ -633,6 +633,10 @@ 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 get_queued_messages(self) -> dict[str, list[WorkflowMessage]]: + """Inspect queued messages without consuming them or copying their payloads.""" + return {source_id: list(messages) for source_id, messages in self._messages.items()} + 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 faecb5e3c1a..0807cf792e7 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -66,7 +66,6 @@ RunnerContext, WorkflowMessage, _get_queued_response_info_events, # pyright: ignore[reportPrivateUsage] - _SupportsQueuedResponseInfoEvents, # pyright: ignore[reportPrivateUsage] ) from ._state import State from ._typing_utils import is_instance_of, try_coerce_to_type @@ -429,6 +428,9 @@ def __init__( self._active_run: weakref.ref[ResponseStream[WorkflowEvent, WorkflowRunResult]] | None = None self._response_batch_validator: _ResponseBatchValidator | None = None self._runner._response_batch_validator = self._validate_response_batch # pyright: ignore[reportPrivateUsage] + self._runner._queued_response_validator = functools.partial( # pyright: ignore[reportPrivateUsage] + self._validate_response_requests, () + ) self._initial_state_initializer: Callable[[State], None] | None = None # Run-scoped pause checkpoint bookkeeping (owned by Workflow, not callers). @@ -1212,27 +1214,25 @@ async def _validate_response_requests(self, requests: Sequence[WorkflowEvent]) - async def _get_response_validation_batches( self, requests: Sequence[WorkflowEvent] ) -> list[_ResponseBatchValidation]: - context = self._runner.context - if isinstance(context, _SupportsQueuedResponseInfoEvents): - queued_requests = await context.get_queued_response_info_events() - else: - # Legacy contexts expose their queues through non-persisting checkpoint snapshots. - checkpoint = await context.build_checkpoint( - self.name, self.graph_signature_hash, self._runner.state, None, 0 - ) - queued_requests = _get_queued_response_info_events(checkpoint.messages) + messages = await self._runner._get_queued_messages() # pyright: ignore[reportPrivateUsage] + queued_requests = _get_queued_response_info_events(messages) + wrapped_requests = self._runner._get_wrapped_response_requests(messages) # pyright: ignore[reportPrivateUsage] batch = [*queued_requests, *requests] + origin_requests = [*batch, *(request for wrapped in wrapped_requests.values() for request in wrapped)] parent_contexts = self._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {}) - def validate_parent_contexts(requests: Sequence[WorkflowEvent]) -> None: - validate_request_contexts(parent_contexts, requests) + def validate_parent_contexts(_requests: Sequence[WorkflowEvent]) -> None: + validate_request_contexts(parent_contexts, origin_requests) validations: list[_ResponseBatchValidation] = [] if self._response_batch_validator is not None: validations.append((self._response_batch_validator, batch)) validations.append((validate_parent_contexts, batch)) for executor in self.executors.values(): - delegated = [request for request in batch if request.source_executor_id == executor.id] + delegated = [ + *(request for request in batch if request.source_executor_id == executor.id), + *wrapped_requests.get(executor.id, []), + ] if delegated: validations.extend( await executor._get_response_validation_batches(delegated) # pyright: ignore[reportPrivateUsage] @@ -1627,10 +1627,6 @@ async def cancel_pending_requests( async def apply_cancellations() -> None: pending = await self._runner.context.get_pending_request_info_events() - validate_request_contexts( - self._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {}), - [request for request_id, request in pending.items() if request_id in selected_ids], - ) computer_executor_ids = { event.source_executor_id for request_id, event in pending.items() @@ -1645,6 +1641,9 @@ async def apply_cancellations() -> None: for request_id, event in pending.items() if event.source_executor_id in computer_executor_ids ) + await self._validate_response_requests([ + request for request_id, request in pending.items() if request_id in selected_ids + ]) cancelled_events = await self._runner.context.cancel_request_info_events(selected_ids) for request_id, request_event in cancelled_events.items(): source_executor_id = request_event.source_executor_id diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index 5611c6cf90a..f23ec064b44 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -717,8 +717,9 @@ def _build_workflow() -> 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" + # Validate both the detached checkpoint candidate and the live dispatch batch. + assert [[event.request_id for event in batch] for batch in validated_requests] == [[request_id], [request_id]] + assert all(batch[0].data == "approve" for batch in validated_requests) async def test_memory_checkpoint_storage_roundtrip_json_native_types(): diff --git a/python/packages/core/tests/workflow/test_sub_workflow.py b/python/packages/core/tests/workflow/test_sub_workflow.py index 474992706ec..e7b03e395b3 100644 --- a/python/packages/core/tests/workflow/test_sub_workflow.py +++ b/python/packages/core/tests/workflow/test_sub_workflow.py @@ -2,6 +2,7 @@ import asyncio import logging +from collections.abc import Sequence from dataclasses import dataclass, field from typing import Any from uuid import uuid4 @@ -279,6 +280,132 @@ def build() -> Workflow: assert not await workflow._runner_context.get_pending_request_info_events() +class _WrappedReplyCoordinator(Executor): + def __init__(self) -> None: + super().__init__(id="coordinator") + self.requests: dict[str, SubWorkflowRequestMessage] = {} + self.outputs: list[tuple[str, str]] = [] + + @handler + async def start(self, email: str, ctx: WorkflowContext[DomainCheckRequest]) -> None: + await ctx.send_message(DomainCheckRequest(domain=email.split("@")[1], email=email)) + + @handler + async def intercept(self, request: SubWorkflowRequestMessage, ctx: WorkflowContext) -> None: + email = request.source_event.data.email + self.requests[email] = request + await ctx.request_info(email, bool, request_id=email) + + @response_handler + async def respond(self, email: str, approved: bool, ctx: WorkflowContext[SubWorkflowResponseMessage]) -> None: + request = self.requests[email] + await ctx.send_message(request.create_response(approved), target_id=request.executor_id) + + @handler + async def queue_replies( + self, replies: list[SubWorkflowResponseMessage], ctx: WorkflowContext[SubWorkflowResponseMessage] + ) -> None: + for reply in replies: + await ctx.send_message(reply) + + @handler + async def collect(self, result: ValidationResult, ctx: WorkflowContext) -> None: + from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY + + self.outputs.append((result.email, ctx.get_state(WORKFLOW_SESSION_ID_KEY))) + + +def _build_wrapped_reply_workflow() -> tuple[Workflow, Workflow, _WrappedReplyCoordinator]: + coordinator = _WrappedReplyCoordinator() + child = WorkflowBuilder(start_executor=EmailDomainValidator()).build() + wrapper = WorkflowExecutor(child, id="child") + parent = ( + WorkflowBuilder(start_executor=coordinator) + .add_edge(coordinator, wrapper) + .add_edge(wrapper, coordinator) + .build() + ) + return parent, child, coordinator + + +@pytest.mark.parametrize("streaming", [False, True]) +@pytest.mark.parametrize("delivery", ["live", "queued_and_new"]) +async def test_wrapped_reply_batches_reject_mixed_parent_origins(delivery: str, streaming: bool) -> None: + from agent_framework._workflows._const import WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, WORKFLOW_SESSION_ID_KEY + + parent, child, coordinator = _build_wrapped_reply_workflow() + emails = ["a@first.example", "b@second.example"] + scopes: list[str] = [] + for email in emails: + await parent.run(email) + scopes.append(parent._runner.state.get(WORKFLOW_SESSION_ID_KEY)) + child_pending = await child._runner_context.get_pending_request_info_events() + parent_pending = await parent._runner_context.get_pending_request_info_events() + origins = parent._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY) + replies = [coordinator.requests[email].create_response(True) for email in emails] + if delivery == "queued_and_new": + await parent._runner_context.send_message( + WorkflowMessage(data=replies[0], source_id=coordinator.id, target_id="child") + ) + + with pytest.raises(ValueError, match="multiple parent workflow scopes"): + message = replies if delivery == "live" else None + responses = {emails[1]: True} if delivery == "queued_and_new" else None + if streaming: + _ = [event async for event in parent.run(message, responses=responses, stream=True)] + else: + await parent.run(message, responses=responses) + + assert coordinator.outputs == [] + assert await child._runner_context.get_pending_request_info_events() == child_pending + assert await parent._runner_context.get_pending_request_info_events() == parent_pending + assert parent._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY) == origins + queued = await parent._runner_context.drain_messages() + assert [message.data.source_event.request_id for batch in queued.values() for message in batch] == [ + reply.source_event.request_id for reply in replies[: 1 if delivery == "queued_and_new" else 2] + ] + + for email, scope in reversed(list(zip(emails, scopes, strict=True))): + await parent.run(responses={email: True}) + assert coordinator.outputs[-1] == (email, scope) + assert not await child._runner_context.get_pending_request_info_events() + assert not await parent._runner_context.get_pending_request_info_events() + assert parent._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY) == {} + + +@pytest.mark.parametrize("delivery", ["live", "queued_and_new"]) +async def test_wrapped_reply_preflight_applies_child_policy(delivery: str) -> None: + parent, child, coordinator = _build_wrapped_reply_workflow() + email = "a@first.example" + await parent.run(email) + parent_pending = await parent._runner_context.get_pending_request_info_events() + child_pending = await child._runner_context.get_pending_request_info_events() + reply = coordinator.requests[email].create_response(True) + + def reject(requests: Sequence[WorkflowEvent]) -> None: + if requests: + raise ValueError("child reply policy") + + child._response_batch_validator = reject + if delivery == "queued_and_new": + await parent._runner_context.send_message( + WorkflowMessage(data=reply, source_id=coordinator.id, target_id="child") + ) + run = parent.run(responses={email: True}) + else: + run = parent.run([reply]) + with pytest.raises(ValueError, match="child reply policy"): + await run + + assert coordinator.outputs == [] + assert await parent._runner_context.get_pending_request_info_events() == parent_pending + assert await child._runner_context.get_pending_request_info_events() == child_pending + queued = await parent._runner_context.drain_messages() + assert [message.data.source_event.request_id for batch in queued.values() for message in batch] == [ + reply.source_event.request_id + ] + + @pytest.mark.parametrize("streaming", [False, True]) async def test_response_validation_preserves_noncopyable_request_payload(streaming: bool) -> None: from threading import Lock diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index 70a18dcc640..99d75f83dd5 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/python/packages/core/tests/workflow/test_workflow.py @@ -139,7 +139,7 @@ def __init__(self) -> None: self.inner = InProcRunnerContext() def __getattr__(self, name: str) -> Any: - if name == "get_queued_response_info_events": + if name in {"get_queued_response_info_events", "get_queued_messages"}: raise AttributeError(name) return getattr(self.inner, name) @@ -203,6 +203,75 @@ def validate(requests: Sequence[WorkflowEvent]) -> None: assert sorted(resumed.get_outputs()) == ["first:yes", "second:yes"] +@pytest.mark.parametrize("legacy", [False, True]) +@pytest.mark.parametrize("mixed", [False, True]) +async def test_cancellation_preflight_checks_queued_reply_origins(legacy: bool, mixed: bool) -> None: + from agent_framework._workflows._const import WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, WORKFLOW_SESSION_ID_KEY + + cancelled: list[str] = [] + handled: list[str] = [] + delivered: list[tuple[str, str]] = [] + + class Requester(Executor): + @handler + async def start(self, message: str, ctx: WorkflowContext) -> None: + for request_id in ("a", "b") if message == "same" else (message,): + await ctx.request_info(request_id, str, request_id=request_id) + + @response_handler + async def respond(self, request: str, response: str, ctx: WorkflowContext[str]) -> None: + handled.append(request) + await ctx.send_message(f"response:{request}") + + async def _cancel_pending_request(self, request_id: str, ctx: WorkflowContext[Any, Any]) -> None: + cancelled.append(request_id) + await ctx.send_message(f"cancelled:{request_id}") + + class Collector(Executor): + @handler + async def collect(self, message: str, ctx: WorkflowContext) -> None: + delivered.append((message, ctx.get_state(WORKFLOW_SESSION_ID_KEY))) + + requester = Requester(id="requester") + workflow = WorkflowBuilder(start_executor=requester).add_edge(requester, Collector(id="collector")).build() + if legacy: + workflow._runner._ctx = cast(RunnerContext, LegacyRunnerContext()) + context = workflow._runner.context + if mixed: + await workflow.run("a") + await workflow.run("b") + else: + await workflow.run("same") + scope = workflow._runner.state.get(WORKFLOW_SESSION_ID_KEY) + await context.send_request_info_response("a", "yes") + checkpoint = await workflow._runner.build_checkpoint() + await workflow._runner.restore_checkpoint(checkpoint) + before = await context.build_checkpoint( + workflow.name, workflow.graph_signature_hash, workflow._runner.state, None, 0 + ) + pending = await context.get_pending_request_info_events() + origins = workflow._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY) + + if mixed: + with pytest.raises(ValueError, match="multiple parent workflow scopes"): + await workflow.cancel_pending_requests(["b"]) + assert cancelled == handled == delivered == [] + assert await context.get_pending_request_info_events() == pending + snapshot = await context.build_checkpoint( + workflow.name, workflow.graph_signature_hash, workflow._runner.state, None, 0 + ) + assert snapshot.messages == before.messages + assert workflow._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY) == origins + assert workflow._runner.state.get(WORKFLOW_SESSION_ID_KEY) == scope + else: + await workflow.cancel_pending_requests(["b"]) + assert cancelled == ["b"] + assert handled == ["a"] + assert sorted(delivered) == [("cancelled:b", scope), ("response:a", scope)] + assert not await context.get_pending_request_info_events() + assert workflow._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY) == {} + + async def test_public_stream_close_immediately_drains_live_executor_without_gc() -> None: started, cancelled, closed, release = (asyncio.Event() for _ in range(4)) cleanup: list[str] = [] diff --git a/skills/agent-framework-builder-python/references/declarative.md b/skills/agent-framework-builder-python/references/declarative.md index 0f938d9f00e..51540d79b22 100644 --- a/skills/agent-framework-builder-python/references/declarative.md +++ b/skills/agent-framework-builder-python/references/declarative.md @@ -32,6 +32,8 @@ maturity. Check current package status and public docstrings before making compa Child responses restore the originating parent scope before invoking downstream bindings, including after checkpoint restoration. Submit responses from different parent scopes in separate runs. The same origin recovery applies to ordinary Core `ctx.request_info()` pauses before invoking a child. + Queued replies participate in response and cancellation preflight. Wrapped child replies are also validated + before live dispatch; do not combine different parent origins in one message or cancellation batch. ## Validation From 7e45edea863e2121a07a8a95ccd2b86664bbe79d Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Tue, 27 Oct 2026 19:50:16 -0400 Subject: [PATCH 69/76] fix: honor Core scopes in manually composed workflows Use the active Core workflow scope for declarative executor lookups while preserving factory and standalone initialization. Cover fresh runs, nested bindings, agent continuity, entry replay, and interleaved approvals. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 8e076c3e-5f45-4049-9cf5-f51ed5971ba8 --- ...thon-declarative-request-scope-recovery.md | 6 + python/packages/declarative/AGENTS.md | 2 + .../_workflows/_mcp_handler.py | 22 +++- .../tests/test_invoke_mcp_tool_executor.py | 123 +++++++++++++++++- .../references/declarative.md | 1 + 5 files changed, 150 insertions(+), 4 deletions(-) diff --git a/docs/decisions/0046-python-declarative-request-scope-recovery.md b/docs/decisions/0046-python-declarative-request-scope-recovery.md index 854d53552a6..8ea3bd80d6c 100644 --- a/docs/decisions/0046-python-declarative-request-scope-recovery.md +++ b/docs/decisions/0046-python-declarative-request-scope-recovery.md @@ -85,6 +85,12 @@ originating parent scope before further nested invocation. Child scope selection occurs only after the child accepts a fresh run, not before checkpoint validation or on response-only continuations. +Declarative executors manually composed with Core `WorkflowBuilder` read the +active Core scope instead of reusing a stale declarative cache. Factory +initialization still selects its existing declarative scopes before the entry +checkpoint; standalone state helpers retain their previous initialization behavior. +Restored request scopes take precedence over inherited metadata from later runs. + Older parent checkpoints without this scope cannot reconstruct a child's original scope if that child had not started. Such parents receive a new isolated scope with a warning; subsequent checkpoints retain it. Existing request-origin diff --git a/python/packages/declarative/AGENTS.md b/python/packages/declarative/AGENTS.md index d6dae130d1b..675120c93c7 100644 --- a/python/packages/declarative/AGENTS.md +++ b/python/packages/declarative/AGENTS.md @@ -31,6 +31,8 @@ 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. +Executors manually composed with Core `WorkflowBuilder` use Core's active run scope, +including inherited child scopes, rather than a cached declarative 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. 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 f2438e77f18..c3a3d16511e 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -40,7 +40,11 @@ from typing import TYPE_CHECKING, Any, ClassVar, Protocol, cast, runtime_checkable import httpx2 as httpx -from agent_framework._workflows._const import WORKFLOW_PARENT_SESSION_ID_KEY, WORKFLOW_SESSION_ID_KEY +from agent_framework._workflows._const import ( + WORKFLOW_AGENT_RUN_KEY, + WORKFLOW_PARENT_SESSION_ID_KEY, + WORKFLOW_SESSION_ID_KEY, +) from ._errors import DeclarativeWorkflowError @@ -66,6 +70,18 @@ def get_or_create_workflow_session_id(state: State) -> str: + # Core runs own the active scope even without the declarative factory initializer. + if state.has(WORKFLOW_AGENT_RUN_KEY): + workflow_session_id = state.get(WORKFLOW_SESSION_ID_KEY) + if workflow_session_id is not None: + if not isinstance(workflow_session_id, str) or not workflow_session_id: + raise ValueError("Invalid MCP workflow session state.") + state.set(_WORKFLOW_SESSION_ID_KEY, workflow_session_id) + return workflow_session_id + return _get_or_create_declarative_session_id(state) + + +def _get_or_create_declarative_session_id(state: State) -> str: agent_session_id = state.get(_WORKFLOW_AGENT_SESSION_ID_KEY) if agent_session_id is not None and state.get(WORKFLOW_PARENT_SESSION_ID_KEY) is None: if not isinstance(agent_session_id, str) or not agent_session_id: @@ -100,7 +116,7 @@ def activate_workflow_session_id(state: State, *, reset_unscoped: bool, is_agent raise ValueError("Invalid parent workflow session scope.") state.set(_WORKFLOW_SESSION_ID_KEY, parent_scope) elif state.get(_WORKFLOW_AGENT_SESSION_ID_KEY) is not None: - get_or_create_workflow_session_id(state) + _get_or_create_declarative_session_id(state) elif is_agent_run: workflow_session_id = state.get(_SESSIONLESS_AGENT_WORKFLOW_SESSION_ID_KEY) if workflow_session_id is None: @@ -111,7 +127,7 @@ def activate_workflow_session_id(state: State, *, reset_unscoped: bool, is_agent state.set(_WORKFLOW_SESSION_ID_KEY, workflow_session_id) elif reset_unscoped: state.set(_WORKFLOW_SESSION_ID_KEY, uuid.uuid4().hex) - state.set(WORKFLOW_SESSION_ID_KEY, get_or_create_workflow_session_id(state)) + state.set(WORKFLOW_SESSION_ID_KEY, _get_or_create_declarative_session_id(state)) def restore_workflow_session_id(state: State, workflow_session_id: str) -> None: 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 1e4faef66a1..c770aa4f2cc 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -28,7 +28,7 @@ reason="PowerFx engine not available (requires dotnet runtime)", ) -from agent_framework import AgentSession, Content, Message # noqa: E402 +from agent_framework import AgentSession, Content, Message, Workflow # noqa: E402 from agent_framework.exceptions import ToolExecutionException # noqa: E402 from agent_framework_declarative._workflows import ( # noqa: E402 @@ -117,6 +117,30 @@ def _yaml(action: dict[str, Any]) -> dict[str, Any]: return {"name": "mcp_test", "actions": [action]} +def _build_manual_mcp_workflow( + mcp_handler: MCPToolHandler, *, depth: int, require_approval: bool = False +) -> tuple[Workflow, Workflow]: + from agent_framework import WorkflowBuilder, WorkflowExecutor + + from agent_framework_declarative._workflows._executors_mcp import InvokeMcpToolActionExecutor + + leaf = WorkflowBuilder( + name="manual_leaf", + start_executor=InvokeMcpToolActionExecutor( + _action(require_approval=require_approval), mcp_tool_handler=mcp_handler + ), + ).build() + root = leaf + for level in range(depth): + root = WorkflowBuilder( + name=f"manual_parent_{level}", + start_executor=WorkflowExecutor( + root, id=f"child_{level}", propagate_request=True, allow_direct_output=True + ), + ).build() + return root, leaf + + # ---------- Builder enforcement -------------------------------------------- @@ -151,6 +175,103 @@ def test_missing_tool_name_fails_validation(self) -> None: class TestFieldForwarding: + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("depth", [0, 1, 2]) + async def test_manual_workflow_scopes_follow_core_run_context(self, streaming: bool, depth: int) -> None: + from agent_framework import InMemoryCheckpointStorage + from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY + + mcp_handler = StubMcpHandler(_ok()) + storage = InMemoryCheckpointStorage() + root, leaf = _build_manual_mcp_workflow(mcp_handler, depth=depth) + expected_scopes: list[str] = [] + entry_id: str | None = None + for index in range(2): + if streaming: + await root.run({}, checkpoint_storage=storage, stream=True).get_final_response() + else: + await root.run({}, checkpoint_storage=storage) + expected_scopes.append(leaf._runner.state.get(WORKFLOW_SESSION_ID_KEY)) + if index == 0: + [entry] = [ + checkpoint + for checkpoint in await storage.list_checkpoints(workflow_name=root.name) + if checkpoint.iteration_count == 0 + ] + entry_id = entry.checkpoint_id + + actual_scopes = [invocation.workflow_session_id for invocation in mcp_handler.invocations] + assert len(set(actual_scopes)) == 2 + assert actual_scopes == expected_scopes + + agent = root.as_agent() + for session in (None, AgentSession()): + for _ in range(2): + if streaming: + await agent.run("agent turn", session=session, stream=True).get_final_response() + else: + await agent.run("agent turn", session=session) + expected_scopes.append(leaf._runner.state.get(WORKFLOW_SESSION_ID_KEY)) + assert expected_scopes[-1] == expected_scopes[-2] + assert len(set(expected_scopes)) == 4 + assert [invocation.workflow_session_id for invocation in mcp_handler.invocations] == expected_scopes + + assert entry_id is not None + root, leaf = _build_manual_mcp_workflow(mcp_handler, depth=depth) + if streaming: + await root.run(checkpoint_id=entry_id, checkpoint_storage=storage, stream=True).get_final_response() + else: + await root.run(checkpoint_id=entry_id, checkpoint_storage=storage) + assert mcp_handler.invocations[-1].workflow_session_id == expected_scopes[0] + assert leaf._runner.state.get(WORKFLOW_SESSION_ID_KEY) == expected_scopes[0] + + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("restore", [False, True]) + async def test_manual_approval_continuations_restore_core_scopes(self, streaming: bool, restore: bool) -> None: + from agent_framework import InMemoryCheckpointStorage, WorkflowEvent + from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY + + from agent_framework_declarative._workflows import ToolApprovalResponse + + mcp_handler = StubMcpHandler(_ok()) + storage = InMemoryCheckpointStorage() + root, leaf = _build_manual_mcp_workflow(mcp_handler, depth=1, require_approval=True) + requests: list[WorkflowEvent[Any]] = [] + for _ in range(2): + if streaming: + await root.run({}, checkpoint_storage=storage, stream=True).get_final_response() + else: + await root.run({}, checkpoint_storage=storage) + [request] = [ + request + for request in (await root._runner_context.get_pending_request_info_events()).values() + if request.request_id not in {previous.request_id for previous in requests} + ] + requests.append(request) + scopes = [request.data.workflow_session_id for request in requests] + assert len(set(scopes)) == 2 + assert mcp_handler.call_count == 0 + if restore: + checkpoint_id = root.get_last_checkpoint_id() + assert checkpoint_id is not None + root, leaf = _build_manual_mcp_workflow(mcp_handler, depth=1, require_approval=True) + await root.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage) + + with pytest.raises(ValueError, match="multiple parent workflow scopes"): + await root.run(responses={request.request_id: ToolApprovalResponse(approved=True) for request in requests}) + assert mcp_handler.call_count == 0 + assert len(await root._runner_context.get_pending_request_info_events()) == 2 + for request, scope in reversed(list(zip(requests, scopes, strict=True))): + responses = {request.request_id: ToolApprovalResponse(approved=True)} + if streaming: + await root.run(responses=responses, stream=True).get_final_response() + else: + await root.run(responses=responses) + assert mcp_handler.invocations[-1].workflow_session_id == scope + assert leaf._runner.state.get(WORKFLOW_SESSION_ID_KEY) == scope + assert not await root._runner_context.get_pending_request_info_events() + assert not await leaf._runner_context.get_pending_request_info_events() + @pytest.mark.asyncio async def test_basic_invocation_forwards_required_fields(self) -> None: handler = StubMcpHandler(_ok()) diff --git a/skills/agent-framework-builder-python/references/declarative.md b/skills/agent-framework-builder-python/references/declarative.md index 51540d79b22..d46ff79e18f 100644 --- a/skills/agent-framework-builder-python/references/declarative.md +++ b/skills/agent-framework-builder-python/references/declarative.md @@ -29,6 +29,7 @@ maturity. Check current package status and public docstrings before making compa - Declarative children wrapped in `WorkflowExecutor` retain MCP session continuity across parent agent turns. Each binding must own a unique child workflow instance; its MCP scope remains isolated from sibling bindings. Replaying a parent entry checkpoint also retains child scopes when rebuilding the same workflow graph. + Declarative executors manually composed with Core `WorkflowBuilder` also follow Core's active run and inherited scopes. Child responses restore the originating parent scope before invoking downstream bindings, including after checkpoint restoration. Submit responses from different parent scopes in separate runs. The same origin recovery applies to ordinary Core `ctx.request_info()` pauses before invoking a child. From f0f190693fa7489989ed4699f2b7dffb0eef4caf Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Thu, 29 Oct 2026 10:37:49 -0400 Subject: [PATCH 70/76] fix(workflows): isolate relay origins and incrementally update request state Preserve independently owned child-boundary origins when coordinators relay the same request ID. Stage individual origin mapping updates to avoid quadratic request fan-out copying while preserving checkpoint and rollback semantics. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 06d37a04-f20d-46d4-ab56-1ef5adda7b2d --- ...thon-declarative-request-scope-recovery.md | 8 ++ python/packages/core/AGENTS.md | 5 +- .../_workflows/_request_context.py | 102 +++++++++----- .../agent_framework/_workflows/_runner.py | 16 ++- .../core/agent_framework/_workflows/_state.py | 72 ++++++++-- .../agent_framework/_workflows/_workflow.py | 13 +- .../_workflows/_workflow_executor.py | 4 +- .../core/tests/workflow/test_state.py | 132 +++++++++++++++++- python/packages/declarative/tests/conftest.py | 67 +++++++++ .../tests/test_invoke_mcp_tool_executor.py | 54 ++++++- 10 files changed, 415 insertions(+), 58 deletions(-) diff --git a/docs/decisions/0046-python-declarative-request-scope-recovery.md b/docs/decisions/0046-python-declarative-request-scope-recovery.md index 8ea3bd80d6c..88ff251e982 100644 --- a/docs/decisions/0046-python-declarative-request-scope-recovery.md +++ b/docs/decisions/0046-python-declarative-request-scope-recovery.md @@ -130,6 +130,14 @@ including for queued checkpoint replies. Core cancellation handlers restore the same metadata. Propagated child responses use this path once; wrapped child messages retain their boundary-specific capture and restoration. +Wrapped origins are owned independently by each child binding, not by the +coordinator's external-request ID. Relaying a child request with that same ID +therefore retains two independently consumable origins across checkpoints. +Capturing or releasing an origin stages only that mapping entry in shared state, +with committed values unchanged until the superstep commits. Discard rolls back +these entry updates. Full-map validation and detached copies remain checkpoint +and response-batch operations, not per-request insertion or removal work. + Preflight includes both propagated responses and wrapped child replies already queued for delivery. Route wrapped replies through the parent graph to validate their child's own policy without treating child scopes as parent-policy scopes. diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 4d69ac2e4bc..f24012d9e21 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -346,7 +346,10 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID retains its run context in checkpointed shared state, including ordinary Core requests before a child binding. Core restores and releases that context before invoking response or cancellation handlers, including queued checkpoint replies. Wrapped child replies use the same helpers at their message boundary; propagated replies - are restored only once. Mixed-parent-scope response batches are rejected before consuming requests, including + are restored only once. Wrapped origins belong to their binding independently of external requests, even when + a coordinator relays the same request ID. Origin insertion and removal stage individual mapping entries; + committed snapshots remain unchanged until commit and discard rolls back pending changes. + Mixed-parent-scope response batches are rejected before consuming requests, including across sibling bindings. Validation includes queued wrapped replies and runs before each live message batch is drained, routing wrapped replies to their child's own policy. Cancellation preflight combines queued replies with the complete selected cancellation set before consuming any requests. diff --git a/python/packages/core/agent_framework/_workflows/_request_context.py b/python/packages/core/agent_framework/_workflows/_request_context.py index 3c9b26dbf51..e2e3622886c 100644 --- a/python/packages/core/agent_framework/_workflows/_request_context.py +++ b/python/packages/core/agent_framework/_workflows/_request_context.py @@ -25,56 +25,89 @@ WORKFLOW_AGENT_RUN_KEY, WORKFLOW_AGENT_SESSION_ID_KEY, ) +_MISSING_ORIGIN = object() -def validate_request_contexts(contexts: Any, requests: Sequence[WorkflowEvent] = ()) -> dict[str, dict[str, Any]]: - """Preflight checkpointed request origins before any responses are consumed.""" - if not isinstance(contexts, dict): - raise ValueError("Invalid parent workflow request contexts.") - contexts = cast(dict[str, dict[str, Any]], contexts) - for request_id, value in contexts.items(): - if not isinstance(request_id, str) or not isinstance(value, dict): - raise ValueError("Invalid parent workflow request context.") - context = cast(dict[str, Any], value) - if set(context) != set(_RUN_CONTEXT_KEYS): - raise ValueError("Invalid parent workflow request context.") - scope = context[WORKFLOW_SESSION_ID_KEY] - inherited_scope = context[WORKFLOW_PARENT_SESSION_ID_KEY] - agent_session = context[WORKFLOW_AGENT_SESSION_ID_KEY] - if ( - not isinstance(scope, str) - or not scope - or (inherited_scope is not None and (not isinstance(inherited_scope, str) or not inherited_scope)) - or not isinstance(context[WORKFLOW_AGENT_RUN_KEY], bool) - or (agent_session is not None and not isinstance(agent_session, str)) - ): - raise ValueError("Invalid parent workflow request context.") - scopes = { - contexts[request.request_id][WORKFLOW_SESSION_ID_KEY] for request in requests if request.request_id in contexts - } +def request_context_state_key(boundary_id: str | None = None) -> str: + """Keep wrapped child origins independent of external requests with the same ID.""" + if boundary_id is None: + return WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY + return f"{WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY}:{boundary_id}" + + +def _validate_origin(value: Any) -> dict[str, Any]: + if not isinstance(value, dict): + raise ValueError("Invalid parent workflow request context.") + context = cast(dict[str, Any], value) + if set(context) != set(_RUN_CONTEXT_KEYS): + raise ValueError("Invalid parent workflow request context.") + scope = context[WORKFLOW_SESSION_ID_KEY] + inherited_scope = context[WORKFLOW_PARENT_SESSION_ID_KEY] + agent_session = context[WORKFLOW_AGENT_SESSION_ID_KEY] + if ( + not isinstance(scope, str) + or not scope + or (inherited_scope is not None and (not isinstance(inherited_scope, str) or not inherited_scope)) + or not isinstance(context[WORKFLOW_AGENT_RUN_KEY], bool) + or (agent_session is not None and not isinstance(agent_session, str)) + ): + raise ValueError("Invalid parent workflow request context.") + return context + + +def validate_request_contexts( + contexts: Any, + requests: Sequence[WorkflowEvent] = (), + *, + boundary_contexts: Sequence[tuple[Any, Sequence[WorkflowEvent]]] = (), +) -> None: + """Preflight all origin owners together before any responses are consumed.""" + scopes: set[str] = set() + for pending_contexts, pending_requests in ((contexts, requests), *boundary_contexts): + if not isinstance(pending_contexts, dict): + raise ValueError("Invalid parent workflow request contexts.") + pending_contexts = cast(dict[str, Any], pending_contexts) + for request_id, value in pending_contexts.items(): + if not isinstance(request_id, str): + raise ValueError("Invalid parent workflow request context.") + _validate_origin(value) + scopes.update( + pending_contexts[request.request_id][WORKFLOW_SESSION_ID_KEY] + for request in pending_requests + if request.request_id in pending_contexts + ) if len(scopes) > 1: raise ValueError("Responses span multiple parent workflow scopes; submit each scope in a separate run.") - return contexts -def capture_request_context(state: State, request_id: str) -> None: +def capture_request_context(state: State, request_id: str, *, boundary_id: str | None = None) -> None: """Associate a request with framework-owned run metadata, without copying its payload.""" - contexts = validate_request_contexts(state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {})) + state_key = request_context_state_key(boundary_id) + existing = state._get_mapping_item(state_key, request_id, _MISSING_ORIGIN) # pyright: ignore[reportPrivateUsage] + if existing is not _MISSING_ORIGIN: + _validate_origin(existing) + return origin = {key: state.get(key) for key in _RUN_CONTEXT_KEYS} if origin[WORKFLOW_SESSION_ID_KEY] is None: logger.warning("Workflow request %s has no active run scope; assigning a new isolated scope.", request_id) origin[WORKFLOW_SESSION_ID_KEY] = uuid.uuid4().hex state.set(WORKFLOW_SESSION_ID_KEY, origin[WORKFLOW_SESSION_ID_KEY]) origin[WORKFLOW_AGENT_RUN_KEY] = origin[WORKFLOW_AGENT_RUN_KEY] is True - contexts.setdefault(request_id, origin) - state.set(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, contexts) + _validate_origin(origin) + state._set_mapping_item(state_key, request_id, origin) # pyright: ignore[reportPrivateUsage] -def restore_request_context(state: State, request_id: str) -> None: +def restore_request_context(state: State, request_id: str, *, boundary_id: str | None = None) -> None: """Restore and release a request's run metadata before its continuation executes.""" - contexts = validate_request_contexts(state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {})) - origin = contexts.pop(request_id, None) - if origin is None: + state_key = request_context_state_key(boundary_id) + if boundary_id is not None and not state.has(state_key): + # Checkpoints predating separate boundary ownership used the external-request map. + state_key = WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY + origin = state._get_mapping_item(state_key, request_id, _MISSING_ORIGIN) # pyright: ignore[reportPrivateUsage] + if origin is not _MISSING_ORIGIN: + _validate_origin(origin) + state._pop_mapping_item(state_key, request_id) # pyright: ignore[reportPrivateUsage] + else: logger.warning( "Workflow request %s has no checkpointed run origin; assigning a new isolated scope.", request_id ) @@ -86,4 +119,3 @@ def restore_request_context(state: State, request_id: str) -> None: } for key in _RUN_CONTEXT_KEYS: state.set(key, origin[key]) - state.set(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, contexts) diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index 61648f9fc4d..df567648244 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -20,7 +20,7 @@ from ._edge_runner import EdgeRunner, create_edge_runner, gather_cancelling_siblings_on_error from ._events import WorkflowEvent from ._executor import Executor -from ._request_context import validate_request_contexts +from ._request_context import request_context_state_key, validate_request_contexts from ._runner_context import ( RunnerContext, WorkflowMessage, @@ -458,11 +458,17 @@ def _validate_checkpoint_tree( if requests and self._response_batch_validator is not None: self._response_batch_validator(requests) wrapped_requests = self._get_wrapped_response_requests(checkpoint.messages) + parent_contexts = checkpoint.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {}) validate_request_contexts( - checkpoint.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {}), - [ - *requests, - *(request for batch in wrapped_requests.values() for request in batch), + parent_contexts, + requests, + boundary_contexts=[ + ( + checkpoint.state.get(request_context_state_key(executor_id), parent_contexts), + wrapped_requests.get(executor_id, []), + ) + for executor_id in self._executors + if executor_id in wrapped_requests or request_context_state_key(executor_id) in checkpoint.state ], ) if EXECUTOR_STATE_KEY in checkpoint.state: diff --git a/python/packages/core/agent_framework/_workflows/_state.py b/python/packages/core/agent_framework/_workflows/_state.py index 136b70de14e..288812f5073 100644 --- a/python/packages/core/agent_framework/_workflows/_state.py +++ b/python/packages/core/agent_framework/_workflows/_state.py @@ -1,8 +1,10 @@ # Copyright (c) Microsoft. All rights reserved. +from __future__ import annotations + import copy from collections.abc import Callable -from typing import Any +from typing import Any, cast class State: @@ -28,6 +30,7 @@ def __init__(self) -> None: """Initialize the state.""" self._committed: dict[str, Any] = {} self._pending: dict[str, Any] = {} + self._pending_mapping_updates: dict[str, dict[str, Any]] = {} def set(self, key: str, value: Any) -> None: """Set a value in the pending state buffer. @@ -44,6 +47,7 @@ def set(self, key: str, value: Any) -> None: in a superstep see the same committed state at the start. """ self._pending[key] = copy.deepcopy(value) + self._pending_mapping_updates.pop(key, None) def get(self, key: str, default: Any = None) -> Any: """Get a value from state, checking pending first then committed. @@ -56,14 +60,51 @@ def get(self, key: str, default: Any = None) -> Any: A deep copy of the value if found, otherwise the default value. Mutate the returned value and pass it to :meth:`set` to update workflow state. """ - if key in self._pending: - value = self._pending[key] - if value is _DeleteSentinel: - return default - return copy.deepcopy(value) - if key in self._committed: - return copy.deepcopy(self._committed[key]) - return default + value = self._get_value(key) + return default if value is _DeleteSentinel else copy.deepcopy(value) + + def _get_value(self, key: str) -> Any: + value = self._pending.get(key, self._committed.get(key, _DeleteSentinel)) + updates = self._pending_mapping_updates.get(key) + if updates: + merged: dict[str, Any] = dict(value) + for item_key, item_value in updates.items(): + if item_value is _DeleteSentinel: + merged.pop(item_key, None) + else: + merged[item_key] = item_value + return merged + return value + + def _get_mapping(self, key: str) -> dict[str, Any]: + value = self._pending.get(key, self._committed.get(key, _DeleteSentinel)) + if value is _DeleteSentinel: + return {} + if not isinstance(value, dict): + raise ValueError(f"Workflow state '{key}' must be a mapping.") + return cast(dict[str, Any], value) + + def _get_mapping_item(self, key: str, item_key: str, default: Any = None) -> Any: + """Read one framework-owned mapping entry without copying its siblings.""" + mapping = self._get_mapping(key) + value = self._pending_mapping_updates.get(key, {}).get(item_key, mapping.get(item_key, _DeleteSentinel)) + return default if value is _DeleteSentinel else copy.deepcopy(value) + + def _set_mapping_item(self, key: str, item_key: str, value: Any) -> None: + """Stage a detached mapping entry while preserving commit/discard semantics.""" + self._get_mapping(key) + copied_value = copy.deepcopy(value) + if not self.has(key): + self._pending[key] = {} + self._pending_mapping_updates.setdefault(key, {})[item_key] = copied_value + + def _pop_mapping_item(self, key: str, item_key: str) -> Any: + """Stage removal of one entry without copying or mutating committed siblings.""" + value = self._get_mapping_item(key, item_key, _DeleteSentinel) + if value is not _DeleteSentinel: + self._pending_mapping_updates.setdefault(key, {})[item_key] = _DeleteSentinel + return value + return None def has(self, key: str) -> bool: """Check if a key exists in pending or committed state.""" @@ -77,7 +118,7 @@ def _validate(self, key: str, validator: Callable[[Any], None]) -> None: The internal validator must not mutate or retain the value. Missing and pending-deleted keys are not validated, matching ``get`` visibility. """ - value = self._pending.get(key, self._committed.get(key, _DeleteSentinel)) + value = self._get_value(key) if value is not _DeleteSentinel: validator(value) @@ -97,11 +138,13 @@ def delete(self, key: str) -> None: elif key in self._pending: # Only exists in pending, safe to just remove del self._pending[key] + self._pending_mapping_updates.pop(key, None) def clear(self) -> None: """Clear both committed and pending state.""" self._committed.clear() self._pending.clear() + self._pending_mapping_updates.clear() def commit(self) -> None: """Commit pending changes to the committed state. @@ -113,11 +156,20 @@ def commit(self) -> None: self._committed.pop(key, None) else: self._committed[key] = value + for key, updates in self._pending_mapping_updates.items(): + mapping = self._committed[key] + for item_key, value in updates.items(): + if value is _DeleteSentinel: + mapping.pop(item_key, None) + else: + mapping[item_key] = value self._pending.clear() + self._pending_mapping_updates.clear() def discard(self) -> None: """Discard all pending changes without committing.""" self._pending.clear() + self._pending_mapping_updates.clear() def export_state(self) -> dict[str, Any]: """Export a deepcopy of the committed state. diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 0807cf792e7..5f034ecdc9b 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -60,7 +60,7 @@ _ResponseBatchValidator, # pyright: ignore[reportPrivateUsage] ) from ._model_utils import DictConvertible -from ._request_context import restore_request_context, validate_request_contexts +from ._request_context import request_context_state_key, restore_request_context, validate_request_contexts from ._runner import RunnerImpl from ._runner_context import ( RunnerContext, @@ -1218,11 +1218,18 @@ async def _get_response_validation_batches( queued_requests = _get_queued_response_info_events(messages) wrapped_requests = self._runner._get_wrapped_response_requests(messages) # pyright: ignore[reportPrivateUsage] batch = [*queued_requests, *requests] - origin_requests = [*batch, *(request for wrapped in wrapped_requests.values() for request in wrapped)] parent_contexts = self._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {}) + boundary_contexts = [ + ( + self._runner.state.get(request_context_state_key(executor_id), parent_contexts), + wrapped_requests.get(executor_id, []), + ) + for executor_id in self.executors + if executor_id in wrapped_requests or self._runner.state.has(request_context_state_key(executor_id)) + ] def validate_parent_contexts(_requests: Sequence[WorkflowEvent]) -> None: - validate_request_contexts(parent_contexts, origin_requests) + validate_request_contexts(parent_contexts, batch, boundary_contexts=boundary_contexts) validations: list[_ResponseBatchValidation] = [] if self._response_batch_validator is not None: diff --git a/python/packages/core/agent_framework/_workflows/_workflow_executor.py b/python/packages/core/agent_framework/_workflows/_workflow_executor.py index da292a4d5db..e6c9651b452 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_executor.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_executor.py @@ -726,7 +726,7 @@ async def _forward_intermediate_output(output: Any) -> None: else: # In a workflow where the parent workflow has an executor that may intercept the # request and handle it directly, a message should be sent. - capture_request_context(ctx.state, request_id) + capture_request_context(ctx.state, request_id, boundary_id=self.id) await ctx.send_message(SubWorkflowRequestMessage(source_event=event, executor_id=self.id)) # Handle final state @@ -786,6 +786,6 @@ async def _handle_response( ) # Propagated replies already restored their context before their response handler. if ctx.request_id != request_id: - restore_request_context(ctx.state, request_id) + restore_request_context(ctx.state, request_id, boundary_id=self.id) self._remove_reemitted_pending_requests(result, pending_requests) await self._process_workflow_result(result, ctx) diff --git a/python/packages/core/tests/workflow/test_state.py b/python/packages/core/tests/workflow/test_state.py index 6b6a52c6084..8973dbbe830 100644 --- a/python/packages/core/tests/workflow/test_state.py +++ b/python/packages/core/tests/workflow/test_state.py @@ -2,7 +2,7 @@ """Unit tests for the State class superstep caching behavior.""" -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest @@ -10,6 +10,136 @@ from agent_framework._workflows._state import State +def test_request_origin_fan_out_copies_only_incremental_entries() -> None: + import copy + + from agent_framework._workflows._const import WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, WORKFLOW_SESSION_ID_KEY + from agent_framework._workflows._request_context import capture_request_context, restore_request_context + + state = State() + state.set(WORKFLOW_SESSION_ID_KEY, "origin") + request_count = 200 + with patch("agent_framework._workflows._state.copy.deepcopy", wraps=copy.deepcopy) as deepcopy: + for index in range(request_count): + capture_request_context(state, str(index)) + state.commit() + for index in range(request_count): + restore_request_context(state, str(index)) + assert state.get(WORKFLOW_SESSION_ID_KEY) == "origin" + copied_entries = sum(len(call.args[0]) for call in deepcopy.call_args_list if isinstance(call.args[0], dict)) + + assert state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY) == {} + assert copied_entries < 40 * request_count + + +def test_incremental_mapping_updates_preserve_snapshots_and_rollback() -> None: + state = State() + state.set("origins", {"first": {"scope": "A"}, "second": {"scope": "B"}}) + state.commit() + snapshot = state.export_state() + update = {"scope": "C"} + + state._set_mapping_item("origins", "third", update) + update["scope"] = "mutated" + removed = state._pop_mapping_item("origins", "first") + removed["scope"] = "mutated" + assert state._get_mapping_item("origins", "third") == {"scope": "C"} + assert state.get("origins") == {"second": {"scope": "B"}, "third": {"scope": "C"}} + assert state.export_state() == snapshot + + state.discard() + assert state.get("origins") == snapshot["origins"] + state._pop_mapping_item("origins", "first") + state._set_mapping_item("origins", "second", {"scope": "D"}) + state.commit() + assert state.export_state()["origins"] == {"second": {"scope": "D"}} + assert snapshot["origins"]["second"] == {"scope": "B"} + + +@pytest.mark.parametrize("operation", ["set", "delete", "clear"]) +def test_whole_value_operations_replace_pending_mapping_updates(operation: str) -> None: + state = State() + state.set("origins", {"first": 1}) + state.commit() + state._set_mapping_item("origins", "second", 2) + + if operation == "set": + state.set("origins", {"replacement": 3}) + expected = {"replacement": 3} + elif operation == "delete": + state.delete("origins") + expected = None + else: + state.clear() + expected = None + assert state.get("origins") == expected + state.commit() + assert state.get("origins") == expected + + +def test_incremental_mapping_updates_support_new_deleted_and_null_entries() -> None: + state = State() + state._set_mapping_item("origins", "first", None) + assert state.has("origins") + state._pop_mapping_item("origins", "first") + assert state.get("origins") == {} + state.commit() + state.delete("origins") + state._set_mapping_item("origins", "second", {"scope": "B"}) + state.commit() + assert state.export_state()["origins"] == {"second": {"scope": "B"}} + + +def test_request_origin_owners_remain_independent_for_the_same_request_id() -> None: + from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY + from agent_framework._workflows._request_context import capture_request_context, restore_request_context + + state = State() + state.set(WORKFLOW_SESSION_ID_KEY, "child-boundary") + capture_request_context(state, "same-id", boundary_id="child") + state.set(WORKFLOW_SESSION_ID_KEY, "coordinator-relay") + capture_request_context(state, "same-id") + state.commit() + + restore_request_context(state, "same-id") + assert state.get(WORKFLOW_SESSION_ID_KEY) == "coordinator-relay" + state.commit() + restore_request_context(state, "same-id", boundary_id="child") + assert state.get(WORKFLOW_SESSION_ID_KEY) == "child-boundary" + + +def test_wrapped_request_restores_legacy_checkpoint_origin() -> None: + from agent_framework._workflows._const import WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, WORKFLOW_SESSION_ID_KEY + from agent_framework._workflows._request_context import capture_request_context, restore_request_context + + legacy = State() + legacy.set(WORKFLOW_SESSION_ID_KEY, "legacy-boundary") + capture_request_context(legacy, "request") + legacy.commit() + restored = State() + restored.import_state(legacy.export_state()) + restored.set(WORKFLOW_SESSION_ID_KEY, "later-run") + + restore_request_context(restored, "request", boundary_id="child") + + assert restored.get(WORKFLOW_SESSION_ID_KEY) == "legacy-boundary" + assert restored.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY) == {} + + +@pytest.mark.parametrize("operation", ["capture", "restore"]) +def test_incremental_request_origin_operations_reject_malformed_entries(operation: str) -> None: + from agent_framework._workflows._const import WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY + from agent_framework._workflows._request_context import capture_request_context, restore_request_context + + state = State() + state.set(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {"request": None}) + state.commit() + action = capture_request_context if operation == "capture" else restore_request_context + with pytest.raises(ValueError, match="Invalid parent workflow request context"): + action(state, "request") + assert state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY) == {"request": None} + + class TestStateBasicOperations: """Tests for basic State get/set/has/delete operations.""" diff --git a/python/packages/declarative/tests/conftest.py b/python/packages/declarative/tests/conftest.py index 083de75660f..d6b7f7859ae 100644 --- a/python/packages/declarative/tests/conftest.py +++ b/python/packages/declarative/tests/conftest.py @@ -2,9 +2,25 @@ """Pytest configuration for declarative tests.""" +from __future__ import annotations + import sys +from collections.abc import Callable +from typing import Any import pytest +from agent_framework import ( + CheckpointStorage, + Executor, + SubWorkflowRequestMessage, + SubWorkflowResponseMessage, + Workflow, + WorkflowBuilder, + WorkflowContext, + WorkflowExecutor, + handler, + response_handler, +) # Skip all tests in this directory on Python 3.14+ because powerfx doesn't support it yet if sys.version_info >= (3, 14): @@ -18,3 +34,54 @@ def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item for item in items: if "declarative" in str(item.fspath): item.add_marker(skip_marker) + + +@pytest.fixture +def coordinator_relay_workflow() -> Callable[[Workflow, CheckpointStorage], Workflow]: + """Build a parent that relays child requests without changing their IDs.""" + + class Requester(Executor): + @handler + async def start(self, message: str, ctx: WorkflowContext) -> None: + await ctx.request_info(message, str) + + @response_handler + async def respond(self, original: str, response: str, ctx: WorkflowContext[Any, dict[str, str]]) -> None: + await ctx.yield_output({"response": response}) + + class Coordinator(Executor): + @handler + async def start(self, message: int, ctx: WorkflowContext[Any]) -> None: + await ctx.send_message({}, target_id="mcp") + await ctx.send_message(str(message), target_id="requester") + + @handler + async def relay(self, request: SubWorkflowRequestMessage, ctx: WorkflowContext) -> None: + await ctx.request_info(request, str, request_id=request.source_event.request_id) + + @response_handler + async def respond( + self, + original: SubWorkflowRequestMessage, + response: str, + ctx: WorkflowContext[SubWorkflowResponseMessage], + ) -> None: + await ctx.send_message(original.create_response(response), target_id="requester") + + @handler + async def collect(self, message: dict[str, str], ctx: WorkflowContext[dict[str, Any]]) -> None: + await ctx.send_message({}, target_id="mcp") + + def build(mcp_workflow: Workflow, storage: CheckpointStorage) -> Workflow: + coordinator = Coordinator(id="coordinator") + requester = WorkflowExecutor(WorkflowBuilder(start_executor=Requester(id="pause")).build(), id="requester") + mcp = WorkflowExecutor(mcp_workflow, id="mcp") + return ( + WorkflowBuilder(start_executor=coordinator, checkpoint_storage=storage) + .add_edge(coordinator, requester) + .add_edge(requester, coordinator) + .add_edge(coordinator, mcp) + .build() + ) + + return build 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 c770aa4f2cc..6a1389fee43 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -9,6 +9,7 @@ """ import sys +from collections.abc import Callable from pathlib import Path from typing import Any from unittest.mock import AsyncMock, MagicMock @@ -28,7 +29,7 @@ reason="PowerFx engine not available (requires dotnet runtime)", ) -from agent_framework import AgentSession, Content, Message, Workflow # noqa: E402 +from agent_framework import AgentSession, CheckpointStorage, Content, Message, Workflow # noqa: E402 from agent_framework.exceptions import ToolExecutionException # noqa: E402 from agent_framework_declarative._workflows import ( # noqa: E402 @@ -304,6 +305,57 @@ 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.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("restore_kind", ["memory", "file", "queued"]) + async def test_same_id_coordinator_relay_preserves_downstream_mcp_scope( + self, + tmp_path: Path, + streaming: bool, + restore_kind: str, + coordinator_relay_workflow: Callable[[Workflow, CheckpointStorage], Workflow], + ) -> None: + mcp_handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=mcp_handler) + storage = _scope_checkpoint_storage(tmp_path) + + parent = coordinator_relay_workflow(factory.create_workflow_from_definition(_yaml(_action())), storage) + requests = [] + for message in (1, 2): + if streaming: + result = await parent.run(message, stream=True).get_final_response() + else: + result = await parent.run(message) + [request] = [event for event in result.get_request_info_events() if event.request_id not in requests] + requests.append(request.request_id) + scopes = [invocation.workflow_session_id for invocation in mcp_handler.invocations] + assert len(scopes) == 2 + assert scopes[0] and scopes[1] and scopes[0] != scopes[1] + if restore_kind == "file": + checkpoint_id = parent.get_last_checkpoint_id() + assert checkpoint_id is not None + parent = coordinator_relay_workflow(factory.create_workflow_from_definition(_yaml(_action())), storage) + await parent.run(checkpoint_id=checkpoint_id) + + for request_id, scope in reversed(list(zip(requests, scopes, strict=True))): + resume_checkpoint_id: str | None = None + responses: dict[str, str] | None = {request_id: "continue"} + if restore_kind == "queued": + await parent._runner_context.send_request_info_response(request_id, "continue") + await parent._runner.create_checkpoint_if_enabled() + resume_checkpoint_id = parent.get_last_checkpoint_id() + assert resume_checkpoint_id is not None + parent = coordinator_relay_workflow(factory.create_workflow_from_definition(_yaml(_action())), storage) + responses = None + if streaming: + await parent.run( + responses=responses, checkpoint_id=resume_checkpoint_id, stream=True + ).get_final_response() + else: + await parent.run(responses=responses, checkpoint_id=resume_checkpoint_id) + assert mcp_handler.invocations[-1].workflow_session_id == scope + assert len(mcp_handler.invocations) == 4 + assert not await parent._runner_context.get_pending_request_info_events() + @pytest.mark.parametrize("streaming", [False, True]) @pytest.mark.parametrize("restore_kind", ["memory", "file", "queued"]) async def test_core_parent_request_restores_origin_before_child( From 5009d68efe26ec0c2eda17f64d82e6bbdf5c3939 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Thu, 29 Oct 2026 11:10:14 -0400 Subject: [PATCH 71/76] fix(workflows): validate selected wrapped reply routes Reuse edge routing decisions during response preflight and dispatch, including asynchronous checkpoint validation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 06d37a04-f20d-46d4-ab56-1ef5adda7b2d --- python/packages/core/AGENTS.md | 7 +- .../_workflows/_edge_runner.py | 90 ++++++++-------- .../agent_framework/_workflows/_executor.py | 2 +- .../agent_framework/_workflows/_runner.py | 22 ++-- .../_workflows/_runner_context.py | 7 ++ .../agent_framework/_workflows/_workflow.py | 2 +- .../_workflows/_workflow_executor.py | 4 +- .../packages/core/tests/workflow/conftest.py | 102 ++++++++++++++++++ .../core/tests/workflow/test_sub_workflow.py | 73 +++++++++++++ 9 files changed, 247 insertions(+), 62 deletions(-) create mode 100644 python/packages/core/tests/workflow/conftest.py diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index f24012d9e21..e7051647f75 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -351,10 +351,13 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID committed snapshots remain unchanged until commit and discard rolls back pending changes. Mixed-parent-scope response batches are rejected before consuming requests, including across sibling bindings. Validation includes queued wrapped replies and runs before each live message - batch is drained, routing wrapped replies to their child's own policy. Cancellation preflight combines + batch is drained, routing wrapped replies to their child's own policy. Preflight uses the edge runner's + actual target/type filtering, conditions, and fan-out selection, and dispatch reuses those decisions + without evaluating predicates again. Fan-in aggregates are not individual wrapped-reply deliveries. + Cancellation preflight combines queued replies with the complete selected cancellation set before consuming any requests. - **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, + replacing live state or queues. `WorkflowExecutor` delegates its async 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. - **Response policy staging** - Live response validation does not copy request payloads. Policies validate without mutation and may return deferred metadata updates; all policies in the workflow tree must succeed diff --git a/python/packages/core/agent_framework/_workflows/_edge_runner.py b/python/packages/core/agent_framework/_workflows/_edge_runner.py index 89d2ecf3a6b..23dc62ae173 100644 --- a/python/packages/core/agent_framework/_workflows/_edge_runner.py +++ b/python/packages/core/agent_framework/_workflows/_edge_runner.py @@ -1,5 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. +from __future__ import annotations + import asyncio import json import logging @@ -7,6 +9,7 @@ from collections import defaultdict from collections.abc import Callable, Coroutine from typing import Any, cast +from uuid import uuid4 from ..exceptions import WorkflowCheckpointException from ..observability import EdgeGroupDeliveryStatus, OtelAttr, create_edge_group_processing_span @@ -64,8 +67,8 @@ def __init__(self, edge_group: EdgeGroup, executors: dict[str, Executor]) -> Non """ self._edge_group = edge_group self._executors = executors + self._delivery_key = str(uuid4()) - @abstractmethod async def send_message( self, message: WorkflowMessage, @@ -83,6 +86,28 @@ async def send_message( bool: True if the message was processed successfully, False if the target executor cannot handle the message. """ + delivery = message._prepared_deliveries.pop(self._delivery_key, None) # pyright: ignore[reportPrivateUsage] + if delivery is None: + delivery = await self._select_targets(message) + processed, target_ids = delivery + source_id = self._edge_group.edges[0].source_id + if len(target_ids) == 1: + await self._execute_on_target(target_ids[0], [source_id], message, state, ctx) + return processed + await gather_cancelling_siblings_on_error( + *(self._execute_on_target(target_id, [source_id], message, state, ctx) for target_id in target_ids) + ) + return processed + + async def _prepare_message(self, message: WorkflowMessage) -> list[str]: + deliveries = message._prepared_deliveries # pyright: ignore[reportPrivateUsage] + if self._delivery_key not in deliveries: + deliveries[self._delivery_key] = await self._select_targets(message) + return deliveries[self._delivery_key][1] + + @abstractmethod + async def _select_targets(self, message: WorkflowMessage) -> tuple[bool, list[str]]: + """Determine delivery targets without executing or buffering the message.""" raise NotImplementedError @property @@ -159,13 +184,8 @@ def __init__(self, edge_group: SingleEdgeGroup | InternalEdgeGroup, executors: d super().__init__(edge_group, executors) self._edge = edge_group.edges[0] - async def send_message( - self, - message: WorkflowMessage, - state: State, - ctx: RunnerContext, - ) -> bool: - """Send a message through the single edge.""" + async def _select_targets(self, message: WorkflowMessage) -> tuple[bool, list[str]]: + """Select the single edge's delivery target without executing it.""" should_execute = False target_id: str | None = None source_id: str | None = None @@ -183,7 +203,7 @@ async def send_message( OtelAttr.EDGE_GROUP_DELIVERED: False, OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TARGET_MISMATCH.value, }) - return False + return False, [] if self._can_handle(self._edge.target_id, message): route_result = await self._edge.should_route(message.data) @@ -202,13 +222,13 @@ async def send_message( OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_CONDITION_FALSE.value, }) # Return True here because message was processed, just condition failed - return True + return True, [] else: span.set_attributes({ OtelAttr.EDGE_GROUP_DELIVERED: False, OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH.value, }) - return False + return False, [] except Exception as e: span.set_attributes({ OtelAttr.EDGE_GROUP_DELIVERED: False, @@ -216,12 +236,10 @@ async def send_message( }) raise e - # Execute outside the span if should_execute and target_id and source_id: - await self._execute_on_target(target_id, [source_id], message, state, ctx) - return True + return True, [target_id] - return False + return False, [] class FanOutEdgeRunner(EdgeRunner): @@ -236,13 +254,8 @@ def __init__(self, edge_group: FanOutEdgeGroup, executors: dict[str, Executor]) Callable[[Any, list[str]], list[str]] | None, getattr(edge_group, "selection_func", None) ) - async def send_message( - self, - message: WorkflowMessage, - state: State, - ctx: RunnerContext, - ) -> bool: - """Send a message through all edges in the fan-out edge group.""" + async def _select_targets(self, message: WorkflowMessage) -> tuple[bool, list[str]]: + """Select fan-out delivery targets without executing them.""" deliverable_edges: list[Edge] = [] single_target_edge: Edge | None = None # Process routing logic within span @@ -269,14 +282,14 @@ async def send_message( OtelAttr.EDGE_GROUP_DELIVERED: False, OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TARGET_MISMATCH.value, }) - return False + return False, [] if not any(self._can_handle(target_id, message) for target_id in self._target_ids): span.set_attributes({ OtelAttr.EDGE_GROUP_DELIVERED: False, OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH.value, }) - return False + return False, [] selection_results = ( self._selection_func(message.data, self._target_ids) if self._selection_func else self._target_ids @@ -313,21 +326,21 @@ async def send_message( OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_CONDITION_FALSE.value, # ruff:ignore[line-too-long] }) # For targeted messages with condition failure, return True (message was processed) - return True + return True, [] else: span.set_attributes({ OtelAttr.EDGE_GROUP_DELIVERED: False, OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH.value, # ruff:ignore[line-too-long] }) # For targeted messages that can't be handled, return False - return False + return False, [] else: span.set_attributes({ OtelAttr.EDGE_GROUP_DELIVERED: False, OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TARGET_MISMATCH.value, }) # For targeted messages not in selection, return False - return False + return False, [] else: # If no target ID, send the message to the selected targets for target_id in selection_results: @@ -355,27 +368,14 @@ async def send_message( }) raise e - # Execute outside the span if single_target_edge: - await self._execute_on_target( - single_target_edge.target_id, - [single_target_edge.source_id], - message, - state, - ctx, - ) - return True + return True, [single_target_edge.target_id] if deliverable_edges: - - async def send_to_edge(edge: Edge) -> None: - await self._execute_on_target(edge.target_id, [edge.source_id], message, state, ctx) - - await gather_cancelling_siblings_on_error(*(send_to_edge(edge) for edge in deliverable_edges)) - return True + return True, [edge.target_id for edge in deliverable_edges] # If we get here, it's a broadcast message with no deliverable edges - return False + return False, [] def _validate_selection_result(self, selection_results: list[str]) -> bool: """Validate the selection results to ensure all IDs are valid target executor IDs.""" @@ -394,6 +394,10 @@ def __init__(self, edge_group: FanInEdgeGroup, executors: dict[str, Executor]) - # Key is the source executor ID, value is a list of messages self._buffer: dict[str, list[WorkflowMessage]] = defaultdict(list) + async def _select_targets(self, message: WorkflowMessage) -> tuple[bool, list[str]]: + # Fan-in delivers an aggregate rather than the wrapped reply itself. + return False, [] + def snapshot_state(self) -> dict[str, Any] | None: """Capture the buffered messages that are still waiting for the remaining sources. diff --git a/python/packages/core/agent_framework/_workflows/_executor.py b/python/packages/core/agent_framework/_workflows/_executor.py index 37421e0dae7..6971d6c3c4b 100644 --- a/python/packages/core/agent_framework/_workflows/_executor.py +++ b/python/packages/core/agent_framework/_workflows/_executor.py @@ -549,7 +549,7 @@ async def _get_response_validation_batches( del requests return [] - def _validate_checkpoint_restore(self, state: dict[str, Any], requests: Sequence[WorkflowEvent] = ()) -> None: + async def _validate_checkpoint_restore(self, state: dict[str, Any], requests: Sequence[WorkflowEvent] = ()) -> None: """Validate a checkpoint snapshot without changing live executor state.""" del state, requests diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index df567648244..5c5624201c1 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -256,7 +256,7 @@ async def _get_queued_messages(self) -> Mapping[str, Sequence[WorkflowMessage]]: ) return checkpoint.messages - def _get_wrapped_response_requests( + async def _get_wrapped_response_requests( self, messages: Mapping[str, Sequence[WorkflowMessage]] ) -> dict[str, list[WorkflowEvent]]: from ._workflow_executor import SubWorkflowResponseMessage @@ -267,13 +267,9 @@ def _get_wrapped_response_requests( if not isinstance(message.data, SubWorkflowResponseMessage): continue # Wrapped replies retain child-local source IDs; route through the parent graph. - 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) - } + targets: set[str] = set() + for runner in self._edge_runner_map.get(message.source_id, []): + targets.update(await runner._prepare_message(message)) # pyright: ignore[reportPrivateUsage] for target_id in targets: requests[target_id].append(message.data.source_event) return dict(requests) @@ -361,7 +357,7 @@ async def restore_from_checkpoint( raise WorkflowCheckpointException(f"Checkpoint {checkpoint_id} not found") checkpoint = copy.deepcopy(checkpoint) - self._validate_checkpoint_tree(checkpoint) + await 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. @@ -421,7 +417,7 @@ async def restore_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None: """ try: checkpoint = copy.deepcopy(checkpoint) - self._validate_checkpoint_tree(checkpoint) + await 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. @@ -437,7 +433,7 @@ 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( + async def _validate_checkpoint_tree( self, checkpoint: WorkflowCheckpoint, delegated_requests: Sequence[WorkflowEvent] = () ) -> None: from ._const import WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY @@ -457,7 +453,7 @@ def _validate_checkpoint_tree( ] if requests and self._response_batch_validator is not None: self._response_batch_validator(requests) - wrapped_requests = self._get_wrapped_response_requests(checkpoint.messages) + wrapped_requests = await self._get_wrapped_response_requests(checkpoint.messages) parent_contexts = checkpoint.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {}) validate_request_contexts( parent_contexts, @@ -478,7 +474,7 @@ def _validate_checkpoint_tree( for target_id, batch in wrapped_requests.items(): delegated_by_executor[target_id].extend(batch) for _, executor, state in self._get_executor_states(checkpoint.state[EXECUTOR_STATE_KEY]): - executor._validate_checkpoint_restore( # pyright: ignore[reportPrivateUsage] + await executor._validate_checkpoint_restore( # pyright: ignore[reportPrivateUsage] state, delegated_by_executor[executor.id] ) diff --git a/python/packages/core/agent_framework/_workflows/_runner_context.py b/python/packages/core/agent_framework/_workflows/_runner_context.py index 71bc526d77c..7e367397e45 100644 --- a/python/packages/core/agent_framework/_workflows/_runner_context.py +++ b/python/packages/core/agent_framework/_workflows/_runner_context.py @@ -9,6 +9,7 @@ from copy import copy from dataclasses import dataclass from enum import Enum +from functools import cached_property from typing import Any, Literal, Protocol, TypeVar, runtime_checkable from .._tools import ToolTypes @@ -52,6 +53,12 @@ class WorkflowMessage: # For response messages, the original request data original_request_info_event: WorkflowEvent[Any] | None = None + @cached_property + def _prepared_deliveries(self) -> dict[str, tuple[bool, list[str]]]: + # Lazy initialization also supports messages unpickled from older checkpoints. + # Per-runner keys prevent reusing decisions in another workflow instance. + return {} + # Backward compatibility properties @property def trace_context(self) -> dict[str, str] | None: diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 5f034ecdc9b..372a0a1316f 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -1216,7 +1216,7 @@ async def _get_response_validation_batches( ) -> list[_ResponseBatchValidation]: messages = await self._runner._get_queued_messages() # pyright: ignore[reportPrivateUsage] queued_requests = _get_queued_response_info_events(messages) - wrapped_requests = self._runner._get_wrapped_response_requests(messages) # pyright: ignore[reportPrivateUsage] + wrapped_requests = await self._runner._get_wrapped_response_requests(messages) # pyright: ignore[reportPrivateUsage] batch = [*queued_requests, *requests] parent_contexts = self._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {}) boundary_contexts = [ diff --git a/python/packages/core/agent_framework/_workflows/_workflow_executor.py b/python/packages/core/agent_framework/_workflows/_workflow_executor.py index e6c9651b452..2feadaf43ae 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_executor.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_executor.py @@ -612,10 +612,10 @@ async def _get_response_validation_batches( ]) @override - def _validate_checkpoint_restore(self, state: dict[str, Any], requests: Sequence[WorkflowEvent] = ()) -> None: + async 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, requests) # pyright: ignore[reportPrivateUsage] + await 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/core/tests/workflow/conftest.py b/python/packages/core/tests/workflow/conftest.py new file mode 100644 index 00000000000..5f3680b4230 --- /dev/null +++ b/python/packages/core/tests/workflow/conftest.py @@ -0,0 +1,102 @@ +# Copyright (c) Microsoft. All rights reserved. + +from collections.abc import Callable +from typing import Any + +import pytest + +from agent_framework import ( + Case, + Default, + Executor, + SubWorkflowRequestMessage, + SubWorkflowResponseMessage, + Workflow, + WorkflowBuilder, + WorkflowContext, + WorkflowExecutor, + handler, + response_handler, +) + + +@pytest.fixture +def selected_reply_workflow() -> Callable[ + [str, bool], tuple[Workflow, dict[str, Workflow], dict[str, SubWorkflowRequestMessage], list[str]] +]: + def build( + routing: str, select_both: bool = False + ) -> tuple[Workflow, dict[str, Workflow], dict[str, SubWorkflowRequestMessage], list[str]]: + requests: dict[str, SubWorkflowRequestMessage] = {} + calls: list[str] = [] + + class Requester(Executor): + @handler + async def start(self, message: str, ctx: WorkflowContext) -> None: + await ctx.request_info(message, bool, request_id="shared") + + @response_handler + async def respond(self, original: str, approved: bool, ctx: WorkflowContext[Any, str]) -> None: + await ctx.yield_output(original) + + class Coordinator(Executor): + @handler + async def start(self, message: str, ctx: WorkflowContext[str]) -> None: + await ctx.send_message(message, target_id=message) + + @handler + async def reply( + self, messages: list[SubWorkflowResponseMessage], ctx: WorkflowContext[SubWorkflowResponseMessage] + ) -> None: + for message in messages: + await ctx.send_message(message) + + @handler + async def intercept(self, request: SubWorkflowRequestMessage, ctx: WorkflowContext) -> None: + requests[request.executor_id] = request + + def select(message: Any, targets: list[str]) -> list[str]: + if isinstance(message, SubWorkflowResponseMessage): + calls.append("selection") + return targets if select_both else ["first"] + return targets + + async def first_condition(message: Any) -> bool: + if isinstance(message, SubWorkflowResponseMessage): + calls.append("first") + return True + + async def second_condition(message: Any) -> bool: + if isinstance(message, SubWorkflowResponseMessage): + calls.append("second") + return select_both + return True + + def first_case(message: Any) -> bool: + if isinstance(message, SubWorkflowResponseMessage): + calls.append("switch") + return True + return message == "first" + + children = { + name: WorkflowBuilder(start_executor=Requester(id="requester")).build() for name in ("first", "second") + } + wrappers = { + name: WorkflowExecutor(child, id=name, allow_direct_output=True) for name, child in children.items() + } + coordinator = Coordinator(id="coordinator") + builder = WorkflowBuilder(start_executor=coordinator) + if routing == "condition": + builder.add_edge(coordinator, wrappers["first"], condition=first_condition) + builder.add_edge(coordinator, wrappers["second"], condition=second_condition) + elif routing == "selection": + builder.add_multi_selection_edge_group(coordinator, list(wrappers.values()), select) + else: + builder.add_switch_case_edge_group( + coordinator, [Case(condition=first_case, target=wrappers["first"]), Default(target=wrappers["second"])] + ) + for wrapper in wrappers.values(): + builder.add_edge(wrapper, coordinator) + return builder.build(), children, requests, calls + + return build diff --git a/python/packages/core/tests/workflow/test_sub_workflow.py b/python/packages/core/tests/workflow/test_sub_workflow.py index e7b03e395b3..4be6d918b93 100644 --- a/python/packages/core/tests/workflow/test_sub_workflow.py +++ b/python/packages/core/tests/workflow/test_sub_workflow.py @@ -406,6 +406,79 @@ def reject(requests: Sequence[WorkflowEvent]) -> None: ] +@pytest.mark.parametrize("routing", ["condition", "selection", "switch"]) +@pytest.mark.parametrize("restore", [False, True]) +@pytest.mark.parametrize("streaming", [False, True]) +async def test_wrapped_reply_preflight_uses_selected_routes( + selected_reply_workflow, routing: str, restore: bool, streaming: bool +) -> None: + from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY + + parent, children, requests, calls = selected_reply_workflow(routing) + await parent.run("first") + first_scope = parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) + await parent.run("second") + assert parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) != first_scope + pending_second = await children["second"]._runner_context.get_pending_request_info_events() + assert requests["first"].source_event.request_id == requests["second"].source_event.request_id == "shared" + reply = requests["first"].create_response(True) + if restore: + await parent._runner_context.send_message(WorkflowMessage(data=reply, source_id="coordinator")) + checkpoint = await parent._runner.build_checkpoint() + storage = InMemoryCheckpointStorage() + checkpoint_id = await storage.save(checkpoint) + run = parent.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage, stream=streaming) + else: + run = parent.run([reply], stream=streaming) + + if streaming: + events = [event async for event in run] + outputs = [event.data for event in events if event.type == "output"] + else: + outputs = (await run).get_outputs() + + assert outputs == ["first"] + assert not await children["first"]._runner_context.get_pending_request_info_events() + remaining = await children["second"]._runner_context.get_pending_request_info_events() + assert {key: event.to_dict() for key, event in remaining.items()} == { + key: event.to_dict() for key, event in pending_second.items() + } + assert parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) == first_scope + assert ( + sorted(calls) == {"condition": ["first", "second"], "selection": ["selection"], "switch": ["switch"]}[routing] + ) + + +@pytest.mark.parametrize("routing", ["condition", "selection"]) +@pytest.mark.parametrize("restore", [False, True]) +async def test_wrapped_reply_selected_mixed_scopes_preserve_queue( + selected_reply_workflow, routing: str, restore: bool +) -> None: + from agent_framework import WorkflowCheckpointException + from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY + + parent, children, requests, _ = selected_reply_workflow(routing, True) + await parent.run("first") + await parent.run("second") + pending = {name: await child._runner_context.get_pending_request_info_events() for name, child in children.items()} + message = WorkflowMessage(data=requests["first"].create_response(True), source_id="coordinator") + await parent._runner_context.send_message(message) + scope = parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) + if restore: + checkpoint = await parent._runner.build_checkpoint() + with pytest.raises(WorkflowCheckpointException) as error: + await parent._runner.restore_checkpoint(checkpoint) + assert isinstance(error.value.__cause__, ValueError) + assert "multiple parent workflow scopes" in str(error.value.__cause__) + else: + with pytest.raises(ValueError, match="multiple parent workflow scopes"): + await parent._runner._run_iteration() + assert parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) == scope + assert await parent._runner_context.get_queued_messages() == {"coordinator": [message]} + for name, child in children.items(): + assert await child._runner_context.get_pending_request_info_events() == pending[name] + + @pytest.mark.parametrize("streaming", [False, True]) async def test_response_validation_preserves_noncopyable_request_payload(streaming: bool) -> None: from threading import Lock From 517ebf12965dd0d5048baff4628721c03ef1d245 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Thu, 29 Oct 2026 11:19:04 -0400 Subject: [PATCH 72/76] fix(workflows): make routing regression types explicit Supply the fixture factory option explicitly, narrow streaming calls through literal overloads, and inspect queues through the runner helper. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 06d37a04-f20d-46d4-ab56-1ef5adda7b2d --- .../core/tests/workflow/test_sub_workflow.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/python/packages/core/tests/workflow/test_sub_workflow.py b/python/packages/core/tests/workflow/test_sub_workflow.py index 4be6d918b93..df17b32a4b9 100644 --- a/python/packages/core/tests/workflow/test_sub_workflow.py +++ b/python/packages/core/tests/workflow/test_sub_workflow.py @@ -414,7 +414,7 @@ async def test_wrapped_reply_preflight_uses_selected_routes( ) -> None: from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY - parent, children, requests, calls = selected_reply_workflow(routing) + parent, children, requests, calls = selected_reply_workflow(routing, False) await parent.run("first") first_scope = parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) await parent.run("second") @@ -422,20 +422,23 @@ async def test_wrapped_reply_preflight_uses_selected_routes( pending_second = await children["second"]._runner_context.get_pending_request_info_events() assert requests["first"].source_event.request_id == requests["second"].source_event.request_id == "shared" reply = requests["first"].create_response(True) + checkpoint_id: str | None = None + storage: InMemoryCheckpointStorage | None = None if restore: await parent._runner_context.send_message(WorkflowMessage(data=reply, source_id="coordinator")) checkpoint = await parent._runner.build_checkpoint() storage = InMemoryCheckpointStorage() checkpoint_id = await storage.save(checkpoint) - run = parent.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage, stream=streaming) - else: - run = parent.run([reply], stream=streaming) + message = None if restore else [reply] if streaming: - events = [event async for event in run] + events = [ + event + async for event in parent.run(message, checkpoint_id=checkpoint_id, checkpoint_storage=storage, stream=True) + ] outputs = [event.data for event in events if event.type == "output"] else: - outputs = (await run).get_outputs() + outputs = (await parent.run(message, checkpoint_id=checkpoint_id, checkpoint_storage=storage)).get_outputs() assert outputs == ["first"] assert not await children["first"]._runner_context.get_pending_request_info_events() @@ -474,7 +477,7 @@ async def test_wrapped_reply_selected_mixed_scopes_preserve_queue( with pytest.raises(ValueError, match="multiple parent workflow scopes"): await parent._runner._run_iteration() assert parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) == scope - assert await parent._runner_context.get_queued_messages() == {"coordinator": [message]} + assert await parent._runner._get_queued_messages() == {"coordinator": [message]} for name, child in children.items(): assert await child._runner_context.get_pending_request_info_events() == pending[name] From 0577d313fb968013dbc2a5a1fe497cd0257e7426 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Thu, 29 Oct 2026 12:44:21 -0400 Subject: [PATCH 73/76] fix(workflows): add state-dependent reply routing regression coverage Reproduce routing decisions cached before earlier handlers run and before checkpointed executor fields are restored. These eight regression cases currently fail; the implementation correction is pending routing-contract guidance. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 06d37a04-f20d-46d4-ab56-1ef5adda7b2d --- .../packages/core/tests/workflow/conftest.py | 46 ++++++++++++--- .../core/tests/workflow/test_sub_workflow.py | 59 ++++++++++++++++++- 2 files changed, 94 insertions(+), 11 deletions(-) diff --git a/python/packages/core/tests/workflow/conftest.py b/python/packages/core/tests/workflow/conftest.py index 5f3680b4230..da126fe0e54 100644 --- a/python/packages/core/tests/workflow/conftest.py +++ b/python/packages/core/tests/workflow/conftest.py @@ -22,10 +22,10 @@ @pytest.fixture def selected_reply_workflow() -> Callable[ - [str, bool], tuple[Workflow, dict[str, Workflow], dict[str, SubWorkflowRequestMessage], list[str]] + [str, bool, bool], tuple[Workflow, dict[str, Workflow], dict[str, SubWorkflowRequestMessage], list[str]] ]: def build( - routing: str, select_both: bool = False + routing: str, select_both: bool = False, stateful: bool = False ) -> tuple[Workflow, dict[str, Workflow], dict[str, SubWorkflowRequestMessage], list[str]]: requests: dict[str, SubWorkflowRequestMessage] = {} calls: list[str] = [] @@ -40,6 +40,16 @@ async def respond(self, original: str, approved: bool, ctx: WorkflowContext[Any, await ctx.yield_output(original) class Coordinator(Executor): + def __init__(self) -> None: + super().__init__(id="coordinator") + self.target = "first" + + async def on_checkpoint_save(self) -> dict[str, Any]: + return {"target": self.target} + + async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: + self.target = state["target"] + @handler async def start(self, message: str, ctx: WorkflowContext[str]) -> None: await ctx.send_message(message, target_id=message) @@ -55,22 +65,41 @@ async def reply( async def intercept(self, request: SubWorkflowRequestMessage, ctx: WorkflowContext) -> None: requests[request.executor_id] = request + @handler + async def queue_ordered( + self, + messages: tuple[str, SubWorkflowResponseMessage], + ctx: WorkflowContext[str | SubWorkflowResponseMessage], + ) -> None: + for message in messages: + await ctx.send_message(message, target_id="first") + + class RoutingWorkflowExecutor(WorkflowExecutor): + @handler + async def apply_routing(self, message: str, ctx: WorkflowContext[Any, Any]) -> None: + if message == "enable-first": + coordinator.target = "first" + else: + await self.process_workflow(message, ctx) + def select(message: Any, targets: list[str]) -> list[str]: if isinstance(message, SubWorkflowResponseMessage): calls.append("selection") - return targets if select_both else ["first"] + return targets if select_both else [coordinator.target if stateful else "first"] return targets async def first_condition(message: Any) -> bool: if isinstance(message, SubWorkflowResponseMessage): calls.append("first") + if stateful: + return coordinator.target == "first" return True async def second_condition(message: Any) -> bool: if isinstance(message, SubWorkflowResponseMessage): calls.append("second") - return select_both - return True + return select_both or (stateful and coordinator.target == "second") + return coordinator.target == "first" if stateful else True def first_case(message: Any) -> bool: if isinstance(message, SubWorkflowResponseMessage): @@ -81,10 +110,9 @@ def first_case(message: Any) -> bool: children = { name: WorkflowBuilder(start_executor=Requester(id="requester")).build() for name in ("first", "second") } - wrappers = { - name: WorkflowExecutor(child, id=name, allow_direct_output=True) for name, child in children.items() - } - coordinator = Coordinator(id="coordinator") + wrapper_type = RoutingWorkflowExecutor if stateful else WorkflowExecutor + wrappers = {name: wrapper_type(child, id=name, allow_direct_output=True) for name, child in children.items()} + coordinator = Coordinator() builder = WorkflowBuilder(start_executor=coordinator) if routing == "condition": builder.add_edge(coordinator, wrappers["first"], condition=first_condition) diff --git a/python/packages/core/tests/workflow/test_sub_workflow.py b/python/packages/core/tests/workflow/test_sub_workflow.py index df17b32a4b9..f72fe3650c9 100644 --- a/python/packages/core/tests/workflow/test_sub_workflow.py +++ b/python/packages/core/tests/workflow/test_sub_workflow.py @@ -414,7 +414,7 @@ async def test_wrapped_reply_preflight_uses_selected_routes( ) -> None: from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY - parent, children, requests, calls = selected_reply_workflow(routing, False) + parent, children, requests, calls = selected_reply_workflow(routing, False, False) await parent.run("first") first_scope = parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) await parent.run("second") @@ -460,7 +460,7 @@ async def test_wrapped_reply_selected_mixed_scopes_preserve_queue( from agent_framework import WorkflowCheckpointException from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY - parent, children, requests, _ = selected_reply_workflow(routing, True) + parent, children, requests, _ = selected_reply_workflow(routing, True, False) await parent.run("first") await parent.run("second") pending = {name: await child._runner_context.get_pending_request_info_events() for name, child in children.items()} @@ -482,6 +482,61 @@ async def test_wrapped_reply_selected_mixed_scopes_preserve_queue( assert await child._runner_context.get_pending_request_info_events() == pending[name] +@pytest.mark.parametrize("routing", ["condition", "selection"]) +@pytest.mark.parametrize("streaming", [False, True]) +async def test_wrapped_reply_routing_observes_earlier_same_edge_handler( + selected_reply_workflow, routing: str, streaming: bool +) -> None: + parent, children, requests, calls = selected_reply_workflow(routing, False, True) + await parent.run("first") + coordinator = parent.get_start_executor() + await coordinator.on_checkpoint_restore({"target": "second"}) + message = ("enable-first", requests["first"].create_response(True)) + + if streaming: + events = [event async for event in parent.run(message, stream=True)] + outputs = [event.data for event in events if event.type == "output"] + else: + outputs = (await parent.run(message)).get_outputs() + + assert await coordinator.on_checkpoint_save() == {"target": "first"} + assert outputs == ["first"] + assert not await children["first"]._runner_context.get_pending_request_info_events() + assert calls == {"condition": ["first"], "selection": ["selection"]}[routing] + + +@pytest.mark.parametrize("routing", ["condition", "selection"]) +@pytest.mark.parametrize("streaming", [False, True]) +async def test_wrapped_reply_routing_uses_checkpointed_executor_fields( + selected_reply_workflow, routing: str, streaming: bool +) -> None: + parent, children, requests, calls = selected_reply_workflow(routing, False, True) + await parent.run("first") + await parent.run("second") + coordinator = parent.get_start_executor() + await coordinator.on_checkpoint_restore({"target": "second"}) + await parent._runner_context.send_message( + WorkflowMessage(data=requests["second"].create_response(True), source_id="coordinator") + ) + storage = InMemoryCheckpointStorage() + checkpoint_id = await storage.save(await parent._runner.build_checkpoint()) + await coordinator.on_checkpoint_restore({"target": "first"}) + + if streaming: + events = [ + event async for event in parent.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage, stream=True) + ] + outputs = [event.data for event in events if event.type == "output"] + else: + outputs = (await parent.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage)).get_outputs() + + assert await coordinator.on_checkpoint_save() == {"target": "second"} + assert outputs == ["second"] + assert not await children["second"]._runner_context.get_pending_request_info_events() + assert set(await children["first"]._runner_context.get_pending_request_info_events()) == {"shared"} + assert sorted(calls) == {"condition": ["first", "second"], "selection": ["selection"]}[routing] + + @pytest.mark.parametrize("streaming", [False, True]) async def test_response_validation_preserves_noncopyable_request_payload(streaming: bool) -> None: from threading import Lock From 1e4d973b4e9065ccdb5650c2a6c9911959b1a3b7 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Thu, 29 Oct 2026 13:44:54 -0400 Subject: [PATCH 74/76] fix(workflows): validate wrapped replies at delivery time Select wrapped reply routes after preceding handlers and checkpoint restoration, validate all selected children before delivery, and retain selected parent scopes through the superstep. Remove message-bound route caches and preserve undelivered messages on rejection. Document the delivery-time validation boundary and checkpoint recovery after partial execution. Preserve ordinary response and structural checkpoint preflight. Validate routing regressions, child-policy staging and the revised rejection boundary without skipping tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 06d37a04-f20d-46d4-ab56-1ef5adda7b2d --- ...thon-declarative-request-scope-recovery.md | 52 +++++--- python/packages/core/AGENTS.md | 17 +-- .../_workflows/_edge_runner.py | 29 +++-- .../_workflows/_request_context.py | 10 +- .../agent_framework/_workflows/_runner.py | 113 ++++++++++++------ .../_workflows/_runner_context.py | 7 -- .../agent_framework/_workflows/_workflow.py | 40 +++++-- .../packages/core/tests/workflow/conftest.py | 4 +- .../core/tests/workflow/test_sub_workflow.py | 97 +++++++++++---- python/packages/declarative/AGENTS.md | 9 +- .../tests/test_invoke_mcp_tool_executor.py | 60 ++++------ .../references/workflow-builder.md | 3 + 12 files changed, 288 insertions(+), 153 deletions(-) diff --git a/docs/decisions/0046-python-declarative-request-scope-recovery.md b/docs/decisions/0046-python-declarative-request-scope-recovery.md index 88ff251e982..b915e4719f4 100644 --- a/docs/decisions/0046-python-declarative-request-scope-recovery.md +++ b/docs/decisions/0046-python-declarative-request-scope-recovery.md @@ -44,9 +44,10 @@ 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. +- Checkpoint restores validate ordinary queued response origins before restoring + shared or executor state, including embedded child checkpoints. Submitted and + ordinary queued responses share batch validation before pending requests are consumed. + Wrapped child replies validate selected routes at delivery, after executor restoration. - 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 @@ -108,11 +109,11 @@ is checkpointed with its owning workflow, including at every embedded boundary. After a child successfully resumes or cancels a request, restore the saved parent context before forwarding outputs or capturing the origins of follow-up requests. Release completed origins, preserve still-pending origins across new inputs, and -reject mixed-parent-scope response batches before consuming any requests, including -batches spanning sibling bindings. Cancellation batches obey the same parent-scope -rule because cancellation can also produce downstream outputs. -Checkpoint preflight applies the same validation -to queued propagated and wrapped replies before replacing live state. +reject mixed-parent-scope ordinary response batches before consuming any requests, +including batches spanning sibling bindings. Cancellation batches obey the same +parent-scope rule because cancellation can also produce downstream outputs. +Checkpoint preflight applies the same validation to queued propagated replies +before replacing live state. Wrapped replies use the delivery boundary below. Older checkpoints without boundary origins cannot reconstruct those ancestors' original scopes. A successful child continuation therefore receives a new isolated @@ -138,11 +139,8 @@ with committed values unchanged until the superstep commits. Discard rolls back these entry updates. Full-map validation and detached copies remain checkpoint and response-batch operations, not per-request insertion or removal work. -Preflight includes both propagated responses and wrapped child replies already -queued for delivery. Route wrapped replies through the parent graph to validate -their child's own policy without treating child scopes as parent-policy scopes. -Repeat validation before each live message batch is drained, so a rejected batch -retains its queued replies. Cancellation validates queued replies together with +Preflight includes propagated responses already queued for delivery. Repeat +validation before each live message batch is drained. Cancellation validates ordinary queued replies together with the complete selected cancellation set before consuming requests or invoking cancellation handlers. Legacy runner contexts use a non-persisting snapshot when optional queued-message inspection is unavailable. @@ -153,7 +151,33 @@ response policies validate without mutation and return deferred updates to the scope metadata they own. Apply these updates only after every policy in the tree has validated. This preserves cross-child rejection atomicity without copying user payloads or changing their identity. Stored checkpoint restoration still preflights -its detached candidate tree before replacing live state. +ordinary responses in its detached candidate tree before replacing live state. + +## State-dependent wrapped-reply routing + +Edge callbacks can read executor fields changed by preceding handlers or restored +from a checkpoint. Evaluating them during queue or detached-checkpoint preflight +and caching their decisions changes normal routing behavior. Copying executors +does not rebind arbitrary callback closures, and speculative handler or restore-hook +execution can have external side effects. An explicit stable/snapshot-aware routing +API would impose a new contract on existing graphs. + +Instead, finish preceding messages from a source before selecting a wrapped +reply's targets. Evaluate each edge callback once, validate all selected targets +together against their own child policies, then deliver that reply. Track selected +parent scopes through the superstep so a later delivery cannot introduce a +different scope after earlier origins have been released. Normal message batches +retain concurrent edge delivery; wrapped replies form delivery barriers, with +their selected targets still executing concurrently. + +This deliberately narrows atomicity for wrapped replies: rejection retains the +rejected and subsequent undelivered messages, but cannot undo earlier handlers, +child continuations, external effects, or already restored checkpoint state. +Recover a partially executed failed run from a known-good checkpoint rather than +retrying its remaining queue. Ordinary response/cancellation batches and detached +structural and ordinary-response checkpoint validation retain their preflight +guarantees. No callback is assumed pure, and no executor is cloned or replayed +speculatively. The proposed deciders are Python code owners; maintainer and engineering-management approval is still required. diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index e7051647f75..7f8bbd78dd5 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -349,16 +349,17 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID are restored only once. Wrapped origins belong to their binding independently of external requests, even when a coordinator relays the same request ID. Origin insertion and removal stage individual mapping entries; committed snapshots remain unchanged until commit and discard rolls back pending changes. - Mixed-parent-scope response batches are rejected before consuming requests, including - across sibling bindings. Validation includes queued wrapped replies and runs before each live message - batch is drained, routing wrapped replies to their child's own policy. Preflight uses the edge runner's - actual target/type filtering, conditions, and fan-out selection, and dispatch reuses those decisions - without evaluating predicates again. Fan-in aggregates are not individual wrapped-reply deliveries. - Cancellation preflight combines - queued replies with the complete selected cancellation set before consuming any requests. + Ordinary response and cancellation batches reject mixed parent scopes before consuming requests. + Wrapped child replies select routes immediately before delivery, after preceding handlers on their source + have completed, and validate all selected targets together using each child's own policy. Predicates run once + per delivery, not during checkpoint or queue preflight. Fan-in aggregates are not individual wrapped replies. + Selected replies cannot introduce a different parent scope later in the same superstep. A rejected reply + and subsequent messages remain queued, but earlier handlers and child continuations are not rolled back. + Recover a partially executed failed run from a known-good checkpoint rather than retrying its remaining queue. - **Nested checkpoint restoration** - Stored and embedded restores preflight a detached checkpoint tree before replacing live state or queues. `WorkflowExecutor` delegates its async 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. + so ordinary queued responses use each child's own policy and migration metadata stays in the candidate snapshot. + Wrapped-reply policies run only at delivery after executor restoration; rejection does not undo that restoration. - **Response policy staging** - Live response validation does not copy request payloads. Policies validate without mutation and may return deferred metadata updates; all policies in the workflow tree must succeed before those updates are applied. This preserves in-process non-copyable payloads and atomic legacy-scope diff --git a/python/packages/core/agent_framework/_workflows/_edge_runner.py b/python/packages/core/agent_framework/_workflows/_edge_runner.py index 23dc62ae173..0448be1f197 100644 --- a/python/packages/core/agent_framework/_workflows/_edge_runner.py +++ b/python/packages/core/agent_framework/_workflows/_edge_runner.py @@ -9,7 +9,6 @@ from collections import defaultdict from collections.abc import Callable, Coroutine from typing import Any, cast -from uuid import uuid4 from ..exceptions import WorkflowCheckpointException from ..observability import EdgeGroupDeliveryStatus, OtelAttr, create_edge_group_processing_span @@ -67,7 +66,6 @@ def __init__(self, edge_group: EdgeGroup, executors: dict[str, Executor]) -> Non """ self._edge_group = edge_group self._executors = executors - self._delivery_key = str(uuid4()) async def send_message( self, @@ -86,9 +84,15 @@ async def send_message( bool: True if the message was processed successfully, False if the target executor cannot handle the message. """ - delivery = message._prepared_deliveries.pop(self._delivery_key, None) # pyright: ignore[reportPrivateUsage] - if delivery is None: - delivery = await self._select_targets(message) + return await self._send_selected_message(message, state, ctx, await self._select_targets(message)) + + async def _send_selected_message( + self, + message: WorkflowMessage, + state: State, + ctx: RunnerContext, + delivery: tuple[bool, list[str]], + ) -> bool: processed, target_ids = delivery source_id = self._edge_group.edges[0].source_id if len(target_ids) == 1: @@ -99,12 +103,6 @@ async def send_message( ) return processed - async def _prepare_message(self, message: WorkflowMessage) -> list[str]: - deliveries = message._prepared_deliveries # pyright: ignore[reportPrivateUsage] - if self._delivery_key not in deliveries: - deliveries[self._delivery_key] = await self._select_targets(message) - return deliveries[self._delivery_key][1] - @abstractmethod async def _select_targets(self, message: WorkflowMessage) -> tuple[bool, list[str]]: """Determine delivery targets without executing or buffering the message.""" @@ -398,6 +396,15 @@ async def _select_targets(self, message: WorkflowMessage) -> tuple[bool, list[st # Fan-in delivers an aggregate rather than the wrapped reply itself. return False, [] + async def _send_selected_message( + self, + message: WorkflowMessage, + state: State, + ctx: RunnerContext, + delivery: tuple[bool, list[str]], + ) -> bool: + return await self.send_message(message, state, ctx) + def snapshot_state(self) -> dict[str, Any] | None: """Capture the buffered messages that are still waiting for the remaining sources. diff --git a/python/packages/core/agent_framework/_workflows/_request_context.py b/python/packages/core/agent_framework/_workflows/_request_context.py index e2e3622886c..e090fe3ca1c 100644 --- a/python/packages/core/agent_framework/_workflows/_request_context.py +++ b/python/packages/core/agent_framework/_workflows/_request_context.py @@ -4,7 +4,7 @@ import logging import uuid -from collections.abc import Sequence +from collections.abc import Collection, Sequence from typing import Any, cast from ._const import ( @@ -60,9 +60,10 @@ def validate_request_contexts( requests: Sequence[WorkflowEvent] = (), *, boundary_contexts: Sequence[tuple[Any, Sequence[WorkflowEvent]]] = (), -) -> None: - """Preflight all origin owners together before any responses are consumed.""" - scopes: set[str] = set() + previous_scopes: Collection[str] = (), +) -> set[str]: + """Validate selected origins together with earlier deliveries in the same batch.""" + scopes = set(previous_scopes) for pending_contexts, pending_requests in ((contexts, requests), *boundary_contexts): if not isinstance(pending_contexts, dict): raise ValueError("Invalid parent workflow request contexts.") @@ -78,6 +79,7 @@ def validate_request_contexts( ) if len(scopes) > 1: raise ValueError("Responses span multiple parent workflow scopes; submit each scope in a separate run.") + return scopes def capture_request_context(state: State, request_id: str, *, boundary_id: str | None = None) -> None: diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index 5c5624201c1..795cb5c3e34 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -90,7 +90,10 @@ def __init__( # Checkpointing related attributes self._previous_checkpoint_id: CheckpointID | None = None self._response_batch_validator: Callable[[Sequence[WorkflowEvent]], None] | None = None - self._queued_response_validator: Callable[[], Awaitable[None]] | None = None + self._queued_response_validator: Callable[[set[str]], Awaitable[None]] | None = None + self._wrapped_response_validator: ( + Callable[[Mapping[str, Sequence[WorkflowEvent]], set[str]], Awaitable[None]] | None + ) = None @property def context(self) -> RunnerContext: @@ -212,7 +215,13 @@ async def _run_iteration(self) -> None: a source by specifying the target, the messages will be delivered to the specified targets in the order they were sent. This is because all messages go through the same edge runner instance which preserves message order. + - Wrapped replies wait for preceding messages from their source, then select and validate all + targets before concurrent delivery. Other sources remain concurrent. """ + from ._workflow_executor import SubWorkflowResponseMessage + + scopes: set[str] = set() + next_message: dict[str, int] = {} async def _deliver_messages(source_executor_id: str, source_messages: list[WorkflowMessage]) -> None: """Outer loop to concurrently deliver messages from all sources to their targets.""" @@ -221,10 +230,12 @@ async def _deliver_message_inner(edge_runner: EdgeRunner, message: WorkflowMessa """Inner loop to deliver a single message through an edge runner.""" return await edge_runner.send_message(message, self._state, self._ctx) - async def _deliver_messages_for_edge_runner(edge_runner: EdgeRunner) -> None: + async def _deliver_messages_for_edge_runner( + edge_runner: EdgeRunner, messages: Sequence[WorkflowMessage] + ) -> None: # Preserve message order per edge runner (and therefore per routed target path) # while still allowing parallelism across different edge runners. - for message in source_messages: + for message in messages: await _deliver_message_inner(edge_runner, message) # Route all messages through normal workflow edges @@ -232,20 +243,75 @@ async def _deliver_messages_for_edge_runner(edge_runner: EdgeRunner) -> None: if not associated_edge_runners: # This is expected for terminal nodes (e.g., EndWorkflow, last action in workflow) logger.debug(f"No outgoing edges found for executor {source_executor_id}; dropping messages.") + next_message[source_executor_id] = len(source_messages) return - tasks = [_deliver_messages_for_edge_runner(edge_runner) for edge_runner in associated_edge_runners] - await gather_cancelling_siblings_on_error(*tasks) + start = 0 + for index, message in enumerate(source_messages): + if not isinstance(message.data, SubWorkflowResponseMessage): + continue + # Earlier handlers must finish before state-dependent reply routes are selected. + next_message[source_executor_id] = index + if index > start: + await gather_cancelling_siblings_on_error( + *( + _deliver_messages_for_edge_runner(runner, source_messages[start:index]) + for runner in associated_edge_runners + ) + ) + deliveries = [ + (runner, await runner._select_targets(message)) # pyright: ignore[reportPrivateUsage] + for runner in associated_edge_runners + ] + targets = {target for _, (_, target_ids) in deliveries for target in target_ids} + if self._wrapped_response_validator is not None: + await self._wrapped_response_validator( + {target: [message.data.source_event] for target in targets}, scopes + ) + next_message[source_executor_id] = index + 1 + await gather_cancelling_siblings_on_error( + *( + runner._send_selected_message( # pyright: ignore[reportPrivateUsage] + message, self._state, self._ctx, delivery + ) + for runner, delivery in deliveries + ) + ) + start = index + 1 + next_message[source_executor_id] = len(source_messages) + if start < len(source_messages): + await gather_cancelling_siblings_on_error( + *( + _deliver_messages_for_edge_runner(runner, source_messages[start:]) + for runner in associated_edge_runners + ) + ) if self._queued_response_validator is not None: - await self._queued_response_validator() + await self._queued_response_validator(scopes) message_batches = await self._ctx.drain_messages() - await gather_cancelling_siblings_on_error( - *( - _deliver_messages(source_executor_id, source_messages) - for source_executor_id, source_messages in message_batches.items() + try: + await gather_cancelling_siblings_on_error( + *( + _deliver_messages(source_executor_id, source_messages) + for source_executor_id, source_messages in message_batches.items() + ) ) - ) + except BaseException: + remaining = [ + message + for source_id, messages in message_batches.items() + for message in messages[next_message.get(source_id, 0) :] + ] + if remaining: + # Siblings have stopped; retain undelivered replies ahead of newly emitted messages. + queued = await self._ctx.drain_messages() + for message in remaining: + await self._ctx.send_message(message) + for messages in queued.values(): + for message in messages: + await self._ctx.send_message(message) + raise async def _get_queued_messages(self) -> Mapping[str, Sequence[WorkflowMessage]]: if isinstance(self._ctx, _SupportsQueuedMessages): @@ -256,24 +322,6 @@ async def _get_queued_messages(self) -> Mapping[str, Sequence[WorkflowMessage]]: ) return checkpoint.messages - async def _get_wrapped_response_requests( - self, messages: Mapping[str, Sequence[WorkflowMessage]] - ) -> dict[str, list[WorkflowEvent]]: - from ._workflow_executor import SubWorkflowResponseMessage - - requests: defaultdict[str, list[WorkflowEvent]] = defaultdict(list) - for batch in messages.values(): - for message in batch: - if not isinstance(message.data, SubWorkflowResponseMessage): - continue - # Wrapped replies retain child-local source IDs; route through the parent graph. - targets: set[str] = set() - for runner in self._edge_runner_map.get(message.source_id, []): - targets.update(await runner._prepare_message(message)) # pyright: ignore[reportPrivateUsage] - for target_id in targets: - requests[target_id].append(message.data.source_event) - return dict(requests) - async def _prepare_checkpoint_state(self) -> None: """Persist executor snapshots into committed shared state. @@ -453,7 +501,6 @@ async def _validate_checkpoint_tree( ] if requests and self._response_batch_validator is not None: self._response_batch_validator(requests) - wrapped_requests = await self._get_wrapped_response_requests(checkpoint.messages) parent_contexts = checkpoint.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {}) validate_request_contexts( parent_contexts, @@ -461,18 +508,16 @@ async def _validate_checkpoint_tree( boundary_contexts=[ ( checkpoint.state.get(request_context_state_key(executor_id), parent_contexts), - wrapped_requests.get(executor_id, []), + (), ) for executor_id in self._executors - if executor_id in wrapped_requests or request_context_state_key(executor_id) in checkpoint.state + if request_context_state_key(executor_id) in checkpoint.state ], ) 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) - for target_id, batch in wrapped_requests.items(): - delegated_by_executor[target_id].extend(batch) for _, executor, state in self._get_executor_states(checkpoint.state[EXECUTOR_STATE_KEY]): await executor._validate_checkpoint_restore( # pyright: ignore[reportPrivateUsage] state, delegated_by_executor[executor.id] diff --git a/python/packages/core/agent_framework/_workflows/_runner_context.py b/python/packages/core/agent_framework/_workflows/_runner_context.py index 7e367397e45..71bc526d77c 100644 --- a/python/packages/core/agent_framework/_workflows/_runner_context.py +++ b/python/packages/core/agent_framework/_workflows/_runner_context.py @@ -9,7 +9,6 @@ from copy import copy from dataclasses import dataclass from enum import Enum -from functools import cached_property from typing import Any, Literal, Protocol, TypeVar, runtime_checkable from .._tools import ToolTypes @@ -53,12 +52,6 @@ class WorkflowMessage: # For response messages, the original request data original_request_info_event: WorkflowEvent[Any] | None = None - @cached_property - def _prepared_deliveries(self) -> dict[str, tuple[bool, list[str]]]: - # Lazy initialization also supports messages unpickled from older checkpoints. - # Per-runner keys prevent reusing decisions in another workflow instance. - return {} - # Backward compatibility properties @property def trace_context(self) -> dict[str, str] | None: diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 372a0a1316f..d738c57d8c6 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -428,8 +428,11 @@ def __init__( self._active_run: weakref.ref[ResponseStream[WorkflowEvent, WorkflowRunResult]] | None = None self._response_batch_validator: _ResponseBatchValidator | None = None self._runner._response_batch_validator = self._validate_response_batch # pyright: ignore[reportPrivateUsage] - self._runner._queued_response_validator = functools.partial( # pyright: ignore[reportPrivateUsage] - self._validate_response_requests, () + self._runner._queued_response_validator = ( # pyright: ignore[reportPrivateUsage] + lambda scopes: self._validate_response_requests((), scopes=scopes) + ) + self._runner._wrapped_response_validator = ( # pyright: ignore[reportPrivateUsage] + lambda wrapped, scopes: self._validate_response_requests((), wrapped_requests=wrapped, scopes=scopes) ) self._initial_state_initializer: Callable[[State], None] | None = None @@ -1204,19 +1207,32 @@ def _validate_response_batch(self, requests: Sequence[WorkflowEvent]) -> None: if migration is not None: migration() - async def _validate_response_requests(self, requests: Sequence[WorkflowEvent]) -> None: - validations = await self._get_response_validation_batches(requests) + async def _validate_response_requests( + self, + requests: Sequence[WorkflowEvent], + *, + wrapped_requests: Mapping[str, Sequence[WorkflowEvent]] | None = None, + scopes: set[str] | None = None, + ) -> None: + validations = await self._get_response_validation_batches( + requests, wrapped_requests=wrapped_requests, scopes=scopes + ) migrations = [validator(batch) for validator, batch in validations] for migration in migrations: if migration is not None: migration() async def _get_response_validation_batches( - self, requests: Sequence[WorkflowEvent] + self, + requests: Sequence[WorkflowEvent], + *, + wrapped_requests: Mapping[str, Sequence[WorkflowEvent]] | None = None, + scopes: set[str] | None = None, ) -> list[_ResponseBatchValidation]: messages = await self._runner._get_queued_messages() # pyright: ignore[reportPrivateUsage] queued_requests = _get_queued_response_info_events(messages) - wrapped_requests = await self._runner._get_wrapped_response_requests(messages) # pyright: ignore[reportPrivateUsage] + if wrapped_requests is None: + wrapped_requests = {} batch = [*queued_requests, *requests] parent_contexts = self._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {}) boundary_contexts = [ @@ -1228,8 +1244,16 @@ async def _get_response_validation_batches( if executor_id in wrapped_requests or self._runner.state.has(request_context_state_key(executor_id)) ] - def validate_parent_contexts(_requests: Sequence[WorkflowEvent]) -> None: - validate_request_contexts(parent_contexts, batch, boundary_contexts=boundary_contexts) + def validate_parent_contexts(_requests: Sequence[WorkflowEvent]) -> Callable[[], None] | None: + selected_scopes = validate_request_contexts( + parent_contexts, + batch, + boundary_contexts=boundary_contexts, + previous_scopes=scopes if scopes is not None else (), + ) + if scopes is not None: + return functools.partial(scopes.update, selected_scopes) + return None validations: list[_ResponseBatchValidation] = [] if self._response_batch_validator is not None: diff --git a/python/packages/core/tests/workflow/conftest.py b/python/packages/core/tests/workflow/conftest.py index da126fe0e54..053bd4fad0d 100644 --- a/python/packages/core/tests/workflow/conftest.py +++ b/python/packages/core/tests/workflow/conftest.py @@ -4,7 +4,6 @@ from typing import Any import pytest - from agent_framework import ( Case, Default, @@ -52,7 +51,8 @@ async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: @handler async def start(self, message: str, ctx: WorkflowContext[str]) -> None: - await ctx.send_message(message, target_id=message) + for target in ("first", "second") if message == "both" else (message,): + await ctx.send_message(target, target_id=target) @handler async def reply( diff --git a/python/packages/core/tests/workflow/test_sub_workflow.py b/python/packages/core/tests/workflow/test_sub_workflow.py index f72fe3650c9..7f0759beb88 100644 --- a/python/packages/core/tests/workflow/test_sub_workflow.py +++ b/python/packages/core/tests/workflow/test_sub_workflow.py @@ -8,8 +8,6 @@ from uuid import uuid4 import pytest -from typing_extensions import Never - from agent_framework import ( Executor, FunctionExecutor, @@ -29,6 +27,7 @@ from agent_framework._workflows._checkpoint import InMemoryCheckpointStorage from agent_framework._workflows._runner_context import InProcRunnerContext, WorkflowMessage from agent_framework._workflows._state import State +from typing_extensions import Never # Test message types @@ -330,7 +329,7 @@ def _build_wrapped_reply_workflow() -> tuple[Workflow, Workflow, _WrappedReplyCo @pytest.mark.parametrize("streaming", [False, True]) @pytest.mark.parametrize("delivery", ["live", "queued_and_new"]) -async def test_wrapped_reply_batches_reject_mixed_parent_origins(delivery: str, streaming: bool) -> None: +async def test_wrapped_reply_delivery_rejects_later_mixed_origins(delivery: str, streaming: bool) -> None: from agent_framework._workflows._const import WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, WORKFLOW_SESSION_ID_KEY parent, child, coordinator = _build_wrapped_reply_workflow() @@ -342,6 +341,7 @@ async def test_wrapped_reply_batches_reject_mixed_parent_origins(delivery: str, child_pending = await child._runner_context.get_pending_request_info_events() parent_pending = await parent._runner_context.get_pending_request_info_events() origins = parent._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY) + checkpoint = await parent._runner.build_checkpoint() replies = [coordinator.requests[email].create_response(True) for email in emails] if delivery == "queued_and_new": await parent._runner_context.send_message( @@ -357,14 +357,32 @@ async def test_wrapped_reply_batches_reject_mixed_parent_origins(delivery: str, await parent.run(message, responses=responses) assert coordinator.outputs == [] - assert await child._runner_context.get_pending_request_info_events() == child_pending - assert await parent._runner_context.get_pending_request_info_events() == parent_pending + expected_child_pending = dict(child_pending) + expected_parent_pending = dict(parent_pending) + if delivery == "live": + del expected_child_pending[replies[0].source_event.request_id] + else: + del expected_parent_pending[emails[1]] + assert await child._runner_context.get_pending_request_info_events() == expected_child_pending + assert await parent._runner_context.get_pending_request_info_events() == expected_parent_pending assert parent._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY) == origins queued = await parent._runner_context.drain_messages() - assert [message.data.source_event.request_id for batch in queued.values() for message in batch] == [ - reply.source_event.request_id for reply in replies[: 1 if delivery == "queued_and_new" else 2] - ] - + assert [ + message.data.source_event.request_id + for batch in queued.values() + for message in batch + if isinstance(message.data, SubWorkflowResponseMessage) + ] == [reply.source_event.request_id for reply in (replies if delivery == "queued_and_new" else replies[1:])] + + await parent._runner.restore_checkpoint(checkpoint) + restored_child = await child._runner_context.get_pending_request_info_events() + restored_parent = await parent._runner_context.get_pending_request_info_events() + assert {key: event.to_dict() for key, event in restored_child.items()} == { + key: event.to_dict() for key, event in child_pending.items() + } + assert {key: event.to_dict() for key, event in restored_parent.items()} == { + key: event.to_dict() for key, event in parent_pending.items() + } for email, scope in reversed(list(zip(emails, scopes, strict=True))): await parent.run(responses={email: True}) assert coordinator.outputs[-1] == (email, scope) @@ -374,7 +392,7 @@ async def test_wrapped_reply_batches_reject_mixed_parent_origins(delivery: str, @pytest.mark.parametrize("delivery", ["live", "queued_and_new"]) -async def test_wrapped_reply_preflight_applies_child_policy(delivery: str) -> None: +async def test_wrapped_reply_delivery_applies_child_policy(delivery: str) -> None: parent, child, coordinator = _build_wrapped_reply_workflow() email = "a@first.example" await parent.run(email) @@ -398,18 +416,23 @@ def reject(requests: Sequence[WorkflowEvent]) -> None: await run assert coordinator.outputs == [] - assert await parent._runner_context.get_pending_request_info_events() == parent_pending + assert await parent._runner_context.get_pending_request_info_events() == ( + parent_pending if delivery == "live" else {} + ) assert await child._runner_context.get_pending_request_info_events() == child_pending queued = await parent._runner_context.drain_messages() - assert [message.data.source_event.request_id for batch in queued.values() for message in batch] == [ - reply.source_event.request_id - ] + assert [ + message.data.source_event.request_id + for batch in queued.values() + for message in batch + if isinstance(message.data, SubWorkflowResponseMessage) + ] == [reply.source_event.request_id] * (2 if delivery == "queued_and_new" else 1) @pytest.mark.parametrize("routing", ["condition", "selection", "switch"]) @pytest.mark.parametrize("restore", [False, True]) @pytest.mark.parametrize("streaming", [False, True]) -async def test_wrapped_reply_preflight_uses_selected_routes( +async def test_wrapped_reply_delivery_uses_selected_routes( selected_reply_workflow, routing: str, restore: bool, streaming: bool ) -> None: from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY @@ -457,7 +480,6 @@ async def test_wrapped_reply_preflight_uses_selected_routes( async def test_wrapped_reply_selected_mixed_scopes_preserve_queue( selected_reply_workflow, routing: str, restore: bool ) -> None: - from agent_framework import WorkflowCheckpointException from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY parent, children, requests, _ = selected_reply_workflow(routing, True, False) @@ -469,19 +491,48 @@ async def test_wrapped_reply_selected_mixed_scopes_preserve_queue( scope = parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) if restore: checkpoint = await parent._runner.build_checkpoint() - with pytest.raises(WorkflowCheckpointException) as error: - await parent._runner.restore_checkpoint(checkpoint) - assert isinstance(error.value.__cause__, ValueError) - assert "multiple parent workflow scopes" in str(error.value.__cause__) - else: - with pytest.raises(ValueError, match="multiple parent workflow scopes"): - await parent._runner._run_iteration() + await parent._runner.restore_checkpoint(checkpoint) + [restored_message] = (await parent._runner._get_queued_messages())["coordinator"] + assert restored_message.data.source_event.to_dict() == message.data.source_event.to_dict() + message = restored_message + pending = { + name: await child._runner_context.get_pending_request_info_events() for name, child in children.items() + } + with pytest.raises(ValueError, match="multiple parent workflow scopes"): + await parent._runner._run_iteration() assert parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) == scope assert await parent._runner._get_queued_messages() == {"coordinator": [message]} for name, child in children.items(): assert await child._runner_context.get_pending_request_info_events() == pending[name] +@pytest.mark.parametrize("routing", ["condition", "selection"]) +async def test_wrapped_reply_validates_all_selected_child_policies_before_delivery( + selected_reply_workflow, routing: str +) -> None: + parent, children, requests, _ = selected_reply_workflow(routing, True, False) + await parent.run("both") + pending = {name: await child._runner_context.get_pending_request_info_events() for name, child in children.items()} + migrations: list[str] = [] + + def stage(requests: Sequence[WorkflowEvent]): + return lambda: migrations.append("first") + + def reject(requests: Sequence[WorkflowEvent]) -> None: + if requests: + raise ValueError("second child policy") + + children["first"]._response_batch_validator = stage + children["second"]._response_batch_validator = reject + with pytest.raises(ValueError, match="second child policy"): + await parent.run([requests["first"].create_response(True)]) + assert migrations == [] + for name, child in children.items(): + assert await child._runner_context.get_pending_request_info_events() == pending[name] + [queued] = (await parent._runner._get_queued_messages())["coordinator"] + assert queued.data.source_event.request_id == "shared" + + @pytest.mark.parametrize("routing", ["condition", "selection"]) @pytest.mark.parametrize("streaming", [False, True]) async def test_wrapped_reply_routing_observes_earlier_same_edge_handler( diff --git a/python/packages/declarative/AGENTS.md b/python/packages/declarative/AGENTS.md index 675120c93c7..ece581e3fbd 100644 --- a/python/packages/declarative/AGENTS.md +++ b/python/packages/declarative/AGENTS.md @@ -109,12 +109,17 @@ 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, +The same validation and migration run before restoring ordinary queued checkpoint responses, 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 +Newly submitted responses are validated together with ordinary queued responses before any pending request is consumed. +Wrapped child replies instead select and validate their actual routes immediately +before delivery, after earlier handlers and checkpoint restoration. All selected +children must validate before that reply is delivered. A later rejected reply does +not roll back earlier continuations or restored state; recover a partially executed +failed run from a known-good checkpoint rather than retrying its remaining queue. 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/tests/test_invoke_mcp_tool_executor.py b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py index 6a1389fee43..25fc8bf62ee 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -31,7 +31,6 @@ from agent_framework import AgentSession, CheckpointStorage, Content, Message, Workflow # noqa: E402 from agent_framework.exceptions import ToolExecutionException # noqa: E402 - from agent_framework_declarative._workflows import ( # noqa: E402 DECLARATIVE_STATE_KEY, DeclarativeWorkflowError, @@ -122,7 +121,6 @@ def _build_manual_mcp_workflow( mcp_handler: MCPToolHandler, *, depth: int, require_approval: bool = False ) -> tuple[Workflow, Workflow]: from agent_framework import WorkflowBuilder, WorkflowExecutor - from agent_framework_declarative._workflows._executors_mcp import InvokeMcpToolActionExecutor leaf = WorkflowBuilder( @@ -231,7 +229,6 @@ async def test_manual_workflow_scopes_follow_core_run_context(self, streaming: b async def test_manual_approval_continuations_restore_core_scopes(self, streaming: bool, restore: bool) -> None: from agent_framework import InMemoryCheckpointStorage, WorkflowEvent from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY - from agent_framework_declarative._workflows import ToolApprovalResponse mcp_handler = StubMcpHandler(_ok()) @@ -473,7 +470,6 @@ async def test_response_only_agent_checkpoint_resume_preserves_workflow_session_ self, pause_kind: str, streaming: bool ) -> None: from agent_framework import InMemoryCheckpointStorage - from agent_framework_declarative._workflows import ExternalInputResponse storage = InMemoryCheckpointStorage() @@ -548,7 +544,6 @@ 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() @@ -608,7 +603,6 @@ 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() @@ -823,7 +817,6 @@ 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() @@ -1152,7 +1145,6 @@ async def test_pending_approval_uses_originating_workflow_session_id( self, restore: bool, agent_sessions: bool, streaming: bool, approved: bool, second_approved: bool ) -> None: from agent_framework import InMemoryCheckpointStorage - from agent_framework_declarative._workflows import ToolApprovalResponse storage = InMemoryCheckpointStorage() @@ -1228,7 +1220,6 @@ async def test_pending_approval_uses_originating_workflow_session_id( @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_response_batch from agent_framework_declarative._workflows._executors_control_flow import JoinExecutor @@ -1947,7 +1938,7 @@ async def restore() -> None: @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( + async def test_wrapped_child_response_validation_uses_restored_state( self, tmp_path: Path, restore_kind: str, depth: int, targeted: bool, mode: str ) -> None: from copy import deepcopy @@ -1955,6 +1946,7 @@ async def test_wrapped_child_response_preflight_preserves_state( from agent_framework import ( Executor, SubWorkflowRequestMessage, + SubWorkflowResponseMessage, WorkflowBuilder, WorkflowContext, WorkflowExecutor, @@ -1963,7 +1955,6 @@ async def test_wrapped_child_response_preflight_preserves_state( 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: @@ -2041,48 +2032,38 @@ async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: 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) + await restore() 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) + with pytest.raises(DeclarativeWorkflowError): + _ = [event async for event in parent._runner.run_until_convergence()] + pending = await leaves[1]._runner_context.get_pending_request_info_events() + assert set(pending) == {child_request.request_id} + assert pending[child_request.request_id].data.workflow_session_id == (42 if mode == "invalid" else None) + queued = await parent._runner._get_queued_messages() + assert [ + message.data.source_event.request_id + for messages in queued.values() + for message in messages + if isinstance(message.data, SubWorkflowResponseMessage) + ] == [child_request.request_id] 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 all(workflow._runner.state.get("restore_marker") == "snapshot" for workflow in workflows) + assert all(not workflow._runner.state.has("uncommitted") for workflow in workflows) + assert recorder.value == "snapshot" + assert recorder.restore_calls == 1 assert encode_checkpoint_value(checkpoint.to_dict()) == checkpoint_before - assert mcp_handler.call_count == (2 if mode in {"migrate", "valid"} else 0) + assert mcp_handler.call_count == (2 if mode in {"migrate", "valid"} else 1) @pytest.mark.parametrize("depth", [1, 2]) @pytest.mark.parametrize("mode", ["missing", "invalid", "queued"]) @@ -2510,7 +2491,6 @@ class TestApprovalFlow: @pytest.mark.asyncio async def test_auth_context_binding_survives_workflow_resume(self, restore: bool, change_context: bool) -> None: from agent_framework import InMemoryCheckpointStorage, WorkflowBuilder - from agent_framework_declarative._workflows import ( ActionTrigger, ExternalInputResponse, diff --git a/skills/agent-framework-builder-python/references/workflow-builder.md b/skills/agent-framework-builder-python/references/workflow-builder.md index c121e97c66e..bf6f4a62d67 100644 --- a/skills/agent-framework-builder-python/references/workflow-builder.md +++ b/skills/agent-framework-builder-python/references/workflow-builder.md @@ -40,6 +40,9 @@ between stages, final output, intermediate output, session ownership, and approv - Sessions are not automatically interchangeable between participants. - Tool approval and workflow human input are different suspension mechanisms. - Workflow streams include framework events in addition to application output. +- Wrapped child replies evaluate routing after preceding handlers and checkpoint restoration, then validate the + selected children before delivery. A later rejection does not roll back earlier execution or restored state; + recover a partially executed failed run from a known-good checkpoint rather than retrying its remaining queue. ## Validation From 4d515ae7b83418208cb96e90a1b4c6ac29277f45 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Thu, 29 Oct 2026 14:06:05 -0400 Subject: [PATCH 75/76] fix(workflows): retain reply scope across continuation steps Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 06d37a04-f20d-46d4-ab56-1ef5adda7b2d --- ...thon-declarative-request-scope-recovery.md | 12 +- python/packages/core/AGENTS.md | 4 +- .../core/agent_framework/_workflows/_const.py | 1 + .../_workflows/_request_context.py | 15 ++- .../agent_framework/_workflows/_runner.py | 10 +- .../agent_framework/_workflows/_workflow.py | 30 +++-- .../packages/core/tests/workflow/conftest.py | 84 ++++++++++++++ .../core/tests/workflow/test_sub_workflow.py | 104 ++++++++++++++++++ python/packages/declarative/AGENTS.md | 2 + .../references/workflow-builder.md | 2 + 10 files changed, 248 insertions(+), 16 deletions(-) diff --git a/docs/decisions/0046-python-declarative-request-scope-recovery.md b/docs/decisions/0046-python-declarative-request-scope-recovery.md index b915e4719f4..6199a27a6ec 100644 --- a/docs/decisions/0046-python-declarative-request-scope-recovery.md +++ b/docs/decisions/0046-python-declarative-request-scope-recovery.md @@ -164,9 +164,15 @@ API would impose a new contract on existing graphs. Instead, finish preceding messages from a source before selecting a wrapped reply's targets. Evaluate each edge callback once, validate all selected targets -together against their own child policies, then deliver that reply. Track selected -parent scopes through the superstep so a later delivery cannot introduce a -different scope after earlier origins have been released. Normal message batches +together against their own child policies, then deliver that reply. Retain the +accepted parent scope in checkpoint-backed shared state while the continuation +has in-flight messages. Later supersteps and mid-run checkpoint replay cannot +introduce a different scope after earlier origins have been released. Clear this +guard at quiescence before capturing the final checkpoint, so a later response-only +run may resume another origin. An accepted new run with no in-flight work also +resets the guard; restoring an in-flight checkpoint preserves it. Validate stored +guard metadata before replacing live state. +Normal message batches retain concurrent edge delivery; wrapped replies form delivery barriers, with their selected targets still executing concurrently. diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 7f8bbd78dd5..fa6d201bf3b 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -353,7 +353,9 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID Wrapped child replies select routes immediately before delivery, after preceding handlers on their source have completed, and validate all selected targets together using each child's own policy. Predicates run once per delivery, not during checkpoint or queue preflight. Fan-in aggregates are not individual wrapped replies. - Selected replies cannot introduce a different parent scope later in the same superstep. A rejected reply + Selected replies cannot introduce a different parent scope while a continuation has in-flight messages. + The accepted scope is checkpoint-backed across supersteps and clears at quiescence before the final + checkpoint, so a later response-only run can select another origin. A rejected reply and subsequent messages remain queued, but earlier handlers and child continuations are not rolled back. Recover a partially executed failed run from a known-good checkpoint rather than retrying its remaining queue. - **Nested checkpoint restoration** - Stored and embedded restores preflight a detached checkpoint tree before diff --git a/python/packages/core/agent_framework/_workflows/_const.py b/python/packages/core/agent_framework/_workflows/_const.py index e36f22392ed..17b97f79360 100644 --- a/python/packages/core/agent_framework/_workflows/_const.py +++ b/python/packages/core/agent_framework/_workflows/_const.py @@ -30,6 +30,7 @@ WORKFLOW_PARENT_SESSION_ID_KEY = "_workflow_parent_session_id" WORKFLOW_AGENT_SESSION_NAMESPACE_KEY = "_workflow_agent_session_namespace" WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY = "_workflow_request_parent_contexts" +WORKFLOW_RESPONSE_SCOPE_KEY = "_workflow_response_scope" # Key used to store collision-free, executor-aware workflow invocation kwargs. # WORKFLOW_RUN_KWARGS_KEY remains in the legacy dict format for compatibility with diff --git a/python/packages/core/agent_framework/_workflows/_request_context.py b/python/packages/core/agent_framework/_workflows/_request_context.py index e090fe3ca1c..5cc2af6852c 100644 --- a/python/packages/core/agent_framework/_workflows/_request_context.py +++ b/python/packages/core/agent_framework/_workflows/_request_context.py @@ -4,7 +4,7 @@ import logging import uuid -from collections.abc import Collection, Sequence +from collections.abc import Collection, Mapping, Sequence from typing import Any, cast from ._const import ( @@ -12,6 +12,7 @@ WORKFLOW_AGENT_SESSION_ID_KEY, WORKFLOW_PARENT_SESSION_ID_KEY, WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, + WORKFLOW_RESPONSE_SCOPE_KEY, WORKFLOW_SESSION_ID_KEY, ) from ._events import WorkflowEvent @@ -28,6 +29,16 @@ _MISSING_ORIGIN = object() +def get_response_scopes(state: State | Mapping[str, Any]) -> set[str]: + """Read the accepted scope retained while a response continuation has in-flight work.""" + scope = state.get(WORKFLOW_RESPONSE_SCOPE_KEY) + if scope is None: + return set() + if not isinstance(scope, str) or not scope: + raise ValueError("Invalid parent workflow response scope.") + return {scope} + + def request_context_state_key(boundary_id: str | None = None) -> str: """Keep wrapped child origins independent of external requests with the same ID.""" if boundary_id is None: @@ -62,7 +73,7 @@ def validate_request_contexts( boundary_contexts: Sequence[tuple[Any, Sequence[WorkflowEvent]]] = (), previous_scopes: Collection[str] = (), ) -> set[str]: - """Validate selected origins together with earlier deliveries in the same batch.""" + """Validate selected origins together with earlier deliveries in the same continuation.""" scopes = set(previous_scopes) for pending_contexts, pending_requests in ((contexts, requests), *boundary_contexts): if not isinstance(pending_contexts, dict): diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index 795cb5c3e34..7ad0f1ec803 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -15,12 +15,12 @@ WorkflowConvergenceException, ) from ._checkpoint import CheckpointID, CheckpointStorage, WorkflowCheckpoint -from ._const import EDGE_STATE_KEY, EXECUTOR_STATE_KEY +from ._const import EDGE_STATE_KEY, EXECUTOR_STATE_KEY, WORKFLOW_RESPONSE_SCOPE_KEY from ._edge import EdgeGroup from ._edge_runner import EdgeRunner, create_edge_runner, gather_cancelling_siblings_on_error from ._events import WorkflowEvent from ._executor import Executor -from ._request_context import request_context_state_key, validate_request_contexts +from ._request_context import get_response_scopes, request_context_state_key, validate_request_contexts from ._runner_context import ( RunnerContext, WorkflowMessage, @@ -174,6 +174,9 @@ async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent]: logger.info(f"Completed superstep {self._iteration}") + # Idle checkpoints must allow a later response run to select a different origin. + if not await self._ctx.has_messages() and self._state.has(WORKFLOW_RESPONSE_SCOPE_KEY): + self._state.delete(WORKFLOW_RESPONSE_SCOPE_KEY) # Commit pending state changes at superstep boundary self._state.commit() committed = True @@ -220,7 +223,7 @@ async def _run_iteration(self) -> None: """ from ._workflow_executor import SubWorkflowResponseMessage - scopes: set[str] = set() + scopes = get_response_scopes(self._state) next_message: dict[str, int] = {} async def _deliver_messages(source_executor_id: str, source_messages: list[WorkflowMessage]) -> None: @@ -505,6 +508,7 @@ async def _validate_checkpoint_tree( validate_request_contexts( parent_contexts, requests, + previous_scopes=get_response_scopes(checkpoint.state), boundary_contexts=[ ( checkpoint.state.get(request_context_state_key(executor_id), parent_contexts), diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index d738c57d8c6..95be9656356 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -40,6 +40,7 @@ WORKFLOW_AGENT_SESSION_NAMESPACE_KEY, WORKFLOW_PARENT_SESSION_ID_KEY, WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, + WORKFLOW_RESPONSE_SCOPE_KEY, WORKFLOW_RUN_KWARGS_KEY, WORKFLOW_SESSION_ID_KEY, ) @@ -60,7 +61,12 @@ _ResponseBatchValidator, # pyright: ignore[reportPrivateUsage] ) from ._model_utils import DictConvertible -from ._request_context import request_context_state_key, restore_request_context, validate_request_contexts +from ._request_context import ( + get_response_scopes, + request_context_state_key, + restore_request_context, + validate_request_contexts, +) from ._runner import RunnerImpl from ._runner_context import ( RunnerContext, @@ -1020,7 +1026,8 @@ async def _run_core( # NOTE: _validate_run_params already enforces that ``message`` is # mutually exclusive with both ``checkpoint_id`` and ``responses``, # so we don't need to re-check those here. - if message is not None and await self._runner.context.has_messages(): + has_messages = await self._runner.context.has_messages() + if message is not None and has_messages: raise RuntimeError( "Cannot start a new run with 'message' while in-flight executor " "messages remain from a prior run. Resume from a checkpoint " @@ -1029,6 +1036,9 @@ async def _run_core( "checkpointing; there is no in-process recovery path." ) + if checkpoint_id is None and not has_messages and self._runner.state.has(WORKFLOW_RESPONSE_SCOPE_KEY): + self._runner.state.delete(WORKFLOW_RESPONSE_SCOPE_KEY) + # Warn (but don't block) when a fresh message or a checkpoint restore begins while the # workflow still has pending request_info events from an unfinished request/response # cycle. A fresh ``message`` does NOT drop those pending requests - they remain pending and @@ -1233,6 +1243,8 @@ async def _get_response_validation_batches( queued_requests = _get_queued_response_info_events(messages) if wrapped_requests is None: wrapped_requests = {} + if scopes is None: + scopes = get_response_scopes(self._runner.state) batch = [*queued_requests, *requests] parent_contexts = self._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {}) boundary_contexts = [ @@ -1244,16 +1256,20 @@ async def _get_response_validation_batches( if executor_id in wrapped_requests or self._runner.state.has(request_context_state_key(executor_id)) ] - def validate_parent_contexts(_requests: Sequence[WorkflowEvent]) -> Callable[[], None] | None: + def validate_parent_contexts(_requests: Sequence[WorkflowEvent]) -> Callable[[], None]: selected_scopes = validate_request_contexts( parent_contexts, batch, boundary_contexts=boundary_contexts, - previous_scopes=scopes if scopes is not None else (), + previous_scopes=scopes, ) - if scopes is not None: - return functools.partial(scopes.update, selected_scopes) - return None + + def accept_scope() -> None: + scopes.update(selected_scopes) + if selected_scopes: + self._runner.state.set(WORKFLOW_RESPONSE_SCOPE_KEY, next(iter(selected_scopes))) + + return accept_scope validations: list[_ResponseBatchValidation] = [] if self._response_batch_validator is not None: diff --git a/python/packages/core/tests/workflow/conftest.py b/python/packages/core/tests/workflow/conftest.py index 053bd4fad0d..959e21c370b 100644 --- a/python/packages/core/tests/workflow/conftest.py +++ b/python/packages/core/tests/workflow/conftest.py @@ -128,3 +128,87 @@ def first_case(message: Any) -> bool: return builder.build(), children, requests, calls return build + + +@pytest.fixture +def staggered_reply_workflow() -> Callable[ + [], tuple[Workflow, dict[str, Workflow], dict[str, SubWorkflowRequestMessage], list[tuple[str, str]]] +]: + def build() -> tuple[Workflow, dict[str, Workflow], dict[str, SubWorkflowRequestMessage], list[tuple[str, str]]]: + requests: dict[str, SubWorkflowRequestMessage] = {} + observed: list[tuple[str, str]] = [] + + class Requester(Executor): + @handler + async def start(self, message: str, ctx: WorkflowContext) -> None: + await ctx.request_info(message, bool, request_id=message) + + @response_handler + async def respond(self, original: str, approved: bool, ctx: WorkflowContext[Any, str]) -> None: + await ctx.yield_output(original) + + class Coordinator(Executor): + @handler + async def start(self, message: str, ctx: WorkflowContext[str]) -> None: + for target in ("first", "second") if message == "both" else (message,): + await ctx.send_message(target, target_id=target) + + @handler + async def intercept(self, request: SubWorkflowRequestMessage, ctx: WorkflowContext) -> None: + requests[request.executor_id] = request + await ctx.request_info(request.executor_id, bool, request_id=request.executor_id) + + @response_handler + async def respond( + self, original: str, approved: bool, ctx: WorkflowContext[SubWorkflowResponseMessage] + ) -> None: + await ctx.send_message(requests[original].create_response(approved), target_id=original) + + @handler + async def stagger( + self, + replies: list[SubWorkflowResponseMessage], + ctx: WorkflowContext[SubWorkflowResponseMessage | tuple[SubWorkflowResponseMessage]], + ) -> None: + await ctx.send_message(replies[0], target_id="first") + await ctx.send_message((replies[1],), target_id="delay") + + class ReplyHop(Executor): + @handler + async def forward( + self, reply: tuple[SubWorkflowResponseMessage], ctx: WorkflowContext[SubWorkflowResponseMessage] + ) -> None: + await ctx.send_message(reply[0]) + + class OutputHop(Executor): + @handler + async def forward(self, message: str, ctx: WorkflowContext[str]) -> None: + await ctx.send_message(message) + + class Recorder(Executor): + @handler + async def record(self, message: str, ctx: WorkflowContext[Any, str]) -> None: + from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY + + observed.append((message, ctx.get_state(WORKFLOW_SESSION_ID_KEY))) + await ctx.yield_output(message) + + children = { + name: WorkflowBuilder(start_executor=Requester(id="requester")).build() for name in ("first", "second") + } + wrappers = {name: WorkflowExecutor(child, id=name) for name, child in children.items()} + coordinator = Coordinator(id="coordinator") + delay = ReplyHop(id="delay") + output = OutputHop(id="output") + recorder = Recorder(id="recorder") + builder = WorkflowBuilder(start_executor=coordinator) + for wrapper in wrappers.values(): + builder.add_edge(coordinator, wrapper).add_edge( + wrapper, coordinator, condition=lambda message: isinstance(message, SubWorkflowRequestMessage) + ).add_edge(wrapper, output) + parent = ( + builder.add_edge(coordinator, delay).add_edge(delay, wrappers["second"]).add_edge(output, recorder).build() + ) + return parent, children, requests, observed + + return build diff --git a/python/packages/core/tests/workflow/test_sub_workflow.py b/python/packages/core/tests/workflow/test_sub_workflow.py index 7f0759beb88..63e6afbbbfd 100644 --- a/python/packages/core/tests/workflow/test_sub_workflow.py +++ b/python/packages/core/tests/workflow/test_sub_workflow.py @@ -588,6 +588,110 @@ async def test_wrapped_reply_routing_uses_checkpointed_executor_fields( assert sorted(calls) == {"condition": ["first", "second"], "selection": ["selection"]}[routing] +@pytest.mark.parametrize("streaming", [False, True]) +@pytest.mark.parametrize("restore", [False, True]) +async def test_staggered_wrapped_replies_reject_cross_scope_continuation( + staggered_reply_workflow, streaming: bool, restore: bool +) -> None: + parent, children, requests, observed = staggered_reply_workflow() + await parent.run("first") + await parent.run("second") + replies = [requests[name].create_response(True) for name in ("first", "second")] + checkpoint_id: str | None = None + storage: InMemoryCheckpointStorage | None = None + if restore: + await parent._runner_context.send_message( + WorkflowMessage(data=replies[0], source_id="coordinator", target_id="first") + ) + await parent._runner_context.send_message( + WorkflowMessage(data=(replies[1],), source_id="coordinator", target_id="delay") + ) + await parent._runner._run_iteration() + storage = InMemoryCheckpointStorage() + checkpoint_id = await storage.save(await parent._runner.build_checkpoint()) + parent, children, _, observed = staggered_reply_workflow() + + with pytest.raises(ValueError, match="multiple parent workflow scopes"): + if streaming: + _ = [ + event + async for event in parent.run( + None if restore else replies, + checkpoint_id=checkpoint_id, + checkpoint_storage=storage, + stream=True, + ) + ] + else: + await parent.run(None if restore else replies, checkpoint_id=checkpoint_id, checkpoint_storage=storage) + assert not await children["first"]._runner_context.get_pending_request_info_events() + assert set(await children["second"]._runner_context.get_pending_request_info_events()) == {"second"} + assert observed == [] + + +@pytest.mark.parametrize("streaming", [False, True]) +async def test_staggered_wrapped_replies_allow_one_parent_scope(staggered_reply_workflow, streaming: bool) -> None: + from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY + + parent, children, requests, observed = staggered_reply_workflow() + await parent.run("both") + scope = parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) + replies = [requests[name].create_response(True) for name in ("first", "second")] + if streaming: + events = [event async for event in parent.run(replies, stream=True)] + outputs = [event.data for event in events if event.type == "output"] + else: + outputs = (await parent.run(replies)).get_outputs() + assert sorted(outputs) == ["first", "second"] + assert sorted(observed) == [("first", scope), ("second", scope)] + for child in children.values(): + assert not await child._runner_context.get_pending_request_info_events() + + +@pytest.mark.parametrize("restore", [False, True]) +async def test_separate_response_runs_can_restore_different_parent_scopes( + staggered_reply_workflow, restore: bool +) -> None: + from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY + + parent, children, _, observed = staggered_reply_workflow() + scopes: list[str] = [] + for name in ("first", "second"): + await parent.run(name) + scopes.append(parent._runner.state.get(WORKFLOW_SESSION_ID_KEY)) + assert scopes[0] != scopes[1] + for name, scope in zip(("first", "second"), scopes, strict=True): + if restore and name == "second": + storage = InMemoryCheckpointStorage() + checkpoint_id = await storage.save(await parent._runner.build_checkpoint()) + result = await parent.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage, responses={name: True}) + else: + result = await parent.run(responses={name: True}) + assert result.get_outputs() == [name] + assert observed[-1] == (name, scope) + assert not await children[name]._runner_context.get_pending_request_info_events() + + +@pytest.mark.parametrize("scope", ["", False, ["first"]]) +async def test_checkpoint_rejects_malformed_response_scope_without_restoring( + staggered_reply_workflow, scope: Any +) -> None: + from agent_framework import WorkflowCheckpointException + from agent_framework._workflows._checkpoint_encoding import encode_checkpoint_value + from agent_framework._workflows._const import WORKFLOW_RESPONSE_SCOPE_KEY + + parent, _, _, _ = staggered_reply_workflow() + await parent.run("first") + checkpoint = await parent._runner.build_checkpoint() + checkpoint.state[WORKFLOW_RESPONSE_SCOPE_KEY] = scope + before = encode_checkpoint_value(parent._runner.state.export_state()) + with pytest.raises(WorkflowCheckpointException) as error: + await parent._runner.restore_checkpoint(checkpoint) + assert isinstance(error.value.__cause__, ValueError) + assert "Invalid parent workflow response scope" in str(error.value.__cause__) + assert encode_checkpoint_value(parent._runner.state.export_state()) == before + + @pytest.mark.parametrize("streaming", [False, True]) async def test_response_validation_preserves_noncopyable_request_payload(streaming: bool) -> None: from threading import Lock diff --git a/python/packages/declarative/AGENTS.md b/python/packages/declarative/AGENTS.md index ece581e3fbd..19607e06674 100644 --- a/python/packages/declarative/AGENTS.md +++ b/python/packages/declarative/AGENTS.md @@ -120,6 +120,8 @@ before delivery, after earlier handlers and checkpoint restoration. All selected children must validate before that reply is delivered. A later rejected reply does not roll back earlier continuations or restored state; recover a partially executed failed run from a known-good checkpoint rather than retrying its remaining queue. +The accepted parent scope persists across supersteps and mid-run checkpoint replay +until all in-flight messages drain; later response-only runs may select other origins. 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/skills/agent-framework-builder-python/references/workflow-builder.md b/skills/agent-framework-builder-python/references/workflow-builder.md index bf6f4a62d67..12d01fdff1d 100644 --- a/skills/agent-framework-builder-python/references/workflow-builder.md +++ b/skills/agent-framework-builder-python/references/workflow-builder.md @@ -43,6 +43,8 @@ between stages, final output, intermediate output, session ownership, and approv - Wrapped child replies evaluate routing after preceding handlers and checkpoint restoration, then validate the selected children before delivery. A later rejection does not roll back earlier execution or restored state; recover a partially executed failed run from a known-good checkpoint rather than retrying its remaining queue. +- Accepted reply scopes persist across supersteps and in-flight checkpoint replay until messages drain. A later + response-only run may resume a different originating scope. ## Validation From 3aad44e54dca8ebccd56ad4063f683c2507c84bb Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Fri, 30 Oct 2026 13:30:22 -0400 Subject: [PATCH 76/76] fix(workflows): avoid cross-edge wrapped reply deadlocks Preserve per-edge delivery ordering while validating wrapped replies only across compatible target paths. Add waiter/child-edge regressions for broadcast, targeted, and independently queued replies. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9f6c9cde-10e6-497b-9dd3-c708cf6f10b3 --- python/packages/core/AGENTS.md | 7 +- .../_workflows/_edge_runner.py | 10 ++ .../agent_framework/_workflows/_runner.py | 92 +++++++++---------- .../packages/core/tests/workflow/conftest.py | 64 +++++++++++++ .../core/tests/workflow/test_sub_workflow.py | 17 +++- 5 files changed, 137 insertions(+), 53 deletions(-) diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index fa6d201bf3b..dcbd2a459cf 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -350,9 +350,10 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID a coordinator relays the same request ID. Origin insertion and removal stage individual mapping entries; committed snapshots remain unchanged until commit and discard rolls back pending changes. Ordinary response and cancellation batches reject mixed parent scopes before consuming requests. - Wrapped child replies select routes immediately before delivery, after preceding handlers on their source - have completed, and validate all selected targets together using each child's own policy. Predicates run once - per delivery, not during checkpoint or queue preflight. Fan-in aggregates are not individual wrapped replies. + Wrapped child replies select routes immediately before delivery, after preceding handlers on compatible + edge runners have completed, without waiting for unrelated outgoing handlers, and validate all selected targets + together using each child's own policy. Predicates run once per delivery, not during checkpoint or queue + preflight. Fan-in aggregates are not individual wrapped replies. Selected replies cannot introduce a different parent scope while a continuation has in-flight messages. The accepted scope is checkpoint-backed across supersteps and clears at quiescence before the final checkpoint, so a later response-only run can select another origin. A rejected reply diff --git a/python/packages/core/agent_framework/_workflows/_edge_runner.py b/python/packages/core/agent_framework/_workflows/_edge_runner.py index 0448be1f197..8824ff10bdf 100644 --- a/python/packages/core/agent_framework/_workflows/_edge_runner.py +++ b/python/packages/core/agent_framework/_workflows/_edge_runner.py @@ -108,6 +108,13 @@ async def _select_targets(self, message: WorkflowMessage) -> tuple[bool, list[st """Determine delivery targets without executing or buffering the message.""" raise NotImplementedError + def _has_compatible_target(self, message: WorkflowMessage) -> bool: + """Check target and handler compatibility without evaluating routing predicates.""" + return any( + (not message.target_id or message.target_id == edge.target_id) and self._can_handle(edge.target_id, message) + for edge in self._edge_group.edges + ) + @property def state_key(self) -> str: """Return a topology-derived key identifying this runner across workflow instances. @@ -396,6 +403,9 @@ async def _select_targets(self, message: WorkflowMessage) -> tuple[bool, list[st # Fan-in delivers an aggregate rather than the wrapped reply itself. return False, [] + def _has_compatible_target(self, message: WorkflowMessage) -> bool: + return False + async def _send_selected_message( self, message: WorkflowMessage, diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index 7ad0f1ec803..86614d1db9b 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -218,8 +218,8 @@ async def _run_iteration(self) -> None: a source by specifying the target, the messages will be delivered to the specified targets in the order they were sent. This is because all messages go through the same edge runner instance which preserves message order. - - Wrapped replies wait for preceding messages from their source, then select and validate all - targets before concurrent delivery. Other sources remain concurrent. + - Wrapped replies wait for preceding messages only on compatible edge runners, then select and + validate all targets before delivery. Unrelated edge runners remain concurrent. """ from ._workflow_executor import SubWorkflowResponseMessage @@ -228,19 +228,6 @@ async def _run_iteration(self) -> None: async def _deliver_messages(source_executor_id: str, source_messages: list[WorkflowMessage]) -> None: """Outer loop to concurrently deliver messages from all sources to their targets.""" - - async def _deliver_message_inner(edge_runner: EdgeRunner, message: WorkflowMessage) -> bool: - """Inner loop to deliver a single message through an edge runner.""" - return await edge_runner.send_message(message, self._state, self._ctx) - - async def _deliver_messages_for_edge_runner( - edge_runner: EdgeRunner, messages: Sequence[WorkflowMessage] - ) -> None: - # Preserve message order per edge runner (and therefore per routed target path) - # while still allowing parallelism across different edge runners. - for message in messages: - await _deliver_message_inner(edge_runner, message) - # Route all messages through normal workflow edges associated_edge_runners = self._edge_runner_map.get(source_executor_id, []) if not associated_edge_runners: @@ -249,46 +236,53 @@ async def _deliver_messages_for_edge_runner( next_message[source_executor_id] = len(source_messages) return - start = 0 + selections: dict[int, dict[EdgeRunner, asyncio.Future[tuple[bool, list[str]]]]] = {} + validated: dict[int, asyncio.Event] = {} for index, message in enumerate(source_messages): - if not isinstance(message.data, SubWorkflowResponseMessage): - continue - # Earlier handlers must finish before state-dependent reply routes are selected. - next_message[source_executor_id] = index - if index > start: - await gather_cancelling_siblings_on_error( - *( - _deliver_messages_for_edge_runner(runner, source_messages[start:index]) - for runner in associated_edge_runners - ) + if isinstance(message.data, SubWorkflowResponseMessage): + selections[index] = { + runner: asyncio.get_running_loop().create_future() + for runner in associated_edge_runners + if runner._has_compatible_target(message) # pyright: ignore[reportPrivateUsage] + } + validated[index] = asyncio.Event() + + async def _deliver_messages_for_edge_runner(edge_runner: EdgeRunner) -> None: + for index, message in enumerate(source_messages): + if index not in selections or edge_runner not in selections[index]: + await edge_runner.send_message(message, self._state, self._ctx) + continue + # Select routes after this edge's earlier handlers, not unrelated outgoing handlers. + delivery = await edge_runner._select_targets(message) # pyright: ignore[reportPrivateUsage] + selections[index][edge_runner].set_result(delivery) + await validated[index].wait() + await edge_runner._send_selected_message( # pyright: ignore[reportPrivateUsage] + message, self._state, self._ctx, delivery ) - deliveries = [ - (runner, await runner._select_targets(message)) # pyright: ignore[reportPrivateUsage] - for runner in associated_edge_runners - ] - targets = {target for _, (_, target_ids) in deliveries for target in target_ids} + + async def _validate_reply(index: int) -> None: + targets: set[str] = set() + for selection in selections[index].values(): + _, target_ids = await selection + targets.update(target_ids) + message = source_messages[index] if self._wrapped_response_validator is not None: await self._wrapped_response_validator( {target: [message.data.source_event] for target in targets}, scopes ) - next_message[source_executor_id] = index + 1 - await gather_cancelling_siblings_on_error( - *( - runner._send_selected_message( # pyright: ignore[reportPrivateUsage] - message, self._state, self._ctx, delivery - ) - for runner, delivery in deliveries - ) - ) - start = index + 1 - next_message[source_executor_id] = len(source_messages) - if start < len(source_messages): - await gather_cancelling_siblings_on_error( - *( - _deliver_messages_for_edge_runner(runner, source_messages[start:]) - for runner in associated_edge_runners - ) - ) + validated[index].set() + # Independent paths may validate out of order; retain the earliest unvalidated reply on failure. + while next_message[source_executor_id] < len(source_messages): + pending_index = next_message[source_executor_id] + if pending_index in validated and not validated[pending_index].is_set(): + break + next_message[source_executor_id] += 1 + + next_message[source_executor_id] = next(iter(selections), len(source_messages)) + await gather_cancelling_siblings_on_error( + *(_validate_reply(index) for index in selections), + *(_deliver_messages_for_edge_runner(runner) for runner in associated_edge_runners), + ) if self._queued_response_validator is not None: await self._queued_response_validator(scopes) diff --git a/python/packages/core/tests/workflow/conftest.py b/python/packages/core/tests/workflow/conftest.py index 959e21c370b..f9b7b3811f3 100644 --- a/python/packages/core/tests/workflow/conftest.py +++ b/python/packages/core/tests/workflow/conftest.py @@ -1,9 +1,11 @@ # Copyright (c) Microsoft. All rights reserved. +import asyncio from collections.abc import Callable from typing import Any import pytest + from agent_framework import ( Case, Default, @@ -130,6 +132,68 @@ def first_case(message: Any) -> bool: return build +@pytest.fixture +def waiting_reply_workflow() -> Callable[[bool, bool], tuple[Workflow, Workflow, list[str]]]: + def build(targeted: bool, preceding_reply: bool) -> tuple[Workflow, Workflow, list[str]]: + started = asyncio.Event() + released = asyncio.Event() + completed: list[str] = [] + requests: list[SubWorkflowRequestMessage] = [] + + class Requester(Executor): + @handler + async def start(self, message: str, ctx: WorkflowContext) -> None: + await ctx.request_info(message, bool, request_id="approval") + + @response_handler + async def respond(self, original: str, approved: bool, ctx: WorkflowContext[Any, str]) -> None: + await started.wait() + released.set() + await ctx.yield_output(original) + + class Waiter(Executor): + @handler + async def wait(self, message: str, ctx: WorkflowContext) -> None: + started.set() + await released.wait() + completed.append(message) + + class ReplyWaiter(Waiter): + @handler + async def receive_reply(self, message: SubWorkflowResponseMessage, ctx: WorkflowContext) -> None: + completed.append("reply") + + class Coordinator(Executor): + @handler + async def start(self, message: str, ctx: WorkflowContext[str]) -> None: + await ctx.send_message(message, target_id="child") + + @handler + async def intercept(self, request: SubWorkflowRequestMessage, ctx: WorkflowContext) -> None: + requests.append(request) + + @handler + async def reply(self, approved: bool, ctx: WorkflowContext[str | SubWorkflowResponseMessage]) -> None: + await ctx.send_message("wait", target_id="waiter") + if preceding_reply: + await ctx.send_message(requests[0].create_response(approved), target_id="waiter") + await ctx.send_message(requests[0].create_response(approved), target_id="child" if targeted else None) + + child = WorkflowBuilder(start_executor=Requester(id="requester")).build() + wrapper = WorkflowExecutor(child, id="child", allow_direct_output=True) + coordinator = Coordinator(id="coordinator") + parent = ( + WorkflowBuilder(start_executor=coordinator) + .add_edge(coordinator, ReplyWaiter(id="waiter") if targeted else Waiter(id="waiter")) + .add_edge(coordinator, wrapper) + .add_edge(wrapper, coordinator) + .build() + ) + return parent, child, completed + + return build + + @pytest.fixture def staggered_reply_workflow() -> Callable[ [], tuple[Workflow, dict[str, Workflow], dict[str, SubWorkflowRequestMessage], list[tuple[str, str]]] diff --git a/python/packages/core/tests/workflow/test_sub_workflow.py b/python/packages/core/tests/workflow/test_sub_workflow.py index 63e6afbbbfd..bb4836b45ed 100644 --- a/python/packages/core/tests/workflow/test_sub_workflow.py +++ b/python/packages/core/tests/workflow/test_sub_workflow.py @@ -8,6 +8,8 @@ from uuid import uuid4 import pytest +from typing_extensions import Never + from agent_framework import ( Executor, FunctionExecutor, @@ -27,7 +29,6 @@ from agent_framework._workflows._checkpoint import InMemoryCheckpointStorage from agent_framework._workflows._runner_context import InProcRunnerContext, WorkflowMessage from agent_framework._workflows._state import State -from typing_extensions import Never # Test message types @@ -556,6 +557,20 @@ async def test_wrapped_reply_routing_observes_earlier_same_edge_handler( assert calls == {"condition": ["first"], "selection": ["selection"]}[routing] +@pytest.mark.parametrize(("targeted", "preceding_reply"), [(False, False), (True, False), (True, True)]) +async def test_wrapped_reply_does_not_wait_for_unrelated_outgoing_handler( + waiting_reply_workflow, targeted: bool, preceding_reply: bool +) -> None: + parent, child, completed = waiting_reply_workflow(targeted, preceding_reply) + await parent.run("start") + + result = await asyncio.wait_for(parent.run(True), timeout=2.0) + + assert result.get_outputs() == ["start"] + assert completed == (["wait", "reply"] if preceding_reply else ["wait"]) + assert not await child._runner_context.get_pending_request_info_events() + + @pytest.mark.parametrize("routing", ["condition", "selection"]) @pytest.mark.parametrize("streaming", [False, True]) async def test_wrapped_reply_routing_uses_checkpointed_executor_fields(