VolcanicAgenticCore/tests/Volcanic.AgenticCore.A2A.Server.Tests/Support/A2ATestServer.cs

using System.Collections.Concurrent;
using A2A;
using Microsoft.AspNetCore.Builder;
using Microsoft.AspNetCore.Hosting;
using Microsoft.AspNetCore.TestHost;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging.Abstractions;
using Volcanic.AgenticCore.A2A.Server.Callers;
using Volcanic.AgenticCore.A2A.Server.Hosting;
using Volcanic.AgenticCore.A2A.Server.Push;
using Volcanic.AgenticCore.Features.Agent;
using Volcanic.AgenticCore.Features.Extensions;
using Volcanic.AgenticCore.Hosting;
using Volcanic.AgenticCore.Testing;

namespace Volcanic.AgenticCore.A2A.Server.Tests.Support;

/// <summary>An in-process A2A server over two workspaces with the FauxProvider, a recording notifier and push handler.</summary>
public sealed class A2ATestServer : IAsyncDisposable
{
    public const string Key = "key-alpha";
    public const string OtherKey = "key-beta";

    private readonly string _root = Path.Combine(Path.GetTempPath(), "a2a-e2e", Guid.NewGuid().ToString("N"));

    public FauxProvider Faux { get; } = new();
    public string Workspace => Path.Combine(_root, "ws1");
    public string OtherWorkspace => Path.Combine(_root, "ws2");
    public RecordingNotifier Notifier { get; } = new();
    public RecordingHandler Push { get; } = new();
    public WebApplication App { get; private set; } = null!;

    public static async Task<A2ATestServer> StartAsync(Action<A2AHostOptions>? configure = null, bool withNotifier = true)
    {
        var host = new A2ATestServer();
        try
        {
            Directory.CreateDirectory(host.Workspace);
            Directory.CreateDirectory(host.OtherWorkspace);
            host.WriteCallers(host.Workspace, $$"""[ { "name": "alpha", "key": "{{Key}}" }, { "name": "gamma", "key": "key-gamma" } ]""");
            host.WriteCallers(host.OtherWorkspace, $$"""[ { "name": "beta", "key": "{{OtherKey}}" } ]""");

            var builder = WebApplication.CreateBuilder();
            builder.WebHost.UseTestServer();
            builder.Services.AddAgenticCore(o =>
            {
                o.AgentDir = Path.Combine(host._root, "agent");
                o.DiscoverExtensionAssemblies = false;
                o.DefaultModel = FauxProvider.ModelKey;
                o.Extensions.Add(ExtensionRegistration.FromInstance(host.Faux.Extension));
                o.ProcessEnvironment.HardenHostProcess = false;
                o.Retry = new RetrySettings { Enabled = false }; // FauxTurn.Error fails the task at once
            });
            builder.Services.AddSingleton<A2APushSender>(sp => new HttpA2APushSender(new HttpClient(host.Push), TimeProvider.System, NullLogger<HttpA2APushSender>.Instance));
            if (withNotifier) builder.Services.AddSingleton<IA2AOwnerNotifier>(host.Notifier);
            builder.Services.AddSingleton<IA2ACallerDirectory>(new WorkspaceCallerDirectory(
                () => [new A2AWorkspace(host.Workspace, "test-ws1", "ws1"), new A2AWorkspace(host.OtherWorkspace, "test-ws2", "ws2")],
                NullLogger<WorkspaceCallerDirectory>.Instance));
            builder.Services.AddA2AServer(o =>
            {
                o.AllowPrivatePushUrls = true;
                configure?.Invoke(o);
            });
            host.App = builder.Build();
            host.App.MapA2AServer();
            await host.App.StartAsync();
            return host;
        }
        catch
        {
            await host.DisposeAsync();
            throw;
        }
    }

    public void WriteCallers(string workspace, string json)
    {
        Directory.CreateDirectory(Path.GetDirectoryName(CallersFile.PathFor(workspace))!);
        File.WriteAllText(CallersFile.PathFor(workspace), json);
    }

    public HttpClient Http(string? key = Key)
    {
        var http = App.GetTestClient();
        if (key is not null) http.DefaultRequestHeaders.Add("X-API-Key", key);
        return http;
    }

    public A2AClient Client(string key = Key) => new(new Uri("http://localhost/a2a"), Http(key));

    public async ValueTask DisposeAsync()
    {
        if (App is not null)
        {
            await App.StopAsync();
            await App.DisposeAsync();
        }
        try { Directory.Delete(_root, true); } catch (Exception e) when (e is IOException or UnauthorizedAccessException) { }
    }

    public sealed class RecordingNotifier : IA2AOwnerNotifier
    {
        public ConcurrentQueue<(string Caller, string Text)> Sent { get; } = new();
        public Task NotifyAsync(A2ACaller caller, string text, CancellationToken cancellationToken)
        {
            Sent.Enqueue((caller.Name, text));
            return Task.CompletedTask;
        }
    }

    public sealed class RecordingHandler : HttpMessageHandler
    {
        public ConcurrentQueue<(Uri Url, string? Token, string Body)> Requests { get; } = new();
        protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
        {
            var token = request.Headers.TryGetValues(A2APushReceiver.TokenHeader, out var v) ? v.Single() : null;
            Requests.Enqueue((request.RequestUri!, token, await request.Content!.ReadAsStringAsync(cancellationToken)));
            return new HttpResponseMessage(System.Net.HttpStatusCode.OK);
        }
    }
}