VolcanicAgenticCore/tests/Volcanic.AgenticCore.Telegram.Tests/A2AClientHostTests.cs

using System.Diagnostics;
using System.Text.Json;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging.Abstractions;
using Volcanic.AgenticCore.A2A.Server;
using Volcanic.AgenticCore.A2A.Server.Agent;
using Volcanic.AgenticCore.A2A.Server.Callers;
using Volcanic.AgenticCore.Features.Extensions;
using Volcanic.AgenticCore.Extensions.A2A;
using Volcanic.AgenticCore.Telegram.A2A;
using Volcanic.AgenticCore.Telegram.Tests.Support;
using Volcanic.AgenticCore.Testing;

namespace Volcanic.AgenticCore.Telegram.Tests;

public class A2AClientHostTests
{
    private static TelegramA2AResultDelivery DeliveryFor(ChatHost host) =>
        new(host.Chats, host.Settings, host.Options, TimeProvider.System, NullLogger<TelegramA2AResultDelivery>.Instance);

    [Fact]
    public async Task ClientDisabled_NoA2aTool()
    {
        await using var host = new ChatHost();
        var chat = await host.ChatAsync();
        Assert.DoesNotContain(chat.Session.Tools, t => t.Name == "a2a");
    }

    [Fact]
    public async Task ClientEnabled_ChatHasA2aTool()
    {
        await using var host = new ChatHost(a2aClient: new A2AClientSettings { Enabled = true });
        var chat = await host.ChatAsync();
        Assert.Contains(chat.Session.Tools, t => t.Name == "a2a");
    }

    [Fact]
    public async Task ResultDelivery_GoesToChatWithTimestampHeader()
    {
        await using var host = new ChatHost(a2aClient: new A2AClientSettings { Enabled = true });
        host.Faux.Enqueue(FauxTurn.Say("Got the result"));
        var delivery = DeliveryFor(host);
        Assert.True(await delivery.DeliverAsync("1001", "[A2A result from \"remote\", task t1 (completed)]\nDone", CancellationToken.None));
        var prompt = host.Faux.Requests.Last().Messages.Last().Text;
        Assert.Contains("[A2A result from \"remote\"", prompt);
        Assert.Matches(@"^\[.*\d{4}", prompt); // the local-time header comes first
        await (await host.ChatAsync()).Renderer.DrainAsync().WaitAsync(TimeSpan.FromSeconds(5));
        Assert.Contains(host.Texts, t => t.Contains("Got the result"));
    }

    [Fact]
    public async Task ResultDelivery_UnknownOrNotAllowedTarget_False()
    {
        await using var host = new ChatHost(a2aClient: new A2AClientSettings { Enabled = true });
        var delivery = DeliveryFor(host);
        Assert.False(await delivery.DeliverAsync("999", "x", CancellationToken.None));
        Assert.False(await delivery.DeliverAsync("not-a-chat", "x", CancellationToken.None));
        Assert.Equal(["1001"], delivery.KnownTargets());
        Assert.Null(delivery.WorkspaceFor("not-a-chat"));
        Assert.Equal(host.WorkspaceOf(), delivery.WorkspaceFor("1001"));
    }

    [Fact]
    public async Task A2aCommand_ListsAgentsAndWatches()
    {
        await using var host = new ChatHost(a2aClient: new A2AClientSettings { Enabled = true });
        var ws = host.WorkspaceOf();
        Directory.CreateDirectory(Path.GetDirectoryName(AgentsFile.PathFor(ws))!);
        File.WriteAllText(AgentsFile.PathFor(ws), """
            [ { "name": "travel-bot", "url": "https://travel.example/", "key": "agent-key-123456789" },
              { "name": "short", "url": "https://short.example/", "key": "short-key" } ]
            """);
        await host.SendAsync("/a2a");
        await host.SettleAsync();
        var reply = host.LastReply();
        Assert.Contains("travel-bot", reply);
        Assert.Contains("https://travel.example/", reply);
        Assert.DoesNotContain("agent-key-123456789", reply);
        Assert.Contains("…6789", reply);
        Assert.DoesNotContain("short-key", reply);
        Assert.Contains("Watches", reply);
        Assert.Contains("A2A server: off", reply);
    }

    [Fact]
    public async Task A2aCommand_BothOff_SaysHowToEnable()
    {
        await using var host = new ChatHost();
        await host.SendAsync("/a2a");
        await host.SettleAsync();
        var reply = host.LastReply();
        Assert.Contains("A2A is off", reply);
        Assert.Contains("A2A__Enabled", reply);
        Assert.Contains("A2AClient__Enabled", reply);
    }

    [Fact]
    public async Task A2aCommand_ServerOnly_ShowsCallersAndClientOff()
    {
        await using var host = new ChatHost(a2a: new A2AHostOptions { Enabled = true });
        await host.SendAsync("/a2a");
        await host.SettleAsync();
        var reply = host.LastReply();
        Assert.Contains("A2A callers", reply);
        Assert.Contains("A2A client: off", reply);
    }

    [Fact]
    public void Setup_ServerOff_NoPushBaseUrlEvenIfConfigured()
    {
        var services = new ServiceCollection();
        var client = new A2AClientSettings { Enabled = true, PushBaseUrl = "https://configured.example" };
        var server = new A2AHostOptions { Enabled = false, PublicUrl = "https://bot.example" };
        TelegramA2AClientSetup.Register(services, client, server, new A2AWatchRegistry(TimeProvider.System, new HttpClientHandler()));
        Assert.Null(client.PushBaseUrl);
        Assert.Null(server.OnPush);
        Assert.Contains(services, d => d.ServiceType == typeof(A2AWatchRegistry));
        Assert.Contains(services, d => d.ImplementationType == typeof(A2AWatchHostedService));
    }

    [Fact]
    public async Task Setup_ServerOn_PushRoutedIntoTheRegistry()
    {
        var services = new ServiceCollection();
        var client = new A2AClientSettings { Enabled = true };
        var server = new A2AHostOptions { Enabled = true, PublicUrl = "https://bot.example" };
        await using var registry = new A2AWatchRegistry(TimeProvider.System, new HttpClientHandler());
        TelegramA2AClientSetup.Register(services, client, server, registry);
        Assert.Equal("https://bot.example", client.PushBaseUrl);
        Assert.NotNull(server.OnPush);
        Assert.False(await server.OnPush!("unknown", "token", JsonDocument.Parse("{}").RootElement, CancellationToken.None));
        Assert.Contains("push at https://bot.example", TelegramA2AClientSetup.Describe(client));
    }

    [Fact]
    public void Setup_ClientOff_RegistersOnlyTheSettings()
    {
        var services = new ServiceCollection();
        var client = new A2AClientSettings { PushBaseUrl = "https://configured.example" };
        var server = new A2AHostOptions { Enabled = true, PublicUrl = "https://bot.example" };
        TelegramA2AClientSetup.Register(services, client, server, new A2AWatchRegistry(TimeProvider.System, new HttpClientHandler()));
        Assert.Null(server.OnPush);
        Assert.Null(client.PushBaseUrl);
        Assert.Equal(typeof(A2AClientSettings), Assert.Single(services).ServiceType);
        Assert.Equal("A2A client: off", TelegramA2AClientSetup.Describe(client));
    }

    [Fact]
    public async Task HostedService_RestoreRunsInTheBackground()
    {
        var settings = new A2AClientSettings { Enabled = true };
        await using var host = new ChatHost(a2aClient: settings);
        var ws = host.WorkspaceOf();
        PendingFile.Upsert(ws, new WatchRecord("w1", ChatHost.Target, "remote", "t1", "c1", null, DateTimeOffset.UtcNow, WatchSources.Poll,
            Undelivered: "[A2A result from \"remote\", task t1 (completed)]\nKept"));
        host.Faux.Enqueue(FauxTurn.Say("Delivered late").WithDelay(TimeSpan.FromSeconds(2)));
        var service = new A2AWatchHostedService(host.A2AWatchRegistry!, settings, DeliveryFor(host), NullLogger<A2AWatchHostedService>.Instance);
        var watch = Stopwatch.StartNew();
        await service.StartAsync(CancellationToken.None);
        Assert.True(watch.Elapsed < TimeSpan.FromSeconds(1), $"StartAsync took {watch.Elapsed}");
        await Wait.UntilAsync(() => host.Texts.Any(t => t.Contains("Delivered late")), "the kept result to be delivered", 15_000);
        await Wait.UntilAsync(() => PendingFile.Read(ws).Count == 0, "the record to be removed");
        await ((IHostedService)service).StopAsync(CancellationToken.None);
    }

    [Fact]
    public async Task BothOn_A2aServerSessionHasNoA2aTool_ChatHasIt()
    {
        var client = new A2AClientSettings { Enabled = true };
        var server = new A2AHostOptions { Enabled = true, PublicUrl = "https://bot.example" };
        await using var registry = new A2AWatchRegistry(TimeProvider.System, new HttpClientHandler());
        await using var host = new ChatHost((h, s) =>
        {
            s.AddLogging();
            s.AddSingleton(_ => h.Chats);
            s.AddSingleton(_ => h.Settings);
            TelegramA2AClientSetup.Register(s, client, server, registry); // as Program.cs: before the server setup
            TelegramA2ASetup.Register(s, h.Options, server);
        }, a2aClient: client);
        Assert.Contains((await host.ChatAsync()).Session.Tools, t => t.Name == "a2a");
        var pool = host.Services.GetRequiredService<A2ASessionPool>();
        var caller = new A2ACaller("travel-bot", "isolated", host.WorkspaceOf(), $"telegram-{Updates.AllowedUser}", Updates.AllowedUser.ToString());
        using var lease = await pool.AcquireAsync(caller, "ctx1", "task1", CancellationToken.None);
        Assert.DoesNotContain(lease.Session.Tools, t => t.Name == "a2a");
        Assert.NotEmpty(lease.Session.Tools);
    }

    [Fact]
    public async Task ClientExtension_CannotBeActivatedHostWide()
    {
        // As if discovered from an extensions folder: its target can't come from DI, so it never loads host-wide.
        await using var host = new ChatHost((_, s) => s.AddSingleton(new A2AClientSettings { Enabled = true }),
            core: o => o.Extensions.Add(ExtensionRegistration.FromType(typeof(A2AClientExtension), "test")));
        var chat = await host.ChatAsync();
        Assert.DoesNotContain(chat.Session.Tools, t => t.Name == "a2a");
        var registration = ExtensionRegistration.FromType(typeof(A2AClientExtension), "test");
        Assert.ThrowsAny<InvalidOperationException>(() => registration.Factory(host.Services));
    }
}