132 lines
4.4 KiB
Python
132 lines
4.4 KiB
Python
"""MCP tools hosted once inside the user-started web service."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from functools import wraps
|
|
import hmac
|
|
import inspect
|
|
import ipaddress
|
|
import logging
|
|
import threading
|
|
from typing import get_type_hints
|
|
|
|
from starlette.responses import JSONResponse
|
|
|
|
from artemis.mcp.tool_server import ArtemisFastMCP
|
|
from artemis.service_registry import TOKEN_HEADER
|
|
|
|
|
|
class ToolExecutor:
|
|
"""One shared tool loop; synchronous legacy HTTP calls cannot block the UI.
|
|
|
|
Device clients keep their loop affinity across MCP connections. Cancellation
|
|
propagates to async tools; dispatched synchronous writes may finish in their
|
|
worker thread, so callers must resolve the original request rather than retry it.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.loop = None
|
|
self.thread = None
|
|
|
|
def start(self):
|
|
ready = threading.Event()
|
|
|
|
def run():
|
|
self.loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(self.loop)
|
|
self.loop.call_soon(ready.set)
|
|
self.loop.run_forever()
|
|
self.loop.close()
|
|
|
|
self.thread = threading.Thread(target=run, name="artemis-shared-mcp", daemon=True)
|
|
self.thread.start()
|
|
ready.wait()
|
|
|
|
async def invoke(self, fn, *args, **kwargs):
|
|
if self.loop is None or not self.loop.is_running():
|
|
raise RuntimeError("ARTEMIS 主程序正在启动或关闭，请稍后重试。")
|
|
|
|
async def dispatch():
|
|
if inspect.iscoroutinefunction(fn):
|
|
return await fn(*args, **kwargs)
|
|
# Legacy tools use synchronous HTTP to call this same web service.
|
|
# Execute them off both event loops, and never call them speculatively
|
|
# to discover whether their return value is awaitable.
|
|
result = await asyncio.to_thread(fn, *args, **kwargs)
|
|
return await result if inspect.isawaitable(result) else result
|
|
|
|
future = asyncio.run_coroutine_threadsafe(dispatch(), self.loop)
|
|
return await asyncio.wrap_future(future)
|
|
|
|
async def stop(self):
|
|
if self.loop is None or not self.loop.is_running():
|
|
return
|
|
|
|
async def cancel_pending():
|
|
pending = [task for task in asyncio.all_tasks() if task is not asyncio.current_task()]
|
|
for task in pending:
|
|
task.cancel()
|
|
if pending:
|
|
await asyncio.gather(*pending, return_exceptions=True)
|
|
|
|
future = asyncio.run_coroutine_threadsafe(cancel_pending(), self.loop)
|
|
try:
|
|
await asyncio.wait_for(asyncio.wrap_future(future), timeout=5)
|
|
except TimeoutError:
|
|
logging.getLogger(__name__).warning("Timed out draining shared MCP tools on shutdown")
|
|
finally:
|
|
self.loop.call_soon_threadsafe(self.loop.stop)
|
|
await asyncio.to_thread(self.thread.join, 2)
|
|
|
|
|
|
def create_shared_mcp():
|
|
from mcp_server.base import mcp as definitions
|
|
from mcp_server import tools
|
|
|
|
executor = ToolExecutor()
|
|
shared = ArtemisFastMCP(
|
|
"artemis",
|
|
instructions=definitions.instructions,
|
|
streamable_http_path="/",
|
|
json_response=True,
|
|
)
|
|
|
|
def wrap(fn):
|
|
@wraps(fn)
|
|
async def call(*args, **kwargs):
|
|
return await executor.invoke(fn, *args, **kwargs)
|
|
|
|
call.__annotations__ = get_type_hints(fn)
|
|
call.__signature__ = inspect.signature(fn, eval_str=True)
|
|
return call
|
|
|
|
for name in tools.__all__:
|
|
shared.add_tool(wrap(getattr(tools, name)))
|
|
return shared, executor
|
|
|
|
|
|
class LocalMcpBoundary:
|
|
"""Only authenticated loopback bridges may enter the shared tool runtime."""
|
|
|
|
def __init__(self, app, token: str):
|
|
self.app = app
|
|
self.token = token
|
|
|
|
async def __call__(self, scope, receive, send):
|
|
if scope["type"] != "http":
|
|
await self.app(scope, receive, send)
|
|
return
|
|
headers = dict(scope.get("headers", []))
|
|
supplied = headers.get(TOKEN_HEADER.lower().encode(), b"").decode("latin-1")
|
|
try:
|
|
local = ipaddress.ip_address((scope.get("client") or ("",))[0]).is_loopback
|
|
except ValueError:
|
|
local = False
|
|
if not local or not hmac.compare_digest(supplied, self.token):
|
|
await JSONResponse({"detail": "MCP 仅接受本机 ARTEMIS 连接。"}, status_code=403)(
|
|
scope, receive, send
|
|
)
|
|
return
|
|
await self.app(scope, receive, send)
|