@neoline/hostbridge 0.2.2 → 1.1.0

package/src/server_control_mcp/daemon.py

@@ -0,0 +1,420 @@
+from __future__ import annotations
+
+import json
+import os
+import signal
+import socket
+import subprocess
+import sys
+import threading
+import time
+from contextlib import suppress
+from dataclasses import asdict
+from pathlib import Path
+from socketserver import StreamRequestHandler, ThreadingMixIn, UnixStreamServer
+from typing import Any
+from urllib.parse import parse_qs, urlparse
+
+from .hosts import HostRegistry, load_hosts
+from .manager import (
+    DEFAULT_COMMAND_TIMEOUT,
+    DEFAULT_CONNECT_TIMEOUT,
+    DEFAULT_MAX_TRANSFER_BYTES,
+    PolicyDeniedError,
+    RemoteCommandError,
+    SessionManager,
+    _new_token,
+)
+from .policy import build_audit_sink, load_policy
+
+DEFAULT_DAEMON_DIR = Path("~/.hostbridge").expanduser()
+DEFAULT_SOCKET_FILE = DEFAULT_DAEMON_DIR / "daemon.sock"
+PID_FILE = DEFAULT_DAEMON_DIR / "daemon.pid"
+LOG_FILE = DEFAULT_DAEMON_DIR / "daemon.log"
+
+_ERROR_CODES = {
+    PolicyDeniedError: "denied",
+    RemoteCommandError: "remote_error",
+}
+
+
+class ThreadingUnixRPCServer(ThreadingMixIn, UnixStreamServer):
+    daemon_threads = True
+    allow_reuse_address = True
+
+
+def _ok(**payload: object) -> dict[str, object]:
+    return {"ok": True, **payload}
+
+
+def _error_code(exc: Exception) -> str:
+    for exc_type, code in _ERROR_CODES.items():
+        if isinstance(exc, exc_type):
+            return code
+    text = str(exc).lower()
+    if "unknown host" in text or "unknown session" in text or "unknown task" in text:
+        return "not_found"
+    return "internal"
+
+
+def _status_for_error(exc: Exception) -> int:
+    code = _error_code(exc)
+    if code == "not_found":
+        return 404
+    if code in {"denied", "remote_error"}:
+        return 409
+    return 500
+
+
+def _error(exc: Exception) -> dict[str, object]:
+    return {"ok": False, "status": _status_for_error(exc), "error": str(exc), "error_code": _error_code(exc)}
+
+
+class HostBridgeDaemonHandler(StreamRequestHandler):
+    """One JSON request per line, one JSON response per line over a Unix socket."""
+
+    manager: SessionManager
+    registry: HostRegistry
+
+    def handle(self) -> None:
+        for raw_line in self.rfile:
+            try:
+                request = json.loads(raw_line.decode("utf-8"))
+                if not isinstance(request, dict):
+                    raise ValueError("request must be a JSON object")
+                response = self.dispatch(request)
+            except Exception as exc:
+                response = _error(exc)
+            self.wfile.write(json.dumps(response, ensure_ascii=False).encode("utf-8") + b"\n")
+            self.wfile.flush()
+
+    def dispatch(self, request: dict[str, Any]) -> dict[str, object]:
+        method = str(request.get("method", "GET")).upper()
+        path = str(request.get("path", "/"))
+        body = request.get("body", {})
+        if body is None:
+            body = {}
+        if not isinstance(body, dict):
+            raise ValueError("request body must be a JSON object")
+        try:
+            return self._dispatch(method, path, body)
+        except Exception as exc:
+            return _error(exc)
+
+    def _dispatch(self, method: str, path: str, body: dict[str, Any]) -> dict[str, object]:
+        parsed = urlparse(path)
+        if method == "GET" and parsed.path == "/health":
+            return _ok(service="hostbridge-daemon", transport="unix-jsonl")
+        if method == "GET" and parsed.path == "/hosts":
+            return _ok(hosts=self.registry.describe())
+        if method == "POST" and parsed.path == "/hosts/reload":
+            registry = _reload_handler_hosts(self)
+            return _ok(hosts=registry.describe())
+        if method == "GET" and parsed.path == "/sessions":
+            return _ok(sessions=self.manager.list_sessions())
+        if method == "POST" and parsed.path == "/sessions":
+            token = body.get("token")
+            token_factory = (lambda: str(token)) if token else None
+            session = self.manager.open_session(
+                str(body.get("host_id", "")),
+                connect_timeout=int(body.get("connect_timeout", DEFAULT_CONNECT_TIMEOUT)),
+                owner_agent_id=_optional_str(body.get("owner_agent_id")),
+                **({"token_factory": token_factory} if token_factory else {}),
+            )
+            descriptor = next(item for item in self.manager.list_sessions() if item["session_id"] == session.id)
+            return _ok(session=descriptor)
+        if method == "POST" and parsed.path == "/sessions/close_all":
+            return _ok(**self.manager.close_all_sessions())
+        if method == "POST" and parsed.path == "/shutdown":
+            threading.Thread(target=self.server.shutdown, daemon=True).start()
+            return _ok(shutting_down=True)
+
+        parts = parsed.path.strip("/").split("/")
+        if len(parts) == 3 and parts[0] == "sessions" and parts[2] == "commands" and method == "POST":
+            token = body.get("token")
+            token_factory = (lambda: str(token)) if token else _new_token
+            command_result = self.manager.run_command(
+                parts[1],
+                str(body.get("command", "")),
+                timeout=int(body.get("timeout", DEFAULT_COMMAND_TIMEOUT)),
+                owner_agent_id=_optional_str(body.get("owner_agent_id")),
+                token_factory=token_factory,
+            )
+            return _ok(result=asdict(command_result))
+        if len(parts) == 3 and parts[0] == "sessions" and parts[2] == "tasks" and method == "POST":
+            task = self.manager.start_task(
+                parts[1],
+                str(body.get("command", "")),
+                owner_agent_id=_optional_str(body.get("owner_agent_id")),
+            )
+            return _ok(task=asdict(task))
+        if len(parts) == 4 and parts[0] == "sessions" and parts[2] == "tasks" and method == "GET":
+            query = parse_qs(parsed.query)
+            owner_agent_id = _optional_str(body.get("owner_agent_id") or query.get("owner_agent_id", [None])[0])
+            return _ok(
+                status=self.manager.task_status(
+                    parts[1],
+                    parts[3],
+                    tail_lines=int(query.get("tail_lines", [80])[0]),
+                    owner_agent_id=owner_agent_id,
+                )
+            )
+        if len(parts) == 4 and parts[0] == "sessions" and parts[2] == "tasks" and method == "DELETE":
+            return _ok(**self.manager.stop_task(parts[1], parts[3], owner_agent_id=_optional_str(body.get("owner_agent_id")), force=bool(body.get("force"))))
+        if len(parts) == 3 and parts[0] == "sessions" and parts[2] == "files" and method == "POST":
+            action = str(body.get("action", ""))
+            owner = _optional_str(body.get("owner_agent_id"))
+            token = body.get("token")
+            token_factory = (lambda: str(token)) if token else None
+            token_factory = token_factory or _new_token
+            if action == "write_text":
+                return _ok(
+                    transfer=self.manager.file_write_text(
+                        parts[1],
+                        str(body.get("remote_path", "")),
+                        str(body.get("content", "")),
+                        mode=str(body.get("mode", "0644")),
+                        max_bytes=int(body.get("max_bytes", DEFAULT_MAX_TRANSFER_BYTES)),
+                        owner_agent_id=owner,
+                        token_factory=token_factory,
+                    )
+                )
+            if action == "read_text":
+                return _ok(
+                    transfer=self.manager.file_read_text(
+                        parts[1],
+                        str(body.get("remote_path", "")),
+                        max_bytes=int(body.get("max_bytes", DEFAULT_MAX_TRANSFER_BYTES)),
+                        owner_agent_id=owner,
+                        token_factory=token_factory,
+                    )
+                )
+            if action == "upload":
+                return _ok(
+                    transfer=self.manager.file_upload(
+                        parts[1],
+                        str(body.get("local_path", "")),
+                        str(body.get("remote_path", "")),
+                        mode=str(body.get("mode", "0644")),
+                        max_bytes=int(body.get("max_bytes", DEFAULT_MAX_TRANSFER_BYTES)),
+                        owner_agent_id=owner,
+                        token_factory=token_factory,
+                    )
+                )
+            if action == "download":
+                return _ok(
+                    transfer=self.manager.file_download(
+                        parts[1],
+                        str(body.get("remote_path", "")),
+                        str(body.get("local_path", "")),
+                        max_bytes=int(body.get("max_bytes", DEFAULT_MAX_TRANSFER_BYTES)),
+                        owner_agent_id=owner,
+                        token_factory=token_factory,
+                    )
+                )
+            return {"ok": False, "status": 400, "error": "unknown file action", "error_code": "bad_request"}
+        if len(parts) == 3 and parts[0] == "sessions" and parts[2] == "write" and method == "POST":
+            result = self.manager.write(parts[1], str(body.get("text", "")), owner_agent_id=_optional_str(body.get("owner_agent_id")))
+            return _ok(**result)
+        if len(parts) == 3 and parts[0] == "sessions" and parts[2] == "read" and method == "POST":
+            result = self.manager.read(
+                parts[1],
+                timeout=float(body.get("timeout", 1.0)),
+                max_chars=int(body.get("max_chars", 20000)),
+                owner_agent_id=_optional_str(body.get("owner_agent_id")),
+            )
+            return _ok(**result)
+        if len(parts) == 2 and parts[0] == "sessions" and method == "DELETE":
+            query = parse_qs(parsed.query)
+            force = bool(body.get("force")) or query.get("force", [""])[0].lower() in {"1", "true", "yes"}
+            return _ok(
+                **self.manager.close_session(
+                    parts[1],
+                    owner_agent_id=_optional_str(body.get("owner_agent_id")),
+                    force=force,
+                )
+            )
+        return {"ok": False, "status": 404, "error": "not found", "error_code": "not_found"}
+
+
+def _optional_str(value: object) -> str | None:
+    if value is None:
+        return None
+    text = str(value).strip()
+    return text or None
+
+
+def _reload_handler_hosts(handler: HostBridgeDaemonHandler) -> HostRegistry:
+    registry = HostRegistry(load_hosts())
+    handler.manager.hosts = registry
+    type(handler).registry = registry
+    return registry
+
+
+def make_handler(manager: SessionManager, registry: HostRegistry | None = None) -> type[HostBridgeDaemonHandler]:
+    class Handler(HostBridgeDaemonHandler):
+        pass
+
+    Handler.manager = manager
+    Handler.registry = registry or manager.hosts
+    return Handler
+
+
+def _max_sessions_per_host_from_env() -> int:
+    raw = os.environ.get("HOSTBRIDGE_MAX_SESSIONS_PER_HOST", "0")
+    with suppress(ValueError):
+        return max(0, int(raw))
+    return 0
+
+
+def build_manager() -> SessionManager:
+    registry = HostRegistry(load_hosts())
+    policy = load_policy()
+    manager = SessionManager(
+        registry,
+        policy=policy,
+        audit_sink=build_audit_sink(policy),
+        max_sessions_per_host=_max_sessions_per_host_from_env(),
+    )
+    manager.start_idle_reaper()
+    return manager
+
+
+def socket_path(path: Path | str | None = None) -> Path:
+    configured = path or os.environ.get("HOSTBRIDGE_DAEMON_SOCKET") or DEFAULT_SOCKET_FILE
+    return Path(configured).expanduser()
+
+
+def _prepare_socket_path(path: Path) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with suppress(OSError):
+        path.parent.chmod(0o700)
+    if path.exists():
+        if _socket_accepts_connections(path):
+            raise RuntimeError(f"hostbridge daemon already running at unix://{path}")
+        path.unlink()
+
+
+def _socket_accepts_connections(path: Path) -> bool:
+    try:
+        with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client:
+            client.settimeout(0.5)
+            client.connect(str(path))
+        return True
+    except OSError:
+        return False
+
+
+def run_daemon(*, socket_file: Path | str | None = None) -> int:
+    manager = build_manager()
+    path = socket_path(socket_file)
+    _prepare_socket_path(path)
+    server = ThreadingUnixRPCServer(str(path), make_handler(manager))
+    PID_FILE.write_text(str(os.getpid()), encoding="utf-8")
+    try:
+        print(f"hostbridge daemon listening on unix://{path}", flush=True)
+        server.serve_forever()
+    finally:
+        manager.close_all_sessions()
+        manager.stop_idle_reaper()
+        server.server_close()
+        with suppress(FileNotFoundError):
+            path.unlink()
+        with suppress(FileNotFoundError):
+            PID_FILE.unlink()
+    return 0
+
+
+def daemon_endpoint() -> str:
+    return f"unix://{socket_path()}"
+
+
+def is_running(*, socket_file: Path | str | None = None) -> bool:
+    try:
+        health = request_json("GET", "/health", timeout=1.0, socket_file=socket_file)
+        return health.get("ok") is True and health.get("service") == "hostbridge-daemon"
+    except Exception:
+        return False
+
+
+def request_json(
+    method: str,
+    path: str,
+    payload: dict[str, object] | None = None,
+    *,
+    timeout: float = 5.0,
+    socket_file: Path | str | None = None,
+) -> dict[str, object]:
+    request = {"method": method.upper(), "path": path, "body": payload or {}}
+    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client:
+        client.settimeout(timeout)
+        client.connect(str(socket_path(socket_file)))
+        client.sendall(json.dumps(request, ensure_ascii=False).encode("utf-8") + b"\n")
+        response = _recv_json_line(client)
+    if not isinstance(response, dict):
+        raise RuntimeError("daemon returned non-object JSON")
+    return response
+
+
+def _recv_json_line(client: socket.socket) -> dict[str, object]:
+    chunks: list[bytes] = []
+    while True:
+        chunk = client.recv(65536)
+        if not chunk:
+            break
+        chunks.append(chunk)
+        if b"\n" in chunk:
+            break
+    line = b"".join(chunks).split(b"\n", 1)[0]
+    if not line:
+        return {}
+    decoded = json.loads(line.decode("utf-8"))
+    if not isinstance(decoded, dict):
+        raise RuntimeError("daemon returned non-object JSON")
+    return decoded
+
+
+def start_background(*, socket_file: Path | str | None = None, quiet: bool = False) -> int:
+    path = socket_path(socket_file)
+    if is_running(socket_file=path):
+        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)
+    log_handle = LOG_FILE.open("a", encoding="utf-8")
+    process = subprocess.Popen(
+        [sys.executable, "-m", "server_control_mcp", "daemon", "start", "--foreground", "--socket", str(path)],
+        stdin=subprocess.DEVNULL,
+        stdout=log_handle,
+        stderr=log_handle,
+        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 not quiet:
+                print(f"hostbridge daemon started at unix://{path} (pid {process.pid})")
+            return 0
+        time.sleep(0.1)
+    if not quiet:
+        print(f"hostbridge daemon failed to start; see {LOG_FILE}", file=sys.stderr)
+    return 1
+
+
+def stop_background(*, socket_file: Path | str | None = None) -> int:
+    path = socket_path(socket_file)
+    if is_running(socket_file=path):
+        request_json("POST", "/shutdown", timeout=2.0, socket_file=path)
+        print("hostbridge daemon stopped")
+        return 0
+    if PID_FILE.exists():
+        with suppress(Exception):
+            os.kill(int(PID_FILE.read_text(encoding="utf-8").strip()), signal.SIGTERM)
+            print("hostbridge daemon process signaled")
+            return 0
+    print("hostbridge daemon is not running")
+    return 0