// Copyright (c) Microsoft. All rights reserved. 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; using CatchAllF = System.Func< Microsoft.Agents.AI.Workflows.PortableValue, // message Microsoft.Agents.AI.Workflows.IWorkflowContext, // context System.Threading.CancellationToken, // cancellation System.Threading.Tasks.ValueTask >; using MessageHandlerF = System.Func< object, // message Microsoft.Agents.AI.Workflows.IWorkflowContext, // context System.Threading.CancellationToken, // cancellation System.Threading.Tasks.ValueTask >; using PortHandlerF = System.Func< Microsoft.Agents.AI.Workflows.ExternalResponse, // message Microsoft.Agents.AI.Workflows.IWorkflowContext, // context System.Threading.CancellationToken, // cancellation System.Threading.Tasks.ValueTask >; namespace Microsoft.Agents.AI.Workflows; /// /// Provides a builder for configuring message type handlers for an . /// public class RouteBuilder { private readonly IExternalRequestContext? _externalRequestContext; private readonly Dictionary _typedHandlers = []; private readonly Dictionary _outputTypes = []; private readonly Dictionary _portHandlers = []; private CatchAllF? _catchAll; internal RouteBuilder(IExternalRequestContext? externalRequestContext) { this._externalRequestContext = externalRequestContext; } internal RouteBuilder AddHandlerInternal(Type messageType, MessageHandlerF handler, Type? outputType, bool overwrite = false) { Throw.IfNull(messageType); Throw.IfNull(handler); if (messageType == typeof(PortableValue)) { throw new InvalidOperationException("Cannot register a handler for PortableValue. Use AddCatchAll() instead."); } Debug.Assert(typeof(CallResult) != outputType, "Must not double-wrap message handlers in the RouteBuilder. " + "Use AddHandlerInternal() or do not wrap user-provided handler."); // Overwrite must be false if the type is not registered. Overwrite must be true if the type is registered. if (this._typedHandlers.ContainsKey(messageType) == overwrite) { this._typedHandlers[messageType] = handler; if (outputType is not null) { this._outputTypes[messageType] = outputType; } else { this._outputTypes.Remove(messageType); } } else if (overwrite) { // overwrite is true, but the type is not registered. throw new ArgumentException($"A handler for message type {messageType.FullName} has not yet been registered (overwrite = true)."); } else if (!overwrite) { throw new ArgumentException($"A handler for message type {messageType.FullName} is already registered (overwrite = false)."); } return this; } internal RouteBuilder AddHandlerUntyped(Type type, Func handler, bool overwrite = false) { Throw.IfNull(handler); return this.AddHandlerInternal(type, WrappedHandlerAsync, outputType: null, overwrite); async ValueTask WrappedHandlerAsync(object message, IWorkflowContext context, CancellationToken cancellationToken) { await handler.Invoke(message, context, cancellationToken).ConfigureAwait(false); return CallResult.ReturnVoid(); } } internal RouteBuilder AddHandlerUntyped(Type type, Func> handler, bool overwrite = false) { Throw.IfNull(handler); return this.AddHandlerInternal(type, WrappedHandlerAsync, outputType: typeof(TResult), overwrite); async ValueTask WrappedHandlerAsync(object message, IWorkflowContext context, CancellationToken cancellationToken) { TResult result = await handler.Invoke(message, context, cancellationToken).ConfigureAwait(false); return CallResult.ReturnResult(result); } } /// /// Registers a port and associated handler for external requests originating from the executor. This generates a PortBinding that can be used to /// submit requests through to the workflow Run call. /// /// The type of request messages that will be sent through this port. /// The type of response messages that will be sent through this port. /// A unique identifier for the port. /// A delegate that processes messages of type within the workflow context. The /// delegate is invoked for each incoming response to requests through this port. /// A representing this port registration providing a means to submit requests. /// Set to replace an existing handler for the specified response; if a port with this id is not /// this will throw. If set to and a handler is registered, this will throw. /// The current instance, enabling fluent configuration of additional handlers or route /// options. /// If a handler is already registered for the specified type, and overwrite is set /// to , or if a handler is not already registered, but overwrite is set to . internal RouteBuilder AddPortHandler(string id, Func handler, out PortBinding portBinding, bool overwrite = false) { if (this._externalRequestContext == null) { throw new InvalidOperationException("An external request context is required to register port handlers."); } RequestPort port = RequestPort.Create(id); IExternalRequestSink sink = this._externalRequestContext!.RegisterPort(port); portBinding = new(port, sink); if (this._portHandlers.ContainsKey(id) == overwrite) { this._portHandlers[id] = InvokeHandlerAsync; } else if (overwrite) { throw new InvalidOperationException($"A handler for port id {id} is not registered (overwrite = true)."); } else { throw new InvalidOperationException($"A handler for port id {id} is already registered (overwrite = false)."); } return this; async ValueTask InvokeHandlerAsync(ExternalResponse response, IWorkflowContext context, CancellationToken cancellationToken) { if (!response.TryGetDataAs(out TResponse? typedResponse)) { throw new InvalidOperationException($"Received response data is not of expected type {typeof(TResponse).FullName} for port {port.Id}."); } await handler(typedResponse, context, cancellationToken).ConfigureAwait(false); return response; } } /// /// Registers a handler for messages of the specified input type in the workflow route. /// /// If a handler for the specified input type already exists and is /// , the existing handler will not be replaced. Handlers are invoked asynchronously and are /// expected to complete their processing before the workflow continues. /// /// A delegate that processes messages of type within the workflow context. The /// delegate is invoked for each incoming message of the specified type. /// Set to replace an existing handler for the specified input type; if no /// handler is registered will throw. If set to and a handler is registered, this will throw. /// The current instance, enabling fluent configuration of additional handlers or route /// options. /// If a handler is already registered for the specified type, and overwrite is set /// to , or if a handler is not already registered, but overwrite is set to . public RouteBuilder AddHandler(Action handler, bool overwrite = false) { Throw.IfNull(handler); return this.AddHandlerInternal(typeof(TInput), WrappedHandlerAsync, outputType: null, overwrite); async ValueTask WrappedHandlerAsync(object message, IWorkflowContext context, CancellationToken cancellationToken) { handler.Invoke((TInput)message, context, cancellationToken); return CallResult.ReturnVoid(); } } /// /// Registers a handler for messages of the specified input type in the workflow route. /// /// If a handler for the specified input type already exists and is /// , the existing handler will not be replaced. Handlers are invoked asynchronously and are /// expected to complete their processing before the workflow continues. /// /// A delegate that processes messages of type within the workflow context. The /// delegate is invoked for each incoming message of the specified type. /// Set to replace an existing handler for the specified input type; if no /// handler is registered will throw. If set to and a handler is registered, this will throw. /// The current instance, enabling fluent configuration of additional handlers or route /// options. /// If a handler is already registered for the specified type, and overwrite is set /// to , or if a handler is not already registered, but overwrite is set to . public RouteBuilder AddHandler(Action handler, bool overwrite = false) { Throw.IfNull(handler); return this.AddHandlerInternal(typeof(TInput), WrappedHandlerAsync, outputType: null, overwrite); async ValueTask WrappedHandlerAsync(object message, IWorkflowContext context, CancellationToken cancellationToken) { handler.Invoke((TInput)message, context); return CallResult.ReturnVoid(); } } /// /// Registers a handler for messages of the specified input type in the workflow route. /// /// If a handler for the specified input type already exists and is /// , the existing handler will not be replaced. Handlers are invoked asynchronously and are /// expected to complete their processing before the workflow continues. /// /// A delegate that processes messages of type within the workflow context. The /// delegate is invoked for each incoming message of the specified type. /// Set to replace an existing handler for the specified input type; if no /// handler is registered will throw. If set to and a handler is registered, this will throw. /// The current instance, enabling fluent configuration of additional handlers or route /// options. /// If a handler is already registered for the specified type, and overwrite is set /// to , or if a handler is not already registered, but overwrite is set to . public RouteBuilder AddHandler(Func handler, bool overwrite = false) { Throw.IfNull(handler); return this.AddHandlerInternal(typeof(TInput), WrappedHandlerAsync, outputType: null, overwrite); async ValueTask WrappedHandlerAsync(object message, IWorkflowContext context, CancellationToken cancellationToken) { await handler.Invoke((TInput)message, context, cancellationToken).ConfigureAwait(false); return CallResult.ReturnVoid(); } } /// /// Registers a handler for messages of the specified input type in the workflow route. /// /// If a handler for the specified input type already exists and is /// , the existing handler will not be replaced. Handlers are invoked asynchronously and are /// expected to complete their processing before the workflow continues. /// /// A delegate that processes messages of type within the workflow context. The /// delegate is invoked for each incoming message of the specified type. /// Set to replace an existing handler for the specified input type; if no /// handler is registered will throw. If set to and a handler is registered, this will throw. /// The current instance, enabling fluent configuration of additional handlers or route /// options. /// If a handler is already registered for the specified type, and overwrite is set /// to , or if a handler is not already registered, but overwrite is set to . public RouteBuilder AddHandler(Func handler, bool overwrite = false) { Throw.IfNull(handler); return this.AddHandlerInternal(typeof(TInput), WrappedHandlerAsync, outputType: null, overwrite); async ValueTask WrappedHandlerAsync(object message, IWorkflowContext context, CancellationToken cancellationToken) { await handler.Invoke((TInput)message, context).ConfigureAwait(false); return CallResult.ReturnVoid(); } } /// /// Registers a handler function for messages of the specified input type in the workflow route. /// /// If a handler for the given input type already exists, setting to /// will replace the existing handler; otherwise, an exception may be thrown. The handler /// receives the input message and workflow context, and returns a result asynchronously. /// The type of input message the handler will process. /// The type of result produced by the handler. /// A function that processes messages of type within the workflow context and returns /// a representing the asynchronous result. /// Set to replace an existing handler for the specified input type; if no /// handler is registered will throw. If set to and a handler is registered, this will throw. /// The current instance, enabling fluent configuration of workflow routes. /// If a handler is already registered for the specified type, and overwrite is set /// to , or if a handler is not already registered, but overwrite is set to . public RouteBuilder AddHandler(Func handler, bool overwrite = false) { Throw.IfNull(handler); return this.AddHandlerInternal(typeof(TInput), WrappedHandlerAsync, outputType: typeof(TResult), overwrite); async ValueTask WrappedHandlerAsync(object message, IWorkflowContext context, CancellationToken cancellationToken) { TResult result = handler.Invoke((TInput)message, context, cancellationToken); return CallResult.ReturnResult(result); } } /// /// Registers a handler function for messages of the specified input type in the workflow route. /// /// If a handler for the given input type already exists, setting to /// will replace the existing handler; otherwise, an exception may be thrown. The handler /// receives the input message and workflow context, and returns a result asynchronously. /// The type of input message the handler will process. /// The type of result produced by the handler. /// A function that processes messages of type within the workflow context and returns /// a representing the asynchronous result. /// Set to replace an existing handler for the specified input type; if no /// handler is registered will throw. If set to and a handler is registered, this will throw. /// The current instance, enabling fluent configuration of workflow routes. /// If a handler is already registered for the specified type, and overwrite is set /// to , or if a handler is not already registered, but overwrite is set to . public RouteBuilder AddHandler(Func handler, bool overwrite = false) { Throw.IfNull(handler); return this.AddHandlerInternal(typeof(TInput), WrappedHandlerAsync, outputType: typeof(TResult), overwrite); async ValueTask WrappedHandlerAsync(object message, IWorkflowContext context, CancellationToken cancellationToken) { TResult result = handler.Invoke((TInput)message, context); return CallResult.ReturnResult(result); } } /// /// Registers a handler function for messages of the specified input type in the workflow route. /// /// If a handler for the given input type already exists, setting to /// will replace the existing handler; otherwise, an exception may be thrown. The handler /// receives the input message and workflow context, and returns a result asynchronously. /// The type of input message the handler will process. /// The type of result produced by the handler. /// A function that processes messages of type within the workflow context and returns /// a representing the asynchronous result. /// Set to replace an existing handler for the specified input type; if no /// handler is registered will throw. If set to and a handler is registered, this will throw. /// The current instance, enabling fluent configuration of workflow routes. /// If a handler is already registered for the specified type, and overwrite is set /// to , or if a handler is not already registered, but overwrite is set to . public RouteBuilder AddHandler(Func> handler, bool overwrite = false) { Throw.IfNull(handler); return this.AddHandlerInternal(typeof(TInput), WrappedHandlerAsync, outputType: typeof(TResult), overwrite); async ValueTask WrappedHandlerAsync(object message, IWorkflowContext context, CancellationToken cancellationToken) { TResult result = await handler((TInput)message, context, cancellationToken).ConfigureAwait(false); return CallResult.ReturnResult(result); } } /// /// Registers a handler function for messages of the specified input type in the workflow route. /// /// If a handler for the given input type already exists, setting to /// will replace the existing handler; otherwise, an exception may be thrown. The handler /// receives the input message and workflow context, and returns a result asynchronously. /// The type of input message the handler will process. /// The type of result produced by the handler. /// A function that processes messages of type within the workflow context and returns /// a representing the asynchronous result. /// Set to replace an existing handler for the specified input type; if no /// handler is registered will throw. If set to and a handler is registered, this will throw. /// The current instance, enabling fluent configuration of workflow routes. /// If a handler is already registered for the specified type, and overwrite is set /// to , or if a handler is not already registered, but overwrite is set to . public RouteBuilder AddHandler(Func> handler, bool overwrite = false) { Throw.IfNull(handler); return this.AddHandlerInternal(typeof(TInput), WrappedHandlerAsync, outputType: typeof(TResult), overwrite); async ValueTask WrappedHandlerAsync(object message, IWorkflowContext context, CancellationToken cancellationToken) { TResult result = await handler.Invoke((TInput)message, context).ConfigureAwait(false); return CallResult.ReturnResult(result); } } private RouteBuilder AddCatchAll(CatchAllF handler, bool overwrite = false) { if (!overwrite && this._catchAll != null) { throw new InvalidOperationException("A catch-all is already registered (overwrite = false)."); } this._catchAll = handler; return this; } /// /// Register a handler function as a catch-all handler: It will be used if not type-matching handler is registered. /// /// If a catch-all handler for already exists, setting to /// will replace the existing handler; otherwise, an exception may be thrown. The handler receives the input message /// wrapped as and workflow context, and returns a result asynchronously. /// A function that processes messages wrapped as within the /// workflow context. The delegate is invoked for each incoming message not otherwise handled. /// Set to replace an existing handler for the specified input type; if no /// handler is registered will throw. If set to and a handler is registered, this will throw. /// The current instance, enabling fluent configuration of workflow routes. /// If a handler is already registered for the specified type, and overwrite is set /// to , or if a handler is not already registered, but overwrite is set to . public RouteBuilder AddCatchAll(Func handler, bool overwrite = false) { Throw.IfNull(handler); return this.AddCatchAll(WrappedHandlerAsync, overwrite); async ValueTask WrappedHandlerAsync(PortableValue message, IWorkflowContext context, CancellationToken cancellationToken) { await handler.Invoke(message, context, cancellationToken).ConfigureAwait(false); return CallResult.ReturnVoid(); } } /// /// Register a handler function as a catch-all handler: It will be used if not type-matching handler is registered. /// /// If a catch-all handler for already exists, setting to /// will replace the existing handler; otherwise, an exception may be thrown. The handler receives the input message /// wrapped as and workflow context, and returns a result asynchronously. /// A function that processes messages wrapped as within the /// workflow context. The delegate is invoked for each incoming message not otherwise handled. /// Set to replace an existing handler for the specified input type; if no /// handler is registered will throw. If set to and a handler is registered, this will throw. /// The current instance, enabling fluent configuration of workflow routes. /// If a handler is already registered for the specified type, and overwrite is set /// to , or if a handler is not already registered, but overwrite is set to . public RouteBuilder AddCatchAll(Func handler, bool overwrite = false) { Throw.IfNull(handler); return this.AddCatchAll(WrappedHandlerAsync, overwrite); async ValueTask WrappedHandlerAsync(PortableValue message, IWorkflowContext context, CancellationToken cancellationToken) { await handler.Invoke(message, context).ConfigureAwait(false); return CallResult.ReturnVoid(); } } /// /// Register a handler function as a catch-all handler: It will be used if not type-matching handler is registered. /// /// If a catch-all handler for already exists, setting to /// will replace the existing handler; otherwise, an exception may be thrown. The handler receives the input message /// wrapped as and workflow context, and returns a result asynchronously. /// A function that processes messages wrapped as within the /// workflow context and returns a representing the asynchronous result. /// Set to replace an existing handler for the specified input type; if no /// handler is registered will throw. If set to and a handler is registered, this will throw. /// The current instance, enabling fluent configuration of workflow routes. /// If a handler is already registered for the specified type, and overwrite is set /// to , or if a handler is not already registered, but overwrite is set to . public RouteBuilder AddCatchAll(Func> handler, bool overwrite = false) { Throw.IfNull(handler); return this.AddCatchAll(WrappedHandlerAsync, overwrite); async ValueTask WrappedHandlerAsync(PortableValue message, IWorkflowContext context, CancellationToken cancellationToken) { TResult result = await handler.Invoke(message, context, cancellationToken).ConfigureAwait(false); return CallResult.ReturnResult(result); } } /// /// Register a handler function as a catch-all handler: It will be used if not type-matching handler is registered. /// /// If a catch-all handler for already exists, setting to /// will replace the existing handler; otherwise, an exception may be thrown. The handler receives the input message /// wrapped as and workflow context, and returns a result asynchronously. /// A function that processes messages wrapped as within the /// workflow context and returns a representing the asynchronous result. /// Set to replace an existing handler for the specified input type; if no /// handler is registered will throw. If set to and a handler is registered, this will throw. /// The current instance, enabling fluent configuration of workflow routes. /// If a handler is already registered for the specified type, and overwrite is set /// to , or if a handler is not already registered, but overwrite is set to . public RouteBuilder AddCatchAll(Func> handler, bool overwrite = false) { Throw.IfNull(handler); return this.AddCatchAll(WrappedHandlerAsync, overwrite); async ValueTask WrappedHandlerAsync(PortableValue message, IWorkflowContext context, CancellationToken cancellationToken) { TResult result = await handler.Invoke(message, context).ConfigureAwait(false); return CallResult.ReturnResult(result); } } /// /// Register a handler function as a catch-all handler: It will be used if not type-matching handler is registered. /// /// If a catch-all handler for already exists, setting to /// will replace the existing handler; otherwise, an exception may be thrown. The handler receives the input message /// wrapped as and workflow context, and returns a result asynchronously. /// A function that processes messages wrapped as within the /// workflow context. The delegate is invoked for each incoming message not otherwise handled. /// Set to replace an existing handler for the specified input type; if no /// handler is registered will throw. If set to and a handler is registered, this will throw. /// The current instance, enabling fluent configuration of workflow routes. /// If a handler is already registered for the specified type, and overwrite is set /// to , or if a handler is not already registered, but overwrite is set to . public RouteBuilder AddCatchAll(Action handler, bool overwrite = false) { Throw.IfNull(handler); return this.AddCatchAll(WrappedHandlerAsync, overwrite); ValueTask WrappedHandlerAsync(PortableValue message, IWorkflowContext ctx, CancellationToken cancellationToken) { handler.Invoke(message, ctx, cancellationToken); return new(CallResult.ReturnVoid()); } } /// /// Register a handler function as a catch-all handler: It will be used if not type-matching handler is registered. /// /// If a catch-all handler for already exists, setting to /// will replace the existing handler; otherwise, an exception may be thrown. The handler receives the input message /// wrapped as and workflow context, and returns a result asynchronously. /// A function that processes messages wrapped as within the /// workflow context. The delegate is invoked for each incoming message not otherwise handled. /// Set to replace an existing handler for the specified input type; if no /// handler is registered will throw. If set to and a handler is registered, this will throw. /// The current instance, enabling fluent configuration of workflow routes. /// If a handler is already registered for the specified type, and overwrite is set /// to , or if a handler is not already registered, but overwrite is set to . public RouteBuilder AddCatchAll(Action handler, bool overwrite = false) { Throw.IfNull(handler); return this.AddCatchAll(WrappedHandlerAsync, overwrite); ValueTask WrappedHandlerAsync(PortableValue message, IWorkflowContext ctx, CancellationToken cancellationToken) { handler.Invoke(message, ctx); return new(CallResult.ReturnVoid()); } } /// /// Register a handler function as a catch-all handler: It will be used if not type-matching handler is registered. /// /// If a catch-all handler for already exists, setting to /// will replace the existing handler; otherwise, an exception may be thrown. The handler receives the input message /// wrapped as and workflow context, and returns a result asynchronously. /// A function that processes messages wrapped as within the /// workflow context and returns a representing the asynchronous result. /// Set to replace an existing handler for the specified input type; if no /// handler is registered will throw. If set to and a handler is registered, this will throw. /// The current instance, enabling fluent configuration of workflow routes. /// If a handler is already registered for the specified type, and overwrite is set /// to , or if a handler is not already registered, but overwrite is set to . public RouteBuilder AddCatchAll(Func handler, bool overwrite = false) { Throw.IfNull(handler); return this.AddCatchAll(WrappedHandlerAsync, overwrite); ValueTask WrappedHandlerAsync(PortableValue message, IWorkflowContext context, CancellationToken cancellationToken) { TResult result = handler.Invoke(message, context, cancellationToken); return new(CallResult.ReturnResult(result)); } } /// /// Register a handler function as a catch-all handler: It will be used if not type-matching handler is registered. /// /// If a catch-all handler for already exists, setting to /// will replace the existing handler; otherwise, an exception may be thrown. The handler receives the input message /// wrapped as and workflow context, and returns a result asynchronously. /// A function that processes messages wrapped as within the /// workflow context and returns a representing the asynchronous result. /// Set to replace an existing handler for the specified input type; if no /// handler is registered will throw. If set to and a handler is registered, this will throw. /// The current instance, enabling fluent configuration of workflow routes. /// If a handler is already registered for the specified type, and overwrite is set /// to , or if a handler is not already registered, but overwrite is set to . public RouteBuilder AddCatchAll(Func handler, bool overwrite = false) { Throw.IfNull(handler); return this.AddCatchAll(WrappedHandlerAsync, overwrite); ValueTask WrappedHandlerAsync(PortableValue message, IWorkflowContext context, CancellationToken cancellationToken) { TResult result = handler.Invoke(message, context); return new(CallResult.ReturnResult(result)); } } private void RegisterPortHandlerRouter() { Dictionary portHandlers = this._portHandlers; this.AddHandler(InvokeHandlerAsync); ValueTask InvokeHandlerAsync(ExternalResponse response, IWorkflowContext context, CancellationToken cancellationToken) { if (portHandlers.TryGetValue(response.PortInfo.PortId, out PortHandlerF? portHandler)) { return portHandler(response, context, cancellationToken); } throw new InvalidOperationException($"Unknown port {response.PortInfo}"); } } internal IEnumerable OutputTypes => this._outputTypes.Values; internal MessageRouter Build() { if (this._portHandlers.Count > 0) { this.RegisterPortHandlerRouter(); } return new(this._typedHandlers, [.. this._outputTypes.Values], this._catchAll); } }