mirror of
https://github.com/agent0ai/agent-zero.git
synced 2026-10-19 15:33:54 +00:00
Send complete SSH commands without blocking the event loop
Paramiko Channel.send can accept only part of a command. Ignoring its byte count truncated large inputs, including heredocs whose closing delimiters never reached the shell. Use sendall through asyncio.to_thread and route shell initialization through the same send path. Close the SSH connection when an in-flight send is cancelled so partial input cannot contaminate later commands. Cancellation intentionally discards that SSH session and its shell state. Add regressions for multi-packet Unicode delivery and blocked-window responsiveness/cancellation, and document the SSH send contract. Validation: 24 focused tests passed. Real Paramiko transport and container OpenSSH probes delivered complete large single/nested heredocs and preserved shell state across terminal, Python and Node.js calls. No local TTY, prompt readiness, tool-dispatch, or platform-specific heredoc workaround changes.
This commit is contained in:
parent
95ba588abc
commit
ad5bebbea5
3 changed files with 81 additions and 2 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
73
tests/test_shell_ssh_io.py
Normal file
73
tests/test_shell_ssh_io.py
Normal file
|
|
@ -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())
|
||||
Loading…
Add table
Add a link
Reference in a new issue