Allow setting env vars inside sandboxes (#1652)

* add env to all sandboxes

* add unit tests

opendevin/runtime/docker/exec_box.py

@@ -62,7 +62,9 @@ class DockerExecBox(Sandbox):
             )
             raise ex
 
-        self.instance_id = sid + str(uuid.uuid4()) if sid is not None else str(uuid.uuid4())
+        self.instance_id = (
+            sid + str(uuid.uuid4()) if sid is not None else str(uuid.uuid4())
+        )
 
         # TODO: this timeout is actually essential - need a better way to set it
         # if it is too short, the container may still waiting for previous
@@ -80,6 +82,7 @@ class DockerExecBox(Sandbox):
         if RUN_AS_DEVIN:
             self.setup_devin_user()
         atexit.register(self.close)
+        super().__init__()
 
     def setup_devin_user(self):
         cmds = [
@@ -89,7 +92,9 @@ class DockerExecBox(Sandbox):
         ]
         for cmd in cmds:
             exit_code, logs = self.container.exec_run(
-                ['/bin/bash', '-c', cmd], workdir=SANDBOX_WORKSPACE_DIR
+                ['/bin/bash', '-c', cmd],
+                workdir=SANDBOX_WORKSPACE_DIR,
+                environment=self._env,
             )
             if exit_code != 0:
                 raise Exception(f'Failed to setup devin user: {logs}')
@@ -109,7 +114,9 @@ class DockerExecBox(Sandbox):
     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=SANDBOX_WORKSPACE_DIR)
+            return container.exec_run(
+                command, workdir=SANDBOX_WORKSPACE_DIR, environment=self._env
+            )
 
         # Use ThreadPoolExecutor to control command and set timeout
         with concurrent.futures.ThreadPoolExecutor() as executor:
@@ -125,7 +132,9 @@ class DockerExecBox(Sandbox):
                 pid = self.get_pid(cmd)
                 if pid is not None:
                     self.container.exec_run(
-                        f'kill -9 {pid}', workdir=SANDBOX_WORKSPACE_DIR
+                        f'kill -9 {pid}',
+                        workdir=SANDBOX_WORKSPACE_DIR,
+                        environment=self._env,
                     )
                 return -1, f'Command: "{cmd}" timed out'
         logs_out = logs.decode('utf-8')
@@ -138,6 +147,7 @@ class DockerExecBox(Sandbox):
         exit_code, logs = self.container.exec_run(
             ['/bin/bash', '-c', f'mkdir -p {sandbox_dest}'],
             workdir=SANDBOX_WORKSPACE_DIR,
+            environment=self._env,
         )
         if exit_code != 0:
             raise Exception(
@@ -173,7 +183,10 @@ class DockerExecBox(Sandbox):
 
     def execute_in_background(self, cmd: str) -> Process:
         result = self.container.exec_run(
-            self.get_exec_cmd(cmd), socket=True, workdir=SANDBOX_WORKSPACE_DIR
+            self.get_exec_cmd(cmd),
+            socket=True,
+            workdir=SANDBOX_WORKSPACE_DIR,
+            environment=self._env,
         )
         result.output._sock.setblocking(0)
         pid = self.get_pid(cmd)
@@ -183,7 +196,7 @@ class DockerExecBox(Sandbox):
         return bg_cmd
 
     def get_pid(self, cmd):
-        exec_result = self.container.exec_run('ps aux')
+        exec_result = self.container.exec_run('ps aux', environment=self._env)
         processes = exec_result.output.decode('utf-8').splitlines()
         cmd = ' '.join(self.get_exec_cmd(cmd))
 
@@ -199,7 +212,9 @@ class DockerExecBox(Sandbox):
         bg_cmd = self.background_commands[id]
         if bg_cmd.pid is not None:
             self.container.exec_run(
-                f'kill -9 {bg_cmd.pid}', workdir=SANDBOX_WORKSPACE_DIR
+                f'kill -9 {bg_cmd.pid}',
+                workdir=SANDBOX_WORKSPACE_DIR,
+                environment=self._env,
             )
         assert isinstance(bg_cmd, DockerProcess)
         bg_cmd.result.output.close()

opendevin/runtime/docker/local_box.py

@@ -33,6 +33,7 @@ class LocalBox(Sandbox):
         self.background_commands: Dict[int, Process] = {}
         self.cur_background_id = 0
         atexit.register(self.cleanup)
+        super().__init__()
 
     def execute(self, cmd: str) -> Tuple[int, str]:
         try:
@@ -43,6 +44,7 @@ class LocalBox(Sandbox):
                 capture_output=True,
                 timeout=self.timeout,
                 cwd=config.get(ConfigType.WORKSPACE_BASE),
+                env=self._env,
             )
             return completed_process.returncode, completed_process.stdout.strip()
         except subprocess.TimeoutExpired:
@@ -55,6 +57,7 @@ class LocalBox(Sandbox):
             shell=True,
             text=True,
             cwd=config.get(ConfigType.WORKSPACE_BASE),
+            env=self._env,
         )
         if res.returncode != 0:
             raise RuntimeError(f'Failed to create directory {sandbox_dest} in sandbox')
@@ -65,6 +68,7 @@ class LocalBox(Sandbox):
                 shell=True,
                 text=True,
                 cwd=config.get(ConfigType.WORKSPACE_BASE),
+                env=self._env,
             )
             if res.returncode != 0:
                 raise RuntimeError(
@@ -76,6 +80,7 @@ class LocalBox(Sandbox):
                 shell=True,
                 text=True,
                 cwd=config.get(ConfigType.WORKSPACE_BASE),
+                env=self._env,
             )
             if res.returncode != 0:
                 raise RuntimeError(

opendevin/runtime/docker/ssh_box.py

@@ -1,4 +1,5 @@
 import atexit
+import json
 import os
 import sys
 import tarfile
@@ -78,7 +79,9 @@ class DockerSSHBox(Sandbox):
             )
             raise ex
 
-        self.instance_id = sid + str(uuid.uuid4()) if sid is not None else str(uuid.uuid4())
+        self.instance_id = (
+            sid + str(uuid.uuid4()) if sid is not None else str(uuid.uuid4())
+        )
 
         # TODO: this timeout is actually essential - need a better way to set it
         # if it is too short, the container may still waiting for previous
@@ -100,6 +103,12 @@ class DockerSSHBox(Sandbox):
         self.setup_user()
         self.start_ssh_session()
         atexit.register(self.close)
+        super().__init__()
+
+    def add_to_env(self, key: str, value: str):
+        super().add_to_env(key, value)
+        # Note: json.dumps gives us nice escaping for free
+        self.execute(f'export {key}={json.dumps(value)}')
 
     def setup_user(self):
         # Make users sudoers passwordless
@@ -107,6 +116,7 @@ class DockerSSHBox(Sandbox):
         exit_code, logs = self.container.exec_run(
             ['/bin/bash', '-c', r"echo '%sudo ALL=(ALL) NOPASSWD:ALL' >> /etc/sudoers"],
             workdir=SANDBOX_WORKSPACE_DIR,
+            environment=self._env,
         )
         if exit_code != 0:
             raise Exception(
@@ -117,12 +127,14 @@ class DockerSSHBox(Sandbox):
         exit_code, logs = self.container.exec_run(
             ['/bin/bash', '-c', 'id -u opendevin'],
             workdir=SANDBOX_WORKSPACE_DIR,
+            environment=self._env,
         )
         if exit_code == 0:
             # User exists, delete it
             exit_code, logs = self.container.exec_run(
                 ['/bin/bash', '-c', 'userdel -r opendevin'],
                 workdir=SANDBOX_WORKSPACE_DIR,
+                environment=self._env,
             )
             if exit_code != 0:
                 raise Exception(f'Failed to remove opendevin user in sandbox: {logs}')
@@ -136,6 +148,7 @@ class DockerSSHBox(Sandbox):
                     f'useradd -rm -d /home/opendevin -s /bin/bash -g root -G sudo -u {USER_ID} opendevin',
                 ],
                 workdir=SANDBOX_WORKSPACE_DIR,
+                environment=self._env,
             )
             if exit_code != 0:
                 raise Exception(f'Failed to create opendevin user in sandbox: {logs}')
@@ -146,6 +159,7 @@ class DockerSSHBox(Sandbox):
                     f"echo 'opendevin:{self._ssh_password}' | chpasswd",
                 ],
                 workdir=SANDBOX_WORKSPACE_DIR,
+                environment=self._env,
             )
             if exit_code != 0:
                 raise Exception(f'Failed to set password in sandbox: {logs}')
@@ -154,6 +168,7 @@ class DockerSSHBox(Sandbox):
             exit_code, logs = self.container.exec_run(
                 ['/bin/bash', '-c', 'chown opendevin:root /home/opendevin'],
                 workdir=SANDBOX_WORKSPACE_DIR,
+                environment=self._env,
             )
             if exit_code != 0:
                 raise Exception(
@@ -162,6 +177,7 @@ class DockerSSHBox(Sandbox):
             exit_code, logs = self.container.exec_run(
                 ['/bin/bash', '-c', f'chown opendevin:root {SANDBOX_WORKSPACE_DIR}'],
                 workdir=SANDBOX_WORKSPACE_DIR,
+                environment=self._env,
             )
             if exit_code != 0:
                 # This is not a fatal error, just a warning
@@ -173,12 +189,14 @@ class DockerSSHBox(Sandbox):
                 # change password for root
                 ['/bin/bash', '-c', f"echo 'root:{self._ssh_password}' | chpasswd"],
                 workdir=SANDBOX_WORKSPACE_DIR,
+                environment=self._env,
             )
             if exit_code != 0:
                 raise Exception(f'Failed to set password for root in sandbox: {logs}')
         exit_code, logs = self.container.exec_run(
             ['/bin/bash', '-c', "echo 'opendevin-sandbox' > /etc/hostname"],
             workdir=SANDBOX_WORKSPACE_DIR,
+            environment=self._env,
         )
 
     def start_ssh_session(self):
@@ -266,6 +284,7 @@ class DockerSSHBox(Sandbox):
         exit_code, logs = self.container.exec_run(
             ['/bin/bash', '-c', f'mkdir -p {sandbox_dest}'],
             workdir=SANDBOX_WORKSPACE_DIR,
+            environment=self._env,
         )
         if exit_code != 0:
             raise Exception(
@@ -301,7 +320,10 @@ class DockerSSHBox(Sandbox):
 
     def execute_in_background(self, cmd: str) -> Process:
         result = self.container.exec_run(
-            self.get_exec_cmd(cmd), socket=True, workdir=SANDBOX_WORKSPACE_DIR
+            self.get_exec_cmd(cmd),
+            socket=True,
+            workdir=SANDBOX_WORKSPACE_DIR,
+            environment=self._env,
         )
         result.output._sock.setblocking(0)
         pid = self.get_pid(cmd)
@@ -311,7 +333,7 @@ class DockerSSHBox(Sandbox):
         return bg_cmd
 
     def get_pid(self, cmd):
-        exec_result = self.container.exec_run('ps aux')
+        exec_result = self.container.exec_run('ps aux', environment=self._env)
         processes = exec_result.output.decode('utf-8').splitlines()
         cmd = ' '.join(self.get_exec_cmd(cmd))
 
@@ -327,7 +349,9 @@ class DockerSSHBox(Sandbox):
         bg_cmd = self.background_commands[id]
         if bg_cmd.pid is not None:
             self.container.exec_run(
-                f'kill -9 {bg_cmd.pid}', workdir=SANDBOX_WORKSPACE_DIR
+                f'kill -9 {bg_cmd.pid}',
+                workdir=SANDBOX_WORKSPACE_DIR,
+                environment=self._env,
             )
         assert isinstance(bg_cmd, DockerProcess)
         bg_cmd.result.output.close()

opendevin/runtime/e2b/sandbox.py

@@ -37,6 +37,7 @@ class E2BBox(Sandbox):
         )
         self.timeout = timeout
         logger.info(f'Started E2B sandbox with ID "{self.sandbox.id}"')
+        super().__init__()
 
     @property
     def filesystem(self):
@@ -74,7 +75,7 @@ class E2BBox(Sandbox):
         return '\n'.join([m.line for m in proc.output_messages])
 
     def execute(self, cmd: str) -> Tuple[int, str]:
-        process = self.sandbox.process.start(cmd)
+        process = self.sandbox.process.start(cmd, envVars=self._env)
         try:
             process_output = process.wait(timeout=self.timeout)
         except TimeoutException:

opendevin/runtime/sandbox.py

@@ -1,3 +1,4 @@
+import os
 from abc import ABC, abstractmethod
 from typing import Dict, Tuple
 
@@ -7,6 +8,16 @@ from opendevin.runtime.plugins.mixin import PluginMixin
 
 class Sandbox(ABC, PluginMixin):
     background_commands: Dict[int, Process] = {}
+    _env: Dict[str, str] = {}
+
+    def __init__(self, **kwargs):
+        for key in os.environ:
+            if key.startswith('SANDBOX_ENV_'):
+                sandbox_key = key.removeprefix('SANDBOX_ENV_')
+                self.add_to_env(sandbox_key, os.environ[key])
+
+    def add_to_env(self, key: str, value: str):
+        self._env[key] = value
 
     @abstractmethod
     def execute(self, cmd: str) -> Tuple[int, str]:

tests/unit/test_sandbox.py

@@ -1,3 +1,4 @@
+import os
 import pathlib
 import tempfile
 from unittest.mock import patch
@@ -5,6 +6,8 @@ from unittest.mock import patch
 import pytest
 
 from opendevin.core import config
+from opendevin.runtime.docker.exec_box import DockerExecBox
+from opendevin.runtime.docker.local_box import LocalBox
 from opendevin.runtime.docker.ssh_box import DockerSSHBox
 
 
@@ -16,6 +19,18 @@ def temp_dir():
         yield temp_dir
 
 
+def test_env_vars(temp_dir):
+    os.environ['SANDBOX_ENV_FOOBAR'] = 'BAZ'
+    for box_class in [DockerSSHBox, DockerExecBox, LocalBox]:
+        box = box_class()
+        box.add_to_env('QUUX', 'abc"def')
+        assert box._env['FOOBAR'] == 'BAZ'
+        assert box._env['QUUX'] == 'abc"def'
+        exit_code, output = box.execute('echo $FOOBAR $QUUX')
+        assert exit_code == 0, 'The exit code should be 0.'
+        assert output.strip() == 'BAZ abc"def', f'Output: {output} for {box_class}'
+
+
 def test_ssh_box_run_as_devin(temp_dir):
     # get a temporary directory
     with patch.dict(