Skip to content

Commit a49bb6b

Browse files
lokitothlokitoth
committed
Fix CheckpointInfo.Parent always null in InProcessRunner (#3796)
Track the last CheckpointInfo in InProcessRunner so that newly created checkpoints reference their parent. When resuming from a checkpoint, the resumed-from checkpoint becomes the parent of the next checkpoint. Adds tests verifying: - First checkpoint has null parent - Subsequent checkpoints chain parents correctly - Checkpoint after resume references the resumed-from checkpoint
1 parent 8ad6663 commit a49bb6b

2 files changed

Lines changed: 183 additions & 4 deletions

File tree

‎dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunner.cs‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -245,6 +245,7 @@ await this.RaiseWorkflowEventAsync(this.StepTracer.Complete(this.RunContext.Next
245245
}
246246

247247
private WorkflowInfo? _workflowInfoCache;
248+
private CheckpointInfo? _lastCheckpointInfo;
248249
private readonly List<CheckpointInfo> _checkpoints = [];
249250
internal async ValueTask CheckpointAsync(CancellationToken cancellationToken = default)
250251
{
@@ -270,10 +271,10 @@ internal async ValueTask CheckpointAsync(CancellationToken cancellationToken = d
270271
RunnerStateData runnerData = await this.RunContext.ExportStateAsync().ConfigureAwait(false);
271272
Dictionary<ScopeKey, PortableValue> stateData = await this.RunContext.StateManager.ExportStateAsync().ConfigureAwait(false);
272273

273-
Checkpoint checkpoint = new(this.StepTracer.StepNumber, this._workflowInfoCache, runnerData, stateData, edgeData);
274-
CheckpointInfo checkpointInfo = await this.CheckpointManager.CommitCheckpointAsync(this.RunId, checkpoint).ConfigureAwait(false);
275-
this.StepTracer.TraceCheckpointCreated(checkpointInfo);
276-
this._checkpoints.Add(checkpointInfo);
274+
Checkpoint checkpoint = new(this.StepTracer.StepNumber, this._workflowInfoCache, runnerData, stateData, edgeData, this._lastCheckpointInfo);
275+
this._lastCheckpointInfo = await this.CheckpointManager.CommitCheckpointAsync(this.RunId, checkpoint).ConfigureAwait(false);
276+
this.StepTracer.TraceCheckpointCreated(this._lastCheckpointInfo);
277+
this._checkpoints.Add(this._lastCheckpointInfo);
277278
}
278279

279280
public async ValueTask RestoreCheckpointAsync(CheckpointInfo checkpointInfo, CancellationToken cancellationToken = default)
@@ -304,6 +305,7 @@ public async ValueTask RestoreCheckpointAsync(CheckpointInfo checkpointInfo, Can
304305
await this.EdgeMap.ImportStateAsync(checkpoint).ConfigureAwait(false);
305306
await Task.WhenAll(executorNotifyTask, republishRequestsTask.AsTask()).ConfigureAwait(false);
306307

308+
this._lastCheckpointInfo = checkpointInfo;
307309
this.StepTracer.Reload(this.StepTracer.StepNumber);
308310
}
309311

‎dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/CheckpointParentTests.cs‎

Lines changed: 177 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,177 @@
1+
// Copyright (c) Microsoft. All rights reserved.
2+
3+
using System.Collections.Generic;
4+
using System.Threading;
5+
using System.Threading.Tasks;
6+
using FluentAssertions;
7+
using Microsoft.Agents.AI.Workflows.Checkpointing;
8+
9+
namespace Microsoft.Agents.AI.Workflows.UnitTests;
10+
11+
/// <summary>
12+
/// Tests for verifying that CheckpointInfo.Parent is properly populated
13+
/// when checkpoints are created during workflow execution (GH #3796).
14+
/// </summary>
15+
public class CheckpointParentTests
16+
{
17+
[Theory]
18+
[InlineData(ExecutionEnvironment.InProcess_Lockstep)]
19+
[InlineData(ExecutionEnvironment.InProcess_OffThread)]
20+
internal async Task Checkpoint_FirstCheckpoint_ShouldHaveNullParentAsync(ExecutionEnvironment environment)
21+
{
22+
// Arrange: A simple two-step workflow that will produce at least one checkpoint.
23+
ForwardMessageExecutor<string> executorA = new("A");
24+
ForwardMessageExecutor<string> executorB = new("B");
25+
26+
Workflow workflow = new WorkflowBuilder(executorA)
27+
.AddEdge(executorA, executorB)
28+
.Build();
29+
30+
CheckpointManager checkpointManager = CheckpointManager.CreateInMemory();
31+
IWorkflowExecutionEnvironment env = environment.ToWorkflowExecutionEnvironment();
32+
33+
// Act
34+
Checkpointed<StreamingRun> checkpointed =
35+
await env.StreamAsync(workflow, "Hello", checkpointManager);
36+
37+
List<CheckpointInfo> checkpoints = [];
38+
await foreach (WorkflowEvent evt in checkpointed.Run.WatchStreamAsync())
39+
{
40+
if (evt is SuperStepCompletedEvent stepEvt && stepEvt.CompletionInfo?.Checkpoint is { } cp)
41+
{
42+
checkpoints.Add(cp);
43+
}
44+
}
45+
46+
// Assert: The first checkpoint should have been created and stored with a null parent.
47+
checkpoints.Should().NotBeEmpty("at least one checkpoint should have been created");
48+
49+
CheckpointInfo firstCheckpoint = checkpoints[0];
50+
Checkpoint storedFirst = await ((ICheckpointManager)checkpointManager)
51+
.LookupCheckpointAsync(firstCheckpoint.RunId, firstCheckpoint);
52+
storedFirst.Parent.Should().BeNull("the first checkpoint should have no parent");
53+
}
54+
55+
[Theory]
56+
[InlineData(ExecutionEnvironment.InProcess_Lockstep)]
57+
[InlineData(ExecutionEnvironment.InProcess_OffThread)]
58+
internal async Task Checkpoint_SubsequentCheckpoints_ShouldChainParentsAsync(ExecutionEnvironment environment)
59+
{
60+
// Arrange: A workflow with a loop that will produce multiple checkpoints.
61+
ForwardMessageExecutor<string> executorA = new("A");
62+
ForwardMessageExecutor<string> executorB = new("B");
63+
64+
// A -> B -> A (loop) to generate multiple supersteps/checkpoints.
65+
Workflow workflow = new WorkflowBuilder(executorA)
66+
.AddEdge(executorA, executorB)
67+
.AddEdge(executorB, executorA)
68+
.Build();
69+
70+
CheckpointManager checkpointManager = CheckpointManager.CreateInMemory();
71+
IWorkflowExecutionEnvironment env = environment.ToWorkflowExecutionEnvironment();
72+
73+
// Act
74+
Checkpointed<StreamingRun> checkpointed =
75+
await env.StreamAsync(workflow, "Hello", checkpointManager);
76+
77+
List<CheckpointInfo> checkpoints = [];
78+
CancellationTokenSource cts = new();
79+
80+
await foreach (WorkflowEvent evt in checkpointed.Run.WatchStreamAsync(cts.Token))
81+
{
82+
if (evt is SuperStepCompletedEvent stepEvt && stepEvt.CompletionInfo?.Checkpoint is { } cp)
83+
{
84+
checkpoints.Add(cp);
85+
if (checkpoints.Count >= 3)
86+
{
87+
cts.Cancel();
88+
}
89+
}
90+
}
91+
92+
// Assert: We should have at least 3 checkpoints
93+
checkpoints.Should().HaveCountGreaterThanOrEqualTo(3);
94+
95+
// Verify the parent chain
96+
Checkpoint stored0 = await ((ICheckpointManager)checkpointManager)
97+
.LookupCheckpointAsync(checkpoints[0].RunId, checkpoints[0]);
98+
stored0.Parent.Should().BeNull("the first checkpoint should have no parent");
99+
100+
Checkpoint stored1 = await ((ICheckpointManager)checkpointManager)
101+
.LookupCheckpointAsync(checkpoints[1].RunId, checkpoints[1]);
102+
stored1.Parent.Should().NotBeNull("the second checkpoint should have a parent");
103+
stored1.Parent.Should().Be(checkpoints[0], "the second checkpoint's parent should be the first checkpoint");
104+
105+
Checkpoint stored2 = await ((ICheckpointManager)checkpointManager)
106+
.LookupCheckpointAsync(checkpoints[2].RunId, checkpoints[2]);
107+
stored2.Parent.Should().NotBeNull("the third checkpoint should have a parent");
108+
stored2.Parent.Should().Be(checkpoints[1], "the third checkpoint's parent should be the second checkpoint");
109+
}
110+
111+
[Theory]
112+
[InlineData(ExecutionEnvironment.InProcess_Lockstep)]
113+
[InlineData(ExecutionEnvironment.InProcess_OffThread)]
114+
internal async Task Checkpoint_AfterResume_ShouldHaveResumedCheckpointAsParentAsync(ExecutionEnvironment environment)
115+
{
116+
// Arrange: A looping workflow that produces checkpoints.
117+
ForwardMessageExecutor<string> executorA = new("A");
118+
ForwardMessageExecutor<string> executorB = new("B");
119+
120+
Workflow workflow = new WorkflowBuilder(executorA)
121+
.AddEdge(executorA, executorB)
122+
.AddEdge(executorB, executorA)
123+
.Build();
124+
125+
CheckpointManager checkpointManager = CheckpointManager.CreateInMemory();
126+
IWorkflowExecutionEnvironment env = environment.ToWorkflowExecutionEnvironment();
127+
128+
// First run: collect a checkpoint to resume from
129+
Checkpointed<StreamingRun> checkpointed =
130+
await env.StreamAsync(workflow, "Hello", checkpointManager);
131+
132+
List<CheckpointInfo> firstRunCheckpoints = [];
133+
CancellationTokenSource cts = new();
134+
await foreach (WorkflowEvent evt in checkpointed.Run.WatchStreamAsync(cts.Token))
135+
{
136+
if (evt is SuperStepCompletedEvent stepEvt && stepEvt.CompletionInfo?.Checkpoint is { } cp)
137+
{
138+
firstRunCheckpoints.Add(cp);
139+
if (firstRunCheckpoints.Count >= 2)
140+
{
141+
cts.Cancel();
142+
}
143+
}
144+
}
145+
146+
firstRunCheckpoints.Should().HaveCountGreaterThanOrEqualTo(2);
147+
CheckpointInfo resumePoint = firstRunCheckpoints[0];
148+
149+
// Dispose the first run to release workflow ownership before resuming.
150+
await checkpointed.DisposeAsync();
151+
152+
// Act: Resume from the first checkpoint
153+
Checkpointed<StreamingRun> resumed =
154+
await env.ResumeStreamAsync(workflow, resumePoint, checkpointManager);
155+
156+
List<CheckpointInfo> resumedCheckpoints = [];
157+
CancellationTokenSource cts2 = new();
158+
await foreach (WorkflowEvent evt in resumed.Run.WatchStreamAsync(cts2.Token))
159+
{
160+
if (evt is SuperStepCompletedEvent stepEvt && stepEvt.CompletionInfo?.Checkpoint is { } cp)
161+
{
162+
resumedCheckpoints.Add(cp);
163+
if (resumedCheckpoints.Count >= 1)
164+
{
165+
cts2.Cancel();
166+
}
167+
}
168+
}
169+
170+
// Assert: The first checkpoint after resume should have the resume point as its parent.
171+
resumedCheckpoints.Should().NotBeEmpty();
172+
Checkpoint storedResumed = await ((ICheckpointManager)checkpointManager)
173+
.LookupCheckpointAsync(resumedCheckpoints[0].RunId, resumedCheckpoints[0]);
174+
storedResumed.Parent.Should().NotBeNull("checkpoint created after resume should have a parent");
175+
storedResumed.Parent.Should().Be(resumePoint, "checkpoint after resume should reference the checkpoint we resumed from");
176+
}
177+
}

0 commit comments

Comments
 (0)