From b2246efa69a60528f5afd273fe671cdf5c9a6eff Mon Sep 17 00:00:00 2001 From: Jacob Alber Date: Thu, 23 Oct 2025 15:32:52 -0400 Subject: [PATCH] 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. --- .../Checkpointing/CheckpointInfoConverter.cs | 49 +++++++++++++++++ .../InMemoryCheckpointManager.cs | 9 ++-- .../JsonConverterDictionarySupportBase.cs | 52 +++++++++++++++++++ .../Checkpointing/JsonMarshaller.cs | 1 + .../Checkpointing/RunCheckpointCache.cs | 25 +++++---- .../Checkpointing/ScopeKeyConverter.cs | 47 ----------------- .../JsonSerializationTests.cs | 34 +++++++++--- 7 files changed, 149 insertions(+), 68 deletions(-) create mode 100644 dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/CheckpointInfoConverter.cs diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/CheckpointInfoConverter.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/CheckpointInfoConverter.cs new file mode 100644 index 0000000000..53c277d822 --- /dev/null +++ b/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; + +/// +/// Provides support for using values as dictionary keys when serializing and deserializing JSON. +/// +internal sealed partial class CheckpointInfoConverter() : JsonConverterDictionarySupportBase +{ + protected override JsonTypeInfo TypeInfo + => WorkflowsJsonUtilities.JsonContext.Default.CheckpointInfo; + + private const string CheckpointInfoPropertyNamePattern = @"^(?(((\|\|)|([^\|]))*))\|(?(((\|\|)|([^\|]))*)?)$"; +#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}"; + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/InMemoryCheckpointManager.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/InMemoryCheckpointManager.cs index aca61fa51e..73d0ce50f4 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/InMemoryCheckpointManager.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/InMemoryCheckpointManager.cs @@ -12,21 +12,22 @@ namespace Microsoft.Agents.AI.Workflows.Checkpointing; /// internal sealed class InMemoryCheckpointManager : ICheckpointManager { - private readonly Dictionary> _store = []; + [JsonInclude] + internal Dictionary> Store { get; } = []; public InMemoryCheckpointManager() { } [JsonConstructor] internal InMemoryCheckpointManager(Dictionary> store) { - this._store = store; + this.Store = store; } private RunCheckpointCache GetRunStore(string runId) { - if (!this._store.TryGetValue(runId, out RunCheckpointCache? runStore)) + if (!this.Store.TryGetValue(runId, out RunCheckpointCache? runStore)) { - runStore = this._store[runId] = new(); + runStore = this.Store[runId] = new(); } return runStore; diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/JsonConverterDictionarySupportBase.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/JsonConverterDictionarySupportBase.cs index 48cbf9c3f5..827f494e16 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/JsonConverterDictionarySupportBase.cs +++ b/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 : 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; diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/JsonMarshaller.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/JsonMarshaller.cs index 6eb7ce9231..a6a69f258f 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/JsonMarshaller.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/JsonMarshaller.cs @@ -18,6 +18,7 @@ internal sealed class JsonMarshaller : IWireMarshaller 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; } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/RunCheckpointCache.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/RunCheckpointCache.cs index 01f98f3588..6dd5da9f9c 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/RunCheckpointCache.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/RunCheckpointCache.cs @@ -8,23 +8,26 @@ namespace Microsoft.Agents.AI.Workflows.Checkpointing; internal sealed class RunCheckpointCache { - private readonly List _checkpointIndex = []; - private readonly Dictionary _cache = []; + [JsonInclude] + internal List CheckpointIndex { get; } = []; + + [JsonInclude] + internal Dictionary Cache { get; } = []; public RunCheckpointCache() { } [JsonConstructor] internal RunCheckpointCache(List checkpointIndex, Dictionary cache) { - this._checkpointIndex = checkpointIndex; - this._cache = cache; + this.CheckpointIndex = checkpointIndex; + this.Cache = cache; } [JsonIgnore] - public IEnumerable Index => this._checkpointIndex; + public IEnumerable 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 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; diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/ScopeKeyConverter.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/ScopeKeyConverter.cs index c63de96bae..550bba445e 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/ScopeKeyConverter.cs +++ b/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); diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/JsonSerializationTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/JsonSerializationTests.cs index 7afbe96abe..1c20a7a091 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/JsonSerializationTests.cs +++ b/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); + } }