Blablablab/audio-classification

1"""2Coding Agent Checkpoint Manager3 4Git-based checkpointing for coding agent sessions. Creates lightweight5commits after each file-modifying tool call, enabling rollback to any6previous step.7 8Uses a dedicated git branch (potato-agent-<session_id>) to avoid9interfering with the user's branches.10"""11 12import logging13import os14import subprocess15import time16from dataclasses import dataclass, field17from typing import Dict, List, Optional18 19logger = logging.getLogger(__name__)20 21 22@dataclass23class Checkpoint:24    """A snapshot of the working directory state."""25    checkpoint_id: str  # git commit hash26    step_index: int27    tool_name: str28    description: str29    timestamp: float30    files_changed: List[str] = field(default_factory=list)31 32    def to_dict(self) -> dict:33        return {34            "checkpoint_id": self.checkpoint_id,35            "step_index": self.step_index,36            "tool_name": self.tool_name,37            "description": self.description,38            "timestamp": self.timestamp,39            "files_changed": self.files_changed,40        }41 42 43class CheckpointManager:44    """Manages git-based checkpoints for a coding agent session."""45 46    def __init__(self, working_dir: str, session_id: str):47        self._working_dir = os.path.abspath(working_dir)48        self._session_id = session_id49        self._branch_name = f"potato-agent-{session_id[:12]}"50        self._checkpoints: List[Checkpoint] = []51        self._initialized = False52 53    @property54    def checkpoints(self) -> List[Checkpoint]:55        return list(self._checkpoints)56 57    def init(self) -> bool:58        """Initialize git repo and create session branch.59 60        Returns True if initialization succeeded.61        """62        if self._initialized:63            return True64 65        # Ensure git repo exists66        if not self._is_git_repo():67            try:68                self._run_git("init")69                self._run_git("add", "-A")70                self._run_git("commit", "--allow-empty", "-m", "[potato] init")71            except Exception as e:72                logger.warning(f"Failed to init git repo: {e}")73                return False74 75        # Create session branch from current HEAD76        try:77            current_branch = self._run_git("rev-parse", "--abbrev-ref", "HEAD").strip()78            self._run_git("checkout", "-b", self._branch_name)79        except subprocess.CalledProcessError:80            # Branch might already exist (session restart)81            try:82                self._run_git("checkout", self._branch_name)83            except subprocess.CalledProcessError as e:84                logger.warning(f"Failed to create/checkout session branch: {e}")85                return False86 87        # Create initial checkpoint88        try:89            self._run_git("add", "-A")90            self._run_git("commit", "--allow-empty", "-m",91                          f"[potato] session start {self._session_id[:8]}")92            commit_hash = self._get_head_hash()93            self._checkpoints.append(Checkpoint(94                checkpoint_id=commit_hash,95                step_index=-1,96                tool_name="init",97                description="Session start",98                timestamp=time.time(),99            ))100        except Exception as e:101            logger.warning(f"Failed to create initial checkpoint: {e}")102 103        self._initialized = True104        logger.info(f"CheckpointManager initialized on branch {self._branch_name}")105        return True106 107    def create_checkpoint(self, step_index: int, tool_name: str,108                          description: str = "") -> Optional[str]:109        """Create a checkpoint after a tool execution.110 111        Returns the commit hash, or None if no changes to commit.112        """113        if not self._initialized:114            if not self.init():115                return None116 117        try:118            # Stage all changes119            self._run_git("add", "-A")120 121            # Check if there are changes to commit122            status = self._run_git("status", "--porcelain")123            if not status.strip():124                # No changes, but still record the checkpoint for rollback125                commit_hash = self._get_head_hash()126                self._checkpoints.append(Checkpoint(127                    checkpoint_id=commit_hash,128                    step_index=step_index,129                    tool_name=tool_name,130                    description=description or f"Step {step_index}: {tool_name}",131                    timestamp=time.time(),132                ))133                return commit_hash134 135            # Get list of changed files136            changed = [137                line.split(None, 1)[-1].strip()138                for line in status.strip().split("\n")139                if line.strip()140            ]141 142            # Commit143            msg = f"[potato] step={step_index} tool={tool_name}"144            if description:145                msg += f" {description}"146            self._run_git("commit", "-m", msg)147 148            commit_hash = self._get_head_hash()149            checkpoint = Checkpoint(150                checkpoint_id=commit_hash,151                step_index=step_index,152                tool_name=tool_name,153                description=description or f"Step {step_index}: {tool_name}",154                timestamp=time.time(),155                files_changed=changed,156            )157            self._checkpoints.append(checkpoint)158 159            logger.debug(f"Created checkpoint {commit_hash[:8]} at step {step_index}")160            return commit_hash161 162        except Exception as e:163            logger.warning(f"Failed to create checkpoint: {e}")164            return None165 166    def rollback_to(self, step_index: int) -> bool:167        """Rollback to the checkpoint at the given step index.168 169        Returns True if rollback succeeded.170        """171        # Find the checkpoint172        target = None173        for cp in self._checkpoints:174            if cp.step_index == step_index:175                target = cp176                break177            if cp.step_index <= step_index:178                target = cp  # Use the latest checkpoint at or before step_index179 180        if not target:181            logger.warning(f"No checkpoint found at or before step {step_index}")182            return False183 184        try:185            self._run_git("reset", "--hard", target.checkpoint_id)186 187            # Truncate checkpoint list188            self._checkpoints = [189                cp for cp in self._checkpoints190                if cp.step_index <= step_index191            ]192 193            logger.info(f"Rolled back to step {step_index} (commit {target.checkpoint_id[:8]})")194            return True195 196        except Exception as e:197            logger.error(f"Rollback failed: {e}")198            return False199 200    def get_diff_between(self, from_step: int, to_step: int) -> str:201        """Get the git diff between two checkpoints."""202        from_cp = self._find_checkpoint(from_step)203        to_cp = self._find_checkpoint(to_step)204        if not from_cp or not to_cp:205            return ""206 207        try:208            return self._run_git("diff", from_cp.checkpoint_id, to_cp.checkpoint_id)209        except Exception:210            return ""211 212    def get_diff_since(self, step_index: int) -> str:213        """Get the diff from a checkpoint to current HEAD."""214        cp = self._find_checkpoint(step_index)215        if not cp:216            return ""217        try:218            return self._run_git("diff", cp.checkpoint_id, "HEAD")219        except Exception:220            return ""221 222    def get_file_at(self, step_index: int, file_path: str) -> Optional[str]:223        """Get file contents at a specific checkpoint."""224        cp = self._find_checkpoint(step_index)225        if not cp:226            return None227        try:228            return self._run_git("show", f"{cp.checkpoint_id}:{file_path}")229        except Exception:230            return None231 232    def list_checkpoints(self) -> List[dict]:233        """Return checkpoint metadata as list of dicts."""234        return [cp.to_dict() for cp in self._checkpoints]235 236    def cleanup(self) -> None:237        """Clean up the session branch."""238        if not self._initialized:239            return240 241        try:242            # Switch back to the original branch243            branches = self._run_git("branch", "--list").strip().split("\n")244            main_branch = None245            for b in branches:246                name = b.strip().lstrip("* ")247                if name and name != self._branch_name:248                    main_branch = name249                    break250 251            if main_branch:252                self._run_git("checkout", main_branch)253                self._run_git("branch", "-D", self._branch_name)254                logger.info(f"Cleaned up session branch {self._branch_name}")255        except Exception as e:256            logger.warning(f"Failed to clean up session branch: {e}")257 258    def _find_checkpoint(self, step_index: int) -> Optional[Checkpoint]:259        for cp in self._checkpoints:260            if cp.step_index == step_index:261                return cp262        return None263 264    def _is_git_repo(self) -> bool:265        try:266            self._run_git("rev-parse", "--git-dir")267            return True268        except (subprocess.CalledProcessError, FileNotFoundError):269            return False270 271    def _get_head_hash(self) -> str:272        return self._run_git("rev-parse", "HEAD").strip()273 274    def _run_git(self, *args) -> str:275        result = subprocess.run(276            ["git"] + list(args),277            cwd=self._working_dir,278            capture_output=True,279            text=True,280            timeout=30,281        )282        if result.returncode != 0:283            raise subprocess.CalledProcessError(284                result.returncode, ["git"] + list(args),285                output=result.stdout, stderr=result.stderr,286            )287        return result.stdout288