fix: InMemoryCheckpointManager is not JSON serializable (#1639)

Checkpointing is used by the WorkflowHostAgent to be able to support resume from a provided thread. When a CheckpointManager is not specified, we use the InMemoryCheckpointManager and serialize its state into the thread's Serialize()ed JsonElement.

At some point InMemoryCheckpointManager became not serializable, breaking this behaviour. This change restores serializability, and adds a test.

dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/CheckpointInfoConverter.cs

@@ -0,0 +1,49 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System.Diagnostics.CodeAnalysis;
+using System.Text.Json;
+using System.Text.Json.Serialization.Metadata;
+using System.Text.RegularExpressions;
+
+namespace Microsoft.Agents.AI.Workflows.Checkpointing;
+
+/// <summary>
+/// Provides support for using <see cref="CheckpointInfo"/> values as dictionary keys when serializing and deserializing JSON.
+/// </summary>
+internal sealed partial class CheckpointInfoConverter() : JsonConverterDictionarySupportBase<CheckpointInfo>
+{
+    protected override JsonTypeInfo<CheckpointInfo> TypeInfo
+        => WorkflowsJsonUtilities.JsonContext.Default.CheckpointInfo;
+
+    private const string CheckpointInfoPropertyNamePattern = @"^(?<runId>(((\|\|)|([^\|]))*))\|(?<checkpointId>(((\|\|)|([^\|]))*)?)$";
+#if NET
+    [GeneratedRegex(CheckpointInfoPropertyNamePattern, RegexOptions.CultureInvariant | RegexOptions.ExplicitCapture)]
+    public static partial Regex CheckpointInfoPropertyNameRegex();
+#else
+    public static Regex CheckpointInfoPropertyNameRegex() => s_scopeKeyPropertyNameRegex;
+    private static readonly Regex s_scopeKeyPropertyNameRegex =
+        new(CheckpointInfoPropertyNamePattern, RegexOptions.Compiled | RegexOptions.CultureInvariant | RegexOptions.ExplicitCapture);
+#endif
+
+    protected override CheckpointInfo Parse(string propertyName)
+    {
+        Match scopeKeyPatternMatch = CheckpointInfoPropertyNameRegex().Match(propertyName);
+        if (!scopeKeyPatternMatch.Success)
+        {
+            throw new JsonException($"Invalid CheckpointInfo property name format. Got '{propertyName}'.");
+        }
+
+        string runId = scopeKeyPatternMatch.Groups["runId"].Value;
+        string checkpointId = scopeKeyPatternMatch.Groups["checkpointId"].Value;
+
+        return new(Unescape(runId)!, Unescape(checkpointId)!);
+    }
+
+    protected override string Stringify([DisallowNull] CheckpointInfo value)
+    {
+        string? runIdEscaped = Escape(value.RunId);
+        string? checkpointIdEscaped = Escape(value.CheckpointId);
+
+        return $"{runIdEscaped}|{checkpointIdEscaped}";
+    }
+}

dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/InMemoryCheckpointManager.cs

@@ -12,21 +12,22 @@ namespace Microsoft.Agents.AI.Workflows.Checkpointing;
 /// </summary>
 internal sealed class InMemoryCheckpointManager : ICheckpointManager
 {
-    private readonly Dictionary<string, RunCheckpointCache<Checkpoint>> _store = [];
+    [JsonInclude]
+    internal Dictionary<string, RunCheckpointCache<Checkpoint>> Store { get; } = [];
 
     public InMemoryCheckpointManager() { }
 
     [JsonConstructor]
     internal InMemoryCheckpointManager(Dictionary<string, RunCheckpointCache<Checkpoint>> store)
     {
-        this._store = store;
+        this.Store = store;
     }
 
     private RunCheckpointCache<Checkpoint> GetRunStore(string runId)
     {
-        if (!this._store.TryGetValue(runId, out RunCheckpointCache<Checkpoint>? runStore))
+        if (!this.Store.TryGetValue(runId, out RunCheckpointCache<Checkpoint>? runStore))
         {
-            runStore = this._store[runId] = new();
+            runStore = this.Store[runId] = new();
         }
 
         return runStore;

dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/JsonConverterDictionarySupportBase.cs

@@ -2,6 +2,7 @@
 
 using System;
 using System.Diagnostics.CodeAnalysis;
+using System.Runtime.CompilerServices;
 using System.Text.Json;
 using System.Text.Json.Serialization.Metadata;
 
@@ -18,6 +19,57 @@ internal abstract class JsonConverterDictionarySupportBase<T> : JsonConverterBas
     protected abstract string Stringify([DisallowNull] T value);
     protected abstract T Parse(string propertyName);
 
+    [return: NotNull]
+    protected static string Escape(string? value, char escapeChar = '|', bool allowNullAndPad = false, [CallerArgumentExpression(nameof(value))] string? componentName = null)
+    {
+        if (!allowNullAndPad && value is null)
+        {
+            throw new JsonException($"Invalid {componentName} '{value}'. Expecting non-null string.");
+        }
+
+        if (value is null)
+        {
+            return string.Empty;
+        }
+
+        string unescaped = escapeChar.ToString();
+        string escaped = new(escapeChar, 2);
+
+        if (allowNullAndPad)
+        {
+            return $"@{value.Replace(unescaped, escaped)}";
+        }
+
+        return $"{value.Replace(unescaped, escaped)}";
+    }
+
+    protected static string? Unescape([DisallowNull] string value, char escapeChar = '|', bool allowNullAndPad = false, [CallerArgumentExpression(nameof(value))] string? componentName = null)
+    {
+        if (value.Length == 0)
+        {
+            if (!allowNullAndPad)
+            {
+                throw new JsonException($"Invalid {componentName} '{value}'. Expecting empty string or a value that is prefixed with '@'.");
+            }
+
+            return null;
+        }
+
+        if (allowNullAndPad && value[0] != '@')
+        {
+            throw new JsonException($"Invalid {componentName} component '{value}'. Expecting empty string or a value that is prefixed with '@'.");
+        }
+
+        if (allowNullAndPad)
+        {
+            value = value.Substring(1);
+        }
+
+        string unescaped = escapeChar.ToString();
+        string escaped = new(escapeChar, 2);
+        return value.Replace(escaped, unescaped);
+    }
+
     public override T ReadAsPropertyName(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options)
     {
         SequencePosition position = reader.Position;

dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/JsonMarshaller.cs

@@ -18,6 +18,7 @@ internal sealed class JsonMarshaller : IWireMarshaller<JsonElement>
         this._internalOptions.Converters.Add(new ExecutorIdentityConverter());
         this._internalOptions.Converters.Add(new ScopeKeyConverter());
         this._internalOptions.Converters.Add(new EdgeIdConverter());
+        this._internalOptions.Converters.Add(new CheckpointInfoConverter());
 
         this._externalOptions = serializerOptions;
     }

dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/RunCheckpointCache.cs

@@ -8,23 +8,26 @@ namespace Microsoft.Agents.AI.Workflows.Checkpointing;
 
 internal sealed class RunCheckpointCache<TStoreObject>
 {
-    private readonly List<CheckpointInfo> _checkpointIndex = [];
-    private readonly Dictionary<CheckpointInfo, TStoreObject> _cache = [];
+    [JsonInclude]
+    internal List<CheckpointInfo> CheckpointIndex { get; } = [];
+
+    [JsonInclude]
+    internal Dictionary<CheckpointInfo, TStoreObject> Cache { get; } = [];
 
     public RunCheckpointCache() { }
 
     [JsonConstructor]
     internal RunCheckpointCache(List<CheckpointInfo> checkpointIndex, Dictionary<CheckpointInfo, TStoreObject> cache)
     {
-        this._checkpointIndex = checkpointIndex;
-        this._cache = cache;
+        this.CheckpointIndex = checkpointIndex;
+        this.Cache = cache;
     }
 
     [JsonIgnore]
-    public IEnumerable<CheckpointInfo> Index => this._checkpointIndex;
+    public IEnumerable<CheckpointInfo> Index => this.CheckpointIndex;
 
-    public bool IsInIndex(CheckpointInfo key) => this._cache.ContainsKey(key);
-    public bool TryGet(CheckpointInfo key, [MaybeNullWhen(false)] out TStoreObject value) => this._cache.TryGetValue(key, out value);
+    public bool IsInIndex(CheckpointInfo key) => this.Cache.ContainsKey(key);
+    public bool TryGet(CheckpointInfo key, [MaybeNullWhen(false)] out TStoreObject value) => this.Cache.TryGetValue(key, out value);
 
     public CheckpointInfo Add(string runId, TStoreObject value)
     {
@@ -45,18 +48,18 @@ internal sealed class RunCheckpointCache<TStoreObject>
             return false;
         }
 
-        this._cache[key] = value;
-        this._checkpointIndex.Add(key);
+        this.Cache[key] = value;
+        this.CheckpointIndex.Add(key);
         return true;
     }
 
     [JsonIgnore]
-    public bool HasCheckpoints => this._checkpointIndex.Count > 0;
+    public bool HasCheckpoints => this.CheckpointIndex.Count > 0;
     public bool TryGetLastCheckpointInfo([NotNullWhen(true)] out CheckpointInfo? checkpointInfo)
     {
         if (this.HasCheckpoints)
         {
-            checkpointInfo = this._checkpointIndex[this._checkpointIndex.Count - 1];
+            checkpointInfo = this.CheckpointIndex[this.CheckpointIndex.Count - 1];
             return true;
         }
         checkpointInfo = default;

dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/ScopeKeyConverter.cs

@@ -1,7 +1,6 @@
 // Copyright (c) Microsoft. All rights reserved.
 
 using System.Diagnostics.CodeAnalysis;
-using System.Runtime.CompilerServices;
 using System.Text.Json;
 using System.Text.Json.Serialization.Metadata;
 using System.Text.RegularExpressions;
@@ -42,52 +41,6 @@ internal sealed partial class ScopeKeyConverter : JsonConverterDictionarySupport
                             Unescape(key)!);
     }
 
-    [return: NotNull]
-    private static string Escape(string? value, bool allowNullAndPad = false, [CallerArgumentExpression(nameof(value))] string componentName = "ScopeKey")
-    {
-        if (!allowNullAndPad && value is null)
-        {
-            throw new JsonException($"Invalid {componentName} '{value}'. Expecting non-null string.");
-        }
-
-        if (value is null)
-        {
-            return string.Empty;
-        }
-
-        if (allowNullAndPad)
-        {
-            return $"@{value.Replace("|", "||")}";
-        }
-
-        return $"{value.Replace("|", "||")}";
-    }
-
-    private static string? Unescape([DisallowNull] string value, bool allowNullAndPad = false, [CallerArgumentExpression(nameof(value))] string componentName = "ScopeKey")
-    {
-        if (value.Length == 0)
-        {
-            if (!allowNullAndPad)
-            {
-                throw new JsonException($"Invalid {componentName} '{value}'. Expecting empty string or a value that is prefixed with '@'.");
-            }
-
-            return null;
-        }
-
-        if (allowNullAndPad && value[0] != '@')
-        {
-            throw new JsonException($"Invalid {componentName} component '{value}'. Expecting empty string or a value that is prefixed with '@'.");
-        }
-
-        if (allowNullAndPad)
-        {
-            value = value.Substring(1);
-        }
-
-        return value.Replace("||", "|");
-    }
-
     protected override string Stringify([DisallowNull] ScopeKey value)
     {
         string? executorIdEscaped = Escape(value.ScopeId.ExecutorId);

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/JsonSerializationTests.cs

@@ -634,13 +634,8 @@ public class JsonSerializationTests
 
     private static CheckpointInfo TestParentCheckpointInfo => new(s_runId, s_parentCheckpointId);
 
-    [Fact]
-    public async Task Test_Checkpoint_JsonRoundTripAsync()
+    private static void ValidateCheckpoint(Checkpoint result, Checkpoint prototype)
     {
-        WorkflowInfo testWorkflowInfo = CreateTestWorkflowInfo();
-        Checkpoint prototype = new(12, testWorkflowInfo, TestRunnerStateData, TestStateData, TestEdgeState, TestParentCheckpointInfo);
-        Checkpoint result = RunJsonRoundtrip(prototype, TestCustomSerializedJsonOptions);
-
         result.Should().Match((Checkpoint checkpoint) => checkpoint.StepNumber == prototype.StepNumber);
 
         result.Parent.Should().Be(prototype.Parent);
@@ -650,4 +645,31 @@ public class JsonSerializationTests
         ValidateStateData(result.StateData, prototype.StateData);
         ValidateEdgeStateData(result.EdgeStateData, prototype.EdgeStateData);
     }
+
+    [Fact]
+    public async Task Test_Checkpoint_JsonRoundTripAsync()
+    {
+        WorkflowInfo testWorkflowInfo = CreateTestWorkflowInfo();
+        Checkpoint prototype = new(12, testWorkflowInfo, TestRunnerStateData, TestStateData, TestEdgeState, TestParentCheckpointInfo);
+        Checkpoint result = RunJsonRoundtrip(prototype, TestCustomSerializedJsonOptions);
+
+        ValidateCheckpoint(result, prototype);
+    }
+
+    [Fact]
+    public async Task Test_InMemoryCheckpointManager_JsonRoundTripAsync()
+    {
+        WorkflowInfo testWorkflowInfo = CreateTestWorkflowInfo();
+        Checkpoint prototype = new(12, testWorkflowInfo, TestRunnerStateData, TestStateData, TestEdgeState, TestParentCheckpointInfo);
+        string runId = Guid.NewGuid().ToString("N");
+
+        InMemoryCheckpointManager manager = new();
+        CheckpointInfo checkpointInfo = await manager.CommitCheckpointAsync(runId, prototype);
+
+        InMemoryCheckpointManager result = RunJsonRoundtrip(manager, TestCustomSerializedJsonOptions);
+
+        Checkpoint? retrievedCheckpoint = await result.LookupCheckpointAsync(runId, checkpointInfo);
+
+        ValidateCheckpoint(retrievedCheckpoint, prototype);
+    }
 }