Agentswarm/agent/git_operations.py

"""Git operations for the agent workspace.

Merged from agent_swarm_v4:
- Adds repository-root discovery (`repo_root`) so per-task child workspaces can still
  operate inside a parent Git checkout (all git commands run at the repo toplevel)
- Keeps the original branch/commit/push API so the protocol stays compatible
"""
import asyncio
import logging
import os
import shutil
import subprocess
import time
from pathlib import Path
from typing import Optional
from urllib.parse import quote, urlsplit, urlunsplit

logger = logging.getLogger(__name__)


class GitOperations:
    """Handles Git operations for agent workspace."""

    def __init__(self, workspace_dir: str, agent_id: str):
        self.workspace_dir = str(Path(workspace_dir))
        self.agent_id = agent_id
        self.result_branch = f"agent-{agent_id}-results"
        self.base_branch = os.getenv("GIT_BASE_BRANCH", "main")

    def _authenticated_repo_url(self, repo_url: str) -> str:
        username = os.getenv("GIT_USERNAME")
        password = os.getenv("GIT_PASSWORD") or os.getenv("GIT_TOKEN")
        if not username or not password:
            return repo_url

        parsed = urlsplit(repo_url)
        if parsed.scheme not in {"http", "https"} or "@" in parsed.netloc:
            return repo_url

        userinfo = f"{quote(username)}:{quote(password)}"
        return urlunsplit((
            parsed.scheme,
            f"{userinfo}@{parsed.netloc}",
            parsed.path,
            parsed.query,
            parsed.fragment,
        ))

    async def repo_root(self) -> Optional[str]:
        result = await self._run_git_command(
            ["git", "rev-parse", "--show-toplevel"],
            cwd=self.workspace_dir,
        )
        if result.returncode != 0:
            return None
        return result.stdout.strip() or None

    async def is_git_workspace(self) -> bool:
        result = await self._run_git_command(
            ["git", "rev-parse", "--is-inside-work-tree"],
            cwd=self.workspace_dir,
        )
        return result.returncode == 0 and result.stdout.strip() == "true"

    async def clone_workspace(self, repo_url: str) -> bool:
        try:
            logger.info(f"Cloning repository from {repo_url}")

            workspace = Path(self.workspace_dir)
            workspace.parent.mkdir(parents=True, exist_ok=True)

            if workspace.exists():
                if await self.is_git_workspace():
                    logger.info("Workspace already contains a Git checkout; refreshing existing clone")
                    return await self._refresh_existing_workspace(repo_url)

                if any(workspace.iterdir()):
                    logger.warning("Workspace directory exists and is not empty; clearing before clone")
                    self._clear_workspace_dir(workspace)

            clone_url = self._authenticated_repo_url(repo_url)
            base_commit = (os.getenv("GIT_BASE_COMMIT") or "").strip()

            if base_commit:
                # SWE-bench/benchmark: work from the EXACT base_commit, not default HEAD. Shallow-fetch
                # just that commit (fast + small even for huge repos) and check it out.
                workspace.mkdir(parents=True, exist_ok=True)
                steps = [
                    ["git", "init", "-q", self.workspace_dir],
                    ["git", "-C", self.workspace_dir, "remote", "add", "origin", clone_url],
                    ["git", "-C", self.workspace_dir, "fetch", "--depth", "1", "origin", base_commit],
                    ["git", "-C", self.workspace_dir, "checkout", "-q", "FETCH_HEAD"],
                ]
                for step in steps:
                    r = await self._run_git_command(step)
                    if r.returncode != 0:
                        logger.error(f"base_commit shallow checkout failed at {step[1:4]}: {r.stderr}")
                        return False
                logger.info(f"Checked out base_commit {base_commit[:12]} (shallow) at {self.workspace_dir}")
            else:
                result = await self._run_git_command([
                    "git", "clone", clone_url, self.workspace_dir
                ])
                if result.returncode != 0:
                    logger.error(f"Git clone failed: {result.stderr}")
                    return False
                logger.info(f"Successfully cloned repository to {self.workspace_dir}")

            await self._configure_git_user()
            await self._ensure_baseline_commit()
            return True

        except Exception as e:
            logger.error(f"Error cloning workspace: {e}")
            return False

    def _clear_workspace_dir(self, workspace: Path):
        for child in workspace.iterdir():
            if child.is_dir() and not child.is_symlink():
                shutil.rmtree(child)
            else:
                child.unlink()

    async def _configure_git_user(self):
        cwd = await self.repo_root() or self.workspace_dir
        await self._run_git_command([
            "git", "config", "user.name", f"Agent {self.agent_id}"
        ], cwd=cwd)
        await self._run_git_command([
            "git", "config", "user.email", f"{self.agent_id}@agent.local"
        ], cwd=cwd)

    async def _refresh_existing_workspace(self, repo_url: str) -> bool:
        cwd = await self.repo_root() or self.workspace_dir
        clone_url = self._authenticated_repo_url(repo_url)
        await self._run_git_command(["git", "remote", "set-url", "origin", clone_url], cwd=cwd)
        await self._configure_git_user()

        fetch = await self._run_git_command(["git", "fetch", "origin"], cwd=cwd)
        if fetch.returncode != 0:
            logger.error(f"Git fetch failed: {fetch.stderr}")
            return False

        base_ref = f"origin/{self.base_branch}"
        base_check = await self._run_git_command(["git", "rev-parse", "--verify", base_ref], cwd=cwd)
        if base_check.returncode == 0:
            reset = await self._run_git_command(["git", "reset", "--hard", base_ref], cwd=cwd)
            if reset.returncode != 0:
                logger.error(f"Git reset failed: {reset.stderr}")
                return False

        clean = await self._run_git_command(["git", "clean", "-fd"], cwd=cwd)
        if clean.returncode != 0:
            logger.error(f"Git clean failed: {clean.stderr}")
            return False

        await self._ensure_baseline_commit()
        return True

    async def _has_commits(self, cwd: str) -> bool:
        """True if the checkout has at least one commit (a valid, born HEAD)."""
        result = await self._run_git_command(
            ["git", "rev-parse", "--verify", "--quiet", "HEAD"], cwd=cwd
        )
        return result.returncode == 0

    async def _ensure_baseline_commit(self) -> bool:
        """Seed an initial commit when the clone is an empty repo (unborn HEAD, no commits).

        An empty remote has no default branch and no commits, so the clone's HEAD is an
        unborn ref. Downstream `git worktree add ... HEAD` / `checkout -B ... HEAD` then fail
        with "invalid reference: HEAD", and the agent silently degrades to "without git push",
        so results never reach the remote (agent_swarm#73). Planting one empty baseline commit
        on the configured base branch gives those operations a valid base and lets the result
        branch push back normally. No-op when the repo already has commits. Returns True if the
        repo ends up with a born HEAD.
        """
        cwd = await self.repo_root() or self.workspace_dir
        if await self._has_commits(cwd):
            return True

        logger.warning(
            "Cloned repository is empty (no commits); seeding baseline commit on '%s' "
            "so task worktrees can be created and pushed (agent_swarm#73)",
            self.base_branch,
        )
        # Land the first commit on the configured base branch even though HEAD is still unborn.
        await self._run_git_command(["git", "checkout", "-B", self.base_branch], cwd=cwd)
        commit = await self._run_git_command(
            ["git", "commit", "--allow-empty", "-m",
             "chore: initialize empty repository baseline"],
            cwd=cwd,
        )
        if commit.returncode != 0:
            logger.error(f"Failed to seed baseline commit for empty repository: {commit.stderr}")
            return False

        logger.info(f"Seeded baseline commit on branch '{self.base_branch}' for empty repository")
        return True

    async def create_result_branch(self, task_id: Optional[str] = None) -> bool:
        try:
            cwd = await self.repo_root() or self.workspace_dir
            if task_id:
                safe_task_id = task_id.replace("/", "-")[:12]
                timestamp = int(time.time())
                self.result_branch = f"agent/{self.agent_id}/{safe_task_id}-{timestamp}"

            await self._run_git_command(["git", "fetch", "origin"], cwd=cwd)

            base_ref = f"origin/{self.base_branch}"
            base_check = await self._run_git_command(["git", "rev-parse", "--verify", base_ref], cwd=cwd)
            if base_check.returncode != 0:
                base_ref = "HEAD"
                logger.warning(f"Base branch origin/{self.base_branch} not found; creating result branch from HEAD")

            result = await self._run_git_command(["git", "checkout", "-B", self.result_branch, base_ref], cwd=cwd)
            if result.returncode != 0:
                logger.error(f"Failed to create branch: {result.stderr}")
                return False

            logger.info(f"Created result branch: {self.result_branch}")
            return True
        except Exception as e:
            logger.error(f"Error creating result branch: {e}")
            return False

    async def add_task_worktree(self, path: str, task_id: str) -> Optional[str]:
        """Create an isolated git worktree at `path` on a fresh per-task result branch.

        Cut from this agent's clone (off origin/<base> or HEAD), the worktree has the full repo
        contents plus its own index/HEAD, so one agent can run several tasks concurrently without
        them clobbering a shared checkout. Returns the branch name, or None on failure.
        """
        try:
            root = await self.repo_root() or self.workspace_dir
            safe_task_id = task_id.replace("/", "-")[:12]
            branch = f"agent/{self.agent_id}/{safe_task_id}-{int(time.time())}"

            await self._run_git_command(["git", "fetch", "origin"], cwd=root)
            base_ref = f"origin/{self.base_branch}"
            if (await self._run_git_command(["git", "rev-parse", "--verify", base_ref], cwd=root)).returncode != 0:
                base_ref = "HEAD"
                logger.warning(f"Base branch origin/{self.base_branch} not found; worktree from HEAD")

            Path(path).parent.mkdir(parents=True, exist_ok=True)
            result = await self._run_git_command(
                ["git", "worktree", "add", "-b", branch, path, base_ref], cwd=root,
            )
            if result.returncode != 0:
                logger.error(f"git worktree add failed: {result.stderr}")
                return None
            logger.info(f"Created task worktree {path} on branch {branch}")
            return branch
        except Exception as e:
            logger.error(f"Error creating task worktree: {e}")
            return None

    async def remove_task_worktree(self, path: str) -> None:
        """Tear down a per-task worktree (best-effort). The branch ref is kept (already pushed)."""
        try:
            root = await self.repo_root() or self.workspace_dir
            await self._run_git_command(["git", "worktree", "remove", "--force", path], cwd=root)
            logger.info(f"Removed task worktree {path}")
        except Exception as e:
            logger.warning(f"Failed to remove task worktree {path}: {e}")

    async def commit_changes(self, message: str) -> Optional[str]:
        try:
            cwd = await self.repo_root() or self.workspace_dir
            status_result = await self._run_git_command(["git", "status", "--porcelain"], cwd=cwd)
            if not status_result.stdout.strip():
                logger.info("No changes to commit")
                return None

            await self._run_git_command(["git", "add", "-A"], cwd=cwd)
            commit_result = await self._run_git_command(["git", "commit", "-m", message], cwd=cwd)
            if commit_result.returncode != 0:
                logger.error(f"Git commit failed: {commit_result.stderr}")
                return None

            sha_result = await self._run_git_command(["git", "rev-parse", "HEAD"], cwd=cwd)
            commit_sha = sha_result.stdout.strip()
            logger.info(f"Committed changes: {commit_sha[:8]} - {message}")
            return commit_sha
        except Exception as e:
            logger.error(f"Error committing changes: {e}")
            return None

    async def push_results(self) -> Optional[str]:
        try:
            cwd = await self.repo_root() or self.workspace_dir
            result = await self._run_git_command(["git", "push", "-u", "origin", self.result_branch], cwd=cwd)
            if result.returncode != 0:
                logger.error(f"Git push failed: {result.stderr}")
                return None
            logger.info(f"Pushed results to branch: {self.result_branch}")
            return self.result_branch
        except Exception as e:
            logger.error(f"Error pushing results: {e}")
            return None

    async def _run_git_command(
        self,
        command: list[str],
        cwd: Optional[str] = None,
    ) -> subprocess.CompletedProcess:
        process = await asyncio.create_subprocess_exec(
            *command,
            cwd=cwd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        stdout, stderr = await process.communicate()
        return subprocess.CompletedProcess(
            args=command,
            returncode=process.returncode,
            stdout=stdout.decode(),
            stderr=stderr.decode(),
        )