@neoline/hostbridge 0.2.2 → 1.1.0

package/src/server_control_mcp/server.py

@@ -1,28 +1,23 @@
 from __future__ import annotations
 
+import os
 import sys
-from dataclasses import asdict
 from typing import Any
+from urllib.parse import quote
 
 from mcp.server.fastmcp import FastMCP
 
-from . import cli
-from .hosts import HostRegistry, load_hosts
+from . import cli, daemon
 from .manager import (
     DEFAULT_COMMAND_TIMEOUT,
     DEFAULT_CONNECT_TIMEOUT,
+    DEFAULT_MAX_TRANSFER_BYTES,
+    MAX_SHELL_TRANSFER_BYTES,
     PolicyDeniedError,
     RemoteCommandError,
-    SessionManager,
 )
-from .policy import build_audit_sink, load_policy
 
 mcp = FastMCP("hostbridge")
-_registry = HostRegistry(load_hosts())
-_policy = load_policy()
-_audit = build_audit_sink(_policy)
-_manager = SessionManager(_registry, policy=_policy, audit_sink=_audit)
-_manager.start_idle_reaper()
 
 
 _ERROR_CODES = {
@@ -40,200 +35,239 @@ def _error_code(exc: Exception) -> str:
     return "internal"
 
 
-def _reload_hosts() -> HostRegistry:
-    """Reload host allowlist without dropping existing remote sessions."""
+def _error(exc: Exception) -> dict[str, object]:
+    return {"ok": False, "error": str(exc), "error_code": _error_code(exc)}
 
-    global _registry
-    _registry = HostRegistry(load_hosts())
-    _manager.hosts = _registry
-    return _registry
 
+def _owner_agent_id() -> str | None:
+    """Return the optional caller identity stamped onto daemon-owned sessions."""
 
-def _error(exc: Exception) -> dict[str, object]:
-    return {"ok": False, "error": str(exc), "error_code": _error_code(exc)}
+    value = os.environ.get("HOSTBRIDGE_AGENT_ID", "").strip()
+    return value or None
+
+
+def _daemon_request(method: str, path: str, payload: dict[str, object] | None = None, *, timeout: float = 5.0) -> dict[str, object]:
+    try:
+        return daemon.request_json(method, path, payload, timeout=timeout)
+    except OSError as exc:
+        # MCP stdout must stay protocol-clean, so auto-start the daemon quietly.
+        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 _ok(**payload: object) -> dict[str, object]:
-    return {"ok": True, **payload}
+def _daemon_payload(**payload: object) -> dict[str, object]:
+    owner = _owner_agent_id()
+    if owner is not None:
+        payload["owner_agent_id"] = owner
+    return payload
+
+
+def _qid(value: str) -> str:
+    return quote(value, safe="")
+
+
+def _safe_transfer_limit(max_bytes: int) -> int:
+    return max(1, min(int(max_bytes), MAX_SHELL_TRANSFER_BYTES))
 
 
 @mcp.tool()
 def hosts_list() -> dict[str, object]:
-    """List approved connection targets and local metadata when available.
-
-    Returns an object with ``hosts``: an array of host descriptors. Safe to call
-    repeatedly; it triggers a config reload from hosts.json.
-    """
+    """List approved connection targets and local metadata when available."""
     try:
-        registry = _reload_hosts()
-        return _ok(hosts=registry.describe())
+        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.
-
-    Returns the same shape as ``hosts_list``. Use this when you know the config
-    file changed and want a fresh descriptor immediately.
-    """
+    """Reload approved connection targets from hosts.json without restarting."""
     try:
-        registry = _reload_hosts()
-        return _ok(hosts=registry.describe())
+        return _daemon_request("POST", "/hosts/reload")
     except Exception as exc:
         return _error(exc)
 
 
 @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.
-
-    Args:
-        host_id: identifier of a host in hosts.json.
-        connect_timeout: seconds to wait for the shell to become ready (5-180).
-
-    Returns the opened session descriptor. Keep the ``session_id`` for later calls.
-    """
+    """Open a persistent interactive shell for one approved host."""
     try:
-        _reload_hosts()
         timeout = max(5, min(int(connect_timeout), 180))
-        session = _manager.open_session(host_id, connect_timeout=timeout)
-        sessions = [item for item in _manager.list_sessions() if item["session_id"] == session.id]
-        return _ok(session=sessions[0])
+        return _daemon_request("POST", "/sessions", _daemon_payload(host_id=host_id, connect_timeout=timeout), timeout=timeout + 5)
     except Exception as exc:
         return _error(exc)
 
 
 @mcp.tool()
 def sessions_list() -> dict[str, object]:
-    """List currently open sessions held by this MCP process."""
-    return _ok(sessions=_manager.list_sessions())
+    """List currently open daemon-owned sessions."""
+    try:
+        return _daemon_request("GET", "/sessions")
+    except Exception as exc:
+        return _error(exc)
 
 
 @mcp.tool()
-def session_close(session_id: str) -> dict[str, object]:
-    """Close one persistent session and free its PTY.
-
-    Args:
-        session_id: value returned by ``session_open``.
-    """
+def session_close(session_id: str, force: bool = False) -> dict[str, object]:
+    """Close one persistent session and free its PTY."""
     try:
-        return _ok(**_manager.close_session(session_id))
+        return _daemon_request("DELETE", f"/sessions/{_qid(session_id)}", _daemon_payload(force=bool(force)))
     except Exception as exc:
         return _error(exc)
 
 
 @mcp.tool()
 def sessions_close_all() -> dict[str, object]:
-    """Close every live session. Use during graceful shutdown or reset.
-
-    Returns ``{"closed": [...session_ids], "count": N}``.
-    """
+    """Close every live session. Use during graceful shutdown or reset."""
     try:
-        return _ok(**_manager.close_all_sessions())
+        return _daemon_request("POST", "/sessions/close_all")
     except Exception as exc:
         return _error(exc)
 
 
 @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.
-
-    Args:
-        session_id: an open session from ``session_open``.
-        command: the shell command string to execute.
-        timeout: max seconds to wait (1 - 86400). On timeout the result has ``timed_out=true``.
-
-    The command is checked against the configured denylist before execution;
-    a blocked command returns ``error_code="denied"``.
-    """
+    """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))
-        result = _manager.run_command(session_id, command, timeout=safe_timeout)
-        return _ok(result=asdict(result))
+        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)
 
 
 @mcp.tool()
 def task_start(session_id: str, command: str) -> dict[str, object]:
-    """Start a long-running background task on the remote host.
-
-    Args:
-        session_id: an open session.
-        command: the shell command to run detached on the remote host.
-
-    Returns a task descriptor with ``task_id`` for polling via ``task_status``.
-    The command is checked against the configured denylist before execution.
-    """
+    """Start a long-running background task on the remote host."""
     try:
-        task = _manager.start_task(session_id, command)
-        return _ok(task=asdict(task))
+        return _daemon_request("POST", f"/sessions/{_qid(session_id)}/tasks", _daemon_payload(command=command))
     except Exception as exc:
         return _error(exc)
 
 
 @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.
-
-    Args:
-        session_id: the session that owns the task.
-        task_id: value returned by ``task_start``.
-        tail_lines: number of trailing log lines to return (1 - 500).
-    """
+    """Poll a background task: state, exit code, and recent log tail."""
     try:
-        return _ok(status=_manager.task_status(session_id, task_id, tail_lines=tail_lines))
+        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)
 
 
 @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.
-
-    Args:
-        session_id: the session that owns the task.
-        task_id: value returned by ``task_start``.
-
-    Returns ``{"task_id": ..., "stopped": bool, "outcome": "killed|already|noop"}``.
-    """
+    """Best-effort terminate a background task by sending SIGTERM to its remote PID."""
     try:
-        return _ok(**_manager.stop_task(session_id, task_id))
+        return _daemon_request("DELETE", f"/sessions/{_qid(session_id)}/tasks/{_qid(task_id)}", _daemon_payload())
     except Exception as exc:
         return _error(exc)
 
 
 @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.
-
-    Args:
-        session_id: an open session.
-        text: the literal line to send. Must not be empty.
-    """
+    """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 _ok(**_manager.write(session_id, text))
+        return _daemon_request("POST", f"/sessions/{_qid(session_id)}/write", _daemon_payload(text=text))
     except Exception as exc:
         return _error(exc)
 
 
 @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.
-
-    Args:
-        session_id: an open session.
-        timeout: max seconds to block waiting for output (0.0 - 10.0).
-        max_chars: cap on returned characters (1 - 120000).
-    """
+    """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 _ok(**_manager.read(session_id, timeout=safe_timeout, max_chars=safe_max))
+        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)
+
+
+@mcp.tool()
+def file_write_text(
+    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)
+
+
+@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)
+
+
+@mcp.tool()
+def file_upload(
+    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)
+
+
+@mcp.tool()
+def file_download(
+    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)