Skip to content

.NET: Add per run / thread feature collection support and improved custom ChatMessageStore support #2354

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 5 commits into
base: main
Choose a base branch
from
Draft

.NET: Add per run / thread feature collection support and improved custom ChatMessageStore support #2354

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
eff5aee
.NET: Add per run / thread feature collection support and improved cu…
westey-m Nov 20, 2025
570bed9
Merge branch 'main' into feature-featurecollections-messagestore
westey-m Nov 21, 2025
9d86adf
.NET: Update AgentFeatureCollections with feedback (#2379)
westey-m Nov 24, 2025
9382526
Merge branch 'main' into feature-featurecollections-messagestore
westey-m Nov 24, 2025
9498c84
Merge branch 'main' into feature-featurecollections-messagestore
westey-m Nov 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions dotnet/samples/GettingStarted/AgentProviders/Agent_With_CustomImplementation/Program.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,10 +28,10 @@ internal sealed class UpperCaseParrotAgent : AIAgent
{
public override string? Name => "UpperCaseParrotAgent";

public override AgentThread GetNewThread()
public override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null)
=> new CustomAgentThread();

public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null)
=> new CustomAgentThread(serializedThread, jsonSerializerOptions);

public override async Task<AgentRunResponse> RunAsync(IEnumerable<ChatMessage> messages, AgentThread? thread = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default)
Expand Down
244 changes: 199 additions & 45 deletions dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyThreadStorage/Program.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,83 +21,237 @@
// Replace this with a vector store implementation of your choice if you want to persist the chat history to disk.
VectorStore vectorStore = new InMemoryVectorStore();

// Create the agent
AIAgent agent = new AzureOpenAIClient(
new Uri(endpoint),
new AzureCliCredential())
.GetChatClient(deploymentName)
.CreateAIAgent(new ChatClientAgentOptions
{
Instructions = "You are good at telling jokes.",
Name = "Joker",
ChatMessageStoreFactory = ctx =>
// Execute various samples showing how to use a custom ChatMessageStore with an agent.
await CustomChatMessageStore_UsingFactory_Async();
await CustomChatMessageStore_UsingFactoryAndExistingExternalId_Async();
await CustomChatMessageStore_PerThread_Async();
await CustomChatMessageStore_PerRun_Async();

// Here we can see how to create a custom ChatMessageStore using a factory method
// provided to the agent via the ChatMessageStoreFactory option.
// This allows us to use a custom chat message store, where the consumer of the agent
// doesn't need to know anything about the storage mechanism used.
async Task CustomChatMessageStore_UsingFactory_Async()
{
Console.WriteLine("\n--- With Factory ---\n");

// Create the agent
AIAgent agent = new AzureOpenAIClient(
new Uri(endpoint),
new AzureCliCredential())
// Use a service that doesn't require storage of chat history in the service itself.
.GetChatClient(deploymentName)
.CreateAIAgent(new ChatClientAgentOptions
{
// Create a new chat message store for this agent that stores the messages in a vector store.
// Each thread must get its own copy of the VectorChatMessageStore, since the store
// also contains the id that the thread is stored under.
return new VectorChatMessageStore(vectorStore, ctx.SerializedState, ctx.JsonSerializerOptions);
}
});
Instructions = "You are good at telling jokes.",
Name = "Joker",
ChatMessageStoreFactory = ctx =>
{
// Create a new chat message store for this agent that stores the messages in a vector store.
// Each thread must get its own copy of the VectorChatMessageStore, since the store
// also contains the id that the thread is stored under.
return new VectorChatMessageStore(vectorStore, ctx.SerializedState, ctx.JsonSerializerOptions, ctx.Features);
}
});

// Start a new thread for the agent conversation.
AgentThread thread = agent.GetNewThread();

// Run the agent with the thread that stores conversation history in the vector store.
Console.WriteLine(await agent.RunAsync("Tell me a joke about a pirate.", thread));

// Serialize the thread state, so it can be stored for later use.
// Since the chat history is stored in the vector store, the serialized thread
// only contains the guid that the messages are stored under in the vector store.
JsonElement serializedThread = thread.Serialize();

Console.WriteLine("\n--- Serialized thread ---\n");
Console.WriteLine(JsonSerializer.Serialize(serializedThread, new JsonSerializerOptions { WriteIndented = true }));

// The serialized thread can now be saved to a database, file, or any other storage mechanism
// and loaded again later.

// Deserialize the thread state after loading from storage.
AgentThread resumedThread = agent.DeserializeThread(serializedThread);

// Run the agent with the thread that stores conversation history in the vector store a second time.
Console.WriteLine(await agent.RunAsync("Now tell the same joke in the voice of a pirate, and add some emojis to the joke.", resumedThread));
}

// Here we can see how to create a custom ChatMessageStore using a factory method
// provided to the agent via the ChatMessageStoreFactory option.
// It also shows how we can pass a custom storage id at runtime to the message store using
// the VectorChatMessageStoreThreadDbKeyFeature.
// Note that not all agents or chat message stores may support this feature.
async Task CustomChatMessageStore_UsingFactoryAndExistingExternalId_Async()
{
Console.WriteLine("\n--- With Factory and Existing External ID ---\n");

// Create the agent
AIAgent agent = new AzureOpenAIClient(
new Uri(endpoint),
new AzureCliCredential())
// Use a service that doesn't require storage of chat history in the service itself.
.GetChatClient(deploymentName)
.CreateAIAgent(new ChatClientAgentOptions
{
Instructions = "You are good at telling jokes.",
Name = "Joker",
ChatMessageStoreFactory = ctx =>
{
// Create a new chat message store for this agent that stores the messages in a vector store.
// Each thread must get its own copy of the VectorChatMessageStore, since the store
// also contains the id that the thread is stored under.
return new VectorChatMessageStore(vectorStore, ctx.SerializedState, ctx.JsonSerializerOptions, ctx.Features);
}
});

// Start a new thread for the agent conversation.
AgentThread thread = agent.GetNewThread();

// Run the agent with the thread that stores conversation history in the vector store.
Console.WriteLine(await agent.RunAsync("Tell me a joke about a pirate.", thread));

// We can access the VectorChatMessageStore via the thread's GetService method if we need to read the key under which threads are stored.
var messageStoreFromFactory = thread.GetService<VectorChatMessageStore>()!;
Console.WriteLine($"\nThread is stored in vector store under key: {messageStoreFromFactory.ThreadDbKey}");

// It's possible to create a new thread that uses the same chat message store id by providing
// the VectorChatMessageStoreThreadDbKeyFeature in the feature collection when creating the new thread.
AgentThread resumedThread = agent.GetNewThread(
new AgentFeatureCollection().WithFeature(new VectorChatMessageStoreThreadDbKeyFeature(messageStoreFromFactory.ThreadDbKey!)));

// Start a new thread for the agent conversation.
AgentThread thread = agent.GetNewThread();
// Run the agent with the thread that stores conversation history in the vector store.
Console.WriteLine(await agent.RunAsync("Now tell the same joke in the voice of a pirate, and add some emojis to the joke.", resumedThread));
}

// Run the agent with the thread that stores conversation history in the vector store.
Console.WriteLine(await agent.RunAsync("Tell me a joke about a pirate.", thread));
// Here we can see how to create a custom ChatMessageStore and pass it to the thread
// when creating a new thread.
async Task CustomChatMessageStore_PerThread_Async()
{
Console.WriteLine("\n--- Per Thread ---\n");

// Serialize the thread state, so it can be stored for later use.
// Since the chat history is stored in the vector store, the serialized thread
// only contains the guid that the messages are stored under in the vector store.
JsonElement serializedThread = thread.Serialize();
// We can also create an agent without a factory that provides a ChatMessageStore.
AIAgent agent = new AzureOpenAIClient(
new Uri(endpoint),
new AzureCliCredential())
// Use a service that doesn't require storage of chat history in the service itself.
.GetChatClient(deploymentName)
.CreateAIAgent(new ChatClientAgentOptions
{
Instructions = "You are good at telling jokes.",
Name = "Joker"
});

Console.WriteLine("\n--- Serialized thread ---\n");
Console.WriteLine(JsonSerializer.Serialize(serializedThread, new JsonSerializerOptions { WriteIndented = true }));
// Instead of using a factory on the agent to create the ChatMessageStore, we can
// create a VectorChatMessageStore ourselves and register it in a feature collection.
// We can then pass the feature collection when creating a new thread.
// We also have the opportunity here to pass any id that we want for storing the chat history in the vector store.
VectorChatMessageStore perThreadMessageStore = new(vectorStore, "chat-history-1");
AgentThread thread = agent.GetNewThread(new AgentFeatureCollection().WithFeature<ChatMessageStore>(perThreadMessageStore));

// The serialized thread can now be saved to a database, file, or any other storage mechanism
// and loaded again later.
Console.WriteLine(await agent.RunAsync("Tell me a joke about a pirate.", thread));

// Deserialize the thread state after loading from storage.
AgentThread resumedThread = agent.DeserializeThread(serializedThread);
// When serializing this thread, we'll see that it has the id from the message store stored in its state.
JsonElement serializedThread = thread.Serialize();

// Run the agent with the thread that stores conversation history in the vector store a second time.
Console.WriteLine(await agent.RunAsync("Now tell the same joke in the voice of a pirate, and add some emojis to the joke.", resumedThread));
Console.WriteLine("\n--- Serialized thread ---\n");
Console.WriteLine(JsonSerializer.Serialize(serializedThread, new JsonSerializerOptions { WriteIndented = true }));
}

// We can access the VectorChatMessageStore via the thread's GetService method if we need to read the key under which threads are stored.
var messageStore = resumedThread.GetService<VectorChatMessageStore>()!;
Console.WriteLine($"\nThread is stored in vector store under key: {messageStore.ThreadDbKey}");
// Here we can see how to create a custom ChatMessageStore for a single run using the Features option
// passed when we run the agent.
// Note that if the agent doesn't support a chat message store, it would be ignored.
async Task CustomChatMessageStore_PerRun_Async()
{
Console.WriteLine("\n--- Per Run ---\n");

// We can also create an agent without a factory that provides a ChatMessageStore.
AIAgent agent = new AzureOpenAIClient(
new Uri(endpoint),
new AzureCliCredential())
// Use a service that doesn't require storage of chat history in the service itself.
.GetChatClient(deploymentName)
.CreateAIAgent(new ChatClientAgentOptions
{
Instructions = "You are good at telling jokes.",
Name = "Joker"
});

// Start a new thread for the agent conversation.
AgentThread thread = agent.GetNewThread();

// Instead of using a factory on the agent to create the ChatMessageStore, we can
// create a VectorChatMessageStore ourselves and register it in a feature collection.
// We can then pass the feature collection to the agent when running it by using the Features option.
// The message store would only be used for the run that it's passed to.
// If the agent doesn't support a message store, it would be ignored.
// We also have the opportunity here to pass any id that we want for storing the chat history in the vector store.
VectorChatMessageStore perRunMessageStore = new(vectorStore, "chat-history-1");
Console.WriteLine(await agent.RunAsync(
"Tell me a joke about a pirate.",
thread,
options: new AgentRunOptions()
{
Features = new AgentFeatureCollection().WithFeature<ChatMessageStore>(perRunMessageStore)
}));

// When serializing this thread, we'll see that it has no messagestore state, since the messagestore was not attached to the thread,
// but just provided for the single run. Note that, depending on the circumstances, the thread may still contain other state, e.g. Memories,
// if an AIContextProvider is attached which adds memory to an agent.
JsonElement serializedThread = thread.Serialize();

Console.WriteLine("\n--- Serialized thread ---\n");
Console.WriteLine(JsonSerializer.Serialize(serializedThread, new JsonSerializerOptions { WriteIndented = true }));
}

namespace SampleApp
{
/// <summary>
/// A feature that allows providing the thread database key for the <see cref="VectorChatMessageStore"/>.
/// </summary>
internal sealed class VectorChatMessageStoreThreadDbKeyFeature(string threadDbKey)
{
public string ThreadDbKey { get; } = threadDbKey;
}

/// <summary>
/// A sample implementation of <see cref="ChatMessageStore"/> that stores chat messages in a vector store.
/// </summary>
internal sealed class VectorChatMessageStore : ChatMessageStore
{
private readonly VectorStore _vectorStore;

public VectorChatMessageStore(VectorStore vectorStore, JsonElement serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null)
public VectorChatMessageStore(VectorStore vectorStore, string threadDbKey)
{
this._vectorStore = vectorStore ?? throw new ArgumentNullException(nameof(vectorStore));
this.ThreadDbKey = threadDbKey ?? throw new ArgumentNullException(nameof(threadDbKey));
}

public VectorChatMessageStore(VectorStore vectorStore, JsonElement serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? features = null)
{
this._vectorStore = vectorStore ?? throw new ArgumentNullException(nameof(vectorStore));

if (serializedStoreState.ValueKind is JsonValueKind.String)
{
// Here we can deserialize the thread id so that we can access the same messages as before the suspension.
this.ThreadDbKey = serializedStoreState.Deserialize<string>();
}
// Here we can deserialize the thread id so that we can access the same messages as before the suspension, or if
// a user provided a ConversationIdAgentFeature in the features collection, we can use that
// or finally we can generate one ourselves.
this.ThreadDbKey = serializedStoreState.ValueKind is JsonValueKind.String
? serializedStoreState.Deserialize<string>()
: features?.TryGet<VectorChatMessageStoreThreadDbKeyFeature>(out var threadDbKeyFeature) is true
? threadDbKeyFeature.ThreadDbKey
: Guid.NewGuid().ToString("N");
}

public string? ThreadDbKey { get; private set; }
public string? ThreadDbKey { get; }

public override async Task AddMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken = default)
{
this.ThreadDbKey ??= Guid.NewGuid().ToString("N");

var collection = this._vectorStore.GetCollection<string, ChatHistoryItem>("ChatHistory");
await collection.EnsureCollectionExistsAsync(cancellationToken);

await collection.UpsertAsync(messages.Select(x => new ChatHistoryItem()
{
Key = this.ThreadDbKey + x.MessageId,
Key = this.ThreadDbKey + (string.IsNullOrWhiteSpace(x.MessageId) ? Guid.NewGuid().ToString("N") : x.MessageId),
Timestamp = DateTimeOffset.UtcNow,
ThreadId = this.ThreadDbKey,
SerializedMessage = JsonSerializer.Serialize(x),
Expand Down
11 changes: 8 additions & 3 deletions dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,8 +55,13 @@ public A2AAgent(A2AClient a2aClient, string? id = null, string? name = null, str
}

/// <inheritdoc/>
public sealed override AgentThread GetNewThread()
=> new A2AAgentThread();
public sealed override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null)
=> new A2AAgentThread()
{
ContextId = featureCollection?.TryGet<ConversationIdAgentFeature>(out var conversationIdFeature) is true
? conversationIdFeature.ConversationId
: null
};

/// <summary>
/// Get a new <see cref="AgentThread"/> instance using an existing context id, to continue that conversation.
Expand All @@ -67,7 +72,7 @@ public AgentThread GetNewThread(string contextId)
=> new A2AAgentThread() { ContextId = contextId };

/// <inheritdoc/>
public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null)
=> new A2AAgentThread(serializedThread, jsonSerializerOptions);

/// <inheritdoc/>
Expand Down
6 changes: 4 additions & 2 deletions dotnet/src/Microsoft.Agents.AI.Abstractions/AIAgent.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,7 @@ public abstract class AIAgent
/// <summary>
/// Creates a new conversation thread that is compatible with this agent.
/// </summary>
/// <param name="featureCollection">An optional feature collection to override or provide additional context or capabilities to the thread where the thread supports these features.</param>
/// <returns>A new <see cref="AgentThread"/> instance ready for use with this agent.</returns>
/// <remarks>
/// <para>
Expand All @@ -121,13 +122,14 @@ public abstract class AIAgent
/// may be deferred until first use to optimize performance.
/// </para>
/// </remarks>
public abstract AgentThread GetNewThread();
public abstract AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null);

/// <summary>
/// Deserializes an agent thread from its JSON serialized representation.
/// </summary>
/// <param name="serializedThread">A <see cref="JsonElement"/> containing the serialized thread state.</param>
/// <param name="jsonSerializerOptions">Optional settings to customize the deserialization process.</param>
/// <param name="featureCollection">An optional feature collection to override or provide additional context or capabilities to the thread where the thread supports these features.</param>
/// <returns>A restored <see cref="AgentThread"/> instance with the state from <paramref name="serializedThread"/>.</returns>
/// <exception cref="ArgumentException">The <paramref name="serializedThread"/> is not in the expected format.</exception>
/// <exception cref="JsonException">The serialized data is invalid or cannot be deserialized.</exception>
Expand All @@ -136,7 +138,7 @@ public abstract class AIAgent
/// allowing conversations to resume across application restarts or be migrated between
/// different agent instances.
/// </remarks>
public abstract AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null);
public abstract AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null);

/// <summary>
/// Run the agent with no message assuming that all required instructions are already provided to the agent or on the thread.
Expand Down
Loading
Loading