ClaudeDo/tests/ClaudeDo.Worker.Tests/Planning/PlanningSessionManagerTests.cs

using ClaudeDo.Data;
using ClaudeDo.Data.Git;
using ClaudeDo.Data.Models;
using ClaudeDo.Data.Repositories;
using ClaudeDo.Worker.Config;
using ClaudeDo.Worker.Planning;
using ClaudeDo.Worker.Tests.Infrastructure;
using TaskStatus = ClaudeDo.Data.Models.TaskStatus;

namespace ClaudeDo.Worker.Tests.Planning;

public sealed class PlanningSessionManagerTests : IDisposable
{
    private readonly DbFixture _db = new();
    private readonly ClaudeDoDbContext _ctx;
    private readonly TaskRepository _tasks;
    private readonly ListRepository _lists;
    private readonly string _rootDir;
    private readonly GitService _git;
    private readonly WorkerConfig _cfg;
    private readonly AppSettingsRepository _settingsRepo;
    private readonly PlanningSessionManager _sut;
    private readonly TaskStateServiceBuilder.Built _built;

    public PlanningSessionManagerTests()
    {
        _ctx = _db.CreateContext();
        _tasks = new TaskRepository(_ctx);
        _lists = new ListRepository(_ctx);
        _rootDir = Path.Combine(Path.GetTempPath(), $"cd_planning_{Guid.NewGuid():N}");
        _git = new GitService();
        _cfg = new WorkerConfig { CentralWorktreeRoot = Path.Combine(_rootDir, "central") };
        _settingsRepo = new AppSettingsRepository(_ctx);
        _settingsRepo.UpdateAsync(new AppSettingsEntity { WorktreeStrategy = "sibling" }).GetAwaiter().GetResult();
        _built = TaskStateServiceBuilder.Build(_db.CreateFactory());
        _sut = new PlanningSessionManager(
            _tasks, _lists, _settingsRepo, _git, _cfg, _rootDir, _built.State, _built.Chain);
    }

    public void Dispose()
    {
        _ctx.Dispose();
        _db.Dispose();
        try { Directory.Delete(_rootDir, recursive: true); } catch { /* ignore */ }
    }

    private async Task<(string listId, string workingDir)> SeedListAsync()
    {
        var listId = Guid.NewGuid().ToString();
        var wd = Path.Combine(Path.GetTempPath(), $"cd_wd_{Guid.NewGuid():N}");
        GitRepoFixture.InitRepoWithInitialCommit(wd);
        await _lists.AddAsync(new ListEntity
        {
            Id = listId,
            Name = "Test",
            WorkingDir = wd,
            CreatedAt = DateTime.UtcNow,
        });
        return (listId, wd);
    }

    private async Task<TaskEntity> SeedManualTaskAsync(string listId)
    {
        var t = new TaskEntity
        {
            Id = Guid.NewGuid().ToString(),
            ListId = listId,
            Title = "Brainstorm auth",
            Description = "- review tokens\n- plan rollout",
            Status = TaskStatus.Manual,
            CreatedAt = DateTime.UtcNow,
            CommitType = "feat",
        };
        await _tasks.AddAsync(t);
        return t;
    }

    [Fact]
    public async Task StartAsync_CreatesSessionFiles_AndTransitionsTaskToPlanning()
    {
        var (listId, wd) = await SeedListAsync();
        var parent = await SeedManualTaskAsync(listId);

        var ctx = await _sut.StartAsync(parent.Id, CancellationToken.None);

        Assert.Equal(parent.Id, ctx.ParentTaskId);
        Assert.Equal(ctx.WorktreePath, ctx.WorkingDir);
        Assert.True(Directory.Exists(ctx.WorktreePath));
        var mcpPath = Path.Combine(ctx.WorktreePath, ".mcp.json");
        Assert.True(File.Exists(mcpPath));
        Assert.True(File.Exists(Path.Combine(ctx.WorktreePath, ".claude", "settings.local.json")));
        Assert.True(File.Exists(ctx.Files.SystemPromptPath));
        Assert.True(File.Exists(ctx.Files.InitialPromptPath));

        var mcp = await File.ReadAllTextAsync(mcpPath);
        Assert.Contains("${CLAUDEDO_PLANNING_TOKEN}", mcp);
        Assert.DoesNotContain(ctx.Token, mcp);

        var initial = await File.ReadAllTextAsync(ctx.Files.InitialPromptPath);
        Assert.Contains("Brainstorm auth", initial);
        Assert.Contains("review tokens", initial);

        var loaded = await _tasks.GetByIdAsync(parent.Id);
        Assert.Equal(TaskStatus.Planning, loaded!.Status);
        Assert.NotNull(loaded.PlanningSessionToken);
    }

    [Fact]
    public async Task StartAsync_TaskNotManual_Throws()
    {
        var (listId, _) = await SeedListAsync();
        var queuedTask = new TaskEntity
        {
            Id = Guid.NewGuid().ToString(),
            ListId = listId,
            Title = "x",
            Status = TaskStatus.Queued,
            CreatedAt = DateTime.UtcNow,
            CommitType = "feat",
        };
        await _tasks.AddAsync(queuedTask);

        await Assert.ThrowsAsync<InvalidOperationException>(() =>
            _sut.StartAsync(queuedTask.Id, CancellationToken.None));
    }

    [Fact]
    public async Task StartAsync_ChildTask_Throws()
    {
        var (listId, _) = await SeedListAsync();
        var parent = await SeedManualTaskAsync(listId);
        await _tasks.SetPlanningStartedAsync(parent.Id, "t");
        var child = await _tasks.CreateChildAsync(parent.Id, "c", null, null, null);

        await Assert.ThrowsAsync<InvalidOperationException>(() =>
            _sut.StartAsync(child.Id, CancellationToken.None));
    }

    [Fact]
    public async Task ResumeAsync_ReturnsExistingSessionDetails()
    {
        var (listId, wd) = await SeedListAsync();
        var parent = await SeedManualTaskAsync(listId);
        var startCtx = await _sut.StartAsync(parent.Id, CancellationToken.None);
        await _tasks.UpdatePlanningSessionIdAsync(parent.Id, "claude-session-42");

        var resumeCtx = await _sut.ResumeAsync(parent.Id, CancellationToken.None);

        Assert.Equal(parent.Id, resumeCtx.ParentTaskId);
        Assert.Equal(resumeCtx.WorktreePath, resumeCtx.WorkingDir);
        Assert.Equal("claude-session-42", resumeCtx.ClaudeSessionId);
        Assert.True(Directory.Exists(resumeCtx.WorktreePath));
        Assert.True(File.Exists(Path.Combine(resumeCtx.WorktreePath, ".mcp.json")));
    }

    [Fact]
    public async Task ResumeAsync_NotPlanning_Throws()
    {
        var (listId, _) = await SeedListAsync();
        var parent = await SeedManualTaskAsync(listId);
        // did not start
        await Assert.ThrowsAsync<InvalidOperationException>(() =>
            _sut.ResumeAsync(parent.Id, CancellationToken.None));
    }

    [Fact]
    public async Task ResumeAsync_NoClaudeSessionId_Throws()
    {
        var (listId, _) = await SeedListAsync();
        var parent = await SeedManualTaskAsync(listId);
        await _sut.StartAsync(parent.Id, CancellationToken.None);
        // UpdatePlanningSessionIdAsync not called

        await Assert.ThrowsAsync<InvalidOperationException>(() =>
            _sut.ResumeAsync(parent.Id, CancellationToken.None));
    }

    [Fact]
    public async Task FinalizeAsync_PromotesDraftsAndMarksPlanningFinalized()
    {
        var (listId, _) = await SeedListAsync();
        var parent = await SeedManualTaskAsync(listId);
        await _sut.StartAsync(parent.Id, CancellationToken.None);
        await _tasks.CreateChildAsync(parent.Id, "c1", null, null, null);
        await _tasks.CreateChildAsync(parent.Id, "c2", null, null, null);

        var count = await _sut.FinalizeAsync(parent.Id, queueAgentTasks: true, CancellationToken.None);

        Assert.Equal(2, count);
        var loaded = await _tasks.GetByIdAsync(parent.Id);
        Assert.Equal(PlanningPhase.Finalized, loaded!.PlanningPhase);
        Assert.NotNull(loaded.PlanningFinalizedAt);
        Assert.Null(loaded.PlanningSessionToken);
    }

    [Fact]
    public async Task GetPendingDraftCountAsync_ReturnsDraftCount()
    {
        var (listId, _) = await SeedListAsync();
        var parent = await SeedManualTaskAsync(listId);
        await _sut.StartAsync(parent.Id, CancellationToken.None);
        await _tasks.CreateChildAsync(parent.Id, "c1", null, null, null);
        await _tasks.CreateChildAsync(parent.Id, "c2", null, null, null);
        await _tasks.CreateChildAsync(parent.Id, "c3", null, null, null);

        var n = await _sut.GetPendingDraftCountAsync(parent.Id, CancellationToken.None);

        Assert.Equal(3, n);
    }

    [Fact]
    public async Task DiscardAsync_DeletesSessionDirAndResetsTask()
    {
        var (listId, _) = await SeedListAsync();
        var parent = await SeedManualTaskAsync(listId);
        var startCtx = await _sut.StartAsync(parent.Id, CancellationToken.None);
        Assert.True(Directory.Exists(startCtx.Files.SessionDirectory));

        await _sut.DiscardAsync(parent.Id, CancellationToken.None);

        Assert.False(Directory.Exists(startCtx.Files.SessionDirectory));
        var loaded = await _tasks.GetByIdAsync(parent.Id);
        Assert.Equal(TaskStatus.Manual, loaded!.Status);
        Assert.Null(loaded.PlanningSessionToken);
    }

    [Fact]
    public async Task DiscardAsync_RemovesWorktreeAndBranch()
    {
        var (listId, wd) = await SeedListAsync();
        var parent = await SeedManualTaskAsync(listId);

        var ctx = await _sut.StartAsync(parent.Id, CancellationToken.None);
        Assert.True(Directory.Exists(ctx.WorktreePath));

        await _sut.DiscardAsync(parent.Id, CancellationToken.None);

        Assert.False(Directory.Exists(ctx.WorktreePath));
        // branch deleted
        var paths = await _git.ListWorktreePathsForBranchAsync(wd, ctx.BranchName);
        Assert.Empty(paths);
    }

    [Fact]
    public async Task StartAsync_ThrowsWhenWorkingDirIsNotGitRepo()
    {
        var listId = Guid.NewGuid().ToString();
        var wd = Path.Combine(Path.GetTempPath(), $"cd_nogit_{Guid.NewGuid():N}");
        Directory.CreateDirectory(wd);
        try
        {
            await _lists.AddAsync(new ListEntity { Id = listId, Name = "NoGit", WorkingDir = wd, CreatedAt = DateTime.UtcNow });
            var t = await SeedManualTaskAsync(listId);
            await Assert.ThrowsAsync<InvalidOperationException>(() => _sut.StartAsync(t.Id, CancellationToken.None));
        }
        finally
        {
            try { Directory.Delete(wd, recursive: true); } catch { /* best effort */ }
        }
    }

    [Fact]
    public async Task StartAsync_SelfHealsWhenBranchAlreadyExists()
    {
        var (listId, wd) = await SeedListAsync();
        var parent = await SeedManualTaskAsync(listId);

        // Pre-create a colliding branch.
        var branch = $"claudedo/planning/{parent.Id.Replace("-", "")}";
        var head = await _git.RevParseHeadAsync(wd);
        var procInfo = new System.Diagnostics.ProcessStartInfo("git") { WorkingDirectory = wd };
        procInfo.ArgumentList.Add("branch");
        procInfo.ArgumentList.Add(branch);
        procInfo.ArgumentList.Add(head);
        var p = System.Diagnostics.Process.Start(procInfo)!;
        p.WaitForExit();
        Assert.True(p.ExitCode == 0, $"git branch setup failed with exit {p.ExitCode}");

        var ctx = await _sut.StartAsync(parent.Id, CancellationToken.None);
        Assert.True(Directory.Exists(ctx.WorktreePath));
    }

    [Fact]
    public async Task ResumeAsync_ReturnsContextWithTokenAndWorktree()
    {
        var (listId, wd) = await SeedListAsync();
        var parent = await SeedManualTaskAsync(listId);

        var startCtx = await _sut.StartAsync(parent.Id, CancellationToken.None);

        // Simulate the claude session capturing its session id.
        await _tasks.UpdatePlanningSessionIdAsync(parent.Id, "session-abc");

        var resumeCtx = await _sut.ResumeAsync(parent.Id, CancellationToken.None);

        Assert.Equal(startCtx.Token, resumeCtx.Token);
        Assert.Equal(startCtx.WorktreePath, resumeCtx.WorktreePath);
        Assert.Equal("session-abc", resumeCtx.ClaudeSessionId);
    }
}