@neoline/hostbridge 2.0.4 → 2.0.5

package/src/server_control_mcp/mux_connection.py

@@ -0,0 +1,326 @@
+"""Lifecycle and bounded resource primitives for multiplexed connections."""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
+from enum import Enum, StrEnum
+
+from .mux_protocol import FrameType, MuxProtocolError
+
+
+class MuxStateError(RuntimeError):
+    """Raised when an operation is invalid for the current lifecycle state."""
+
+
+class StreamState(StrEnum):
+    OPENING = "OPENING"
+    OPEN = "OPEN"
+    LOCAL_EOF = "LOCAL_EOF"
+    REMOTE_EOF = "REMOTE_EOF"
+    DRAINING = "DRAINING"
+    CLOSING = "CLOSING"
+    CLOSED = "CLOSED"
+    RESET = "RESET"
+
+
+class TerminalKind(StrEnum):
+    CLOSED = "CLOSED"
+    RESET = "RESET"
+
+
+class ConnectionState(StrEnum):
+    DISCONNECTED = "DISCONNECTED"
+    CONNECTING = "CONNECTING"
+    HANDSHAKING = "HANDSHAKING"
+    READY = "READY"
+    DRAINING = "DRAINING"
+    CLOSED = "CLOSED"
+    FAILED = "FAILED"
+
+
+class FrameDisposition(Enum):
+    UNKNOWN_STREAM = "unknown_stream"
+    DISCARD_LATE_DATA = "discard_late_data"
+    IGNORE_TERMINAL = "ignore_terminal"
+    PROTOCOL_ERROR = "protocol_error"
+
+
+class ByteCreditWindow:
+    def __init__(self, *, initial_credit: int, maximum_credit: int) -> None:
+        _positive_or_zero(initial_credit, "initial_credit")
+        _positive(maximum_credit, "maximum_credit")
+        if initial_credit > maximum_credit:
+            raise ValueError("initial credit exceeds maximum credit")
+        self._available = initial_credit
+        self._maximum = maximum_credit
+        self._closed = False
+
+    @property
+    def available(self) -> int:
+        return self._available
+
+    @property
+    def closed(self) -> bool:
+        return self._closed
+
+    def consume(self, amount: int) -> None:
+        _positive(amount, "amount")
+        if self._closed:
+            raise MuxStateError("credit window is closed")
+        if amount > self._available:
+            raise MuxStateError("insufficient byte credit")
+        self._available -= amount
+
+    def grant(self, amount: int) -> None:
+        _positive(amount, "amount")
+        if self._closed:
+            return
+        if self._available + amount > self._maximum:
+            raise MuxProtocolError("credit grant exceeds negotiated maximum")
+        self._available += amount
+
+    def close(self) -> None:
+        self._closed = True
+
+
+class StreamLifecycle:
+    def __init__(self, *, open_initiator: bool) -> None:
+        if not isinstance(open_initiator, bool):
+            raise TypeError("open_initiator must be a boolean")
+        self.open_initiator = open_initiator
+        self.state = StreamState.OPENING
+        self.terminal_kind: TerminalKind | None = None
+        self._opened = False
+        self._local_eof = False
+        self._remote_eof = False
+        self._close_sent = False
+        self._close_received = False
+
+    @property
+    def opened(self) -> bool:
+        return self._opened
+
+    @property
+    def local_eof(self) -> bool:
+        return self._local_eof
+
+    @property
+    def remote_eof(self) -> bool:
+        return self._remote_eof
+
+    @property
+    def close_sent(self) -> bool:
+        return self._close_sent
+
+    @property
+    def close_received(self) -> bool:
+        return self._close_received
+
+    @property
+    def terminal(self) -> bool:
+        return self.state in {StreamState.CLOSED, StreamState.RESET}
+
+    def mark_open(self) -> None:
+        self._require(StreamState.OPENING)
+        self._opened = True
+        self.state = StreamState.OPEN
+
+    def send_eof(self) -> None:
+        if self._local_eof:
+            return
+        self._require_active("send EOF")
+        self._local_eof = True
+        self._update_eof_state()
+
+    def receive_eof(self) -> None:
+        if self._remote_eof:
+            return
+        self._require_active("receive EOF")
+        self._remote_eof = True
+        self._update_eof_state()
+
+    def validate_send_data(self) -> None:
+        if self.state not in {StreamState.OPEN, StreamState.REMOTE_EOF}:
+            raise MuxStateError(f"cannot send DATA while stream is {self.state.value}")
+
+    def validate_receive_data(self) -> None:
+        if self.state not in {StreamState.OPEN, StreamState.LOCAL_EOF}:
+            raise MuxStateError(f"cannot receive DATA while stream is {self.state.value}")
+
+    def send_close(self) -> None:
+        if self._close_sent:
+            return
+        if not self.open_initiator:
+            raise MuxStateError("only the OPEN initiator may send CLOSE")
+        self._require(StreamState.DRAINING)
+        self._close_sent = True
+        self.state = StreamState.CLOSING
+
+    def receive_close(self) -> None:
+        if self._close_received:
+            return
+        if self.open_initiator:
+            raise MuxStateError("OPEN initiator cannot receive CLOSE")
+        self._require(StreamState.DRAINING)
+        self._close_received = True
+        self.state = StreamState.CLOSING
+
+    def send_close_ack(self) -> None:
+        if self.state is StreamState.CLOSED:
+            return
+        if self.open_initiator or not self._close_received:
+            raise MuxStateError("CLOSE_ACK requires a received CLOSE on the acceptor")
+        self._finish_closed()
+
+    def receive_close_ack(self) -> None:
+        if self.state is StreamState.CLOSED:
+            return
+        if not self.open_initiator or not self._close_sent:
+            raise MuxStateError("CLOSE_ACK requires a sent CLOSE on the initiator")
+        self._finish_closed()
+
+    def send_reset(self) -> None:
+        self._finish_reset()
+
+    def receive_reset(self) -> None:
+        self._finish_reset()
+
+    def _update_eof_state(self) -> None:
+        if self._local_eof and self._remote_eof:
+            self.state = StreamState.DRAINING
+        elif self._local_eof:
+            self.state = StreamState.LOCAL_EOF
+        else:
+            self.state = StreamState.REMOTE_EOF
+
+    def _require_active(self, action: str) -> None:
+        if self.state not in {StreamState.OPEN, StreamState.LOCAL_EOF, StreamState.REMOTE_EOF}:
+            raise MuxStateError(f"cannot {action} while stream is {self.state.value}")
+
+    def _require(self, state: StreamState) -> None:
+        if self.state is not state:
+            raise MuxStateError(f"stream state must be {state.value}, found {self.state.value}")
+
+    def _finish_closed(self) -> None:
+        self.state = StreamState.CLOSED
+        self.terminal_kind = TerminalKind.CLOSED
+
+    def _finish_reset(self) -> None:
+        if self.state in {StreamState.CLOSED, StreamState.RESET}:
+            return
+        self.state = StreamState.RESET
+        self.terminal_kind = TerminalKind.RESET
+
+
+class ConnectionLifecycle:
+    def __init__(self) -> None:
+        self.state = ConnectionState.DISCONNECTED
+
+    def start_connecting(self) -> None:
+        self._require_any(ConnectionState.DISCONNECTED, ConnectionState.CLOSED, ConnectionState.FAILED)
+        self.state = ConnectionState.CONNECTING
+
+    @property
+    def terminal(self) -> bool:
+        return self.state in {ConnectionState.CLOSED, ConnectionState.FAILED}
+
+    @property
+    def draining(self) -> bool:
+        return self.state is ConnectionState.DRAINING
+
+    def start_handshake(self) -> None:
+        self._require(ConnectionState.CONNECTING)
+        self.state = ConnectionState.HANDSHAKING
+
+    def mark_ready(self) -> None:
+        self._require(ConnectionState.HANDSHAKING)
+        self.state = ConnectionState.READY
+
+    def start_draining(self) -> None:
+        if self.state is ConnectionState.DRAINING:
+            return
+        self._require(ConnectionState.READY)
+        self.state = ConnectionState.DRAINING
+
+    def mark_closed(self) -> None:
+        if self.state is ConnectionState.CLOSED:
+            return
+        self._require_any(ConnectionState.CONNECTING, ConnectionState.HANDSHAKING, ConnectionState.DRAINING)
+        self.state = ConnectionState.CLOSED
+
+    def mark_failed(self) -> None:
+        if self.state is ConnectionState.FAILED:
+            return
+        if self.state is ConnectionState.CLOSED:
+            raise MuxStateError("cannot fail connection while CLOSED")
+        self.state = ConnectionState.FAILED
+
+    def _require(self, state: ConnectionState) -> None:
+        if self.state is not state:
+            raise MuxStateError(f"connection state must be {state.value}, found {self.state.value}")
+
+    def _require_any(self, *states: ConnectionState) -> None:
+        if self.state not in states:
+            allowed = ", ".join(state.value for state in states)
+            raise MuxStateError(f"connection state must be one of {allowed}, found {self.state.value}")
+
+
+@dataclass(frozen=True, slots=True)
+class _Tombstone:
+    terminal_kind: TerminalKind
+    expires_at: float
+
+
+class StreamTombstones:
+    def __init__(
+        self,
+        *,
+        retention_seconds: float,
+        max_entries: int = 4096,
+        clock: Callable[[], float] = time.monotonic,
+    ) -> None:
+        if retention_seconds <= 0:
+            raise ValueError("retention_seconds must be positive")
+        _positive(max_entries, "max_entries")
+        self._retention_seconds = float(retention_seconds)
+        self._max_entries = max_entries
+        self._clock = clock
+        self._entries: dict[int, _Tombstone] = {}
+
+    def add(self, stream_id: int, terminal_kind: TerminalKind) -> None:
+        if not isinstance(stream_id, int) or isinstance(stream_id, bool) or stream_id <= 0:
+            raise ValueError("stream_id must be a positive integer")
+        if not isinstance(terminal_kind, TerminalKind):
+            raise TypeError("terminal_kind must be a TerminalKind")
+        self.expire()
+        if stream_id not in self._entries and len(self._entries) >= self._max_entries:
+            raise MuxStateError("tombstone capacity is exhausted")
+        self._entries[stream_id] = _Tombstone(terminal_kind, self._clock() + self._retention_seconds)
+
+    def classify(self, stream_id: int, frame_type: FrameType) -> FrameDisposition:
+        self.expire()
+        if stream_id not in self._entries:
+            return FrameDisposition.UNKNOWN_STREAM
+        if frame_type is FrameType.OPEN:
+            return FrameDisposition.PROTOCOL_ERROR
+        if frame_type is FrameType.DATA:
+            return FrameDisposition.DISCARD_LATE_DATA
+        return FrameDisposition.IGNORE_TERMINAL
+
+    def expire(self) -> None:
+        now = self._clock()
+        expired = [stream_id for stream_id, tombstone in self._entries.items() if tombstone.expires_at <= now]
+        for stream_id in expired:
+            del self._entries[stream_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 _positive_or_zero(value: int, name: str) -> None:
+    if not isinstance(value, int) or isinstance(value, bool) or value < 0:
+        raise ValueError(f"{name} must be a non-negative integer")

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)