[BREAKING] .NET: Decouple Checkpointing from Run/StreamAsync APIs (#4037)

* [BREAKING] refactor: Decouple Checkpointing and Execution APIs

With this change, Checkpointing becomes an property of an IWorkflowExecutionEnvironment. This lets environments that are tightly-coupled to their CheckpointManager avoid needing to present APIs that would not work (e.g. taking in an InMemory CheckpointManager for Durable Tasks, for example)

* refactor: Normalize IsCheckpointingEnabled naming

dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.IntegrationTests/Framework/WorkflowHarness.cs

@@ -45,7 +45,7 @@ internal sealed class WorkflowHarness(Workflow workflow, string runId)
     public async Task<WorkflowEvents> RunWorkflowAsync<TInput>(TInput input, bool useJson = false) where TInput : notnull
     {
         Console.WriteLine("RUNNING WORKFLOW...");
-        Checkpointed<StreamingRun> run = await InProcessExecution.StreamAsync(workflow, input, this.GetCheckpointManager(useJson), runId);
+        StreamingRun run = await InProcessExecution.StreamAsync(workflow, input, this.GetCheckpointManager(useJson), runId);
         IReadOnlyList<WorkflowEvent> workflowEvents = await MonitorAndDisposeWorkflowRunAsync(run).ToArrayAsync();
         this._lastCheckpoint = workflowEvents.OfType<SuperStepCompletedEvent>().LastOrDefault()?.CompletionInfo?.Checkpoint;
         return new WorkflowEvents(workflowEvents);
@@ -55,7 +55,7 @@ internal sealed class WorkflowHarness(Workflow workflow, string runId)
     {
         Console.WriteLine("\nRESUMING WORKFLOW...");
         Assert.NotNull(this._lastCheckpoint);
-        Checkpointed<StreamingRun> run = await InProcessExecution.ResumeStreamAsync(workflow, this._lastCheckpoint, this.GetCheckpointManager());
+        StreamingRun run = await InProcessExecution.ResumeStreamAsync(workflow, this._lastCheckpoint, this.GetCheckpointManager());
         IReadOnlyList<WorkflowEvent> workflowEvents = await MonitorAndDisposeWorkflowRunAsync(run, response).ToArrayAsync();
         this._lastCheckpoint = workflowEvents.OfType<SuperStepCompletedEvent>().LastOrDefault()?.CompletionInfo?.Checkpoint;
         return new WorkflowEvents(workflowEvents);
@@ -97,19 +97,19 @@ internal sealed class WorkflowHarness(Workflow workflow, string runId)
         return this._checkpointManager;
     }
 
-    private static async IAsyncEnumerable<WorkflowEvent> MonitorAndDisposeWorkflowRunAsync(Checkpointed<StreamingRun> run, ExternalResponse? response = null)
+    private static async IAsyncEnumerable<WorkflowEvent> MonitorAndDisposeWorkflowRunAsync(StreamingRun run, ExternalResponse? response = null)
     {
         await using IAsyncDisposable disposeRun = run;
 
         if (response is not null)
         {
-            await run.Run.SendResponseAsync(response).ConfigureAwait(false);
+            await run.SendResponseAsync(response).ConfigureAwait(false);
         }
 
         bool exitLoop = false;
         bool hasRequest = false;
 
-        await foreach (WorkflowEvent workflowEvent in run.Run.WatchStreamAsync().ConfigureAwait(false))
+        await foreach (WorkflowEvent workflowEvent in run.WatchStreamAsync().ConfigureAwait(false))
         {
             switch (workflowEvent)
             {

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

@@ -9,6 +9,7 @@ using System.Text.Json;
 using System.Text.RegularExpressions;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.Agents.AI.Workflows.InProc;
 using Microsoft.Extensions.AI;
 
 #pragma warning disable SYSLIB1045 // Use GeneratedRegex
@@ -389,7 +390,7 @@ public class AgentWorkflowBuilderTests
     {
         StringBuilder sb = new();
 
-        IWorkflowExecutionEnvironment environment = executionEnvironment.ToWorkflowExecutionEnvironment();
+        InProcessExecutionEnvironment environment = executionEnvironment.ToWorkflowExecutionEnvironment();
         await using StreamingRun run = await environment.StreamAsync(workflow, input);
         await run.TrySendMessageAsync(new TurnToken(emitEvents: true));
 

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

@@ -5,6 +5,7 @@ using System.Threading;
 using System.Threading.Tasks;
 using FluentAssertions;
 using Microsoft.Agents.AI.Workflows.Checkpointing;
+using Microsoft.Agents.AI.Workflows.InProc;
 
 namespace Microsoft.Agents.AI.Workflows.UnitTests;
 
@@ -28,14 +29,14 @@ public class CheckpointParentTests
             .Build();
 
         CheckpointManager checkpointManager = CheckpointManager.CreateInMemory();
-        IWorkflowExecutionEnvironment env = environment.ToWorkflowExecutionEnvironment();
+        InProcessExecutionEnvironment env = environment.ToWorkflowExecutionEnvironment();
 
         // Act
-        Checkpointed<StreamingRun> checkpointed =
-            await env.StreamAsync(workflow, "Hello", checkpointManager);
+        StreamingRun run =
+            await env.WithCheckpointing(checkpointManager).StreamAsync(workflow, "Hello");
 
         List<CheckpointInfo> checkpoints = [];
-        await foreach (WorkflowEvent evt in checkpointed.Run.WatchStreamAsync())
+        await foreach (WorkflowEvent evt in run.WatchStreamAsync())
         {
             if (evt is SuperStepCompletedEvent stepEvt && stepEvt.CompletionInfo?.Checkpoint is { } cp)
             {
@@ -68,16 +69,15 @@ public class CheckpointParentTests
             .Build();
 
         CheckpointManager checkpointManager = CheckpointManager.CreateInMemory();
-        IWorkflowExecutionEnvironment env = environment.ToWorkflowExecutionEnvironment();
+        InProcessExecutionEnvironment env = environment.ToWorkflowExecutionEnvironment();
 
         // Act
-        await using Checkpointed<StreamingRun> checkpointed =
-            await env.StreamAsync(workflow, "Hello", checkpointManager);
+        await using StreamingRun run = await env.WithCheckpointing(checkpointManager).StreamAsync(workflow, "Hello");
 
         List<CheckpointInfo> checkpoints = [];
         using CancellationTokenSource cts = new();
 
-        await foreach (WorkflowEvent evt in checkpointed.Run.WatchStreamAsync(cts.Token))
+        await foreach (WorkflowEvent evt in run.WatchStreamAsync(cts.Token))
         {
             if (evt is SuperStepCompletedEvent stepEvt && stepEvt.CompletionInfo?.Checkpoint is { } cp)
             {
@@ -123,15 +123,14 @@ public class CheckpointParentTests
             .Build();
 
         CheckpointManager checkpointManager = CheckpointManager.CreateInMemory();
-        IWorkflowExecutionEnvironment env = environment.ToWorkflowExecutionEnvironment();
+        InProcessExecutionEnvironment env = environment.ToWorkflowExecutionEnvironment();
 
         // First run: collect a checkpoint to resume from
-        await using Checkpointed<StreamingRun> checkpointed =
-            await env.StreamAsync(workflow, "Hello", checkpointManager);
+        await using StreamingRun run = await env.WithCheckpointing(checkpointManager).StreamAsync(workflow, "Hello");
 
         List<CheckpointInfo> firstRunCheckpoints = [];
         using CancellationTokenSource cts = new();
-        await foreach (WorkflowEvent evt in checkpointed.Run.WatchStreamAsync(cts.Token))
+        await foreach (WorkflowEvent evt in run.WatchStreamAsync(cts.Token))
         {
             if (evt is SuperStepCompletedEvent stepEvt && stepEvt.CompletionInfo?.Checkpoint is { } cp)
             {
@@ -147,15 +146,14 @@ public class CheckpointParentTests
         CheckpointInfo resumePoint = firstRunCheckpoints[0];
 
         // Dispose the first run to release workflow ownership before resuming.
-        await checkpointed.DisposeAsync();
+        await run.DisposeAsync();
 
         // Act: Resume from the first checkpoint
-        Checkpointed<StreamingRun> resumed =
-            await env.ResumeStreamAsync(workflow, resumePoint, checkpointManager);
+        StreamingRun resumed = await env.WithCheckpointing(checkpointManager).ResumeStreamAsync(workflow, resumePoint);
 
         List<CheckpointInfo> resumedCheckpoints = [];
         using CancellationTokenSource cts2 = new();
-        await foreach (WorkflowEvent evt in resumed.Run.WatchStreamAsync(cts2.Token))
+        await foreach (WorkflowEvent evt in resumed.WatchStreamAsync(cts2.Token))
         {
             if (evt is SuperStepCompletedEvent stepEvt && stepEvt.CompletionInfo?.Checkpoint is { } cp)
             {

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/ExecutionExtensions.cs

@@ -1,12 +1,13 @@
 // Copyright (c) Microsoft. All rights reserved.
 
 using System;
+using Microsoft.Agents.AI.Workflows.InProc;
 
 namespace Microsoft.Agents.AI.Workflows.UnitTests;
 
 internal static class ExecutionExtensions
 {
-    public static IWorkflowExecutionEnvironment ToWorkflowExecutionEnvironment(this ExecutionEnvironment environment)
+    public static InProcessExecutionEnvironment ToWorkflowExecutionEnvironment(this ExecutionEnvironment environment)
     {
         return environment switch
         {

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/InProcessStateTests.cs

@@ -131,11 +131,11 @@ public partial class InProcessStateTests
                 .AddEdge(writer, validator, MaxTurns(4))
                 .AddEdge(validator, writer, MaxTurns(4)).Build();
 
-        Checkpointed<Run> checkpointed = await InProcessExecution.RunAsync<TurnToken>(workflow, new(), CheckpointManager.Default);
+        Run checkpointed = await InProcessExecution.RunAsync<TurnToken>(workflow, new(), CheckpointManager.Default);
 
         checkpointed.Checkpoints.Should().HaveCount(4);
 
-        RunStatus status = await checkpointed.Run.GetStatusAsync();
+        RunStatus status = await checkpointed.GetStatusAsync();
         status.Should().Be(RunStatus.Idle);
 
         writer.Completed.Should().BeTrue();

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/05_Simple_Workflow_Checkpointing.cs

@@ -6,12 +6,13 @@ using System.IO;
 using System.Threading;
 using System.Threading.Tasks;
 using FluentAssertions;
+using Microsoft.Agents.AI.Workflows.InProc;
 
 namespace Microsoft.Agents.AI.Workflows.Sample;
 
 internal static class Step5EntryPoint
 {
-    public static async ValueTask<string> RunAsync(TextWriter writer, Func<string, int> userGuessCallback, IWorkflowExecutionEnvironment environment, bool rehydrateToRestore = false, CheckpointManager? checkpointManager = null)
+    public static async ValueTask<string> RunAsync(TextWriter writer, Func<string, int> userGuessCallback, InProcessExecutionEnvironment environment, bool rehydrateToRestore = false, CheckpointManager? checkpointManager = null)
     {
         Dictionary<CheckpointInfo, (NumberSignal signal, string? prompt)> checkpointedOutputs = [];
 
@@ -22,14 +23,14 @@ internal static class Step5EntryPoint
 
         Workflow workflow = Step4EntryPoint.CreateWorkflowInstance(out JudgeExecutor judge);
 
-        Checkpointed<StreamingRun> checkpointed =
-            await environment.StreamAsync(workflow, NumberSignal.Init, checkpointManager)
+        StreamingRun handle =
+            await environment.WithCheckpointing(checkpointManager)
+                             .StreamAsync(workflow, NumberSignal.Init)
                              .ConfigureAwait(false);
 
         List<CheckpointInfo> checkpoints = [];
         CancellationTokenSource cancellationSource = new();
 
-        StreamingRun handle = checkpointed.Run;
         string? result = await RunStreamToHaltOrMaxStepAsync(maxStep: 6).ConfigureAwait(false);
 
         result.Should().BeNull();
@@ -42,13 +43,13 @@ internal static class Step5EntryPoint
         {
             await handle.DisposeAsync().ConfigureAwait(false);
 
-            checkpointed = await environment.ResumeStreamAsync(workflow, targetCheckpoint, checkpointManager, cancellationToken: CancellationToken.None)
-                                            .ConfigureAwait(false);
-            handle = checkpointed.Run;
+            handle = await environment.WithCheckpointing(checkpointManager)
+                                      .ResumeStreamAsync(workflow, targetCheckpoint, CancellationToken.None)
+                                      .ConfigureAwait(false);
         }
         else
         {
-            await checkpointed.RestoreCheckpointAsync(checkpoints[2], CancellationToken.None).ConfigureAwait(false);
+            await handle.RestoreCheckpointAsync(checkpoints[2], CancellationToken.None).ConfigureAwait(false);
         }
 
         (signal, prompt) = checkpointedOutputs[targetCheckpoint];

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/13_Subworkflow_Checkpointing.cs

@@ -48,10 +48,9 @@ internal static class Step13EntryPoint
         return session;
     }
 
-    public static async ValueTask<CheckpointInfo> RunAsync(TextWriter writer, string input, IWorkflowExecutionEnvironment environment, CheckpointManager checkpointManager, CheckpointInfo? resumeFrom)
+    public static async ValueTask<CheckpointInfo> RunAsync(TextWriter writer, string input, IWorkflowExecutionEnvironment environment, CheckpointInfo? resumeFrom)
     {
-        await using Checkpointed<StreamingRun> checkpointed = await BeginAsync();
-        StreamingRun run = checkpointed.Run;
+        await using StreamingRun run = await BeginAsync();
 
         await run.TrySendMessageAsync(new TurnToken());
 
@@ -80,16 +79,16 @@ internal static class Step13EntryPoint
 
         return lastCheckpoint!;
 
-        async ValueTask<Checkpointed<StreamingRun>> BeginAsync()
+        async ValueTask<StreamingRun> BeginAsync()
         {
             if (resumeFrom == null)
             {
-                return await environment.StreamAsync(WorkflowInstance, input, checkpointManager);
+                return await environment.StreamAsync(WorkflowInstance, input);
             }
 
-            Checkpointed<StreamingRun> checkpointed = await environment.ResumeStreamAsync(WorkflowInstance, resumeFrom, checkpointManager);
-            await checkpointed.Run.TrySendMessageAsync(input);
-            return checkpointed;
+            StreamingRun run = await environment.ResumeStreamAsync(WorkflowInstance, resumeFrom);
+            await run.TrySendMessageAsync(input);
+            return run;
         }
     }
 }

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/SampleSmokeTest.cs

@@ -7,6 +7,7 @@ using System.Linq;
 using System.Text.Json;
 using System.Threading.Tasks;
 using FluentAssertions;
+using Microsoft.Agents.AI.Workflows.InProc;
 using Microsoft.Agents.AI.Workflows.Sample;
 
 namespace Microsoft.Agents.AI.Workflows.UnitTests;
@@ -378,9 +379,9 @@ public class SampleSmokeTest
     [InlineData(ExecutionEnvironment.InProcess_Concurrent)]
     internal async Task Test_RunSample_Step13Async(ExecutionEnvironment environment)
     {
-        IWorkflowExecutionEnvironment executionEnvironment = environment.ToWorkflowExecutionEnvironment();
-
         CheckpointManager checkpointManager = CheckpointManager.CreateInMemory();
+        InProcessExecutionEnvironment executionEnvironment = environment.ToWorkflowExecutionEnvironment().WithCheckpointing(checkpointManager);
+
         CheckpointInfo? resumeFrom = null;
 
         await RunAndValidateAsync(1);
@@ -393,7 +394,7 @@ public class SampleSmokeTest
             using StringWriter writer = new();
             string input = $"[{step}] Hello, World!";
 
-            resumeFrom = await Step13EntryPoint.RunAsync(writer, input, executionEnvironment, checkpointManager, resumeFrom);
+            resumeFrom = await Step13EntryPoint.RunAsync(writer, input, executionEnvironment, resumeFrom);
 
             string result = writer.ToString();
             string[] lines = result.Split([Environment.NewLine], StringSplitOptions.RemoveEmptyEntries);

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/TestRunContext.cs

@@ -144,7 +144,7 @@ public class TestRunContext : IRunnerContext
     public Dictionary<string, Executor> Executors { get; set; } = [];
     public string StartingExecutorId { get; set; } = string.Empty;
 
-    public bool WithCheckpointing => false;
+    public bool IsCheckpointingEnabled => false;
     public bool ConcurrentRunsEnabled => false;
 
     WorkflowTelemetryContext IRunnerContext.TelemetryContext => WorkflowTelemetryContext.Disabled;