fix: Update WorkflowHostAgent post Checkpointing implementation (#1498)

WorkflowHostAgent was initially implemented before Checkpointing was available. This meant that in order to support resuming, the WorkflowHostAgent needed to keep the runs around, which broke it when stricter rules about concurrent sharing of workflows during execution were introduced.

This change updates the hosting logic to release the underlying StreamingRun when the RunStreamingAsync or RunAsync are invoked, in favour of keeping the checkpointing information in the WorkflowThread to enable resumption.

dotnet/src/Microsoft.Agents.AI.Workflows/CheckpointManager.cs

@@ -20,7 +20,7 @@ public sealed class CheckpointManager : ICheckpointManager
         return new CheckpointManagerImpl<TStoreObject>(marshaller, store);
     }
 
-    private CheckpointManager(ICheckpointManager impl)
+    internal CheckpointManager(ICheckpointManager impl)
     {
         this._impl = impl;
     }

dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs

@@ -3,8 +3,6 @@
 using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
-using System.Diagnostics;
-using System.Linq;
 using System.Runtime.CompilerServices;
 using System.Text.Json;
 using System.Threading;
@@ -16,18 +14,19 @@ namespace Microsoft.Agents.AI.Workflows;
 
 internal sealed class WorkflowHostAgent : AIAgent
 {
-    private readonly Workflow<List<ChatMessage>> _workflow;
+    private readonly Workflow _workflow;
     private readonly string? _id;
+    private readonly CheckpointManager? _checkpointManager;
 
     private readonly ConcurrentDictionary<string, string> _assignedRunIds = [];
-    private readonly Dictionary<string, StreamingRun> _runningWorkflows = [];
 
-    public WorkflowHostAgent(Workflow<List<ChatMessage>> workflow, string? id = null, string? name = null)
+    public WorkflowHostAgent(Workflow<List<ChatMessage>> workflow, string? id = null, string? name = null, CheckpointManager? checkpointManager = null)
     {
         this._workflow = Throw.IfNull(workflow);
 
         this._id = id;
         this.Name = name;
+        this._checkpointManager = checkpointManager;
     }
 
     public override string? Name { get; }
@@ -45,59 +44,10 @@ internal sealed class WorkflowHostAgent : AIAgent
         return result;
     }
 
-    public override AgentThread GetNewThread() => new WorkflowThread(this.Id, this.Name, this.GenerateNewId());
+    public override AgentThread GetNewThread() => new WorkflowThread(this._workflow, this.GenerateNewId(), this._checkpointManager);
 
     public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
-        => new WorkflowThread(serializedThread, jsonSerializerOptions);
-
-    private async
-    IAsyncEnumerable<AgentRunResponseUpdate> InvokeStageAsync(
-        WorkflowThread conversation,
-        [EnumeratorCancellation] CancellationToken cancellationToken = default)
-    {
-        string runId = conversation.RunId;
-        List<ChatMessage> messages = conversation.MessageStore.GetFromBookmark().ToList();
-
-        try
-        {
-            // technically there is a race condition here between assigning the ID, and checking if it exists
-            // in the case of new threads.
-            if (!this._runningWorkflows.TryGetValue(runId, out StreamingRun? run))
-            {
-                run = await InProcessExecution.StreamAsync(this._workflow, messages, cancellationToken: cancellationToken)
-                                              .ConfigureAwait(false);
-                this._runningWorkflows[runId] = run;
-            }
-            else
-            {
-                bool sentMessages = await run.TrySendMessageAsync(messages).ConfigureAwait(false);
-                Debug.Assert(sentMessages, "Hosted workflow is required to take List<ChatMessage> as input.");
-            }
-
-            await run.TrySendMessageAsync(new TurnToken(emitEvents: true)).ConfigureAwait(false);
-            await foreach (WorkflowEvent evt in run.WatchStreamAsync(blockOnPendingRequest: false, cancellationToken)
-                                               .ConfigureAwait(false)
-                                               .WithCancellation(cancellationToken))
-            {
-                switch (evt)
-                {
-                    case AgentRunUpdateEvent agentUpdate:
-                        yield return agentUpdate.Update;
-                        break;
-                    case RequestInfoEvent requestInfo:
-                        FunctionCallContent fcContent = requestInfo.Request.ToFunctionCall();
-                        AgentRunResponseUpdate update = conversation.CreateUpdate(fcContent);
-                        yield return update;
-                        break;
-                }
-            }
-        }
-        finally
-        {
-            // Do we want to try to undo the step, and not update the bookmark?
-            conversation.MessageStore.UpdateBookmark();
-        }
-    }
+        => new WorkflowThread(this._workflow, serializedThread, this._checkpointManager, jsonSerializerOptions);
 
     private async ValueTask<WorkflowThread> UpdateThreadAsync(IEnumerable<ChatMessage> messages, AgentThread? thread = null, CancellationToken cancellationToken = default)
     {
@@ -122,14 +72,14 @@ internal sealed class WorkflowHostAgent : AIAgent
         WorkflowThread workflowThread = await this.UpdateThreadAsync(messages, thread, cancellationToken).ConfigureAwait(false);
         MessageMerger merger = new();
 
-        await foreach (AgentRunResponseUpdate update in this.InvokeStageAsync(workflowThread, cancellationToken)
-                                                            .ConfigureAwait(false)
-                                                            .WithCancellation(cancellationToken))
+        await foreach (AgentRunResponseUpdate update in workflowThread.InvokeStageAsync(cancellationToken)
+                                                                      .ConfigureAwait(false)
+                                                                      .WithCancellation(cancellationToken))
         {
             merger.AddUpdate(update);
         }
 
-        return merger.ComputeMerged(workflowThread.ResponseId, this.Id, this.Name);
+        return merger.ComputeMerged(workflowThread.LastResponseId!, this.Id, this.Name);
     }
 
     public override async
@@ -140,9 +90,9 @@ internal sealed class WorkflowHostAgent : AIAgent
         [EnumeratorCancellation] CancellationToken cancellationToken = default)
     {
         WorkflowThread workflowThread = await this.UpdateThreadAsync(messages, thread, cancellationToken).ConfigureAwait(false);
-        await foreach (AgentRunResponseUpdate update in this.InvokeStageAsync(workflowThread, cancellationToken)
-                                                            .ConfigureAwait(false)
-                                                            .WithCancellation(cancellationToken))
+        await foreach (AgentRunResponseUpdate update in workflowThread.InvokeStageAsync(cancellationToken)
+                                                                      .ConfigureAwait(false)
+                                                                      .WithCancellation(cancellationToken))
         {
             yield return update;
         }

dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowMessageStore.cs

@@ -1,11 +1,11 @@
 // Copyright (c) Microsoft. All rights reserved.
 
-using System;
 using System.Collections.Generic;
 using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
+using Microsoft.Shared.Diagnostics;
 
 namespace Microsoft.Agents.AI.Workflows;
 
@@ -18,22 +18,22 @@ internal sealed class WorkflowMessageStore : ChatMessageStore
     {
     }
 
-    public WorkflowMessageStore(JsonElement serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null)
+    public WorkflowMessageStore(StoreState state)
     {
-        if (serializedStoreState.ValueKind is not JsonValueKind.Object)
-        {
-            throw new ArgumentException("The provided JsonElement must be a json object", nameof(serializedStoreState));
-        }
+        this.ImportStoreState(Throw.IfNull(state));
+    }
 
-        StoreState? state =
-            serializedStoreState.Deserialize(
-                AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(StoreState))) as StoreState;
+    private void ImportStoreState(StoreState state, bool clearMessages = false)
+    {
+        if (clearMessages)
+        {
+            this._chatMessages.Clear();
+        }
 
         if (state?.Messages is not null)
         {
             this._chatMessages.AddRange(state.Messages);
         }
-
         this._bookmark = state?.Bookmark ?? 0;
     }
 
@@ -66,13 +66,11 @@ internal sealed class WorkflowMessageStore : ChatMessageStore
 
     public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
     {
-        StoreState state = new()
-        {
-            Bookmark = this._bookmark,
-            Messages = this._chatMessages,
-        };
+        StoreState state = this.ExportStoreState();
 
         return JsonSerializer.SerializeToElement(state,
             WorkflowsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(StoreState)));
     }
+
+    internal StoreState ExportStoreState() => new() { Bookmark = this._bookmark, Messages = this._chatMessages };
 }

dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowThread.cs

@@ -1,7 +1,13 @@
 // Copyright (c) Microsoft. All rights reserved.
 
 using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Runtime.CompilerServices;
 using System.Text.Json;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Agents.AI.Workflows.Checkpointing;
 using Microsoft.Extensions.AI;
 using Microsoft.Shared.Diagnostics;
 
@@ -9,26 +15,71 @@ namespace Microsoft.Agents.AI.Workflows;
 
 internal sealed class WorkflowThread : AgentThread
 {
-    public WorkflowThread(string workflowId, string? workflowName, string runId)
+    private readonly CheckpointManager _checkpointManager;
+    private readonly InMemoryCheckpointManager? _inMemoryCheckpointManager;
+    private readonly Workflow _workflow;
+
+    public WorkflowThread(Workflow workflow, string runId, CheckpointManager? checkpointManager = null)
     {
-        this.MessageStore = new();
+        this._workflow = Throw.IfNull(workflow);
+
+        // If the user provided an external checkpoint manager, use that, otherwise rely on an in-memory one.
+        // TODO: Implement persist-only-last functionality for in-memory checkpoint manager, to avoid unbounded
+        // memory growth.
+        this._checkpointManager = checkpointManager ?? new(this._inMemoryCheckpointManager = new());
+
         this.RunId = Throw.IfNullOrEmpty(runId);
+        this.MessageStore = new WorkflowMessageStore();
     }
 
-    public WorkflowThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
+    public WorkflowThread(Workflow workflow, JsonElement serializedThread, CheckpointManager? checkpointManager = null, JsonSerializerOptions? jsonSerializerOptions = null)
     {
-        throw new NotImplementedException("Pending Checkpointing work.");
+        this._workflow = Throw.IfNull(workflow);
+
+        JsonMarshaller marshaller = new(jsonSerializerOptions);
+        ThreadState threadState = marshaller.Marshal<ThreadState>(serializedThread);
+
+        this._inMemoryCheckpointManager = threadState.CheckpointManager;
+        if (this._inMemoryCheckpointManager is not null && checkpointManager is not null)
+        {
+            // The thread was externalized with an in-memory checkpoint manager, but the caller is providing an external one.
+            throw new ArgumentException("Cannot provide an external checkpoint manager when deserializing a thread that " +
+                "was serialized with an in-memory checkpoint manager.", nameof(checkpointManager));
+        }
+        else if (this._inMemoryCheckpointManager is null && checkpointManager is null)
+        {
+            // The thread was externalized without an in-memory checkpoint manager, and the caller is not providing an external one.
+            throw new ArgumentException("An external checkpoint manager must be provided when deserializing a thread that " +
+                "was serialized without an in-memory checkpoint manager.", nameof(checkpointManager));
+        }
+        else
+        {
+            this._checkpointManager = checkpointManager ?? new(this._inMemoryCheckpointManager!);
+        }
+
+        this.RunId = threadState.RunId;
+        this.LastCheckpoint = threadState.LastCheckpoint;
+        this.MessageStore = new WorkflowMessageStore(threadState.MessageStoreState);
     }
 
-    public string RunId { get; }
-    public int Halts { get; }
+    public CheckpointInfo? LastCheckpoint { get; set; }
 
-    public string ResponseId => $"{this.RunId}@{this.Halts}";
+    protected override Task MessagesReceivedAsync(IEnumerable<ChatMessage> newMessages, CancellationToken cancellationToken = default)
+        => this.MessageStore.AddMessagesAsync(newMessages, cancellationToken);
 
     public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
-        => throw new NotImplementedException("Pending Checkpointing work.");
+    {
+        JsonMarshaller marshaller = new(jsonSerializerOptions);
+        ThreadState info = new(
+            this.RunId,
+            this.LastCheckpoint,
+            this.MessageStore.ExportStoreState(),
+            this._inMemoryCheckpointManager);
 
-    public AgentRunResponseUpdate CreateUpdate(params AIContent[] parts)
+        return marshaller.Marshal(info);
+    }
+
+    public AgentRunResponseUpdate CreateUpdate(string responseId, params AIContent[] parts)
     {
         Throw.IfNullOrEmpty(parts);
 
@@ -36,6 +87,8 @@ internal sealed class WorkflowThread : AgentThread
         {
             CreatedAt = DateTimeOffset.UtcNow,
             MessageId = Guid.NewGuid().ToString("N"),
+            Role = ChatRole.Assistant,
+            ResponseId = responseId
         };
 
         this.MessageStore.AddMessages(update.ToChatMessage());
@@ -43,6 +96,89 @@ internal sealed class WorkflowThread : AgentThread
         return update;
     }
 
+    private async ValueTask<Checkpointed<StreamingRun>> CreateOrResumeRunAsync(List<ChatMessage> messages, CancellationToken cancellationToken = default)
+    {
+        if (this.LastCheckpoint is not null)
+        {
+            Checkpointed<StreamingRun> checkpointed =
+                await InProcessExecution.ResumeStreamAsync(this._workflow,
+                                                           this.LastCheckpoint,
+                                                           this._checkpointManager,
+                                                           this.RunId,
+                                                           cancellationToken)
+                                        .ConfigureAwait(false);
+
+            await checkpointed.Run.TrySendMessageAsync(messages).ConfigureAwait(false);
+            return checkpointed;
+        }
+
+        return await InProcessExecution.StreamAsync(this._workflow,
+                                                    messages,
+                                                    this._checkpointManager,
+                                                    this.RunId,
+                                                    cancellationToken)
+                                       .ConfigureAwait(false);
+    }
+
+    internal async
+    IAsyncEnumerable<AgentRunResponseUpdate> InvokeStageAsync(
+        [EnumeratorCancellation] CancellationToken cancellationToken = default)
+    {
+        try
+        {
+            this.LastResponseId = Guid.NewGuid().ToString("N");
+            List<ChatMessage> messages = this.MessageStore.GetFromBookmark().ToList();
+
+#pragma warning disable CA2007 // Analyzer misfiring and not seeing .ConfigureAwait(false) below.
+            await using Checkpointed<StreamingRun> checkpointed =
+                await this.CreateOrResumeRunAsync(messages, cancellationToken).ConfigureAwait(false);
+#pragma warning restore CA2007
+
+            StreamingRun run = checkpointed.Run;
+            await run.TrySendMessageAsync(new TurnToken(emitEvents: true)).ConfigureAwait(false);
+            await foreach (WorkflowEvent evt in run.WatchStreamAsync(blockOnPendingRequest: false, cancellationToken)
+                                               .ConfigureAwait(false)
+                                               .WithCancellation(cancellationToken))
+            {
+                switch (evt)
+                {
+                    case AgentRunUpdateEvent agentUpdate:
+                        yield return agentUpdate.Update;
+                        break;
+                    case RequestInfoEvent requestInfo:
+                        FunctionCallContent fcContent = requestInfo.Request.ToFunctionCall();
+                        AgentRunResponseUpdate update = this.CreateUpdate(this.LastResponseId, fcContent);
+                        yield return update;
+                        break;
+                    case SuperStepCompletedEvent stepCompleted:
+                        this.LastCheckpoint = stepCompleted.CompletionInfo?.Checkpoint;
+                        break;
+                }
+            }
+        }
+        finally
+        {
+            // Do we want to try to undo the step, and not update the bookmark?
+            this.MessageStore.UpdateBookmark();
+        }
+    }
+
+    public string? LastResponseId { get; set; }
+
+    public string RunId { get; }
+
     /// <inheritdoc/>
     public WorkflowMessageStore MessageStore { get; }
+
+    internal sealed class ThreadState(
+        string runId,
+        CheckpointInfo? lastCheckpoint,
+        WorkflowMessageStore.StoreState messageStoreState,
+        InMemoryCheckpointManager? checkpointManager = null)
+    {
+        public string RunId { get; } = runId;
+        public CheckpointInfo? LastCheckpoint { get; } = lastCheckpoint;
+        public WorkflowMessageStore.StoreState MessageStoreState { get; } = messageStoreState;
+        public InMemoryCheckpointManager? CheckpointManager { get; } = checkpointManager;
+    }
 }

dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowsJsonUtilities.cs

@@ -7,7 +7,6 @@ using System.Text.Json.Serialization;
 using Microsoft.Agents.AI.Workflows.Checkpointing;
 using Microsoft.Agents.AI.Workflows.Execution;
 using Microsoft.Extensions.AI;
-using static Microsoft.Agents.AI.Workflows.WorkflowMessageStore;
 
 namespace Microsoft.Agents.AI.Workflows;
 
@@ -80,7 +79,8 @@ internal static partial class WorkflowsJsonUtilities
     [JsonSerializable(typeof(EdgeConnection))]
 
     // Workflow-as-Agent
-    [JsonSerializable(typeof(StoreState))]
+    [JsonSerializable(typeof(WorkflowMessageStore.StoreState))]
+    [JsonSerializable(typeof(WorkflowThread.ThreadState))]
 
     // Message Types
     [JsonSerializable(typeof(ChatMessage))]
@@ -91,7 +91,7 @@ internal static partial class WorkflowsJsonUtilities
     // Event Types
     //[JsonSerializable(typeof(WorkflowEvent))]
     //   Currently cannot be serialized because it includes Exceptions.
-    //   We'll need a way to marshal this correct in the AgentRuntime case.
+    //   We'll need a way to marshal this correctly in the AgentRuntime case.
     //   For now this is okay, because we never serialize WorkflowEvents into
     //   checkpoints.
     [JsonSerializable(typeof(JsonElement))]