.NET: [BREAKING] Add GetNewThread overload for taking a ChatMessageStore at thread creation time (#2429)

* Add GetNewThread overload for taking a ChatMessageStore at thread creation time and move message store creation to thread creation time.

* Fix formatting.

* Apply suggestions from code review

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

* Apply suggestion from @rogerbarreto

Co-authored-by: Roger Barreto <19890735+rogerbarreto@users.noreply.github.com>

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Roger Barreto <19890735+rogerbarreto@users.noreply.github.com>
This commit is contained in:
westey
2025-11-25 10:49:01 +00:00
committed by GitHub
Unverified
parent a89c15d6e6
commit 112410ecd6
4 changed files with 282 additions and 47 deletions
@@ -289,6 +289,7 @@ public sealed partial 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 })
};
@@ -316,6 +317,34 @@ public sealed partial class ChatClientAgent : AIAgent
AIContextProvider = this._agentOptions?.AIContextProviderFactory?.Invoke(new() { SerializedState = default, JsonSerializerOptions = null })
};
/// <summary>
/// Creates a new agent thread instance using an existing <see cref="ChatMessageStore"/> to continue a conversation.
/// </summary>
/// <param name="chatMessageStore">The <see cref="ChatMessageStore"/> instance to use for managing the conversation's message history.</param>
/// <returns>
/// A new <see cref="AgentThread"/> instance configured to work with the provided <paramref name="chatMessageStore"/>.
/// </returns>
/// <remarks>
/// <para>
/// This method creates threads that do not support server-side conversation storage.
/// Some AI services require server-side conversation storage to function properly, and creating a thread
/// with a <see cref="ChatMessageStore"/> may not be compatible with these services.
/// </para>
/// <para>
/// Where a service requires server-side conversation storage, use <see cref="GetNewThread(string)"/>.
/// </para>
/// <para>
/// If the agent detects, during the first run, that the underlying AI service requires server-side conversation storage,
/// the thread will throw an exception to indicate that it cannot continue using the provided <see cref="ChatMessageStore"/>.
/// </para>
/// </remarks>
public AgentThread GetNewThread(ChatMessageStore chatMessageStore)
=> new ChatClientAgentThread()
{
MessageStore = Throw.IfNull(chatMessageStore),
AIContextProvider = this._agentOptions?.AIContextProviderFactory?.Invoke(new() { SerializedState = default, JsonSerializerOptions = null })
};
/// <inheritdoc/>
public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
{
@@ -459,10 +459,74 @@ public partial class ChatClientAgentTests
}
/// <summary>
/// Verify that RunAsync doesn't use the ChatMessageStore factory when the chat client returns a conversation id.
/// Verify that RunAsync uses the default InMemoryChatMessageStore when the chat client returns no conversation id.
/// </summary>
[Fact]
public async Task RunAsyncIgnoresChatMessageStoreWhenConversationIdReturnedByChatClientAsync()
public async Task RunAsyncUsesDefaultInMemoryChatMessageStoreWhenNoConversationIdReturnedByChatClientAsync()
{
// 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")]));
ChatClientAgent agent = new(mockService.Object, options: new()
{
Instructions = "test instructions",
});
// Act
ChatClientAgentThread? thread = agent.GetNewThread() as ChatClientAgentThread;
await agent.RunAsync([new(ChatRole.User, "test")], thread);
// Assert
var messageStore = Assert.IsType<InMemoryChatMessageStore>(thread!.MessageStore);
Assert.Equal(2, messageStore.Count);
Assert.Equal("test", messageStore[0].Text);
Assert.Equal("response", messageStore[1].Text);
}
/// <summary>
/// Verify that RunAsync uses the ChatMessageStore factory when the chat client returns no conversation id.
/// </summary>
[Fact]
public async Task RunAsyncUsesChatMessageStoreFactoryWhenProvidedAndNoConversationIdReturnedByChatClientAsync()
{
// 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<ChatMessageStore> mockChatMessageStore = new();
Mock<Func<ChatClientAgentOptions.ChatMessageStoreFactoryContext, ChatMessageStore>> mockFactory = new();
mockFactory.Setup(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>())).Returns(mockChatMessageStore.Object);
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<ChatMessageStore>(thread!.MessageStore, exactMatch: false);
mockChatMessageStore.Verify(s => s.AddMessagesAsync(It.Is<IEnumerable<ChatMessage>>(x => x.Count() == 2), It.IsAny<CancellationToken>()), Times.Once);
mockFactory.Verify(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>()), Times.Once);
}
/// <summary>
/// Verify that RunAsync throws when a ChatMessageStore Factory is provided and the chat client returns a conversation id.
/// </summary>
[Fact]
public async Task RunAsyncThrowsWhenChatMessageStoreFactoryProvidedAndConversationIdReturnedByChatClientAsync()
{
// Arrange
Mock<IChatClient> mockService = new();
@@ -479,13 +543,10 @@ public partial class ChatClientAgentTests
ChatMessageStoreFactory = mockFactory.Object
});
// Act
// Act & Assert
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);
var exception = await Assert.ThrowsAsync<InvalidOperationException>(() => agent.RunAsync([new(ChatRole.User, "test")], thread));
Assert.Equal("Only the ConversationId or MessageStore may be set, but not both and switching from one to another is not supported.", exception.Message);
}
/// <summary>
@@ -1914,10 +1975,10 @@ public partial class ChatClientAgentTests
}
/// <summary>
/// Verify that RunStreamingAsync doesn't use the ChatMessageStore factory when the chat client returns a conversation id.
/// Verify that RunStreamingAsync throws when a ChatMessageStore factory is provided and the chat client returns a conversation id.
/// </summary>
[Fact]
public async Task RunStreamingAsyncIgnoresChatMessageStoreWhenConversationIdReturnedByChatClientAsync()
public async Task RunStreamingAsyncThrowsWhenChatMessageStoreFactoryProvidedAndConversationIdReturnedByChatClientAsync()
{
// Arrange
Mock<IChatClient> mockService = new();
@@ -1939,13 +2000,10 @@ public partial class ChatClientAgentTests
ChatMessageStoreFactory = mockFactory.Object
});
// Act
// Act & Assert
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);
var exception = await Assert.ThrowsAsync<InvalidOperationException>(async () => await agent.RunStreamingAsync([new(ChatRole.User, "test")], thread).ToListAsync());
Assert.Equal("Only the ConversationId or MessageStore may be set, but not both and switching from one to another is not supported.", exception.Message);
}
/// <summary>
@@ -2074,37 +2132,6 @@ public partial class ChatClientAgentTests
#endregion
#region GetNewThread Tests
[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.IsType<ChatClientAgentThread>(thread);
var typedThread = (ChatClientAgentThread)thread;
Assert.Same(mockContextProvider.Object, typedThread.AIContextProvider);
}
#endregion
#region Background Responses Tests
[Theory]
@@ -0,0 +1,79 @@
// Copyright (c) Microsoft. All rights reserved.
using System.Text.Json;
using Microsoft.Extensions.AI;
using Moq;
namespace Microsoft.Agents.AI.UnitTests.ChatClient;
/// <summary>
/// Contains unit tests for the ChatClientAgent.DeserializeThread methods.
/// </summary>
public class ChatClientAgent_DeserializeThreadTests
{
[Fact]
public void DeserializeThread_UsesAIContextProviderFactory_IfProvided()
{
// 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;
}
});
var json = JsonSerializer.Deserialize("""
{
"aiContextProviderState": ["CP1"]
}
""", TestJsonSerializerContext.Default.JsonElement);
// Act
var thread = agent.DeserializeThread(json);
// Assert
Assert.True(factoryCalled, "AIContextProviderFactory was not called.");
Assert.IsType<ChatClientAgentThread>(thread);
var typedThread = (ChatClientAgentThread)thread;
Assert.Same(mockContextProvider.Object, typedThread.AIContextProvider);
}
[Fact]
public void DeserializeThread_UsesChatMessageStoreFactory_IfProvided()
{
// Arrange
var mockChatClient = new Mock<IChatClient>();
var mockMessageStore = new Mock<ChatMessageStore>();
var factoryCalled = false;
var agent = new ChatClientAgent(mockChatClient.Object, new ChatClientAgentOptions
{
Instructions = "Test instructions",
ChatMessageStoreFactory = _ =>
{
factoryCalled = true;
return mockMessageStore.Object;
}
});
var json = JsonSerializer.Deserialize("""
{
"storeState": { }
}
""", TestJsonSerializerContext.Default.JsonElement);
// Act
var thread = agent.DeserializeThread(json);
// Assert
Assert.True(factoryCalled, "ChatMessageStoreFactory was not called.");
Assert.IsType<ChatClientAgentThread>(thread);
var typedThread = (ChatClientAgentThread)thread;
Assert.Same(mockMessageStore.Object, typedThread.MessageStore);
}
}
@@ -0,0 +1,100 @@
// Copyright (c) Microsoft. All rights reserved.
using Microsoft.Extensions.AI;
using Moq;
namespace Microsoft.Agents.AI.UnitTests.ChatClient;
/// <summary>
/// Contains unit tests for the ChatClientAgent.GetNewThread methods.
/// </summary>
public class ChatClientAgent_GetNewThreadTests
{
[Fact]
public void GetNewThread_UsesAIContextProviderFactory_IfProvided()
{
// 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.IsType<ChatClientAgentThread>(thread);
var typedThread = (ChatClientAgentThread)thread;
Assert.Same(mockContextProvider.Object, typedThread.AIContextProvider);
}
[Fact]
public void GetNewThread_UsesChatMessageStoreFactory_IfProvided()
{
// Arrange
var mockChatClient = new Mock<IChatClient>();
var mockMessageStore = new Mock<ChatMessageStore>();
var factoryCalled = false;
var agent = new ChatClientAgent(mockChatClient.Object, new ChatClientAgentOptions
{
Instructions = "Test instructions",
ChatMessageStoreFactory = _ =>
{
factoryCalled = true;
return mockMessageStore.Object;
}
});
// Act
var thread = agent.GetNewThread();
// Assert
Assert.True(factoryCalled, "ChatMessageStoreFactory was not called.");
Assert.IsType<ChatClientAgentThread>(thread);
var typedThread = (ChatClientAgentThread)thread;
Assert.Same(mockMessageStore.Object, typedThread.MessageStore);
}
[Fact]
public void GetNewThread_UsesChatMessageStore_FromTypedOverload()
{
// Arrange
var mockChatClient = new Mock<IChatClient>();
var mockMessageStore = new Mock<ChatMessageStore>();
var agent = new ChatClientAgent(mockChatClient.Object);
// Act
var thread = agent.GetNewThread(mockMessageStore.Object);
// Assert
Assert.IsType<ChatClientAgentThread>(thread);
var typedThread = (ChatClientAgentThread)thread;
Assert.Same(mockMessageStore.Object, typedThread.MessageStore);
}
[Fact]
public void GetNewThread_UsesConversationId_FromTypedOverload()
{
// Arrange
var mockChatClient = new Mock<IChatClient>();
const string TestConversationId = "test_conversation_id";
var agent = new ChatClientAgent(mockChatClient.Object);
// Act
var thread = agent.GetNewThread(TestConversationId);
// Assert
Assert.IsType<ChatClientAgentThread>(thread);
var typedThread = (ChatClientAgentThread)thread;
Assert.Equal(TestConversationId, typedThread.ConversationId);
}
}