feat: Make user configurable

docker_agent_sandbox/sandbox.py

@@ -4,20 +4,16 @@ from __future__ import annotations
 
 import io
 import select as _select
-import socket as _socket
 import tarfile
-from contextlib import ExitStack
-from select import select as _original_select
 from pathlib import Path
 from typing import TYPE_CHECKING
-from unittest.mock import patch
+import docker
+import docker.errors
+from loguru import logger
+
 
 _original_poll = getattr(_select, "poll", None)
 
-import docker
-import docker.errors
-from docker.utils.socket import consume_socket_output, demux_adaptor, frames_iter
-from loguru import logger
 
 if TYPE_CHECKING:
     import docker.models.containers
@@ -50,6 +46,8 @@ class DockerSandbox:
         cpu_limit: float = 8,
         memory_limit: str = "16g",
         command: str | None = None,
+        uid: int | None = None,
+        gid: int | None = None,
     ) -> None:
         self.container_name = container_name
         self._image = image
@@ -65,9 +63,21 @@ class DockerSandbox:
         self._nano_cpus = int(cpu_limit * 1e9)
         self._memory_limit = memory_limit
         self._command = command
+        self._uid = uid
+        self._gid = gid
         self._client: docker.DockerClient = docker.from_env()
         self._container: docker.models.containers.Container | None = None
 
+    @property
+    def _user(self) -> str | None:
+        """Docker ``user`` spec ("uid" or "uid:gid"), or None to use the image default."""
+        if self._uid is None and self._gid is None:
+            return None
+        uid = "" if self._uid is None else str(self._uid)
+        if self._gid is None:
+            return uid
+        return f"{uid}:{self._gid}"
+
     def __enter__(self) -> "DockerSandbox":
         return self
 
@@ -135,9 +145,13 @@ class DockerSandbox:
             run_kwargs["security_opt"] = self._security_opt
         run_kwargs["nano_cpus"] = self._nano_cpus
         run_kwargs["mem_limit"] = self._memory_limit
+        if self._user is not None:
+            run_kwargs["user"] = self._user
 
         try:
-            self._container = self._client.containers.run(self._image, self._command, **run_kwargs)
+            self._container = self._client.containers.run(
+                self._image, self._command, **run_kwargs
+            )
         except docker.errors.ImageNotFound:
             raise RuntimeError(
                 f"Image {self._image!r} not found locally. "
@@ -247,13 +261,20 @@ class DockerSandbox:
     # Command execution
     # ------------------------------------------------------------------
 
-    def exec(self, command: str, timeout: int = 120, kill_grace: int = 5) -> tuple[int, str]:
+    def exec(
+        self, command: str, timeout: int = 120, kill_grace: int = 5
+    ) -> tuple[int, str]:
         if self._container is None:
             return 1, "Sandbox container is not running."
 
         wrapped = [
-            "timeout", "--kill-after", f"{kill_grace}s", f"{timeout}s",
+            "timeout",
-            "bash", "-c", command,
+            "--kill-after",
+            f"{kill_grace}s",
+            f"{timeout}s",
+            "bash",
+            "-c",
+            command,
         ]
         kwargs = {}
         if self._working_dir is not None: