@neoline/hostbridge 2.0.4 → 2.0.5

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