Agent Zero
727d86614e
P0 (7): Auth-bypass removed, migrations fixed, plugin-upload disabled, RLS FORCE+WITH CHECK, plugin double-registration fixed, persistent volume, domain removed P1 (11): User/tenant model, Redis centralized, worker separated, transactional outbox, XSS fixed, DMS chunked streaming, permissions unified, password reset, metrics secured, config/docs fixed, cross-tenant FK P2 (4): Contact model normalized, cross-imports reduced 94%, commands+state machines for contacts/dms/mail/calendar, SPA path-traversal 8 new migrations, 99 unit tests, 13 commands, 8 contracts, 72 files changed
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.contracts 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)
|