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); + } }