@neoline/hostbridge 2.0.3 → 2.0.5

package/src/server_control_mcp/client.py

@@ -1,34 +1,24 @@
-"""Typed synchronous client for the HostBridge daemon protocol."""
+"""Typed synchronous client for the HostBridge multiplexed daemon protocol."""
 
 from __future__ import annotations
 
-import base64
-import binascii
 import hashlib
-import json
-import socket
 import threading
+import time
 import uuid
-from collections.abc import Mapping
+from collections.abc import Iterable, Mapping
+from contextlib import suppress
 from dataclasses import dataclass
 from pathlib import Path
+from types import TracebackType
 
-from .protocol import (
-    MAX_FRAME_BYTES,
-    BinaryFrame,
-    ErrorInfo,
-    FrameFlags,
-    FrameSequenceValidator,
-    ProtocolError,
-    Request,
-    Response,
-    encode_frame,
-    read_frame,
-)
+from . import __version__
+from .mux_protocol import MUX_PROTOCOL_VERSION
+from .mux_records import MuxRecordChannel, MuxRecordDecoder, MuxRecordError
+from .mux_rpc import MuxCapacityError, MuxConnectionClosed, MuxRemoteError
+from .mux_sync_client import MuxSyncClient, MuxSyncStream
 from .runtime import runtime_paths
-from .transports.base import ShellReadResult, TransferResult
-
-MAX_CONTROL_LINE_BYTES = 4 * 1024 * 1024
+from .transports.base import ShellReadResult, ShellTerminalResult, TransferResult
 
 
 def _required_text(value: object, field: str) -> str:
@@ -37,23 +27,6 @@ def _required_text(value: object, field: str) -> str:
     return value
 
 
-def _encode_inline_bytes(value: bytes | None) -> dict[str, str] | None:
-    if value is None:
-        return None
-    if not isinstance(value, bytes):
-        raise TypeError("stdin must be bytes or None")
-    return {"encoding": "base64", "data": base64.b64encode(value).decode("ascii")}
-
-
-def _decode_inline_bytes(value: object, field: str) -> bytes:
-    if not isinstance(value, dict) or value.get("encoding") != "base64" or not isinstance(value.get("data"), str):
-        raise HostBridgeError("protocol_mismatch", f"response field {field} must contain base64 data")
-    try:
-        return base64.b64decode(value["data"], validate=True)
-    except (binascii.Error, ValueError) as exc:
-        raise HostBridgeError("protocol_mismatch", f"response field {field} contains invalid base64") from exc
-
-
 @dataclass(frozen=True, slots=True)
 class SessionInfo:
     session_id: str
@@ -80,21 +53,70 @@ class ExecResult:
     stderr: bytes
     timed_out: bool = False
 
-    @classmethod
-    def from_dict(cls, value: Mapping[str, object]) -> ExecResult:
-        exit_code = value.get("exit_code")
-        if exit_code is not None and (not isinstance(exit_code, int) or isinstance(exit_code, bool)):
-            raise HostBridgeError("protocol_mismatch", "response field exit_code must be an integer or null")
-        timed_out = value.get("timed_out", False)
-        if not isinstance(timed_out, bool):
-            raise HostBridgeError("protocol_mismatch", "response field timed_out must be a boolean")
-        return cls(
-            session_id=_required_text(value.get("session_id"), "session_id"),
-            exit_code=exit_code,
-            stdout=_decode_inline_bytes(value.get("stdout"), "stdout"),
-            stderr=_decode_inline_bytes(value.get("stderr"), "stderr"),
-            timed_out=timed_out,
-        )
+
+class HostBridgeShell:
+    def __init__(self, stream: MuxSyncStream) -> None:
+        self._stream = stream
+        self._state_lock = threading.Lock()
+        self._closed = False
+        self._input_closed = False
+        self._terminal: ShellTerminalResult | None = None
+
+    @property
+    def closed(self) -> bool:
+        with self._state_lock:
+            return self._closed
+
+    @property
+    def exit_code(self) -> int | None:
+        with self._state_lock:
+            return None if self._terminal is None else self._terminal.exit_code
+
+    @property
+    def exit_signal(self) -> str | None:
+        with self._state_lock:
+            return None if self._terminal is None else self._terminal.signal
+
+    def receive(self, max_bytes: int = 64 * 1024) -> bytes | None:
+        if max_bytes <= 0:
+            raise ValueError("max_bytes must be positive")
+        with self._state_lock:
+            if self._closed:
+                return None
+        chunk = self._stream.receive_unbounded(max_bytes)
+        if chunk is not None:
+            return chunk
+        result = self._stream.finish(timeout=1)
+        try:
+            terminal = ShellTerminalResult.from_mapping(result)
+        except ValueError as exc:
+            raise HostBridgeError("protocol_mismatch", f"invalid shell terminal result: {exc}") from exc
+        with self._state_lock:
+            self._terminal = terminal
+            self._closed = True
+        return None
+
+    def send(self, data: bytes) -> None:
+        if not isinstance(data, bytes):
+            raise TypeError("shell data must be bytes")
+        with self._state_lock:
+            if self._closed or self._input_closed:
+                raise HostBridgeError("connection_lost", "shell input is closed")
+        self._stream.send(data)
+
+    def send_eof(self) -> None:
+        with self._state_lock:
+            if self._closed or self._input_closed:
+                return
+            self._input_closed = True
+        self._stream.send_eof()
+
+    def close(self) -> None:
+        with self._state_lock:
+            if self._closed:
+                return
+            self._closed = True
+        self._stream.reset("shell output consumer closed", timeout=1)
 
 
 class HostBridgeError(RuntimeError):
@@ -111,10 +133,6 @@ class HostBridgeError(RuntimeError):
         self.retryable = retryable
         self.details = details or {}
 
-    @classmethod
-    def from_info(cls, error: ErrorInfo) -> HostBridgeError:
-        return cls(error.code, error.message, retryable=error.retryable, details=error.details)
-
 
 class HostBridgeClient:
     def __init__(
@@ -127,6 +145,25 @@ class HostBridgeClient:
             raise ValueError("default_timeout must be positive")
         self.socket_path = runtime_paths(socket_path).socket
         self.default_timeout = float(default_timeout)
+        self._mux = MuxSyncClient(self.socket_path, default_timeout=self.default_timeout)
+
+    @property
+    def closed(self) -> bool:
+        return self._mux.closed
+
+    def close(self) -> None:
+        self._mux.close()
+
+    def __enter__(self) -> HostBridgeClient:
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        self.close()
 
     def request(
         self,
@@ -135,28 +172,56 @@ class HostBridgeClient:
         *,
         timeout: float | None = None,
     ) -> dict[str, object]:
-        request = Request.new(method, dict(params))
-        response = self._exchange(request, timeout=self.default_timeout if timeout is None else timeout)
-        if not response.ok:
-            if response.error is None:
-                raise HostBridgeError("protocol_mismatch", "failed response omitted error details")
-            raise HostBridgeError.from_info(response.error)
-        return response.result
+        try:
+            return self._mux.request(method, params, timeout=timeout)
+        except BaseException as exc:
+            raise _translate_error(exc) from exc
 
     def hello(self) -> dict[str, object]:
-        return self.request("system.hello", {})
+        result = self.request("system.hello", {})
+        if (
+            result.get("service") != "hostbridge"
+            or result.get("protocol") != MUX_PROTOCOL_VERSION
+            or result.get("version") != __version__
+        ):
+            raise HostBridgeError(
+                "protocol_mismatch",
+                f"HostBridge daemon is not compatible with runtime {__version__}",
+            )
+        return result
+
+    def list_hosts(self) -> list[dict[str, object]]:
+        hosts = self.request("host.list", {}).get("hosts")
+        if not isinstance(hosts, list) or any(not isinstance(host, dict) for host in hosts):
+            raise HostBridgeError("protocol_mismatch", "host list response must contain objects")
+        return hosts
+
+    def tunnel_health(self) -> dict[str, object]:
+        result = self.request("tunnel.health", {})
+        hosts = result.get("hosts")
+        if not isinstance(hosts, list) or any(not isinstance(host, dict) for host in hosts):
+            raise HostBridgeError("protocol_mismatch", "tunnel health response must contain host objects")
+        return result
+
+    def list_sessions(self) -> list[dict[str, object]]:
+        sessions = self.request("session.list", {}).get("sessions")
+        if not isinstance(sessions, list) or any(not isinstance(session, dict) for session in sessions):
+            raise HostBridgeError("protocol_mismatch", "session list response must contain objects")
+        return sessions
 
     def open_session(
         self,
         host_id: str,
         *,
         owner: str | None = None,
-        timeout: float | None = None,
+        connect_timeout: float = 30,
     ) -> SessionInfo:
+        if connect_timeout <= 0:
+            raise ValueError("connect_timeout must be positive")
         result = self.request(
             "session.open",
-            {"host_id": host_id, "owner": owner},
-            timeout=timeout,
+            {"host_id": host_id, "owner": owner, "connect_timeout": connect_timeout},
+            timeout=connect_timeout + 5,
         )
         return SessionInfo.from_dict(result)
 
@@ -170,26 +235,76 @@ class HostBridgeClient:
         output_limit: int | None = None,
         owner: str | None = None,
     ) -> ExecResult:
-        params: dict[str, object] = {
-            "session_id": session_id,
-            "command": command,
-            "stdin": _encode_inline_bytes(stdin),
-        }
+        if stdin is not None and not isinstance(stdin, bytes):
+            raise TypeError("stdin must be bytes or None")
+        chunks: Iterable[bytes] = () if stdin is None else (stdin,)
+        return self.exec_stream(
+            session_id,
+            command,
+            chunks,
+            timeout=timeout,
+            output_limit=output_limit,
+            owner=owner,
+        )
+
+    def exec_stream(
+        self,
+        session_id: str,
+        command: str,
+        stdin: Iterable[bytes],
+        *,
+        timeout: float | None = None,
+        output_limit: int | None = None,
+        owner: str | None = None,
+        cancel_event: threading.Event | None = None,
+    ) -> ExecResult:
+        params: dict[str, object] = {"session_id": session_id, "command": command}
         if timeout is not None:
             params["timeout"] = timeout
         if output_limit is not None:
             params["output_limit"] = output_limit
         if owner is not None:
             params["owner"] = owner
-        result = self.request("exec.start", params, timeout=timeout)
-        return ExecResult.from_dict(result)
+        deadline = _deadline(self.default_timeout if timeout is None else timeout)
+        stream = self._open_stream("exec.start", params, deadline)
+        decoder = MuxRecordDecoder()
+        stdout = bytearray()
+        stderr = bytearray()
+        try:
+            self._send_input(stream, stdin, deadline, cancel_event=cancel_event, label="exec stdin")
+            while chunk := stream.receive(timeout=_remaining(deadline)):
+                for record in decoder.feed(chunk):
+                    target = stdout if record.channel is MuxRecordChannel.STDOUT else stderr
+                    target.extend(record.payload)
+            if decoder.buffered_bytes:
+                raise HostBridgeError("protocol_mismatch", "exec output ended with an incomplete binary record")
+            result = stream.finish(timeout=_remaining(deadline))
+        except BaseException as exc:
+            self._reset_stream(stream, str(exc))
+            raise _translate_error(exc) from exc
+        exit_code = result.get("exit_code")
+        if exit_code is not None and (not isinstance(exit_code, int) or isinstance(exit_code, bool)):
+            raise HostBridgeError("protocol_mismatch", "response field exit_code must be an integer or null")
+        timed_out = result.get("timed_out", False)
+        if not isinstance(timed_out, bool):
+            raise HostBridgeError("protocol_mismatch", "response field timed_out must be a boolean")
+        return ExecResult(
+            _required_text(result.get("session_id"), "session_id"),
+            exit_code,
+            bytes(stdout),
+            bytes(stderr),
+            timed_out,
+        )
 
-    def close_session(self, session_id: str, *, owner: str | None = None, force: bool = False) -> None:
-        params: dict[str, object] = {"session_id": session_id, "force": force}
+    def close_session(self, session_id: str, *, owner: str | None = None) -> None:
+        params: dict[str, object] = {"session_id": session_id}
         if owner is not None:
             params["owner"] = owner
         self.request("session.close", params)
 
+    def release_mock_ssh_session(self, session_id: str) -> None:
+        self.request("session.release_mock_ssh", {"session_id": session_id})
+
     def write_text(
         self,
         session_id: str,
@@ -209,8 +324,7 @@ class HostBridgeClient:
         }
         if owner is not None:
             params["owner"] = owner
-        result = self.request("file.write_text", params)
-        return self._transfer_result(result)
+        return self._transfer_result(self.request("file.write_text", params))
 
     def read_text(
         self,
@@ -220,28 +334,55 @@ class HostBridgeClient:
         max_bytes: int = 4 * 1024 * 1024,
         owner: str | None = None,
     ) -> str:
-        params: dict[str, object] = {"session_id": session_id, "remote_path": remote_path, "max_bytes": max_bytes}
+        params: dict[str, object] = {
+            "session_id": session_id,
+            "remote_path": remote_path,
+            "max_bytes": max_bytes,
+        }
         if owner is not None:
             params["owner"] = owner
-        result = self.request("file.read_text", params)
-        content = result.get("content")
+        content = self.request("file.read_text", params).get("content")
         if not isinstance(content, str):
             raise HostBridgeError("protocol_mismatch", "text response content must be a string")
         return content
 
     def shell_write(self, session_id: str, data: bytes, *, owner: str | None = None) -> int:
-        params: dict[str, object] = {
-            "session_id": session_id,
-            "data": _encode_inline_bytes(data),
-        }
+        if not isinstance(data, bytes):
+            raise TypeError("shell data must be bytes")
+        params: dict[str, object] = {"session_id": session_id, "max_bytes": max(1, len(data))}
         if owner is not None:
             params["owner"] = owner
-        result = self.request("shell.write", params)
+        deadline = _deadline(self.default_timeout)
+        stream = self._open_stream("shell.write", params, deadline)
+        try:
+            self._send_input(stream, (data,), deadline, cancel_event=None, label="shell input")
+            while stream.receive(timeout=_remaining(deadline)) is not None:
+                pass
+            result = stream.finish(timeout=_remaining(deadline))
+        except BaseException as exc:
+            self._reset_stream(stream, str(exc))
+            raise _translate_error(exc) from exc
         written = result.get("written")
         if not isinstance(written, int) or isinstance(written, bool) or written < 0:
             raise HostBridgeError("protocol_mismatch", "shell write response contains invalid byte count")
         return written
 
+    def open_shell(
+        self,
+        session_id: str,
+        *,
+        owner: str | None = None,
+        max_bytes: int = 64 * 1024,
+    ) -> HostBridgeShell:
+        if max_bytes <= 0:
+            raise ValueError("max_bytes must be positive")
+        params: dict[str, object] = {"session_id": session_id, "max_bytes": max_bytes}
+        if owner is not None:
+            params["owner"] = owner
+        deadline = _deadline(self.default_timeout)
+        stream = self._open_stream("shell.attach", params, deadline)
+        return HostBridgeShell(stream)
+
     def shell_read(
         self,
         session_id: str,
@@ -250,18 +391,24 @@ class HostBridgeClient:
         max_bytes: int = 4096,
         owner: str | None = None,
     ) -> ShellReadResult:
-        params: dict[str, object] = {
-            "session_id": session_id,
-            "timeout": timeout,
-            "max_bytes": max_bytes,
-        }
+        params: dict[str, object] = {"session_id": session_id, "timeout": timeout, "max_bytes": max_bytes}
         if owner is not None:
             params["owner"] = owner
-        result = self.request("shell.read", params, timeout=max(self.default_timeout, timeout + 1))
+        deadline = _deadline(max(self.default_timeout, timeout + 1))
+        stream = self._open_stream("shell.read", params, deadline)
+        output = bytearray()
+        try:
+            stream.send_eof(timeout=_remaining(deadline))
+            while chunk := stream.receive(timeout=_remaining(deadline)):
+                output.extend(chunk)
+            result = stream.finish(timeout=_remaining(deadline))
+        except BaseException as exc:
+            self._reset_stream(stream, str(exc))
+            raise _translate_error(exc) from exc
         alive = result.get("alive")
         if not isinstance(alive, bool):
             raise HostBridgeError("protocol_mismatch", "shell read response contains invalid alive flag")
-        return ShellReadResult(_decode_inline_bytes(result.get("data"), "data"), alive)
+        return ShellReadResult(bytes(output), alive)
 
     def upload_file(
         self,
@@ -274,54 +421,40 @@ class HostBridgeClient:
         timeout: float | None = None,
         cancel_event: threading.Event | None = None,
         owner: str | None = None,
+        max_bytes: int | None = None,
     ) -> TransferResult:
         self._validate_chunk_size(chunk_size)
+        self._validate_optional_byte_limit(max_bytes)
         source = Path(local_path)
         expected_size = source.stat().st_size
-        stream_id = uuid.uuid4().hex
+        if max_bytes is not None and expected_size > max_bytes:
+            raise HostBridgeError("resource_limit", f"upload exceeds {max_bytes} bytes")
         params: dict[str, object] = {
             "session_id": session_id,
             "remote_path": remote_path,
-            "stream_id": stream_id,
             "mode": mode,
             "expected_size": expected_size,
         }
         if owner is not None:
             params["owner"] = owner
-        request = Request.new("transfer.upload", params)
+        deadline = _deadline(self.default_timeout if timeout is None else timeout)
+        stream = self._open_stream("transfer.upload", params, deadline)
         digest = hashlib.sha256()
-        sequence = 0
-        request_timeout = self.default_timeout if timeout is None else timeout
+
+        def chunks() -> Iterable[bytes]:
+            with source.open("rb") as handle:
+                while chunk := handle.read(chunk_size):
+                    digest.update(chunk)
+                    yield chunk
+
         try:
-            with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection:
-                connection.settimeout(request_timeout)
-                connection.connect(str(self.socket_path))
-                connection.sendall(request.to_json().encode("utf-8") + b"\n")
-                with source.open("rb") as local_file:
-                    while True:
-                        if cancel_event is not None and cancel_event.is_set():
-                            payload = json.dumps(
-                                {"code": "cancelled", "message": "upload cancelled by client"},
-                                separators=(",", ":"),
-                            ).encode("utf-8")
-                            connection.sendall(
-                                encode_frame(BinaryFrame(stream_id, sequence, FrameFlags.CANCEL, payload))
-                            )
-                            break
-                        chunk = local_file.read(chunk_size)
-                        if not chunk:
-                            connection.sendall(encode_frame(BinaryFrame(stream_id, sequence, FrameFlags.EOF)))
-                            break
-                        digest.update(chunk)
-                        connection.sendall(encode_frame(BinaryFrame(stream_id, sequence, FrameFlags.DATA, chunk)))
-                        sequence += 1
-                with connection.makefile("rb") as reader:
-                    response = self._read_response(reader, request)
-        except TimeoutError as exc:
-            raise HostBridgeError("timeout", "HostBridge upload timed out", retryable=True) from exc
-        except OSError as exc:
-            raise HostBridgeError("connection_lost", "HostBridge upload connection failed", retryable=True) from exc
-        result = self._transfer_result(response.result)
+            self._send_input(stream, chunks(), deadline, cancel_event=cancel_event, label="upload")
+            while stream.receive(timeout=_remaining(deadline)) is not None:
+                pass
+            result = self._transfer_result(stream.finish(timeout=_remaining(deadline)))
+        except BaseException as exc:
+            self._reset_stream(stream, str(exc))
+            raise _translate_error(exc) from exc
         if result.bytes_transferred != expected_size or result.sha256 != digest.hexdigest():
             raise HostBridgeError("integrity_error", "HostBridge upload integrity check failed")
         return result
@@ -335,69 +468,82 @@ class HostBridgeClient:
         chunk_size: int = 256 * 1024,
         timeout: float | None = None,
         owner: str | None = None,
+        max_bytes: int | None = None,
     ) -> TransferResult:
         self._validate_chunk_size(chunk_size)
+        self._validate_optional_byte_limit(max_bytes)
         destination = Path(local_path)
         destination.parent.mkdir(parents=True, exist_ok=True)
         temporary = destination.with_name(f".{destination.name}.hostbridge-{uuid.uuid4().hex}.tmp")
-        stream_id = uuid.uuid4().hex
         params: dict[str, object] = {
             "session_id": session_id,
             "remote_path": remote_path,
-            "stream_id": stream_id,
             "chunk_size": chunk_size,
         }
         if owner is not None:
             params["owner"] = owner
-        request = Request.new("transfer.download", params)
+        deadline = _deadline(self.default_timeout if timeout is None else timeout)
+        stream = self._open_stream("transfer.download", params, deadline)
         digest = hashlib.sha256()
         transferred = 0
-        validator = FrameSequenceValidator()
-        request_timeout = self.default_timeout if timeout is None else timeout
         try:
-            with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection:
-                connection.settimeout(request_timeout)
-                connection.connect(str(self.socket_path))
-                connection.sendall(request.to_json().encode("utf-8") + b"\n")
-                with connection.makefile("rb") as reader:
-                    response = self._read_response(reader, request)
-                    if response.result.get("stream_id") != stream_id:
-                        raise HostBridgeError("protocol_mismatch", "download stream_id does not match")
-                    with temporary.open("wb") as local_file:
-                        while True:
-                            frame = read_frame(reader)
-                            validator.accept(frame)
-                            if frame.stream_id != stream_id:
-                                raise HostBridgeError("protocol_mismatch", "download frame stream_id does not match")
-                            if frame.flags & FrameFlags.CANCEL:
-                                self._raise_cancel_frame(frame)
-                            if frame.flags & FrameFlags.DATA:
-                                local_file.write(frame.payload)
-                                digest.update(frame.payload)
-                                transferred += len(frame.payload)
-                            if frame.flags & FrameFlags.EOF:
-                                metadata = self._frame_metadata(frame.payload)
-                                break
-        except HostBridgeError:
-            temporary.unlink(missing_ok=True)
-            raise
-        except TimeoutError as exc:
-            temporary.unlink(missing_ok=True)
-            raise HostBridgeError("timeout", "HostBridge download timed out", retryable=True) from exc
-        except (OSError, ProtocolError) as exc:
-            temporary.unlink(missing_ok=True)
-            raise HostBridgeError("connection_lost", "HostBridge download connection failed", retryable=True) from exc
-        expected = self._transfer_result(metadata)
-        if expected.bytes_transferred != transferred or expected.sha256 != digest.hexdigest():
+            stream.send_eof(timeout=_remaining(deadline))
+            with temporary.open("wb") as handle:
+                while chunk := stream.receive(timeout=_remaining(deadline)):
+                    transferred += len(chunk)
+                    if max_bytes is not None and transferred > max_bytes:
+                        raise HostBridgeError("resource_limit", f"download exceeds {max_bytes} bytes")
+                    handle.write(chunk)
+                    digest.update(chunk)
+            result = self._transfer_result(stream.finish(timeout=_remaining(deadline)))
+            if result.bytes_transferred != transferred or result.sha256 != digest.hexdigest():
+                raise HostBridgeError("integrity_error", "HostBridge download integrity check failed")
+            temporary.replace(destination)
+            return result
+        except BaseException as exc:
             temporary.unlink(missing_ok=True)
-            raise HostBridgeError("integrity_error", "HostBridge download integrity check failed")
-        temporary.replace(destination)
-        return expected
+            self._reset_stream(stream, str(exc))
+            raise _translate_error(exc) from exc
+
+    def _open_stream(self, method: str, params: Mapping[str, object], deadline: float) -> MuxSyncStream:
+        try:
+            return self._mux.open_stream(method, params, timeout=_remaining(deadline))
+        except BaseException as exc:
+            raise _translate_error(exc) from exc
+
+    @staticmethod
+    def _send_input(
+        stream: MuxSyncStream,
+        chunks: Iterable[bytes],
+        deadline: float,
+        *,
+        cancel_event: threading.Event | None,
+        label: str,
+    ) -> None:
+        for chunk in chunks:
+            if cancel_event is not None and cancel_event.is_set():
+                raise HostBridgeError("cancelled", f"{label} cancelled")
+            if not isinstance(chunk, bytes):
+                raise TypeError(f"{label} chunks must be bytes")
+            stream.send(chunk, timeout=_remaining(deadline))
+        if cancel_event is not None and cancel_event.is_set():
+            raise HostBridgeError("cancelled", f"{label} cancelled")
+        stream.send_eof(timeout=_remaining(deadline))
+
+    @staticmethod
+    def _reset_stream(stream: MuxSyncStream, message: str) -> None:
+        with suppress(BaseException):
+            stream.reset(message or "client operation failed", timeout=1)
 
     @staticmethod
     def _validate_chunk_size(chunk_size: int) -> None:
-        if not isinstance(chunk_size, int) or isinstance(chunk_size, bool) or not 4096 <= chunk_size <= MAX_FRAME_BYTES:
-            raise ValueError(f"chunk_size must be between 4096 and {MAX_FRAME_BYTES} bytes")
+        if not isinstance(chunk_size, int) or isinstance(chunk_size, bool) or not 4096 <= chunk_size <= 1024 * 1024:
+            raise ValueError("chunk_size must be between 4096 and 1048576 bytes")
+
+    @staticmethod
+    def _validate_optional_byte_limit(max_bytes: int | None) -> None:
+        if max_bytes is not None and (not isinstance(max_bytes, int) or isinstance(max_bytes, bool) or max_bytes <= 0):
+            raise ValueError("max_bytes must be a positive integer or None")
 
     @staticmethod
     def _transfer_result(value: Mapping[str, object]) -> TransferResult:
@@ -409,70 +555,31 @@ class HostBridgeClient:
             raise HostBridgeError("protocol_mismatch", "transfer response contains invalid SHA-256")
         return TransferResult(byte_count, sha256.lower())
 
-    def _read_response(self, reader, request: Request) -> Response:
-        raw_response = reader.readline(MAX_CONTROL_LINE_BYTES + 1)
-        if not raw_response:
-            raise HostBridgeError("connection_lost", "HostBridge daemon closed before responding", retryable=True)
-        if len(raw_response) > MAX_CONTROL_LINE_BYTES:
-            raise HostBridgeError("protocol_mismatch", "HostBridge daemon response exceeds the control limit")
-        try:
-            response = Response.from_json(raw_response.decode("utf-8"))
-        except (ProtocolError, UnicodeDecodeError) as exc:
-            raise HostBridgeError("protocol_mismatch", "HostBridge daemon returned an invalid response") from exc
-        if response.request_id != request.request_id:
-            raise HostBridgeError("protocol_mismatch", "HostBridge daemon response request_id does not match")
-        if not response.ok:
-            if response.error is None:
-                raise HostBridgeError("protocol_mismatch", "failed response omitted error details")
-            raise HostBridgeError.from_info(response.error)
-        return response
-
-    @staticmethod
-    def _frame_metadata(payload: bytes) -> dict[str, object]:
-        try:
-            value = json.loads(payload.decode("utf-8"))
-        except (UnicodeDecodeError, json.JSONDecodeError) as exc:
-            raise HostBridgeError("protocol_mismatch", "download EOF contains invalid metadata") from exc
-        if not isinstance(value, dict):
-            raise HostBridgeError("protocol_mismatch", "download EOF metadata must be an object")
-        return value
-
-    @staticmethod
-    def _raise_cancel_frame(frame: BinaryFrame) -> None:
-        metadata = HostBridgeClient._frame_metadata(frame.payload)
-        code = metadata.get("code")
-        message = metadata.get("message")
-        raise HostBridgeError(
-            code if isinstance(code, str) else "cancelled",
-            message if isinstance(message, str) else "HostBridge transfer cancelled",
-        )
-
-    def _exchange(self, request: Request, *, timeout: float) -> Response:
-        if timeout <= 0:
-            raise ValueError("timeout must be positive")
-        try:
-            with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection:
-                connection.settimeout(timeout)
-                connection.connect(str(self.socket_path))
-                connection.sendall(request.to_json().encode("utf-8") + b"\n")
-                with connection.makefile("rb") as reader:
-                    raw_response = reader.readline(MAX_CONTROL_LINE_BYTES + 1)
-        except TimeoutError as exc:
-            raise HostBridgeError("timeout", "HostBridge daemon request timed out", retryable=True) from exc
-        except OSError as exc:
-            raise HostBridgeError("connection_lost", "HostBridge daemon connection failed", retryable=True) from exc
-        if not raw_response:
-            raise HostBridgeError("connection_lost", "HostBridge daemon closed before responding", retryable=True)
-        if len(raw_response) > MAX_CONTROL_LINE_BYTES:
-            raise HostBridgeError("protocol_mismatch", "HostBridge daemon response exceeds the control limit")
-        return self._read_response_from_bytes(raw_response, request)
 
-    @staticmethod
-    def _read_response_from_bytes(raw_response: bytes, request: Request) -> Response:
-        try:
-            response = Response.from_json(raw_response.decode("utf-8"))
-        except (ProtocolError, UnicodeDecodeError) as exc:
-            raise HostBridgeError("protocol_mismatch", "HostBridge daemon returned an invalid response") from exc
-        if response.request_id != request.request_id:
-            raise HostBridgeError("protocol_mismatch", "HostBridge daemon response request_id does not match")
-        return response
+def _deadline(timeout: float) -> float:
+    if timeout <= 0:
+        raise ValueError("timeout must be positive")
+    return time.monotonic() + timeout
+
+
+def _remaining(deadline: float) -> float:
+    remaining = deadline - time.monotonic()
+    if remaining <= 0:
+        raise HostBridgeError("timeout", "HostBridge operation timed out", retryable=True)
+    return remaining
+
+
+def _translate_error(error: BaseException) -> BaseException:
+    if isinstance(error, MuxRecordError):
+        return HostBridgeError("protocol_mismatch", str(error))
+    if isinstance(error, HostBridgeError | TypeError | ValueError):
+        return error
+    if isinstance(error, MuxRemoteError):
+        return HostBridgeError(error.code, str(error), retryable=error.retryable, details=error.details)
+    if isinstance(error, MuxCapacityError):
+        return HostBridgeError("busy", str(error), retryable=True)
+    if isinstance(error, MuxConnectionClosed | RuntimeError):
+        return HostBridgeError("connection_lost", str(error), retryable=True)
+    if isinstance(error, TimeoutError):
+        return HostBridgeError("timeout", str(error), retryable=True)
+    return error