Skip to content

Commit 670f7ea

Browse files
author
Jacob Alber
committed
fix: Fix test logic for Handoff to correctly use checkpointing for multiturn
1 parent 6746e68 commit 670f7ea

1 file changed

Lines changed: 56 additions & 71 deletions

File tree

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

Lines changed: 56 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -147,7 +147,7 @@ from i in Enumerable.Range(1, numAgents)
147147
for (int iter = 0; iter < 3; iter++)
148148
{
149149
const string UserInput = "abc";
150-
(string updateText, List<ChatMessage>? result) = await RunWorkflowAsync(workflow, [new ChatMessage(ChatRole.User, UserInput)]);
150+
(string updateText, List<ChatMessage>? result, _) = await RunWorkflowAsync(workflow, [new ChatMessage(ChatRole.User, UserInput)]);
151151

152152
Assert.NotNull(result);
153153
Assert.Equal(numAgents + 1, result.Count);
@@ -225,7 +225,7 @@ public async Task BuildConcurrent_AgentsRunInParallelAsync()
225225
barrier.Value = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
226226
remaining.Value = 2;
227227

228-
(string updateText, List<ChatMessage>? result) = await RunWorkflowAsync(workflow, [new ChatMessage(ChatRole.User, "abc")]);
228+
(string updateText, List<ChatMessage>? result, _) = await RunWorkflowAsync(workflow, [new ChatMessage(ChatRole.User, "abc")]);
229229
Assert.NotEmpty(updateText);
230230
Assert.NotNull(result);
231231

@@ -258,7 +258,7 @@ public async Task Handoffs_NoTransfers_ResponseServedByOriginalAgentAsync()
258258
}), description: "nop"))
259259
.Build();
260260

261-
(string updateText, List<ChatMessage>? result) = await RunWorkflowAsync(workflow, [new ChatMessage(ChatRole.User, "abc")]);
261+
(string updateText, List<ChatMessage>? result, _) = await RunWorkflowAsync(workflow, [new ChatMessage(ChatRole.User, "abc")]);
262262

263263
Assert.Equal("Hello from agent1", updateText);
264264
Assert.NotNull(result);
@@ -296,7 +296,7 @@ public async Task Handoffs_OneTransfer_ResponseServedBySecondAgentAsync()
296296
.WithHandoff(initialAgent, nextAgent)
297297
.Build();
298298

299-
(string updateText, List<ChatMessage>? result) = await RunWorkflowAsync(workflow, [new ChatMessage(ChatRole.User, "abc")]);
299+
(string updateText, List<ChatMessage>? result, _) = await RunWorkflowAsync(workflow, [new ChatMessage(ChatRole.User, "abc")]);
300300

301301
Assert.Equal("Hello from agent2", updateText);
302302
Assert.NotNull(result);
@@ -406,7 +406,7 @@ public async Task Handoffs_TwoTransfers_HandoffTargetsDoNotReceiveHandoffFunctio
406406
.WithHandoff(secondAgent, thirdAgent)
407407
.Build();
408408

409-
(string updateText, _) = await RunWorkflowAsync(workflow, [new ChatMessage(ChatRole.User, "abc")]);
409+
(string updateText, _, _) = await RunWorkflowAsync(workflow, [new ChatMessage(ChatRole.User, "abc")]);
410410

411411
Assert.Contains("Hello from agent3", updateText);
412412

@@ -604,7 +604,7 @@ public async Task Handoffs_TwoTransfers_ResponseServedByThirdAgentAsync()
604604
.WithHandoff(secondAgent, thirdAgent)
605605
.Build();
606606

607-
(string updateText, List<ChatMessage>? result) = await RunWorkflowAsync(workflow, [new ChatMessage(ChatRole.User, "abc")]);
607+
(string updateText, List<ChatMessage>? result, _) = await RunWorkflowAsync(workflow, [new ChatMessage(ChatRole.User, "abc")]);
608608

609609
Assert.Equal("Hello from agent3", updateText);
610610
Assert.NotNull(result);
@@ -651,7 +651,7 @@ public async Task BuildGroupChat_AgentsRunInOrderAsync(int maxIterations)
651651
for (int iter = 0; iter < 3; iter++)
652652
{
653653
const string UserInput = "abc";
654-
(string updateText, List<ChatMessage>? result) = await RunWorkflowAsync(workflow, [new ChatMessage(ChatRole.User, UserInput)]);
654+
(string updateText, List<ChatMessage>? result, _) = await RunWorkflowAsync(workflow, [new ChatMessage(ChatRole.User, UserInput)]);
655655

656656
Assert.NotNull(result);
657657
Assert.Equal(maxIterations + 1, result.Count);
@@ -705,17 +705,15 @@ public async Task Handoffs_ReturnToPrevious_DisabledByDefault_SecondTurnRoutesVi
705705
.WithHandoff(coordinator, specialist)
706706
.Build();
707707

708-
var environment = InProcessExecution.Lockstep;
709-
string sessionId = Guid.NewGuid().ToString("N");
708+
CheckpointManager checkpointManager = CheckpointManager.CreateInMemory();
709+
const ExecutionEnvironment Environment = ExecutionEnvironment.InProcess_Lockstep;
710710

711711
// Turn 1: coordinator hands off to specialist
712-
(_, List<ChatMessage>? turn1Result) = await RunWorkflowAsync(workflow, [new ChatMessage(ChatRole.User, "book an appointment")], environment, sessionId);
712+
WorkflowRunResult result = await RunWorkflowCheckpointedAsync(workflow, [new ChatMessage(ChatRole.User, "book an appointment")], Environment, checkpointManager);
713713
Assert.Equal(1, coordinatorCallCount);
714714

715715
// Turn 2: without ReturnToPrevious, coordinator should be invoked again
716-
Assert.NotNull(turn1Result);
717-
turn1Result.Add(new ChatMessage(ChatRole.User, "my id is 12345"));
718-
_ = await RunWorkflowAsync(workflow, turn1Result, environment, sessionId);
716+
_ = await RunWorkflowCheckpointedAsync(workflow, [new ChatMessage(ChatRole.User, "my id is 12345")], Environment, checkpointManager, result.LastCheckpoint);
719717
Assert.Equal(2, coordinatorCallCount);
720718
}
721719

@@ -744,18 +742,16 @@ public async Task Handoffs_ReturnToPrevious_Enabled_SecondTurnRoutesDirectlyToSp
744742
.EnableReturnToPrevious()
745743
.Build();
746744

747-
var environment = InProcessExecution.Lockstep;
748-
string sessionId = Guid.NewGuid().ToString("N");
745+
CheckpointManager checkpointManager = CheckpointManager.CreateInMemory();
746+
const ExecutionEnvironment Environment = ExecutionEnvironment.InProcess_Lockstep;
749747

750748
// Turn 1: coordinator hands off to specialist
751-
(_, List<ChatMessage>? turn1Result) = await RunWorkflowAsync(workflow, [new ChatMessage(ChatRole.User, "book an appointment")], environment, sessionId);
749+
WorkflowRunResult result = await RunWorkflowCheckpointedAsync(workflow, [new ChatMessage(ChatRole.User, "book an appointment")], Environment, checkpointManager);
752750
Assert.Equal(1, coordinatorCallCount);
753751
Assert.Equal(1, specialistCallCount);
754752

755753
// Turn 2: with ReturnToPrevious, specialist should be invoked directly, coordinator should NOT be called again
756-
Assert.NotNull(turn1Result);
757-
turn1Result.Add(new ChatMessage(ChatRole.User, "my id is 12345"));
758-
_ = await RunWorkflowAsync(workflow, turn1Result, environment, sessionId);
754+
_ = await RunWorkflowCheckpointedAsync(workflow, [new ChatMessage(ChatRole.User, "my id is 12345")], Environment, checkpointManager, result.LastCheckpoint);
759755
Assert.Equal(1, coordinatorCallCount); // coordinator NOT called again
760756
Assert.Equal(2, specialistCallCount); // specialist called again
761757
}
@@ -782,11 +778,8 @@ public async Task Handoffs_ReturnToPrevious_Enabled_BeforeAnyHandoff_RoutesViaIn
782778
.EnableReturnToPrevious()
783779
.Build();
784780

785-
var environment = InProcessExecution.Lockstep;
786-
string sessionId = Guid.NewGuid().ToString("N");
787-
788781
// First turn with no prior handoff: should route to initial (coordinator) agent
789-
_ = await RunWorkflowAsync(workflow, [new ChatMessage(ChatRole.User, "hello")], environment, sessionId);
782+
_ = await RunWorkflowAsync(workflow, [new ChatMessage(ChatRole.User, "hello")]);
790783
Assert.Equal(1, coordinatorCallCount);
791784
}
792785

@@ -825,81 +818,73 @@ public async Task Handoffs_ReturnToPrevious_Enabled_AfterHandoffBackToCoordinato
825818
.EnableReturnToPrevious()
826819
.Build();
827820

828-
var environment = InProcessExecution.Lockstep;
829-
string sessionId = Guid.NewGuid().ToString("N");
821+
CheckpointManager checkpointManager = CheckpointManager.CreateInMemory();
822+
const ExecutionEnvironment Environment = ExecutionEnvironment.InProcess_Lockstep;
830823

831824
// Turn 1: coordinator → specialist → coordinator (specialist hands back)
832-
(_, List<ChatMessage>? turn1Result) = await RunWorkflowAsync(workflow, [new ChatMessage(ChatRole.User, "book an appointment")], environment, sessionId);
825+
WorkflowRunResult result = await RunWorkflowCheckpointedAsync(workflow, [new ChatMessage(ChatRole.User, "book an appointment")], Environment, checkpointManager);
833826
Assert.Equal(2, coordinatorCallCount); // called twice: initial handoff + receiving handback
834827
Assert.Equal(1, specialistCallCount); // specialist called once, then handed back
835828

836829
// Turn 2: after handoff back to coordinator, should route to coordinator (not specialist)
837-
Assert.NotNull(turn1Result);
838-
turn1Result.Add(new ChatMessage(ChatRole.User, "never mind"));
839-
_ = await RunWorkflowAsync(workflow, turn1Result, environment, sessionId);
830+
_ = await RunWorkflowCheckpointedAsync(workflow, [new ChatMessage(ChatRole.User, "never mind")], Environment, checkpointManager, result.LastCheckpoint);
840831
Assert.Equal(3, coordinatorCallCount); // coordinator called again on turn 2
841832
Assert.Equal(1, specialistCallCount); // specialist NOT called
842833
}
843834

844-
private static async Task<(string UpdateText, List<ChatMessage>? Result)> RunWorkflowAsync(
845-
Workflow workflow, List<ChatMessage> input, InProcessExecutionEnvironment environment, string? sessionId = null)
846-
{
847-
StringBuilder sb = new();
848-
849-
await using StreamingRun run = await environment.RunStreamingAsync(workflow, input, sessionId);
850-
await run.TrySendMessageAsync(new TurnToken(emitEvents: true));
835+
private sealed record WorkflowRunResult(string UpdateText, List<ChatMessage>? Result, CheckpointInfo? LastCheckpoint);
851836

852-
WorkflowOutputEvent? output = null;
853-
await foreach (WorkflowEvent evt in run.WatchStreamAsync().ConfigureAwait(false))
854-
{
855-
if (evt is AgentResponseUpdateEvent executorComplete)
856-
{
857-
sb.Append(executorComplete.Data);
858-
}
859-
else if (evt is WorkflowOutputEvent e)
860-
{
861-
output = e;
862-
break;
863-
}
864-
else if (evt is WorkflowErrorEvent errorEvent)
865-
{
866-
Assert.Fail($"Workflow execution failed with error: {errorEvent.Exception}");
867-
}
868-
}
837+
private static Task<WorkflowRunResult> RunWorkflowCheckpointedAsync(
838+
Workflow workflow, List<ChatMessage> input, ExecutionEnvironment executionEnvironment, CheckpointManager checkpointManager, CheckpointInfo? fromCheckpoint = null)
839+
{
840+
InProcessExecutionEnvironment environment = executionEnvironment.ToWorkflowExecutionEnvironment()
841+
.WithCheckpointing(checkpointManager);
869842

870-
return (sb.ToString(), output?.As<List<ChatMessage>>());
843+
return RunWorkflowCheckpointedAsync(workflow, input, environment, fromCheckpoint);
871844
}
872845

873-
private static async Task<(string UpdateText, List<ChatMessage>? Result)> RunWorkflowAsync(
874-
Workflow workflow, List<ChatMessage> input, ExecutionEnvironment executionEnvironment = ExecutionEnvironment.InProcess_Lockstep)
846+
private static async Task<WorkflowRunResult> RunWorkflowCheckpointedAsync(
847+
Workflow workflow, List<ChatMessage> input, InProcessExecutionEnvironment environment, CheckpointInfo? fromCheckpoint = null)
875848
{
876-
StringBuilder sb = new();
849+
await using StreamingRun run =
850+
fromCheckpoint != null ? await environment.ResumeStreamingAsync(workflow, fromCheckpoint)
851+
: await environment.OpenStreamingAsync(workflow);
877852

878-
InProcessExecutionEnvironment environment = executionEnvironment.ToWorkflowExecutionEnvironment();
879-
await using StreamingRun run = await environment.RunStreamingAsync(workflow, input);
853+
await run.TrySendMessageAsync(input);
880854
await run.TrySendMessageAsync(new TurnToken(emitEvents: true));
881855

856+
StringBuilder sb = new();
882857
WorkflowOutputEvent? output = null;
858+
CheckpointInfo? lastCheckpoint = null;
883859
await foreach (WorkflowEvent evt in run.WatchStreamAsync().ConfigureAwait(false))
884860
{
885-
if (evt is AgentResponseUpdateEvent executorComplete)
886-
{
887-
sb.Append(executorComplete.Data);
888-
}
889-
else if (evt is WorkflowOutputEvent e)
861+
switch (evt)
890862
{
891-
output = e;
892-
break;
893-
}
894-
else if (evt is WorkflowErrorEvent errorEvent)
895-
{
896-
Assert.Fail($"Workflow execution failed with error: {errorEvent.Exception}");
863+
case AgentResponseUpdateEvent executorComplete:
864+
sb.Append(executorComplete.Data);
865+
break;
866+
867+
case WorkflowOutputEvent e:
868+
output = e;
869+
break;
870+
871+
case WorkflowErrorEvent errorEvent:
872+
Assert.Fail($"Workflow execution failed with error: {errorEvent.Exception}");
873+
break;
874+
875+
case SuperStepCompletedEvent stepCompleted:
876+
lastCheckpoint = stepCompleted.CompletionInfo?.Checkpoint;
877+
break;
897878
}
898879
}
899880

900-
return (sb.ToString(), output?.As<List<ChatMessage>>());
881+
return new(sb.ToString(), output?.As<List<ChatMessage>>(), lastCheckpoint);
901882
}
902883

884+
private static Task<WorkflowRunResult> RunWorkflowAsync(
885+
Workflow workflow, List<ChatMessage> input, ExecutionEnvironment executionEnvironment = ExecutionEnvironment.InProcess_Lockstep)
886+
=> RunWorkflowCheckpointedAsync(workflow, input, executionEnvironment.ToWorkflowExecutionEnvironment());
887+
903888
private sealed class DoubleEchoAgentWithBarrier(string name, StrongBox<TaskCompletionSource<bool>> barrier, StrongBox<int> remaining) : DoubleEchoAgent(name)
904889
{
905890
protected override async IAsyncEnumerable<AgentResponseUpdate> RunCoreStreamingAsync(

0 commit comments

Comments
 (0)