ClaudeDo/tests/ClaudeDo.Worker.Tests/External/HandoffMcpToolsTests.cs

using ClaudeDo.Data;
using ClaudeDo.Data.Models;
using ClaudeDo.Data.Repositories;
using ClaudeDo.Worker.External;
using ClaudeDo.Worker.Hub;
using ClaudeDo.Worker.Runner;
using ClaudeDo.Worker.Tests.Infrastructure;
using ClaudeDo.Worker.Usage;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.Logging.Abstractions;
using TaskStatus = ClaudeDo.Data.Models.TaskStatus;

namespace ClaudeDo.Worker.Tests.External;

public sealed class HandoffMcpToolsTests : IDisposable
{
    private readonly DbFixture _db = new();
    private readonly ClaudeDoDbContext _ctx;
    private readonly TaskRepository _tasks;
    private readonly ListRepository _lists;
    private readonly CapturingHubContext _hubContext = new();
    private readonly string _tempDir;

    public HandoffMcpToolsTests()
    {
        _ctx = _db.CreateContext();
        _tasks = new TaskRepository(_ctx);
        _lists = new ListRepository(_ctx);
        _tempDir = Path.Combine(Path.GetTempPath(), $"cd_handoffmcp_{Guid.NewGuid():N}");
        Directory.CreateDirectory(_tempDir);
    }

    public void Dispose()
    {
        _ctx.Dispose();
        _db.Dispose();
        try { Directory.Delete(_tempDir, true); } catch { }
    }

    // No usage snapshot -> HandlerUsageGuard never withholds, which is the behaviour every test
    // below assumes. The withholding path has its own tests in HandlerUsageGuardTests.
    private HandoffMcpTools BuildSut(HandoffRoundTracker? tracker = null)
    {
        var rounds = tracker ?? new HandoffRoundTracker();
        var guard = new HandlerUsageGuard(
            _db.CreateFactory(), new UsageState(),
            new HandlerRoundState(Path.Combine(_tempDir, "handler-rounds.json")), rounds,
            new HubBroadcaster(_hubContext), NullLogger<HandlerUsageGuard>.Instance);
        return new HandoffMcpTools(_tasks, new HubBroadcaster(_hubContext), rounds, _db.CreateFactory(), guard);
    }

    /// The cap the SUT reads comes from app_settings, so a test that wants a different cap has to
    /// write it there. GetAsync is AsNoTracking, hence the separate tracked query.
    private async Task SetMaxHandoffsAsync(int value)
    {
        await new AppSettingsRepository(_ctx).GetAsync();
        var row = await _ctx.AppSettings.FirstAsync(s => s.Id == AppSettingsEntity.SingletonId);
        row.ListHandlerMaxHandoffs = value;
        await _ctx.SaveChangesAsync();
    }

    private async Task<TaskEntity> SeedTaskAsync(string listId, TaskStatus status = TaskStatus.Idle, string title = "t")
    {
        var task = new TaskEntity
        {
            Id = Guid.NewGuid().ToString(), ListId = listId, Title = title,
            Status = status, CreatedAt = DateTime.UtcNow, CommitType = "chore",
        };
        await _tasks.AddAsync(task);
        return task;
    }

    private async Task<string> SeedListAsync()
    {
        var listId = Guid.NewGuid().ToString();
        await _lists.AddAsync(new ListEntity { Id = listId, Name = "L", CreatedAt = DateTime.UtcNow });
        return listId;
    }

    [Fact]
    public async Task HandoffListHandler_ValidIds_BroadcastsAndReturnsCount()
    {
        var listId = await SeedListAsync();
        var handlerTask = await SeedTaskAsync(listId, title: "List handler: Alpha");
        var survivor = await SeedTaskAsync(listId, TaskStatus.WaitingForReview, title: "Survivor");

        var sut = BuildSut();
        var result = await sut.HandoffListHandler(handlerTask.Id, new[] { survivor.Id }, cancellationToken: CancellationToken.None);

        Assert.True(result.Requested);
        Assert.Equal(handlerTask.Id, result.TaskId);
        Assert.Equal(1, result.SurvivingCount);
        Assert.Equal("wait", result.NextPhase);

        var call = Assert.Single(_hubContext.Proxy.Calls);
        Assert.Equal("HandoffRequested", call.Method);
        Assert.Equal(handlerTask.Id, call.Args[0]);
        Assert.Equal("wait", call.Args[2]);
    }

    [Fact]
    public async Task HandoffListHandler_NextPhaseDefaultsToWait()
    {
        var listId = await SeedListAsync();
        var handlerTask = await SeedTaskAsync(listId);
        var survivor = await SeedTaskAsync(listId, TaskStatus.WaitingForReview, title: "Survivor");

        var sut = BuildSut();
        var result = await sut.HandoffListHandler(handlerTask.Id, new[] { survivor.Id }, cancellationToken: CancellationToken.None);

        Assert.Equal("wait", result.NextPhase);
    }

    [Theory]
    [InlineData("wait")]
    [InlineData("merge")]
    [InlineData("wait_final")]
    [InlineData("merge_final")]
    public async Task HandoffListHandler_AcceptsEveryKnownPhase(string phase)
    {
        var listId = await SeedListAsync();
        var handlerTask = await SeedTaskAsync(listId);
        var survivor = await SeedTaskAsync(listId, TaskStatus.WaitingForReview, title: "Survivor");

        var sut = BuildSut();
        var result = await sut.HandoffListHandler(handlerTask.Id, new[] { survivor.Id }, phase, CancellationToken.None);

        Assert.Equal(phase, result.NextPhase);
    }

    [Fact]
    public async Task HandoffListHandler_UnknownPhase_Throws()
    {
        var listId = await SeedListAsync();
        var handlerTask = await SeedTaskAsync(listId);
        var survivor = await SeedTaskAsync(listId, TaskStatus.WaitingForReview, title: "Survivor");

        var sut = BuildSut();
        var ex = await Assert.ThrowsAsync<ArgumentException>(() =>
            sut.HandoffListHandler(handlerTask.Id, new[] { survivor.Id }, "phase-9", CancellationToken.None));
        Assert.Contains("phase-9", ex.Message);
    }

    [Fact]
    public async Task HandoffListHandler_EmptySurvivingIds_Throws()
    {
        var listId = await SeedListAsync();
        var handlerTask = await SeedTaskAsync(listId);

        var sut = BuildSut();
        await Assert.ThrowsAsync<InvalidOperationException>(() =>
            sut.HandoffListHandler(handlerTask.Id, Array.Empty<string>(), cancellationToken: CancellationToken.None));
    }

    [Fact]
    public async Task HandoffListHandler_UnknownHandlerTask_Throws()
    {
        var sut = BuildSut();
        await Assert.ThrowsAsync<InvalidOperationException>(() =>
            sut.HandoffListHandler("missing", new[] { "x" }, cancellationToken: CancellationToken.None));
    }

    [Fact]
    public async Task HandoffListHandler_UnknownSurvivingTask_Throws()
    {
        var listId = await SeedListAsync();
        var handlerTask = await SeedTaskAsync(listId);

        var sut = BuildSut();
        await Assert.ThrowsAsync<InvalidOperationException>(() =>
            sut.HandoffListHandler(handlerTask.Id, new[] { "missing" }, cancellationToken: CancellationToken.None));
    }

    [Fact]
    public async Task HandoffListHandler_NormalWaitMergeWaitFinalMergeFinalSequence_PassesUntouched()
    {
        var listId = await SeedListAsync();
        var handlerTask = await SeedTaskAsync(listId);
        var survivor = await SeedTaskAsync(listId, TaskStatus.WaitingForReview, title: "Survivor");
        var tracker = new HandoffRoundTracker();
        var survivors = new[] { survivor.Id };

        foreach (var phase in new[] { "wait", "merge", "wait_final", "merge_final" })
        {
            var sut = BuildSut(tracker);
            var result = await sut.HandoffListHandler(handlerTask.Id, survivors, phase, CancellationToken.None);
            Assert.False(result.Coerced);
            Assert.Equal(phase, result.NextPhase);
        }
    }

    [Fact]
    public async Task HandoffListHandler_PastCap_CoercesNonFinalPhaseToItsFinalCounterpart()
    {
        var listId = await SeedListAsync();
        var handlerTask = await SeedTaskAsync(listId);
        var survivor = await SeedTaskAsync(listId, TaskStatus.WaitingForReview, title: "Survivor");
        var tracker = new HandoffRoundTracker();
        var survivors = new[] { survivor.Id };

        foreach (var phase in new[] { "wait", "merge", "wait", "merge" })
        {
            var sut = BuildSut(tracker);
            var result = await sut.HandoffListHandler(handlerTask.Id, survivors, phase, CancellationToken.None);
            Assert.False(result.Coerced);
        }

        var fifthSut = BuildSut(tracker);
        var fifth = await fifthSut.HandoffListHandler(handlerTask.Id, survivors, "wait", CancellationToken.None);

        Assert.True(fifth.Coerced);
        Assert.Equal("wait_final", fifth.NextPhase);
    }

    [Fact]
    public async Task HandoffListHandler_AfterMergeFinal_AnyFurtherHandoffThrows()
    {
        var listId = await SeedListAsync();
        var handlerTask = await SeedTaskAsync(listId);
        var survivor = await SeedTaskAsync(listId, TaskStatus.WaitingForReview, title: "Survivor");
        var tracker = new HandoffRoundTracker();
        var survivors = new[] { survivor.Id };

        var firstSut = BuildSut(tracker);
        await firstSut.HandoffListHandler(handlerTask.Id, survivors, "merge_final", CancellationToken.None);

        var secondSut = BuildSut(tracker);
        await Assert.ThrowsAsync<InvalidOperationException>(() =>
            secondSut.HandoffListHandler(handlerTask.Id, survivors, "wait", CancellationToken.None));
    }

    [Theory]
    [InlineData(1, 3)]
    [InlineData(2, 2)]
    [InlineData(4, 0)]
    public void HandoffRoundTracker_GetState_ReportsCountAfterNRegistrations(int registrations, int expectedRemaining)
    {
        var tracker = new HandoffRoundTracker();
        var taskId = Guid.NewGuid().ToString();

        for (var i = 0; i < registrations; i++)
            tracker.Register(taskId, "wait");

        var (count, sealedState) = tracker.GetState(taskId);
        Assert.Equal(registrations, count);
        Assert.Equal(expectedRemaining, HandoffRoundTracker.DefaultMaxHandoffs - count);
        Assert.False(sealedState);
    }

    [Fact]
    public void HandoffRoundTracker_GetState_UnregisteredTask_ReturnsZero()
    {
        var tracker = new HandoffRoundTracker();
        var (count, sealedState) = tracker.GetState(Guid.NewGuid().ToString());
        Assert.Equal(0, count);
        Assert.False(sealedState);
    }

    [Fact]
    public void HandoffRoundTracker_Register_DoesNotMutateStateVisibleFromGetState_BeyondItsOwnCall()
    {
        var tracker = new HandoffRoundTracker();
        var taskId = Guid.NewGuid().ToString();

        tracker.Register(taskId, "wait");
        var (beforeCount, _) = tracker.GetState(taskId);
        Assert.Equal(1, beforeCount);

        // GetState itself must be read-only -- calling it repeatedly must not advance the round.
        tracker.GetState(taskId);
        tracker.GetState(taskId);
        var (afterCount, _) = tracker.GetState(taskId);
        Assert.Equal(beforeCount, afterCount);
    }

    [Fact]
    public void HandoffRoundTracker_RejectedHandoff_ErrorNamesTheNextAction()
    {
        var tracker = new HandoffRoundTracker();
        var taskId = Guid.NewGuid().ToString();
        tracker.Register(taskId, "merge_final");

        var result = tracker.Register(taskId, "wait");

        Assert.True(result.Rejected);
        Assert.Contains("new handler run", result.Error, StringComparison.OrdinalIgnoreCase);
    }

    [Fact]
    public async Task HandoffListHandler_RaisedCapSetting_AllowsMoreRoundsBeforeCoercing()
    {
        await SetMaxHandoffsAsync(8);

        var listId = await SeedListAsync();
        var handlerTask = await SeedTaskAsync(listId);
        var survivor = await SeedTaskAsync(listId, TaskStatus.WaitingForReview, title: "Survivor");
        var tracker = new HandoffRoundTracker();
        var survivors = new[] { survivor.Id };

        // The 5th handoff is where the default cap of 4 would coerce; with 8 configured it must not.
        foreach (var phase in new[] { "wait", "merge", "wait", "merge", "wait", "merge", "wait", "merge" })
        {
            var sut = BuildSut(tracker);
            var result = await sut.HandoffListHandler(handlerTask.Id, survivors, phase, CancellationToken.None);
            Assert.False(result.Coerced);
            Assert.Equal(phase, result.NextPhase);
        }

        var ninthSut = BuildSut(tracker);
        var ninth = await ninthSut.HandoffListHandler(handlerTask.Id, survivors, "wait", CancellationToken.None);

        Assert.True(ninth.Coerced);
        Assert.Equal("wait_final", ninth.NextPhase);
    }

    [Fact]
    public async Task HandoffListHandler_CapOfZero_IsClampedToOne_NotUnbounded()
    {
        // 0 must not read as "no cap" — that would remove the only bound on a runaway chain.
        await SetMaxHandoffsAsync(0);

        var listId = await SeedListAsync();
        var handlerTask = await SeedTaskAsync(listId);
        var survivor = await SeedTaskAsync(listId, TaskStatus.WaitingForReview, title: "Survivor");
        var tracker = new HandoffRoundTracker();
        var survivors = new[] { survivor.Id };

        var firstSut = BuildSut(tracker);
        var first = await firstSut.HandoffListHandler(handlerTask.Id, survivors, "wait", CancellationToken.None);
        Assert.False(first.Coerced);

        var secondSut = BuildSut(tracker);
        var second = await secondSut.HandoffListHandler(handlerTask.Id, survivors, "merge", CancellationToken.None);
        Assert.True(second.Coerced);
        Assert.Equal("merge_final", second.NextPhase);
    }

    /// GetHandoffStatus must report against the SAME cap HandoffListHandler enforces. Reading it
    /// off the constant would understate a raised budget and make the tool lie about how much room
    /// the chain has left.
    [Fact]
    public async Task GetHandoffStatus_RaisedCap_ReportsConfiguredRemaining()
    {
        await SetMaxHandoffsAsync(40);

        var listId = await SeedListAsync();
        var handlerTask = await SeedTaskAsync(listId);
        var survivor = await SeedTaskAsync(listId, TaskStatus.WaitingForReview, title: "Survivor");
        var tracker = new HandoffRoundTracker();
        var survivors = new[] { survivor.Id };

        await BuildSut(tracker).HandoffListHandler(handlerTask.Id, survivors, "wait", CancellationToken.None);
        await BuildSut(tracker).HandoffListHandler(handlerTask.Id, survivors, "merge", CancellationToken.None);

        var status = await BuildSut(tracker).GetHandoffStatus(handlerTask.Id, CancellationToken.None);

        Assert.Equal(2, status.Round);
        Assert.Equal(38, status.Remaining);
        Assert.False(status.Final);
    }

    [Fact]
    public async Task HandoffListHandler_TwoHandlerTasks_DoNotShareACounter()
    {
        var listId = await SeedListAsync();
        var handlerA = await SeedTaskAsync(listId, title: "Handler A");
        var handlerB = await SeedTaskAsync(listId, title: "Handler B");
        var survivorA = await SeedTaskAsync(listId, TaskStatus.WaitingForReview, title: "Survivor A");
        var survivorB = await SeedTaskAsync(listId, TaskStatus.WaitingForReview, title: "Survivor B");
        var tracker = new HandoffRoundTracker();

        foreach (var phase in new[] { "wait", "merge", "wait", "merge" })
        {
            var sut = BuildSut(tracker);
            await sut.HandoffListHandler(handlerA.Id, new[] { survivorA.Id }, phase, CancellationToken.None);
        }

        var sutB = BuildSut(tracker);
        var resultB = await sutB.HandoffListHandler(handlerB.Id, new[] { survivorB.Id }, "wait", CancellationToken.None);

        Assert.False(resultB.Coerced);
        Assert.Equal("wait", resultB.NextPhase);
    }

    [Fact]
    public async Task GetHandoffStatus_NeverHandedOff_ReportsZeroRoundFullBudget()
    {
        var listId = await SeedListAsync();
        var handlerTask = await SeedTaskAsync(listId);

        var sut = BuildSut();
        var status = await sut.GetHandoffStatus(handlerTask.Id, CancellationToken.None);

        Assert.Equal(handlerTask.Id, status.TaskId);
        Assert.Equal(handlerTask.Number, status.Number);
        Assert.Equal(0, status.Round);
        Assert.Equal(HandoffRoundTracker.DefaultMaxHandoffs, status.Remaining);
        Assert.False(status.Final);
    }

    [Fact]
    public async Task GetHandoffStatus_AfterHandoffs_ReflectsTrackerState()
    {
        var listId = await SeedListAsync();
        var handlerTask = await SeedTaskAsync(listId);
        var survivor = await SeedTaskAsync(listId, TaskStatus.WaitingForReview, title: "Survivor");
        var tracker = new HandoffRoundTracker();
        var survivors = new[] { survivor.Id };

        await BuildSut(tracker).HandoffListHandler(handlerTask.Id, survivors, "wait", CancellationToken.None);
        await BuildSut(tracker).HandoffListHandler(handlerTask.Id, survivors, "merge", CancellationToken.None);

        var status = await BuildSut(tracker).GetHandoffStatus(handlerTask.Id, CancellationToken.None);

        Assert.Equal(2, status.Round);
        Assert.Equal(HandoffRoundTracker.DefaultMaxHandoffs - 2, status.Remaining);
        Assert.False(status.Final);
    }

    [Fact]
    public async Task GetHandoffStatus_AfterMergeFinal_ReportsFinal()
    {
        var listId = await SeedListAsync();
        var handlerTask = await SeedTaskAsync(listId);
        var survivor = await SeedTaskAsync(listId, TaskStatus.WaitingForReview, title: "Survivor");
        var tracker = new HandoffRoundTracker();
        var survivors = new[] { survivor.Id };

        await BuildSut(tracker).HandoffListHandler(handlerTask.Id, survivors, "merge_final", CancellationToken.None);

        var status = await BuildSut(tracker).GetHandoffStatus(handlerTask.Id, CancellationToken.None);

        Assert.True(status.Final);
        Assert.Equal(1, status.Round);
    }

    [Fact]
    public async Task GetHandoffStatus_UnknownTask_Throws()
    {
        var sut = BuildSut();
        await Assert.ThrowsAsync<InvalidOperationException>(() =>
            sut.GetHandoffStatus("missing", CancellationToken.None));
    }
}