ClaudeDo/tests/ClaudeDo.Worker.Tests/Runner/ContinueAsyncExceptionTests.cs

using ClaudeDo.Data;
using ClaudeDo.Data.Models;
using ClaudeDo.Data.Repositories;
using ClaudeDo.Worker.Config;
using ClaudeDo.Worker.Hub;
using ClaudeDo.Worker.Runner;
using ClaudeDo.Worker.Skills;
using ClaudeDo.Worker.Tests.Infrastructure;
using Microsoft.Extensions.Logging.Abstractions;
using TaskStatus = ClaudeDo.Data.Models.TaskStatus;

namespace ClaudeDo.Worker.Tests.Runner;

/// <summary>
/// Verifies that ContinueAsync wraps RunOnceAsync exceptions so the task
/// is never left stuck in Running status on an unexpected error.
/// </summary>
public sealed class ContinueAsyncExceptionTests : IDisposable
{
    private readonly DbFixture _db = new();
    private readonly string _tempDir;
    private readonly WorkerConfig _cfg;

    public ContinueAsyncExceptionTests()
    {
        _tempDir = Path.Combine(Path.GetTempPath(), $"cd_continue_{Guid.NewGuid():N}");
        Directory.CreateDirectory(_tempDir);
        _cfg = new WorkerConfig { SandboxRoot = _tempDir, LogRoot = _tempDir };
    }

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

    private TaskRunner BuildRunner(
        IClaudeProcess claude, ClaudeDoDbContext ctx,
        ISessionSkillSeeder? skillSeeder = null, TaskRunTokenRegistry? tokens = null, WorkerConfig? cfg = null)
    {
        var dbFactory = _db.CreateFactory();
        var broadcaster = new HubBroadcaster(new CapturingHubContext());
        var state = TaskStateServiceBuilder.Build(dbFactory).State;
        var effectiveCfg = cfg ?? _cfg;
        var wt = new WorktreeManager(new ClaudeDo.Data.Git.GitService(), dbFactory, effectiveCfg, NullLogger<WorktreeManager>.Instance);
        return new TaskRunner(claude, dbFactory, broadcaster, wt, new ClaudeArgsBuilder(), effectiveCfg,
            NullLogger<TaskRunner>.Instance, state, tokens ?? new TaskRunTokenRegistry(), new AttachmentStore(),
            skillSeeder ?? new FakeSessionSkillSeeder(), new FakeTranscriptUsageReader());
    }

    private async Task<string> SeedContinuableTaskAsync(string sessionId)
    {
        string listId, taskId;
        using var ctx = _db.CreateContext();
        listId = Guid.NewGuid().ToString();
        ctx.Lists.Add(new ListEntity { Id = listId, Name = "L", WorkingDir = null, CreatedAt = DateTime.UtcNow });

        taskId = Guid.NewGuid().ToString();
        ctx.Tasks.Add(new TaskEntity
        {
            Id = taskId,
            ListId = listId,
            Title = "Continue me",
            Status = TaskStatus.WaitingForReview,
            CreatedAt = DateTime.UtcNow,
        });
        await ctx.SaveChangesAsync();

        await new TaskRunRepository(ctx).AddAsync(new TaskRunEntity
        {
            Id = Guid.NewGuid().ToString(),
            TaskId = taskId,
            RunNumber = 1,
            IsRetry = false,
            Prompt = "original prompt",
            SessionId = sessionId,
            StartedAt = DateTime.UtcNow.AddMinutes(-5),
            FinishedAt = DateTime.UtcNow.AddMinutes(-1),
            ExitCode = 0,
            ResultMarkdown = "first result",
        });
        return taskId;
    }

    // Bug: SeedAsync used to run after the Running claim but before ContinueAsync's own
    // try/catch started, so a throw here propagated straight out of ContinueAsync to the
    // caller (OverrideSlotService.RunContinueInSlotAsync), which only logs — leaving the task
    // stuck Running forever.
    [Fact]
    public async Task ContinueAsync_SkillSeederThrowsAfterClaim_MarksTaskFailed_NotStuckRunning()
    {
        var taskId = await SeedContinuableTaskAsync("sess-seed-throws");
        var throwingSeeder = new FakeSessionSkillSeeder { ThrowOnSeed = new InvalidOperationException("skill seed exploded") };

        using var ctx2 = _db.CreateContext();
        var runner = BuildRunner(new FakeClaudeProcess(), ctx2, skillSeeder: throwingSeeder);

        await runner.ContinueAsync(taskId, "please continue", "slot-1", CancellationToken.None);

        using var verify = _db.CreateContext();
        var task = await new TaskRepository(verify).GetByIdAsync(taskId);
        Assert.NotNull(task);
        Assert.Equal(TaskStatus.Failed, task.Status);
    }

    // Bug: SetupMcpConfigAsync registers the per-run MCP token *before* it can fail (writing
    // the config file to disk). When it threw after Register but before the try/catch existed,
    // the token was never unregistered — a leak in TaskRunTokenRegistry on top of the stuck
    // Running task.
    [Fact]
    public async Task ContinueAsync_McpSetupThrowsAfterTokenRegistered_MarksTaskFailed_AndUnregistersToken()
    {
        var taskId = await SeedContinuableTaskAsync("sess-mcp-throws");

        // SetupMcpConfigAsync calls Directory.CreateDirectory(cfg.LogRoot) right after
        // registering the token; pointing LogRoot at an existing file makes that call throw.
        var badLogRoot = Path.Combine(_tempDir, "logroot_is_a_file");
        await File.WriteAllTextAsync(badLogRoot, "not a directory");
        var cfg = new WorkerConfig { SandboxRoot = _tempDir, LogRoot = badLogRoot };
        var tokens = new TaskRunTokenRegistry();

        using var ctx2 = _db.CreateContext();
        var runner = BuildRunner(new FakeClaudeProcess(), ctx2, tokens: tokens, cfg: cfg);

        await runner.ContinueAsync(taskId, "please continue", "slot-1", CancellationToken.None);

        using var verify = _db.CreateContext();
        var task = await new TaskRepository(verify).GetByIdAsync(taskId);
        Assert.NotNull(task);
        Assert.Equal(TaskStatus.Failed, task.Status);
        Assert.False(tokens.HasAnyTokenFor(taskId));
    }

    [Fact]
    public async Task ContinueAsync_UnhandledException_MarksTaskFailed_NotStuckRunning()
    {
        var dbFactory = _db.CreateFactory();
        string listId, taskId;

        using (var ctx = _db.CreateContext())
        {
            listId = Guid.NewGuid().ToString();
            ctx.Lists.Add(new ListEntity { Id = listId, Name = "L", WorkingDir = null, CreatedAt = DateTime.UtcNow });

            taskId = Guid.NewGuid().ToString();
            ctx.Tasks.Add(new TaskEntity
            {
                Id = taskId,
                ListId = listId,
                Title = "Continue me",
                Status = TaskStatus.WaitingForReview,
                CreatedAt = DateTime.UtcNow,
            });
            await ctx.SaveChangesAsync();

            // A prior run with a session ID is required for ContinueAsync to proceed.
            await new TaskRunRepository(ctx).AddAsync(new TaskRunEntity
            {
                Id = Guid.NewGuid().ToString(),
                TaskId = taskId,
                RunNumber = 1,
                IsRetry = false,
                Prompt = "original prompt",
                SessionId = "sess-continue-test",
                StartedAt = DateTime.UtcNow.AddMinutes(-5),
                FinishedAt = DateTime.UtcNow.AddMinutes(-1),
                ExitCode = 0,
                ResultMarkdown = "first result",
            });
        }

        // This process throws a non-cancellation exception to simulate an unexpected failure.
        var throwingProcess = new ThrowingClaudeProcess(new InvalidOperationException("disk full"));

        using var ctx2 = _db.CreateContext();
        var runner = BuildRunner(throwingProcess, ctx2);

        // ContinueAsync must not propagate the exception and must leave the task in Failed.
        await runner.ContinueAsync(taskId, "please continue", "slot-1", CancellationToken.None);

        using var verify = _db.CreateContext();
        var task = await new TaskRepository(verify).GetByIdAsync(taskId);
        Assert.NotNull(task);
        Assert.Equal(TaskStatus.Failed, task.Status);
    }

    [Fact]
    public async Task ContinueAsync_WiresMcpConfigAndAllowedTools_LikeAFreshRun()
    {
        string listId, taskId;

        using (var ctx = _db.CreateContext())
        {
            listId = Guid.NewGuid().ToString();
            ctx.Lists.Add(new ListEntity { Id = listId, Name = "L", WorkingDir = null, CreatedAt = DateTime.UtcNow });

            taskId = Guid.NewGuid().ToString();
            ctx.Tasks.Add(new TaskEntity
            {
                Id = taskId,
                ListId = listId,
                Title = "Continue me",
                Status = TaskStatus.WaitingForReview,
                CreatedAt = DateTime.UtcNow,
            });
            await ctx.SaveChangesAsync();

            await new TaskRunRepository(ctx).AddAsync(new TaskRunEntity
            {
                Id = Guid.NewGuid().ToString(),
                TaskId = taskId,
                RunNumber = 1,
                IsRetry = false,
                Prompt = "original prompt",
                SessionId = "sess-continue-mcp-test",
                StartedAt = DateTime.UtcNow.AddMinutes(-5),
                FinishedAt = DateTime.UtcNow.AddMinutes(-1),
                ExitCode = 0,
                ResultMarkdown = "first result",
            });
        }

        IReadOnlyList<string>? capturedArgs = null;
        var claude = new FakeClaudeProcess((_, _, args, _, _) =>
        {
            capturedArgs = args;
            return Task.FromResult(new RunResult { ExitCode = 0, SessionId = "sess-continue-mcp-test", ResultMarkdown = "ok" });
        });

        using var ctx2 = _db.CreateContext();
        var runner = BuildRunner(claude, ctx2);

        await runner.ContinueAsync(taskId, "please continue", "slot-1", CancellationToken.None);

        Assert.NotNull(capturedArgs);
        Assert.Contains("--mcp-config", capturedArgs!);
        Assert.Contains("--allowedTools", capturedArgs!);
    }

    private sealed class ThrowingClaudeProcess : IClaudeProcess
    {
        private readonly Exception _ex;
        public ThrowingClaudeProcess(Exception ex) => _ex = ex;

        public Task<RunResult> RunAsync(
            IReadOnlyList<string> arguments, string prompt, string workingDirectory,
            Func<string, Task> onStdoutLine, CancellationToken ct)
            => throw _ex;
    }
}