wehub-resource-sync
412 lines
16 KiB
Python
412 lines
16 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import Awaitable, Callable, Iterable
|
|
from contextlib import AbstractAsyncContextManager
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
from ..logger import logger
|
|
from .server import MCPServer
|
|
|
|
|
|
@dataclass
|
|
class _ServerCommand:
|
|
action: str
|
|
timeout_seconds: float | None
|
|
future: asyncio.Future[None]
|
|
|
|
|
|
class _ServerWorker:
|
|
def __init__(
|
|
self,
|
|
server: MCPServer,
|
|
connect_timeout_seconds: float | None,
|
|
cleanup_timeout_seconds: float | None,
|
|
) -> None:
|
|
self._server = server
|
|
self._connect_timeout_seconds = connect_timeout_seconds
|
|
self._cleanup_timeout_seconds = cleanup_timeout_seconds
|
|
self._queue: asyncio.Queue[_ServerCommand] = asyncio.Queue()
|
|
self._task = asyncio.create_task(self._run())
|
|
|
|
@property
|
|
def is_done(self) -> bool:
|
|
return self._task.done()
|
|
|
|
async def connect(self) -> None:
|
|
await self._submit("connect", self._connect_timeout_seconds)
|
|
|
|
async def cleanup(self) -> None:
|
|
await self._submit("cleanup", self._cleanup_timeout_seconds)
|
|
|
|
async def _submit(self, action: str, timeout_seconds: float | None) -> None:
|
|
loop = asyncio.get_running_loop()
|
|
future: asyncio.Future[None] = loop.create_future()
|
|
await self._queue.put(
|
|
_ServerCommand(action=action, timeout_seconds=timeout_seconds, future=future)
|
|
)
|
|
await future
|
|
|
|
async def _run(self) -> None:
|
|
while True:
|
|
command = await self._queue.get()
|
|
should_exit = command.action == "cleanup"
|
|
try:
|
|
if command.action == "connect":
|
|
await _run_with_timeout_in_task(self._server.connect, command.timeout_seconds)
|
|
elif command.action == "cleanup":
|
|
await _run_with_timeout_in_task(self._server.cleanup, command.timeout_seconds)
|
|
else:
|
|
raise ValueError(f"Unknown command: {command.action}")
|
|
if not command.future.cancelled():
|
|
command.future.set_result(None)
|
|
except BaseException as exc:
|
|
if not command.future.cancelled():
|
|
command.future.set_exception(exc)
|
|
if should_exit:
|
|
return
|
|
|
|
|
|
async def _run_with_timeout_in_task(
|
|
func: Callable[[], Awaitable[Any]], timeout_seconds: float | None
|
|
) -> None:
|
|
# Use an in-task timeout to preserve task affinity for MCP cleanup.
|
|
# asyncio.wait_for creates a new Task on Python < 3.11, which breaks
|
|
# libraries that require connect/cleanup in the same task (e.g. AnyIO cancel scopes).
|
|
if timeout_seconds is None:
|
|
await func()
|
|
return
|
|
timeout_context = getattr(asyncio, "timeout", None)
|
|
if timeout_context is not None:
|
|
async with timeout_context(timeout_seconds):
|
|
await func()
|
|
return
|
|
task = asyncio.current_task()
|
|
if task is None:
|
|
await asyncio.wait_for(func(), timeout=timeout_seconds)
|
|
return
|
|
timed_out = False
|
|
loop = asyncio.get_running_loop()
|
|
|
|
def _cancel() -> None:
|
|
nonlocal timed_out
|
|
timed_out = True
|
|
task.cancel()
|
|
|
|
handle = loop.call_later(timeout_seconds, _cancel)
|
|
try:
|
|
await func()
|
|
except asyncio.CancelledError as exc:
|
|
if timed_out:
|
|
raise asyncio.TimeoutError() from exc
|
|
raise
|
|
finally:
|
|
handle.cancel()
|
|
|
|
|
|
class MCPServerManager(AbstractAsyncContextManager["MCPServerManager"]):
|
|
"""Manage MCP server lifecycles and expose only connected servers.
|
|
|
|
Use this helper to keep MCP connect/cleanup on the same task and avoid
|
|
run failures when a server is unavailable. The manager will attempt to
|
|
connect each server and then expose the connected subset via
|
|
`active_servers`.
|
|
|
|
Basic usage:
|
|
async with MCPServerManager([server_a, server_b]) as manager:
|
|
agent = Agent(
|
|
name="Assistant",
|
|
instructions="...",
|
|
mcp_servers=manager.active_servers,
|
|
)
|
|
|
|
FastAPI lifespan example:
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
async with MCPServerManager([server_a, server_b]) as manager:
|
|
app.state.mcp_manager = manager
|
|
yield
|
|
|
|
app = FastAPI(lifespan=lifespan)
|
|
|
|
Important behaviors:
|
|
- `active_servers` only includes servers that connected successfully.
|
|
`failed_servers` holds the failures and `errors` maps servers to errors.
|
|
- `drop_failed_servers=True` removes failed servers from `active_servers`
|
|
(recommended). If False, `active_servers` will still include all servers.
|
|
- `strict=True` raises on the first connection failure. If False, failures
|
|
are recorded and the run can proceed with the remaining servers.
|
|
- `reconnect(failed_only=True)` retries failed servers and refreshes
|
|
`active_servers`.
|
|
- `connect_in_parallel=True` uses a dedicated worker task per server to
|
|
allow concurrent connects while preserving task affinity for cleanup.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
servers: Iterable[MCPServer],
|
|
*,
|
|
connect_timeout_seconds: float | None = 10.0,
|
|
cleanup_timeout_seconds: float | None = 10.0,
|
|
drop_failed_servers: bool = True,
|
|
strict: bool = False,
|
|
suppress_cancelled_error: bool = True,
|
|
connect_in_parallel: bool = False,
|
|
) -> None:
|
|
self._all_servers = list(servers)
|
|
self._active_servers = list(servers)
|
|
self.connect_timeout_seconds = connect_timeout_seconds
|
|
self.cleanup_timeout_seconds = cleanup_timeout_seconds
|
|
self.drop_failed_servers = drop_failed_servers
|
|
self.strict = strict
|
|
self.suppress_cancelled_error = suppress_cancelled_error
|
|
self.connect_in_parallel = connect_in_parallel
|
|
self._workers: dict[MCPServer, _ServerWorker] = {}
|
|
|
|
self.failed_servers: list[MCPServer] = []
|
|
self._failed_server_set: set[MCPServer] = set()
|
|
self._connected_servers: set[MCPServer] = set()
|
|
self.errors: dict[MCPServer, BaseException] = {}
|
|
|
|
@property
|
|
def active_servers(self) -> list[MCPServer]:
|
|
"""Return the active MCP servers after connection attempts."""
|
|
return list(self._active_servers)
|
|
|
|
@property
|
|
def all_servers(self) -> list[MCPServer]:
|
|
"""Return all MCP servers managed by this instance."""
|
|
return list(self._all_servers)
|
|
|
|
async def __aenter__(self) -> MCPServerManager:
|
|
await self.connect_all()
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc_val, exc_tb) -> bool | None:
|
|
await self.cleanup_all()
|
|
return None
|
|
|
|
async def connect_all(self) -> list[MCPServer]:
|
|
"""Connect all servers in order and return the active list."""
|
|
previous_connected_servers = set(self._connected_servers)
|
|
previous_active_servers = list(self._active_servers)
|
|
self.failed_servers = []
|
|
self._failed_server_set = set()
|
|
self.errors = {}
|
|
|
|
servers_to_connect = self._servers_to_connect(self._all_servers)
|
|
connected_servers: list[MCPServer] = []
|
|
try:
|
|
if self.connect_in_parallel:
|
|
await self._connect_all_parallel(servers_to_connect)
|
|
else:
|
|
for server in servers_to_connect:
|
|
await self._attempt_connect(server)
|
|
if server not in self._failed_server_set:
|
|
connected_servers.append(server)
|
|
except BaseException:
|
|
if self.connect_in_parallel:
|
|
await self._cleanup_servers(servers_to_connect)
|
|
else:
|
|
servers_to_cleanup = self._unique_servers(
|
|
[*connected_servers, *self.failed_servers]
|
|
)
|
|
await self._cleanup_servers(servers_to_cleanup)
|
|
if self.drop_failed_servers:
|
|
self._active_servers = [
|
|
server for server in self._all_servers if server in previous_connected_servers
|
|
]
|
|
else:
|
|
self._active_servers = previous_active_servers
|
|
raise
|
|
|
|
self._refresh_active_servers()
|
|
|
|
return self._active_servers
|
|
|
|
async def reconnect(self, *, failed_only: bool = True) -> list[MCPServer]:
|
|
"""Reconnect servers and return the active list.
|
|
|
|
Args:
|
|
failed_only: If True, only retry servers that previously failed.
|
|
If False, cleanup and retry all servers.
|
|
"""
|
|
if failed_only:
|
|
servers_to_retry = self._unique_servers(self.failed_servers)
|
|
else:
|
|
await self.cleanup_all()
|
|
servers_to_retry = list(self._all_servers)
|
|
self.failed_servers = []
|
|
self._failed_server_set = set()
|
|
self.errors = {}
|
|
|
|
servers_to_retry = self._servers_to_connect(servers_to_retry)
|
|
try:
|
|
if self.connect_in_parallel:
|
|
await self._connect_all_parallel(servers_to_retry)
|
|
else:
|
|
for server in servers_to_retry:
|
|
await self._attempt_connect(server)
|
|
finally:
|
|
self._refresh_active_servers()
|
|
return self._active_servers
|
|
|
|
async def cleanup_all(self) -> None:
|
|
"""Cleanup all servers in reverse order."""
|
|
for server in reversed(self._all_servers):
|
|
try:
|
|
await self._cleanup_server(server)
|
|
except asyncio.CancelledError as exc:
|
|
if not self.suppress_cancelled_error:
|
|
raise
|
|
logger.debug("Cleanup cancelled for MCP server '%s': %s", server.name, exc)
|
|
self.errors[server] = exc
|
|
except Exception as exc:
|
|
logger.exception("Failed to cleanup MCP server '%s': %s", server.name, exc)
|
|
self.errors[server] = exc
|
|
|
|
async def _run_with_timeout(
|
|
self, func: Callable[[], Awaitable[Any]], timeout_seconds: float | None
|
|
) -> None:
|
|
await _run_with_timeout_in_task(func, timeout_seconds)
|
|
|
|
async def _attempt_connect(
|
|
self, server: MCPServer, *, raise_on_error: bool | None = None
|
|
) -> None:
|
|
if raise_on_error is None:
|
|
raise_on_error = self.strict
|
|
try:
|
|
await self._run_connect(server)
|
|
self._connected_servers.add(server)
|
|
if server in self.failed_servers:
|
|
self._remove_failed_server(server)
|
|
self.errors.pop(server, None)
|
|
except asyncio.CancelledError as exc:
|
|
if not self.suppress_cancelled_error:
|
|
raise
|
|
self._record_failure(server, exc, phase="connect")
|
|
except Exception as exc:
|
|
self._record_failure(server, exc, phase="connect")
|
|
if raise_on_error:
|
|
raise
|
|
except BaseException as exc:
|
|
self._record_failure(server, exc, phase="connect")
|
|
raise
|
|
|
|
def _refresh_active_servers(self) -> None:
|
|
if self.drop_failed_servers:
|
|
failed = set(self._failed_server_set)
|
|
self._active_servers = [server for server in self._all_servers if server not in failed]
|
|
else:
|
|
self._active_servers = list(self._all_servers)
|
|
|
|
def _record_failure(self, server: MCPServer, exc: BaseException, phase: str) -> None:
|
|
logger.exception("Failed to %s MCP server '%s': %s", phase, server.name, exc)
|
|
if server not in self._failed_server_set:
|
|
self.failed_servers.append(server)
|
|
self._failed_server_set.add(server)
|
|
self.errors[server] = exc
|
|
|
|
async def _run_connect(self, server: MCPServer) -> None:
|
|
if self.connect_in_parallel:
|
|
worker = self._get_worker(server)
|
|
await worker.connect()
|
|
else:
|
|
await self._run_with_timeout(server.connect, self.connect_timeout_seconds)
|
|
|
|
async def _cleanup_server(self, server: MCPServer) -> None:
|
|
if self.connect_in_parallel and server in self._workers:
|
|
worker = self._workers[server]
|
|
if worker.is_done:
|
|
self._workers.pop(server, None)
|
|
self._connected_servers.discard(server)
|
|
return
|
|
try:
|
|
await worker.cleanup()
|
|
finally:
|
|
self._workers.pop(server, None)
|
|
self._connected_servers.discard(server)
|
|
return
|
|
try:
|
|
await self._run_with_timeout(server.cleanup, self.cleanup_timeout_seconds)
|
|
finally:
|
|
self._connected_servers.discard(server)
|
|
|
|
async def _cleanup_servers(self, servers: Iterable[MCPServer]) -> None:
|
|
for server in reversed(list(servers)):
|
|
try:
|
|
await self._cleanup_server(server)
|
|
except asyncio.CancelledError as exc:
|
|
if not self.suppress_cancelled_error:
|
|
raise
|
|
logger.debug("Cleanup cancelled for MCP server '%s': %s", server.name, exc)
|
|
self.errors[server] = exc
|
|
except Exception as exc:
|
|
logger.exception("Failed to cleanup MCP server '%s': %s", server.name, exc)
|
|
self.errors[server] = exc
|
|
|
|
async def _connect_all_parallel(self, servers: list[MCPServer]) -> None:
|
|
tasks = [
|
|
asyncio.create_task(self._attempt_connect(server, raise_on_error=False))
|
|
for server in servers
|
|
]
|
|
results = await asyncio.gather(*tasks, return_exceptions=True)
|
|
if not self.suppress_cancelled_error:
|
|
for result in results:
|
|
if isinstance(result, asyncio.CancelledError):
|
|
raise result
|
|
for result in results:
|
|
if isinstance(result, BaseException) and not isinstance(result, asyncio.CancelledError):
|
|
raise result
|
|
if self.strict and self.failed_servers:
|
|
first_failure = None
|
|
if self.suppress_cancelled_error:
|
|
for server in self.failed_servers:
|
|
error = self.errors.get(server)
|
|
if error is None or isinstance(error, asyncio.CancelledError):
|
|
continue
|
|
first_failure = server
|
|
break
|
|
else:
|
|
first_failure = self.failed_servers[0]
|
|
if first_failure is not None:
|
|
error = self.errors.get(first_failure)
|
|
if error is not None:
|
|
raise error
|
|
raise RuntimeError(f"Failed to connect MCP server '{first_failure.name}'")
|
|
|
|
def _get_worker(self, server: MCPServer) -> _ServerWorker:
|
|
worker = self._workers.get(server)
|
|
if worker is None or worker.is_done:
|
|
worker = _ServerWorker(
|
|
server=server,
|
|
connect_timeout_seconds=self.connect_timeout_seconds,
|
|
cleanup_timeout_seconds=self.cleanup_timeout_seconds,
|
|
)
|
|
self._workers[server] = worker
|
|
return worker
|
|
|
|
def _remove_failed_server(self, server: MCPServer) -> None:
|
|
if server in self._failed_server_set:
|
|
self._failed_server_set.remove(server)
|
|
self.failed_servers = [
|
|
failed_server for failed_server in self.failed_servers if failed_server != server
|
|
]
|
|
|
|
def _servers_to_connect(self, servers: Iterable[MCPServer]) -> list[MCPServer]:
|
|
unique = self._unique_servers(servers)
|
|
if not self._connected_servers:
|
|
return unique
|
|
return [server for server in unique if server not in self._connected_servers]
|
|
|
|
@staticmethod
|
|
def _unique_servers(servers: Iterable[MCPServer]) -> list[MCPServer]:
|
|
seen: set[MCPServer] = set()
|
|
unique: list[MCPServer] = []
|
|
for server in servers:
|
|
if server not in seen:
|
|
seen.add(server)
|
|
unique.append(server)
|
|
return unique
|