Add FunctionInvokingChatClient.CurrentContext (#5786)

* Add `FunctionInvokingChatClient.CurrentContext`

* PR feedback

Author: MackinnonBuck

src/Libraries/Microsoft.Extensions.AI/ChatCompletion/FunctionInvokingChatClient.cs

@@ -40,6 +40,9 @@ namespace Microsoft.Extensions.AI;
 /// </remarks>
 public partial class FunctionInvokingChatClient : DelegatingChatClient
 {
+    /// <summary>The <see cref="FunctionInvocationContext"/> for the current function invocation.</summary>
+    private static readonly AsyncLocal<FunctionInvocationContext?> _currentContext = new();
+
     /// <summary>The logger to use for logging information about function invocation.</summary>
     private readonly ILogger _logger;
 
@@ -50,6 +53,18 @@ public partial class FunctionInvokingChatClient : DelegatingChatClient
     /// <summary>Maximum number of roundtrips allowed to the inner client.</summary>
     private int? _maximumIterationsPerRequest;
 
+    /// <summary>
+    /// Gets or sets the <see cref="FunctionInvocationContext"/> for the current function invocation.
+    /// </summary>
+    /// <remarks>
+    /// This value flows across async calls.
+    /// </remarks>
+    public static FunctionInvocationContext? CurrentContext
+    {
+        get => _currentContext.Value;
+        set => _currentContext.Value = value;
+    }
+
     /// <summary>
     /// Initializes a new instance of the <see cref="FunctionInvokingChatClient"/> class.
     /// </summary>
@@ -661,6 +676,7 @@ FunctionResultContent CreateFunctionResultContent(FunctionInvocationResult resul
         object? result = null;
         try
         {
+            CurrentContext = context;
             result = await context.Function.InvokeAsync(context.CallContent.Arguments, cancellationToken).ConfigureAwait(false);
         }
         catch (Exception e)

test/Libraries/Microsoft.Extensions.AI.Tests/ChatCompletion/FunctionInvokingChatClientTests.cs

@@ -553,6 +553,93 @@ public async Task SupportsConsecutiveStreamingUpdatesWithFunctionCalls()
         Assert.Equal("OK bye", singleUpdateContent.Text);
     }
 
+    [Fact]
+    public async Task CanAccesssFunctionInvocationContextFromFunctionCall()
+    {
+        var invocationContexts = new List<FunctionInvokingChatClient.FunctionInvocationContext>();
+        var function = AIFunctionFactory.Create(async (int i) =>
+        {
+            // The context should propogate across async calls
+            await Task.Yield();
+
+            var context = FunctionInvokingChatClient.CurrentContext!;
+            invocationContexts.Add(context);
+
+            if (i == 42)
+            {
+                context.Terminate = true;
+            }
+
+            return $"Result {i}";
+        }, "Func1");
+
+        var options = new ChatOptions
+        {
+            Tools = [function],
+        };
+
+        // The invocation loop should terminate after the second function call
+        List<ChatMessage> planBeforeTermination =
+        [
+            new ChatMessage(ChatRole.User, "hello"),
+            new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("callId1", "Func1", new Dictionary<string, object?> { ["i"] = 41 })]),
+            new ChatMessage(ChatRole.Tool, [new FunctionResultContent("callId1", "Func1", result: "Result 41")]),
+            new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("callId2", "Func1", new Dictionary<string, object?> { ["i"] = 42 })]),
+            new ChatMessage(ChatRole.Tool, [new FunctionResultContent("callId2", "Func1", result: "Result 42")]),
+        ];
+
+        // The full plan should never be fulfilled
+        List<ChatMessage> plan =
+        [
+            .. planBeforeTermination,
+            new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("callId3", "Func1", new Dictionary<string, object?> { ["i"] = 43 })]),
+            new ChatMessage(ChatRole.Tool, [new FunctionResultContent("callId3", "Func1", result: "Result 43")]),
+            new ChatMessage(ChatRole.Assistant, "world"),
+        ];
+
+        await InvokeAsync(() => InvokeAndAssertAsync(options, plan, expected: [
+            .. planBeforeTermination,
+
+            // The last message is the one returned by the chat client
+            // This message's content should contain the last function call before the termination
+            new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("callId2", "Func1", new Dictionary<string, object?> { ["i"] = 42 })]),
+        ]));
+
+        await InvokeAsync(() => InvokeAndAssertStreamingAsync(options, plan, expected: [
+            .. planBeforeTermination,
+
+            // The last message is the one returned by the chat client
+            // When streaming, function call content is removed from this message
+            new ChatMessage(ChatRole.Assistant, []),
+        ]));
+
+        // The current context should be null outside the async call stack for the function invocation
+        Assert.Null(FunctionInvokingChatClient.CurrentContext);
+
+        async Task InvokeAsync(Func<Task<List<ChatMessage>>> work)
+        {
+            invocationContexts.Clear();
+
+            var chatMessages = await work();
+
+            Assert.Collection(invocationContexts,
+                c => AssertInvocationContext(c, iteration: 0, terminate: false),
+                c => AssertInvocationContext(c, iteration: 1, terminate: true));
+
+            void AssertInvocationContext(FunctionInvokingChatClient.FunctionInvocationContext context, int iteration, bool terminate)
+            {
+                Assert.NotNull(context);
+                Assert.Same(chatMessages, context.ChatMessages);
+                Assert.Same(function, context.Function);
+                Assert.Equal("Func1", context.CallContent.Name);
+                Assert.Equal(0, context.FunctionCallIndex);
+                Assert.Equal(1, context.FunctionCount);
+                Assert.Equal(iteration, context.Iteration);
+                Assert.Equal(terminate, context.Terminate);
+            }
+        }
+    }
+
     private static async Task<List<ChatMessage>> InvokeAndAssertAsync(
         ChatOptions options,
         List<ChatMessage> plan,