@neoline/hostbridge 2.0.4 → 2.0.5

package/src/server_control_mcp/mux_service.py

@@ -0,0 +1,506 @@
+"""HostBridge service operations exposed over multiplexed RPC streams."""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+from collections.abc import AsyncGenerator, Awaitable, Callable
+from dataclasses import asdict
+from typing import Any
+
+from . import __version__
+from .async_lifecycle import finish_task_before_cancellation
+from .mux_protocol import MUX_PROTOCOL_VERSION
+from .mux_records import MAX_RECORD_PAYLOAD_BYTES, MuxRecord, MuxRecordChannel, encode_record
+from .mux_rpc import MuxRpcFailure
+from .mux_stream import MuxRpcStream
+from .runtime_services import AsyncHostBridgeServices
+from .services import ServiceError
+from .tunnel_manager import HostTunnelManager, TunnelError
+
+
+class HostBridgeMuxService:
+    def __init__(self, services: AsyncHostBridgeServices, *, tunnel_manager: HostTunnelManager | None = None) -> None:
+        self.services = services
+        self.tunnel_manager = tunnel_manager
+        self._unary_methods = {
+            "system.hello",
+            "system.health",
+            "tunnel.health",
+            "system.shutdown",
+            "host.list",
+            "session.open",
+            "session.list",
+            "session.close",
+            "session.release_mock_ssh",
+            "session.close_all",
+            "task.start",
+            "task.status",
+            "task.cancel",
+            "file.write_text",
+            "file.read_text",
+        }
+        self._stream_handlers: dict[
+            str,
+            Callable[[dict[str, object], MuxRpcStream], Awaitable[dict[str, object]]],
+        ] = {
+            "exec.start": self._exec,
+            "transfer.upload": self._upload,
+            "transfer.download": self._download,
+            "shell.write": self._shell_write,
+            "shell.read": self._shell_read,
+            "shell.attach": self._shell_attach,
+        }
+
+    async def handle_unary_async(
+        self,
+        method: str,
+        params: dict[str, object],
+    ) -> dict[str, object]:
+        try:
+            return await self._dispatch_unary(method, params)
+        except asyncio.CancelledError:
+            raise
+        except MuxRpcFailure:
+            raise
+        except ServiceError as exc:
+            raise _service_failure(exc) from exc
+        except (TypeError, ValueError) as exc:
+            raise MuxRpcFailure("invalid_request", str(exc)) from exc
+
+    async def handle_stream(
+        self,
+        method: str,
+        params: dict[str, object],
+        stream: MuxRpcStream,
+        client_connection_id: str,
+    ) -> dict[str, object]:
+        handler = None if method == "tcp.connect" else self._stream_handlers.get(method)
+        if handler is None:
+            if method == "tcp.connect":
+                try:
+                    return await self._tcp_connect(params, stream, client_connection_id)
+                except asyncio.CancelledError:
+                    raise
+                except MuxRpcFailure:
+                    raise
+                except TunnelError as exc:
+                    raise MuxRpcFailure(exc.code, str(exc), retryable=exc.retryable) from exc
+                except (TypeError, ValueError) as exc:
+                    raise MuxRpcFailure("invalid_request", str(exc)) from exc
+            if method in self._unary_methods:
+                raise MuxRpcFailure("invalid_request", f"method {method} does not accept a binary stream")
+            raise MuxRpcFailure("not_found", f"unknown method {method}")
+        try:
+            return await handler(params, stream)
+        except asyncio.CancelledError:
+            raise
+        except MuxRpcFailure:
+            raise
+        except ServiceError as exc:
+            raise _service_failure(exc) from exc
+        except TunnelError as exc:
+            raise MuxRpcFailure(exc.code, str(exc), retryable=exc.retryable) from exc
+        except (TypeError, ValueError) as exc:
+            raise MuxRpcFailure("invalid_request", str(exc)) from exc
+
+    async def _dispatch_unary(self, method: str, params: dict[str, object]) -> dict[str, object]:
+        if method in self._stream_handlers or method == "tcp.connect":
+            raise MuxRpcFailure("invalid_request", f"method {method} requires a binary stream")
+        if method not in self._unary_methods:
+            raise MuxRpcFailure("not_found", f"unknown method {method}")
+        if method == "system.hello":
+            return {
+                "service": "hostbridge",
+                "protocol": MUX_PROTOCOL_VERSION,
+                "version": __version__,
+                "capabilities": ["sessions", "exec", "binary_streams", "persistent_mux"],
+            }
+        if method == "system.health":
+            return {"status": "ok", "protocol": MUX_PROTOCOL_VERSION}
+        if method == "host.list":
+            return {"hosts": self.services.hosts.describe()}
+        if method == "tunnel.health":
+            return self._tunnel_health(params)
+        if method == "system.shutdown":
+            return await self.services.close_all(force=True)
+        if method == "session.open":
+            return await self._session_open(params)
+        if method == "session.list":
+            return {"sessions": await self.services.list_sessions()}
+        if method == "session.close":
+            return await self._session_close(params)
+        if method == "session.release_mock_ssh":
+            return await self._session_release_mock_ssh(params)
+        if method == "session.close_all":
+            return await self.services.close_all(owner=_optional_owner(params))
+        if method == "task.start":
+            task = await self.services.start_task(
+                _required_string(params, "session_id"),
+                _required_string(params, "command"),
+                owner=_optional_owner(params),
+            )
+            return asdict(task)
+        if method == "task.status":
+            return await self.services.task_status(
+                _required_string(params, "session_id"),
+                _required_string(params, "task_id"),
+                owner=_optional_owner(params),
+                tail_lines=_integer(params.get("tail_lines", 80), "tail_lines"),
+            )
+        if method == "task.cancel":
+            return await self.services.cancel_task(
+                _required_string(params, "session_id"),
+                _required_string(params, "task_id"),
+                owner=_optional_owner(params),
+            )
+        if method == "file.write_text":
+            return await self._file_write_text(params)
+        return await self._file_read_text(params)
+
+    def _tunnel_health(self, params: dict[str, object]) -> dict[str, object]:
+        if params:
+            raise ValueError("tunnel.health does not accept parameters")
+        manager = self.tunnel_manager
+        if manager is None:
+            raise MuxRpcFailure("not_found", "TCP tunnel service is unavailable")
+        hosts = []
+        for description in self.services.hosts.describe():
+            host_id = description.get("id")
+            if not isinstance(host_id, str):
+                raise MuxRpcFailure("internal", "host registry returned an invalid host id")
+            hosts.append(manager.health(host_id))
+        return {"hosts": hosts}
+
+    async def _session_open(self, params: dict[str, object]) -> dict[str, object]:
+        connect_timeout = params.get("connect_timeout")
+        if connect_timeout is not None and (
+            not isinstance(connect_timeout, int | float)
+            or isinstance(connect_timeout, bool)
+            or not 0 < connect_timeout <= 180
+        ):
+            raise ValueError("connect_timeout must be between 0 and 180 seconds")
+        session = await self.services.open_session(
+            _required_string(params, "host_id"),
+            owner=_optional_owner(params),
+            connect_timeout=float(connect_timeout) if connect_timeout is not None else None,
+        )
+        return session.describe()
+
+    async def _session_close(self, params: dict[str, object]) -> dict[str, object]:
+        if "force" in params:
+            raise ValueError("force close is not available over the public protocol")
+        session_id = _required_string(params, "session_id")
+        await self.services.close_session(session_id, owner=_optional_owner(params))
+        return {"session_id": session_id, "closed": True}
+
+    async def _session_release_mock_ssh(self, params: dict[str, object]) -> dict[str, object]:
+        unknown = sorted(set(params) - {"session_id"})
+        if unknown:
+            raise ValueError(f"unknown session.release_mock_ssh fields: {', '.join(unknown)}")
+        session_id = _required_string(params, "session_id")
+        await self.services.close_mock_ssh_session(session_id)
+        return {"session_id": session_id, "closed": True}
+
+    async def _file_write_text(self, params: dict[str, object]) -> dict[str, object]:
+        content = params.get("content")
+        if not isinstance(content, str):
+            raise ValueError("content must be a string")
+        result = await self.services.write_text(
+            _required_string(params, "session_id"),
+            _required_string(params, "remote_path"),
+            content,
+            mode=_integer(params.get("mode", 0o644), "mode"),
+            max_bytes=_positive_integer(params.get("max_bytes", 4 * 1024 * 1024), "max_bytes"),
+            owner=_optional_owner(params),
+        )
+        return _transfer_result(result.bytes_transferred, result.sha256)
+
+    async def _file_read_text(self, params: dict[str, object]) -> dict[str, object]:
+        content = await self.services.read_text(
+            _required_string(params, "session_id"),
+            _required_string(params, "remote_path"),
+            max_bytes=_positive_integer(params.get("max_bytes", 4 * 1024 * 1024), "max_bytes"),
+            owner=_optional_owner(params),
+        )
+        return {"content": content, "bytes_transferred": len(content.encode("utf-8"))}
+
+    async def _exec(
+        self,
+        params: dict[str, object],
+        stream: MuxRpcStream,
+    ) -> dict[str, object]:
+        timeout = _positive_number(params.get("timeout", 60), "timeout")
+        output_limit = _positive_integer(params.get("output_limit", 1_000_000), "output_limit")
+        session_id = _required_string(params, "session_id")
+        command = _required_string(params, "command")
+        owner = _optional_owner(params)
+        await stream.accept()
+        result = await self.services.exec(
+            session_id,
+            command,
+            stdin=_receive_chunks(stream),
+            timeout=timeout,
+            output_limit=output_limit,
+            owner=owner,
+        )
+        await _send_records(stream, MuxRecordChannel.STDOUT, result.stdout)
+        await _send_records(stream, MuxRecordChannel.STDERR, result.stderr)
+        return {
+            "session_id": result.session_id,
+            "exit_code": result.exit_code,
+            "timed_out": result.timed_out,
+        }
+
+    async def _upload(
+        self,
+        params: dict[str, object],
+        stream: MuxRpcStream,
+    ) -> dict[str, object]:
+        expected_size_value = params.get("expected_size")
+        expected_size = (
+            None if expected_size_value is None else _non_negative_integer(expected_size_value, "expected_size")
+        )
+        expected_sha256 = params.get("expected_sha256")
+        if expected_sha256 is not None and not isinstance(expected_sha256, str):
+            raise ValueError("expected_sha256 must be a string or null")
+        session_id = _required_string(params, "session_id")
+        remote_path = _required_string(params, "remote_path")
+        mode = _integer(params.get("mode", 0o644), "mode")
+        owner = _optional_owner(params)
+        await stream.accept()
+        result = await self.services.upload(
+            session_id,
+            _receive_chunks(stream),
+            remote_path,
+            mode=mode,
+            expected_size=expected_size,
+            expected_sha256=expected_sha256,
+            owner=owner,
+        )
+        return _transfer_result(result.bytes_transferred, result.sha256)
+
+    async def _download(
+        self,
+        params: dict[str, object],
+        stream: MuxRpcStream,
+    ) -> dict[str, object]:
+        chunk_size = _positive_integer(params.get("chunk_size", 256 * 1024), "chunk_size")
+        if not 4096 <= chunk_size <= 1024 * 1024:
+            raise ValueError("chunk_size must be between 4096 and 1048576 bytes")
+        session_id = _required_string(params, "session_id")
+        remote_path = _required_string(params, "remote_path")
+        owner = _optional_owner(params)
+        await stream.accept()
+        digest = hashlib.sha256()
+        transferred = 0
+        async for chunk in self.services.download(
+            session_id,
+            remote_path,
+            chunk_size=chunk_size,
+            owner=owner,
+        ):
+            await stream.send(chunk)
+            digest.update(chunk)
+            transferred += len(chunk)
+        return _transfer_result(transferred, digest.hexdigest())
+
+    async def _shell_write(
+        self,
+        params: dict[str, object],
+        stream: MuxRpcStream,
+    ) -> dict[str, object]:
+        max_bytes = _positive_integer(params.get("max_bytes", 1024 * 1024), "max_bytes")
+        session_id = _required_string(params, "session_id")
+        owner = _optional_owner(params)
+        await stream.accept()
+        data = bytearray()
+        async for chunk in _receive_chunks(stream):
+            data.extend(chunk)
+            if len(data) > max_bytes:
+                raise ServiceError("resource_limit", f"shell write exceeds {max_bytes} bytes")
+        written = await self.services.shell_write(session_id, bytes(data), owner=owner)
+        return {"written": written}
+
+    async def _shell_read(
+        self,
+        params: dict[str, object],
+        stream: MuxRpcStream,
+    ) -> dict[str, object]:
+        timeout = _non_negative_number(params.get("timeout", 0.2), "timeout")
+        max_bytes = _positive_integer(params.get("max_bytes", 4096), "max_bytes")
+        session_id = _required_string(params, "session_id")
+        owner = _optional_owner(params)
+        await stream.accept()
+        result = await self.services.shell_read(
+            session_id,
+            timeout=timeout,
+            max_bytes=max_bytes,
+            owner=owner,
+        )
+        await stream.send(result.data)
+        return {"alive": result.alive, "bytes": len(result.data)}
+
+    async def _shell_attach(
+        self,
+        params: dict[str, object],
+        stream: MuxRpcStream,
+    ) -> dict[str, object]:
+        max_bytes = _positive_integer(params.get("max_bytes", 64 * 1024), "max_bytes")
+        session_id = _required_string(params, "session_id")
+        owner = _optional_owner(params)
+        await stream.accept()
+        transferred = 0
+        attachment = await self.services.attach_shell(
+            session_id,
+            _receive_chunks(stream),
+            max_bytes=max_bytes,
+            owner=owner,
+        )
+        async for chunk in attachment:
+            await stream.send(chunk)
+            transferred += len(chunk)
+        terminal = await attachment.finish()
+        return {**terminal.to_mapping(), "bytes": transferred}
+
+    async def _tcp_connect(
+        self,
+        params: dict[str, object],
+        stream: MuxRpcStream,
+        client_connection_id: str,
+    ) -> dict[str, object]:
+        manager = self.tunnel_manager
+        if manager is None:
+            raise MuxRpcFailure("not_found", "TCP tunnel service is unavailable")
+        allowed = {"host_id", "destination_host", "destination_port"}
+        unknown = sorted(set(params) - allowed)
+        if unknown:
+            raise ValueError(f"unknown tcp.connect fields: {', '.join(unknown)}")
+        host_id = _required_string(params, "host_id")
+        destination_host = _required_string(params, "destination_host")
+        destination_port = _positive_integer(params.get("destination_port"), "destination_port")
+        handle = await manager.open_stream(
+            host_id,
+            client_connection_id,
+            stream.stream_id,
+            destination_host,
+            destination_port,
+        )
+        await stream.accept()
+        bytes_sent = 0
+        bytes_received = 0
+        tasks: list[asyncio.Future[Any]] = []
+        clean = False
+
+        async def upstream() -> None:
+            nonlocal bytes_sent
+            while chunk := await stream.receive():
+                await handle.send(chunk)
+                bytes_sent += len(chunk)
+            await handle.send_eof()
+
+        async def downstream() -> None:
+            nonlocal bytes_received
+            while chunk := await handle.receive():
+                await stream.send(chunk)
+                bytes_received += len(chunk)
+
+        try:
+            bridge = asyncio.gather(
+                asyncio.create_task(upstream(), name=f"hostbridge-tcp-upstream-{stream.stream_id}"),
+                asyncio.create_task(downstream(), name=f"hostbridge-tcp-downstream-{stream.stream_id}"),
+            )
+            terminal = asyncio.create_task(
+                handle.wait_closed(),
+                name=f"hostbridge-tcp-terminal-{stream.stream_id}",
+            )
+            tasks = [bridge, terminal]
+            done, _pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+            if terminal in done:
+                await terminal
+                raise TunnelError("connection_lost", "destination stream closed before both directions drained")
+            await bridge
+            clean = True
+            return {"bytes_sent": bytes_sent, "bytes_received": bytes_received}
+        finally:
+
+            async def finish_bridge() -> None:
+                for task in tasks:
+                    if not task.done():
+                        task.cancel()
+                if tasks:
+                    await asyncio.gather(*tasks, return_exceptions=True)
+                if clean:
+                    await handle.close()
+                else:
+                    await handle.reset("daemon client tunnel stream closed")
+
+            cleanup = asyncio.create_task(
+                finish_bridge(),
+                name=f"hostbridge-tcp-cleanup-{stream.stream_id}",
+            )
+            await finish_task_before_cancellation(cleanup)
+
+
+async def _send_records(stream: MuxRpcStream, channel: MuxRecordChannel, data: bytes) -> None:
+    for offset in range(0, len(data), MAX_RECORD_PAYLOAD_BYTES):
+        await stream.send(encode_record(MuxRecord(channel, data[offset : offset + MAX_RECORD_PAYLOAD_BYTES])))
+
+
+async def _receive_chunks(stream: MuxRpcStream) -> AsyncGenerator[bytes, None]:
+    while chunk := await stream.receive():
+        yield chunk
+
+
+def _service_failure(error: ServiceError) -> MuxRpcFailure:
+    return MuxRpcFailure(error.code, str(error), retryable=error.retryable, details=error.details)
+
+
+def _required_string(params: dict[str, object], name: str) -> str:
+    value = params.get(name)
+    if not isinstance(value, str) or not value:
+        raise ValueError(f"{name} must be a non-empty string")
+    return value
+
+
+def _optional_owner(params: dict[str, object]) -> str | None:
+    owner = params.get("owner")
+    if owner is not None and (not isinstance(owner, str) or not owner):
+        raise ValueError("owner must be a non-empty string or null")
+    return owner
+
+
+def _integer(value: object, name: str) -> int:
+    if not isinstance(value, int) or isinstance(value, bool):
+        raise ValueError(f"{name} must be an integer")
+    return value
+
+
+def _positive_integer(value: object, name: str) -> int:
+    parsed = _integer(value, name)
+    if parsed <= 0:
+        raise ValueError(f"{name} must be positive")
+    return parsed
+
+
+def _non_negative_integer(value: object, name: str) -> int:
+    parsed = _integer(value, name)
+    if parsed < 0:
+        raise ValueError(f"{name} must be non-negative")
+    return parsed
+
+
+def _positive_number(value: object, name: str) -> float:
+    if not isinstance(value, int | float) or isinstance(value, bool) or value <= 0:
+        raise ValueError(f"{name} must be positive")
+    return float(value)
+
+
+def _non_negative_number(value: object, name: str) -> float:
+    if not isinstance(value, int | float) or isinstance(value, bool) or value < 0:
+        raise ValueError(f"{name} must be non-negative")
+    return float(value)
+
+
+def _transfer_result(bytes_transferred: int, sha256: str) -> dict[str, object]:
+    return {"bytes_transferred": bytes_transferred, "sha256": sha256.lower()}