Fix InMemoryChatMessageStore serialization bug. (microsoft#2185)

Author: westey-m

dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatMessageStore.cs

@@ -97,8 +97,9 @@ public InMemoryChatMessageStore(IChatReducer? chatReducer, JsonElement serialize
 
         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 override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptio
             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 />

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 async Task SerializeAndDeserializeConstructorRoundtripsAsync()
         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 async Task GetMessagesAsync_WithReducer_ButWrongTrigger_DoesNotInvokeRedu
         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;
+    }
 }

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;