@neoline/hostbridge 2.0.3 → 2.0.4

package/src/server_control_mcp/mock_ssh_exec.py

@@ -0,0 +1,210 @@
+from __future__ import annotations
+
+import asyncio
+import queue
+import sys
+import threading
+from contextlib import suppress
+from functools import partial
+from typing import cast
+
+import asyncssh
+
+from .client import 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 open_task
+            if session_id is not None:
+                await asyncio.to_thread(self._session_factory.close, session_id, reusable=False)
+            self._set_exit_status(process, 255)
+            return
+        assert session_id is not None
+        reusable = False
+        try:
+            reusable = await self._run(process, session_id)
+        except asyncio.CancelledError:
+            self._set_exit_status(process, 255)
+            return
+        finally:
+            await asyncio.to_thread(self._session_factory.close, session_id, reusable=reusable)
+
+    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) -> None:
+            while True:
+                try:
+                    items.put_nowait(item)
+                    return
+                except queue.Full:
+                    if consumer.done():
+                        await consumer
+                    await asyncio.sleep(0.005)
+
+        try:
+            while not process.stdin.at_eof():
+                remaining = staging_deadline - loop.time()
+                if remaining <= 0:
+                    await put(HostBridgeError("timeout", "timed out waiting for mock-ssh exec stdin"))
+                    break
+                try:
+                    data = await asyncio.wait_for(process.stdin.read(64 * 1024), timeout=remaining)
+                except TimeoutError:
+                    await put(HostBridgeError("timeout", "timed out waiting for mock-ssh exec stdin"))
+                    break
+                except BaseException as exc:
+                    await put(exc)
+                    break
+                if not data:
+                    break
+                await put(_to_bytes(data))
+            await put(end)
+            return await consumer
+        except asyncio.CancelledError:
+            cancelled.set()
+            await put(end)
+            with suppress(BaseException):
+                await asyncio.shield(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:
+        loop = asyncio.get_running_loop()
+        stop = asyncio.Event()
+        input_done = asyncio.Event()
+
+        async def pump_input() -> None:
+            try:
+                while not process.stdin.at_eof():
+                    data = await process.stdin.read(4096)
+                    if not data:
+                        break
+                    await loop.run_in_executor(
+                        None,
+                        partial(self._client.shell_write, session_id, _to_bytes(data), owner="mock-ssh"),
+                    )
+            finally:
+                input_done.set()
+
+        async def pump_output() -> None:
+            try:
+                idle_after_input = 0
+                while not stop.is_set():
+                    result = await loop.run_in_executor(
+                        None,
+                        lambda: self._client.shell_read(
+                            session_id,
+                            timeout=0.2,
+                            max_bytes=4096,
+                            owner="mock-ssh",
+                        ),
+                    )
+                    if result.data:
+                        idle_after_input = 0
+                        process.stdout.write(result.data)
+                    elif input_done.is_set():
+                        idle_after_input += 1
+                        if idle_after_input >= 5:
+                            stop.set()
+                            break
+                    if not result.alive:
+                        stop.set()
+                        break
+            except Exception as exc:
+                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"))
+                stop.set()
+
+        tasks = [asyncio.create_task(pump_input()), asyncio.create_task(pump_output())]
+        await stop.wait()
+        for task in tasks:
+            task.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)
+        self._set_exit_status(process, 0)

package/src/server_control_mcp/mock_ssh_session.py

@@ -0,0 +1,68 @@
+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.close_session(session_id, owner="mock-ssh")
+                return
+            except HostBridgeError as exc:
+                if exc.code == "not_found":
+                    return
+                last_error = exc
+            except Exception as exc:
+                last_error = exc
+        assert last_error is not None
+        raise last_error