ClaudeDo/src/ClaudeDo.Worker/Planning/PlanningMergeOrchestrator.cs

using System.Collections.Concurrent;
using ClaudeDo.Data;
using ClaudeDo.Data.Git;
using ClaudeDo.Data.Models;
using ClaudeDo.Worker.Hub;
using ClaudeDo.Worker.Lifecycle;
using ClaudeDo.Worker.State;
using Microsoft.EntityFrameworkCore;
using ModelContextProtocol;
using TaskStatus = ClaudeDo.Data.Models.TaskStatus;

namespace ClaudeDo.Worker.Planning;

/// <summary>A unit-merge conflict currently paused, driven by an MCP session rather than the UI.</summary>
public sealed record ExternalPlanningMergeConflict(string PlanningTaskId, string SubtaskId);

public sealed class PlanningMergeOrchestrator : IActiveMergeState
{
    private readonly IDbContextFactory<ClaudeDoDbContext> _dbFactory;
    private readonly TaskMergeService _merge;
    private readonly PlanningAggregator _aggregator;
    private readonly HubBroadcaster _broadcaster;
    private readonly GitService _git;
    private readonly ITaskStateService _state;
    private readonly ILogger<PlanningMergeOrchestrator> _logger;

    private sealed class State
    {
        public required string TargetBranch { get; init; }
        public required Queue<string> RemainingSubtaskIds { get; init; }
        public required bool IsPlanning { get; init; }
        public required string WorkingDir { get; init; }
        /// <summary>True when this unit merge was started by an MCP tool call (e.g. review_task
        /// from a running Claude session) rather than a direct UI action. Conflicts on such a
        /// merge must not auto-pop the in-app resolver — the driving session owns resolution.</summary>
        public required bool ExternallyDriven { get; init; }
        public string? CurrentSubtaskId { get; set; }
    }

    private readonly ConcurrentDictionary<string, State> _states = new();

    public PlanningMergeOrchestrator(
        IDbContextFactory<ClaudeDoDbContext> dbFactory,
        TaskMergeService merge,
        PlanningAggregator aggregator,
        HubBroadcaster broadcaster,
        GitService git,
        ITaskStateService state,
        ILogger<PlanningMergeOrchestrator> logger)
    {
        _dbFactory = dbFactory;
        _merge = merge;
        _aggregator = aggregator;
        _broadcaster = broadcaster;
        _git = git;
        _state = state;
        _logger = logger;
    }

    public async Task StartAsync(
        string parentTaskId, string targetBranch, CancellationToken ct, bool externallyDriven = false,
        IProgress<ProgressNotificationValue>? progress = null)
    {
        string workingDir;
        List<TaskEntity> children;
        bool isPlanning;
        bool parentHasWorktree;

        using (var ctx = _dbFactory.CreateDbContext())
        {
            var parent = await ctx.Tasks
                .Include(t => t.List)
                .Include(t => t.Worktree)
                .Include(t => t.Children).ThenInclude(c => c.Worktree)
                .SingleOrDefaultAsync(t => t.Id == parentTaskId, ct)
                ?? throw new KeyNotFoundException($"Planning task '{parentTaskId}' not found.");
            workingDir = parent.List.WorkingDir
                ?? throw new InvalidOperationException("List has no working directory.");
            children = parent.Children.OrderBy(c => c.SortOrder).ToList();
            isPlanning = parent.PlanningPhase != PlanningPhase.None;
            parentHasWorktree = parent.Worktree is { State: WorktreeState.Active };
        }

        if (isPlanning)
        {
            foreach (var c in children)
            {
                if (c.Status != TaskStatus.Done)
                    throw new InvalidOperationException($"subtask {c.Id} is not Done (status {c.Status})");
                if (c.Worktree is null)
                    throw new InvalidOperationException($"subtask {c.Id} has no worktree");
                if (c.Worktree.State != WorktreeState.Active && c.Worktree.State != WorktreeState.Merged)
                    throw new InvalidOperationException(
                        $"subtask {c.Id} worktree state is {c.Worktree.State}");
            }
        }

        if (await _git.IsMidMergeAsync(workingDir, ct))
            throw new InvalidOperationException(
                "repo is mid-merge; use AbortPlanningMerge to reset the repository, then Approve again");
        if (await _git.HasChangesAsync(workingDir, includeUntracked: false, ct))
            throw new InvalidOperationException("working tree has uncommitted changes");

        var idsToMerge = new List<string>();
        if (!isPlanning && parentHasWorktree)
            idsToMerge.Add(parentTaskId);
        idsToMerge.AddRange(
            children
                .Where(c => c.Status == TaskStatus.Done && c.Worktree is { State: WorktreeState.Active })
                .Select(c => c.Id));

        var queue = new Queue<string>(idsToMerge);

        var state = new State
        {
            TargetBranch = targetBranch,
            RemainingSubtaskIds = queue,
            IsPlanning = isPlanning,
            WorkingDir = workingDir,
            ExternallyDriven = externallyDriven,
        };
        if (!_states.TryAdd(parentTaskId, state))
            throw new InvalidOperationException($"Merge already in progress for {parentTaskId}.");

        await _broadcaster.PlanningMergeStarted(parentTaskId, targetBranch);
        await DrainAsync(parentTaskId, ct, progress);
    }

    /// <summary>True when a unit merge for this parent is paused on a conflict (in-memory state).</summary>
    public bool HasActiveMerge(string parentTaskId) =>
        _states.TryGetValue(parentTaskId, out var s) && s.CurrentSubtaskId is not null;

    /// <summary>Externally-driven unit merges currently paused on a conflict, so the UI can
    /// recover this on reconnect instead of relying solely on the one-shot broadcast. Checked
    /// against <see cref="GitService.IsMidMergeAsync"/> rather than trusting the in-memory flag
    /// alone, so a stale entry (e.g. the repo was already resolved through another path) never
    /// reports an external merge that no longer exists.</summary>
    public async Task<IReadOnlyList<ExternalPlanningMergeConflict>> GetActiveExternalConflictsAsync(CancellationToken ct)
    {
        var result = new List<ExternalPlanningMergeConflict>();
        foreach (var (planningTaskId, state) in _states)
        {
            if (!state.ExternallyDriven || state.CurrentSubtaskId is null) continue;
            if (!await _git.IsMidMergeAsync(state.WorkingDir, ct)) continue;
            result.Add(new ExternalPlanningMergeConflict(planningTaskId, state.CurrentSubtaskId));
        }
        return result;
    }

    public async Task ContinueAsync(
        string planningTaskId, CancellationToken ct, IProgress<ProgressNotificationValue>? progress = null)
    {
        if (!_states.TryGetValue(planningTaskId, out var state) || state.CurrentSubtaskId is null)
            throw new InvalidOperationException(
                "no in-progress merge to continue; if the worker was restarted during a conflict, use AbortPlanningMerge to reset the repository");

        var current = state.CurrentSubtaskId;
        var result = await _merge.ContinueMergeAsync(current, ct, progress);

        if (result.Status == TaskMergeService.StatusConflict)
        {
            await _broadcaster.PlanningMergeConflict(planningTaskId, current, result.ConflictFiles, state.ExternallyDriven);
            return;
        }

        if (result.Status != TaskMergeService.StatusMerged)
        {
            _logger.LogWarning(
                "Planning continue blocked on subtask {Subtask}: {Msg}",
                current, result.ErrorMessage);
            _states.TryRemove(planningTaskId, out _);
            await _broadcaster.PlanningMergeAborted(planningTaskId);
            return;
        }

        await _broadcaster.PlanningSubtaskMerged(planningTaskId, current);

        state.CurrentSubtaskId = null;
        await DrainAsync(planningTaskId, ct, progress);
    }

    public async Task AbortAsync(string planningTaskId, CancellationToken ct)
    {
        if (!_states.TryGetValue(planningTaskId, out var state) || state.CurrentSubtaskId is null)
        {
            // No in-memory state — worker may have been restarted while a conflict was paused.
            // Check whether the list repo is still mid-merge and abort it directly.
            await AbortStatelessAsync(planningTaskId, ct);
            return;
        }

        await _merge.AbortMergeAsync(state.CurrentSubtaskId, ct);
        _states.TryRemove(planningTaskId, out _);
        await _broadcaster.PlanningMergeAborted(planningTaskId);
    }

    private async Task AbortStatelessAsync(string planningTaskId, CancellationToken ct)
    {
        string? workingDir;
        await using (var ctx = _dbFactory.CreateDbContext())
        {
            workingDir = await ctx.Tasks
                .Where(t => t.Id == planningTaskId)
                .Select(t => t.List.WorkingDir)
                .FirstOrDefaultAsync(ct);
        }

        if (string.IsNullOrWhiteSpace(workingDir) || !await _git.IsMidMergeAsync(workingDir, ct))
            throw new InvalidOperationException("no in-progress merge to abort");

        await _git.MergeAbortAsync(workingDir, ct);
        _logger.LogInformation(
            "Stateless abort of mid-merge for planning task {ParentId} (post-restart recovery)",
            planningTaskId);
        await _broadcaster.PlanningMergeAborted(planningTaskId);
        // Parent remains WaitingForReview — Approve will restart the unit merge from scratch.
    }

    private async Task DrainAsync(
        string planningTaskId, CancellationToken ct, IProgress<ProgressNotificationValue>? progress = null)
    {
        if (!_states.TryGetValue(planningTaskId, out var state)) return;

        var keepState = false;
        try
        {
            while (state.RemainingSubtaskIds.TryDequeue(out var subtaskId))
            {
                state.CurrentSubtaskId = subtaskId;
                var result = await _merge.MergeAsync(
                    subtaskId,
                    state.TargetBranch,
                    removeWorktree: true,
                    commitMessage: "",   // blank -> TaskMergeService builds the conventional default
                    leaveConflictsInTree: true,
                    ct,
                    progress);

                if (result.Status == TaskMergeService.StatusConflict)
                {
                    await _broadcaster.PlanningMergeConflict(
                        planningTaskId, subtaskId, result.ConflictFiles, state.ExternallyDriven);
                    keepState = true;
                    return;
                }

                if (result.Status != TaskMergeService.StatusMerged)
                {
                    _logger.LogWarning(
                        "Planning merge blocked on subtask {Subtask}: {Msg}",
                        subtaskId, result.ErrorMessage);
                    await _broadcaster.PlanningMergeAborted(planningTaskId);
                    return; // keepState stays false → finally removes the state entry
                }

                await _broadcaster.PlanningSubtaskMerged(planningTaskId, subtaskId);
            }

            state.CurrentSubtaskId = null;
            var finalized = await FinalizeParentDoneAsync(planningTaskId, state.IsPlanning, ct);
            if (finalized)
                await _broadcaster.PlanningCompleted(planningTaskId);
        }
        finally
        {
            if (!keepState) _states.TryRemove(planningTaskId, out _);
        }
    }

    private async Task<bool> FinalizeParentDoneAsync(string parentTaskId, bool isPlanning, CancellationToken ct)
    {
        var result = await _state.ApproveReviewAsync(parentTaskId, ct);
        if (!result.Ok)
        {
            // ApproveReviewAsync requires WaitingForReview. For improvement parents whose own
            // worktree is in the merge queue, TaskMergeService.ApproveIfWaitingForReviewAsync
            // already approved the parent during the drain — check for that expected path.
            await using var ctx = _dbFactory.CreateDbContext();
            var current = await ctx.Tasks
                .Where(t => t.Id == parentTaskId)
                .Select(t => (TaskStatus?)t.Status)
                .FirstOrDefaultAsync(ct);

            if (current != TaskStatus.Done)
            {
                // Parent was cancelled or moved to an unexpected state during the merge drain.
                // Do not overwrite — the external transition takes precedence.
                _logger.LogWarning(
                    "Unit-merge drain completed but parent {ParentTaskId} could not be finalized (status: {Status}): {Reason}",
                    parentTaskId, current, result.Reason);
                return false;
            }
        }

        // Only planning builds an integration branch via the aggregator; skip cleanup otherwise.
        if (isPlanning)
        {
            try { await _aggregator.CleanupIntegrationBranchAsync(parentTaskId, ct); }
            catch (Exception ex) { _logger.LogWarning(ex, "integration branch cleanup failed"); }
        }

        return true;
    }
}