@neoline/hostbridge 2.0.4 → 2.0.5

package/src/server_control_mcp/mux_stream.py

@@ -0,0 +1,283 @@
+"""Bounded stream state shared by multiplexed RPC clients and servers."""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Callable
+from contextlib import suppress
+from typing import Protocol
+
+from .mux_connection import ByteCreditWindow, StreamLifecycle, StreamState
+from .mux_protocol import SequenceTracker
+
+
+class MuxStreamOwner(Protocol):
+    async def _stream_accept(self, stream: MuxRpcStream) -> None: ...
+
+    async def _stream_send_data(self, stream: MuxRpcStream, data: bytes) -> None: ...
+
+    async def _stream_send_eof(self, stream: MuxRpcStream) -> None: ...
+
+    async def _stream_release_credit(self, stream: MuxRpcStream, amount: int) -> None: ...
+
+    async def _stream_reset(self, stream: MuxRpcStream, message: str) -> None: ...
+
+
+class AsyncByteCreditWindow:
+    def __init__(self, credit: int, *, on_stall: Callable[[], None] | None = None) -> None:
+        if not isinstance(credit, int) or isinstance(credit, bool) or credit <= 0:
+            raise ValueError("credit must be a positive integer")
+        self._available = credit
+        self._maximum = credit
+        self._on_stall = on_stall
+        self._stall_count = 0
+        self._closed_error: BaseException | None = None
+        self._condition = asyncio.Condition()
+
+    @property
+    def available(self) -> int:
+        return self._available
+
+    @property
+    def maximum(self) -> int:
+        return self._maximum
+
+    @property
+    def stall_count(self) -> int:
+        return self._stall_count
+
+    async def acquire(self, amount: int) -> None:
+        _positive(amount, "amount")
+        if amount > self._maximum:
+            raise ValueError("amount exceeds maximum credit")
+        async with self._condition:
+            if self._closed_error is None and self._available < amount:
+                self._stall_count += 1
+                if self._on_stall is not None:
+                    with suppress(Exception):
+                        self._on_stall()
+            await self._condition.wait_for(lambda: self._closed_error is not None or self._available >= amount)
+            if self._closed_error is not None:
+                raise self._closed_error
+            self._available -= amount
+
+    async def grant(self, amount: int) -> None:
+        _positive(amount, "amount")
+        async with self._condition:
+            if self._closed_error is not None:
+                return
+            if self._available + amount > self._maximum:
+                raise ValueError("credit grant exceeds negotiated maximum")
+            self._available += amount
+            self._condition.notify_all()
+
+    async def close(self, error: BaseException) -> None:
+        async with self._condition:
+            if self._closed_error is None:
+                self._closed_error = error
+            self._condition.notify_all()
+
+
+class _ReceiveBuffer:
+    def __init__(self) -> None:
+        self._buffer = bytearray()
+        self._eof = False
+        self._error: BaseException | None = None
+        self._condition = asyncio.Condition()
+
+    @property
+    def buffered_bytes(self) -> int:
+        return len(self._buffer)
+
+    async def feed_data(self, data: bytes) -> None:
+        async with self._condition:
+            if self._eof or self._error is not None:
+                return
+            self._buffer.extend(data)
+            self._condition.notify_all()
+
+    async def feed_eof(self) -> None:
+        async with self._condition:
+            self._eof = True
+            self._condition.notify_all()
+
+    async def fail(self, error: BaseException) -> None:
+        async with self._condition:
+            self._buffer.clear()
+            self._error = error
+            self._condition.notify_all()
+
+    async def read(self, max_bytes: int) -> bytes | None:
+        _positive(max_bytes, "max_bytes")
+        async with self._condition:
+            await self._condition.wait_for(lambda: bool(self._buffer) or self._eof or self._error is not None)
+            if self._error is not None:
+                raise self._error
+            if self._buffer:
+                size = min(max_bytes, len(self._buffer))
+                data = bytes(self._buffer[:size])
+                del self._buffer[:size]
+                return data
+            return None
+
+
+class MuxRpcStream:
+    def __init__(
+        self,
+        owner: MuxStreamOwner,
+        stream_id: int,
+        sequences: SequenceTracker,
+        *,
+        send_credit: int,
+        receive_credit: int,
+        open_initiator: bool,
+        result_expected: bool = True,
+        accepted: bool = True,
+        on_send_stall: Callable[[], None] | None = None,
+    ) -> None:
+        self._owner = owner
+        self.stream_id = stream_id
+        self._sequences = sequences
+        self._send_window = AsyncByteCreditWindow(send_credit, on_stall=on_send_stall)
+        self._receive_window = ByteCreditWindow(
+            initial_credit=receive_credit,
+            maximum_credit=receive_credit,
+        )
+        self._receive_buffer = _ReceiveBuffer()
+        self._result: asyncio.Future[dict[str, object]] | None = (
+            asyncio.get_running_loop().create_future() if result_expected else None
+        )
+        self._lifecycle = StreamLifecycle(open_initiator=open_initiator)
+        if accepted:
+            self._lifecycle.mark_open()
+
+    @property
+    def buffered_bytes(self) -> int:
+        return self._receive_buffer.buffered_bytes
+
+    @property
+    def closed(self) -> bool:
+        return self._lifecycle.terminal
+
+    @property
+    def state(self) -> StreamState:
+        return self._lifecycle.state
+
+    @property
+    def local_eof(self) -> bool:
+        return self._lifecycle.local_eof
+
+    @property
+    def accepted(self) -> bool:
+        return self._lifecycle.opened
+
+    @property
+    def remote_eof(self) -> bool:
+        return self._lifecycle.remote_eof
+
+    @property
+    def close_sent(self) -> bool:
+        return self._lifecycle.close_sent
+
+    @property
+    def close_received(self) -> bool:
+        return self._lifecycle.close_received
+
+    @property
+    def receive_credit(self) -> int:
+        return self._receive_window.available
+
+    async def accept(self) -> None:
+        if self.accepted:
+            return
+        if self.closed:
+            raise RuntimeError("stream is closed")
+        self._lifecycle.mark_open()
+        await self._owner._stream_accept(self)
+
+    async def send(self, data: bytes) -> None:
+        if not isinstance(data, bytes):
+            raise TypeError("stream data must be bytes")
+        if not data:
+            return
+        if self.closed:
+            raise RuntimeError("stream is closed")
+        if not self.accepted:
+            raise RuntimeError("stream is not accepted")
+        self._lifecycle.validate_send_data()
+        await self._owner._stream_send_data(self, data)
+
+    async def send_eof(self) -> None:
+        if self.closed or self.local_eof:
+            return
+        if not self.accepted:
+            raise RuntimeError("stream is not accepted")
+        self._lifecycle.send_eof()
+        await self._owner._stream_send_eof(self)
+
+    async def receive(self, max_bytes: int = 64 * 1024) -> bytes | None:
+        if not self.accepted:
+            raise RuntimeError("stream is not accepted")
+        data = await self._receive_buffer.read(max_bytes)
+        if data is not None:
+            await self._owner._stream_release_credit(self, len(data))
+        return data
+
+    async def finish(self) -> dict[str, object]:
+        if self._result is None:
+            raise RuntimeError("server-owned streams do not receive a terminal result")
+        return await self._result
+
+    async def reset(self, message: str = "stream reset") -> None:
+        if self.closed:
+            return
+        await self._owner._stream_reset(self, message)
+
+    async def _accept_data(self, data: bytes) -> None:
+        self._lifecycle.validate_receive_data()
+        self._receive_window.consume(len(data))
+        await self._receive_buffer.feed_data(data)
+
+    async def _accept_eof(self) -> None:
+        self._lifecycle.receive_eof()
+        await self._receive_buffer.feed_eof()
+
+    def _send_close(self) -> None:
+        self._lifecycle.send_close()
+
+    def _receive_close(self) -> None:
+        self._lifecycle.receive_close()
+
+    async def _release_receive_credit(self, amount: int) -> None:
+        self._receive_window.grant(amount)
+
+    async def _grant_send_credit(self, amount: int) -> None:
+        await self._send_window.grant(amount)
+
+    async def _fail(self, error: BaseException) -> None:
+        if self.closed:
+            return
+        self._lifecycle.send_reset()
+        self._receive_window.close()
+        await self._send_window.close(error)
+        await self._receive_buffer.fail(error)
+        if self._result is not None and not self._result.done():
+            self._result.set_exception(error)
+            self._result.exception()
+
+    async def _finish(self, result: dict[str, object]) -> None:
+        if self.closed:
+            return
+        if self._lifecycle.open_initiator:
+            self._lifecycle.receive_close_ack()
+        else:
+            self._lifecycle.send_close_ack()
+        self._receive_window.close()
+        await self._send_window.close(RuntimeError("stream is closed"))
+        if self._result is not None and not self._result.done():
+            self._result.set_result(result)
+
+
+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")

package/src/server_control_mcp/mux_sync_client.py

@@ -0,0 +1,224 @@
+"""Thread-safe synchronous facade for the asyncio multiplexed RPC client."""
+
+from __future__ import annotations
+
+import asyncio
+import threading
+from collections.abc import Coroutine, Mapping
+from concurrent.futures import TimeoutError as FutureTimeoutError
+from pathlib import Path
+from types import TracebackType
+from typing import Any, TypeVar
+
+from .mux_rpc import MuxConnectionClosed, MuxRpcClient
+from .mux_stream import MuxRpcStream
+
+T = TypeVar("T")
+
+
+class MuxSyncStream:
+    def __init__(self, owner: MuxSyncClient, stream: MuxRpcStream) -> None:
+        self._owner = owner
+        self._stream = stream
+
+    @property
+    def stream_id(self) -> int:
+        return self._stream.stream_id
+
+    @property
+    def closed(self) -> bool:
+        return self._stream.closed
+
+    def send(self, data: bytes, *, timeout: float | None = None) -> None:
+        self._owner._submit(self._stream.send(data), timeout=timeout)
+
+    def send_eof(self, *, timeout: float | None = None) -> None:
+        self._owner._submit(self._stream.send_eof(), timeout=timeout)
+
+    def receive(self, max_bytes: int = 64 * 1024, *, timeout: float | None = None) -> bytes | None:
+        return self._owner._submit(self._stream.receive(max_bytes), timeout=timeout)
+
+    def receive_unbounded(self, max_bytes: int = 64 * 1024) -> bytes | None:
+        return self._owner._submit_unbounded(self._stream.receive(max_bytes))
+
+    def finish(self, *, timeout: float | None = None) -> dict[str, object]:
+        return self._owner._submit(self._stream.finish(), timeout=timeout)
+
+    def reset(self, message: str = "stream reset", *, timeout: float | None = None) -> None:
+        self._owner._submit(self._stream.reset(message), timeout=timeout)
+
+
+class MuxSyncClient:
+    def __init__(
+        self,
+        socket_path: Path | str,
+        *,
+        default_timeout: float = 30.0,
+        max_pending: int = 256,
+    ) -> None:
+        if default_timeout <= 0:
+            raise ValueError("default_timeout must be positive")
+        if not isinstance(max_pending, int) or isinstance(max_pending, bool) or max_pending <= 0:
+            raise ValueError("max_pending must be a positive integer")
+        self.socket_path = Path(socket_path)
+        self.default_timeout = float(default_timeout)
+        self._max_pending = max_pending
+        self._loop = asyncio.new_event_loop()
+        self._client: MuxRpcClient | None = None
+        self._connect_lock: asyncio.Lock | None = None
+        self._state_lock = threading.Lock()
+        self._started = threading.Event()
+        self._closed = False
+        self._thread = threading.Thread(target=self._run_loop, name="hostbridge-mux-client", daemon=True)
+        self._thread.start()
+        self._started.wait()
+
+    @property
+    def closed(self) -> bool:
+        with self._state_lock:
+            return self._closed
+
+    @property
+    def worker_thread_alive(self) -> bool:
+        return self._thread.is_alive()
+
+    def request(
+        self,
+        method: str,
+        params: Mapping[str, object],
+        *,
+        timeout: float | None = None,
+    ) -> dict[str, object]:
+        return self._submit(self._request(method, params), timeout=timeout)
+
+    def open_stream(
+        self,
+        method: str,
+        params: Mapping[str, object],
+        *,
+        timeout: float | None = None,
+        receive_window: int | None = None,
+    ) -> MuxSyncStream:
+        stream = self._submit(
+            self._open_stream(method, params, receive_window=receive_window),
+            timeout=timeout,
+        )
+        return MuxSyncStream(self, stream)
+
+    def close(self) -> None:
+        with self._state_lock:
+            if self._closed:
+                return
+            self._closed = True
+        if self._thread.is_alive():
+            future = asyncio.run_coroutine_threadsafe(self._close_client(), self._loop)
+            try:
+                future.result(timeout=self.default_timeout)
+            finally:
+                self._loop.call_soon_threadsafe(self._loop.stop)
+                if threading.get_ident() != self._thread.ident:
+                    self._thread.join(timeout=self.default_timeout)
+
+    def __enter__(self) -> MuxSyncClient:
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        self.close()
+
+    async def _request(self, method: str, params: Mapping[str, object]) -> dict[str, object]:
+        client = await self._ensure_client()
+        try:
+            return await client.request(method, params)
+        except MuxConnectionClosed:
+            await self._discard_client(client)
+            raise
+
+    async def _open_stream(
+        self,
+        method: str,
+        params: Mapping[str, object],
+        *,
+        receive_window: int | None,
+    ) -> MuxRpcStream:
+        client = await self._ensure_client()
+        try:
+            return await client.open_stream(method, params, receive_window=receive_window)
+        except MuxConnectionClosed:
+            await self._discard_client(client)
+            raise
+
+    async def _ensure_client(self) -> MuxRpcClient:
+        current = self._client
+        if current is not None and not current.closed:
+            return current
+        if self._connect_lock is None:
+            self._connect_lock = asyncio.Lock()
+        async with self._connect_lock:
+            current = self._client
+            if current is not None and not current.closed:
+                return current
+            replacement = await MuxRpcClient.connect_unix(self.socket_path, max_pending=self._max_pending)
+            self._client = replacement
+        if current is not None:
+            await current.close()
+        return replacement
+
+    async def _discard_client(self, client: MuxRpcClient) -> None:
+        if self._connect_lock is None:
+            self._connect_lock = asyncio.Lock()
+        async with self._connect_lock:
+            if self._client is client:
+                self._client = None
+        await client.close()
+
+    async def _close_client(self) -> None:
+        if self._client is not None:
+            await self._client.close()
+            self._client = None
+
+    def _submit(self, coroutine: Coroutine[Any, Any, T], *, timeout: float | None) -> T:
+        effective_timeout = self.default_timeout if timeout is None else timeout
+        if effective_timeout <= 0:
+            coroutine.close()
+            raise ValueError("timeout must be positive")
+        with self._state_lock:
+            if self._closed:
+                coroutine.close()
+                raise RuntimeError("multiplexed client is closed")
+            future = asyncio.run_coroutine_threadsafe(coroutine, self._loop)
+        try:
+            return future.result(timeout=effective_timeout)
+        except FutureTimeoutError as exc:
+            future.cancel()
+            raise TimeoutError("multiplexed client operation timed out") from exc
+
+    def _submit_unbounded(self, coroutine: Coroutine[Any, Any, T]) -> T:
+        with self._state_lock:
+            if self._closed:
+                coroutine.close()
+                raise RuntimeError("multiplexed client is closed")
+            future = asyncio.run_coroutine_threadsafe(coroutine, self._loop)
+        try:
+            return future.result()
+        except BaseException:
+            future.cancel()
+            raise
+
+    def _run_loop(self) -> None:
+        asyncio.set_event_loop(self._loop)
+        self._started.set()
+        try:
+            self._loop.run_forever()
+        finally:
+            pending = asyncio.all_tasks(self._loop)
+            for task in pending:
+                task.cancel()
+            if pending:
+                self._loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
+            self._loop.run_until_complete(self._loop.shutdown_asyncgens())
+            self._loop.close()

package/src/server_control_mcp/remote_agent_bundle.py

@@ -0,0 +1,56 @@
+"""Build the deterministic, self-contained remote tunnel agent zipapp."""
+
+from __future__ import annotations
+
+import io
+import zipfile
+from pathlib import Path
+
+_PACKAGE_MODULES = (
+    "async_lifecycle.py",
+    "mux_connection.py",
+    "mux_protocol.py",
+    "mux_records.py",
+    "mux_rpc.py",
+    "mux_stream.py",
+    "remote_mux_agent.py",
+)
+_ZIP_TIMESTAMP = (1980, 1, 1, 0, 0, 0)
+_PACKAGE_INIT = b'"""Self-contained HostBridge remote tunnel agent."""\n'
+_ENTRYPOINT = b"""from __future__ import annotations
+
+import asyncio
+import os
+import sys
+
+from server_control_mcp.remote_mux_agent import run_stdio_agent
+
+try:
+    exit_code = asyncio.run(run_stdio_agent())
+except Exception as exc:
+    nonce = sys.argv[1].encode("ascii")
+    error = type(exc).__name__.encode("ascii", "replace")[:64]
+    os.write(1, b"\\0HBFAIL2:" + nonce + b":" + error + b"\\0")
+    raise
+raise SystemExit(exit_code)
+"""
+
+
+def build_remote_agent_bundle() -> bytes:
+    package_dir = Path(__file__).resolve().parent
+    files: list[tuple[str, bytes]] = [
+        ("__main__.py", _ENTRYPOINT),
+        ("server_control_mcp/__init__.py", _PACKAGE_INIT),
+    ]
+    for name in _PACKAGE_MODULES:
+        files.append((f"server_control_mcp/{name}", (package_dir / name).read_bytes()))
+
+    output = io.BytesIO()
+    with zipfile.ZipFile(output, mode="w", compression=zipfile.ZIP_STORED, strict_timestamps=True) as archive:
+        for path, content in files:
+            info = zipfile.ZipInfo(path, date_time=_ZIP_TIMESTAMP)
+            info.compress_type = zipfile.ZIP_STORED
+            info.create_system = 3
+            info.external_attr = 0o100644 << 16
+            archive.writestr(info, content)
+    return output.getvalue()