.NET: Add AIContextProvider support (#691)

* Add AIContextProvider support

* Address feedback.

* Address PR comments.

* Switch to valuetask and remove parallel calls for AIContextProvider

* Remove Model from ModelInvokingAsync method name

* Remove agent thread id again and remove it from context provider interface

* Add AIContextProvider serialization support to AgentThread and update sample to show this feature

* Address PR comments

* Improve memory sample

* Update sample comment.

* Remove AggregateAIContextProvider for now since it makes too many assumptions.  We can include it later as a sample if needed.

* Update AIContextProviders to have an Invoked method instead of MessagesAddingAsync.

* Remove unused using.

* Address PR comments.

* Address PR comment.

* Update comment.

* Update comment

* Address PR comments.

dotnet/tests/Microsoft.Extensions.AI.Agents.Abstractions.UnitTests/AIAgentTests.cs

@@ -235,7 +235,7 @@ public class AIAgentTests
 
         await MockAgent.NotifyThreadOfNewMessagesAsync(threadMock.Object, messages, cancellationToken);
 
-        threadMock.Protected().Verify("OnNewMessagesAsync", Times.Once(), messages, cancellationToken);
+        threadMock.Protected().Verify("MessagesReceivedAsync", Times.Once(), messages, cancellationToken);
     }
 
     #region GetService Method Tests

dotnet/tests/Microsoft.Extensions.AI.Agents.Abstractions.UnitTests/AIContextProviderTests.cs

@@ -0,0 +1,56 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System.Collections.Generic;
+using System.Collections.ObjectModel;
+using System.Text.Json;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.Extensions.AI.Agents.Abstractions.UnitTests;
+
+public class AIContextProviderTests
+{
+    [Fact]
+    public async Task InvokedAsync_ReturnsCompletedTaskAsync()
+    {
+        var provider = new TestAIContextProvider();
+        var messages = new ReadOnlyCollection<ChatMessage>(new List<ChatMessage>());
+        var task = provider.InvokedAsync(new(messages));
+        Assert.Equal(default, task);
+    }
+
+    [Fact]
+    public async Task SerializeAsync_ReturnsEmptyElementAsync()
+    {
+        var provider = new TestAIContextProvider();
+        var actual = await provider.SerializeAsync();
+        Assert.Equal(default, actual);
+    }
+
+    [Fact]
+    public async Task DeserializeAsync_ReturnsCompletedTaskAsync()
+    {
+        var provider = new TestAIContextProvider();
+        var element = default(JsonElement);
+        var actual = provider.DeserializeAsync(element);
+        Assert.Equal(default, actual);
+    }
+
+    private sealed class TestAIContextProvider : AIContextProvider
+    {
+        public override ValueTask<AIContext> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default)
+        {
+            return default;
+        }
+
+        public override async ValueTask<JsonElement?> SerializeAsync(JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+        {
+            return await base.SerializeAsync(jsonSerializerOptions, cancellationToken);
+        }
+
+        public override async ValueTask DeserializeAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+        {
+            await base.DeserializeAsync(serializedState, jsonSerializerOptions, cancellationToken);
+        }
+    }
+}

dotnet/tests/Microsoft.Extensions.AI.Agents.Abstractions.UnitTests/AIContextTests.cs

@@ -0,0 +1,58 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System.Collections.Generic;
+
+namespace Microsoft.Extensions.AI.Agents.Abstractions.UnitTests;
+
+/// <summary>
+/// Unit tests for <see cref="AIContext"/>.
+/// </summary>
+public class AIContextTests
+{
+    [Fact]
+    public void SetInstructionsRoundtrips()
+    {
+        var context = new AIContext
+        {
+            Instructions = "Test Instructions"
+        };
+
+        Assert.Equal("Test Instructions", context.Instructions);
+    }
+
+    [Fact]
+    public void SetMessagesRoundtrips()
+    {
+        var context = new AIContext
+        {
+            Messages = new List<ChatMessage>
+            {
+                new(ChatRole.User, "Hello"),
+                new(ChatRole.Assistant, "Hi there!")
+            }
+        };
+
+        Assert.NotNull(context.Messages);
+        Assert.Equal(2, context.Messages.Count);
+        Assert.Equal("Hello", context.Messages[0].Text);
+        Assert.Equal("Hi there!", context.Messages[1].Text);
+    }
+
+    [Fact]
+    public void SetAIFunctionsRoundtrips()
+    {
+        var context = new AIContext
+        {
+            Tools = new List<AITool>
+            {
+                AIFunctionFactory.Create(() => "Function1", "Function1", "Description1"),
+                AIFunctionFactory.Create(() => "Function2", "Function2", "Description2"),
+            }
+        };
+
+        Assert.NotNull(context.Tools);
+        Assert.Equal(2, context.Tools.Count);
+        Assert.Equal("Function1", context.Tools[0].Name);
+        Assert.Equal("Function2", context.Tools[1].Name);
+    }
+}

dotnet/tests/Microsoft.Extensions.AI.Agents.Abstractions.UnitTests/AgentThreadTests.cs

@@ -8,6 +8,8 @@ using System.Threading;
 using System.Threading.Tasks;
 using Moq;
 
+#pragma warning disable CA1861 // Avoid constant arrays as arguments
+
 namespace Microsoft.Extensions.AI.Agents.Abstractions.UnitTests;
 
 public class AgentThreadTests
@@ -102,7 +104,7 @@ public class AgentThreadTests
         };
 
         // Act
-        await thread.OnNewMessagesAsync(messages, CancellationToken.None);
+        await thread.MessagesReceivedAsync(messages, CancellationToken.None);
         Assert.Equal("thread-123", thread.ConversationId);
         Assert.Null(thread.MessageStore);
     }
@@ -120,7 +122,7 @@ public class AgentThreadTests
         };
 
         // Act
-        await thread.OnNewMessagesAsync(messages, CancellationToken.None);
+        await thread.MessagesReceivedAsync(messages, CancellationToken.None);
 
         // Assert
         Assert.Equal(2, store.Count);
@@ -173,6 +175,26 @@ public class AgentThreadTests
         Assert.Null(thread.MessageStore);
     }
 
+    [Fact]
+    public async Task VerifyDeserializeWithAIContextProviderAsync()
+    {
+        // Arrange
+        var json = JsonSerializer.Deserialize("""
+            {
+                "aiContextProviderState": ["CP1"]
+            }
+            """, TestJsonSerializerContext.Default.JsonElement);
+        Mock<AIContextProvider> mockProvider = new();
+        var thread = new AgentThread() { AIContextProvider = mockProvider.Object };
+
+        // Act
+        await thread.DeserializeAsync(json);
+
+        // Assert
+        Assert.Null(thread.MessageStore);
+        mockProvider.Verify(m => m.DeserializeAsync(It.Is<JsonElement>(e => e.ValueKind == JsonValueKind.Array && e.GetArrayLength() == 1), It.IsAny<JsonSerializerOptions?>(), It.IsAny<CancellationToken>()), Times.Once);
+    }
+
     [Fact]
     public async Task DeserializeWithInvalidJsonThrowsAsync()
     {
@@ -245,6 +267,31 @@ public class AgentThreadTests
         Assert.Equal("TestContent", textContent.GetProperty("text").GetString());
     }
 
+    [Fact]
+    public async Task VerifyThreadSerializationWithWithAIContextProviderAsync()
+    {
+        // Arrange
+        Mock<AIContextProvider> mockProvider = new();
+        var providerStateElement = JsonSerializer.SerializeToElement(new[] { "CP1" }, TestJsonSerializerContext.Default.StringArray);
+        mockProvider
+            .Setup(m => m.SerializeAsync(It.IsAny<JsonSerializerOptions?>(), It.IsAny<CancellationToken>()))
+            .ReturnsAsync(providerStateElement);
+
+        var thread = new AgentThread();
+        thread.AIContextProvider = mockProvider.Object;
+
+        // Act
+        var json = await thread.SerializeAsync();
+
+        // Assert
+        Assert.Equal(JsonValueKind.Object, json.ValueKind);
+        Assert.True(json.TryGetProperty("aiContextProviderState", out var providerStateProperty));
+        Assert.Equal(JsonValueKind.Array, providerStateProperty.ValueKind);
+        Assert.Single(providerStateProperty.EnumerateArray());
+        Assert.Equal("CP1", providerStateProperty.EnumerateArray().First().GetString());
+        mockProvider.Verify(m => m.SerializeAsync(It.IsAny<JsonSerializerOptions?>(), It.IsAny<CancellationToken>()), Times.Once);
+    }
+
     /// <summary>
     /// Verify thread serialization to JSON with custom options.
     /// </summary>

dotnet/tests/Microsoft.Extensions.AI.Agents.Abstractions.UnitTests/TestJsonSerializerContext.cs

@@ -17,4 +17,5 @@ namespace Microsoft.Extensions.AI.Agents.Abstractions.UnitTests;
 [JsonSerializable(typeof(Animal))]
 [JsonSerializable(typeof(JsonElement))]
 [JsonSerializable(typeof(Dictionary<string, object?>))]
+[JsonSerializable(typeof(string[]))]
 internal sealed partial class TestJsonSerializerContext : JsonSerializerContext;

dotnet/tests/Microsoft.Extensions.AI.Agents.UnitTests/ChatCompletion/ChatClientAgentOptionsTests.cs

@@ -22,6 +22,7 @@ public class ChatClientAgentOptionsTests
         Assert.Null(options.Description);
         Assert.Null(options.ChatOptions);
         Assert.Null(options.ChatMessageStoreFactory);
+        Assert.Null(options.AIContextProviderFactory);
     }
 
     [Fact]
@@ -39,6 +40,7 @@ public class ChatClientAgentOptionsTests
         Assert.Null(options.Instructions);
         Assert.Null(options.Description);
         Assert.Null(options.ChatOptions);
+        Assert.Null(options.AIContextProviderFactory);
     }
 
     [Fact]
@@ -163,11 +165,13 @@ public class ChatClientAgentOptionsTests
         const string Description = "Test description";
         var tools = new List<AITool> { AIFunctionFactory.Create(() => "test") };
         static IChatMessageStore ChatMessageStoreFactory() => new Mock<IChatMessageStore>().Object;
+        static AIContextProvider AIContextProviderFactory() => new Mock<AIContextProvider>().Object;
 
         var original = new ChatClientAgentOptions(Instructions, Name, Description, tools)
         {
             Id = "test-id",
-            ChatMessageStoreFactory = ChatMessageStoreFactory
+            ChatMessageStoreFactory = ChatMessageStoreFactory,
+            AIContextProviderFactory = AIContextProviderFactory
         };
 
         // Act
@@ -180,6 +184,7 @@ public class ChatClientAgentOptionsTests
         Assert.Equal(original.Instructions, clone.Instructions);
         Assert.Equal(original.Description, clone.Description);
         Assert.Same(original.ChatMessageStoreFactory, clone.ChatMessageStoreFactory);
+        Assert.Same(original.AIContextProviderFactory, clone.AIContextProviderFactory);
 
         // ChatOptions should be cloned, not the same reference
         Assert.NotSame(original.ChatOptions, clone.ChatOptions);
@@ -209,5 +214,7 @@ public class ChatClientAgentOptionsTests
         Assert.Equal(original.Instructions, clone.Instructions);
         Assert.Equal(original.Description, clone.Description);
         Assert.Null(clone.ChatOptions);
+        Assert.Null(clone.ChatMessageStoreFactory);
+        Assert.Null(clone.AIContextProviderFactory);
     }
 }

dotnet/tests/Microsoft.Extensions.AI.Agents.UnitTests/ChatCompletion/ChatClientAgentTests.cs

@@ -10,6 +10,8 @@ namespace Microsoft.Extensions.AI.Agents.UnitTests.ChatCompletion;
 
 public class ChatClientAgentTests
 {
+    #region Constructor Tests
+
     /// <summary>
     /// Verify the invocation and response of <see cref="ChatClientAgent"/>.
     /// </summary>
@@ -38,6 +40,10 @@ public class ChatClientAgentTests
         Assert.Equal("AgentInvokedChatClient", agent.ChatClient.GetType().Name);
     }
 
+    #endregion
+
+    #region RunAsync Tests
+
     /// <summary>
     /// Verify the invocation and response of <see cref="ChatClientAgent"/> using <see cref="IChatClient"/>.
     /// </summary>
@@ -390,6 +396,176 @@ public class ChatClientAgentTests
         await Assert.ThrowsAsync<InvalidOperationException>(() => agent.RunAsync([new(ChatRole.User, "test")], thread));
     }
 
+    /// <summary>
+    /// Verify that RunAsync sets the ConversationId on the thread when the service returns one.
+    /// </summary>
+    [Fact]
+    public async Task RunAsyncSetsConversationIdOnThreadWhenReturnedByChatClientAsync()
+    {
+        // Arrange
+        Mock<IChatClient> mockService = new();
+        mockService.Setup(
+            s => s.GetResponseAsync(
+                It.IsAny<IEnumerable<ChatMessage>>(),
+                It.IsAny<ChatOptions>(),
+                It.IsAny<CancellationToken>())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")]) { ConversationId = "ConvId" });
+        ChatClientAgent agent = new(mockService.Object, options: new() { Instructions = "test instructions" });
+        AgentThread thread = new();
+
+        // Act
+        await agent.RunAsync([new(ChatRole.User, "test")], thread);
+
+        // Assert
+        Assert.Equal("ConvId", thread.ConversationId);
+    }
+
+    /// <summary>
+    /// Verify that RunAsync invokes any provided AIContextProvider and uses the result.
+    /// </summary>
+    [Fact]
+    public async Task RunAsyncInvokesAIContextProviderAndUsesResultAsync()
+    {
+        // Arrange
+        ChatMessage[] requestMessages = [new(ChatRole.User, "user message")];
+        ChatMessage[] responseMessages = [new(ChatRole.Assistant, "response")];
+        Mock<IChatClient> mockService = new();
+        List<ChatMessage> capturedMessages = [];
+        string capturedInstructions = string.Empty;
+        List<AITool> capturedTools = [];
+        mockService
+            .Setup(s => s.GetResponseAsync(
+                It.IsAny<IEnumerable<ChatMessage>>(),
+                It.IsAny<ChatOptions>(),
+                It.IsAny<CancellationToken>()))
+            .Callback<IEnumerable<ChatMessage>, ChatOptions, CancellationToken>((msgs, opts, ct) =>
+            {
+                capturedMessages.AddRange(msgs);
+                capturedInstructions = opts.Instructions ?? string.Empty;
+                if (opts.Tools != null)
+                {
+                    capturedTools.AddRange(opts.Tools);
+                }
+            })
+            .ReturnsAsync(new ChatResponse(responseMessages));
+
+        var mockProvider = new Mock<AIContextProvider>();
+        mockProvider
+            .Setup(p => p.InvokingAsync(It.IsAny<AIContextProvider.InvokingContext>(), It.IsAny<CancellationToken>()))
+            .ReturnsAsync(new AIContext
+            {
+                Messages = [new(ChatRole.System, "context provider message")],
+                Instructions = "context provider instructions",
+                Tools = [AIFunctionFactory.Create(() => { }, "context provider function")]
+            });
+        mockProvider
+            .Setup(p => p.InvokedAsync(It.IsAny<AIContextProvider.InvokedContext>(), It.IsAny<CancellationToken>()))
+            .Returns(new ValueTask());
+
+        ChatClientAgent agent = new(mockService.Object, options: new() { Instructions = "base instructions", AIContextProviderFactory = () => mockProvider.Object, ChatOptions = new() { Tools = [AIFunctionFactory.Create(() => { }, "base function")] } });
+
+        // Act
+        await agent.RunAsync(requestMessages);
+
+        // Assert
+        // Should contain: base instructions, context message, user message, base function, context function
+        Assert.Equal(2, capturedMessages.Count);
+        Assert.Equal("base instructions\ncontext provider instructions", capturedInstructions);
+        Assert.Equal("context provider message", capturedMessages[0].Text);
+        Assert.Equal(ChatRole.System, capturedMessages[0].Role);
+        Assert.Equal("user message", capturedMessages[1].Text);
+        Assert.Equal(ChatRole.User, capturedMessages[1].Role);
+        Assert.Equal(2, capturedTools.Count);
+        Assert.Contains(capturedTools, t => t.Name == "base function");
+        Assert.Contains(capturedTools, t => t.Name == "context provider function");
+        mockProvider.Verify(p => p.InvokingAsync(It.IsAny<AIContextProvider.InvokingContext>(), It.IsAny<CancellationToken>()), Times.Once);
+        mockProvider.Verify(p => p.InvokedAsync(It.Is<AIContextProvider.InvokedContext>(x => x.RequestMessages == requestMessages && x.ResponseMessages == responseMessages && x.InvokeException == null), It.IsAny<CancellationToken>()), Times.Once);
+    }
+
+    /// <summary>
+    /// Verify that RunAsync invokes any provided AIContextProvider when the downstream GetResponse call fails.
+    /// </summary>
+    [Fact]
+    public async Task RunAsyncInvokesAIContextProviderWhenGetResponseFailsAsync()
+    {
+        // Arrange
+        ChatMessage[] requestMessages = [new(ChatRole.User, "user message")];
+        ChatMessage[] responseMessages = [new(ChatRole.Assistant, "response")];
+        Mock<IChatClient> mockService = new();
+        mockService
+            .Setup(s => s.GetResponseAsync(
+                It.IsAny<IEnumerable<ChatMessage>>(),
+                It.IsAny<ChatOptions>(),
+                It.IsAny<CancellationToken>()))
+            .Throws(new InvalidOperationException("downstream failure"));
+
+        var mockProvider = new Mock<AIContextProvider>();
+        mockProvider
+            .Setup(p => p.InvokingAsync(It.IsAny<AIContextProvider.InvokingContext>(), It.IsAny<CancellationToken>()))
+            .ReturnsAsync(new AIContext());
+        mockProvider
+            .Setup(p => p.InvokedAsync(It.IsAny<AIContextProvider.InvokedContext>(), It.IsAny<CancellationToken>()))
+            .Returns(new ValueTask());
+
+        ChatClientAgent agent = new(mockService.Object, options: new() { Instructions = "base instructions", AIContextProviderFactory = () => mockProvider.Object, ChatOptions = new() { Tools = [AIFunctionFactory.Create(() => { }, "base function")] } });
+
+        // Act
+        await Assert.ThrowsAsync<InvalidOperationException>(() => agent.RunAsync(requestMessages));
+
+        // Assert
+        mockProvider.Verify(p => p.InvokingAsync(It.IsAny<AIContextProvider.InvokingContext>(), It.IsAny<CancellationToken>()), Times.Once);
+        mockProvider.Verify(p => p.InvokedAsync(It.Is<AIContextProvider.InvokedContext>(x => x.RequestMessages == requestMessages && x.ResponseMessages == null && x.InvokeException is InvalidOperationException), It.IsAny<CancellationToken>()), Times.Once);
+    }
+
+    /// <summary>
+    /// Verify that RunAsync invokes any provided AIContextProvider and succeeds even when the AIContext is empty.
+    /// </summary>
+    [Fact]
+    public async Task RunAsyncInvokesAIContextProviderAndSucceedsWithEmptyAIContextAsync()
+    {
+        // Arrange
+        Mock<IChatClient> mockService = new();
+        List<ChatMessage> capturedMessages = [];
+        string capturedInstructions = string.Empty;
+        List<AITool> capturedTools = [];
+        mockService
+            .Setup(s => s.GetResponseAsync(
+                It.IsAny<IEnumerable<ChatMessage>>(),
+                It.IsAny<ChatOptions>(),
+                It.IsAny<CancellationToken>()))
+            .Callback<IEnumerable<ChatMessage>, ChatOptions, CancellationToken>((msgs, opts, ct) =>
+            {
+                capturedMessages.AddRange(msgs);
+                capturedInstructions = opts.Instructions ?? string.Empty;
+                if (opts.Tools != null)
+                {
+                    capturedTools.AddRange(opts.Tools);
+                }
+            })
+            .ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")]));
+
+        var mockProvider = new Mock<AIContextProvider>();
+        mockProvider
+            .Setup(p => p.InvokingAsync(It.IsAny<AIContextProvider.InvokingContext>(), It.IsAny<CancellationToken>()))
+            .ReturnsAsync(new AIContext());
+
+        ChatClientAgent agent = new(mockService.Object, options: new() { Instructions = "base instructions", AIContextProviderFactory = () => mockProvider.Object, ChatOptions = new() { Tools = [AIFunctionFactory.Create(() => { }, "base function")] } });
+
+        // Act
+        await agent.RunAsync([new(ChatRole.User, "user message")]);
+
+        // Assert
+        // Should contain: base instructions, user message, base function
+        Assert.Single(capturedMessages);
+        Assert.Equal("base instructions", capturedInstructions);
+        Assert.Equal("user message", capturedMessages[0].Text);
+        Assert.Equal(ChatRole.User, capturedMessages[0].Role);
+        Assert.Single(capturedTools);
+        Assert.Contains(capturedTools, t => t.Name == "base function");
+        mockProvider.Verify(p => p.InvokingAsync(It.IsAny<AIContextProvider.InvokingContext>(), It.IsAny<CancellationToken>()), Times.Once);
+    }
+
+    #endregion
+
     #region Property Override Tests
 
     /// <summary>
@@ -1524,6 +1700,61 @@ public class ChatClientAgentTests
 
     #endregion
 
+    #region GetNewThread Tests
+
+    [Fact]
+    public void GetNewThreadUsesChatMessageStoreFactoryIfProvided()
+    {
+        // Arrange
+        var mockChatClient = new Mock<IChatClient>();
+        var mockStore = new Mock<IChatMessageStore>();
+        var factoryCalled = false;
+
+        var agent = new ChatClientAgent(mockChatClient.Object, new ChatClientAgentOptions
+        {
+            Instructions = "Test instructions",
+            ChatMessageStoreFactory = () =>
+            {
+                factoryCalled = true;
+                return mockStore.Object;
+            }
+        });
+
+        // Act
+        var thread = agent.GetNewThread();
+
+        // Assert
+        Assert.True(factoryCalled, "ChatMessageStoreFactory was not called.");
+        Assert.Same(mockStore.Object, thread.MessageStore);
+    }
+
+    [Fact]
+    public void GetNewThreadUsesAIContextProviderFactoryIfProvided()
+    {
+        // Arrange
+        var mockChatClient = new Mock<IChatClient>();
+        var mockContextProvider = new Mock<AIContextProvider>();
+        var factoryCalled = false;
+        var agent = new ChatClientAgent(mockChatClient.Object, new ChatClientAgentOptions
+        {
+            Instructions = "Test instructions",
+            AIContextProviderFactory = () =>
+            {
+                factoryCalled = true;
+                return mockContextProvider.Object;
+            }
+        });
+
+        // Act
+        var thread = agent.GetNewThread();
+
+        // Assert
+        Assert.True(factoryCalled, "AIContextProviderFactory was not called.");
+        Assert.Same(mockContextProvider.Object, thread.AIContextProvider);
+    }
+
+    #endregion
+
     private static async IAsyncEnumerable<T> ToAsyncEnumerableAsync<T>(IEnumerable<T> values)
     {
         await Task.Yield();

dotnet/tests/Microsoft.Extensions.AI.Agents.UnitTests/TestJsonSerializerContext.cs

@@ -0,0 +1,14 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System.Text.Json;
+using System.Text.Json.Serialization;
+
+namespace Microsoft.Extensions.AI.Agents.UnitTests;
+
+[JsonSourceGenerationOptions(
+    PropertyNamingPolicy = JsonKnownNamingPolicy.CamelCase,
+    DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull,
+    UseStringEnumConverter = true)]
+[JsonSerializable(typeof(JsonElement))]
+[JsonSerializable(typeof(string))]
+internal sealed partial class TestJsonSerializerContext : JsonSerializerContext;