From 5902bcb10af45a1230e7cc68a6bf175f24b2438b Mon Sep 17 00:00:00 2001 From: Jacob Alber Date: Tue, 7 Oct 2025 20:28:44 -0400 Subject: [PATCH] .NET: [BREAKING] Propagate CancellationToken into Workflow Executors and message handlers (#1280) * feat: Propagate CancellationToken to Executors * Also adds cancellation propagation to `Executor`-accessible APIs * Adds registrators for cancellable handlers to `RouteBuilder` * [BREAKING]: Adds `CancellationToken` to `IMessageHandler.HandleAsync` * test: Re-enable Concurrent Orchestration test * refactor: Delete unused IInputCoordinator * refactor: Remove superfluous argument qualifications --- .../Agents/CustomAgentExecutors/Program.cs | 24 +- .../WorkflowAsAnAgent/WorkflowHelper.cs | 14 +- .../CheckpointAndRehydrate/WorkflowHelper.cs | 24 +- .../CheckpointAndResume/WorkflowHelper.cs | 24 +- .../WorkflowHelper.cs | 12 +- .../Concurrent/Concurrent/Program.cs | 14 +- .../Workflows/Concurrent/MapReduce/Program.cs | 37 +-- .../01_EdgeCondition/Program.cs | 20 +- .../ConditionalEdges/02_SwitchCase/Program.cs | 26 +- .../03_MultiSelection/Program.cs | 42 +-- .../HumanInTheLoopBasic/WorkflowHelper.cs | 8 +- .../GettingStarted/Workflows/Loop/Program.cs | 16 +- .../ApplicationInsights/Program.cs | 9 +- .../Observability/AspireDashboard/Program.cs | 9 +- .../Workflows/SharedStates/Program.cs | 19 +- .../01_ExecutorsAndEdges/Program.cs | 8 +- .../_Foundational/02_Streaming/Program.cs | 8 +- .../Extensions/AgentProviderExtensions.cs | 6 +- .../Extensions/IWorkflowContextExtensions.cs | 30 +-- .../Interpreter/DeclarativeActionExecutor.cs | 8 +- .../Interpreter/DeclarativeWorkflowContext.cs | 45 ++-- .../DeclarativeWorkflowExecutor.cs | 11 +- .../Interpreter/DelegateActionExecutor.cs | 11 +- .../Kit/ActionExecutor.cs | 4 +- .../Kit/IWorkflowContextExtensions.cs | 7 +- .../Kit/RootExecutor.cs | 12 +- .../ObjectModel/ClearAllVariablesExecutor.cs | 2 +- .../ObjectModel/ConditionGroupExecutor.cs | 2 +- .../ObjectModel/CreateConversationExecutor.cs | 2 +- .../ObjectModel/ForeachExecutor.cs | 10 +- .../ObjectModel/QuestionExecutor.cs | 10 +- .../ObjectModel/ResetVariableExecutor.cs | 2 +- .../ObjectModel/SendActivityExecutor.cs | 2 +- .../PowerFx/WorkflowFormulaState.cs | 4 +- .../AgentWorkflowBuilder.cs | 125 +++++---- .../AggregatingExecutor.cs | 6 +- .../AsyncBarrier.cs | 6 +- .../AsyncCoordinator.cs | 7 +- .../ChatProtocolExecutor.cs | 16 +- .../Execution/AsyncRunHandle.cs | 35 ++- .../Execution/AsyncRunHandleExtensions.cs | 20 +- .../Execution/CallResult.cs | 19 +- .../Execution/IInputCoordinator.cs | 11 - .../Execution/IRunEventStream.cs | 4 +- .../Execution/IRunnerContext.cs | 9 +- .../Execution/ISuperStepJoinContext.cs | 6 +- .../Execution/ISuperStepRunner.cs | 8 +- .../Execution/InputWaiter.cs | 6 +- .../Execution/LockstepRunEventStream.cs | 6 +- .../Execution/MessageRouter.cs | 9 +- .../Execution/StreamingRunEventStream.cs | 16 +- .../Microsoft.Agents.AI.Workflows/Executor.cs | 24 +- .../FunctionExecutor.cs | 4 +- .../IWorkflowContext.cs | 86 +++++- .../InProc/InProcessRunner.cs | 47 ++-- .../InProc/InProcessRunnerContext.cs | 48 ++-- .../Reflection/IMessageHandler.cs | 9 +- .../Reflection/MessageHandlerInfo.cs | 29 ++- .../RouteBuilder.cs | 245 ++++++++++++++++-- .../src/Microsoft.Agents.AI.Workflows/Run.cs | 6 +- .../Specialized/AIAgentHostExecutor.cs | 8 +- .../Specialized/RequestInfoExecutor.cs | 17 +- .../Specialized/WorkflowHostExecutor.cs | 30 +-- .../StreamingRun.cs | 6 +- .../ObjectModel/WorkflowActionExecutorTest.cs | 5 +- .../AgentWorkflowBuilderTests.cs | 11 +- .../InProcessStateTests.cs | 4 +- .../ReflectionSmokeTest.cs | 7 +- .../Sample/01_Simple_Workflow_Sequential.cs | 7 +- .../Sample/02_Simple_Workflow_Condition.cs | 15 +- .../Sample/03_Simple_Workflow_Loop.cs | 10 +- .../Sample/08_Subworkflow_Simple.cs | 4 +- .../SpecializedExecutorSmokeTests.cs | 14 +- .../TestRunContext.cs | 40 +-- 74 files changed, 910 insertions(+), 557 deletions(-) delete mode 100644 dotnet/src/Microsoft.Agents.AI.Workflows/Execution/IInputCoordinator.cs diff --git a/dotnet/samples/GettingStarted/Workflows/Agents/CustomAgentExecutors/Program.cs b/dotnet/samples/GettingStarted/Workflows/Agents/CustomAgentExecutors/Program.cs index 600182974c..437453a865 100644 --- a/dotnet/samples/GettingStarted/Workflows/Agents/CustomAgentExecutors/Program.cs +++ b/dotnet/samples/GettingStarted/Workflows/Agents/CustomAgentExecutors/Program.cs @@ -134,17 +134,17 @@ internal sealed class SloganWriterExecutor this._thread = this._agent.GetNewThread(); } - public async ValueTask HandleAsync(string message, IWorkflowContext context) + public async ValueTask HandleAsync(string message, IWorkflowContext context, CancellationToken cancellationToken = default) { - var result = await this._agent.RunAsync(message, this._thread); + var result = await this._agent.RunAsync(message, this._thread, cancellationToken: cancellationToken); var sloganResult = JsonSerializer.Deserialize(result.Text) ?? throw new InvalidOperationException("Failed to deserialize slogan result."); - await context.AddEventAsync(new SloganGeneratedEvent(sloganResult)); + await context.AddEventAsync(new SloganGeneratedEvent(sloganResult), cancellationToken); return sloganResult; } - public async ValueTask HandleAsync(FeedbackResult message, IWorkflowContext context) + public async ValueTask HandleAsync(FeedbackResult message, IWorkflowContext context, CancellationToken cancellationToken = default) { var feedbackMessage = $""" Here is the feedback on your previous slogan: @@ -155,10 +155,10 @@ internal sealed class SloganWriterExecutor Please use this feedback to improve your slogan. """; - var result = await this._agent.RunAsync(feedbackMessage, this._thread); + var result = await this._agent.RunAsync(feedbackMessage, this._thread, cancellationToken: cancellationToken); var sloganResult = JsonSerializer.Deserialize(result.Text) ?? throw new InvalidOperationException("Failed to deserialize slogan result."); - await context.AddEventAsync(new SloganGeneratedEvent(sloganResult)); + await context.AddEventAsync(new SloganGeneratedEvent(sloganResult), cancellationToken); return sloganResult; } } @@ -205,7 +205,7 @@ internal sealed class FeedbackExecutor : ReflectingExecutor, I this._thread = this._agent.GetNewThread(); } - public async ValueTask HandleAsync(SloganResult message, IWorkflowContext context) + public async ValueTask HandleAsync(SloganResult message, IWorkflowContext context, CancellationToken cancellationToken = default) { var sloganMessage = $""" Here is a slogan for the task '{message.Task}': @@ -213,24 +213,24 @@ internal sealed class FeedbackExecutor : ReflectingExecutor, I Please provide feedback on this slogan, including comments, a rating from 1 to 10, and suggested actions for improvement. """; - var response = await this._agent.RunAsync(sloganMessage, this._thread); + var response = await this._agent.RunAsync(sloganMessage, this._thread, cancellationToken: cancellationToken); var feedback = JsonSerializer.Deserialize(response.Text) ?? throw new InvalidOperationException("Failed to deserialize feedback."); - await context.AddEventAsync(new FeedbackEvent(feedback)); + await context.AddEventAsync(new FeedbackEvent(feedback), cancellationToken); if (feedback.Rating >= this.MinimumRating) { - await context.YieldOutputAsync($"The following slogan was accepted:\n\n{message.Slogan}"); + await context.YieldOutputAsync($"The following slogan was accepted:\n\n{message.Slogan}", cancellationToken); return; } if (this._attempts >= this.MaxAttempts) { - await context.YieldOutputAsync($"The slogan was rejected after {this.MaxAttempts} attempts. Final slogan:\n\n{message.Slogan}"); + await context.YieldOutputAsync($"The slogan was rejected after {this.MaxAttempts} attempts. Final slogan:\n\n{message.Slogan}", cancellationToken); return; } - await context.SendMessageAsync(feedback); + await context.SendMessageAsync(feedback, cancellationToken: cancellationToken); this._attempts++; } } diff --git a/dotnet/samples/GettingStarted/Workflows/Agents/WorkflowAsAnAgent/WorkflowHelper.cs b/dotnet/samples/GettingStarted/Workflows/Agents/WorkflowAsAnAgent/WorkflowHelper.cs index 1da5e6932e..3e2261918b 100644 --- a/dotnet/samples/GettingStarted/Workflows/Agents/WorkflowAsAnAgent/WorkflowHelper.cs +++ b/dotnet/samples/GettingStarted/Workflows/Agents/WorkflowAsAnAgent/WorkflowHelper.cs @@ -51,13 +51,15 @@ internal static class WorkflowHelper /// /// The user message to process /// Workflow context for accessing workflow services and adding events - public async ValueTask HandleAsync(List message, IWorkflowContext context) + /// The to monitor for cancellation requests. + /// The default is . + public async ValueTask HandleAsync(List message, IWorkflowContext context, CancellationToken cancellationToken = default) { // Broadcast the message to all connected agents. Receiving agents will queue // the message but will not start processing until they receive a turn token. - await context.SendMessageAsync(message); + await context.SendMessageAsync(message, cancellationToken: cancellationToken); // Broadcast the turn token to kick off the agents. - await context.SendMessageAsync(new TurnToken(emitEvents: true)); + await context.SendMessageAsync(new TurnToken(emitEvents: true), cancellationToken: cancellationToken); } } @@ -75,14 +77,16 @@ internal static class WorkflowHelper /// /// The message from the agent /// Workflow context for accessing workflow services and adding events - public async ValueTask HandleAsync(ChatMessage message, IWorkflowContext context) + /// The to monitor for cancellation requests. + /// The default is . + public async ValueTask HandleAsync(ChatMessage message, IWorkflowContext context, CancellationToken cancellationToken = default) { this._messages.Add(message); if (this._messages.Count == 2) { var formattedMessages = string.Join(Environment.NewLine, this._messages.Select(m => $"{m.Text}")); - await context.YieldOutputAsync(formattedMessages); + await context.YieldOutputAsync(formattedMessages, cancellationToken); } } } diff --git a/dotnet/samples/GettingStarted/Workflows/Checkpoint/CheckpointAndRehydrate/WorkflowHelper.cs b/dotnet/samples/GettingStarted/Workflows/Checkpoint/CheckpointAndRehydrate/WorkflowHelper.cs index f911fa8a05..cb568fca35 100644 --- a/dotnet/samples/GettingStarted/Workflows/Checkpoint/CheckpointAndRehydrate/WorkflowHelper.cs +++ b/dotnet/samples/GettingStarted/Workflows/Checkpoint/CheckpointAndRehydrate/WorkflowHelper.cs @@ -69,20 +69,20 @@ internal sealed class GuessNumberExecutor() : ReflectingExecutor (this.LowerBound + this.UpperBound) / 2; - public async ValueTask HandleAsync(NumberSignal message, IWorkflowContext context) + public async ValueTask HandleAsync(NumberSignal message, IWorkflowContext context, CancellationToken cancellationToken = default) { switch (message) { case NumberSignal.Init: - await context.SendMessageAsync(this.NextGuess).ConfigureAwait(false); + await context.SendMessageAsync(this.NextGuess, cancellationToken: cancellationToken).ConfigureAwait(false); break; case NumberSignal.Above: this.UpperBound = this.NextGuess - 1; - await context.SendMessageAsync(this.NextGuess).ConfigureAwait(false); + await context.SendMessageAsync(this.NextGuess, cancellationToken: cancellationToken).ConfigureAwait(false); break; case NumberSignal.Below: this.LowerBound = this.NextGuess + 1; - await context.SendMessageAsync(this.NextGuess).ConfigureAwait(false); + await context.SendMessageAsync(this.NextGuess, cancellationToken: cancellationToken).ConfigureAwait(false); break; } } @@ -92,14 +92,14 @@ internal sealed class GuessNumberExecutor() : ReflectingExecutor protected override ValueTask OnCheckpointingAsync(IWorkflowContext context, CancellationToken cancellationToken = default) => - context.QueueStateUpdateAsync(StateKey, (this.LowerBound, this.UpperBound)); + context.QueueStateUpdateAsync(StateKey, (this.LowerBound, this.UpperBound), cancellationToken: cancellationToken); /// /// Restore the state of the executor from a checkpoint. /// This must be overridden to restore any state that was saved during checkpointing. /// protected override async ValueTask OnCheckpointRestoredAsync(IWorkflowContext context, CancellationToken cancellationToken = default) => - (this.LowerBound, this.UpperBound) = await context.ReadStateAsync<(int, int)>(StateKey).ConfigureAwait(false); + (this.LowerBound, this.UpperBound) = await context.ReadStateAsync<(int, int)>(StateKey, cancellationToken: cancellationToken).ConfigureAwait(false); } /// @@ -120,20 +120,20 @@ internal sealed class JudgeExecutor() : ReflectingExecutor("Judge this._targetNumber = targetNumber; } - public async ValueTask HandleAsync(int message, IWorkflowContext context) + public async ValueTask HandleAsync(int message, IWorkflowContext context, CancellationToken cancellationToken = default) { this._tries++; if (message == this._targetNumber) { - await context.YieldOutputAsync($"{this._targetNumber} found in {this._tries} tries!").ConfigureAwait(false); + await context.YieldOutputAsync($"{this._targetNumber} found in {this._tries} tries!", cancellationToken: cancellationToken).ConfigureAwait(false); } else if (message < this._targetNumber) { - await context.SendMessageAsync(NumberSignal.Below).ConfigureAwait(false); + await context.SendMessageAsync(NumberSignal.Below, cancellationToken: cancellationToken).ConfigureAwait(false); } else { - await context.SendMessageAsync(NumberSignal.Above).ConfigureAwait(false); + await context.SendMessageAsync(NumberSignal.Above, cancellationToken: cancellationToken).ConfigureAwait(false); } } @@ -142,12 +142,12 @@ internal sealed class JudgeExecutor() : ReflectingExecutor("Judge /// This must be overridden to save any state that is needed to resume the executor. /// protected override ValueTask OnCheckpointingAsync(IWorkflowContext context, CancellationToken cancellationToken = default) => - context.QueueStateUpdateAsync(StateKey, this._tries); + context.QueueStateUpdateAsync(StateKey, this._tries, cancellationToken: cancellationToken); /// /// Restore the state of the executor from a checkpoint. /// This must be overridden to restore any state that was saved during checkpointing. /// protected override async ValueTask OnCheckpointRestoredAsync(IWorkflowContext context, CancellationToken cancellationToken = default) => - this._tries = await context.ReadStateAsync(StateKey).ConfigureAwait(false); + this._tries = await context.ReadStateAsync(StateKey, cancellationToken: cancellationToken).ConfigureAwait(false); } diff --git a/dotnet/samples/GettingStarted/Workflows/Checkpoint/CheckpointAndResume/WorkflowHelper.cs b/dotnet/samples/GettingStarted/Workflows/Checkpoint/CheckpointAndResume/WorkflowHelper.cs index 8bfabfbf4d..ba6268c924 100644 --- a/dotnet/samples/GettingStarted/Workflows/Checkpoint/CheckpointAndResume/WorkflowHelper.cs +++ b/dotnet/samples/GettingStarted/Workflows/Checkpoint/CheckpointAndResume/WorkflowHelper.cs @@ -69,20 +69,20 @@ internal sealed class GuessNumberExecutor() : ReflectingExecutor (this.LowerBound + this.UpperBound) / 2; - public async ValueTask HandleAsync(NumberSignal message, IWorkflowContext context) + public async ValueTask HandleAsync(NumberSignal message, IWorkflowContext context, CancellationToken cancellationToken = default) { switch (message) { case NumberSignal.Init: - await context.SendMessageAsync(this.NextGuess).ConfigureAwait(false); + await context.SendMessageAsync(this.NextGuess, cancellationToken: cancellationToken).ConfigureAwait(false); break; case NumberSignal.Above: this.UpperBound = this.NextGuess - 1; - await context.SendMessageAsync(this.NextGuess).ConfigureAwait(false); + await context.SendMessageAsync(this.NextGuess, cancellationToken: cancellationToken).ConfigureAwait(false); break; case NumberSignal.Below: this.LowerBound = this.NextGuess + 1; - await context.SendMessageAsync(this.NextGuess).ConfigureAwait(false); + await context.SendMessageAsync(this.NextGuess, cancellationToken: cancellationToken).ConfigureAwait(false); break; } } @@ -92,14 +92,14 @@ internal sealed class GuessNumberExecutor() : ReflectingExecutor protected override ValueTask OnCheckpointingAsync(IWorkflowContext context, CancellationToken cancellationToken = default) => - context.QueueStateUpdateAsync(StateKey, (this.LowerBound, this.UpperBound)); + context.QueueStateUpdateAsync(StateKey, (this.LowerBound, this.UpperBound), cancellationToken: cancellationToken); /// /// Restore the state of the executor from a checkpoint. /// This must be overridden to restore any state that was saved during checkpointing. /// protected override async ValueTask OnCheckpointRestoredAsync(IWorkflowContext context, CancellationToken cancellationToken = default) => - (this.LowerBound, this.UpperBound) = await context.ReadStateAsync<(int, int)>(StateKey).ConfigureAwait(false); + (this.LowerBound, this.UpperBound) = await context.ReadStateAsync<(int, int)>(StateKey, cancellationToken: cancellationToken).ConfigureAwait(false); } /// @@ -120,20 +120,20 @@ internal sealed class JudgeExecutor() : ReflectingExecutor("Judge this._targetNumber = targetNumber; } - public async ValueTask HandleAsync(int message, IWorkflowContext context) + public async ValueTask HandleAsync(int message, IWorkflowContext context, CancellationToken cancellationToken = default) { this._tries++; if (message == this._targetNumber) { - await context.YieldOutputAsync($"{this._targetNumber} found in {this._tries} tries!").ConfigureAwait(false); + await context.YieldOutputAsync($"{this._targetNumber} found in {this._tries} tries!", cancellationToken).ConfigureAwait(false); } else if (message < this._targetNumber) { - await context.SendMessageAsync(NumberSignal.Below).ConfigureAwait(false); + await context.SendMessageAsync(NumberSignal.Below, cancellationToken: cancellationToken).ConfigureAwait(false); } else { - await context.SendMessageAsync(NumberSignal.Above).ConfigureAwait(false); + await context.SendMessageAsync(NumberSignal.Above, cancellationToken: cancellationToken).ConfigureAwait(false); } } @@ -142,12 +142,12 @@ internal sealed class JudgeExecutor() : ReflectingExecutor("Judge /// This must be overridden to save any state that is needed to resume the executor. /// protected override ValueTask OnCheckpointingAsync(IWorkflowContext context, CancellationToken cancellationToken = default) => - context.QueueStateUpdateAsync(StateKey, this._tries); + context.QueueStateUpdateAsync(StateKey, this._tries, cancellationToken: cancellationToken); /// /// Restore the state of the executor from a checkpoint. /// This must be overridden to restore any state that was saved during checkpointing. /// protected override async ValueTask OnCheckpointRestoredAsync(IWorkflowContext context, CancellationToken cancellationToken = default) => - this._tries = await context.ReadStateAsync(StateKey).ConfigureAwait(false); + this._tries = await context.ReadStateAsync(StateKey, cancellationToken: cancellationToken).ConfigureAwait(false); } diff --git a/dotnet/samples/GettingStarted/Workflows/Checkpoint/CheckpointWithHumanInTheLoop/WorkflowHelper.cs b/dotnet/samples/GettingStarted/Workflows/Checkpoint/CheckpointWithHumanInTheLoop/WorkflowHelper.cs index 42f77ae387..9fabfee873 100644 --- a/dotnet/samples/GettingStarted/Workflows/Checkpoint/CheckpointWithHumanInTheLoop/WorkflowHelper.cs +++ b/dotnet/samples/GettingStarted/Workflows/Checkpoint/CheckpointWithHumanInTheLoop/WorkflowHelper.cs @@ -69,21 +69,21 @@ internal sealed class JudgeExecutor() : ReflectingExecutor("Judge this._targetNumber = targetNumber; } - public async ValueTask HandleAsync(int message, IWorkflowContext context) + public async ValueTask HandleAsync(int message, IWorkflowContext context, CancellationToken cancellationToken = default) { this._tries++; if (message == this._targetNumber) { - await context.YieldOutputAsync($"{this._targetNumber} found in {this._tries} tries!") + await context.YieldOutputAsync($"{this._targetNumber} found in {this._tries} tries!", cancellationToken) .ConfigureAwait(false); } else if (message < this._targetNumber) { - await context.SendMessageAsync(new SignalWithNumber(NumberSignal.Below, message)).ConfigureAwait(false); + await context.SendMessageAsync(new SignalWithNumber(NumberSignal.Below, message), cancellationToken: cancellationToken).ConfigureAwait(false); } else { - await context.SendMessageAsync(new SignalWithNumber(NumberSignal.Above, message)).ConfigureAwait(false); + await context.SendMessageAsync(new SignalWithNumber(NumberSignal.Above, message), cancellationToken: cancellationToken).ConfigureAwait(false); } } @@ -92,12 +92,12 @@ internal sealed class JudgeExecutor() : ReflectingExecutor("Judge /// This must be overridden to save any state that is needed to resume the executor. /// protected override ValueTask OnCheckpointingAsync(IWorkflowContext context, CancellationToken cancellationToken = default) => - context.QueueStateUpdateAsync(StateKey, this._tries); + context.QueueStateUpdateAsync(StateKey, this._tries, cancellationToken: cancellationToken); /// /// Restore the state of the executor from a checkpoint. /// This must be overridden to restore any state that was saved during checkpointing. /// protected override async ValueTask OnCheckpointRestoredAsync(IWorkflowContext context, CancellationToken cancellationToken = default) => - this._tries = await context.ReadStateAsync(StateKey).ConfigureAwait(false); + this._tries = await context.ReadStateAsync(StateKey, cancellationToken: cancellationToken).ConfigureAwait(false); } diff --git a/dotnet/samples/GettingStarted/Workflows/Concurrent/Concurrent/Program.cs b/dotnet/samples/GettingStarted/Workflows/Concurrent/Concurrent/Program.cs index d43474ece0..44d0f0fe77 100644 --- a/dotnet/samples/GettingStarted/Workflows/Concurrent/Concurrent/Program.cs +++ b/dotnet/samples/GettingStarted/Workflows/Concurrent/Concurrent/Program.cs @@ -82,14 +82,16 @@ internal sealed class ConcurrentStartExecutor() : /// /// The user message to process /// Workflow context for accessing workflow services and adding events + /// The to monitor for cancellation requests. + /// The default is . /// A task representing the asynchronous operation - public async ValueTask HandleAsync(string message, IWorkflowContext context) + public async ValueTask HandleAsync(string message, IWorkflowContext context, CancellationToken cancellationToken = default) { // Broadcast the message to all connected agents. Receiving agents will queue // the message but will not start processing until they receive a turn token. - await context.SendMessageAsync(new ChatMessage(ChatRole.User, message)); + await context.SendMessageAsync(new ChatMessage(ChatRole.User, message), cancellationToken: cancellationToken); // Broadcast the turn token to kick off the agents. - await context.SendMessageAsync(new TurnToken(emitEvents: true)); + await context.SendMessageAsync(new TurnToken(emitEvents: true), cancellationToken: cancellationToken); } } @@ -107,15 +109,17 @@ internal sealed class ConcurrentAggregationExecutor() : /// /// The message from the agent /// Workflow context for accessing workflow services and adding events + /// The to monitor for cancellation requests. + /// The default is . /// A task representing the asynchronous operation - public async ValueTask HandleAsync(ChatMessage message, IWorkflowContext context) + public async ValueTask HandleAsync(ChatMessage message, IWorkflowContext context, CancellationToken cancellationToken = default) { this._messages.Add(message); if (this._messages.Count == 2) { var formattedMessages = string.Join(Environment.NewLine, this._messages.Select(m => $"{m.AuthorName}: {m.Text}")); - await context.YieldOutputAsync(formattedMessages); + await context.YieldOutputAsync(formattedMessages, cancellationToken); } } } diff --git a/dotnet/samples/GettingStarted/Workflows/Concurrent/MapReduce/Program.cs b/dotnet/samples/GettingStarted/Workflows/Concurrent/MapReduce/Program.cs index a0f059a00e..759f602563 100644 --- a/dotnet/samples/GettingStarted/Workflows/Concurrent/MapReduce/Program.cs +++ b/dotnet/samples/GettingStarted/Workflows/Concurrent/MapReduce/Program.cs @@ -5,6 +5,7 @@ using System.Collections.Generic; using System.IO; using System.Linq; using System.Text.Json; +using System.Threading; using System.Threading.Tasks; using Microsoft.Agents.AI.Workflows; using Microsoft.Agents.AI.Workflows.Reflection; @@ -138,7 +139,7 @@ internal sealed class Split(string[] mapperIds, string id) : /// /// Tokenize input and assign contiguous index ranges to each mapper via shared state. /// - public async ValueTask HandleAsync(string message, IWorkflowContext context) + public async ValueTask HandleAsync(string message, IWorkflowContext context, CancellationToken cancellationToken = default) { // Ensure temp directory exists Directory.CreateDirectory(MapReduceConstants.TempDir); @@ -147,7 +148,7 @@ internal sealed class Split(string[] mapperIds, string id) : var wordList = Preprocess(message); // Store the tokenized words once so that all mappers can read by index - await context.QueueStateUpdateAsync(MapReduceConstants.DataToProcessKey, wordList, scopeName: MapReduceConstants.StateScope); + await context.QueueStateUpdateAsync(MapReduceConstants.DataToProcessKey, wordList, scopeName: MapReduceConstants.StateScope, cancellationToken); // Divide indices into contiguous slices for each mapper var mapperCount = this._mapperIds.Length; @@ -160,10 +161,10 @@ internal sealed class Split(string[] mapperIds, string id) : var endIndex = i < mapperCount - 1 ? startIndex + chunkSize : wordList.Length; // Save the indices under the mapper's Id - await context.QueueStateUpdateAsync(this._mapperIds[i], (startIndex, endIndex), scopeName: MapReduceConstants.StateScope); + await context.QueueStateUpdateAsync(this._mapperIds[i], (startIndex, endIndex), scopeName: MapReduceConstants.StateScope, cancellationToken); // Notify the mapper that data is ready - await context.SendMessageAsync(new SplitComplete(), targetId: this._mapperIds[i]); + await context.SendMessageAsync(new SplitComplete(), targetId: this._mapperIds[i], cancellationToken); } // Process all the chunks @@ -192,10 +193,10 @@ internal sealed class Mapper(string id) : ReflectingExecutor(id), IMessa /// /// Read the assigned slice, emit (word, 1) pairs, and persist to disk. /// - public async ValueTask HandleAsync(SplitComplete message, IWorkflowContext context) + public async ValueTask HandleAsync(SplitComplete message, IWorkflowContext context, CancellationToken cancellationToken = default) { - var dataToProcess = await context.ReadStateAsync(MapReduceConstants.DataToProcessKey, scopeName: MapReduceConstants.StateScope); - var chunk = await context.ReadStateAsync<(int start, int end)>(this.Id, scopeName: MapReduceConstants.StateScope); + var dataToProcess = await context.ReadStateAsync(MapReduceConstants.DataToProcessKey, scopeName: MapReduceConstants.StateScope, cancellationToken); + var chunk = await context.ReadStateAsync<(int start, int end)>(this.Id, scopeName: MapReduceConstants.StateScope, cancellationToken); var results = dataToProcess![chunk.start..chunk.end] .Select(word => (word, 1)) @@ -204,9 +205,9 @@ internal sealed class Mapper(string id) : ReflectingExecutor(id), IMessa // Write this mapper's results as simple text lines for easy debugging var filePath = Path.Combine(MapReduceConstants.TempDir, $"map_results_{this.Id}.txt"); var lines = results.Select(r => $"{r.word}: {r.Item2}"); - await File.WriteAllLinesAsync(filePath, lines); + await File.WriteAllLinesAsync(filePath, lines, cancellationToken); - await context.SendMessageAsync(new MapComplete(filePath)); + await context.SendMessageAsync(new MapComplete(filePath), cancellationToken: cancellationToken); } } @@ -224,7 +225,7 @@ internal sealed class Shuffler(string[] reducerIds, string[] mapperIds, string i /// /// Aggregate mapper outputs and write one partition file per reducer. /// - public async ValueTask HandleAsync(MapComplete message, IWorkflowContext context) + public async ValueTask HandleAsync(MapComplete message, IWorkflowContext context, CancellationToken cancellationToken = default) { this._mapResults.Add(message); @@ -241,9 +242,9 @@ internal sealed class Shuffler(string[] reducerIds, string[] mapperIds, string i // Write one grouped partition for reducer index and notify that reducer var filePath = Path.Combine(MapReduceConstants.TempDir, $"shuffle_results_{index}.txt"); var lines = chunk.Select(kvp => $"{kvp.key}: {JsonSerializer.Serialize(kvp.values)}"); - await File.WriteAllLinesAsync(filePath, lines); + await File.WriteAllLinesAsync(filePath, lines, cancellationToken); - await context.SendMessageAsync(new ShuffleComplete(filePath, this._reducerIds[index])); + await context.SendMessageAsync(new ShuffleComplete(filePath, this._reducerIds[index]), cancellationToken: cancellationToken); } var tasks = chunks.Select((chunk, i) => ProcessChunkAsync(chunk, i)); @@ -318,7 +319,7 @@ internal sealed class Reducer(string id) : ReflectingExecutor(id), IMes /// /// Read one shuffle partition and reduce it to totals. /// - public async ValueTask HandleAsync(ShuffleComplete message, IWorkflowContext context) + public async ValueTask HandleAsync(ShuffleComplete message, IWorkflowContext context, CancellationToken cancellationToken = default) { if (message.ReducerId != this.Id) { @@ -327,7 +328,7 @@ internal sealed class Reducer(string id) : ReflectingExecutor(id), IMes } // Read grouped values from the shuffle output - var lines = await File.ReadAllLinesAsync(message.FilePath); + var lines = await File.ReadAllLinesAsync(message.FilePath, cancellationToken); // Sum values per key. Values are serialized JSON arrays like [1, 1, ...] var reducedResults = new Dictionary(); @@ -345,9 +346,9 @@ internal sealed class Reducer(string id) : ReflectingExecutor(id), IMes // Persist our partition totals var filePath = Path.Combine(MapReduceConstants.TempDir, $"reduced_results_{this.Id}.txt"); var outputLines = reducedResults.Select(kvp => $"{kvp.Key}: {kvp.Value}"); - await File.WriteAllLinesAsync(filePath, outputLines); + await File.WriteAllLinesAsync(filePath, outputLines, cancellationToken); - await context.SendMessageAsync(new ReduceComplete(filePath)); + await context.SendMessageAsync(new ReduceComplete(filePath), cancellationToken: cancellationToken); } } @@ -361,10 +362,10 @@ internal sealed class CompletionExecutor(string id) : /// /// Collect reducer output file paths and yield final output. /// - public async ValueTask HandleAsync(List message, IWorkflowContext context) + public async ValueTask HandleAsync(List message, IWorkflowContext context, CancellationToken cancellationToken = default) { var filePaths = message.ConvertAll(r => r.FilePath); - await context.YieldOutputAsync(filePaths); + await context.YieldOutputAsync(filePaths, cancellationToken); } } diff --git a/dotnet/samples/GettingStarted/Workflows/ConditionalEdges/01_EdgeCondition/Program.cs b/dotnet/samples/GettingStarted/Workflows/ConditionalEdges/01_EdgeCondition/Program.cs index 65ac2aa852..6298337d79 100644 --- a/dotnet/samples/GettingStarted/Workflows/ConditionalEdges/01_EdgeCondition/Program.cs +++ b/dotnet/samples/GettingStarted/Workflows/ConditionalEdges/01_EdgeCondition/Program.cs @@ -160,7 +160,7 @@ internal sealed class SpamDetectionExecutor : ReflectingExecutor HandleAsync(ChatMessage message, IWorkflowContext context) + public async ValueTask HandleAsync(ChatMessage message, IWorkflowContext context, CancellationToken cancellationToken = default) { // Generate a random email ID and store the email content to the shared state var newEmail = new Email @@ -168,10 +168,10 @@ internal sealed class SpamDetectionExecutor : ReflectingExecutor(response.Text); detectionResult!.EmailId = newEmail.EmailId; @@ -205,7 +205,7 @@ internal sealed class EmailAssistantExecutor : ReflectingExecutor HandleAsync(DetectionResult message, IWorkflowContext context) + public async ValueTask HandleAsync(DetectionResult message, IWorkflowContext context, CancellationToken cancellationToken = default) { if (message.IsSpam) { @@ -213,11 +213,11 @@ internal sealed class EmailAssistantExecutor : ReflectingExecutor(message.EmailId, scopeName: EmailStateConstants.EmailStateScope) + var email = await context.ReadStateAsync(message.EmailId, scopeName: EmailStateConstants.EmailStateScope, cancellationToken) ?? throw new InvalidOperationException("Email not found."); // Invoke the agent - var response = await this._emailAssistantAgent.RunAsync(email.EmailContent); + var response = await this._emailAssistantAgent.RunAsync(email.EmailContent, cancellationToken: cancellationToken); var emailResponse = JsonSerializer.Deserialize(response.Text); return emailResponse!; @@ -232,8 +232,8 @@ internal sealed class SendEmailExecutor() : ReflectingExecutor /// Simulate the sending of an email. /// - public async ValueTask HandleAsync(EmailResponse message, IWorkflowContext context) => - await context.YieldOutputAsync($"Email sent: {message.Response}"); + public async ValueTask HandleAsync(EmailResponse message, IWorkflowContext context, CancellationToken cancellationToken = default) => + await context.YieldOutputAsync($"Email sent: {message.Response}", cancellationToken); } /// @@ -244,11 +244,11 @@ internal sealed class HandleSpamExecutor() : ReflectingExecutor /// Simulate the handling of a spam message. /// - public async ValueTask HandleAsync(DetectionResult message, IWorkflowContext context) + public async ValueTask HandleAsync(DetectionResult message, IWorkflowContext context, CancellationToken cancellationToken = default) { if (message.IsSpam) { - await context.YieldOutputAsync($"Email marked as spam: {message.Reason}"); + await context.YieldOutputAsync($"Email marked as spam: {message.Reason}", cancellationToken); } else { diff --git a/dotnet/samples/GettingStarted/Workflows/ConditionalEdges/02_SwitchCase/Program.cs b/dotnet/samples/GettingStarted/Workflows/ConditionalEdges/02_SwitchCase/Program.cs index 4f985f47a2..17de5311c5 100644 --- a/dotnet/samples/GettingStarted/Workflows/ConditionalEdges/02_SwitchCase/Program.cs +++ b/dotnet/samples/GettingStarted/Workflows/ConditionalEdges/02_SwitchCase/Program.cs @@ -185,7 +185,7 @@ internal sealed class SpamDetectionExecutor : ReflectingExecutor HandleAsync(ChatMessage message, IWorkflowContext context) + public async ValueTask HandleAsync(ChatMessage message, IWorkflowContext context, CancellationToken cancellationToken = default) { // Generate a random email ID and store the email content var newEmail = new Email @@ -193,10 +193,10 @@ internal sealed class SpamDetectionExecutor : ReflectingExecutor(response.Text); detectionResult!.EmailId = newEmail.EmailId; @@ -230,7 +230,7 @@ internal sealed class EmailAssistantExecutor : ReflectingExecutor HandleAsync(DetectionResult message, IWorkflowContext context) + public async ValueTask HandleAsync(DetectionResult message, IWorkflowContext context, CancellationToken cancellationToken = default) { if (message.spamDecision == SpamDecision.Spam) { @@ -238,10 +238,10 @@ internal sealed class EmailAssistantExecutor : ReflectingExecutor(message.EmailId, scopeName: EmailStateConstants.EmailStateScope); + var email = await context.ReadStateAsync(message.EmailId, scopeName: EmailStateConstants.EmailStateScope, cancellationToken); // Invoke the agent - var response = await this._emailAssistantAgent.RunAsync(email!.EmailContent); + var response = await this._emailAssistantAgent.RunAsync(email!.EmailContent, cancellationToken: cancellationToken); var emailResponse = JsonSerializer.Deserialize(response.Text); return emailResponse!; @@ -256,8 +256,8 @@ internal sealed class SendEmailExecutor() : ReflectingExecutor /// Simulate the sending of an email. /// - public async ValueTask HandleAsync(EmailResponse message, IWorkflowContext context) => - await context.YieldOutputAsync($"Email sent: {message.Response}").ConfigureAwait(false); + public async ValueTask HandleAsync(EmailResponse message, IWorkflowContext context, CancellationToken cancellationToken = default) => + await context.YieldOutputAsync($"Email sent: {message.Response}", cancellationToken).ConfigureAwait(false); } /// @@ -268,11 +268,11 @@ internal sealed class HandleSpamExecutor() : ReflectingExecutor /// Simulate the handling of a spam message. /// - public async ValueTask HandleAsync(DetectionResult message, IWorkflowContext context) + public async ValueTask HandleAsync(DetectionResult message, IWorkflowContext context, CancellationToken cancellationToken = default) { if (message.spamDecision == SpamDecision.Spam) { - await context.YieldOutputAsync($"Email marked as spam: {message.Reason}").ConfigureAwait(false); + await context.YieldOutputAsync($"Email marked as spam: {message.Reason}", cancellationToken).ConfigureAwait(false); } else { @@ -289,12 +289,12 @@ internal sealed class HandleUncertainExecutor() : ReflectingExecutor /// Simulate the handling of an uncertain spam decision. /// - public async ValueTask HandleAsync(DetectionResult message, IWorkflowContext context) + public async ValueTask HandleAsync(DetectionResult message, IWorkflowContext context, CancellationToken cancellationToken = default) { if (message.spamDecision == SpamDecision.Uncertain) { - var email = await context.ReadStateAsync(message.EmailId, scopeName: EmailStateConstants.EmailStateScope); - await context.YieldOutputAsync($"Email marked as uncertain: {message.Reason}. Email content: {email?.EmailContent}"); + var email = await context.ReadStateAsync(message.EmailId, scopeName: EmailStateConstants.EmailStateScope, cancellationToken); + await context.YieldOutputAsync($"Email marked as uncertain: {message.Reason}. Email content: {email?.EmailContent}", cancellationToken); } else { diff --git a/dotnet/samples/GettingStarted/Workflows/ConditionalEdges/03_MultiSelection/Program.cs b/dotnet/samples/GettingStarted/Workflows/ConditionalEdges/03_MultiSelection/Program.cs index 987dd7ffd6..cb3914fa64 100644 --- a/dotnet/samples/GettingStarted/Workflows/ConditionalEdges/03_MultiSelection/Program.cs +++ b/dotnet/samples/GettingStarted/Workflows/ConditionalEdges/03_MultiSelection/Program.cs @@ -241,7 +241,7 @@ internal sealed class EmailAnalysisExecutor : ReflectingExecutor HandleAsync(ChatMessage message, IWorkflowContext context) + public async ValueTask HandleAsync(ChatMessage message, IWorkflowContext context, CancellationToken cancellationToken = default) { // Generate a random email ID and store the email content var newEmail = new Email @@ -249,10 +249,10 @@ internal sealed class EmailAnalysisExecutor : ReflectingExecutor(response.Text); AnalysisResult!.EmailId = newEmail.EmailId; @@ -287,7 +287,7 @@ internal sealed class EmailAssistantExecutor : ReflectingExecutor HandleAsync(AnalysisResult message, IWorkflowContext context) + public async ValueTask HandleAsync(AnalysisResult message, IWorkflowContext context, CancellationToken cancellationToken = default) { if (message.spamDecision == SpamDecision.Spam) { @@ -295,10 +295,10 @@ internal sealed class EmailAssistantExecutor : ReflectingExecutor(message.EmailId, scopeName: EmailStateConstants.EmailStateScope); + var email = await context.ReadStateAsync(message.EmailId, scopeName: EmailStateConstants.EmailStateScope, cancellationToken); // Invoke the agent - var response = await this._emailAssistantAgent.RunAsync(email!.EmailContent); + var response = await this._emailAssistantAgent.RunAsync(email!.EmailContent, cancellationToken: cancellationToken); var emailResponse = JsonSerializer.Deserialize(response.Text); return emailResponse!; @@ -313,8 +313,8 @@ internal sealed class SendEmailExecutor() : ReflectingExecutor /// Simulate the sending of an email. /// - public async ValueTask HandleAsync(EmailResponse message, IWorkflowContext context) => - await context.YieldOutputAsync($"Email sent: {message.Response}"); + public async ValueTask HandleAsync(EmailResponse message, IWorkflowContext context, CancellationToken cancellationToken = default) => + await context.YieldOutputAsync($"Email sent: {message.Response}", cancellationToken); } /// @@ -325,11 +325,11 @@ internal sealed class HandleSpamExecutor() : ReflectingExecutor /// Simulate the handling of a spam message. /// - public async ValueTask HandleAsync(AnalysisResult message, IWorkflowContext context) + public async ValueTask HandleAsync(AnalysisResult message, IWorkflowContext context, CancellationToken cancellationToken = default) { if (message.spamDecision == SpamDecision.Spam) { - await context.YieldOutputAsync($"Email marked as spam: {message.Reason}"); + await context.YieldOutputAsync($"Email marked as spam: {message.Reason}", cancellationToken); } else { @@ -346,12 +346,12 @@ internal sealed class HandleUncertainExecutor() : ReflectingExecutor /// Simulate the handling of an uncertain spam decision. /// - public async ValueTask HandleAsync(AnalysisResult message, IWorkflowContext context) + public async ValueTask HandleAsync(AnalysisResult message, IWorkflowContext context, CancellationToken cancellationToken = default) { if (message.spamDecision == SpamDecision.Uncertain) { - var email = await context.ReadStateAsync(message.EmailId, scopeName: EmailStateConstants.EmailStateScope); - await context.YieldOutputAsync($"Email marked as uncertain: {message.Reason}. Email content: {email?.EmailContent}"); + var email = await context.ReadStateAsync(message.EmailId, scopeName: EmailStateConstants.EmailStateScope, cancellationToken); + await context.YieldOutputAsync($"Email marked as uncertain: {message.Reason}. Email content: {email?.EmailContent}", cancellationToken); } else { @@ -385,13 +385,13 @@ internal sealed class EmailSummaryExecutor : ReflectingExecutor HandleAsync(AnalysisResult message, IWorkflowContext context) + public async ValueTask HandleAsync(AnalysisResult message, IWorkflowContext context, CancellationToken cancellationToken = default) { // Read the email content from the shared states - var email = await context.ReadStateAsync(message.EmailId, scopeName: EmailStateConstants.EmailStateScope); + var email = await context.ReadStateAsync(message.EmailId, scopeName: EmailStateConstants.EmailStateScope, cancellationToken); // Invoke the agent - var response = await this._emailSummaryAgent.RunAsync(email!.EmailContent); + var response = await this._emailSummaryAgent.RunAsync(email!.EmailContent, cancellationToken: cancellationToken); var emailSummary = JsonSerializer.Deserialize(response.Text); message.EmailSummary = emailSummary!.Summary; @@ -410,17 +410,17 @@ internal sealed class DatabaseEvent(string message) : WorkflowEvent(message) { } /// internal sealed class DatabaseAccessExecutor() : ReflectingExecutor("DatabaseAccessExecutor"), IMessageHandler { - public async ValueTask HandleAsync(AnalysisResult message, IWorkflowContext context) + public async ValueTask HandleAsync(AnalysisResult message, IWorkflowContext context, CancellationToken cancellationToken = default) { // 1. Save the email content - await context.ReadStateAsync(message.EmailId, scopeName: EmailStateConstants.EmailStateScope); - await Task.Delay(100); // Simulate database access delay + await context.ReadStateAsync(message.EmailId, scopeName: EmailStateConstants.EmailStateScope, cancellationToken); + await Task.Delay(100, cancellationToken); // Simulate database access delay // 2. Save the analysis result - await Task.Delay(100); // Simulate database access delay + await Task.Delay(100, cancellationToken); // Simulate database access delay // Not using the `WorkflowCompletedEvent` because this is not the end of the workflow. // The end of the workflow is signaled by the `SendEmailExecutor` or the `HandleUnknownExecutor`. - await context.AddEventAsync(new DatabaseEvent($"Email {message.EmailId} saved to database.")); + await context.AddEventAsync(new DatabaseEvent($"Email {message.EmailId} saved to database."), cancellationToken); } } diff --git a/dotnet/samples/GettingStarted/Workflows/HumanInTheLoop/HumanInTheLoopBasic/WorkflowHelper.cs b/dotnet/samples/GettingStarted/Workflows/HumanInTheLoop/HumanInTheLoopBasic/WorkflowHelper.cs index 8f0f8c7b35..5fcc6cf765 100644 --- a/dotnet/samples/GettingStarted/Workflows/HumanInTheLoop/HumanInTheLoopBasic/WorkflowHelper.cs +++ b/dotnet/samples/GettingStarted/Workflows/HumanInTheLoop/HumanInTheLoopBasic/WorkflowHelper.cs @@ -53,21 +53,21 @@ internal sealed class JudgeExecutor() : ReflectingExecutor("Judge this._targetNumber = targetNumber; } - public async ValueTask HandleAsync(int message, IWorkflowContext context) + public async ValueTask HandleAsync(int message, IWorkflowContext context, CancellationToken cancellationToken = default) { this._tries++; if (message == this._targetNumber) { - await context.YieldOutputAsync($"{this._targetNumber} found in {this._tries} tries!") + await context.YieldOutputAsync($"{this._targetNumber} found in {this._tries} tries!", cancellationToken) .ConfigureAwait(false); } else if (message < this._targetNumber) { - await context.SendMessageAsync(NumberSignal.Below).ConfigureAwait(false); + await context.SendMessageAsync(NumberSignal.Below, cancellationToken: cancellationToken).ConfigureAwait(false); } else { - await context.SendMessageAsync(NumberSignal.Above).ConfigureAwait(false); + await context.SendMessageAsync(NumberSignal.Above, cancellationToken: cancellationToken).ConfigureAwait(false); } } } diff --git a/dotnet/samples/GettingStarted/Workflows/Loop/Program.cs b/dotnet/samples/GettingStarted/Workflows/Loop/Program.cs index 020c193ca0..a5f4168767 100644 --- a/dotnet/samples/GettingStarted/Workflows/Loop/Program.cs +++ b/dotnet/samples/GettingStarted/Workflows/Loop/Program.cs @@ -83,20 +83,20 @@ internal sealed class GuessNumberExecutor : ReflectingExecutor (this.LowerBound + this.UpperBound) / 2; - public async ValueTask HandleAsync(NumberSignal message, IWorkflowContext context) + public async ValueTask HandleAsync(NumberSignal message, IWorkflowContext context, CancellationToken cancellationToken = default) { switch (message) { case NumberSignal.Init: - await context.SendMessageAsync(this.NextGuess).ConfigureAwait(false); + await context.SendMessageAsync(this.NextGuess, cancellationToken: cancellationToken).ConfigureAwait(false); break; case NumberSignal.Above: this.UpperBound = this.NextGuess - 1; - await context.SendMessageAsync(this.NextGuess).ConfigureAwait(false); + await context.SendMessageAsync(this.NextGuess, cancellationToken: cancellationToken).ConfigureAwait(false); break; case NumberSignal.Below: this.LowerBound = this.NextGuess + 1; - await context.SendMessageAsync(this.NextGuess).ConfigureAwait(false); + await context.SendMessageAsync(this.NextGuess, cancellationToken: cancellationToken).ConfigureAwait(false); break; } } @@ -120,21 +120,21 @@ internal sealed class JudgeExecutor : ReflectingExecutor, IMessag this._targetNumber = targetNumber; } - public async ValueTask HandleAsync(int message, IWorkflowContext context) + public async ValueTask HandleAsync(int message, IWorkflowContext context, CancellationToken cancellationToken = default) { this._tries++; if (message == this._targetNumber) { - await context.YieldOutputAsync($"{this._targetNumber} found in {this._tries} tries!") + await context.YieldOutputAsync($"{this._targetNumber} found in {this._tries} tries!", cancellationToken) .ConfigureAwait(false); } else if (message < this._targetNumber) { - await context.SendMessageAsync(NumberSignal.Below).ConfigureAwait(false); + await context.SendMessageAsync(NumberSignal.Below, cancellationToken: cancellationToken).ConfigureAwait(false); } else { - await context.SendMessageAsync(NumberSignal.Above).ConfigureAwait(false); + await context.SendMessageAsync(NumberSignal.Above, cancellationToken: cancellationToken).ConfigureAwait(false); } } } diff --git a/dotnet/samples/GettingStarted/Workflows/Observability/ApplicationInsights/Program.cs b/dotnet/samples/GettingStarted/Workflows/Observability/ApplicationInsights/Program.cs index 0ff7a50e8b..6c3debd003 100644 --- a/dotnet/samples/GettingStarted/Workflows/Observability/ApplicationInsights/Program.cs +++ b/dotnet/samples/GettingStarted/Workflows/Observability/ApplicationInsights/Program.cs @@ -76,8 +76,10 @@ internal sealed class UppercaseExecutor() : ReflectingExecutor /// The input text to convert /// Workflow context for accessing workflow services and adding events + /// The to monitor for cancellation requests. + /// The default is . /// The input text converted to uppercase - public async ValueTask HandleAsync(string message, IWorkflowContext context) => + public async ValueTask HandleAsync(string message, IWorkflowContext context, CancellationToken cancellationToken = default) => message.ToUpperInvariant(); // The return value will be sent as a message along an edge to subsequent executors } @@ -91,6 +93,9 @@ internal sealed class ReverseTextExecutor() : ReflectingExecutor /// The input text to reverse /// Workflow context for accessing workflow services and adding events + /// The to monitor for cancellation requests. + /// The default is . /// The input text reversed - public async ValueTask HandleAsync(string message, IWorkflowContext context) => new string(message.Reverse().ToArray()); + public async ValueTask HandleAsync(string message, IWorkflowContext context, CancellationToken cancellationToken = default) + => new(message.Reverse().ToArray()); } diff --git a/dotnet/samples/GettingStarted/Workflows/Observability/AspireDashboard/Program.cs b/dotnet/samples/GettingStarted/Workflows/Observability/AspireDashboard/Program.cs index c02d6bc4a0..ddd88e7709 100644 --- a/dotnet/samples/GettingStarted/Workflows/Observability/AspireDashboard/Program.cs +++ b/dotnet/samples/GettingStarted/Workflows/Observability/AspireDashboard/Program.cs @@ -78,8 +78,10 @@ internal sealed class UppercaseExecutor() : ReflectingExecutor /// The input text to convert /// Workflow context for accessing workflow services and adding events + /// The to monitor for cancellation requests. + /// The default is . /// The input text converted to uppercase - public async ValueTask HandleAsync(string message, IWorkflowContext context) => + public async ValueTask HandleAsync(string message, IWorkflowContext context, CancellationToken cancellationToken = default) => message.ToUpperInvariant(); // The return value will be sent as a message along an edge to subsequent executors } @@ -93,6 +95,9 @@ internal sealed class ReverseTextExecutor() : ReflectingExecutor /// The input text to reverse /// Workflow context for accessing workflow services and adding events + /// The to monitor for cancellation requests. + /// The default is . /// The input text reversed - public async ValueTask HandleAsync(string message, IWorkflowContext context) => new string(message.Reverse().ToArray()); + public async ValueTask HandleAsync(string message, IWorkflowContext context, CancellationToken cancellationToken = default) + => new(message.Reverse().ToArray()); } diff --git a/dotnet/samples/GettingStarted/Workflows/SharedStates/Program.cs b/dotnet/samples/GettingStarted/Workflows/SharedStates/Program.cs index 8812d2689e..246bb776c9 100644 --- a/dotnet/samples/GettingStarted/Workflows/SharedStates/Program.cs +++ b/dotnet/samples/GettingStarted/Workflows/SharedStates/Program.cs @@ -54,13 +54,13 @@ internal static class FileContentStateConstants internal sealed class FileReadExecutor() : ReflectingExecutor("FileReadExecutor"), IMessageHandler { - public async ValueTask HandleAsync(string message, IWorkflowContext context) + public async ValueTask HandleAsync(string message, IWorkflowContext context, CancellationToken cancellationToken = default) { // Read file content from embedded resource string fileContent = Resources.Read(message); // Store file content in a shared state for access by other executors string fileID = Guid.NewGuid().ToString("N"); - await context.QueueStateUpdateAsync(fileID, fileContent, scopeName: FileContentStateConstants.FileContentStateScope); + await context.QueueStateUpdateAsync(fileID, fileContent, scopeName: FileContentStateConstants.FileContentStateScope, cancellationToken); return fileID; } @@ -74,10 +74,10 @@ internal sealed class FileStats internal sealed class WordCountingExecutor() : ReflectingExecutor("WordCountingExecutor"), IMessageHandler { - public async ValueTask HandleAsync(string message, IWorkflowContext context) + public async ValueTask HandleAsync(string message, IWorkflowContext context, CancellationToken cancellationToken = default) { // Retrieve the file content from the shared state - var fileContent = await context.ReadStateAsync(message, scopeName: FileContentStateConstants.FileContentStateScope) + var fileContent = await context.ReadStateAsync(message, scopeName: FileContentStateConstants.FileContentStateScope, cancellationToken) ?? throw new InvalidOperationException("File content state not found"); int wordCount = fileContent.Split([' ', '\n', '\r'], StringSplitOptions.RemoveEmptyEntries).Length; @@ -86,12 +86,13 @@ internal sealed class WordCountingExecutor() : ReflectingExecutor("ParagraphCountingExecutor"), IMessageHandler +internal sealed class ParagraphCountingExecutor() : ReflectingExecutor("ParagraphCountingExecutor"), + IMessageHandler { - public async ValueTask HandleAsync(string message, IWorkflowContext context) + public async ValueTask HandleAsync(string message, IWorkflowContext context, CancellationToken cancellationToken = default) { // Retrieve the file content from the shared state - var fileContent = await context.ReadStateAsync(message, scopeName: FileContentStateConstants.FileContentStateScope) + var fileContent = await context.ReadStateAsync(message, scopeName: FileContentStateConstants.FileContentStateScope, cancellationToken) ?? throw new InvalidOperationException("File content state not found"); int paragraphCount = fileContent.Split(['\n', '\r'], StringSplitOptions.RemoveEmptyEntries).Length; @@ -104,7 +105,7 @@ internal sealed class AggregationExecutor() : ReflectingExecutor _messages = []; - public async ValueTask HandleAsync(FileStats message, IWorkflowContext context) + public async ValueTask HandleAsync(FileStats message, IWorkflowContext context, CancellationToken cancellationToken = default) { this._messages.Add(message); @@ -113,7 +114,7 @@ internal sealed class AggregationExecutor() : ReflectingExecutor m.ParagraphCount); var totalWordCount = this._messages.Sum(m => m.WordCount); - await context.YieldOutputAsync($"Total Paragraphs: {totalParagraphCount}, Total Words: {totalWordCount}"); + await context.YieldOutputAsync($"Total Paragraphs: {totalParagraphCount}, Total Words: {totalWordCount}", cancellationToken); } } } diff --git a/dotnet/samples/GettingStarted/Workflows/_Foundational/01_ExecutorsAndEdges/Program.cs b/dotnet/samples/GettingStarted/Workflows/_Foundational/01_ExecutorsAndEdges/Program.cs index b3712eed11..ba70329383 100644 --- a/dotnet/samples/GettingStarted/Workflows/_Foundational/01_ExecutorsAndEdges/Program.cs +++ b/dotnet/samples/GettingStarted/Workflows/_Foundational/01_ExecutorsAndEdges/Program.cs @@ -51,8 +51,10 @@ internal sealed class UppercaseExecutor() : ReflectingExecutor /// The input text to convert /// Workflow context for accessing workflow services and adding events + /// The to monitor for cancellation requests. + /// The default is . /// The input text converted to uppercase - public async ValueTask HandleAsync(string message, IWorkflowContext context) => + public async ValueTask HandleAsync(string message, IWorkflowContext context, CancellationToken cancellationToken = default) => message.ToUpperInvariant(); // The return value will be sent as a message along an edge to subsequent executors } @@ -66,8 +68,10 @@ internal sealed class ReverseTextExecutor() : ReflectingExecutor /// The input text to reverse /// Workflow context for accessing workflow services and adding events + /// The to monitor for cancellation requests. + /// The default is . /// The input text reversed - public async ValueTask HandleAsync(string message, IWorkflowContext context) + public async ValueTask HandleAsync(string message, IWorkflowContext context, CancellationToken cancellationToken = default) { // Because we do not suppress it, the returned result will be yielded as an output from this executor. return string.Concat(message.Reverse()); diff --git a/dotnet/samples/GettingStarted/Workflows/_Foundational/02_Streaming/Program.cs b/dotnet/samples/GettingStarted/Workflows/_Foundational/02_Streaming/Program.cs index c3c1370b84..5ffd64aedb 100644 --- a/dotnet/samples/GettingStarted/Workflows/_Foundational/02_Streaming/Program.cs +++ b/dotnet/samples/GettingStarted/Workflows/_Foundational/02_Streaming/Program.cs @@ -50,8 +50,10 @@ internal sealed class UppercaseExecutor() : ReflectingExecutor /// The input text to convert /// Workflow context for accessing workflow services and adding events + /// The to monitor for cancellation requests. + /// The default is . /// The input text converted to uppercase - public async ValueTask HandleAsync(string message, IWorkflowContext context) => + public async ValueTask HandleAsync(string message, IWorkflowContext context, CancellationToken cancellationToken = default) => message.ToUpperInvariant(); // The return value will be sent as a message along an edge to subsequent executors } @@ -65,8 +67,10 @@ internal sealed class ReverseTextExecutor() : ReflectingExecutor /// The input text to reverse /// Workflow context for accessing workflow services and adding events + /// The to monitor for cancellation requests. + /// The default is . /// The input text reversed - public async ValueTask HandleAsync(string message, IWorkflowContext context) + public async ValueTask HandleAsync(string message, IWorkflowContext context, CancellationToken cancellationToken) { // Because we do not suppress it, the returned result will be yielded as an output from this executor. return string.Concat(message.Reverse()); diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Extensions/AgentProviderExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Extensions/AgentProviderExtensions.cs index 5b74c8f7a7..2de25ec7dc 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Extensions/AgentProviderExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Extensions/AgentProviderExtensions.cs @@ -69,7 +69,7 @@ internal static class AgentProviderExtensions if (autoSend) { - await context.AddEventAsync(new AgentRunUpdateEvent(executorId, update)).ConfigureAwait(false); + await context.AddEventAsync(new AgentRunUpdateEvent(executorId, update), cancellationToken).ConfigureAwait(false); } } @@ -77,7 +77,7 @@ internal static class AgentProviderExtensions if (autoSend) { - await context.AddEventAsync(new AgentRunResponseEvent(executorId, response)).ConfigureAwait(false); + await context.AddEventAsync(new AgentRunResponseEvent(executorId, response), cancellationToken).ConfigureAwait(false); } if (autoSend && !isWorkflowConversation && workflowConversationId is not null) @@ -103,7 +103,7 @@ internal static class AgentProviderExtensions { conversationId = assignValue; - await context.QueueConversationUpdateAsync(conversationId).ConfigureAwait(false); + await context.QueueConversationUpdateAsync(conversationId, cancellationToken: cancellationToken).ConfigureAwait(false); } } } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Extensions/IWorkflowContextExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Extensions/IWorkflowContextExtensions.cs index 4e2815f218..aa38ee0b7a 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Extensions/IWorkflowContextExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Extensions/IWorkflowContextExtensions.cs @@ -14,23 +14,23 @@ namespace Microsoft.Agents.AI.Workflows.Declarative.Extensions; internal static class IWorkflowContextExtensions { - public static ValueTask RaiseInvocationEventAsync(this IWorkflowContext context, DialogAction action, string? priorEventId = null) => - context.AddEventAsync(new DeclarativeActionInvokedEvent(action, priorEventId)); + public static ValueTask RaiseInvocationEventAsync(this IWorkflowContext context, DialogAction action, string? priorEventId = null, CancellationToken cancellationToken = default) => + context.AddEventAsync(new DeclarativeActionInvokedEvent(action, priorEventId), cancellationToken); - public static ValueTask RaiseCompletionEventAsync(this IWorkflowContext context, DialogAction action) => - context.AddEventAsync(new DeclarativeActionCompletedEvent(action)); + public static ValueTask RaiseCompletionEventAsync(this IWorkflowContext context, DialogAction action, CancellationToken cancellationToken = default) => + context.AddEventAsync(new DeclarativeActionCompletedEvent(action), cancellationToken); public static ValueTask SendResultMessageAsync(this IWorkflowContext context, string id, object? result = null, CancellationToken cancellationToken = default) => - context.SendMessageAsync(new ActionExecutorResult(id, result)); + context.SendMessageAsync(new ActionExecutorResult(id, result), cancellationToken: cancellationToken); - public static ValueTask QueueStateResetAsync(this IWorkflowContext context, PropertyPath variablePath) => - context.QueueStateUpdateAsync(Throw.IfNull(variablePath.VariableName), UnassignedValue.Instance, Throw.IfNull(variablePath.NamespaceAlias)); + public static ValueTask QueueStateResetAsync(this IWorkflowContext context, PropertyPath variablePath, CancellationToken cancellationToken = default) => + context.QueueStateUpdateAsync(Throw.IfNull(variablePath.VariableName), UnassignedValue.Instance, Throw.IfNull(variablePath.NamespaceAlias), cancellationToken); - public static ValueTask QueueStateUpdateAsync(this IWorkflowContext context, PropertyPath variablePath, TValue? value) => - context.QueueStateUpdateAsync(Throw.IfNull(variablePath.VariableName), value, Throw.IfNull(variablePath.NamespaceAlias)); + public static ValueTask QueueStateUpdateAsync(this IWorkflowContext context, PropertyPath variablePath, TValue? value, CancellationToken cancellationToken = default) => + context.QueueStateUpdateAsync(Throw.IfNull(variablePath.VariableName), value, Throw.IfNull(variablePath.NamespaceAlias), cancellationToken); - public static ValueTask QueueSystemUpdateAsync(this IWorkflowContext context, string key, TValue? value) => - DeclarativeContext(context).QueueSystemUpdateAsync(key, value); + public static ValueTask QueueSystemUpdateAsync(this IWorkflowContext context, string key, TValue? value, CancellationToken cancellationToken = default) => + DeclarativeContext(context).QueueSystemUpdateAsync(key, value, cancellationToken); public static FormulaValue ReadState(this IWorkflowContext context, PropertyPath variablePath) => context.ReadState(Throw.IfNull(variablePath.VariableName), Throw.IfNull(variablePath.NamespaceAlias)); @@ -38,18 +38,18 @@ internal static class IWorkflowContextExtensions public static FormulaValue ReadState(this IWorkflowContext context, string key, string? scopeName = null) => DeclarativeContext(context).State.Get(key, scopeName); - public static async ValueTask QueueConversationUpdateAsync(this IWorkflowContext context, string conversationId, bool isExternal = false) + public static async ValueTask QueueConversationUpdateAsync(this IWorkflowContext context, string conversationId, bool isExternal = false, CancellationToken cancellationToken = default) { RecordValue conversation = (RecordValue)context.ReadState(SystemScope.Names.Conversation, VariableScopeNames.System); if (isExternal) { conversation.UpdateField("Id", FormulaValue.New(conversationId)); - await context.QueueSystemUpdateAsync(SystemScope.Names.Conversation, conversation).ConfigureAwait(false); - await context.QueueSystemUpdateAsync(SystemScope.Names.ConversationId, FormulaValue.New(conversationId)).ConfigureAwait(false); + await context.QueueSystemUpdateAsync(SystemScope.Names.Conversation, conversation, cancellationToken).ConfigureAwait(false); + await context.QueueSystemUpdateAsync(SystemScope.Names.ConversationId, FormulaValue.New(conversationId), cancellationToken).ConfigureAwait(false); } - await context.AddEventAsync(new ConversationUpdateEvent(conversationId) { IsWorkflow = isExternal }).ConfigureAwait(false); + await context.AddEventAsync(new ConversationUpdateEvent(conversationId) { IsWorkflow = isExternal }, cancellationToken).ConfigureAwait(false); } public static bool IsWorkflowConversation( diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeActionExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeActionExecutor.cs index 95604846ec..3b0169bc00 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeActionExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeActionExecutor.cs @@ -61,7 +61,7 @@ internal abstract class DeclarativeActionExecutor : Executor - public override async ValueTask HandleAsync(ActionExecutorResult message, IWorkflowContext context) + public override async ValueTask HandleAsync(ActionExecutorResult message, IWorkflowContext context, CancellationToken cancellationToken = default) { if (this.Model.Disabled) { @@ -69,7 +69,7 @@ internal abstract class DeclarativeActionExecutor : Executor? TraceContext => this.Source.TraceContext; /// - public ValueTask AddEventAsync(WorkflowEvent workflowEvent) => this.Source.AddEventAsync(workflowEvent); + public ValueTask AddEventAsync(WorkflowEvent workflowEvent, CancellationToken cancellationToken = default) + => this.Source.AddEventAsync(workflowEvent, cancellationToken); /// - public ValueTask YieldOutputAsync(object output) => this.Source.YieldOutputAsync(output); + public ValueTask YieldOutputAsync(object output, CancellationToken cancellationToken = default) + => this.Source.YieldOutputAsync(output, cancellationToken); /// public ValueTask RequestHaltAsync() => this.Source.RequestHaltAsync(); /// - public async ValueTask QueueClearScopeAsync(string? scopeName = null) + public async ValueTask QueueClearScopeAsync(string? scopeName = null, CancellationToken cancellationToken = default) { if (scopeName is not null) { @@ -50,12 +53,12 @@ internal sealed class DeclarativeWorkflowContext : IWorkflowContext // Copy keys to array to avoid modifying collection during enumeration. foreach (string key in this.State.Keys(scopeName).ToArray()) { - await this.UpdateStateAsync(key, UnassignedValue.Instance, scopeName).ConfigureAwait(false); + await this.UpdateStateAsync(key, UnassignedValue.Instance, scopeName, cancellationToken: cancellationToken).ConfigureAwait(false); } } else { - await this.Source.QueueClearScopeAsync(scopeName).ConfigureAwait(false); + await this.Source.QueueClearScopeAsync(scopeName, cancellationToken).ConfigureAwait(false); } this.State.Bind(); @@ -63,20 +66,20 @@ internal sealed class DeclarativeWorkflowContext : IWorkflowContext } /// - public async ValueTask QueueStateUpdateAsync(string key, T? value, string? scopeName = null) + public async ValueTask QueueStateUpdateAsync(string key, T? value, string? scopeName = null, CancellationToken cancellationToken = default) { - await this.UpdateStateAsync(key, value, scopeName).ConfigureAwait(false); + await this.UpdateStateAsync(key, value, scopeName, cancellationToken: cancellationToken).ConfigureAwait(false); this.State.Bind(); } - public async ValueTask QueueSystemUpdateAsync(string key, TValue? value) + public async ValueTask QueueSystemUpdateAsync(string key, TValue? value, CancellationToken cancellationToken = default) { - await this.UpdateStateAsync(key, value, VariableScopeNames.System, allowSystem: true).ConfigureAwait(false); + await this.UpdateStateAsync(key, value, VariableScopeNames.System, allowSystem: true, cancellationToken).ConfigureAwait(false); this.State.Bind(); } /// - public async ValueTask ReadStateAsync(string key, string? scopeName = null) + public async ValueTask ReadStateAsync(string key, string? scopeName = null, CancellationToken cancellationToken = default) { bool isManagedScope = scopeName is not null && // null scope cannot be managed @@ -86,21 +89,23 @@ internal sealed class DeclarativeWorkflowContext : IWorkflowContext { // Not a managed scope, just pass through. This is valid when a declarative // workflow has been ejected to code (where DeclarativeWorkflowContext is also utilized). - _ when !isManagedScope => await this.Source.ReadStateAsync(key, scopeName).ConfigureAwait(false), + _ when !isManagedScope => await this.Source.ReadStateAsync(key, scopeName, cancellationToken).ConfigureAwait(false), // Retrieve formula values directly from the managed state to avoid conversion. _ when typeof(TValue) == typeof(FormulaValue) => (TValue?)(object?)this.State.Get(key, scopeName), // Retrieve native types from the source context to avoid conversion. - _ => await this.Source.ReadStateAsync(key, scopeName).ConfigureAwait(false), + _ => await this.Source.ReadStateAsync(key, scopeName, cancellationToken).ConfigureAwait(false), }; } /// - public ValueTask> ReadStateKeysAsync(string? scopeName = null) => this.Source.ReadStateKeysAsync(scopeName); + public ValueTask> ReadStateKeysAsync(string? scopeName = null, CancellationToken cancellationToken = default) + => this.Source.ReadStateKeysAsync(scopeName, cancellationToken); /// - public ValueTask SendMessageAsync(object message, string? targetId = null) => this.Source.SendMessageAsync(message, targetId); + public ValueTask SendMessageAsync(object message, string? targetId = null, CancellationToken cancellationToken = default) + => this.Source.SendMessageAsync(message, targetId, cancellationToken); - private ValueTask UpdateStateAsync(string key, T? value, string? scopeName, bool allowSystem = true) + private ValueTask UpdateStateAsync(string key, T? value, string? scopeName, bool allowSystem = true, CancellationToken cancellationToken = default) { bool isManagedScope = scopeName is not null && // null scope cannot be managed @@ -110,7 +115,7 @@ internal sealed class DeclarativeWorkflowContext : IWorkflowContext { // Not a managed scope, just pass through. This is valid when a declarative // workflow has been ejected to code (where DeclarativeWorkflowContext is also utilized). - return this.Source.QueueStateUpdateAsync(key, value, scopeName); + return this.Source.QueueStateUpdateAsync(key, value, scopeName, cancellationToken); } if (!ManagedScopes.Contains(scopeName!) && !allowSystem) @@ -134,7 +139,7 @@ internal sealed class DeclarativeWorkflowContext : IWorkflowContext { this.State.Set(key, FormulaValue.NewBlank(), scopeName); } - return this.Source.QueueStateUpdateAsync(key, UnassignedValue.Instance, scopeName); + return this.Source.QueueStateUpdateAsync(key, UnassignedValue.Instance, scopeName, cancellationToken); } ValueTask QueueFormulaStateAsync(FormulaValue formulaValue) @@ -143,7 +148,7 @@ internal sealed class DeclarativeWorkflowContext : IWorkflowContext { this.State.Set(key, formulaValue, scopeName); } - return this.Source.QueueStateUpdateAsync(key, formulaValue.ToObject(), scopeName); + return this.Source.QueueStateUpdateAsync(key, formulaValue.ToObject(), scopeName, cancellationToken); } ValueTask QueueDataValueStateAsync(DataValue dataValue) @@ -153,7 +158,7 @@ internal sealed class DeclarativeWorkflowContext : IWorkflowContext FormulaValue formulaValue = dataValue.ToFormula(); this.State.Set(key, formulaValue, scopeName); } - return this.Source.QueueStateUpdateAsync(key, dataValue.ToObject(), scopeName); + return this.Source.QueueStateUpdateAsync(key, dataValue.ToObject(), scopeName, cancellationToken); } ValueTask QueueNativeStateAsync(object? rawValue) @@ -163,7 +168,7 @@ internal sealed class DeclarativeWorkflowContext : IWorkflowContext FormulaValue formulaValue = rawValue.ToFormula(); this.State.Set(key, formulaValue, scopeName); } - return this.Source.QueueStateUpdateAsync(key, rawValue, scopeName); + return this.Source.QueueStateUpdateAsync(key, rawValue, scopeName, cancellationToken); } } } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowExecutor.cs index e30bb94a01..c228533f17 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowExecutor.cs @@ -1,6 +1,7 @@ // Copyright (c) Microsoft. All rights reserved. using System; +using System.Threading; using System.Threading.Tasks; using Microsoft.Agents.AI.Workflows.Declarative.Extensions; using Microsoft.Agents.AI.Workflows.Declarative.PowerFx; @@ -24,7 +25,7 @@ internal sealed class DeclarativeWorkflowExecutor( return default; } - public override async ValueTask HandleAsync(TInput message, IWorkflowContext context) + public override async ValueTask HandleAsync(TInput message, IWorkflowContext context, CancellationToken cancellationToken = default) { // No state to restore if we're starting from the beginning. state.SetInitialized(); @@ -35,13 +36,13 @@ internal sealed class DeclarativeWorkflowExecutor( string? conversationId = options.ConversationId; if (string.IsNullOrWhiteSpace(conversationId)) { - conversationId = await options.AgentProvider.CreateConversationAsync(cancellationToken: default).ConfigureAwait(false); + conversationId = await options.AgentProvider.CreateConversationAsync(cancellationToken).ConfigureAwait(false); } - await declarativeContext.QueueConversationUpdateAsync(conversationId, isExternal: true).ConfigureAwait(false); + await declarativeContext.QueueConversationUpdateAsync(conversationId, isExternal: true, cancellationToken).ConfigureAwait(false); - await options.AgentProvider.CreateMessageAsync(conversationId, input, cancellationToken: default).ConfigureAwait(false); + await options.AgentProvider.CreateMessageAsync(conversationId, input, cancellationToken).ConfigureAwait(false); await declarativeContext.SetLastMessageAsync(input).ConfigureAwait(false); - await context.SendResultMessageAsync(this.Id).ConfigureAwait(false); + await context.SendResultMessageAsync(this.Id, cancellationToken: cancellationToken).ConfigureAwait(false); } } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DelegateActionExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DelegateActionExecutor.cs index 21f1f1aa4d..e05f0908af 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DelegateActionExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DelegateActionExecutor.cs @@ -1,6 +1,7 @@ // Copyright (c) Microsoft. All rights reserved. using System.Diagnostics; +using System.Threading; using System.Threading.Tasks; using Microsoft.Agents.AI.Workflows.Declarative.Extensions; using Microsoft.Agents.AI.Workflows.Declarative.Kit; @@ -11,11 +12,11 @@ namespace Microsoft.Agents.AI.Workflows.Declarative.Interpreter; internal sealed class DelegateActionExecutor(string actionId, WorkflowFormulaState state, DelegateAction? action = null, bool emitResult = true) : DelegateActionExecutor(actionId, state, action, emitResult) { - public override ValueTask HandleAsync(ActionExecutorResult message, IWorkflowContext context) + public override ValueTask HandleAsync(ActionExecutorResult message, IWorkflowContext context, CancellationToken cancellationToken) { Debug.WriteLine($"RESULT #{this.Id} - {message.Result ?? "(null)"}"); - return base.HandleAsync(message, context); + return base.HandleAsync(message, context, cancellationToken); } } @@ -39,16 +40,16 @@ internal class DelegateActionExecutor : Executor, IResettabl return default; } - public override async ValueTask HandleAsync(TMessage message, IWorkflowContext context) + public override async ValueTask HandleAsync(TMessage message, IWorkflowContext context, CancellationToken cancellationToken = default) { if (this._action is not null) { - await this._action.Invoke(new DeclarativeWorkflowContext(context, this._state), message, default).ConfigureAwait(false); + await this._action.Invoke(new DeclarativeWorkflowContext(context, this._state), message, cancellationToken).ConfigureAwait(false); } if (this._emitResult) { - await context.SendResultMessageAsync(this.Id).ConfigureAwait(false); + await context.SendResultMessageAsync(this.Id, cancellationToken: cancellationToken).ConfigureAwait(false); } } } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/ActionExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/ActionExecutor.cs index e92c6b1b7d..a73eb70a73 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/ActionExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/ActionExecutor.cs @@ -73,12 +73,12 @@ public abstract class ActionExecutor : Executor, IResettable } /// - public override async ValueTask HandleAsync(TMessage message, IWorkflowContext context) + public override async ValueTask HandleAsync(TMessage message, IWorkflowContext context, CancellationToken cancellationToken) { object? result = await this.ExecuteAsync(new DeclarativeWorkflowContext(context, this._session.State), message, cancellationToken: default).ConfigureAwait(false); Debug.WriteLine($"RESULT #{this.Id} - {result ?? "(null)"}"); - await context.SendResultMessageAsync(this.Id, result).ConfigureAwait(false); + await context.SendResultMessageAsync(this.Id, result, cancellationToken).ConfigureAwait(false); } /// diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/IWorkflowContextExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/IWorkflowContextExtensions.cs index 185068d550..e33f32a1a3 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/IWorkflowContextExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/IWorkflowContextExtensions.cs @@ -132,7 +132,7 @@ public static class IWorkflowContextExtensions /// The converted value public static async ValueTask ConvertValueAsync(this IWorkflowContext context, VariableType targetType, string key, string? scopeName = null, CancellationToken cancellationToken = default) { - object? sourceValue = await context.ReadStateAsync(key, scopeName).ConfigureAwait(false); + object? sourceValue = await context.ReadStateAsync(key, scopeName, cancellationToken).ConfigureAwait(false); return sourceValue.ConvertType(targetType); } @@ -143,10 +143,11 @@ public static class IWorkflowContextExtensions /// The workflow execution context used to restore persisted state prior to formatting. /// The key of the state value. /// An optional name that specifies the scope to read.If null, the default scope is used. + /// A token that propagates notification when operation should be canceled. /// The evaluated list expression - public static async ValueTask?> ReadListAsync(this IWorkflowContext context, string key, string? scopeName = null) + public static async ValueTask?> ReadListAsync(this IWorkflowContext context, string key, string? scopeName = null, CancellationToken cancellationToken = default) { - object? value = await context.ReadStateAsync(key, scopeName).ConfigureAwait(false); + object? value = await context.ReadStateAsync(key, scopeName, cancellationToken).ConfigureAwait(false); return value.AsList(); } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/RootExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/RootExecutor.cs index a1ebd09f5b..87957b8b9d 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/RootExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Kit/RootExecutor.cs @@ -55,23 +55,23 @@ public abstract class RootExecutor : Executor, IResettableExecut } /// - public override async ValueTask HandleAsync(TInput message, IWorkflowContext context) + public override async ValueTask HandleAsync(TInput message, IWorkflowContext context, CancellationToken cancellationToken) { DeclarativeWorkflowContext declarativeContext = new(context, this._state); - await this.ExecuteAsync(message, declarativeContext, cancellationToken: default).ConfigureAwait(false); + await this.ExecuteAsync(message, declarativeContext, cancellationToken).ConfigureAwait(false); ChatMessage input = (this._inputTransform ?? DefaultInputTransform).Invoke(message); if (string.IsNullOrWhiteSpace(this._conversationId)) { - this._conversationId = await this._agentProvider.CreateConversationAsync(cancellationToken: default).ConfigureAwait(false); + this._conversationId = await this._agentProvider.CreateConversationAsync(cancellationToken).ConfigureAwait(false); } - await declarativeContext.QueueConversationUpdateAsync(this._conversationId, isExternal: true).ConfigureAwait(false); + await declarativeContext.QueueConversationUpdateAsync(this._conversationId, isExternal: true, cancellationToken).ConfigureAwait(false); - await this._agentProvider.CreateMessageAsync(this._conversationId, input, cancellationToken: default).ConfigureAwait(false); + await this._agentProvider.CreateMessageAsync(this._conversationId, input, cancellationToken).ConfigureAwait(false); await declarativeContext.SetLastMessageAsync(input).ConfigureAwait(false); - await declarativeContext.SendMessageAsync(new ActionExecutorResult(this.Id)).ConfigureAwait(false); + await declarativeContext.SendMessageAsync(new ActionExecutorResult(this.Id), cancellationToken: cancellationToken).ConfigureAwait(false); } /// diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/ClearAllVariablesExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/ClearAllVariablesExecutor.cs index 419e112d0a..62f0dbac75 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/ClearAllVariablesExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/ClearAllVariablesExecutor.cs @@ -28,7 +28,7 @@ internal sealed class ClearAllVariablesExecutor(ClearAllVariables model, Workflo if (scope is not null) { - await context.QueueClearScopeAsync(scope).ConfigureAwait(false); + await context.QueueClearScopeAsync(scope, cancellationToken).ConfigureAwait(false); Debug.WriteLine( $""" STATE: {this.GetType().Name} [{this.Id}] diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/ConditionGroupExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/ConditionGroupExecutor.cs index 65a6510b18..b935b6e185 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/ConditionGroupExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/ConditionGroupExecutor.cs @@ -69,5 +69,5 @@ internal sealed class ConditionGroupExecutor : DeclarativeActionExecutor - await context.RaiseCompletionEventAsync(this.Model).ConfigureAwait(false); + await context.RaiseCompletionEventAsync(this.Model, cancellationToken).ConfigureAwait(false); } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/CreateConversationExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/CreateConversationExecutor.cs index 1da428f931..6ed1d4db32 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/CreateConversationExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/CreateConversationExecutor.cs @@ -17,7 +17,7 @@ internal sealed class CreateConversationExecutor(CreateConversation model, Workf { string conversationId = await agentProvider.CreateConversationAsync(cancellationToken).ConfigureAwait(false); await this.AssignAsync(this.Model.ConversationId?.Path, FormulaValue.New(conversationId), context).ConfigureAwait(false); - await context.QueueConversationUpdateAsync(conversationId).ConfigureAwait(false); + await context.QueueConversationUpdateAsync(conversationId, cancellationToken: cancellationToken).ConfigureAwait(false); return default; } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/ForeachExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/ForeachExecutor.cs index 41c9e7cb5a..3130a29b85 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/ForeachExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/ForeachExecutor.cs @@ -68,11 +68,11 @@ internal sealed class ForeachExecutor : DeclarativeActionExecutor { FormulaValue value = this._values[this._index]; - await context.QueueStateUpdateAsync(Throw.IfNull(this.Model.Value), value).ConfigureAwait(false); + await context.QueueStateUpdateAsync(Throw.IfNull(this.Model.Value), value, cancellationToken).ConfigureAwait(false); if (this.Model.Index is not null) { - await context.QueueStateUpdateAsync(this.Model.Index.Path, FormulaValue.New(this._index)).ConfigureAwait(false); + await context.QueueStateUpdateAsync(this.Model.Index.Path, FormulaValue.New(this._index), cancellationToken).ConfigureAwait(false); } this._index++; @@ -83,15 +83,15 @@ internal sealed class ForeachExecutor : DeclarativeActionExecutor { try { - await context.QueueStateResetAsync(Throw.IfNull(this.Model.Value)).ConfigureAwait(false); + await context.QueueStateResetAsync(Throw.IfNull(this.Model.Value), cancellationToken).ConfigureAwait(false); if (this.Model.Index is not null) { - await context.QueueStateResetAsync(this.Model.Index).ConfigureAwait(false); + await context.QueueStateResetAsync(this.Model.Index, cancellationToken).ConfigureAwait(false); } } finally { - await context.RaiseCompletionEventAsync(this.Model).ConfigureAwait(false); + await context.RaiseCompletionEventAsync(this.Model, cancellationToken).ConfigureAwait(false); } } } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/QuestionExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/QuestionExecutor.cs index 1d34869441..be82ca4935 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/QuestionExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/QuestionExecutor.cs @@ -75,7 +75,7 @@ internal sealed class QuestionExecutor(Question model, WorkflowFormulaState stat { int count = await this._promptCount.ReadAsync(context).ConfigureAwait(false); InputRequest inputRequest = new(this.FormatPrompt(this.Model.Prompt)); - await context.SendMessageAsync(inputRequest).ConfigureAwait(false); + await context.SendMessageAsync(inputRequest, cancellationToken: cancellationToken).ConfigureAwait(false); await this._promptCount.WriteAsync(context, count + 1).ConfigureAwait(false); } @@ -85,7 +85,7 @@ internal sealed class QuestionExecutor(Question model, WorkflowFormulaState stat if (string.IsNullOrWhiteSpace(message.Value)) { string unrecognizedResponse = this.FormatPrompt(this.Model.UnrecognizedPrompt); - await context.AddEventAsync(new MessageActivityEvent(unrecognizedResponse.Trim())).ConfigureAwait(false); + await context.AddEventAsync(new MessageActivityEvent(unrecognizedResponse.Trim()), cancellationToken).ConfigureAwait(false); } else { @@ -97,7 +97,7 @@ internal sealed class QuestionExecutor(Question model, WorkflowFormulaState stat else { string invalidResponse = this.FormatPrompt(this.Model.InvalidPrompt); - await context.AddEventAsync(new MessageActivityEvent(invalidResponse.Trim())).ConfigureAwait(false); + await context.AddEventAsync(new MessageActivityEvent(invalidResponse.Trim()), cancellationToken).ConfigureAwait(false); } } @@ -115,7 +115,7 @@ internal sealed class QuestionExecutor(Question model, WorkflowFormulaState stat public async ValueTask CompleteAsync(IWorkflowContext context, ActionExecutorResult message, CancellationToken cancellationToken) { - await context.RaiseCompletionEventAsync(this.Model).ConfigureAwait(false); + await context.RaiseCompletionEventAsync(this.Model, cancellationToken).ConfigureAwait(false); } private async ValueTask PromptAsync(IWorkflowContext context, CancellationToken cancellationToken) @@ -128,7 +128,7 @@ internal sealed class QuestionExecutor(Question model, WorkflowFormulaState stat DataValue defaultValue = this.Evaluator.GetValue(defaultValueExpression).Value; await this.AssignAsync(this.Model.Variable?.Path, defaultValue.ToFormula(), context).ConfigureAwait(false); string defaultValueResponse = this.FormatPrompt(this.Model.DefaultValueResponse); - await context.AddEventAsync(new MessageActivityEvent(defaultValueResponse.Trim())).ConfigureAwait(false); + await context.AddEventAsync(new MessageActivityEvent(defaultValueResponse.Trim()), cancellationToken).ConfigureAwait(false); await context.SendResultMessageAsync(this.Id, result: null, cancellationToken).ConfigureAwait(false); } else diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/ResetVariableExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/ResetVariableExecutor.cs index da625b11ae..eb679fa4b0 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/ResetVariableExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/ResetVariableExecutor.cs @@ -17,7 +17,7 @@ internal sealed class ResetVariableExecutor(ResetVariable model, WorkflowFormula protected override async ValueTask ExecuteAsync(IWorkflowContext context, CancellationToken cancellationToken = default) { Throw.IfNull(this.Model.Variable, $"{nameof(this.Model)}.{nameof(model.Variable)}"); - await context.QueueStateResetAsync(this.Model.Variable).ConfigureAwait(false); + await context.QueueStateResetAsync(this.Model.Variable, cancellationToken).ConfigureAwait(false); Debug.WriteLine( $""" STATE: {this.GetType().Name} [{this.Id}] diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/SendActivityExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/SendActivityExecutor.cs index 69ed8bb970..9af463f865 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/SendActivityExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/SendActivityExecutor.cs @@ -18,7 +18,7 @@ internal sealed class SendActivityExecutor(SendActivity model, WorkflowFormulaSt { string activityText = this.Engine.Format(messageActivity.Text).Trim(); - await context.AddEventAsync(new MessageActivityEvent(activityText.Trim())).ConfigureAwait(false); + await context.AddEventAsync(new MessageActivityEvent(activityText.Trim()), cancellationToken).ConfigureAwait(false); } return default; diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PowerFx/WorkflowFormulaState.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PowerFx/WorkflowFormulaState.cs index ee58a3c63c..3d810d073b 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PowerFx/WorkflowFormulaState.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PowerFx/WorkflowFormulaState.cs @@ -72,10 +72,10 @@ internal sealed class WorkflowFormulaState async Task ReadScopeAsync(string scopeName) { - HashSet keys = await context.ReadStateKeysAsync(scopeName).ConfigureAwait(false); + HashSet keys = await context.ReadStateKeysAsync(scopeName, cancellationToken).ConfigureAwait(false); foreach (string key in keys) { - object? value = await context.ReadStateAsync(key, scopeName).ConfigureAwait(false); + object? value = await context.ReadStateAsync(key, scopeName, cancellationToken).ConfigureAwait(false); if (value is null or UnassignedValue) { value = FormulaValue.NewBlank(); diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/AgentWorkflowBuilder.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/AgentWorkflowBuilder.cs index b21edee04c..ed91c14e5f 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/AgentWorkflowBuilder.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/AgentWorkflowBuilder.cs @@ -150,12 +150,12 @@ public static partial class AgentWorkflowBuilder protected override RouteBuilder ConfigureRoutes(RouteBuilder routeBuilder) => routeBuilder - .AddHandler((message, context) => this._pendingMessages.Add(new(ChatRole.User, message))) - .AddHandler((message, context) => this._pendingMessages.Add(message)) - .AddHandler>((messages, _) => this._pendingMessages.AddRange(messages)) - .AddHandler((messages, _) => this._pendingMessages.AddRange(messages)) // TODO: Remove once https://github.com/microsoft/agent-framework/issues/782 is addressed - .AddHandler>((messages, _) => this._pendingMessages.AddRange(messages)) // TODO: Remove once https://github.com/microsoft/agent-framework/issues/782 is addressed - .AddHandler(async (token, context) => + .AddHandler((message, _, __) => this._pendingMessages.Add(new(ChatRole.User, message))) + .AddHandler((message, _, __) => this._pendingMessages.Add(message)) + .AddHandler>((messages, _, __) => this._pendingMessages.AddRange(messages)) + .AddHandler((messages, _, __) => this._pendingMessages.AddRange(messages)) // TODO: Remove once https://github.com/microsoft/agent-framework/issues/782 is addressed + .AddHandler>((messages, _, __) => this._pendingMessages.AddRange(messages)) // TODO: Remove once https://github.com/microsoft/agent-framework/issues/782 is addressed + .AddHandler(async (token, context, cancellationToken) => { List messages = [.. this._pendingMessages]; this._pendingMessages.Clear(); @@ -163,12 +163,12 @@ public static partial class AgentWorkflowBuilder List? roleChanged = ChangeAssistantToUserForOtherParticipants(agent.DisplayName, messages); List updates = []; - await foreach (var update in agent.RunStreamingAsync(messages).ConfigureAwait(false)) + await foreach (var update in agent.RunStreamingAsync(messages, cancellationToken: cancellationToken).ConfigureAwait(false)) { updates.Add(update); if (token.EmitEvents is true) { - await context.AddEventAsync(new AgentRunUpdateEvent(this.Id, update)).ConfigureAwait(false); + await context.AddEventAsync(new AgentRunUpdateEvent(this.Id, update), cancellationToken).ConfigureAwait(false); } } @@ -181,8 +181,8 @@ public static partial class AgentWorkflowBuilder messages.AddRange(updates.ToAgentRunResponse().Messages); - await context.SendMessageAsync(messages).ConfigureAwait(false); - await context.SendMessageAsync(token).ConfigureAwait(false); + await context.SendMessageAsync(messages, cancellationToken: cancellationToken).ConfigureAwait(false); + await context.SendMessageAsync(token, cancellationToken: cancellationToken).ConfigureAwait(false); }); public ValueTask ResetAsync() @@ -199,7 +199,7 @@ public static partial class AgentWorkflowBuilder private sealed class OutputMessagesExecutor() : ChatProtocolExecutor("OutputMessages"), IResettableExecutor { protected override ValueTask TakeTurnAsync(List messages, IWorkflowContext context, bool? emitEvents, CancellationToken cancellationToken = default) - => context.YieldOutputAsync(messages); + => context.YieldOutputAsync(messages, cancellationToken); ValueTask IResettableExecutor.ResetAsync() => this.ResetAsync(); } @@ -209,10 +209,10 @@ public static partial class AgentWorkflowBuilder { protected override RouteBuilder ConfigureRoutes(RouteBuilder routeBuilder) => routeBuilder - .AddHandler((message, context) => context.SendMessageAsync(new ChatMessage(ChatRole.User, message))) - .AddHandler((message, context) => context.SendMessageAsync(message)) - .AddHandler>((messages, context) => context.SendMessageAsync(messages)) - .AddHandler((turnToken, context) => context.SendMessageAsync(turnToken)); + .AddHandler((message, context, cancellationToken) => context.SendMessageAsync(new ChatMessage(ChatRole.User, message), cancellationToken: cancellationToken)) + .AddHandler((message, context, cancellationToken) => context.SendMessageAsync(message, cancellationToken: cancellationToken)) + .AddHandler>((messages, context, cancellationToken) => context.SendMessageAsync(messages, cancellationToken: cancellationToken)) + .AddHandler((turnToken, context, cancellationToken) => context.SendMessageAsync(turnToken, cancellationToken: cancellationToken)); public ValueTask ResetAsync() => default; } @@ -224,7 +224,7 @@ public static partial class AgentWorkflowBuilder private sealed class BatchChatMessagesToListExecutor(string id) : ChatProtocolExecutor(id), IResettableExecutor { protected override ValueTask TakeTurnAsync(List messages, IWorkflowContext context, bool? emitEvents, CancellationToken cancellationToken = default) - => context.SendMessageAsync(messages); + => context.SendMessageAsync(messages, cancellationToken: cancellationToken); ValueTask IResettableExecutor.ResetAsync() => this.ResetAsync(); } @@ -256,7 +256,7 @@ public static partial class AgentWorkflowBuilder } protected override RouteBuilder ConfigureRoutes(RouteBuilder routeBuilder) => - routeBuilder.AddHandler>(async (messages, context) => + routeBuilder.AddHandler>(async (messages, context, cancellationToken) => { // TODO: https://github.com/microsoft/agent-framework/issues/784 // This locking should not be necessary. @@ -273,7 +273,7 @@ public static partial class AgentWorkflowBuilder var results = this._allResults; this._allResults = new List>(this._expectedInputs); - await context.YieldOutputAsync(this._aggregator(results)).ConfigureAwait(false); + await context.YieldOutputAsync(this._aggregator(results), cancellationToken).ConfigureAwait(false); } }); @@ -457,16 +457,17 @@ public static partial class AgentWorkflowBuilder protected override RouteBuilder ConfigureRoutes(RouteBuilder routeBuilder) => routeBuilder - .AddHandler((message, context) => this._pendingMessages.Add(new(ChatRole.User, message))) - .AddHandler((message, context) => this._pendingMessages.Add(message)) - .AddHandler>((messages, _) => this._pendingMessages.AddRange(messages)) - .AddHandler((messages, _) => this._pendingMessages.AddRange(messages)) // TODO: Remove once https://github.com/microsoft/agent-framework/issues/782 is addressed - .AddHandler>((messages, _) => this._pendingMessages.AddRange(messages)) // TODO: Remove once https://github.com/microsoft/agent-framework/issues/782 is addressed - .AddHandler(async (token, context) => + .AddHandler((message, context, _) => this._pendingMessages.Add(new(ChatRole.User, message))) + .AddHandler((message, context, _) => this._pendingMessages.Add(message)) + .AddHandler>((messages, _, __) => this._pendingMessages.AddRange(messages)) + .AddHandler((messages, _, __) => this._pendingMessages.AddRange(messages)) // TODO: Remove once https://github.com/microsoft/agent-framework/issues/782 is addressed + .AddHandler>((messages, _, __) => this._pendingMessages.AddRange(messages)) // TODO: Remove once https://github.com/microsoft/agent-framework/issues/782 is addressed + .AddHandler(async (token, context, cancellationToken) => { var messages = new List(this._pendingMessages); this._pendingMessages.Clear(); - await context.SendMessageAsync(new HandoffState(token, null, messages)).ConfigureAwait(false); + await context.SendMessageAsync(new HandoffState(token, null, messages), cancellationToken: cancellationToken) + .ConfigureAwait(false); }); public ValueTask ResetAsync() @@ -480,8 +481,8 @@ public static partial class AgentWorkflowBuilder private sealed class EndHandoffsExecutor() : Executor("HandoffEnd"), IResettableExecutor { protected override RouteBuilder ConfigureRoutes(RouteBuilder routeBuilder) => - routeBuilder.AddHandler((handoff, context) => - context.YieldOutputAsync(handoff.Messages)); + routeBuilder.AddHandler((handoff, context, cancellationToken) => + context.YieldOutputAsync(handoff.Messages, cancellationToken)); public ValueTask ResetAsync() => default; } @@ -534,7 +535,7 @@ public static partial class AgentWorkflowBuilder }); protected override RouteBuilder ConfigureRoutes(RouteBuilder routeBuilder) => - routeBuilder.AddHandler(async (handoffState, context) => + routeBuilder.AddHandler(async (handoffState, context, cancellationToken) => { string? requestedHandoff = null; List updates = []; @@ -542,24 +543,31 @@ public static partial class AgentWorkflowBuilder List? roleChanges = ChangeAssistantToUserForOtherParticipants(this._agent.DisplayName, allMessages); - await foreach (var update in this._agent.RunStreamingAsync(allMessages, options: this._agentOptions).ConfigureAwait(false)) + await foreach (var update in this._agent.RunStreamingAsync(allMessages, + options: this._agentOptions, + cancellationToken: cancellationToken) + .ConfigureAwait(false)) { - await AddUpdateAsync(update).ConfigureAwait(false); + await AddUpdateAsync(update, cancellationToken).ConfigureAwait(false); foreach (var c in update.Contents) { if (c is FunctionCallContent fcc && this._handoffFunctionNames.Contains(fcc.Name)) { requestedHandoff = fcc.Name; - await AddUpdateAsync(new AgentRunResponseUpdate - { - AgentId = this._agent.Id, - AuthorName = this._agent.DisplayName, - Contents = [new FunctionResultContent(fcc.CallId, "Transferred.")], - CreatedAt = DateTimeOffset.UtcNow, - MessageId = Guid.NewGuid().ToString("N"), - Role = ChatRole.Tool, - }).ConfigureAwait(false); + await AddUpdateAsync( + new AgentRunResponseUpdate + { + AgentId = this._agent.Id, + AuthorName = this._agent.DisplayName, + Contents = [new FunctionResultContent(fcc.CallId, "Transferred.")], + CreatedAt = DateTimeOffset.UtcNow, + MessageId = Guid.NewGuid().ToString("N"), + Role = ChatRole.Tool, + }, + cancellationToken + ) + .ConfigureAwait(false); } } } @@ -568,14 +576,14 @@ public static partial class AgentWorkflowBuilder ResetUserToAssistantForChangedRoles(roleChanges); - await context.SendMessageAsync(new HandoffState(handoffState.TurnToken, requestedHandoff, allMessages)).ConfigureAwait(false); + await context.SendMessageAsync(new HandoffState(handoffState.TurnToken, requestedHandoff, allMessages), cancellationToken: cancellationToken).ConfigureAwait(false); - async Task AddUpdateAsync(AgentRunResponseUpdate update) + async Task AddUpdateAsync(AgentRunResponseUpdate update, CancellationToken cancellationToken) { updates.Add(update); if (handoffState.TurnToken.EmitEvents is true) { - await context.AddEventAsync(new AgentRunUpdateEvent(this.Id, update)).ConfigureAwait(false); + await context.AddEventAsync(new AgentRunUpdateEvent(this.Id, update), cancellationToken).ConfigureAwait(false); } } }); @@ -623,7 +631,8 @@ public static partial class AgentWorkflowBuilder /// Selects the next agent to participate in the group chat based on the provided chat history and team. /// /// The chat history to consider. - /// The to monitor for cancellation requests. The default is . + /// The to monitor for cancellation requests. + /// The default is . /// The next to speak. This agent must be part of the chat. protected internal abstract ValueTask SelectNextAgentAsync( IReadOnlyList history, @@ -633,7 +642,8 @@ public static partial class AgentWorkflowBuilder /// Filters the chat history before it's passed to the next agent. /// /// The chat history to filter. - /// The to monitor for cancellation requests. The default is . + /// The to monitor for cancellation requests. + /// The default is . /// The filtered chat history. protected internal virtual ValueTask> UpdateHistoryAsync( IReadOnlyList history, @@ -644,7 +654,8 @@ public static partial class AgentWorkflowBuilder /// Determines whether the group chat should be terminated based on the provided chat history and iteration count. /// /// The chat history to consider. - /// The to monitor for cancellation requests. The default is . + /// The to monitor for cancellation requests. + /// The default is . /// A indicating whether the chat should be terminated. protected internal virtual ValueTask ShouldTerminateAsync( IReadOnlyList history, @@ -789,35 +800,35 @@ public static partial class AgentWorkflowBuilder private GroupChatManager? _manager; protected override RouteBuilder ConfigureRoutes(RouteBuilder routeBuilder) => routeBuilder - .AddHandler((message, context) => this._pendingMessages.Add(new(ChatRole.User, message))) - .AddHandler((message, context) => this._pendingMessages.Add(message)) - .AddHandler>((messages, _) => this._pendingMessages.AddRange(messages)) - .AddHandler((messages, _) => this._pendingMessages.AddRange(messages)) // TODO: Remove once https://github.com/microsoft/agent-framework/issues/782 is addressed - .AddHandler>((messages, _) => this._pendingMessages.AddRange(messages)) // TODO: Remove once https://github.com/microsoft/agent-framework/issues/782 is addressed - .AddHandler(async (token, context) => + .AddHandler((message, context, _) => this._pendingMessages.Add(new(ChatRole.User, message))) + .AddHandler((message, context, _) => this._pendingMessages.Add(message)) + .AddHandler>((messages, _, __) => this._pendingMessages.AddRange(messages)) + .AddHandler((messages, _, __) => this._pendingMessages.AddRange(messages)) // TODO: Remove once https://github.com/microsoft/agent-framework/issues/782 is addressed + .AddHandler>((messages, _, __) => this._pendingMessages.AddRange(messages)) // TODO: Remove once https://github.com/microsoft/agent-framework/issues/782 is addressed + .AddHandler(async (token, context, cancellationToken) => { List messages = [.. this._pendingMessages]; this._pendingMessages.Clear(); this._manager ??= this._managerFactory(this._agents); - if (!await this._manager.ShouldTerminateAsync(messages).ConfigureAwait(false)) + if (!await this._manager.ShouldTerminateAsync(messages, cancellationToken).ConfigureAwait(false)) { - var filtered = await this._manager.UpdateHistoryAsync(messages).ConfigureAwait(false); + var filtered = await this._manager.UpdateHistoryAsync(messages, cancellationToken).ConfigureAwait(false); messages = filtered is null || ReferenceEquals(filtered, messages) ? messages : [.. filtered]; - if (await this._manager.SelectNextAgentAsync(messages).ConfigureAwait(false) is AIAgent nextAgent && + if (await this._manager.SelectNextAgentAsync(messages, cancellationToken).ConfigureAwait(false) is AIAgent nextAgent && this._agentMap.TryGetValue(nextAgent, out var executor)) { this._manager.IterationCount++; - await context.SendMessageAsync(messages, executor.Id).ConfigureAwait(false); - await context.SendMessageAsync(token, executor.Id).ConfigureAwait(false); + await context.SendMessageAsync(messages, executor.Id, cancellationToken).ConfigureAwait(false); + await context.SendMessageAsync(token, executor.Id, cancellationToken).ConfigureAwait(false); return; } } this._manager = null; - await context.YieldOutputAsync(messages).ConfigureAwait(false); + await context.YieldOutputAsync(messages, cancellationToken).ConfigureAwait(false); }); public ValueTask ResetAsync() diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/AggregatingExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/AggregatingExecutor.cs index c6cb36abe3..06f6ce243a 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/AggregatingExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/AggregatingExecutor.cs @@ -28,7 +28,7 @@ public class AggregatingExecutor(string id, private TAggregate? _runningAggregate; /// - public override ValueTask HandleAsync(TInput message, IWorkflowContext context) + public override ValueTask HandleAsync(TInput message, IWorkflowContext context, CancellationToken cancellationToken = default) { this._runningAggregate = aggregator(this._runningAggregate, message); return new(this._runningAggregate); @@ -37,7 +37,7 @@ public class AggregatingExecutor(string id, /// protected internal override async ValueTask OnCheckpointingAsync(IWorkflowContext context, CancellationToken cancellationToken = default) { - await context.QueueStateUpdateAsync(AggregateStateKey, this._runningAggregate).ConfigureAwait(false); + await context.QueueStateUpdateAsync(AggregateStateKey, this._runningAggregate, cancellationToken: cancellationToken).ConfigureAwait(false); await base.OnCheckpointingAsync(context, cancellationToken).ConfigureAwait(false); } @@ -47,6 +47,6 @@ public class AggregatingExecutor(string id, { await base.OnCheckpointRestoredAsync(context, cancellationToken).ConfigureAwait(false); - this._runningAggregate = await context.ReadStateAsync(AggregateStateKey).ConfigureAwait(false); + this._runningAggregate = await context.ReadStateAsync(AggregateStateKey, cancellationToken: cancellationToken).ConfigureAwait(false); } } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/AsyncBarrier.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/AsyncBarrier.cs index 99f8a1fa62..2cd0e84347 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/AsyncBarrier.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/AsyncBarrier.cs @@ -10,7 +10,7 @@ internal sealed class AsyncBarrier() { private readonly InitLocked> _completionSource = new(); - public async ValueTask JoinAsync(CancellationToken cancellation = default) + public async ValueTask JoinAsync(CancellationToken cancellationToken = default) { this._completionSource.Init(() => new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously)); TaskCompletionSource completionSource = this._completionSource.Get()!; @@ -19,10 +19,10 @@ internal sealed class AsyncBarrier() // should not cancel the entire barrier. TaskCompletionSource cancellationSource = new(); - using CancellationTokenRegistration registration = cancellation.Register(() => cancellationSource.SetResult(new())); + using CancellationTokenRegistration registration = cancellationToken.Register(() => cancellationSource.SetResult(new())); await Task.WhenAny(completionSource.Task, cancellationSource.Task).ConfigureAwait(false); - return !cancellation.IsCancellationRequested; + return !cancellationToken.IsCancellationRequested; } public bool ReleaseBarrier() diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/AsyncCoordinator.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/AsyncCoordinator.cs index 835a78ab2a..e1478396c8 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/AsyncCoordinator.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/AsyncCoordinator.cs @@ -12,12 +12,13 @@ internal sealed class AsyncCoordinator /// /// Wait for the Coordination owner to mark the next coordination point, then continue execution. /// - /// A cancellation token that can be used to cancel the wait. + /// The to monitor for cancellation requests. + /// The default is . /// /// A task that represents the asynchronous operation. The task result is /// if the wait was completed; otherwise, for example, if the wait was cancelled, . /// - public async ValueTask WaitForCoordinationAsync(CancellationToken cancellation = default) + public async ValueTask WaitForCoordinationAsync(CancellationToken cancellationToken = default) { // There is a chance that we might get a stale barrier that is getting released if there is a // release happening concurrently with this call. This is by design, and should be considered @@ -26,7 +27,7 @@ internal sealed class AsyncCoordinator ?? Interlocked.CompareExchange(ref this._coordinationBarrier, new(), null) ?? this._coordinationBarrier!; // Re-read after setting - return await actualBarrier.JoinAsync(cancellation).ConfigureAwait(false); + return await actualBarrier.JoinAsync(cancellationToken).ConfigureAwait(false); } /// diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/ChatProtocolExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/ChatProtocolExecutor.cs index afcc7e52cc..1ede5dbdea 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/ChatProtocolExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/ChatProtocolExecutor.cs @@ -30,19 +30,19 @@ internal abstract class ChatProtocolExecutor(string id, ChatProtocolExecutorOpti { if (this._stringMessageChatRole.HasValue) { - routeBuilder = routeBuilder.AddHandler((message, _) => this._pendingMessages.Add(new(this._stringMessageChatRole.Value, message))); + routeBuilder = routeBuilder.AddHandler((message, _, __) => this._pendingMessages.Add(new(this._stringMessageChatRole.Value, message))); } - return routeBuilder.AddHandler((message, _) => this._pendingMessages.Add(message)) - .AddHandler>((messages, _) => this._pendingMessages.AddRange(messages)) + return routeBuilder.AddHandler((message, _, __) => this._pendingMessages.Add(message)) + .AddHandler>((messages, _, __) => this._pendingMessages.AddRange(messages)) .AddHandler(this.TakeTurnAsync); } - public async ValueTask TakeTurnAsync(TurnToken token, IWorkflowContext context) + public async ValueTask TakeTurnAsync(TurnToken token, IWorkflowContext context, CancellationToken cancellationToken = default) { - await this.TakeTurnAsync(this._pendingMessages, context, token.EmitEvents).ConfigureAwait(false); + await this.TakeTurnAsync(this._pendingMessages, context, token.EmitEvents, cancellationToken).ConfigureAwait(false); this._pendingMessages = []; - await context.SendMessageAsync(token).ConfigureAwait(false); + await context.SendMessageAsync(token, cancellationToken: cancellationToken).ConfigureAwait(false); } protected abstract ValueTask TakeTurnAsync(List messages, IWorkflowContext context, bool? emitEvents, CancellationToken cancellationToken = default); @@ -54,7 +54,7 @@ internal abstract class ChatProtocolExecutor(string id, ChatProtocolExecutorOpti if (this._pendingMessages.Count > 0) { JsonElement messagesValue = this._pendingMessages.Serialize(); - messagesTask = context.QueueStateUpdateAsync(PendingMessagesStateKey, messagesValue).AsTask(); + messagesTask = context.QueueStateUpdateAsync(PendingMessagesStateKey, messagesValue, cancellationToken: cancellationToken).AsTask(); } await messagesTask.ConfigureAwait(false); @@ -62,7 +62,7 @@ internal abstract class ChatProtocolExecutor(string id, ChatProtocolExecutorOpti protected internal override async ValueTask OnCheckpointRestoredAsync(IWorkflowContext context, CancellationToken cancellationToken = default) { - JsonElement? messagesValue = await context.ReadStateAsync(PendingMessagesStateKey).ConfigureAwait(false); + JsonElement? messagesValue = await context.ReadStateAsync(PendingMessagesStateKey, cancellationToken: cancellationToken).ConfigureAwait(false); if (messagesValue.HasValue) { List messages = messagesValue.Value.DeserializeMessages(); diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/AsyncRunHandle.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/AsyncRunHandle.cs index 8fe6edce2a..de1a1dfc7f 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/AsyncRunHandle.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/AsyncRunHandle.cs @@ -44,19 +44,14 @@ internal sealed class AsyncRunHandle : ICheckpointingHandle, IAsyncDisposable } } - //private readonly AsyncCoordinator _waitForResponseCoordinator = new(); - - //public ValueTask WaitForNextInputAsync(CancellationToken cancellation = default) - // => this._waitForResponseCoordinator.WaitForCoordinationAsync(cancellation); - public string RunId => this._stepRunner.RunId; public IReadOnlyList Checkpoints => this._checkpointingHandle.Checkpoints; - public ValueTask GetStatusAsync(CancellationToken cancellation = default) - => this._eventStream.GetStatusAsync(cancellation); + public ValueTask GetStatusAsync(CancellationToken cancellationToken = default) + => this._eventStream.GetStatusAsync(cancellationToken); - public async IAsyncEnumerable TakeEventStreamAsync(bool blockOnPendingRequest, [EnumeratorCancellation] CancellationToken cancellation = default) + public async IAsyncEnumerable TakeEventStreamAsync(bool blockOnPendingRequest, [EnumeratorCancellation] CancellationToken cancellationToken = default) { //Debug.Assert(breakOnHalt); // Enforce single active enumerator (this runs when enumeration begins) @@ -68,7 +63,7 @@ internal sealed class AsyncRunHandle : ICheckpointingHandle, IAsyncDisposable CancellationTokenSource? linked = null; try { - linked = CancellationTokenSource.CreateLinkedTokenSource(cancellation, this._endRunSource.Token); + linked = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, this._endRunSource.Token); var token = linked.Token; // Build the inner stream before the loop so synchronous exceptions still release the gate @@ -92,21 +87,21 @@ internal sealed class AsyncRunHandle : ICheckpointingHandle, IAsyncDisposable } } - public ValueTask IsValidInputTypeAsync(CancellationToken cancellation = default) - => this._stepRunner.IsValidInputTypeAsync(cancellation); + public ValueTask IsValidInputTypeAsync(CancellationToken cancellationToken = default) + => this._stepRunner.IsValidInputTypeAsync(cancellationToken); - public async ValueTask EnqueueMessageAsync(T message, CancellationToken cancellation = default) + public async ValueTask EnqueueMessageAsync(T message, CancellationToken cancellationToken = default) { if (message is ExternalResponse response) { // EnqueueResponseAsync handles signaling - await this.EnqueueResponseAsync(response, cancellation) + await this.EnqueueResponseAsync(response, cancellationToken) .ConfigureAwait(false); return true; } - bool result = await this._stepRunner.EnqueueMessageAsync(message, cancellation) + bool result = await this._stepRunner.EnqueueMessageAsync(message, cancellationToken) .ConfigureAwait(false); // Signal the run loop that new input is available @@ -115,7 +110,7 @@ internal sealed class AsyncRunHandle : ICheckpointingHandle, IAsyncDisposable return result; } - public async ValueTask EnqueueMessageUntypedAsync([NotNull] object message, Type? declaredType = null, CancellationToken cancellation = default) + public async ValueTask EnqueueMessageUntypedAsync([NotNull] object message, Type? declaredType = null, CancellationToken cancellationToken = default) { if (declaredType?.IsInstanceOfType(message) == false) { @@ -125,7 +120,7 @@ internal sealed class AsyncRunHandle : ICheckpointingHandle, IAsyncDisposable if (declaredType != null && typeof(ExternalResponse).IsAssignableFrom(declaredType)) { // EnqueueResponseAsync handles signaling - await this.EnqueueResponseAsync((ExternalResponse)message, cancellation) + await this.EnqueueResponseAsync((ExternalResponse)message, cancellationToken) .ConfigureAwait(false); return true; @@ -133,13 +128,13 @@ internal sealed class AsyncRunHandle : ICheckpointingHandle, IAsyncDisposable else if (declaredType == null && message is ExternalResponse response) { // EnqueueResponseAsync handles signaling - await this.EnqueueResponseAsync(response, cancellation) + await this.EnqueueResponseAsync(response, cancellationToken) .ConfigureAwait(false); return true; } - bool result = await this._stepRunner.EnqueueMessageUntypedAsync(message, declaredType ?? message.GetType(), cancellation) + bool result = await this._stepRunner.EnqueueMessageUntypedAsync(message, declaredType ?? message.GetType(), cancellationToken) .ConfigureAwait(false); // Signal the run loop that new input is available @@ -148,9 +143,9 @@ internal sealed class AsyncRunHandle : ICheckpointingHandle, IAsyncDisposable return result; } - public async ValueTask EnqueueResponseAsync(ExternalResponse response, CancellationToken cancellation = default) + public async ValueTask EnqueueResponseAsync(ExternalResponse response, CancellationToken cancellationToken = default) { - await this._stepRunner.EnqueueResponseAsync(response, cancellation).ConfigureAwait(false); + await this._stepRunner.EnqueueResponseAsync(response, cancellationToken).ConfigureAwait(false); // Signal the run loop that new input is available this.SignalInputToRunLoop(); diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/AsyncRunHandleExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/AsyncRunHandleExtensions.cs index 65b60cd6a8..c7ac339a0c 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/AsyncRunHandleExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/AsyncRunHandleExtensions.cs @@ -14,33 +14,33 @@ internal static class AsyncRunHandleExtensions return new Checkpointed(run, runHandle); } - public static async ValueTask EnqueueAndStreamAsync(this AsyncRunHandle runHandle, TInput input, CancellationToken cancellation = default) + public static async ValueTask EnqueueAndStreamAsync(this AsyncRunHandle runHandle, TInput input, CancellationToken cancellationToken = default) { - await runHandle.EnqueueMessageAsync(input, cancellation).ConfigureAwait(false); + await runHandle.EnqueueMessageAsync(input, cancellationToken).ConfigureAwait(false); return new(runHandle); } - public static async ValueTask EnqueueUntypedAndStreamAsync(this AsyncRunHandle runHandle, object input, CancellationToken cancellation = default) + public static async ValueTask EnqueueUntypedAndStreamAsync(this AsyncRunHandle runHandle, object input, CancellationToken cancellationToken = default) { - await runHandle.EnqueueMessageUntypedAsync(input, cancellation: cancellation).ConfigureAwait(false); + await runHandle.EnqueueMessageUntypedAsync(input, cancellationToken: cancellationToken).ConfigureAwait(false); return new(runHandle); } - public static async ValueTask EnqueueAndRunAsync(this AsyncRunHandle runHandle, TInput input, CancellationToken cancellation = default) + public static async ValueTask EnqueueAndRunAsync(this AsyncRunHandle runHandle, TInput input, CancellationToken cancellationToken = default) { - await runHandle.EnqueueMessageAsync(input, cancellation).ConfigureAwait(false); + await runHandle.EnqueueMessageAsync(input, cancellationToken).ConfigureAwait(false); Run run = new(runHandle); - await run.RunToNextHaltAsync(cancellation).ConfigureAwait(false); + await run.RunToNextHaltAsync(cancellationToken).ConfigureAwait(false); return run; } - public static async ValueTask EnqueueUntypedAndRunAsync(this AsyncRunHandle runHandle, object input, CancellationToken cancellation = default) + public static async ValueTask EnqueueUntypedAndRunAsync(this AsyncRunHandle runHandle, object input, CancellationToken cancellationToken = default) { - await runHandle.EnqueueMessageUntypedAsync(input, cancellation: cancellation).ConfigureAwait(false); + await runHandle.EnqueueMessageUntypedAsync(input, cancellationToken: cancellationToken).ConfigureAwait(false); Run run = new(runHandle); - await run.RunToNextHaltAsync(cancellation).ConfigureAwait(false); + await run.RunToNextHaltAsync(cancellationToken).ConfigureAwait(false); return run; } } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/CallResult.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/CallResult.cs index 655aa1597a..952b48cd6c 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/CallResult.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/CallResult.cs @@ -1,6 +1,7 @@ // Copyright (c) Microsoft. All rights reserved. using System; +using System.Threading; using Microsoft.Shared.Diagnostics; namespace Microsoft.Agents.AI.Workflows.Execution; @@ -26,16 +27,22 @@ internal sealed class CallResult /// public Exception? Exception { get; init; } + /// + /// Indicated whether the call was cancelled (e.g., via a ). + /// + public bool IsCancelled { get; init; } + /// /// Indicates whether the call was successful. A call is considered successful if it returned /// without throwing an exception. /// - public bool IsSuccess => this.Exception is null; + public bool IsSuccess => this.Exception is null && !this.IsCancelled; - private CallResult(bool isVoid = false) + private CallResult(bool isVoid = false, bool isCancelled = false) { // Private constructor to enforce use of static methods. this.IsVoid = isVoid; + this.IsCancelled = isCancelled; } /// @@ -51,6 +58,14 @@ internal sealed class CallResult /// A indicating the result of the call. public static CallResult ReturnVoid() => new(isVoid: true); + /// + /// Create a indicating that the call was cancelled. + /// + /// A boolean specifying whether the call was void (was not expected to return + /// a value). + /// A indicating the result of the call. + public static CallResult Cancelled(bool wasVoid) => new(wasVoid, isCancelled: true); + /// /// Create a indicating that an exception was raised during the call. /// diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/IInputCoordinator.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/IInputCoordinator.cs deleted file mode 100644 index ed9f4bf22d..0000000000 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/IInputCoordinator.cs +++ /dev/null @@ -1,11 +0,0 @@ -// Copyright (c) Microsoft. All rights reserved. - -using System.Threading; -using System.Threading.Tasks; - -namespace Microsoft.Agents.AI.Workflows.Execution; - -internal interface IInputCoordinator -{ - ValueTask WaitForNextInputAsync(CancellationToken cancellation = default); -} diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/IRunEventStream.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/IRunEventStream.cs index 5ffce6ce36..dfc35c7566 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/IRunEventStream.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/IRunEventStream.cs @@ -15,7 +15,7 @@ internal interface IRunEventStream : IAsyncDisposable // this cannot be cancelled ValueTask StopAsync(); - ValueTask GetStatusAsync(CancellationToken cancellation = default); + ValueTask GetStatusAsync(CancellationToken cancellationToken = default); - IAsyncEnumerable TakeEventStreamAsync(bool blockOnPendingRequest, CancellationToken cancellation = default); + IAsyncEnumerable TakeEventStreamAsync(bool blockOnPendingRequest, CancellationToken cancellationToken = default); } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/IRunnerContext.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/IRunnerContext.cs index b3a988306c..f3fc762336 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/IRunnerContext.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/IRunnerContext.cs @@ -1,16 +1,17 @@ // Copyright (c) Microsoft. All rights reserved. using System.Collections.Generic; +using System.Threading; using System.Threading.Tasks; namespace Microsoft.Agents.AI.Workflows.Execution; internal interface IRunnerContext : IExternalRequestSink, ISuperStepJoinContext { - ValueTask AddEventAsync(WorkflowEvent workflowEvent); - ValueTask SendMessageAsync(string sourceId, object message, string? targetId = null); + ValueTask AddEventAsync(WorkflowEvent workflowEvent, CancellationToken cancellationToken = default); + ValueTask SendMessageAsync(string sourceId, object message, string? targetId = null, CancellationToken cancellationToken = default); - ValueTask AdvanceAsync(); + ValueTask AdvanceAsync(CancellationToken cancellationToken = default); IWorkflowContext Bind(string executorId, Dictionary? traceContext = null); - ValueTask EnsureExecutorAsync(string executorId, IStepTracer? tracer); + ValueTask EnsureExecutorAsync(string executorId, IStepTracer? tracer, CancellationToken cancellationToken = default); } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/ISuperStepJoinContext.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/ISuperStepJoinContext.cs index 1e27fe6b4d..ed4ded17d6 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/ISuperStepJoinContext.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/ISuperStepJoinContext.cs @@ -10,8 +10,8 @@ internal interface ISuperStepJoinContext { bool WithCheckpointing { get; } - ValueTask ForwardWorkflowEventAsync(WorkflowEvent workflowEvent, CancellationToken cancellation = default); - ValueTask SendMessageAsync(string senderId, [DisallowNull] TMessage message, CancellationToken cancellation = default); + ValueTask ForwardWorkflowEventAsync(WorkflowEvent workflowEvent, CancellationToken cancellationToken = default); + ValueTask SendMessageAsync(string senderId, [DisallowNull] TMessage message, CancellationToken cancellationToken = default); - ValueTask AttachSuperstepAsync(ISuperStepRunner superStepRunner, CancellationToken cancellation = default); + ValueTask AttachSuperstepAsync(ISuperStepRunner superStepRunner, CancellationToken cancellationToken = default); } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/ISuperStepRunner.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/ISuperStepRunner.cs index 7384dfc4d0..a7923a7d9b 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/ISuperStepRunner.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/ISuperStepRunner.cs @@ -15,11 +15,11 @@ internal interface ISuperStepRunner bool HasUnservicedRequests { get; } bool HasUnprocessedMessages { get; } - ValueTask EnqueueResponseAsync(ExternalResponse response, CancellationToken cancellation = default); + ValueTask EnqueueResponseAsync(ExternalResponse response, CancellationToken cancellationToken = default); - ValueTask IsValidInputTypeAsync(CancellationToken cancellation = default); - ValueTask EnqueueMessageAsync(T message, CancellationToken cancellation = default); - ValueTask EnqueueMessageUntypedAsync(object message, Type declaredType, CancellationToken cancellation = default); + ValueTask IsValidInputTypeAsync(CancellationToken cancellationToken = default); + ValueTask EnqueueMessageAsync(T message, CancellationToken cancellationToken = default); + ValueTask EnqueueMessageUntypedAsync(object message, Type declaredType, CancellationToken cancellationToken = default); ConcurrentEventSink OutgoingEvents { get; } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/InputWaiter.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/InputWaiter.cs index f86f3721bb..d50f284f48 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/InputWaiter.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/InputWaiter.cs @@ -33,10 +33,10 @@ internal sealed class InputWaiter : IDisposable } } - public Task WaitForInputAsync(CancellationToken cancellation = default) => this.WaitForInputAsync(null, cancellation); + public Task WaitForInputAsync(CancellationToken cancellationToken = default) => this.WaitForInputAsync(null, cancellationToken); - public async Task WaitForInputAsync(TimeSpan? timeout = null, CancellationToken cancellation = default) + public async Task WaitForInputAsync(TimeSpan? timeout = null, CancellationToken cancellationToken = default) { - await this._inputSignal.WaitAsync(timeout ?? TimeSpan.FromMilliseconds(-1), cancellation).ConfigureAwait(false); + await this._inputSignal.WaitAsync(timeout ?? TimeSpan.FromMilliseconds(-1), cancellationToken).ConfigureAwait(false); } } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/LockstepRunEventStream.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/LockstepRunEventStream.cs index a66f0af978..b47a692113 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/LockstepRunEventStream.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/LockstepRunEventStream.cs @@ -22,7 +22,7 @@ internal sealed class LockstepRunEventStream : IRunEventStream private readonly ISuperStepRunner _stepRunner; - public ValueTask GetStatusAsync(CancellationToken cancellation = default) => new(this.RunStatus); + public ValueTask GetStatusAsync(CancellationToken cancellationToken = default) => new(this.RunStatus); public LockstepRunEventStream(ISuperStepRunner stepRunner) { @@ -36,7 +36,7 @@ internal sealed class LockstepRunEventStream : IRunEventStream // No-op for lockstep execution } - public async IAsyncEnumerable TakeEventStreamAsync(bool blockOnPendingRequest, [EnumeratorCancellation] CancellationToken cancellation = default) + public async IAsyncEnumerable TakeEventStreamAsync(bool blockOnPendingRequest, [EnumeratorCancellation] CancellationToken cancellationToken = default) { #if NET ObjectDisposedException.ThrowIf(Volatile.Read(ref this._isDisposed) == 1, this); @@ -47,7 +47,7 @@ internal sealed class LockstepRunEventStream : IRunEventStream } #endif - CancellationTokenSource linkedSource = CancellationTokenSource.CreateLinkedTokenSource(this._stopCancellation.Token, cancellation); + CancellationTokenSource linkedSource = CancellationTokenSource.CreateLinkedTokenSource(this._stopCancellation.Token, cancellationToken); ConcurrentQueue eventSink = []; diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/MessageRouter.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/MessageRouter.cs index 0a23105c74..10ce345ad8 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/MessageRouter.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/MessageRouter.cs @@ -4,6 +4,7 @@ using System; using System.Collections.Generic; using System.Diagnostics.CodeAnalysis; using System.Linq; +using System.Threading; using System.Threading.Tasks; using Microsoft.Agents.AI.Workflows.Checkpointing; using Microsoft.Shared.Diagnostics; @@ -11,12 +12,14 @@ using CatchAllF = System.Func< Microsoft.Agents.AI.Workflows.PortableValue, // message Microsoft.Agents.AI.Workflows.IWorkflowContext, // context + System.Threading.CancellationToken, // cancellation System.Threading.Tasks.ValueTask >; using MessageHandlerF = System.Func< object, // message Microsoft.Agents.AI.Workflows.IWorkflowContext, // context + System.Threading.CancellationToken, // cancellation System.Threading.Tasks.ValueTask >; @@ -56,7 +59,7 @@ internal sealed class MessageRouter public HashSet DefaultOutputTypes { get; } - public async ValueTask RouteMessageAsync(object message, IWorkflowContext context, bool requireRoute = false) + public async ValueTask RouteMessageAsync(object message, IWorkflowContext context, bool requireRoute = false, CancellationToken cancellationToken = default) { Throw.IfNull(message); @@ -74,13 +77,13 @@ internal sealed class MessageRouter { if (this._typedHandlers.TryGetValue(message.GetType(), out MessageHandlerF? handler)) { - result = await handler(message, context).ConfigureAwait(false); + result = await handler(message, context, cancellationToken).ConfigureAwait(false); } else if (this.HasCatchAll) { portableValue ??= new PortableValue(message); - result = await this._catchAllFunc(portableValue, context).ConfigureAwait(false); + result = await this._catchAllFunc(portableValue, context, cancellationToken).ConfigureAwait(false); } } catch (Exception e) diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/StreamingRunEventStream.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/StreamingRunEventStream.cs index 0f941ea733..718ebcd11c 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/StreamingRunEventStream.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/StreamingRunEventStream.cs @@ -50,10 +50,10 @@ internal sealed class StreamingRunEventStream : IRunEventStream } } - private async Task RunLoopAsync(CancellationToken cancellation) + private async Task RunLoopAsync(CancellationToken cancellationToken) { using CancellationTokenSource errorSource = new(); - CancellationTokenSource linkedSource = CancellationTokenSource.CreateLinkedTokenSource(errorSource.Token, cancellation); + CancellationTokenSource linkedSource = CancellationTokenSource.CreateLinkedTokenSource(errorSource.Token, cancellationToken); // Subscribe to events - they will flow directly to the channel as they're raised this._stepRunner.OutgoingEvents.EventRaised += OnEventRaisedAsync; @@ -62,7 +62,7 @@ internal sealed class StreamingRunEventStream : IRunEventStream { // Wait for the first input before starting // The consumer will call EnqueueMessageAsync which signals the run loop - await this._inputWaiter.WaitForInputAsync(cancellation: linkedSource.Token).ConfigureAwait(false); + await this._inputWaiter.WaitForInputAsync(cancellationToken: linkedSource.Token).ConfigureAwait(false); this._runStatus = RunStatus.Running; @@ -134,7 +134,7 @@ internal sealed class StreamingRunEventStream : IRunEventStream public async IAsyncEnumerable TakeEventStreamAsync( bool blockOnPendingRequest, - [EnumeratorCancellation] CancellationToken cancellation = default) + [EnumeratorCancellation] CancellationToken cancellationToken = default) { // Get the current epoch - we'll only respond to completion signals from this epoch or later int myEpoch = Volatile.Read(ref this._completionEpoch) + 1; @@ -143,7 +143,7 @@ internal sealed class StreamingRunEventStream : IRunEventStream // Note: When cancellation is requested, ReadAllAsync may throw OperationCanceledException // or may complete the enumeration. We check IsCancellationRequested explicitly at superstep // boundaries to ensure clean cancellation. - await foreach (WorkflowEvent evt in this._eventChannel.Reader.ReadAllAsync(cancellation).ConfigureAwait(false)) + await foreach (WorkflowEvent evt in this._eventChannel.Reader.ReadAllAsync(cancellationToken).ConfigureAwait(false)) { // Filter out internal signals used for run loop coordination if (evt is InternalHaltSignal completionSignal) @@ -156,7 +156,7 @@ internal sealed class StreamingRunEventStream : IRunEventStream // Check for cancellation at superstep boundaries (before processing completion signal) // This allows consumers to stop reading events cleanly between supersteps - if (cancellation.IsCancellationRequested) + if (cancellationToken.IsCancellationRequested) { yield break; } @@ -186,7 +186,7 @@ internal sealed class StreamingRunEventStream : IRunEventStream yield break; } - if (cancellation.IsCancellationRequested) + if (cancellationToken.IsCancellationRequested) { yield break; } @@ -195,7 +195,7 @@ internal sealed class StreamingRunEventStream : IRunEventStream } } - public ValueTask GetStatusAsync(CancellationToken cancellation = default) + public ValueTask GetStatusAsync(CancellationToken cancellationToken = default) { // Thread-safe read of status (enum is read atomically on most platforms) return new ValueTask(this._runStatus); diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Executor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Executor.cs index e1df4fb37e..cd97fcaea4 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Executor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Executor.cs @@ -90,10 +90,12 @@ public abstract class Executor : IIdentified /// The "declared" type of the message (captured when it was being sent). This is /// used to enable routing messages as their base types, in absence of true polymorphic type routing. /// The workflow context in which the executor executes. + /// The to monitor for cancellation requests. + /// The default is . /// A ValueTask representing the asynchronous operation, wrapping the output from the executor. /// No handler found for the message type. /// An exception is generated while handling the message. - public async ValueTask ExecuteAsync(object message, TypeId messageType, IWorkflowContext context) + public async ValueTask ExecuteAsync(object message, TypeId messageType, IWorkflowContext context, CancellationToken cancellationToken = default) { using var activity = s_activitySource.StartActivity(ActivityNames.ExecutorProcess, ActivityKind.Internal); activity?.SetTag(Tags.ExecutorId, this.Id) @@ -101,9 +103,9 @@ public abstract class Executor : IIdentified .SetTag(Tags.MessageType, messageType.TypeName) .CreateSourceLinks(context.TraceContext); - await context.AddEventAsync(new ExecutorInvokedEvent(this.Id, message)).ConfigureAwait(false); + await context.AddEventAsync(new ExecutorInvokedEvent(this.Id, message), cancellationToken).ConfigureAwait(false); - CallResult? result = await this.Router.RouteMessageAsync(message, context, requireRoute: true) + CallResult? result = await this.Router.RouteMessageAsync(message, context, requireRoute: true, cancellationToken) .ConfigureAwait(false); ExecutorEvent executionResult; @@ -116,7 +118,7 @@ public abstract class Executor : IIdentified executionResult = new ExecutorFailedEvent(this.Id, result.Exception); } - await context.AddEventAsync(executionResult).ConfigureAwait(false); + await context.AddEventAsync(executionResult, cancellationToken).ConfigureAwait(false); if (result is null) { @@ -137,11 +139,11 @@ public abstract class Executor : IIdentified // If we had a real return type, raise it as a SendMessage; TODO: Should we have a way to disable this behaviour? if (result.Result is not null && this.Options.AutoSendMessageHandlerResultObject) { - await context.SendMessageAsync(result.Result).ConfigureAwait(false); + await context.SendMessageAsync(result.Result, cancellationToken: cancellationToken).ConfigureAwait(false); } if (result.Result is not null && this.Options.AutoYieldOutputHandlerResultObject) { - await context.YieldOutputAsync(result.Result).ConfigureAwait(false); + await context.YieldOutputAsync(result.Result, cancellationToken).ConfigureAwait(false); } return result.Result; @@ -152,7 +154,8 @@ public abstract class Executor : IIdentified /// /// The workflow context. /// A ValueTask representing the asynchronous operation. - /// The to monitor for cancellation requests. The default is . + /// The to monitor for cancellation requests. + /// The default is . protected internal virtual ValueTask OnCheckpointingAsync(IWorkflowContext context, CancellationToken cancellationToken = default) => default; /// @@ -160,7 +163,8 @@ public abstract class Executor : IIdentified /// /// The workflow context. /// A ValueTask representing the asynchronous operation. - /// The to monitor for cancellation requests. The default is . + /// The to monitor for cancellation requests. + /// The default is . protected internal virtual ValueTask OnCheckpointRestoredAsync(IWorkflowContext context, CancellationToken cancellationToken = default) => default; /// @@ -210,7 +214,7 @@ public abstract class Executor(string id, ExecutorOptions? options = nul routeBuilder.AddHandler(this.HandleAsync); /// - public abstract ValueTask HandleAsync(TInput message, IWorkflowContext context); + public abstract ValueTask HandleAsync(TInput message, IWorkflowContext context, CancellationToken cancellationToken = default); } /// @@ -229,5 +233,5 @@ public abstract class Executor(string id, ExecutorOptions? opti routeBuilder.AddHandler(this.HandleAsync); /// - public abstract ValueTask HandleAsync(TInput message, IWorkflowContext context); + public abstract ValueTask HandleAsync(TInput message, IWorkflowContext context, CancellationToken cancellationToken = default); } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/FunctionExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/FunctionExecutor.cs index 0db15f42bf..dd692165c7 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/FunctionExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/FunctionExecutor.cs @@ -29,7 +29,7 @@ public class FunctionExecutor(string id, } /// - public override ValueTask HandleAsync(TInput message, IWorkflowContext context) => handlerAsync(message, context, default); + public override ValueTask HandleAsync(TInput message, IWorkflowContext context, CancellationToken cancellationToken) => handlerAsync(message, context, cancellationToken); /// /// Creates a new instance of the class. @@ -65,7 +65,7 @@ public class FunctionExecutor(string id, } /// - public override ValueTask HandleAsync(TInput message, IWorkflowContext context) => handlerAsync(message, context, default); + public override ValueTask HandleAsync(TInput message, IWorkflowContext context, CancellationToken cancellationToken) => handlerAsync(message, context, cancellationToken); /// /// Creates a new instance of the class. diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/IWorkflowContext.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/IWorkflowContext.cs index eed3b15d72..efc074c2d4 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/IWorkflowContext.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/IWorkflowContext.cs @@ -1,6 +1,7 @@ // Copyright (c) Microsoft. All rights reserved. using System.Collections.Generic; +using System.Threading; using System.Threading.Tasks; namespace Microsoft.Agents.AI.Workflows; @@ -15,8 +16,10 @@ public interface IWorkflowContext /// end of the current SuperStep. /// /// The event to be raised. + /// The to monitor for cancellation requests. + /// The default is . /// A representing the asynchronous operation. - ValueTask AddEventAsync(WorkflowEvent workflowEvent); + ValueTask AddEventAsync(WorkflowEvent workflowEvent, CancellationToken cancellationToken = default); /// /// Queues a message to be sent to connected executors. The message will be sent during the next SuperStep. @@ -25,8 +28,22 @@ public interface IWorkflowContext /// An optional identifier of the target executor. If null, the message is sent to all connected /// executors. If the target executor is not connected from this executor via an edge, it will still not receive the /// message. + /// The to monitor for cancellation requests. + /// The default is . /// A representing the asynchronous operation. - ValueTask SendMessageAsync(object message, string? targetId = null); + ValueTask SendMessageAsync(object message, string? targetId = null, CancellationToken cancellationToken = default); + +#if NET // What's the right way to do this so we do not make life a misery for netstandard2.0 targets? + // What's the value if they have to still write `cancellationToken: cancellationToken` to skip the targetId parameter? + // TODO: Remove this? (Maybe not: NET will eventually be the only target framework, right?) + /// + /// Queues a message to be sent to connected executors. The message will be sent during the next SuperStep. + /// + /// The message to be sent. + /// The to monitor for cancellation requests. + /// A representing the asynchronous operation. + ValueTask SendMessageAsync(object message, CancellationToken cancellationToken) => this.SendMessageAsync(message, null, cancellationToken); +#endif /// /// Adds an output value to the workflow's output queue. These outputs will be bubbled out of the workflow using the @@ -37,8 +54,10 @@ public interface IWorkflowContext /// types of registered message handlers are considered output types, unless otherwise specified using . /// /// The output value to be returned. + /// The to monitor for cancellation requests. + /// The default is . /// A representing the asynchronous operation. - ValueTask YieldOutputAsync(object output); + ValueTask YieldOutputAsync(object output, CancellationToken cancellationToken = default); /// /// Adds a request to "halt" workflow execution at the end of the current SuperStep. @@ -54,15 +73,32 @@ public interface IWorkflowContext /// The key of the state value. /// An optional name that specifies the scope to read.If null, the default scope is /// used. + /// The to monitor for cancellation requests. + /// The default is . /// A representing the asynchronous operation. - ValueTask ReadStateAsync(string key, string? scopeName = null); + ValueTask ReadStateAsync(string key, string? scopeName = null, CancellationToken cancellationToken = default); + +#if NET // See above for musings about this construction + /// + /// Reads a state value from the workflow's state store. If no scope is provided, the executor's + /// default scope is used. + /// + /// The type of the state value. + /// The key of the state value. + /// The to monitor for cancellation requests. + /// A representing the asynchronous operation. + ValueTask ReadStateAsync(string key, CancellationToken cancellationToken) => this.ReadStateAsync(key, null, cancellationToken); + +#endif /// /// Asynchronously reads all state keys within the specified scope. /// /// An optional name that specifies the scope to read. If null, the default scope is /// used. - ValueTask> ReadStateKeysAsync(string? scopeName = null); + /// The to monitor for cancellation requests. + /// The default is . + ValueTask> ReadStateKeysAsync(string? scopeName = null, CancellationToken cancellationToken = default); /// /// Asynchronously updates the state of a queue entry identified by the specified key and optional scope. @@ -77,8 +113,27 @@ public interface IWorkflowContext /// implementation. /// An optional name that specifies the scope to update. If null, the default scope is /// used. + /// The to monitor for cancellation requests. + /// The default is . /// A ValueTask that represents the asynchronous update operation. - ValueTask QueueStateUpdateAsync(string key, T? value, string? scopeName = null); + ValueTask QueueStateUpdateAsync(string key, T? value, string? scopeName = null, CancellationToken cancellationToken = default); + +#if NET // See above for musings about this construction + /// + /// Asynchronously updates the state of a queue entry identified by the specified key and optional scope. + /// + /// + /// Subsequent reads by this executor will result in the new value of the state. Other executors will only see + /// the new state starting from the next SuperStep. + /// + /// The type of the value to associate with the queue entry. + /// The unique identifier for the queue entry to update. Cannot be null or empty. + /// The value to set for the queue entry. If null, the entry's state may be cleared or reset depending on + /// implementation. + /// The to monitor for cancellation requests. + /// A ValueTask that represents the asynchronous update operation. + ValueTask QueueStateUpdateAsync(string key, T? value, CancellationToken cancellationToken) => this.QueueStateUpdateAsync(key, value, null, cancellationToken); +#endif /// /// Asynchronously clears all state entries within the specified scope. @@ -90,8 +145,25 @@ public interface IWorkflowContext /// see the cleared state starting from the next SuperStep. /// /// An optional name that specifies the scope to clear. If null, the default scope is used. + /// The to monitor for cancellation requests. + /// The default is . /// A ValueTask that represents the asynchronous clear operation. - ValueTask QueueClearScopeAsync(string? scopeName = null); + ValueTask QueueClearScopeAsync(string? scopeName = null, CancellationToken cancellationToken = default); + +#if NET // See above for musings about this construction + /// + /// Asynchronously clears all state entries within the specified scope. + /// + /// This semantically equivalent to retrieving all keys in the scope and deleting them one-by-one. + /// + /// + /// Subsequent reads by this executor will not find any entries in the cleared scope. Other executors will only + /// see the cleared state starting from the next SuperStep. + /// + /// The to monitor for cancellation requests. + /// A ValueTask that represents the asynchronous clear operation. + ValueTask QueueClearScopeAsync(CancellationToken cancellationToken) => this.QueueClearScopeAsync(null, cancellationToken); +#endif /// /// The trace context associated with the current message about to be processed by the executor, if any. diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunner.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunner.cs index 6f4cef55c5..6d7a4bd691 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunner.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunner.cs @@ -46,14 +46,14 @@ internal sealed class InProcessRunner : ISuperStepRunner, ICheckpointingHandle public string StartExecutorId { get; } private readonly HashSet _knownValidInputTypes; - public async ValueTask IsValidInputTypeAsync(Type messageType, CancellationToken cancellation = default) + public async ValueTask IsValidInputTypeAsync(Type messageType, CancellationToken cancellationToken = default) { if (this._knownValidInputTypes.Contains(messageType)) { return true; } - Executor startingExecutor = await this.RunContext.EnsureExecutorAsync(this.Workflow.StartExecutorId, tracer: null).ConfigureAwait(false); + Executor startingExecutor = await this.RunContext.EnsureExecutorAsync(this.Workflow.StartExecutorId, tracer: null, cancellationToken).ConfigureAwait(false); if (startingExecutor.CanHandle(messageType)) { this._knownValidInputTypes.Add(messageType); @@ -63,10 +63,10 @@ internal sealed class InProcessRunner : ISuperStepRunner, ICheckpointingHandle return false; } - public ValueTask IsValidInputTypeAsync(CancellationToken cancellation = default) - => this.IsValidInputTypeAsync(typeof(T), cancellation); + public ValueTask IsValidInputTypeAsync(CancellationToken cancellationToken = default) + => this.IsValidInputTypeAsync(typeof(T), cancellationToken); - public async ValueTask EnqueueMessageUntypedAsync(object message, Type declaredType, CancellationToken cancellation = default) + public async ValueTask EnqueueMessageUntypedAsync(object message, Type declaredType, CancellationToken cancellationToken = default) { this.RunContext.CheckEnded(); Throw.IfNull(message); @@ -78,7 +78,7 @@ internal sealed class InProcessRunner : ISuperStepRunner, ICheckpointingHandle // Check that the type of the incoming message is compatible with the starting executor's // input type. - if (!await this.IsValidInputTypeAsync(declaredType, cancellation).ConfigureAwait(false)) + if (!await this.IsValidInputTypeAsync(declaredType, cancellationToken).ConfigureAwait(false)) { return false; } @@ -87,13 +87,13 @@ internal sealed class InProcessRunner : ISuperStepRunner, ICheckpointingHandle return true; } - public ValueTask EnqueueMessageAsync(T message, CancellationToken cancellation = default) - => this.EnqueueMessageUntypedAsync(Throw.IfNull(message), typeof(T), cancellation); + public ValueTask EnqueueMessageAsync(T message, CancellationToken cancellationToken = default) + => this.EnqueueMessageUntypedAsync(Throw.IfNull(message), typeof(T), cancellationToken); - public ValueTask EnqueueMessageAsync(object message, CancellationToken cancellation = default) - => this.EnqueueMessageUntypedAsync(Throw.IfNull(message), message.GetType(), cancellation); + public ValueTask EnqueueMessageUntypedAsync(object message, CancellationToken cancellationToken = default) + => this.EnqueueMessageUntypedAsync(Throw.IfNull(message), message.GetType(), cancellationToken); - ValueTask ISuperStepRunner.EnqueueResponseAsync(ExternalResponse response, CancellationToken cancellation) + ValueTask ISuperStepRunner.EnqueueResponseAsync(ExternalResponse response, CancellationToken cancellationToken) { // TODO: Check that there exists a corresponding input port? return this.RunContext.AddExternalResponseAsync(response); @@ -110,13 +110,13 @@ internal sealed class InProcessRunner : ISuperStepRunner, ICheckpointingHandle private ValueTask RaiseWorkflowEventAsync(WorkflowEvent workflowEvent) => this.OutgoingEvents.EnqueueAsync(workflowEvent); - public ValueTask BeginStreamAsync(ExecutionMode mode, CancellationToken cancellation = default) + public ValueTask BeginStreamAsync(ExecutionMode mode, CancellationToken cancellationToken = default) { this.RunContext.CheckEnded(); return new(new AsyncRunHandle(this, this, mode)); } - public async ValueTask ResumeStreamAsync(ExecutionMode mode, CheckpointInfo fromCheckpoint, CancellationToken cancellation = default) + public async ValueTask ResumeStreamAsync(ExecutionMode mode, CheckpointInfo fromCheckpoint, CancellationToken cancellationToken = default) { this.RunContext.CheckEnded(); Throw.IfNull(fromCheckpoint); @@ -125,7 +125,7 @@ internal sealed class InProcessRunner : ISuperStepRunner, ICheckpointingHandle throw new InvalidOperationException("This runner was not configured with a CheckpointManager, so it cannot restore checkpoints."); } - await this.RestoreCheckpointAsync(fromCheckpoint, cancellation).ConfigureAwait(false); + await this.RestoreCheckpointAsync(fromCheckpoint, cancellationToken).ConfigureAwait(false); return new AsyncRunHandle(this, this, mode); } @@ -142,7 +142,7 @@ internal sealed class InProcessRunner : ISuperStepRunner, ICheckpointingHandle return false; } - StepContext currentStep = await this.RunContext.AdvanceAsync().ConfigureAwait(false); + StepContext currentStep = await this.RunContext.AdvanceAsync(cancellationToken).ConfigureAwait(false); if (currentStep.HasMessages || this.RunContext.HasQueuedExternalDeliveries || @@ -150,7 +150,7 @@ internal sealed class InProcessRunner : ISuperStepRunner, ICheckpointingHandle { try { - await this.RunSuperstepAsync(currentStep).ConfigureAwait(false); + await this.RunSuperstepAsync(currentStep, cancellationToken).ConfigureAwait(false); } catch (OperationCanceledException) { } @@ -165,9 +165,9 @@ internal sealed class InProcessRunner : ISuperStepRunner, ICheckpointingHandle return false; } - private async ValueTask DeliverMessagesAsync(string receiverId, ConcurrentQueue envelopes) + private async ValueTask DeliverMessagesAsync(string receiverId, ConcurrentQueue envelopes, CancellationToken cancellationToken) { - Executor executor = await this.RunContext.EnsureExecutorAsync(receiverId, this.StepTracer).ConfigureAwait(false); + Executor executor = await this.RunContext.EnsureExecutorAsync(receiverId, this.StepTracer, cancellationToken).ConfigureAwait(false); this.StepTracer.TraceActivated(receiverId); while (envelopes.TryDequeue(out var envelope)) @@ -175,19 +175,20 @@ internal sealed class InProcessRunner : ISuperStepRunner, ICheckpointingHandle await executor.ExecuteAsync( envelope.Message, envelope.MessageType, - this.RunContext.Bind(receiverId, envelope.TraceContext) + this.RunContext.Bind(receiverId, envelope.TraceContext), + cancellationToken ).ConfigureAwait(false); } } - private async ValueTask RunSuperstepAsync(StepContext currentStep) + private async ValueTask RunSuperstepAsync(StepContext currentStep, CancellationToken cancellationToken) { await this.RaiseWorkflowEventAsync(this.StepTracer.Advance(currentStep)).ConfigureAwait(false); // Deliver the messages and queue the next step List receiverTasks = currentStep.QueuedMessages.Keys - .Select(receiverId => this.DeliverMessagesAsync(receiverId, currentStep.MessagesFor(receiverId)).AsTask()) + .Select(receiverId => this.DeliverMessagesAsync(receiverId, currentStep.MessagesFor(receiverId), cancellationToken).AsTask()) .ToList(); // TODO: Should we let the user specify that they want strictly turn-based execution of the edges, vs. concurrent? @@ -202,12 +203,12 @@ internal sealed class InProcessRunner : ISuperStepRunner, ICheckpointingHandle List subworkflowTasks = new(); foreach (ISuperStepRunner subworkflowRunner in this.RunContext.JoinedSubworkflowRunners) { - subworkflowTasks.Add(subworkflowRunner.RunSuperStepAsync(CancellationToken.None).AsTask()); + subworkflowTasks.Add(subworkflowRunner.RunSuperStepAsync(cancellationToken).AsTask()); } await Task.WhenAll(subworkflowTasks).ConfigureAwait(false); - await this.CheckpointAsync().ConfigureAwait(false); + await this.CheckpointAsync(cancellationToken).ConfigureAwait(false); await this.RaiseWorkflowEventAsync(this.StepTracer.Complete(this.RunContext.NextStepHasActions, this.RunContext.HasUnservicedRequests)) .ConfigureAwait(false); diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunnerContext.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunnerContext.cs index cd3f763ae8..f707c72be0 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunnerContext.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunnerContext.cs @@ -57,7 +57,7 @@ internal sealed class InProcessRunnerContext : IRunnerContext this.OutgoingEvents = outgoingEvents; } - public async ValueTask EnsureExecutorAsync(string executorId, IStepTracer? tracer) + public async ValueTask EnsureExecutorAsync(string executorId, IStepTracer? tracer, CancellationToken cancellationToken = default) { this.CheckEnded(); Task executorTask = this._executors.GetOrAdd(executorId, CreateExecutorAsync); @@ -88,9 +88,9 @@ internal sealed class InProcessRunnerContext : IRunnerContext return await executorTask.ConfigureAwait(false); } - public async ValueTask> GetStartingExecutorInputTypesAsync(CancellationToken cancellation = default) + public async ValueTask> GetStartingExecutorInputTypesAsync(CancellationToken cancellationToken = default) { - Executor startingExecutor = await this.EnsureExecutorAsync(this._workflow.StartExecutorId, tracer: null) + Executor startingExecutor = await this.EnsureExecutorAsync(this._workflow.StartExecutorId, tracer: null, cancellationToken) .ConfigureAwait(false); return startingExecutor.InputTypes; @@ -145,7 +145,7 @@ internal sealed class InProcessRunnerContext : IRunnerContext public bool HasUnservicedRequests => !this._externalRequests.IsEmpty || this._joinedSubworkflowRunners.Any(joinedRunner => joinedRunner.HasUnservicedRequests); - public async ValueTask AdvanceAsync() + public async ValueTask AdvanceAsync(CancellationToken cancellationToken = default) { this.CheckEnded(); @@ -159,7 +159,7 @@ internal sealed class InProcessRunnerContext : IRunnerContext return Interlocked.Exchange(ref this._nextStep, new StepContext()); } - public ValueTask AddEventAsync(WorkflowEvent workflowEvent) + public ValueTask AddEventAsync(WorkflowEvent workflowEvent, CancellationToken cancellationToken = default) { this.CheckEnded(); return this.OutgoingEvents.EnqueueAsync(workflowEvent); @@ -168,7 +168,7 @@ internal sealed class InProcessRunnerContext : IRunnerContext private static readonly string s_namespace = typeof(IWorkflowContext).Namespace!; private static readonly ActivitySource s_activitySource = new(s_namespace); - public async ValueTask SendMessageAsync(string sourceId, object message, string? targetId = null) + public async ValueTask SendMessageAsync(string sourceId, object message, string? targetId = null, CancellationToken cancellationToken = default) { using Activity? activity = s_activitySource.StartActivity(ActivityNames.MessageSend, ActivityKind.Producer); // Create a carrier for trace context propagation @@ -231,19 +231,19 @@ internal sealed class InProcessRunnerContext : IRunnerContext OutputFilter outputFilter, Dictionary? traceContext) : IWorkflowContext { - public ValueTask AddEventAsync(WorkflowEvent workflowEvent) => RunnerContext.AddEventAsync(workflowEvent); + public ValueTask AddEventAsync(WorkflowEvent workflowEvent, CancellationToken cancellationToken = default) => RunnerContext.AddEventAsync(workflowEvent, cancellationToken); - public ValueTask SendMessageAsync(object message, string? targetId = null) + public ValueTask SendMessageAsync(object message, string? targetId = null, CancellationToken cancellationToken = default) { - return RunnerContext.SendMessageAsync(ExecutorId, message, targetId); + return RunnerContext.SendMessageAsync(ExecutorId, message, targetId, cancellationToken); } - public async ValueTask YieldOutputAsync(object output) + public async ValueTask YieldOutputAsync(object output, CancellationToken cancellationToken = default) { RunnerContext.CheckEnded(); Throw.IfNull(output); - Executor sourceExecutor = await RunnerContext.EnsureExecutorAsync(ExecutorId, tracer: null).ConfigureAwait(false); + Executor sourceExecutor = await RunnerContext.EnsureExecutorAsync(ExecutorId, tracer: null, cancellationToken).ConfigureAwait(false); if (!sourceExecutor.CanOutput(output.GetType())) { throw new InvalidOperationException($"Cannot output object of type {output.GetType().Name}. Expecting one of [{string.Join(", ", sourceExecutor.OutputTypes)}]."); @@ -251,22 +251,22 @@ internal sealed class InProcessRunnerContext : IRunnerContext if (outputFilter.CanOutput(ExecutorId, output)) { - await this.AddEventAsync(new WorkflowOutputEvent(output, ExecutorId)).ConfigureAwait(false); + await this.AddEventAsync(new WorkflowOutputEvent(output, ExecutorId), cancellationToken).ConfigureAwait(false); } } public ValueTask RequestHaltAsync() => this.AddEventAsync(new RequestHaltEvent()); - public ValueTask ReadStateAsync(string key, string? scopeName = null) + public ValueTask ReadStateAsync(string key, string? scopeName = null, CancellationToken cancellationToken = default) => RunnerContext.StateManager.ReadStateAsync(ExecutorId, scopeName, key); - public ValueTask> ReadStateKeysAsync(string? scopeName = null) + public ValueTask> ReadStateKeysAsync(string? scopeName = null, CancellationToken cancellationToken = default) => RunnerContext.StateManager.ReadKeysAsync(ExecutorId, scopeName); - public ValueTask QueueStateUpdateAsync(string key, T? value, string? scopeName = null) + public ValueTask QueueStateUpdateAsync(string key, T? value, string? scopeName = null, CancellationToken cancellationToken = default) => RunnerContext.StateManager.WriteStateAsync(ExecutorId, scopeName, key, value); - public ValueTask QueueClearScopeAsync(string? scopeName = null) + public ValueTask QueueClearScopeAsync(string? scopeName = null, CancellationToken cancellationToken = default) => RunnerContext.StateManager.ClearStateAsync(ExecutorId, scopeName); public IReadOnlyDictionary? TraceContext => traceContext; @@ -274,7 +274,7 @@ internal sealed class InProcessRunnerContext : IRunnerContext public bool WithCheckpointing { get; } - internal Task PrepareForCheckpointAsync(CancellationToken cancellation = default) + internal Task PrepareForCheckpointAsync(CancellationToken cancellationToken = default) { this.CheckEnded(); @@ -283,7 +283,7 @@ internal sealed class InProcessRunnerContext : IRunnerContext async Task InvokeCheckpointingAsync(Task executorTask) { Executor executor = await executorTask.ConfigureAwait(false); - await executor.OnCheckpointingAsync(this.Bind(executor.Id), cancellation).ConfigureAwait(false); + await executor.OnCheckpointingAsync(this.Bind(executor.Id), cancellationToken).ConfigureAwait(false); } } @@ -320,7 +320,7 @@ internal sealed class InProcessRunnerContext : IRunnerContext { foreach (string requestId in this._externalRequests.Keys) { - await this.AddEventAsync(new RequestInfoEvent(this._externalRequests[requestId])) + await this.AddEventAsync(new RequestInfoEvent(this._externalRequests[requestId]), cancellationToken) .ConfigureAwait(false); } } @@ -386,7 +386,7 @@ internal sealed class InProcessRunnerContext : IRunnerContext public IEnumerable JoinedSubworkflowRunners => this._joinedSubworkflowRunners; - public ValueTask AttachSuperstepAsync(ISuperStepRunner superStepRunner, CancellationToken cancellation = default) + public ValueTask AttachSuperstepAsync(ISuperStepRunner superStepRunner, CancellationToken cancellationToken = default) { // This needs to be a thread-safe ordered collection because we can potentially instantiate executors // in parallel, which means multiple sub-workflows could be attaching at the same time. @@ -394,9 +394,9 @@ internal sealed class InProcessRunnerContext : IRunnerContext return default; } - ValueTask ISuperStepJoinContext.ForwardWorkflowEventAsync(WorkflowEvent workflowEvent, CancellationToken cancellation) - => this.AddEventAsync(workflowEvent); + ValueTask ISuperStepJoinContext.ForwardWorkflowEventAsync(WorkflowEvent workflowEvent, CancellationToken cancellationToken) + => this.AddEventAsync(workflowEvent, cancellationToken); - ValueTask ISuperStepJoinContext.SendMessageAsync(string senderId, [DisallowNull] TMessage message, CancellationToken cancellation) - => this.SendMessageAsync(senderId, Throw.IfNull(message)); + ValueTask ISuperStepJoinContext.SendMessageAsync(string senderId, [DisallowNull] TMessage message, CancellationToken cancellationToken) + => this.SendMessageAsync(senderId, Throw.IfNull(message), cancellationToken: cancellationToken); } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Reflection/IMessageHandler.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Reflection/IMessageHandler.cs index 1903e49d44..3b18379907 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Reflection/IMessageHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Reflection/IMessageHandler.cs @@ -1,5 +1,6 @@ // Copyright (c) Microsoft. All rights reserved. +using System.Threading; using System.Threading.Tasks; namespace Microsoft.Agents.AI.Workflows.Reflection; @@ -15,8 +16,10 @@ public interface IMessageHandler /// /// The message to handle. /// The execution context. + /// The to monitor for cancellation requests. + /// The default is . /// A task that represents the asynchronous operation. - ValueTask HandleAsync(TMessage message, IWorkflowContext context); + ValueTask HandleAsync(TMessage message, IWorkflowContext context, CancellationToken cancellationToken = default); } /// @@ -32,6 +35,8 @@ public interface IMessageHandler /// /// The message to handle. /// The execution context. + /// The to monitor for cancellation requests. + /// The default is . /// A task that represents the asynchronous operation. - ValueTask HandleAsync(TMessage message, IWorkflowContext context); + ValueTask HandleAsync(TMessage message, IWorkflowContext context, CancellationToken cancellationToken = default); } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Reflection/MessageHandlerInfo.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Reflection/MessageHandlerInfo.cs index e00675c473..f63a43b4a8 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Reflection/MessageHandlerInfo.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Reflection/MessageHandlerInfo.cs @@ -5,6 +5,7 @@ using System.Diagnostics; using System.Diagnostics.CodeAnalysis; using System.Linq; using System.Reflection; +using System.Threading; using System.Threading.Tasks; using Microsoft.Agents.AI.Workflows.Execution; @@ -26,14 +27,19 @@ internal readonly struct MessageHandlerInfo this.HandlerInfo = handlerInfo; ParameterInfo[] parameters = handlerInfo.GetParameters(); - if (parameters.Length != 2) + if (parameters.Length != 3) { - throw new ArgumentException("Handler method must have exactly two parameters: TMessage and IExecutionContext.", nameof(handlerInfo)); + throw new ArgumentException("Handler method must have exactly three parameters: TMessage, IWorkflowContext, and CancellationToken.", nameof(handlerInfo)); } if (parameters[1].ParameterType != typeof(IWorkflowContext)) { - throw new ArgumentException("Handler method's second parameter must be of type IExecutionContext.", nameof(handlerInfo)); + throw new ArgumentException("Handler method's second parameter must be of type IWorkflowContext.", nameof(handlerInfo)); + } + + if (parameters[2].ParameterType != typeof(CancellationToken)) + { + throw new ArgumentException("Handler method's third parameter must be of type CancellationToken.", nameof(handlerInfo)); } this.InType = parameters[0].ParameterType; @@ -61,17 +67,17 @@ internal readonly struct MessageHandlerInfo } } - public static Func> Bind(Func handlerAsync, bool checkType, Type? resultType = null, Func>? unwrapper = null) + public static Func> Bind(Func handlerAsync, bool checkType, Type? resultType = null, Func>? unwrapper = null) { return InvokeHandlerAsync; - async ValueTask InvokeHandlerAsync(object message, IWorkflowContext workflowContext) + async ValueTask InvokeHandlerAsync(object message, IWorkflowContext workflowContext, CancellationToken cancellationToken) { bool expectingVoid = resultType is null || resultType == typeof(void); try { - object? maybeValueTask = handlerAsync(message, workflowContext); + object? maybeValueTask = handlerAsync(message, workflowContext, cancellationToken); if (expectingVoid) { @@ -109,6 +115,11 @@ internal readonly struct MessageHandlerInfo return CallResult.ReturnResult(result); } + catch (OperationCanceledException) + { + // If the operation was canceled, return a canceled CallResult. + return CallResult.Cancelled(wasVoid: expectingVoid); + } catch (Exception ex) { // If the handler throws an exception, return it in the CallResult. @@ -117,7 +128,7 @@ internal readonly struct MessageHandlerInfo } } - public Func> Bind< + public Func> Bind< [DynamicallyAccessedMembers( ReflectionDemands.RuntimeInterfaceDiscoveryAndInvocation) ] TExecutor @@ -128,9 +139,9 @@ internal readonly struct MessageHandlerInfo MethodInfo handlerMethod = this.HandlerInfo; return Bind(InvokeHandler, checkType, this.OutType, this.Unwrapper); - object? InvokeHandler(object message, IWorkflowContext workflowContext) + object? InvokeHandler(object message, IWorkflowContext workflowContext, CancellationToken cancellationToken) { - return handlerMethod.Invoke(executor, [message, workflowContext]); + return handlerMethod.Invoke(executor, [message, workflowContext, cancellationToken]); } } } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/RouteBuilder.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/RouteBuilder.cs index 965de0cdaa..99cfdb6992 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/RouteBuilder.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/RouteBuilder.cs @@ -3,6 +3,7 @@ using System; using System.Collections.Generic; using System.Diagnostics; +using System.Threading; using System.Threading.Tasks; using Microsoft.Agents.AI.Workflows.Execution; using Microsoft.Shared.Diagnostics; @@ -10,12 +11,14 @@ using CatchAllF = System.Func< Microsoft.Agents.AI.Workflows.PortableValue, // message Microsoft.Agents.AI.Workflows.IWorkflowContext, // context + System.Threading.CancellationToken, // cancellation System.Threading.Tasks.ValueTask >; using MessageHandlerF = System.Func< object, // message Microsoft.Agents.AI.Workflows.IWorkflowContext, // context + System.Threading.CancellationToken, // cancellation System.Threading.Tasks.ValueTask >; @@ -73,32 +76,58 @@ public class RouteBuilder return this; } - internal RouteBuilder AddHandlerUntyped(Type type, Func handler, bool overwrite = false) + internal RouteBuilder AddHandlerUntyped(Type type, Func handler, bool overwrite = false) { Throw.IfNull(handler); return this.AddHandlerInternal(type, WrappedHandlerAsync, outputType: null, overwrite); - async ValueTask WrappedHandlerAsync(object msg, IWorkflowContext ctx) + async ValueTask WrappedHandlerAsync(object message, IWorkflowContext context, CancellationToken cancellationToken) { - await handler.Invoke(msg, ctx).ConfigureAwait(false); + await handler.Invoke(message, context, cancellationToken).ConfigureAwait(false); return CallResult.ReturnVoid(); } } - internal RouteBuilder AddHandlerUntyped(Type type, Func> handler, bool overwrite = false) + internal RouteBuilder AddHandlerUntyped(Type type, Func> handler, bool overwrite = false) { Throw.IfNull(handler); return this.AddHandlerInternal(type, WrappedHandlerAsync, outputType: typeof(TResult), overwrite); - async ValueTask WrappedHandlerAsync(object msg, IWorkflowContext ctx) + async ValueTask WrappedHandlerAsync(object message, IWorkflowContext context, CancellationToken cancellationToken) { - TResult result = await handler.Invoke(msg, ctx).ConfigureAwait(false); + TResult result = await handler.Invoke(message, context, cancellationToken).ConfigureAwait(false); return CallResult.ReturnResult(result); } } + /// + /// Registers a handler for messages of the specified input type in the workflow route. + /// + /// If a handler for the specified input type already exists and is + /// , the existing handler will not be replaced. Handlers are invoked asynchronously and are + /// expected to complete their processing before the workflow continues. + /// + /// A delegate that processes messages of type within the workflow context. The + /// delegate is invoked for each incoming message of the specified type. + /// to replace any existing handler for the specified input type; otherwise, to preserve the existing handler. + /// The current instance, enabling fluent configuration of additional handlers or route + /// options. + public RouteBuilder AddHandler(Action handler, bool overwrite = false) + { + Throw.IfNull(handler); + + return this.AddHandlerInternal(typeof(TInput), WrappedHandlerAsync, outputType: null, overwrite); + + async ValueTask WrappedHandlerAsync(object message, IWorkflowContext context, CancellationToken cancellationToken) + { + handler.Invoke((TInput)message, context, cancellationToken); + return CallResult.ReturnVoid(); + } + } + /// /// Registers a handler for messages of the specified input type in the workflow route. /// @@ -118,9 +147,35 @@ public class RouteBuilder return this.AddHandlerInternal(typeof(TInput), WrappedHandlerAsync, outputType: null, overwrite); - async ValueTask WrappedHandlerAsync(object msg, IWorkflowContext ctx) + async ValueTask WrappedHandlerAsync(object message, IWorkflowContext context, CancellationToken cancellationToken) { - handler.Invoke((TInput)msg, ctx); + handler.Invoke((TInput)message, context); + return CallResult.ReturnVoid(); + } + } + + /// + /// Registers a handler for messages of the specified input type in the workflow route. + /// + /// If a handler for the specified input type already exists and is + /// , the existing handler will not be replaced. Handlers are invoked asynchronously and are + /// expected to complete their processing before the workflow continues. + /// + /// A delegate that processes messages of type within the workflow context. The + /// delegate is invoked for each incoming message of the specified type. + /// to replace any existing handler for the specified input type; otherwise, to preserve the existing handler. + /// The current instance, enabling fluent configuration of additional handlers or route + /// options. + public RouteBuilder AddHandler(Func handler, bool overwrite = false) + { + Throw.IfNull(handler); + + return this.AddHandlerInternal(typeof(TInput), WrappedHandlerAsync, outputType: null, overwrite); + + async ValueTask WrappedHandlerAsync(object message, IWorkflowContext context, CancellationToken cancellationToken) + { + await handler.Invoke((TInput)message, context, cancellationToken).ConfigureAwait(false); return CallResult.ReturnVoid(); } } @@ -144,13 +199,39 @@ public class RouteBuilder return this.AddHandlerInternal(typeof(TInput), WrappedHandlerAsync, outputType: null, overwrite); - async ValueTask WrappedHandlerAsync(object msg, IWorkflowContext ctx) + async ValueTask WrappedHandlerAsync(object message, IWorkflowContext context, CancellationToken cancellationToken) { - await handler.Invoke((TInput)msg, ctx).ConfigureAwait(false); + await handler.Invoke((TInput)message, context).ConfigureAwait(false); return CallResult.ReturnVoid(); } } + /// + /// Registers a handler function for messages of the specified input type in the workflow route. + /// + /// If a handler for the given input type already exists, setting to + /// will replace the existing handler; otherwise, an exception may be thrown. The handler + /// receives the input message and workflow context, and returns a result asynchronously. + /// The type of input message the handler will process. + /// The type of result produced by the handler. + /// A function that processes messages of type within the workflow context and returns + /// a representing the asynchronous result. + /// to replace any existing handler for the input type; otherwise, to + /// preserve existing handlers. + /// The current instance, enabling fluent configuration of workflow routes. + public RouteBuilder AddHandler(Func handler, bool overwrite = false) + { + Throw.IfNull(handler); + + return this.AddHandlerInternal(typeof(TInput), WrappedHandlerAsync, outputType: typeof(TResult), overwrite); + + async ValueTask WrappedHandlerAsync(object message, IWorkflowContext context, CancellationToken cancellationToken) + { + TResult result = handler.Invoke((TInput)message, context, cancellationToken); + return CallResult.ReturnResult(result); + } + } + /// /// Registers a handler function for messages of the specified input type in the workflow route. /// @@ -170,9 +251,35 @@ public class RouteBuilder return this.AddHandlerInternal(typeof(TInput), WrappedHandlerAsync, outputType: typeof(TResult), overwrite); - async ValueTask WrappedHandlerAsync(object msg, IWorkflowContext ctx) + async ValueTask WrappedHandlerAsync(object message, IWorkflowContext context, CancellationToken cancellationToken) { - TResult result = handler.Invoke((TInput)msg, ctx); + TResult result = handler.Invoke((TInput)message, context); + return CallResult.ReturnResult(result); + } + } + + /// + /// Registers a handler function for messages of the specified input type in the workflow route. + /// + /// If a handler for the given input type already exists, setting to + /// will replace the existing handler; otherwise, an exception may be thrown. The handler + /// receives the input message and workflow context, and returns a result asynchronously. + /// The type of input message the handler will process. + /// The type of result produced by the handler. + /// A function that processes messages of type within the workflow context and returns + /// a representing the asynchronous result. + /// to replace any existing handler for the input type; otherwise, to + /// preserve existing handlers. + /// The current instance, enabling fluent configuration of workflow routes. + public RouteBuilder AddHandler(Func> handler, bool overwrite = false) + { + Throw.IfNull(handler); + + return this.AddHandlerInternal(typeof(TInput), WrappedHandlerAsync, outputType: typeof(TResult), overwrite); + + async ValueTask WrappedHandlerAsync(object message, IWorkflowContext context, CancellationToken cancellationToken) + { + TResult result = await handler.Invoke((TInput)message, context, cancellationToken).ConfigureAwait(false); return CallResult.ReturnResult(result); } } @@ -196,9 +303,9 @@ public class RouteBuilder return this.AddHandlerInternal(typeof(TInput), WrappedHandlerAsync, outputType: typeof(TResult), overwrite); - async ValueTask WrappedHandlerAsync(object msg, IWorkflowContext ctx) + async ValueTask WrappedHandlerAsync(object message, IWorkflowContext context, CancellationToken cancellationToken) { - TResult result = await handler.Invoke((TInput)msg, ctx).ConfigureAwait(false); + TResult result = await handler.Invoke((TInput)message, context).ConfigureAwait(false); return CallResult.ReturnResult(result); } } @@ -215,6 +322,30 @@ public class RouteBuilder return this; } + /// + /// Register a handler function as a catch-all handler: It will be used if not type-matching handler is registered. + /// + /// If a catch-all handler for already exists, setting to + /// will replace the existing handler; otherwise, an exception may be thrown. The handler receives the input message + /// wrapped as and workflow context, and returns a result asynchronously. + /// A function that processes messages wrapped as within the + /// workflow context. The delegate is invoked for each incoming message not otherwise handled. + /// to replace any existing handler for the input type; otherwise, to + /// preserve existing handlers. + /// The current instance, enabling fluent configuration of workflow routes. + public RouteBuilder AddCatchAll(Func handler, bool overwrite = false) + { + Throw.IfNull(handler); + + return this.AddCatchAll(WrappedHandlerAsync, overwrite); + + async ValueTask WrappedHandlerAsync(PortableValue message, IWorkflowContext context, CancellationToken cancellationToken) + { + await handler.Invoke(message, context, cancellationToken).ConfigureAwait(false); + return CallResult.ReturnVoid(); + } + } + /// /// Register a handler function as a catch-all handler: It will be used if not type-matching handler is registered. /// @@ -232,13 +363,37 @@ public class RouteBuilder return this.AddCatchAll(WrappedHandlerAsync, overwrite); - async ValueTask WrappedHandlerAsync(PortableValue message, IWorkflowContext ctx) + async ValueTask WrappedHandlerAsync(PortableValue message, IWorkflowContext context, CancellationToken cancellationToken) { - await handler.Invoke(message, ctx).ConfigureAwait(false); + await handler.Invoke(message, context).ConfigureAwait(false); return CallResult.ReturnVoid(); } } + /// + /// Register a handler function as a catch-all handler: It will be used if not type-matching handler is registered. + /// + /// If a catch-all handler for already exists, setting to + /// will replace the existing handler; otherwise, an exception may be thrown. The handler receives the input message + /// wrapped as and workflow context, and returns a result asynchronously. + /// A function that processes messages wrapped as within the + /// workflow context and returns a representing the asynchronous result. + /// to replace any existing handler for the input type; otherwise, to + /// preserve existing handlers. + /// The current instance, enabling fluent configuration of workflow routes. + public RouteBuilder AddCatchAll(Func> handler, bool overwrite = false) + { + Throw.IfNull(handler); + + return this.AddCatchAll(WrappedHandlerAsync, overwrite); + + async ValueTask WrappedHandlerAsync(PortableValue message, IWorkflowContext context, CancellationToken cancellationToken) + { + TResult result = await handler.Invoke(message, context, cancellationToken).ConfigureAwait(false); + return CallResult.ReturnResult(result); + } + } + /// /// Register a handler function as a catch-all handler: It will be used if not type-matching handler is registered. /// @@ -256,13 +411,37 @@ public class RouteBuilder return this.AddCatchAll(WrappedHandlerAsync, overwrite); - async ValueTask WrappedHandlerAsync(PortableValue message, IWorkflowContext ctx) + async ValueTask WrappedHandlerAsync(PortableValue message, IWorkflowContext context, CancellationToken cancellationToken) { - TResult result = await handler.Invoke(message, ctx).ConfigureAwait(false); + TResult result = await handler.Invoke(message, context).ConfigureAwait(false); return CallResult.ReturnResult(result); } } + /// + /// Register a handler function as a catch-all handler: It will be used if not type-matching handler is registered. + /// + /// If a catch-all handler for already exists, setting to + /// will replace the existing handler; otherwise, an exception may be thrown. The handler receives the input message + /// wrapped as and workflow context, and returns a result asynchronously. + /// A function that processes messages wrapped as within the + /// workflow context. The delegate is invoked for each incoming message not otherwise handled. + /// to replace any existing handler for the input type; otherwise, to + /// preserve existing handlers. + /// The current instance, enabling fluent configuration of workflow routes. + public RouteBuilder AddCatchAll(Action handler, bool overwrite = false) + { + Throw.IfNull(handler); + + return this.AddCatchAll(WrappedHandlerAsync, overwrite); + + ValueTask WrappedHandlerAsync(PortableValue message, IWorkflowContext ctx, CancellationToken cancellationToken) + { + handler.Invoke(message, ctx, cancellationToken); + return new(CallResult.ReturnVoid()); + } + } + /// /// Register a handler function as a catch-all handler: It will be used if not type-matching handler is registered. /// @@ -280,13 +459,37 @@ public class RouteBuilder return this.AddCatchAll(WrappedHandlerAsync, overwrite); - ValueTask WrappedHandlerAsync(PortableValue message, IWorkflowContext ctx) + ValueTask WrappedHandlerAsync(PortableValue message, IWorkflowContext ctx, CancellationToken cancellationToken) { handler.Invoke(message, ctx); return new(CallResult.ReturnVoid()); } } + /// + /// Register a handler function as a catch-all handler: It will be used if not type-matching handler is registered. + /// + /// If a catch-all handler for already exists, setting to + /// will replace the existing handler; otherwise, an exception may be thrown. The handler receives the input message + /// wrapped as and workflow context, and returns a result asynchronously. + /// A function that processes messages wrapped as within the + /// workflow context and returns a representing the asynchronous result. + /// to replace any existing handler for the input type; otherwise, to + /// preserve existing handlers. + /// The current instance, enabling fluent configuration of workflow routes. + public RouteBuilder AddCatchAll(Func handler, bool overwrite = false) + { + Throw.IfNull(handler); + + return this.AddCatchAll(WrappedHandlerAsync, overwrite); + + ValueTask WrappedHandlerAsync(PortableValue message, IWorkflowContext context, CancellationToken cancellationToken) + { + TResult result = handler.Invoke(message, context, cancellationToken); + return new(CallResult.ReturnResult(result)); + } + } + /// /// Register a handler function as a catch-all handler: It will be used if not type-matching handler is registered. /// @@ -304,9 +507,9 @@ public class RouteBuilder return this.AddCatchAll(WrappedHandlerAsync, overwrite); - ValueTask WrappedHandlerAsync(PortableValue message, IWorkflowContext ctx) + ValueTask WrappedHandlerAsync(PortableValue message, IWorkflowContext context, CancellationToken cancellationToken) { - TResult result = handler.Invoke(message, ctx); + TResult result = handler.Invoke(message, context); return new(CallResult.ReturnResult(result)); } } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Run.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Run.cs index dec02fbd46..3dfa4f271e 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Run.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Run.cs @@ -43,8 +43,8 @@ public sealed class Run : IAsyncDisposable /// /// Gets the current execution status of the workflow run. /// - public ValueTask GetStatusAsync(CancellationToken cancellation = default) - => this._runHandle.GetStatusAsync(cancellation); + public ValueTask GetStatusAsync(CancellationToken cancellationToken = default) + => this._runHandle.GetStatusAsync(cancellationToken); /// /// Gets all events emitted by the workflow. @@ -113,7 +113,7 @@ public sealed class Run : IAsyncDisposable { foreach (object? message in messages) { - await this._runHandle.EnqueueMessageUntypedAsync(message, cancellation: cancellationToken).ConfigureAwait(false); + await this._runHandle.EnqueueMessageUntypedAsync(message, cancellationToken: cancellationToken).ConfigureAwait(false); } } else diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/AIAgentHostExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/AIAgentHostExecutor.cs index e50e0f803a..836399c5c1 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/AIAgentHostExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/AIAgentHostExecutor.cs @@ -30,7 +30,7 @@ internal sealed class AIAgentHostExecutor : ChatProtocolExecutor if (this._thread is not null) { JsonElement threadValue = this._thread.Serialize(); - threadTask = context.QueueStateUpdateAsync(ThreadStateKey, threadValue).AsTask(); + threadTask = context.QueueStateUpdateAsync(ThreadStateKey, threadValue, cancellationToken: cancellationToken).AsTask(); } Task baseTask = base.OnCheckpointingAsync(context, cancellationToken).AsTask(); @@ -40,7 +40,7 @@ internal sealed class AIAgentHostExecutor : ChatProtocolExecutor protected internal override async ValueTask OnCheckpointRestoredAsync(IWorkflowContext context, CancellationToken cancellationToken = default) { - JsonElement? threadValue = await context.ReadStateAsync(ThreadStateKey).ConfigureAwait(false); + JsonElement? threadValue = await context.ReadStateAsync(ThreadStateKey, cancellationToken: cancellationToken).ConfigureAwait(false); if (threadValue.HasValue) { this._thread = this._agent.DeserializeThread(threadValue.Value); @@ -67,7 +67,7 @@ internal sealed class AIAgentHostExecutor : ChatProtocolExecutor if (emitEvents ?? this._emitEvents) { - await context.AddEventAsync(new AgentRunUpdateEvent(this.Id, update)).ConfigureAwait(false); + await context.AddEventAsync(new AgentRunUpdateEvent(this.Id, update), cancellationToken).ConfigureAwait(false); } // TODO: FunctionCall request handling, and user info request handling. @@ -100,7 +100,7 @@ internal sealed class AIAgentHostExecutor : ChatProtocolExecutor currentStreamingMessage.Contents = updates; updates = []; - await context.SendMessageAsync(currentStreamingMessage).ConfigureAwait(false); + await context.SendMessageAsync(currentStreamingMessage, cancellationToken: cancellationToken).ConfigureAwait(false); } currentStreamingMessage = null; diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/RequestInfoExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/RequestInfoExecutor.cs index 6226b00db9..afb07507f9 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/RequestInfoExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/RequestInfoExecutor.cs @@ -3,6 +3,7 @@ using System; using System.Collections.Generic; using System.Diagnostics; +using System.Threading; using System.Threading.Tasks; using Microsoft.Agents.AI.Workflows.Execution; using Microsoft.Shared.Diagnostics; @@ -55,7 +56,7 @@ internal sealed class RequestInfoExecutor : Executor internal void AttachRequestSink(IExternalRequestSink requestSink) => this.RequestSink = Throw.IfNull(requestSink); - public async ValueTask HandleCatchAllAsync(PortableValue message, IWorkflowContext context) + public async ValueTask HandleCatchAllAsync(PortableValue message, IWorkflowContext context, CancellationToken cancellationToken) { Throw.IfNull(message); @@ -70,13 +71,13 @@ internal sealed class RequestInfoExecutor : Executor } else if (message.Is(out ExternalRequest? request)) { - return await this.HandleAsync(request, context).ConfigureAwait(false); + return await this.HandleAsync(request, context, cancellationToken).ConfigureAwait(false); } return null; } - public async ValueTask HandleAsync(ExternalRequest message, IWorkflowContext context) + public async ValueTask HandleAsync(ExternalRequest message, IWorkflowContext context, CancellationToken cancellationToken = default) { Debug.Assert(this._allowWrapped); Throw.IfNull(message); @@ -100,7 +101,7 @@ internal sealed class RequestInfoExecutor : Executor return request; } - public async ValueTask HandleAsync(object message, IWorkflowContext context) + public async ValueTask HandleAsync(object message, IWorkflowContext context, CancellationToken cancellationToken = default) { Throw.IfNull(message); Debug.Assert(this.Port.Request.IsInstanceOfType(message)); @@ -111,7 +112,7 @@ internal sealed class RequestInfoExecutor : Executor return request; } - public async ValueTask HandleAsync(ExternalResponse message, IWorkflowContext context) + public async ValueTask HandleAsync(ExternalResponse message, IWorkflowContext context, CancellationToken cancellationToken = default) { Throw.IfNull(message); Throw.IfNull(message.Data); @@ -127,14 +128,14 @@ internal sealed class RequestInfoExecutor : Executor if (this._allowWrapped && this._wrappedRequests.TryGetValue(message.RequestId, out ExternalRequest? originalRequest)) { - await context.SendMessageAsync(originalRequest.RewrapResponse(message)).ConfigureAwait(false); + await context.SendMessageAsync(originalRequest.RewrapResponse(message), cancellationToken: cancellationToken).ConfigureAwait(false); } else { - await context.SendMessageAsync(message).ConfigureAwait(false); + await context.SendMessageAsync(message, cancellationToken: cancellationToken).ConfigureAwait(false); } - await context.SendMessageAsync(data).ConfigureAwait(false); + await context.SendMessageAsync(data, cancellationToken: cancellationToken).ConfigureAwait(false); return message; } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/WorkflowHostExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/WorkflowHostExecutor.cs index 78f79148e7..4bcffafef9 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/WorkflowHostExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/WorkflowHostExecutor.cs @@ -44,22 +44,22 @@ internal class WorkflowHostExecutor : Executor, IAsyncDisposable return routeBuilder.AddCatchAll(this.QueueExternalMessageAsync); } - private async ValueTask QueueExternalMessageAsync(PortableValue portableValue, IWorkflowContext context) + private async ValueTask QueueExternalMessageAsync(PortableValue portableValue, IWorkflowContext context, CancellationToken cancellationToken) { if (portableValue.Is(out ExternalResponse? response)) { response = this.CheckAndUnqualifyResponse(response); - await this.EnsureRunSendMessageAsync(response).ConfigureAwait(false); + await this.EnsureRunSendMessageAsync(response, cancellationToken: cancellationToken).ConfigureAwait(false); } else { InProcessRunner runner = await this.EnsureRunnerAsync().ConfigureAwait(false); - IEnumerable validInputTypes = await runner.RunContext.GetStartingExecutorInputTypesAsync().ConfigureAwait(false); + IEnumerable validInputTypes = await runner.RunContext.GetStartingExecutorInputTypesAsync(cancellationToken).ConfigureAwait(false); foreach (Type candidateType in validInputTypes) { if (portableValue.IsType(candidateType, out object? message)) { - await this.EnsureRunSendMessageAsync(message, candidateType).ConfigureAwait(false); + await this.EnsureRunSendMessageAsync(message, candidateType, cancellationToken: cancellationToken).ConfigureAwait(false); return; } } @@ -87,7 +87,7 @@ internal class WorkflowHostExecutor : Executor, IAsyncDisposable return this._activeRunner; } - internal async ValueTask EnsureRunSendMessageAsync(object? incomingMessage = null, Type? incomingMessageType = null, bool resume = false, CancellationToken cancellation = default) + internal async ValueTask EnsureRunSendMessageAsync(object? incomingMessage = null, Type? incomingMessageType = null, bool resume = false, CancellationToken cancellationToken = default) { Debug.Assert(this._joinContext != null, "Must attach to a join context before starting the run."); @@ -114,20 +114,20 @@ internal class WorkflowHostExecutor : Executor, IAsyncDisposable throw new InvalidOperationException("No checkpoints available to resume from."); } - runHandle = await activeRunner.ResumeStreamAsync(ExecutionMode.Subworkflow, lastCheckpoint!, cancellation) + runHandle = await activeRunner.ResumeStreamAsync(ExecutionMode.Subworkflow, lastCheckpoint!, cancellationToken) .ConfigureAwait(false); if (incomingMessage != null) { - await runHandle.EnqueueUntypedAndRunAsync(incomingMessage, cancellation).ConfigureAwait(false); + await runHandle.EnqueueUntypedAndRunAsync(incomingMessage, cancellationToken).ConfigureAwait(false); } } else if (incomingMessage != null) { - runHandle = await activeRunner.BeginStreamAsync(ExecutionMode.Subworkflow, cancellation) + runHandle = await activeRunner.BeginStreamAsync(ExecutionMode.Subworkflow, cancellationToken) .ConfigureAwait(false); - await runHandle.EnqueueUntypedAndRunAsync(incomingMessage, cancellation).ConfigureAwait(false); + await runHandle.EnqueueUntypedAndRunAsync(incomingMessage, cancellationToken).ConfigureAwait(false); } else { @@ -136,14 +136,14 @@ internal class WorkflowHostExecutor : Executor, IAsyncDisposable } else { - runHandle = await activeRunner.BeginStreamAsync(ExecutionMode.Subworkflow, cancellation).ConfigureAwait(false); + runHandle = await activeRunner.BeginStreamAsync(ExecutionMode.Subworkflow, cancellationToken).ConfigureAwait(false); - await runHandle.EnqueueMessageUntypedAsync(Throw.IfNull(incomingMessage), cancellation: cancellation).ConfigureAwait(false); + await runHandle.EnqueueMessageUntypedAsync(Throw.IfNull(incomingMessage), cancellationToken: cancellationToken).ConfigureAwait(false); } this._run = new(runHandle); - await this._joinContext.AttachSuperstepAsync(activeRunner, cancellation).ConfigureAwait(false); + await this._joinContext.AttachSuperstepAsync(activeRunner, cancellationToken).ConfigureAwait(false); activeRunner.OutgoingEvents.EventRaised += this.ForwardWorkflowEventAsync; return this._run; @@ -228,7 +228,7 @@ internal class WorkflowHostExecutor : Executor, IAsyncDisposable protected internal override async ValueTask OnCheckpointingAsync(IWorkflowContext context, CancellationToken cancellationToken = default) { - await context.QueueStateUpdateAsync(nameof(CheckpointManager), this._checkpointManager).ConfigureAwait(false); + await context.QueueStateUpdateAsync(nameof(CheckpointManager), this._checkpointManager, cancellationToken: cancellationToken).ConfigureAwait(false); await base.OnCheckpointingAsync(context, cancellationToken).ConfigureAwait(false); } @@ -237,7 +237,7 @@ internal class WorkflowHostExecutor : Executor, IAsyncDisposable { await base.OnCheckpointRestoredAsync(context, cancellationToken).ConfigureAwait(false); - InMemoryCheckpointManager manager = await context.ReadStateAsync(nameof(InMemoryCheckpointManager)).ConfigureAwait(false) ?? new(); + InMemoryCheckpointManager manager = await context.ReadStateAsync(nameof(InMemoryCheckpointManager), cancellationToken: cancellationToken).ConfigureAwait(false) ?? new(); if (this._checkpointManager == manager) { // We are restoring in the context of the same run; not need to rebuild the entire execution stack. @@ -249,7 +249,7 @@ internal class WorkflowHostExecutor : Executor, IAsyncDisposable await this.ResetAsync().ConfigureAwait(false); } - StreamingRun run = await this.EnsureRunSendMessageAsync(cancellation: cancellationToken).ConfigureAwait(false); + StreamingRun run = await this.EnsureRunSendMessageAsync(cancellationToken: cancellationToken).ConfigureAwait(false); } private async ValueTask ResetAsync() diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/StreamingRun.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/StreamingRun.cs index aa70307c1e..ad6727fc54 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/StreamingRun.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/StreamingRun.cs @@ -31,8 +31,8 @@ public sealed class StreamingRun : IAsyncDisposable /// /// Gets the current execution status of the workflow run. /// - public ValueTask GetStatusAsync(CancellationToken cancellation = default) - => this._runHandle.GetStatusAsync(cancellation); + public ValueTask GetStatusAsync(CancellationToken cancellationToken = default) + => this._runHandle.GetStatusAsync(cancellationToken); /// /// Asynchronously sends the specified response to the external system and signals completion of the current @@ -67,7 +67,7 @@ public sealed class StreamingRun : IAsyncDisposable /// progresses. The stream completes when a is encountered. Events are /// delivered in the order they are raised. /// A that can be used to cancel the streaming operation. If cancellation is - /// requested, the stream will end and no further events will be yielded. + /// requested, the stream will end and no further events will be yielded, but this will not cancel the workflow execution. /// An asynchronous stream of objects representing significant workflow state changes. /// The stream ends when the workflow completes or when cancellation is requested. public IAsyncEnumerable WatchStreamAsync( diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/ObjectModel/WorkflowActionExecutorTest.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/ObjectModel/WorkflowActionExecutorTest.cs index 62f9231e68..738b26b6f9 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/ObjectModel/WorkflowActionExecutorTest.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/ObjectModel/WorkflowActionExecutorTest.cs @@ -2,6 +2,7 @@ using System; using System.Linq; +using System.Threading; using System.Threading.Tasks; using Microsoft.Agents.AI.Workflows.Declarative.Extensions; using Microsoft.Agents.AI.Workflows.Declarative.Interpreter; @@ -70,7 +71,7 @@ public abstract class WorkflowActionExecutorTest(ITestOutputHelper output) : Wor internal sealed class TestWorkflowExecutor() : Executor("test_workflow") { - public override async ValueTask HandleAsync(WorkflowFormulaState message, IWorkflowContext context) => - await context.SendMessageAsync(new ActionExecutorResult(this.Id)).ConfigureAwait(false); + public override async ValueTask HandleAsync(WorkflowFormulaState message, IWorkflowContext context, CancellationToken cancellationToken) => + await context.SendMessageAsync(new ActionExecutorResult(this.Id), cancellationToken: cancellationToken).ConfigureAwait(false); } } diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/AgentWorkflowBuilderTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/AgentWorkflowBuilderTests.cs index 8a534269f0..e86df11d6f 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/AgentWorkflowBuilderTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/AgentWorkflowBuilderTests.cs @@ -6,6 +6,7 @@ using System.Linq; using System.Runtime.CompilerServices; using System.Text; using System.Text.Json; +using System.Text.RegularExpressions; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; @@ -159,7 +160,7 @@ public class AgentWorkflowBuilderTests private sealed class DoubleEchoAgentThread() : InMemoryAgentThread(); - [Fact(Skip = "issue #1109")] + [Fact] public async Task BuildConcurrent_AgentsRunInParallelAsync() { StrongBox> barrier = new(); @@ -182,10 +183,10 @@ public class AgentWorkflowBuilderTests // TODO: https://github.com/microsoft/agent-framework/issues/784 // These asserts are flaky until we guarantee message delivery order. - //Assert.Single(Regex.Matches(updateText, "agent1")); - //Assert.Single(Regex.Matches(updateText, "agent2")); - //Assert.Equal(4, Regex.Matches(updateText, "abc").Count); - //Assert.Equal(2, result.Count); + Assert.Single(Regex.Matches(updateText, "agent1")); + Assert.Single(Regex.Matches(updateText, "agent2")); + Assert.Equal(4, Regex.Matches(updateText, "abc").Count); + Assert.Equal(2, result.Count); } } diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/InProcessStateTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/InProcessStateTests.cs index debde5fc95..014c51b3c0 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/InProcessStateTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/InProcessStateTests.cs @@ -43,12 +43,12 @@ public class InProcessStateTests return async (turn, context, cancellation) => { - TState? state = await context.ReadStateAsync(stateKey.Key, stateKey.ScopeId.ScopeName) + TState? state = await context.ReadStateAsync(stateKey.Key, stateKey.ScopeId.ScopeName, cancellation) .ConfigureAwait(false); state = action(state); - await context.QueueStateUpdateAsync(stateKey.Key, state, stateKey.ScopeId.ScopeName); + await context.QueueStateUpdateAsync(stateKey.Key, state, stateKey.ScopeId.ScopeName, cancellation); return turn.Next; }; diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/ReflectionSmokeTest.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/ReflectionSmokeTest.cs index dd8470b11c..5027028387 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/ReflectionSmokeTest.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/ReflectionSmokeTest.cs @@ -1,6 +1,7 @@ // Copyright (c) Microsoft. All rights reserved. using System; +using System.Threading; using System.Threading.Tasks; using Microsoft.Agents.AI.Workflows.Execution; using Microsoft.Agents.AI.Workflows.Reflection; @@ -21,7 +22,7 @@ public class BaseTestExecutor(string id) : ReflectingExecutor( public class DefaultHandler() : BaseTestExecutor(nameof(DefaultHandler)), IMessageHandler { - public ValueTask HandleAsync(object message, IWorkflowContext context) + public ValueTask HandleAsync(object message, IWorkflowContext context, CancellationToken cancellationToken = default) { this.OnInvokedHandler(); return this.Handler(message, context); @@ -36,7 +37,7 @@ public class DefaultHandler() : BaseTestExecutor(nameof(DefaultH public class TypedHandler() : BaseTestExecutor>(nameof(TypedHandler)), IMessageHandler { - public ValueTask HandleAsync(TInput message, IWorkflowContext context) + public ValueTask HandleAsync(TInput message, IWorkflowContext context, CancellationToken cancellationToken = default) { this.OnInvokedHandler(); return this.Handler(message, context); @@ -51,7 +52,7 @@ public class TypedHandler() : BaseTestExecutor>(nam public class TypedHandlerWithOutput() : BaseTestExecutor>(nameof(TypedHandlerWithOutput)), IMessageHandler { - public ValueTask HandleAsync(TInput message, IWorkflowContext context) + public ValueTask HandleAsync(TInput message, IWorkflowContext context, CancellationToken cancellationToken) { this.OnInvokedHandler(); return this.Handler(message, context); diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/01_Simple_Workflow_Sequential.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/01_Simple_Workflow_Sequential.cs index 2e89dabd9e..7ec83b4d03 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/01_Simple_Workflow_Sequential.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/01_Simple_Workflow_Sequential.cs @@ -3,6 +3,7 @@ using System; using System.IO; using System.Linq; +using System.Threading; using System.Threading.Tasks; using Microsoft.Agents.AI.Workflows.InProc; using Microsoft.Agents.AI.Workflows.Reflection; @@ -44,17 +45,17 @@ internal static class Step1EntryPoint internal sealed class UppercaseExecutor() : ReflectingExecutor("UppercaseExecutor"), IMessageHandler { - public async ValueTask HandleAsync(string message, IWorkflowContext context) => + public async ValueTask HandleAsync(string message, IWorkflowContext context, CancellationToken cancellationToken = default) => message.ToUpperInvariant(); } internal sealed class ReverseTextExecutor() : ReflectingExecutor("ReverseTextExecutor"), IMessageHandler { - public async ValueTask HandleAsync(string message, IWorkflowContext context) + public async ValueTask HandleAsync(string message, IWorkflowContext context, CancellationToken cancellationToken = default) { string result = string.Concat(message.Reverse()); - await context.YieldOutputAsync(result).ConfigureAwait(false); + await context.YieldOutputAsync(result, cancellationToken).ConfigureAwait(false); return result; } } diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/02_Simple_Workflow_Condition.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/02_Simple_Workflow_Condition.cs index dd02f93594..5b9a0712ba 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/02_Simple_Workflow_Condition.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/02_Simple_Workflow_Condition.cs @@ -3,6 +3,7 @@ using System; using System.IO; using System.Linq; +using System.Threading; using System.Threading.Tasks; using Microsoft.Agents.AI.Workflows.InProc; using Microsoft.Agents.AI.Workflows.Reflection; @@ -56,7 +57,7 @@ internal static class Step2EntryPoint internal sealed class DetectSpamExecutor(string id, params string[] spamKeywords) : ReflectingExecutor(id), IMessageHandler { - public async ValueTask HandleAsync(string message, IWorkflowContext context) => + public async ValueTask HandleAsync(string message, IWorkflowContext context, CancellationToken cancellationToken = default) => spamKeywords.Any(keyword => message.IndexOf(keyword, StringComparison.OrdinalIgnoreCase) >= 0); } @@ -64,7 +65,7 @@ internal sealed class RespondToMessageExecutor(string id) : ReflectingExecutor (this.LowerBound + this.UpperBound) / 2; private int _currGuess = -1; - public async ValueTask HandleAsync(NumberSignal message, IWorkflowContext context) + public async ValueTask HandleAsync(NumberSignal message, IWorkflowContext context, CancellationToken cancellationToken = default) { switch (message) { case NumberSignal.Matched: - await context.YieldOutputAsync($"Guessed the number: {this._currGuess}") + await context.YieldOutputAsync($"Guessed the number: {this._currGuess}", cancellationToken) .ConfigureAwait(false); break; @@ -106,7 +106,7 @@ internal sealed class JudgeExecutor : ReflectingExecutor, IMessag this._targetNumber = targetNumber; } - public async ValueTask HandleAsync(int message, IWorkflowContext context) + public async ValueTask HandleAsync(int message, IWorkflowContext context, CancellationToken cancellationToken = default) { this.Tries = this.Tries is int tries ? tries + 1 : 1; @@ -118,12 +118,12 @@ internal sealed class JudgeExecutor : ReflectingExecutor, IMessag protected internal override ValueTask OnCheckpointingAsync(IWorkflowContext context, CancellationToken cancellationToken = default) { - return context.QueueStateUpdateAsync("TryCount", this.Tries); + return context.QueueStateUpdateAsync("TryCount", this.Tries, cancellationToken: cancellationToken); } protected internal override async ValueTask OnCheckpointRestoredAsync(IWorkflowContext context, CancellationToken cancellationToken = default) { - this.Tries = await context.ReadStateAsync("TryCount").ConfigureAwait(false) ?? 0; + this.Tries = await context.ReadStateAsync("TryCount", cancellationToken: cancellationToken).ConfigureAwait(false) ?? 0; } public ValueTask ResetAsync() diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/08_Subworkflow_Simple.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/08_Subworkflow_Simple.cs index d865a6275c..35b739fa44 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/08_Subworkflow_Simple.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/08_Subworkflow_Simple.cs @@ -56,7 +56,7 @@ internal static class Step8EntryPoint return results; } - private static ValueTask ProcessTextAsync(TextProcessingRequest request, IWorkflowContext context, CancellationToken cancellation = default) + private static ValueTask ProcessTextAsync(TextProcessingRequest request, IWorkflowContext context, CancellationToken cancellationToken = default) { int wordCount = 0; int charCount = 0; @@ -67,7 +67,7 @@ internal static class Step8EntryPoint charCount = request.Text.Length; } - return context.YieldOutputAsync(new TextProcessingResult(request.TaskId, request.Text, wordCount, charCount)); + return context.YieldOutputAsync(new TextProcessingResult(request.TaskId, request.Text, wordCount, charCount), cancellationToken); } private sealed class TextProcessingOrchestrator() : Executor("TextOrchestrator") diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/SpecializedExecutorSmokeTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/SpecializedExecutorSmokeTests.cs index 6b883ebd1e..9aea98f068 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/SpecializedExecutorSmokeTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/SpecializedExecutorSmokeTests.cs @@ -115,28 +115,28 @@ public class SpecializedExecutorSmokeTests { public List> Updates { get; } = []; - public ValueTask AddEventAsync(WorkflowEvent workflowEvent) => + public ValueTask AddEventAsync(WorkflowEvent workflowEvent, CancellationToken cancellationToken = default) => default; - public ValueTask YieldOutputAsync(object output) => + public ValueTask YieldOutputAsync(object output, CancellationToken cancellationToken = default) => default; public ValueTask RequestHaltAsync() => default; - public ValueTask QueueClearScopeAsync(string? scopeName = null) => + public ValueTask QueueClearScopeAsync(string? scopeName = null, CancellationToken cancellationToken = default) => default; - public ValueTask QueueStateUpdateAsync(string key, T? value, string? scopeName = null) => + public ValueTask QueueStateUpdateAsync(string key, T? value, string? scopeName = null, CancellationToken cancellationToken = default) => default; - public ValueTask ReadStateAsync(string key, string? scopeName = null) => + public ValueTask ReadStateAsync(string key, string? scopeName = null, CancellationToken cancellationToken = default) => throw new NotImplementedException(); - public ValueTask> ReadStateKeysAsync(string? scopeName = null) => + public ValueTask> ReadStateKeysAsync(string? scopeName = null, CancellationToken cancellationToken = default) => throw new NotImplementedException(); - public ValueTask SendMessageAsync(object message, string? targetId = null) + public ValueTask SendMessageAsync(object message, string? targetId = null, CancellationToken cancellationToken = default) { if (message is List messages) { diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/TestRunContext.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/TestRunContext.cs index c996bdb4e1..17935892a6 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/TestRunContext.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/TestRunContext.cs @@ -15,36 +15,36 @@ public class TestRunContext : IRunnerContext TestRunContext runnerContext, IReadOnlyDictionary? traceContext) : IWorkflowContext { - public ValueTask AddEventAsync(WorkflowEvent workflowEvent) - => runnerContext.AddEventAsync(workflowEvent); + public ValueTask AddEventAsync(WorkflowEvent workflowEvent, CancellationToken cancellationToken = default) + => runnerContext.AddEventAsync(workflowEvent, cancellationToken); - public ValueTask YieldOutputAsync(object output) - => this.AddEventAsync(new WorkflowOutputEvent(output, executorId)); + public ValueTask YieldOutputAsync(object output, CancellationToken cancellationToken = default) + => this.AddEventAsync(new WorkflowOutputEvent(output, executorId), cancellationToken); public ValueTask RequestHaltAsync() => this.AddEventAsync(new RequestHaltEvent()); - public ValueTask QueueClearScopeAsync(string? scopeName = null) + public ValueTask QueueClearScopeAsync(string? scopeName = null, CancellationToken cancellationToken = default) => default; - public ValueTask QueueStateUpdateAsync(string key, T? value, string? scopeName = null) + public ValueTask QueueStateUpdateAsync(string key, T? value, string? scopeName = null, CancellationToken cancellationToken = default) => default; - public ValueTask ReadStateAsync(string key, string? scopeName = null) + public ValueTask ReadStateAsync(string key, string? scopeName = null, CancellationToken cancellationToken = default) => new(default(T?)); - public ValueTask> ReadStateKeysAsync(string? scopeName = null) + public ValueTask> ReadStateKeysAsync(string? scopeName = null, CancellationToken cancellationToken = default) => new([]); - public ValueTask SendMessageAsync(object message, string? targetId = null) - => runnerContext.SendMessageAsync(executorId, message, targetId); + public ValueTask SendMessageAsync(object message, string? targetId = null, CancellationToken cancellationToken = default) + => runnerContext.SendMessageAsync(executorId, message, targetId, cancellationToken); public IReadOnlyDictionary? TraceContext => traceContext; } public List Events { get; } = []; - public ValueTask AddEventAsync(WorkflowEvent workflowEvent) + public ValueTask AddEventAsync(WorkflowEvent workflowEvent, CancellationToken cancellationToken) { this.Events.Add(workflowEvent); return default; @@ -61,7 +61,7 @@ public class TestRunContext : IRunnerContext } internal Dictionary> QueuedMessages { get; } = []; - public ValueTask SendMessageAsync(string sourceId, object message, string? targetId = null) + public ValueTask SendMessageAsync(string sourceId, object message, string? targetId = null, CancellationToken cancellationToken = default) { if (!this.QueuedMessages.TryGetValue(sourceId, out List? deliveryQueue)) { @@ -72,7 +72,7 @@ public class TestRunContext : IRunnerContext return default; } - ValueTask IRunnerContext.AdvanceAsync() => + ValueTask IRunnerContext.AdvanceAsync(CancellationToken cancellationToken) => throw new NotImplementedException(); public Dictionary Executors { get; set; } = []; @@ -80,10 +80,10 @@ public class TestRunContext : IRunnerContext public bool WithCheckpointing => throw new NotSupportedException(); - ValueTask IRunnerContext.EnsureExecutorAsync(string executorId, IStepTracer? tracer) => + ValueTask IRunnerContext.EnsureExecutorAsync(string executorId, IStepTracer? tracer, CancellationToken cancellationToken) => new(this.Executors[executorId]); - public ValueTask> GetStartingExecutorInputTypesAsync(CancellationToken cancellation = default) + public ValueTask> GetStartingExecutorInputTypesAsync(CancellationToken cancellationToken = default) { if (this.Executors.TryGetValue(this.StartingExecutorId, out Executor? executor)) { @@ -93,11 +93,11 @@ public class TestRunContext : IRunnerContext throw new InvalidOperationException($"No executor with ID '{this.StartingExecutorId}' is registered in this context."); } - public ValueTask ForwardWorkflowEventAsync(WorkflowEvent workflowEvent, CancellationToken cancellation = default) - => this.AddEventAsync(workflowEvent); + public ValueTask ForwardWorkflowEventAsync(WorkflowEvent workflowEvent, CancellationToken cancellationToken = default) + => this.AddEventAsync(workflowEvent, cancellationToken); - public ValueTask SendMessageAsync(string senderId, [System.Diagnostics.CodeAnalysis.DisallowNull] TMessage message, CancellationToken cancellation = default) - => this.SendMessageAsync(senderId, message, cancellation); + public ValueTask SendMessageAsync(string senderId, [System.Diagnostics.CodeAnalysis.DisallowNull] TMessage message, CancellationToken cancellationToken = default) + => this.SendMessageAsync(senderId, message, cancellationToken); - ValueTask ISuperStepJoinContext.AttachSuperstepAsync(ISuperStepRunner superStepRunner, CancellationToken cancellation) => default; + ValueTask ISuperStepJoinContext.AttachSuperstepAsync(ISuperStepRunner superStepRunner, CancellationToken cancellationToken) => default; }