diff --git a/plugins/_code_execution/AGENTS.md b/plugins/_code_execution/AGENTS.md index 0e5de63bc..0cc78b305 100644 --- a/plugins/_code_execution/AGENTS.md +++ b/plugins/_code_execution/AGENTS.md @@ -21,6 +21,7 @@ - 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. - POSIX PTY I/O is nonblocking. The existing write/drain interface buffers partial writes and yields while the terminal is full; cancellation discards the shell only while POSIX input bytes remain unsent. Fully drained sends retain the shell and its state. +- SSH command writes must send every byte without blocking the event loop. Cancelling an in-flight SSH send closes its connection so a partial command cannot contaminate subsequent input. - TTY destruction uses best-effort kill, never nested event-loop execution; explicit close retains awaited cleanup. - Windows cleanup uses the native process methods without looking up POSIX-only signals; the pywinpty adapter passes SIGTERM for forced termination. POSIX close retains process-group TERM-then-KILL escalation. - Local and SSH session wrappers must synchronously release their owned process or connection resources when discarded. diff --git a/plugins/_code_execution/helpers/shell_ssh.py b/plugins/_code_execution/helpers/shell_ssh.py index 9092b9bad..6e8dfb5e2 100644 --- a/plugins/_code_execution/helpers/shell_ssh.py +++ b/plugins/_code_execution/helpers/shell_ssh.py @@ -86,7 +86,7 @@ class SSHInteractiveSession: initial_command = f"unset PROMPT_COMMAND PS0; stty -echo; {PAGER_DISABLE_COMMAND}" if self.cwd: initial_command = f"cd {self.cwd}; {initial_command}" - self.shell.send(f"{initial_command}\n".encode()) + await self.send_command(initial_command) # wait for initial prompt/output to settle while True: @@ -123,7 +123,12 @@ class SSHInteractiveSession: command = command + "\n" self.last_command = command.encode() self.trimmed_command_length = 0 - self.shell.send(self.last_command) + try: + await asyncio.to_thread(self.shell.sendall, self.last_command) + except asyncio.CancelledError: + # Stop the worker and discard a potentially incomplete command. + await self.close() + raise def is_terminated(self) -> bool: if not self.shell: diff --git a/tests/test_shell_ssh_io.py b/tests/test_shell_ssh_io.py new file mode 100644 index 000000000..829dbca50 --- /dev/null +++ b/tests/test_shell_ssh_io.py @@ -0,0 +1,73 @@ +"""Exercise Paramiko's real packet and flow-control behavior without a server.""" + +import asyncio +import threading +import time +from types import SimpleNamespace + +import paramiko +import pytest + +from plugins._code_execution.helpers.shell_ssh import SSHInteractiveSession + + +def make_session(window=2**21): + received = bytearray() + + def receive(message): + message = paramiko.Message(message.asbytes()) + if message.get_byte() == b"\x5e": # SSH_MSG_CHANNEL_DATA + message.get_int() + received.extend(message.get_binary()) + + channel = paramiko.Channel(0) + channel.transport = SimpleNamespace( + _send_user_message=receive, _sanitize_packet_size=lambda size: size + ) + channel._set_remote_channel(0, window, 32768) + session = object.__new__(SSHInteractiveSession) + session.shell = channel + session.client = SimpleNamespace(close=lambda: None) + return session, received + + +def test_ssh_sends_complete_nested_heredoc_across_packets(): + command = "bash <<'OUTER'\ncat <<'INNER'\n" + "héllo 🐇\n" * 20000 + "INNER\nOUTER" + + async def run(): + session, received = make_session() + try: + await session.send_command(command) + assert bytes(received) == (command + "\n").encode() + assert not session.shell.closed + finally: + await session.close() + + asyncio.run(run()) + + +def test_cancel_blocked_ssh_send_closes_channel_without_stalling_loop(): + async def run(): + session, received = make_session(window=8) + # Also bound a regression that blocks the event loop in sendall(). + watchdog = threading.Timer(2, session.shell.close) + watchdog.start() + pending = asyncio.create_task(session.send_command("x" * 100000)) + started = time.monotonic() + try: + while not received and not pending.done(): + await asyncio.sleep(0.01) + assert time.monotonic() - started < 1 + assert bytes(received) == b"x" * 8 + assert not pending.done() + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(pending, 1) + assert session.shell.closed + finally: + await session.close() + watchdog.cancel() + await asyncio.gather(pending, return_exceptions=True) + + # Loop shutdown also waits for the cancelled send's worker to finish. + asyncio.run(run())