worktree.py

"""Git worktree lifecycle for the Conscious multi-agent system (Tier 3, Phase 2).

Each agent gets an isolated ``git worktree`` on the shared repo: a separate
working directory sharing the repo's ``.git`` object store. No clone, no
duplication. N agents on one repo with zero file clobbering.

Layout (TIER3_PLAN.md §12)::

    <workspace_sandbox>/
    ├── .brain/                       # canonical brain (shared)
    ├── .git/                         # shared object store
    ├── ... (main working tree)
    └── .worktrees/
        └── <agent-id>/               # one worktree per agent
            ├── .brain/ -> ../../.brain  # symlink to shared brain
            └── ... (agent's isolated working tree)

Tier 1 safety (non-negotiable): worktree operations never touch the remote
(no push/fetch), so no token is needed. The main repo's ``origin`` is already
token-free (Tier 1's ``_strip_remote_token``). Pushes from worktrees go
through ``github_integration.push_to_remote`` — token only in argv.

This module is stdlib-only (subprocess + pathlib). It does NOT import
github_integration (avoids a circular import; the merge path returns a
result that the caller — conscious_routes — uses to decide whether to push).

Security: all user-controlled inputs (agent_id, filenames, resolution keys)
are validated against strict allowlists + path-containment checks before any
filesystem operation. See ``_validate_agent_id`` and ``_safe_join``.
"""
from __future__ import annotations

import os
import re
import subprocess
from pathlib import Path
from typing import Any

# ---------------------------------------------------------------------------
# security helpers
# ---------------------------------------------------------------------------

# Agent IDs are generated by conscious_db._gen_id() as uuid4().hex[:16],
# so they're always 16 lowercase hex chars. But we accept the broader
# cuid-style ids the TS preview generates too. Reject anything with /, ..,
# spaces, or shell metachars — this closes the worktree-creation /
# branch-name / rm -rf path-traversal vectors (audit C6/C7/H1/H2).
_AGENT_ID_RE = re.compile(r"^[A-Za-z0-9_-]{1,64}$")


def _validate_agent_id(agent_id: str) -> str:
    """Reject agent IDs that could escape the worktree path or inject refnames.

    Raises ValueError if the id contains anything outside [A-Za-z0-9_-] or is
    longer than 64 chars. This is the single chokepoint that closes:
      - worktree creation at attacker-chosen paths (audit H1)
      - branch name injection via ``conscious/<agent_id>`` (audit H2)
      - ``rm -rf`` / ``shutil.rmtree`` on arbitrary dirs (audit C6/C7)
    """
    if not agent_id or not _AGENT_ID_RE.match(agent_id):
        raise ValueError(f"invalid agent_id (rejected by allowlist): {agent_id!r}")
    if ".." in agent_id:
        raise ValueError(f"agent_id contains '..': {agent_id!r}")
    return agent_id


def _safe_join(base: Path, user_input: str) -> Path:
    """Join ``user_input`` onto ``base`` and verify the result stays inside ``base``.

    Closes the path-traversal vector where ``inputs.filename = "../../etc/evil"``
    or a conflict-resolution key escapes the sandbox. Rejects:
      - absolute paths (``/etc/passwd``)
      - ``..`` segments that resolve outside ``base``
      - NUL bytes
    """
    if not user_input or "\x00" in user_input:
        raise ValueError("invalid path: empty or contains NUL byte")
    if os.path.isabs(user_input):
        raise ValueError(f"absolute paths not allowed: {user_input!r}")
    target = (base / user_input).resolve()
    base_resolved = base.resolve()
    # the target must be inside base (base itself is ok for dirs)
    try:
        target.relative_to(base_resolved)
    except ValueError:
        raise ValueError(
            f"path escapes sandbox: {user_input!r} resolves outside {base_resolved}")
    return target


def _sanitize_git_stderr(stderr: str) -> str:
    """Redact potential tokens / auth URLs from git stderr before surfacing.

    Guards against the case where git echoes a configured remote URL (which
    could briefly contain a token during the clone window). Patterns:
      - ``https://<token>@host`` → ``https://***@host``
      - ``ghp_<36 chars>`` → ``ghp_***``
      - ``github_pat_<82 chars>`` → ``github_pat_***``
      - ``hf_<30+ chars>`` → ``hf_***``
      - ``AKIA<16 chars>`` → ``AKIA***``
    """
    import re as _re
    s = stderr
    s = _re.sub(r"https://[^@/\s]+@", "https://***@", s)
    s = _re.sub(r"\bghp_[A-Za-z0-9]{36}\b", "ghp_***", s)
    s = _re.sub(r"\bgithub_pat_[A-Za-z0-9_]{82}\b", "github_pat_***", s)
    s = _re.sub(r"\bhf_[A-Za-z0-9]{30,}\b", "hf_***", s)
    s = _re.sub(r"\bAKIA[0-9A-Z]{16}\b", "AKIA***", s)
    return s[:400]


# ---------------------------------------------------------------------------
# low-level git helper
# ---------------------------------------------------------------------------

def _git(sandbox: str | Path, *args: str, timeout: int = 60) -> str:
    """Run a git command in ``sandbox``. Returns stdout. Raises on failure.

    Security: stderr is sanitized via ``_sanitize_git_stderr`` before being
    included in the error message, so tokens / auth URLs never leak into
    exceptions (which can flow into HTTP 500 bodies).
    """
    cmd = ["git", "-C", str(sandbox)] + list(args)
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
    if result.returncode != 0:
        # sanitize BOTH the args (in case an auth URL was passed) and stderr
        safe_args = _sanitize_git_stderr(" ".join(args))
        safe_stderr = _sanitize_git_stderr(result.stderr.strip())
        raise RuntimeError(
            f"git {safe_args} failed (exit {result.returncode}): "
            f"{safe_stderr or _sanitize_git_stderr(result.stdout.strip())}"
        )
    return result.stdout


def _is_git_repo(sandbox: Path) -> bool:
    try:
        _git(sandbox, "rev-parse", "--git-dir")
        return True
    except Exception:
        return False


# ---------------------------------------------------------------------------
# main repo init
# ---------------------------------------------------------------------------

def init_main_repo(sandbox: Path, base_branch: str = "main") -> None:
    """Ensure the sandbox is a git repo with at least one commit on ``base_branch``.

    Idempotent. If the sandbox isn't a repo, init one + create an initial
    commit (empty or a README) so worktrees have a base to branch from.
    If it's already a repo, ensure there's at least one commit (worktrees
    can't be created from an empty repo).
    """
    sandbox.mkdir(parents=True, exist_ok=True)
    if not _is_git_repo(sandbox):
        _git(sandbox, "init", "-b", base_branch)
        # minimal identity so the initial commit succeeds in headless Spaces
        try:
            _git(sandbox, "config", "user.email", "conscious@loom.local")
            _git(sandbox, "config", "user.name", "Conscious")
        except Exception:
            pass
    # ensure at least one commit exists
    try:
        _git(sandbox, "rev-parse", "HEAD")
    except Exception:
        readme = sandbox / "README.md"
        if not readme.exists():
            readme.write_text(
                f"# {sandbox.name}\n\nShared workspace for the Conscious multi-agent system.\n",
                encoding="utf-8",
            )
        _git(sandbox, "add", "README.md")
        _git(sandbox, "commit", "-m", "initial commit", "--allow-empty")
    # rename to base_branch if currently on something else (e.g. master)
    try:
        current = _git(sandbox, "rev-parse", "--abbrev-ref", "HEAD").strip()
        if current != base_branch:
            _git(sandbox, "branch", "-M", base_branch)
    except Exception:
        pass


# ---------------------------------------------------------------------------
# worktree lifecycle
# ---------------------------------------------------------------------------

def worktree_path_for(sandbox: Path, agent_id: str) -> Path:
    _validate_agent_id(agent_id)
    return sandbox / ".worktrees" / agent_id


def branch_name_for(agent_id: str) -> str:
    _validate_agent_id(agent_id)
    return f"conscious/{agent_id}"


def create_worktree(sandbox: Path, agent_id: str,
                    base_branch: str = "main") -> Path:
    """Create a git worktree for ``agent_id`` at ``<sandbox>/.worktrees/<agent_id>``.

    Creates a new branch ``conscious/<agent_id>`` off ``base_branch``.
    Sets up the ``.brain/`` symlink so the worktree shares the canonical brain.
    Returns the worktree path.

    Raises if the worktree already exists or the branch already exists.
    The caller (conscious_routes) is responsible for rolling back the agent
    DB row if this raises.
    """
    _validate_agent_id(agent_id)  # defense-in-depth (worktree_path_for also checks)
    init_main_repo(sandbox, base_branch)
    wt = worktree_path_for(sandbox, agent_id)
    if wt.exists():
        raise RuntimeError(f"worktree already exists: {wt}")
    branch = branch_name_for(agent_id)
    # `git worktree add -b <branch> <path> <base>` — base must be a valid ref
    _git(sandbox, "worktree", "add", "-b", branch, str(wt), base_branch)
    # symlink .brain/ → ../../.brain (relative, survives path layout)
    brain_link = wt / ".brain"
    if not brain_link.exists():
        try:
            os.symlink("../../.brain", str(brain_link))
        except OSError:
            # symlink creation can fail on some filesystems; not fatal —
            # the agent can still read .brain/ via the main sandbox path.
            pass
    return wt


def remove_worktree(sandbox: Path, agent_id: str) -> None:
    """Remove an agent's worktree + delete its branch. Idempotent."""
    _validate_agent_id(agent_id)  # prevents rm -rf on arbitrary paths
    wt = worktree_path_for(sandbox, agent_id)
    branch = branch_name_for(agent_id)
    if wt.exists():
        try:
            _git(sandbox, "worktree", "remove", "--force", str(wt))
        except Exception:
            # if the worktree is locked or busy, prune the metadata
            try:
                _git(sandbox, "worktree", "prune")
            except Exception:
                pass
            # last resort: rm the dir
            import shutil
            shutil.rmtree(wt, ignore_errors=True)
    try:
        _git(sandbox, "branch", "-D", branch)
    except Exception:
        pass  # branch may not exist (already deleted / never created)


def list_worktrees(sandbox: Path) -> list[dict]:
    """List all worktrees. Returns [{path, branch, head, bare}]."""
    if not _is_git_repo(sandbox):
        return []
    try:
        out = _git(sandbox, "worktree", "list", "--porcelain")
    except Exception:
        return []
    items: list[dict] = []
    cur: dict | None = None
    for line in out.splitlines():
        if not line.strip():
            if cur:
                items.append(cur)
                cur = None
            continue
        key, _, val = line.partition(" ")
        if key == "worktree":
            cur = {"path": val}
        elif cur is not None:
            if key == "HEAD":
                cur["head"] = val
            elif key == "branch":
                cur["branch"] = val
            elif key == "bare":
                cur["bare"] = val == "true"
    if cur:
        items.append(cur)
    return items


def prune_worktrees(sandbox: Path) -> None:
    """Prune orphaned worktree metadata (worktrees whose dirs were deleted externally)."""
    if not _is_git_repo(sandbox):
        return
    try:
        _git(sandbox, "worktree", "prune")
    except Exception:
        pass


# ---------------------------------------------------------------------------
# merge agent branch → main
# ---------------------------------------------------------------------------

def agent_branch_status(sandbox: Path, agent_id: str,
                        base_branch: str = "main") -> dict:
    """Return how the agent's branch diverges from base: {ahead, behind, files_changed, diff_summary}."""
    branch = branch_name_for(agent_id)
    try:
        ahead_behind = _git(
            sandbox, "rev-list", "--left-right", "--count",
            f"{base_branch}...{branch}").strip().split()
        ahead = int(ahead_behind[1]) if len(ahead_behind) >= 2 else 0
        behind = int(ahead_behind[0]) if len(ahead_behind) >= 2 else 0
    except Exception:
        ahead = behind = 0
    try:
        diff = _git(sandbox, "diff", "--name-status", f"{base_branch}..{branch}").strip()
        files = [line.split("\t", 1) for line in diff.splitlines() if line.strip()]
    except Exception:
        files = []
    return {
        "branch": branch,
        "ahead": ahead,
        "behind": behind,
        "files_changed": [{"status": f[0], "path": f[1] if len(f) > 1 else ""}
                          for f in files],
    }


def merge_agent_branch(sandbox: Path, agent_id: str,
                       base_branch: str = "main") -> dict:
    """Merge ``conscious/<agent_id>`` into ``base_branch``.

    Returns one of:
      - ``{"status": "clean", "merged_files": [...], "merge_commit": "<sha>"}``
      - ``{"status": "conflicts", "conflicts": [{"file", "ours", "theirs"}], "aborted": True}``

    On conflict: the merge is **aborted** (``git merge --abort``) so the main
    branch stays clean. The caller (conscious_routes) creates one proposal per
    conflicting file for the orchestrator to resolve. The orchestrator then
    calls ``apply_resolution_and_merge`` to write the resolved content onto the
    agent's branch and retry the merge.

    This is the Q3 "merging is a judgment call, not a deterministic function"
    invariant from TIER3_PLAN §1 — conflicts surface to the orchestrator, never
    auto-resolved.
    """
    branch = branch_name_for(agent_id)
    # ensure we're on base_branch in the main worktree
    _git(sandbox, "checkout", base_branch)
    # attempt the merge
    result = subprocess.run(
        ["git", "-C", str(sandbox), "merge", "--no-ff", "--no-commit", "-m",
         f"[conscious] merge agent {agent_id}", branch],
        capture_output=True, text=True, timeout=60,
    )
    if result.returncode == 0:
        # clean merge — finalize the commit
        _git(sandbox, "commit", "--no-edit")
        sha = _git(sandbox, "rev-parse", "HEAD").strip()
        files = _git(sandbox, "diff", "--name-only", "HEAD~1..HEAD").strip().splitlines()
        return {"status": "clean", "merged_files": files, "merge_commit": sha}
    # conflict — extract conflicting files + abort
    conflicts: list[dict] = []
    try:
        status = _git(sandbox, "status", "--porcelain")
        for line in status.splitlines():
            if line.startswith(("UU", "AA", "DU", "UD", "AU", "UA", "DD")):
                file = line[3:].strip()
                # path-containment check before reading (audit M4)
                try:
                    ours_path = _safe_join(sandbox, file)
                    theirs_path = _safe_join(worktree_path_for(sandbox, agent_id), file)
                    ours = _safe_read(ours_path)
                    theirs = _safe_read(theirs_path)
                except ValueError:
                    # file path escapes sandbox — skip it (don't read arbitrary files)
                    ours = "(path skipped — outside sandbox)"
                    theirs = "(path skipped — outside sandbox)"
                conflicts.append({"file": file, "ours": ours, "theirs": theirs})
    except Exception:
        pass
    # abort the merge so base_branch stays clean
    try:
        _git(sandbox, "merge", "--abort")
    except Exception:
        pass
    return {"status": "conflicts", "conflicts": conflicts, "aborted": True}


def apply_resolution_and_merge(sandbox: Path, agent_id: str,
                               resolutions: dict[str, str],
                               base_branch: str = "main") -> dict:
    """Apply the orchestrator's resolutions and complete the merge.

    Resolves directly on the main worktree: starts the merge (which will
    conflict again), overwrites each conflicted file with the orchestrator's
    chosen content, stages it, then completes the merge commit. This is the
    standard "resolve conflicts during merge" flow.
    """
    branch = branch_name_for(agent_id)
    _git(sandbox, "checkout", base_branch)
    # start the merge (will conflict — expected)
    import subprocess
    subprocess.run(
        ["git", "-C", str(sandbox), "merge", "--no-ff", "--no-commit", "-m",
         f"[conscious] merge agent {agent_id} (resolved)", branch],
        capture_output=True, text=True, timeout=60)
    # write resolutions to the main worktree — each file path validated via _safe_join
    for file, content in resolutions.items():
        target = _safe_join(sandbox, file)  # path-containment check (audit C4)
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_text(content, encoding="utf-8")
        _git(sandbox, "add", "--", file)
    # complete the merge commit
    try:
        _git(sandbox, "commit", "--no-edit")
    except Exception:
        try:
            _git(sandbox, "commit", "--no-edit", "--allow-empty")
        except Exception:
            try:
                _git(sandbox, "merge", "--abort")
            except Exception:
                pass
            return {"status": "conflicts", "conflicts": [], "aborted": True,
                    "error": "could not complete the merge commit after resolution"}
    sha = _git(sandbox, "rev-parse", "HEAD").strip()
    files = _git(sandbox, "diff", "--name-only", "HEAD~1..HEAD").strip().splitlines()
    return {"status": "clean", "merged_files": files, "merge_commit": sha}


def _safe_read(p: Path) -> str:
    try:
        return p.read_text(encoding="utf-8")[:4000]
    except Exception:
        return ""


# ---------------------------------------------------------------------------
# main-branch file listing (for smoke-test verification)
# ---------------------------------------------------------------------------

def list_files_on_branch(sandbox: Path, branch: str = "main",
                         prefix: str = "") -> list[str]:
    """List files on ``branch`` (optionally under ``prefix``). Excludes .git/.brain/.worktrees."""
    try:
        out = _git(sandbox, "ls-tree", "-r", "--name-only", branch)
    except Exception:
        return []
    files = []
    for line in out.splitlines():
        if not line.strip():
            continue
        if line.startswith((".git/", ".worktrees/", ".brain/")):
            continue
        if prefix and not line.startswith(prefix):
            continue
        files.append(line)
    return files


def read_file_on_branch(sandbox: Path, file: str,
                        branch: str = "main") -> str:
    """Read a file's content at ``branch`` HEAD."""
    try:
        return _git(sandbox, "show", f"{branch}:{file}")
    except Exception:
        return ""