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

using System.Diagnostics;
using ClaudeDo.Data;
using ClaudeDo.Data.Models;
using ClaudeDo.Data.Repositories;
using ClaudeDo.Worker.External;
using ClaudeDo.Worker.Tests.Infrastructure;
using ModelContextProtocol;
using TaskStatus = ClaudeDo.Data.Models.TaskStatus;

namespace ClaudeDo.Worker.Tests.External;

public sealed class TaskWaitMcpToolsTests : IDisposable
{
    private readonly DbFixture _db = new();
    private readonly ClaudeDoDbContext _ctx;
    private readonly TaskRepository _tasks;
    private readonly ListRepository _lists;

    public TaskWaitMcpToolsTests()
    {
        _ctx = _db.CreateContext();
        _tasks = new TaskRepository(_ctx);
        _lists = new ListRepository(_ctx);
    }

    public void Dispose() { _ctx.Dispose(); _db.Dispose(); }

    private TaskWaitMcpTools BuildSut() => new(_db.CreateFactory());

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

    [Fact]
    public async Task WaitForTaskChange_AlreadyOutOfQueuedRunning_ReturnsImmediately()
    {
        var task = await SeedTaskAsync(TaskStatus.WaitingForReview);
        var sut = BuildSut();
        var sw = Stopwatch.StartNew();

        var result = await sut.WaitForTaskChange([task.Id], timeoutSeconds: 30, cancellationToken: CancellationToken.None);

        sw.Stop();
        Assert.False(result.TimedOut);
        Assert.Equal(task.Id, Assert.Single(result.Changed).TaskId);
        Assert.Equal("WaitingForReview", result.Changed[0].Status);
        Assert.True(sw.Elapsed < TimeSpan.FromSeconds(2), $"took {sw.Elapsed}");
    }

    [Fact]
    public async Task WaitForTaskChange_UnknownId_ReturnsImmediatelyAsNotFound()
    {
        var sut = BuildSut();
        var sw = Stopwatch.StartNew();

        var result = await sut.WaitForTaskChange(["missing-id"], timeoutSeconds: 30, cancellationToken: CancellationToken.None);

        sw.Stop();
        Assert.False(result.TimedOut);
        Assert.Equal("NotFound", Assert.Single(result.Changed).Status);
        Assert.True(sw.Elapsed < TimeSpan.FromSeconds(2), $"took {sw.Elapsed}");
    }

    [Fact]
    public async Task WaitForTaskChange_StatusChangesWhileWaiting_ReturnsBeforeTimeout()
    {
        var task = await SeedTaskAsync(TaskStatus.Running);
        var sut = BuildSut();
        var sw = Stopwatch.StartNew();

        var waitTask = sut.WaitForTaskChange([task.Id], timeoutSeconds: 10, cancellationToken: CancellationToken.None);

        await Task.Delay(150);
        // Simulate the status change a broadcast would announce, via a separate context
        // (mirrors what TaskStateService does from a different scope/process).
        await using (var writeCtx = _db.CreateContext())
        {
            var writeRepo = new TaskRepository(writeCtx);
            var loaded = await writeRepo.GetByIdAsync(task.Id);
            loaded!.Status = TaskStatus.Done;
            await writeRepo.UpdateAsync(loaded);
        }

        var result = await waitTask;
        sw.Stop();

        Assert.False(result.TimedOut);
        Assert.Equal("Done", Assert.Single(result.Changed).Status);
        Assert.True(sw.Elapsed < TimeSpan.FromSeconds(5), $"took {sw.Elapsed}");
    }

    [Fact]
    public async Task WaitForTaskChange_NoChange_TimesOut()
    {
        var task = await SeedTaskAsync(TaskStatus.Running);
        var sut = BuildSut();
        var sw = Stopwatch.StartNew();

        var result = await sut.WaitForTaskChange([task.Id], timeoutSeconds: 1, cancellationToken: CancellationToken.None);

        sw.Stop();
        Assert.True(result.TimedOut);
        Assert.Empty(result.Changed);
        Assert.True(sw.Elapsed >= TimeSpan.FromMilliseconds(900), $"took {sw.Elapsed}");
    }

    [Fact]
    public async Task WaitForTaskChange_EmptyTaskIds_Throws()
    {
        var sut = BuildSut();
        await Assert.ThrowsAsync<ArgumentException>(() =>
            sut.WaitForTaskChange([], timeoutSeconds: 5, cancellationToken: CancellationToken.None));
    }

    [Fact]
    public async Task WaitForTaskChange_WaitingForChildren_DefaultBehavior_ReturnsImmediately()
    {
        var task = await SeedTaskAsync(TaskStatus.WaitingForChildren);
        var sut = BuildSut();
        var sw = Stopwatch.StartNew();

        var result = await sut.WaitForTaskChange([task.Id], timeoutSeconds: 30, cancellationToken: CancellationToken.None);

        sw.Stop();
        Assert.False(result.TimedOut);
        Assert.Equal("WaitingForChildren", Assert.Single(result.Changed).Status);
        Assert.True(sw.Elapsed < TimeSpan.FromSeconds(2), $"took {sw.Elapsed}");
    }

    [Fact]
    public async Task WaitForTaskChange_TreatWaitingForChildrenAsBusy_DoesNotReturnImmediately_TimesOut()
    {
        var task = await SeedTaskAsync(TaskStatus.WaitingForChildren);
        var sut = BuildSut();
        var sw = Stopwatch.StartNew();

        var result = await sut.WaitForTaskChange(
            [task.Id], timeoutSeconds: 1, treatWaitingForChildrenAsBusy: true, cancellationToken: CancellationToken.None);

        sw.Stop();
        Assert.True(result.TimedOut);
        Assert.Empty(result.Changed);
        Assert.True(sw.Elapsed >= TimeSpan.FromMilliseconds(900), $"took {sw.Elapsed}");
    }

    [Fact]
    public async Task WaitForTaskChange_TreatWaitingForChildrenAsBusy_ReturnsWhenParentReachesWaitingForReview()
    {
        var task = await SeedTaskAsync(TaskStatus.WaitingForChildren);
        var sut = BuildSut();
        var sw = Stopwatch.StartNew();

        var waitTask = sut.WaitForTaskChange(
            [task.Id], timeoutSeconds: 30, treatWaitingForChildrenAsBusy: true, cancellationToken: CancellationToken.None);

        await Task.Delay(150);
        await using (var writeCtx = _db.CreateContext())
        {
            var writeRepo = new TaskRepository(writeCtx);
            var loaded = await writeRepo.GetByIdAsync(task.Id);
            loaded!.Status = TaskStatus.WaitingForReview;
            await writeRepo.UpdateAsync(loaded);
        }

        var result = await waitTask;
        sw.Stop();

        // Bound is generous (well under timeoutSeconds) purely to prove this didn't just
        // coincidentally land on the timeout path -- correctness is already covered by
        // Assert.False(result.TimedOut) above; this is not a performance assertion.
        Assert.False(result.TimedOut);
        Assert.Equal("WaitingForReview", Assert.Single(result.Changed).Status);
        Assert.True(sw.Elapsed < TimeSpan.FromSeconds(25), $"took {sw.Elapsed}");
    }

    [Fact]
    public async Task WaitForTaskChange_TreatWaitingForChildrenAsBusy_UnknownId_StillReturnsImmediatelyAsNotFound()
    {
        var sut = BuildSut();
        var sw = Stopwatch.StartNew();

        var result = await sut.WaitForTaskChange(
            ["missing-id"], timeoutSeconds: 30, treatWaitingForChildrenAsBusy: true, cancellationToken: CancellationToken.None);

        sw.Stop();
        Assert.False(result.TimedOut);
        Assert.Equal("NotFound", Assert.Single(result.Changed).Status);
        Assert.True(sw.Elapsed < TimeSpan.FromSeconds(2), $"took {sw.Elapsed}");
    }

    [Fact]
    public async Task WaitForTaskChange_QueuedWithBlockedByTaskId_KeepsWaiting_BecauseTheChainResolvesItself()
    {
        // A planning-chain link clears automatically once the predecessor goes terminal, so it
        // must NOT short-circuit the wait -- otherwise waiting on a queued fan-out returns
        // instantly for every chained child and the caller burns a turn per poll.
        var predecessor = await SeedTaskAsync(TaskStatus.Queued);
        var task = await SeedTaskAsync(TaskStatus.Queued);
        task.BlockedByTaskId = predecessor.Id;
        await _tasks.UpdateAsync(task);
        var sut = BuildSut();

        var result = await sut.WaitForTaskChange([task.Id], timeoutSeconds: 1, cancellationToken: CancellationToken.None);

        Assert.True(result.TimedOut);
        Assert.Empty(result.Changed);
    }

    [Fact]
    public async Task WaitForTaskChange_QueuedWithUnmetDependsOn_ReportsBlockedImmediately_InsteadOfTimingOut()
    {
        var dependency = await SeedTaskAsync(TaskStatus.Idle);
        var task = await SeedTaskAsync(TaskStatus.Queued);
        await _ctx.SeedDependsOnAsync(task.Id, dependency.Id);
        var sut = BuildSut();
        var sw = Stopwatch.StartNew();

        var result = await sut.WaitForTaskChange([task.Id], timeoutSeconds: 30, cancellationToken: CancellationToken.None);

        sw.Stop();
        Assert.False(result.TimedOut);
        var change = Assert.Single(result.Changed);
        Assert.Equal("Blocked", change.Status);
        Assert.Contains(dependency.Id, change.BlockedReason);
        Assert.True(sw.Elapsed < TimeSpan.FromSeconds(2), $"took {sw.Elapsed}");
    }

    [Fact]
    public async Task WaitForTaskChange_QueuedWithDoneDependsOn_IsNotBlocked_StillWaitsAsBusy()
    {
        var dependency = await SeedTaskAsync(TaskStatus.Done);
        var task = await SeedTaskAsync(TaskStatus.Queued);
        await _ctx.SeedDependsOnAsync(task.Id, dependency.Id);
        var sut = BuildSut();
        var sw = Stopwatch.StartNew();

        var result = await sut.WaitForTaskChange([task.Id], timeoutSeconds: 1, cancellationToken: CancellationToken.None);

        sw.Stop();
        Assert.True(result.TimedOut);
        Assert.Empty(result.Changed);
        Assert.True(sw.Elapsed >= TimeSpan.FromMilliseconds(900), $"took {sw.Elapsed}");
    }

    [Fact]
    public async Task WaitForTaskChange_UnknownNumberHashForm_ReturnsImmediatelyAsNotFound()
    {
        var sut = BuildSut();
        var sw = Stopwatch.StartNew();

        var result = await sut.WaitForTaskChange(["#999999"], timeoutSeconds: 30, cancellationToken: CancellationToken.None);

        sw.Stop();
        Assert.False(result.TimedOut);
        Assert.Equal("NotFound", Assert.Single(result.Changed).Status);
        Assert.True(sw.Elapsed < TimeSpan.FromSeconds(2), $"took {sw.Elapsed}");
    }

    [Fact]
    public async Task WaitForTaskChange_UnknownBareNumber_ReturnsImmediatelyAsNotFound()
    {
        var sut = BuildSut();
        var sw = Stopwatch.StartNew();

        var result = await sut.WaitForTaskChange(["999999"], timeoutSeconds: 30, cancellationToken: CancellationToken.None);

        sw.Stop();
        Assert.False(result.TimedOut);
        Assert.Equal("NotFound", Assert.Single(result.Changed).Status);
        Assert.True(sw.Elapsed < TimeSpan.FromSeconds(2), $"took {sw.Elapsed}");
    }

    [Fact]
    public async Task WaitForTaskChange_MixOfValidAndUnknownNumber_ReportsBothCorrectly()
    {
        var task = await SeedTaskAsync(TaskStatus.WaitingForReview);
        var sut = BuildSut();
        var sw = Stopwatch.StartNew();

        var result = await sut.WaitForTaskChange(
            [$"#{task.Number}", "#999999"], timeoutSeconds: 30, cancellationToken: CancellationToken.None);

        sw.Stop();
        Assert.False(result.TimedOut);
        Assert.Equal(2, result.Changed.Count);
        Assert.Contains(result.Changed, c => c.TaskId == task.Id && c.Status == "WaitingForReview");
        Assert.Contains(result.Changed, c => c.TaskId == "#999999" && c.Status == "NotFound");
        Assert.True(sw.Elapsed < TimeSpan.FromSeconds(2), $"took {sw.Elapsed}");
    }

    [Fact]
    public void MaxTimeoutSeconds_StaysComfortablyUnderMcpToolTimeout()
    {
        // ClaudeProcess / InteractiveLaunchSpecService set MCP_TOOL_TIMEOUT=930000ms.
        Assert.True(TaskWaitMcpTools.MaxTimeoutSeconds < 930);
    }

    [Fact]
    public async Task WaitForTaskChange_LongWait_ReportsProgressWellUnderClaudeCodeIdleTimeout()
    {
        var original = TaskWaitMcpTools.ProgressReportInterval;
        TaskWaitMcpTools.ProgressReportInterval = TimeSpan.FromMilliseconds(200);
        try
        {
            var task = await SeedTaskAsync(TaskStatus.Running);
            var sut = BuildSut();
            var reports = new List<ProgressNotificationValue>();
            var progress = new Progress<ProgressNotificationValue>(reports.Add);

            var result = await sut.WaitForTaskChange(
                [task.Id], timeoutSeconds: 1, progress: progress, cancellationToken: CancellationToken.None);

            Assert.True(result.TimedOut);
            // Progress<T> marshals via the SynchronizationContext captured at construction; give
            // any queued callbacks a beat to run before asserting on `reports`.
            await Task.Delay(200);
            Assert.NotEmpty(reports);
        }
        finally
        {
            TaskWaitMcpTools.ProgressReportInterval = original;
        }
    }

    [Fact]
    public async Task WaitForTaskChange_NoProgressHandlerPassed_DoesNotThrow()
    {
        var task = await SeedTaskAsync(TaskStatus.Running);
        var sut = BuildSut();

        var result = await sut.WaitForTaskChange(
            [task.Id], timeoutSeconds: 1, progress: null, cancellationToken: CancellationToken.None);

        Assert.True(result.TimedOut);
    }
}