@neoline/hostbridge 2.0.3 → 2.0.5

package/src/server_control_mcp/transports/ssh.py

@@ -4,15 +4,62 @@ from __future__ import annotations
 
 import asyncio
 import hashlib
+import math
 import uuid
-from collections.abc import Iterable, Sequence
+from collections import deque
+from collections.abc import AsyncGenerator, Sequence
 from contextlib import suppress
 from dataclasses import dataclass
-from pathlib import Path
 
 import asyncssh
 
-from .base import ExecResult, TransferResult, TransportCapabilities
+from ..async_lifecycle import (
+    finish_task_before_cancellation,
+    finish_task_despite_cancellation,
+    stop_input_after_peer_eof,
+)
+from .base import (
+    ByteStream,
+    ExecResult,
+    ShellAttachment,
+    ShellReadResult,
+    ShellTerminalResult,
+    TransferResult,
+    TransportCapabilities,
+    TransportError,
+    iterate_byte_stream,
+)
+
+
+class _CombinedOutputBuffer:
+    def __init__(self, limit: int) -> None:
+        self._limit = limit
+        self._segments: deque[tuple[bool, bytes]] = deque()
+        self._size = 0
+
+    def append(self, chunk: bytes, *, stderr: bool) -> None:
+        if not chunk:
+            return
+        self._segments.append((stderr, bytes(chunk)))
+        self._size += len(chunk)
+        excess = self._size - self._limit
+        while excess > 0:
+            segment_stderr, segment = self._segments[0]
+            if len(segment) <= excess:
+                self._segments.popleft()
+                self._size -= len(segment)
+                excess -= len(segment)
+            else:
+                self._segments[0] = (segment_stderr, segment[excess:])
+                self._size -= excess
+                excess = 0
+
+    def outputs(self) -> tuple[bytes, bytes]:
+        stdout = bytearray()
+        stderr = bytearray()
+        for is_stderr, segment in self._segments:
+            (stderr if is_stderr else stdout).extend(segment)
+        return bytes(stdout), bytes(stderr)
 
 
 @dataclass(frozen=True, slots=True)
@@ -21,21 +68,41 @@ class SshConnectionConfig:
     port: int = 22
     username: str | None = None
     client_keys: Sequence[str] = ()
-    known_hosts: str | None = None
+    known_hosts: str | None | tuple[()] = ()
+    tunnel: str | None = None
+    password: str | None = None
+
 
+DEFAULT_SFTP_CLEANUP_TIMEOUT_SECONDS = 5.0
 
-class SshTransport:
+
+class AsyncSshConnection:
     capabilities = TransportCapabilities(
         separate_stderr=True,
         native_binary_transfer=True,
         parallel_channels=8,
     )
 
-    def __init__(self, connection: asyncssh.SSHClientConnection) -> None:
+    def __init__(
+        self,
+        connection: asyncssh.SSHClientConnection,
+        *,
+        sftp_cleanup_timeout: float = DEFAULT_SFTP_CLEANUP_TIMEOUT_SECONDS,
+    ) -> None:
+        if (
+            not isinstance(sftp_cleanup_timeout, (int, float))
+            or isinstance(sftp_cleanup_timeout, bool)
+            or not math.isfinite(sftp_cleanup_timeout)
+            or sftp_cleanup_timeout <= 0
+        ):
+            raise ValueError("sftp_cleanup_timeout must be a finite positive number")
         self._connection = connection
+        self._sftp_cleanup_timeout = float(sftp_cleanup_timeout)
+        self._shells: dict[str, asyncssh.SSHClientProcess[bytes]] = {}
+        self._shell_creation_lock = asyncio.Lock()
 
     @classmethod
-    async def connect(cls, config: SshConnectionConfig) -> SshTransport:
+    async def connect(cls, config: SshConnectionConfig, *, connect_timeout: float | None = None) -> AsyncSshConnection:
         if not config.host.strip():
             raise ValueError("SSH host must not be empty")
         if not 1 <= config.port <= 65535:
@@ -46,6 +113,9 @@ class SshTransport:
             username=config.username,
             client_keys=list(config.client_keys),
             known_hosts=config.known_hosts,
+            tunnel=config.tunnel,
+            password=config.password,
+            connect_timeout=connect_timeout,
             encoding=None,
         )
         return cls(connection)
@@ -54,7 +124,7 @@ class SshTransport:
         self,
         command: str,
         *,
-        stdin: Iterable[bytes] | None,
+        stdin: ByteStream | None,
         timeout: float,
         output_limit: int,
     ) -> ExecResult:
@@ -64,121 +134,181 @@ class SshTransport:
             raise ValueError("timeout must be positive")
         if output_limit <= 0:
             raise ValueError("output_limit must be positive")
-        process = await self._connection.create_process(command, encoding=None)
-        stdout_task = asyncio.create_task(self._read_stream(process.stdout, output_limit))
-        stderr_task = asyncio.create_task(self._read_stream(process.stderr, output_limit))
+        open_task: asyncio.Future[asyncssh.SSHClientProcess[bytes]] = asyncio.ensure_future(
+            self._connection.create_process(command, encoding=None)
+        )
+        try:
+            process = await open_task
+        except asyncio.CancelledError:
+            if not open_task.done():
+                open_task.cancel()
+            await asyncio.gather(open_task, return_exceptions=True)
+            raise
+        output = _CombinedOutputBuffer(output_limit)
+        stdout_task = asyncio.create_task(self._read_stream(process.stdout, output, stderr=False))
+        stderr_task = asyncio.create_task(self._read_stream(process.stderr, output, stderr=True))
         input_task = asyncio.create_task(self._write_stdin(process, stdin))
+        wait_task = asyncio.create_task(process.wait_closed())
+        tasks = (stdout_task, stderr_task, input_task, wait_task)
         try:
-            await asyncio.wait_for(process.wait_closed(), timeout=timeout)
-            await input_task
-            stdout, stderr = await asyncio.gather(stdout_task, stderr_task)
+            await asyncio.wait_for(asyncio.gather(asyncio.shield(wait_task), input_task), timeout=timeout)
+            await asyncio.gather(stdout_task, stderr_task)
+            stdout, stderr = output.outputs()
             return ExecResult(process.exit_status, stdout, stderr)
         except TimeoutError:
-            process.terminate()
+            cleanup = asyncio.create_task(self._abort_process(process, tasks))
+            await finish_task_before_cancellation(cleanup)
+            stdout, stderr = output.outputs()
+            return ExecResult(None, stdout, stderr, timed_out=True)
+        except BaseException:
+            cleanup = asyncio.create_task(self._abort_process(process, tasks))
+            await finish_task_despite_cancellation(cleanup)
+            raise
+
+    @staticmethod
+    async def _abort_process(
+        process: asyncssh.SSHClientProcess[bytes],
+        tasks: tuple[asyncio.Task[None], asyncio.Task[None], asyncio.Task[None], asyncio.Task[None]],
+    ) -> None:
+        process.terminate()
+        process.close()
+        wait_task = tasks[-1]
+        try:
+            await asyncio.wait_for(asyncio.shield(wait_task), timeout=1)
+        except Exception:
+            process.kill()
+            process.close()
             with suppress(Exception):
-                await process.wait_closed()
-            for task in (stdout_task, stderr_task, input_task):
+                await asyncio.wait_for(asyncio.shield(wait_task), timeout=1)
+        for task in tasks:
+            if not task.done():
                 task.cancel()
-            await asyncio.gather(stdout_task, stderr_task, input_task, return_exceptions=True)
-            return ExecResult(None, b"", b"", timed_out=True)
+        await asyncio.gather(*tasks, return_exceptions=True)
 
     @staticmethod
-    async def _write_stdin(process: asyncssh.SSHClientProcess[bytes], stdin: Iterable[bytes] | None) -> None:
+    async def _write_stdin(process: asyncssh.SSHClientProcess[bytes], stdin: ByteStream | None) -> None:
         if stdin is not None:
-            for chunk in stdin:
-                if not isinstance(chunk, bytes):
-                    raise TypeError("stdin chunks must be bytes")
+            async for chunk in iterate_byte_stream(stdin, label="stdin"):
                 process.stdin.write(chunk)
                 await process.stdin.drain()
         process.stdin.write_eof()
 
     @staticmethod
-    async def _read_stream(stream, output_limit: int) -> bytes:
-        output = bytearray()
+    async def _read_stream(stream, output: _CombinedOutputBuffer, *, stderr: bool) -> None:
         while True:
             chunk = await stream.read(64 * 1024)
             if not chunk:
                 break
-            output.extend(chunk)
-            if len(output) > output_limit:
-                del output[: len(output) - output_limit]
-        return bytes(output)
+            output.append(bytes(chunk), stderr=stderr)
 
-    async def upload_file(
+    async def upload(
         self,
-        local_path: Path | str,
+        chunks: ByteStream,
         remote_path: str,
         *,
-        chunk_size: int = 256 * 1024,
-        mode: int = 0o644,
+        mode: int,
+        expected_size: int | None,
+        expected_sha256: str | None,
     ) -> TransferResult:
-        source = Path(local_path)
-        self._validate_chunk_size(chunk_size)
         self._validate_remote_path(remote_path)
         temporary = f"{remote_path}.hostbridge-{uuid.uuid4().hex}.tmp"
         digest = hashlib.sha256()
         transferred = 0
         sftp = await self._connection.start_sftp_client()
         try:
-            local_file = source.open("rb")
-            try:
-                async with sftp.open(temporary, "wb") as remote_file:
-                    while True:
-                        chunk = await asyncio.to_thread(local_file.read, chunk_size)
-                        if not chunk:
-                            break
-                        await remote_file.write(chunk)
-                        digest.update(chunk)
-                        transferred += len(chunk)
-            finally:
-                local_file.close()
+            async with sftp.open(temporary, "wb", encoding=None) as remote_file:
+                async for chunk in iterate_byte_stream(chunks, label="upload"):
+                    await remote_file.write(chunk)
+                    digest.update(chunk)
+                    transferred += len(chunk)
+            if expected_size is not None and transferred != expected_size:
+                raise TransportError(f"upload size mismatch: expected {expected_size}, received {transferred}")
+            actual_sha256 = digest.hexdigest()
+            if expected_sha256 is not None and actual_sha256.lower() != expected_sha256.lower():
+                raise TransportError(
+                    f"upload SHA-256 mismatch: expected {expected_sha256.lower()}, received {actual_sha256}"
+                )
             await sftp.chmod(temporary, mode)
             await sftp.rename(temporary, remote_path)
-        except Exception:
-            with suppress(Exception):
-                await sftp.remove(temporary)
+        except BaseException as primary_error:
+            cleanup = asyncio.create_task(self._finish_sftp(sftp, temporary=temporary))
+            try:
+                await finish_task_despite_cancellation(cleanup)
+            except BaseException as cleanup_error:
+                primary_error.add_note(f"SFTP temporary cleanup failed: {cleanup_error}")
             raise
-        finally:
-            sftp.exit()
+        cleanup = asyncio.create_task(self._finish_sftp(sftp))
+        await finish_task_before_cancellation(cleanup)
         return TransferResult(transferred, digest.hexdigest())
 
-    async def download_file(
-        self,
-        remote_path: str,
-        local_path: Path | str,
-        *,
-        chunk_size: int = 256 * 1024,
-    ) -> TransferResult:
-        destination = Path(local_path)
+    async def _finish_sftp(self, sftp: asyncssh.SFTPClient, *, temporary: str | None = None) -> None:
+        removed = await self._close_sftp_channel(sftp, temporary=temporary)
+        if temporary is None or removed:
+            return
+
+        replacement: asyncssh.SFTPClient | None = None
+        try:
+            async with asyncio.timeout(self._sftp_cleanup_timeout):
+                replacement = await self._connection.start_sftp_client()
+        except Exception as exc:
+            self._abort_connection()
+            raise TransportError(f"could not reopen SFTP to remove temporary file {temporary!r}") from exc
+
+        if await self._close_sftp_channel(replacement, temporary=temporary):
+            return
+        self._abort_connection()
+        raise TransportError(f"could not remove SFTP temporary file {temporary!r}")
+
+    async def _close_sftp_channel(self, sftp: asyncssh.SFTPClient, *, temporary: str | None) -> bool:
+        exit_sent = False
+        removed = temporary is None
+        try:
+            async with asyncio.timeout(self._sftp_cleanup_timeout):
+                if temporary is not None:
+                    try:
+                        await sftp.remove(temporary)
+                    except asyncssh.SFTPNoSuchFile:
+                        removed = True
+                    else:
+                        removed = True
+                sftp.exit()
+                exit_sent = True
+                with suppress(Exception):
+                    await sftp.wait_closed()
+        except Exception:
+            pass
+        finally:
+            if not exit_sent:
+                with suppress(Exception):
+                    sftp.exit()
+        return removed
+
+    def _abort_connection(self) -> None:
+        try:
+            self._connection.abort()
+        except Exception as exc:
+            try:
+                self._connection.close()
+            except Exception as fallback_error:
+                fallback_error.add_note(f"SSH connection abort failed: {exc}")
+                raise TransportError(
+                    "could not terminate SSH connection after SFTP cleanup failure"
+                ) from fallback_error
+
+    async def download(self, remote_path: str, *, chunk_size: int) -> AsyncGenerator[bytes, None]:
         self._validate_chunk_size(chunk_size)
         self._validate_remote_path(remote_path)
-        temporary = destination.with_name(f".{destination.name}.hostbridge-{uuid.uuid4().hex}.tmp")
-        digest = hashlib.sha256()
-        transferred = 0
         sftp = await self._connection.start_sftp_client()
         try:
-            destination.parent.mkdir(parents=True, exist_ok=True)
-            local_file = temporary.open("wb")
-            try:
-                async with sftp.open(remote_path, "rb") as remote_file:
-                    while True:
-                        raw_chunk = await remote_file.read(chunk_size)
-                        if not raw_chunk:
-                            break
-                        chunk = raw_chunk.encode("utf-8") if isinstance(raw_chunk, str) else bytes(raw_chunk)
-                        await asyncio.to_thread(local_file.write, chunk)
-                        digest.update(chunk)
-                        transferred += len(chunk)
-            finally:
-                local_file.close()
-            temporary.chmod(0o600)
-            temporary.replace(destination)
-        except Exception:
-            temporary.unlink(missing_ok=True)
-            raise
+            async with sftp.open(remote_path, "rb", encoding=None) as remote_file:
+                while True:
+                    raw_chunk = await remote_file.read(chunk_size)
+                    if not raw_chunk:
+                        break
+                    yield raw_chunk.encode("utf-8") if isinstance(raw_chunk, str) else bytes(raw_chunk)
         finally:
-            sftp.exit()
-        return TransferResult(transferred, digest.hexdigest())
+            cleanup = asyncio.create_task(self._finish_sftp(sftp))
+            await finish_task_before_cancellation(cleanup)
 
     @staticmethod
     def _validate_chunk_size(chunk_size: int) -> None:
@@ -190,6 +320,107 @@ class SshTransport:
         if not isinstance(remote_path, str) or not remote_path.strip() or "\x00" in remote_path:
             raise ValueError("remote_path must be a non-empty path without NUL bytes")
 
+    async def shell_write(self, shell_id: str, data: bytes) -> None:
+        if not isinstance(data, bytes):
+            raise TypeError("shell data must be bytes")
+        process = await self._ensure_shell(shell_id)
+        process.stdin.write(data)
+        await process.stdin.drain()
+
+    async def shell_read(self, shell_id: str, *, timeout: float | None, max_bytes: int) -> ShellReadResult:
+        if timeout is not None and timeout < 0:
+            raise ValueError("timeout must be non-negative")
+        if max_bytes <= 0:
+            raise ValueError("max_bytes must be positive")
+        process = await self._ensure_shell(shell_id)
+        if timeout is None:
+            data = await process.stdout.read(max_bytes)
+        else:
+            try:
+                data = await asyncio.wait_for(process.stdout.read(max_bytes), timeout=timeout)
+            except TimeoutError:
+                data = b""
+        return ShellReadResult(bytes(data), process.exit_status is None and not self._connection.is_closed())
+
+    async def attach_shell(self, shell_id: str, stdin: ByteStream, *, max_bytes: int) -> ShellAttachment:
+        process = await self._ensure_shell(shell_id)
+        terminal = asyncio.get_running_loop().create_future()
+
+        async def send_input() -> None:
+            async for chunk in iterate_byte_stream(stdin, label="shell input"):
+                if len(chunk) > max_bytes:
+                    raise TransportError(f"shell input chunk exceeds {max_bytes} bytes")
+                process.stdin.write(chunk)
+                await process.stdin.drain()
+            process.stdin.write_eof()
+
+        async def output() -> AsyncGenerator[bytes, None]:
+            input_task = asyncio.create_task(send_input(), name=f"hostbridge-native-shell-input-{shell_id}")
+            try:
+                while data := await process.stdout.read(max_bytes):
+                    yield bytes(data)
+                await stop_input_after_peer_eof(input_task)
+                await process.wait_closed()
+                terminal.set_result(_ssh_shell_terminal(process))
+            except BaseException as exc:
+                if not terminal.done():
+                    terminal.set_exception(exc)
+                    terminal.exception()
+                raise
+            finally:
+                if not input_task.done():
+                    input_task.cancel()
+                cleanup = asyncio.gather(input_task, return_exceptions=True)
+                await finish_task_before_cancellation(cleanup)
+
+        return ShellAttachment(output(), terminal)
+
+    async def _ensure_shell(self, shell_id: str) -> asyncssh.SSHClientProcess[bytes]:
+        if not shell_id:
+            raise ValueError("shell_id must not be empty")
+        async with self._shell_creation_lock:
+            shell = self._shells.get(shell_id)
+            if shell is None or shell.exit_status is not None:
+                shell = await self._connection.create_process(term_type="xterm", encoding=None)
+                self._shells[shell_id] = shell
+            return shell
+
+    async def shell_close(self, shell_id: str) -> None:
+        async with self._shell_creation_lock:
+            shell = self._shells.pop(shell_id, None)
+        if shell is None:
+            return
+        cleanup = asyncio.create_task(self._close_shell_process(shell))
+        await finish_task_before_cancellation(cleanup)
+
+    @staticmethod
+    async def _close_shell_process(shell: asyncssh.SSHClientProcess[bytes]) -> None:
+        shell.terminate()
+        shell.close()
+        try:
+            await asyncio.wait_for(shell.wait_closed(), timeout=1)
+        except TimeoutError:
+            shell.kill()
+            shell.close()
+            with suppress(Exception):
+                await asyncio.wait_for(shell.wait_closed(), timeout=1)
+        except Exception:
+            pass
+
     async def close(self) -> None:
+        await asyncio.gather(*(self.shell_close(shell_id) for shell_id in tuple(self._shells)), return_exceptions=True)
         self._connection.close()
         await self._connection.wait_closed()
+
+
+def _ssh_shell_terminal(process: asyncssh.SSHClientProcess[bytes]) -> ShellTerminalResult:
+    exit_status = process.exit_status
+    raw_signal = process.exit_signal
+    signal_name: str | None = None
+    if isinstance(raw_signal, tuple) and raw_signal and isinstance(raw_signal[0], str):
+        signal_name = raw_signal[0]
+    elif isinstance(raw_signal, str):
+        signal_name = raw_signal
+    if signal_name:
+        return ShellTerminalResult(None, signal_name)
+    return ShellTerminalResult(exit_status, None)