ClaudeDo/src/ClaudeDo.Worker/Planning/PlanningChainCoordinator.cs

using ClaudeDo.Data;
using ClaudeDo.Data.Models;
using ClaudeDo.Worker.State;
using Microsoft.EntityFrameworkCore;
using TaskStatus = ClaudeDo.Data.Models.TaskStatus;

namespace ClaudeDo.Worker.Planning;

public sealed class PlanningChainCoordinator
{
    private readonly IDbContextFactory<ClaudeDoDbContext> _dbFactory;
    private readonly Func<ITaskStateService> _state;

    public PlanningChainCoordinator(
        IDbContextFactory<ClaudeDoDbContext> dbFactory,
        Func<ITaskStateService> state)
    {
        _dbFactory = dbFactory;
        _state = state;
    }

    // Sets up a sequential chain over a planning parent's children.
    // - First non-terminal child gets BlockedByTaskId=null.
    // - Each subsequent non-terminal child gets BlockedByTaskId=<predecessor>,
    //   so the picker skips them until the predecessor finishes.
    // - When enqueue is true, each non-terminal child is also set to Status=Queued
    //   (the user-driven "Queue plan"). When false (finalize), children are left
    //   Idle and only the blocked-by links are established, so nothing runs until
    //   the user queues the plan.
    // - Terminal children (Done/Failed/Cancelled) are left untouched; they are
    //   skipped when computing predecessors so a re-run on a partially executed
    //   chain leaves history alone but still reshapes the tail.
    // - Running children abort the operation — the chain cannot be reshaped while
    //   one of its members is mid-flight.
    // Returns the number of children placed in the chain.
    internal async Task<int> SetupChainAsync(string parentTaskId, bool enqueue, CancellationToken ct = default)
    {
        await using var ctx = await _dbFactory.CreateDbContextAsync(ct);
        var parent = await ctx.Tasks.FirstOrDefaultAsync(t => t.Id == parentTaskId, ct)
            ?? throw new InvalidOperationException($"Task {parentTaskId} not found.");

        var children = await ctx.Tasks
            .Where(t => t.ParentTaskId == parentTaskId)
            .OrderBy(t => t.SortOrder).ThenBy(t => t.CreatedAt)
            .ToListAsync(ct);
        if (children.Count == 0)
            throw new InvalidOperationException("Parent has no subtasks.");

        var running = children.FirstOrDefault(c => c.Status == TaskStatus.Running);
        if (running is not null)
            throw new InvalidOperationException(
                $"Child {running.Id} is running; cannot reshape chain.");

        // Re-shape over Idle and Queued children only; leave Done/Failed/Cancelled
        // (terminal) results in place.
        var sequenceable = children
            .Where(c => c.Status == TaskStatus.Idle || c.Status == TaskStatus.Queued)
            .ToList();

        var state = _state();
        for (int i = 0; i < sequenceable.Count; i++)
        {
            if (enqueue)
                await state.EnqueueAsync(sequenceable[i].Id, ct);
            if (i == 0)
                await state.UnblockAsync(sequenceable[i].Id, ct);
            else
                await state.BlockOnAsync(sequenceable[i].Id, sequenceable[i - 1].Id, ct);
        }

        return sequenceable.Count;
    }

    // User-triggered "send plan to queue". Only valid once the plan is finalized
    // (children are "Planned"); otherwise the children are still drafts.
    public async Task<int> QueuePlanAsync(string parentTaskId, CancellationToken ct = default)
    {
        await using var ctx = await _dbFactory.CreateDbContextAsync(ct);
        var phase = await ctx.Tasks.AsNoTracking()
            .Where(t => t.Id == parentTaskId)
            .Select(t => (PlanningPhase?)t.PlanningPhase)
            .FirstOrDefaultAsync(ct);

        if (phase is null)
            throw new InvalidOperationException($"Task {parentTaskId} not found.");
        if (phase != PlanningPhase.Finalized)
            throw new InvalidOperationException("Plan must be finalized before it can be queued.");

        return await SetupChainAsync(parentTaskId, enqueue: true, ct);
    }

    public async Task<string?> OnChildFinishedAsync(
        string childTaskId, TaskStatus finalStatus, CancellationToken ct = default)
    {
        var nextId = await FindSuccessorAsync(childTaskId, ct);
        if (nextId is null) return null;

        if (finalStatus == TaskStatus.Done)
        {
            await _state().UnblockAsync(nextId, ct);
            return nextId;
        }

        // Child failed or was cancelled: cancel the immediate successor so the chain is
        // not left wedged. If it's cancellable, CancelAsync's own OnChildTerminalAsync
        // callback recurses into this method for its successor, cascading cancellation
        // through the rest of the chain. If it's not (e.g. it was parked back to Idle
        // out of band), CancelAsync is a no-op and nothing will ever call back for it —
        // keep walking the chain ourselves so the tail isn't left wedged forever.
        var predecessorId = nextId;
        while (true)
        {
            var result = await _state().CancelAsync(predecessorId, DateTime.UtcNow, ct);
            if (result.Ok) return null;

            var following = await FindSuccessorAsync(predecessorId, ct);
            if (following is null) return null;
            predecessorId = following;
        }
    }

    // The successor is whichever sibling explicitly blocks on this task.
    private async Task<string?> FindSuccessorAsync(string taskId, CancellationToken ct)
    {
        await using var ctx = await _dbFactory.CreateDbContextAsync(ct);
        return await ctx.Tasks
            .AsNoTracking()
            .Where(t => t.BlockedByTaskId == taskId)
            .OrderBy(t => t.SortOrder).ThenBy(t => t.CreatedAt)
            .Select(t => t.Id)
            .FirstOrDefaultAsync(ct);
    }
}