@neoline/hostbridge 2.0.2 → 2.0.4

package/src/server_control_mcp/daemon.py

@@ -6,6 +6,7 @@ import hashlib
 import json
 import os
 import socket
+import stat
 import subprocess
 import sys
 import threading
@@ -14,10 +15,11 @@ from contextlib import suppress
 from dataclasses import asdict
 from pathlib import Path
 from socketserver import StreamRequestHandler, ThreadingMixIn, UnixStreamServer
+from typing import Any
 
+from . import __version__
 from .client import HostBridgeClient, HostBridgeError
-from .config import load_v1_config
-from .hosts import HostRegistry, load_hosts
+from .hosts import HostRegistry, configured_hosts_path, load_hosts
 from .policy import build_audit_sink, load_policy
 from .protocol import (
     PROTOCOL_VERSION,
@@ -32,17 +34,65 @@ from .protocol import (
     read_frame,
 )
 from .runtime import runtime_paths
+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
 
 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)
@@ -120,6 +170,7 @@ class HostBridgeV1Dispatcher:
         return {
             "service": "hostbridge",
             "protocol": PROTOCOL_VERSION,
+            "version": __version__,
             "capabilities": ["sessions", "exec", "explicit_stdin"],
         }
 
@@ -128,28 +179,38 @@ class HostBridgeV1Dispatcher:
 
     def system_shutdown(self, params: dict[str, object]) -> dict[str, object]:
         self.services.stop_idle_reaper()
-        return self.services.close_all()
+        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]:
-        session = self.services.open_session(_required_param(params, "host_id"), owner=_optional_owner(params))
+        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]:
-        force = params.get("force", False)
-        if not isinstance(force, bool):
-            raise ServiceError("invalid_request", "force must be a boolean")
+        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), force=force)
+        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()
+        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)
@@ -263,6 +324,12 @@ 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:
@@ -272,6 +339,8 @@ class HostBridgeV1Handler(StreamRequestHandler):
                 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
@@ -327,6 +396,49 @@ class HostBridgeV1Handler(StreamRequestHandler):
             {"bytes_transferred": result.bytes_transferred, "sha256": result.sha256},
         )
 
+    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")
@@ -377,17 +489,36 @@ def build_v1_services(
     *,
     connect_timeout: int = DEFAULT_CONNECT_TIMEOUT,
 ) -> HostBridgeServices:
-    if config_path is not None:
-        load_v1_config(config_path)
-    registry = HostRegistry(load_hosts(config_path))
+    resolved_config = configured_hosts_path(config_path)
+    registry = HostRegistry(load_hosts(resolved_config))
     policy = load_policy()
 
-    def transport_factory(host):
-        return PtyTransport.connect(host, connect_timeout=connect_timeout)
+    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(
         registry,
         transport_factory=transport_factory,
+        default_connect_timeout=connect_timeout,
         policy=policy,
         audit_sink=build_audit_sink(policy),
     )
@@ -404,21 +535,22 @@ def socket_path(path: Path | str | None = None) -> Path:
     return runtime_paths(path).socket
 
 
-def pid_file() -> Path:
-    return runtime_paths().pid
+def pid_file(socket_file: Path | str | None = None) -> Path:
+    return runtime_paths(socket_file).pid
 
 
-def log_file() -> Path:
-    return runtime_paths().log
+def log_file(socket_file: Path | str | None = None) -> Path:
+    return runtime_paths(socket_file).log
 
 
 def _prepare_socket_path(path: Path) -> None:
-    path.parent.mkdir(parents=True, exist_ok=True)
-    log_file().parent.mkdir(parents=True, exist_ok=True)
-    pid_file().parent.mkdir(parents=True, exist_ok=True)
-    with suppress(OSError):
-        path.parent.chmod(0o700)
-    if path.exists():
+    runtime_paths(path).ensure_directory()
+    if os.path.lexists(path):
+        metadata = path.lstat()
+        if stat.S_ISLNK(metadata.st_mode) or not stat.S_ISSOCK(metadata.st_mode):
+            raise RuntimeError(f"refusing to remove non-socket runtime path: {path}")
+        if metadata.st_uid != os.getuid():
+            raise PermissionError(f"runtime socket is not owned by the current user: {path}")
         if _socket_accepts_connections(path):
             raise RuntimeError(f"hostbridge daemon already running at unix://{path}")
         path.unlink()
@@ -440,24 +572,34 @@ def run_daemon(
     config_path: Path | str | None = None,
 ) -> int:
     services = build_v1_services(config_path)
-    path = socket_path(socket_file)
-    _prepare_socket_path(path)
-    server = ThreadingUnixRPCServer(
-        str(path),
-        make_v1_handler(HostBridgeV1Dispatcher(services)),
-    )
-    pid_file().write_text(str(os.getpid()), encoding="utf-8")
+    paths = runtime_paths(socket_file)
+    path = paths.socket
+    server: ThreadingUnixRPCServer | None = None
+    owns_runtime_paths = False
     try:
+        _prepare_socket_path(path)
+        server = ThreadingUnixRPCServer(
+            str(path),
+            make_v1_handler(HostBridgeV1Dispatcher(services)),
+            bind_and_activate=False,
+        )
+        server.server_bind()
+        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()
     finally:
-        services.close_all()
+        services.close_all(force=True)
         services.stop_idle_reaper()
-        server.server_close()
-        with suppress(FileNotFoundError):
-            path.unlink()
-        with suppress(FileNotFoundError):
-            pid_file().unlink()
+        if server is not None:
+            server.server_close()
+        if owns_runtime_paths:
+            with suppress(FileNotFoundError):
+                path.unlink()
+            with suppress(FileNotFoundError):
+                paths.pid.unlink()
     return 0
 
 
@@ -466,11 +608,23 @@ def daemon_endpoint() -> str:
 
 
 def is_running(*, socket_file: Path | str | None = None) -> bool:
+    return _daemon_is_compatible(_probe_daemon(socket_path(socket_file)))
+
+
+def _probe_daemon(path: Path) -> dict[str, object] | None:
     try:
-        hello = HostBridgeClient(socket_file, default_timeout=1.0).hello()
-        return hello.get("service") == "hostbridge" and hello.get("protocol") == PROTOCOL_VERSION
+        return HostBridgeClient(path, default_timeout=1.0).request("system.hello", {})
     except HostBridgeError:
-        return False
+        return None
+
+
+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("version") == __version__
+    )
 
 
 def start_background(
@@ -479,19 +633,33 @@ def start_background(
     config_path: Path | str | None = None,
     quiet: bool = False,
 ) -> int:
-    path = socket_path(socket_file)
-    if is_running(socket_file=path):
+    paths = runtime_paths(socket_file)
+    path = paths.socket
+    hello = _probe_daemon(path)
+    if _daemon_is_compatible(hello):
         if not quiet:
             print(f"hostbridge daemon already running at unix://{path}")
         return 0
-    path.parent.mkdir(parents=True, exist_ok=True)
-    with suppress(OSError):
-        path.parent.chmod(0o700)
+    if hello is not None and hello.get("service") == "hostbridge":
+        try:
+            HostBridgeClient(path, default_timeout=2.0).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 not quiet:
+                print(f"hostbridge incompatible daemon did not stop at unix://{path}", file=sys.stderr)
+            return 1
+    paths.ensure_directory()
     command = [sys.executable, "-m", "server_control_mcp", "daemon"]
     if config_path is not None:
         command.extend(["--config", str(config_path)])
     command.extend(["start", "--foreground", "--socket", str(path)])
-    with log_file().open("a", encoding="utf-8") as log_handle:
+    with open_private_append(paths.log) as log_handle:
         process = subprocess.Popen(
             command,
             stdin=subprocess.DEVNULL,
@@ -500,24 +668,46 @@ def start_background(
             start_new_session=True,
             env=os.environ.copy(),
         )
-    pid_file().write_text(str(process.pid), encoding="utf-8")
     deadline = time.time() + 5
     while time.time() < deadline:
         if is_running(socket_file=path):
+            if _read_daemon_pid(paths.pid) != process.pid:
+                _terminate_background_process(process)
             if not quiet:
                 print(f"hostbridge daemon started at unix://{path} (pid {process.pid})")
             return 0
         if process.poll() is not None:
             break
         time.sleep(0.1)
+    _terminate_background_process(process)
     if not quiet:
-        print(f"hostbridge daemon failed to start; see {log_file()}", file=sys.stderr)
+        print(f"hostbridge daemon failed to start; see {paths.log}", file=sys.stderr)
     return 1
 
 
+def _read_daemon_pid(path: Path) -> int | None:
+    try:
+        value = int(path.read_text(encoding="ascii").strip())
+    except (FileNotFoundError, OSError, UnicodeError, ValueError):
+        return None
+    return value if value > 0 else None
+
+
+def _terminate_background_process(process: subprocess.Popen[bytes]) -> None:
+    if process.poll() is not None:
+        return
+    process.terminate()
+    try:
+        process.wait(timeout=2)
+    except subprocess.TimeoutExpired:
+        process.kill()
+        process.wait(timeout=2)
+
+
 def stop_background(*, socket_file: Path | str | None = None) -> int:
     path = socket_path(socket_file)
-    if is_running(socket_file=path):
+    hello = _probe_daemon(path)
+    if hello is not None and hello.get("service") == "hostbridge":
         HostBridgeClient(path, default_timeout=2.0).request("system.shutdown", {})
         print("hostbridge daemon stopped")
         return 0

package/src/server_control_mcp/doctor.py

@@ -60,8 +60,10 @@ def _version_diagnostic() -> DiagnosticResult:
         "npm": npm.get("version"),
         "python": pyproject.get("project", {}).get("version"),
     }
-    status = "pass" if len(set(versions.values())) == 1 and __version__ == "2.0.2" else "fail"
-    return DiagnosticResult("version", status, "version sources agree" if status == "pass" else "version drift detected", versions)
+    status = "pass" if len(set(versions.values())) == 1 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:
@@ -69,7 +71,9 @@ def _config_diagnostic(config_path: Path | str) -> DiagnosticResult:
         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)})
+    return DiagnosticResult(
+        "config", "pass", "schema v1 configuration is valid", {"path": str(config.path), "hosts": len(config.hosts)}
+    )
 
 
 def _runtime_diagnostic() -> DiagnosticResult:
@@ -93,7 +97,7 @@ 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
+    valid = hello.get("service") == "hostbridge" and hello.get("protocol") == 1 and hello.get("version") == __version__
     return DiagnosticResult(
         "daemon",
         "pass" if valid else "fail",