.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
This commit is contained in:
Jacob Alber
2025-10-07 20:28:44 -04:00
committed by GitHub
Unverified
parent eb049c43a6
commit 5902bcb10a
74 changed files with 910 additions and 557 deletions
@@ -134,17 +134,17 @@ internal sealed class SloganWriterExecutor
this._thread = this._agent.GetNewThread();
}
public async ValueTask<SloganResult> HandleAsync(string message, IWorkflowContext context)
public async ValueTask<SloganResult> 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<SloganResult>(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<SloganResult> HandleAsync(FeedbackResult message, IWorkflowContext context)
public async ValueTask<SloganResult> 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<SloganResult>(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<FeedbackExecutor>, 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<FeedbackExecutor>, 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<FeedbackResult>(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++;
}
}
@@ -51,13 +51,15 @@ internal static class WorkflowHelper
/// </summary>
/// <param name="message">The user message to process</param>
/// <param name="context">Workflow context for accessing workflow services and adding events</param>
public async ValueTask HandleAsync(List<ChatMessage> message, IWorkflowContext context)
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
public async ValueTask HandleAsync(List<ChatMessage> 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
/// </summary>
/// <param name="message">The message from the agent</param>
/// <param name="context">Workflow context for accessing workflow services and adding events</param>
public async ValueTask HandleAsync(ChatMessage message, IWorkflowContext context)
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
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);
}
}
}
@@ -69,20 +69,20 @@ internal sealed class GuessNumberExecutor() : ReflectingExecutor<GuessNumberExec
private int NextGuess => (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<GuessNumberExec
/// This must be overridden to save any state that is needed to resume the executor.
/// </summary>
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);
/// <summary>
/// Restore the state of the executor from a checkpoint.
/// This must be overridden to restore any state that was saved during checkpointing.
/// </summary>
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);
}
/// <summary>
@@ -120,20 +120,20 @@ internal sealed class JudgeExecutor() : ReflectingExecutor<JudgeExecutor>("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<JudgeExecutor>("Judge
/// This must be overridden to save any state that is needed to resume the executor.
/// </summary>
protected override ValueTask OnCheckpointingAsync(IWorkflowContext context, CancellationToken cancellationToken = default) =>
context.QueueStateUpdateAsync(StateKey, this._tries);
context.QueueStateUpdateAsync(StateKey, this._tries, cancellationToken: cancellationToken);
/// <summary>
/// Restore the state of the executor from a checkpoint.
/// This must be overridden to restore any state that was saved during checkpointing.
/// </summary>
protected override async ValueTask OnCheckpointRestoredAsync(IWorkflowContext context, CancellationToken cancellationToken = default) =>
this._tries = await context.ReadStateAsync<int>(StateKey).ConfigureAwait(false);
this._tries = await context.ReadStateAsync<int>(StateKey, cancellationToken: cancellationToken).ConfigureAwait(false);
}
@@ -69,20 +69,20 @@ internal sealed class GuessNumberExecutor() : ReflectingExecutor<GuessNumberExec
private int NextGuess => (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<GuessNumberExec
/// This must be overridden to save any state that is needed to resume the executor.
/// </summary>
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);
/// <summary>
/// Restore the state of the executor from a checkpoint.
/// This must be overridden to restore any state that was saved during checkpointing.
/// </summary>
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);
}
/// <summary>
@@ -120,20 +120,20 @@ internal sealed class JudgeExecutor() : ReflectingExecutor<JudgeExecutor>("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<JudgeExecutor>("Judge
/// This must be overridden to save any state that is needed to resume the executor.
/// </summary>
protected override ValueTask OnCheckpointingAsync(IWorkflowContext context, CancellationToken cancellationToken = default) =>
context.QueueStateUpdateAsync(StateKey, this._tries);
context.QueueStateUpdateAsync(StateKey, this._tries, cancellationToken: cancellationToken);
/// <summary>
/// Restore the state of the executor from a checkpoint.
/// This must be overridden to restore any state that was saved during checkpointing.
/// </summary>
protected override async ValueTask OnCheckpointRestoredAsync(IWorkflowContext context, CancellationToken cancellationToken = default) =>
this._tries = await context.ReadStateAsync<int>(StateKey).ConfigureAwait(false);
this._tries = await context.ReadStateAsync<int>(StateKey, cancellationToken: cancellationToken).ConfigureAwait(false);
}
@@ -69,21 +69,21 @@ internal sealed class JudgeExecutor() : ReflectingExecutor<JudgeExecutor>("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<JudgeExecutor>("Judge
/// This must be overridden to save any state that is needed to resume the executor.
/// </summary>
protected override ValueTask OnCheckpointingAsync(IWorkflowContext context, CancellationToken cancellationToken = default) =>
context.QueueStateUpdateAsync(StateKey, this._tries);
context.QueueStateUpdateAsync(StateKey, this._tries, cancellationToken: cancellationToken);
/// <summary>
/// Restore the state of the executor from a checkpoint.
/// This must be overridden to restore any state that was saved during checkpointing.
/// </summary>
protected override async ValueTask OnCheckpointRestoredAsync(IWorkflowContext context, CancellationToken cancellationToken = default) =>
this._tries = await context.ReadStateAsync<int>(StateKey).ConfigureAwait(false);
this._tries = await context.ReadStateAsync<int>(StateKey, cancellationToken: cancellationToken).ConfigureAwait(false);
}
@@ -82,14 +82,16 @@ internal sealed class ConcurrentStartExecutor() :
/// </summary>
/// <param name="message">The user message to process</param>
/// <param name="context">Workflow context for accessing workflow services and adding events</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>A task representing the asynchronous operation</returns>
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() :
/// </summary>
/// <param name="message">The message from the agent</param>
/// <param name="context">Workflow context for accessing workflow services and adding events</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>A task representing the asynchronous operation</returns>
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);
}
}
}
@@ -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) :
/// <summary>
/// Tokenize input and assign contiguous index ranges to each mapper via shared state.
/// </summary>
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<Mapper>(id), IMessa
/// <summary>
/// Read the assigned slice, emit (word, 1) pairs, and persist to disk.
/// </summary>
public async ValueTask HandleAsync(SplitComplete message, IWorkflowContext context)
public async ValueTask HandleAsync(SplitComplete message, IWorkflowContext context, CancellationToken cancellationToken = default)
{
var dataToProcess = await context.ReadStateAsync<string[]>(MapReduceConstants.DataToProcessKey, scopeName: MapReduceConstants.StateScope);
var chunk = await context.ReadStateAsync<(int start, int end)>(this.Id, scopeName: MapReduceConstants.StateScope);
var dataToProcess = await context.ReadStateAsync<string[]>(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<Mapper>(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
/// <summary>
/// Aggregate mapper outputs and write one partition file per reducer.
/// </summary>
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<Reducer>(id), IMes
/// <summary>
/// Read one shuffle partition and reduce it to totals.
/// </summary>
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<Reducer>(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<string, int>();
@@ -345,9 +346,9 @@ internal sealed class Reducer(string id) : ReflectingExecutor<Reducer>(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) :
/// <summary>
/// Collect reducer output file paths and yield final output.
/// </summary>
public async ValueTask HandleAsync(List<ReduceComplete> message, IWorkflowContext context)
public async ValueTask HandleAsync(List<ReduceComplete> message, IWorkflowContext context, CancellationToken cancellationToken = default)
{
var filePaths = message.ConvertAll(r => r.FilePath);
await context.YieldOutputAsync(filePaths);
await context.YieldOutputAsync(filePaths, cancellationToken);
}
}
@@ -160,7 +160,7 @@ internal sealed class SpamDetectionExecutor : ReflectingExecutor<SpamDetectionEx
this._spamDetectionAgent = spamDetectionAgent;
}
public async ValueTask<DetectionResult> HandleAsync(ChatMessage message, IWorkflowContext context)
public async ValueTask<DetectionResult> 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<SpamDetectionEx
EmailId = Guid.NewGuid().ToString("N"),
EmailContent = message.Text
};
await context.QueueStateUpdateAsync(newEmail.EmailId, newEmail, scopeName: EmailStateConstants.EmailStateScope);
await context.QueueStateUpdateAsync(newEmail.EmailId, newEmail, scopeName: EmailStateConstants.EmailStateScope, cancellationToken);
// Invoke the agent
var response = await this._spamDetectionAgent.RunAsync(message);
var response = await this._spamDetectionAgent.RunAsync(message, cancellationToken: cancellationToken);
var detectionResult = JsonSerializer.Deserialize<DetectionResult>(response.Text);
detectionResult!.EmailId = newEmail.EmailId;
@@ -205,7 +205,7 @@ internal sealed class EmailAssistantExecutor : ReflectingExecutor<EmailAssistant
this._emailAssistantAgent = emailAssistantAgent;
}
public async ValueTask<EmailResponse> HandleAsync(DetectionResult message, IWorkflowContext context)
public async ValueTask<EmailResponse> HandleAsync(DetectionResult message, IWorkflowContext context, CancellationToken cancellationToken = default)
{
if (message.IsSpam)
{
@@ -213,11 +213,11 @@ internal sealed class EmailAssistantExecutor : ReflectingExecutor<EmailAssistant
}
// Retrieve the email content from the shared state
var email = await context.ReadStateAsync<Email>(message.EmailId, scopeName: EmailStateConstants.EmailStateScope)
var email = await context.ReadStateAsync<Email>(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<EmailResponse>(response.Text);
return emailResponse!;
@@ -232,8 +232,8 @@ internal sealed class SendEmailExecutor() : ReflectingExecutor<SendEmailExecutor
/// <summary>
/// Simulate the sending of an email.
/// </summary>
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);
}
/// <summary>
@@ -244,11 +244,11 @@ internal sealed class HandleSpamExecutor() : ReflectingExecutor<HandleSpamExecut
/// <summary>
/// Simulate the handling of a spam message.
/// </summary>
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
{
@@ -185,7 +185,7 @@ internal sealed class SpamDetectionExecutor : ReflectingExecutor<SpamDetectionEx
this._spamDetectionAgent = spamDetectionAgent;
}
public async ValueTask<DetectionResult> HandleAsync(ChatMessage message, IWorkflowContext context)
public async ValueTask<DetectionResult> 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<SpamDetectionEx
EmailId = Guid.NewGuid().ToString("N"),
EmailContent = message.Text
};
await context.QueueStateUpdateAsync(newEmail.EmailId, newEmail, scopeName: EmailStateConstants.EmailStateScope);
await context.QueueStateUpdateAsync(newEmail.EmailId, newEmail, scopeName: EmailStateConstants.EmailStateScope, cancellationToken);
// Invoke the agent
var response = await this._spamDetectionAgent.RunAsync(message);
var response = await this._spamDetectionAgent.RunAsync(message, cancellationToken: cancellationToken);
var detectionResult = JsonSerializer.Deserialize<DetectionResult>(response.Text);
detectionResult!.EmailId = newEmail.EmailId;
@@ -230,7 +230,7 @@ internal sealed class EmailAssistantExecutor : ReflectingExecutor<EmailAssistant
this._emailAssistantAgent = emailAssistantAgent;
}
public async ValueTask<EmailResponse> HandleAsync(DetectionResult message, IWorkflowContext context)
public async ValueTask<EmailResponse> HandleAsync(DetectionResult message, IWorkflowContext context, CancellationToken cancellationToken = default)
{
if (message.spamDecision == SpamDecision.Spam)
{
@@ -238,10 +238,10 @@ internal sealed class EmailAssistantExecutor : ReflectingExecutor<EmailAssistant
}
// Retrieve the email content from the context
var email = await context.ReadStateAsync<Email>(message.EmailId, scopeName: EmailStateConstants.EmailStateScope);
var email = await context.ReadStateAsync<Email>(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<EmailResponse>(response.Text);
return emailResponse!;
@@ -256,8 +256,8 @@ internal sealed class SendEmailExecutor() : ReflectingExecutor<SendEmailExecutor
/// <summary>
/// Simulate the sending of an email.
/// </summary>
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);
}
/// <summary>
@@ -268,11 +268,11 @@ internal sealed class HandleSpamExecutor() : ReflectingExecutor<HandleSpamExecut
/// <summary>
/// Simulate the handling of a spam message.
/// </summary>
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<HandleUncer
/// <summary>
/// Simulate the handling of an uncertain spam decision.
/// </summary>
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<Email>(message.EmailId, scopeName: EmailStateConstants.EmailStateScope);
await context.YieldOutputAsync($"Email marked as uncertain: {message.Reason}. Email content: {email?.EmailContent}");
var email = await context.ReadStateAsync<Email>(message.EmailId, scopeName: EmailStateConstants.EmailStateScope, cancellationToken);
await context.YieldOutputAsync($"Email marked as uncertain: {message.Reason}. Email content: {email?.EmailContent}", cancellationToken);
}
else
{
@@ -241,7 +241,7 @@ internal sealed class EmailAnalysisExecutor : ReflectingExecutor<EmailAnalysisEx
this._emailAnalysisAgent = emailAnalysisAgent;
}
public async ValueTask<AnalysisResult> HandleAsync(ChatMessage message, IWorkflowContext context)
public async ValueTask<AnalysisResult> 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<EmailAnalysisEx
EmailId = Guid.NewGuid().ToString("N"),
EmailContent = message.Text
};
await context.QueueStateUpdateAsync(newEmail.EmailId, newEmail, scopeName: EmailStateConstants.EmailStateScope);
await context.QueueStateUpdateAsync(newEmail.EmailId, newEmail, scopeName: EmailStateConstants.EmailStateScope, cancellationToken);
// Invoke the agent
var response = await this._emailAnalysisAgent.RunAsync(message);
var response = await this._emailAnalysisAgent.RunAsync(message, cancellationToken: cancellationToken);
var AnalysisResult = JsonSerializer.Deserialize<AnalysisResult>(response.Text);
AnalysisResult!.EmailId = newEmail.EmailId;
@@ -287,7 +287,7 @@ internal sealed class EmailAssistantExecutor : ReflectingExecutor<EmailAssistant
this._emailAssistantAgent = emailAssistantAgent;
}
public async ValueTask<EmailResponse> HandleAsync(AnalysisResult message, IWorkflowContext context)
public async ValueTask<EmailResponse> HandleAsync(AnalysisResult message, IWorkflowContext context, CancellationToken cancellationToken = default)
{
if (message.spamDecision == SpamDecision.Spam)
{
@@ -295,10 +295,10 @@ internal sealed class EmailAssistantExecutor : ReflectingExecutor<EmailAssistant
}
// Retrieve the email content from the context
var email = await context.ReadStateAsync<Email>(message.EmailId, scopeName: EmailStateConstants.EmailStateScope);
var email = await context.ReadStateAsync<Email>(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<EmailResponse>(response.Text);
return emailResponse!;
@@ -313,8 +313,8 @@ internal sealed class SendEmailExecutor() : ReflectingExecutor<SendEmailExecutor
/// <summary>
/// Simulate the sending of an email.
/// </summary>
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);
}
/// <summary>
@@ -325,11 +325,11 @@ internal sealed class HandleSpamExecutor() : ReflectingExecutor<HandleSpamExecut
/// <summary>
/// Simulate the handling of a spam message.
/// </summary>
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<HandleUncer
/// <summary>
/// Simulate the handling of an uncertain spam decision.
/// </summary>
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<Email>(message.EmailId, scopeName: EmailStateConstants.EmailStateScope);
await context.YieldOutputAsync($"Email marked as uncertain: {message.Reason}. Email content: {email?.EmailContent}");
var email = await context.ReadStateAsync<Email>(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<EmailSummaryExec
this._emailSummaryAgent = emailSummaryAgent;
}
public async ValueTask<AnalysisResult> HandleAsync(AnalysisResult message, IWorkflowContext context)
public async ValueTask<AnalysisResult> HandleAsync(AnalysisResult message, IWorkflowContext context, CancellationToken cancellationToken = default)
{
// Read the email content from the shared states
var email = await context.ReadStateAsync<Email>(message.EmailId, scopeName: EmailStateConstants.EmailStateScope);
var email = await context.ReadStateAsync<Email>(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<EmailSummary>(response.Text);
message.EmailSummary = emailSummary!.Summary;
@@ -410,17 +410,17 @@ internal sealed class DatabaseEvent(string message) : WorkflowEvent(message) { }
/// </summary>
internal sealed class DatabaseAccessExecutor() : ReflectingExecutor<DatabaseAccessExecutor>("DatabaseAccessExecutor"), IMessageHandler<AnalysisResult>
{
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<Email>(message.EmailId, scopeName: EmailStateConstants.EmailStateScope);
await Task.Delay(100); // Simulate database access delay
await context.ReadStateAsync<Email>(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);
}
}
@@ -53,21 +53,21 @@ internal sealed class JudgeExecutor() : ReflectingExecutor<JudgeExecutor>("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);
}
}
}
@@ -83,20 +83,20 @@ internal sealed class GuessNumberExecutor : ReflectingExecutor<GuessNumberExecut
private int NextGuess => (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<JudgeExecutor>, 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);
}
}
}
@@ -76,8 +76,10 @@ internal sealed class UppercaseExecutor() : ReflectingExecutor<UppercaseExecutor
/// </summary>
/// <param name="message">The input text to convert</param>
/// <param name="context">Workflow context for accessing workflow services and adding events</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>The input text converted to uppercase</returns>
public async ValueTask<string> HandleAsync(string message, IWorkflowContext context) =>
public async ValueTask<string> 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<ReverseTextExec
/// </summary>
/// <param name="message">The input text to reverse</param>
/// <param name="context">Workflow context for accessing workflow services and adding events</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>The input text reversed</returns>
public async ValueTask<string> HandleAsync(string message, IWorkflowContext context) => new string(message.Reverse().ToArray());
public async ValueTask<string> HandleAsync(string message, IWorkflowContext context, CancellationToken cancellationToken = default)
=> new(message.Reverse().ToArray());
}
@@ -78,8 +78,10 @@ internal sealed class UppercaseExecutor() : ReflectingExecutor<UppercaseExecutor
/// </summary>
/// <param name="message">The input text to convert</param>
/// <param name="context">Workflow context for accessing workflow services and adding events</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>The input text converted to uppercase</returns>
public async ValueTask<string> HandleAsync(string message, IWorkflowContext context) =>
public async ValueTask<string> 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<ReverseTextExec
/// </summary>
/// <param name="message">The input text to reverse</param>
/// <param name="context">Workflow context for accessing workflow services and adding events</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>The input text reversed</returns>
public async ValueTask<string> HandleAsync(string message, IWorkflowContext context) => new string(message.Reverse().ToArray());
public async ValueTask<string> HandleAsync(string message, IWorkflowContext context, CancellationToken cancellationToken = default)
=> new(message.Reverse().ToArray());
}
@@ -54,13 +54,13 @@ internal static class FileContentStateConstants
internal sealed class FileReadExecutor() : ReflectingExecutor<FileReadExecutor>("FileReadExecutor"), IMessageHandler<string, string>
{
public async ValueTask<string> HandleAsync(string message, IWorkflowContext context)
public async ValueTask<string> 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>("WordCountingExecutor"), IMessageHandler<string, FileStats>
{
public async ValueTask<FileStats> HandleAsync(string message, IWorkflowContext context)
public async ValueTask<FileStats> HandleAsync(string message, IWorkflowContext context, CancellationToken cancellationToken = default)
{
// Retrieve the file content from the shared state
var fileContent = await context.ReadStateAsync<string>(message, scopeName: FileContentStateConstants.FileContentStateScope)
var fileContent = await context.ReadStateAsync<string>(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<WordCountingEx
}
}
internal sealed class ParagraphCountingExecutor() : ReflectingExecutor<ParagraphCountingExecutor>("ParagraphCountingExecutor"), IMessageHandler<string, FileStats>
internal sealed class ParagraphCountingExecutor() : ReflectingExecutor<ParagraphCountingExecutor>("ParagraphCountingExecutor"),
IMessageHandler<string, FileStats>
{
public async ValueTask<FileStats> HandleAsync(string message, IWorkflowContext context)
public async ValueTask<FileStats> HandleAsync(string message, IWorkflowContext context, CancellationToken cancellationToken = default)
{
// Retrieve the file content from the shared state
var fileContent = await context.ReadStateAsync<string>(message, scopeName: FileContentStateConstants.FileContentStateScope)
var fileContent = await context.ReadStateAsync<string>(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<AggregationExec
{
private readonly List<FileStats> _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<AggregationExec
// Aggregate the results from both executors
var totalParagraphCount = this._messages.Sum(m => 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);
}
}
}
@@ -51,8 +51,10 @@ internal sealed class UppercaseExecutor() : ReflectingExecutor<UppercaseExecutor
/// </summary>
/// <param name="message">The input text to convert</param>
/// <param name="context">Workflow context for accessing workflow services and adding events</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>The input text converted to uppercase</returns>
public async ValueTask<string> HandleAsync(string message, IWorkflowContext context) =>
public async ValueTask<string> 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<ReverseTextExec
/// </summary>
/// <param name="message">The input text to reverse</param>
/// <param name="context">Workflow context for accessing workflow services and adding events</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>The input text reversed</returns>
public async ValueTask<string> HandleAsync(string message, IWorkflowContext context)
public async ValueTask<string> 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());
@@ -50,8 +50,10 @@ internal sealed class UppercaseExecutor() : ReflectingExecutor<UppercaseExecutor
/// </summary>
/// <param name="message">The input text to convert</param>
/// <param name="context">Workflow context for accessing workflow services and adding events</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>The input text converted to uppercase</returns>
public async ValueTask<string> HandleAsync(string message, IWorkflowContext context) =>
public async ValueTask<string> 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<ReverseTextExec
/// </summary>
/// <param name="message">The input text to reverse</param>
/// <param name="context">Workflow context for accessing workflow services and adding events</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>The input text reversed</returns>
public async ValueTask<string> HandleAsync(string message, IWorkflowContext context)
public async ValueTask<string> 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());
@@ -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);
}
}
}
@@ -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<TValue>(this IWorkflowContext context, PropertyPath variablePath, TValue? value) =>
context.QueueStateUpdateAsync(Throw.IfNull(variablePath.VariableName), value, Throw.IfNull(variablePath.NamespaceAlias));
public static ValueTask QueueStateUpdateAsync<TValue>(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<TValue>(this IWorkflowContext context, string key, TValue? value) =>
DeclarativeContext(context).QueueSystemUpdateAsync(key, value);
public static ValueTask QueueSystemUpdateAsync<TValue>(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(
@@ -61,7 +61,7 @@ internal abstract class DeclarativeActionExecutor : Executor<ActionExecutorResul
}
/// <inheritdoc/>
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<ActionExecutorResul
return;
}
await context.RaiseInvocationEventAsync(this.Model, message.ExecutorId).ConfigureAwait(false);
await context.RaiseInvocationEventAsync(this.Model, message.ExecutorId, cancellationToken).ConfigureAwait(false);
try
{
@@ -78,7 +78,7 @@ internal abstract class DeclarativeActionExecutor : Executor<ActionExecutorResul
if (this.EmitResultEvent)
{
await context.SendResultMessageAsync(this.Id, result).ConfigureAwait(false);
await context.SendResultMessageAsync(this.Id, result, cancellationToken).ConfigureAwait(false);
}
}
catch (DeclarativeActionException exception)
@@ -95,7 +95,7 @@ internal abstract class DeclarativeActionExecutor : Executor<ActionExecutorResul
{
if (this.IsDiscreteAction)
{
await context.RaiseCompletionEventAsync(this.Model).ConfigureAwait(false);
await context.RaiseCompletionEventAsync(this.Model, cancellationToken).ConfigureAwait(false);
}
}
}
@@ -3,6 +3,7 @@
using System.Collections.Frozen;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Agents.AI.Workflows.Declarative.Extensions;
using Microsoft.Agents.AI.Workflows.Declarative.Kit;
@@ -32,16 +33,18 @@ internal sealed class DeclarativeWorkflowContext : IWorkflowContext
public IReadOnlyDictionary<string, string>? TraceContext => this.Source.TraceContext;
/// <inheritdoc/>
public ValueTask AddEventAsync(WorkflowEvent workflowEvent) => this.Source.AddEventAsync(workflowEvent);
public ValueTask AddEventAsync(WorkflowEvent workflowEvent, CancellationToken cancellationToken = default)
=> this.Source.AddEventAsync(workflowEvent, cancellationToken);
/// <inheritdoc/>
public ValueTask YieldOutputAsync(object output) => this.Source.YieldOutputAsync(output);
public ValueTask YieldOutputAsync(object output, CancellationToken cancellationToken = default)
=> this.Source.YieldOutputAsync(output, cancellationToken);
/// <inheritdoc/>
public ValueTask RequestHaltAsync() => this.Source.RequestHaltAsync();
/// <inheritdoc/>
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
}
/// <inheritdoc/>
public async ValueTask QueueStateUpdateAsync<T>(string key, T? value, string? scopeName = null)
public async ValueTask QueueStateUpdateAsync<T>(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<TValue>(string key, TValue? value)
public async ValueTask QueueSystemUpdateAsync<TValue>(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();
}
/// <inheritdoc/>
public async ValueTask<TValue?> ReadStateAsync<TValue>(string key, string? scopeName = null)
public async ValueTask<TValue?> ReadStateAsync<TValue>(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<TValue>(key, scopeName).ConfigureAwait(false),
_ when !isManagedScope => await this.Source.ReadStateAsync<TValue>(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<TValue>(key, scopeName).ConfigureAwait(false),
_ => await this.Source.ReadStateAsync<TValue>(key, scopeName, cancellationToken).ConfigureAwait(false),
};
}
/// <inheritdoc/>
public ValueTask<HashSet<string>> ReadStateKeysAsync(string? scopeName = null) => this.Source.ReadStateKeysAsync(scopeName);
public ValueTask<HashSet<string>> ReadStateKeysAsync(string? scopeName = null, CancellationToken cancellationToken = default)
=> this.Source.ReadStateKeysAsync(scopeName, cancellationToken);
/// <inheritdoc/>
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<T>(string key, T? value, string? scopeName, bool allowSystem = true)
private ValueTask UpdateStateAsync<T>(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);
}
}
}
@@ -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<TInput>(
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<TInput>(
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);
}
}
@@ -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<ActionExecutorResult>? action = null, bool emitResult = true)
: DelegateActionExecutor<ActionExecutorResult>(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<TMessage> : Executor<TMessage>, 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);
}
}
}
@@ -73,12 +73,12 @@ public abstract class ActionExecutor<TMessage> : Executor<TMessage>, IResettable
}
/// <inheritdoc/>
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);
}
/// <summary>
@@ -132,7 +132,7 @@ public static class IWorkflowContextExtensions
/// <returns>The converted value</returns>
public static async ValueTask<object?> ConvertValueAsync(this IWorkflowContext context, VariableType targetType, string key, string? scopeName = null, CancellationToken cancellationToken = default)
{
object? sourceValue = await context.ReadStateAsync<object>(key, scopeName).ConfigureAwait(false);
object? sourceValue = await context.ReadStateAsync<object>(key, scopeName, cancellationToken).ConfigureAwait(false);
return sourceValue.ConvertType(targetType);
}
@@ -143,10 +143,11 @@ public static class IWorkflowContextExtensions
/// <param name="context">The workflow execution context used to restore persisted state prior to formatting.</param>
/// <param name="key">The key of the state value.</param>
/// <param name = "scopeName" > An optional name that specifies the scope to read.If null, the default scope is used.</param>
/// <param name="cancellationToken">A token that propagates notification when operation should be canceled.</param>
/// <returns>The evaluated list expression</returns>
public static async ValueTask<IList<TElement>?> ReadListAsync<TElement>(this IWorkflowContext context, string key, string? scopeName = null)
public static async ValueTask<IList<TElement>?> ReadListAsync<TElement>(this IWorkflowContext context, string key, string? scopeName = null, CancellationToken cancellationToken = default)
{
object? value = await context.ReadStateAsync<object>(key, scopeName).ConfigureAwait(false);
object? value = await context.ReadStateAsync<object>(key, scopeName, cancellationToken).ConfigureAwait(false);
return value.AsList<TElement>();
}
@@ -55,23 +55,23 @@ public abstract class RootExecutor<TInput> : Executor<TInput>, IResettableExecut
}
/// <inheritdoc/>
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);
}
/// <summary>
@@ -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}]
@@ -69,5 +69,5 @@ internal sealed class ConditionGroupExecutor : DeclarativeActionExecutor<Conditi
}
public async ValueTask DoneAsync(IWorkflowContext context, ActionExecutorResult _, CancellationToken cancellationToken) =>
await context.RaiseCompletionEventAsync(this.Model).ConfigureAwait(false);
await context.RaiseCompletionEventAsync(this.Model, cancellationToken).ConfigureAwait(false);
}
@@ -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;
}
@@ -68,11 +68,11 @@ internal sealed class ForeachExecutor : DeclarativeActionExecutor<Foreach>
{
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<Foreach>
{
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);
}
}
}
@@ -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
@@ -17,7 +17,7 @@ internal sealed class ResetVariableExecutor(ResetVariable model, WorkflowFormula
protected override async ValueTask<object?> 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}]
@@ -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;
@@ -72,10 +72,10 @@ internal sealed class WorkflowFormulaState
async Task ReadScopeAsync(string scopeName)
{
HashSet<string> keys = await context.ReadStateKeysAsync(scopeName).ConfigureAwait(false);
HashSet<string> keys = await context.ReadStateKeysAsync(scopeName, cancellationToken).ConfigureAwait(false);
foreach (string key in keys)
{
object? value = await context.ReadStateAsync<object>(key, scopeName).ConfigureAwait(false);
object? value = await context.ReadStateAsync<object>(key, scopeName, cancellationToken).ConfigureAwait(false);
if (value is null or UnassignedValue)
{
value = FormulaValue.NewBlank();
@@ -150,12 +150,12 @@ public static partial class AgentWorkflowBuilder
protected override RouteBuilder ConfigureRoutes(RouteBuilder routeBuilder) =>
routeBuilder
.AddHandler<string>((message, context) => this._pendingMessages.Add(new(ChatRole.User, message)))
.AddHandler<ChatMessage>((message, context) => this._pendingMessages.Add(message))
.AddHandler<IEnumerable<ChatMessage>>((messages, _) => this._pendingMessages.AddRange(messages))
.AddHandler<ChatMessage[]>((messages, _) => this._pendingMessages.AddRange(messages)) // TODO: Remove once https://github.com/microsoft/agent-framework/issues/782 is addressed
.AddHandler<List<ChatMessage>>((messages, _) => this._pendingMessages.AddRange(messages)) // TODO: Remove once https://github.com/microsoft/agent-framework/issues/782 is addressed
.AddHandler<TurnToken>(async (token, context) =>
.AddHandler<string>((message, _, __) => this._pendingMessages.Add(new(ChatRole.User, message)))
.AddHandler<ChatMessage>((message, _, __) => this._pendingMessages.Add(message))
.AddHandler<IEnumerable<ChatMessage>>((messages, _, __) => this._pendingMessages.AddRange(messages))
.AddHandler<ChatMessage[]>((messages, _, __) => this._pendingMessages.AddRange(messages)) // TODO: Remove once https://github.com/microsoft/agent-framework/issues/782 is addressed
.AddHandler<List<ChatMessage>>((messages, _, __) => this._pendingMessages.AddRange(messages)) // TODO: Remove once https://github.com/microsoft/agent-framework/issues/782 is addressed
.AddHandler<TurnToken>(async (token, context, cancellationToken) =>
{
List<ChatMessage> messages = [.. this._pendingMessages];
this._pendingMessages.Clear();
@@ -163,12 +163,12 @@ public static partial class AgentWorkflowBuilder
List<ChatMessage>? roleChanged = ChangeAssistantToUserForOtherParticipants(agent.DisplayName, messages);
List<AgentRunResponseUpdate> 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<ChatMessage> 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<string>((message, context) => context.SendMessageAsync(new ChatMessage(ChatRole.User, message)))
.AddHandler<ChatMessage>((message, context) => context.SendMessageAsync(message))
.AddHandler<List<ChatMessage>>((messages, context) => context.SendMessageAsync(messages))
.AddHandler<TurnToken>((turnToken, context) => context.SendMessageAsync(turnToken));
.AddHandler<string>((message, context, cancellationToken) => context.SendMessageAsync(new ChatMessage(ChatRole.User, message), cancellationToken: cancellationToken))
.AddHandler<ChatMessage>((message, context, cancellationToken) => context.SendMessageAsync(message, cancellationToken: cancellationToken))
.AddHandler<List<ChatMessage>>((messages, context, cancellationToken) => context.SendMessageAsync(messages, cancellationToken: cancellationToken))
.AddHandler<TurnToken>((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<ChatMessage> 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<List<ChatMessage>>(async (messages, context) =>
routeBuilder.AddHandler<List<ChatMessage>>(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<List<ChatMessage>>(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<string>((message, context) => this._pendingMessages.Add(new(ChatRole.User, message)))
.AddHandler<ChatMessage>((message, context) => this._pendingMessages.Add(message))
.AddHandler<IEnumerable<ChatMessage>>((messages, _) => this._pendingMessages.AddRange(messages))
.AddHandler<ChatMessage[]>((messages, _) => this._pendingMessages.AddRange(messages)) // TODO: Remove once https://github.com/microsoft/agent-framework/issues/782 is addressed
.AddHandler<List<ChatMessage>>((messages, _) => this._pendingMessages.AddRange(messages)) // TODO: Remove once https://github.com/microsoft/agent-framework/issues/782 is addressed
.AddHandler<TurnToken>(async (token, context) =>
.AddHandler<string>((message, context, _) => this._pendingMessages.Add(new(ChatRole.User, message)))
.AddHandler<ChatMessage>((message, context, _) => this._pendingMessages.Add(message))
.AddHandler<IEnumerable<ChatMessage>>((messages, _, __) => this._pendingMessages.AddRange(messages))
.AddHandler<ChatMessage[]>((messages, _, __) => this._pendingMessages.AddRange(messages)) // TODO: Remove once https://github.com/microsoft/agent-framework/issues/782 is addressed
.AddHandler<List<ChatMessage>>((messages, _, __) => this._pendingMessages.AddRange(messages)) // TODO: Remove once https://github.com/microsoft/agent-framework/issues/782 is addressed
.AddHandler<TurnToken>(async (token, context, cancellationToken) =>
{
var messages = new List<ChatMessage>(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<HandoffState>((handoff, context) =>
context.YieldOutputAsync(handoff.Messages));
routeBuilder.AddHandler<HandoffState>((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<HandoffState>(async (handoffState, context) =>
routeBuilder.AddHandler<HandoffState>(async (handoffState, context, cancellationToken) =>
{
string? requestedHandoff = null;
List<AgentRunResponseUpdate> updates = [];
@@ -542,24 +543,31 @@ public static partial class AgentWorkflowBuilder
List<ChatMessage>? 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.
/// </summary>
/// <param name="history">The chat history to consider.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>The next <see cref="AIAgent"/> to speak. This agent must be part of the chat.</returns>
protected internal abstract ValueTask<AIAgent> SelectNextAgentAsync(
IReadOnlyList<ChatMessage> history,
@@ -633,7 +642,8 @@ public static partial class AgentWorkflowBuilder
/// Filters the chat history before it's passed to the next agent.
/// </summary>
/// <param name="history">The chat history to filter.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>The filtered chat history.</returns>
protected internal virtual ValueTask<IEnumerable<ChatMessage>> UpdateHistoryAsync(
IReadOnlyList<ChatMessage> 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.
/// </summary>
/// <param name="history">The chat history to consider.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>A <see cref="bool"/> indicating whether the chat should be terminated.</returns>
protected internal virtual ValueTask<bool> ShouldTerminateAsync(
IReadOnlyList<ChatMessage> history,
@@ -789,35 +800,35 @@ public static partial class AgentWorkflowBuilder
private GroupChatManager? _manager;
protected override RouteBuilder ConfigureRoutes(RouteBuilder routeBuilder) => routeBuilder
.AddHandler<string>((message, context) => this._pendingMessages.Add(new(ChatRole.User, message)))
.AddHandler<ChatMessage>((message, context) => this._pendingMessages.Add(message))
.AddHandler<IEnumerable<ChatMessage>>((messages, _) => this._pendingMessages.AddRange(messages))
.AddHandler<ChatMessage[]>((messages, _) => this._pendingMessages.AddRange(messages)) // TODO: Remove once https://github.com/microsoft/agent-framework/issues/782 is addressed
.AddHandler<List<ChatMessage>>((messages, _) => this._pendingMessages.AddRange(messages)) // TODO: Remove once https://github.com/microsoft/agent-framework/issues/782 is addressed
.AddHandler<TurnToken>(async (token, context) =>
.AddHandler<string>((message, context, _) => this._pendingMessages.Add(new(ChatRole.User, message)))
.AddHandler<ChatMessage>((message, context, _) => this._pendingMessages.Add(message))
.AddHandler<IEnumerable<ChatMessage>>((messages, _, __) => this._pendingMessages.AddRange(messages))
.AddHandler<ChatMessage[]>((messages, _, __) => this._pendingMessages.AddRange(messages)) // TODO: Remove once https://github.com/microsoft/agent-framework/issues/782 is addressed
.AddHandler<List<ChatMessage>>((messages, _, __) => this._pendingMessages.AddRange(messages)) // TODO: Remove once https://github.com/microsoft/agent-framework/issues/782 is addressed
.AddHandler<TurnToken>(async (token, context, cancellationToken) =>
{
List<ChatMessage> 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()
@@ -28,7 +28,7 @@ public class AggregatingExecutor<TInput, TAggregate>(string id,
private TAggregate? _runningAggregate;
/// <inheritdoc/>
public override ValueTask<TAggregate?> HandleAsync(TInput message, IWorkflowContext context)
public override ValueTask<TAggregate?> 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<TInput, TAggregate>(string id,
/// <inheritdoc/>
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<TInput, TAggregate>(string id,
{
await base.OnCheckpointRestoredAsync(context, cancellationToken).ConfigureAwait(false);
this._runningAggregate = await context.ReadStateAsync<TAggregate>(AggregateStateKey).ConfigureAwait(false);
this._runningAggregate = await context.ReadStateAsync<TAggregate>(AggregateStateKey, cancellationToken: cancellationToken).ConfigureAwait(false);
}
}
@@ -10,7 +10,7 @@ internal sealed class AsyncBarrier()
{
private readonly InitLocked<TaskCompletionSource<object>> _completionSource = new();
public async ValueTask<bool> JoinAsync(CancellationToken cancellation = default)
public async ValueTask<bool> JoinAsync(CancellationToken cancellationToken = default)
{
this._completionSource.Init(() => new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously));
TaskCompletionSource<object> completionSource = this._completionSource.Get()!;
@@ -19,10 +19,10 @@ internal sealed class AsyncBarrier()
// should not cancel the entire barrier.
TaskCompletionSource<object> 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()
@@ -12,12 +12,13 @@ internal sealed class AsyncCoordinator
/// <summary>
/// Wait for the Coordination owner to mark the next coordination point, then continue execution.
/// </summary>
/// <param name="cancellation">A cancellation token that can be used to cancel the wait.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>
/// A task that represents the asynchronous operation. The task result is <see langword="true"/>
/// if the wait was completed; otherwise, for example, if the wait was cancelled, <see langword="false"/>.
/// </returns>
public async ValueTask<bool> WaitForCoordinationAsync(CancellationToken cancellation = default)
public async ValueTask<bool> 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);
}
/// <summary>
@@ -30,19 +30,19 @@ internal abstract class ChatProtocolExecutor(string id, ChatProtocolExecutorOpti
{
if (this._stringMessageChatRole.HasValue)
{
routeBuilder = routeBuilder.AddHandler<string>((message, _) => this._pendingMessages.Add(new(this._stringMessageChatRole.Value, message)));
routeBuilder = routeBuilder.AddHandler<string>((message, _, __) => this._pendingMessages.Add(new(this._stringMessageChatRole.Value, message)));
}
return routeBuilder.AddHandler<ChatMessage>((message, _) => this._pendingMessages.Add(message))
.AddHandler<List<ChatMessage>>((messages, _) => this._pendingMessages.AddRange(messages))
return routeBuilder.AddHandler<ChatMessage>((message, _, __) => this._pendingMessages.Add(message))
.AddHandler<List<ChatMessage>>((messages, _, __) => this._pendingMessages.AddRange(messages))
.AddHandler<TurnToken>(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<ChatMessage> 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<JsonElement?>(PendingMessagesStateKey).ConfigureAwait(false);
JsonElement? messagesValue = await context.ReadStateAsync<JsonElement?>(PendingMessagesStateKey, cancellationToken: cancellationToken).ConfigureAwait(false);
if (messagesValue.HasValue)
{
List<ChatMessage> messages = messagesValue.Value.DeserializeMessages();
@@ -44,19 +44,14 @@ internal sealed class AsyncRunHandle : ICheckpointingHandle, IAsyncDisposable
}
}
//private readonly AsyncCoordinator _waitForResponseCoordinator = new();
//public ValueTask<bool> WaitForNextInputAsync(CancellationToken cancellation = default)
// => this._waitForResponseCoordinator.WaitForCoordinationAsync(cancellation);
public string RunId => this._stepRunner.RunId;
public IReadOnlyList<CheckpointInfo> Checkpoints => this._checkpointingHandle.Checkpoints;
public ValueTask<RunStatus> GetStatusAsync(CancellationToken cancellation = default)
=> this._eventStream.GetStatusAsync(cancellation);
public ValueTask<RunStatus> GetStatusAsync(CancellationToken cancellationToken = default)
=> this._eventStream.GetStatusAsync(cancellationToken);
public async IAsyncEnumerable<WorkflowEvent> TakeEventStreamAsync(bool blockOnPendingRequest, [EnumeratorCancellation] CancellationToken cancellation = default)
public async IAsyncEnumerable<WorkflowEvent> 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<bool> IsValidInputTypeAsync<T>(CancellationToken cancellation = default)
=> this._stepRunner.IsValidInputTypeAsync<T>(cancellation);
public ValueTask<bool> IsValidInputTypeAsync<T>(CancellationToken cancellationToken = default)
=> this._stepRunner.IsValidInputTypeAsync<T>(cancellationToken);
public async ValueTask<bool> EnqueueMessageAsync<T>(T message, CancellationToken cancellation = default)
public async ValueTask<bool> EnqueueMessageAsync<T>(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<bool> EnqueueMessageUntypedAsync([NotNull] object message, Type? declaredType = null, CancellationToken cancellation = default)
public async ValueTask<bool> 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();
@@ -14,33 +14,33 @@ internal static class AsyncRunHandleExtensions
return new Checkpointed<TRunType>(run, runHandle);
}
public static async ValueTask<StreamingRun> EnqueueAndStreamAsync<TInput>(this AsyncRunHandle runHandle, TInput input, CancellationToken cancellation = default)
public static async ValueTask<StreamingRun> EnqueueAndStreamAsync<TInput>(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<StreamingRun> EnqueueUntypedAndStreamAsync(this AsyncRunHandle runHandle, object input, CancellationToken cancellation = default)
public static async ValueTask<StreamingRun> 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<Run> EnqueueAndRunAsync<TInput>(this AsyncRunHandle runHandle, TInput input, CancellationToken cancellation = default)
public static async ValueTask<Run> EnqueueAndRunAsync<TInput>(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<Run> EnqueueUntypedAndRunAsync(this AsyncRunHandle runHandle, object input, CancellationToken cancellation = default)
public static async ValueTask<Run> 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;
}
}
@@ -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
/// </summary>
public Exception? Exception { get; init; }
/// <summary>
/// Indicated whether the call was cancelled (e.g., via a <see cref="CancellationToken"/>).
/// </summary>
public bool IsCancelled { get; init; }
/// <summary>
/// Indicates whether the call was successful. A call is considered successful if it returned
/// without throwing an exception.
/// </summary>
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;
}
/// <summary>
@@ -51,6 +58,14 @@ internal sealed class CallResult
/// <returns>A <see cref="CallResult"/> indicating the result of the call.</returns>
public static CallResult ReturnVoid() => new(isVoid: true);
/// <summary>
/// Create a <see cref="CallResult"/> indicating that the call was cancelled.
/// </summary>
/// <param name="wasVoid">A boolean specifying whether the call was void (was not expected to return
/// a value).</param>
/// <returns>A <see cref="CallResult"/> indicating the result of the call.</returns>
public static CallResult Cancelled(bool wasVoid) => new(wasVoid, isCancelled: true);
/// <summary>
/// Create a <see cref="CallResult"/> indicating that an exception was raised during the call.
/// </summary>
@@ -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<bool> WaitForNextInputAsync(CancellationToken cancellation = default);
}
@@ -15,7 +15,7 @@ internal interface IRunEventStream : IAsyncDisposable
// this cannot be cancelled
ValueTask StopAsync();
ValueTask<RunStatus> GetStatusAsync(CancellationToken cancellation = default);
ValueTask<RunStatus> GetStatusAsync(CancellationToken cancellationToken = default);
IAsyncEnumerable<WorkflowEvent> TakeEventStreamAsync(bool blockOnPendingRequest, CancellationToken cancellation = default);
IAsyncEnumerable<WorkflowEvent> TakeEventStreamAsync(bool blockOnPendingRequest, CancellationToken cancellationToken = default);
}
@@ -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<StepContext> AdvanceAsync();
ValueTask<StepContext> AdvanceAsync(CancellationToken cancellationToken = default);
IWorkflowContext Bind(string executorId, Dictionary<string, string>? traceContext = null);
ValueTask<Executor> EnsureExecutorAsync(string executorId, IStepTracer? tracer);
ValueTask<Executor> EnsureExecutorAsync(string executorId, IStepTracer? tracer, CancellationToken cancellationToken = default);
}
@@ -10,8 +10,8 @@ internal interface ISuperStepJoinContext
{
bool WithCheckpointing { get; }
ValueTask ForwardWorkflowEventAsync(WorkflowEvent workflowEvent, CancellationToken cancellation = default);
ValueTask SendMessageAsync<TMessage>(string senderId, [DisallowNull] TMessage message, CancellationToken cancellation = default);
ValueTask ForwardWorkflowEventAsync(WorkflowEvent workflowEvent, CancellationToken cancellationToken = default);
ValueTask SendMessageAsync<TMessage>(string senderId, [DisallowNull] TMessage message, CancellationToken cancellationToken = default);
ValueTask AttachSuperstepAsync(ISuperStepRunner superStepRunner, CancellationToken cancellation = default);
ValueTask AttachSuperstepAsync(ISuperStepRunner superStepRunner, CancellationToken cancellationToken = default);
}
@@ -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<bool> IsValidInputTypeAsync<T>(CancellationToken cancellation = default);
ValueTask<bool> EnqueueMessageAsync<T>(T message, CancellationToken cancellation = default);
ValueTask<bool> EnqueueMessageUntypedAsync(object message, Type declaredType, CancellationToken cancellation = default);
ValueTask<bool> IsValidInputTypeAsync<T>(CancellationToken cancellationToken = default);
ValueTask<bool> EnqueueMessageAsync<T>(T message, CancellationToken cancellationToken = default);
ValueTask<bool> EnqueueMessageUntypedAsync(object message, Type declaredType, CancellationToken cancellationToken = default);
ConcurrentEventSink OutgoingEvents { get; }
@@ -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);
}
}
@@ -22,7 +22,7 @@ internal sealed class LockstepRunEventStream : IRunEventStream
private readonly ISuperStepRunner _stepRunner;
public ValueTask<RunStatus> GetStatusAsync(CancellationToken cancellation = default) => new(this.RunStatus);
public ValueTask<RunStatus> 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<WorkflowEvent> TakeEventStreamAsync(bool blockOnPendingRequest, [EnumeratorCancellation] CancellationToken cancellation = default)
public async IAsyncEnumerable<WorkflowEvent> 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<WorkflowEvent> eventSink = [];
@@ -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<Microsoft.Agents.AI.Workflows.Execution.CallResult>
>;
using MessageHandlerF =
System.Func<
object, // message
Microsoft.Agents.AI.Workflows.IWorkflowContext, // context
System.Threading.CancellationToken, // cancellation
System.Threading.Tasks.ValueTask<Microsoft.Agents.AI.Workflows.Execution.CallResult>
>;
@@ -56,7 +59,7 @@ internal sealed class MessageRouter
public HashSet<Type> DefaultOutputTypes { get; }
public async ValueTask<CallResult?> RouteMessageAsync(object message, IWorkflowContext context, bool requireRoute = false)
public async ValueTask<CallResult?> 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)
@@ -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<WorkflowEvent> 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<RunStatus> GetStatusAsync(CancellationToken cancellation = default)
public ValueTask<RunStatus> GetStatusAsync(CancellationToken cancellationToken = default)
{
// Thread-safe read of status (enum is read atomically on most platforms)
return new ValueTask<RunStatus>(this._runStatus);
@@ -90,10 +90,12 @@ public abstract class Executor : IIdentified
/// <param name="messageType">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.</param>
/// <param name="context">The workflow context in which the executor executes.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>A ValueTask representing the asynchronous operation, wrapping the output from the executor.</returns>
/// <exception cref="NotSupportedException">No handler found for the message type.</exception>
/// <exception cref="TargetInvocationException">An exception is generated while handling the message.</exception>
public async ValueTask<object?> ExecuteAsync(object message, TypeId messageType, IWorkflowContext context)
public async ValueTask<object?> 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
/// </summary>
/// <param name="context">The workflow context.</param>
/// <returns>A ValueTask representing the asynchronous operation.</returns>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
protected internal virtual ValueTask OnCheckpointingAsync(IWorkflowContext context, CancellationToken cancellationToken = default) => default;
/// <summary>
@@ -160,7 +163,8 @@ public abstract class Executor : IIdentified
/// </summary>
/// <param name="context">The workflow context.</param>
/// <returns>A ValueTask representing the asynchronous operation.</returns>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
protected internal virtual ValueTask OnCheckpointRestoredAsync(IWorkflowContext context, CancellationToken cancellationToken = default) => default;
/// <summary>
@@ -210,7 +214,7 @@ public abstract class Executor<TInput>(string id, ExecutorOptions? options = nul
routeBuilder.AddHandler<TInput>(this.HandleAsync);
/// <inheritdoc/>
public abstract ValueTask HandleAsync(TInput message, IWorkflowContext context);
public abstract ValueTask HandleAsync(TInput message, IWorkflowContext context, CancellationToken cancellationToken = default);
}
/// <summary>
@@ -229,5 +233,5 @@ public abstract class Executor<TInput, TOutput>(string id, ExecutorOptions? opti
routeBuilder.AddHandler<TInput, TOutput>(this.HandleAsync);
/// <inheritdoc/>
public abstract ValueTask<TOutput> HandleAsync(TInput message, IWorkflowContext context);
public abstract ValueTask<TOutput> HandleAsync(TInput message, IWorkflowContext context, CancellationToken cancellationToken = default);
}
@@ -29,7 +29,7 @@ public class FunctionExecutor<TInput>(string id,
}
/// <inheritdoc/>
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);
/// <summary>
/// Creates a new instance of the <see cref="FunctionExecutor{TInput}"/> class.
@@ -65,7 +65,7 @@ public class FunctionExecutor<TInput, TOutput>(string id,
}
/// <inheritdoc/>
public override ValueTask<TOutput> HandleAsync(TInput message, IWorkflowContext context) => handlerAsync(message, context, default);
public override ValueTask<TOutput> HandleAsync(TInput message, IWorkflowContext context, CancellationToken cancellationToken) => handlerAsync(message, context, cancellationToken);
/// <summary>
/// Creates a new instance of the <see cref="FunctionExecutor{TInput,TOutput}"/> class.
@@ -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.
/// </summary>
/// <param name="workflowEvent">The event to be raised.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>A <see cref="ValueTask"/> representing the asynchronous operation.</returns>
ValueTask AddEventAsync(WorkflowEvent workflowEvent);
ValueTask AddEventAsync(WorkflowEvent workflowEvent, CancellationToken cancellationToken = default);
/// <summary>
/// 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
/// <param name="targetId">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.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>A <see cref="ValueTask"/> representing the asynchronous operation.</returns>
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?)
/// <summary>
/// Queues a message to be sent to connected executors. The message will be sent during the next SuperStep.
/// </summary>
/// <param name="message">The message to be sent.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.</param>
/// <returns>A <see cref="ValueTask"/> representing the asynchronous operation.</returns>
ValueTask SendMessageAsync(object message, CancellationToken cancellationToken) => this.SendMessageAsync(message, null, cancellationToken);
#endif
/// <summary>
/// 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 <see cref="ExecutorOptions"/>.
/// </remarks>
/// <param name="output">The output value to be returned.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>A <see cref="ValueTask"/> representing the asynchronous operation.</returns>
ValueTask YieldOutputAsync(object output);
ValueTask YieldOutputAsync(object output, CancellationToken cancellationToken = default);
/// <summary>
/// Adds a request to "halt" workflow execution at the end of the current SuperStep.
@@ -54,15 +73,32 @@ public interface IWorkflowContext
/// <param name="key">The key of the state value.</param>
/// <param name = "scopeName" > An optional name that specifies the scope to read.If null, the default scope is
/// used.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>A <see cref="ValueTask{T}"/> representing the asynchronous operation.</returns>
ValueTask<T?> ReadStateAsync<T>(string key, string? scopeName = null);
ValueTask<T?> ReadStateAsync<T>(string key, string? scopeName = null, CancellationToken cancellationToken = default);
#if NET // See above for musings about this construction
/// <summary>
/// Reads a state value from the workflow's state store. If no scope is provided, the executor's
/// default scope is used.
/// </summary>
/// <typeparam name="T">The type of the state value.</typeparam>
/// <param name="key">The key of the state value.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.</param>
/// <returns>A <see cref="ValueTask{T}"/> representing the asynchronous operation.</returns>
ValueTask<T?> ReadStateAsync<T>(string key, CancellationToken cancellationToken) => this.ReadStateAsync<T>(key, null, cancellationToken);
#endif
/// <summary>
/// Asynchronously reads all state keys within the specified scope.
/// </summary>
/// <param name="scopeName">An optional name that specifies the scope to read. If null, the default scope is
/// used.</param>
ValueTask<HashSet<string>> ReadStateKeysAsync(string? scopeName = null);
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
ValueTask<HashSet<string>> ReadStateKeysAsync(string? scopeName = null, CancellationToken cancellationToken = default);
/// <summary>
/// Asynchronously updates the state of a queue entry identified by the specified key and optional scope.
@@ -77,8 +113,27 @@ public interface IWorkflowContext
/// implementation.</param>
/// <param name="scopeName">An optional name that specifies the scope to update. If null, the default scope is
/// used.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>A ValueTask that represents the asynchronous update operation.</returns>
ValueTask QueueStateUpdateAsync<T>(string key, T? value, string? scopeName = null);
ValueTask QueueStateUpdateAsync<T>(string key, T? value, string? scopeName = null, CancellationToken cancellationToken = default);
#if NET // See above for musings about this construction
/// <summary>
/// Asynchronously updates the state of a queue entry identified by the specified key and optional scope.
/// </summary>
/// <remarks>
/// 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.
/// </remarks>
/// <typeparam name="T">The type of the value to associate with the queue entry.</typeparam>
/// <param name="key">The unique identifier for the queue entry to update. Cannot be null or empty.</param>
/// <param name="value">The value to set for the queue entry. If null, the entry's state may be cleared or reset depending on
/// implementation.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.</param>
/// <returns>A ValueTask that represents the asynchronous update operation.</returns>
ValueTask QueueStateUpdateAsync<T>(string key, T? value, CancellationToken cancellationToken) => this.QueueStateUpdateAsync(key, value, null, cancellationToken);
#endif
/// <summary>
/// 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.
/// </remarks>
/// <param name="scopeName">An optional name that specifies the scope to clear. If null, the default scope is used.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>A ValueTask that represents the asynchronous clear operation.</returns>
ValueTask QueueClearScopeAsync(string? scopeName = null);
ValueTask QueueClearScopeAsync(string? scopeName = null, CancellationToken cancellationToken = default);
#if NET // See above for musings about this construction
/// <summary>
/// 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.
/// </summary>
/// <remarks>
/// 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.
/// </remarks>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.</param>
/// <returns>A ValueTask that represents the asynchronous clear operation.</returns>
ValueTask QueueClearScopeAsync(CancellationToken cancellationToken) => this.QueueClearScopeAsync(null, cancellationToken);
#endif
/// <summary>
/// The trace context associated with the current message about to be processed by the executor, if any.
@@ -46,14 +46,14 @@ internal sealed class InProcessRunner : ISuperStepRunner, ICheckpointingHandle
public string StartExecutorId { get; }
private readonly HashSet<Type> _knownValidInputTypes;
public async ValueTask<bool> IsValidInputTypeAsync(Type messageType, CancellationToken cancellation = default)
public async ValueTask<bool> 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<bool> IsValidInputTypeAsync<T>(CancellationToken cancellation = default)
=> this.IsValidInputTypeAsync(typeof(T), cancellation);
public ValueTask<bool> IsValidInputTypeAsync<T>(CancellationToken cancellationToken = default)
=> this.IsValidInputTypeAsync(typeof(T), cancellationToken);
public async ValueTask<bool> EnqueueMessageUntypedAsync(object message, Type declaredType, CancellationToken cancellation = default)
public async ValueTask<bool> 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<bool> EnqueueMessageAsync<T>(T message, CancellationToken cancellation = default)
=> this.EnqueueMessageUntypedAsync(Throw.IfNull(message), typeof(T), cancellation);
public ValueTask<bool> EnqueueMessageAsync<T>(T message, CancellationToken cancellationToken = default)
=> this.EnqueueMessageUntypedAsync(Throw.IfNull(message), typeof(T), cancellationToken);
public ValueTask<bool> EnqueueMessageAsync(object message, CancellationToken cancellation = default)
=> this.EnqueueMessageUntypedAsync(Throw.IfNull(message), message.GetType(), cancellation);
public ValueTask<bool> 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<AsyncRunHandle> BeginStreamAsync(ExecutionMode mode, CancellationToken cancellation = default)
public ValueTask<AsyncRunHandle> BeginStreamAsync(ExecutionMode mode, CancellationToken cancellationToken = default)
{
this.RunContext.CheckEnded();
return new(new AsyncRunHandle(this, this, mode));
}
public async ValueTask<AsyncRunHandle> ResumeStreamAsync(ExecutionMode mode, CheckpointInfo fromCheckpoint, CancellationToken cancellation = default)
public async ValueTask<AsyncRunHandle> 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<MessageEnvelope> envelopes)
private async ValueTask DeliverMessagesAsync(string receiverId, ConcurrentQueue<MessageEnvelope> 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<Task> 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<Task> 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);
@@ -57,7 +57,7 @@ internal sealed class InProcessRunnerContext : IRunnerContext
this.OutgoingEvents = outgoingEvents;
}
public async ValueTask<Executor> EnsureExecutorAsync(string executorId, IStepTracer? tracer)
public async ValueTask<Executor> EnsureExecutorAsync(string executorId, IStepTracer? tracer, CancellationToken cancellationToken = default)
{
this.CheckEnded();
Task<Executor> executorTask = this._executors.GetOrAdd(executorId, CreateExecutorAsync);
@@ -88,9 +88,9 @@ internal sealed class InProcessRunnerContext : IRunnerContext
return await executorTask.ConfigureAwait(false);
}
public async ValueTask<IEnumerable<Type>> GetStartingExecutorInputTypesAsync(CancellationToken cancellation = default)
public async ValueTask<IEnumerable<Type>> 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<StepContext> AdvanceAsync()
public async ValueTask<StepContext> 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<string, string>? 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<T?> ReadStateAsync<T>(string key, string? scopeName = null)
public ValueTask<T?> ReadStateAsync<T>(string key, string? scopeName = null, CancellationToken cancellationToken = default)
=> RunnerContext.StateManager.ReadStateAsync<T>(ExecutorId, scopeName, key);
public ValueTask<HashSet<string>> ReadStateKeysAsync(string? scopeName = null)
public ValueTask<HashSet<string>> ReadStateKeysAsync(string? scopeName = null, CancellationToken cancellationToken = default)
=> RunnerContext.StateManager.ReadKeysAsync(ExecutorId, scopeName);
public ValueTask QueueStateUpdateAsync<T>(string key, T? value, string? scopeName = null)
public ValueTask QueueStateUpdateAsync<T>(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<string, string>? 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<Executor> 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<ISuperStepRunner> 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<TMessage>(string senderId, [DisallowNull] TMessage message, CancellationToken cancellation)
=> this.SendMessageAsync(senderId, Throw.IfNull(message));
ValueTask ISuperStepJoinContext.SendMessageAsync<TMessage>(string senderId, [DisallowNull] TMessage message, CancellationToken cancellationToken)
=> this.SendMessageAsync(senderId, Throw.IfNull(message), cancellationToken: cancellationToken);
}
@@ -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<TMessage>
/// </summary>
/// <param name="message">The message to handle.</param>
/// <param name="context">The execution context.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>A task that represents the asynchronous operation.</returns>
ValueTask HandleAsync(TMessage message, IWorkflowContext context);
ValueTask HandleAsync(TMessage message, IWorkflowContext context, CancellationToken cancellationToken = default);
}
/// <summary>
@@ -32,6 +35,8 @@ public interface IMessageHandler<TMessage, TResult>
/// </summary>
/// <param name="message">The message to handle.</param>
/// <param name="context">The execution context.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.
/// The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>A task that represents the asynchronous operation.</returns>
ValueTask<TResult> HandleAsync(TMessage message, IWorkflowContext context);
ValueTask<TResult> HandleAsync(TMessage message, IWorkflowContext context, CancellationToken cancellationToken = default);
}
@@ -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<object, IWorkflowContext, ValueTask<CallResult>> Bind(Func<object, IWorkflowContext, object?> handlerAsync, bool checkType, Type? resultType = null, Func<object, ValueTask<object?>>? unwrapper = null)
public static Func<object, IWorkflowContext, CancellationToken, ValueTask<CallResult>> Bind(Func<object, IWorkflowContext, CancellationToken, object?> handlerAsync, bool checkType, Type? resultType = null, Func<object, ValueTask<object?>>? unwrapper = null)
{
return InvokeHandlerAsync;
async ValueTask<CallResult> InvokeHandlerAsync(object message, IWorkflowContext workflowContext)
async ValueTask<CallResult> 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<object, IWorkflowContext, ValueTask<CallResult>> Bind<
public Func<object, IWorkflowContext, CancellationToken, ValueTask<CallResult>> 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]);
}
}
}
@@ -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<Microsoft.Agents.AI.Workflows.Execution.CallResult>
>;
using MessageHandlerF =
System.Func<
object, // message
Microsoft.Agents.AI.Workflows.IWorkflowContext, // context
System.Threading.CancellationToken, // cancellation
System.Threading.Tasks.ValueTask<Microsoft.Agents.AI.Workflows.Execution.CallResult>
>;
@@ -73,32 +76,58 @@ public class RouteBuilder
return this;
}
internal RouteBuilder AddHandlerUntyped(Type type, Func<object, IWorkflowContext, ValueTask> handler, bool overwrite = false)
internal RouteBuilder AddHandlerUntyped(Type type, Func<object, IWorkflowContext, CancellationToken, ValueTask> handler, bool overwrite = false)
{
Throw.IfNull(handler);
return this.AddHandlerInternal(type, WrappedHandlerAsync, outputType: null, overwrite);
async ValueTask<CallResult> WrappedHandlerAsync(object msg, IWorkflowContext ctx)
async ValueTask<CallResult> 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<TResult>(Type type, Func<object, IWorkflowContext, ValueTask<TResult>> handler, bool overwrite = false)
internal RouteBuilder AddHandlerUntyped<TResult>(Type type, Func<object, IWorkflowContext, CancellationToken, ValueTask<TResult>> handler, bool overwrite = false)
{
Throw.IfNull(handler);
return this.AddHandlerInternal(type, WrappedHandlerAsync, outputType: typeof(TResult), overwrite);
async ValueTask<CallResult> WrappedHandlerAsync(object msg, IWorkflowContext ctx)
async ValueTask<CallResult> 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);
}
}
/// <summary>
/// Registers a handler for messages of the specified input type in the workflow route.
/// </summary>
/// <remarks>If a handler for the specified input type already exists and <paramref name="overwrite"/> is
/// <see langword="false"/>, the existing handler will not be replaced. Handlers are invoked asynchronously and are
/// expected to complete their processing before the workflow continues.</remarks>
/// <typeparam name="TInput"></typeparam>
/// <param name="handler">A delegate that processes messages of type <typeparamref name="TInput"/> within the workflow context. The
/// delegate is invoked for each incoming message of the specified type.</param>
/// <param name="overwrite"><see langword="true"/> to replace any existing handler for the specified input type; otherwise, <see
/// langword="false"/> to preserve the existing handler.</param>
/// <returns>The current <see cref="RouteBuilder"/> instance, enabling fluent configuration of additional handlers or route
/// options.</returns>
public RouteBuilder AddHandler<TInput>(Action<TInput, IWorkflowContext, CancellationToken> handler, bool overwrite = false)
{
Throw.IfNull(handler);
return this.AddHandlerInternal(typeof(TInput), WrappedHandlerAsync, outputType: null, overwrite);
async ValueTask<CallResult> WrappedHandlerAsync(object message, IWorkflowContext context, CancellationToken cancellationToken)
{
handler.Invoke((TInput)message, context, cancellationToken);
return CallResult.ReturnVoid();
}
}
/// <summary>
/// Registers a handler for messages of the specified input type in the workflow route.
/// </summary>
@@ -118,9 +147,35 @@ public class RouteBuilder
return this.AddHandlerInternal(typeof(TInput), WrappedHandlerAsync, outputType: null, overwrite);
async ValueTask<CallResult> WrappedHandlerAsync(object msg, IWorkflowContext ctx)
async ValueTask<CallResult> WrappedHandlerAsync(object message, IWorkflowContext context, CancellationToken cancellationToken)
{
handler.Invoke((TInput)msg, ctx);
handler.Invoke((TInput)message, context);
return CallResult.ReturnVoid();
}
}
/// <summary>
/// Registers a handler for messages of the specified input type in the workflow route.
/// </summary>
/// <remarks>If a handler for the specified input type already exists and <paramref name="overwrite"/> is
/// <see langword="false"/>, the existing handler will not be replaced. Handlers are invoked asynchronously and are
/// expected to complete their processing before the workflow continues.</remarks>
/// <typeparam name="TInput"></typeparam>
/// <param name="handler">A delegate that processes messages of type <typeparamref name="TInput"/> within the workflow context. The
/// delegate is invoked for each incoming message of the specified type.</param>
/// <param name="overwrite"><see langword="true"/> to replace any existing handler for the specified input type; otherwise, <see
/// langword="false"/> to preserve the existing handler.</param>
/// <returns>The current <see cref="RouteBuilder"/> instance, enabling fluent configuration of additional handlers or route
/// options.</returns>
public RouteBuilder AddHandler<TInput>(Func<TInput, IWorkflowContext, CancellationToken, ValueTask> handler, bool overwrite = false)
{
Throw.IfNull(handler);
return this.AddHandlerInternal(typeof(TInput), WrappedHandlerAsync, outputType: null, overwrite);
async ValueTask<CallResult> 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<CallResult> WrappedHandlerAsync(object msg, IWorkflowContext ctx)
async ValueTask<CallResult> 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();
}
}
/// <summary>
/// Registers a handler function for messages of the specified input type in the workflow route.
/// </summary>
/// <remarks>If a handler for the given input type already exists, setting <paramref name="overwrite"/> to
/// <see langword="true"/> 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.</remarks>
/// <typeparam name="TInput">The type of input message the handler will process.</typeparam>
/// <typeparam name="TResult">The type of result produced by the handler.</typeparam>
/// <param name="handler">A function that processes messages of type <typeparamref name="TInput"/> within the workflow context and returns
/// a <see cref="ValueTask{TResult}"/> representing the asynchronous result.</param>
/// <param name="overwrite"><see langword="true"/> to replace any existing handler for the input type; otherwise, <see langword="false"/> to
/// preserve existing handlers.</param>
/// <returns>The current <see cref="RouteBuilder"/> instance, enabling fluent configuration of workflow routes.</returns>
public RouteBuilder AddHandler<TInput, TResult>(Func<TInput, IWorkflowContext, CancellationToken, TResult> handler, bool overwrite = false)
{
Throw.IfNull(handler);
return this.AddHandlerInternal(typeof(TInput), WrappedHandlerAsync, outputType: typeof(TResult), overwrite);
async ValueTask<CallResult> WrappedHandlerAsync(object message, IWorkflowContext context, CancellationToken cancellationToken)
{
TResult result = handler.Invoke((TInput)message, context, cancellationToken);
return CallResult.ReturnResult(result);
}
}
/// <summary>
/// Registers a handler function for messages of the specified input type in the workflow route.
/// </summary>
@@ -170,9 +251,35 @@ public class RouteBuilder
return this.AddHandlerInternal(typeof(TInput), WrappedHandlerAsync, outputType: typeof(TResult), overwrite);
async ValueTask<CallResult> WrappedHandlerAsync(object msg, IWorkflowContext ctx)
async ValueTask<CallResult> 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);
}
}
/// <summary>
/// Registers a handler function for messages of the specified input type in the workflow route.
/// </summary>
/// <remarks>If a handler for the given input type already exists, setting <paramref name="overwrite"/> to
/// <see langword="true"/> 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.</remarks>
/// <typeparam name="TInput">The type of input message the handler will process.</typeparam>
/// <typeparam name="TResult">The type of result produced by the handler.</typeparam>
/// <param name="handler">A function that processes messages of type <typeparamref name="TInput"/> within the workflow context and returns
/// a <see cref="ValueTask{TResult}"/> representing the asynchronous result.</param>
/// <param name="overwrite"><see langword="true"/> to replace any existing handler for the input type; otherwise, <see langword="false"/> to
/// preserve existing handlers.</param>
/// <returns>The current <see cref="RouteBuilder"/> instance, enabling fluent configuration of workflow routes.</returns>
public RouteBuilder AddHandler<TInput, TResult>(Func<TInput, IWorkflowContext, CancellationToken, ValueTask<TResult>> handler, bool overwrite = false)
{
Throw.IfNull(handler);
return this.AddHandlerInternal(typeof(TInput), WrappedHandlerAsync, outputType: typeof(TResult), overwrite);
async ValueTask<CallResult> 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<CallResult> WrappedHandlerAsync(object msg, IWorkflowContext ctx)
async ValueTask<CallResult> 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;
}
/// <summary>
/// Register a handler function as a catch-all handler: It will be used if not type-matching handler is registered.
/// </summary>
/// <remarks>If a catch-all handler for already exists, setting <paramref name="overwrite"/> to <see langword="true"/>
/// will replace the existing handler; otherwise, an exception may be thrown. The handler receives the input message
/// wrapped as <see cref="PortableValue"/> and workflow context, and returns a result asynchronously.</remarks>
/// <param name="handler">A function that processes messages wrapped as <see cref="PortableValue"/> within the
/// workflow context. The delegate is invoked for each incoming message not otherwise handled.</param>
/// <param name="overwrite"><see langword="true"/> to replace any existing handler for the input type; otherwise, <see langword="false"/> to
/// preserve existing handlers.</param>
/// <returns>The current <see cref="RouteBuilder"/> instance, enabling fluent configuration of workflow routes.</returns>
public RouteBuilder AddCatchAll(Func<PortableValue, IWorkflowContext, CancellationToken, ValueTask> handler, bool overwrite = false)
{
Throw.IfNull(handler);
return this.AddCatchAll(WrappedHandlerAsync, overwrite);
async ValueTask<CallResult> WrappedHandlerAsync(PortableValue message, IWorkflowContext context, CancellationToken cancellationToken)
{
await handler.Invoke(message, context, cancellationToken).ConfigureAwait(false);
return CallResult.ReturnVoid();
}
}
/// <summary>
/// Register a handler function as a catch-all handler: It will be used if not type-matching handler is registered.
/// </summary>
@@ -232,13 +363,37 @@ public class RouteBuilder
return this.AddCatchAll(WrappedHandlerAsync, overwrite);
async ValueTask<CallResult> WrappedHandlerAsync(PortableValue message, IWorkflowContext ctx)
async ValueTask<CallResult> WrappedHandlerAsync(PortableValue message, IWorkflowContext context, CancellationToken cancellationToken)
{
await handler.Invoke(message, ctx).ConfigureAwait(false);
await handler.Invoke(message, context).ConfigureAwait(false);
return CallResult.ReturnVoid();
}
}
/// <summary>
/// Register a handler function as a catch-all handler: It will be used if not type-matching handler is registered.
/// </summary>
/// <remarks>If a catch-all handler for already exists, setting <paramref name="overwrite"/> to <see langword="true"/>
/// will replace the existing handler; otherwise, an exception may be thrown. The handler receives the input message
/// wrapped as <see cref="PortableValue"/> and workflow context, and returns a result asynchronously.</remarks>
/// <param name="handler">A function that processes messages wrapped as <see cref="PortableValue"/> within the
/// workflow context and returns a <see cref="ValueTask{TResult}"/> representing the asynchronous result.</param>
/// <param name="overwrite"><see langword="true"/> to replace any existing handler for the input type; otherwise, <see langword="false"/> to
/// preserve existing handlers.</param>
/// <returns>The current <see cref="RouteBuilder"/> instance, enabling fluent configuration of workflow routes.</returns>
public RouteBuilder AddCatchAll<TResult>(Func<PortableValue, IWorkflowContext, CancellationToken, ValueTask<TResult>> handler, bool overwrite = false)
{
Throw.IfNull(handler);
return this.AddCatchAll(WrappedHandlerAsync, overwrite);
async ValueTask<CallResult> WrappedHandlerAsync(PortableValue message, IWorkflowContext context, CancellationToken cancellationToken)
{
TResult result = await handler.Invoke(message, context, cancellationToken).ConfigureAwait(false);
return CallResult.ReturnResult(result);
}
}
/// <summary>
/// Register a handler function as a catch-all handler: It will be used if not type-matching handler is registered.
/// </summary>
@@ -256,13 +411,37 @@ public class RouteBuilder
return this.AddCatchAll(WrappedHandlerAsync, overwrite);
async ValueTask<CallResult> WrappedHandlerAsync(PortableValue message, IWorkflowContext ctx)
async ValueTask<CallResult> 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);
}
}
/// <summary>
/// Register a handler function as a catch-all handler: It will be used if not type-matching handler is registered.
/// </summary>
/// <remarks>If a catch-all handler for already exists, setting <paramref name="overwrite"/> to <see langword="true"/>
/// will replace the existing handler; otherwise, an exception may be thrown. The handler receives the input message
/// wrapped as <see cref="PortableValue"/> and workflow context, and returns a result asynchronously.</remarks>
/// <param name="handler">A function that processes messages wrapped as <see cref="PortableValue"/> within the
/// workflow context. The delegate is invoked for each incoming message not otherwise handled.</param>
/// <param name="overwrite"><see langword="true"/> to replace any existing handler for the input type; otherwise, <see langword="false"/> to
/// preserve existing handlers.</param>
/// <returns>The current <see cref="RouteBuilder"/> instance, enabling fluent configuration of workflow routes.</returns>
public RouteBuilder AddCatchAll(Action<PortableValue, IWorkflowContext, CancellationToken> handler, bool overwrite = false)
{
Throw.IfNull(handler);
return this.AddCatchAll(WrappedHandlerAsync, overwrite);
ValueTask<CallResult> WrappedHandlerAsync(PortableValue message, IWorkflowContext ctx, CancellationToken cancellationToken)
{
handler.Invoke(message, ctx, cancellationToken);
return new(CallResult.ReturnVoid());
}
}
/// <summary>
/// Register a handler function as a catch-all handler: It will be used if not type-matching handler is registered.
/// </summary>
@@ -280,13 +459,37 @@ public class RouteBuilder
return this.AddCatchAll(WrappedHandlerAsync, overwrite);
ValueTask<CallResult> WrappedHandlerAsync(PortableValue message, IWorkflowContext ctx)
ValueTask<CallResult> WrappedHandlerAsync(PortableValue message, IWorkflowContext ctx, CancellationToken cancellationToken)
{
handler.Invoke(message, ctx);
return new(CallResult.ReturnVoid());
}
}
/// <summary>
/// Register a handler function as a catch-all handler: It will be used if not type-matching handler is registered.
/// </summary>
/// <remarks>If a catch-all handler for already exists, setting <paramref name="overwrite"/> to <see langword="true"/>
/// will replace the existing handler; otherwise, an exception may be thrown. The handler receives the input message
/// wrapped as <see cref="PortableValue"/> and workflow context, and returns a result asynchronously.</remarks>
/// <param name="handler">A function that processes messages wrapped as <see cref="PortableValue"/> within the
/// workflow context and returns a <see cref="ValueTask{TResult}"/> representing the asynchronous result.</param>
/// <param name="overwrite"><see langword="true"/> to replace any existing handler for the input type; otherwise, <see langword="false"/> to
/// preserve existing handlers.</param>
/// <returns>The current <see cref="RouteBuilder"/> instance, enabling fluent configuration of workflow routes.</returns>
public RouteBuilder AddCatchAll<TResult>(Func<PortableValue, IWorkflowContext, CancellationToken, TResult> handler, bool overwrite = false)
{
Throw.IfNull(handler);
return this.AddCatchAll(WrappedHandlerAsync, overwrite);
ValueTask<CallResult> WrappedHandlerAsync(PortableValue message, IWorkflowContext context, CancellationToken cancellationToken)
{
TResult result = handler.Invoke(message, context, cancellationToken);
return new(CallResult.ReturnResult(result));
}
}
/// <summary>
/// Register a handler function as a catch-all handler: It will be used if not type-matching handler is registered.
/// </summary>
@@ -304,9 +507,9 @@ public class RouteBuilder
return this.AddCatchAll(WrappedHandlerAsync, overwrite);
ValueTask<CallResult> WrappedHandlerAsync(PortableValue message, IWorkflowContext ctx)
ValueTask<CallResult> WrappedHandlerAsync(PortableValue message, IWorkflowContext context, CancellationToken cancellationToken)
{
TResult result = handler.Invoke(message, ctx);
TResult result = handler.Invoke(message, context);
return new(CallResult.ReturnResult(result));
}
}
@@ -43,8 +43,8 @@ public sealed class Run : IAsyncDisposable
/// <summary>
/// Gets the current execution status of the workflow run.
/// </summary>
public ValueTask<RunStatus> GetStatusAsync(CancellationToken cancellation = default)
=> this._runHandle.GetStatusAsync(cancellation);
public ValueTask<RunStatus> GetStatusAsync(CancellationToken cancellationToken = default)
=> this._runHandle.GetStatusAsync(cancellationToken);
/// <summary>
/// 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
@@ -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<JsonElement?>(ThreadStateKey).ConfigureAwait(false);
JsonElement? threadValue = await context.ReadStateAsync<JsonElement?>(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;
@@ -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<ExternalRequest?> HandleCatchAllAsync(PortableValue message, IWorkflowContext context)
public async ValueTask<ExternalRequest?> 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<ExternalRequest> HandleAsync(ExternalRequest message, IWorkflowContext context)
public async ValueTask<ExternalRequest> 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<ExternalRequest> HandleAsync(object message, IWorkflowContext context)
public async ValueTask<ExternalRequest> 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<ExternalResponse?> HandleAsync(ExternalResponse message, IWorkflowContext context)
public async ValueTask<ExternalResponse?> 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;
}
@@ -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<Type> validInputTypes = await runner.RunContext.GetStartingExecutorInputTypesAsync().ConfigureAwait(false);
IEnumerable<Type> 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<StreamingRun> EnsureRunSendMessageAsync(object? incomingMessage = null, Type? incomingMessageType = null, bool resume = false, CancellationToken cancellation = default)
internal async ValueTask<StreamingRun> 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<InMemoryCheckpointManager>(nameof(InMemoryCheckpointManager)).ConfigureAwait(false) ?? new();
InMemoryCheckpointManager manager = await context.ReadStateAsync<InMemoryCheckpointManager>(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()
@@ -31,8 +31,8 @@ public sealed class StreamingRun : IAsyncDisposable
/// <summary>
/// Gets the current execution status of the workflow run.
/// </summary>
public ValueTask<RunStatus> GetStatusAsync(CancellationToken cancellation = default)
=> this._runHandle.GetStatusAsync(cancellation);
public ValueTask<RunStatus> GetStatusAsync(CancellationToken cancellationToken = default)
=> this._runHandle.GetStatusAsync(cancellationToken);
/// <summary>
/// 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 <see cref="RequestHaltEvent"/> is encountered. Events are
/// delivered in the order they are raised.</remarks>
/// <param name="cancellationToken">A <see cref="CancellationToken"/> that can be used to cancel the streaming operation. If cancellation is
/// requested, the stream will end and no further events will be yielded.</param>
/// requested, the stream will end and no further events will be yielded, but this will not cancel the workflow execution.</param>
/// <returns>An asynchronous stream of <see cref="WorkflowEvent"/> objects representing significant workflow state changes.
/// The stream ends when the workflow completes or when cancellation is requested.</returns>
public IAsyncEnumerable<WorkflowEvent> WatchStreamAsync(
@@ -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<WorkflowFormulaState>("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);
}
}
@@ -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<TaskCompletionSource<bool>> 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);
}
}
@@ -43,12 +43,12 @@ public class InProcessStateTests
return
async (turn, context, cancellation) =>
{
TState? state = await context.ReadStateAsync<TState>(stateKey.Key, stateKey.ScopeId.ScopeName)
TState? state = await context.ReadStateAsync<TState>(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;
};
@@ -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<TActual>(string id) : ReflectingExecutor<TActual>(
public class DefaultHandler() : BaseTestExecutor<DefaultHandler>(nameof(DefaultHandler)), IMessageHandler<object>
{
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<DefaultHandler>(nameof(DefaultH
public class TypedHandler<TInput>() : BaseTestExecutor<TypedHandler<TInput>>(nameof(TypedHandler<TInput>)), IMessageHandler<TInput>
{
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<TInput>() : BaseTestExecutor<TypedHandler<TInput>>(nam
public class TypedHandlerWithOutput<TInput, TResult>() : BaseTestExecutor<TypedHandlerWithOutput<TInput, TResult>>(nameof(TypedHandlerWithOutput<TInput, TResult>)), IMessageHandler<TInput, TResult>
{
public ValueTask<TResult> HandleAsync(TInput message, IWorkflowContext context)
public ValueTask<TResult> HandleAsync(TInput message, IWorkflowContext context, CancellationToken cancellationToken)
{
this.OnInvokedHandler();
return this.Handler(message, context);
@@ -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>("UppercaseExecutor"), IMessageHandler<string, string>
{
public async ValueTask<string> HandleAsync(string message, IWorkflowContext context) =>
public async ValueTask<string> HandleAsync(string message, IWorkflowContext context, CancellationToken cancellationToken = default) =>
message.ToUpperInvariant();
}
internal sealed class ReverseTextExecutor() : ReflectingExecutor<ReverseTextExecutor>("ReverseTextExecutor"), IMessageHandler<string, string>
{
public async ValueTask<string> HandleAsync(string message, IWorkflowContext context)
public async ValueTask<string> 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;
}
}
@@ -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<DetectSpamExecutor>(id), IMessageHandler<string, bool>
{
public async ValueTask<bool> HandleAsync(string message, IWorkflowContext context) =>
public async ValueTask<bool> 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<R
{
public const string ActionResult = "Message processed successfully.";
public async ValueTask HandleAsync(bool message, IWorkflowContext context)
public async ValueTask HandleAsync(bool message, IWorkflowContext context, CancellationToken cancellationToken = default)
{
if (message)
{
@@ -72,9 +73,9 @@ internal sealed class RespondToMessageExecutor(string id) : ReflectingExecutor<R
throw new InvalidOperationException("Received a spam message that should not be getting a reply.");
}
await Task.Delay(1000).ConfigureAwait(false); // Simulate some processing delay
await Task.Delay(1000, cancellationToken).ConfigureAwait(false); // Simulate some processing delay
await context.YieldOutputAsync(ActionResult)
await context.YieldOutputAsync(ActionResult, cancellationToken)
.ConfigureAwait(false);
}
}
@@ -83,7 +84,7 @@ internal sealed class RemoveSpamExecutor(string id) : ReflectingExecutor<RemoveS
{
public const string ActionResult = "Spam message removed.";
public async ValueTask HandleAsync(bool message, IWorkflowContext context)
public async ValueTask HandleAsync(bool message, IWorkflowContext context, CancellationToken cancellationToken = default)
{
if (!message)
{
@@ -91,9 +92,9 @@ internal sealed class RemoveSpamExecutor(string id) : ReflectingExecutor<RemoveS
throw new InvalidOperationException("Received a non-spam message that should not be getting removed.");
}
await Task.Delay(1000).ConfigureAwait(false); // Simulate some processing delay
await Task.Delay(1000, cancellationToken).ConfigureAwait(false); // Simulate some processing delay
await context.YieldOutputAsync(ActionResult)
await context.YieldOutputAsync(ActionResult, cancellationToken)
.ConfigureAwait(false);
}
}
@@ -73,12 +73,12 @@ internal sealed class GuessNumberExecutor : ReflectingExecutor<GuessNumberExecut
private int NextGuess => (this.LowerBound + this.UpperBound) / 2;
private int _currGuess = -1;
public async ValueTask<int> HandleAsync(NumberSignal message, IWorkflowContext context)
public async ValueTask<int> 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<JudgeExecutor>, IMessag
this._targetNumber = targetNumber;
}
public async ValueTask<NumberSignal> HandleAsync(int message, IWorkflowContext context)
public async ValueTask<NumberSignal> 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<JudgeExecutor>, 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<int?>("TryCount").ConfigureAwait(false) ?? 0;
this.Tries = await context.ReadStateAsync<int?>("TryCount", cancellationToken: cancellationToken).ConfigureAwait(false) ?? 0;
}
public ValueTask ResetAsync()
@@ -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")
@@ -115,28 +115,28 @@ public class SpecializedExecutorSmokeTests
{
public List<List<ChatMessage>> 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<T>(string key, T? value, string? scopeName = null) =>
public ValueTask QueueStateUpdateAsync<T>(string key, T? value, string? scopeName = null, CancellationToken cancellationToken = default) =>
default;
public ValueTask<T?> ReadStateAsync<T>(string key, string? scopeName = null) =>
public ValueTask<T?> ReadStateAsync<T>(string key, string? scopeName = null, CancellationToken cancellationToken = default) =>
throw new NotImplementedException();
public ValueTask<HashSet<string>> ReadStateKeysAsync(string? scopeName = null) =>
public ValueTask<HashSet<string>> 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<ChatMessage> messages)
{
@@ -15,36 +15,36 @@ public class TestRunContext : IRunnerContext
TestRunContext runnerContext,
IReadOnlyDictionary<string, string>? 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<T>(string key, T? value, string? scopeName = null)
public ValueTask QueueStateUpdateAsync<T>(string key, T? value, string? scopeName = null, CancellationToken cancellationToken = default)
=> default;
public ValueTask<T?> ReadStateAsync<T>(string key, string? scopeName = null)
public ValueTask<T?> ReadStateAsync<T>(string key, string? scopeName = null, CancellationToken cancellationToken = default)
=> new(default(T?));
public ValueTask<HashSet<string>> ReadStateKeysAsync(string? scopeName = null)
public ValueTask<HashSet<string>> 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<string, string>? TraceContext => traceContext;
}
public List<WorkflowEvent> 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<string, List<MessageEnvelope>> 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<MessageEnvelope>? deliveryQueue))
{
@@ -72,7 +72,7 @@ public class TestRunContext : IRunnerContext
return default;
}
ValueTask<StepContext> IRunnerContext.AdvanceAsync() =>
ValueTask<StepContext> IRunnerContext.AdvanceAsync(CancellationToken cancellationToken) =>
throw new NotImplementedException();
public Dictionary<string, Executor> Executors { get; set; } = [];
@@ -80,10 +80,10 @@ public class TestRunContext : IRunnerContext
public bool WithCheckpointing => throw new NotSupportedException();
ValueTask<Executor> IRunnerContext.EnsureExecutorAsync(string executorId, IStepTracer? tracer) =>
ValueTask<Executor> IRunnerContext.EnsureExecutorAsync(string executorId, IStepTracer? tracer, CancellationToken cancellationToken) =>
new(this.Executors[executorId]);
public ValueTask<IEnumerable<Type>> GetStartingExecutorInputTypesAsync(CancellationToken cancellation = default)
public ValueTask<IEnumerable<Type>> 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<TMessage>(string senderId, [System.Diagnostics.CodeAnalysis.DisallowNull] TMessage message, CancellationToken cancellation = default)
=> this.SendMessageAsync(senderId, message, cancellation);
public ValueTask SendMessageAsync<TMessage>(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;
}