fix(sandbox): SSH box .execute misalign issue (All-Hands-AI#1304)

* fix weird behavior of pxssh where the output would not flush correctly

* make ssh box can handle exit_code properly as well

* Update opendevin/sandbox/docker/ssh_box.py

Co-authored-by: Robert Brennan <[email protected]>

* fix typo

---------

Co-authored-by: Robert Brennan <[email protected]>

opendevin/sandbox/docker/ssh_box.py

@@ -205,6 +205,7 @@ def read_logs(self, id) -> str:
         return bg_cmd.read_logs()
 
     def execute(self, cmd: str) -> Tuple[int, str]:
+        cmd = cmd.strip()
         # use self.ssh
         self.ssh.sendline(cmd)
         success = self.ssh.prompt(timeout=self.timeout)
@@ -217,13 +218,45 @@ def execute(self, cmd: str) -> Tuple[int, str]:
             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()
+        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()')
+            self.ssh.sendline('\n')
+            timeout_not_reached = self.ssh.prompt(timeout=1)
+            if not timeout_not_reached:
+                logger.debug('TIMEOUT REACHED')
+                break
+            logger.debug('WAITING FOR .before')
+            output = self.ssh.before.decode('utf-8').strip()
+            logger.debug(f'WAITING FOR END OF command output ({bool(output)}): {output}')
+            if output == '':
+                break
+            command_output += output
+        command_output = command_output.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
+        while not exit_code.startswith('echo $?'):
+            self.ssh.prompt()
+            exit_code = self.ssh.before.decode('utf-8')
+            logger.debug(f'WAITING FOR exit code: {exit_code}')
         exit_code = int(exit_code.lstrip('echo $?').strip())
         return exit_code, command_output