Help needed: custom ChatMessageStore not available until after RunStreamingAsync

[hexbit2]

Help needed: custom ChatMessageStore not available until after RunStreamingAsync

Runtime & packages

• net9.0
• Microsoft.Agents.AI.OpenAI 1.0.0-preview.251110.2
• Azure.AI.OpenAI 2.5.0-beta.1

How we configure the agent

We create agents through an AgentFactory that wires up a custom PostgreSQL-backed chat history store:

// doccore/DocCore.Services/Services/AgentFactory.cs
var azureClient = _clientProvider.GetClient(modelKey);

var chatClient = azureClient
    .GetChatClient(modelConfig.DeploymentName)
    .AsIChatClient();

var agentName = name ?? profile.Key;
var toolArray = tools?
    .Cast<AITool>()
    .ToArray();

var options = toolArray is null
    ? new ChatClientAgentOptions()
    : new ChatClientAgentOptions(profile.Instructions, agentName, description: null, tools: toolArray);

options.Instructions ??= profile.Instructions;
options.Name ??= agentName;
options.ChatMessageStoreFactory = ctx => new PostgresChatMessageStore(
    _conversationService,
    _tenantContext,
    ctx.SerializedState,
    ctx.JsonSerializerOptions ?? new JsonSerializerOptions());

return chatClient.CreateAIAgent(options);

AzureOpenAIClientProvider caches AzureOpenAIClient instances for each configured deployment so every request reuses the same underlying AzureOpenAIClient.

Our store (simplified):

// doccore/DocCore.Services/Agents/PostgresChatMessageStore.cs
public sealed class PostgresChatMessageStore : ChatMessageStore
{
    private readonly IConversationService _conversationService;
    private readonly ITenantContext _tenantContext;
    private Guid? _conversationId;

    public PostgresChatMessageStore(
        IConversationService conversationService,
        ITenantContext tenantContext,
        JsonElement serializedState,
        JsonSerializerOptions? jsonSerializerOptions = null)
    {
        _conversationService = conversationService;
        _tenantContext = tenantContext;

        if (serializedState.ValueKind != JsonValueKind.Undefined &&
            serializedState.ValueKind != JsonValueKind.Null)
        {
            var state = serializedState.Deserialize<StoreState>(jsonSerializerOptions ?? JsonSerializerOptions.Default);
            if (state != null && Guid.TryParse(state.ConversationId, out var parsedId))
            {
                _conversationId = parsedId;
            }
        }
    }

    public void AttachConversation(Guid conversationId) => _conversationId = conversationId;

    public override async Task AddMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken ct)
    {
        var tenantId = EnsureTenantId();
        var conversationId = await EnsureConversationAsync(ct);
        // … map ChatMessage -> CreateConversationMessageDto and persist via ConversationService …
    }

    public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) =>
        JsonSerializer.SerializeToElement(new StoreState { ConversationId = _conversationId?.ToString() },
            jsonSerializerOptions ?? JsonSerializerOptions.Default);
}

The store implements Serialize/Deserialize and exposes an AttachConversation(Guid conversationId) helper so we can link the agent thread to an existing conversation row in our database.

How a request is processed

In the SignalR hub we:

1. Pre-create a conversation in PostgreSQL and obtain its ID.
2. Call var agent = _agentFactory.GetAgent(...);
3. Call var thread = agent.GetNewThread();
4. Goal: fetch the PostgresChatMessageStore immediately and call store.AttachConversation(conversationId) before invoking the LLM, so that the very first user + assistant turns append to that pre-created conversation.
5. Call await foreach (var chunk in agent.RunStreamingAsync(message, thread: thread)) { ... }

The problem

• Step 3 returns a ChatClientAgentThread, but thread.MessageStore is null.
• Because the store is null, we cannot attach the existing conversation ID. When the user sends the first message, our store creates a new conversation instead of appending to the one we prepared.
• As soon as RunStreamingAsync finishes, thread.MessageStore now contains our PostgresChatMessageStore (factory clearly ran) and message persistence works from that point forward—but it’s too late to associate the first turn with the pre-created conversation.

What we’ve already checked

• Verified the ChatMessageStoreFactory lambda executes: PostgresChatMessageStore.AddMessagesAsync is called after the run and the serialized state is captured correctly.
• Tried manually new-ing up a store inside the hub when MessageStore is null. That works but defeats the purpose of the factory and raises questions about how the SDK expects lifecycle/state to be managed.
• Looked for overloads such as GetNewThread(serializedState) or guidance on setting ChatClientAgentThread.ConversationId; could not find any.
• Observed that when the framework restores from serialized state, the store is populated as expected—the gap only happens on the very first run.

Questions for the Agent Framework team

1. Is it expected that ChatClientAgentThread.MessageStore is null until after the first call to RunAsync/RunStreamingAsync? If so, how should we attach our existing conversation ID before that first call?
2. Is there a recommended API/pattern for associating an existing custom store instance (or serialized state) with a newly created thread?
3. Should we be setting ChatClientAgentThread.ConversationId ourselves to prime the thread, or is there another hook we should be using?

Goal

We need to link a new agent thread to an already existing conversation record before sending the very first user message, so that no duplicate conversation rows are created in our history tables. Any guidance on the intended lifecycle would be greatly appreciated.*** End Patch

[hexbit2]

Quick update on the ChatMessageStore issue:

• I implemented the workaround suggested in discussion #1908.

• Before the first request, I deserialize an empty thread so the framework instantiates the store immediately:

  var blankState = JsonSerializer.Deserialize<JsonElement>("{}");
  var thread = agent.DeserializeThread(blankState);
  if (thread is not ChatClientAgentThread)
  {
      thread = agent.GetNewThread();
  }
  
  if (thread is ChatClientAgentThread chatThread &&
      chatThread.MessageStore is PostgresChatMessageStore store)
  {
      store.AttachConversation(conversationGuid);
  }

• With that shim in place, thread.MessageStore is populated before RunStreamingAsync, letting us attach the pre-existing conversation ID and avoid duplicate rows.

Could you confirm whether this behavior (store being null until after the first run) is expected, or if we should treat this as a bug? An official way to prime the thread/store without the workaround would be ideal.

[westey-m]

Hi @hexbit2, thanks for raising this issue. I do think we have a gap for your scenario where you have a pre-existing id that you want to use, which is created using some external id generator.

For now, the hack you did should work, but I think we can provide a better solution.

I've created a work item to track this and provide a solution: #2325
Please take a look and feel free to comment and add thoughts.

[hexbit2]

Thank you !