Agent Zero
317d5c81f8
- New plugin app/plugins/builtins/mcp_client/ with server config CRUD
- Models: McpServerConfig with TenantMixin (name, url, api_token, enabled)
- Routes: GET/POST/PATCH/DELETE /api/v1/mcp-client/servers
- Routes: GET /servers/{id}/tools, POST /servers/{id}/execute
- client.py: async MCP client using httpx for external server calls
- tool_registry_integration.py: registers external MCP tools in AI tool registry
- Migration: 0001_initial.sql for mcp_server_configs table
- Frontend: mcpClient.ts API client with React Query hooks
- Frontend: MCP Client settings UI in SettingsMcp.tsx (server CRUD, tool viewing)
- Tests: 8 tests covering CRUD, auth, tool registry integration
95 lines
3.6 KiB
Python
95 lines
3.6 KiB
Python
"""Integrates external MCP server tools into the AI Assistant tool_registry.
|
|
|
|
When the MCP Client plugin activates, it loads all enabled MCP server configs,
|
|
fetches their tool lists, and registers each tool in the global ToolRegistry.
|
|
Agents can then call external MCP tools like native tools.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import uuid
|
|
from typing import Any
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.plugins.builtins.ai_assistant.tool_registry import get_tool_registry
|
|
from app.plugins.builtins.mcp_client.client import McpClient
|
|
from app.plugins.builtins.mcp_client.models import McpServerConfig as McpServerConfigModel
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
PLUGIN_NAME = "mcp_client"
|
|
|
|
|
|
def _make_tool_name(server_name: str, tool_name: str) -> str:
|
|
"""Generate a unique tool name: mcp__{server}__{tool}."""
|
|
safe_server = server_name.replace(" ", "_").replace("-", "_").lower()
|
|
return f"mcp__{safe_server}__{tool_name}"
|
|
|
|
|
|
def _make_handler(server_cfg: McpServerConfigModel, tool_name: str):
|
|
"""Create an async handler that calls the external MCP server."""
|
|
|
|
async def _handler(arguments: dict[str, Any], context: dict[str, Any]) -> str:
|
|
client = McpClient(base_url=server_cfg.url, api_token=server_cfg.api_token)
|
|
try:
|
|
resp = await client.execute_tool(tool_name, arguments)
|
|
if resp.success:
|
|
import json
|
|
result = resp.result if resp.result is not None else {}
|
|
return json.dumps(result) if isinstance(result, dict) else str(result)
|
|
return f"Error: {resp.error or 'Unknown error'}"
|
|
except Exception as exc:
|
|
logger.exception("MCP tool execution failed: %s/%s", server_cfg.name, tool_name)
|
|
return f"Error: {exc}"
|
|
|
|
return _handler
|
|
|
|
|
|
async def sync_external_tools(db: AsyncSession, tenant_id: uuid.UUID) -> int:
|
|
"""Fetch tools from all enabled MCP servers and register them in the tool registry.
|
|
|
|
Returns the number of tools registered.
|
|
"""
|
|
registry = get_tool_registry()
|
|
# Unregister previous tools from this plugin
|
|
registry.unregister_plugin(PLUGIN_NAME)
|
|
|
|
stmt = select(McpServerConfigModel).where(
|
|
McpServerConfigModel.tenant_id == tenant_id,
|
|
McpServerConfigModel.enabled.is_(True),
|
|
)
|
|
result = await db.execute(stmt)
|
|
configs = result.scalars().all()
|
|
|
|
count = 0
|
|
for cfg in configs:
|
|
try:
|
|
client = McpClient(base_url=cfg.url, api_token=cfg.api_token, timeout=10.0)
|
|
tools_resp = await client.list_tools()
|
|
for tool in tools_resp.tools:
|
|
tool_name = _make_tool_name(cfg.name, tool.name)
|
|
registry.register(
|
|
name=tool_name,
|
|
description=f"[MCP:{cfg.name}] {tool.description}",
|
|
parameters=tool.parameters if isinstance(tool.parameters, dict) else {},
|
|
handler=_make_handler(cfg, tool.name),
|
|
plugin_name=PLUGIN_NAME,
|
|
required_permission="mcp-client:read",
|
|
category="mcp-external",
|
|
)
|
|
count += 1
|
|
logger.info("Registered %d tools from MCP server %s", len(tools_resp.tools), cfg.name)
|
|
except Exception as exc:
|
|
logger.warning("Failed to sync tools from MCP server %s: %s", cfg.name, exc)
|
|
|
|
return count
|
|
|
|
|
|
def unregister_all_external_tools() -> None:
|
|
"""Remove all MCP client tools from the registry."""
|
|
registry = get_tool_registry()
|
|
registry.unregister_plugin(PLUGIN_NAME)
|