Agent Zero
217 lines
8.0 KiB
Python
217 lines
8.0 KiB
Python
"""MCP Client plugin routes — CRUD for server configs, tool listing, execution."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import uuid
|
|
from datetime import UTC, datetime
|
|
from typing import Any
|
|
|
|
from fastapi import APIRouter, Depends, HTTPException
|
|
from sqlalchemy import select, update
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.core.db import get_db
|
|
from app.core.visibility import apply_visibility_filter
|
|
from app.deps import require_permission
|
|
from app.plugins.builtins.mcp_client.client import McpClient
|
|
from app.plugins.builtins.mcp_client.models import McpServerConfig as McpServerConfigModel
|
|
from app.plugins.builtins.mcp_client.schemas import (
|
|
McpServerConfigCreate,
|
|
McpServerConfigResponse,
|
|
McpServerConfigUpdate,
|
|
McpServerExecuteRequest,
|
|
McpServerExecuteResponse,
|
|
McpServerToolsResponse,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
router = APIRouter(prefix="/api/v1/mcp-client", tags=["mcp-client"])
|
|
|
|
|
|
def _config_to_response(cfg: McpServerConfigModel) -> McpServerConfigResponse:
|
|
return McpServerConfigResponse(
|
|
id=str(cfg.id),
|
|
name=cfg.name,
|
|
url=cfg.url,
|
|
api_token=cfg.api_token,
|
|
enabled=cfg.enabled,
|
|
description=cfg.description,
|
|
last_connected_at=cfg.last_connected_at,
|
|
created_by=str(cfg.created_by) if cfg.created_by else None,
|
|
created_at=cfg.created_at,
|
|
updated_at=cfg.updated_at,
|
|
)
|
|
|
|
|
|
@router.get("/servers", response_model=list[McpServerConfigResponse])
|
|
async def list_mcp_servers(
|
|
db: AsyncSession = Depends(get_db),
|
|
current_user: dict[str, Any] = Depends(require_permission("mcp-client:read")),
|
|
) -> list[McpServerConfigResponse]:
|
|
"""List all configured MCP servers for the current tenant."""
|
|
tenant_id = uuid.UUID(current_user["tenant_id"])
|
|
user_id = uuid.UUID(current_user["user_id"])
|
|
is_system_admin = current_user.get("is_system_admin", False)
|
|
stmt = select(McpServerConfigModel).where(McpServerConfigModel.tenant_id == tenant_id)
|
|
stmt = await apply_visibility_filter(
|
|
db, stmt, "mcp_server_config", McpServerConfigModel, user_id, tenant_id, is_system_admin
|
|
)
|
|
result = await db.execute(stmt)
|
|
configs = result.scalars().all()
|
|
return [_config_to_response(c) for c in configs]
|
|
|
|
|
|
@router.post("/servers", response_model=McpServerConfigResponse, status_code=201)
|
|
async def create_mcp_server(
|
|
body: McpServerConfigCreate,
|
|
db: AsyncSession = Depends(get_db),
|
|
current_user: dict[str, Any] = Depends(require_permission("mcp-client:write")),
|
|
) -> McpServerConfigResponse:
|
|
"""Create a new MCP server configuration."""
|
|
cfg = McpServerConfigModel(
|
|
tenant_id=uuid.UUID(current_user["tenant_id"]),
|
|
name=body.name,
|
|
url=body.url,
|
|
api_token=body.api_token,
|
|
enabled=body.enabled,
|
|
description=body.description,
|
|
created_by=uuid.UUID(current_user["user_id"]),
|
|
owner_id=uuid.UUID(current_user["user_id"]),
|
|
)
|
|
db.add(cfg)
|
|
await db.commit()
|
|
await db.refresh(cfg)
|
|
return _config_to_response(cfg)
|
|
|
|
|
|
@router.patch("/servers/{server_id}", response_model=McpServerConfigResponse)
|
|
async def update_mcp_server(
|
|
server_id: str,
|
|
body: McpServerConfigUpdate,
|
|
db: AsyncSession = Depends(get_db),
|
|
current_user: dict[str, Any] = Depends(require_permission("mcp-client:write")),
|
|
) -> McpServerConfigResponse:
|
|
"""Update an existing MCP server configuration."""
|
|
try:
|
|
sid = uuid.UUID(server_id)
|
|
except (ValueError, TypeError):
|
|
raise HTTPException(400, detail={"detail": "Invalid server_id", "code": "invalid_id"}) from None
|
|
|
|
stmt = select(McpServerConfigModel).where(
|
|
McpServerConfigModel.id == sid,
|
|
McpServerConfigModel.tenant_id == uuid.UUID(current_user["tenant_id"]),
|
|
)
|
|
result = await db.execute(stmt)
|
|
cfg = result.scalar_one_or_none()
|
|
if not cfg:
|
|
raise HTTPException(404, detail={"detail": "MCP server not found", "code": "not_found"})
|
|
|
|
update_data = body.model_dump(exclude_unset=True)
|
|
for key, val in update_data.items():
|
|
setattr(cfg, key, val)
|
|
|
|
await db.commit()
|
|
await db.refresh(cfg)
|
|
return _config_to_response(cfg)
|
|
|
|
|
|
@router.delete("/servers/{server_id}", status_code=204)
|
|
async def delete_mcp_server(
|
|
server_id: str,
|
|
db: AsyncSession = Depends(get_db),
|
|
current_user: dict[str, Any] = Depends(require_permission("mcp-client:write")),
|
|
):
|
|
"""Delete an MCP server configuration."""
|
|
try:
|
|
sid = uuid.UUID(server_id)
|
|
except (ValueError, TypeError):
|
|
raise HTTPException(400, detail={"detail": "Invalid server_id", "code": "invalid_id"}) from None
|
|
|
|
stmt = select(McpServerConfigModel).where(
|
|
McpServerConfigModel.id == sid,
|
|
McpServerConfigModel.tenant_id == uuid.UUID(current_user["tenant_id"]),
|
|
)
|
|
result = await db.execute(stmt)
|
|
cfg = result.scalar_one_or_none()
|
|
if not cfg:
|
|
raise HTTPException(404, detail={"detail": "MCP server not found", "code": "not_found"})
|
|
|
|
await db.delete(cfg)
|
|
await db.commit()
|
|
|
|
|
|
@router.get("/servers/{server_id}/tools", response_model=McpServerToolsResponse)
|
|
async def list_server_tools(
|
|
server_id: str,
|
|
db: AsyncSession = Depends(get_db),
|
|
current_user: dict[str, Any] = Depends(require_permission("mcp-client:read")),
|
|
) -> McpServerToolsResponse:
|
|
"""List tools available on a specific external MCP server."""
|
|
try:
|
|
sid = uuid.UUID(server_id)
|
|
except (ValueError, TypeError):
|
|
raise HTTPException(400, detail={"detail": "Invalid server_id", "code": "invalid_id"}) from None
|
|
|
|
stmt = select(McpServerConfigModel).where(
|
|
McpServerConfigModel.id == sid,
|
|
McpServerConfigModel.tenant_id == uuid.UUID(current_user["tenant_id"]),
|
|
)
|
|
result = await db.execute(stmt)
|
|
cfg = result.scalar_one_or_none()
|
|
if not cfg:
|
|
raise HTTPException(404, detail={"detail": "MCP server not found", "code": "not_found"})
|
|
if not cfg.enabled:
|
|
raise HTTPException(400, detail={"detail": "MCP server is disabled", "code": "server_disabled"})
|
|
|
|
client = McpClient(base_url=cfg.url, api_token=cfg.api_token)
|
|
try:
|
|
tools_resp = await client.list_tools()
|
|
# Update last_connected_at
|
|
await db.execute(
|
|
update(McpServerConfigModel).where(McpServerConfigModel.id == sid).values(last_connected_at=datetime.now(UTC))
|
|
)
|
|
await db.commit()
|
|
return tools_resp
|
|
except Exception as exc:
|
|
logger.exception("Failed to list tools from MCP server %s", cfg.name)
|
|
raise HTTPException(502, detail={"detail": f"Failed to connect: {exc}", "code": "connection_failed"}) from exc
|
|
|
|
|
|
@router.post("/servers/{server_id}/execute", response_model=McpServerExecuteResponse)
|
|
async def execute_server_tool(
|
|
server_id: str,
|
|
body: McpServerExecuteRequest,
|
|
db: AsyncSession = Depends(get_db),
|
|
current_user: dict[str, Any] = Depends(require_permission("mcp-client:write")),
|
|
) -> McpServerExecuteResponse:
|
|
"""Execute a tool on a specific external MCP server."""
|
|
try:
|
|
sid = uuid.UUID(server_id)
|
|
except (ValueError, TypeError):
|
|
raise HTTPException(400, detail={"detail": "Invalid server_id", "code": "invalid_id"}) from None
|
|
|
|
stmt = select(McpServerConfigModel).where(
|
|
McpServerConfigModel.id == sid,
|
|
McpServerConfigModel.tenant_id == uuid.UUID(current_user["tenant_id"]),
|
|
)
|
|
result = await db.execute(stmt)
|
|
cfg = result.scalar_one_or_none()
|
|
if not cfg:
|
|
raise HTTPException(404, detail={"detail": "MCP server not found", "code": "not_found"})
|
|
if not cfg.enabled:
|
|
raise HTTPException(400, detail={"detail": "MCP server is disabled", "code": "server_disabled"})
|
|
|
|
client = McpClient(base_url=cfg.url, api_token=cfg.api_token)
|
|
try:
|
|
return await client.execute_tool(body.tool_name, body.arguments)
|
|
except Exception as exc:
|
|
logger.exception("Failed to execute tool on MCP server %s", cfg.name)
|
|
return McpServerExecuteResponse(
|
|
server_name=cfg.name,
|
|
tool=body.tool_name,
|
|
success=False,
|
|
error=str(exc),
|
|
)
|