.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-08 00:28:44 +00:00
committed by GitHub
parent eb049c43a6
commit 5902bcb10a
74 changed files with 910 additions and 557 deletions
@@ -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(