@neoline/hostbridge 2.0.4 → 2.0.5

package/src/server_control_mcp/tunnel_pty_agent.py

@@ -0,0 +1,909 @@
+"""Persistent raw-binary PTY tunnel provider and bootstrap transport."""
+
+from __future__ import annotations
+
+import asyncio
+import errno
+import hashlib
+import os
+import re
+import secrets
+import shlex
+import signal
+from collections.abc import AsyncGenerator, Callable
+from contextlib import suppress
+from re import Pattern
+
+from .async_lifecycle import finish_task_before_cancellation, stop_input_after_peer_eof
+from .hosts import ForwardingPolicy, HostConfig
+from .mux_records import MuxRecordChannel, MuxRecordDecoder
+from .mux_rpc import (
+    DEFAULT_CONNECTION_WINDOW_BYTES,
+    DEFAULT_STREAM_WINDOW_BYTES,
+    MuxRpcClient,
+)
+from .mux_stream import MuxRpcStream
+from .remote_agent_bundle import build_remote_agent_bundle
+from .transports.base import (
+    ByteStream,
+    ExecResult,
+    ShellAttachment,
+    ShellReadResult,
+    ShellTerminalResult,
+    TransferResult,
+    TransportCapabilities,
+    iterate_byte_stream,
+)
+from .tunnel_manager import TunnelBackendError
+
+MAX_BOOTSTRAP_BUFFER_BYTES = 1024 * 1024
+MAX_PTY_WRITE_BUFFER_BYTES = 2 * 1024 * 1024
+PTY_IO_CHUNK_BYTES = 64 * 1024
+PROCESS_TERM_TIMEOUT_SECONDS = 1.0
+MAX_BOOTSTRAP_ERROR_BYTES = 64
+
+_BOOTSTRAP = r"""import hashlib,os,runpy,sys,tempfile,termios,tty
+q=sys.argv[1].encode();n=int(sys.argv[2]);h=sys.argv[3];t=termios.tcgetattr(0);p=None;f=-1
+try:
+ tty.setraw(0);e=os.open(os.devnull,os.O_WRONLY);os.dup2(e,2);os.close(e);os.write(1,b"\0HBRAW2:"+q+b"\0");f,p=tempfile.mkstemp(prefix="hostbridge-agent-",suffix=".pyz");d=hashlib.sha256()
+ while n:
+  b=os.read(0,min(n,65536))
+  if not b:raise EOFError()
+  os.write(f,b);d.update(b);n-=len(b)
+ os.close(f);f=-1
+ if d.hexdigest()!=h:raise ValueError("hash mismatch")
+ os.write(1,b"\0HBOK2:"+q+b"\0");runpy.run_path(p,run_name="__main__")
+finally:
+ if f>=0:os.close(f)
+ if p:
+  try:os.unlink(p)
+  except FileNotFoundError:pass
+ termios.tcsetattr(0,termios.TCSANOW,t)
+"""
+
+
+class _AsyncPtyWriter:
+    def __init__(self, owner: AsyncPtyByteStream) -> None:
+        self._owner = owner
+
+    def write(self, data: bytes) -> None:
+        self._owner.write(data)
+
+    async def drain(self) -> None:
+        await self._owner.drain()
+
+    def close(self) -> None:
+        self._owner.close_nowait()
+
+    async def wait_closed(self) -> None:
+        await asyncio.shield(self._owner.closed_future)
+
+    def is_closing(self) -> bool:
+        return self._owner.closed
+
+
+class _PtyHandshakeReader(asyncio.StreamReader):
+    def __init__(self, reader: asyncio.StreamReader, nonce: bytes) -> None:
+        super().__init__()
+        self._reader = reader
+        self._failure_prefix = b"\x00HBFAIL2:" + nonce + b":"
+        self._pending = bytearray()
+        self._handshake_complete = False
+
+    async def read(self, n: int = -1) -> bytes:
+        if self._handshake_complete:
+            if self._pending:
+                size = len(self._pending) if n < 0 else min(n, len(self._pending))
+                data = bytes(self._pending[:size])
+                del self._pending[:size]
+                return data
+            return await self._reader.read(n)
+        while True:
+            if self._pending.startswith(b"HBM2"):
+                self._handshake_complete = True
+                size = len(self._pending) if n < 0 else min(n, len(self._pending))
+                data = bytes(self._pending[:size])
+                del self._pending[:size]
+                return data
+            if self._pending.startswith(self._failure_prefix):
+                terminator = self._pending.find(b"\x00", len(self._failure_prefix))
+                if terminator >= 0:
+                    raw_error = bytes(self._pending[len(self._failure_prefix) : terminator])
+                    error = raw_error.decode("ascii", "replace") or "unknown error"
+                    raise TunnelBackendError(f"remote agent bootstrap failed: {error}")
+                if len(self._pending) > len(self._failure_prefix) + MAX_BOOTSTRAP_ERROR_BYTES:
+                    raise TunnelBackendError("remote agent returned an oversized bootstrap failure")
+            elif not (b"HBM2".startswith(self._pending) or self._failure_prefix.startswith(self._pending)):
+                raise TunnelBackendError("remote agent returned invalid bytes before the mux handshake")
+
+            data = await self._reader.read(64 * 1024)
+            if not data:
+                raise TunnelBackendError("remote agent closed before the mux handshake")
+            self._pending.extend(data)
+
+
+class AsyncPtyByteStream:
+    def __init__(
+        self,
+        master_fd: int,
+        *,
+        max_bootstrap_buffer: int = MAX_BOOTSTRAP_BUFFER_BYTES,
+        max_write_buffer: int = MAX_PTY_WRITE_BUFFER_BYTES,
+    ) -> None:
+        if master_fd < 0:
+            raise ValueError("master_fd must be non-negative")
+        if max_bootstrap_buffer <= 0 or max_write_buffer <= 0:
+            raise ValueError("PTY buffer limits must be positive")
+        self._loop = asyncio.get_running_loop()
+        self._fd = master_fd
+        self._max_bootstrap_buffer = max_bootstrap_buffer
+        self._max_write_buffer = max_write_buffer
+        self._bootstrap_buffer = bytearray()
+        self._write_buffer = bytearray()
+        self._data_ready = asyncio.Event()
+        self._write_drained = asyncio.Event()
+        self._write_drained.set()
+        self._reader: asyncio.StreamReader | None = None
+        self._handed_off = False
+        self._eof = False
+        self._closed = False
+        self.closed_future: asyncio.Future[None] = self._loop.create_future()
+        os.set_blocking(master_fd, False)
+        self._loop.add_reader(master_fd, self._read_ready)
+
+    @property
+    def closed(self) -> bool:
+        return self._closed
+
+    async def read_until(self, pattern: bytes | Pattern[bytes], *, timeout: float) -> bytes:
+        if self._handed_off:
+            raise RuntimeError("PTY byte stream has already been handed off")
+        if timeout <= 0:
+            raise ValueError("timeout must be positive")
+        compiled = re.compile(re.escape(pattern)) if isinstance(pattern, bytes) else pattern
+        try:
+            async with asyncio.timeout(timeout):
+                while True:
+                    match = compiled.search(self._bootstrap_buffer)
+                    if match is not None:
+                        consumed = bytes(self._bootstrap_buffer[: match.end()])
+                        del self._bootstrap_buffer[: match.end()]
+                        return consumed
+                    if self._eof or self._closed:
+                        raise EOFError("PTY closed before bootstrap marker")
+                    if len(self._bootstrap_buffer) > self._max_bootstrap_buffer:
+                        raise TunnelBackendError("PTY bootstrap output exceeded its buffer limit")
+                    self._data_ready.clear()
+                    await self._data_ready.wait()
+        except TimeoutError as exc:
+            raise TunnelBackendError("PTY bootstrap timed out") from exc
+
+    def handoff(self) -> tuple[asyncio.StreamReader, _AsyncPtyWriter]:
+        if self._handed_off:
+            raise RuntimeError("PTY byte stream has already been handed off")
+        self._handed_off = True
+        self._reader = asyncio.StreamReader(limit=self._max_bootstrap_buffer)
+        if self._bootstrap_buffer:
+            self._reader.feed_data(bytes(self._bootstrap_buffer))
+            self._bootstrap_buffer.clear()
+        if self._eof or self._closed:
+            self._reader.feed_eof()
+        return self._reader, _AsyncPtyWriter(self)
+
+    def write(self, data: bytes) -> None:
+        if not isinstance(data, bytes):
+            raise TypeError("PTY writes must be bytes")
+        if not data:
+            return
+        if self._closed:
+            raise ConnectionResetError("PTY byte stream is closed")
+        if len(self._write_buffer) + len(data) > self._max_write_buffer:
+            raise BufferError("PTY write buffer limit exceeded")
+        self._write_drained.clear()
+        self._write_buffer.extend(data)
+        self._flush_write()
+
+    async def drain(self) -> None:
+        while self._write_buffer and not self._closed:
+            await self._write_drained.wait()
+        if self._closed and self._write_buffer:
+            raise ConnectionResetError("PTY closed before buffered data was written")
+
+    async def close(self) -> None:
+        self.close_nowait()
+        await asyncio.shield(self.closed_future)
+
+    def close_nowait(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        self._remove_watchers()
+        with suppress(OSError):
+            os.close(self._fd)
+        self._eof = True
+        self._data_ready.set()
+        self._write_drained.set()
+        if self._reader is not None:
+            self._reader.feed_eof()
+        if not self.closed_future.done():
+            self.closed_future.set_result(None)
+
+    def _read_ready(self) -> None:
+        while not self._closed:
+            try:
+                data = os.read(self._fd, PTY_IO_CHUNK_BYTES)
+            except BlockingIOError:
+                return
+            except OSError as exc:
+                if exc.errno == errno.EIO:
+                    self._mark_eof()
+                    return
+                else:
+                    self.close_nowait()
+                    return
+            if not data:
+                self._mark_eof()
+                return
+            if self._reader is None:
+                self._bootstrap_buffer.extend(data)
+            else:
+                self._reader.feed_data(data)
+            self._data_ready.set()
+
+    def _flush_write(self) -> None:
+        while self._write_buffer and not self._closed:
+            try:
+                written = os.write(self._fd, self._write_buffer)
+            except BlockingIOError:
+                self._loop.add_writer(self._fd, self._write_ready)
+                return
+            except OSError:
+                self.close_nowait()
+                return
+            del self._write_buffer[:written]
+        if not self._write_buffer:
+            with suppress(Exception):
+                self._loop.remove_writer(self._fd)
+            self._write_drained.set()
+
+    def _write_ready(self) -> None:
+        self._flush_write()
+
+    def _remove_watchers(self) -> None:
+        with suppress(Exception):
+            self._loop.remove_reader(self._fd)
+        with suppress(Exception):
+            self._loop.remove_writer(self._fd)
+
+    def _mark_eof(self) -> None:
+        if self._eof:
+            return
+        self._eof = True
+        with suppress(Exception):
+            self._loop.remove_reader(self._fd)
+        self._data_ready.set()
+        if self._reader is not None:
+            self._reader.feed_eof()
+
+
+class PtyAgentTunnelEndpoint:
+    def __init__(self, stream: MuxRpcStream, closed_handler: Callable[[int], None]) -> None:
+        self._stream = stream
+        self.provider_stream_id = stream.stream_id
+        self._closed_handler = closed_handler
+        self._watch_task = asyncio.create_task(
+            self._watch_terminal(),
+            name=f"hostbridge-pty-agent-stream-{stream.stream_id}",
+        )
+
+    @property
+    def closed(self) -> bool:
+        return self._stream.closed
+
+    @property
+    def queued_bytes(self) -> int:
+        return self._stream.buffered_bytes
+
+    async def send(self, data: bytes) -> None:
+        await self._stream.send(data)
+
+    async def send_eof(self) -> None:
+        await self._stream.send_eof()
+
+    async def receive(self, max_bytes: int = PTY_IO_CHUNK_BYTES) -> bytes | None:
+        return await self._stream.receive(max_bytes)
+
+    async def close(self) -> None:
+        if not self._stream.closed:
+            await self._stream.reset("PTY tunnel stream closed")
+        await self._wait_watcher()
+
+    async def reset(self, reason: str) -> None:
+        if not self._stream.closed:
+            await self._stream.reset(reason)
+        await self._wait_watcher()
+
+    async def _watch_terminal(self) -> None:
+        try:
+            await self._stream.finish()
+        except BaseException:
+            pass
+        finally:
+            self._closed_handler(self.provider_stream_id)
+
+    async def _wait_watcher(self) -> None:
+        if self._watch_task is asyncio.current_task():
+            return
+        await asyncio.gather(self._watch_task, return_exceptions=True)
+
+
+class PtyAgentHostRuntime:
+    provider_name = "pty_agent"
+
+    @property
+    def capabilities(self) -> TransportCapabilities:
+        return TransportCapabilities(
+            separate_stderr=True,
+            native_binary_transfer=True,
+            parallel_channels=self.host.forwarding.max_streams,
+        )
+
+    def __init__(
+        self,
+        host: HostConfig,
+        *,
+        generation: int,
+        connect_timeout: float = 30,
+        bundle_builder: Callable[[], bytes] = build_remote_agent_bundle,
+    ) -> None:
+        if generation <= 0:
+            raise ValueError("generation must be positive")
+        if connect_timeout <= 0:
+            raise ValueError("connect_timeout must be positive")
+        self.host = host
+        self.generation = generation
+        self.connect_timeout = float(connect_timeout)
+        self._bundle_builder = bundle_builder
+        self._process: asyncio.subprocess.Process | None = None
+        self._byte_stream: AsyncPtyByteStream | None = None
+        self._client: MuxRpcClient | None = None
+        self._monitor_task: asyncio.Task[None] | None = None
+        self._closed_future: asyncio.Future[BaseException | None] | None = None
+        self._endpoints: dict[int, PtyAgentTunnelEndpoint] = {}
+        self._closing = False
+        self._close_lock = asyncio.Lock()
+        self._agent_pid: int | None = None
+        self._process_returncode: int | None = None
+
+    @property
+    def agent_pid(self) -> int | None:
+        return self._agent_pid
+
+    @property
+    def process_returncode(self) -> int | None:
+        process = self._process
+        return self._process_returncode if process is None else process.returncode
+
+    @property
+    def queued_bytes(self) -> int:
+        return sum(endpoint.queued_bytes for endpoint in self._endpoints.values())
+
+    @property
+    def window_stalls(self) -> int:
+        client = self._client
+        return 0 if client is None else client.window_stalls
+
+    @property
+    def protocol_failures(self) -> int:
+        client = self._client
+        return 0 if client is None else client.protocol_failures
+
+    async def start(self) -> None:
+        if self._process is not None:
+            return
+        self._closed_future = asyncio.get_running_loop().create_future()
+        master_fd, slave_fd = os.openpty()
+        try:
+            process = await asyncio.create_subprocess_exec(
+                self.host.command,
+                *self.host.args,
+                stdin=slave_fd,
+                stdout=slave_fd,
+                stderr=slave_fd,
+                start_new_session=True,
+            )
+        except BaseException:
+            os.close(master_fd)
+            raise
+        finally:
+            os.close(slave_fd)
+        self._process = process
+        self._agent_pid = process.pid
+        self._byte_stream = AsyncPtyByteStream(master_fd)
+        try:
+            nonce = await self._bootstrap(self._byte_stream)
+            reader, writer = self._byte_stream.handoff()
+            client = MuxRpcClient(
+                _PtyHandshakeReader(reader, nonce),
+                writer,  # type: ignore[arg-type]
+                max_pending=self.host.forwarding.max_streams,
+                connection_window=DEFAULT_CONNECTION_WINDOW_BYTES,
+                stream_window=DEFAULT_STREAM_WINDOW_BYTES,
+            )
+            self._client = client
+            async with asyncio.timeout(self.connect_timeout):
+                await client._start()
+            self._monitor_task = asyncio.create_task(
+                self._monitor_process(),
+                name=f"hostbridge-pty-agent-monitor-{self.host.id}-{self.generation}",
+            )
+        except BaseException as exc:
+            await self.close()
+            if isinstance(exc, asyncio.CancelledError):
+                raise
+            if isinstance(exc, TunnelBackendError):
+                raise
+            raise TunnelBackendError(f"PTY agent bootstrap failed: {exc}") from exc
+
+    async def open_stream(
+        self,
+        host: str,
+        port: int,
+        policy: ForwardingPolicy,
+    ) -> PtyAgentTunnelEndpoint:
+        client = self._client
+        if self._closing or client is None or client.closed:
+            raise TunnelBackendError("PTY agent runtime is not connected")
+        if policy != self.host.forwarding:
+            raise TunnelBackendError("PTY agent stream policy does not match the host runtime")
+        try:
+            stream = await client.open_stream(
+                "tcp.connect",
+                {
+                    "host": host,
+                    "port": port,
+                    "allow_hostnames": list(policy.allow_hostnames),
+                    "allow_cidrs": list(policy.allow_cidrs),
+                    "allow_ports": list(policy.allow_ports),
+                    "connect_timeout": min(self.connect_timeout, 180),
+                },
+            )
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            raise TunnelBackendError(f"PTY agent destination open failed: {exc}") from exc
+        endpoint = PtyAgentTunnelEndpoint(stream, self._endpoint_closed)
+        self._endpoints[stream.stream_id] = endpoint
+        return endpoint
+
+    async def exec(
+        self,
+        command: str,
+        *,
+        stdin: ByteStream | None,
+        timeout: float,
+        output_limit: int,
+    ) -> ExecResult:
+        client = self._client
+        if self._closing or client is None or client.closed:
+            raise TunnelBackendError("PTY agent runtime is not connected")
+        try:
+            stream = await client.open_stream(
+                "exec.start",
+                {
+                    "command": command,
+                    "timeout": timeout,
+                    "output_limit": output_limit,
+                },
+            )
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            raise TunnelBackendError(f"PTY agent command open failed: {exc}") from exc
+
+        stdout = bytearray()
+        stderr = bytearray()
+        decoder = MuxRecordDecoder()
+
+        async def send_input() -> None:
+            if stdin is not None:
+                async for chunk in iterate_byte_stream(stdin, label="stdin"):
+                    await stream.send(chunk)
+            await stream.send_eof()
+
+        async def receive_output() -> None:
+            while chunk := await stream.receive():
+                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 TunnelBackendError("PTY agent command returned an incomplete output record")
+
+        tasks = [
+            asyncio.create_task(send_input(), name=f"hostbridge-pty-exec-input-{stream.stream_id}"),
+            asyncio.create_task(receive_output(), name=f"hostbridge-pty-exec-output-{stream.stream_id}"),
+        ]
+        try:
+            await asyncio.gather(*tasks)
+            result = await stream.finish()
+        except BaseException as exc:
+            for task in tasks:
+                if not task.done():
+                    task.cancel()
+            await asyncio.gather(*tasks, return_exceptions=True)
+            await _reset_mux_stream(stream, str(exc) or "PTY agent command failed")
+            raise
+        exit_code = result.get("exit_code")
+        if exit_code is not None and (not isinstance(exit_code, int) or isinstance(exit_code, bool)):
+            raise TunnelBackendError("PTY agent command returned an invalid exit status")
+        timed_out = result.get("timed_out")
+        if not isinstance(timed_out, bool):
+            raise TunnelBackendError("PTY agent command returned an invalid timeout status")
+        return ExecResult(exit_code, bytes(stdout), bytes(stderr), timed_out=timed_out)
+
+    async def upload(
+        self,
+        chunks: ByteStream,
+        remote_path: str,
+        *,
+        mode: int,
+        expected_size: int | None,
+        expected_sha256: str | None,
+    ) -> TransferResult:
+        client = self._require_client()
+        try:
+            stream = await client.open_stream(
+                "transfer.upload",
+                {
+                    "remote_path": remote_path,
+                    "mode": mode,
+                    "expected_size": expected_size,
+                    "expected_sha256": expected_sha256,
+                    "max_bytes": self.host.limits.max_transfer_bytes,
+                },
+            )
+            async for chunk in iterate_byte_stream(chunks, label="upload"):
+                await stream.send(chunk)
+            await stream.send_eof()
+            while await stream.receive() is not None:
+                pass
+            result = await stream.finish()
+        except asyncio.CancelledError:
+            if "stream" in locals():
+                await _reset_mux_stream(stream, "PTY agent upload cancelled")
+            raise
+        except Exception as exc:
+            if "stream" in locals():
+                await _reset_mux_stream(stream, str(exc) or "PTY agent upload failed")
+            raise TunnelBackendError(f"PTY agent upload failed: {exc}") from exc
+        return _transfer_result(result)
+
+    async def download(self, remote_path: str, *, chunk_size: int) -> AsyncGenerator[bytes, None]:
+        client = self._require_client()
+        stream: MuxRpcStream | None = None
+        completed = False
+        try:
+            stream = await client.open_stream(
+                "transfer.download",
+                {
+                    "remote_path": remote_path,
+                    "chunk_size": chunk_size,
+                    "max_bytes": self.host.limits.max_transfer_bytes,
+                },
+            )
+            await stream.send_eof()
+            while chunk := await stream.receive(chunk_size):
+                yield chunk
+            _transfer_result(await stream.finish())
+            completed = True
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            raise TunnelBackendError(f"PTY agent download failed: {exc}") from exc
+        finally:
+            if stream is not None and not completed:
+                await _reset_mux_stream(stream, "PTY agent download closed before completion")
+
+    async def shell_write(self, shell_id: str, data: bytes) -> None:
+        stream: MuxRpcStream | None = None
+        try:
+            stream = await self._require_client().open_stream(
+                "shell.write",
+                {"shell_id": shell_id, "max_bytes": self.host.limits.max_shell_bytes},
+            )
+            await stream.send(data)
+            await stream.send_eof()
+            while await stream.receive() is not None:
+                pass
+            result = await stream.finish()
+            if result.get("written") != len(data):
+                raise TunnelBackendError("PTY agent shell write returned an invalid byte count")
+        except asyncio.CancelledError:
+            if stream is not None:
+                await _reset_mux_stream(stream, "PTY agent shell write cancelled")
+            raise
+        except Exception as exc:
+            if stream is not None:
+                await _reset_mux_stream(stream, str(exc) or "PTY agent shell write failed")
+            if isinstance(exc, TunnelBackendError):
+                raise
+            raise TunnelBackendError(f"PTY agent shell write failed: {exc}") from exc
+
+    async def attach_shell(self, shell_id: str, stdin: ByteStream, *, max_bytes: int) -> ShellAttachment:
+        stream = await self._require_client().open_stream(
+            "shell.attach",
+            {"shell_id": shell_id, "max_bytes": max_bytes},
+        )
+        completed = False
+        terminal = asyncio.get_running_loop().create_future()
+
+        async def send_input() -> None:
+            async for chunk in iterate_byte_stream(stdin, label="shell input"):
+                await stream.send(chunk)
+            await stream.send_eof()
+
+        async def output() -> AsyncGenerator[bytes, None]:
+            nonlocal completed
+            input_task = asyncio.create_task(send_input(), name=f"hostbridge-pty-shell-input-{stream.stream_id}")
+            transferred = 0
+            try:
+                while chunk := await stream.receive(max_bytes):
+                    transferred += len(chunk)
+                    yield chunk
+                result = await stream.finish()
+                await stop_input_after_peer_eof(input_task)
+                if result.get("bytes") != transferred:
+                    raise TunnelBackendError("PTY agent attached shell returned an invalid byte count")
+                try:
+                    shell_terminal = ShellTerminalResult.from_mapping(result)
+                except ValueError as exc:
+                    raise TunnelBackendError("PTY agent attached shell returned invalid terminal metadata") from exc
+                terminal.set_result(shell_terminal)
+                completed = True
+            except BaseException as exc:
+                if not terminal.done():
+                    terminal.set_exception(exc)
+                    terminal.exception()
+                raise
+            finally:
+
+                async def finish_attachment() -> None:
+                    if not input_task.done():
+                        input_task.cancel()
+                    await asyncio.gather(input_task, return_exceptions=True)
+                    if not completed:
+                        await _reset_mux_stream(stream, "PTY agent attached shell closed before completion")
+
+                cleanup = asyncio.create_task(
+                    finish_attachment(),
+                    name=f"hostbridge-pty-shell-cleanup-{stream.stream_id}",
+                )
+                await finish_task_before_cancellation(cleanup)
+
+        return ShellAttachment(output(), terminal)
+
+    async def shell_read(self, shell_id: str, *, timeout: float | None, max_bytes: int) -> ShellReadResult:
+        stream: MuxRpcStream | None = None
+        data = bytearray()
+        try:
+            params: dict[str, object] = {"shell_id": shell_id, "max_bytes": max_bytes}
+            if timeout is not None:
+                params["timeout"] = timeout
+            stream = await self._require_client().open_stream(
+                "shell.read",
+                params,
+            )
+            await stream.send_eof()
+            while chunk := await stream.receive():
+                data.extend(chunk)
+            result = await stream.finish()
+        except asyncio.CancelledError:
+            if stream is not None:
+                await _reset_mux_stream(stream, "PTY agent shell read cancelled")
+            raise
+        except Exception as exc:
+            if stream is not None:
+                await _reset_mux_stream(stream, str(exc) or "PTY agent shell read failed")
+            raise TunnelBackendError(f"PTY agent shell read failed: {exc}") from exc
+        alive = result.get("alive")
+        if not isinstance(alive, bool) or result.get("bytes") != len(data):
+            raise TunnelBackendError("PTY agent shell read returned invalid metadata")
+        return ShellReadResult(bytes(data), alive)
+
+    async def shell_close(self, shell_id: str) -> None:
+        stream = await self._require_client().open_stream("shell.close", {"shell_id": shell_id})
+        try:
+            await stream.send_eof()
+            while await stream.receive() is not None:
+                pass
+            result = await stream.finish()
+            if result.get("closed") is not True:
+                raise TunnelBackendError("PTY agent shell close returned invalid metadata")
+        except BaseException:
+            await _reset_mux_stream(stream, "PTY agent shell close failed")
+            raise
+
+    def _require_client(self) -> MuxRpcClient:
+        client = self._client
+        if self._closing or client is None or client.closed:
+            raise TunnelBackendError("PTY agent runtime is not connected")
+        return client
+
+    async def wait_closed(self) -> BaseException | None:
+        if self._closed_future is None:
+            raise TunnelBackendError("PTY agent runtime has not started")
+        return await asyncio.shield(self._closed_future)
+
+    async def close(self) -> None:
+        async with self._close_lock:
+            self._closing = True
+            monitor_is_caller = self._monitor_task is asyncio.current_task()
+            cleanup = asyncio.create_task(self._finish_close(monitor_is_caller))
+            await finish_task_before_cancellation(cleanup)
+
+    def abort(self) -> None:
+        self._closing = True
+        error = TunnelBackendError("PTY agent runtime cleanup timed out")
+        byte_stream = self._byte_stream
+        abort_error: Exception | None = None
+        if byte_stream is not None:
+            try:
+                byte_stream.close_nowait()
+            except Exception as exc:
+                abort_error = exc
+        process = self._process
+        if process is not None and process.returncode is None:
+            try:
+                _signal_process_tree(process, signal.SIGKILL)
+                abort_error = None
+            except Exception as exc:
+                try:
+                    process.kill()
+                    abort_error = None
+                except Exception as fallback_error:
+                    fallback_error.add_note(f"process group kill failed: {exc}")
+                    abort_error = fallback_error
+        if self._closed_future is not None and not self._closed_future.done():
+            with suppress(Exception):
+                self._closed_future.set_result(error)
+        if abort_error is not None:
+            raise abort_error
+
+    async def _finish_close(self, monitor_is_caller: bool) -> None:
+        endpoints = list(self._endpoints.values())
+        await asyncio.gather(
+            *(endpoint.reset("PTY agent runtime closed") for endpoint in endpoints),
+            return_exceptions=True,
+        )
+        for endpoint in endpoints:
+            self._endpoints.pop(endpoint.provider_stream_id, None)
+        client = self._client
+        if client is not None:
+            await client.close()
+            if self._client is client:
+                self._client = None
+        byte_stream = self._byte_stream
+        if byte_stream is not None:
+            await byte_stream.close()
+            if self._byte_stream is byte_stream:
+                self._byte_stream = None
+        await self._terminate_process()
+        monitor = self._monitor_task
+        if monitor is not None and not monitor_is_caller:
+            monitor.cancel()
+            await asyncio.gather(monitor, return_exceptions=True)
+            if self._monitor_task is monitor:
+                self._monitor_task = None
+        if self._closed_future is not None and not self._closed_future.done():
+            self._closed_future.set_result(None)
+
+    async def _bootstrap(self, byte_stream: AsyncPtyByteStream) -> bytes:
+        for step in self.host.login_steps:
+            await byte_stream.read_until(re.compile(step.expect.encode("utf-8")), timeout=self.connect_timeout)
+            if step.send_secret is not None:
+                value = self.host.secrets.get(step.send_secret)
+                if value is None:
+                    raise TunnelBackendError(f"login step references missing secret {step.send_secret!r}")
+            else:
+                value = step.send or ""
+            byte_stream.write(value.encode("utf-8") + b"\r")
+            await byte_stream.drain()
+        if self.host.login_steps:
+            ready = self.host.ready_expect or r"(?m)[^\r\n]*[#$]\s*$"
+            await byte_stream.read_until(re.compile(ready.encode("utf-8")), timeout=self.connect_timeout)
+
+        bundle = self._bundle_builder()
+        nonce = secrets.token_hex(32)
+        digest = hashlib.sha256(bundle).hexdigest()
+        command = (f"python3 -u -c {shlex.quote(f'exec({_BOOTSTRAP!r})')} {nonce} {len(bundle)} {digest}\r").encode(
+            "ascii"
+        )
+        byte_stream.write(command)
+        await byte_stream.drain()
+        ready_marker = b"\x00HBRAW2:" + nonce.encode("ascii") + b"\x00"
+        await byte_stream.read_until(ready_marker, timeout=self.connect_timeout)
+        byte_stream.write(bundle)
+        await byte_stream.drain()
+        verified_marker = b"\x00HBOK2:" + nonce.encode("ascii") + b"\x00"
+        await byte_stream.read_until(verified_marker, timeout=self.connect_timeout)
+        return nonce.encode("ascii")
+
+    async def _monitor_process(self) -> None:
+        process = self._process
+        if process is None:
+            return
+        returncode = await process.wait()
+        self._process_returncode = returncode
+        if self._closing:
+            return
+        error = TunnelBackendError(f"PTY agent process exited with status {returncode}")
+        if self._closed_future is not None and not self._closed_future.done():
+            self._closed_future.set_result(error)
+        client = self._client
+        if client is not None:
+            await client.close()
+
+    async def _terminate_process(self) -> None:
+        process = self._process
+        if process is None:
+            return
+        if process.returncode is None:
+            _signal_process_tree(process, signal.SIGTERM)
+            try:
+                async with asyncio.timeout(PROCESS_TERM_TIMEOUT_SECONDS):
+                    await process.wait()
+            except TimeoutError:
+                _signal_process_tree(process, signal.SIGKILL)
+                await process.wait()
+        else:
+            await process.wait()
+        self._process_returncode = process.returncode
+
+    def _endpoint_closed(self, stream_id: int) -> None:
+        self._endpoints.pop(stream_id, None)
+
+
+def pty_agent_provider_factory(
+    host: HostConfig,
+    provider_name: str,
+    generation: int,
+) -> PtyAgentHostRuntime:
+    if provider_name != "pty_agent":
+        raise TunnelBackendError(f"PTY agent factory cannot create provider {provider_name!r}")
+    return PtyAgentHostRuntime(host, generation=generation)
+
+
+def _signal_process_tree(process: asyncio.subprocess.Process, signal_number: signal.Signals) -> None:
+    try:
+        os.killpg(process.pid, signal_number)
+    except ProcessLookupError:
+        return
+    except PermissionError:
+        with suppress(ProcessLookupError):
+            process.send_signal(signal_number)
+
+
+async def _reset_mux_stream(stream: MuxRpcStream, reason: str) -> None:
+    cleanup = asyncio.create_task(stream.reset(reason), name=f"hostbridge-pty-stream-reset-{stream.stream_id}")
+    while not cleanup.done():
+        try:
+            await asyncio.shield(cleanup)
+        except asyncio.CancelledError:
+            continue
+        except BaseException:
+            break
+    await asyncio.gather(cleanup, return_exceptions=True)
+
+
+def _transfer_result(result: dict[str, object]) -> TransferResult:
+    bytes_transferred = result.get("bytes_transferred")
+    sha256 = result.get("sha256")
+    if (
+        not isinstance(bytes_transferred, int)
+        or isinstance(bytes_transferred, bool)
+        or bytes_transferred < 0
+        or not isinstance(sha256, str)
+        or len(sha256) != 64
+    ):
+        raise TunnelBackendError("PTY agent returned an invalid transfer result")
+    return TransferResult(bytes_transferred, sha256)