.NET: Fix bug where ChatClientAgent throws when providing a ChatMessageStore with a service that requries service storage (#1229)

* Fix bug where ChatClientAgent throws when providing a ChatMessageStore with a service that requries service storage

* Update dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
westey
2025-10-07 09:31:34 +01:00
committed by GitHub
Unverified
parent 7506046f43
commit 74905af778
2 changed files with 124 additions and 22 deletions
@@ -322,7 +322,6 @@ public sealed class ChatClientAgent : AIAgent
public override AgentThread GetNewThread()
=> new ChatClientAgentThread
{
MessageStore = this._agentOptions?.ChatMessageStoreFactory?.Invoke(new() { SerializedState = default, JsonSerializerOptions = null }),
AIContextProvider = this._agentOptions?.AIContextProviderFactory?.Invoke(new() { SerializedState = default, JsonSerializerOptions = null })
};
@@ -2,6 +2,7 @@
using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.AI;
@@ -420,6 +421,66 @@ public class ChatClientAgentTests
Assert.Equal("ConvId", thread.ConversationId);
}
/// <summary>
/// Verify that RunAsync uses the ChatMessageStore factory when the chat client returns no conversation id.
/// </summary>
[Fact]
public async Task RunAsyncUsesChatMessageStoreWhenNoConversationIdReturnedByChatClientAsync()
{
// 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")]));
Mock<Func<ChatClientAgentOptions.ChatMessageStoreFactoryContext, ChatMessageStore>> mockFactory = new();
mockFactory.Setup(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>())).Returns(new InMemoryChatMessageStore());
ChatClientAgent agent = new(mockService.Object, options: new()
{
Instructions = "test instructions",
ChatMessageStoreFactory = mockFactory.Object
});
// Act
ChatClientAgentThread? thread = agent.GetNewThread() as ChatClientAgentThread;
await agent.RunAsync([new(ChatRole.User, "test")], thread);
// Assert
Assert.IsType<InMemoryChatMessageStore>(thread!.MessageStore);
mockFactory.Verify(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>()), Times.Once);
}
/// <summary>
/// Verify that RunAsync doesn't use the ChatMessageStore factory when the chat client returns a conversation id.
/// </summary>
[Fact]
public async Task RunAsyncIgnoresChatMessageStoreWhenConversationIdReturnedByChatClientAsync()
{
// 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" });
Mock<Func<ChatClientAgentOptions.ChatMessageStoreFactoryContext, ChatMessageStore>> mockFactory = new();
mockFactory.Setup(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>())).Returns(new InMemoryChatMessageStore());
ChatClientAgent agent = new(mockService.Object, options: new()
{
Instructions = "test instructions",
ChatMessageStoreFactory = mockFactory.Object
});
// Act
ChatClientAgentThread? thread = agent.GetNewThread() as ChatClientAgentThread;
await agent.RunAsync([new(ChatRole.User, "test")], thread);
// Assert
Assert.Equal("ConvId", thread!.ConversationId);
mockFactory.Verify(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>()), Times.Never);
}
/// <summary>
/// Verify that RunAsync invokes any provided AIContextProvider and uses the result.
/// </summary>
@@ -1721,38 +1782,80 @@ public class ChatClientAgentTests
Times.Once);
}
#endregion
#region GetNewThread Tests
/// <summary>
/// Verify that RunStreamingAsync uses the ChatMessageStore factory when the chat client returns no conversation id.
/// </summary>
[Fact]
public void GetNewThreadUsesChatMessageStoreFactoryIfProvided()
public async Task RunStreamingAsyncUsesChatMessageStoreWhenNoConversationIdReturnedByChatClientAsync()
{
// Arrange
var mockChatClient = new Mock<IChatClient>();
var mockStore = new Mock<ChatMessageStore>();
var factoryCalled = false;
var agent = new ChatClientAgent(mockChatClient.Object, new ChatClientAgentOptions
Mock<IChatClient> mockService = new();
ChatResponseUpdate[] returnUpdates =
[
new ChatResponseUpdate(role: ChatRole.Assistant, content: "wh"),
new ChatResponseUpdate(role: ChatRole.Assistant, content: "at?"),
];
mockService.Setup(
s => s.GetStreamingResponseAsync(
It.IsAny<IEnumerable<ChatMessage>>(),
It.IsAny<ChatOptions>(),
It.IsAny<CancellationToken>())).Returns(ToAsyncEnumerableAsync(returnUpdates));
Mock<Func<ChatClientAgentOptions.ChatMessageStoreFactoryContext, ChatMessageStore>> mockFactory = new();
mockFactory.Setup(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>())).Returns(new InMemoryChatMessageStore());
ChatClientAgent agent = new(mockService.Object, options: new()
{
Instructions = "Test instructions",
ChatMessageStoreFactory = _ =>
{
factoryCalled = true;
return mockStore.Object;
}
Instructions = "test instructions",
ChatMessageStoreFactory = mockFactory.Object
});
// Act
var thread = agent.GetNewThread();
ChatClientAgentThread? thread = agent.GetNewThread() as ChatClientAgentThread;
await agent.RunStreamingAsync([new(ChatRole.User, "test")], thread).ToListAsync();
// Assert
Assert.True(factoryCalled, "ChatMessageStoreFactory was not called.");
Assert.IsType<ChatClientAgentThread>(thread);
var typedThread = (ChatClientAgentThread)thread;
Assert.Same(mockStore.Object, typedThread.MessageStore);
Assert.IsType<InMemoryChatMessageStore>(thread!.MessageStore);
mockFactory.Verify(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>()), Times.Once);
}
/// <summary>
/// Verify that RunStreamingAsync doesn't use the ChatMessageStore factory when the chat client returns a conversation id.
/// </summary>
[Fact]
public async Task RunStreamingAsyncIgnoresChatMessageStoreWhenConversationIdReturnedByChatClientAsync()
{
// Arrange
Mock<IChatClient> mockService = new();
ChatResponseUpdate[] returnUpdates =
[
new ChatResponseUpdate(role: ChatRole.Assistant, content: "wh") { ConversationId = "ConvId" },
new ChatResponseUpdate(role: ChatRole.Assistant, content: "at?") { ConversationId = "ConvId" },
];
mockService.Setup(
s => s.GetStreamingResponseAsync(
It.IsAny<IEnumerable<ChatMessage>>(),
It.IsAny<ChatOptions>(),
It.IsAny<CancellationToken>())).Returns(ToAsyncEnumerableAsync(returnUpdates));
Mock<Func<ChatClientAgentOptions.ChatMessageStoreFactoryContext, ChatMessageStore>> mockFactory = new();
mockFactory.Setup(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>())).Returns(new InMemoryChatMessageStore());
ChatClientAgent agent = new(mockService.Object, options: new()
{
Instructions = "test instructions",
ChatMessageStoreFactory = mockFactory.Object
});
// Act
ChatClientAgentThread? thread = agent.GetNewThread() as ChatClientAgentThread;
await agent.RunStreamingAsync([new(ChatRole.User, "test")], thread).ToListAsync();
// Assert
Assert.Equal("ConvId", thread!.ConversationId);
mockFactory.Verify(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>()), Times.Never);
}
#endregion
#region GetNewThread Tests
[Fact]
public void GetNewThreadUsesAIContextProviderFactoryIfProvided()
{