From f8093389aa4f48ddd8ddc369246943df7356d57d Mon Sep 17 00:00:00 2001
From: Daniel Cazzulino <daniel@cazzulino.com>
Date: Tue, 14 Oct 2025 14:57:02 -0300
Subject: [PATCH] Add service-driven ChatMessageStoreFactory for agents

In order to extend the configured agents, in addition to providing the configureOptions delegate when adding agents, the user can now also rely on auto-wiring of services exported as ChatMessageStoreFactory to customize message persistence.

The export can be keyed to a specific agent (takes priority) or without a key (fallback). This allows granular wiring without resorting to more code.
---
 src/Agents/ChatMessageStoreFactory.cs | 24 +++++++++++
 src/Agents/ConfigurableAIAgent.cs     |  9 +++++
 src/Tests/ConfigurableAgentTests.cs   | 57 ++++++++++++++++++++++++++-
 3 files changed, 89 insertions(+), 1 deletion(-)
 create mode 100644 src/Agents/ChatMessageStoreFactory.cs

diff --git a/src/Agents/ChatMessageStoreFactory.cs b/src/Agents/ChatMessageStoreFactory.cs
new file mode 100644
index 0000000..0bf6cb8
--- /dev/null
+++ b/src/Agents/ChatMessageStoreFactory.cs
@@ -0,0 +1,24 @@
+using Microsoft.Agents.AI;
+using static Microsoft.Agents.AI.ChatClientAgentOptions;
+
+namespace Devlooped.Agents.AI;
+
+/// <summary>
+/// An implementation of a <see cref="ChatMessageStore"/> factory as a class that can provide 
+/// the functionality to <see cref="ChatClientAgentOptions.ChatMessageStoreFactory"/> and integrates 
+/// more easily into a service collection.
+/// </summary>
+/// <remarks>
+/// The <see cref="ChatMessageStore"/> is a key extensibility point in Microsoft.Agents.AI, allowing 
+/// storage and retrieval of chat messages.
+/// </remarks>
+public abstract class ChatMessageStoreFactory
+{
+    /// <summary>
+    /// Provides the implementation of <see cref="ChatClientAgentOptions.ChatMessageStoreFactory"/> 
+    /// to provide message persistence.
+    /// </summary>
+    /// <param name="context">The context to potentially hydrate state from.</param>
+    /// <returns>The message store that will handle chat messages.</returns>
+    public abstract ChatMessageStore CreateStore(ChatMessageStoreFactoryContext context);
+}
diff --git a/src/Agents/ConfigurableAIAgent.cs b/src/Agents/ConfigurableAIAgent.cs
index eaa9f34..b0fe1c5 100644
--- a/src/Agents/ConfigurableAIAgent.cs
+++ b/src/Agents/ConfigurableAIAgent.cs
@@ -98,6 +98,15 @@ public override IAsyncEnumerable<AgentRunResponseUpdate> RunStreamingAsync(IEnum
                 options.AIContextProviderFactory = contextFactory.CreateProvider;
         }
 
+        if (options.ChatMessageStoreFactory is null)
+        {
+            var storeFactory = services.GetKeyedService<ChatMessageStoreFactory>(name) ??
+                services.GetService<ChatMessageStoreFactory>();
+
+            if (storeFactory is not null)
+                options.ChatMessageStoreFactory = storeFactory.CreateStore;
+        }
+
         LogConfigured(name);
 
         return (new ChatClientAgent(client, options, services.GetRequiredService<ILoggerFactory>(), services), options, client);
diff --git a/src/Tests/ConfigurableAgentTests.cs b/src/Tests/ConfigurableAgentTests.cs
index 66784c8..392a259 100644
--- a/src/Tests/ConfigurableAgentTests.cs
+++ b/src/Tests/ConfigurableAgentTests.cs
@@ -131,8 +131,63 @@ public void AssignsContextProviderFromService()
         var options = agent.GetService<ChatClientAgentOptions>();
 
         Assert.NotNull(options?.AIContextProviderFactory);
-        Assert.Same(context, options?.AIContextProviderFactory?.Invoke(new ChatClientAgentOptions.AIContextProviderFactoryContext()));
+        Assert.Same(context, options?.AIContextProviderFactory?.Invoke(new()));
+    }
+
+    [Fact]
+    public void AssignsMessageStoreFactoryFromKeyedService()
+    {
+        var builder = new HostApplicationBuilder();
+        var context = Mock.Of<ChatMessageStore>();
+
+        builder.Services.AddKeyedSingleton<ChatMessageStoreFactory>("bot",
+            Mock.Of<ChatMessageStoreFactory>(x
+                => x.CreateStore(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>()) == context));
+
+        builder.Configuration.AddInMemoryCollection(new Dictionary<string, string?>
+        {
+            ["ai:clients:chat:modelid"] = "gpt-4.1-nano",
+            ["ai:clients:chat:apikey"] = "sk-asdfasdf",
+            ["ai:agents:bot:client"] = "chat",
+            ["ai:agents:bot:options:temperature"] = "0.5",
+        });
+
+        builder.AddAIAgents();
+
+        var app = builder.Build();
+        var agent = app.Services.GetRequiredKeyedService<AIAgent>("bot");
+        var options = agent.GetService<ChatClientAgentOptions>();
+
+        Assert.NotNull(options?.ChatMessageStoreFactory);
+        Assert.Same(context, options?.ChatMessageStoreFactory?.Invoke(new()));
     }
 
+    [Fact]
+    public void AssignsMessageStoreFactoryFromService()
+    {
+        var builder = new HostApplicationBuilder();
+        var context = Mock.Of<ChatMessageStore>();
+
+        builder.Services.AddSingleton<ChatMessageStoreFactory>(
+            Mock.Of<ChatMessageStoreFactory>(x
+                => x.CreateStore(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>()) == context));
+
+        builder.Configuration.AddInMemoryCollection(new Dictionary<string, string?>
+        {
+            ["ai:clients:chat:modelid"] = "gpt-4.1-nano",
+            ["ai:clients:chat:apikey"] = "sk-asdfasdf",
+            ["ai:agents:bot:client"] = "chat",
+            ["ai:agents:bot:options:temperature"] = "0.5",
+        });
+
+        builder.AddAIAgents();
+
+        var app = builder.Build();
+        var agent = app.Services.GetRequiredKeyedService<AIAgent>("bot");
+        var options = agent.GetService<ChatClientAgentOptions>();
+
+        Assert.NotNull(options?.ChatMessageStoreFactory);
+        Assert.Same(context, options?.ChatMessageStoreFactory?.Invoke(new()));
+    }
 }