.NET: [BREAKING] Add session statebag to use for state storage instead of inside providers (#3737)

* Add a StateBag to AgentSession and pass Agent and AgentSession to AIContextProvider and ChatHistoryProviders

* Convert all AIContextProviders to use the statebag

* Update InMemoryChatHistoryProvider to use StateBag

* Update Comsos and Workflow ChatHistoryProviders

* Update 3rd party chat history storage sample.

* Remove serialize method from providers

* Replacing provider factories with properties

* Remove Providers from Session and flatten state bag serialization

* Update samples to use getservice on agent

* Updated additional session types to serialize statebag

* Fix regression

* Address PR comments

* Address PR comments.

* Fix formatting

* Fix unit tests

* Remove InMemoryAgentSession since it is not required anymore.

* Address PR comments

* Convert sessions for A2AAgent, ChatClientAgent, CopilotStudioAgent and GithubCopilotAgent to use regular json serialization.

* Fix durable agent session jso usgae

* Add jso to InMemory and Workflow ChatHistoryProviders

* Update InMemoryChatHistoryProvider to use an options class for it's many optional settings.

* Apply suggestions from code review

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

* Address PR feedback

* Fix verification bug.

* Improve state bag thread safety

* Address PR comments and fix unit tests

* Address PR comments

* Fix unit test

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
westey
2026-02-10 12:03:37 +00:00
committed by GitHub
co-authored by Copilot
parent 977c3adfb2
commit b12ff578af
93 changed files with 2990 additions and 3178 deletions
@@ -34,16 +34,21 @@ AIAgent agent = new AzureOpenAIClient(
{
ChatOptions = new() { Instructions = "You are good at telling jokes." },
Name = "Joker",
AIContextProviderFactory = (ctx, ct) => new ValueTask<AIContextProvider>(new ChatHistoryMemoryProvider(
AIContextProvider = new ChatHistoryMemoryProvider(
vectorStore,
collectionName: "chathistory",
vectorDimensions: 3072,
// Configure the scope values under which chat messages will be stored.
// In this case, we are using a fixed user ID and a unique session ID for each new session.
storageScope: new() { UserId = "UID1", SessionId = Guid.NewGuid().ToString() },
// Configure the scope which would be used to search for relevant prior messages.
// In this case, we are searching for any messages for the user across all sessions.
searchScope: new() { UserId = "UID1" }))
// Callback to configure the initial state of the ChatHistoryMemoryProvider.
// The ChatHistoryMemoryProvider stores its state in the AgentSession and this callback
// will be called whenever the ChatHistoryMemoryProvider cannot find existing state in the session,
// typically the first time it is used with a new session.
session => new ChatHistoryMemoryProvider.State(
// Configure the scope values under which chat messages will be stored.
// In this case, we are using a fixed user ID and a unique session ID for each new session.
storageScope: new() { UserId = "UID1", SessionId = Guid.NewGuid().ToString() },
// Configure the scope which would be used to search for relevant prior messages.
// In this case, we are searching for any messages for the user across all sessions.
searchScope: new() { UserId = "UID1" }))
});
// Start a new session for the agent conversation.
@@ -31,20 +31,21 @@ AIAgent agent = new AzureOpenAIClient(
.AsAIAgent(new ChatClientAgentOptions()
{
ChatOptions = new() { Instructions = "You are a friendly travel assistant. Use known memories about the user when responding, and do not invent details." },
AIContextProviderFactory = (ctx, ct) => new ValueTask<AIContextProvider>(ctx.SerializedState.ValueKind is not JsonValueKind.Null and not JsonValueKind.Undefined
// If each session should have its own Mem0 scope, you can create a new id per session here:
// ? new Mem0Provider(mem0HttpClient, new Mem0ProviderScope() { ThreadId = Guid.NewGuid().ToString() })
// In this case we are storing memories scoped by application and user instead so that memories are retained across threads.
? new Mem0Provider(mem0HttpClient, new Mem0ProviderScope() { ApplicationId = "getting-started-agents", UserId = "sample-user" })
// For cases where we are restoring from serialized state:
: new Mem0Provider(mem0HttpClient, ctx.SerializedState, ctx.JsonSerializerOptions))
// The stateInitializer can be used to customize the Mem0 scope per session and it will be called each time a session
// is encountered by the Mem0Provider that does not already have Mem0Provider state stored on the session.
// If each session should have its own Mem0 scope, you can create a new id per session via the stateInitializer, e.g.:
// new Mem0Provider(mem0HttpClient, stateInitializer: _ => new(new Mem0ProviderScope() { ThreadId = Guid.NewGuid().ToString() }))
// In our case we are storing memories scoped by application and user instead so that memories are retained across threads.
AIContextProvider = new Mem0Provider(mem0HttpClient, stateInitializer: _ => new(new Mem0ProviderScope() { ApplicationId = "getting-started-agents", UserId = "sample-user" }))
});
AgentSession session = await agent.CreateSessionAsync();
// Clear any existing memories for this scope to demonstrate fresh behavior.
Mem0Provider mem0Provider = session.GetService<Mem0Provider>()!;
await mem0Provider.ClearStoredMemoriesAsync();
// Note that the ClearStoredMemoriesAsync method will clear memories
// using the scope stored in the session, or provided via the stateInitializer.
Mem0Provider mem0Provider = agent.GetService<Mem0Provider>()!;
await mem0Provider.ClearStoredMemoriesAsync(session);
Console.WriteLine(await agent.RunAsync("Hi there! My name is Taylor and I'm planning a hiking trip to Patagonia in November.", session));
Console.WriteLine(await agent.RunAsync("I'm travelling with my sister and we love finding scenic viewpoints.", session));
@@ -33,7 +33,7 @@ ChatClient chatClient = new AzureOpenAIClient(
AIAgent agent = chatClient.AsAIAgent(new ChatClientAgentOptions()
{
ChatOptions = new() { Instructions = "You are a friendly assistant. Always address the user by their name." },
AIContextProviderFactory = (ctx, ct) => new ValueTask<AIContextProvider>(new UserInfoMemory(chatClient.AsIChatClient(), ctx.SerializedState, ctx.JsonSerializerOptions))
AIContextProvider = new UserInfoMemory(chatClient.AsIChatClient())
});
// Create a new session for the conversation.
@@ -55,10 +55,10 @@ Console.WriteLine("\n>> Use deserialized session with previously created memorie
var deserializedSession = await agent.DeserializeSessionAsync(sesionElement);
Console.WriteLine(await agent.RunAsync("What is my name and age?", deserializedSession));
Console.WriteLine("\n>> Read memories from memory component\n");
Console.WriteLine("\n>> Read memories using memory component\n");
// It's possible to access the memory component via the session's GetService method.
var userInfo = deserializedSession.GetService<UserInfoMemory>()?.UserInfo;
// It's possible to access the memory component via the agent's GetService method.
var userInfo = agent.GetService<UserInfoMemory>()?.GetUserInfo(deserializedSession);
// Output the user info that was captured by the memory component.
Console.WriteLine($"MEMORY - User Name: {userInfo?.UserName}");
@@ -66,12 +66,12 @@ Console.WriteLine($"MEMORY - User Age: {userInfo?.UserAge}");
Console.WriteLine("\n>> Use new session with previously created memories\n");
// It is also possible to set the memories in a memory component on an individual session.
// It is also possible to set the memories using a memory component on an individual session.
// This is useful if we want to start a new session, but have it share the same memories as a previous session.
var newSession = await agent.CreateSessionAsync();
if (userInfo is not null && newSession.GetService<UserInfoMemory>() is UserInfoMemory newSessionMemory)
if (userInfo is not null && agent.GetService<UserInfoMemory>() is UserInfoMemory newSessionMemory)
{
newSessionMemory.UserInfo = userInfo;
newSessionMemory.SetUserInfo(newSession, userInfo);
}
// Invoke the agent and output the text result.
@@ -86,28 +86,27 @@ namespace SampleApp
internal sealed class UserInfoMemory : AIContextProvider
{
private readonly IChatClient _chatClient;
private readonly Func<AgentSession?, UserInfo> _stateInitializer;
public UserInfoMemory(IChatClient chatClient, UserInfo? userInfo = null)
public UserInfoMemory(IChatClient chatClient, Func<AgentSession?, UserInfo>? stateInitializer = null)
{
this._chatClient = chatClient;
this.UserInfo = userInfo ?? new UserInfo();
this._stateInitializer = stateInitializer ?? (_ => new UserInfo());
}
public UserInfoMemory(IChatClient chatClient, JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null)
{
this._chatClient = chatClient;
public UserInfo GetUserInfo(AgentSession session)
=> session.StateBag.GetValue<UserInfo>(nameof(UserInfoMemory)) ?? new UserInfo();
this.UserInfo = serializedState.ValueKind == JsonValueKind.Object ?
serializedState.Deserialize<UserInfo>(jsonSerializerOptions)! :
new UserInfo();
}
public UserInfo UserInfo { get; set; }
public void SetUserInfo(AgentSession session, UserInfo userInfo)
=> session.StateBag.SetValue(nameof(UserInfoMemory), userInfo);
protected override async ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
{
var userInfo = context.Session?.StateBag.GetValue<UserInfo>(nameof(UserInfoMemory))
?? this._stateInitializer.Invoke(context.Session);
// Try and extract the user name and age from the message if we don't have it already and it's a user message.
if ((this.UserInfo.UserName is null || this.UserInfo.UserAge is null) && context.RequestMessages.Any(x => x.Role == ChatRole.User))
if ((userInfo.UserName is null || userInfo.UserAge is null) && context.RequestMessages.Any(x => x.Role == ChatRole.User))
{
var result = await this._chatClient.GetResponseAsync<UserInfo>(
context.RequestMessages,
@@ -117,36 +116,36 @@ namespace SampleApp
},
cancellationToken: cancellationToken);
this.UserInfo.UserName ??= result.Result.UserName;
this.UserInfo.UserAge ??= result.Result.UserAge;
userInfo.UserName ??= result.Result.UserName;
userInfo.UserAge ??= result.Result.UserAge;
}
context.Session?.StateBag.SetValue(nameof(UserInfoMemory), userInfo);
}
protected override ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
{
var userInfo = context.Session?.StateBag.GetValue<UserInfo>(nameof(UserInfoMemory))
?? this._stateInitializer.Invoke(context.Session);
StringBuilder instructions = new();
// If we don't already know the user's name and age, add instructions to ask for them, otherwise just provide what we have to the context.
instructions
.AppendLine(
this.UserInfo.UserName is null ?
userInfo.UserName is null ?
"Ask the user for their name and politely decline to answer any questions until they provide it." :
$"The user's name is {this.UserInfo.UserName}.")
$"The user's name is {userInfo.UserName}.")
.AppendLine(
this.UserInfo.UserAge is null ?
userInfo.UserAge is null ?
"Ask the user for their age and politely decline to answer any questions until they provide it." :
$"The user's age is {this.UserInfo.UserAge}.");
$"The user's age is {userInfo.UserAge}.");
return new ValueTask<AIContext>(new AIContext
{
Instructions = instructions.ToString()
});
}
public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
{
return JsonSerializer.SerializeToElement(this.UserInfo, jsonSerializerOptions);
}
}
internal sealed class UserInfo