@neoline/hostbridge 2.0.4 → 2.0.5

package/src/server_control_mcp/daemon.py

@@ -1,530 +1,62 @@
 from __future__ import annotations
 
-import base64
-import binascii
-import hashlib
-import json
+import asyncio
 import os
+import shlex
+import signal
 import socket
 import stat
-import subprocess
+import subprocess  # nosec B404
 import sys
-import threading
 import time
 from contextlib import suppress
-from dataclasses import asdict
+from dataclasses import dataclass
 from pathlib import Path
-from socketserver import StreamRequestHandler, ThreadingMixIn, UnixStreamServer
-from typing import Any
 
 from . import __version__
 from .client import HostBridgeClient, HostBridgeError
 from .hosts import HostRegistry, configured_hosts_path, load_hosts
+from .mux_daemon import MuxDaemonServer
+from .mux_protocol import MUX_PROTOCOL_VERSION
+from .mux_service import HostBridgeMuxService
 from .policy import build_audit_sink, load_policy
-from .protocol import (
-    PROTOCOL_VERSION,
-    BinaryFrame,
-    ErrorInfo,
-    FrameFlags,
-    FrameSequenceValidator,
-    ProtocolError,
-    Request,
-    Response,
-    encode_frame,
-    read_frame,
-)
 from .runtime import runtime_paths
+from .runtime_services import AsyncHostBridgeServices
 from .secure_files import open_private_append, replace_private_file
-from .services import HostBridgeServices, ServiceError
-from .transports.pty import PtyTransport
-from .transports.ssh import SshConnectionConfig, SshTransport
+from .tunnel_manager import HostTunnelManager
+from .tunnel_providers import tunnel_provider_factory
 
-MAX_V1_CONTROL_LINE_BYTES = 4 * 1024 * 1024
 DEFAULT_CONNECT_TIMEOUT = 30
-DEFAULT_MAX_DAEMON_CONNECTIONS = 64
-DEFAULT_DAEMON_CONNECTION_IDLE_TIMEOUT = 30.0
-
-
-class ThreadingUnixRPCServer(ThreadingMixIn, UnixStreamServer):
-    daemon_threads = True
-    allow_reuse_address = True
-
-    def __init__(
-        self,
-        server_address: str,
-        request_handler_class: type[StreamRequestHandler],
-        bind_and_activate: bool = True,
-        *,
-        max_connections: int = DEFAULT_MAX_DAEMON_CONNECTIONS,
-        connection_idle_timeout: float = DEFAULT_DAEMON_CONNECTION_IDLE_TIMEOUT,
-    ) -> None:
-        if not isinstance(max_connections, int) or isinstance(max_connections, bool) or max_connections <= 0:
-            raise ValueError("max_connections must be a positive integer")
-        if connection_idle_timeout <= 0:
-            raise ValueError("connection_idle_timeout must be positive")
-        self._connection_slots = threading.BoundedSemaphore(max_connections)
-        self._connection_idle_timeout = float(connection_idle_timeout)
-        super().__init__(server_address, request_handler_class, bind_and_activate)
-
-    def get_request(self) -> tuple[socket.socket, Any]:
-        request, client_address = super().get_request()
-        request.settimeout(self._connection_idle_timeout)
-        return request, client_address
-
-    def process_request(
-        self,
-        request: socket.socket | tuple[bytes, socket.socket],
-        client_address: Any,
-    ) -> None:
-        self._connection_slots.acquire()
-        try:
-            super().process_request(request, client_address)
-        except BaseException:
-            self._connection_slots.release()
-            raise
-
-    def process_request_thread(
-        self,
-        request: socket.socket | tuple[bytes, socket.socket],
-        client_address: Any,
-    ) -> None:
-        try:
-            super().process_request_thread(request, client_address)
-        finally:
-            self._connection_slots.release()
-
-
-def _required_param(params: dict[str, object], name: str) -> str:
-    value = params.get(name)
-    if not isinstance(value, str) or not value.strip():
-        raise ServiceError("invalid_request", f"{name} must be a non-empty string")
-    return value
-
-
-def _optional_owner(params: dict[str, object]) -> str | None:
-    owner = params.get("owner")
-    if owner is not None and not isinstance(owner, str):
-        raise ServiceError("invalid_request", "owner must be a string or null")
-    return owner
-
-
-def _decode_inline_stdin(value: object) -> list[bytes] | None:
-    if value is None:
-        return None
-    if not isinstance(value, dict) or value.get("encoding") != "base64" or not isinstance(value.get("data"), str):
-        raise ServiceError("invalid_request", "stdin must contain base64 data")
-    try:
-        return [base64.b64decode(value["data"], validate=True)]
-    except (binascii.Error, ValueError) as exc:
-        raise ServiceError("invalid_request", "stdin contains invalid base64") from exc
-
-
-def _encoded_bytes(value: bytes) -> dict[str, str]:
-    return {"encoding": "base64", "data": base64.b64encode(value).decode("ascii")}
-
-
-class HostBridgeV1Dispatcher:
-    def __init__(self, services: HostBridgeServices) -> None:
-        self.services = services
-        self.handlers = {
-            "system.hello": self.system_hello,
-            "system.health": self.system_health,
-            "system.shutdown": self.system_shutdown,
-            "host.list": self.host_list,
-            "session.open": self.session_open,
-            "session.list": self.session_list,
-            "session.close": self.session_close,
-            "session.close_all": self.session_close_all,
-            "exec.start": self.exec_start,
-            "task.start": self.task_start,
-            "task.status": self.task_status,
-            "task.cancel": self.task_cancel,
-            "file.write_text": self.file_write_text,
-            "file.read_text": self.file_read_text,
-            "shell.write": self.shell_write,
-            "shell.read": self.shell_read,
-        }
-
-    def dispatch(self, request: Request) -> Response:
-        handler = self.handlers.get(request.method)
-        if handler is None:
-            return Response.failure(
-                request.request_id,
-                ErrorInfo("not_found", f"unknown method {request.method}"),
-            )
-        try:
-            return Response.success(request.request_id, handler(request.params))
-        except ServiceError as exc:
-            return Response.failure(
-                request.request_id,
-                ErrorInfo(exc.code, str(exc), retryable=exc.retryable, details=exc.details),
-            )
-        except (TypeError, ValueError) as exc:
-            return Response.failure(request.request_id, ErrorInfo("invalid_request", str(exc)))
-        except Exception:
-            return Response.failure(
-                request.request_id, ErrorInfo("internal", "internal HostBridge error", retryable=True)
-            )
-
-    def system_hello(self, params: dict[str, object]) -> dict[str, object]:
-        return {
-            "service": "hostbridge",
-            "protocol": PROTOCOL_VERSION,
-            "version": __version__,
-            "capabilities": ["sessions", "exec", "explicit_stdin"],
-        }
-
-    def system_health(self, params: dict[str, object]) -> dict[str, object]:
-        return {"status": "ok", "protocol": PROTOCOL_VERSION}
-
-    def system_shutdown(self, params: dict[str, object]) -> dict[str, object]:
-        self.services.stop_idle_reaper()
-        return self.services.close_all(force=True)
-
-    def host_list(self, params: dict[str, object]) -> dict[str, object]:
-        return {"hosts": self.services.hosts.describe()}
-
-    def session_open(self, params: dict[str, object]) -> dict[str, object]:
-        connect_timeout = params.get("connect_timeout")
-        if connect_timeout is not None and (
-            not isinstance(connect_timeout, int | float)
-            or isinstance(connect_timeout, bool)
-            or not 0 < connect_timeout <= 180
-        ):
-            raise ServiceError("invalid_request", "connect_timeout must be between 0 and 180 seconds")
-        session = self.services.open_session(
-            _required_param(params, "host_id"),
-            owner=_optional_owner(params),
-            connect_timeout=float(connect_timeout) if connect_timeout is not None else None,
-        )
-        return session.describe()
-
-    def session_list(self, params: dict[str, object]) -> dict[str, object]:
-        return {"sessions": self.services.list_sessions()}
-
-    def session_close(self, params: dict[str, object]) -> dict[str, object]:
-        if "force" in params:
-            raise ServiceError("invalid_request", "force close is not available over the public protocol")
-        session_id = _required_param(params, "session_id")
-        self.services.close_session(session_id, owner=_optional_owner(params))
-        return {"session_id": session_id, "closed": True}
-
-    def session_close_all(self, params: dict[str, object]) -> dict[str, object]:
-        return self.services.close_all(owner=_optional_owner(params))
-
-    def exec_start(self, params: dict[str, object]) -> dict[str, object]:
-        timeout = params.get("timeout", 60)
-        output_limit = params.get("output_limit", 1_000_000)
-        if not isinstance(timeout, int | float) or isinstance(timeout, bool) or timeout <= 0:
-            raise ServiceError("invalid_request", "timeout must be positive")
-        if not isinstance(output_limit, int) or isinstance(output_limit, bool) or output_limit <= 0:
-            raise ServiceError("invalid_request", "output_limit must be a positive integer")
-        result = self.services.exec(
-            _required_param(params, "session_id"),
-            _required_param(params, "command"),
-            stdin=_decode_inline_stdin(params.get("stdin")),
-            timeout=float(timeout),
-            output_limit=output_limit,
-            owner=_optional_owner(params),
-        )
-        return {
-            "session_id": result.session_id,
-            "exit_code": result.exit_code,
-            "stdout": _encoded_bytes(result.stdout),
-            "stderr": _encoded_bytes(result.stderr),
-            "timed_out": result.timed_out,
-        }
-
-    def task_start(self, params: dict[str, object]) -> dict[str, object]:
-        task = self.services.start_task(
-            _required_param(params, "session_id"),
-            _required_param(params, "command"),
-            owner=_optional_owner(params),
-        )
-        return asdict(task)
-
-    def task_status(self, params: dict[str, object]) -> dict[str, object]:
-        tail_lines = params.get("tail_lines", 80)
-        if not isinstance(tail_lines, int) or isinstance(tail_lines, bool):
-            raise ServiceError("invalid_request", "tail_lines must be an integer")
-        return self.services.task_status(
-            _required_param(params, "session_id"),
-            _required_param(params, "task_id"),
-            owner=_optional_owner(params),
-            tail_lines=tail_lines,
-        )
-
-    def task_cancel(self, params: dict[str, object]) -> dict[str, object]:
-        return self.services.cancel_task(
-            _required_param(params, "session_id"),
-            _required_param(params, "task_id"),
-            owner=_optional_owner(params),
-        )
-
-    def file_write_text(self, params: dict[str, object]) -> dict[str, object]:
-        content = params.get("content")
-        mode = params.get("mode", 0o644)
-        max_bytes = params.get("max_bytes", 4 * 1024 * 1024)
-        if not isinstance(content, str):
-            raise ServiceError("invalid_request", "content must be a string")
-        if not isinstance(mode, int) or isinstance(mode, bool):
-            raise ServiceError("invalid_request", "mode must be an integer")
-        if not isinstance(max_bytes, int) or isinstance(max_bytes, bool) or max_bytes <= 0:
-            raise ServiceError("invalid_request", "max_bytes must be a positive integer")
-        result = self.services.write_text(
-            _required_param(params, "session_id"),
-            _required_param(params, "remote_path"),
-            content,
-            mode=mode,
-            max_bytes=max_bytes,
-            owner=_optional_owner(params),
-        )
-        return {"bytes_transferred": result.bytes_transferred, "sha256": result.sha256}
-
-    def file_read_text(self, params: dict[str, object]) -> dict[str, object]:
-        max_bytes = params.get("max_bytes", 4 * 1024 * 1024)
-        if not isinstance(max_bytes, int) or isinstance(max_bytes, bool) or max_bytes <= 0:
-            raise ServiceError("invalid_request", "max_bytes must be a positive integer")
-        content = self.services.read_text(
-            _required_param(params, "session_id"),
-            _required_param(params, "remote_path"),
-            max_bytes=max_bytes,
-            owner=_optional_owner(params),
-        )
-        return {"content": content, "bytes_transferred": len(content.encode("utf-8"))}
-
-    def shell_write(self, params: dict[str, object]) -> dict[str, object]:
-        chunks = _decode_inline_stdin(params.get("data"))
-        if chunks is None:
-            raise ServiceError("invalid_request", "data must contain base64 bytes")
-        written = self.services.shell_write(
-            _required_param(params, "session_id"),
-            b"".join(chunks),
-            owner=_optional_owner(params),
-        )
-        return {"written": written}
-
-    def shell_read(self, params: dict[str, object]) -> dict[str, object]:
-        timeout = params.get("timeout", 0.2)
-        max_bytes = params.get("max_bytes", 4096)
-        if not isinstance(timeout, int | float) or isinstance(timeout, bool) or timeout < 0:
-            raise ServiceError("invalid_request", "timeout must be non-negative")
-        if not isinstance(max_bytes, int) or isinstance(max_bytes, bool) or max_bytes <= 0:
-            raise ServiceError("invalid_request", "max_bytes must be a positive integer")
-        result = self.services.shell_read(
-            _required_param(params, "session_id"),
-            timeout=float(timeout),
-            max_bytes=max_bytes,
-            owner=_optional_owner(params),
-        )
-        return {"data": _encoded_bytes(result.data), "alive": result.alive}
-
+_DAEMON_STOP_GRACE_SECONDS = 5.0
+_DAEMON_TERM_SECONDS = 2.0
+_DAEMON_KILL_SECONDS = 2.0
+_PROCESS_POLL_INTERVAL_SECONDS = 0.05
 
-class HostBridgeV1Handler(StreamRequestHandler):
-    dispatcher: HostBridgeV1Dispatcher
 
-    def handle(self) -> None:
-        try:
-            self._handle_requests()
-        except OSError:
-            return
-
-    def _handle_requests(self) -> None:
-        while raw_line := self.rfile.readline(MAX_V1_CONTROL_LINE_BYTES + 1):
-            request_id = "invalid"
-            try:
-                if len(raw_line) > MAX_V1_CONTROL_LINE_BYTES:
-                    raise ProtocolError("request exceeds the control limit")
-                request = Request.from_json(raw_line.decode("utf-8"))
-                request_id = request.request_id
-                if request.method == "transfer.upload":
-                    response = self._handle_upload(request)
-                elif request.method == "exec.stream":
-                    response = self._handle_exec_stream(request)
-                elif request.method == "transfer.download":
-                    self._handle_download(request)
-                    continue
-                else:
-                    response = self.dispatcher.dispatch(request)
-            except ServiceError as exc:
-                response = Response.failure(
-                    request_id,
-                    ErrorInfo(exc.code, str(exc), retryable=exc.retryable, details=exc.details),
-                )
-            except (ProtocolError, UnicodeDecodeError) as exc:
-                response = Response.failure(request_id, ErrorInfo("invalid_request", str(exc)))
-            self.wfile.write(response.to_json().encode("utf-8") + b"\n")
-            self.wfile.flush()
-            if request_id != "invalid" and request.method == "system.shutdown":
-                threading.Thread(target=self.server.shutdown, daemon=True).start()
-
-    def _handle_upload(self, request: Request) -> Response:
-        params = request.params
-        stream_id = _required_param(params, "stream_id")
-        expected_size = params.get("expected_size")
-        mode = params.get("mode", 0o644)
-        if not isinstance(expected_size, int) or isinstance(expected_size, bool) or expected_size < 0:
-            raise ServiceError("invalid_request", "expected_size must be a non-negative integer")
-        if not isinstance(mode, int) or isinstance(mode, bool):
-            raise ServiceError("invalid_request", "mode must be an integer")
-        validator = FrameSequenceValidator()
-
-        def chunks():
-            while True:
-                frame = read_frame(self.rfile)
-                validator.accept(frame)
-                if frame.stream_id != stream_id:
-                    raise ServiceError("invalid_request", "upload frame stream_id does not match")
-                if frame.flags & FrameFlags.CANCEL:
-                    raise ServiceError("cancelled", "upload cancelled")
-                if frame.flags & FrameFlags.DATA:
-                    yield frame.payload
-                if frame.flags & FrameFlags.EOF:
-                    return
-
-        result = self.dispatcher.services.upload(
-            _required_param(params, "session_id"),
-            chunks(),
-            _required_param(params, "remote_path"),
-            mode=mode,
-            expected_size=expected_size,
-            expected_sha256=None,
-            owner=_optional_owner(params),
-        )
-        return Response.success(
-            request.request_id,
-            {"bytes_transferred": result.bytes_transferred, "sha256": result.sha256},
-        )
+@dataclass(frozen=True, slots=True)
+class DaemonRuntime:
+    services: AsyncHostBridgeServices
+    tunnel_manager: HostTunnelManager
 
-    def _handle_exec_stream(self, request: Request) -> Response:
-        params = request.params
-        stream_id = _required_param(params, "stream_id")
-        timeout = params.get("timeout", 60)
-        output_limit = params.get("output_limit", 1_000_000)
-        if not isinstance(timeout, int | float) or isinstance(timeout, bool) or timeout <= 0:
-            raise ServiceError("invalid_request", "timeout must be positive")
-        if not isinstance(output_limit, int) or isinstance(output_limit, bool) or output_limit <= 0:
-            raise ServiceError("invalid_request", "output_limit must be a positive integer")
-        validator = FrameSequenceValidator()
-
-        def chunks():
-            while True:
-                frame = read_frame(self.rfile)
-                validator.accept(frame)
-                if frame.stream_id != stream_id:
-                    raise ServiceError("invalid_request", "exec frame stream_id does not match")
-                if frame.flags & FrameFlags.CANCEL:
-                    raise ServiceError("cancelled", "exec stdin cancelled")
-                if frame.flags & FrameFlags.DATA:
-                    yield frame.payload
-                if frame.flags & FrameFlags.EOF:
-                    return
-
-        result = self.dispatcher.services.exec(
-            _required_param(params, "session_id"),
-            _required_param(params, "command"),
-            stdin=chunks(),
-            timeout=float(timeout),
-            output_limit=output_limit,
-            owner=_optional_owner(params),
-        )
-        return Response.success(
-            request.request_id,
-            {
-                "session_id": result.session_id,
-                "exit_code": result.exit_code,
-                "stdout": _encoded_bytes(result.stdout),
-                "stderr": _encoded_bytes(result.stderr),
-                "timed_out": result.timed_out,
-            },
-        )
 
-    def _handle_download(self, request: Request) -> None:
-        params = request.params
-        stream_id = _required_param(params, "stream_id")
-        chunk_size = params.get("chunk_size", 256 * 1024)
-        if not isinstance(chunk_size, int) or isinstance(chunk_size, bool) or not 4096 <= chunk_size <= 1024 * 1024:
-            raise ServiceError("invalid_request", "chunk_size must be between 4096 and 1048576 bytes")
-        response = Response.success(request.request_id, {"stream_id": stream_id})
-        self.wfile.write(response.to_json().encode("utf-8") + b"\n")
-        self.wfile.flush()
-        digest = hashlib.sha256()
-        transferred = 0
-        sequence = 0
-        try:
-            chunks = self.dispatcher.services.download(
-                _required_param(params, "session_id"),
-                _required_param(params, "remote_path"),
-                chunk_size=chunk_size,
-                owner=_optional_owner(params),
-            )
-            for chunk in chunks:
-                digest.update(chunk)
-                transferred += len(chunk)
-                self.wfile.write(encode_frame(BinaryFrame(stream_id, sequence, FrameFlags.DATA, chunk)))
-                self.wfile.flush()
-                sequence += 1
-            metadata = json.dumps(
-                {"bytes_transferred": transferred, "sha256": digest.hexdigest()},
-                separators=(",", ":"),
-            ).encode("utf-8")
-            self.wfile.write(encode_frame(BinaryFrame(stream_id, sequence, FrameFlags.EOF, metadata)))
-            self.wfile.flush()
-        except ServiceError as exc:
-            metadata = json.dumps({"code": exc.code, "message": str(exc)}, separators=(",", ":")).encode("utf-8")
-            self.wfile.write(encode_frame(BinaryFrame(stream_id, sequence, FrameFlags.CANCEL, metadata)))
-            self.wfile.flush()
-
-
-def make_v1_handler(dispatcher: HostBridgeV1Dispatcher) -> type[HostBridgeV1Handler]:
-    class Handler(HostBridgeV1Handler):
-        pass
-
-    Handler.dispatcher = dispatcher
-    return Handler
-
-
-def build_v1_services(
+def build_runtime(
     config_path: Path | str | None = None,
     *,
     connect_timeout: int = DEFAULT_CONNECT_TIMEOUT,
-) -> HostBridgeServices:
+) -> DaemonRuntime:
     resolved_config = configured_hosts_path(config_path)
     registry = HostRegistry(load_hosts(resolved_config))
     policy = load_policy()
-
-    def transport_factory(host, session_connect_timeout):
-        native_ssh = (
-            host.ssh is not None
-            and not host.login_steps
-            and not host.ssh.extra_args
-            and host.transport in {"auto", "ssh"}
-        )
-        if native_ssh:
-            assert host.ssh is not None
-            return SshTransport.connect(
-                SshConnectionConfig(
-                    host=host.ssh.host,
-                    port=host.ssh.port,
-                    username=host.ssh.username,
-                    client_keys=([str(host.ssh.identity_file)] if host.ssh.identity_file else ()),
-                    tunnel=host.ssh.proxy_jump,
-                    password=host.secrets.get("password"),
-                ),
-                connect_timeout=session_connect_timeout,
-            )
-        return PtyTransport.connect(host, connect_timeout=session_connect_timeout)
-
-    services = HostBridgeServices(
+    tunnel_manager = HostTunnelManager(registry, tunnel_provider_factory)
+    services = AsyncHostBridgeServices(
         registry,
-        transport_factory=transport_factory,
+        tunnel_manager,
         default_connect_timeout=connect_timeout,
         policy=policy,
+        idle_timeout_seconds=policy.idle_timeout_seconds,
         audit_sink=build_audit_sink(policy),
     )
-    if policy.enabled and policy.idle_timeout_seconds > 0:
-        services.start_idle_reaper(timeout=policy.idle_timeout_seconds)
-    return services
+    return DaemonRuntime(services, tunnel_manager)
 
 
 def daemon_dir(path: Path | str | None = None) -> Path:
@@ -566,43 +98,111 @@ def _socket_accepts_connections(path: Path) -> bool:
         return False
 
 
-def run_daemon(
+async def run_daemon_async(
     *,
     socket_file: Path | str | None = None,
     config_path: Path | str | None = None,
 ) -> int:
-    services = build_v1_services(config_path)
+    runtime = build_runtime(config_path)
+    services = runtime.services
+    tunnel_manager = runtime.tunnel_manager
     paths = runtime_paths(socket_file)
     path = paths.socket
-    server: ThreadingUnixRPCServer | None = None
+    shutdown = asyncio.Event()
+    server: MuxDaemonServer | None = None
     owns_runtime_paths = False
+
+    async def request_closed(method: str) -> None:
+        if method == "system.shutdown":
+            shutdown.set()
+
     try:
         _prepare_socket_path(path)
-        server = ThreadingUnixRPCServer(
-            str(path),
-            make_v1_handler(HostBridgeV1Dispatcher(services)),
-            bind_and_activate=False,
+        service = HostBridgeMuxService(services, tunnel_manager=tunnel_manager)
+        server = MuxDaemonServer(
+            path,
+            service.handle_unary_async,
+            stream_handler=service.handle_stream,
+            request_closed_handler=request_closed,
         )
-        server.server_bind()
+        await server.start()
         owns_runtime_paths = True
-        server.server_activate()
-        path.chmod(0o600)
         replace_private_file(paths.pid, str(os.getpid()).encode("ascii"))
         print(f"hostbridge daemon listening on unix://{path}", flush=True)
-        server.serve_forever()
+        await shutdown.wait()
     finally:
-        services.close_all(force=True)
-        services.stop_idle_reaper()
-        if server is not None:
-            server.server_close()
-        if owns_runtime_paths:
-            with suppress(FileNotFoundError):
-                path.unlink()
-            with suppress(FileNotFoundError):
-                paths.pid.unlink()
+        primary_error = sys.exception()
+        cleanup = asyncio.create_task(
+            _cleanup_daemon_resources(
+                server,
+                services,
+                tunnel_manager,
+                path=path,
+                pid_path=paths.pid,
+                owns_runtime_paths=owns_runtime_paths,
+            ),
+            name="hostbridge-daemon-cleanup",
+        )
+        cancellation: asyncio.CancelledError | None = None
+        while not cleanup.done():
+            try:
+                await asyncio.shield(cleanup)
+            except asyncio.CancelledError as exc:
+                cancellation = exc
+        cleanup_errors = cleanup.result()
+        if cancellation is not None:
+            for error in cleanup_errors:
+                cancellation.add_note(f"shutdown cleanup failed: {error!r}")
+            if primary_error is None:
+                raise cancellation
+            primary_error.add_note(f"shutdown was cancelled during cleanup: {cancellation!r}")
+        if cleanup_errors:
+            if primary_error is None:
+                raise cleanup_errors[0]
+            for error in cleanup_errors:
+                primary_error.add_note(f"shutdown cleanup failed: {error!r}")
     return 0
 
 
+async def _cleanup_daemon_resources(
+    server: MuxDaemonServer | None,
+    services: AsyncHostBridgeServices,
+    tunnel_manager: HostTunnelManager,
+    *,
+    path: Path,
+    pid_path: Path,
+    owns_runtime_paths: bool,
+) -> list[BaseException]:
+    cleanup_errors: list[BaseException] = []
+    if server is not None:
+        try:
+            await server.close()
+        except BaseException as exc:
+            cleanup_errors.append(exc)
+    try:
+        await services.close_all(force=True)
+    except BaseException as exc:
+        cleanup_errors.append(exc)
+    try:
+        await tunnel_manager.close()
+    except BaseException as exc:
+        cleanup_errors.append(exc)
+    if owns_runtime_paths:
+        with suppress(FileNotFoundError):
+            path.unlink()
+        with suppress(FileNotFoundError):
+            pid_path.unlink()
+    return cleanup_errors
+
+
+def run_daemon(
+    *,
+    socket_file: Path | str | None = None,
+    config_path: Path | str | None = None,
+) -> int:
+    return asyncio.run(run_daemon_async(socket_file=socket_file, config_path=config_path))
+
+
 def daemon_endpoint() -> str:
     return f"unix://{socket_path()}"
 
@@ -613,7 +213,8 @@ def is_running(*, socket_file: Path | str | None = None) -> bool:
 
 def _probe_daemon(path: Path) -> dict[str, object] | None:
     try:
-        return HostBridgeClient(path, default_timeout=1.0).request("system.hello", {})
+        with HostBridgeClient(path, default_timeout=1.0) as client:
+            return client.request("system.hello", {})
     except HostBridgeError:
         return None
 
@@ -622,7 +223,7 @@ def _daemon_is_compatible(hello: dict[str, object] | None) -> bool:
     return bool(
         hello is not None
         and hello.get("service") == "hostbridge"
-        and hello.get("protocol") == PROTOCOL_VERSION
+        and hello.get("protocol") == MUX_PROTOCOL_VERSION
         and hello.get("version") == __version__
     )
 
@@ -641,16 +242,23 @@ def start_background(
             print(f"hostbridge daemon already running at unix://{path}")
         return 0
     if hello is not None and hello.get("service") == "hostbridge":
+        daemon_pid = _read_daemon_pid(paths.pid)
+        owned_daemon = daemon_pid is not None and _process_matches_daemon(daemon_pid, path)
         try:
-            HostBridgeClient(path, default_timeout=2.0).request("system.shutdown", {})
+            with HostBridgeClient(path, default_timeout=2.0) as client:
+                client.request("system.shutdown", {})
         except HostBridgeError:
             if not quiet:
                 print(f"hostbridge incompatible daemon could not be stopped at unix://{path}", file=sys.stderr)
             return 1
-        stop_deadline = time.time() + 5
-        while _socket_accepts_connections(path) and time.time() < stop_deadline:
-            time.sleep(0.05)
-        if _socket_accepts_connections(path):
+        if owned_daemon and daemon_pid is not None:
+            stopped = _stop_daemon_process(daemon_pid, path)
+        else:
+            stop_deadline = time.monotonic() + _DAEMON_STOP_GRACE_SECONDS
+            while _socket_accepts_connections(path) and time.monotonic() < stop_deadline:
+                time.sleep(_PROCESS_POLL_INTERVAL_SECONDS)
+            stopped = not _socket_accepts_connections(path)
+        if not stopped:
             if not quiet:
                 print(f"hostbridge incompatible daemon did not stop at unix://{path}", file=sys.stderr)
             return 1
@@ -660,7 +268,8 @@ def start_background(
         command.extend(["--config", str(config_path)])
     command.extend(["start", "--foreground", "--socket", str(path)])
     with open_private_append(paths.log) as log_handle:
-        process = subprocess.Popen(
+        # The command is the current HostBridge module argv.
+        process = subprocess.Popen(  # nosec B603
             command,
             stdin=subprocess.DEVNULL,
             stdout=log_handle,
@@ -693,6 +302,89 @@ def _read_daemon_pid(path: Path) -> int | None:
     return value if value > 0 else None
 
 
+def _process_is_alive(pid: int) -> bool:
+    try:
+        os.kill(pid, 0)
+    except ProcessLookupError:
+        return False
+    except PermissionError:
+        return True
+    return True
+
+
+def _process_matches_daemon(pid: int, path: Path) -> bool:
+    try:
+        process = subprocess.run(  # nosec B603 B607
+            ["/bin/ps", "-p", str(pid), "-o", "uid=", "-o", "command="],
+            check=False,
+            capture_output=True,
+            text=True,
+            timeout=1,
+        )
+    except (OSError, subprocess.TimeoutExpired):
+        return False
+    if process.returncode != 0:
+        return False
+    fields = process.stdout.strip().split(maxsplit=1)
+    if len(fields) != 2:
+        return False
+    try:
+        process_uid = int(fields[0])
+        argv = shlex.split(fields[1])
+    except (ValueError, UnicodeError):
+        return False
+    if process_uid != os.getuid():
+        return False
+    try:
+        module_index = next(
+            index for index in range(len(argv) - 2) if argv[index : index + 3] == ["-m", "server_control_mcp", "daemon"]
+        )
+        socket_index = argv.index("--socket", module_index + 3)
+    except (StopIteration, ValueError):
+        return False
+    expected_socket = os.path.abspath(os.path.expanduser(str(path)))
+    actual_socket = os.path.abspath(os.path.expanduser(argv[socket_index + 1])) if socket_index + 1 < len(argv) else ""
+    return (
+        "start" in argv[module_index + 3 : socket_index]
+        and "--foreground" in argv[module_index + 3 : socket_index]
+        and actual_socket == expected_socket
+    )
+
+
+def _wait_for_process_exit(pid: int, timeout: float) -> bool:
+    deadline = time.monotonic() + max(0.0, timeout)
+    while _process_is_alive(pid):
+        remaining = deadline - time.monotonic()
+        if remaining <= 0:
+            return False
+        time.sleep(min(_PROCESS_POLL_INTERVAL_SECONDS, remaining))
+    return True
+
+
+def _stop_daemon_process(pid: int, path: Path) -> bool:
+    if _wait_for_process_exit(pid, _DAEMON_STOP_GRACE_SECONDS):
+        return True
+    if not _process_matches_daemon(pid, path):
+        return False
+    try:
+        os.kill(pid, signal.SIGTERM)
+    except ProcessLookupError:
+        return True
+    except OSError:
+        return False
+    if _wait_for_process_exit(pid, _DAEMON_TERM_SECONDS):
+        return True
+    if not _process_matches_daemon(pid, path):
+        return not _process_is_alive(pid)
+    try:
+        os.kill(pid, signal.SIGKILL)
+    except ProcessLookupError:
+        return True
+    except OSError:
+        return False
+    return _wait_for_process_exit(pid, _DAEMON_KILL_SECONDS)
+
+
 def _terminate_background_process(process: subprocess.Popen[bytes]) -> None:
     if process.poll() is not None:
         return
@@ -705,11 +397,26 @@ def _terminate_background_process(process: subprocess.Popen[bytes]) -> None:
 
 
 def stop_background(*, socket_file: Path | str | None = None) -> int:
-    path = socket_path(socket_file)
+    paths = runtime_paths(socket_file)
+    path = paths.socket
+    daemon_pid = _read_daemon_pid(paths.pid)
+    owned_daemon = daemon_pid is not None and _process_matches_daemon(daemon_pid, path)
     hello = _probe_daemon(path)
     if hello is not None and hello.get("service") == "hostbridge":
-        HostBridgeClient(path, default_timeout=2.0).request("system.shutdown", {})
+        with HostBridgeClient(path, default_timeout=2.0) as client:
+            client.request("system.shutdown", {})
+    elif not owned_daemon:
+        print("hostbridge daemon is not running")
+        return 0
+    if owned_daemon and daemon_pid is not None:
+        stopped = _stop_daemon_process(daemon_pid, path)
+    else:
+        stop_deadline = time.monotonic() + _DAEMON_STOP_GRACE_SECONDS
+        while _socket_accepts_connections(path) and time.monotonic() < stop_deadline:
+            time.sleep(_PROCESS_POLL_INTERVAL_SECONDS)
+        stopped = not _socket_accepts_connections(path)
+    if stopped:
         print("hostbridge daemon stopped")
         return 0
-    print("hostbridge daemon is not running")
-    return 0
+    print(f"hostbridge daemon did not stop at unix://{path}", file=sys.stderr)
+    return 1