agent-framework/dotnet/src/Microsoft.Agents.AI.Workflows/CheckpointManager.cs

// Copyright (c) Microsoft. All rights reserved.

using System.Collections.Generic;
using System.Text.Json;
using System.Threading.Tasks;
using Microsoft.Agents.AI.Workflows.Checkpointing;

namespace Microsoft.Agents.AI.Workflows;

/// <summary>
/// A manager for storing and retrieving workflow execution checkpoints.
/// </summary>
public sealed class CheckpointManager : ICheckpointManager
{
    private readonly ICheckpointManager _impl;

    private static CheckpointManagerImpl<TStoreObject> CreateImpl<TStoreObject>(
        IWireMarshaller<TStoreObject> marshaller,
        ICheckpointStore<TStoreObject> store)
    {
        return new CheckpointManagerImpl<TStoreObject>(marshaller, store);
    }

    internal CheckpointManager(ICheckpointManager impl)
    {
        this._impl = impl;
    }

    /// <summary>
    /// Creates a new instance of <see cref="ICheckpointManager"/> that uses the specified marshaller and store.
    /// </summary>
    /// <returns></returns>
    public static CheckpointManager CreateInMemory() => new(new InMemoryCheckpointManager());

    /// <summary>
    /// Gets the default in-memory checkpoint manager instance.
    /// </summary>
    public static CheckpointManager Default { get; } = CreateInMemory();

    /// <summary>
    /// Creates a new instance of the CheckpointManager that uses JSON serialization for checkpoint data.
    /// </summary>
    /// <param name="store">The checkpoint store to use for persisting and retrieving checkpoint data as JSON elements. Cannot be null.</param>
    /// <param name="customOptions">Optional custom JSON serializer options to use for serialization and deserialization. Must be provided if
    /// using custom types in messages or state.</param>
    /// <returns>A CheckpointManager instance configured to serialize checkpoint data as JSON.</returns>
    public static CheckpointManager CreateJson(ICheckpointStore<JsonElement> store, JsonSerializerOptions? customOptions = null)
    {
        JsonMarshaller marshaller = new(customOptions);
        return new(CreateImpl(marshaller, store));
    }

    ValueTask<CheckpointInfo> ICheckpointManager.CommitCheckpointAsync(string sessionId, Checkpoint checkpoint)
        => this._impl.CommitCheckpointAsync(sessionId, checkpoint);

    ValueTask<Checkpoint> ICheckpointManager.LookupCheckpointAsync(string sessionId, CheckpointInfo checkpointInfo)
        => this._impl.LookupCheckpointAsync(sessionId, checkpointInfo);

    ValueTask<IEnumerable<CheckpointInfo>> ICheckpointManager.RetrieveIndexAsync(string sessionId, CheckpointInfo? withParent)
        => this._impl.RetrieveIndexAsync(sessionId, withParent);
}