@neoline/hostbridge 2.0.3 → 2.0.5

package/src/server_control_mcp/mux_daemon.py

@@ -0,0 +1,186 @@
+"""Async Unix daemon primitives for persistent multiplexed HostBridge RPC."""
+
+from __future__ import annotations
+
+import asyncio
+import os
+from collections.abc import Awaitable, Callable
+from contextlib import suppress
+from pathlib import Path
+from typing import Any
+from uuid import uuid4
+
+from .mux_rpc import (
+    DEFAULT_CONNECTION_WINDOW_BYTES,
+    DEFAULT_STREAM_WINDOW_BYTES,
+    MuxRpcRequestClosedHandler,
+    MuxRpcServerConnection,
+)
+from .mux_stream import MuxRpcStream
+
+AsyncRpcHandler = Callable[[str, dict[str, object]], Awaitable[dict[str, object]]]
+MuxDaemonStreamHandler = Callable[
+    [str, dict[str, object], MuxRpcStream, str],
+    Awaitable[dict[str, object]],
+]
+
+
+class MuxDaemonServer:
+    def __init__(
+        self,
+        socket_path: Path | str,
+        handler: AsyncRpcHandler,
+        *,
+        stream_handler: MuxDaemonStreamHandler | None = None,
+        request_closed_handler: MuxRpcRequestClosedHandler | None = None,
+        max_handlers_per_connection: int = 256,
+        drain_timeout: float = 5.0,
+    ) -> None:
+        _positive(max_handlers_per_connection, "max_handlers_per_connection")
+        if drain_timeout <= 0:
+            raise ValueError("drain_timeout must be positive")
+        self.socket_path = Path(socket_path)
+        self._handler = handler
+        self._stream_handler = stream_handler
+        self._request_closed_handler = request_closed_handler
+        self._max_handlers_per_connection = max_handlers_per_connection
+        self._drain_timeout = float(drain_timeout)
+        self._server: asyncio.AbstractServer | None = None
+        self._tasks: set[asyncio.Task[None]] = set()
+        self._writers: set[asyncio.StreamWriter] = set()
+        self._connections: set[MuxRpcServerConnection] = set()
+        self._closing = False
+        self._close_task: asyncio.Task[None] | None = None
+        self.accepted_connections = 0
+
+    @property
+    def active_connections(self) -> int:
+        return len(self._tasks)
+
+    @property
+    def active_handlers(self) -> int:
+        return sum(connection.active_requests for connection in self._connections)
+
+    async def start(self) -> None:
+        if self._server is not None:
+            raise RuntimeError("multiplexed daemon is already started")
+        if self._closing:
+            raise RuntimeError("multiplexed daemon is closed")
+        self.socket_path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
+        self._server = await asyncio.start_unix_server(self._accept, path=self.socket_path)
+        os.chmod(self.socket_path, 0o600)
+
+    async def close(self) -> None:
+        task = self._close_task
+        if task is None:
+            self._closing = True
+            task = asyncio.create_task(
+                self._close_owned_resources(),
+                name="hostbridge-mux-daemon-close",
+            )
+            self._close_task = task
+        await asyncio.shield(task)
+
+    async def _close_owned_resources(self) -> None:
+        listener = self._server
+        if self._server is not None:
+            self._server.close()
+            self._server = None
+        connections = list(self._connections)
+        if connections:
+            drain_tasks = [asyncio.create_task(connection.start_draining()) for connection in connections]
+            await _wait_tasks_bounded(drain_tasks, self._drain_timeout)
+        pending = set(self._tasks)
+        if pending:
+            _, pending = await asyncio.wait(pending, timeout=self._drain_timeout)
+        for writer in list(self._writers):
+            writer.close()
+        if self._writers:
+            waiters = [asyncio.create_task(writer.wait_closed()) for writer in list(self._writers)]
+            await _wait_tasks_bounded(waiters, self._drain_timeout)
+        for task in pending:
+            task.cancel()
+            task.add_done_callback(_consume_task_result)
+        if listener is not None:
+            await listener.wait_closed()
+        self.socket_path.unlink(missing_ok=True)
+
+    def _accept(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
+        if self._closing:
+            writer.close()
+            return
+        self.accepted_connections += 1
+        client_connection_id = uuid4().hex
+        self._writers.add(writer)
+
+        async def run_stream_handler(
+            method: str,
+            params: dict[str, object],
+            stream: MuxRpcStream,
+        ) -> dict[str, object]:
+            return await self._run_stream_handler(
+                method,
+                params,
+                stream,
+                client_connection_id,
+            )
+
+        connection = MuxRpcServerConnection(
+            reader,
+            writer,
+            self._run_handler,
+            max_handlers=self._max_handlers_per_connection,
+            stream_handler=run_stream_handler if self._stream_handler is not None else None,
+            request_closed_handler=self._request_closed_handler,
+            connection_window=DEFAULT_CONNECTION_WINDOW_BYTES,
+            stream_window=DEFAULT_STREAM_WINDOW_BYTES,
+            cleanup_timeout=self._drain_timeout,
+        )
+        self._connections.add(connection)
+        task = asyncio.create_task(
+            connection.serve(),
+            name=f"hostbridge-mux-daemon-{self.accepted_connections}",
+        )
+        self._tasks.add(task)
+
+        def finished(done: asyncio.Task[None]) -> None:
+            _consume_task_result(done)
+            self._tasks.discard(done)
+            self._writers.discard(writer)
+            self._connections.discard(connection)
+
+        task.add_done_callback(finished)
+
+    async def _run_handler(self, method: str, params: dict[str, object]) -> dict[str, object]:
+        return await self._handler(method, params)
+
+    async def _run_stream_handler(
+        self,
+        method: str,
+        params: dict[str, object],
+        stream: MuxRpcStream,
+        client_connection_id: str,
+    ) -> dict[str, object]:
+        handler = self._stream_handler
+        if handler is None:
+            raise RuntimeError("multiplexed daemon stream handler is unavailable")
+        return await handler(method, params, stream, client_connection_id)
+
+
+def _positive(value: int, name: str) -> None:
+    if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
+        raise ValueError(f"{name} must be a positive integer")
+
+
+def _consume_task_result(task: asyncio.Task[object]) -> None:
+    with suppress(BaseException):
+        task.result()
+
+
+async def _wait_tasks_bounded(tasks: list[asyncio.Task[Any]], timeout: float) -> None:
+    done, stalled = await asyncio.wait(tasks, timeout=timeout)
+    for task in done:
+        _consume_task_result(task)
+    for task in stalled:
+        task.cancel()
+        task.add_done_callback(_consume_task_result)

package/src/server_control_mcp/mux_protocol.py

@@ -0,0 +1,279 @@
+"""Binary framing primitives for the HostBridge multiplexed protocol."""
+
+from __future__ import annotations
+
+import json
+import struct
+from collections.abc import Mapping
+from dataclasses import dataclass
+from enum import IntEnum
+
+MUX_MAGIC = b"HBM2"
+MUX_PROTOCOL_VERSION = 2
+MAX_FRAME_PAYLOAD_BYTES = 256 * 1024
+MAX_CONTROL_PAYLOAD_BYTES = 64 * 1024
+MAX_SEQUENCE = 0xFFFFFFFF
+MAX_STREAM_ID = 0xFFFFFFFFFFFFFFFF
+
+_FRAME_HEADER = struct.Struct("!4sBBHQII")
+FRAME_HEADER_SIZE = _FRAME_HEADER.size
+_WINDOW_UPDATE = struct.Struct("!Q")
+
+
+class MuxProtocolError(ValueError):
+    """Raised when a multiplexed frame violates the wire contract."""
+
+
+class FrameType(IntEnum):
+    HELLO = 1
+    HELLO_OK = 2
+    GOAWAY = 3
+    PING = 4
+    PONG = 5
+    OPEN = 6
+    OPEN_OK = 7
+    DATA = 8
+    WINDOW_UPDATE = 9
+    EOF = 10
+    CLOSE = 11
+    CLOSE_ACK = 12
+    RESET = 13
+
+
+class StreamParity(IntEnum):
+    EVEN = 0
+    ODD = 1
+
+
+_CONNECTION_FRAME_TYPES = {
+    FrameType.HELLO,
+    FrameType.HELLO_OK,
+    FrameType.GOAWAY,
+    FrameType.PING,
+    FrameType.PONG,
+}
+_STREAM_FRAME_TYPES = {
+    FrameType.OPEN,
+    FrameType.OPEN_OK,
+    FrameType.DATA,
+    FrameType.EOF,
+    FrameType.CLOSE,
+    FrameType.CLOSE_ACK,
+    FrameType.RESET,
+}
+_JSON_CONTROL_FRAME_TYPES = {
+    FrameType.HELLO,
+    FrameType.HELLO_OK,
+    FrameType.GOAWAY,
+    FrameType.OPEN,
+    FrameType.OPEN_OK,
+    FrameType.CLOSE,
+    FrameType.RESET,
+}
+
+
+@dataclass(frozen=True, slots=True)
+class MuxFrame:
+    frame_type: FrameType
+    stream_id: int
+    sequence: int
+    flags: int = 0
+    payload: bytes = b""
+
+
+def encode_frame(frame: MuxFrame) -> bytes:
+    _validate_frame(frame)
+    return (
+        _FRAME_HEADER.pack(
+            MUX_MAGIC,
+            MUX_PROTOCOL_VERSION,
+            int(frame.frame_type),
+            frame.flags,
+            frame.stream_id,
+            frame.sequence,
+            len(frame.payload),
+        )
+        + frame.payload
+    )
+
+
+class MuxCodec:
+    """Incrementally decode complete frames from an arbitrary byte stream."""
+
+    def __init__(self) -> None:
+        self._buffer = bytearray()
+
+    @property
+    def buffered_bytes(self) -> int:
+        return len(self._buffer)
+
+    def feed(self, data: bytes | bytearray) -> list[MuxFrame]:
+        if not isinstance(data, bytes | bytearray):
+            raise TypeError("mux input must be bytes-like")
+        self._buffer.extend(data)
+        frames: list[MuxFrame] = []
+        offset = 0
+        while len(self._buffer) - offset >= FRAME_HEADER_SIZE:
+            frame_type, flags, stream_id, sequence, payload_length = _decode_header(self._buffer, offset)
+            frame_length = FRAME_HEADER_SIZE + payload_length
+            if len(self._buffer) - offset < frame_length:
+                break
+            payload_start = offset + FRAME_HEADER_SIZE
+            payload = bytes(self._buffer[payload_start : offset + frame_length])
+            frame = MuxFrame(frame_type, stream_id, sequence, flags=flags, payload=payload)
+            _validate_frame(frame)
+            frames.append(frame)
+            offset += frame_length
+        if offset:
+            del self._buffer[:offset]
+        return frames
+
+
+def encode_control_payload(value: Mapping[str, object]) -> bytes:
+    if not isinstance(value, Mapping):
+        raise MuxProtocolError("control payload must be a JSON object")
+    try:
+        payload = json.dumps(dict(value), separators=(",", ":"), sort_keys=True).encode("utf-8")
+    except (TypeError, ValueError) as exc:
+        raise MuxProtocolError("control payload must be a JSON object") from exc
+    if len(payload) > MAX_CONTROL_PAYLOAD_BYTES:
+        raise MuxProtocolError(f"control payload exceeds {MAX_CONTROL_PAYLOAD_BYTES} bytes")
+    return payload
+
+
+def decode_control_payload(payload: bytes) -> dict[str, object]:
+    if not isinstance(payload, bytes):
+        raise TypeError("control payload must be bytes")
+    if len(payload) > MAX_CONTROL_PAYLOAD_BYTES:
+        raise MuxProtocolError(f"control payload exceeds {MAX_CONTROL_PAYLOAD_BYTES} bytes")
+    try:
+        value = json.loads(payload.decode("utf-8"))
+    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
+        raise MuxProtocolError("control payload must contain valid JSON") from exc
+    if not isinstance(value, dict):
+        raise MuxProtocolError("control payload must be a JSON object")
+    return value
+
+
+def encode_window_update(credit: int) -> bytes:
+    _validate_credit(credit)
+    return _WINDOW_UPDATE.pack(credit)
+
+
+def decode_window_update(payload: bytes) -> int:
+    if not isinstance(payload, bytes):
+        raise TypeError("window update payload must be bytes")
+    if len(payload) != _WINDOW_UPDATE.size:
+        raise MuxProtocolError("window update payload must contain an 8-byte credit")
+    (credit,) = _WINDOW_UPDATE.unpack(payload)
+    _validate_credit(credit)
+    return credit
+
+
+class StreamIdAllocator:
+    def __init__(self, parity: StreamParity, *, first: int | None = None) -> None:
+        if not isinstance(parity, StreamParity):
+            raise TypeError("stream parity must be a StreamParity")
+        candidate = (1 if parity is StreamParity.ODD else 2) if first is None else first
+        _validate_stream_id(candidate, allow_zero=False)
+        if candidate % 2 != int(parity):
+            raise ValueError("first stream ID does not match allocator parity")
+        self._next = candidate
+
+    def allocate(self) -> int:
+        if self._next > MAX_STREAM_ID:
+            raise MuxProtocolError("stream ID space is exhausted")
+        stream_id = self._next
+        self._next += 2
+        return stream_id
+
+
+class SequenceTracker:
+    """Track independent inbound and outbound sequence spaces."""
+
+    def __init__(self, *, next_inbound: int = 0, next_outbound: int = 0) -> None:
+        _validate_sequence(next_inbound)
+        _validate_sequence(next_outbound)
+        self._next_inbound = next_inbound
+        self._next_outbound = next_outbound
+
+    def accept_inbound(self, sequence: int) -> None:
+        if self._next_inbound > MAX_SEQUENCE:
+            raise MuxProtocolError("inbound sequence space is exhausted")
+        _validate_sequence(sequence)
+        if sequence != self._next_inbound:
+            raise MuxProtocolError(f"expected sequence {self._next_inbound}, received {sequence}")
+        self._next_inbound += 1
+
+    def next_outbound(self) -> int:
+        if self._next_outbound > MAX_SEQUENCE:
+            raise MuxProtocolError("outbound sequence space is exhausted")
+        sequence = self._next_outbound
+        self._next_outbound += 1
+        return sequence
+
+
+def _decode_header(data: bytes | bytearray, offset: int = 0) -> tuple[FrameType, int, int, int, int]:
+    magic, version, raw_type, flags, stream_id, sequence, payload_length = _FRAME_HEADER.unpack_from(data, offset)
+    if magic != MUX_MAGIC:
+        raise MuxProtocolError("invalid mux frame magic")
+    if version != MUX_PROTOCOL_VERSION:
+        raise MuxProtocolError(f"unsupported mux protocol version: {version}")
+    try:
+        frame_type = FrameType(raw_type)
+    except ValueError as exc:
+        raise MuxProtocolError(f"unsupported mux frame type: {raw_type}") from exc
+    if flags != 0:
+        raise MuxProtocolError("unsupported mux frame flags")
+    if payload_length > MAX_FRAME_PAYLOAD_BYTES:
+        raise MuxProtocolError(f"frame payload length exceeds {MAX_FRAME_PAYLOAD_BYTES} bytes")
+    return frame_type, flags, stream_id, sequence, payload_length
+
+
+def _validate_frame(frame: MuxFrame) -> None:
+    if not isinstance(frame, MuxFrame):
+        raise TypeError("frame must be a MuxFrame")
+    if not isinstance(frame.frame_type, FrameType):
+        raise MuxProtocolError("invalid mux frame type")
+    if not isinstance(frame.flags, int) or isinstance(frame.flags, bool) or frame.flags != 0:
+        raise MuxProtocolError("unsupported mux frame flags")
+    _validate_stream_id(frame.stream_id, allow_zero=True)
+    _validate_sequence(frame.sequence)
+    if not isinstance(frame.payload, bytes):
+        raise MuxProtocolError("frame payload must be bytes")
+    if len(frame.payload) > MAX_FRAME_PAYLOAD_BYTES:
+        raise MuxProtocolError(f"frame payload exceeds {MAX_FRAME_PAYLOAD_BYTES} bytes")
+    if frame.frame_type in _CONNECTION_FRAME_TYPES and frame.stream_id != 0:
+        raise MuxProtocolError(f"{frame.frame_type.name} requires connection stream ID 0")
+    if frame.frame_type in _STREAM_FRAME_TYPES and frame.stream_id == 0:
+        raise MuxProtocolError(f"{frame.frame_type.name} requires a nonzero stream ID")
+    if frame.frame_type is FrameType.DATA and not frame.payload:
+        raise MuxProtocolError("DATA requires a non-empty payload")
+    if frame.frame_type is FrameType.EOF and frame.payload:
+        raise MuxProtocolError(f"{frame.frame_type.name} must not contain a payload")
+    if frame.frame_type is FrameType.CLOSE_ACK and frame.payload:
+        decode_control_payload(frame.payload)
+    if frame.frame_type is FrameType.WINDOW_UPDATE:
+        decode_window_update(frame.payload)
+    if frame.frame_type in _JSON_CONTROL_FRAME_TYPES:
+        decode_control_payload(frame.payload)
+    if frame.frame_type in {FrameType.PING, FrameType.PONG} and len(frame.payload) > 8:
+        raise MuxProtocolError(f"{frame.frame_type.name} payload exceeds 8 bytes")
+
+
+def _validate_stream_id(stream_id: int, *, allow_zero: bool) -> None:
+    if not isinstance(stream_id, int) or isinstance(stream_id, bool):
+        raise MuxProtocolError("stream ID must be an integer")
+    minimum = 0 if allow_zero else 1
+    if not minimum <= stream_id <= MAX_STREAM_ID:
+        raise MuxProtocolError("stream ID is outside the unsigned 64-bit range")
+
+
+def _validate_sequence(sequence: int) -> None:
+    if not isinstance(sequence, int) or isinstance(sequence, bool) or not 0 <= sequence <= MAX_SEQUENCE:
+        raise MuxProtocolError("sequence is outside the unsigned 32-bit range")
+
+
+def _validate_credit(credit: int) -> None:
+    if not isinstance(credit, int) or isinstance(credit, bool) or not 1 <= credit <= MAX_STREAM_ID:
+        raise MuxProtocolError("credit must be an unsigned nonzero 64-bit integer")

package/src/server_control_mcp/mux_records.py

@@ -0,0 +1,71 @@
+"""Binary application records carried inside multiplexed DATA frames."""
+
+from __future__ import annotations
+
+import struct
+from dataclasses import dataclass
+from enum import IntEnum
+
+MAX_RECORD_PAYLOAD_BYTES = 256 * 1024
+_HEADER = struct.Struct("!BI")
+
+
+class MuxRecordError(ValueError):
+    pass
+
+
+class MuxRecordChannel(IntEnum):
+    STDOUT = 1
+    STDERR = 2
+
+
+@dataclass(frozen=True, slots=True)
+class MuxRecord:
+    channel: MuxRecordChannel
+    payload: bytes
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.channel, MuxRecordChannel):
+            raise TypeError("record channel must be a MuxRecordChannel")
+        if not isinstance(self.payload, bytes):
+            raise TypeError("record payload must be bytes")
+
+
+def encode_record(record: MuxRecord) -> bytes:
+    if not isinstance(record, MuxRecord):
+        raise TypeError("record must be a MuxRecord")
+    if len(record.payload) > MAX_RECORD_PAYLOAD_BYTES:
+        raise MuxRecordError(f"record payload exceeds maximum of {MAX_RECORD_PAYLOAD_BYTES} bytes")
+    return _HEADER.pack(int(record.channel), len(record.payload)) + record.payload
+
+
+class MuxRecordDecoder:
+    def __init__(self) -> None:
+        self._buffer = bytearray()
+
+    @property
+    def buffered_bytes(self) -> int:
+        return len(self._buffer)
+
+    def feed(self, data: bytes) -> list[MuxRecord]:
+        if not isinstance(data, bytes):
+            raise TypeError("record data must be bytes")
+        if data:
+            self._buffer.extend(data)
+
+        records: list[MuxRecord] = []
+        while len(self._buffer) >= _HEADER.size:
+            channel_value, payload_length = _HEADER.unpack_from(self._buffer)
+            try:
+                channel = MuxRecordChannel(channel_value)
+            except ValueError as exc:
+                raise MuxRecordError(f"unknown record channel {channel_value}") from exc
+            if payload_length > MAX_RECORD_PAYLOAD_BYTES:
+                raise MuxRecordError(f"record payload exceeds maximum of {MAX_RECORD_PAYLOAD_BYTES} bytes")
+            record_length = _HEADER.size + payload_length
+            if len(self._buffer) < record_length:
+                break
+            payload = bytes(self._buffer[_HEADER.size : record_length])
+            del self._buffer[:record_length]
+            records.append(MuxRecord(channel, payload))
+        return records