Isolate Code Execution in a Container and Limit Side Effects to a Directory

[SanderGi]

Currently code execution with run_python_repl works via exec meaning it has the same file/network/etc. permissions as the python process. There are good reasons to not containerize the entire process (such as tools controlling external software, and also it has to use API keys somewhere). I think it would make sense to isolate just the code execution tool and perhaps give the agent a specific working directory that limits the scope of the files/directories it can create/delete/modify to just that and any subfolders. If we add this as an optional security layer when Docker is installed, this is straightforward to accomplish with a bind volume. Here's a minimal example that allows executing bash/python commands and create files within a playground directory while capturing stdout+stderr+exitcode:

import os, shutil, subprocess
from pathlib import Path
from io import BytesIO

DEFAULT_PLAYGROUND_PATH = os.path.join(
    os.path.dirname(__file__), "agent-playground"
)
DOCKERFILE_PATH = os.path.join(os.path.dirname(__file__), "agent_playground.Dockerfile")
DOCKER_CONTAINER_TAG = "biomni:agent-playground"

class AgentPlayground:
    """
    Isolates a working director for agent use (download files from the internet, run arbitrary commands, and read files)
    Requires docker to be running. By default uses an image with python3.9, git, and bash pre-installed.
    """

    def __init__(
        self,
        playground_path=DEFAULT_PLAYGROUND_PATH,
        base_image="python:3.9-slim-bullseye",
        install="RUN apt-get update && apt-get install -y --no-install-recommends git bash && rm -rf /var/lib/apt/lists/*",
        hide_build_output=True,
    ):
        """Instantiates playground_path as the working directory for the agent"""

        self.playground_path = playground_path

        os.makedirs(self.playground_path, exist_ok=True)

        with open(DOCKERFILE_PATH, "w") as f:
            f.write(f"FROM {base_image}\n{install}\nWORKDIR /app")

        subprocess.run(["docker", "build", "-f", DOCKERFILE_PATH, ".", "-t", DOCKER_CONTAINER_TAG], check=True, capture_output=hide_build_output)

    def clean(self):
        """Clears the content of playground_path"""

        shutil.rmtree(self.playground_path)
        os.makedirs(self.playground_path)

    def _resolve_relative_path(self, relative_path: str):
        path = os.path.join(self.playground_path, relative_path)
        assert Path(path).is_relative_to(self.playground_path), "must not access file outside playground"
        return path

    def open(self, relative_path: str, mode: str):
        """Returns an open file handle for any path inside playground_path"""

        return open(self._resolve_relative_path(relative_path), mode)

    def run(self, cmd: list[str], timeout_seconds=120):
        """Runs an arbitrary command with playground_path as the working directory."""

        res = subprocess.run(["docker", "run", "--mount", f"type=bind,source={self.playground_path},target=/app", DOCKER_CONTAINER_TAG, *cmd], capture_output=True, timeout=timeout_seconds)
        return res.stdout.decode(), res.stderr.decode(), res.returncode


if __name__ == "__main__":
    # Example usage
    playground = AgentPlayground()
    playground.clean()

    with playground.open("test.py", "w") as f:
        f.write("print('hello world!')")

    out, err, returncode = playground.run(["python", "test.py"])
    print("OUTPUT:")
    print(out)
    if returncode != 0:
        print(f"ERROR (exitcode={returncode}):")
        print(err)

There are plenty of other ways to achieve this too if there is some issue with this approach I have not considered. Let me know if this would make sense to submit a PR for.