@neoline/hostbridge 2.0.4 → 2.0.5

package/src/server_control_mcp/transports/ssh.py

@@ -4,20 +4,31 @@ from __future__ import annotations
 
 import asyncio
 import hashlib
-import threading
+import math
 import uuid
 from collections import deque
-from collections.abc import AsyncGenerator, Awaitable, Iterable, Iterator, Sequence
+from collections.abc import AsyncGenerator, Sequence
 from contextlib import suppress
 from dataclasses import dataclass
-from pathlib import Path
-from typing import TypeVar
 
 import asyncssh
 
-from .base import ExecResult, ShellReadResult, TransferResult, TransportCapabilities, TransportError
-
-T = TypeVar("T")
+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:
@@ -62,18 +73,36 @@ class SshConnectionConfig:
     password: str | None = None
 
 
-class _AsyncSshTransport:
+DEFAULT_SFTP_CLEANUP_TIMEOUT_SECONDS = 5.0
+
+
+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, *, connect_timeout: float | None = None) -> _AsyncSshTransport:
+    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:
@@ -95,7 +124,7 @@ class _AsyncSshTransport:
         self,
         command: str,
         *,
-        stdin: Iterable[bytes] | None,
+        stdin: ByteStream | None,
         timeout: float,
         output_limit: int,
     ) -> ExecResult:
@@ -109,15 +138,8 @@ class _AsyncSshTransport:
             self._connection.create_process(command, encoding=None)
         )
         try:
-            process = await asyncio.shield(open_task)
+            process = await open_task
         except asyncio.CancelledError:
-            self._connection.abort()
-            close_task = asyncio.create_task(self._connection.wait_closed())
-            while not close_task.done():
-                try:
-                    await asyncio.shield(close_task)
-                except asyncio.CancelledError:
-                    continue
             if not open_task.done():
                 open_task.cancel()
             await asyncio.gather(open_task, return_exceptions=True)
@@ -134,17 +156,13 @@ class _AsyncSshTransport:
             stdout, stderr = output.outputs()
             return ExecResult(process.exit_status, stdout, stderr)
         except TimeoutError:
-            await self._abort_process(process, tasks)
+            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))
-            while not cleanup.done():
-                try:
-                    await asyncio.shield(cleanup)
-                except asyncio.CancelledError:
-                    continue
-            await cleanup
+            await finish_task_despite_cancellation(cleanup)
             raise
 
     @staticmethod
@@ -157,10 +175,10 @@ class _AsyncSshTransport:
         wait_task = tasks[-1]
         try:
             await asyncio.wait_for(asyncio.shield(wait_task), timeout=1)
-        except BaseException:
+        except Exception:
             process.kill()
             process.close()
-            with suppress(BaseException):
+            with suppress(Exception):
                 await asyncio.wait_for(asyncio.shield(wait_task), timeout=1)
         for task in tasks:
             if not task.done():
@@ -168,15 +186,9 @@ class _AsyncSshTransport:
         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:
-            iterator = iter(stdin)
-            while True:
-                present, chunk = await asyncio.to_thread(_next_chunk, iterator)
-                if not present:
-                    break
-                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()
@@ -189,81 +201,100 @@ class _AsyncSshTransport:
                 break
             output.append(bytes(chunk), stderr=stderr)
 
-    async def upload_file(
-        self,
-        local_path: Path | str,
-        remote_path: str,
-        *,
-        chunk_size: int = 256 * 1024,
-        mode: int = 0o644,
-    ) -> 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()
-            await sftp.chmod(temporary, mode)
-            await sftp.rename(temporary, remote_path)
-        except Exception:
-            with suppress(Exception):
-                await sftp.remove(temporary)
-            raise
-        finally:
-            sftp.exit()
-        return TransferResult(transferred, digest.hexdigest())
-
     async def upload(
         self,
-        chunks: Iterable[bytes],
+        chunks: ByteStream,
         remote_path: str,
         *,
         mode: int,
         expected_size: int | None,
+        expected_sha256: str | None,
     ) -> TransferResult:
         self._validate_remote_path(remote_path)
         temporary = f"{remote_path}.hostbridge-{uuid.uuid4().hex}.tmp"
         digest = hashlib.sha256()
         transferred = 0
-        iterator = iter(chunks)
         sftp = await self._connection.start_sftp_client()
         try:
             async with sftp.open(temporary, "wb", encoding=None) as remote_file:
-                while True:
-                    present, chunk = await asyncio.to_thread(_next_chunk, iterator)
-                    if not present:
-                        break
-                    if not isinstance(chunk, bytes):
-                        raise TypeError("upload chunks must be bytes")
+                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 BaseException:
-            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 _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)
@@ -276,45 +307,8 @@ class _AsyncSshTransport:
                         break
                     yield raw_chunk.encode("utf-8") if isinstance(raw_chunk, str) else bytes(raw_chunk)
         finally:
-            sftp.exit()
-
-    async def download_file(
-        self,
-        remote_path: str,
-        local_path: Path | str,
-        *,
-        chunk_size: int = 256 * 1024,
-    ) -> TransferResult:
-        destination = Path(local_path)
-        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
-        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:
@@ -326,154 +320,107 @@ class _AsyncSshTransport:
         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, data: bytes) -> None:
+    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()
+        process = await self._ensure_shell(shell_id)
         process.stdin.write(data)
         await process.stdin.drain()
 
-    async def shell_read(self, *, timeout: float, max_bytes: int) -> ShellReadResult:
-        if timeout < 0:
+    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()
-        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 _ensure_shell(self) -> asyncssh.SSHClientProcess[bytes]:
-        shell = getattr(self, "_shell", None)
-        if shell is None or shell.exit_status is not None:
-            shell = await self._connection.create_process(term_type="xterm", encoding=None)
-            self._shell = shell
-        return shell
-
-    async def close(self) -> None:
-        shell = getattr(self, "_shell", None)
-        if shell is not None:
-            shell.terminate()
-            shell.close()
-            with suppress(Exception):
-                await asyncio.wait_for(shell.wait_closed(), timeout=1)
-        self._connection.close()
-        await self._connection.wait_closed()
-
-
-def _next_chunk(iterator: Iterator[bytes]) -> tuple[bool, bytes]:
-    try:
-        return True, next(iterator)
-    except StopIteration:
-        return False, b""
-
-
-class _EventLoopThread:
-    def __init__(self) -> None:
-        self.loop = asyncio.new_event_loop()
-        self.thread = threading.Thread(target=self._run, name="hostbridge-ssh", daemon=True)
-        self.started = threading.Event()
-        self.thread.start()
-        self.started.wait()
-
-    def _run(self) -> None:
-        asyncio.set_event_loop(self.loop)
-        self.started.set()
-        self.loop.run_forever()
-        self.loop.run_until_complete(self.loop.shutdown_asyncgens())
-        self.loop.close()
-
-    def run(self, awaitable: Awaitable[T]) -> T:
-        if not self.thread.is_alive():
-            raise TransportError("SSH transport is closed")
-        return asyncio.run_coroutine_threadsafe(_await_value(awaitable), self.loop).result()
-
-    def stop(self) -> None:
-        if self.thread.is_alive():
-            self.loop.call_soon_threadsafe(self.loop.stop)
-            self.thread.join(timeout=2)
-
-
-class SshTransport:
-    capabilities = _AsyncSshTransport.capabilities
-
-    def __init__(self, core: _AsyncSshTransport, runner: _EventLoopThread) -> None:
-        self._core = core
-        self._runner = runner
-        self._close_lock = threading.Lock()
-        self._closed = False
-
-    @classmethod
-    def connect(cls, config: SshConnectionConfig, *, connect_timeout: float | None = None) -> SshTransport:
-        runner = _EventLoopThread()
-        try:
-            core = runner.run(_AsyncSshTransport.connect(config, connect_timeout=connect_timeout))
-        except BaseException:
-            runner.stop()
-            raise
-        return cls(core, runner)
-
-    def exec(
-        self,
-        command: str,
-        *,
-        stdin: Iterable[bytes] | None,
-        timeout: float,
-        output_limit: int,
-    ) -> ExecResult:
-        return self._run(self._core.exec(command, stdin=stdin, timeout=timeout, output_limit=output_limit))
-
-    def upload(
-        self,
-        chunks: Iterable[bytes],
-        remote_path: str,
-        *,
-        mode: int,
-        expected_size: int | None,
-    ) -> TransferResult:
-        return self._run(self._core.upload(chunks, remote_path, mode=mode, expected_size=expected_size))
-
-    def download(self, remote_path: str, *, chunk_size: int) -> Iterator[bytes]:
-        stream = self._core.download(remote_path, chunk_size=chunk_size)
-
-        def chunks() -> Iterator[bytes]:
+        process = await self._ensure_shell(shell_id)
+        if timeout is None:
+            data = await process.stdout.read(max_bytes)
+        else:
             try:
-                while True:
-                    try:
-                        yield self._run(stream.__anext__())
-                    except StopAsyncIteration:
-                        return
-            finally:
-                self._run(stream.aclose())
-
-        return chunks()
+                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())
 
-    def shell_write(self, data: bytes) -> None:
-        self._run(self._core.shell_write(data))
+    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()
 
-    def shell_read(self, *, timeout: float, max_bytes: int) -> ShellReadResult:
-        return self._run(self._core.shell_read(timeout=timeout, max_bytes=max_bytes))
+        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()
 
-    def close(self) -> None:
-        with self._close_lock:
-            if self._closed:
-                return
-            self._closed = True
+        async def output() -> AsyncGenerator[bytes, None]:
+            input_task = asyncio.create_task(send_input(), name=f"hostbridge-native-shell-input-{shell_id}")
             try:
-                self._runner.run(self._core.close())
+                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:
-                self._runner.stop()
+                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)
 
-    def _run(self, awaitable: Awaitable[T]) -> T:
+    @staticmethod
+    async def _close_shell_process(shell: asyncssh.SSHClientProcess[bytes]) -> None:
+        shell.terminate()
+        shell.close()
         try:
-            return self._runner.run(awaitable)
-        except (StopAsyncIteration, TypeError, ValueError, TransportError):
-            raise
-        except Exception as exc:
-            raise TransportError(str(exc) or exc.__class__.__name__) from exc
+            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()
 
 
-async def _await_value(awaitable: Awaitable[T]) -> T:
-    return await awaitable
+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)