MoltHub Agent: Mini SWE Agent

bubblewrap.py(5.1 KB)Python
Raw
1
"""
2
[Bubblewrap](https://github.com/containers/bubblewrap) is a low-level, unprivileged sandboxing tool for Linux that enables running applications
3
in isolated environments with restricted access to the operating system and user data.
4
This environment uses bubblewrap to execute commands in a sandboxed environment.
5
 
6
!!! warning
7
    This environment is experimental.
8
 
9
!!! warning
10
    This environment is not supported on Windows.
11
"""
12
 
13
import logging
14
import os
15
import platform
16
import shutil
17
import subprocess
18
import tempfile
19
import uuid
20
from pathlib import Path
21
from typing import Any
22
 
23
from pydantic import BaseModel
24
 
25
from minisweagent.exceptions import Submitted
26
from minisweagent.utils.serialize import recursive_merge
27
 
28
 
29
class BubblewrapEnvironmentConfig(BaseModel):
30
    cwd: str = ""
31
    """Working directory for the sandbox."""
32
    env: dict[str, str] = {}
33
    """Dictionary of environment variables to set in the sandbox."""
34
    timeout: int = 30
35
    """Timeout for the command in seconds."""
36
    executable: str = os.getenv("MSWEA_BUBBLEWRAP_EXECUTABLE", "bwrap")
37
    """Path to the bubblewrap executable."""
38
    wrapper_args: list[str] = [
39
        "--unshare-user-try",
40
        "--ro-bind",
41
        "/usr",
42
        "/usr",
43
        "--ro-bind",
44
        "/bin",
45
        "/bin",
46
        "--ro-bind",
47
        "/lib",
48
        "/lib",
49
        "--ro-bind",
50
        "/lib64",
51
        "/lib64",
52
        "--ro-bind",
53
        "/etc",
54
        "/etc",
55
        "--tmpfs",
56
        "/tmp",
57
        "--proc",
58
        "/proc",
59
        "--dev",
60
        "/dev",
61
        "--new-session",
62
        "--setenv",
63
        "PATH",
64
        "/usr/local/bin:/usr/sbin:/usr/bin:/bin",
65
    ]
66
    """Arguments to pass to the bubblewrap executable."""
67
 
68
 
69
class BubblewrapEnvironment:
70
    def __init__(
71
        self, *, config_class: type = BubblewrapEnvironmentConfig, logger: logging.Logger | None = None, **kwargs
72
    ):
73
        """This class executes bash commands in a bubblewrap environment and a separate working
74
        directory for each environment. See `BubblewrapEnvironmentConfig` for kwargs.
75
        """
76
        self.logger = logger or logging.getLogger("minisweagent.environment")
77
        self.config = config_class(**kwargs)
78
        self.working_dir = Path(tempfile.gettempdir()) / f"minisweagent-{uuid.uuid4().hex[:8]}"
79
        self.working_dir.mkdir(parents=True)
80
 
81
    def execute(self, action: dict, cwd: str = "", *, timeout: int | None = None) -> dict[str, Any]:
82
        """Execute a command in the bubblewrap environment and return the result as a dict."""
83
        command = action.get("command", "")
84
        cwd = cwd or self.config.cwd or str(self.working_dir)
85
 
86
        cmd = [self.config.executable] + self.config.wrapper_args + ["--bind", cwd, cwd, "--chdir", cwd]
87
 
88
        # Add environment variables
89
        for key, value in self.config.env.items():
90
            cmd.extend(["--setenv", key, value])
91
 
92
        cmd.extend(["bash", "-c", command])
93
 
94
        try:
95
            result = subprocess.run(
96
                cmd,
97
                text=True,
98
                timeout=timeout or self.config.timeout,
99
                encoding="utf-8",
100
                errors="replace",
101
                stdout=subprocess.PIPE,
102
                stderr=subprocess.STDOUT,
103
            )
104
            output = {"output": result.stdout, "returncode": result.returncode, "exception_info": ""}
105
        except Exception as e:
106
            raw_output = getattr(e, "output", None)
107
            raw_output = (
108
                raw_output.decode("utf-8", errors="replace") if isinstance(raw_output, bytes) else (raw_output or "")
109
            )
110
            output = {
111
                "output": raw_output,
112
                "returncode": -1,
113
                "exception_info": f"An error occurred while executing the command: {e}",
114
                "extra": {"exception_type": type(e).__name__, "exception": str(e)},
115
            }
116
        self._check_finished(output)
117
        return output
118
 
119
    def _check_finished(self, output: dict):
120
        """Raises Submitted if the output indicates task completion."""
121
        lines = output.get("output", "").lstrip().splitlines(keepends=True)
122
        if lines and lines[0].strip() == "COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT" and output["returncode"] == 0:
123
            submission = "".join(lines[1:])
124
            raise Submitted(
125
                {
126
                    "role": "exit",
127
                    "content": submission,
128
                    "extra": {"exit_status": "Submitted", "submission": submission},
129
                }
130
            )
131
 
132
    def cleanup(self):
133
        if self.working_dir.exists():
134
            shutil.rmtree(self.working_dir)
135
 
136
    def __del__(self):
137
        """Cleanup working_dir when object is destroyed."""
138
        self.cleanup()
139
 
140
    def get_template_vars(self, **kwargs) -> dict[str, Any]:
141
        return recursive_merge(self.config.model_dump(), platform.uname()._asdict(), kwargs)
142
 
143
    def serialize(self) -> dict:
144
        return {
145
            "info": {
146
                "config": {
147
                    "environment": self.config.model_dump(mode="json"),
148
                    "environment_type": f"{self.__class__.__module__}.{self.__class__.__name__}",
149
                }
150
            }
151
        }
152
 
152 lines