637 lines
21 KiB
Python
637 lines
21 KiB
Python
"""MCP Service - Unified business logic and state management for MCP.
|
|
|
|
Responsibilities:
|
|
- Server configuration CRUD operations
|
|
- Configuration synchronization (Database <-> Cache)
|
|
- Unified entry point for Agent tool retrieval (auto-filtering disabled_tools)
|
|
- MCP Client and Tools management (formerly in agents/common/mcp.py)
|
|
"""
|
|
|
|
import asyncio
|
|
import re
|
|
import traceback
|
|
from collections.abc import Callable
|
|
from typing import Any, cast
|
|
|
|
from langchain_mcp_adapters.client import MultiServerMCPClient
|
|
from sqlalchemy import func, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
from yuxi.storage.postgres.models_business import MCPServer
|
|
from yuxi.utils import logger
|
|
|
|
# =============================================================================
|
|
# === Global Cache & State ===
|
|
# =============================================================================
|
|
|
|
# Global Lock for MCP state
|
|
_mcp_lock = asyncio.Lock()
|
|
|
|
# Global MCP tools cache
|
|
_mcp_tools_cache: dict[str, list[Callable[..., Any]]] = {}
|
|
|
|
# MCP tools statistics (for reporting enabled/disabled counts)
|
|
_mcp_tools_stats: dict[str, dict[str, int]] = {}
|
|
|
|
# MCP Server configurations (Runtime cache, loaded from DB)
|
|
MCP_SERVERS: dict[str, dict[str, Any]] = {}
|
|
_UNSET = object()
|
|
|
|
# Default MCP Server configurations (Imported to DB on first run)
|
|
_DEFAULT_MCP_SERVERS = {
|
|
"sequentialthinking": {
|
|
"url": "https://remote.mcpservers.org/sequentialthinking/mcp",
|
|
"transport": "streamable_http",
|
|
"description": "顺序思考工具，帮助 AI 将复杂问题分解为多个步骤",
|
|
"icon": "🧠",
|
|
"tags": ["内置", "AI"],
|
|
},
|
|
"mcp-server-chart": {
|
|
"command": "npx",
|
|
"args": ["-y", "@antv/mcp-server-chart"],
|
|
"transport": "stdio",
|
|
"description": "图表生成工具，支持生成各类图表（柱状图、折线图、饼图等）",
|
|
"icon": "📊",
|
|
"tags": ["内置", "图表"],
|
|
},
|
|
}
|
|
|
|
# =============================================================================
|
|
# === Core Logic (Moved from agents/common/mcp.py) ===
|
|
# =============================================================================
|
|
|
|
|
|
async def load_mcp_servers_from_db() -> None:
|
|
"""Load all enabled MCP server configurations from database to MCP_SERVERS cache."""
|
|
global MCP_SERVERS
|
|
|
|
# Delayed import to avoid circular references
|
|
from yuxi.storage.postgres.manager import pg_manager
|
|
|
|
try:
|
|
async with pg_manager.get_async_session_context() as session:
|
|
result = await session.execute(select(MCPServer).filter(MCPServer.enabled == 1))
|
|
servers = result.scalars().all()
|
|
|
|
async with _mcp_lock:
|
|
MCP_SERVERS.clear()
|
|
for server in servers:
|
|
MCP_SERVERS[server.name] = server.to_mcp_config()
|
|
|
|
logger.info(f"Loaded {len(MCP_SERVERS)} MCP servers from database: {list(MCP_SERVERS.keys())}")
|
|
except Exception as e:
|
|
logger.error(f"Failed to load MCP servers from database: {e}")
|
|
|
|
|
|
async def sync_mcp_server_to_cache(name: str, config: dict[str, Any] | None) -> None:
|
|
"""Sync a single MCP server configuration to cache.
|
|
|
|
Args:
|
|
name: Server name
|
|
config: Server configuration, or None to remove from cache
|
|
"""
|
|
global MCP_SERVERS
|
|
|
|
async with _mcp_lock:
|
|
if config is None:
|
|
MCP_SERVERS.pop(name, None)
|
|
logger.info(f"Removed MCP server '{name}' from cache")
|
|
else:
|
|
MCP_SERVERS[name] = config
|
|
logger.info(f"Synced MCP server '{name}' to cache")
|
|
|
|
# Clear tools cache for this server
|
|
_mcp_tools_cache.pop(name, None)
|
|
|
|
|
|
async def init_mcp_servers() -> None:
|
|
"""Initialize MCP server configurations.
|
|
|
|
On first run, if database is empty, import default configurations.
|
|
Then load configurations from database to MCP_SERVERS cache.
|
|
Also ensures all built-in MCP servers are present in the database.
|
|
"""
|
|
# Delayed import to avoid circular references
|
|
from yuxi.storage.postgres.manager import pg_manager
|
|
|
|
try:
|
|
async with pg_manager.get_async_session_context() as session:
|
|
# Check if database has MCP configurations
|
|
result = await session.execute(select(func.count(MCPServer.name)))
|
|
count = result.scalar()
|
|
|
|
if count == 0:
|
|
# Database is empty, import default configurations
|
|
logger.info("No MCP servers in database, importing default configurations...")
|
|
for name, config in _DEFAULT_MCP_SERVERS.items():
|
|
server = MCPServer(
|
|
name=name,
|
|
description=config.get("description"),
|
|
transport=config["transport"],
|
|
url=config.get("url"),
|
|
command=config.get("command"),
|
|
args=config.get("args"),
|
|
env=config.get("env"),
|
|
headers=config.get("headers"),
|
|
timeout=config.get("timeout"),
|
|
sse_read_timeout=config.get("sse_read_timeout"),
|
|
tags=config.get("tags"),
|
|
icon=config.get("icon"),
|
|
enabled=1,
|
|
created_by="system",
|
|
updated_by="system",
|
|
)
|
|
session.add(server)
|
|
await session.commit()
|
|
logger.info(f"Imported {len(_DEFAULT_MCP_SERVERS)} default MCP servers to database")
|
|
else:
|
|
# Ensure all built-in MCP servers exist in database
|
|
for name, config in _DEFAULT_MCP_SERVERS.items():
|
|
result = await session.execute(select(MCPServer).filter(MCPServer.name == name))
|
|
existing = result.scalar_one_or_none()
|
|
if not existing:
|
|
server = MCPServer(
|
|
name=name,
|
|
description=config.get("description"),
|
|
transport=config["transport"],
|
|
url=config.get("url"),
|
|
command=config.get("command"),
|
|
args=config.get("args"),
|
|
env=config.get("env"),
|
|
headers=config.get("headers"),
|
|
timeout=config.get("timeout"),
|
|
sse_read_timeout=config.get("sse_read_timeout"),
|
|
tags=config.get("tags"),
|
|
icon=config.get("icon"),
|
|
enabled=1,
|
|
created_by="system",
|
|
updated_by="system",
|
|
)
|
|
session.add(server)
|
|
logger.info(f"Added built-in MCP server '{name}' to database")
|
|
# Commit if any new servers were added (check session state)
|
|
if session.new:
|
|
await session.commit()
|
|
|
|
# Load configurations from database to cache
|
|
await load_mcp_servers_from_db()
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize MCP servers: {e}, traceback: {traceback.format_exc()}")
|
|
|
|
|
|
async def get_mcp_client(
|
|
server_configs: dict[str, Any] | None = None,
|
|
) -> MultiServerMCPClient | None:
|
|
"""Initializes an MCP client with the given server configurations."""
|
|
try:
|
|
client = MultiServerMCPClient(server_configs) # pyright: ignore[reportArgumentType]
|
|
logger.info(f"Initialized MCP client with servers: {list(server_configs.keys())}")
|
|
return client
|
|
except Exception as e:
|
|
logger.error("Failed to initialize MCP client: {}", e)
|
|
return None
|
|
|
|
|
|
def to_camel_case(s: str) -> str:
|
|
"""Convert string to lowerCamelCase."""
|
|
|
|
# Handle - and _
|
|
s = re.sub(r"[-_]+(.)", lambda m: m.group(1).upper(), s)
|
|
# Lowercase first letter
|
|
if len(s) > 0:
|
|
s = s[0].lower() + s[1:]
|
|
return s
|
|
|
|
|
|
async def get_mcp_tools(
|
|
server_name: str,
|
|
additional_servers: dict[str, dict] = None,
|
|
disabled_tools: list[str] = None,
|
|
cache: bool = True,
|
|
force_refresh: bool = False,
|
|
raise_on_error: bool = False,
|
|
) -> list[Callable[..., Any]]:
|
|
"""Get MCP tools for a specific server.
|
|
|
|
Architecture:
|
|
1. Fetching: Connects to MCP server to get ALL tools.
|
|
2. Caching: Stores the FULL, UNFILTERED list of tools in `_mcp_tools_cache`.
|
|
3. Filtering: Filters the return value based on `disabled_tools` argument.
|
|
|
|
Args:
|
|
server_name: Server name
|
|
additional_servers: Additional server configurations
|
|
disabled_tools: List of tool names to filter out from the RETURN value (does not affect cache)
|
|
cache: Whether to use/update the cache (default: True)
|
|
force_refresh: Whether to force a refresh from the server (default: False)
|
|
"""
|
|
global _mcp_tools_cache
|
|
|
|
# 1. Prepare Server Config
|
|
async with _mcp_lock:
|
|
mcp_servers = MCP_SERVERS | (additional_servers or {})
|
|
|
|
all_processed_tools = []
|
|
|
|
# 2. Check Cache / Fetch Strategy
|
|
# If we have it in cache and don't need to force refresh, use cache.
|
|
if not force_refresh and cache and server_name in _mcp_tools_cache:
|
|
all_processed_tools = _mcp_tools_cache[server_name]
|
|
else:
|
|
# Need to fetch from server
|
|
try:
|
|
assert server_name in mcp_servers, f"Server {server_name} not found in ({list(mcp_servers.keys())})"
|
|
|
|
# Extract connection config
|
|
server_config = mcp_servers[server_name]
|
|
client_config = {k: v for k, v in server_config.items() if k not in ("disabled_tools",)}
|
|
|
|
client = await get_mcp_client({server_name: client_config})
|
|
if client is None:
|
|
return []
|
|
|
|
# Get ALL tools (Raw)
|
|
raw_tools = cast(list[Any], await client.get_tools())
|
|
|
|
# Render IDs for ALL tools
|
|
server_cc = to_camel_case(server_name)
|
|
for tool in raw_tools:
|
|
# Render unique ID rule: mcp__[camelCaseServer]__[camelCaseTool]
|
|
original_name = tool.name
|
|
tool_cc = to_camel_case(original_name)
|
|
unique_id = f"mcp__{server_cc}__{tool_cc}"
|
|
|
|
# Use metadata to store
|
|
if tool.metadata is None:
|
|
tool.metadata = {}
|
|
tool.metadata["id"] = unique_id
|
|
|
|
all_processed_tools.append(tool)
|
|
|
|
# Update Cache (Store the FULL list)
|
|
if cache:
|
|
_mcp_tools_cache[server_name] = all_processed_tools
|
|
|
|
# Update Stats
|
|
# Stats should reflect the GLOBAL configuration state
|
|
# (How many are disabled in the stored config, not the transient arg)
|
|
global_config_disabled = mcp_servers.get(server_name, {}).get("disabled_tools") or []
|
|
enabled_count = len([t for t in all_processed_tools if t.name not in global_config_disabled])
|
|
|
|
_mcp_tools_stats[server_name] = {
|
|
"total": len(all_processed_tools),
|
|
"enabled": enabled_count,
|
|
"disabled": len(all_processed_tools) - enabled_count,
|
|
}
|
|
|
|
logger.info(f"Refreshed MCP tools cache for '{server_name}': {len(all_processed_tools)} tools loaded.")
|
|
|
|
except AssertionError as e:
|
|
logger.warning(f"[assert] Failed to load tools from MCP server '{server_name}': {e}")
|
|
if raise_on_error:
|
|
raise
|
|
return []
|
|
except Exception as e:
|
|
logger.error(
|
|
f"Failed to load tools from MCP server '{server_name}': {e}, traceback: {traceback.format_exc()}"
|
|
)
|
|
if raise_on_error:
|
|
raise
|
|
return []
|
|
|
|
# 3. Filtering (Apply to Return Value Only)
|
|
if disabled_tools:
|
|
filtered_tools = [t for t in all_processed_tools if t.name not in disabled_tools]
|
|
logger.debug(
|
|
f"Returning {len(filtered_tools)}/{len(all_processed_tools)} tools for '{server_name}' "
|
|
f"(filtered {len(disabled_tools)} by argument)"
|
|
)
|
|
return filtered_tools
|
|
|
|
return all_processed_tools
|
|
|
|
|
|
async def get_tools_from_all_servers() -> list[Callable[..., Any]]:
|
|
"""Get all tools from all configured MCP servers."""
|
|
all_tools = []
|
|
for server_name in MCP_SERVERS.keys():
|
|
tools = await get_mcp_tools(server_name)
|
|
all_tools.extend(tools)
|
|
return all_tools
|
|
|
|
|
|
def add_mcp_server(name: str, config: dict[str, Any]) -> None:
|
|
"""Add a new MCP server configuration."""
|
|
MCP_SERVERS[name] = config
|
|
# Clear client to force reinitialization with new config
|
|
clear_mcp_cache()
|
|
|
|
|
|
def clear_mcp_cache() -> None:
|
|
"""Clear the MCP tools cache (useful for testing)."""
|
|
global _mcp_tools_cache, _mcp_tools_stats
|
|
_mcp_tools_cache = {}
|
|
_mcp_tools_stats = {}
|
|
|
|
|
|
def clear_mcp_server_tools_cache(server_name: str) -> None:
|
|
"""Clear the tools cache for a specific MCP server."""
|
|
global _mcp_tools_cache, _mcp_tools_stats
|
|
_mcp_tools_cache.pop(server_name, None)
|
|
_mcp_tools_stats.pop(server_name, None)
|
|
logger.info(f"Cleared tools cache for MCP server '{server_name}'")
|
|
|
|
|
|
def get_mcp_tools_stats(server_name: str) -> dict[str, int] | None:
|
|
"""Get tools statistics for a MCP server.
|
|
|
|
Returns:
|
|
dict with 'total', 'enabled', 'disabled' counts, or None if not available
|
|
"""
|
|
return _mcp_tools_stats.get(server_name)
|
|
|
|
|
|
# =============================================================================
|
|
# === Server Config CRUD (Existing in mcp_service.py) ===
|
|
# =============================================================================
|
|
|
|
|
|
async def get_mcp_server(db: AsyncSession, name: str) -> MCPServer | None:
|
|
"""Get single server configuration."""
|
|
result = await db.execute(select(MCPServer).filter(MCPServer.name == name))
|
|
return result.scalar_one_or_none()
|
|
|
|
|
|
async def get_all_mcp_servers(db: AsyncSession) -> list[MCPServer]:
|
|
"""Get all server configurations."""
|
|
result = await db.execute(select(MCPServer))
|
|
return list(result.scalars().all())
|
|
|
|
|
|
async def create_mcp_server(
|
|
db: AsyncSession,
|
|
name: str,
|
|
transport: str,
|
|
url: str = None,
|
|
command: str = None,
|
|
args: list = None,
|
|
env: dict = None,
|
|
description: str = None,
|
|
headers: dict = None,
|
|
timeout: int = None,
|
|
sse_read_timeout: int = None,
|
|
tags: list = None,
|
|
icon: str = None,
|
|
created_by: str = None,
|
|
) -> MCPServer:
|
|
"""Create server."""
|
|
# Check if name exists
|
|
existing = await get_mcp_server(db, name)
|
|
if existing:
|
|
raise ValueError(f"Server name '{name}' already exists")
|
|
|
|
server = MCPServer(
|
|
name=name,
|
|
description=description,
|
|
transport=transport,
|
|
url=url,
|
|
command=command,
|
|
args=args,
|
|
env=env,
|
|
headers=headers,
|
|
timeout=timeout,
|
|
sse_read_timeout=sse_read_timeout,
|
|
tags=tags,
|
|
icon=icon,
|
|
enabled=1,
|
|
created_by=created_by,
|
|
updated_by=created_by,
|
|
)
|
|
db.add(server)
|
|
await db.commit()
|
|
await db.refresh(server)
|
|
|
|
# Sync to cache
|
|
await sync_mcp_server_to_cache(name, server.to_mcp_config())
|
|
|
|
logger.info(f"Created MCP server '{name}'")
|
|
return server
|
|
|
|
|
|
async def update_mcp_server(
|
|
db: AsyncSession,
|
|
name: str,
|
|
description: str = None,
|
|
transport: str = None,
|
|
url: str = None,
|
|
command: str = None,
|
|
args: list = None,
|
|
env: Any = _UNSET,
|
|
headers: dict = None,
|
|
timeout: int = None,
|
|
sse_read_timeout: int = None,
|
|
tags: list = None,
|
|
icon: str = None,
|
|
updated_by: str = None,
|
|
) -> MCPServer:
|
|
"""Update server configuration."""
|
|
server = await get_mcp_server(db, name)
|
|
if not server:
|
|
raise ValueError(f"Server '{name}' does not exist")
|
|
|
|
if description is not None:
|
|
server.description = description
|
|
if transport is not None:
|
|
server.transport = transport
|
|
if url is not None:
|
|
server.url = url
|
|
if command is not None:
|
|
server.command = command
|
|
if args is not None:
|
|
server.args = args
|
|
if env is not _UNSET:
|
|
server.env = env
|
|
if headers is not None:
|
|
server.headers = headers
|
|
if timeout is not None:
|
|
server.timeout = timeout
|
|
if sse_read_timeout is not None:
|
|
server.sse_read_timeout = sse_read_timeout
|
|
if tags is not None:
|
|
server.tags = tags
|
|
if icon is not None:
|
|
server.icon = icon
|
|
if updated_by is not None:
|
|
server.updated_by = updated_by
|
|
|
|
await db.commit()
|
|
await db.refresh(server)
|
|
|
|
# Sync to cache (if enabled)
|
|
if server.enabled:
|
|
await sync_mcp_server_to_cache(name, server.to_mcp_config())
|
|
|
|
logger.info(f"Updated MCP server '{name}'")
|
|
return server
|
|
|
|
|
|
async def delete_mcp_server(db: AsyncSession, name: str) -> bool:
|
|
"""Delete server."""
|
|
server = await get_mcp_server(db, name)
|
|
if not server:
|
|
return False
|
|
|
|
await db.delete(server)
|
|
await db.commit()
|
|
|
|
# Remove from cache
|
|
await sync_mcp_server_to_cache(name, None)
|
|
|
|
logger.info(f"Deleted MCP server '{name}'")
|
|
return True
|
|
|
|
|
|
# =============================================================================
|
|
# === Tool Management ===
|
|
# =============================================================================
|
|
|
|
|
|
async def toggle_server_enabled(db: AsyncSession, name: str, updated_by: str = None) -> tuple[bool, MCPServer]:
|
|
"""Toggle server enabled status."""
|
|
server = await get_mcp_server(db, name)
|
|
if not server:
|
|
raise ValueError(f"Server '{name}' does not exist")
|
|
|
|
server.enabled = 0 if server.enabled else 1
|
|
if updated_by is not None:
|
|
server.updated_by = updated_by
|
|
await db.commit()
|
|
|
|
# Sync to cache
|
|
is_enabled = bool(server.enabled)
|
|
server_config = server.to_mcp_config() if is_enabled else None
|
|
await sync_mcp_server_to_cache(name, server_config)
|
|
|
|
logger.info(f"Toggled MCP server '{name}' enabled={is_enabled}")
|
|
return is_enabled, server
|
|
|
|
|
|
async def toggle_tool_enabled(
|
|
db: AsyncSession,
|
|
server_name: str,
|
|
tool_name: str,
|
|
updated_by: str = None,
|
|
) -> tuple[bool, MCPServer]:
|
|
"""Toggle single tool enabled status.
|
|
|
|
Args:
|
|
db: Database session
|
|
server_name: Server name
|
|
tool_name: Tool name
|
|
updated_by: Updater
|
|
|
|
Returns:
|
|
(enabled, server): Tool enabled status and updated server object
|
|
"""
|
|
server = await get_mcp_server(db, server_name)
|
|
if not server:
|
|
raise ValueError(f"Server '{server_name}' does not exist")
|
|
|
|
disabled_tools = list(server.disabled_tools or [])
|
|
|
|
if tool_name in disabled_tools:
|
|
disabled_tools.remove(tool_name)
|
|
enabled = True
|
|
else:
|
|
disabled_tools.append(tool_name)
|
|
enabled = False
|
|
|
|
server.disabled_tools = disabled_tools
|
|
if updated_by is not None:
|
|
server.updated_by = updated_by
|
|
await db.commit()
|
|
|
|
# Clear tool cache (re-filtered on next fetch)
|
|
clear_mcp_server_tools_cache(server_name)
|
|
|
|
logger.info(f"Toggled tool '{tool_name}' for server '{server_name}' enabled={enabled}")
|
|
return enabled, server
|
|
|
|
|
|
# =============================================================================
|
|
# === Unified Entry Points (Wrappers) ===
|
|
# =============================================================================
|
|
|
|
|
|
def get_mcp_server_names() -> list[str]:
|
|
"""Get list of loaded MCP server names.
|
|
|
|
Returns a copy of keys to avoid runtime modification issues during iteration.
|
|
"""
|
|
return list(MCP_SERVERS.keys())
|
|
|
|
|
|
async def get_enabled_mcp_tools(server_name: str) -> list:
|
|
"""Get MCP server tools (auto-filtering disabled_tools).
|
|
|
|
Unified entry point for Agents, automatically:
|
|
1. Gets server config from cache
|
|
2. Gets all tools
|
|
3. Filters out disabled_tools
|
|
|
|
Args:
|
|
server_name: Server name
|
|
|
|
Returns:
|
|
List of enabled tools
|
|
"""
|
|
config = MCP_SERVERS.get(server_name)
|
|
if not config:
|
|
logger.warning(f"MCP server '{server_name}' not found in cache")
|
|
return []
|
|
|
|
disabled_tools = config.get("disabled_tools") or []
|
|
return await get_mcp_tools(server_name, disabled_tools=disabled_tools)
|
|
|
|
|
|
async def get_servers_config(names: list[str]) -> dict[str, dict[str, Any]]:
|
|
"""Batch get server configurations.
|
|
|
|
Args:
|
|
names: List of server names
|
|
|
|
Returns:
|
|
{name: config} dictionary, containing only found servers
|
|
"""
|
|
return {name: MCP_SERVERS[name] for name in names if name in MCP_SERVERS}
|
|
|
|
|
|
async def get_all_mcp_tools(server_name: str, raise_on_error: bool = False) -> list:
|
|
"""Get all tools of an MCP server (no filtering).
|
|
|
|
For management UI to display tool list, supports viewing all tools and their enabled status.
|
|
Does NOT update the global tools cache to avoid polluting agent's filtered view.
|
|
|
|
Args:
|
|
server_name: Server name
|
|
raise_on_error: Whether to raise an exception on error instead of returning empty list
|
|
|
|
Returns:
|
|
List of all tools (unfiltered)
|
|
"""
|
|
config = MCP_SERVERS.get(server_name)
|
|
if not config:
|
|
logger.warning(f"MCP server '{server_name}' not found in cache")
|
|
if raise_on_error:
|
|
raise ValueError(f"MCP server '{server_name}' not found in cache")
|
|
return []
|
|
|
|
# Get all tools (no filtering, force refresh, no cache update)
|
|
return await get_mcp_tools(
|
|
server_name,
|
|
disabled_tools=[],
|
|
cache=False,
|
|
force_refresh=True,
|
|
raise_on_error=raise_on_error,
|
|
)
|