From af615706a689f7d01e7e9ba69376b9a0e4287d13 Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Tue, 15 Sep 2026 15:50:10 +0000 Subject: [PATCH 1/7] fix: preserve dynamic function middleware Keep harness callbacks consistent as tool collections change. Preserve existing function wrappers and request-local composition without coupling reusable options to agent instances. --- .../FunctionInvocationDelegatingAgent.cs | 146 +++++- ...ocationDelegatingAgentBuilderExtensions.cs | 5 + .../FunctionInvocationDelegatingAgentTests.cs | 416 ++++++++++++++++++ .../core/test_function_invocation_logic.py | 73 ++- 4 files changed, 624 insertions(+), 16 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgent.cs b/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgent.cs index 13f15ff1e9f..93804ddf726 100644 --- a/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgent.cs @@ -2,7 +2,8 @@ using System; using System.Collections.Generic; -using System.Linq; +using System.Collections.ObjectModel; +using System.Runtime.CompilerServices; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; @@ -28,7 +29,7 @@ protected override IAsyncEnumerable RunCoreStreamingAsync(I => this.InnerAgent.RunStreamingAsync(messages, session, this.AgentRunOptionsWithFunctionMiddleware(options), cancellationToken); // Decorate options to add the middleware function - private AgentRunOptions? AgentRunOptionsWithFunctionMiddleware(AgentRunOptions? options) + private ChatClientAgentRunOptions AgentRunOptionsWithFunctionMiddleware(AgentRunOptions? options) { if (options is null || options.GetType() == typeof(AgentRunOptions)) { @@ -40,6 +41,10 @@ protected override IAsyncEnumerable RunCoreStreamingAsync(I AdditionalProperties = options?.AdditionalProperties, }; } + else if (options is ChatClientAgentRunOptions chatOptions) + { + options = chatOptions.Clone(); + } if (options is not ChatClientAgentRunOptions aco) { @@ -47,28 +52,129 @@ protected override IAsyncEnumerable RunCoreStreamingAsync(I } var originalFactory = aco.ChatClientFactory; - aco.ChatClientFactory = chatClient => + aco.ChatClientFactory = chatClient => new FunctionMiddlewarePreservingChatClient(chatClient, this).Build(originalFactory); + + return aco; + } + + /// + /// Preserves the function middleware chain when tools are added or replaced during a run. + /// + private sealed class FunctionMiddlewarePreservingChatClient(IChatClient innerClient, FunctionInvocationDelegatingAgent middleware) : DelegatingChatClient(innerClient) + { + private readonly FunctionInvocationDelegatingAgent _middleware = middleware; + private FunctionInvocationDelegatingAgent[] _middlewareChain = [middleware]; + + internal IChatClient Build(Func? originalFactory) { - var builder = chatClient.AsBuilder(); + var builder = this.AsBuilder(); if (originalFactory is not null) { builder.Use(originalFactory); } - return builder.ConfigureOptions(co - => co.Tools = co.Tools?.Select(tool => tool is AIFunction aiFunction - ? new MiddlewareEnabledFunction(this.InnerAgent, aiFunction, this._delegateFunc) - : tool) - .ToList()) - .Build(); - }; + var pipeline = builder.Build(); + var chain = new List(); + for (var client = pipeline.GetService(); + client is not null && !ReferenceEquals(client, this); + client = client.InnerClient.GetService()) + { + chain.Add(client._middleware); + } + + chain.Add(this._middleware); + // Initialize only this request's decorator, never the shared client or caller's options. + this._middlewareChain = [.. chain]; + return pipeline; + } + + public override async Task GetResponseAsync( + IEnumerable messages, ChatOptions? options = null, CancellationToken cancellationToken = default) + => await this.InnerClient.GetResponseAsync(messages, this.ConfigureOptions(options), cancellationToken).ConfigureAwait(false); + + public override async IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, ChatOptions? options = null, [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + await foreach (var update in this.InnerClient.GetStreamingResponseAsync(messages, this.ConfigureOptions(options), cancellationToken).ConfigureAwait(false)) + { + yield return update; + } + } + + private ChatOptions ConfigureOptions(ChatOptions? options) + { + options = options?.Clone() ?? new(); + if (options.Tools is { } tools) + { + options.Tools = new MiddlewareEnabledTools(tools, this._middlewareChain); + } + + return options; + } + } + + private sealed class MiddlewareEnabledTools : Collection + { + private readonly FunctionInvocationDelegatingAgent[] _middleware; + + internal MiddlewareEnabledTools(IList tools, FunctionInvocationDelegatingAgent[] middleware) + { + this._middleware = middleware; + foreach (var tool in tools) + { + this.Add(tool); + } + } + + internal void ApplyTo(ChatOptions options) + { + if (options.Tools is { } tools && + (tools is not MiddlewareEnabledTools existing || !ReferenceEquals(existing._middleware, this._middleware))) + { + options.Tools = new MiddlewareEnabledTools(tools, this._middleware); + } + } + + protected override void InsertItem(int index, AITool item) => base.InsertItem(index, this.Wrap(item)); + + protected override void SetItem(int index, AITool item) => base.SetItem(index, this.Wrap(item)); + + private AITool Wrap(AITool tool) + { + if (tool is AIFunction function) + { + foreach (var middleware in this._middleware) + { + function = MiddlewareEnabledFunction.Wrap(function, middleware); + } + + return function; + } - return options; + return tool; + } } - private sealed class MiddlewareEnabledFunction(AIAgent innerAgent, AIFunction innerFunction, Func>, CancellationToken, ValueTask> next) : DelegatingAIFunction(innerFunction) + private sealed class MiddlewareEnabledFunction(AIFunction innerFunction, FunctionInvocationDelegatingAgent middleware) : DelegatingAIFunction(innerFunction) { + private readonly FunctionInvocationDelegatingAgent _middleware = middleware; + + internal static AIFunction Wrap(AIFunction function, FunctionInvocationDelegatingAgent middleware) + { + for (var existing = function.GetService(); + existing is not null; + existing = existing.InnerFunction.GetService()) + { + if (ReferenceEquals(existing._middleware, middleware)) + { + return function; + } + } + + return new MiddlewareEnabledFunction(function, middleware); + } + protected override async ValueTask InvokeCoreAsync(AIFunctionArguments arguments, CancellationToken cancellationToken) { var context = FunctionInvokingChatClient.CurrentContext @@ -79,7 +185,19 @@ private sealed class MiddlewareEnabledFunction(AIAgent innerAgent, AIFunction in CallContent = new(string.Empty, this.InnerFunction.Name, new Dictionary(arguments)), }; - return await next(innerAgent, context, CoreLogicAsync, cancellationToken).ConfigureAwait(false); + var tools = context.Options?.Tools as MiddlewareEnabledTools; + try + { + return await this._middleware._delegateFunc(this._middleware.InnerAgent, context, CoreLogicAsync, cancellationToken).ConfigureAwait(false); + } + finally + { + // A function or middleware can replace the entire collection during invocation. + if (tools is not null && context.Options is { } options) + { + tools.ApplyTo(options); + } + } ValueTask CoreLogicAsync(FunctionInvocationContext ctx, CancellationToken cancellationToken) => base.InvokeCoreAsync(ctx.Arguments, cancellationToken); diff --git a/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgentBuilderExtensions.cs b/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgentBuilderExtensions.cs index 5ff23f600c7..02d5ab64c32 100644 --- a/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgentBuilderExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgentBuilderExtensions.cs @@ -30,6 +30,11 @@ public static class FunctionInvocationDelegatingAgentBuilderExtensions /// unless it intends to completely replace the function's behavior. /// /// + /// The callbacks also apply to functions added to or replaced in the current + /// collection during execution. Replacing the collection within a function or callback preserves the + /// callbacks for subsequent invocations. + /// + /// /// The inner agent or the pipeline wrapping it must include a . If one does not exist, /// the added to the pipline by this method will throw an exception when it is invoked. /// diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/FunctionInvocationDelegatingAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/FunctionInvocationDelegatingAgentTests.cs index d39edde983d..809bdb4a367 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/FunctionInvocationDelegatingAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/FunctionInvocationDelegatingAgentTests.cs @@ -3,6 +3,7 @@ using System; using System.Collections.Generic; using System.Linq; +using System.Text.Json; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; @@ -287,6 +288,389 @@ public async Task RunAsync_WithMultipleFunctionCalls_InvokesMiddlewareForEachAsy #endregion + [Theory] + [InlineData(false, "Add")] + [InlineData(true, "Add")] + [InlineData(false, "Insert")] + [InlineData(true, "Insert")] + [InlineData(false, "Set")] + [InlineData(true, "Set")] + [InlineData(false, "Replace")] + [InlineData(true, "Replace")] + public async Task RunAsync_DynamicallyAddedFunction_InvokesMiddlewareAsync(bool streaming, string operation) + { + // Arrange + var invokedFunctions = new List(); + var functionExecuted = false; + var dynamicFunction = AIFunctionFactory.Create(() => + { + functionExecuted = true; + return "Function result"; + }, "DynamicFunction"); + var loader = AIFunctionFactory.Create(() => + { + var options = FunctionInvokingChatClient.CurrentContext!.Options!; + switch (operation) + { + case "Add": + options.Tools!.Add(dynamicFunction); + break; + case "Insert": + options.Tools!.Insert(0, dynamicFunction); + break; + case "Set": + options.Tools![0] = dynamicFunction; + break; + case "Replace": + options.Tools = [dynamicFunction]; + break; + default: + throw new ArgumentOutOfRangeException(nameof(operation)); + } + + return "Function added"; + }, "LoadFunction"); + + var responses = new Queue( + [ + new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("load", loader.Name)])), + new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("invoke", dynamicFunction.Name)])), + new(new ChatMessage(ChatRole.Assistant, "Complete")), + ]); + var mockChatClient = CreateMockChatClient(responses); + var agent = new ChatClientAgent(mockChatClient.Object, tools: [loader]) + .AsBuilder() + .Use((agent, context, next, cancellationToken) => + { + invokedFunctions.Add(context.Function.Name); + return context.Function.Name == dynamicFunction.Name + ? new ValueTask("Function handled by middleware") + : next(context, cancellationToken); + }) + .Build(); + + // Act + var response = streaming + ? await agent.RunStreamingAsync("Run the functions").ToAgentResponseAsync() + : await agent.RunAsync("Run the functions"); + + // Assert + Assert.Equal([loader.Name, dynamicFunction.Name], invokedFunctions); + Assert.False(functionExecuted); + Assert.Equal("Function handled by middleware", response.Messages + .SelectMany(m => m.Contents) + .OfType() + .Single(r => r.CallId == "invoke").Result); + Assert.Empty(responses); + } + + [Theory] + [InlineData(false, false)] + [InlineData(true, false)] + [InlineData(false, true)] + [InlineData(true, true)] + public async Task RunAsync_DynamicFunction_PreservesInvocationOrderAsync(bool streaming, bool replaceTools) + => await RunDynamicFunctionPreservingInvocationOrderAsync(streaming, replaceTools, recreateOptions: false); + + [Theory] + [InlineData(false, false)] + [InlineData(true, false)] + [InlineData(false, true)] + [InlineData(true, true)] + public async Task RunAsync_RecreatedOptions_PreservesFunctionMiddlewareAsync(bool streaming, bool replaceTools) + => await RunDynamicFunctionPreservingInvocationOrderAsync(streaming, replaceTools, recreateOptions: true); + + private static async Task RunDynamicFunctionPreservingInvocationOrderAsync(bool streaming, bool replaceTools, bool recreateOptions) + { + // Arrange + var executionOrder = new List(); + var function = AIFunctionFactory.Create(() => + { + executionOrder.Add("Function"); + return "Function result"; + }, "TestFunction"); + var loader = AIFunctionFactory.Create(() => + { + var options = FunctionInvokingChatClient.CurrentContext!.Options!; + if (replaceTools) + { + options.Tools = [function]; + } + else + { + options.Tools!.Add(function); + options.Tools[options.Tools.Count - 1] = options.Tools[options.Tools.Count - 1]; + } + + return "Function added"; + }, "LoadFunction"); + + var responses = new Queue(); + var mockChatClient = CreateMockChatClient(responses); + using var functionClient = new TrackingFunctionInvokingChatClient(mockChatClient.Object, executionOrder, function.Name); + async ValueTask InvokeAsync(FunctionInvocationContext context, CancellationToken cancellationToken) + { + if (context.Function.Name == function.Name) + { + executionOrder.Add("Invoker-Pre"); + } + + var result = await context.Function.InvokeAsync(context.Arguments, cancellationToken); + if (context.Function.Name == function.Name) + { + executionOrder.Add("Invoker-Post"); + } + + return result; + } + + functionClient.FunctionInvoker = InvokeAsync; + var innerAgent = new ChatClientAgent(functionClient, new ChatClientAgentOptions { UseProvidedChatClientAsIs = true }); + var first = innerAgent.AsBuilder().Use(async (agent, context, next, cancellationToken) => + { + Assert.Same(innerAgent, agent); + if (recreateOptions) + { + Assert.Equal(0.5f, context.Options?.Temperature); + } + + if (context.Function.Name == function.Name) + { + executionOrder.Add("First-Pre"); + } + + var result = await next(context, cancellationToken); + if (context.Function.Name == function.Name) + { + executionOrder.Add("First-Post"); + } + + return result; + }).Build(); + var nextAgent = recreateOptions + ? new AnonymousDelegatingAIAgent( + first, + (messages, session, options, agent, cancellationToken) => + agent.RunAsync(messages, session, RecreateOptions(options), cancellationToken), + (messages, session, options, agent, cancellationToken) => + agent.RunStreamingAsync(messages, session, RecreateOptions(options), cancellationToken)) + : first; + var decorated = nextAgent.AsBuilder().Use(async (agent, context, next, cancellationToken) => + { + Assert.Same(nextAgent, agent); + if (context.Function.Name == function.Name) + { + executionOrder.Add("Second-Pre"); + } + + var result = await next(context, cancellationToken); + if (context.Function.Name == function.Name) + { + executionOrder.Add("Second-Post"); + } + + return result; + }).Build(); + + var options = new ChatClientAgentRunOptions(new ChatOptions { Tools = [function] }); + var expectedOrder = new[] + { + "Override-Pre", "Invoker-Pre", "First-Pre", "Second-Pre", + "Function", "Second-Post", "First-Post", "Invoker-Post", "Override-Post", + }; + + // Act + for (int run = 0; run < 3; run++) + { + executionOrder.Clear(); + if (run > 0) + { + options.ChatOptions!.Tools = [loader]; + responses.Enqueue(new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("load", loader.Name)]))); + } + + responses.Enqueue(new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("invoke", function.Name)]))); + responses.Enqueue(new(new ChatMessage(ChatRole.Assistant, "Complete"))); + var response = streaming + ? await decorated.RunStreamingAsync("Run the functions", options: options).ToAgentResponseAsync() + : await decorated.RunAsync("Run the functions", options: options); + + // Assert + Assert.Equal(expectedOrder, executionOrder); + Assert.Equal("Function result", Assert.IsType(response.Messages.SelectMany(m => m.Contents) + .OfType().Single(r => r.CallId == "invoke").Result).GetString()); + Assert.Empty(responses); + Assert.Equal(InvokeAsync, functionClient.FunctionInvoker); + } + + static ChatClientAgentRunOptions RecreateOptions(AgentRunOptions? options) + { + var original = Assert.IsType(options); + var factory = Assert.IsType>(original.ChatClientFactory); + return new ChatClientAgentRunOptions(original.ChatOptions?.Clone()) + { + ChatClientFactory = client => factory(new ConfigureOptionsChatClient(client, options => options.Temperature = 0.5f)), + }; + } + } + + [Theory] + [InlineData(false, false)] + [InlineData(true, false)] + [InlineData(false, true)] + [InlineData(true, true)] + public async Task RunAsync_DynamicFunction_PreservesApprovalAsync(bool streaming, bool approved) + { + // Arrange + var invocationCount = 0; + var observedFunctions = new List(); + var function = AIFunctionFactory.Create(() => + { + invocationCount++; + return "Function result"; + }, "TestFunction", "Function requiring approval"); + var approvalFunction = new ApprovalRequiredAIFunction(function); + var loader = AIFunctionFactory.Create(() => + { + var tools = FunctionInvokingChatClient.CurrentContext!.Options!.Tools!; + tools.Add(approvalFunction); + Assert.Equal(function.Name, tools[tools.Count - 1].Name); + Assert.Same(approvalFunction, tools[tools.Count - 1].GetService()); + Assert.Equal(function.JsonSchema.GetRawText(), Assert.IsAssignableFrom(tools[tools.Count - 1]).JsonSchema.GetRawText()); + return "Function added"; + }, "LoadFunction"); + + var responses = new Queue( + [ + new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("load", loader.Name)])), + new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("invoke", function.Name)])), + ]); + var mockChatClient = CreateMockChatClient(responses); + var agent = new ChatClientAgent(mockChatClient.Object, tools: [loader]).AsBuilder() + .Use((agent, context, next, cancellationToken) => + { + observedFunctions.Add(context.Function.Name); + return next(context, cancellationToken); + }).Build(); + var session = await agent.CreateSessionAsync(); + + // Act + var response = streaming + ? await agent.RunStreamingAsync("Run the functions", session).ToAgentResponseAsync() + : await agent.RunAsync("Run the functions", session); + + // Assert + var request = Assert.Single(response.Messages.SelectMany(m => m.Contents).OfType()); + Assert.Equal(0, invocationCount); + Assert.Equal([loader.Name], observedFunctions); + Assert.Empty(responses); + + // Act + responses.Enqueue(new(new ChatMessage(ChatRole.Assistant, "Complete"))); + var options = new ChatClientAgentRunOptions(new ChatOptions { Tools = [approvalFunction] }); + var approvalMessage = new ChatMessage(ChatRole.User, [request.CreateResponse(approved)]); + var resumedResponse = streaming + ? await agent.RunStreamingAsync(approvalMessage, session, options).ToAgentResponseAsync() + : await agent.RunAsync(approvalMessage, session, options); + + // Assert + Assert.Equal(approved ? 1 : 0, invocationCount); + Assert.Equal(approved ? [loader.Name, function.Name] : new[] { loader.Name }, observedFunctions); + Assert.Single(resumedResponse.Messages.SelectMany(m => m.Contents) + .OfType(), result => result.CallId == "invoke"); + Assert.Empty(responses); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task RunAsync_SharedClientAndOptions_IsolatesFunctionMiddlewareAsync(bool streaming) + { + // Arrange + var loadersEntered = 0; + var gate = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + using var cancellation = new CancellationTokenSource(TimeSpan.FromSeconds(30)); + using var registration = cancellation.Token.Register(() => gate.TrySetCanceled()); + var function = AIFunctionFactory.Create(() => "Function result", "TestFunction"); + var loader = AIFunctionFactory.Create(async () => + { + if (Interlocked.Increment(ref loadersEntered) == 2) + { + gate.TrySetResult(true); + } + + await gate.Task; + FunctionInvokingChatClient.CurrentContext!.Options!.Tools!.Add(function); + return "Function added"; + }, "LoadFunction"); + var mockChatClient = new Mock(); + ChatResponse GetResponse(IEnumerable messages) + { + var results = messages.SelectMany(m => m.Contents).OfType().ToArray(); + if (!results.Any(r => r.CallId == "load")) + { + return new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("load", loader.Name)])); + } + + return results.Any(r => r.CallId == "invoke") + ? new(new ChatMessage(ChatRole.Assistant, "Complete")) + : new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("invoke", function.Name)])); + } + + mockChatClient.Setup(c => c.GetResponseAsync( + It.IsAny>(), It.IsAny(), It.IsAny())) + .Returns((IEnumerable messages, ChatOptions? options, CancellationToken ct) => Task.FromResult(GetResponse(messages))); + mockChatClient.Setup(c => c.GetStreamingResponseAsync( + It.IsAny>(), It.IsAny(), It.IsAny())) + .Returns((IEnumerable messages, ChatOptions? options, CancellationToken ct) => GetResponse(messages).ToChatResponseUpdates().ToAsyncEnumerable()); + + using var client = new FunctionInvokingChatClient(mockChatClient.Object); + var innerAgent = new ChatClientAgent(client, new ChatClientAgentOptions { UseProvidedChatClientAsIs = true }); + var firstInvocations = new List(); + var secondInvocations = new List(); + var first = innerAgent.AsBuilder().Use((agent, context, next, cancellationToken) => + { + firstInvocations.Add(context.Function.Name); + return context.Function.Name == function.Name ? new ValueTask("First") : next(context, cancellationToken); + }).Build(); + var second = innerAgent.AsBuilder().Use((agent, context, next, cancellationToken) => + { + secondInvocations.Add(context.Function.Name); + return context.Function.Name == function.Name ? new ValueTask("Second") : next(context, cancellationToken); + }).Build(); + Func factory = static client => client; + var options = new ChatClientAgentRunOptions(new ChatOptions { Tools = [loader] }) { ChatClientFactory = factory }; + + // Act + Task RunAsync(AIAgent agent, AgentRunOptions? runOptions = null) => streaming + ? agent.RunStreamingAsync("Run the functions", options: runOptions ?? options, cancellationToken: cancellation.Token).ToAgentResponseAsync(cancellation.Token) + : agent.RunAsync("Run the functions", options: runOptions ?? options, cancellationToken: cancellation.Token); + var responses = await Task.WhenAll(RunAsync(first), RunAsync(second)); + + // Assert + Assert.Equal([loader.Name, function.Name], firstInvocations); + Assert.Equal([loader.Name, function.Name], secondInvocations); + Assert.Equal("First", responses[0].Messages.SelectMany(m => m.Contents) + .OfType().Single(r => r.CallId == "invoke").Result); + Assert.Equal("Second", responses[1].Messages.SelectMany(m => m.Contents) + .OfType().Single(r => r.CallId == "invoke").Result); + Assert.Same(factory, options.ChatClientFactory); + Assert.Same(loader, Assert.Single(options.ChatOptions!.Tools!)); + Assert.Null(client.FunctionInvoker); + + // Act & Assert + foreach (var response in new[] { await RunAsync(innerAgent), await RunAsync(innerAgent, options.Clone()) }) + { + Assert.Equal("Function result", Assert.IsType(response.Messages.SelectMany(m => m.Contents) + .OfType().Single(r => r.CallId == "invoke").Result).GetString()); + } + + Assert.Equal([loader.Name, function.Name], firstInvocations); + Assert.Equal([loader.Name, function.Name], secondInvocations); + Assert.Same(factory, options.ChatClientFactory); + Assert.Same(loader, Assert.Single(options.ChatOptions!.Tools!)); + } + #region Context Validation Tests /// @@ -994,6 +1378,38 @@ public async Task RunAsync_WithBaseAgentRunOptions_PreservesAllOriginalOptionsAs #endregion + private static Mock CreateMockChatClient(Queue responses) + { + var mockChatClient = new Mock(); + mockChatClient.Setup(c => c.GetResponseAsync( + It.IsAny>(), It.IsAny(), It.IsAny())) + .ReturnsAsync(() => responses.Dequeue()); + mockChatClient.Setup(c => c.GetStreamingResponseAsync( + It.IsAny>(), It.IsAny(), It.IsAny())) + .Returns(() => responses.Dequeue().ToChatResponseUpdates().ToAsyncEnumerable()); + return mockChatClient; + } + + private sealed class TrackingFunctionInvokingChatClient(IChatClient innerClient, List executionOrder, string functionName) + : FunctionInvokingChatClient(innerClient) + { + protected override async ValueTask InvokeFunctionAsync(FunctionInvocationContext context, CancellationToken cancellationToken) + { + if (context.Function.Name == functionName) + { + executionOrder.Add("Override-Pre"); + } + + var result = await base.InvokeFunctionAsync(context, cancellationToken); + if (context.Function.Name == functionName) + { + executionOrder.Add("Override-Post"); + } + + return result; + } + } + /// /// Creates a mock IChatClient with predefined responses for testing. /// diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index d829aaa4c31..0ce63a34160 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -19,6 +19,8 @@ ChatResponse, ChatResponseUpdate, Content, + FunctionInvocationContext, + FunctionMiddleware, Message, ResponseStream, SupportsChatGetResponse, @@ -36,8 +38,6 @@ included_token_count, ) from agent_framework._middleware import ( - FunctionInvocationContext, - FunctionMiddleware, FunctionMiddlewarePipeline, MiddlewareFailure, MiddlewareTermination, @@ -7493,6 +7493,75 @@ def load_math(ctx: FunctionInvocationContext) -> str: assert exec_counter == 1 +@pytest.mark.parametrize("streaming", [False, True], ids=["non-streaming", "streaming"]) +@pytest.mark.parametrize("allowed", [False, True], ids=["blocked", "allowed"]) +@pytest.mark.parametrize("max_iterations", [3], indirect=True) +async def test_add_tools_respects_function_middleware( + chat_client_base: SupportsChatGetResponse, + streaming: bool, + allowed: bool, +) -> None: + """Late-added tools remain subject to the agent's function middleware.""" + observed_tools: list[str] = [] + target_calls: list[str] = [] + + class PolicyMiddleware(FunctionMiddleware): + async def process(self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]) -> None: + observed_tools.append(context.function.name) + if context.function.name == "late_tool" and not allowed: + context.result = "blocked by policy" + return + await call_next() + + @tool(name="late_tool", approval_mode="never_require") + def late_tool() -> str: + target_calls.append("invoked") + return "target completed" + + @tool(name="load_tool", approval_mode="never_require") + def load_tool(context: FunctionInvocationContext) -> str: + context.add_tools(late_tool) + return "target loaded" + + responses = [ + _pte_function_call_response("1", "load_tool"), + _pte_function_call_response("2", "late_tool"), + _pte_text_response(), + ] + if streaming: + chat_client_base.streaming_responses = [ # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + [ChatResponseUpdate(role="assistant", contents=message.contents) for message in response.messages] + for response in responses + ] + else: + chat_client_base.run_responses = responses # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + + agent = Agent(client=chat_client_base, tools=[load_tool], middleware=[PolicyMiddleware()]) + expected_results = [("1", "target loaded"), ("2", "target completed" if allowed else "blocked by policy")] + if streaming: + stream = agent.run("Load and call the tool.", stream=True) + updates = [update async for update in stream] + response = await stream.get_final_response() + assert [ + (content.call_id, content.result) + for update in updates + for content in update.contents + if content.type == "function_result" + ] == expected_results + else: + response = await agent.run("Load and call the tool.") + + assert observed_tools == ["load_tool", "late_tool"] + assert target_calls == (["invoked"] if allowed else []) + assert [ + (content.call_id, content.result) + for message in response.messages + for content in message.contents + if content.type == "function_result" + ] == expected_results + assert response.messages[-1].text == "done" + + async def test_add_tools_with_approval_required_tool(chat_client_base: SupportsChatGetResponse): @tool(name="secure_tool", approval_mode="always_require") def secure_tool(value: str) -> str: From d9eb7fd9fdeda22c1f58d95b08672b87966a7512 Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Tue, 15 Sep 2026 18:16:08 +0000 Subject: [PATCH 2/7] fix: compose harness middleware across clients Keep callback composition independent of client service discovery. Use request-local execution scopes without changing reusable options. Cover scope restoration and tool repair after collection updates. --- .../FunctionInvocationDelegatingAgent.cs | 48 ++-- .../FunctionInvocationDelegatingAgentTests.cs | 271 +++++++++++++++++- 2 files changed, 298 insertions(+), 21 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgent.cs b/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgent.cs index 93804ddf726..40dac4d2989 100644 --- a/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgent.cs @@ -62,8 +62,8 @@ private ChatClientAgentRunOptions AgentRunOptionsWithFunctionMiddleware(AgentRun /// private sealed class FunctionMiddlewarePreservingChatClient(IChatClient innerClient, FunctionInvocationDelegatingAgent middleware) : DelegatingChatClient(innerClient) { + private static readonly AsyncLocal s_currentScope = new(); private readonly FunctionInvocationDelegatingAgent _middleware = middleware; - private FunctionInvocationDelegatingAgent[] _middlewareChain = [middleware]; internal IChatClient Build(Func? originalFactory) { @@ -74,44 +74,56 @@ internal IChatClient Build(Func? originalFactory) builder.Use(originalFactory); } - var pipeline = builder.Build(); - var chain = new List(); - for (var client = pipeline.GetService(); - client is not null && !ReferenceEquals(client, this); - client = client.InnerClient.GetService()) - { - chain.Add(client._middleware); - } - - chain.Add(this._middleware); - // Initialize only this request's decorator, never the shared client or caller's options. - this._middlewareChain = [.. chain]; - return pipeline; + return builder.Build(); } public override async Task GetResponseAsync( IEnumerable messages, ChatOptions? options = null, CancellationToken cancellationToken = default) - => await this.InnerClient.GetResponseAsync(messages, this.ConfigureOptions(options), cancellationToken).ConfigureAwait(false); + { + var scope = this.CreateScope(s_currentScope.Value); + s_currentScope.Value = scope; + return await this.InnerClient.GetResponseAsync(messages, ConfigureOptions(options, scope), cancellationToken).ConfigureAwait(false); + } public override async IAsyncEnumerable GetStreamingResponseAsync( IEnumerable messages, ChatOptions? options = null, [EnumeratorCancellation] CancellationToken cancellationToken = default) { - await foreach (var update in this.InnerClient.GetStreamingResponseAsync(messages, this.ConfigureOptions(options), cancellationToken).ConfigureAwait(false)) + var scope = this.CreateScope(s_currentScope.Value); + s_currentScope.Value = scope; + await foreach (var update in this.InnerClient.GetStreamingResponseAsync(messages, ConfigureOptions(options, scope), cancellationToken).ConfigureAwait(false)) { yield return update; + + // Resume this request's scope after the consumer's execution context. + s_currentScope.Value = scope; } } - private ChatOptions ConfigureOptions(ChatOptions? options) + private MiddlewareScope CreateScope(MiddlewareScope? previous) + { + var runContext = CurrentRunContext; + // A nested agent run must not inherit the calling agent's callbacks. + return new(runContext, previous is not null && ReferenceEquals(previous.RunContext, runContext) + ? [.. previous.Middleware, this._middleware] + : [this._middleware]); + } + + private static ChatOptions ConfigureOptions(ChatOptions? options, MiddlewareScope scope) { options = options?.Clone() ?? new(); if (options.Tools is { } tools) { - options.Tools = new MiddlewareEnabledTools(tools, this._middlewareChain); + options.Tools = new MiddlewareEnabledTools(tools, scope.Middleware); } return options; } + + private sealed class MiddlewareScope(AgentRunContext? runContext, FunctionInvocationDelegatingAgent[] middleware) + { + internal AgentRunContext? RunContext { get; } = runContext; + internal FunctionInvocationDelegatingAgent[] Middleware { get; } = middleware; + } } private sealed class MiddlewareEnabledTools : Collection diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/FunctionInvocationDelegatingAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/FunctionInvocationDelegatingAgentTests.cs index 809bdb4a367..fff734b761b 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/FunctionInvocationDelegatingAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/FunctionInvocationDelegatingAgentTests.cs @@ -380,7 +380,15 @@ public async Task RunAsync_DynamicFunction_PreservesInvocationOrderAsync(bool st public async Task RunAsync_RecreatedOptions_PreservesFunctionMiddlewareAsync(bool streaming, bool replaceTools) => await RunDynamicFunctionPreservingInvocationOrderAsync(streaming, replaceTools, recreateOptions: true); - private static async Task RunDynamicFunctionPreservingInvocationOrderAsync(bool streaming, bool replaceTools, bool recreateOptions) + [Theory] + [InlineData(false, false)] + [InlineData(true, false)] + [InlineData(false, true)] + [InlineData(true, true)] + public async Task RunAsync_OpaqueChatClient_PreservesFunctionMiddlewareAsync(bool streaming, bool replaceTools) + => await RunDynamicFunctionPreservingInvocationOrderAsync(streaming, replaceTools, recreateOptions: true, hideServices: true); + + private static async Task RunDynamicFunctionPreservingInvocationOrderAsync(bool streaming, bool replaceTools, bool recreateOptions, bool hideServices = false) { // Arrange var executionOrder = new List(); @@ -503,13 +511,17 @@ private static async Task RunDynamicFunctionPreservingInvocationOrderAsync(bool Assert.Equal(InvokeAsync, functionClient.FunctionInvoker); } - static ChatClientAgentRunOptions RecreateOptions(AgentRunOptions? options) + ChatClientAgentRunOptions RecreateOptions(AgentRunOptions? options) { var original = Assert.IsType(options); var factory = Assert.IsType>(original.ChatClientFactory); return new ChatClientAgentRunOptions(original.ChatOptions?.Clone()) { - ChatClientFactory = client => factory(new ConfigureOptionsChatClient(client, options => options.Temperature = 0.5f)), + ChatClientFactory = client => + { + var pipeline = factory(new ConfigureOptionsChatClient(client, options => options.Temperature = 0.5f)); + return hideServices ? new OpaqueChatClient(pipeline) : pipeline; + }, }; } } @@ -671,6 +683,253 @@ Task RunAsync(AIAgent agent, AgentRunOptions? runOptions = null) Assert.Same(loader, Assert.Single(options.ChatOptions!.Tools!)); } + [Theory] + [InlineData(false, false)] + [InlineData(false, true)] + [InlineData(true, false)] + [InlineData(true, true)] + public async Task RunAsync_NestedAgent_IsolatesFunctionMiddlewareAsync(bool streaming, bool nestedStreaming) + { + // Arrange + var innerInvocations = new List(); + var outerInvocations = new List(); + var innerFunction = AIFunctionFactory.Create(() => "Inner result", "InnerFunction"); + var innerResponses = new Queue( + [ + new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("inner", innerFunction.Name)])), + new(new ChatMessage(ChatRole.Assistant, "Inner complete")), + ]); + var innerAgent = new ChatClientAgent(CreateMockChatClient(innerResponses).Object, tools: [innerFunction]).AsBuilder() + .Use((agent, context, next, cancellationToken) => + { + innerInvocations.Add(context.Function.Name); + return next(context, cancellationToken); + }).Build(); + + var outerFunction = AIFunctionFactory.Create(async () => + { + await Task.Yield(); + var response = nestedStreaming + ? await innerAgent.RunStreamingAsync("Run the inner function").ToAgentResponseAsync() + : await innerAgent.RunAsync("Run the inner function"); + return response.Text; + }, "OuterFunction"); + var outerResponses = new Queue( + [ + new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("outer", outerFunction.Name)])), + new(new ChatMessage(ChatRole.Assistant, "Outer complete")), + ]); + var outerAgent = new ChatClientAgent(CreateMockChatClient(outerResponses).Object, tools: [outerFunction]).AsBuilder() + .Use((agent, context, next, cancellationToken) => + { + outerInvocations.Add(context.Function.Name); + return next(context, cancellationToken); + }).Build(); + + // Act + var response = streaming + ? await outerAgent.RunStreamingAsync("Run the outer function").ToAgentResponseAsync() + : await outerAgent.RunAsync("Run the outer function"); + + // Assert + Assert.Equal([innerFunction.Name], innerInvocations); + Assert.Equal([outerFunction.Name], outerInvocations); + Assert.Equal("Inner complete", Assert.IsType(response.Messages.SelectMany(m => m.Contents) + .OfType().Single(r => r.CallId == "outer").Result).GetString()); + Assert.Empty(innerResponses); + Assert.Empty(outerResponses); + } + + [Theory] + [InlineData(false, "Completed")] + [InlineData(false, "Faulted")] + [InlineData(false, "Canceled")] + [InlineData(true, "Completed")] + [InlineData(true, "Faulted")] + [InlineData(true, "Canceled")] + [InlineData(true, "Disposed")] + public async Task ChatClient_RequestScope_RestoresCallerContextAsync(bool streaming, string outcome) + { + // Arrange + var firstInvocations = new List(); + var secondInvocations = new List(); + var first = await CaptureClientAsync(firstInvocations); + var second = await CaptureClientAsync(secondInvocations); + using var firstClient = first.Client; + using var secondClient = second.Client; + var firstFunction = AIFunctionFactory.Create(() => "First result", "FirstFunction"); + var secondFunction = AIFunctionFactory.Create(() => "Second result", "SecondFunction"); + var firstOptions = new ChatOptions { Tools = [firstFunction] }; + using var cancellation = new CancellationTokenSource(); + cancellation.Cancel(); + var token = outcome == "Canceled" ? cancellation.Token : CancellationToken.None; + Exception? expectedFailure = outcome switch + { + "Faulted" => new InvalidOperationException("Request failed"), + "Canceled" => new OperationCanceledException(token), + _ => null, + }; + if (expectedFailure is not null) + { + first.Mock.Setup(c => c.GetResponseAsync( + It.IsAny>(), It.IsAny(), It.IsAny())) + .ThrowsAsync(expectedFailure); + first.Mock.Setup(c => c.GetStreamingResponseAsync( + It.IsAny>(), It.IsAny(), It.IsAny())) + .Returns(FailingResponseAsync(expectedFailure)); + } + else + { + if (outcome == "Completed") + { + first.Responses.Enqueue(new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("first", firstFunction.Name)]))); + } + + first.Responses.Enqueue(new(new ChatMessage(ChatRole.Assistant, "Complete"))); + } + + var callerContext = AIAgent.CurrentRunContext; + var secondRequestCount = 0; + + // Act + if (outcome == "Disposed") + { + await using var iterator = firstClient.GetStreamingResponseAsync([new(ChatRole.User, "First")], firstOptions).GetAsyncEnumerator(); + Assert.True(await iterator.MoveNextAsync()); + await VerifySecondClientAsync(); + } + else + { + var observedFailure = false; + try + { + if (streaming) + { + await foreach (var _ in firstClient.GetStreamingResponseAsync([new(ChatRole.User, "First")], firstOptions, token)) + { + } + } + else + { + await firstClient.GetResponseAsync([new(ChatRole.User, "First")], firstOptions, token); + } + } + catch (InvalidOperationException exception) when (ReferenceEquals(exception, expectedFailure)) + { + observedFailure = true; + } + catch (OperationCanceledException) when (outcome == "Canceled") + { + observedFailure = true; + } + + Assert.Equal(expectedFailure is not null, observedFailure); + } + + // Assert + await VerifySecondClientAsync(); + Assert.Equal(outcome == "Completed" ? new[] { firstFunction.Name } : [], firstInvocations); + + async Task VerifySecondClientAsync() + { + // Invoke clients directly under the same run context so a new agent run cannot mask a leaked scope. + Assert.Same(callerContext, AIAgent.CurrentRunContext); + second.Responses.Enqueue(new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("second", secondFunction.Name)]))); + second.Responses.Enqueue(new(new ChatMessage(ChatRole.Assistant, "Complete"))); + var response = await secondClient.GetResponseAsync( + [new(ChatRole.User, "Second")], new ChatOptions { Tools = [secondFunction] }); + + Assert.Equal("Second result", Assert.IsType(Assert.Single(response.Messages + .SelectMany(m => m.Contents).OfType()).Result).GetString()); + Assert.DoesNotContain(secondFunction.Name, firstInvocations); + Assert.Equal(++secondRequestCount, secondInvocations.Count); + Assert.All(secondInvocations, name => Assert.Equal(secondFunction.Name, name)); + Assert.Same(callerContext, AIAgent.CurrentRunContext); + } + + static async Task<(IChatClient Client, Mock Mock, Queue Responses)> CaptureClientAsync(List invocations) + { + var responses = new Queue([new(new ChatMessage(ChatRole.Assistant, "Initialized"))]); + var mock = CreateMockChatClient(responses); + IChatClient? capturedClient = null; + var agent = new ChatClientAgent(mock.Object).AsBuilder() + .Use((agent, context, next, cancellationToken) => + { + invocations.Add(context.Function.Name); + return next(context, cancellationToken); + }).Build(); + await agent.RunAsync("Initialize", options: new ChatClientAgentRunOptions + { + ChatClientFactory = client => capturedClient = client, + }); + return (Assert.IsAssignableFrom(capturedClient), mock, responses); + } + + static async IAsyncEnumerable FailingResponseAsync(Exception exception) + { + await Task.Yield(); + await Task.FromException(exception); + yield break; + } + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task RunAsync_ToolReplacementBeforeFailure_PreservesFunctionMiddlewareAsync(bool streaming) + { + // Arrange + var invocations = new List(); + var functionExecuted = false; + var function = AIFunctionFactory.Create(() => + { + functionExecuted = true; + return "Function result"; + }, "DynamicFunction"); + var expectedFailure = new InvalidOperationException("Loader failed after updating tools"); + string LoadFunction() + { + FunctionInvokingChatClient.CurrentContext!.Options!.Tools = [function]; + throw expectedFailure; + } + + var loader = AIFunctionFactory.Create(LoadFunction); + var responses = new Queue( + [ + new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("load", loader.Name)])), + new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("invoke", function.Name)])), + new(new ChatMessage(ChatRole.Assistant, "Complete")), + ]); + using var client = new FunctionInvokingChatClient(CreateMockChatClient(responses).Object) + { + MaximumConsecutiveErrorsPerRequest = 1, + }; + var agent = new ChatClientAgent(client, new ChatClientAgentOptions + { + UseProvidedChatClientAsIs = true, + ChatOptions = new() { Tools = [loader] }, + }).AsBuilder().Use((agent, context, next, cancellationToken) => + { + invocations.Add(context.Function.Name); + return context.Function.Name == function.Name + ? new ValueTask("Handled by middleware") + : next(context, cancellationToken); + }).Build(); + + // Act + var response = streaming + ? await agent.RunStreamingAsync("Run the functions").ToAgentResponseAsync() + : await agent.RunAsync("Run the functions"); + + // Assert + Assert.Equal([loader.Name, function.Name], invocations); + Assert.False(functionExecuted); + var results = response.Messages.SelectMany(m => m.Contents).OfType().ToArray(); + Assert.Same(expectedFailure, Assert.Single(results, result => result.CallId == "load").Exception); + Assert.Equal("Handled by middleware", Assert.Single(results, result => result.CallId == "invoke").Result); + Assert.Empty(responses); + } + #region Context Validation Tests /// @@ -1390,6 +1649,12 @@ private static Mock CreateMockChatClient(Queue respon return mockChatClient; } + private sealed class OpaqueChatClient(IChatClient innerClient) : DelegatingChatClient(innerClient) + { + public override object? GetService(Type serviceType, object? serviceKey = null) + => serviceKey is null && serviceType.IsInstanceOfType(this) ? this : null; + } + private sealed class TrackingFunctionInvokingChatClient(IChatClient innerClient, List executionOrder, string functionName) : FunctionInvokingChatClient(innerClient) { From 20087cc932017b132945fe1783515641c34bd239 Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Wed, 16 Sep 2026 10:17:34 +0000 Subject: [PATCH 3/7] fix: retain middleware through options clones Keep harness callback composition consistent across function-loop iterations by retaining the full chain on function wrappers. Cover real client-driven clones, repeated calls and approval flows. --- .../FunctionInvocationDelegatingAgent.cs | 33 +-- ...ocationDelegatingAgentBuilderExtensions.cs | 3 +- .../FunctionInvocationDelegatingAgentTests.cs | 215 ++++++++++++++++++ 3 files changed, 236 insertions(+), 15 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgent.cs b/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgent.cs index 40dac4d2989..e775b6847b6 100644 --- a/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgent.cs @@ -128,23 +128,23 @@ private sealed class MiddlewareScope(AgentRunContext? runContext, FunctionInvoca private sealed class MiddlewareEnabledTools : Collection { - private readonly FunctionInvocationDelegatingAgent[] _middleware; - internal MiddlewareEnabledTools(IList tools, FunctionInvocationDelegatingAgent[] middleware) { - this._middleware = middleware; + this.MiddlewareChain = middleware; foreach (var tool in tools) { this.Add(tool); } } - internal void ApplyTo(ChatOptions options) + internal FunctionInvocationDelegatingAgent[] MiddlewareChain { get; } + + internal static void ApplyTo(ChatOptions options, FunctionInvocationDelegatingAgent[] middleware) { if (options.Tools is { } tools && - (tools is not MiddlewareEnabledTools existing || !ReferenceEquals(existing._middleware, this._middleware))) + (tools is not MiddlewareEnabledTools existing || !ReferenceEquals(existing.MiddlewareChain, middleware))) { - options.Tools = new MiddlewareEnabledTools(tools, this._middleware); + options.Tools = new MiddlewareEnabledTools(tools, middleware); } } @@ -156,9 +156,9 @@ private AITool Wrap(AITool tool) { if (tool is AIFunction function) { - foreach (var middleware in this._middleware) + foreach (var middleware in this.MiddlewareChain) { - function = MiddlewareEnabledFunction.Wrap(function, middleware); + function = MiddlewareEnabledFunction.Wrap(function, middleware, this.MiddlewareChain); } return function; @@ -168,11 +168,15 @@ private AITool Wrap(AITool tool) } } - private sealed class MiddlewareEnabledFunction(AIFunction innerFunction, FunctionInvocationDelegatingAgent middleware) : DelegatingAIFunction(innerFunction) + private sealed class MiddlewareEnabledFunction( + AIFunction innerFunction, + FunctionInvocationDelegatingAgent middleware, + FunctionInvocationDelegatingAgent[] middlewareChain) : DelegatingAIFunction(innerFunction) { private readonly FunctionInvocationDelegatingAgent _middleware = middleware; + private readonly FunctionInvocationDelegatingAgent[] _middlewareChain = middlewareChain; - internal static AIFunction Wrap(AIFunction function, FunctionInvocationDelegatingAgent middleware) + internal static AIFunction Wrap(AIFunction function, FunctionInvocationDelegatingAgent middleware, FunctionInvocationDelegatingAgent[] middlewareChain) { for (var existing = function.GetService(); existing is not null; @@ -184,7 +188,7 @@ internal static AIFunction Wrap(AIFunction function, FunctionInvocationDelegatin } } - return new MiddlewareEnabledFunction(function, middleware); + return new MiddlewareEnabledFunction(function, middleware, middlewareChain); } protected override async ValueTask InvokeCoreAsync(AIFunctionArguments arguments, CancellationToken cancellationToken) @@ -197,7 +201,8 @@ internal static AIFunction Wrap(AIFunction function, FunctionInvocationDelegatin CallContent = new(string.Empty, this.InnerFunction.Name, new Dictionary(arguments)), }; - var tools = context.Options?.Tools as MiddlewareEnabledTools; + // Function wrappers survive ChatOptions.Clone even when the middleware-aware collection does not. + var middlewareChain = (context.Options?.Tools as MiddlewareEnabledTools)?.MiddlewareChain ?? this._middlewareChain; try { return await this._middleware._delegateFunc(this._middleware.InnerAgent, context, CoreLogicAsync, cancellationToken).ConfigureAwait(false); @@ -205,9 +210,9 @@ internal static AIFunction Wrap(AIFunction function, FunctionInvocationDelegatin finally { // A function or middleware can replace the entire collection during invocation. - if (tools is not null && context.Options is { } options) + if (context.Options is { } options) { - tools.ApplyTo(options); + MiddlewareEnabledTools.ApplyTo(options, middlewareChain); } } diff --git a/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgentBuilderExtensions.cs b/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgentBuilderExtensions.cs index 02d5ab64c32..996adc835f5 100644 --- a/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgentBuilderExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgentBuilderExtensions.cs @@ -32,7 +32,8 @@ public static class FunctionInvocationDelegatingAgentBuilderExtensions /// /// The callbacks also apply to functions added to or replaced in the current /// collection during execution. Replacing the collection within a function or callback preserves the - /// callbacks for subsequent invocations. + /// callbacks for subsequent invocations, including when the function-calling client clones the options + /// between iterations. /// /// /// The inner agent or the pipeline wrapping it must include a . If one does not exist, diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/FunctionInvocationDelegatingAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/FunctionInvocationDelegatingAgentTests.cs index fff734b761b..0893a1aa847 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/FunctionInvocationDelegatingAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/FunctionInvocationDelegatingAgentTests.cs @@ -364,6 +364,221 @@ public async Task RunAsync_DynamicallyAddedFunction_InvokesMiddlewareAsync(bool Assert.Empty(responses); } + [Theory] + [InlineData(false, "None", false)] + [InlineData(true, "None", false)] + [InlineData(false, "None", true)] + [InlineData(true, "None", true)] + [InlineData(false, "RequiredTool", false)] + [InlineData(true, "RequiredTool", false)] + [InlineData(false, "RequiredTool", true)] + [InlineData(true, "RequiredTool", true)] + [InlineData(false, "ConversationId", false)] + [InlineData(true, "ConversationId", false)] + [InlineData(false, "ConversationId", true)] + [InlineData(true, "ConversationId", true)] + public async Task RunAsync_InterIterationOptionsClone_PreservesDynamicFunctionMiddlewareAsync( + bool streaming, string cloneTrigger, bool replaceTools) + => await VerifyInterIterationOptionsCloneAsync(streaming, cloneTrigger, replaceTools); + + [Theory] + [InlineData(false, "RequiredTool")] + [InlineData(true, "RequiredTool")] + [InlineData(false, "ConversationId")] + [InlineData(true, "ConversationId")] + public async Task RunAsync_InterIterationOptionsClone_AllowedFunctionExecutesAsync(bool streaming, string cloneTrigger) + => await VerifyInterIterationOptionsCloneAsync(streaming, cloneTrigger, replaceTools: true, allowExecution: true); + + [Theory] + [InlineData(false, "RequiredTool")] + [InlineData(true, "RequiredTool")] + [InlineData(false, "ConversationId")] + [InlineData(true, "ConversationId")] + public async Task RunAsync_InterIterationOptionsClone_LoaderFailurePreservesMiddlewareAsync(bool streaming, string cloneTrigger) + => await VerifyInterIterationOptionsCloneAsync(streaming, cloneTrigger, replaceTools: true, loaderThrows: true); + + private static async Task VerifyInterIterationOptionsCloneAsync( + bool streaming, string cloneTrigger, bool replaceTools, bool allowExecution = false, bool loaderThrows = false) + { + // Arrange + var invocations = new List(); + var invocationCount = 0; + var loaderFailure = new InvalidOperationException("Loader failed after updating tools"); + ChatOptions? initialOptions = null; + var function = AIFunctionFactory.Create(() => + { + invocationCount++; + return "Function result"; + }, "DynamicFunction"); + var secondLoader = AIFunctionFactory.Create(() => + { + var options = FunctionInvokingChatClient.CurrentContext!.Options!; + if (cloneTrigger == "None") + { + Assert.Same(initialOptions, options); + } + else + { + Assert.NotSame(initialOptions, options); + } + + if (replaceTools) + { + options.Tools = [function]; + } + else + { + options.Tools!.Add(function); + } + + if (loaderThrows) + { + throw loaderFailure; + } + + return "Function added"; + }, "SecondLoader"); + var firstLoader = AIFunctionFactory.Create(() => + { + initialOptions = FunctionInvokingChatClient.CurrentContext!.Options!; + initialOptions.Tools!.Add(secondLoader); + return "Second loader added"; + }, "FirstLoader"); + var conversationId = cloneTrigger == "ConversationId" ? "conversation" : null; + var responses = new Queue( + [ + new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("first", firstLoader.Name)])) { ConversationId = conversationId }, + new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("second", secondLoader.Name)])) { ConversationId = conversationId }, + new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("invoke", function.Name)])) { ConversationId = conversationId }, + new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("invoke-again", function.Name)])) { ConversationId = conversationId }, + new(new ChatMessage(ChatRole.Assistant, "Complete")) { ConversationId = conversationId }, + ]); + using var client = new FunctionInvokingChatClient(CreateMockChatClient(responses).Object) + { + MaximumConsecutiveErrorsPerRequest = 1, + }; + var innerAgent = new ChatClientAgent(client, tools: [firstLoader]); + var first = innerAgent.AsBuilder().Use(async (agent, context, next, cancellationToken) => + { + invocations.Add($"First-Pre-{context.Function.Name}"); + var result = await next(context, cancellationToken); + invocations.Add($"First-Post-{context.Function.Name}"); + return result; + }).Build(); + var decorated = first.AsBuilder().Use(async (agent, context, next, cancellationToken) => + { + invocations.Add($"Second-Pre-{context.Function.Name}"); + var result = context.Function.Name == function.Name && !allowExecution + ? "Handled by middleware" + : await next(context, cancellationToken); + invocations.Add($"Second-Post-{context.Function.Name}"); + return result; + }).Build(); + var options = new ChatClientAgentRunOptions(new ChatOptions + { + ToolMode = cloneTrigger == "RequiredTool" ? ChatToolMode.RequireAny : null, + }); + + // Act + var response = streaming + ? await decorated.RunStreamingAsync("Run the functions", options: options).ToAgentResponseAsync() + : await decorated.RunAsync("Run the functions", options: options); + + // Assert + Assert.Equal(new[] { firstLoader.Name, secondLoader.Name, function.Name, function.Name }.SelectMany(name => + loaderThrows && name == secondLoader.Name + ? new[] { $"First-Pre-{name}", $"Second-Pre-{name}" } + : new[] { $"First-Pre-{name}", $"Second-Pre-{name}", $"Second-Post-{name}", $"First-Post-{name}" }), invocations); + Assert.Equal(allowExecution ? 2 : 0, invocationCount); + var results = response.Messages.SelectMany(m => m.Contents).OfType().ToArray(); + foreach (var callId in new[] { "invoke", "invoke-again" }) + { + var result = Assert.Single(results, result => result.CallId == callId).Result; + if (allowExecution) + { + Assert.Equal("Function result", Assert.IsType(result).GetString()); + } + else + { + Assert.Equal("Handled by middleware", result); + } + } + + Assert.Same(loaderThrows ? loaderFailure : null, Assert.Single(results, result => result.CallId == "second").Exception); + Assert.Empty(responses); + } + + [Theory] + [InlineData(false, false)] + [InlineData(true, false)] + [InlineData(false, true)] + [InlineData(true, true)] + public async Task RunAsync_InterIterationOptionsClone_PreservesApprovalAsync(bool streaming, bool approved) + { + // Arrange + var invocationCount = 0; + var middlewareInvocations = new List(); + ChatOptions? initialOptions = null; + var function = new ApprovalRequiredAIFunction(AIFunctionFactory.Create(() => + { + invocationCount++; + return "Function result"; + }, "DynamicFunction")); + var secondLoader = AIFunctionFactory.Create(() => + { + var options = FunctionInvokingChatClient.CurrentContext!.Options!; + Assert.NotSame(initialOptions, options); + options.Tools!.Add(function); + return "Function added"; + }, "SecondLoader"); + var firstLoader = AIFunctionFactory.Create(() => + { + initialOptions = FunctionInvokingChatClient.CurrentContext!.Options!; + initialOptions.Tools!.Add(secondLoader); + return "Second loader added"; + }, "FirstLoader"); + var responses = new Queue( + [ + new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("first", firstLoader.Name)])), + new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("second", secondLoader.Name)])), + new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("invoke", function.Name)])), + ]); + var agent = new ChatClientAgent(CreateMockChatClient(responses).Object, tools: [firstLoader]).AsBuilder() + .Use((agent, context, next, cancellationToken) => + { + middlewareInvocations.Add(context.Function.Name); + return next(context, cancellationToken); + }).Build(); + var session = await agent.CreateSessionAsync(); + var runOptions = new ChatClientAgentRunOptions(new ChatOptions { ToolMode = ChatToolMode.RequireAny }); + + // Act + var response = streaming + ? await agent.RunStreamingAsync("Run the functions", session, runOptions).ToAgentResponseAsync() + : await agent.RunAsync("Run the functions", session, runOptions); + + // Assert + var request = Assert.Single(response.Messages.SelectMany(m => m.Contents).OfType()); + Assert.Equal(0, invocationCount); + Assert.Equal([firstLoader.Name, secondLoader.Name], middlewareInvocations); + Assert.Empty(responses); + + // Act + responses.Enqueue(new(new ChatMessage(ChatRole.Assistant, "Complete"))); + var approval = new ChatMessage(ChatRole.User, [request.CreateResponse(approved)]); + var resumedOptions = new ChatClientAgentRunOptions(new ChatOptions { Tools = [function] }); + var resumedResponse = streaming + ? await agent.RunStreamingAsync(approval, session, resumedOptions).ToAgentResponseAsync() + : await agent.RunAsync(approval, session, resumedOptions); + + // Assert + Assert.Equal(approved ? 1 : 0, invocationCount); + Assert.Equal(approved ? [firstLoader.Name, secondLoader.Name, function.Name] : new[] { firstLoader.Name, secondLoader.Name }, middlewareInvocations); + Assert.Single(resumedResponse.Messages.SelectMany(m => m.Contents) + .OfType(), result => result.CallId == "invoke"); + Assert.Empty(responses); + } + [Theory] [InlineData(false, false)] [InlineData(true, false)] From 33dd41ef46411fc0f4721c1660a1c78284aed265 Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Wed, 16 Sep 2026 11:00:50 +0000 Subject: [PATCH 4/7] test: extend dynamic tool middleware coverage Cover Python harness callback consistency across progressive tool loading, continuation changes, failures, approvals and reused options. --- .../core/test_function_invocation_logic.py | 331 ++++++++++++++++++ 1 file changed, 331 insertions(+) diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index 0ce63a34160..b575afe4ad2 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -15,12 +15,14 @@ from agent_framework import ( Agent, AgentSession, + ChatContext, ChatOptions, ChatResponse, ChatResponseUpdate, Content, FunctionInvocationContext, FunctionMiddleware, + FunctionTool, Message, ResponseStream, SupportsChatGetResponse, @@ -7268,6 +7270,21 @@ def _pte_text_response(text: str = "done") -> ChatResponse: return ChatResponse(messages=Message(role="assistant", contents=[text])) +def _pte_set_responses(client: SupportsChatGetResponse, responses: list[ChatResponse], *, streaming: bool) -> None: + if streaming: + client.streaming_responses = [ # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + [ + ChatResponseUpdate( + role="assistant", contents=message.contents, conversation_id=response.conversation_id + ) + for message in response.messages + ] + for response in responses + ] + else: + client.run_responses = responses # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + + @tool(name="factorial", approval_mode="never_require") def _pte_factorial(n: int) -> int: """Compute the factorial of n.""" @@ -7562,6 +7579,320 @@ def load_tool(context: FunctionInvocationContext) -> str: assert response.messages[-1].text == "done" +@pytest.mark.parametrize("streaming", [False, True], ids=["non-streaming", "streaming"]) +@pytest.mark.parametrize("transition", ["none", "required", "conversation"]) +@pytest.mark.parametrize("replace_tools", [False, True], ids=["append", "replace"]) +@pytest.mark.parametrize("outcome", ["blocked", "allowed", "loader-failed"]) +@pytest.mark.parametrize("max_iterations", [5], indirect=True) +async def test_second_generation_tools_preserve_middleware_across_iterations( + chat_client_base: SupportsChatGetResponse, + streaming: bool, + transition: str, + replace_tools: bool, + outcome: str, +) -> None: + """Tool-list updates and continuation changes preserve agent and run-level callbacks.""" + invocations: list[str] = [] + target_calls: list[str] = [] + model_options: list[tuple[Any, Any, list[str]]] = [] + streamed_results: list[tuple[str | None, Any]] | None = None + loader_failure = RuntimeError("loader failed after updating tools") + + class AgentLevelFunctionMiddleware(FunctionMiddleware): + async def process(self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]) -> None: + invocations.append(f"agent-before-{context.function.name}") + await call_next() + invocations.append(f"agent-after-{context.function.name}") + + class RunLevelFunctionMiddleware(FunctionMiddleware): + async def process(self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]) -> None: + invocations.append(f"run-before-{context.function.name}") + if context.function.name == "late_tool" and outcome != "allowed": + context.result = "blocked by policy" + else: + await call_next() + invocations.append(f"run-after-{context.function.name}") + + @chat_middleware + async def observe_options(context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None: + assert context.options is not None + model_options.append(( + context.options.get("tool_choice"), + context.options.get("conversation_id"), + [item.name for item in context.options.get("tools", []) if isinstance(item, FunctionTool)], + )) + await call_next() + + @tool(name="late_tool", approval_mode="never_require") + def late_tool() -> str: + target_calls.append("invoked") + return "target completed" + + @tool(name="second_loader", approval_mode="never_require") + def second_loader(context: FunctionInvocationContext) -> str: + assert context.tools is not None + if replace_tools: + context.tools[:] = [late_tool] + else: + context.add_tools(late_tool) + if outcome == "loader-failed": + raise loader_failure + return "target loaded" + + @tool(name="first_loader", approval_mode="never_require") + def first_loader(context: FunctionInvocationContext) -> str: + context.add_tools(second_loader) + return "second loader loaded" + + responses = [ + _pte_function_call_response("1", "first_loader"), + _pte_function_call_response("2", "second_loader"), + _pte_function_call_response("3", "late_tool"), + _pte_function_call_response("4", "late_tool"), + _pte_text_response(), + ] + if transition == "conversation": + for index, chat_response in enumerate(responses, start=1): + chat_response.conversation_id = f"conversation-{index}" + _pte_set_responses(chat_client_base, responses, streaming=streaming) + caller_tools = [first_loader] + caller_options: ChatOptions = {"tool_choice": "required" if transition == "required" else "auto"} + agent = Agent( + client=chat_client_base, tools=caller_tools, middleware=[AgentLevelFunctionMiddleware(), observe_options] + ) + run_middleware = [RunLevelFunctionMiddleware()] + session = AgentSession() + if streaming: + stream = agent.run( + "Load and call tools.", stream=True, session=session, options=caller_options, middleware=run_middleware + ) + updates = [update async for update in stream] + response = await stream.get_final_response() + streamed_results = [ + (content.call_id, content.result) + for update in updates + for content in update.contents + if content.type == "function_result" + ] + else: + response = await agent.run( + "Load and call tools.", session=session, options=caller_options, middleware=run_middleware + ) + + expected_invocations: list[str] = [] + for name in ["first_loader", "second_loader", "late_tool", "late_tool"]: + expected_invocations.extend([f"agent-before-{name}", f"run-before-{name}"]) + if name != "second_loader" or outcome != "loader-failed": + expected_invocations.extend([f"run-after-{name}", f"agent-after-{name}"]) + assert invocations == expected_invocations + assert target_calls == (["invoked", "invoked"] if outcome == "allowed" else []) + results = [ + content for message in response.messages for content in message.contents if content.type == "function_result" + ] + assert [content.call_id for content in results] == ["1", "2", "3", "4"] + assert results[0].result == "second loader loaded" + if outcome == "loader-failed": + assert results[1].exception == str(loader_failure) + assert results[1].result == "Error: Function failed." + else: + assert results[1].result == "target loaded" + expected_target_result = "target completed" if outcome == "allowed" else "blocked by policy" + assert [content.result for content in results[2:]] == [expected_target_result, expected_target_result] + if streaming: + assert streamed_results == [(content.call_id, content.result) for content in results] + assert response.messages[-1].text == "done" + assert [item[0] for item in model_options] == ( + ["required", None, None, None, None] if transition == "required" else ["auto"] * 5 + ) + assert [item[1] for item in model_options] == ( + [None, "conversation-1", "conversation-2", "conversation-3", "conversation-4"] + if transition == "conversation" + else [None] * 5 + ) + final_tools = ["late_tool"] if replace_tools else ["first_loader", "second_loader", "late_tool"] + assert [item[2] for item in model_options] == [ + ["first_loader"], + ["first_loader", "second_loader"], + final_tools, + final_tools, + final_tools, + ] + assert caller_tools == [first_loader] + assert caller_options == {"tool_choice": "required" if transition == "required" else "auto"} + assert session.service_session_id == ("conversation-5" if transition == "conversation" else None) + + +@pytest.mark.parametrize("streaming", [False, True], ids=["non-streaming", "streaming"]) +@pytest.mark.parametrize("approved", [False, True], ids=["rejected", "approved"]) +@pytest.mark.parametrize("max_iterations", [4], indirect=True) +async def test_second_generation_tools_preserve_middleware_on_approval_resume( + chat_client_base: SupportsChatGetResponse, streaming: bool, approved: bool +) -> None: + """A second-generation tool requires approval, then retains middleware on resume.""" + invocations: list[str] = [] + target_calls: list[str] = [] + + class RecordingMiddleware(FunctionMiddleware): + async def process(self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]) -> None: + invocations.append(context.function.name) + await call_next() + + @tool(name="guarded_tool", approval_mode="always_require") + def guarded_tool() -> str: + target_calls.append("invoked") + return "target completed" + + @tool(name="second_loader", approval_mode="never_require") + def second_loader(context: FunctionInvocationContext) -> str: + context.add_tools(guarded_tool) + return "target loaded" + + @tool(name="first_loader", approval_mode="never_require") + def first_loader(context: FunctionInvocationContext) -> str: + context.add_tools(second_loader) + return "second loader loaded" + + _pte_set_responses( + chat_client_base, + [ + _pte_function_call_response("1", "first_loader"), + _pte_function_call_response("2", "second_loader"), + _pte_function_call_response("3", "guarded_tool"), + ], + streaming=streaming, + ) + session = AgentSession() + agent = Agent(client=chat_client_base, tools=[first_loader], middleware=[RecordingMiddleware()]) + if streaming: + first_stream = agent.run( + "Load guarded tool.", stream=True, session=session, options={"tool_choice": "required"} + ) + first_updates = [update async for update in first_stream] + first_response = await first_stream.get_final_response() + assert not any( + content.type == "function_result" and content.call_id == "3" + for update in first_updates + for content in update.contents + ) + else: + first_response = await agent.run("Load guarded tool.", session=session, options={"tool_choice": "required"}) + + requests = [ + content + for message in first_response.messages + for content in message.contents + if content.type == "function_approval_request" + ] + assert len(requests) == 1 + assert invocations == ["first_loader", "second_loader"] + assert target_calls == [] + + _pte_set_responses(chat_client_base, [_pte_text_response()], streaming=streaming) + approval = Message(role="user", contents=[requests[0].to_function_approval_response(approved=approved)]) + expected_result = "target completed" if approved else "Error: Tool call invocation was rejected by user." + if streaming: + resumed_stream = agent.run(approval, stream=True, session=session, tools=[guarded_tool]) + resumed_updates = [update async for update in resumed_stream] + resumed_response = await resumed_stream.get_final_response() + assert [ + (content.call_id, content.result) + for update in resumed_updates + for content in update.contents + if content.type == "function_result" + ] == [("3", expected_result)] + else: + resumed_response = await agent.run(approval, session=session, tools=[guarded_tool]) + + assert invocations == ["first_loader", "second_loader"] + (["guarded_tool"] if approved else []) + assert target_calls == (["invoked"] if approved else []) + assert [ + (content.call_id, content.result) + for message in resumed_response.messages + for content in message.contents + if content.type == "function_result" + ] == [("3", expected_result)] + assert resumed_response.messages[-1].text == "done" + + +@pytest.mark.parametrize("streaming", [False, True], ids=["non-streaming", "streaming"]) +@pytest.mark.parametrize("max_iterations", [4], indirect=True) +async def test_second_generation_tools_keep_shared_options_and_middleware_isolated( + chat_client_base: SupportsChatGetResponse, streaming: bool +) -> None: + """Shared clients/options do not transfer dynamic tools or run middleware between agents.""" + target_calls: list[str] = [] + invocations: list[tuple[str, str]] = [] + + class RunPolicy(FunctionMiddleware): + def __init__(self, label: str, *, allowed: bool) -> None: + self.label = label + self.allowed = allowed + + async def process(self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]) -> None: + invocations.append((self.label, context.function.name)) + if context.function.name == "late_tool" and not self.allowed: + context.result = "blocked by policy" + return + await call_next() + + @tool(name="late_tool", approval_mode="never_require") + def late_tool() -> str: + target_calls.append("invoked") + return "target completed" + + @tool(name="second_loader", approval_mode="never_require") + def second_loader(context: FunctionInvocationContext) -> str: + context.add_tools(late_tool) + return "target loaded" + + @tool(name="first_loader", approval_mode="never_require") + def first_loader(context: FunctionInvocationContext) -> str: + assert context.tools == [first_loader] + context.add_tools(second_loader) + return "second loader loaded" + + caller_tools = [first_loader] + options: ChatOptions = {"tools": caller_tools, "tool_choice": "required"} + first_agent = Agent(client=chat_client_base) + second_agent = Agent(client=chat_client_base) + middleware: list[RunPolicy] + for agent, run_options, middleware, expected_result in [ + (first_agent, options, [RunPolicy("first", allowed=False)], "blocked by policy"), + (second_agent, options, [RunPolicy("second", allowed=True)], "target completed"), + (first_agent, options.copy(), [], "target completed"), + ]: + _pte_set_responses( + chat_client_base, + [ + _pte_function_call_response("1", "first_loader"), + _pte_function_call_response("2", "second_loader"), + _pte_function_call_response("3", "late_tool"), + _pte_text_response(), + ], + streaming=streaming, + ) + if streaming: + stream = agent.run("Load and call.", stream=True, options=run_options, middleware=middleware) + async for _ in stream: + pass + response = await stream.get_final_response() + else: + response = await agent.run("Load and call.", options=run_options, middleware=middleware) + assert [ + content.result + for message in response.messages + for content in message.contents + if content.type == "function_result" and content.call_id == "3" + ] == [expected_result] + assert options == {"tools": [first_loader], "tool_choice": "required"} + assert options["tools"] is caller_tools + + assert target_calls == ["invoked", "invoked"] + assert invocations == [ + (label, name) for label in ["first", "second"] for name in ["first_loader", "second_loader", "late_tool"] + ] + + async def test_add_tools_with_approval_required_tool(chat_client_base: SupportsChatGetResponse): @tool(name="secure_tool", approval_mode="always_require") def secure_tool(value: str) -> str: From db96d480a75d2033d9dd2d394652a28aa8122c2c Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Wed, 16 Sep 2026 11:14:50 +0000 Subject: [PATCH 5/7] test: simplify streaming response collection --- .../core/tests/core/test_function_invocation_logic.py | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index b575afe4ad2..38fac38e014 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -7556,15 +7556,7 @@ def load_tool(context: FunctionInvocationContext) -> str: agent = Agent(client=chat_client_base, tools=[load_tool], middleware=[PolicyMiddleware()]) expected_results = [("1", "target loaded"), ("2", "target completed" if allowed else "blocked by policy")] if streaming: - stream = agent.run("Load and call the tool.", stream=True) - updates = [update async for update in stream] - response = await stream.get_final_response() - assert [ - (content.call_id, content.result) - for update in updates - for content in update.contents - if content.type == "function_result" - ] == expected_results + response = await agent.run("Load and call the tool.", stream=True).get_final_response() else: response = await agent.run("Load and call the tool.") From 123cf1835ae2506b53eedbc7d9d977c20e6ee0ae Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Wed, 16 Sep 2026 12:35:28 +0000 Subject: [PATCH 6/7] fix: preserve harness pipeline boundaries Build callback chains around factory results instead of sharing them between active client calls. Track active invocation callbacks across opaque function decorators while preserving explicit continuations. Cover replacement factories, captured clients and failure cleanup. --- .../FunctionInvocationDelegatingAgent.cs | 86 +++--- ...ocationDelegatingAgentBuilderExtensions.cs | 2 +- .../FunctionInvocationDelegatingAgentTests.cs | 278 ++++++++++++++++-- 3 files changed, 313 insertions(+), 53 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgent.cs b/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgent.cs index e775b6847b6..f1bd7193904 100644 --- a/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgent.cs @@ -52,7 +52,7 @@ private ChatClientAgentRunOptions AgentRunOptionsWithFunctionMiddleware(AgentRun } var originalFactory = aco.ChatClientFactory; - aco.ChatClientFactory = chatClient => new FunctionMiddlewarePreservingChatClient(chatClient, this).Build(originalFactory); + aco.ChatClientFactory = chatClient => FunctionMiddlewarePreservingChatClient.Build(chatClient, originalFactory, this); return aco; } @@ -60,69 +60,66 @@ private ChatClientAgentRunOptions AgentRunOptionsWithFunctionMiddleware(AgentRun /// /// Preserves the function middleware chain when tools are added or replaced during a run. /// - private sealed class FunctionMiddlewarePreservingChatClient(IChatClient innerClient, FunctionInvocationDelegatingAgent middleware) : DelegatingChatClient(innerClient) + private sealed class FunctionMiddlewarePreservingChatClient( + IChatClient innerClient, FunctionInvocationDelegatingAgent[] middlewareChain) : DelegatingChatClient(innerClient) { - private static readonly AsyncLocal s_currentScope = new(); - private readonly FunctionInvocationDelegatingAgent _middleware = middleware; + private static readonly AsyncLocal s_buildScope = new(); + private readonly FunctionInvocationDelegatingAgent[] _middlewareChain = middlewareChain; - internal IChatClient Build(Func? originalFactory) + internal static FunctionMiddlewarePreservingChatClient Build( + IChatClient chatClient, Func? originalFactory, FunctionInvocationDelegatingAgent middleware) { - var builder = this.AsBuilder(); + var previous = s_buildScope.Value; + var scope = previous is not null && ReferenceEquals(previous.RunContext, CurrentRunContext) + ? previous + : new PipelineBuildScope(CurrentRunContext); + scope.Middleware.Insert(0, middleware); + s_buildScope.Value = scope; + try + { + var builder = chatClient.AsBuilder(); + if (originalFactory is not null) + { + builder.Use(originalFactory); + } - if (originalFactory is not null) + return new FunctionMiddlewarePreservingChatClient(builder.Build(), [.. scope.Middleware]); + } + finally { - builder.Use(originalFactory); + // Factory composition is synchronous; restore its construction scope before returning. + s_buildScope.Value = previous; } - - return builder.Build(); } public override async Task GetResponseAsync( IEnumerable messages, ChatOptions? options = null, CancellationToken cancellationToken = default) - { - var scope = this.CreateScope(s_currentScope.Value); - s_currentScope.Value = scope; - return await this.InnerClient.GetResponseAsync(messages, ConfigureOptions(options, scope), cancellationToken).ConfigureAwait(false); - } + => await this.InnerClient.GetResponseAsync(messages, this.ConfigureOptions(options), cancellationToken).ConfigureAwait(false); public override async IAsyncEnumerable GetStreamingResponseAsync( IEnumerable messages, ChatOptions? options = null, [EnumeratorCancellation] CancellationToken cancellationToken = default) { - var scope = this.CreateScope(s_currentScope.Value); - s_currentScope.Value = scope; - await foreach (var update in this.InnerClient.GetStreamingResponseAsync(messages, ConfigureOptions(options, scope), cancellationToken).ConfigureAwait(false)) + await foreach (var update in this.InnerClient.GetStreamingResponseAsync(messages, this.ConfigureOptions(options), cancellationToken).ConfigureAwait(false)) { yield return update; - - // Resume this request's scope after the consumer's execution context. - s_currentScope.Value = scope; } } - private MiddlewareScope CreateScope(MiddlewareScope? previous) - { - var runContext = CurrentRunContext; - // A nested agent run must not inherit the calling agent's callbacks. - return new(runContext, previous is not null && ReferenceEquals(previous.RunContext, runContext) - ? [.. previous.Middleware, this._middleware] - : [this._middleware]); - } - - private static ChatOptions ConfigureOptions(ChatOptions? options, MiddlewareScope scope) + private ChatOptions ConfigureOptions(ChatOptions? options) { options = options?.Clone() ?? new(); if (options.Tools is { } tools) { - options.Tools = new MiddlewareEnabledTools(tools, scope.Middleware); + options.Tools = new MiddlewareEnabledTools(tools, this._middlewareChain); } return options; } - private sealed class MiddlewareScope(AgentRunContext? runContext, FunctionInvocationDelegatingAgent[] middleware) + private sealed class PipelineBuildScope(AgentRunContext? runContext) { internal AgentRunContext? RunContext { get; } = runContext; - internal FunctionInvocationDelegatingAgent[] Middleware { get; } = middleware; + internal List Middleware { get; } = []; } } @@ -173,6 +170,7 @@ private sealed class MiddlewareEnabledFunction( FunctionInvocationDelegatingAgent middleware, FunctionInvocationDelegatingAgent[] middlewareChain) : DelegatingAIFunction(innerFunction) { + private static readonly AsyncLocal s_invocationScope = new(); private readonly FunctionInvocationDelegatingAgent _middleware = middleware; private readonly FunctionInvocationDelegatingAgent[] _middlewareChain = middlewareChain; @@ -201,6 +199,18 @@ internal static AIFunction Wrap(AIFunction function, FunctionInvocationDelegatin CallContent = new(string.Empty, this.InnerFunction.Name, new Dictionary(arguments)), }; + var previous = s_invocationScope.Value; + for (var active = previous; active is not null; active = active.Parent) + { + if (ReferenceEquals(active.Context, context) && ReferenceEquals(active.Middleware, this._middleware)) + { + return await CoreLogicAsync(context, cancellationToken).ConfigureAwait(false); + } + } + + // An opaque decorator can reach an already active callback for the same invocation. + s_invocationScope.Value = new(context, this._middleware, previous); + // Function wrappers survive ChatOptions.Clone even when the middleware-aware collection does not. var middlewareChain = (context.Options?.Tools as MiddlewareEnabledTools)?.MiddlewareChain ?? this._middlewareChain; try @@ -219,5 +229,13 @@ internal static AIFunction Wrap(AIFunction function, FunctionInvocationDelegatin ValueTask CoreLogicAsync(FunctionInvocationContext ctx, CancellationToken cancellationToken) => base.InvokeCoreAsync(ctx.Arguments, cancellationToken); } + + private sealed class InvocationScope( + FunctionInvocationContext context, FunctionInvocationDelegatingAgent middleware, InvocationScope? parent) + { + internal FunctionInvocationContext Context { get; } = context; + internal FunctionInvocationDelegatingAgent Middleware { get; } = middleware; + internal InvocationScope? Parent { get; } = parent; + } } } diff --git a/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgentBuilderExtensions.cs b/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgentBuilderExtensions.cs index 996adc835f5..de207fdabcb 100644 --- a/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgentBuilderExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgentBuilderExtensions.cs @@ -33,7 +33,7 @@ public static class FunctionInvocationDelegatingAgentBuilderExtensions /// The callbacks also apply to functions added to or replaced in the current /// collection during execution. Replacing the collection within a function or callback preserves the /// callbacks for subsequent invocations, including when the function-calling client clones the options - /// between iterations. + /// between iterations. A per-request chat-client factory can replace the client without removing the callbacks. /// /// /// The inner agent or the pipeline wrapping it must include a . If one does not exist, diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/FunctionInvocationDelegatingAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/FunctionInvocationDelegatingAgentTests.cs index 0893a1aa847..a4cd5363066 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/FunctionInvocationDelegatingAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/FunctionInvocationDelegatingAgentTests.cs @@ -955,6 +955,235 @@ public async Task RunAsync_NestedAgent_IsolatesFunctionMiddlewareAsync(bool stre Assert.Empty(outerResponses); } + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task RunAsync_ReplacingChatClientFactory_PreservesMiddlewareAsync(bool streaming) + { + // Arrange + var invocations = new List(); + var functionExecuted = false; + var function = AIFunctionFactory.Create(() => + { + functionExecuted = true; + return "Function result"; + }, "DynamicFunction"); + var loader = AIFunctionFactory.Create(() => + { + FunctionInvokingChatClient.CurrentContext!.Options!.Tools!.Add(function); + return "Function added"; + }, "LoadFunction"); + var responses = new Queue( + [ + new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("load", loader.Name)])), + new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("invoke", function.Name)])), + new(new ChatMessage(ChatRole.Assistant, "Complete")), + ]); + using var replacement = new FunctionInvokingChatClient(CreateMockChatClient(responses).Object); + var originalClient = new Mock(); + var first = new ChatClientAgent(originalClient.Object, tools: [loader]).AsBuilder() + .Use((agent, context, next, cancellationToken) => + { + invocations.Add($"First-{context.Function.Name}"); + return next(context, cancellationToken); + }).Build(); + var agent = first.AsBuilder().Use((agent, context, next, cancellationToken) => + { + invocations.Add($"Second-{context.Function.Name}"); + return context.Function.Name == function.Name + ? new ValueTask("Handled by middleware") + : next(context, cancellationToken); + }).Build(); + var options = new ChatClientAgentRunOptions { ChatClientFactory = _ => replacement }; + + // Act + var response = streaming + ? await agent.RunStreamingAsync("Run", options: options).ToAgentResponseAsync() + : await agent.RunAsync("Run", options: options); + + // Assert + Assert.Equal(["First-LoadFunction", "Second-LoadFunction", "First-DynamicFunction", "Second-DynamicFunction"], invocations); + Assert.False(functionExecuted); + Assert.Equal("Handled by middleware", response.Messages.SelectMany(m => m.Contents) + .OfType().Single(r => r.CallId == "invoke").Result); + Assert.Empty(responses); + originalClient.Verify(c => c.GetResponseAsync( + It.IsAny>(), It.IsAny(), It.IsAny()), Times.Never); + originalClient.Verify(c => c.GetStreamingResponseAsync( + It.IsAny>(), It.IsAny(), It.IsAny()), Times.Never); + } + + [Theory] + [InlineData(false, false)] + [InlineData(false, true)] + [InlineData(true, false)] + [InlineData(true, true)] + public async Task RunAsync_IndependentCapturedClient_DoesNotInheritMiddlewareAsync(bool streaming, bool nestedStreaming) + { + // Arrange + var innerInvocations = new List(); + var captured = await CaptureClientAsync(innerInvocations); + using var innerClient = captured.Client; + var innerFunction = AIFunctionFactory.Create(() => "Inner result", "InnerFunction"); + captured.Responses.Enqueue(new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("inner", innerFunction.Name)]))); + captured.Responses.Enqueue(new(new ChatMessage(ChatRole.Assistant, "Complete"))); + var outerFunction = AIFunctionFactory.Create(async () => + { + var options = new ChatOptions { Tools = [innerFunction] }; + if (nestedStreaming) + { + await foreach (var _ in innerClient.GetStreamingResponseAsync([new(ChatRole.User, "Inner")], options)) + { + } + } + else + { + await innerClient.GetResponseAsync([new(ChatRole.User, "Inner")], options); + } + + return "Outer result"; + }, "OuterFunction"); + var outerInvocations = new List(); + var responses = new Queue( + [ + new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("outer", outerFunction.Name)])), + new(new ChatMessage(ChatRole.Assistant, "Complete")), + ]); + var agent = new ChatClientAgent(CreateMockChatClient(responses).Object, tools: [outerFunction]).AsBuilder() + .Use((agent, context, next, cancellationToken) => + { + outerInvocations.Add(context.Function.Name); + return next(context, cancellationToken); + }).Build(); + + // Act + if (streaming) + { + await agent.RunStreamingAsync("Run").ToAgentResponseAsync(); + } + else + { + await agent.RunAsync("Run"); + } + + // Assert + Assert.Equal([outerFunction.Name], outerInvocations); + Assert.Equal([innerFunction.Name], innerInvocations); + Assert.Empty(responses); + Assert.Empty(captured.Responses); + } + + [Theory] + [InlineData(false, false)] + [InlineData(true, false)] + [InlineData(false, true)] + [InlineData(true, true)] + public async Task RunAsync_OpaqueFunctionDecorator_DoesNotDuplicateMiddlewareAsync(bool streaming, bool invokeTwice) + { + // Arrange + var invocations = new List(); + var invocationCount = 0; + var replaced = false; + var function = AIFunctionFactory.Create(() => ++invocationCount, "TestFunction"); + var responses = new Queue( + [ + new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("first", function.Name)])), + new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("second", function.Name)])), + new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("third", function.Name)])), + new(new ChatMessage(ChatRole.Assistant, "Complete")), + ]); + var first = new ChatClientAgent(CreateMockChatClient(responses).Object, tools: [function]).AsBuilder() + .Use(async (agent, context, next, cancellationToken) => + { + invocations.Add("First"); + if (!replaced) + { + replaced = true; + context.Options!.Tools = [new OpaqueFunction(context.Function)]; + } + + var result = await next(context, cancellationToken); + return invokeTwice ? await next(context, cancellationToken) : result; + }).Build(); + var agent = first.AsBuilder().Use((agent, context, next, cancellationToken) => + { + invocations.Add("Second"); + return next(context, cancellationToken); + }).Build(); + + // Act + if (streaming) + { + await agent.RunStreamingAsync("Run").ToAgentResponseAsync(); + } + else + { + await agent.RunAsync("Run"); + } + + // Assert + string[] expectedPerInvocation = invokeTwice ? ["First", "Second", "Second"] : ["First", "Second"]; + Assert.Equal(Enumerable.Range(0, 3).SelectMany(_ => expectedPerInvocation), invocations); + Assert.Equal(invokeTwice ? 6 : 3, invocationCount); + Assert.Empty(responses); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task ChatClient_FailedFactory_DoesNotLeakMiddlewareAsync(bool returnsNull) + { + // Arrange + var invocations = new List(); + var function = AIFunctionFactory.Create(() => "Function result", "TestFunction"); + var responses = new Queue( + [ + new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("invoke", function.Name)])), + new(new ChatMessage(ChatRole.Assistant, "Complete")), + ]); + using var client = new FunctionInvokingChatClient(CreateMockChatClient(responses).Object); + var expectedFailure = new InvalidOperationException("Factory failed"); + var failingFactory = await CaptureFactoryAsync("Failed", _ => returnsNull ? null! : throw expectedFailure); + var successfulFactory = await CaptureFactoryAsync("Successful", null); + + // Act + if (returnsNull) + { + Assert.Throws(() => failingFactory(client)); + } + else + { + Assert.Same(expectedFailure, Assert.Throws(() => failingFactory(client))); + } + + using var pipeline = successfulFactory(client); + await pipeline.GetResponseAsync([new(ChatRole.User, "Run")], new ChatOptions { Tools = [function] }); + + // Assert + Assert.Equal(["Successful"], invocations); + Assert.Empty(responses); + + async Task> CaptureFactoryAsync(string name, Func? factory) + { + Func? capturedFactory = null; + var capturingAgent = new AnonymousDelegatingAIAgent( + new ChatClientAgent(client), + (messages, session, options, agent, cancellationToken) => + { + capturedFactory = Assert.IsType(options).ChatClientFactory; + return Task.FromResult(new AgentResponse(new ChatMessage(ChatRole.Assistant, "Captured"))); + }, + runStreamingFunc: null); + var agent = capturingAgent.AsBuilder().Use((agent, context, next, cancellationToken) => + { + invocations.Add(name); + return next(context, cancellationToken); + }).Build(); + await agent.RunAsync("Capture", options: new ChatClientAgentRunOptions { ChatClientFactory = factory }); + return Assert.IsType>(capturedFactory); + } + } + [Theory] [InlineData(false, "Completed")] [InlineData(false, "Faulted")] @@ -1062,24 +1291,6 @@ async Task VerifySecondClientAsync() Assert.Same(callerContext, AIAgent.CurrentRunContext); } - static async Task<(IChatClient Client, Mock Mock, Queue Responses)> CaptureClientAsync(List invocations) - { - var responses = new Queue([new(new ChatMessage(ChatRole.Assistant, "Initialized"))]); - var mock = CreateMockChatClient(responses); - IChatClient? capturedClient = null; - var agent = new ChatClientAgent(mock.Object).AsBuilder() - .Use((agent, context, next, cancellationToken) => - { - invocations.Add(context.Function.Name); - return next(context, cancellationToken); - }).Build(); - await agent.RunAsync("Initialize", options: new ChatClientAgentRunOptions - { - ChatClientFactory = client => capturedClient = client, - }); - return (Assert.IsAssignableFrom(capturedClient), mock, responses); - } - static async IAsyncEnumerable FailingResponseAsync(Exception exception) { await Task.Yield(); @@ -1852,6 +2063,37 @@ public async Task RunAsync_WithBaseAgentRunOptions_PreservesAllOriginalOptionsAs #endregion + private static async Task<(IChatClient Client, Mock Mock, Queue Responses)> CaptureClientAsync(List invocations) + { + var responses = new Queue([new(new ChatMessage(ChatRole.Assistant, "Initialized"))]); + var mock = CreateMockChatClient(responses); + IChatClient? capturedClient = null; + var capturingAgent = new AnonymousDelegatingAIAgent( + new ChatClientAgent(mock.Object), + (messages, session, options, innerAgent, cancellationToken) => + { + var runOptions = Assert.IsType(options?.Clone()); + var factory = Assert.IsType>(runOptions.ChatClientFactory); + runOptions.ChatClientFactory = client => capturedClient = factory(client); + return innerAgent.RunAsync(messages, session, runOptions, cancellationToken); + }, + runStreamingFunc: null); + var agent = capturingAgent.AsBuilder() + .Use((agent, context, next, cancellationToken) => + { + invocations.Add(context.Function.Name); + return next(context, cancellationToken); + }).Build(); + await agent.RunAsync("Initialize"); + return (Assert.IsAssignableFrom(capturedClient), mock, responses); + } + + private sealed class OpaqueFunction(AIFunction innerFunction) : DelegatingAIFunction(innerFunction) + { + public override object? GetService(Type serviceType, object? serviceKey = null) + => serviceKey is null && serviceType.IsInstanceOfType(this) ? this : null; + } + private static Mock CreateMockChatClient(Queue responses) { var mockChatClient = new Mock(); From f46c0171b51ca6a0ff58c1d3d905e7fb8e680ae7 Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Fri, 18 Sep 2026 12:23:19 +0100 Subject: [PATCH 7/7] docs(dotnet): explain function callback flow Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../FunctionInvocationDelegatingAgent.cs | 30 ++++++++++++++----- 1 file changed, 23 insertions(+), 7 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgent.cs b/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgent.cs index f1bd7193904..5fdadda9264 100644 --- a/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgent.cs @@ -11,7 +11,7 @@ namespace Microsoft.Agents.AI; /// -/// Internal agent decorator that adds function invocation middleware logic. +/// Internal agent wrapper that gives callbacks control over function calls. /// internal sealed class FunctionInvocationDelegatingAgent : DelegatingAIAgent { @@ -28,11 +28,12 @@ protected override Task RunCoreAsync(IEnumerable mes protected override IAsyncEnumerable RunCoreStreamingAsync(IEnumerable messages, AgentSession? session = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default) => this.InnerAgent.RunStreamingAsync(messages, session, this.AgentRunOptionsWithFunctionMiddleware(options), cancellationToken); - // Decorate options to add the middleware function + // Work on a per-run copy so adding callback support does not change options that the caller may reuse. private ChatClientAgentRunOptions AgentRunOptionsWithFunctionMiddleware(AgentRunOptions? options) { if (options is null || options.GetType() == typeof(AgentRunOptions)) { + // Plain agent options cannot hold a chat-client factory, so copy their shared values to chat-specific options. options = new ChatClientAgentRunOptions() { ResponseFormat = options?.ResponseFormat, @@ -52,6 +53,7 @@ private ChatClientAgentRunOptions AgentRunOptionsWithFunctionMiddleware(AgentRun } var originalFactory = aco.ChatClientFactory; + // Apply the original factory before adding our wrapper so even a replacement client keeps these callbacks. aco.ChatClientFactory = chatClient => FunctionMiddlewarePreservingChatClient.Build(chatClient, originalFactory, this); return aco; @@ -63,6 +65,7 @@ private ChatClientAgentRunOptions AgentRunOptionsWithFunctionMiddleware(AgentRun private sealed class FunctionMiddlewarePreservingChatClient( IChatClient innerClient, FunctionInvocationDelegatingAgent[] middlewareChain) : DelegatingChatClient(innerClient) { + // Concurrent runs must not combine the callbacks collected while their clients are built. private static readonly AsyncLocal s_buildScope = new(); private readonly FunctionInvocationDelegatingAgent[] _middlewareChain = middlewareChain; @@ -70,9 +73,11 @@ internal static FunctionMiddlewarePreservingChatClient Build( IChatClient chatClient, Func? originalFactory, FunctionInvocationDelegatingAgent middleware) { var previous = s_buildScope.Value; + // Nested agent wrappers build one client synchronously. Share their list only within the same run. var scope = previous is not null && ReferenceEquals(previous.RunContext, CurrentRunContext) ? previous : new PipelineBuildScope(CurrentRunContext); + // Function wrapping reverses this list, so insert at the front to keep callbacks in registration order. scope.Middleware.Insert(0, middleware); s_buildScope.Value = scope; try @@ -87,7 +92,7 @@ internal static FunctionMiddlewarePreservingChatClient Build( } finally { - // Factory composition is synchronous; restore its construction scope before returning. + // Restore the previous list so a later client build cannot reuse callbacks from this one. s_buildScope.Value = previous; } } @@ -107,9 +112,11 @@ public override async IAsyncEnumerable GetStreamingResponseA private ChatOptions ConfigureOptions(ChatOptions? options) { + // Each request gets its own options object, leaving caller-owned options unchanged. options = options?.Clone() ?? new(); if (options.Tools is { } tools) { + // This collection also wraps functions that are added or replaced later in the same run. options.Tools = new MiddlewareEnabledTools(tools, this._middlewareChain); } @@ -128,6 +135,7 @@ private sealed class MiddlewareEnabledTools : Collection internal MiddlewareEnabledTools(IList tools, FunctionInvocationDelegatingAgent[] middleware) { this.MiddlewareChain = middleware; + // Add also wraps the functions already present, just as it will wrap functions added later. foreach (var tool in tools) { this.Add(tool); @@ -138,6 +146,7 @@ internal MiddlewareEnabledTools(IList tools, FunctionInvocationDelegatin internal static void ApplyTo(ChatOptions options, FunctionInvocationDelegatingAgent[] middleware) { + // Keep the current collection only when it already uses this exact callback list. if (options.Tools is { } tools && (tools is not MiddlewareEnabledTools existing || !ReferenceEquals(existing.MiddlewareChain, middleware))) { @@ -145,6 +154,7 @@ internal static void ApplyTo(ChatOptions options, FunctionInvocationDelegatingAg } } + // Route later additions and replacements through the same callback wrapping. protected override void InsertItem(int index, AITool item) => base.InsertItem(index, this.Wrap(item)); protected override void SetItem(int index, AITool item) => base.SetItem(index, this.Wrap(item)); @@ -170,12 +180,14 @@ private sealed class MiddlewareEnabledFunction( FunctionInvocationDelegatingAgent middleware, FunctionInvocationDelegatingAgent[] middlewareChain) : DelegatingAIFunction(innerFunction) { + // Keep active callbacks with the current asynchronous operation instead of sharing process-wide state. private static readonly AsyncLocal s_invocationScope = new(); private readonly FunctionInvocationDelegatingAgent _middleware = middleware; private readonly FunctionInvocationDelegatingAgent[] _middlewareChain = middlewareChain; internal static AIFunction Wrap(AIFunction function, FunctionInvocationDelegatingAgent middleware, FunctionInvocationDelegatingAgent[] middlewareChain) { + // Do not add the callback again when the function exposes an existing wrapper for it. for (var existing = function.GetService(); existing is not null; existing = existing.InnerFunction.GetService()) @@ -192,7 +204,8 @@ internal static AIFunction Wrap(AIFunction function, FunctionInvocationDelegatin protected override async ValueTask InvokeCoreAsync(AIFunctionArguments arguments, CancellationToken cancellationToken) { var context = FunctionInvokingChatClient.CurrentContext - ?? new FunctionInvocationContext() // When there is no ambient context, create a new one to hold the arguments + // Direct function calls have no context from FunctionInvokingChatClient, so create the values the callback expects. + ?? new FunctionInvocationContext() { Arguments = arguments, Function = this.InnerFunction, @@ -200,6 +213,7 @@ internal static AIFunction Wrap(AIFunction function, FunctionInvocationDelegatin }; var previous = s_invocationScope.Value; + // A custom function wrapper can lead back to this callback during the same call. Run it only once. for (var active = previous; active is not null; active = active.Parent) { if (ReferenceEquals(active.Context, context) && ReferenceEquals(active.Middleware, this._middleware)) @@ -208,10 +222,11 @@ internal static AIFunction Wrap(AIFunction function, FunctionInvocationDelegatin } } - // An opaque decorator can reach an already active callback for the same invocation. + // Record the callback while it runs, including through wrappers that do not expose their inner function. s_invocationScope.Value = new(context, this._middleware, previous); - // Function wrappers survive ChatOptions.Clone even when the middleware-aware collection does not. + // ChatOptions.Clone keeps function wrappers but copies the tool list into a plain collection. + // In that case, use the callback list saved on this function. var middlewareChain = (context.Options?.Tools as MiddlewareEnabledTools)?.MiddlewareChain ?? this._middlewareChain; try { @@ -219,13 +234,14 @@ internal static AIFunction Wrap(AIFunction function, FunctionInvocationDelegatin } finally { - // A function or middleware can replace the entire collection during invocation. + // A function or callback can replace the entire tool list. Wrap the replacement before the next call. if (context.Options is { } options) { MiddlewareEnabledTools.ApplyTo(options, middlewareChain); } } + // Continue with the next function wrapper using any arguments changed by the callback. ValueTask CoreLogicAsync(FunctionInvocationContext ctx, CancellationToken cancellationToken) => base.InvokeCoreAsync(ctx.Arguments, cancellationToken); }