@neoline/hostbridge 2.0.2 → 2.0.4

package/src/server_control_mcp/server.py

@@ -4,6 +4,7 @@ import os
 import sys
 from collections.abc import Callable
 from dataclasses import asdict
+from pathlib import Path
 from typing import Any, TypeVar
 
 from mcp.server.fastmcp import FastMCP
@@ -21,6 +22,8 @@ T = TypeVar("T")
 
 
 def _client() -> HostBridgeClient:
+    if daemon.start_background(quiet=True) != 0:
+        raise RuntimeError(f"hostbridge daemon failed to start; see {daemon.log_file()}")
     return HostBridgeClient()
 
 
@@ -48,14 +51,18 @@ def _request(
     owner = _owner()
     if owner is not None:
         payload.setdefault("owner", owner)
+    return _after_daemon_start(lambda: _client().request(method, payload, timeout=timeout))
+
+
+def _after_daemon_start(operation: Callable[[], T]) -> T:
     try:
-        return _client().request(method, payload, timeout=timeout)
+        return operation()
     except HostBridgeError as exc:
-        if exc.code != "connection_lost" or os.environ.get("HOSTBRIDGE_NPM_LAUNCHER"):
+        if exc.code != "connection_lost":
             raise
         if daemon.start_background(quiet=True) != 0:
             raise RuntimeError(f"hostbridge daemon failed to start; see {daemon.log_file()}") from exc
-        return _client().request(method, payload, timeout=timeout)
+        return operation()
 
 
 def _mode(value: str | int) -> int:
@@ -78,9 +85,11 @@ def hosts_list() -> dict[str, object]:
 @mcp.tool()
 def session_open(host_id: str, connect_timeout: int = DEFAULT_CONNECT_TIMEOUT) -> dict[str, object]:
     """Open a persistent session for an allowlisted host."""
-    _ = max(1, min(int(connect_timeout), 180))
+    safe_timeout = max(1, min(int(connect_timeout), 180))
     return _call(
-        lambda: _client().open_session(host_id, owner=_owner()),
+        lambda: _after_daemon_start(
+            lambda: _client().open_session(host_id, owner=_owner(), connect_timeout=safe_timeout)
+        ),
         lambda value: {"ok": True, "session": asdict(value)},
     )
 
@@ -92,10 +101,10 @@ def sessions_list() -> dict[str, object]:
 
 
 @mcp.tool()
-def session_close(session_id: str, force: bool = False) -> dict[str, object]:
+def session_close(session_id: str) -> dict[str, object]:
     """Close a session."""
     return _call(
-        lambda: _client().close_session(session_id, owner=_owner(), force=bool(force)),
+        lambda: _client().close_session(session_id, owner=_owner()),
         lambda _: {"ok": True, "session_id": session_id, "closed": True},
     )
 
@@ -204,9 +213,18 @@ 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]:
     """Stream a local file to the remote host."""
-    del max_bytes
+    safe_limit = _limit(max_bytes)
+    source = Path(local_path)
+    try:
+        source_size = source.stat().st_size
+    except OSError as exc:
+        return _error(exc)
+    if source_size > safe_limit:
+        return _error(HostBridgeError("resource_limit", f"upload exceeds {safe_limit} bytes"))
     return _call(
-        lambda: _client().upload_file(session_id, local_path, remote_path, mode=_mode(mode), owner=_owner()),
+        lambda: _client().upload_file(
+            session_id, source, remote_path, mode=_mode(mode), owner=_owner(), max_bytes=safe_limit
+        ),
         lambda value: {"ok": True, "transfer": asdict(value)},
     )
 
@@ -216,16 +234,16 @@ def file_download(
     session_id: str, remote_path: str, local_path: str, max_bytes: int = DEFAULT_MAX_TRANSFER_BYTES
 ) -> dict[str, object]:
     """Stream a remote file to a local path."""
-    del max_bytes
+    safe_limit = _limit(max_bytes)
     return _call(
-        lambda: _client().download_file(session_id, remote_path, local_path, owner=_owner()),
+        lambda: _client().download_file(session_id, remote_path, local_path, owner=_owner(), max_bytes=safe_limit),
         lambda value: {"ok": True, "transfer": asdict(value)},
     )
 
 
 def main(argv: list[str] | None = None) -> Any:
     args = sys.argv[1:] if argv is None else argv
-    if args and (args[0] in cli.SETUP_COMMANDS or args[0] in {"-h", "--help"}):
+    if args:
         return cli.main(args)
     return mcp.run(transport="stdio")