.NET: Fix InMemoryCheckpointManager serialization

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 @@ public JsonMarshaller(JsonSerializerOptions? serializerOptions = null)
         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 @@ public bool Add(CheckpointInfo key, TStoreObject value)
             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 @@ protected override ScopeKey Parse(string propertyName)
                             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 void Test_ExecutorStateData_JsonRoundTrip()
 
     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 async Task Test_Checkpoint_JsonRoundTripAsync()
         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);
+    }
 }