Spaces:
Sleeping
Sleeping
File size: 2,430 Bytes
881aff7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | """Synchronous MCP client for Flask (stdio subprocess)."""
from __future__ import annotations
import asyncio
import atexit
import os
import threading
from typing import Any
from backend.mcp_client import PaperMCPClient
_client: PaperMCPClient | None = None
_connect_lock = threading.Lock()
_shutdown_registered = False
_loop: asyncio.AbstractEventLoop | None = None
_loop_thread: threading.Thread | None = None
_loop_ready = threading.Event()
def use_mcp() -> bool:
return os.getenv("USE_MCP", "1").lower() in {"1", "true", "yes"}
def _start_loop_thread() -> asyncio.AbstractEventLoop:
global _loop, _loop_thread
if _loop is not None and _loop.is_running():
return _loop
def _runner() -> None:
global _loop
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
_loop = loop
_loop_ready.set()
loop.run_forever()
_loop_ready.clear()
_loop_thread = threading.Thread(target=_runner, daemon=True, name="mcp-async-loop")
_loop_thread.start()
_loop_ready.wait()
if _loop is None:
raise RuntimeError("MCP event loop failed to start.")
return _loop
def _run_async(coro, *, timeout: float = 600):
loop = _start_loop_thread()
future = asyncio.run_coroutine_threadsafe(coro, loop)
return future.result(timeout=timeout)
def ensure_mcp_connected() -> PaperMCPClient:
global _client, _shutdown_registered
if _client is not None and _client._session is not None:
return _client
with _connect_lock:
if _client is not None and _client._session is not None:
return _client
client = PaperMCPClient()
_run_async(client.connect())
_client = client
if not _shutdown_registered:
atexit.register(shutdown_mcp)
_shutdown_registered = True
return _client
def shutdown_mcp() -> None:
global _client, _loop
if _client is None:
return
try:
_run_async(_client.close(), timeout=30)
except Exception:
pass
_client = None
if _loop is not None and _loop.is_running():
_loop.call_soon_threadsafe(_loop.stop)
def call_mcp_tool(name: str, arguments: dict[str, Any] | None = None) -> Any:
client = ensure_mcp_connected()
return _run_async(client.call_tool(name, arguments))
|