@neoline/hostbridge 2.0.3 → 2.0.5

package/src/server_control_mcp/mock_ssh_exec.py

@@ -0,0 +1,299 @@
+from __future__ import annotations
+
+import asyncio
+import queue
+import sys
+import threading
+from contextlib import suppress
+from typing import cast
+
+import asyncssh
+
+from .async_lifecycle import finish_task_despite_cancellation
+from .client import ExecResult, HostBridgeClient, HostBridgeError
+from .mock_ssh_session import HostBridgeSessionFactory
+
+
+def _to_bytes(chunk: str | bytes) -> bytes:
+    if isinstance(chunk, bytes):
+        return chunk
+    return chunk.encode("utf-8")
+
+
+class HostBridgeSSHProcess:
+    def __init__(
+        self, client: HostBridgeClient, session_factory: HostBridgeSessionFactory, timeout: int, output_limit: int
+    ) -> None:
+        self._client = client
+        self._session_factory = session_factory
+        self._timeout = timeout
+        self._output_limit = output_limit
+
+    async def __call__(self, process: asyncssh.SSHServerProcess[bytes]) -> None:
+        open_task = asyncio.create_task(asyncio.to_thread(self._session_factory.open))
+        try:
+            session_id = await asyncio.shield(open_task)
+        except asyncio.CancelledError:
+            session_id = None
+            with suppress(BaseException):
+                session_id = await finish_task_despite_cancellation(open_task)
+            if session_id is not None:
+                try:
+                    await self._close_session(session_id, reusable=False)
+                except Exception as exc:
+                    self._report_shell_error(process, exc)
+            self._set_exit_status(process, 255)
+            return
+        if session_id is None:
+            self._set_exit_status(process, 255)
+            return
+        reusable = False
+        try:
+            reusable = await self._run(process, session_id)
+        except asyncio.CancelledError:
+            self._set_exit_status(process, 255)
+            return
+        finally:
+            try:
+                await self._close_session(session_id, reusable=reusable)
+            except Exception as exc:
+                self._report_shell_error(process, exc)
+                self._set_exit_status(process, 255)
+
+    async def _close_session(self, session_id: str, *, reusable: bool) -> None:
+        close_task = asyncio.create_task(asyncio.to_thread(self._session_factory.close, session_id, reusable=reusable))
+        await finish_task_despite_cancellation(close_task)
+
+    async def _run(self, process: asyncssh.SSHServerProcess[bytes], session_id: str) -> bool:
+        command = process.command
+        if not command:
+            await self._run_shell(process, session_id)
+            return False
+
+        try:
+            result = await self._exec_stream(process, session_id, command)
+            if result.stdout:
+                process.stdout.write(result.stdout)
+            if result.stderr:
+                process.stderr.write(result.stderr)
+            if result.timed_out:
+                self._set_exit_status(process, 124)
+                return False
+            if result.exit_code is None:
+                self._set_exit_status(process, 255)
+                return False
+            self._set_exit_status(process, int(result.exit_code))
+            return True
+        except Exception as exc:
+            print(f"hostbridge mock-ssh exec failed: {exc}", file=sys.stderr, flush=True)
+            with suppress(Exception):
+                process.stderr.write(f"hostbridge mock-ssh error: {exc}\n".encode("utf-8", errors="replace"))
+            status = 124 if isinstance(exc, HostBridgeError) and exc.code == "timeout" else 1
+            self._set_exit_status(process, status)
+            return False
+
+    async def _exec_stream(
+        self,
+        process: asyncssh.SSHServerProcess[bytes],
+        session_id: str,
+        command: str,
+    ):
+        loop = asyncio.get_running_loop()
+        items: queue.Queue[bytes | BaseException | object] = queue.Queue(maxsize=8)
+        end = object()
+        cancelled = threading.Event()
+
+        def chunks():
+            while True:
+                item = items.get()
+                if item is end:
+                    return
+                if isinstance(item, BaseException):
+                    raise item
+                yield cast(bytes, item)
+
+        consumer = loop.run_in_executor(
+            None,
+            lambda: self._client.exec_stream(
+                session_id,
+                command,
+                chunks(),
+                timeout=self._timeout,
+                output_limit=self._output_limit,
+                owner="mock-ssh",
+                cancel_event=cancelled,
+            ),
+        )
+        staging_deadline = loop.time() + self._timeout
+
+        async def put(item: bytes | BaseException | object) -> bool:
+            while True:
+                if consumer.done():
+                    await consumer
+                    return False
+                try:
+                    items.put_nowait(item)
+                    return True
+                except queue.Full:
+                    await asyncio.sleep(0.005)
+
+        try:
+            while not process.stdin.at_eof():
+                remaining = staging_deadline - loop.time()
+                if remaining <= 0:
+                    if not await put(HostBridgeError("timeout", "timed out waiting for mock-ssh exec stdin")):
+                        return await consumer
+                    break
+                read_task = asyncio.create_task(process.stdin.read(64 * 1024))
+                try:
+                    completed, _ = await asyncio.wait(
+                        (read_task, consumer),
+                        timeout=remaining,
+                        return_when=asyncio.FIRST_COMPLETED,
+                    )
+                except BaseException:
+                    read_task.cancel()
+                    await asyncio.gather(read_task, return_exceptions=True)
+                    raise
+                if consumer in completed:
+                    read_task.cancel()
+                    await asyncio.gather(read_task, return_exceptions=True)
+                    return await consumer
+                if read_task not in completed:
+                    read_task.cancel()
+                    await asyncio.gather(read_task, return_exceptions=True)
+                    if not await put(HostBridgeError("timeout", "timed out waiting for mock-ssh exec stdin")):
+                        return await consumer
+                    break
+                try:
+                    data = read_task.result()
+                except asyncio.CancelledError:
+                    raise
+                except BaseException as exc:
+                    if not await put(exc):
+                        return await consumer
+                    break
+                if not data:
+                    break
+                if not await put(_to_bytes(data)):
+                    return await consumer
+            if not await put(end):
+                return await consumer
+            return await asyncio.shield(consumer)
+        except asyncio.CancelledError:
+            cancelled.set()
+            close_succeeded = False
+
+            def observe_consumer(future: asyncio.Future[ExecResult]) -> None:
+                try:
+                    future.result()
+                except asyncio.CancelledError:
+                    return
+                except Exception as exc:
+                    self._report_shell_error(process, exc)
+
+            try:
+                await self._close_session(session_id, reusable=False)
+                close_succeeded = True
+            except Exception as exc:
+                self._report_shell_error(process, exc)
+            with suppress(queue.Full):
+                items.put_nowait(end)
+            if close_succeeded:
+                with suppress(BaseException):
+                    await finish_task_despite_cancellation(consumer)
+            else:
+                consumer.add_done_callback(observe_consumer)
+            raise
+
+    @staticmethod
+    def _set_exit_status(process: asyncssh.SSHServerProcess[bytes], status: int) -> None:
+        with suppress(Exception):
+            process.exit(status)
+
+    async def _run_shell(self, process: asyncssh.SSHServerProcess[bytes], session_id: str) -> None:
+        stop = asyncio.Event()
+        failed = False
+        terminal_failed = False
+        exit_code: int | None = None
+        exit_signal: str | None = None
+        open_task = asyncio.create_task(asyncio.to_thread(self._client.open_shell, session_id, owner="mock-ssh"))
+        try:
+            shell = await asyncio.shield(open_task)
+        except asyncio.CancelledError:
+            with suppress(BaseException):
+                shell = await finish_task_despite_cancellation(open_task)
+                close_task = asyncio.create_task(asyncio.to_thread(shell.close))
+                await finish_task_despite_cancellation(close_task)
+            raise
+
+        async def pump_input() -> None:
+            nonlocal failed
+            try:
+                while not process.stdin.at_eof():
+                    data = await process.stdin.read(4096)
+                    if not data:
+                        break
+                    await asyncio.to_thread(shell.send, _to_bytes(data))
+            except Exception as exc:
+                failed = True
+                self._report_shell_error(process, exc)
+            finally:
+                try:
+                    await asyncio.to_thread(shell.send_eof)
+                except Exception as exc:
+                    failed = True
+                    self._report_shell_error(process, exc)
+                    stop.set()
+
+        async def pump_output() -> None:
+            nonlocal exit_code, exit_signal, failed, terminal_failed
+            try:
+                while data := await asyncio.to_thread(shell.receive):
+                    process.stdout.write(data)
+                exit_code = shell.exit_code
+                exit_signal = shell.exit_signal
+            except Exception as exc:
+                failed = True
+                terminal_failed = True
+                self._report_shell_error(process, exc)
+            finally:
+                stop.set()
+
+        tasks = [asyncio.create_task(pump_input()), asyncio.create_task(pump_output())]
+        try:
+            await stop.wait()
+        finally:
+            try:
+                close_task = asyncio.create_task(asyncio.to_thread(shell.close))
+                await finish_task_despite_cancellation(close_task)
+            except Exception as exc:
+                failed = True
+                self._report_shell_error(process, exc)
+            finally:
+                for task in tasks:
+                    task.cancel()
+                await asyncio.gather(*tasks, return_exceptions=True)
+        if terminal_failed:
+            self._set_exit_status(process, 255)
+        elif failed:
+            self._set_exit_status(process, 1)
+        elif isinstance(exit_code, int) and not isinstance(exit_code, bool) and 0 <= exit_code <= 255:
+            self._set_exit_status(process, exit_code)
+        elif isinstance(exit_signal, str) and exit_signal.strip():
+            self._set_exit_signal(process, exit_signal)
+        else:
+            self._set_exit_status(process, 255)
+
+    @staticmethod
+    def _report_shell_error(process: asyncssh.SSHServerProcess[bytes], exc: Exception) -> None:
+        print(f"hostbridge mock-ssh shell failed: {exc}", file=sys.stderr, flush=True)
+        with suppress(Exception):
+            process.stderr.write(f"hostbridge mock-ssh shell error: {exc}\n".encode("utf-8", errors="replace"))
+
+    @classmethod
+    def _set_exit_signal(cls, process: asyncssh.SSHServerProcess[bytes], signal_name: str) -> None:
+        try:
+            process.exit_with_signal(signal_name)
+        except Exception:
+            cls._set_exit_status(process, 255)

package/src/server_control_mcp/mock_ssh_session.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import threading
+
+from .client import HostBridgeClient, HostBridgeError
+
+
+class HostBridgeSessionFactory:
+    def __init__(self, client: HostBridgeClient, host_id: str, connect_timeout: int = 60) -> None:
+        self._client = client
+        self._host_id = host_id
+        self._connect_timeout = connect_timeout
+        self._lock = threading.Lock()
+        self._leased: set[str] = set()
+        self._shutdown = False
+
+    def open(self) -> str:
+        with self._lock:
+            if self._shutdown:
+                raise RuntimeError("mock-ssh session factory is shut down")
+        session = self._client.open_session(
+            self._host_id,
+            owner="mock-ssh",
+            connect_timeout=self._connect_timeout,
+        )
+        session_id = session.session_id
+        with self._lock:
+            if not self._shutdown:
+                self._leased.add(session_id)
+                return session_id
+        self._close_backend(session_id)
+        raise RuntimeError("mock-ssh session factory is shut down")
+
+    def close(self, session_id: str, *, reusable: bool = True) -> None:
+        with self._lock:
+            if session_id not in self._leased:
+                return
+        self._close_backend(session_id)
+        with self._lock:
+            self._leased.discard(session_id)
+
+    def shutdown(self) -> None:
+        with self._lock:
+            self._shutdown = True
+            session_ids = list(self._leased)
+        last_error: Exception | None = None
+        for session_id in session_ids:
+            try:
+                self.close(session_id, reusable=False)
+            except Exception as exc:
+                last_error = exc
+        if last_error is not None:
+            raise last_error
+
+    def _close_backend(self, session_id: str) -> None:
+        last_error: Exception | None = None
+        for _ in range(2):
+            try:
+                self._client.release_mock_ssh_session(session_id)
+                return
+            except HostBridgeError as exc:
+                if exc.code == "not_found":
+                    return
+                last_error = exc
+            except Exception as exc:
+                last_error = exc
+        if last_error is None:
+            raise RuntimeError("backend close failed without an error")
+        raise last_error