From 69e1ab0409c8430d36515ab656892852e6fe957b Mon Sep 17 00:00:00 2001 From: Jacob Alber Date: Wed, 15 Oct 2025 14:18:28 -0400 Subject: [PATCH] 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. --- .../CheckpointManager.cs | 2 +- .../WorkflowHostAgent.cs | 76 ++------- .../WorkflowMessageStore.cs | 28 ++-- .../WorkflowThread.cs | 154 +++++++++++++++++- .../WorkflowsJsonUtilities.cs | 6 +- 5 files changed, 175 insertions(+), 91 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/CheckpointManager.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/CheckpointManager.cs index 9334158b5b..c50283e728 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/CheckpointManager.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/CheckpointManager.cs @@ -20,7 +20,7 @@ public sealed class CheckpointManager : ICheckpointManager return new CheckpointManagerImpl(marshaller, store); } - private CheckpointManager(ICheckpointManager impl) + internal CheckpointManager(ICheckpointManager impl) { this._impl = impl; } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs index 965a23a22f..8fa289d7ea 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs +++ b/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> _workflow; + private readonly Workflow _workflow; private readonly string? _id; + private readonly CheckpointManager? _checkpointManager; private readonly ConcurrentDictionary _assignedRunIds = []; - private readonly Dictionary _runningWorkflows = []; - public WorkflowHostAgent(Workflow> workflow, string? id = null, string? name = null) + public WorkflowHostAgent(Workflow> 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 InvokeStageAsync( - WorkflowThread conversation, - [EnumeratorCancellation] CancellationToken cancellationToken = default) - { - string runId = conversation.RunId; - List 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 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 UpdateThreadAsync(IEnumerable 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; } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowMessageStore.cs index 1fffb502fc..39c83bcadf 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowMessageStore.cs +++ b/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 }; } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowThread.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowThread.cs index 8db2eabe93..61dc0ab337 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowThread.cs +++ b/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(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 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> CreateOrResumeRunAsync(List messages, CancellationToken cancellationToken = default) + { + if (this.LastCheckpoint is not null) + { + Checkpointed 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 InvokeStageAsync( + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + try + { + this.LastResponseId = Guid.NewGuid().ToString("N"); + List messages = this.MessageStore.GetFromBookmark().ToList(); + +#pragma warning disable CA2007 // Analyzer misfiring and not seeing .ConfigureAwait(false) below. + await using Checkpointed 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; } + /// 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; + } } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowsJsonUtilities.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowsJsonUtilities.cs index 33cf002517..f64c1a6b80 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowsJsonUtilities.cs +++ b/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))]