Simplify ssh box to fix multi-line command issues and add unit tests

opendevin/sandbox/docker/ssh_box.py

@@ -216,19 +216,6 @@ def execute(self, cmd: str) -> Tuple[int, str]:
  return -1, f'Command: "{cmd}" timed out. Sending SIGINT to the process: {command_output}'
  command_output = self.ssh.before.decode('utf-8').strip()
 
- # NOTE: there's some weird behavior with the prompt (it may come AFTER the command output)
- # so we need to check if the command is in the output
- n_tries = 5
- while not command_output.startswith(cmd) and n_tries > 0:
- self.ssh.prompt()
- command_output = self.ssh.before.decode('utf-8').strip()
- time.sleep(0.5)
- n_tries -= 1
- if n_tries == 0 and not command_output.startswith(cmd):
- raise Exception(
- f'Something went wrong with the SSH sanbox, cannot get output for command [{cmd}] after 5 retries'
- )
- logger.debug(f'Command output GOT SO FAR: {command_output}')
  # once out, make sure that we have *every* output, we while loop until we get an empty output
  while True:
  logger.debug('WAITING FOR .prompt()')

tests/unit/test_sandbox.py

@@ -0,0 +1,102 @@
+import os
+import pathlib
+import tempfile
+
+
+def test_ssh_box_run_as_devin():
+ # get a temporary directory
+ with tempfile.TemporaryDirectory() as temp_dir:
+ pathlib.Path().mkdir(parents=True, exist_ok=True)
+ os.environ['WORKSPACE_BASE'] = temp_dir
+ os.environ['RUN_AS_DEVIN'] = 'true'
+ os.environ['SANDBOX_TYPE'] = 'ssh'
+ from opendevin.sandbox.docker.ssh_box import DockerSSHBox
+ ssh_box = DockerSSHBox()
+
+ # test the ssh box
+ exit_code, output = ssh_box.execute('ls -l')
+ assert exit_code == 0, 'The exit code should be 0.'
+ assert output.strip() == 'total 0'
+
+ exit_code, output = ssh_box.execute('mkdir test')
+ assert exit_code == 0, 'The exit code should be 0.'
+ assert output.strip() == ''
+
+ exit_code, output = ssh_box.execute('ls -l')
+ assert exit_code == 0, 'The exit code should be 0.'
+ assert 'opendevin' in output, "The output should contain username 'opendevin'"
+ assert 'test' in output, 'The output should contain the test directory'
+
+ exit_code, output = ssh_box.execute('touch test/foo.txt')
+ assert exit_code == 0, 'The exit code should be 0.'
+ assert output.strip() == ''
+
+ exit_code, output = ssh_box.execute('ls -l test')
+ assert exit_code == 0, 'The exit code should be 0.'
+ assert 'foo.txt' in output, 'The output should contain the foo.txt file'
+
+ # cleanup the environment variables
+ del os.environ['WORKSPACE_BASE']
+ del os.environ['SANDBOX_TYPE']
+
+def test_ssh_box_multi_line_cmd_run_as_devin():
+ # get a temporary directory
+ with tempfile.TemporaryDirectory() as temp_dir:
+ pathlib.Path().mkdir(parents=True, exist_ok=True)
+ os.environ['WORKSPACE_BASE'] = temp_dir
+ os.environ['RUN_AS_DEVIN'] = 'true'
+ os.environ['SANDBOX_TYPE'] = 'ssh'
+ from opendevin.sandbox.docker.ssh_box import DockerSSHBox
+ ssh_box = DockerSSHBox()
+
+ # test the ssh box
+ exit_code, output = ssh_box.execute('pwd\nls -l')
+ assert exit_code == 0, 'The exit code should be 0.'
+ expected_lines = ['/workspacels -l', 'total 0']
+ assert output.strip().splitlines() == expected_lines
+
+ # cleanup the environment variables
+ del os.environ['WORKSPACE_BASE']
+ del os.environ['SANDBOX_TYPE']
+
+
+def test_ssh_box_stateful_cmd_run_as_devin():
+ # get a temporary directory
+ with tempfile.TemporaryDirectory() as temp_dir:
+ pathlib.Path().mkdir(parents=True, exist_ok=True)
+ os.environ['WORKSPACE_BASE'] = temp_dir
+ os.environ['RUN_AS_DEVIN'] = 'true'
+ os.environ['SANDBOX_TYPE'] = 'ssh'
+ from opendevin.sandbox.docker.ssh_box import DockerSSHBox
+ ssh_box = DockerSSHBox()
+
+ # test the ssh box
+ exit_code, output = ssh_box.execute('mkdir test')
+ assert exit_code == 0, 'The exit code should be 0.'
+ assert output.strip() == ''
+
+ exit_code, output = ssh_box.execute('cd test')
+ assert exit_code == 0, 'The exit code should be 0.'
+ assert output.strip() == ''
+
+ exit_code, output = ssh_box.execute('pwd')
+ assert exit_code == 0, 'The exit code should be 0.'
+ assert output.strip() == '/workspace/test'
+
+ # cleanup the environment variables
+ del os.environ['WORKSPACE_BASE']
+ del os.environ['SANDBOX_TYPE']
+
+def test_ssh_box_failed_cmd_run_as_devin():
+ # get a temporary directory
+ with tempfile.TemporaryDirectory() as temp_dir:
+ pathlib.Path().mkdir(parents=True, exist_ok=True)
+ os.environ['WORKSPACE_BASE'] = temp_dir
+ os.environ['RUN_AS_DEVIN'] = 'true'
+ os.environ['SANDBOX_TYPE'] = 'ssh'
+ from opendevin.sandbox.docker.ssh_box import DockerSSHBox
+ ssh_box = DockerSSHBox()
+
+ # test the ssh box with a command that fails
+ exit_code, output = ssh_box.execute('non_existing_command')
+ assert exit_code != 0, 'The exit code should not be 0 for a failed command.'