From 4ffcdf5246ba6f31bd27a44f41af41adf1c8cefd Mon Sep 17 00:00:00 2001 From: frdel Date: Mon, 20 Jul 2026 17:04:05 +0200 Subject: [PATCH] Handle shell/SSH process exits as completion Detect and treat terminated local/SSH/TTY shells as definitive command completion. Add is_terminated and get_exit_code helpers to LocalInteractiveSession, SSHInteractiveSession, and TTYSession; expand _is_closed_pty_error to recognize exited TTY processes. CodeExecution now reports a shell-exit prompt, recreates terminated sessions lazily before the next command, and returns immediately when a shell has exited. Docs and README updated to describe strict-mode/exit behavior, and tests were added/updated to cover the new termination detection and handling. --- plugins/_code_execution/AGENTS.md | 1 + plugins/_code_execution/README.md | 1 + .../_code_execution/helpers/shell_local.py | 10 +- plugins/_code_execution/helpers/shell_ssh.py | 23 ++++ .../_code_execution/helpers/tty_session.py | 10 ++ .../prompts/fw.code.shell_exit.md | 1 + .../tools/code_execution_tool.py | 35 ++++- tests/test_code_execution_pager.py | 122 +++++++++++++++++- 8 files changed, 199 insertions(+), 4 deletions(-) create mode 100644 plugins/_code_execution/prompts/fw.code.shell_exit.md diff --git a/plugins/_code_execution/AGENTS.md b/plugins/_code_execution/AGENTS.md index 8b6a38dbb..015852174 100644 --- a/plugins/_code_execution/AGENTS.md +++ b/plugins/_code_execution/AGENTS.md @@ -15,6 +15,7 @@ - Keep session concurrency, timeout, streaming, and reset behavior predictable. - Execute multi-line terminal input as one current-shell compound so intermediate prompts cannot mark queued work complete; preserve `cd`, exports, and other shell state. +- Treat local process exit and SSH channel termination as definitive command completion even when no final prompt is emitted; recreate terminated sessions before their next command. - Terminal reset/close must not hang on foreground commands or shells that ignore SIGTERM. - Explicitly target local versus SSH execution runtimes. - Do not hardcode secrets, SSH credentials, or local user paths. diff --git a/plugins/_code_execution/README.md b/plugins/_code_execution/README.md index 54f27f2cd..0ffe5a8d0 100644 --- a/plugins/_code_execution/README.md +++ b/plugins/_code_execution/README.md @@ -24,6 +24,7 @@ This plugin provides the code execution tool used by agents for development task - Can open SSH interactive sessions instead of local shells when configured. - **Streaming output** - Continuously reads shell output, updates the current log item, and detects progress while commands are running. + - Detects local shell exit and SSH channel termination when strict mode, `exit`, or a lost connection prevents a final prompt from appearing. - **Long-running work** - Keeps normal command execution responsive while giving the `output` runtime longer polling windows for builds, installs, servers, tests, and training jobs. - **Safety around running sessions** diff --git a/plugins/_code_execution/helpers/shell_local.py b/plugins/_code_execution/helpers/shell_local.py index 515f30916..99a56efbe 100644 --- a/plugins/_code_execution/helpers/shell_local.py +++ b/plugins/_code_execution/helpers/shell_local.py @@ -57,6 +57,14 @@ class LocalInteractiveSession: raise Exception("Shell not connected") self.full_output = "" await self.session.sendline(command) + + def is_terminated(self) -> bool: + return self.session is None or self.session.is_terminated() + + def get_exit_code(self) -> int | None: + if not self.session: + return None + return self.session.get_exit_code() async def read_output(self, timeout: float = 0, reset_full_output: bool = False) -> Tuple[str, Optional[str]]: if not self.session: @@ -75,4 +83,4 @@ class LocalInteractiveSession: if not partial_output: return clean_full_output, None - return clean_full_output, partial_output \ No newline at end of file + return clean_full_output, partial_output diff --git a/plugins/_code_execution/helpers/shell_ssh.py b/plugins/_code_execution/helpers/shell_ssh.py index 52a23bade..99945f7d8 100644 --- a/plugins/_code_execution/helpers/shell_ssh.py +++ b/plugins/_code_execution/helpers/shell_ssh.py @@ -34,6 +34,7 @@ class SSHInteractiveSession: self.last_command = b"" self.trimmed_command_length = 0 # Initialize trimmed_command_length self.cwd = cwd + self._exit_code: int | None = None async def connect(self, keepalive_interval: int = 5): """ @@ -67,6 +68,7 @@ class SSHInteractiveSession: # invoke interactive shell self.shell = self.client.invoke_shell(width=100, height=50) + self._exit_code = None # disable systemd/OSC prompt metadata and disable local echo initial_command = f"unset PROMPT_COMMAND PS0; stty -echo; {PAGER_DISABLE_COMMAND}" @@ -110,6 +112,27 @@ class SSHInteractiveSession: self.last_command = command.encode() self.trimmed_command_length = 0 self.shell.send(self.last_command) + + def is_terminated(self) -> bool: + if not self.shell: + return True + try: + transport = self.client.get_transport() + if not transport or not transport.is_active(): + return True + return self.shell.closed or self.shell.exit_status_ready() + except Exception: + return True + + def get_exit_code(self) -> int | None: + if self._exit_code is not None: + return self._exit_code + try: + if self.shell and self.shell.exit_status_ready(): + self._exit_code = self.shell.recv_exit_status() + except Exception: + return None + return self._exit_code async def read_output( self, timeout: float = 0, reset_full_output: bool = False diff --git a/plugins/_code_execution/helpers/tty_session.py b/plugins/_code_execution/helpers/tty_session.py index 5d3e3d018..6176961ed 100644 --- a/plugins/_code_execution/helpers/tty_session.py +++ b/plugins/_code_execution/helpers/tty_session.py @@ -176,6 +176,16 @@ class TTYSession: raise RuntimeError("TTYSpawn is not started") return await self._proc.wait() + def is_terminated(self) -> bool: + """Return whether the managed shell process has exited.""" + return self._proc is None or getattr(self._proc, "returncode", None) is not None + + def get_exit_code(self) -> int | None: + """Return the managed shell exit code when it is already available.""" + if self._proc is None: + return None + return getattr(self._proc, "returncode", None) + def kill(self): """Force-kill the running child process. diff --git a/plugins/_code_execution/prompts/fw.code.shell_exit.md b/plugins/_code_execution/prompts/fw.code.shell_exit.md new file mode 100644 index 000000000..09a6af568 --- /dev/null +++ b/plugins/_code_execution/prompts/fw.code.shell_exit.md @@ -0,0 +1 @@ +Terminal shell exited{{status}}. The command has finished; a new shell will be created before the next command. diff --git a/plugins/_code_execution/tools/code_execution_tool.py b/plugins/_code_execution/tools/code_execution_tool.py index 6b6e58e75..022d5975a 100644 --- a/plugins/_code_execution/tools/code_execution_tool.py +++ b/plugins/_code_execution/tools/code_execution_tool.py @@ -17,8 +17,10 @@ from plugins._code_execution.helpers.shell_ssh import SSHInteractiveSession def _is_closed_pty_error(exc: BaseException) -> bool: - if isinstance(exc, RuntimeError) and "TTYSpawn PTY is closed" in str(exc): - return True + if isinstance(exc, RuntimeError): + message = str(exc) + if "TTYSpawn PTY is closed" in message or "TTYSpawn process has exited" in message: + return True if isinstance(exc, OSError) and exc.errno in (errno.EBADF, errno.EIO, errno.EINVAL): return True cause = getattr(exc, "__cause__", None) @@ -190,6 +192,11 @@ class CodeExecution(Tool): if response := await self.handle_running_session(cfg, session): return response + # A strict-mode command can terminate the persistent shell itself. + # Recreate such a session lazily before accepting the next command. + if self.state.shells[session].session.is_terminated(): + await self.prepare_state(cfg, reset=True, session=session) + # try again on lost connection for i in range(2): try: @@ -295,6 +302,26 @@ class CodeExecution(Tool): last_output_time = now got_output = True + # ``set -e``, ``exit``, or a lost SSH channel can end the managed + # shell without ever producing another prompt. Treat that process + # or channel termination as a definitive command end. + shell = self.state.shells[session].session + if shell.is_terminated(): + exit_code = shell.get_exit_code() + status = f" with exit code {exit_code}" if exit_code is not None else "" + sysinfo = self.agent.read_prompt( + "fw.code.shell_exit.md", status=status + ) + response = self.agent.read_prompt("fw.code.info.md", info=sysinfo) + if truncated_output: + response = truncated_output + "\n\n" + response + PrintStyle.warning(sysinfo) + heading = self.get_heading_from_output(truncated_output, 0, True) + self.log.update(content=prefix + response, heading=heading) + self.mark_session_idle(session) + return response + + if partial_output: # Check for shell prompt at the end of output last_lines = ( truncated_output.splitlines()[-3:] if truncated_output else [] @@ -411,6 +438,10 @@ class CodeExecution(Tool): await self.set_progress(truncated_output) heading = self.get_heading_from_output(truncated_output, 0) + if self.state.shells[session].session.is_terminated(): + self.mark_session_idle(session) + return None + last_lines = ( truncated_output.splitlines()[-3:] if truncated_output else [] ) diff --git a/tests/test_code_execution_pager.py b/tests/test_code_execution_pager.py index a5ef70b19..19d51758e 100644 --- a/tests/test_code_execution_pager.py +++ b/tests/test_code_execution_pager.py @@ -5,10 +5,17 @@ code execution tool: without user input they block forever and spin at 100% CPU. """ import asyncio +from types import SimpleNamespace from plugins._code_execution.helpers import shell_local, shell_ssh from plugins._code_execution.helpers.tty_session import TTYSession -from plugins._code_execution.tools.code_execution_tool import _group_multiline_command +from plugins._code_execution.tools.code_execution_tool import ( + CodeExecution, + ShellWrap, + State, + _group_multiline_command, + _is_closed_pty_error, +) def test_local_env_disables_pagers_and_preserves_existing(): @@ -44,6 +51,10 @@ def test_multiline_terminal_commands_are_one_current_shell_compound(): ) +def test_exited_tty_process_is_a_recoverable_closed_session(): + assert _is_closed_pty_error(RuntimeError("TTYSpawn process has exited")) + + def test_tty_close_kills_term_resistant_process(): async def run(): session = TTYSession("bash -lc 'trap \"\" TERM; sleep 30'") @@ -52,3 +63,112 @@ def test_tty_close_kills_term_resistant_process(): assert session._proc is None asyncio.run(run()) + + +def test_tty_reports_strict_mode_shell_exit(): + async def run(): + session = TTYSession("/bin/bash --noprofile --norc -i") + await session.start() + await session.read_full_until_idle(idle_timeout=0.05, total_timeout=1) + await session.sendline("{\nset -euo pipefail\nfalse\nprintf 'unreachable\\n'\n}") + + exit_code = await asyncio.wait_for(session.wait(), timeout=5) + + assert exit_code != 0 + assert session.is_terminated() + assert session.get_exit_code() == exit_code + await session.close() + + asyncio.run(run()) + + +def test_ssh_session_reports_channel_exit_status(): + class FakeChannel: + closed = False + + @staticmethod + def exit_status_ready(): + return True + + @staticmethod + def recv_exit_status(): + return 7 + + session = object.__new__(shell_ssh.SSHInteractiveSession) + session.shell = FakeChannel() + session.client = SimpleNamespace( + get_transport=lambda: SimpleNamespace(is_active=lambda: True) + ) + session._exit_code = None + + assert session.is_terminated() + assert session.get_exit_code() == 7 + + +def test_code_execution_returns_immediately_when_shell_exits(): + class FinishedSession: + async def read_output(self, timeout=0, reset_full_output=False): + return "nothing to commit, working tree clean\n", "nothing to commit, working tree clean\n" + + @staticmethod + def is_terminated(): + return True + + @staticmethod + def get_exit_code(): + return 1 + + class FakeAgent: + agent_name = "test" + + async def handle_intervention(self): + return None + + @staticmethod + def read_prompt(name, **kwargs): + if name == "fw.code.shell_exit.md": + return f"Terminal shell exited{kwargs['status']}. The command has finished." + if name == "fw.code.info.md": + return f"[SYSTEM: {kwargs['info']}]" + raise AssertionError(f"Unexpected prompt: {name}") + + async def run(): + session = FinishedSession() + state = State( + ssh_enabled=False, + shells={0: ShellWrap(id=0, session=session, running=True)}, + ) + tool = CodeExecution( + FakeAgent(), + "code_execution_tool", + "", + {"runtime": "terminal", "session": 0}, + "", + None, + ) + updates = [] + tool.log = SimpleNamespace(update=lambda **kwargs: updates.append(kwargs)) + + async def prepare_state(*args, **kwargs): + return state + + async def set_progress(content): + return None + + tool.prepare_state = prepare_state + tool.set_progress = set_progress + tool.fix_full_output = lambda output: output + + response = await tool.get_terminal_output( + {"prompt_patterns": [], "dialog_patterns": []}, + session=0, + sleep_time=0, + ) + + assert "nothing to commit" in response + assert "exit code 1" in response + assert "command has finished" in response + assert not state.shells[0].running + assert updates[-1]["heading"].endswith(" icon://done_all") + + asyncio.run(run())