feat: Replace Docker monkey-patch with timeout

docker_agent_sandbox/sandbox.py

@@ -247,12 +247,15 @@ class DockerSandbox:
     # Command execution
     # ------------------------------------------------------------------
 
-    def exec(self, command: str, timeout: int = 120) -> tuple[int, str]:
+    def exec(self, command: str, timeout: int = 120, kill_grace: int = 5) -> tuple[int, str]:
         """
         Run *command* inside the container via the low-level exec API.
 
         Returns ``(exit_code, combined_stdout_stderr)``.
-        Timeout is enforced via a socket-level timeout on the exec socket.
+
+        Exit code 124 indicates a timeout (137 if escalation to SIGKILL was
+        needed). *command* is passed as a single argv element to ``bash -c``,
+        so shell metacharacters within it need no extra escaping.
         """
         if self._container is None:
             return 1, "Sandbox container is not running."
@@ -261,60 +264,37 @@ class DockerSandbox:
         if self._working_dir is not None:
             create_kwargs["workdir"] = self._working_dir
 
-        # TODO(fragile): timeout enforcement relies on private docker-py internals
+        wrapped = [
-        # (frames_iter, demux_adaptor, consume_socket_output from docker.utils.socket)
+            "timeout",
-        # and monkey-patches select.select / select.poll for the duration of the read
+            "--kill-after",
-        # — not thread-safe if multiple exec() calls run concurrently. Replace when
+            f"{kill_grace}s",
-        # docker-py adds native per-call timeout support.
+            f"{timeout}s",
-        # See https://github.com/docker/docker-py/issues/2651
+            "bash",
-        #
+            "-c",
-        # On Linux docker-py uses select.poll (not select.select), so both are patched.
+            command,
+        ]
+
         try:
             exec_id = self._client.api.exec_create(
                 self._container.id,
-                ["bash", "-c", command],
+                wrapped,
                 stdout=True,
                 stderr=True,
                 **create_kwargs,
             )
             sock = self._client.api.exec_start(exec_id["Id"], socket=True)
-            sock._sock.settimeout(timeout)
 
-            timeout_ms = timeout * 1000
+            gen = (demux_adaptor(*frame) for frame in frames_iter(sock, tty=False))
-
+            stdout, stderr = consume_socket_output(gen, demux=True)
-            class _PollWithTimeout:
-                def __init__(self):
-                    self._inner = _original_poll()
-
-                def register(self, fd, eventmask):
-                    return self._inner.register(fd, eventmask)
-
-                def poll(self, *args):
-                    result = self._inner.poll(timeout_ms)
-                    if not result:
-                        raise _socket.timeout(f"timed out after {timeout}s")
-                    return result
-
-            with ExitStack() as stack:
-                stack.enter_context(patch.object(
-                    _select, "select",
-                    new=lambda rlist, wlist, xlist: _original_select(
-                        rlist, wlist, xlist, timeout
-                    ),
-                ))
-                if _original_poll is not None:
-                    stack.enter_context(
-                        patch.object(_select, "poll", new=_PollWithTimeout)
-                    )
-                gen = (demux_adaptor(*frame) for frame in frames_iter(sock, tty=False))
-                stdout, stderr = consume_socket_output(gen, demux=True)
 
             exit_code = self._client.api.exec_inspect(exec_id["Id"])["ExitCode"]
             if exit_code is None:
                 exit_code = 0
-            output = (stdout or b"") + (stderr or b"")
+
-            return exit_code, output.decode("utf-8", errors="replace")
+            output = ((stdout or b"") + (stderr or b"")).decode("utf-8", errors="replace")
-        except _socket.timeout:
+
-            return 124, f"Command timed out after {timeout}s"
+            if exit_code in (124, 137):
+                return exit_code, output + f"\n[command timed out after {timeout}s]"
+            return exit_code, output
         except Exception as exc:
             return 1, f"Error running command in container: {exc}"