[BREAKING] .NET: Workflow Off-Thread Execution Mode (#1233)

* Updates to async run loop.

* fix: Workflow Onwership can be release by nonowner

* fix: Incorrect handling of blockOnPending in StreamingRun

Depending on whether we are running in streaming on non-streaming mode, we may be using the StreamingRun in different ways. Unfortunately, the only place we can really know what is the actual state of execution is in the RunEventStream implementations.

This resulted in blocking where blocking was unneeded and occasionally not-blocking when blocking was needed.

The fix is to move the logic of handling this blocking into RunEventStream implementations.

* fix: Fix cleanup on error and end run

This ensures we clean up the background resources correctly.

* fix: Ensure we let the run loop proceed when shutting down

* fix: Add timeout for Input Waiting

* fix: Make the samples properly clean up `Run`s and `StreamingRun`s

* fix: Simplify Declarative Workflow Run disposal pattern

* Also fixes missing .Disposal() in Integration tests

---------

Co-authored-by: Ben Thomas <ben.thomas@microsoft.com>

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

@@ -40,7 +40,7 @@ internal sealed class WorkflowHarness(Workflow workflow, string runId)
     {
         Console.WriteLine("RUNNING WORKFLOW...");
         Checkpointed<StreamingRun> run = await InProcessExecution.StreamAsync(workflow, input, this._checkpointManager, runId);
-        IReadOnlyList<WorkflowEvent> workflowEvents = await MonitorWorkflowRunAsync(run).ToArrayAsync();
+        IReadOnlyList<WorkflowEvent> workflowEvents = await MonitorAndDisposeWorkflowRunAsync(run).ToArrayAsync();
         this.LastCheckpoint = workflowEvents.OfType<SuperStepCompletedEvent>().LastOrDefault()?.CompletionInfo?.Checkpoint;
         return new WorkflowEvents(workflowEvents);
     }
@@ -50,7 +50,7 @@ internal sealed class WorkflowHarness(Workflow workflow, string runId)
         Console.WriteLine("RESUMING WORKFLOW...");
         Assert.NotNull(this.LastCheckpoint);
         Checkpointed<StreamingRun> run = await InProcessExecution.ResumeStreamAsync(workflow, this.LastCheckpoint, this._checkpointManager, runId);
-        IReadOnlyList<WorkflowEvent> workflowEvents = await MonitorWorkflowRunAsync(run, response).ToArrayAsync();
+        IReadOnlyList<WorkflowEvent> workflowEvents = await MonitorAndDisposeWorkflowRunAsync(run, response).ToArrayAsync();
         return new WorkflowEvents(workflowEvents);
     }
 
@@ -75,8 +75,10 @@ internal sealed class WorkflowHarness(Workflow workflow, string runId)
         return new WorkflowHarness(workflow, runId);
     }
 
-    private static async IAsyncEnumerable<WorkflowEvent> MonitorWorkflowRunAsync(Checkpointed<StreamingRun> run, InputResponse? response = null)
+    private static async IAsyncEnumerable<WorkflowEvent> MonitorAndDisposeWorkflowRunAsync(Checkpointed<StreamingRun> run, InputResponse? response = null)
     {
+        await using IAsyncDisposable disposeRun = run;
+
         await foreach (WorkflowEvent workflowEvent in run.Run.WatchStreamAsync().ConfigureAwait(false))
         {
             bool exitLoop = false;
@@ -93,7 +95,6 @@ internal sealed class WorkflowHarness(Workflow workflow, string runId)
                     }
                     else
                     {
-                        await run.Run.EndRunAsync().ConfigureAwait(false);
                         exitLoop = true;
                     }
                     break;

dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/DeclarativeWorkflowTest.cs

@@ -259,7 +259,7 @@ public sealed class DeclarativeWorkflowTest(ITestOutputHelper output) : Workflow
 
         Workflow workflow = DeclarativeWorkflowBuilder.Build<TInput>(yamlReader, workflowContext);
 
-        StreamingRun run = await InProcessExecution.StreamAsync(workflow, workflowInput);
+        await using StreamingRun run = await InProcessExecution.StreamAsync(workflow, workflowInput);
 
         this.WorkflowEvents = run.WatchStreamAsync().ToEnumerable().ToList();
         foreach (WorkflowEvent workflowEvent in this.WorkflowEvents)

dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/ObjectModel/WorkflowActionExecutorTest.cs

@@ -29,7 +29,7 @@ public abstract class WorkflowActionExecutorTest(ITestOutputHelper output) : Wor
         TestWorkflowExecutor workflowExecutor = new();
         WorkflowBuilder workflowBuilder = new(workflowExecutor);
         workflowBuilder.AddEdge(workflowExecutor, executor);
-        StreamingRun run = await InProcessExecution.StreamAsync(workflowBuilder.Build(), this.State);
+        await using StreamingRun run = await InProcessExecution.StreamAsync(workflowBuilder.Build(), this.State);
         WorkflowEvent[] events = await run.WatchStreamAsync().ToArrayAsync();
         Assert.Contains(events, e => e is DeclarativeActionInvokedEvent);
         Assert.Contains(events, e => e is DeclarativeActionCompletedEvent);

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

@@ -385,31 +385,24 @@ public class AgentWorkflowBuilderTests
     {
         StringBuilder sb = new();
 
-        StreamingRun run = await InProcessExecution.StreamAsync(workflow, input);
-        try
-        {
-            await run.TrySendMessageAsync(new TurnToken(emitEvents: true));
+        await using StreamingRun run = await InProcessExecution.StreamAsync(workflow, input);
+        await run.TrySendMessageAsync(new TurnToken(emitEvents: true));
 
-            WorkflowOutputEvent? output = null;
-            await foreach (WorkflowEvent evt in run.WatchStreamAsync().ConfigureAwait(false))
+        WorkflowOutputEvent? output = null;
+        await foreach (WorkflowEvent evt in run.WatchStreamAsync().ConfigureAwait(false))
+        {
+            if (evt is AgentRunUpdateEvent executorComplete)
             {
-                if (evt is AgentRunUpdateEvent executorComplete)
-                {
-                    sb.Append(executorComplete.Data);
-                }
-                else if (evt is WorkflowOutputEvent e)
-                {
-                    output = e;
-                    break;
-                }
+                sb.Append(executorComplete.Data);
             }
+            else if (evt is WorkflowOutputEvent e)
+            {
+                output = e;
+                break;
+            }
+        }
 
-            return (sb.ToString(), output?.As<List<ChatMessage>>());
-        }
-        finally
-        {
-            await run.EndRunAsync();
-        }
+        return (sb.ToString(), output?.As<List<ChatMessage>>());
     }
 
     private sealed class DoubleEchoAgentWithBarrier(string name, StrongBox<TaskCompletionSource<bool>> barrier, StrongBox<int> remaining) : DoubleEchoAgent(name)

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

@@ -0,0 +1,20 @@
+// 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 InProcessExecutionEnvironment GetEnvironment(this ExecutionMode executionMode)
+    {
+        return executionMode switch
+        {
+            ExecutionMode.OffThread => InProcessExecution.OffThread,
+            ExecutionMode.Lockstep => InProcessExecution.Lockstep,
+            ExecutionMode.Subworkflow => throw new NotSupportedException(),
+            _ => throw new InvalidOperationException($"Unknown execution mode {executionMode}")
+        };
+    }
+}

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

@@ -163,9 +163,25 @@ public class InProcessStateTests
                 .AddFanOutEdge(forward, targets: [testExecutor, testExecutor2])
                 .Build();
 
-        var act = async () => await InProcessExecution.RunAsync(workflow, new TurnToken());
+        Run runWithFailure = await InProcessExecution.RunAsync(workflow, new TurnToken());
 
-        var result = await act.Should()
-                              .ThrowAsync("multiple writers to the same shared scope key");
+        bool hadFailure = false;
+        foreach (WorkflowEvent evt in runWithFailure.NewEvents)
+        {
+            if (evt is WorkflowErrorEvent errorEvent)
+            {
+                hadFailure.Should().BeFalse("There can be only one!");
+                hadFailure = true;
+
+                errorEvent.Data.Should().BeOfType<InvalidOperationException>()
+                                        .Subject.Message.Should().Contain("TestKey");
+            }
+        }
+
+        hadFailure.Should().BeTrue();
+
+        //var act = async () => await InProcessExecution.RunAsync(workflow, new TurnToken());
+        //var result = await act.Should()
+        //                      .ThrowAsync("multiple writers to the same shared scope key");
     }
 }

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/01_Simple_Workflow_Sequential.cs

@@ -1,9 +1,12 @@
 // Copyright (c) Microsoft. All rights reserved.
 
+using System;
 using System.IO;
 using System.Linq;
 using System.Threading.Tasks;
+using Microsoft.Agents.AI.Workflows.InProc;
 using Microsoft.Agents.AI.Workflows.Reflection;
+using Microsoft.Agents.AI.Workflows.UnitTests;
 
 namespace Microsoft.Agents.AI.Workflows.Sample;
 
@@ -23,9 +26,11 @@ internal static class Step1EntryPoint
         }
     }
 
-    public static async ValueTask RunAsync(TextWriter writer)
+    public static async ValueTask RunAsync(TextWriter writer, ExecutionMode executionMode)
     {
-        StreamingRun run = await InProcessExecution.StreamAsync(WorkflowInstance, "Hello, World!").ConfigureAwait(false);
+        InProcessExecutionEnvironment env = executionMode.GetEnvironment();
+
+        StreamingRun run = await env.StreamAsync(WorkflowInstance, "Hello, World!").ConfigureAwait(false);
 
         await foreach (WorkflowEvent evt in run.WatchStreamAsync().ConfigureAwait(false))
         {

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/01a_Simple_Workflow_Sequential.cs

@@ -2,16 +2,18 @@
 
 using System.IO;
 using System.Threading.Tasks;
-
+using Microsoft.Agents.AI.Workflows.InProc;
+using Microsoft.Agents.AI.Workflows.UnitTests;
 using static Microsoft.Agents.AI.Workflows.Sample.Step1EntryPoint;
 
 namespace Microsoft.Agents.AI.Workflows.Sample;
 
 internal static class Step1aEntryPoint
 {
-    public static async ValueTask RunAsync(TextWriter writer)
+    public static async ValueTask RunAsync(TextWriter writer, ExecutionMode executionMode)
     {
-        Run run = await InProcessExecution.RunAsync(WorkflowInstance, "Hello, World!").ConfigureAwait(false);
+        InProcessExecutionEnvironment env = executionMode.GetEnvironment();
+        Run run = await env.RunAsync(WorkflowInstance, "Hello, World!").ConfigureAwait(false);
 
         Assert.Equal(RunStatus.Idle, await run.GetStatusAsync());
 

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/02_Simple_Workflow_Condition.cs

@@ -4,7 +4,9 @@ using System;
 using System.IO;
 using System.Linq;
 using System.Threading.Tasks;
+using Microsoft.Agents.AI.Workflows.InProc;
 using Microsoft.Agents.AI.Workflows.Reflection;
+using Microsoft.Agents.AI.Workflows.UnitTests;
 
 namespace Microsoft.Agents.AI.Workflows.Sample;
 
@@ -28,9 +30,10 @@ internal static class Step2EntryPoint
         }
     }
 
-    public static async ValueTask<string> RunAsync(TextWriter writer, string input = "This is a spam message.")
+    public static async ValueTask<string> RunAsync(TextWriter writer, ExecutionMode executionMode, string input = "This is a spam message.")
     {
-        StreamingRun handle = await InProcessExecution.StreamAsync(WorkflowInstance, input).ConfigureAwait(false);
+        InProcessExecutionEnvironment env = executionMode.GetEnvironment();
+        StreamingRun handle = await env.StreamAsync(WorkflowInstance, input).ConfigureAwait(false);
         await foreach (WorkflowEvent evt in handle.WatchStreamAsync().ConfigureAwait(false))
         {
             switch (evt)

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/03_Simple_Workflow_Loop.cs

@@ -4,7 +4,9 @@ using System;
 using System.IO;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.Agents.AI.Workflows.InProc;
 using Microsoft.Agents.AI.Workflows.Reflection;
+using Microsoft.Agents.AI.Workflows.UnitTests;
 
 namespace Microsoft.Agents.AI.Workflows.Sample;
 
@@ -25,9 +27,10 @@ internal static class Step3EntryPoint
         }
     }
 
-    public static async ValueTask<string> RunAsync(TextWriter writer)
+    public static async ValueTask<string> RunAsync(TextWriter writer, ExecutionMode executionMode)
     {
-        StreamingRun run = await InProcessExecution.StreamAsync(WorkflowInstance, NumberSignal.Init).ConfigureAwait(false);
+        InProcessExecutionEnvironment env = executionMode.GetEnvironment();
+        StreamingRun run = await env.StreamAsync(WorkflowInstance, NumberSignal.Init).ConfigureAwait(false);
 
         await foreach (WorkflowEvent evt in run.WatchStreamAsync().ConfigureAwait(false))
         {

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/04_Simple_Workflow_ExternalRequest.cs

@@ -4,6 +4,8 @@ using System;
 using System.Collections.Generic;
 using System.IO;
 using System.Threading.Tasks;
+using Microsoft.Agents.AI.Workflows.InProc;
+using Microsoft.Agents.AI.Workflows.UnitTests;
 
 namespace Microsoft.Agents.AI.Workflows.Sample;
 
@@ -37,13 +39,14 @@ internal static class Step4EntryPoint
         }
     }
 
-    public static async ValueTask<string> RunAsync(TextWriter writer, Func<string, int> userGuessCallback)
+    public static async ValueTask<string> RunAsync(TextWriter writer, Func<string, int> userGuessCallback, ExecutionMode executionMode)
     {
         NumberSignal signal = NumberSignal.Init;
         string? prompt = UpdatePrompt(null, signal);
 
         Workflow workflow = WorkflowInstance;
-        StreamingRun handle = await InProcessExecution.StreamAsync(workflow, NumberSignal.Init).ConfigureAwait(false);
+        InProcessExecutionEnvironment env = executionMode.GetEnvironment();
+        StreamingRun handle = await env.StreamAsync(workflow, NumberSignal.Init).ConfigureAwait(false);
 
         List<ExternalRequest> requests = [];
         await foreach (WorkflowEvent evt in handle.WatchStreamAsync().ConfigureAwait(false))

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

@@ -6,12 +6,14 @@ using System.IO;
 using System.Threading;
 using System.Threading.Tasks;
 using FluentAssertions;
+using Microsoft.Agents.AI.Workflows.InProc;
+using Microsoft.Agents.AI.Workflows.UnitTests;
 
 namespace Microsoft.Agents.AI.Workflows.Sample;
 
 internal static class Step5EntryPoint
 {
-    public static async ValueTask<string> RunAsync(TextWriter writer, Func<string, int> userGuessCallback, bool rehydrateToRestore = false, CheckpointManager? checkpointManager = null)
+    public static async ValueTask<string> RunAsync(TextWriter writer, Func<string, int> userGuessCallback, ExecutionMode executionMode, bool rehydrateToRestore = false, CheckpointManager? checkpointManager = null)
     {
         Dictionary<CheckpointInfo, (NumberSignal signal, string? prompt)> checkpointedOutputs = [];
 
@@ -21,9 +23,11 @@ internal static class Step5EntryPoint
         checkpointManager ??= CheckpointManager.Default;
 
         Workflow workflow = Step4EntryPoint.CreateWorkflowInstance(out JudgeExecutor judge);
+
+        InProcessExecutionEnvironment env = executionMode.GetEnvironment();
         Checkpointed<StreamingRun> checkpointed =
-            await InProcessExecution.StreamAsync(workflow, NumberSignal.Init, checkpointManager)
-                                    .ConfigureAwait(false);
+            await env.StreamAsync(workflow, NumberSignal.Init, checkpointManager)
+                     .ConfigureAwait(false);
 
         List<CheckpointInfo> checkpoints = [];
         CancellationTokenSource cancellationSource = new();
@@ -40,10 +44,10 @@ internal static class Step5EntryPoint
         Console.WriteLine($"Restoring to checkpoint {targetCheckpoint} from run {targetCheckpoint.RunId}");
         if (rehydrateToRestore)
         {
-            await handle.EndRunAsync().ConfigureAwait(false);
+            await handle.DisposeAsync().ConfigureAwait(false);
 
-            checkpointed = await InProcessExecution.ResumeStreamAsync(workflow, targetCheckpoint, checkpointManager, runId: handle.RunId, cancellationToken: CancellationToken.None)
-                                                   .ConfigureAwait(false);
+            checkpointed = await env.ResumeStreamAsync(workflow, targetCheckpoint, checkpointManager, runId: handle.RunId, cancellationToken: CancellationToken.None)
+                                    .ConfigureAwait(false);
             handle = checkpointed.Run;
         }
         else
@@ -112,20 +116,21 @@ internal static class Step5EntryPoint
                         {
                             Console.WriteLine($"*** Max step {maxStep} reached, cancelling.");
                             cancellationSource.Cancel();
+                            return null;
                         }
-                        else
-                        {
-                            Console.WriteLine($"*** Processing {requests.Count} queued requests.");
-                            foreach (ExternalRequest request in requests)
-                            {
-                                ExternalResponse response = ExecuteExternalRequest(request, userGuessCallback, prompt);
-                                Console.WriteLine($"!!! Sending response: {response}");
-                                await handle.SendResponseAsync(response).ConfigureAwait(false);
-                            }
 
-                            requests.Clear();
-                            Console.WriteLine("*** Completed processing requests.");
+                        Console.WriteLine($"*** Processing {requests.Count} queued requests.");
+                        foreach (ExternalRequest request in requests)
+                        {
+                            ExternalResponse response = ExecuteExternalRequest(request, userGuessCallback, prompt);
+                            Console.WriteLine($"!!! Sending response: {response}");
+                            await handle.SendResponseAsync(response).ConfigureAwait(false);
                         }
+
+                        requests.Clear();
+
+                        Console.WriteLine("*** Completed processing requests.");
+
                         break;
 
                     case ExecutorCompletedEvent executorCompleteEvt:

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/06_GroupChat_Workflow.cs

@@ -10,6 +10,8 @@ using System.Text;
 using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.Agents.AI.Workflows.InProc;
+using Microsoft.Agents.AI.Workflows.UnitTests;
 using Microsoft.Extensions.AI;
 
 namespace Microsoft.Agents.AI.Workflows.Sample;
@@ -22,12 +24,13 @@ internal static class Step6EntryPoint
             .AddParticipants(new HelloAgent(), new EchoAgent())
             .Build();
 
-    public static async ValueTask RunAsync(TextWriter writer, int maxSteps = 2)
+    public static async ValueTask RunAsync(TextWriter writer, ExecutionMode executionMode, int maxSteps = 2)
     {
         Workflow workflow = CreateWorkflow(maxSteps);
 
-        StreamingRun run = await InProcessExecution.StreamAsync(workflow, Array.Empty<ChatMessage>())
-                                                   .ConfigureAwait(false);
+        InProcessExecutionEnvironment env = executionMode.GetEnvironment();
+        StreamingRun run = await env.StreamAsync(workflow, Array.Empty<ChatMessage>())
+                                    .ConfigureAwait(false);
         await run.TrySendMessageAsync(new TurnToken(emitEvents: true));
 
         await foreach (WorkflowEvent evt in run.WatchStreamAsync().ConfigureAwait(false))

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/07_GroupChat_Workflow_HostAsAgent.cs

@@ -1,6 +1,5 @@
 // Copyright (c) Microsoft. All rights reserved.
 
-using System;
 using System.Collections.Generic;
 using System.IO;
 using System.Threading.Tasks;
@@ -26,7 +25,6 @@ internal static class Step7EntryPoint
                                    ?? update.AgentId
                                    ?? update.Role.ToString()
                                    ?? ChatRole.Assistant.ToString()}: {update.Text}";
-            Console.WriteLine(updateText);
             writer.WriteLine(updateText);
         }
     }

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/08_Subworkflow_Simple.cs

@@ -7,6 +7,8 @@ using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using FluentAssertions;
+using Microsoft.Agents.AI.Workflows.InProc;
+using Microsoft.Agents.AI.Workflows.UnitTests;
 
 namespace Microsoft.Agents.AI.Workflows.Sample;
 
@@ -26,7 +28,7 @@ internal static class Step8EntryPoint
             "   Spaces   around   text   ",
         ];
 
-    public static async ValueTask<List<TextProcessingResult>> RunAsync(TextWriter writer, List<string> textsToProcess)
+    public static async ValueTask<List<TextProcessingResult>> RunAsync(TextWriter writer, ExecutionMode executionMode, List<string> textsToProcess)
     {
         Func<TextProcessingRequest, IWorkflowContext, CancellationToken, ValueTask> processTextAsyncFunc = ProcessTextAsync;
         ExecutorIsh processText = processTextAsyncFunc.AsExecutor("TextProcessor");
@@ -41,7 +43,8 @@ internal static class Step8EntryPoint
             .AddEdge(textProcessor, orchestrator)
             .Build();
 
-        Run workflowRun = await InProcessExecution.RunAsync(workflow, textsToProcess);
+        InProcessExecutionEnvironment env = executionMode.GetEnvironment();
+        Run workflowRun = await env.RunAsync(workflow, textsToProcess);
 
         RunStatus status = await workflowRun.GetStatusAsync();
         status.Should().Be(RunStatus.Idle);

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/09_Subworkflow_ExternalRequest.cs

@@ -7,6 +7,8 @@ using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using FluentAssertions;
+using Microsoft.Agents.AI.Workflows.InProc;
+using Microsoft.Agents.AI.Workflows.UnitTests;
 
 namespace Microsoft.Agents.AI.Workflows.Sample;
 
@@ -170,12 +172,13 @@ internal static class Step9EntryPoint
                              .Select(request => Part2FinishedResponses[request.Id])
                              .OrderBy(request => request.Id)];
 
-    public static async ValueTask<List<RequestFinished>> RunAsync(TextWriter writer)
+    public static async ValueTask<List<RequestFinished>> RunAsync(TextWriter writer, ExecutionMode executionMode)
     {
         RunStatus runStatus;
         List<RequestFinished> results = [];
 
-        Run workflowRun = await InProcessExecution.RunAsync(WorkflowInstance, RequestsToProcess.ToList());
+        InProcessExecutionEnvironment env = executionMode.GetEnvironment();
+        Run workflowRun = await env.RunAsync(WorkflowInstance, RequestsToProcess.ToList());
 
         RunStatus part1Status = ExpectedResponsesPart2.Length > 0 ? RunStatus.PendingRequests : RunStatus.Idle;
         runStatus = await workflowRun.GetStatusAsync();

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

@@ -13,12 +13,14 @@ namespace Microsoft.Agents.AI.Workflows.UnitTests;
 
 public class SampleSmokeTest
 {
-    [Fact]
-    public async Task Test_RunSample_Step1Async()
+    [Theory]
+    [InlineData(ExecutionMode.Lockstep)]
+    [InlineData(ExecutionMode.OffThread)]
+    internal async Task Test_RunSample_Step1Async(ExecutionMode executionMode)
     {
         using StringWriter writer = new();
 
-        await Step1EntryPoint.RunAsync(writer);
+        await Step1EntryPoint.RunAsync(writer, executionMode);
 
         string result = writer.ToString();
         string[] lines = result.Split([Environment.NewLine], StringSplitOptions.RemoveEmptyEntries);
@@ -31,12 +33,14 @@ public class SampleSmokeTest
         );
     }
 
-    [Fact]
-    public async Task Test_RunSample_Step1aAsync()
+    [Theory]
+    [InlineData(ExecutionMode.Lockstep)]
+    [InlineData(ExecutionMode.OffThread)]
+    internal async Task Test_RunSample_Step1aAsync(ExecutionMode executionMode)
     {
         using StringWriter writer = new();
 
-        await Step1aEntryPoint.RunAsync(writer);
+        await Step1aEntryPoint.RunAsync(writer, executionMode);
 
         string result = writer.ToString();
         string[] lines = result.Split([Environment.NewLine], StringSplitOptions.RemoveEmptyEntries);
@@ -49,32 +53,38 @@ public class SampleSmokeTest
         );
     }
 
-    [Fact]
-    public async Task Test_RunSample_Step2Async()
+    [Theory]
+    [InlineData(ExecutionMode.Lockstep)]
+    [InlineData(ExecutionMode.OffThread)]
+    internal async Task Test_RunSample_Step2Async(ExecutionMode executionMode)
     {
         using StringWriter writer = new();
 
-        string spamResult = await Step2EntryPoint.RunAsync(writer);
+        string spamResult = await Step2EntryPoint.RunAsync(writer, executionMode);
 
         Assert.Equal(RemoveSpamExecutor.ActionResult, spamResult);
 
-        string nonSpamResult = await Step2EntryPoint.RunAsync(writer, "This is a valid message.");
+        string nonSpamResult = await Step2EntryPoint.RunAsync(writer, executionMode, "This is a valid message.");
 
         Assert.Equal(RespondToMessageExecutor.ActionResult, nonSpamResult);
     }
 
-    [Fact]
-    public async Task Test_RunSample_Step3Async()
+    [Theory]
+    [InlineData(ExecutionMode.Lockstep)]
+    [InlineData(ExecutionMode.OffThread)]
+    internal async Task Test_RunSample_Step3Async(ExecutionMode executionMode)
     {
         using StringWriter writer = new();
 
-        string guessResult = await Step3EntryPoint.RunAsync(writer);
+        string guessResult = await Step3EntryPoint.RunAsync(writer, executionMode);
 
         Assert.Equal("Guessed the number: 42", guessResult);
     }
 
-    [Fact]
-    public async Task Test_RunSample_Step4Async()
+    [Theory]
+    [InlineData(ExecutionMode.Lockstep)]
+    [InlineData(ExecutionMode.OffThread)]
+    internal async Task Test_RunSample_Step4Async(ExecutionMode executionMode)
     {
         using StringWriter writer = new();
 
@@ -83,12 +93,14 @@ public class SampleSmokeTest
             ("Your guess was too high. Try again.", 23),
             ("Your guess was too low. Try again.", 42));
 
-        string guessResult = await Step4EntryPoint.RunAsync(writer, userGuessCallback: responder.InvokeNext);
+        string guessResult = await Step4EntryPoint.RunAsync(writer, userGuessCallback: responder.InvokeNext, executionMode);
         Assert.Equal("You guessed correctly! You Win!", guessResult);
     }
 
-    [Fact]
-    public async Task Test_RunSample_Step5Async()
+    [Theory]
+    [InlineData(ExecutionMode.Lockstep)]
+    [InlineData(ExecutionMode.OffThread)]
+    internal async Task Test_RunSample_Step5Async(ExecutionMode executionMode)
     {
         using StringWriter writer = new();
 
@@ -102,12 +114,14 @@ public class SampleSmokeTest
             ("Your guess was too low. Try again.", 42)
          );
 
-        string guessResult = await Step5EntryPoint.RunAsync(writer, userGuessCallback: responder.InvokeNext);
+        string guessResult = await Step5EntryPoint.RunAsync(writer, userGuessCallback: responder.InvokeNext, executionMode);
         Assert.Equal("You guessed correctly! You Win!", guessResult);
     }
 
-    [Fact]
-    public async Task Test_RunSample_Step5aAsync()
+    [Theory]
+    [InlineData(ExecutionMode.Lockstep)]
+    [InlineData(ExecutionMode.OffThread)]
+    internal async Task Test_RunSample_Step5aAsync(ExecutionMode executionMode)
     {
         using StringWriter writer = new();
 
@@ -121,12 +135,14 @@ public class SampleSmokeTest
             ("Your guess was too low. Try again.", 42)
          );
 
-        string guessResult = await Step5EntryPoint.RunAsync(writer, userGuessCallback: responder.InvokeNext, rehydrateToRestore: true);
+        string guessResult = await Step5EntryPoint.RunAsync(writer, userGuessCallback: responder.InvokeNext, executionMode, rehydrateToRestore: true);
         Assert.Equal("You guessed correctly! You Win!", guessResult);
     }
 
-    [Fact]
-    public async Task Test_RunSample_Step5bAsync()
+    [Theory]
+    [InlineData(ExecutionMode.Lockstep)]
+    [InlineData(ExecutionMode.OffThread)]
+    internal async Task Test_RunSample_Step5bAsync(ExecutionMode executionMode)
     {
         using StringWriter writer = new();
 
@@ -144,16 +160,18 @@ public class SampleSmokeTest
         options.MakeReadOnly();
 
         CheckpointManager memoryJsonManager = CheckpointManager.CreateJson(new InMemoryJsonStore(), options);
-        string guessResult = await Step5EntryPoint.RunAsync(writer, userGuessCallback: responder.InvokeNext, rehydrateToRestore: true, checkpointManager: memoryJsonManager);
+        string guessResult = await Step5EntryPoint.RunAsync(writer, userGuessCallback: responder.InvokeNext, executionMode, rehydrateToRestore: true, checkpointManager: memoryJsonManager);
         Assert.Equal("You guessed correctly! You Win!", guessResult);
     }
 
-    [Fact]
-    public async Task Test_RunSample_Step6Async()
+    [Theory]
+    [InlineData(ExecutionMode.Lockstep)]
+    [InlineData(ExecutionMode.OffThread)]
+    internal async Task Test_RunSample_Step6Async(ExecutionMode executionMode)
     {
         using StringWriter writer = new();
 
-        await Step6EntryPoint.RunAsync(writer);
+        await Step6EntryPoint.RunAsync(writer, executionMode);
 
         string result = writer.ToString();
         string[] lines = result.Split([Environment.NewLine], StringSplitOptions.RemoveEmptyEntries);
@@ -180,8 +198,10 @@ public class SampleSmokeTest
         );
     }
 
-    [Fact]
-    public async Task Test_RunSample_Step8Async()
+    [Theory]
+    [InlineData(ExecutionMode.Lockstep)]
+    [InlineData(ExecutionMode.OffThread)]
+    internal async Task Test_RunSample_Step8Async(ExecutionMode executionMode)
     {
         List<string> textsToProcess = [
             "Hello world! This is a simple test.",
@@ -194,7 +214,7 @@ public class SampleSmokeTest
 
         using StringWriter writer = new();
 
-        List<TextProcessingResult> results = await Step8EntryPoint.RunAsync(writer, textsToProcess);
+        List<TextProcessingResult> results = await Step8EntryPoint.RunAsync(writer, executionMode, textsToProcess);
         Assert.Equal(textsToProcess.Count, results.Count);
 
         Assert.Collection(results,
@@ -216,11 +236,13 @@ public class SampleSmokeTest
         }
     }
 
-    [Fact]
-    public async Task Test_RunSample_Step9Async()
+    [Theory]
+    [InlineData(ExecutionMode.Lockstep)]
+    [InlineData(ExecutionMode.OffThread)]
+    internal async Task Test_RunSample_Step9Async(ExecutionMode executionMode)
     {
         using StringWriter writer = new();
-        _ = await Step9EntryPoint.RunAsync(writer);
+        _ = await Step9EntryPoint.RunAsync(writer, executionMode);
     }
 }