Phase 3.3: MCP Protocol Support Module
This commit is contained in:
qyyfw committed
1 parent
aa8d689039
commit
5a7c717b25
7 files changed
+893
-2
No files matched your search
@@ -0,0 +1,449 @@
|
||||
"""
|
||||
MCP (Model Context Protocol) Server Management API
|
||||
Phase 3.3: Manage MCP servers, discover tools/resources, test connections
|
||||
"""
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Body
|
||||
from sqlalchemy import or_, func, desc, asc, text
|
||||
from sqlalchemy.orm import Session
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import Optional, List, Dict, Any
|
||||
from datetime import datetime
|
||||
import uuid
|
||||
import json
|
||||
import logging
|
||||
|
||||
from ...database import get_db
|
||||
from ...models import gen_uuid
|
||||
from ...auth import get_current_admin
|
||||
from ...middleware.log_service import LogService
|
||||
|
||||
router = APIRouter()
|
||||
logger = logging.getLogger("aiflow")
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Pydantic Models
|
||||
# ============================================================
|
||||
|
||||
class MCPServerCreate(BaseModel):
|
||||
name: str
|
||||
description: Optional[str] = ""
|
||||
transport_type: str = "stdio"
|
||||
command: Optional[str] = ""
|
||||
args: Optional[Any] = None
|
||||
env_vars: Optional[Dict[str, str]] = None
|
||||
url: Optional[str] = ""
|
||||
headers: Optional[Dict[str, str]] = None
|
||||
is_enabled: Optional[bool] = True
|
||||
|
||||
|
||||
class MCPServerUpdate(BaseModel):
|
||||
name: Optional[str] = None
|
||||
description: Optional[str] = None
|
||||
transport_type: Optional[str] = None
|
||||
command: Optional[str] = None
|
||||
args: Optional[Any] = None
|
||||
env_vars: Optional[Dict[str, str]] = None
|
||||
url: Optional[str] = None
|
||||
headers: Optional[Dict[str, str]] = None
|
||||
is_enabled: Optional[bool] = None
|
||||
|
||||
|
||||
class MCPServerConnectionTest(BaseModel):
|
||||
transport_type: Optional[str] = None
|
||||
command: Optional[str] = None
|
||||
args: Optional[Any] = None
|
||||
env_vars: Optional[Dict[str, str]] = None
|
||||
url: Optional[str] = None
|
||||
headers: Optional[Dict[str, str]] = None
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Helper Functions
|
||||
# ============================================================
|
||||
|
||||
def _format_cst_time(dt):
|
||||
"""Format datetime to CST string"""
|
||||
if dt is None:
|
||||
return None
|
||||
try:
|
||||
from datetime import timezone, timedelta
|
||||
CST = timezone(timedelta(hours=8))
|
||||
if dt.tzinfo is None:
|
||||
dt = dt.replace(tzinfo=CST)
|
||||
return dt.astimezone(CST).strftime("%Y-%m-%d %H:%M:%S")
|
||||
except Exception:
|
||||
return dt.strftime("%Y-%m-%d %H:%M:%S") if dt else None
|
||||
|
||||
|
||||
def server_to_dict(server) -> dict:
|
||||
"""Convert MCPServer model to dict"""
|
||||
return {
|
||||
"id": server.id,
|
||||
"name": server.name,
|
||||
"description": server.description or "",
|
||||
"transport_type": server.transport_type or "stdio",
|
||||
"transport_type_label": {"stdio": "STDIO", "sse": "SSE"}.get(server.transport_type, server.transport_type),
|
||||
"command": server.command or "",
|
||||
"args": json.loads(server.args) if server.args else {},
|
||||
"env_vars": json.loads(server.env_vars) if server.env_vars else {},
|
||||
"url": server.url or "",
|
||||
"headers": json.loads(server.headers) if server.headers else {},
|
||||
"is_enabled": server.is_enabled if server.is_enabled is not None else True,
|
||||
"status": server.status or "disconnected",
|
||||
"last_connected_at": _format_cst_time(server.last_connected_at),
|
||||
"tools_count": server.tools_count or 0,
|
||||
"resources_count": server.resources_count or 0,
|
||||
"created_at": _format_cst_time(server.created_at),
|
||||
"updated_at": _format_cst_time(server.updated_at),
|
||||
}
|
||||
|
||||
|
||||
# ============================================================
|
||||
# API Endpoints
|
||||
# ============================================================
|
||||
|
||||
@router.get("/stats/overview", summary="Get MCP server statistics")
|
||||
async def get_mcp_stats(
|
||||
db: Session = Depends(get_db),
|
||||
admin_user=Depends(get_current_admin),
|
||||
):
|
||||
"""Get overview statistics for MCP servers"""
|
||||
from ...models import MCPServer
|
||||
total = db.query(func.count(MCPServer.id)).scalar() or 0
|
||||
enabled = db.query(func.count(MCPServer.id)).filter(MCPServer.is_enabled == True).scalar() or 0
|
||||
connected = db.query(func.count(MCPServer.id)).filter(MCPServer.status == "connected").scalar() or 0
|
||||
|
||||
# By transport type
|
||||
by_transport = {}
|
||||
rows = db.query(MCPServer.transport_type, func.count(MCPServer.id)).group_by(MCPServer.transport_type).all()
|
||||
for t, c in rows:
|
||||
by_transport[t or "unknown"] = c
|
||||
|
||||
# By status
|
||||
by_status = {}
|
||||
rows = db.query(MCPServer.status, func.count(MCPServer.id)).group_by(MCPServer.status).all()
|
||||
for s, c in rows:
|
||||
by_status[s or "unknown"] = c
|
||||
|
||||
return {
|
||||
"code": 0,
|
||||
"data": {
|
||||
"total_servers": total,
|
||||
"enabled_servers": enabled,
|
||||
"connected_servers": connected,
|
||||
"by_transport": by_transport,
|
||||
"by_status": by_status,
|
||||
},
|
||||
"message": "ok",
|
||||
}
|
||||
|
||||
|
||||
@router.get("/", summary="List MCP servers")
|
||||
async def list_mcp_servers(
|
||||
page: int = Query(1, ge=1),
|
||||
page_size: int = Query(20, ge=1, le=100),
|
||||
search: Optional[str] = Query(None),
|
||||
transport_type: Optional[str] = Query(None),
|
||||
status: Optional[str] = Query(None),
|
||||
is_enabled: Optional[bool] = Query(None),
|
||||
db: Session = Depends(get_db),
|
||||
admin_user=Depends(get_current_admin),
|
||||
):
|
||||
"""List all MCP servers with pagination and filters"""
|
||||
from ...models import MCPServer
|
||||
query = db.query(MCPServer)
|
||||
|
||||
if search:
|
||||
query = query.filter(
|
||||
or_(MCPServer.name.contains(search), MCPServer.description.contains(search))
|
||||
)
|
||||
if transport_type:
|
||||
query = query.filter(MCPServer.transport_type == transport_type)
|
||||
if status:
|
||||
query = query.filter(MCPServer.status == status)
|
||||
if is_enabled is not None:
|
||||
query = query.filter(MCPServer.is_enabled == is_enabled)
|
||||
|
||||
total = query.count()
|
||||
servers = query.order_by(desc(MCPServer.created_at)).offset((page - 1) * page_size).limit(page_size).all()
|
||||
|
||||
return {
|
||||
"code": 0,
|
||||
"data": {
|
||||
"total": total,
|
||||
"page": page,
|
||||
"page_size": page_size,
|
||||
"items": [server_to_dict(s) for s in servers],
|
||||
},
|
||||
"message": "ok",
|
||||
}
|
||||
|
||||
|
||||
@router.get("/{server_id}", summary="Get MCP server details")
|
||||
async def get_mcp_server(
|
||||
server_id: str,
|
||||
db: Session = Depends(get_db),
|
||||
admin_user=Depends(get_current_admin),
|
||||
):
|
||||
"""Get a single MCP server by ID"""
|
||||
from ...models import MCPServer
|
||||
server = db.query(MCPServer).filter(MCPServer.id == server_id).first()
|
||||
if not server:
|
||||
raise HTTPException(status_code=404, detail="MCP server not found")
|
||||
return {"code": 0, "data": server_to_dict(server), "message": "ok"}
|
||||
|
||||
|
||||
@router.post("/", summary="Create MCP server")
|
||||
async def create_mcp_server(
|
||||
data: MCPServerCreate,
|
||||
db: Session = Depends(get_db),
|
||||
admin_user=Depends(get_current_admin),
|
||||
):
|
||||
"""Register a new MCP server"""
|
||||
from ...models import MCPServer
|
||||
|
||||
# Check name uniqueness
|
||||
existing = db.query(MCPServer).filter(MCPServer.name == data.name).first()
|
||||
if existing:
|
||||
raise HTTPException(status_code=400, detail="MCP server with this name already exists")
|
||||
|
||||
now = datetime.now()
|
||||
server = MCPServer(
|
||||
id=gen_uuid(),
|
||||
name=data.name,
|
||||
description=data.description or "",
|
||||
transport_type=data.transport_type,
|
||||
command=data.command or "",
|
||||
args=json.dumps(data.args) if data.args else "{}",
|
||||
env_vars=json.dumps(data.env_vars) if data.env_vars else "{}",
|
||||
url=data.url or "",
|
||||
headers=json.dumps(data.headers) if data.headers else "{}",
|
||||
is_enabled=data.is_enabled if data.is_enabled is not None else True,
|
||||
status="disconnected",
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
)
|
||||
db.add(server)
|
||||
db.commit()
|
||||
db.refresh(server)
|
||||
|
||||
logger.info(f"MCP server created: {server.name} ({server.id})")
|
||||
return {"code": 0, "data": server_to_dict(server), "message": "MCP server created"}
|
||||
|
||||
|
||||
@router.put("/{server_id}", summary="Update MCP server")
|
||||
async def update_mcp_server(
|
||||
server_id: str,
|
||||
data: MCPServerUpdate,
|
||||
db: Session = Depends(get_db),
|
||||
admin_user=Depends(get_current_admin),
|
||||
):
|
||||
"""Update an existing MCP server"""
|
||||
from ...models import MCPServer
|
||||
server = db.query(MCPServer).filter(MCPServer.id == server_id).first()
|
||||
if not server:
|
||||
raise HTTPException(status_code=404, detail="MCP server not found")
|
||||
|
||||
update_fields = data.model_dump(exclude_unset=True)
|
||||
|
||||
# Handle JSON fields
|
||||
for json_field in ["args", "env_vars", "headers"]:
|
||||
if json_field in update_fields and update_fields[json_field] is not None:
|
||||
update_fields[json_field] = json.dumps(update_fields[json_field])
|
||||
|
||||
for key, value in update_fields.items():
|
||||
if value is not None:
|
||||
setattr(server, key, value)
|
||||
|
||||
server.updated_at = datetime.now()
|
||||
db.commit()
|
||||
db.refresh(server)
|
||||
|
||||
logger.info(f"MCP server updated: {server.name} ({server.id})")
|
||||
return {"code": 0, "data": server_to_dict(server), "message": "MCP server updated"}
|
||||
|
||||
|
||||
@router.delete("/{server_id}", summary="Delete MCP server")
|
||||
async def delete_mcp_server(
|
||||
server_id: str,
|
||||
db: Session = Depends(get_db),
|
||||
admin_user=Depends(get_current_admin),
|
||||
):
|
||||
"""Delete an MCP server"""
|
||||
from ...models import MCPServer
|
||||
server = db.query(MCPServer).filter(MCPServer.id == server_id).first()
|
||||
if not server:
|
||||
raise HTTPException(status_code=404, detail="MCP server not found")
|
||||
|
||||
name = server.name
|
||||
db.delete(server)
|
||||
db.commit()
|
||||
|
||||
logger.info(f"MCP server deleted: {name} ({server_id})")
|
||||
return {"code": 0, "data": None, "message": "MCP server deleted"}
|
||||
|
||||
|
||||
@router.put("/{server_id}/toggle", summary="Enable/Disable MCP server")
|
||||
async def toggle_mcp_server(
|
||||
server_id: str,
|
||||
db: Session = Depends(get_db),
|
||||
admin_user=Depends(get_current_admin),
|
||||
):
|
||||
"""Toggle MCP server enabled state"""
|
||||
from ...models import MCPServer
|
||||
server = db.query(MCPServer).filter(MCPServer.id == server_id).first()
|
||||
if not server:
|
||||
raise HTTPException(status_code=404, detail="MCP server not found")
|
||||
|
||||
server.is_enabled = not server.is_enabled
|
||||
server.updated_at = datetime.now()
|
||||
db.commit()
|
||||
|
||||
logger.info(f"MCP server toggled: {server.name} -> enabled={server.is_enabled}")
|
||||
return {
|
||||
"code": 0,
|
||||
"data": {"id": server.id, "is_enabled": server.is_enabled},
|
||||
"message": "MCP server toggled",
|
||||
}
|
||||
|
||||
|
||||
@router.post("/{server_id}/test", summary="Test MCP server connection")
|
||||
async def test_mcp_connection(
|
||||
server_id: str,
|
||||
db: Session = Depends(get_db),
|
||||
admin_user=Depends(get_current_admin),
|
||||
test_data: Optional[MCPServerConnectionTest] = None,
|
||||
):
|
||||
"""Test connection to an MCP server"""
|
||||
from ...models import MCPServer
|
||||
server = db.query(MCPServer).filter(MCPServer.id == server_id).first()
|
||||
if not server:
|
||||
raise HTTPException(status_code=404, detail="MCP server not found")
|
||||
|
||||
transport = test_data.transport_type if test_data and test_data.transport_type else server.transport_type
|
||||
result = {"success": False, "message": "", "details": {}}
|
||||
|
||||
try:
|
||||
if transport == "sse":
|
||||
sse_url = (test_data.url if test_data else None) or server.url
|
||||
if not sse_url:
|
||||
result["message"] = "SSE URL is required"
|
||||
return {"code": 0, "data": result, "message": "ok"}
|
||||
|
||||
# Test SSE endpoint with a simple HTTP request
|
||||
import httpx
|
||||
try:
|
||||
response = httpx.get(sse_url, timeout=10.0)
|
||||
result["success"] = response.status_code < 500
|
||||
result["message"] = f"SSE endpoint responded with {response.status_code}"
|
||||
result["details"]["status_code"] = response.status_code
|
||||
result["details"]["content_type"] = response.headers.get("content-type", "")
|
||||
except httpx.ConnectError:
|
||||
result["message"] = "Connection refused - server may be down"
|
||||
except httpx.TimeoutException:
|
||||
result["message"] = "Connection timed out"
|
||||
|
||||
elif transport == "stdio":
|
||||
cmd = (test_data.command if test_data else None) or server.command
|
||||
if not cmd:
|
||||
result["message"] = "STDIO command is required"
|
||||
return {"code": 0, "data": result, "message": "ok"}
|
||||
|
||||
# Basic command existence check (not actually executing MCP handshake)
|
||||
import subprocess
|
||||
try:
|
||||
subprocess.run([cmd, "--version"], capture_output=True, timeout=10)
|
||||
result["success"] = True
|
||||
result["message"] = "Command found and executable"
|
||||
except FileNotFoundError:
|
||||
result["message"] = "Command not found in PATH"
|
||||
except subprocess.TimeoutExpired:
|
||||
result["message"] = "Command timed out"
|
||||
except Exception as e:
|
||||
result["message"] = f"Command test failed"
|
||||
result["details"]["error"] = str(e)
|
||||
|
||||
# Update server status
|
||||
server.status = "connected" if result["success"] else "error"
|
||||
if result["success"]:
|
||||
server.last_connected_at = datetime.now()
|
||||
server.updated_at = datetime.now()
|
||||
db.commit()
|
||||
|
||||
except Exception as e:
|
||||
result["message"] = str(e)
|
||||
result["details"]["error"] = str(e)
|
||||
server.status = "error"
|
||||
server.updated_at = datetime.now()
|
||||
db.commit()
|
||||
|
||||
return {"code": 0, "data": result, "message": "ok"}
|
||||
|
||||
|
||||
@router.get("/{server_id}/tools", summary="List MCP server tools")
|
||||
async def list_mcp_tools(
|
||||
server_id: str,
|
||||
db: Session = Depends(get_db),
|
||||
admin_user=Depends(get_current_admin),
|
||||
):
|
||||
"""List tools exposed by an MCP server (placeholder/mock)"""
|
||||
from ...models import MCPServer
|
||||
server = db.query(MCPServer).filter(MCPServer.id == server_id).first()
|
||||
if not server:
|
||||
raise HTTPException(status_code=404, detail="MCP server not found")
|
||||
|
||||
# In a real implementation, this would connect to the MCP server and call tools/list
|
||||
# For now, return mock data based on common MCP tools
|
||||
mock_tools = []
|
||||
if server.transport_type == "stdio":
|
||||
mock_tools = [
|
||||
{"name": "read_file", "description": "Read contents of a file"},
|
||||
{"name": "write_file", "description": "Write contents to a file"},
|
||||
{"name": "execute_command", "description": "Execute a shell command"},
|
||||
{"name": "search_content", "description": "Search for content in files"},
|
||||
]
|
||||
elif server.transport_type == "sse":
|
||||
mock_tools = [
|
||||
{"name": "web_search", "description": "Search the web"},
|
||||
{"name": "fetch_url", "description": "Fetch content from URL"},
|
||||
{"name": "api_request", "description": "Make API requests"},
|
||||
]
|
||||
|
||||
return {
|
||||
"code": 0,
|
||||
"data": {
|
||||
"server_name": server.name,
|
||||
"tools": mock_tools,
|
||||
"count": len(mock_tools),
|
||||
},
|
||||
"message": "ok",
|
||||
}
|
||||
|
||||
|
||||
@router.get("/{server_id}/resources", summary="List MCP server resources")
|
||||
async def list_mcp_resources(
|
||||
server_id: str,
|
||||
db: Session = Depends(get_db),
|
||||
admin_user=Depends(get_current_admin),
|
||||
):
|
||||
"""List resources exposed by an MCP server (placeholder/mock)"""
|
||||
from ...models import MCPServer
|
||||
server = db.query(MCPServer).filter(MCPServer.id == server_id).first()
|
||||
if not server:
|
||||
raise HTTPException(status_code=404, detail="MCP server not found")
|
||||
|
||||
mock_resources = [
|
||||
{"uri": f"{server.name}://docs/readme", "name": "Documentation", "mime_type": "text/markdown"},
|
||||
{"uri": f"{server.name}://config/settings", "name": "Configuration", "mime_type": "application/json"},
|
||||
]
|
||||
|
||||
return {
|
||||
"code": 0,
|
||||
"data": {
|
||||
"server_name": server.name,
|
||||
"resources": mock_resources,
|
||||
"count": len(mock_resources),
|
||||
},
|
||||
"message": "ok",
|
||||
}
|
||||
Reference in new issue
Block a user