@neoline/hostbridge 2.0.3 → 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/policy.py

@@ -18,12 +18,16 @@ from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Protocol
 
+from .secure_files import open_private_append, read_private_text
+
 DEFAULT_POLICY_FILE = Path("~/.hostbridge/policy.json").expanduser()
 LEGACY_POLICY_FILE = Path("~/.server_control_mcp/policy.json").expanduser()
 DEFAULT_AUDIT_LOG = Path("~/.hostbridge/audit.jsonl").expanduser()
 LEGACY_AUDIT_LOG = Path("~/.server_control_mcp/audit.jsonl").expanduser()
 DEFAULT_IDLE_TIMEOUT_SECONDS = 1800
 DEFAULT_REAPER_INTERVAL_SECONDS = 60
+DEFAULT_AUDIT_MAX_BYTES = 10 * 1024 * 1024
+DEFAULT_AUDIT_BACKUP_COUNT = 3
 
 DEFAULT_DENYLIST_PATTERNS = (
     r"\brm\s+-rf\s+/(?:\s|$)",
@@ -46,16 +50,37 @@ class _NullAudit:
 
 
 class _JsonlAudit:
-    def __init__(self, path: Path) -> None:
+    def __init__(self, path: Path, *, max_bytes: int, backup_count: int) -> None:
         self._path = path
+        self._max_bytes = max_bytes
+        self._backup_count = backup_count
         self._lock = threading.Lock()
 
     def write(self, event: dict[str, Any]) -> None:
         line = json.dumps(event, ensure_ascii=False, sort_keys=True) + "\n"
+        encoded_size = len(line.encode("utf-8"))
         with self._lock:
-            self._path.parent.mkdir(parents=True, exist_ok=True)
-            with self._path.open("a", encoding="utf-8") as handle:
+            current_size = self._path.stat().st_size if self._path.exists() else 0
+            if current_size and current_size + encoded_size > self._max_bytes:
+                self._rotate()
+            with open_private_append(self._path) as handle:
                 handle.write(line)
+                handle.flush()
+
+    def _rotate(self) -> None:
+        if self._backup_count == 0:
+            self._path.unlink(missing_ok=True)
+            return
+        oldest = self._backup_path(self._backup_count)
+        oldest.unlink(missing_ok=True)
+        for index in range(self._backup_count - 1, 0, -1):
+            source = self._backup_path(index)
+            if source.exists():
+                source.replace(self._backup_path(index + 1))
+        self._path.replace(self._backup_path(1))
+
+    def _backup_path(self, index: int) -> Path:
+        return self._path.with_name(f"{self._path.name}.{index}")
 
 
 @dataclass(frozen=True, slots=True)
@@ -70,6 +95,8 @@ class PolicyConfig:
     idle_timeout_seconds: int = DEFAULT_IDLE_TIMEOUT_SECONDS
     command_denylist: tuple[re.Pattern[str], ...] = field(default_factory=tuple)
     audit_log_path: Path | None = None
+    audit_max_bytes: int = DEFAULT_AUDIT_MAX_BYTES
+    audit_backup_count: int = DEFAULT_AUDIT_BACKUP_COUNT
     enabled: bool = True
 
     @classmethod
@@ -90,7 +117,9 @@ class PolicyConfig:
         idle_timeout = int(
             env.get(
                 "HOSTBRIDGE_IDLE_TIMEOUT",
-                env.get("SERVER_CONTROL_MCP_IDLE_TIMEOUT", raw.get("idle_timeout_seconds", DEFAULT_IDLE_TIMEOUT_SECONDS)),
+                env.get(
+                    "SERVER_CONTROL_MCP_IDLE_TIMEOUT", raw.get("idle_timeout_seconds", DEFAULT_IDLE_TIMEOUT_SECONDS)
+                ),
             )
         )
         idle_timeout = max(0, idle_timeout)
@@ -104,17 +133,43 @@ class PolicyConfig:
         if env_extra:
             patterns.extend(p for p in env_extra.split(os.pathsep) if p)
 
-        default_audit = DEFAULT_AUDIT_LOG if DEFAULT_AUDIT_LOG.exists() or not LEGACY_AUDIT_LOG.exists() else LEGACY_AUDIT_LOG
+        default_audit = (
+            DEFAULT_AUDIT_LOG if DEFAULT_AUDIT_LOG.exists() or not LEGACY_AUDIT_LOG.exists() else LEGACY_AUDIT_LOG
+        )
         audit_path_str = env.get(
             "HOSTBRIDGE_AUDIT_LOG",
             env.get("SERVER_CONTROL_MCP_AUDIT_LOG", str(raw.get("audit_log_path") or default_audit)),
         )
         audit_path = Path(audit_path_str).expanduser() if audit_path_str else None
+        audit_max_bytes = int(
+            env.get(
+                "HOSTBRIDGE_AUDIT_MAX_BYTES",
+                env.get(
+                    "SERVER_CONTROL_MCP_AUDIT_MAX_BYTES",
+                    raw.get("audit_max_bytes", DEFAULT_AUDIT_MAX_BYTES),
+                ),
+            )
+        )
+        audit_backup_count = int(
+            env.get(
+                "HOSTBRIDGE_AUDIT_BACKUP_COUNT",
+                env.get(
+                    "SERVER_CONTROL_MCP_AUDIT_BACKUP_COUNT",
+                    raw.get("audit_backup_count", DEFAULT_AUDIT_BACKUP_COUNT),
+                ),
+            )
+        )
+        if audit_max_bytes <= 0:
+            raise ValueError("audit_max_bytes must be positive")
+        if audit_backup_count < 0:
+            raise ValueError("audit_backup_count must be non-negative")
 
         return cls(
             idle_timeout_seconds=idle_timeout,
             command_denylist=tuple(re.compile(pattern) for pattern in patterns),
             audit_log_path=audit_path,
+            audit_max_bytes=audit_max_bytes,
+            audit_backup_count=audit_backup_count,
             enabled=True,
         )
 
@@ -130,7 +185,9 @@ class PolicyConfig:
 
 
 def load_policy(config_path: Path | str | None = None) -> PolicyConfig:
-    if (os.environ.get("HOSTBRIDGE_DISABLE_POLICY", "") or os.environ.get("SERVER_CONTROL_MCP_DISABLE_POLICY", "")).lower() in (
+    if (
+        os.environ.get("HOSTBRIDGE_DISABLE_POLICY", "") or os.environ.get("SERVER_CONTROL_MCP_DISABLE_POLICY", "")
+    ).lower() in (
         "1",
         "true",
         "yes",
@@ -141,24 +198,39 @@ def load_policy(config_path: Path | str | None = None) -> PolicyConfig:
         config_path
         or os.environ.get("HOSTBRIDGE_POLICY_FILE")
         or os.environ.get("SERVER_CONTROL_MCP_POLICY_FILE")
-        or (DEFAULT_POLICY_FILE if DEFAULT_POLICY_FILE.exists() or not LEGACY_POLICY_FILE.exists() else LEGACY_POLICY_FILE)
+        or (
+            DEFAULT_POLICY_FILE
+            if DEFAULT_POLICY_FILE.exists() or not LEGACY_POLICY_FILE.exists()
+            else LEGACY_POLICY_FILE
+        )
     ).expanduser()
 
     raw: dict[str, Any] = {}
-    if path.exists():
+    try:
+        policy_text = read_private_text(path)
+    except FileNotFoundError:
+        policy_text = None
+    except OSError as exc:
+        raise ValueError(f"failed to read HostBridge policy file {path}") from exc
+    if policy_text is not None:
         try:
-            data = json.loads(path.read_text(encoding="utf-8"))
-            if isinstance(data, dict):
-                raw = data
-        except (OSError, json.JSONDecodeError):
-            raw = {}
+            data = json.loads(policy_text)
+        except json.JSONDecodeError as exc:
+            raise ValueError(f"HostBridge policy file {path} is not valid JSON") from exc
+        if not isinstance(data, dict):
+            raise ValueError(f"HostBridge policy file {path} must contain a JSON object")
+        raw = data
     return PolicyConfig.from_dict(raw)
 
 
 def build_audit_sink(policy: PolicyConfig) -> AuditSink:
     if not policy.enabled or policy.audit_log_path is None:
         return _NullAudit()
-    return _JsonlAudit(policy.audit_log_path)
+    return _JsonlAudit(
+        policy.audit_log_path,
+        max_bytes=policy.audit_max_bytes,
+        backup_count=policy.audit_backup_count,
+    )
 
 
 def audit_event(