@neoline/hostbridge 1.4.4 → 2.0.0

package/src/server_control_mcp/server.py

@@ -2,281 +2,225 @@ from __future__ import annotations
 
 import os
 import sys
-from typing import Any
-from urllib.parse import quote
+from collections.abc import Callable
+from dataclasses import asdict
+from typing import Any, TypeVar
 
 from mcp.server.fastmcp import FastMCP
 
 from . import cli, daemon
-from .manager import (
-    DEFAULT_COMMAND_TIMEOUT,
-    DEFAULT_CONNECT_TIMEOUT,
-    DEFAULT_MAX_TRANSFER_BYTES,
-    MAX_SHELL_TRANSFER_BYTES,
-    PolicyDeniedError,
-    RemoteCommandError,
-)
+from .client import HostBridgeClient, HostBridgeError
+
+DEFAULT_COMMAND_TIMEOUT = 60
+DEFAULT_CONNECT_TIMEOUT = 30
+DEFAULT_MAX_TRANSFER_BYTES = 100 * 1024 * 1024
+MAX_TRANSFER_BYTES = 1024 * 1024 * 1024
 
 mcp = FastMCP("hostbridge")
+T = TypeVar("T")
 
 
-_ERROR_CODES = {
-    PolicyDeniedError: "denied",
-    RemoteCommandError: "remote_error",
-}
+def _client() -> HostBridgeClient:
+    return HostBridgeClient()
 
 
-def _error_code(exc: Exception) -> str:
-    for exc_type, code in _ERROR_CODES.items():
-        if isinstance(exc, exc_type):
-            return code
-    if "unknown host" in str(exc).lower() or "unknown session" in str(exc).lower() or "unknown task" in str(exc).lower():
-        return "not_found"
-    return "internal"
+def _owner() -> str | None:
+    value = os.environ.get("HOSTBRIDGE_AGENT_ID", "").strip()
+    return value or None
 
 
 def _error(exc: Exception) -> dict[str, object]:
-    return {"ok": False, "error": str(exc), "error_code": _error_code(exc)}
+    code = exc.code if isinstance(exc, HostBridgeError) else "internal"
+    return {"ok": False, "error": str(exc), "error_code": code}
 
 
-def _owner_agent_id() -> str | None:
-    """Return the optional caller identity stamped onto daemon-owned sessions."""
-
-    value = os.environ.get("HOSTBRIDGE_AGENT_ID", "").strip()
-    return value or None
+def _call(operation: Callable[[], T], wrap: Callable[[T], dict[str, object]]) -> dict[str, object]:
+    try:
+        return wrap(operation())
+    except Exception as exc:
+        return _error(exc)
 
 
-def _daemon_request(method: str, path: str, payload: dict[str, object] | None = None, *, timeout: float = 5.0) -> dict[str, object]:
+def _request(
+    method: str, params: dict[str, object] | None = None, *, timeout: float | None = None
+) -> dict[str, object]:
+    payload = dict(params or {})
+    owner = _owner()
+    if owner is not None:
+        payload.setdefault("owner", owner)
     try:
-        return daemon.request_json(method, path, payload, timeout=timeout)
-    except OSError as exc:
-        if os.environ.get("HOSTBRIDGE_NPM_LAUNCHER"):
-            socket_path = daemon.socket_path()
-            raise RuntimeError(
-                "hostbridge daemon is not reachable at "
-                f"unix://{socket_path}. Start it outside the MCP sandbox, for example: "
-                f"HOSTBRIDGE_DAEMON_SOCKET={socket_path} npx -y @neoline/hostbridge daemon start --foreground"
-            ) from exc
-        # MCP stdout must stay protocol-clean, so auto-start the daemon quietly.
+        return _client().request(method, payload, timeout=timeout)
+    except HostBridgeError as exc:
+        if exc.code != "connection_lost" or os.environ.get("HOSTBRIDGE_NPM_LAUNCHER"):
+            raise
         if daemon.start_background(quiet=True) != 0:
             raise RuntimeError(f"hostbridge daemon failed to start; see {daemon.log_file()}") from exc
-        return daemon.request_json(method, path, payload, timeout=timeout)
-
-
-def _daemon_payload(**payload: object) -> dict[str, object]:
-    owner = _owner_agent_id()
-    if owner is not None:
-        payload["owner_agent_id"] = owner
-    return payload
+        return _client().request(method, payload, timeout=timeout)
 
 
-def _qid(value: str) -> str:
-    return quote(value, safe="")
+def _mode(value: str | int) -> int:
+    mode = int(value, 8) if isinstance(value, str) else int(value)
+    if not 0 <= mode <= 0o7777:
+        raise ValueError("mode must be between 0000 and 7777")
+    return mode
 
 
-def _safe_transfer_limit(max_bytes: int) -> int:
-    return max(1, min(int(max_bytes), MAX_SHELL_TRANSFER_BYTES))
+def _limit(value: int) -> int:
+    return max(1, min(int(value), MAX_TRANSFER_BYTES))
 
 
 @mcp.tool()
 def hosts_list() -> dict[str, object]:
-    """List approved connection targets and local metadata when available."""
-    try:
-        return _daemon_request("POST", "/hosts/reload")
-    except Exception as exc:
-        return _error(exc)
-
-
-@mcp.tool()
-def hosts_reload() -> dict[str, object]:
-    """Reload approved connection targets from hosts.json without restarting."""
-    try:
-        return _daemon_request("POST", "/hosts/reload")
-    except Exception as exc:
-        return _error(exc)
+    """List configured allowlisted hosts."""
+    return _call(lambda: _request("host.list"), lambda value: {"ok": True, **value})
 
 
 @mcp.tool()
 def session_open(host_id: str, connect_timeout: int = DEFAULT_CONNECT_TIMEOUT) -> dict[str, object]:
-    """Open a persistent interactive shell for one approved host."""
-    try:
-        timeout = max(5, min(int(connect_timeout), 180))
-        return _daemon_request("POST", "/sessions", _daemon_payload(host_id=host_id, connect_timeout=timeout), timeout=timeout + 5)
-    except Exception as exc:
-        return _error(exc)
+    """Open a persistent session for an allowlisted host."""
+    _ = max(1, min(int(connect_timeout), 180))
+    return _call(
+        lambda: _client().open_session(host_id, owner=_owner()),
+        lambda value: {"ok": True, "session": asdict(value)},
+    )
 
 
 @mcp.tool()
 def sessions_list() -> dict[str, object]:
-    """List currently open daemon-owned sessions."""
-    try:
-        return _daemon_request("GET", "/sessions")
-    except Exception as exc:
-        return _error(exc)
+    """List live sessions."""
+    return _call(lambda: _request("session.list"), lambda value: {"ok": True, **value})
 
 
 @mcp.tool()
 def session_close(session_id: str, force: bool = False) -> dict[str, object]:
-    """Close one persistent session and free its PTY."""
-    try:
-        return _daemon_request("DELETE", f"/sessions/{_qid(session_id)}", _daemon_payload(force=bool(force)))
-    except Exception as exc:
-        return _error(exc)
+    """Close a session."""
+    return _call(
+        lambda: _client().close_session(session_id, owner=_owner(), force=bool(force)),
+        lambda _: {"ok": True, "session_id": session_id, "closed": True},
+    )
 
 
 @mcp.tool()
 def sessions_close_all() -> dict[str, object]:
-    """Close every live session. Use during graceful shutdown or reset."""
-    try:
-        return _daemon_request("POST", "/sessions/close_all")
-    except Exception as exc:
-        return _error(exc)
+    """Close all sessions owned by this caller."""
+    return _call(lambda: _request("session.close_all"), lambda value: {"ok": True, **value})
 
 
 @mcp.tool()
 def command_run(session_id: str, command: str, timeout: int = DEFAULT_COMMAND_TIMEOUT) -> dict[str, object]:
-    """Run a foreground command on an open session and wait for completion or timeout."""
-    try:
-        safe_timeout = max(1, min(int(timeout), 24 * 60 * 60))
-        return _daemon_request(
-            "POST",
-            f"/sessions/{_qid(session_id)}/commands",
-            _daemon_payload(command=command, timeout=safe_timeout),
-            timeout=safe_timeout + 5,
-        )
-    except Exception as exc:
-        return _error(exc)
+    """Run a foreground command with bounded output."""
+    safe_timeout = max(1, min(int(timeout), 24 * 60 * 60))
+    return _call(
+        lambda: _client().exec(session_id, command, timeout=safe_timeout, owner=_owner()),
+        lambda value: {
+            "ok": True,
+            "result": {
+                "session_id": value.session_id,
+                "exit_code": value.exit_code,
+                "stdout": value.stdout.decode("utf-8", errors="replace"),
+                "stderr": value.stderr.decode("utf-8", errors="replace"),
+                "timed_out": value.timed_out,
+            },
+        },
+    )
 
 
 @mcp.tool()
 def task_start(session_id: str, command: str) -> dict[str, object]:
-    """Start a long-running background task on the remote host."""
-    try:
-        return _daemon_request("POST", f"/sessions/{_qid(session_id)}/tasks", _daemon_payload(command=command))
-    except Exception as exc:
-        return _error(exc)
+    """Start a long-running command."""
+    return _call(
+        lambda: _request("task.start", {"session_id": session_id, "command": command}),
+        lambda value: {"ok": True, "task": value},
+    )
 
 
 @mcp.tool()
 def task_status(session_id: str, task_id: str, tail_lines: int = 80) -> dict[str, object]:
-    """Poll a background task: state, exit code, and recent log tail."""
-    try:
-        return _daemon_request(
-            "GET",
-            f"/sessions/{_qid(session_id)}/tasks/{_qid(task_id)}?tail_lines={int(tail_lines)}",
-            _daemon_payload(),
-        )
-    except Exception as exc:
-        return _error(exc)
+    """Get task state and recent output."""
+    return _call(
+        lambda: _request(
+            "task.status",
+            {"session_id": session_id, "task_id": task_id, "tail_lines": max(0, min(int(tail_lines), 10000))},
+        ),
+        lambda value: {"ok": True, "status": value},
+    )
 
 
 @mcp.tool()
 def task_stop(session_id: str, task_id: str) -> dict[str, object]:
-    """Best-effort terminate a background task by sending SIGTERM to its remote PID."""
-    try:
-        return _daemon_request("DELETE", f"/sessions/{_qid(session_id)}/tasks/{_qid(task_id)}", _daemon_payload())
-    except Exception as exc:
-        return _error(exc)
+    """Cancel a background task."""
+    return _call(
+        lambda: _request("task.cancel", {"session_id": session_id, "task_id": task_id}),
+        lambda value: {"ok": True, **value},
+    )
 
 
 @mcp.tool()
 def session_write(session_id: str, text: str) -> dict[str, object]:
-    """Send one raw line to a session for rare interactive operations; prefer command_run otherwise."""
-    try:
-        if not text.strip():
-            raise RemoteCommandError("text must not be empty")
-        return _daemon_request("POST", f"/sessions/{_qid(session_id)}/write", _daemon_payload(text=text))
-    except Exception as exc:
-        return _error(exc)
+    """Write UTF-8 data to an interactive shell."""
+    if not text:
+        return _error(ValueError("text must not be empty"))
+    return _call(
+        lambda: _client().shell_write(session_id, text.encode("utf-8"), owner=_owner()),
+        lambda value: {"ok": True, "written": value},
+    )
 
 
 @mcp.tool()
 def session_read(session_id: str, timeout: float = 1.0, max_chars: int = 20000) -> dict[str, object]:
-    """Read currently available raw output from a session after session_write or interactive commands."""
-    try:
-        safe_timeout = max(0.0, min(float(timeout), 10.0))
-        safe_max = max(1, min(int(max_chars), 120000))
-        return _daemon_request(
-            "POST",
-            f"/sessions/{_qid(session_id)}/read",
-            _daemon_payload(timeout=safe_timeout, max_chars=safe_max),
-            timeout=safe_timeout + 5,
-        )
-    except Exception as exc:
-        return _error(exc)
+    """Read available UTF-8 output from an interactive shell."""
+    safe_timeout = max(0.0, min(float(timeout), 10.0))
+    safe_max = max(1, min(int(max_chars), 120000))
+    return _call(
+        lambda: _client().shell_read(session_id, timeout=safe_timeout, max_bytes=safe_max, owner=_owner()),
+        lambda value: {"ok": True, "text": value.data.decode("utf-8", errors="replace"), "alive": value.alive},
+    )
 
 
 @mcp.tool()
 def file_write_text(
-    session_id: str,
-    remote_path: str,
-    content: str,
-    mode: str = "0644",
-    max_bytes: int = DEFAULT_MAX_TRANSFER_BYTES,
+    session_id: str, remote_path: str, content: str, mode: str = "0644", max_bytes: int = DEFAULT_MAX_TRANSFER_BYTES
 ) -> dict[str, object]:
-    """Write a UTF-8 text file through the active shell session; supports up to 100MiB via chunked base64."""
-    try:
-        safe_max = _safe_transfer_limit(max_bytes)
-        return _daemon_request(
-            "POST",
-            f"/sessions/{_qid(session_id)}/files",
-            _daemon_payload(action="write_text", remote_path=remote_path, content=content, mode=mode, max_bytes=safe_max),
-        )
-    except Exception as exc:
-        return _error(exc)
+    """Write a UTF-8 remote file atomically."""
+    return _call(
+        lambda: _client().write_text(
+            session_id, remote_path, content, mode=_mode(mode), max_bytes=_limit(max_bytes), owner=_owner()
+        ),
+        lambda value: {"ok": True, "transfer": asdict(value)},
+    )
 
 
 @mcp.tool()
 def file_read_text(session_id: str, remote_path: str, max_bytes: int = DEFAULT_MAX_TRANSFER_BYTES) -> dict[str, object]:
-    """Read a UTF-8 text file through the active shell session; supports up to 100MiB via chunked base64."""
-    try:
-        safe_max = _safe_transfer_limit(max_bytes)
-        return _daemon_request(
-            "POST",
-            f"/sessions/{_qid(session_id)}/files",
-            _daemon_payload(action="read_text", remote_path=remote_path, max_bytes=safe_max),
-        )
-    except Exception as exc:
-        return _error(exc)
+    """Read a bounded UTF-8 remote file."""
+    return _call(
+        lambda: _client().read_text(session_id, remote_path, max_bytes=_limit(max_bytes), owner=_owner()),
+        lambda value: {"ok": True, "content": value},
+    )
 
 
 @mcp.tool()
 def file_upload(
-    session_id: str,
-    local_path: str,
-    remote_path: str,
-    mode: str = "0644",
-    max_bytes: int = DEFAULT_MAX_TRANSFER_BYTES,
+    session_id: str, local_path: str, remote_path: str, mode: str = "0644", max_bytes: int = DEFAULT_MAX_TRANSFER_BYTES
 ) -> dict[str, object]:
-    """Upload a local file through the active shell session; supports up to 100MiB via chunked base64."""
-    try:
-        safe_max = _safe_transfer_limit(max_bytes)
-        return _daemon_request(
-            "POST",
-            f"/sessions/{_qid(session_id)}/files",
-            _daemon_payload(action="upload", local_path=local_path, remote_path=remote_path, mode=mode, max_bytes=safe_max),
-        )
-    except Exception as exc:
-        return _error(exc)
+    """Stream a local file to the remote host."""
+    del max_bytes
+    return _call(
+        lambda: _client().upload_file(session_id, local_path, remote_path, mode=_mode(mode), owner=_owner()),
+        lambda value: {"ok": True, "transfer": asdict(value)},
+    )
 
 
 @mcp.tool()
 def file_download(
-    session_id: str,
-    remote_path: str,
-    local_path: str,
-    max_bytes: int = DEFAULT_MAX_TRANSFER_BYTES,
+    session_id: str, remote_path: str, local_path: str, max_bytes: int = DEFAULT_MAX_TRANSFER_BYTES
 ) -> dict[str, object]:
-    """Download a remote file through the active shell session; supports up to 100MiB via chunked base64."""
-    try:
-        safe_max = _safe_transfer_limit(max_bytes)
-        return _daemon_request(
-            "POST",
-            f"/sessions/{_qid(session_id)}/files",
-            _daemon_payload(action="download", remote_path=remote_path, local_path=local_path, max_bytes=safe_max),
-        )
-    except Exception as exc:
-        return _error(exc)
+    """Stream a remote file to a local path."""
+    del max_bytes
+    return _call(
+        lambda: _client().download_file(session_id, remote_path, local_path, owner=_owner()),
+        lambda value: {"ok": True, "transfer": asdict(value)},
+    )
 
 
 def main(argv: list[str] | None = None) -> Any: