feat(sandbox): Support sshd-based stateful docker session (All-Hands-…

…AI#847)

* support sshd-based stateful docker session

* use .getLogger to avoid same logging message to get printed twice

* update poetry lock for dependency

* fix ruff

* bump docker image version with sshd

* set-up random user password and only allow localhost connection for sandbox

* fix poetry

* move apt install up

opendevin/sandbox/Dockerfile

@@ -14,4 +14,8 @@ RUN apt-get update && apt-get install -y \
     python3-venv \
     python3-dev \
     build-essential \
+    openssh-server \
+    sudo \
     && rm -rf /var/lib/apt/lists/*
+
+RUN service ssh start

opendevin/sandbox/Makefile

@@ -1,7 +1,7 @@
 DOCKER_BUILD_REGISTRY=ghcr.io
 DOCKER_BUILD_ORG=opendevin
 DOCKER_BUILD_REPO=sandbox
-DOCKER_BUILD_TAG=v0.1.0
+DOCKER_BUILD_TAG=v0.1.1
 FULL_IMAGE=$(DOCKER_BUILD_REGISTRY)/$(DOCKER_BUILD_ORG)/$(DOCKER_BUILD_REPO):$(DOCKER_BUILD_TAG)
 
 LATEST_FULL_IMAGE=$(DOCKER_BUILD_REGISTRY)/$(DOCKER_BUILD_ORG)/$(DOCKER_BUILD_REPO):latest

opendevin/sandbox/sandbox.py

@@ -4,11 +4,11 @@
 import sys
 import time
 import uuid
+from pexpect import pxssh
 from collections import namedtuple
 from typing import Dict, List, Tuple
 
 import docker
-import concurrent.futures
 
 from opendevin import config
 from opendevin.logging import opendevin_logger as logger
@@ -132,23 +132,50 @@ def __init__(
 
         if not self.is_container_running():
             self.restart_docker_container()
+        # set up random user password
+        self._ssh_password = str(uuid.uuid4())
         if RUN_AS_DEVIN:
             self.setup_devin_user()
+            self.start_ssh_session()
+        else:
+            # TODO: implement ssh into root
+            raise NotImplementedError(
+                'Running as root is not supported at the moment.')
         atexit.register(self.cleanup)
 
     def setup_devin_user(self):
         exit_code, logs = self.container.exec_run(
-            [
-                '/bin/bash',
-                '-c',
-                f'useradd --shell /bin/bash -u {USER_ID} -o -c "" -m devin',
-            ],
+            ['/bin/bash', '-c',
+                f'useradd -rm -d /home/opendevin -s /bin/bash -g root -G sudo -u {USER_ID} opendevin'],
+            workdir='/workspace',
+        )
+        exit_code, logs = self.container.exec_run(
+            ['/bin/bash', '-c',
+                f"echo 'opendevin:{self._ssh_password}' | chpasswd"],
+            workdir='/workspace',
+        )
+        exit_code, logs = self.container.exec_run(
+            ['/bin/bash', '-c', "echo 'opendevin-sandbox' > /etc/hostname"],
             workdir='/workspace',
         )
 
+    def start_ssh_session(self):
+        # start ssh session at the background
+        self.ssh = pxssh.pxssh()
+        hostname = 'localhost'
+        username = 'opendevin'
+        self.ssh.login(hostname, username, self._ssh_password, port=2222)
+
+        # Fix: https://github.com/pexpect/pexpect/issues/669
+        self.ssh.sendline("bind 'set enable-bracketed-paste off'")
+        self.ssh.prompt()
+        # cd to workspace
+        self.ssh.sendline('cd /workspace')
+        self.ssh.prompt()
+
     def get_exec_cmd(self, cmd: str) -> List[str]:
         if RUN_AS_DEVIN:
-            return ['su', 'devin', '-c', cmd]
+            return ['su', 'opendevin', '-c', cmd]
         else:
             return ['/bin/bash', '-c', cmd]
 
@@ -159,26 +186,27 @@ def read_logs(self, id) -> str:
         return bg_cmd.read_logs()
 
     def execute(self, cmd: str) -> Tuple[int, str]:
-        # TODO: each execute is not stateful! We need to keep track of the current working directory
-        def run_command(container, command):
-            return container.exec_run(command, workdir='/workspace')
-
-        # Use ThreadPoolExecutor to control command and set timeout
-        with concurrent.futures.ThreadPoolExecutor() as executor:
-            future = executor.submit(
-                run_command, self.container, self.get_exec_cmd(cmd)
-            )
-            try:
-                exit_code, logs = future.result(timeout=self.timeout)
-            except concurrent.futures.TimeoutError:
-                logger.exception(
-                    'Command timed out, killing process...', exc_info=False)
-                pid = self.get_pid(cmd)
-                if pid is not None:
-                    self.container.exec_run(
-                        f'kill -9 {pid}', workdir='/workspace')
-                return -1, f'Command: "{cmd}" timed out'
-        return exit_code, logs.decode('utf-8')
+        # use self.ssh
+        self.ssh.sendline(cmd)
+        success = self.ssh.prompt(timeout=self.timeout)
+        if not success:
+            logger.exception(
+                'Command timed out, killing process...', exc_info=False)
+            # send a SIGINT to the process
+            self.ssh.sendintr()
+            self.ssh.prompt()
+            command_output = self.ssh.before.decode(
+                'utf-8').lstrip(cmd).strip()
+            return -1, f'Command: "{cmd}" timed out. Sending SIGINT to the process: {command_output}'
+        command_output = self.ssh.before.decode('utf-8').lstrip(cmd).strip()
+
+        # get the exit code
+        self.ssh.sendline('echo $?')
+        self.ssh.prompt()
+        exit_code = self.ssh.before.decode('utf-8')
+        # remove the echo $? itself
+        exit_code = int(exit_code.lstrip('echo $?').strip())
+        return exit_code, command_output
 
     def execute_in_background(self, cmd: str) -> BackgroundCommand:
         result = self.container.exec_run(
@@ -267,10 +295,12 @@ def restart_docker_container(self):
             # start the container
             self.container = docker_client.containers.run(
                 self.container_image,
-                command='tail -f /dev/null',
+                # only allow connections from localhost
+                command="/usr/sbin/sshd -D -p 2222 -o 'ListenAddress=127.0.0.1'",
                 network_mode='host',
                 working_dir='/workspace',
                 name=self.container_name,
+                hostname='opendevin_sandbox',
                 detach=True,
                 volumes={self.workspace_dir: {
                     'bind': '/workspace', 'mode': 'rw'}},

pyproject.toml

@@ -23,6 +23,7 @@ types-toml = "*"
 numpy = "*"
 json-repair = "*"
 playwright = "*"
+pexpect = "*"
 
 [tool.poetry.group.llama-index.dependencies]
 llama-index = "*"