@neoline/hostbridge 1.4.3 → 2.0.0

package/src/server_control_mcp/doctor.py

@@ -0,0 +1,108 @@
+"""Deterministic HostBridge installation diagnostics."""
+
+from __future__ import annotations
+
+import json
+import tomllib
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Protocol
+
+from . import __version__
+from .config import load_v1_config
+from .runtime import runtime_paths
+
+
+class DoctorClient(Protocol):
+    def hello(self) -> dict[str, object]:
+        raise NotImplementedError
+
+
+@dataclass(frozen=True, slots=True)
+class DiagnosticResult:
+    name: str
+    status: str
+    message: str
+    details: dict[str, object]
+
+
+@dataclass(frozen=True, slots=True)
+class DoctorReport:
+    results: tuple[DiagnosticResult, ...]
+
+    @property
+    def ok(self) -> bool:
+        return all(result.status != "fail" for result in self.results)
+
+
+def collect_diagnostics(
+    config_path: Path | str,
+    *,
+    client: DoctorClient,
+    executable_paths: list[str],
+) -> DoctorReport:
+    results = [
+        _version_diagnostic(),
+        _config_diagnostic(config_path),
+        _runtime_diagnostic(),
+        _daemon_diagnostic(client),
+        _executable_diagnostic(executable_paths),
+    ]
+    return DoctorReport(tuple(results))
+
+
+def _version_diagnostic() -> DiagnosticResult:
+    root = Path(__file__).parents[2]
+    npm = json.loads((root / "package.json").read_text(encoding="utf-8"))
+    pyproject = tomllib.loads((root / "pyproject.toml").read_text(encoding="utf-8"))
+    versions = {
+        "runtime": __version__,
+        "npm": npm.get("version"),
+        "python": pyproject.get("project", {}).get("version"),
+    }
+    status = "pass" if len(set(versions.values())) == 1 and __version__ == "2.0.0" else "fail"
+    return DiagnosticResult("version", status, "version sources agree" if status == "pass" else "version drift detected", versions)
+
+
+def _config_diagnostic(config_path: Path | str) -> DiagnosticResult:
+    try:
+        config = load_v1_config(config_path)
+    except (OSError, ValueError) as exc:
+        return DiagnosticResult("config", "fail", str(exc), {"path": str(config_path)})
+    return DiagnosticResult("config", "pass", "schema v1 configuration is valid", {"path": str(config.path), "hosts": len(config.hosts)})
+
+
+def _runtime_diagnostic() -> DiagnosticResult:
+    paths = runtime_paths()
+    try:
+        paths.ensure_directory()
+    except OSError as exc:
+        return DiagnosticResult("runtime", "fail", str(exc), {"directory": str(paths.directory)})
+    mode = paths.directory.stat().st_mode & 0o777
+    status = "pass" if mode == 0o700 else "fail"
+    return DiagnosticResult(
+        "runtime",
+        status,
+        "runtime directory is private" if status == "pass" else "runtime directory must use mode 0700",
+        {"directory": str(paths.directory), "mode": format(mode, "04o")},
+    )
+
+
+def _daemon_diagnostic(client: DoctorClient) -> DiagnosticResult:
+    try:
+        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
+    return DiagnosticResult(
+        "daemon",
+        "pass" if valid else "fail",
+        "daemon protocol v1 is healthy" if valid else "daemon returned an incompatible handshake",
+        dict(hello),
+    )
+
+
+def _executable_diagnostic(paths: list[str]) -> DiagnosticResult:
+    if not paths:
+        return DiagnosticResult("executable", "fail", "hostbridge executable was not found", {"paths": []})
+    return DiagnosticResult("executable", "pass", "hostbridge executable found", {"paths": paths})

package/src/server_control_mcp/hosts.py

@@ -34,11 +34,12 @@ class HostConfig:
     secrets: dict[str, str] = field(default_factory=dict)
     login_steps: list[LoginStep] = field(default_factory=list)
     ready_expect: str | None = None
+    transport: str = "auto"
 
     @classmethod
-    def script(cls, host_id: str, label: str, script: Path | str) -> HostConfig:
+    def script(cls, host_id: str, label: str, script: Path | str, *, transport: str = "auto") -> HostConfig:
         path = Path(script).expanduser()
-        return cls(host_id, label, "/bin/sh", [str(path)], "script", path)
+        return cls(host_id, label, "/bin/sh", [str(path)], "script", path, transport=transport)
 
     @classmethod
     def command_line(
@@ -52,6 +53,7 @@ class HostConfig:
         secrets: dict[str, str] | None = None,
         login_steps: list[LoginStep] | None = None,
         ready_expect: str | None = None,
+        transport: str = "auto",
     ) -> HostConfig:
         return cls(
             host_id,
@@ -63,6 +65,7 @@ class HostConfig:
             secrets or {},
             login_steps or [],
             ready_expect,
+            transport,
         )
 
     def describe(self) -> dict[str, object]:
@@ -72,6 +75,7 @@ class HostConfig:
             "connection_type": self.connection_type,
             "command": self.command,
             "args": list(self.args),
+            "transport": self.transport,
         }
         if self.script_path is not None:
             info.update(
@@ -130,16 +134,40 @@ def _parse_host(item: object, path: Path, index: int) -> HostConfig:
     secrets = _parse_secrets(item.get("secrets", {}), path, index)
     login_steps = _parse_login_steps(item.get("login_steps", []), path, index)
     ready_expect = _parse_ready_expect(item.get("ready_expect"), path, index)
+    transport = str(item.get("transport", "auto")).strip().lower()
+    if transport not in {"auto", "pty", "ssh"}:
+        raise ValueError(f"{path} hosts[{index}].transport must be one of: auto, pty, ssh")
 
     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)
-        return HostConfig(host.id, host.label, host.command, host.args, host.connection_type, host.script_path, secrets, login_steps, ready_expect)
+        host = HostConfig.script(host_id, label, script, transport=transport)
+        return HostConfig(
+            host.id,
+            host.label,
+            host.command,
+            host.args,
+            host.connection_type,
+            host.script_path,
+            secrets,
+            login_steps,
+            ready_expect,
+            transport,
+        )
 
     if "ssh" in item:
-        return HostConfig.command_line(host_id, label, "ssh", _parse_ssh_args(item["ssh"], path, index), connection_type="ssh", secrets=secrets, login_steps=login_steps, ready_expect=ready_expect)
+        return HostConfig.command_line(
+            host_id,
+            label,
+            "ssh",
+            _parse_ssh_args(item["ssh"], path, index),
+            connection_type="ssh",
+            secrets=secrets,
+            login_steps=login_steps,
+            ready_expect=ready_expect,
+            transport=transport,
+        )
 
     command = str(item.get("command", "")).strip()
     args = item.get("args", [])
@@ -147,7 +175,16 @@ def _parse_host(item: object, path: Path, index: int) -> HostConfig:
         raise ValueError(f"{path} hosts[{index}].command must not be empty")
     if not isinstance(args, list) or not all(isinstance(arg, str) for arg in args):
         raise ValueError(f"{path} hosts[{index}].args must be a list of strings")
-    return HostConfig.command_line(host_id, label, command, args, secrets=secrets, login_steps=login_steps, ready_expect=ready_expect)
+    return HostConfig.command_line(
+        host_id,
+        label,
+        command,
+        args,
+        secrets=secrets,
+        login_steps=login_steps,
+        ready_expect=ready_expect,
+        transport=transport,
+    )
 
 
 def _parse_ready_expect(raw_ready_expect: object, path: Path, index: int) -> str | None:
@@ -196,7 +233,15 @@ def _parse_ssh_args(raw_ssh: object, path: Path, index: int) -> list[str]:
     args: list[str] = []
     port = raw_ssh.get("port")
     if port not in (None, ""):
-        args.extend(["-p", str(int(port))])
+        if isinstance(port, bool) or not isinstance(port, int | str):
+            raise ValueError(f"{path} hosts[{index}].ssh.port must be an integer")
+        try:
+            parsed_port = int(port)
+        except ValueError as exc:
+            raise ValueError(f"{path} hosts[{index}].ssh.port must be an integer") from exc
+        if not 1 <= parsed_port <= 65535:
+            raise ValueError(f"{path} hosts[{index}].ssh.port must be between 1 and 65535")
+        args.extend(["-p", str(parsed_port)])
     identity_file = str(raw_ssh.get("identity_file", "")).strip()
     if identity_file:
         args.extend(["-i", str(Path(identity_file).expanduser())])

package/src/server_control_mcp/mock_ssh.py

@@ -3,19 +3,40 @@ from __future__ import annotations
 import asyncio
 import json
 import posixpath
+import re
 import shlex
 import sys
+import tempfile
 import time
-from contextlib import suppress
 from dataclasses import dataclass, field
+from functools import partial
 from pathlib import Path
+from typing import Any, cast
 
 import asyncssh
 
-from .daemon import build_manager
-from .manager import DEFAULT_COMMAND_TIMEOUT, DEFAULT_MAX_TRANSFER_BYTES, RemoteCommandError, SessionManager, _new_token
+from .client import HostBridgeClient, HostBridgeError
 
 DEFAULT_PORT_RANGE = range(2222, 2300)
+DEFAULT_COMMAND_TIMEOUT = 60
+DEFAULT_MAX_TRANSFER_BYTES = 1024 * 1024 * 1024
+EXEC_STDIN_PROBE_TIMEOUT = 0.05
+STDIN_READING_COMMANDS = [
+    re.compile(r"\bcat\s*(?:>|>>)"),
+    re.compile(r"\btee(?:\s|$)"),
+    re.compile(r"(?:^|\s)(?:/bin/)?(?:sh|bash|dash|zsh)\b[^\n;]*\s-s(?:\s|$)"),
+    re.compile(r"(?:^|\s)(?:python[\d.]*|node)(?:\s+\S+)*\s+-($|\s)"),
+    re.compile(r"(?:^|\s)tar\s+(?:-[A-Za-z]*x[A-Za-z]*|[A-Za-z]*x[A-Za-z]*)(?:\s|$)"),
+    re.compile(r"(?:^|\s)tar\s+.*(?:^|\s)f\s*-(?:\s|$)"),
+    re.compile(r"(?:^|\s)dd\s+.*\bof="),
+    re.compile(r"(?:^|\s)scp(?:\s+\S+)*\s+-[A-Za-z]*(?:t|f)[A-Za-z]*(?:\s|$)"),
+]
+
+
+def _to_bytes(chunk: str | bytes) -> bytes:
+    if isinstance(chunk, bytes):
+        return chunk
+    return chunk.encode("utf-8")
 
 
 class HostBridgeSSHServer(asyncssh.SSHServer):
@@ -36,8 +57,8 @@ class HostBridgeSSHServer(asyncssh.SSHServer):
 
 
 class HostBridgeSSHProcess:
-    def __init__(self, manager: SessionManager, session_factory: "HostBridgeSessionFactory", timeout: int, output_limit: int) -> None:
-        self._manager = manager
+    def __init__(self, client: HostBridgeClient, session_factory: HostBridgeSessionFactory, timeout: int, output_limit: int) -> None:
+        self._client = client
         self._session_factory = session_factory
         self._timeout = timeout
         self._output_limit = output_limit
@@ -57,26 +78,57 @@ class HostBridgeSSHProcess:
 
         loop = asyncio.get_running_loop()
         try:
+            stdin_data = await self._read_exec_stdin(process)
             result = await loop.run_in_executor(
                 None,
-                lambda: self._manager.run_command(
+                lambda: self._client.exec(
                     session_id,
                     command,
+                    stdin=stdin_data,
                     timeout=self._timeout,
-                    token_factory=_new_token,
                     output_limit=self._output_limit,
+                    owner="mock-ssh",
                 ),
             )
-            if result.output:
-                process.stdout.write(result.output)
-                if not result.output.endswith("\n"):
+            if result.stdout:
+                output = result.stdout.decode("utf-8", errors="replace")
+                process.stdout.write(output)
+                if not output.endswith("\n"):
                     process.stdout.write("\n")
+            if result.stderr:
+                process.stderr.write(result.stderr.decode("utf-8", errors="replace"))
             process.exit(124 if result.timed_out else int(result.exit_code or 0))
         except Exception as exc:
             print(f"hostbridge mock-ssh exec failed: {exc}", file=sys.stderr, flush=True)
             process.stderr.write(f"hostbridge mock-ssh error: {exc}\n")
             process.exit(1)
 
+    async def _read_exec_stdin(self, process: asyncssh.SSHServerProcess[str]) -> bytes | None:
+        if process.stdin.at_eof():
+            return b""
+        try:
+            first_chunk = await asyncio.wait_for(process.stdin.read(4096), timeout=EXEC_STDIN_PROBE_TIMEOUT)
+        except TimeoutError:
+            return None
+        if not first_chunk:
+            return b""
+
+        first_chunk_bytes = _to_bytes(first_chunk)
+        total = len(first_chunk_bytes)
+        if total > self._output_limit:
+            raise HostBridgeError("resource_limit", f"mock-ssh exec stdin exceeds limit ({self._output_limit} bytes)")
+        chunks = [first_chunk_bytes]
+        while not process.stdin.at_eof():
+            chunk = await asyncio.wait_for(process.stdin.read(4096), timeout=self._timeout)
+            if not chunk:
+                break
+            chunk_bytes = _to_bytes(chunk)
+            total += len(chunk_bytes)
+            if total > self._output_limit:
+                raise HostBridgeError("resource_limit", f"mock-ssh exec stdin exceeds limit ({self._output_limit} bytes)")
+            chunks.append(chunk_bytes)
+        return b"".join(chunks)
+
     async def _run_shell(self, process: asyncssh.SSHServerProcess[str], session_id: str) -> None:
         loop = asyncio.get_running_loop()
         stop = asyncio.Event()
@@ -90,7 +142,12 @@ class HostBridgeSSHProcess:
                         break
                     await loop.run_in_executor(
                         None,
-                        lambda chunk=data: self._manager.send(session_id, chunk, owner_agent_id="mock-ssh"),
+                        partial(
+                            self._client.shell_write,
+                            session_id,
+                            _to_bytes(data),
+                            owner="mock-ssh",
+                        ),
                     )
             finally:
                 input_done.set()
@@ -101,23 +158,22 @@ class HostBridgeSSHProcess:
                 while not stop.is_set():
                     result = await loop.run_in_executor(
                         None,
-                        lambda: self._manager.read(
+                        lambda: self._client.shell_read(
                             session_id,
                             timeout=0.2,
-                            max_chars=4096,
-                            owner_agent_id="mock-ssh",
+                            max_bytes=4096,
+                            owner="mock-ssh",
                         ),
                     )
-                    output = str(result.get("output") or "")
-                    if output:
+                    if result.data:
                         idle_after_input = 0
-                        process.stdout.write(output)
+                        process.stdout.write(result.data.decode("utf-8", errors="replace"))
                     elif input_done.is_set():
                         idle_after_input += 1
                         if idle_after_input >= 5:
                             stop.set()
                             break
-                    if not result.get("alive", True):
+                    if not result.alive:
                         stop.set()
                         break
             except Exception as exc:
@@ -143,36 +199,39 @@ class HostBridgeSFTPHandle:
     mode: int = 0o644
 
 
+@dataclass(frozen=True, slots=True)
+class _TextCommandResult:
+    output: str
+    exit_code: int | None
+    timed_out: bool
+
+
 class HostBridgeSessionFactory:
-    def __init__(self, manager: SessionManager, host_id: str, connect_timeout: int) -> None:
-        self._manager = manager
+    def __init__(self, client: HostBridgeClient, host_id: str, connect_timeout: int = 60) -> None:
+        self._client = client
         self._host_id = host_id
         self._connect_timeout = connect_timeout
 
     def open(self) -> str:
-        session = self._manager.open_session(
-            self._host_id,
-            connect_timeout=self._connect_timeout,
-            owner_agent_id="mock-ssh",
-        )
-        return session.id
+        session = self._client.open_session(self._host_id, owner="mock-ssh")
+        return session.session_id
 
     def close(self, session_id: str) -> None:
         with _suppress_error():
-            self._manager.close_session(session_id, owner_agent_id="mock-ssh", force=True)
+            self._client.close_session(session_id, owner="mock-ssh", force=True)
 
 
 class HostBridgeSFTPServer(asyncssh.SFTPServer):
     def __init__(
         self,
         chan: asyncssh.SSHServerChannel,
-        manager: SessionManager,
+        client: HostBridgeClient,
         session_factory: HostBridgeSessionFactory,
         timeout: int,
         max_bytes: int,
     ) -> None:
         super().__init__(chan)
-        self._manager = manager
+        self._client = client
         self._session_factory = session_factory
         self._session_id = session_factory.open()
         self._timeout = timeout
@@ -278,17 +337,25 @@ class HostBridgeSFTPServer(asyncssh.SFTPServer):
             self._checked_run(f"python3 -c {shlex.quote(script)} {shlex.quote(remote_path)} {atime} {mtime}")
 
     def _download(self, remote_path: str) -> bytes:
+        temporary_path: Path | None = None
         try:
-            result = self._manager.file_read_bytes(
+            with tempfile.NamedTemporaryFile(prefix="hostbridge-sftp-", delete=False) as temporary:
+                temporary_path = Path(temporary.name)
+            result = self._client.download_file(
                 self._session_id,
                 remote_path,
-                max_bytes=self._max_bytes,
-                token_factory=_new_token,
-                owner_agent_id="mock-ssh",
+                temporary_path,
+                timeout=self._timeout,
+                owner="mock-ssh",
             )
-        except RemoteCommandError as exc:
+            if result.bytes_transferred > self._max_bytes:
+                raise HostBridgeError("resource_limit", f"file exceeds mock-ssh limit ({self._max_bytes} bytes)")
+            return temporary_path.read_bytes()
+        except HostBridgeError as exc:
             raise self._sftp_error(exc) from exc
-        return bytes(result["data_bytes"])
+        finally:
+            if temporary_path is not None:
+                temporary_path.unlink(missing_ok=True)
 
     def _upload(self, handle: HostBridgeSFTPHandle) -> None:
         size = len(handle.data)
@@ -299,19 +366,27 @@ class HostBridgeSFTPServer(asyncssh.SFTPServer):
             data.extend(b"\x00" * (size - len(data)))
         for offset, chunk in sorted(handle.writes.items()):
             data[offset : offset + len(chunk)] = chunk
+        if len(data) > self._max_bytes:
+            raise asyncssh.SFTPFailure(f"file exceeds mock-ssh limit ({self._max_bytes} bytes)")
+        temporary_path: Path | None = None
         try:
-            self._manager.file_write_bytes(
+            with tempfile.NamedTemporaryFile(prefix="hostbridge-sftp-", delete=False) as temporary:
+                temporary.write(data)
+                temporary_path = Path(temporary.name)
+            self._client.upload_file(
                 self._session_id,
+                temporary_path,
                 handle.path,
-                bytes(data),
-                mode=format(handle.mode & 0o7777, "04o"),
-                max_bytes=self._max_bytes,
-                token_factory=_new_token,
-                owner_agent_id="mock-ssh",
+                mode=handle.mode & 0o7777,
+                timeout=self._timeout,
+                owner="mock-ssh",
             )
-        except RemoteCommandError as exc:
+        except HostBridgeError as exc:
             print(f"hostbridge mock-ssh sftp upload failed: {exc}", file=sys.stderr, flush=True)
             raise self._sftp_error(exc) from exc
+        finally:
+            if temporary_path is not None:
+                temporary_path.unlink(missing_ok=True)
 
     def _stat(self, remote_path: str, *, follow: bool) -> asyncssh.SFTPAttrs:
         script = r"""
@@ -362,7 +437,7 @@ print(json.dumps(items))
         ]
 
     def _attrs(self, data: dict[str, object]) -> asyncssh.SFTPAttrs:
-        permissions = int(data.get("permissions") or 0)
+        permissions = self._int_value(data.get("permissions"), 0)
         raw_type = data.get("type")
         file_type = {
             "file": asyncssh.FILEXFER_TYPE_REGULAR,
@@ -371,39 +446,52 @@ print(json.dumps(items))
         }.get(str(raw_type), asyncssh.FILEXFER_TYPE_UNKNOWN)
         return asyncssh.SFTPAttrs(
             type=file_type,
-            size=int(data.get("size") or 0),
-            uid=int(data.get("uid") or 0),
-            gid=int(data.get("gid") or 0),
+            size=self._int_value(data.get("size"), 0),
+            uid=self._int_value(data.get("uid"), 0),
+            gid=self._int_value(data.get("gid"), 0),
             permissions=permissions,
-            atime=int(data.get("atime") or time.time()),
-            mtime=int(data.get("mtime") or time.time()),
+            atime=self._int_value(data.get("atime"), int(time.time())),
+            mtime=self._int_value(data.get("mtime"), int(time.time())),
         )
 
+    @staticmethod
+    def _int_value(value: object, default: int) -> int:
+        if value is None:
+            return default
+        if isinstance(value, bool) or not isinstance(value, int | float | str):
+            raise asyncssh.SFTPFailure("remote metadata contains a non-numeric field")
+        try:
+            return int(value)
+        except ValueError as exc:
+            raise asyncssh.SFTPFailure("remote metadata contains an invalid numeric field") from exc
+
     def _checked_run(self, command: str):
         try:
-            result = self._manager.run_command(
+            result = self._client.exec(
                 self._session_id,
                 command,
                 timeout=self._timeout,
-                token_factory=_new_token,
-                owner_agent_id="mock-ssh",
+                owner="mock-ssh",
             )
-        except RemoteCommandError as exc:
+        except HostBridgeError as exc:
             raise self._sftp_error(exc) from exc
+        output = result.stdout.decode("utf-8", errors="replace")
         if result.timed_out:
             print(f"hostbridge mock-ssh sftp command timed out: {command}", file=sys.stderr, flush=True)
             raise asyncssh.SFTPFailure("remote command timed out")
         if result.exit_code not in (0, None):
-            error = self._sftp_error(RemoteCommandError(result.output or f"remote command exited {result.exit_code}"))
+            error = self._sftp_error(
+                HostBridgeError("remote_error", output or f"remote command exited {result.exit_code}")
+            )
             if isinstance(error, asyncssh.SFTPNoSuchFile):
                 raise error
             print(
-                f"hostbridge mock-ssh sftp command failed exit={result.exit_code}: {command}\n{result.output}",
+                f"hostbridge mock-ssh sftp command failed exit={result.exit_code}: {command}\n{output}",
                 file=sys.stderr,
                 flush=True,
             )
             raise error
-        return result
+        return _TextCommandResult(output, result.exit_code, result.timed_out)
 
     def _path(self, path: bytes) -> str:
         text = path.decode("utf-8", "surrogateescape")
@@ -416,7 +504,7 @@ print(json.dumps(items))
             raise asyncssh.SFTPInvalidHandle("invalid HostBridge SFTP handle")
         return file_obj
 
-    def _sftp_error(self, exc: RemoteCommandError) -> Exception:
+    def _sftp_error(self, exc: HostBridgeError) -> Exception:
         text = str(exc)
         lowered = text.lower()
         if "no such file" in lowered or "not found" in lowered or "filenotfounderror" in lowered:
@@ -595,49 +683,51 @@ async def serve_mock_ssh(
     command_timeout: int = DEFAULT_COMMAND_TIMEOUT,
     output_limit: int = DEFAULT_MAX_TRANSFER_BYTES,
 ) -> int:
+    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",
+        )
     key_dir = key_dir or default_key_dir(host_id)
     ssh_host_alias = ssh_host_alias or default_ssh_host_alias(host_id)
     host_key_path, client_key_path, public_key = ensure_keypair(key_dir)
-    manager = build_manager(config_path)
-    session_factory = HostBridgeSessionFactory(manager, host_id, connect_timeout)
-    try:
-        server, selected_port = await create_server_on_available_port(
-            server_factory=lambda: HostBridgeSSHServer(public_key),
+    client = HostBridgeClient()
+    client.hello()
+    session_factory = HostBridgeSessionFactory(client, host_id, connect_timeout)
+    server, selected_port = await create_server_on_available_port(
+        server_factory=lambda: HostBridgeSSHServer(public_key),
+        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),
+    )
+    installed_config_path: Path | None = None
+    if install_ssh_config:
+        installed_config_path = install_ssh_config_block(
+            host_alias=ssh_host_alias,
             listen_host=listen_host,
-            port=port,
-            server_host_keys=[str(host_key_path)],
-            process_factory=HostBridgeSSHProcess(manager, session_factory, command_timeout, output_limit),
-            sftp_factory=lambda chan: HostBridgeSFTPServer(chan, manager, session_factory, command_timeout, output_limit),
+            port=selected_port,
+            client_key_path=client_key_path,
+            config_path=ssh_config_path,
         )
-        installed_config_path: Path | None = None
-        if install_ssh_config:
-            installed_config_path = install_ssh_config_block(
-                host_alias=ssh_host_alias,
-                listen_host=listen_host,
-                port=selected_port,
-                client_key_path=client_key_path,
-                config_path=ssh_config_path,
-            )
-        print(f"hostbridge mock-ssh forwarding host '{host_id}'", flush=True)
-        print(f"listening on ssh://{listen_host}:{selected_port}", flush=True)
-        print(f"ssh host alias: {ssh_host_alias}", flush=True)
-        if installed_config_path is not None:
-            print(f"ssh config: {installed_config_path}", flush=True)
-        print(f"identity file: {client_key_path}", flush=True)
-        print("supports SSH exec, shell, SFTP, and SCP; press Ctrl-C to stop", flush=True)
-        await server.wait_closed()
-    finally:
-        manager.close_all_sessions()
-        manager.stop_idle_reaper()
+    print(f"hostbridge mock-ssh forwarding host '{host_id}'", flush=True)
+    print(f"listening on ssh://{listen_host}:{selected_port}", flush=True)
+    print(f"ssh host alias: {ssh_host_alias}", flush=True)
+    if installed_config_path is not None:
+        print(f"ssh config: {installed_config_path}", flush=True)
+    print(f"identity file: {client_key_path}", flush=True)
+    print("supports SSH exec, shell, SFTP, and SCP; press Ctrl-C to stop", flush=True)
+    await server.wait_closed()
     return 0
 
 
 def run_mock_ssh(**kwargs: object) -> int:
     try:
-        return asyncio.run(serve_mock_ssh(**kwargs))
+        return asyncio.run(cast(Any, serve_mock_ssh)(**kwargs))
     except KeyboardInterrupt:
         return 130
-    except RemoteCommandError as exc:
+    except HostBridgeError as exc:
         print(f"hostbridge mock-ssh failed: {exc}")
         return 1