104 lines
3.3 KiB
Python
104 lines
3.3 KiB
Python
"""Multi-server MCP client using fastmcp composition.
|
|
|
|
Composes the three domain-scoped MCP servers into namespaced configurations
|
|
that agents connect to as a single client.
|
|
"""
|
|
|
|
import os
|
|
from contextlib import asynccontextmanager
|
|
from typing import Any
|
|
|
|
from fastmcp import Client
|
|
|
|
from agents.shared.parser import (
|
|
parse_prompt_result,
|
|
parse_resource_result,
|
|
parse_tool_result,
|
|
)
|
|
|
|
SERVER_MODULES = {
|
|
"shared": "mcp_servers.shared",
|
|
"ops": "mcp_servers.ops",
|
|
"passenger": "mcp_servers.passenger",
|
|
}
|
|
|
|
AGENT_PROFILES = {
|
|
"fce": ["shared", "ops", "passenger"],
|
|
"handover": ["shared", "ops"],
|
|
}
|
|
|
|
_FORWARDED_ENV_KEYS = (
|
|
"LLM_PROVIDER", "GROQ_API_KEY", "GROQ_MODEL",
|
|
"ANTHROPIC_API_KEY", "ANTHROPIC_MODEL",
|
|
"OPENAI_API_KEY", "OPENAI_BASE_URL", "OPENAI_MODEL",
|
|
"AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_DEFAULT_REGION",
|
|
"BEDROCK_MODEL_ID", "USE_BEDROCK",
|
|
"PATH",
|
|
)
|
|
|
|
|
|
def _subprocess_env() -> dict:
|
|
"""Env vars forwarded to MCP server subprocesses (LLM config + active scenario)."""
|
|
from mcp_servers.data.scenarios.manager import scenario_manager
|
|
|
|
env = {k: os.environ[k] for k in _FORWARDED_ENV_KEYS if k in os.environ}
|
|
env["ACTIVE_SCENARIO"] = scenario_manager.active_id
|
|
return env
|
|
|
|
|
|
def _server_config(module: str) -> dict:
|
|
return {"command": "uv", "args": ["run", "python", "-m", module], "env": _subprocess_env()}
|
|
|
|
|
|
class MCPMultiClient:
|
|
"""Manages connections to multiple MCP servers via fastmcp Client."""
|
|
|
|
def __init__(self) -> None:
|
|
self._clients: dict[str, Client] = {}
|
|
|
|
async def connect(self, server_names: list[str]) -> None:
|
|
for name in server_names:
|
|
if name not in SERVER_MODULES:
|
|
raise ValueError(f"Unknown server: {name}. Available: {list(SERVER_MODULES.keys())}")
|
|
config = {"mcpServers": {"default": _server_config(SERVER_MODULES[name])}}
|
|
client = Client(config)
|
|
await client.__aenter__()
|
|
self._clients[name] = client
|
|
|
|
async def close(self) -> None:
|
|
for client in self._clients.values():
|
|
try:
|
|
await client.__aexit__(None, None, None)
|
|
except (Exception, BaseException):
|
|
pass
|
|
self._clients.clear()
|
|
|
|
def _client(self, server: str) -> Client:
|
|
client = self._clients.get(server)
|
|
if not client:
|
|
raise ValueError(f"Not connected to server: {server}")
|
|
return client
|
|
|
|
async def call_tool(self, server: str, tool_name: str, arguments: dict) -> Any:
|
|
result = await self._client(server).call_tool(tool_name, arguments)
|
|
return parse_tool_result(result)
|
|
|
|
async def read_resource(self, server: str, uri: str) -> Any:
|
|
result = await self._client(server).read_resource(uri)
|
|
return parse_resource_result(result)
|
|
|
|
async def get_prompt(self, server: str, prompt_name: str, arguments: dict) -> str:
|
|
result = await self._client(server).get_prompt(prompt_name, arguments)
|
|
return parse_prompt_result(result)
|
|
|
|
|
|
@asynccontextmanager
|
|
async def connect_servers(server_names: list[str]):
|
|
"""Context manager for multi-server MCP connections."""
|
|
client = MCPMultiClient()
|
|
try:
|
|
await client.connect(server_names)
|
|
yield client
|
|
finally:
|
|
await client.close()
|