File size: 5,407 Bytes
aceb1b2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
"""
Coding Agent Sandbox Manager

Manages isolated working directories for coding agent sessions.
Supports three modes:
- worktree: git worktree (lightweight copy, requires git repo)
- docker: Docker container with mounted workspace
- direct: No isolation (works directly in working_dir)
"""

import logging
import os
import shutil
import subprocess
import uuid
from typing import Optional

logger = logging.getLogger(__name__)


class SandboxManager:
    """Manages sandboxed working directories for agent sessions."""

    def __init__(self, mode: str = "worktree", base_dir: str = "."):
        """Initialize the sandbox manager.

        Args:
            mode: Sandbox mode — "worktree", "docker", or "direct"
            base_dir: Base directory for creating sandboxes
        """
        if mode not in ("worktree", "docker", "direct"):
            raise ValueError(f"Invalid sandbox mode: {mode}. Must be worktree, docker, or direct.")
        self._mode = mode
        self._base_dir = os.path.abspath(base_dir)
        self._sandbox_dir: Optional[str] = None
        self._session_id: Optional[str] = None
        self._worktree_branch: Optional[str] = None

    @property
    def working_dir(self) -> str:
        """The working directory for the agent."""
        return self._sandbox_dir or self._base_dir

    @property
    def mode(self) -> str:
        return self._mode

    def create(self, session_id: str) -> str:
        """Create a sandbox for the given session.

        Returns:
            The working directory path.
        """
        self._session_id = session_id

        if self._mode == "worktree":
            return self._create_worktree(session_id)
        elif self._mode == "docker":
            return self._create_docker(session_id)
        else:  # direct
            self._sandbox_dir = self._base_dir
            return self._base_dir

    def cleanup(self) -> None:
        """Clean up the sandbox."""
        if self._mode == "worktree":
            self._cleanup_worktree()
        elif self._mode == "docker":
            self._cleanup_docker()
        # direct mode: nothing to clean up

    def _create_worktree(self, session_id: str) -> str:
        """Create a git worktree for isolation."""
        # Check if base_dir is a git repo
        try:
            subprocess.run(
                ["git", "rev-parse", "--git-dir"],
                cwd=self._base_dir, capture_output=True, check=True,
            )
        except (subprocess.CalledProcessError, FileNotFoundError):
            logger.warning(
                f"Directory {self._base_dir} is not a git repo. "
                f"Falling back to direct mode."
            )
            self._mode = "direct"
            self._sandbox_dir = self._base_dir
            return self._base_dir

        # Create worktree in a temp location
        branch_name = f"potato-agent-{session_id[:8]}"
        worktree_dir = os.path.join(
            os.path.dirname(self._base_dir),
            f".potato-sandbox-{session_id[:8]}",
        )

        try:
            # Create a new branch from HEAD
            subprocess.run(
                ["git", "branch", branch_name, "HEAD"],
                cwd=self._base_dir, capture_output=True, check=True,
            )

            # Create worktree
            subprocess.run(
                ["git", "worktree", "add", worktree_dir, branch_name],
                cwd=self._base_dir, capture_output=True, check=True,
            )

            self._sandbox_dir = worktree_dir
            self._worktree_branch = branch_name
            logger.info(f"Created git worktree sandbox at {worktree_dir}")
            return worktree_dir

        except subprocess.CalledProcessError as e:
            logger.warning(f"Failed to create worktree: {e}. Falling back to direct mode.")
            self._mode = "direct"
            self._sandbox_dir = self._base_dir
            return self._base_dir

    def _cleanup_worktree(self) -> None:
        """Remove the git worktree and branch."""
        if not self._sandbox_dir or self._sandbox_dir == self._base_dir:
            return

        try:
            # Remove worktree
            subprocess.run(
                ["git", "worktree", "remove", self._sandbox_dir, "--force"],
                cwd=self._base_dir, capture_output=True,
            )
            logger.info(f"Removed worktree at {self._sandbox_dir}")
        except Exception as e:
            logger.warning(f"Failed to remove worktree: {e}")
            # Manual cleanup
            if os.path.exists(self._sandbox_dir):
                shutil.rmtree(self._sandbox_dir, ignore_errors=True)

        # Clean up the branch
        if self._worktree_branch:
            try:
                subprocess.run(
                    ["git", "branch", "-D", self._worktree_branch],
                    cwd=self._base_dir, capture_output=True,
                )
            except Exception:
                pass

    def _create_docker(self, session_id: str) -> str:
        """Create a Docker container for maximum isolation."""
        # For Phase 4 — placeholder
        logger.warning("Docker sandbox not yet implemented. Using direct mode.")
        self._mode = "direct"
        self._sandbox_dir = self._base_dir
        return self._base_dir

    def _cleanup_docker(self) -> None:
        """Remove Docker container."""
        pass  # Phase 4