VolcanicAgenticCore/tests/Volcanic.AgenticCore.A2A.Server.Tests/SessionPoolTests.cs

using A2A;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Extensions.Options;
using Microsoft.Extensions.Time.Testing;
using Volcanic.AgenticCore.A2A.Server.Agent;
using Volcanic.AgenticCore.A2A.Server.Callers;
using Volcanic.AgenticCore.A2A.Server.Push;
using Volcanic.AgenticCore.A2A.Server.Tasks;
using Volcanic.AgenticCore.Features.Agent;
using Volcanic.AgenticCore.Features.Extensions;
using Volcanic.AgenticCore.Features.Processes;
using Volcanic.AgenticCore.Hosting;
using Volcanic.AgenticCore.Testing;

namespace Volcanic.AgenticCore.A2A.Server.Tests;

public sealed class SessionPoolTests : IAsyncDisposable
{
    private readonly string _root = Path.Combine(Path.GetTempPath(), "a2a-pool", Guid.NewGuid().ToString("N"));
    private readonly FauxProvider _faux = new();
    private readonly ServiceProvider _services;
    private readonly FakeTimeProvider _time = new(DateTimeOffset.Parse("2026-10-07T10:00:00Z"));
    private readonly A2AHostOptions _options = new() { SessionIdleTimeout = TimeSpan.FromMinutes(30) };
    private readonly A2ASessionPool _pool;
    private readonly A2ACaller _caller;

    public SessionPoolTests()
    {
        Directory.CreateDirectory(_root);
        _caller = new A2ACaller("alpha", "isolated", Path.Combine(_root, "ws"), "p", "1");
        var services = new ServiceCollection();
        services.AddAgenticCore(o =>
        {
            o.AgentDir = Path.Combine(_root, ".agent");
            o.DiscoverExtensionAssemblies = false;
            o.DefaultModel = FauxProvider.ModelKey;
            o.Extensions.Add(ExtensionRegistration.FromInstance(_faux.Extension));
        });
        _services = services.BuildServiceProvider();
        _pool = new A2ASessionPool(_services.GetRequiredService<IAgentSessionFactory>(), Options.Create(_options), _time,
            NullLogger<A2ASessionPool>.Instance);
    }

    public async ValueTask DisposeAsync()
    {
        await _pool.DisposeAsync();
        await _services.DisposeAsync();
        try { Directory.Delete(_root, true); } catch (Exception) { }
    }

    private A2AAgentHandler Handler() =>
        new(_caller, new WorkspaceTaskStore(_caller, _pool, A2APushSender.Disabled, _time, NullLogger<WorkspaceTaskStore>.Instance),
            _pool, Options.Create(_options), SecretRedactor.Empty, new A2AFileDownloader(new HttpClient()), TimeProvider.System,
            NullLogger<A2AAgentHandler>.Instance);

    private static RequestContext Context(string text, string taskId = "t1", string contextId = "c1") => new()
    {
        Message = new Message { Role = Role.User, MessageId = Guid.NewGuid().ToString("N"), Parts = [Part.FromText(text)] },
        TaskId = taskId,
        ContextId = contextId,
        StreamingResponse = false,
    };

    [Fact]
    public void Pool_Track_MakesTaskLiveBeforeAcquire()
    {
        var tracking = _pool.Track("t1", out var run);
        Assert.True(_pool.IsLive("t1"));
        Assert.Null(run.Entry);
        tracking.Dispose();
        Assert.True(_pool.IsLive("t1")); // the SDK may still be applying the final event
        _time.Advance(A2ASessionPool.FinishedGrace);
        Assert.False(_pool.IsLive("t1"));
    }

    [Fact]
    public void Pool_Track_FinishedRunIsReplacedByContinuation()
    {
        var first = _pool.Track("t1", out _);
        first.Dispose();
        using var second = _pool.Track("t1", out var run);
        _time.Advance(A2ASessionPool.FinishedGrace);
        Assert.True(_pool.IsLive("t1"));
        Assert.Null(run.FinishedAt);
    }

    [Fact]
    public void Pool_Track_UnfinishedRunOfSameTask_Refused()
    {
        using var first = _pool.Track("t1", out _);
        var ex = Assert.Throws<A2AException>(() => _pool.Track("t1", out _));
        Assert.Equal(A2AErrorCode.UnsupportedOperation, ex.ErrorCode);
        Assert.True(_pool.IsRunning("t1"));
    }

    [Fact]
    public async Task Pool_AbortAsync_FinishedRun_IsNoOp()
    {
        _pool.Track("t1", out var run).Dispose();
        await _pool.AbortAsync("t1"); // the run's token source is disposed; must not throw
        Assert.NotNull(run.FinishedAt);
    }

    [Fact]
    public async Task Pool_Acquire_SetsTrackedEntry()
    {
        using var tracking = _pool.Track("t1", out var run);
        using var lease = await _pool.AcquireAsync(_caller, "c1", "t1", CancellationToken.None);
        Assert.Same(lease.Session, run.Entry!.Session);
        Assert.True(_pool.IsBusy(A2ASessionPool.ScopeFor(_caller, "c1")));
    }

    [Fact]
    public async Task Pool_SecondAcquire_SameContext_Throws()
    {
        using var lease = await _pool.AcquireAsync(_caller, "c1", "t1", CancellationToken.None);
        var ex = await Assert.ThrowsAsync<A2AException>(() => _pool.AcquireAsync(_caller, "c1", "t2", CancellationToken.None));
        Assert.Equal(A2AErrorCode.UnsupportedOperation, ex.ErrorCode);
    }

    [Fact]
    public async Task Pool_AbortAsync_CancelsRunToken()
    {
        using var tracking = _pool.Track("t1", out var run);
        await _pool.AbortAsync("t1");
        Assert.True(run.Cancel.IsCancellationRequested);
    }

    [Fact]
    public async Task Pool_EvictedEntry_IsReopenedOnAcquire()
    {
        AgentSession first;
        using (var lease = await _pool.AcquireAsync(_caller, "c1", "t1", CancellationToken.None)) first = lease.Session;
        _time.Advance(TimeSpan.FromMinutes(31));
        await _pool.EvictIdleAsync();
        Assert.Equal(0, _pool.OpenCount);

        _faux.Enqueue(FauxTurn.Say("hello again"));
        using var again = await _pool.AcquireAsync(_caller, "c1", "t2", CancellationToken.None);
        Assert.NotSame(first, again.Session);
        await again.Session.PromptAsync("hi"); // a disposed session would throw ObjectDisposedException
    }

    [Fact]
    public async Task Pool_Eviction_SkipsBusyAndRecentSessions()
    {
        using var busy = await _pool.AcquireAsync(_caller, "busy", "t1", CancellationToken.None);
        using (await _pool.AcquireAsync(_caller, "idle", "t2", CancellationToken.None)) { }
        await _pool.EvictIdleAsync();
        Assert.Equal(2, _pool.OpenCount);
        _time.Advance(TimeSpan.FromMinutes(31));
        await _pool.EvictIdleAsync();
        Assert.Equal(1, _pool.OpenCount);
        Assert.True(_pool.IsBusy(A2ASessionPool.ScopeFor(_caller, "busy")));
    }

    [Fact]
    public async Task Handler_TaskIsLive_ForEveryQueuedEvent()
    {
        _faux.Enqueue(FauxTurn.Say("done here"));
        var queue = new AgentEventQueue();
        var exec = Task.Run(async () =>
        {
            try { await Handler().ExecuteAsync(Context("go"), queue, CancellationToken.None); }
            finally { queue.Complete(); }
        });
        var liveness = new List<bool>();
        var states = new List<TaskState>();
        await foreach (var e in queue)
        {
            liveness.Add(_pool.IsLive("t1"));
            if (e.StatusUpdate is { } s) states.Add(s.Status.State);
            if (e.Task is { } t) states.Add(t.Status.State);
        }
        await exec;
        Assert.NotEmpty(liveness);
        Assert.All(liveness, Assert.True);
        Assert.Equal(TaskState.Submitted, states[0]);
        Assert.Equal(TaskState.Completed, states[^1]);
        _time.Advance(A2ASessionPool.FinishedGrace);
        Assert.False(_pool.IsLive("t1"));
    }

    [Fact]
    public async Task Handler_HopLimit_FailsWhileTracked()
    {
        _options.MaxHops = -1;
        var queue = new AgentEventQueue();
        var exec = Task.Run(async () =>
        {
            try { await Handler().ExecuteAsync(Context("go"), queue, CancellationToken.None); }
            finally { queue.Complete(); }
        });
        TaskState? last = null;
        var allLive = true;
        await foreach (var e in queue)
        {
            allLive &= _pool.IsLive("t1");
            if (e.StatusUpdate is { } s) last = s.Status.State;
        }
        await exec;
        Assert.True(allLive);
        Assert.Equal(TaskState.Failed, last);
    }

    [Fact]
    public async Task Handler_Cancel_EndsCanceled()
    {
        _faux.Enqueue(FauxTurn.Say("slow").WithDelay(TimeSpan.FromSeconds(30)));
        var handler = Handler();
        var queue = new AgentEventQueue();
        var exec = Task.Run(async () =>
        {
            try { await handler.ExecuteAsync(Context("go"), queue, CancellationToken.None); }
            finally { queue.Complete(); }
        });
        var states = new List<TaskState>();
        var reader = Task.Run(async () =>
        {
            await foreach (var e in queue)
                if (e.StatusUpdate is { } s) lock (states) states.Add(s.Status.State);
        });
        var scope = A2ASessionPool.ScopeFor(_caller, "c1");
        while (!_pool.IsBusy(scope)) await Task.Delay(10);
        await Task.Delay(100); // let the model call start
        await handler.CancelAsync(Context("go"), new AgentEventQueue(), CancellationToken.None);
        await reader.WaitAsync(TimeSpan.FromSeconds(10));
        await exec.WaitAsync(TimeSpan.FromSeconds(10));
        Assert.Equal(TaskState.Canceled, states[^1]);
        Assert.DoesNotContain(TaskState.Completed, states);
    }

    [Fact]
    public async Task Handler_RequestTokenCancelledBeforeStart_EndsCanceledWithoutHanging()
    {
        using var cts = new CancellationTokenSource();
        await cts.CancelAsync();
        var queue = new AgentEventQueue();
        var exec = Task.Run(async () =>
        {
            try { await Handler().ExecuteAsync(Context("go"), queue, cts.Token); }
            finally { queue.Complete(); }
        });
        var states = new List<TaskState>();
        await foreach (var e in queue)
            if (e.StatusUpdate is { } s) states.Add(s.Status.State);
        await exec.WaitAsync(TimeSpan.FromSeconds(10));
        Assert.Equal([TaskState.Canceled], states);
    }
}