@neoline/hostbridge 2.0.4 → 2.0.5

package/src/server_control_mcp/server.py

@@ -1,8 +1,11 @@
 from __future__ import annotations
 
+import asyncio
 import os
 import sys
-from collections.abc import Callable
+import threading
+from collections.abc import AsyncIterator, Callable
+from contextlib import asynccontextmanager
 from dataclasses import asdict
 from pathlib import Path
 from typing import Any, TypeVar
@@ -17,14 +20,39 @@ DEFAULT_CONNECT_TIMEOUT = 30
 DEFAULT_MAX_TRANSFER_BYTES = 100 * 1024 * 1024
 MAX_TRANSFER_BYTES = 1024 * 1024 * 1024
 
-mcp = FastMCP("hostbridge")
+_client_instance: HostBridgeClient | None = None
+_client_lock = threading.Lock()
+
+
+@asynccontextmanager
+async def _mcp_lifespan(_server: FastMCP[Any]) -> AsyncIterator[None]:
+    try:
+        yield None
+    finally:
+        await asyncio.to_thread(_close_client)
+
+
+mcp = FastMCP("hostbridge", lifespan=_mcp_lifespan)
 T = TypeVar("T")
 
 
 def _client() -> HostBridgeClient:
+    global _client_instance
     if daemon.start_background(quiet=True) != 0:
         raise RuntimeError(f"hostbridge daemon failed to start; see {daemon.log_file()}")
-    return HostBridgeClient()
+    with _client_lock:
+        if _client_instance is None or _client_instance.closed:
+            _client_instance = HostBridgeClient()
+        return _client_instance
+
+
+def _close_client() -> None:
+    global _client_instance
+    with _client_lock:
+        client = _client_instance
+        _client_instance = None
+    if client is not None:
+        client.close()
 
 
 def _owner() -> str | None:
@@ -51,18 +79,7 @@ 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 operation()
-    except HostBridgeError as exc:
-        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 operation()
+    return _client().request(method, payload, timeout=timeout)
 
 
 def _mode(value: str | int) -> int:
@@ -87,9 +104,7 @@ def session_open(host_id: str, connect_timeout: int = DEFAULT_CONNECT_TIMEOUT) -
     """Open a persistent session for an allowlisted host."""
     safe_timeout = max(1, min(int(connect_timeout), 180))
     return _call(
-        lambda: _after_daemon_start(
-            lambda: _client().open_session(host_id, owner=_owner(), connect_timeout=safe_timeout)
-        ),
+        lambda: _client().open_session(host_id, owner=_owner(), connect_timeout=safe_timeout),
         lambda value: {"ok": True, "session": asdict(value)},
     )