@neoline/hostbridge 2.0.4 → 2.0.5

package/src/server_control_mcp/mock_ssh_tunnel.py

@@ -1,309 +1,64 @@
+"""AsyncSSH direct-tcpip adaptation for persistent binary mux streams."""
+
 from __future__ import annotations
 
 import asyncio
-import base64
-import binascii
 import sys
-import uuid
 from collections import deque
+from collections.abc import Awaitable, Callable
 from contextlib import suppress
-from dataclasses import dataclass
-from enum import StrEnum
 from typing import Protocol
 
 import asyncssh
 
-from .client import HostBridgeClient, HostBridgeError
-from .remote_tunnel_agent import build_remote_tunnel_command
-
-TUNNEL_PROTOCOL_VERSION = 1
-UPSTREAM_FRAME_BYTES = 2048
-DOWNSTREAM_FRAME_BYTES = 16 * 1024
-TUNNEL_READ_BYTES = 64 * 1024
-TUNNEL_WRITE_BATCH_BYTES = 32 * 1024
-TUNNEL_HIGH_WATER_BYTES = 128 * 1024
+TUNNEL_HIGH_WATER_BYTES = 256 * 1024
 TUNNEL_LOW_WATER_BYTES = 64 * 1024
-TUNNEL_MAX_LINE_BYTES = 128 * 1024
-
-
-class TunnelProtocolError(ValueError):
-    pass
-
-
-class Direction(StrEnum):
-    UPSTREAM = "U"
-    DOWNSTREAM = "D"
-
-
-class FrameKind(StrEnum):
-    READY = "READY"
-    DATA = "DATA"
-    EOF = "EOF"
-    ERROR = "ERROR"
-    CLOSED = "CLOSED"
-
-
-@dataclass(frozen=True, slots=True)
-class TunnelFrame:
-    direction: Direction
-    sequence: int
-    kind: FrameKind
-    payload: bytes = b""
-    status: int | None = None
-
-
-class TunnelCodec:
-    def __init__(self, prefix: str, *, inbound_direction: Direction) -> None:
-        if not prefix or any(character in prefix for character in "\r\n:"):
-            raise ValueError("tunnel prefix must be non-empty and must not contain CR, LF, or colon")
-        try:
-            prefix_bytes = prefix.encode("ascii")
-        except UnicodeEncodeError as exc:
-            raise ValueError("tunnel prefix must be ASCII") from exc
-        self._prefix = prefix_bytes
-        self._marker = prefix_bytes + b":V"
-        self._inbound_direction = inbound_direction
-        self._next_sequence = 0
-        self._buffer = bytearray()
-
-    def encode(self, frame: TunnelFrame) -> bytes:
-        self._validate_frame(frame)
-        status = b"-" if frame.status is None else str(frame.status).encode("ascii")
-        payload = base64.b64encode(frame.payload)
-        return (
-            b":".join(
-                (
-                    self._prefix,
-                    f"V{TUNNEL_PROTOCOL_VERSION}".encode("ascii"),
-                    frame.direction.value.encode("ascii"),
-                    str(frame.sequence).encode("ascii"),
-                    frame.kind.value.encode("ascii"),
-                    status,
-                    payload,
-                )
-            )
-            + b"\n"
-        )
-
-    def feed(self, data: bytes) -> list[TunnelFrame]:
-        self._buffer.extend(data)
-        frames: list[TunnelFrame] = []
-        while b"\n" in self._buffer:
-            raw_line, _, remainder = self._buffer.partition(b"\n")
-            self._buffer = bytearray(remainder)
-            line = raw_line.rstrip(b"\r")
-            marker = line.find(self._marker)
-            if marker < 0:
-                continue
-            frames.append(self._decode_line(bytes(line[marker:])))
-        if len(self._buffer) > TUNNEL_MAX_LINE_BYTES:
-            raise TunnelProtocolError("tunnel frame exceeds maximum line size")
-        return frames
-
-    def _decode_line(self, line: bytes) -> TunnelFrame:
-        fields = line.split(b":", 6)
-        if len(fields) != 7:
-            raise TunnelProtocolError("malformed tunnel frame")
-        prefix, version, direction_raw, sequence_raw, kind_raw, status_raw, payload_raw = fields
-        if prefix != self._prefix:
-            raise TunnelProtocolError("tunnel frame prefix mismatch")
-        if version != f"V{TUNNEL_PROTOCOL_VERSION}".encode("ascii"):
-            raise TunnelProtocolError("unsupported tunnel protocol version")
-        try:
-            direction = Direction(direction_raw.decode("ascii"))
-        except (UnicodeDecodeError, ValueError) as exc:
-            raise TunnelProtocolError("invalid tunnel frame direction") from exc
-        if direction is not self._inbound_direction:
-            raise TunnelProtocolError("unexpected tunnel frame direction")
-        try:
-            sequence = int(sequence_raw)
-        except ValueError as exc:
-            raise TunnelProtocolError("invalid tunnel frame sequence") from exc
-        if sequence != self._next_sequence:
-            raise TunnelProtocolError(f"tunnel frame sequence expected {self._next_sequence}, received {sequence}")
-        try:
-            kind = FrameKind(kind_raw.decode("ascii"))
-        except (UnicodeDecodeError, ValueError) as exc:
-            raise TunnelProtocolError("invalid tunnel frame kind") from exc
-        if status_raw == b"-":
-            status = None
-        else:
-            try:
-                status = int(status_raw)
-            except ValueError as exc:
-                raise TunnelProtocolError("invalid tunnel frame status") from exc
-        try:
-            payload = base64.b64decode(payload_raw, validate=True)
-        except (binascii.Error, ValueError) as exc:
-            raise TunnelProtocolError("invalid tunnel frame base64 payload") from exc
-        frame = TunnelFrame(direction, sequence, kind, payload, status)
-        self._validate_frame(frame)
-        self._next_sequence += 1
-        return frame
 
-    @staticmethod
-    def _validate_frame(frame: TunnelFrame) -> None:
-        if not isinstance(frame.direction, Direction):
-            raise TunnelProtocolError("invalid tunnel frame direction")
-        if not isinstance(frame.sequence, int) or isinstance(frame.sequence, bool) or frame.sequence < 0:
-            raise TunnelProtocolError("invalid tunnel frame sequence")
-        if not isinstance(frame.kind, FrameKind):
-            raise TunnelProtocolError("invalid tunnel frame kind")
-        if not isinstance(frame.payload, bytes):
-            raise TunnelProtocolError("tunnel frame payload must be bytes")
-        if frame.kind is FrameKind.CLOSED:
-            if not isinstance(frame.status, int) or isinstance(frame.status, bool) or not 0 <= frame.status <= 255:
-                raise TunnelProtocolError("CLOSED tunnel frame requires an exit status")
-        elif frame.status is not None:
-            raise TunnelProtocolError("only CLOSED tunnel frames may contain a status")
-        if frame.kind not in {FrameKind.DATA, FrameKind.ERROR} and frame.payload:
-            raise TunnelProtocolError(f"{frame.kind.value} tunnel frame must not contain a payload")
-
-
-class TunnelState(StrEnum):
-    CONNECTING = "CONNECTING"
-    OPEN = "OPEN"
-    LOCAL_EOF = "LOCAL_EOF"
-    REMOTE_EOF = "REMOTE_EOF"
-    DRAINING = "DRAINING"
-    ABORTING = "ABORTING"
-    CLOSED = "CLOSED"
-
-
-class TunnelLifecycle:
-    def __init__(self) -> None:
-        self.state = TunnelState.CONNECTING
-        self.reusable = False
-
-    def mark_ready(self) -> None:
-        self._require(TunnelState.CONNECTING)
-        self.state = TunnelState.OPEN
-
-    def mark_local_eof(self) -> None:
-        if self.state is TunnelState.OPEN:
-            self.state = TunnelState.LOCAL_EOF
-        elif self.state is TunnelState.REMOTE_EOF:
-            self.state = TunnelState.DRAINING
-        else:
-            self._invalid("local EOF")
-
-    def mark_remote_eof(self) -> None:
-        if self.state is TunnelState.OPEN:
-            self.state = TunnelState.REMOTE_EOF
-        elif self.state is TunnelState.LOCAL_EOF:
-            self.state = TunnelState.DRAINING
-        else:
-            self._invalid("remote EOF")
-
-    def mark_closed(self, status: int) -> None:
-        self._require(TunnelState.DRAINING)
-        if not isinstance(status, int) or isinstance(status, bool) or not 0 <= status <= 255:
-            raise TunnelProtocolError("terminal tunnel status is invalid")
-        self.state = TunnelState.CLOSED
-        self.reusable = status == 0
 
-    def abort(self) -> None:
-        if self.state is not TunnelState.CLOSED:
-            self.state = TunnelState.ABORTING
-            self.reusable = False
+class MuxTunnelStream(Protocol):
+    stream_id: int
 
-    def finish_abort(self) -> None:
-        self._require(TunnelState.ABORTING)
-        self.state = TunnelState.CLOSED
+    async def send(self, data: bytes) -> None: ...
 
-    def _require(self, expected: TunnelState) -> None:
-        if self.state is not expected:
-            raise TunnelProtocolError(f"tunnel state must be {expected.value}, found {self.state.value}")
+    async def send_eof(self) -> None: ...
 
-    def _invalid(self, event: str) -> None:
-        raise TunnelProtocolError(f"cannot accept {event} while tunnel state is {self.state.value}")
+    async def receive(self, max_bytes: int = 64 * 1024) -> bytes | None: ...
 
+    async def finish(self) -> dict[str, object]: ...
 
-class SessionPool(Protocol):
-    def open(self) -> str: ...
+    async def reset(self, message: str = "stream reset") -> None: ...
 
-    def close(self, session_id: str, *, reusable: bool = True) -> None: ...
 
+class MuxTunnelClient(Protocol):
+    @property
+    def closed(self) -> bool: ...
 
-class TunnelLease:
-    def __init__(self, pool: SessionPool, session_id: str) -> None:
-        self._pool = pool
-        self.session_id = session_id
-        self._close_task: asyncio.Task[None] | None = None
+    async def open_stream(self, method: str, params: dict[str, object]) -> MuxTunnelStream: ...
 
-    @classmethod
-    async def acquire(cls, pool: SessionPool) -> TunnelLease:
-        open_task = asyncio.create_task(asyncio.to_thread(pool.open))
-        try:
-            session_id = await asyncio.shield(open_task)
-        except asyncio.CancelledError:
-            await cls._wait_ignoring_cancellation(open_task)
-            if not open_task.cancelled() and open_task.exception() is None:
-                lease = cls(pool, open_task.result())
-                await lease.close(reusable=False)
-            raise
-        return cls(pool, session_id)
+    async def close(self) -> None: ...
 
-    async def close(self, *, reusable: bool) -> None:
-        if self._close_task is None:
-            self._close_task = asyncio.create_task(
-                asyncio.to_thread(self._pool.close, self.session_id, reusable=reusable)
-            )
-        await self._wait_ignoring_cancellation(self._close_task)
-        await self._close_task
 
-    @staticmethod
-    async def _wait_ignoring_cancellation(task: asyncio.Task[object]) -> None:
-        while not task.done():
-            try:
-                await asyncio.shield(task)
-            except asyncio.CancelledError:
-                continue
-
-
-class _TunnelFrameStream:
-    def __init__(self, client: HostBridgeClient, lease: TunnelLease, prefix: str) -> None:
-        self._client = client
-        self._lease = lease
-        self._codec = TunnelCodec(prefix, inbound_direction=Direction.DOWNSTREAM)
-        self._pending: deque[TunnelFrame] = deque()
-
-    async def next(self) -> TunnelFrame:
-        while not self._pending:
-            result = await asyncio.to_thread(
-                self._client.shell_read,
-                self._lease.session_id,
-                timeout=0.2,
-                max_bytes=TUNNEL_READ_BYTES,
-                owner="mock-ssh",
-            )
-            if result.data:
-                try:
-                    self._pending.extend(self._codec.feed(result.data))
-                except TunnelProtocolError as exc:
-                    raise HostBridgeError("protocol_mismatch", str(exc)) from exc
-            if not result.alive and not self._pending:
-                raise HostBridgeError(
-                    "connection_lost",
-                    "remote TCP tunnel closed without a terminal frame",
-                    retryable=True,
-                )
-        return self._pending.popleft()
+MuxClientFactory = Callable[[], Awaitable[MuxTunnelClient]]
 
 
 class TunnelConnector:
     def __init__(
         self,
-        client: HostBridgeClient,
-        session_pool: SessionPool,
+        host_id: str,
+        client_factory: MuxClientFactory,
         *,
-        connect_timeout: int,
+        connect_timeout: float,
     ) -> None:
+        if not isinstance(host_id, str) or not host_id:
+            raise ValueError("host_id must be a non-empty string")
         if connect_timeout <= 0:
             raise ValueError("connect_timeout must be positive")
-        self._client = client
-        self._session_pool = session_pool
-        self._connect_timeout = connect_timeout
+        self._host_id = host_id
+        self._client_factory = client_factory
+        self._connect_timeout = float(connect_timeout)
+        self._client: MuxTunnelClient | None = None
+        self._lock = asyncio.Lock()
+        self._closing = False
 
     @staticmethod
     def validate_destination(dest_host: str, dest_port: int) -> None:
@@ -312,109 +67,103 @@ class TunnelConnector:
         if not isinstance(dest_port, int) or isinstance(dest_port, bool) or not 1 <= dest_port <= 65535:
             raise ValueError("TCP forwarding destination port must be between 1 and 65535")
 
+    async def start(self) -> None:
+        await self._get_client()
+
     async def connect(self, dest_host: str, dest_port: int) -> asyncssh.SSHTCPSession[bytes]:
         self.validate_destination(dest_host, dest_port)
+        client: MuxTunnelClient | None = None
         try:
-            lease = await TunnelLease.acquire(self._session_pool)
-        except asyncio.CancelledError:
-            raise
-        except Exception as exc:
-            raise asyncssh.ChannelOpenError(asyncssh.OPEN_CONNECT_FAILED, str(exc)) from exc
-
-        prefix = f"__HB_TUNNEL_{uuid.uuid4().hex}__"
-        frames = _TunnelFrameStream(self._client, lease, prefix)
-        lifecycle = TunnelLifecycle()
-        try:
-            await asyncio.to_thread(
-                self._client.shell_write,
-                lease.session_id,
-                build_remote_tunnel_command(dest_host, dest_port, prefix),
-                owner="mock-ssh",
-            )
-            first = await asyncio.wait_for(frames.next(), timeout=self._connect_timeout)
-            if first.kind is FrameKind.ERROR:
-                raise HostBridgeError(
-                    "connection_lost",
-                    first.payload.decode("utf-8", errors="replace") or "remote TCP connection failed",
-                    retryable=True,
-                )
-            if first.kind is not FrameKind.READY:
-                raise HostBridgeError(
-                    "protocol_mismatch",
-                    f"remote TCP tunnel returned {first.kind.value} before READY",
+            client = await self._get_client()
+            async with asyncio.timeout(self._connect_timeout):
+                stream = await client.open_stream(
+                    "tcp.connect",
+                    {
+                        "host_id": self._host_id,
+                        "destination_host": dest_host,
+                        "destination_port": dest_port,
+                    },
                 )
-            lifecycle.mark_ready()
         except asyncio.CancelledError:
-            lifecycle.abort()
-            await lease.close(reusable=False)
-            lifecycle.finish_abort()
             raise
-        except Exception as exc:
-            lifecycle.abort()
-            await lease.close(reusable=False)
-            lifecycle.finish_abort()
-            raise asyncssh.ChannelOpenError(asyncssh.OPEN_CONNECT_FAILED, str(exc)) from exc
-        return DirectTCPIPSession(
-            self._client,
-            lease,
-            prefix,
-            frames,
-            lifecycle,
-            dest_host=dest_host,
-            dest_port=dest_port,
-        )
+        except (Exception, TimeoutError) as exc:
+            if client is not None and client.closed:
+                await self._discard_client(client)
+            message = str(exc) or "TCP tunnel stream open timed out"
+            raise asyncssh.ChannelOpenError(asyncssh.OPEN_CONNECT_FAILED, message) from exc
+        return DirectTCPIPSession(stream, dest_host=dest_host, dest_port=dest_port)
+
+    async def close(self) -> None:
+        async with self._lock:
+            self._closing = True
+            client = self._client
+            if client is not None:
+                await client.close()
+                if self._client is client:
+                    self._client = None
+
+    async def _get_client(self) -> MuxTunnelClient:
+        async with self._lock:
+            if self._closing:
+                raise RuntimeError("TCP tunnel connector is closed")
+            if self._client is not None and not self._client.closed:
+                return self._client
+            stale = self._client
+            async with asyncio.timeout(self._connect_timeout):
+                client = await self._client_factory()
+            self._client = client
+        if stale is not None and stale is not client:
+            with suppress(Exception):
+                await stale.close()
+        return client
+
+    async def _discard_client(self, client: MuxTunnelClient) -> None:
+        async with self._lock:
+            if self._client is client:
+                self._client = None
 
 
 _INPUT_EOF = object()
 
 
 class DirectTCPIPSession(asyncssh.SSHTCPSession[bytes]):
-    def __init__(
-        self,
-        client: HostBridgeClient,
-        lease: TunnelLease,
-        prefix: str,
-        frames: _TunnelFrameStream,
-        lifecycle: TunnelLifecycle,
-        *,
-        dest_host: str,
-        dest_port: int,
-    ) -> None:
-        self._client = client
-        self._lease = lease
-        self._codec = TunnelCodec(prefix, inbound_direction=Direction.DOWNSTREAM)
-        self._frames = frames
-        self._lifecycle = lifecycle
+    def __init__(self, stream: MuxTunnelStream, *, dest_host: str, dest_port: int) -> None:
+        self._stream = stream
         self._dest_host = dest_host
         self._dest_port = dest_port
         self._channel: asyncssh.SSHTCPChannel[bytes] | None = None
-        self._input: asyncio.Queue[bytes | object] = asyncio.Queue()
+        self._input: deque[bytes | object] = deque()
+        self._input_ready = asyncio.Event()
         self._input_buffered_bytes = 0
         self._input_paused = False
-        self._upstream_sequence = 0
+        self._input_eof = False
         self._write_ready = asyncio.Event()
         self._write_ready.set()
         self._runner: asyncio.Task[None] | None = None
-        self._orphan_cleanup: asyncio.Task[None] | None = None
+        self._cleanup_task: asyncio.Task[None] | None = None
         self._finished = asyncio.Event()
-
-    @property
-    def state(self) -> TunnelState:
-        return self._lifecycle.state
+        self._aborted = False
+        self._reset_sent = False
 
     def connection_made(self, chan: asyncssh.SSHTCPChannel[bytes]) -> None:
         self._channel = chan
 
     def session_started(self) -> None:
-        self._runner = asyncio.create_task(self._run())
+        if self._runner is not None or self._finished.is_set():
+            return
+        self._runner = asyncio.create_task(
+            self._run(),
+            name=f"hostbridge-mock-ssh-tcp-{self._stream.stream_id}",
+        )
         self._runner.add_done_callback(self._runner_done)
 
     def data_received(self, data: bytes, datatype: asyncssh.DataType) -> None:  # noqa: ARG002
-        if not data:
+        if not data or self._input_eof or self._finished.is_set():
             return
         payload = bytes(data)
-        self._input.put_nowait(payload)
+        self._input.append(payload)
         self._input_buffered_bytes += len(payload)
+        self._input_ready.set()
         if (
             self._input_buffered_bytes >= TUNNEL_HIGH_WATER_BYTES
             and not self._input_paused
@@ -424,7 +173,10 @@ class DirectTCPIPSession(asyncssh.SSHTCPSession[bytes]):
             self._input_paused = True
 
     def eof_received(self) -> bool:
-        self._input.put_nowait(_INPUT_EOF)
+        if not self._input_eof and not self._finished.is_set():
+            self._input_eof = True
+            self._input.append(_INPUT_EOF)
+            self._input_ready.set()
         return True
 
     def connection_lost(self, exc: Exception | None) -> None:  # noqa: ARG002
@@ -437,41 +189,36 @@ class DirectTCPIPSession(asyncssh.SSHTCPSession[bytes]):
         self._write_ready.set()
 
     def abort(self) -> None:
-        self._lifecycle.abort()
+        if self._finished.is_set():
+            return
+        self._aborted = True
         if self._runner is not None and self._runner is not asyncio.current_task() and not self._runner.done():
             self._runner.cancel()
-        elif self._runner is None:
-            self._schedule_orphan_cleanup()
+        elif self._runner is None and self._cleanup_task is None:
+            self._cleanup_task = asyncio.create_task(self._finish_orphaned_session())
 
     async def wait_finished(self) -> None:
         await self._finished.wait()
 
     def _runner_done(self, task: asyncio.Task[None]) -> None:  # noqa: ARG002
-        if not self._finished.is_set():
-            self._schedule_orphan_cleanup()
-
-    def _schedule_orphan_cleanup(self) -> None:
-        if self._orphan_cleanup is None:
-            self._orphan_cleanup = asyncio.create_task(self._finish_orphaned_session())
+        if not self._finished.is_set() and self._cleanup_task is None:
+            self._cleanup_task = asyncio.create_task(self._finish_orphaned_session())
 
     async def _finish_orphaned_session(self) -> None:
-        self._lifecycle.abort()
-        await self._lease.close(reusable=False)
-        if self._lifecycle.state is TunnelState.ABORTING:
-            self._lifecycle.finish_abort()
-        self._close_channel()
-        self._finished.set()
+        await self._reset_stream()
+        self._finish_channel()
 
     async def _run(self) -> None:
         input_task = asyncio.create_task(self._pump_input())
         output_task = asyncio.create_task(self._pump_output())
+        clean = False
         try:
             await asyncio.gather(input_task, output_task)
+            await self._stream.finish()
+            clean = True
         except asyncio.CancelledError:
-            self._lifecycle.abort()
             raise
         except Exception as exc:
-            self._lifecycle.abort()
             print(
                 f"hostbridge mock-ssh TCP forwarding to {self._dest_host}:{self._dest_port} closed: {exc}",
                 file=sys.stderr,
@@ -482,111 +229,61 @@ class DirectTCPIPSession(asyncssh.SSHTCPSession[bytes]):
                 if not task.done():
                     task.cancel()
             await asyncio.gather(input_task, output_task, return_exceptions=True)
-            await self._lease.close(reusable=self._lifecycle.reusable)
-            if self._lifecycle.state is TunnelState.ABORTING:
-                self._lifecycle.finish_abort()
-            self._close_channel()
-            self._finished.set()
+            if not clean:
+                await self._reset_stream()
+            self._finish_channel()
 
     async def _pump_input(self) -> None:
         while True:
-            item = await self._input.get()
+            item = await self._next_input()
             if item is _INPUT_EOF:
-                await self._send_eof()
+                await self._stream.send_eof()
                 return
-
             if not isinstance(item, bytes):
-                raise RuntimeError("TCP tunnel input queue contained an invalid item")
-            batch = bytearray(item)
-            saw_eof = False
-            while not self._input.empty():
-                queued = self._input.get_nowait()
-                if queued is _INPUT_EOF:
-                    saw_eof = True
-                    break
-                if not isinstance(queued, bytes):
-                    raise RuntimeError("TCP tunnel input queue contained an invalid item")
-                batch.extend(queued)
+                raise TypeError("tunnel input queue contains a non-bytes item")
             try:
-                for offset in range(0, len(batch), TUNNEL_WRITE_BATCH_BYTES):
-                    await self._send_data(bytes(batch[offset : offset + TUNNEL_WRITE_BATCH_BYTES]))
+                await self._stream.send(item)
             finally:
-                self._input_buffered_bytes -= len(batch)
+                self._input_buffered_bytes -= len(item)
                 self._resume_input_if_ready()
-            if saw_eof:
-                await self._send_eof()
-                return
 
-    async def _send_data(self, payload: bytes) -> None:
-        frames = []
-        for offset in range(0, len(payload), UPSTREAM_FRAME_BYTES):
-            frames.append(
-                TunnelFrame(
-                    Direction.UPSTREAM,
-                    self._upstream_sequence,
-                    FrameKind.DATA,
-                    payload=payload[offset : offset + UPSTREAM_FRAME_BYTES],
-                )
-            )
-            self._upstream_sequence += 1
-        await self._write_frames(frames)
-
-    async def _send_eof(self) -> None:
-        frame = TunnelFrame(Direction.UPSTREAM, self._upstream_sequence, FrameKind.EOF)
-        self._upstream_sequence += 1
-        await self._write_frames([frame])
-        self._lifecycle.mark_local_eof()
-
-    async def _write_frames(self, frames: list[TunnelFrame]) -> None:
-        payload = b"".join(self._codec.encode(frame) for frame in frames)
-        await asyncio.to_thread(
-            self._client.shell_write,
-            self._lease.session_id,
-            payload,
-            owner="mock-ssh",
-        )
+    async def _next_input(self) -> bytes | object:
+        while not self._input:
+            self._input_ready.clear()
+            if not self._input:
+                await self._input_ready.wait()
+        item = self._input.popleft()
+        if not self._input:
+            self._input_ready.clear()
+        return item
 
     async def _pump_output(self) -> None:
         if self._channel is None:
             raise RuntimeError("TCP tunnel channel was not initialized")
         while True:
-            frame = await self._frames.next()
-            if frame.kind is FrameKind.DATA:
-                if self._lifecycle.state not in {TunnelState.OPEN, TunnelState.LOCAL_EOF}:
-                    raise HostBridgeError("protocol_mismatch", "remote TCP data arrived after EOF")
-                await self._write_ready.wait()
-                self._channel.write(frame.payload)
-            elif frame.kind is FrameKind.EOF:
-                self._lifecycle.mark_remote_eof()
+            await self._write_ready.wait()
+            chunk = await self._stream.receive()
+            if chunk is None:
                 self._channel.write_eof()
-            elif frame.kind is FrameKind.ERROR:
-                raise HostBridgeError(
-                    "connection_lost",
-                    frame.payload.decode("utf-8", errors="replace") or "remote TCP tunnel failed",
-                    retryable=True,
-                )
-            elif frame.kind is FrameKind.CLOSED:
-                assert frame.status is not None
-                self._lifecycle.mark_closed(frame.status)
-                if frame.status != 0:
-                    raise HostBridgeError(
-                        "connection_lost",
-                        f"remote TCP tunnel exited with status {frame.status}",
-                        retryable=True,
-                    )
                 return
-            else:
-                raise HostBridgeError(
-                    "protocol_mismatch",
-                    f"unexpected {frame.kind.value} frame after READY",
-                )
+            self._channel.write(chunk)
+
+    async def _reset_stream(self) -> None:
+        if self._reset_sent:
+            return
+        self._reset_sent = True
+        with suppress(Exception):
+            await self._stream.reset("mock-ssh TCP channel closed")
 
     def _resume_input_if_ready(self) -> None:
         if self._input_paused and self._input_buffered_bytes <= TUNNEL_LOW_WATER_BYTES and self._channel is not None:
             self._channel.resume_reading()
             self._input_paused = False
 
-    def _close_channel(self) -> None:
+    def _finish_channel(self) -> None:
+        if self._finished.is_set():
+            return
+        self._finished.set()
         if self._channel is not None:
             with suppress(Exception):
                 self._channel.close()