From f273ca7353fdfce3efa470108c105e5a1a523c73 Mon Sep 17 00:00:00 2001
From: westey <164392973+westey-m@users.noreply.github.com>
Date: Thu, 13 Nov 2025 15:43:01 +0000
Subject: [PATCH] Fix InMemoryChatMessageStore serialization bug. (#2185)

---
 .../InMemoryChatMessageStore.cs               |  6 ++--
 .../InMemoryChatMessageStoreTests.cs          | 30 +++++++++++++++++++
 .../TestJsonSerializerContext.cs              |  1 +
 3 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatMessageStore.cs
index 17d1cdce93..79d303207c 100644
--- a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatMessageStore.cs
+++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatMessageStore.cs
@@ -97,8 +97,9 @@ public sealed class InMemoryChatMessageStore : ChatMessageStore, IList<ChatMessa
 
         if (serializedStoreState.ValueKind is JsonValueKind.Object)
         {
+            var jso = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions;
             var state = serializedStoreState.Deserialize(
-                AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(StoreState))) as StoreState;
+                jso.GetTypeInfo(typeof(StoreState))) as StoreState;
             if (state?.Messages is { } messages)
             {
                 this._messages = messages;
@@ -164,7 +165,8 @@ public sealed class InMemoryChatMessageStore : ChatMessageStore, IList<ChatMessa
             Messages = this._messages,
         };
 
-        return JsonSerializer.SerializeToElement(state, AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(StoreState)));
+        var jso = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions;
+        return JsonSerializer.SerializeToElement(state, jso.GetTypeInfo(typeof(StoreState)));
     }
 
     /// <inheritdoc />
diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatMessageStoreTests.cs
index 4c793d17f4..824fb62f6d 100644
--- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatMessageStoreTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatMessageStoreTests.cs
@@ -3,7 +3,9 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Text.Encodings.Web;
 using System.Text.Json;
+using System.Text.Json.Serialization.Metadata;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
@@ -114,6 +116,29 @@ public class InMemoryChatMessageStoreTests
         Assert.Equal("B", newStore[1].Text);
     }
 
+    [Fact]
+    public async Task SerializeAndDeserializeConstructorRoundtripsWithCustomAIContentAsync()
+    {
+        JsonSerializerOptions options = new(TestJsonSerializerContext.Default.Options)
+        {
+            TypeInfoResolver = JsonTypeInfoResolver.Combine(AgentAbstractionsJsonUtilities.DefaultOptions.TypeInfoResolver, TestJsonSerializerContext.Default),
+            Encoder = JavaScriptEncoder.UnsafeRelaxedJsonEscaping,
+        };
+        options.AddAIContentType<TestAIContent>(typeDiscriminatorId: "testContent");
+
+        var store = new InMemoryChatMessageStore
+        {
+            new ChatMessage(ChatRole.User, [new TestAIContent("foo data")]),
+        };
+
+        var jsonElement = store.Serialize(options);
+        var newStore = new InMemoryChatMessageStore(jsonElement, options);
+
+        Assert.Single(newStore);
+        var actualTestAIContent = Assert.IsType<TestAIContent>(newStore[0].Contents[0]);
+        Assert.Equal("foo data", actualTestAIContent.TestData);
+    }
+
     [Fact]
     public async Task SerializeAndDeserializeWorksWithExperimentalContentTypesAsync()
     {
@@ -558,4 +583,9 @@ public class InMemoryChatMessageStoreTests
         Assert.Equal("Hello", result[0].Text);
         reducerMock.Verify(r => r.ReduceAsync(It.IsAny<IEnumerable<ChatMessage>>(), It.IsAny<CancellationToken>()), Times.Never);
     }
+
+    public class TestAIContent(string testData) : AIContent
+    {
+        public string TestData => testData;
+    }
 }
diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/TestJsonSerializerContext.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/TestJsonSerializerContext.cs
index b7c553d348..ec343504ab 100644
--- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/TestJsonSerializerContext.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/TestJsonSerializerContext.cs
@@ -22,4 +22,5 @@ namespace Microsoft.Agents.AI.Abstractions.UnitTests;
 [JsonSerializable(typeof(InMemoryAgentThread.InMemoryAgentThreadState))]
 [JsonSerializable(typeof(ServiceIdAgentThread.ServiceIdAgentThreadState))]
 [JsonSerializable(typeof(ServiceIdAgentThreadTests.EmptyObject))]
+[JsonSerializable(typeof(InMemoryChatMessageStoreTests.TestAIContent))]
 internal sealed partial class TestJsonSerializerContext : JsonSerializerContext;