gdmcode 0.1.0__py3-none-any.whl

src/runtime/branch_farm.py

@@ -0,0 +1,372 @@
+"""Branch Farm — parallel exploration of N independent solution approaches.
+
+Workflow::
+
+    farm = BranchFarm(git_workflow, config=FarmConfig(n_branches=3))
+    best = farm.run("Fix the failing auth tests")
+
+Each approach runs in its own ``git worktree`` via a ``ProcessPoolExecutor``
+so workers have completely isolated filesystems.  The winning branch is
+squash- or cherry-pick-merged back onto the originating branch.
+
+Requires ``autonomy_level >= 4`` to run unattended.
+"""
+from __future__ import annotations
+
+import logging
+import time
+import uuid
+from concurrent.futures import ProcessPoolExecutor, as_completed, Future
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+from src.exceptions import DirtyWorkingTreeError, MergeConflictError
+from src.git_workflow import GitWorkflow, WorktreeInfo
+
+__all__ = [
+    "BranchFarm",
+    "BranchResult",
+    "BranchScore",
+    "FarmConfig",
+    "SCORE_WEIGHTS",
+]
+
+log = logging.getLogger(__name__)
+
+_FARM_BRANCH_PREFIX = "gdm/farm"
+_WORKTREE_ROOT = Path(".context-memory/worktrees")
+
+SCORE_WEIGHTS: dict[str, float] = {
+    "test_pass_rate":    0.40,
+    "lint_clean":        0.20,
+    "security_clean":    0.20,
+    "diff_size_inv":     0.10,
+    "conflict_risk_inv": 0.10,
+}
+
+
+# ---------------------------------------------------------------------------
+# Data classes
+# ---------------------------------------------------------------------------
+
+@dataclass
+class FarmConfig:
+    """Configuration for a BranchFarm run."""
+    n_branches: int = 3
+    max_workers: int = 3
+    timeout_secs: float = 300.0
+    autonomy_level: int = 4
+    merge_strategy: str = "squash"   # "squash" | "cherry-pick" | "patch"
+
+
+@dataclass
+class BranchResult:
+    """Outcome from one worker branch."""
+    branch: str
+    worktree: WorktreeInfo
+    ok: bool
+    error: str | None = None
+    tests_passed: int = 0
+    tests_total: int = 0
+    lint_errors: int = 0
+    security_findings: int = 0
+    diff_lines_added: int = 0
+    diff_lines_removed: int = 0
+    changed_files: int = 0
+    conflict_risk: float = 0.0   # 0.0–1.0; populated by dry-run merge analysis
+    head_sha: str | None = None  # HEAD sha on the farm branch after the run
+
+
+@dataclass
+class BranchScore:
+    """Composite quality score for one branch result."""
+    branch: str
+    tests_passed: int = 0
+    tests_total: int = 0
+    lint_errors: int = 0
+    security_findings: int = 0
+    diff_lines_added: int = 0
+    diff_lines_removed: int = 0
+    changed_files: int = 0
+    conflict_risk: float = 0.0   # 0.0–1.0
+    composite: float = 0.0       # weighted sum; higher = better
+
+
+# ---------------------------------------------------------------------------
+# Module-level worker function (must be pickleable for ProcessPoolExecutor)
+# ---------------------------------------------------------------------------
+
+def _worker_run(
+    worktree_path: Path,
+    branch: str,
+    task: str,
+    extra: dict[str, Any] | None = None,
+) -> BranchResult:
+    """Top-level pickleable entry point for each farm worker subprocess.
+
+    Override for testing by setting *extra["_worker_fn"]*.  In production
+    this calls ``_run_gdm_in_worktree()``.
+    """
+    if extra and "_worker_fn" in extra:
+        fn = extra["_worker_fn"]
+        return fn(worktree_path, branch, task)
+    return _run_gdm_in_worktree(worktree_path, branch, task)
+
+
+def _run_gdm_in_worktree(
+    worktree_path: Path,
+    branch: str,
+    task: str,
+) -> BranchResult:
+    """Run gdm inside *worktree_path* and return scored result.
+
+    Subprocesses do NOT have access to the parent's in-memory state; all
+    information must come through function arguments and the filesystem.
+    """
+    # Lazy imports to avoid circular dependencies inside the subprocess
+    import subprocess  # noqa: PLC0415
+
+    from src.git_workflow import GitWorkflow
+
+    wf = GitWorkflow(worktree_path)
+    wt_info = WorktreeInfo(
+        branch=branch,
+        path=worktree_path,
+        db_path=worktree_path / ".context-memory" / "gdm.db",
+    )
+
+    # --- run pytest to collect pass/fail counts ---
+    pytest_result = subprocess.run(
+        ["python", "-m", "pytest", "--tb=no", "-q"],
+        cwd=str(worktree_path),
+        capture_output=True,
+        text=True,
+    )
+    tests_passed = 0
+    tests_total = 0
+    for line in pytest_result.stdout.splitlines():
+        import re
+        m = re.search(r"(\d+) passed", line)
+        if m:
+            tests_passed = int(m.group(1))
+        m2 = re.search(r"(\d+) failed", line)
+        if m2:
+            tests_total += int(m2.group(1))
+    tests_total += tests_passed
+
+    # --- diff size vs base HEAD ---
+    diff_result = subprocess.run(
+        ["git", "diff", "--stat", "HEAD"],
+        cwd=str(worktree_path),
+        capture_output=True,
+        text=True,
+    )
+    lines_added = 0
+    lines_removed = 0
+    changed_files = 0
+    for line in diff_result.stdout.splitlines():
+        import re
+        m = re.search(r"(\d+) insertion", line)
+        if m:
+            lines_added = int(m.group(1))
+        m = re.search(r"(\d+) deletion", line)
+        if m:
+            lines_removed = int(m.group(1))
+        m = re.search(r"(\d+) file", line)
+        if m:
+            changed_files = int(m.group(1))
+
+    # --- current HEAD sha ---
+    sha_result = subprocess.run(
+        ["git", "rev-parse", "HEAD"],
+        cwd=str(worktree_path),
+        capture_output=True,
+        text=True,
+    )
+    head_sha = sha_result.stdout.strip() or None
+
+    return BranchResult(
+        branch=branch,
+        worktree=wt_info,
+        ok=pytest_result.returncode == 0,
+        tests_passed=tests_passed,
+        tests_total=tests_total,
+        diff_lines_added=lines_added,
+        diff_lines_removed=lines_removed,
+        changed_files=changed_files,
+        head_sha=head_sha,
+    )
+
+
+# ---------------------------------------------------------------------------
+# BranchFarm
+# ---------------------------------------------------------------------------
+
+class BranchFarm:
+    """Orchestrates parallel branch exploration and selects the best result.
+
+    Args:
+        git_workflow: GitWorkflow for the project root.
+        config: FarmConfig tuning parameters.
+    """
+
+    def __init__(
+        self,
+        git_workflow: GitWorkflow,
+        config: FarmConfig | None = None,
+        *,
+        _worker_fn: Any = None,    # injectable for tests
+    ) -> None:
+        self._wf = git_workflow
+        self._config = config or FarmConfig()
+        self._worker_fn = _worker_fn   # override _run_gdm_in_worktree in tests
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def run(self, task: str) -> BranchResult:
+        """Guard → create → dispatch → score → merge → cleanup."""
+        self._require_autonomy()
+        task_id = uuid.uuid4().hex[:8]
+        branches: list[WorktreeInfo] = []
+        try:
+            branches = self.create_branches(self._config.n_branches, task_id)
+            results = self._dispatch(branches, task)
+            scores = self.evaluate_results(results)
+            best_score = scores[0]
+            best_result = next(r for r in results if r.branch == best_score.branch)
+            self._merge_best(best_result)
+            return best_result
+        finally:
+            self._cleanup(branches)
+
+    def create_branches(self, n: int, task_id: str) -> list[WorktreeInfo]:
+        """Guard dirty tree, then create N worktrees on fresh farm branches."""
+        if not self._wf.is_clean():
+            raise DirtyWorkingTreeError()
+        worktrees: list[WorktreeInfo] = []
+        root = self._wf._root
+        worktree_root = root / _WORKTREE_ROOT
+        for i in range(n):
+            branch = f"{_FARM_BRANCH_PREFIX}-{task_id}-{i}"
+            path = worktree_root / f"{task_id}-{i}"
+            wt = self._wf.create_worktree(branch, path)
+            worktrees.append(wt)
+            log.info("Created worktree %s on branch %s", path, branch)
+        return worktrees
+
+    def evaluate_results(self, results: list[BranchResult]) -> list[BranchScore]:
+        """Score each result and return sorted list (best first)."""
+        scores: list[BranchScore] = []
+        for r in results:
+            if not r.ok:
+                scores.append(BranchScore(branch=r.branch, composite=-1.0))
+                continue
+            test_rate = r.tests_passed / max(r.tests_total, 1)
+            lint_score = 1.0 / (1.0 + r.lint_errors)
+            sec_score = 1.0 / (1.0 + r.security_findings)
+            diff_score = 1.0 / (1.0 + r.diff_lines_added + r.diff_lines_removed)
+            risk_score = 1.0 - r.conflict_risk
+            composite = (
+                SCORE_WEIGHTS["test_pass_rate"] * test_rate
+                + SCORE_WEIGHTS["lint_clean"] * lint_score
+                + SCORE_WEIGHTS["security_clean"] * sec_score
+                + SCORE_WEIGHTS["diff_size_inv"] * diff_score
+                + SCORE_WEIGHTS["conflict_risk_inv"] * risk_score
+            )
+            scores.append(BranchScore(
+                branch=r.branch,
+                tests_passed=r.tests_passed,
+                tests_total=r.tests_total,
+                lint_errors=r.lint_errors,
+                security_findings=r.security_findings,
+                diff_lines_added=r.diff_lines_added,
+                diff_lines_removed=r.diff_lines_removed,
+                changed_files=r.changed_files,
+                conflict_risk=r.conflict_risk,
+                composite=composite,
+            ))
+        # Primary: composite desc; secondary: smaller diff (tiebreak)
+        return sorted(
+            scores,
+            key=lambda s: (s.composite, -(s.diff_lines_added + s.diff_lines_removed)),
+            reverse=True,
+        )
+
+    def pick_best(self, results: list[BranchResult]) -> BranchResult:
+        """Return the BranchResult with the highest composite score."""
+        scores = self.evaluate_results(results)
+        best_branch = scores[0].branch
+        return next(r for r in results if r.branch == best_branch)
+
+    # ------------------------------------------------------------------
+    # Private helpers
+    # ------------------------------------------------------------------
+
+    def _require_autonomy(self) -> None:
+        if self._config.autonomy_level < 4:
+            raise PermissionError(
+                f"Branch farm requires autonomy_level >= 4 "
+                f"(current: {self._config.autonomy_level}). "
+                "Set autonomy_level=4 in FarmConfig or approve the run."
+            )
+
+    def _dispatch(
+        self,
+        worktrees: list[WorktreeInfo],
+        task: str,
+    ) -> list[BranchResult]:
+        """Submit all workers to a ProcessPoolExecutor, collect results."""
+        results: list[BranchResult] = []
+        extra = {"_worker_fn": self._worker_fn} if self._worker_fn else None
+
+        with ProcessPoolExecutor(max_workers=self._config.max_workers) as pool:
+            futures: dict[Future[BranchResult], WorktreeInfo] = {
+                pool.submit(_worker_run, wt.path, wt.branch, task, extra): wt
+                for wt in worktrees
+            }
+            for future in as_completed(futures, timeout=self._config.timeout_secs):
+                wt = futures[future]
+                try:
+                    result = future.result()
+                except Exception as exc:
+                    log.error("Worker %s failed: %s", wt.branch, exc)
+                    result = BranchResult(
+                        branch=wt.branch,
+                        worktree=wt,
+                        ok=False,
+                        error=str(exc),
+                    )
+                results.append(result)
+        return results
+
+    def _merge_best(self, best: BranchResult) -> None:
+        """Dry-run check then merge the winning branch back to HEAD."""
+        dry_run = self._wf.merge_tree_dry_run(best.branch)
+        if "<<<<<<<" in dry_run:
+            raise MergeConflictError(best.branch, "Use --no-farm-auto-merge to resolve manually.")
+
+        strategy = self._config.merge_strategy
+        if strategy == "squash":
+            self._wf.squash_merge(
+                best.branch,
+                message=f"[gdm-farm] squash-merge {best.branch} (tests={best.tests_passed}/{best.tests_total})",
+            )
+        elif strategy == "cherry-pick":
+            if best.head_sha:
+                self._wf.cherry_pick(best.head_sha)
+            else:
+                raise MergeConflictError(best.branch, "No HEAD sha available for cherry-pick.")
+        else:
+            raise ValueError(f"Unknown merge_strategy: {strategy!r}")
+
+    def _cleanup(self, worktrees: list[WorktreeInfo]) -> None:
+        """Remove all worktrees and farm branches (best-effort)."""
+        for wt in worktrees:
+            try:
+                self._wf.remove_worktree(wt.path)
+                self._wf.delete_branch(wt.branch)
+            except Exception as exc:
+                log.warning("Cleanup failed for %s: %s", wt.path, exc)