AirPlan-V2/lib/air_runtime/worktree.py

"""
Worktree 隔离并行模块 — V2 P1-15 修复。
区域级冲突检测 + git worktree 隔离，允许同文件不同区域任务安全并行。
"""

from __future__ import annotations

import enum
import subprocess
from dataclasses import dataclass
from pathlib import Path


class ConflictLevel(enum.Enum):
    NONE = "none"       # 无文件重叠，直接并行
    SOFT = "soft"       # 同文件不同区域，worktree 隔离并行
    HARD = "hard"       # 同文件同区域，必须串行


@dataclass
class Conflict:
    task_a: str
    task_b: str
    level: ConflictLevel  # NONE / SOFT / HARD
    reason: str = ""

    def to_dict(self) -> dict:
        return {
            "taskA": self.task_a,
            "taskB": self.task_b,
            "level": self.level.value,
            "reason": self.reason,
        }


@dataclass
class MergeResult:
    task_id: str
    successful: bool
    conflicts: list[str]


class RegionConflictDetector:
    """区域级写集冲突检测。

    区域定义（按优先级）：
    1. todo.md 中的区域标注 [region: xxx]
    2. 函数/类边界（AST 分析或标记注释）
    3. 行号区间（从 diff 或任务元数据获取）
    """

    def detect(self, task_a_write_set: list[str], task_b_write_set: list[str],
               regions_a: dict[str, list[tuple[int, int]]] | None = None,
               regions_b: dict[str, list[tuple[int, int]]] | None = None) -> ConflictLevel:
        file_overlap = set(task_a_write_set) & set(task_b_write_set)
        if not file_overlap:
            return ConflictLevel.NONE

        if regions_a is None or regions_b is None:
            return ConflictLevel.SOFT

        for fpath in file_overlap:
            ra = regions_a.get(fpath, [(0, 999999)])
            rb = regions_b.get(fpath, [(0, 999999)])
            for a_start, a_end in ra:
                for b_start, b_end in rb:
                    if a_start < b_end and b_start < a_end:
                        return ConflictLevel.HARD
        return ConflictLevel.SOFT

    def detect_batch(self, task_write_sets: dict[str, list[str]],
                     task_regions: dict[str, dict[str, list[tuple[int, int]]]] | None = None,
                     ) -> list[Conflict]:
        """批量检测：接收 {task_id: [file, ...]} 映射，返回所有冲突对。

        逐对调用 detect()，仅返回 level != NONE 的冲突。
        """
        conflicts: list[Conflict] = []
        task_ids = list(task_write_sets.keys())
        for i in range(len(task_ids)):
            for j in range(i + 1, len(task_ids)):
                a, b = task_ids[i], task_ids[j]
                ra = task_regions.get(a) if task_regions else None
                rb = task_regions.get(b) if task_regions else None
                level = self.detect(task_write_sets[a], task_write_sets[b], ra, rb)
                if level != ConflictLevel.NONE:
                    overlap = sorted(set(task_write_sets[a]) & set(task_write_sets[b]))
                    reason = f"file overlap: {', '.join(overlap)}" if overlap else ""
                    if level == ConflictLevel.HARD:
                        reason = f"region overlap: {', '.join(overlap)}" if overlap else "same region"
                    conflicts.append(Conflict(task_a=a, task_b=b, level=level, reason=reason))
        return conflicts

    def extract_regions_from_task(self, file_path: Path, region_markers: list[tuple[int, int]]) -> dict[str, list[tuple[int, int]]]:
        """从任务的区域标注提取行号区间。"""
        return {str(file_path): region_markers}


class WorktreeIsolation:
    """为 SOFT 冲突任务创建 git worktree 隔离，完成后合并回主分支。"""

    def __init__(self, repo_root: Path):
        self._repo_root = repo_root

    def create_worktree(self, task_id: str, base_ref: str = "HEAD") -> Path:
        branch = f"air-{task_id}"
        wt_path = self._repo_root.parent / f"{self._repo_root.name}-air-{task_id}"

        subprocess.run(
            ["git", "-C", str(self._repo_root), "worktree", "add", "-b", branch, str(wt_path), base_ref],
            check=True, capture_output=True, text=True,
        )
        return wt_path

    def merge_back(self, task_id: str, wt_path: Path) -> MergeResult:
        branch = f"air-{task_id}"
        try:
            subprocess.run(
                ["git", "-C", str(self._repo_root), "merge", "--no-ff", branch],
                check=True, capture_output=True, text=True,
            )
            return MergeResult(task_id=task_id, successful=True, conflicts=[])
        except subprocess.CalledProcessError as exc:
            conflicts = self._parse_conflicts(exc.stderr)
            self._abort_merge()
            return MergeResult(task_id=task_id, successful=False, conflicts=conflicts)

    def cleanup(self, task_id: str, wt_path: Path) -> None:
        branch = f"air-{task_id}"
        try:
            subprocess.run(
                ["git", "-C", str(self._repo_root), "worktree", "remove", str(wt_path), "--force"],
                check=True, capture_output=True, text=True,
            )
            subprocess.run(
                ["git", "-C", str(self._repo_root), "branch", "-D", branch],
                check=True, capture_output=True, text=True,
            )
        except subprocess.CalledProcessError:
            pass

    def _abort_merge(self) -> None:
        try:
            subprocess.run(
                ["git", "-C", str(self._repo_root), "merge", "--abort"],
                check=True, capture_output=True, text=True,
            )
        except subprocess.CalledProcessError:
            pass

    @staticmethod
    def _parse_conflicts(stderr: str) -> list[str]:
        conflicts: list[str] = []
        for line in stderr.splitlines():
            if "CONFLICT" in line:
                conflicts.append(line.strip())
        return conflicts