@neoline/hostbridge 2.0.4 → 2.0.5

package/src/server_control_mcp/doctor.py

@@ -10,6 +10,7 @@ from typing import Protocol
 
 from . import __version__
 from .config import load_v1_config
+from .mux_protocol import MUX_PROTOCOL_VERSION
 from .runtime import runtime_paths
 
 
@@ -17,6 +18,9 @@ class DoctorClient(Protocol):
     def hello(self) -> dict[str, object]:
         raise NotImplementedError
 
+    def tunnel_health(self) -> dict[str, object]:
+        raise NotImplementedError
+
 
 @dataclass(frozen=True, slots=True)
 class DiagnosticResult:
@@ -46,6 +50,7 @@ def collect_diagnostics(
         _config_diagnostic(config_path),
         _runtime_diagnostic(),
         _daemon_diagnostic(client),
+        _tunnel_diagnostic(client),
         _executable_diagnostic(executable_paths),
     ]
     return DoctorReport(tuple(results))
@@ -97,15 +102,42 @@ def _daemon_diagnostic(client: DoctorClient) -> DiagnosticResult:
         hello = client.hello()
     except Exception as exc:
         return DiagnosticResult("daemon", "fail", f"daemon unavailable: {exc}", {})
-    valid = hello.get("service") == "hostbridge" and hello.get("protocol") == 1 and hello.get("version") == __version__
+    valid = (
+        hello.get("service") == "hostbridge"
+        and hello.get("protocol") == MUX_PROTOCOL_VERSION
+        and hello.get("version") == __version__
+    )
     return DiagnosticResult(
         "daemon",
         "pass" if valid else "fail",
-        "daemon protocol v1 is healthy" if valid else "daemon returned an incompatible handshake",
+        f"daemon protocol v{MUX_PROTOCOL_VERSION} is healthy" if valid else "daemon returned an incompatible handshake",
         dict(hello),
     )
 
 
+def _tunnel_diagnostic(client: DoctorClient) -> DiagnosticResult:
+    try:
+        result = client.tunnel_health()
+    except Exception as exc:
+        return DiagnosticResult("tunnel", "fail", f"tunnel diagnostics unavailable: {exc}", {})
+    hosts = result.get("hosts")
+    if not isinstance(hosts, list) or any(not isinstance(host, dict) for host in hosts):
+        return DiagnosticResult("tunnel", "fail", "daemon returned invalid tunnel diagnostics", dict(result))
+    enabled = [host for host in hosts if host.get("enabled") is True]
+    failed = [host for host in enabled if host.get("state") == "failed"]
+    inactive = [host for host in enabled if host.get("state") != "ready"]
+    if failed:
+        status = "fail"
+        message = f"{len(failed)} tunnel provider runtime(s) failed"
+    elif inactive:
+        status = "warn"
+        message = f"{len(inactive)} tunnel provider runtime(s) configured but not handshaken"
+    else:
+        status = "pass"
+        message = "tunnel provider runtimes are ready"
+    return DiagnosticResult("tunnel", status, message, dict(result))
+
+
 def _executable_diagnostic(paths: list[str]) -> DiagnosticResult:
     if not paths:
         return DiagnosticResult("executable", "fail", "hostbridge executable was not found", {"paths": []})

package/src/server_control_mcp/hosts.py

@@ -4,6 +4,7 @@ import os
 import re
 from collections.abc import Iterable, Sequence
 from dataclasses import dataclass, field
+from ipaddress import ip_address, ip_network
 from pathlib import Path
 
 from .config import load_v1_config
@@ -42,6 +43,17 @@ class SshHostConfig:
     extra_args: tuple[str, ...] = ()
 
 
+@dataclass(frozen=True, slots=True)
+class ForwardingPolicy:
+    enabled: bool = False
+    allow_cidrs: tuple[str, ...] = ()
+    allow_hostnames: tuple[str, ...] = ()
+    allow_ports: tuple[int, ...] = ()
+    max_streams: int = 128
+    idle_timeout_seconds: float = 300
+    stream_deadline_seconds: float = 24 * 60 * 60
+
+
 @dataclass(frozen=True, slots=True)
 class HostConfig:
     """A single approved connection target."""
@@ -58,6 +70,7 @@ class HostConfig:
     transport: str = "auto"
     limits: HostLimits = field(default_factory=HostLimits)
     ssh: SshHostConfig | None = None
+    forwarding: ForwardingPolicy = field(default_factory=ForwardingPolicy)
 
     @classmethod
     def script(
@@ -68,10 +81,19 @@ class HostConfig:
         *,
         transport: str = "auto",
         limits: HostLimits | None = None,
+        forwarding: ForwardingPolicy | None = None,
     ) -> HostConfig:
         path = Path(script).expanduser()
         return cls(
-            host_id, label, "/bin/sh", [str(path)], "script", path, transport=transport, limits=limits or HostLimits()
+            host_id,
+            label,
+            "/bin/sh",
+            [str(path)],
+            "script",
+            path,
+            transport=transport,
+            limits=limits or HostLimits(),
+            forwarding=forwarding or ForwardingPolicy(),
         )
 
     @classmethod
@@ -89,6 +111,7 @@ class HostConfig:
         transport: str = "auto",
         limits: HostLimits | None = None,
         ssh: SshHostConfig | None = None,
+        forwarding: ForwardingPolicy | None = None,
     ) -> HostConfig:
         return cls(
             host_id,
@@ -103,6 +126,7 @@ class HostConfig:
             transport,
             limits or HostLimits(),
             ssh,
+            forwarding or ForwardingPolicy(),
         )
 
     def describe(self) -> dict[str, object]:
@@ -113,6 +137,12 @@ class HostConfig:
             "command": self.command,
             "args": list(self.args),
             "transport": self.transport,
+            "forwarding": {
+                "enabled": self.forwarding.enabled,
+                "max_streams": self.forwarding.max_streams,
+                "idle_timeout_seconds": self.forwarding.idle_timeout_seconds,
+                "stream_deadline_seconds": self.forwarding.stream_deadline_seconds,
+            },
         }
         if self.script_path is not None:
             info.update(
@@ -181,12 +211,20 @@ def _parse_host(item: object, path: Path, index: int) -> HostConfig:
     if transport == "ssh" and "ssh" not in item:
         raise ValueError(f"{path} hosts[{index}].transport ssh requires an ssh connection")
     limits = _parse_limits(item.get("limits"), path, index)
+    forwarding = _parse_forwarding(item.get("forwarding"), path, index)
 
     if "script" in item:
         script = str(item.get("script", "")).strip()
         if not script:
             raise ValueError(f"{path} hosts[{index}].script must not be empty")
-        host = HostConfig.script(host_id, label, script, transport=transport, limits=limits)
+        host = HostConfig.script(
+            host_id,
+            label,
+            script,
+            transport=transport,
+            limits=limits,
+            forwarding=forwarding,
+        )
         return HostConfig(
             host.id,
             host.label,
@@ -199,6 +237,7 @@ def _parse_host(item: object, path: Path, index: int) -> HostConfig:
             ready_expect,
             transport,
             limits,
+            forwarding=forwarding,
         )
 
     if "ssh" in item:
@@ -219,6 +258,7 @@ def _parse_host(item: object, path: Path, index: int) -> HostConfig:
             transport=transport,
             limits=limits,
             ssh=ssh_config,
+            forwarding=forwarding,
         )
 
     command = str(item.get("command", "")).strip()
@@ -237,6 +277,100 @@ def _parse_host(item: object, path: Path, index: int) -> HostConfig:
         ready_expect=ready_expect,
         transport=transport,
         limits=limits,
+        forwarding=forwarding,
+    )
+
+
+def _parse_forwarding(raw: object, path: Path, index: int) -> ForwardingPolicy:
+    if raw is None:
+        return ForwardingPolicy()
+    if not isinstance(raw, dict):
+        raise ValueError(f"{path} hosts[{index}].forwarding must be an object")
+    allowed_fields = {
+        "enabled",
+        "allow_cidrs",
+        "allow_hostnames",
+        "allow_ports",
+        "max_streams",
+        "idle_timeout_seconds",
+        "stream_deadline_seconds",
+    }
+    unknown = sorted(set(raw) - allowed_fields)
+    if unknown:
+        raise ValueError(f"{path} hosts[{index}].forwarding contains unknown field {unknown[0]!r}")
+
+    enabled = raw.get("enabled", False)
+    if not isinstance(enabled, bool):
+        raise ValueError(f"{path} hosts[{index}].forwarding.enabled must be a boolean")
+
+    raw_cidrs = raw.get("allow_cidrs", [])
+    if not isinstance(raw_cidrs, list) or not all(isinstance(value, str) for value in raw_cidrs):
+        raise ValueError(f"{path} hosts[{index}].forwarding.allow_cidrs must be a list of CIDR strings")
+    cidrs: list[str] = []
+    for value in raw_cidrs:
+        try:
+            network = ip_network(value, strict=True)
+        except ValueError as exc:
+            raise ValueError(f"{path} hosts[{index}].forwarding contains invalid CIDR {value!r}") from exc
+        cidrs.append(str(network))
+
+    raw_hostnames = raw.get("allow_hostnames", [])
+    if not isinstance(raw_hostnames, list) or not all(isinstance(value, str) for value in raw_hostnames):
+        raise ValueError(f"{path} hosts[{index}].forwarding.allow_hostnames must be a list of hostnames")
+    hostnames: list[str] = []
+    for value in raw_hostnames:
+        hostname = value.strip().rstrip(".").lower()
+        if not _valid_exact_hostname(hostname):
+            raise ValueError(f"{path} hosts[{index}].forwarding contains invalid hostname {value!r}")
+        hostnames.append(hostname)
+
+    raw_ports = raw.get("allow_ports", [])
+    if not isinstance(raw_ports, list) or any(
+        not isinstance(port, int) or isinstance(port, bool) or not 1 <= port <= 65535 for port in raw_ports
+    ):
+        raise ValueError(f"{path} hosts[{index}].forwarding.allow_ports must contain valid TCP ports")
+
+    max_streams = raw.get("max_streams", 128)
+    if not isinstance(max_streams, int) or isinstance(max_streams, bool) or not 1 <= max_streams <= 4096:
+        raise ValueError(f"{path} hosts[{index}].forwarding.max_streams must be between 1 and 4096")
+    idle_timeout = raw.get("idle_timeout_seconds", 300)
+    if not isinstance(idle_timeout, (int, float)) or isinstance(idle_timeout, bool) or not 0 < idle_timeout <= 86400:
+        raise ValueError(f"{path} hosts[{index}].forwarding.idle_timeout_seconds must be between 0 and 86400")
+    stream_deadline = raw.get("stream_deadline_seconds", 24 * 60 * 60)
+    if (
+        not isinstance(stream_deadline, (int, float))
+        or isinstance(stream_deadline, bool)
+        or not 0 < stream_deadline <= 7 * 24 * 60 * 60
+    ):
+        raise ValueError(f"{path} hosts[{index}].forwarding.stream_deadline_seconds must be between 0 and 604800")
+
+    return ForwardingPolicy(
+        enabled=enabled,
+        allow_cidrs=tuple(dict.fromkeys(cidrs)),
+        allow_hostnames=tuple(dict.fromkeys(hostnames)),
+        allow_ports=tuple(dict.fromkeys(raw_ports)),
+        max_streams=max_streams,
+        idle_timeout_seconds=float(idle_timeout),
+        stream_deadline_seconds=float(stream_deadline),
+    )
+
+
+def _valid_exact_hostname(hostname: str) -> bool:
+    if not hostname or len(hostname) > 253 or "*" in hostname or "\x00" in hostname:
+        return False
+    try:
+        ip_address(hostname)
+    except ValueError:
+        pass
+    else:
+        return False
+    labels = hostname.split(".")
+    return all(
+        1 <= len(label) <= 63
+        and label[0].isalnum()
+        and label[-1].isalnum()
+        and all(character.isalnum() or character == "-" for character in label)
+        for label in labels
     )
 
 

package/src/server_control_mcp/mock_ssh.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import sys
 from collections.abc import Awaitable
 from pathlib import Path
 from typing import Any, cast
@@ -12,11 +13,14 @@ from .mock_ssh_exec import HostBridgeSSHProcess
 from .mock_ssh_session import HostBridgeSessionFactory
 from .mock_ssh_sftp import HostBridgeSFTPServer
 from .mock_ssh_tunnel import TunnelConnector
+from .mux_rpc import MuxRpcClient
 from .secure_files import create_private_file, ensure_private_directory, replace_private_file, validate_private_file
 
 DEFAULT_PORT_RANGE = range(2222, 2300)
 DEFAULT_COMMAND_TIMEOUT = 60
 DEFAULT_MAX_TRANSFER_BYTES = 1024 * 1024 * 1024
+DEFAULT_SFTP_MAX_OPEN_HANDLES = 16
+DEFAULT_SFTP_MAX_TEMP_BYTES = 1024 * 1024 * 1024
 
 
 class HostBridgeSSHServer(asyncssh.SSHServer):
@@ -225,7 +229,7 @@ async def serve_mock_ssh(
     if config_path is not None:
         raise HostBridgeError(
             "invalid_request",
-            "mock-ssh no longer accepts a separate config path; start the v1 daemon with the intended config",
+            "mock-ssh uses the daemon's config; restart the daemon with the intended config",
         )
     key_dir = key_dir or default_key_dir(host_id)
     ssh_host_alias = ssh_host_alias or default_ssh_host_alias(host_id)
@@ -233,16 +237,30 @@ async def serve_mock_ssh(
     client = HostBridgeClient()
     client.hello()
     session_factory = HostBridgeSessionFactory(client, host_id, connect_timeout)
-    tunnel_connector = TunnelConnector(client, session_factory, connect_timeout=connect_timeout)
-    server, selected_port = await create_server_on_available_port(
-        server_factory=lambda: HostBridgeSSHServer(public_key, tunnel_connector),
-        listen_host=listen_host,
-        port=port,
-        server_host_keys=[str(host_key_path)],
-        process_factory=HostBridgeSSHProcess(client, session_factory, command_timeout, output_limit),
-        sftp_factory=lambda chan: HostBridgeSFTPServer(chan, client, session_factory, command_timeout, output_limit),
-    )
+
+    async def connect_mux_client() -> MuxRpcClient:
+        return await MuxRpcClient.connect_unix(client.socket_path)
+
+    tunnel_connector = TunnelConnector(host_id, connect_mux_client, connect_timeout=connect_timeout)
+    server: asyncssh.SSHAcceptor | None = None
     try:
+        await tunnel_connector.start()
+        server, selected_port = await create_server_on_available_port(
+            server_factory=lambda: HostBridgeSSHServer(public_key, tunnel_connector),
+            listen_host=listen_host,
+            port=port,
+            server_host_keys=[str(host_key_path)],
+            process_factory=HostBridgeSSHProcess(client, session_factory, command_timeout, output_limit),
+            sftp_factory=lambda chan: HostBridgeSFTPServer(
+                chan,
+                client,
+                session_factory,
+                command_timeout,
+                output_limit,
+                DEFAULT_SFTP_MAX_OPEN_HANDLES,
+                DEFAULT_SFTP_MAX_TEMP_BYTES,
+            ),
+        )
         installed_config_path: Path | None = None
         if install_ssh_config:
             installed_config_path = install_ssh_config_block(
@@ -262,11 +280,31 @@ async def serve_mock_ssh(
         await server.wait_closed()
         return 0
     finally:
+        primary_error = sys.exception()
+        cleanup_errors: list[BaseException] = []
+        if server is not None:
+            try:
+                server.close()
+                await server.wait_closed()
+            except BaseException as exc:
+                cleanup_errors.append(exc)
+        try:
+            await tunnel_connector.close()
+        except BaseException as exc:
+            cleanup_errors.append(exc)
+        try:
+            await asyncio.to_thread(session_factory.shutdown)
+        except BaseException as exc:
+            cleanup_errors.append(exc)
         try:
-            server.close()
-            await server.wait_closed()
-        finally:
-            session_factory.shutdown()
+            await asyncio.to_thread(client.close)
+        except BaseException as exc:
+            cleanup_errors.append(exc)
+        if cleanup_errors:
+            if primary_error is None:
+                raise cleanup_errors[0]
+            for error in cleanup_errors:
+                primary_error.add_note(f"mock-ssh cleanup failed: {error!r}")
 
 
 def run_mock_ssh(**kwargs: object) -> int:

package/src/server_control_mcp/mock_ssh_exec.py

@@ -5,12 +5,12 @@ import queue
 import sys
 import threading
 from contextlib import suppress
-from functools import partial
 from typing import cast
 
 import asyncssh
 
-from .client import HostBridgeClient, HostBridgeError
+from .async_lifecycle import finish_task_despite_cancellation
+from .client import ExecResult, HostBridgeClient, HostBridgeError
 from .mock_ssh_session import HostBridgeSessionFactory
 
 
@@ -36,12 +36,17 @@ class HostBridgeSSHProcess:
         except asyncio.CancelledError:
             session_id = None
             with suppress(BaseException):
-                session_id = await open_task
+                session_id = await finish_task_despite_cancellation(open_task)
             if session_id is not None:
-                await asyncio.to_thread(self._session_factory.close, session_id, reusable=False)
+                try:
+                    await self._close_session(session_id, reusable=False)
+                except Exception as exc:
+                    self._report_shell_error(process, exc)
+            self._set_exit_status(process, 255)
+            return
+        if session_id is None:
             self._set_exit_status(process, 255)
             return
-        assert session_id is not None
         reusable = False
         try:
             reusable = await self._run(process, session_id)
@@ -49,7 +54,15 @@ class HostBridgeSSHProcess:
             self._set_exit_status(process, 255)
             return
         finally:
-            await asyncio.to_thread(self._session_factory.close, session_id, reusable=reusable)
+            try:
+                await self._close_session(session_id, reusable=reusable)
+            except Exception as exc:
+                self._report_shell_error(process, exc)
+                self._set_exit_status(process, 255)
+
+    async def _close_session(self, session_id: str, *, reusable: bool) -> None:
+        close_task = asyncio.create_task(asyncio.to_thread(self._session_factory.close, session_id, reusable=reusable))
+        await finish_task_despite_cancellation(close_task)
 
     async def _run(self, process: asyncssh.SSHServerProcess[bytes], session_id: str) -> bool:
         command = process.command
@@ -113,40 +126,84 @@ class HostBridgeSSHProcess:
         )
         staging_deadline = loop.time() + self._timeout
 
-        async def put(item: bytes | BaseException | object) -> None:
+        async def put(item: bytes | BaseException | object) -> bool:
             while True:
+                if consumer.done():
+                    await consumer
+                    return False
                 try:
                     items.put_nowait(item)
-                    return
+                    return True
                 except queue.Full:
-                    if consumer.done():
-                        await consumer
                     await asyncio.sleep(0.005)
 
         try:
             while not process.stdin.at_eof():
                 remaining = staging_deadline - loop.time()
                 if remaining <= 0:
-                    await put(HostBridgeError("timeout", "timed out waiting for mock-ssh exec stdin"))
+                    if not await put(HostBridgeError("timeout", "timed out waiting for mock-ssh exec stdin")):
+                        return await consumer
                     break
+                read_task = asyncio.create_task(process.stdin.read(64 * 1024))
                 try:
-                    data = await asyncio.wait_for(process.stdin.read(64 * 1024), timeout=remaining)
-                except TimeoutError:
-                    await put(HostBridgeError("timeout", "timed out waiting for mock-ssh exec stdin"))
+                    completed, _ = await asyncio.wait(
+                        (read_task, consumer),
+                        timeout=remaining,
+                        return_when=asyncio.FIRST_COMPLETED,
+                    )
+                except BaseException:
+                    read_task.cancel()
+                    await asyncio.gather(read_task, return_exceptions=True)
+                    raise
+                if consumer in completed:
+                    read_task.cancel()
+                    await asyncio.gather(read_task, return_exceptions=True)
+                    return await consumer
+                if read_task not in completed:
+                    read_task.cancel()
+                    await asyncio.gather(read_task, return_exceptions=True)
+                    if not await put(HostBridgeError("timeout", "timed out waiting for mock-ssh exec stdin")):
+                        return await consumer
                     break
+                try:
+                    data = read_task.result()
+                except asyncio.CancelledError:
+                    raise
                 except BaseException as exc:
-                    await put(exc)
+                    if not await put(exc):
+                        return await consumer
                     break
                 if not data:
                     break
-                await put(_to_bytes(data))
-            await put(end)
-            return await consumer
+                if not await put(_to_bytes(data)):
+                    return await consumer
+            if not await put(end):
+                return await consumer
+            return await asyncio.shield(consumer)
         except asyncio.CancelledError:
             cancelled.set()
-            await put(end)
-            with suppress(BaseException):
-                await asyncio.shield(consumer)
+            close_succeeded = False
+
+            def observe_consumer(future: asyncio.Future[ExecResult]) -> None:
+                try:
+                    future.result()
+                except asyncio.CancelledError:
+                    return
+                except Exception as exc:
+                    self._report_shell_error(process, exc)
+
+            try:
+                await self._close_session(session_id, reusable=False)
+                close_succeeded = True
+            except Exception as exc:
+                self._report_shell_error(process, exc)
+            with suppress(queue.Full):
+                items.put_nowait(end)
+            if close_succeeded:
+                with suppress(BaseException):
+                    await finish_task_despite_cancellation(consumer)
+            else:
+                consumer.add_done_callback(observe_consumer)
             raise
 
     @staticmethod
@@ -155,56 +212,88 @@ class HostBridgeSSHProcess:
             process.exit(status)
 
     async def _run_shell(self, process: asyncssh.SSHServerProcess[bytes], session_id: str) -> None:
-        loop = asyncio.get_running_loop()
         stop = asyncio.Event()
-        input_done = asyncio.Event()
+        failed = False
+        terminal_failed = False
+        exit_code: int | None = None
+        exit_signal: str | None = None
+        open_task = asyncio.create_task(asyncio.to_thread(self._client.open_shell, session_id, owner="mock-ssh"))
+        try:
+            shell = await asyncio.shield(open_task)
+        except asyncio.CancelledError:
+            with suppress(BaseException):
+                shell = await finish_task_despite_cancellation(open_task)
+                close_task = asyncio.create_task(asyncio.to_thread(shell.close))
+                await finish_task_despite_cancellation(close_task)
+            raise
 
         async def pump_input() -> None:
+            nonlocal failed
             try:
                 while not process.stdin.at_eof():
                     data = await process.stdin.read(4096)
                     if not data:
                         break
-                    await loop.run_in_executor(
-                        None,
-                        partial(self._client.shell_write, session_id, _to_bytes(data), owner="mock-ssh"),
-                    )
+                    await asyncio.to_thread(shell.send, _to_bytes(data))
+            except Exception as exc:
+                failed = True
+                self._report_shell_error(process, exc)
             finally:
-                input_done.set()
+                try:
+                    await asyncio.to_thread(shell.send_eof)
+                except Exception as exc:
+                    failed = True
+                    self._report_shell_error(process, exc)
+                    stop.set()
 
         async def pump_output() -> None:
+            nonlocal exit_code, exit_signal, failed, terminal_failed
             try:
-                idle_after_input = 0
-                while not stop.is_set():
-                    result = await loop.run_in_executor(
-                        None,
-                        lambda: self._client.shell_read(
-                            session_id,
-                            timeout=0.2,
-                            max_bytes=4096,
-                            owner="mock-ssh",
-                        ),
-                    )
-                    if result.data:
-                        idle_after_input = 0
-                        process.stdout.write(result.data)
-                    elif input_done.is_set():
-                        idle_after_input += 1
-                        if idle_after_input >= 5:
-                            stop.set()
-                            break
-                    if not result.alive:
-                        stop.set()
-                        break
+                while data := await asyncio.to_thread(shell.receive):
+                    process.stdout.write(data)
+                exit_code = shell.exit_code
+                exit_signal = shell.exit_signal
             except Exception as exc:
-                print(f"hostbridge mock-ssh shell failed: {exc}", file=sys.stderr, flush=True)
-                with suppress(Exception):
-                    process.stderr.write(f"hostbridge mock-ssh shell error: {exc}\n".encode("utf-8", errors="replace"))
+                failed = True
+                terminal_failed = True
+                self._report_shell_error(process, exc)
+            finally:
                 stop.set()
 
         tasks = [asyncio.create_task(pump_input()), asyncio.create_task(pump_output())]
-        await stop.wait()
-        for task in tasks:
-            task.cancel()
-        await asyncio.gather(*tasks, return_exceptions=True)
-        self._set_exit_status(process, 0)
+        try:
+            await stop.wait()
+        finally:
+            try:
+                close_task = asyncio.create_task(asyncio.to_thread(shell.close))
+                await finish_task_despite_cancellation(close_task)
+            except Exception as exc:
+                failed = True
+                self._report_shell_error(process, exc)
+            finally:
+                for task in tasks:
+                    task.cancel()
+                await asyncio.gather(*tasks, return_exceptions=True)
+        if terminal_failed:
+            self._set_exit_status(process, 255)
+        elif failed:
+            self._set_exit_status(process, 1)
+        elif isinstance(exit_code, int) and not isinstance(exit_code, bool) and 0 <= exit_code <= 255:
+            self._set_exit_status(process, exit_code)
+        elif isinstance(exit_signal, str) and exit_signal.strip():
+            self._set_exit_signal(process, exit_signal)
+        else:
+            self._set_exit_status(process, 255)
+
+    @staticmethod
+    def _report_shell_error(process: asyncssh.SSHServerProcess[bytes], exc: Exception) -> None:
+        print(f"hostbridge mock-ssh shell failed: {exc}", file=sys.stderr, flush=True)
+        with suppress(Exception):
+            process.stderr.write(f"hostbridge mock-ssh shell error: {exc}\n".encode("utf-8", errors="replace"))
+
+    @classmethod
+    def _set_exit_signal(cls, process: asyncssh.SSHServerProcess[bytes], signal_name: str) -> None:
+        try:
+            process.exit_with_signal(signal_name)
+        except Exception:
+            cls._set_exit_status(process, 255)