diff --git a/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgent.cs b/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgent.cs index 13f15ff1e9f..5fdadda9264 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; @@ -10,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 { @@ -27,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 - private AgentRunOptions? AgentRunOptionsWithFunctionMiddleware(AgentRunOptions? options) + // 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, @@ -40,6 +42,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,42 +53,205 @@ protected override IAsyncEnumerable RunCoreStreamingAsync(I } var originalFactory = aco.ChatClientFactory; - aco.ChatClientFactory = chatClient => + // 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; + } + + /// + /// Preserves the function middleware chain when tools are added or replaced during a run. + /// + 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; + + 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 + { + var builder = chatClient.AsBuilder(); + if (originalFactory is not null) + { + builder.Use(originalFactory); + } + + return new FunctionMiddlewarePreservingChatClient(builder.Build(), [.. scope.Middleware]); + } + finally + { + // Restore the previous list so a later client build cannot reuse callbacks from this one. + s_buildScope.Value = previous; + } + } + + 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) { - var builder = chatClient.AsBuilder(); + await foreach (var update in this.InnerClient.GetStreamingResponseAsync(messages, this.ConfigureOptions(options), cancellationToken).ConfigureAwait(false)) + { + yield return update; + } + } + + 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); + } + + return options; + } + + private sealed class PipelineBuildScope(AgentRunContext? runContext) + { + internal AgentRunContext? RunContext { get; } = runContext; + internal List Middleware { get; } = []; + } + } + + 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); + } + } + + internal FunctionInvocationDelegatingAgent[] MiddlewareChain { get; } - if (originalFactory is not null) + 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))) { - builder.Use(originalFactory); + options.Tools = new MiddlewareEnabledTools(tools, middleware); } + } - return builder.ConfigureOptions(co - => co.Tools = co.Tools?.Select(tool => tool is AIFunction aiFunction - ? new MiddlewareEnabledFunction(this.InnerAgent, aiFunction, this._delegateFunc) - : tool) - .ToList()) - .Build(); - }; + // Route later additions and replacements through the same callback wrapping. + protected override void InsertItem(int index, AITool item) => base.InsertItem(index, this.Wrap(item)); - return options; + 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.MiddlewareChain) + { + function = MiddlewareEnabledFunction.Wrap(function, middleware, this.MiddlewareChain); + } + + return function; + } + + return tool; + } } - private sealed class MiddlewareEnabledFunction(AIAgent innerAgent, AIFunction innerFunction, Func>, CancellationToken, ValueTask> next) : DelegatingAIFunction(innerFunction) + private sealed class MiddlewareEnabledFunction( + AIFunction innerFunction, + 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()) + { + if (ReferenceEquals(existing._middleware, middleware)) + { + return function; + } + } + + return new MiddlewareEnabledFunction(function, middleware, middlewareChain); + } + 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, CallContent = new(string.Empty, this.InnerFunction.Name, new Dictionary(arguments)), }; - return await next(innerAgent, context, CoreLogicAsync, cancellationToken).ConfigureAwait(false); + 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)) + { + return await CoreLogicAsync(context, cancellationToken).ConfigureAwait(false); + } + } + // Record the callback while it runs, including through wrappers that do not expose their inner function. + s_invocationScope.Value = new(context, this._middleware, previous); + + // 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 + { + return await this._middleware._delegateFunc(this._middleware.InnerAgent, context, CoreLogicAsync, cancellationToken).ConfigureAwait(false); + } + finally + { + // 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); } + + 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 5ff23f600c7..de207fdabcb 100644 --- a/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgentBuilderExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI/FunctionInvocationDelegatingAgentBuilderExtensions.cs @@ -30,6 +30,12 @@ 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, including when the function-calling client clones the options + /// 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, /// 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..a4cd5363066 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,1074 @@ 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, "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)] + [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); + + [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(); + 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); + } + + ChatClientAgentRunOptions RecreateOptions(AgentRunOptions? options) + { + var original = Assert.IsType(options); + var factory = Assert.IsType>(original.ChatClientFactory); + return new ChatClientAgentRunOptions(original.ChatOptions?.Clone()) + { + ChatClientFactory = client => + { + var pipeline = factory(new ConfigureOptionsChatClient(client, options => options.Temperature = 0.5f)); + return hideServices ? new OpaqueChatClient(pipeline) : pipeline; + }, + }; + } + } + + [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!)); + } + + [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)] + [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")] + [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 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 /// @@ -994,6 +2063,75 @@ 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(); + 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 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) + { + 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..38fac38e014 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -15,10 +15,14 @@ from agent_framework import ( Agent, AgentSession, + ChatContext, ChatOptions, ChatResponse, ChatResponseUpdate, Content, + FunctionInvocationContext, + FunctionMiddleware, + FunctionTool, Message, ResponseStream, SupportsChatGetResponse, @@ -36,8 +40,6 @@ included_token_count, ) from agent_framework._middleware import ( - FunctionInvocationContext, - FunctionMiddleware, FunctionMiddlewarePipeline, MiddlewareFailure, MiddlewareTermination, @@ -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.""" @@ -7493,6 +7510,381 @@ 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: + response = await agent.run("Load and call the tool.", stream=True).get_final_response() + 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" + + +@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: