agent-framework/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/RequestInfoExecutor.cs

// 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;

namespace Microsoft.Agents.AI.Workflows.Specialized;

internal sealed class RequestPortOptions;

internal sealed class RequestInfoExecutor : Executor
{
    private const string WrappedRequestsStateKey = nameof(WrappedRequestsStateKey);
    private readonly Dictionary<string, ExternalRequest> _wrappedRequests = [];
    private RequestPort Port { get; }
    private IExternalRequestSink? RequestSink { get; set; }

    private static ExecutorOptions DefaultOptions => new()
    {
        // We need to be able to return the ExternalRequest/Result objects so they can be bubbled up
        // through the event system, but we do not want to forward the Request message.
        AutoSendMessageHandlerResultObject = false,
        AutoYieldOutputHandlerResultObject = false
    };

    private readonly bool _allowWrapped;
    public RequestInfoExecutor(RequestPort port, bool allowWrapped = true) : base(port.Id, DefaultOptions)
    {
        this.Port = port;

        this._allowWrapped = allowWrapped;
    }

    protected override ProtocolBuilder ConfigureProtocol(ProtocolBuilder protocolBuilder)
    {
        return protocolBuilder.ConfigureRoutes(ConfigureRoutes)
                              .SendsMessage<ExternalRequest>()
                              .SendsMessageType(this.Port.Response);

        void ConfigureRoutes(RouteBuilder routeBuilder)
        {
            routeBuilder = routeBuilder
                // Handle incoming requests (as raw request payloads)
                .AddHandlerUntyped(this.Port.Request, this.HandleAsync)
                .AddCatchAll(this.HandleCatchAllAsync);

            if (this._allowWrapped)
            {
                routeBuilder = routeBuilder
                    .AddHandler<ExternalRequest, ExternalRequest>(this.HandleAsync);
            }

            routeBuilder
                // Handle incoming responses (as wrapped Response object)
                .AddHandler<ExternalResponse, ExternalResponse?>(this.HandleAsync);
        }
    }

    internal void AttachRequestSink(IExternalRequestSink requestSink) => this.RequestSink = Throw.IfNull(requestSink);

    public async ValueTask<ExternalRequest?> HandleCatchAllAsync(PortableValue message, IWorkflowContext context, CancellationToken cancellationToken)
    {
        Throw.IfNull(message);

        object? maybeRequest = message.AsType(this.Port.Request);
        if (maybeRequest != null)
        {
            Debug.Assert(this.Port.Request.IsInstanceOfType(maybeRequest));

            ExternalRequest request = ExternalRequest.Create(this.Port, maybeRequest!);
            await this.RequestSink!.PostAsync(request).ConfigureAwait(false);
            return request;
        }
        else if (message.Is(out ExternalRequest? request))
        {
            return await this.HandleAsync(request, context, cancellationToken).ConfigureAwait(false);
        }

        return null;
    }

    public async ValueTask<ExternalRequest> HandleAsync(ExternalRequest message, IWorkflowContext context, CancellationToken cancellationToken = default)
    {
        Debug.Assert(this._allowWrapped);
        Throw.IfNull(message);

        if (!message.Data.IsType(this.Port.Request, out var requestData))
        {
            throw new InvalidOperationException($"Message type {message.Data.TypeId} could not be interpreted as a value of Request Type {this.Port.Request}");
        }

        if (!message.PortInfo.ResponseType.IsMatchPolymorphic(this.Port.Response))
        {
            throw new InvalidOperationException($"Response type {this.Port.Response} is not a valid response for original request, whose expected response is {message.PortInfo.ResponseType}");
        }

        ExternalRequest request = ExternalRequest.Create(this.Port, requestData, message.RequestId);

        this._wrappedRequests.Add(message.RequestId, message);

        await this.RequestSink!.PostAsync(request).ConfigureAwait(false);

        return request;
    }

    public async ValueTask<ExternalRequest> HandleAsync(object message, IWorkflowContext context, CancellationToken cancellationToken = default)
    {
        Throw.IfNull(message);
        Debug.Assert(this.Port.Request.IsInstanceOfType(message));

        ExternalRequest request = ExternalRequest.Create(this.Port, message);
        await this.RequestSink!.PostAsync(request).ConfigureAwait(false);

        return request;
    }

    public async ValueTask<ExternalResponse?> HandleAsync(ExternalResponse message, IWorkflowContext context, CancellationToken cancellationToken = default)
    {
        if (!this.Port.IsResponsePort(message))
        {
            return null;
        }

        if (!message.Data.IsType(this.Port.Response, out object? data))
        {
            throw this.Port.CreateExceptionForType(message);
        }

        if (this._allowWrapped && this._wrappedRequests.TryGetValue(message.RequestId, out ExternalRequest? originalRequest))
        {
            await context.SendMessageAsync(originalRequest.RewrapResponse(message), cancellationToken: cancellationToken).ConfigureAwait(false);
            this._wrappedRequests.Remove(message.RequestId);
        }
        else
        {
            await context.SendMessageAsync(message, cancellationToken: cancellationToken).ConfigureAwait(false);
            await context.SendMessageAsync(data, cancellationToken: cancellationToken).ConfigureAwait(false);
        }

        return message;
    }

    protected internal override async ValueTask OnCheckpointingAsync(IWorkflowContext context, CancellationToken cancellationToken = default)
    {
        await context.QueueStateUpdateAsync(WrappedRequestsStateKey,
                                            new Dictionary<string, ExternalRequest>(this._wrappedRequests, StringComparer.Ordinal),
                                            cancellationToken: cancellationToken).ConfigureAwait(false);
        await base.OnCheckpointingAsync(context, cancellationToken).ConfigureAwait(false);
    }

    protected internal override async ValueTask OnCheckpointRestoredAsync(IWorkflowContext context, CancellationToken cancellationToken = default)
    {
        await base.OnCheckpointRestoredAsync(context, cancellationToken).ConfigureAwait(false);

        this._wrappedRequests.Clear();

        Dictionary<string, ExternalRequest> wrappedRequests =
            await context.ReadStateAsync<Dictionary<string, ExternalRequest>>(WrappedRequestsStateKey, cancellationToken: cancellationToken)
                         .ConfigureAwait(false) ?? [];

        foreach (KeyValuePair<string, ExternalRequest> wrappedRequest in wrappedRequests)
        {
            this._wrappedRequests[wrappedRequest.Key] = wrappedRequest.Value;
        }
    }
}