Preserve command output on timeout

src/swerex/exceptions.py

@@ -25,7 +25,18 @@ def __init__(self, message: str, *, extra_info: dict[str, Any] = None):
         self.extra_info = extra_info
 
 
-class CommandTimeoutError(SwerexException, RuntimeError, TimeoutError): ...
+class CommandTimeoutError(SwerexException, RuntimeError, TimeoutError):
+    """Raised when a command times out, but includes any output collected before the timeout.
+
+    Attributes:
+        partial_output (str): The output that was collected before the timeout occurred
+    """
+    def __init__(self, message: str, partial_output: str = "", *, extra_info: dict[str, Any] = None):
+        super().__init__(message)
+        self.partial_output = partial_output
+        if extra_info is None:
+            extra_info = {}
+        self.extra_info = extra_info
 
 
 class NoExitCodeError(SwerexException, RuntimeError): ...

src/swerex/runtime/local.py

@@ -217,12 +217,12 @@ async def run(self, action: BashAction | BashInterruptAction) -> BashObservation
 
         Raises:
             SessionNotInitializedError: If the shell is not initialized.
-            CommandTimeoutError: If the command times out.
+            CommandTimeoutError: If the command times out. The error will contain any output collected before the timeout.
             NonZeroExitCodeError: If the command has a non-zero exit code and `action.check` is True.
             NoExitCodeError: If we cannot get the exit code of the command.
 
         Returns:
-            BashObservation: The observation of the command.
+            BashObservation: The observation of the command, including any partial output if a timeout occurred.
         """
         if self.shell is None:
             msg = "shell not initialized"
@@ -231,7 +231,11 @@ async def run(self, action: BashAction | BashInterruptAction) -> BashObservation
             return await self.interrupt(action)
         if action.is_interactive_command or action.is_interactive_quit:
             return await self._run_interactive(action)
-        r = await self._run_normal(action)
+        try:
+            r = await self._run_normal(action)
+        except CommandTimeoutError as e:
+            # Return partial output on timeout
+            return BashObservation(output=e.partial_output, exit_code=None, failure_reason="timeout")
         if action.check == "raise" and r.exit_code != 0:
             msg = (
                 f"Command {action.command!r} failed with exit code {r.exit_code}. " f"Here is the output:\n{r.output!r}"
@@ -309,8 +313,9 @@ async def _run_normal(self, action: BashAction) -> BashObservation:
             expect_index = self.shell.expect(expect_strings, timeout=action.timeout)  # type: ignore
             matched_expect_string = expect_strings[expect_index]
         except pexpect.TIMEOUT as e:
+            partial_output = _strip_control_chars(self.shell.before).strip()  # type: ignore
             msg = f"timeout after {action.timeout} seconds while running command {action.command!r}"
-            raise CommandTimeoutError(msg) from e
+            raise CommandTimeoutError(msg, partial_output=partial_output) from e
         output: str = _strip_control_chars(self.shell.before).strip()  # type: ignore
 
         # Part 3: Get the exit code