1261 lines
56 KiB
Python
1261 lines
56 KiB
Python
#!/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
|
|
"""
|
|
Doris MCP SSE Server Implementation
|
|
|
|
Implements a standard MCP SSE server based on MCP's SseServerTransport,
|
|
supports bidirectional communication with clients, and integrates with the existing Doris-MCP-Server.
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import uuid
|
|
import logging
|
|
import time
|
|
from typing import Any, Optional
|
|
from fastapi import FastAPI, Request
|
|
from fastapi.responses import JSONResponse
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from sse_starlette.sse import EventSourceResponse
|
|
|
|
# Get logger
|
|
logger = logging.getLogger("doris-mcp-sse")
|
|
|
|
class DorisMCPSseServer:
|
|
"""Doris MCP SSE Server Implementation"""
|
|
|
|
def __init__(self, mcp_server, app: FastAPI):
|
|
"""
|
|
Initialize the Doris MCP SSE server
|
|
|
|
Args:
|
|
mcp_server: FastMCP server instance
|
|
app: FastAPI application instance
|
|
"""
|
|
self.mcp_server = mcp_server
|
|
|
|
# Ensure app is a FastAPI instance
|
|
if not isinstance(app, FastAPI):
|
|
logger.warning("Passed application is not a FastAPI instance, will use the existing FastAPI instance")
|
|
|
|
self.app = app
|
|
|
|
# Add CORS middleware
|
|
self.app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["http://localhost:3100"], # Specify frontend domain
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
expose_headers=["*"]
|
|
)
|
|
|
|
# Client session management
|
|
self.client_sessions = {}
|
|
|
|
# Set up SSE routes
|
|
self.setup_sse_routes()
|
|
|
|
# Register startup event
|
|
@self.app.on_event("startup")
|
|
async def startup_event():
|
|
# Start session cleanup task
|
|
asyncio.create_task(self.cleanup_idle_sessions())
|
|
# Start task to send periodic status updates
|
|
asyncio.create_task(self.send_periodic_updates())
|
|
|
|
def setup_sse_routes(self):
|
|
"""Set up SSE related routes"""
|
|
|
|
@self.app.get("/health")
|
|
async def health_check():
|
|
"""Health check endpoint"""
|
|
try:
|
|
# Use direct health check logic
|
|
return {
|
|
"status": "healthy",
|
|
"timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
|
|
"server": "Doris MCP Server"
|
|
}
|
|
except Exception as e:
|
|
return {
|
|
"status": "error",
|
|
"error": str(e),
|
|
"timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
|
|
}
|
|
|
|
@self.app.get("/status")
|
|
async def status():
|
|
"""Get server status"""
|
|
try:
|
|
# Get tool list
|
|
tools = await self.mcp_server.list_tools()
|
|
tool_names = [tool.name if hasattr(tool, 'name') else str(tool) for tool in tools]
|
|
logger.info(f"Getting tool list, currently registered tools: {tool_names}")
|
|
|
|
# Get resource list
|
|
resources = await self.mcp_server.list_resources()
|
|
resource_names = [res.name if hasattr(res, 'name') else str(res) for res in resources]
|
|
|
|
# Get prompt template list
|
|
prompts = await self.mcp_server.list_prompts()
|
|
prompt_names = [prompt.name if hasattr(prompt, 'name') else str(prompt) for prompt in prompts]
|
|
|
|
return {
|
|
"status": "running",
|
|
"name": self.mcp_server.name,
|
|
"mode": "mcp_sse",
|
|
"clients": len(self.client_sessions),
|
|
"tools": tool_names,
|
|
"resources": resource_names,
|
|
"prompts": prompt_names
|
|
}
|
|
except Exception as e:
|
|
logger.error(f"Error getting status: {str(e)}")
|
|
return {
|
|
"status": "error",
|
|
"error": str(e)
|
|
}
|
|
|
|
@self.app.get("/sse")
|
|
async def mcp_sse_init(request: Request):
|
|
"""SSE service entry point, establishes client connection (New Endpoint)"""
|
|
# Generate session ID
|
|
session_id = str(uuid.uuid4())
|
|
logger.info(f"New SSE connection [Session ID: {session_id}] at /sse")
|
|
|
|
# Create client session
|
|
self.client_sessions[session_id] = {
|
|
"client_id": request.headers.get("X-Client-ID", f"client_{str(uuid.uuid4())[:8]}"),
|
|
"created_at": time.time(),
|
|
"last_active": time.time(),
|
|
"queue": asyncio.Queue()
|
|
}
|
|
|
|
# Immediately put endpoint information into the queue
|
|
endpoint_data = f"/mcp/messages?session_id={session_id}"
|
|
await self.client_sessions[session_id]["queue"].put({
|
|
"event": "endpoint",
|
|
"data": endpoint_data
|
|
})
|
|
|
|
# Create event generator
|
|
async def event_generator():
|
|
try:
|
|
while True:
|
|
# Use timeout to get new messages, to detect client disconnect
|
|
try:
|
|
message = await asyncio.wait_for(
|
|
self.client_sessions[session_id]["queue"].get(),
|
|
timeout=30
|
|
)
|
|
|
|
# Check if it's a close command
|
|
if isinstance(message, dict) and message.get("event") == "close":
|
|
logger.info(f"Received close command [Session ID: {session_id}]")
|
|
break
|
|
|
|
# Return message
|
|
if isinstance(message, dict):
|
|
if "event" in message:
|
|
# If event field exists, it's a system event
|
|
event_type = message["event"]
|
|
event_data = message["data"]
|
|
yield {
|
|
"event": event_type,
|
|
"data": event_data
|
|
}
|
|
else:
|
|
# Otherwise it's a normal message, use message event
|
|
yield {
|
|
"event": "message",
|
|
"data": json.dumps(message)
|
|
}
|
|
elif isinstance(message, str):
|
|
# If it's a string, send directly
|
|
yield {
|
|
"event": "message",
|
|
"data": message
|
|
}
|
|
else:
|
|
# Other types, convert to JSON
|
|
yield {
|
|
"event": "message",
|
|
"data": json.dumps(message)
|
|
}
|
|
except asyncio.TimeoutError:
|
|
# Send ping to keep connection alive
|
|
yield {
|
|
"event": "ping",
|
|
"data": "keepalive"
|
|
}
|
|
continue
|
|
except asyncio.CancelledError:
|
|
# Connection cancelled
|
|
logger.info(f"SSE connection cancelled [Session ID: {session_id}]")
|
|
except Exception as e:
|
|
# Other error occurred
|
|
logger.error(f"SSE event generator error [Session ID: {session_id}]: {str(e)}")
|
|
finally:
|
|
# Clean up session
|
|
if session_id in self.client_sessions:
|
|
logger.info(f"Cleaning up session [Session ID: {session_id}]")
|
|
del self.client_sessions[session_id]
|
|
|
|
# Return standard SSE response
|
|
return EventSourceResponse(
|
|
event_generator(),
|
|
headers={
|
|
"Access-Control-Allow-Origin": "*",
|
|
"Access-Control-Allow-Credentials": "true",
|
|
"Access-Control-Allow-Methods": "*",
|
|
"Access-Control-Allow-Headers": "*",
|
|
"Access-Control-Expose-Headers": "*",
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive"
|
|
}
|
|
)
|
|
|
|
@self.app.options("/mcp/messages")
|
|
async def mcp_messages_options(request: Request):
|
|
"""Handle preflight requests"""
|
|
return JSONResponse(
|
|
{},
|
|
headers={
|
|
"Access-Control-Allow-Origin": "*",
|
|
"Access-Control-Allow-Credentials": "true",
|
|
"Access-Control-Allow-Methods": "*",
|
|
"Access-Control-Allow-Headers": "*",
|
|
"Access-Control-Expose-Headers": "*"
|
|
}
|
|
)
|
|
|
|
@self.app.post("/mcp/messages")
|
|
async def mcp_messages_handler(request: Request):
|
|
"""Handle client message requests, using class method"""
|
|
return await self.mcp_message(request)
|
|
|
|
async def cleanup_idle_sessions(self):
|
|
"""Clean up idle client sessions"""
|
|
while True:
|
|
await asyncio.sleep(60) # Check every minute
|
|
current_time = time.time()
|
|
|
|
# Find sessions idle for over 5 minutes
|
|
idle_sessions = []
|
|
for session_id, session in self.client_sessions.items():
|
|
if current_time - session["last_active"] > 300: # 5 minutes
|
|
idle_sessions.append(session_id)
|
|
|
|
# Close and remove idle sessions
|
|
for session_id in idle_sessions:
|
|
try:
|
|
# Send close message
|
|
await self.client_sessions[session_id]["queue"].put({"event": "close"})
|
|
# Clean up session
|
|
logger.info(f"Cleaned up idle session: {session_id}")
|
|
except Exception as e:
|
|
logger.error(f"Error cleaning up session: {str(e)}")
|
|
finally:
|
|
# Ensure session is removed
|
|
if session_id in self.client_sessions:
|
|
del self.client_sessions[session_id]
|
|
|
|
async def send_periodic_updates(self):
|
|
"""Periodically send status updates to all clients"""
|
|
while True:
|
|
try:
|
|
# Send status update every 5 seconds
|
|
await asyncio.sleep(5)
|
|
|
|
# If no clients are connected, skip this update
|
|
if not self.client_sessions:
|
|
continue
|
|
|
|
# Get current status
|
|
status_data = {
|
|
"timestamp": time.time(),
|
|
"clients_count": len(self.client_sessions),
|
|
"server_status": "running"
|
|
}
|
|
|
|
# Send status update to all clients
|
|
await self.broadcast_status_update(status_data)
|
|
except Exception as e:
|
|
logger.error(f"Error sending periodic updates: {str(e)}")
|
|
# Wait a bit after error before continuing
|
|
await asyncio.sleep(1)
|
|
|
|
async def broadcast_status_update(self, status_data):
|
|
"""Broadcast status update to all clients
|
|
|
|
Args:
|
|
status_data: Status data
|
|
"""
|
|
logger.debug(f"Broadcasting status update: {status_data}")
|
|
message = {
|
|
"jsonrpc": "2.0",
|
|
"method": "notifications/status",
|
|
"params": {
|
|
"type": "status_update",
|
|
"data": status_data
|
|
}
|
|
}
|
|
await self.broadcast_message(message)
|
|
|
|
async def broadcast_visualization_data(self, visualization_data):
|
|
"""Broadcast visualization data to all clients
|
|
|
|
Args:
|
|
visualization_data: Visualization data, should include type field
|
|
"""
|
|
if not visualization_data or not isinstance(visualization_data, dict) or "type" not in visualization_data:
|
|
logger.warning(f"Invalid visualization data: {visualization_data}")
|
|
return
|
|
|
|
logger.info(f"Broadcasting visualization data: {visualization_data['type']}")
|
|
message = {
|
|
"jsonrpc": "2.0",
|
|
"method": "notifications/visualization",
|
|
"params": {
|
|
"type": "visualization",
|
|
"data": visualization_data
|
|
}
|
|
}
|
|
await self.broadcast_message(message)
|
|
|
|
async def send_visualization_data(self, session_id, visualization_data):
|
|
"""Send visualization data to a specific client
|
|
|
|
Args:
|
|
session_id: Session ID
|
|
visualization_data: Visualization data, should include type field
|
|
"""
|
|
if not visualization_data or not isinstance(visualization_data, dict) or "type" not in visualization_data:
|
|
logger.warning(f"Invalid visualization data: {visualization_data}")
|
|
return
|
|
|
|
if session_id not in self.client_sessions:
|
|
logger.warning(f"Session does not exist: {session_id}")
|
|
return
|
|
|
|
logger.info(f"Sending visualization data to session {session_id}: {visualization_data['type']}")
|
|
message = {
|
|
"jsonrpc": "2.0",
|
|
"method": "notifications/visualization",
|
|
"params": {
|
|
"type": "visualization",
|
|
"data": visualization_data
|
|
}
|
|
}
|
|
await self.client_sessions[session_id]["queue"].put(message)
|
|
|
|
async def send_tool_result(self, session_id, tool_name, result_data, is_final=True):
|
|
"""Send tool execution result to the client
|
|
|
|
Args:
|
|
session_id: Session ID
|
|
tool_name: Tool name
|
|
result_data: Result data
|
|
is_final: Whether it is the final result
|
|
"""
|
|
if session_id not in self.client_sessions:
|
|
logger.warning(f"Session does not exist: {session_id}")
|
|
return
|
|
|
|
logger.info(f"Sending tool result to session {session_id}: {tool_name}")
|
|
message = {
|
|
"jsonrpc": "2.0",
|
|
"method": "notifications/tool_result",
|
|
"params": {
|
|
"type": "tool_result",
|
|
"tool": tool_name,
|
|
"result": result_data,
|
|
"is_final": is_final
|
|
}
|
|
}
|
|
await self.client_sessions[session_id]["queue"].put(message)
|
|
|
|
async def broadcast_message(self, message):
|
|
"""Broadcast a message to all active sessions
|
|
|
|
Args:
|
|
message: Message to broadcast
|
|
"""
|
|
# If no clients are connected, return immediately
|
|
if not self.client_sessions:
|
|
return
|
|
|
|
# Create a copy of the session ID list so the original dictionary can be safely modified during iteration
|
|
session_ids = list(self.client_sessions.keys())
|
|
|
|
# Send message to all sessions
|
|
for session_id in session_ids:
|
|
try:
|
|
if session_id in self.client_sessions: # Check again, as session might have been removed during iteration
|
|
await self.client_sessions[session_id]["queue"].put(message)
|
|
except Exception as e:
|
|
logger.error(f"Error sending message to session {session_id}: {str(e)}")
|
|
|
|
async def get_status(self):
|
|
"""Get server status"""
|
|
return {
|
|
"status": "running",
|
|
"name": self.mcp_server.name,
|
|
"mode": "mcp_sse",
|
|
"clients": len(self.client_sessions)
|
|
}
|
|
|
|
async def mcp_message(self, request: Request):
|
|
"""Endpoint to receive client messages"""
|
|
try:
|
|
# Parse request parameters
|
|
session_id = self._get_session_id(request)
|
|
|
|
# Check if session exists
|
|
if not session_id or session_id not in self.client_sessions:
|
|
logger.warning(f"Session does not exist: {session_id}")
|
|
return JSONResponse(
|
|
{"jsonrpc": "2.0", "error": {"code": -32000, "message": "Session does not exist or has expired"}},
|
|
status_code=401
|
|
)
|
|
|
|
# Update session last active time
|
|
self.client_sessions[session_id]["last_active"] = time.time()
|
|
|
|
# Get request body
|
|
try:
|
|
body = await request.json()
|
|
logger.info(f"Received message [Session ID: {session_id}]: {json.dumps(body)}")
|
|
|
|
# Process message
|
|
message_id = body.get("id", str(uuid.uuid4()))
|
|
|
|
# Handle JSON-RPC 2.0 formatted commands
|
|
if "jsonrpc" not in body or body.get("jsonrpc") != "2.0" or "method" not in body:
|
|
return JSONResponse(
|
|
{"jsonrpc": "2.0", "id": message_id, "error": {"code": -32600, "message": "Invalid request, must be JSON-RPC 2.0 format"}},
|
|
status_code=400
|
|
)
|
|
|
|
# Get method and parameters
|
|
method = body.get("method")
|
|
params = body.get("params", {})
|
|
|
|
# Special handling for JSON-RPC formatted commands
|
|
if method == "initialize":
|
|
# Initialization request
|
|
logger.info(f"Processing initialize command [Session ID: {session_id}]")
|
|
response = {
|
|
"jsonrpc": "2.0",
|
|
"id": message_id,
|
|
"result": {
|
|
"protocolVersion": "2024-11-05",
|
|
"name": self.mcp_server.name,
|
|
"instructions": "This is an MCP server for Apache Doris database",
|
|
"serverInfo": {
|
|
"version": "0.1.0",
|
|
"name": "Doris MCP Server"
|
|
},
|
|
"capabilities": {
|
|
"tools": {
|
|
"supportsStreaming": True,
|
|
"supportsProgress": True
|
|
},
|
|
"resources": {
|
|
"supportsStreaming": False
|
|
},
|
|
"prompts": {
|
|
"supported": True
|
|
}
|
|
}
|
|
}
|
|
}
|
|
await self.client_sessions[session_id]["queue"].put(response)
|
|
return JSONResponse(
|
|
{"status": "success"},
|
|
headers={
|
|
"Access-Control-Allow-Origin": "*",
|
|
"Access-Control-Allow-Credentials": "true",
|
|
"Access-Control-Allow-Methods": "*",
|
|
"Access-Control-Allow-Headers": "*",
|
|
"Access-Control-Expose-Headers": "*"
|
|
}
|
|
)
|
|
|
|
elif method == "mcp/listOfferings":
|
|
# List all available features
|
|
logger.info(f"Processing listOfferings command [Session ID: {session_id}]")
|
|
|
|
# Get tool list
|
|
tools = await self.mcp_server.list_tools()
|
|
tools_json = [
|
|
{
|
|
"name": tool.name if hasattr(tool, "name") else str(tool),
|
|
"description": tool.description if hasattr(tool, "description") else "",
|
|
"inputSchema": tool.parameters if hasattr(tool, "parameters") else {
|
|
"type": "object",
|
|
"properties": {},
|
|
"required": []
|
|
}
|
|
}
|
|
for tool in tools
|
|
]
|
|
|
|
# Get resource list
|
|
resources = await self.mcp_server.list_resources()
|
|
resources_json = [res.model_dump() if hasattr(res, "model_dump") else res for res in resources]
|
|
|
|
# Get prompt template list
|
|
prompts = await self.mcp_server.list_prompts()
|
|
prompts_json = [prompt.model_dump() if hasattr(prompt, "model_dump") else prompt for prompt in prompts]
|
|
|
|
# Build response
|
|
response = {
|
|
"jsonrpc": "2.0",
|
|
"id": message_id,
|
|
"result": {
|
|
"tools": tools_json,
|
|
"resources": resources_json,
|
|
"prompts": prompts_json
|
|
}
|
|
}
|
|
await self.client_sessions[session_id]["queue"].put(response)
|
|
return JSONResponse(
|
|
{"status": "success"},
|
|
headers={
|
|
"Access-Control-Allow-Origin": "*",
|
|
"Access-Control-Allow-Credentials": "true",
|
|
"Access-Control-Allow-Methods": "*",
|
|
"Access-Control-Allow-Headers": "*",
|
|
"Access-Control-Expose-Headers": "*"
|
|
}
|
|
)
|
|
|
|
elif method == "mcp/listTools" or method == "tools/list":
|
|
# List all tools
|
|
logger.info(f"Processing listTools command [Session ID: {session_id}]")
|
|
tools = await self.mcp_server.list_tools()
|
|
tools_json = [
|
|
{
|
|
"name": tool.name if hasattr(tool, "name") else str(tool),
|
|
"description": tool.description if hasattr(tool, "description") else "",
|
|
"inputSchema": tool.parameters if hasattr(tool, "parameters") else {
|
|
"type": "object",
|
|
"properties": {},
|
|
"required": []
|
|
}
|
|
}
|
|
for tool in tools
|
|
]
|
|
response = {
|
|
"jsonrpc": "2.0",
|
|
"id": message_id,
|
|
"result": {
|
|
"tools": tools_json
|
|
}
|
|
}
|
|
await self.client_sessions[session_id]["queue"].put(response)
|
|
return JSONResponse(
|
|
{"status": "success"},
|
|
headers={
|
|
"Access-Control-Allow-Origin": "*",
|
|
"Access-Control-Allow-Credentials": "true",
|
|
"Access-Control-Allow-Methods": "*",
|
|
"Access-Control-Allow-Headers": "*",
|
|
"Access-Control-Expose-Headers": "*"
|
|
}
|
|
)
|
|
|
|
elif method == "mcp/listResources":
|
|
# List all resources
|
|
logger.info(f"Processing listResources command [Session ID: {session_id}]")
|
|
resources = await self.mcp_server.list_resources()
|
|
resources_json = [res.model_dump() if hasattr(res, "model_dump") else res for res in resources]
|
|
response = {
|
|
"jsonrpc": "2.0",
|
|
"id": message_id,
|
|
"result": {
|
|
"resources": resources_json
|
|
}
|
|
}
|
|
await self.client_sessions[session_id]["queue"].put(response)
|
|
return JSONResponse(
|
|
{"status": "success"},
|
|
headers={
|
|
"Access-Control-Allow-Origin": "*",
|
|
"Access-Control-Allow-Credentials": "true",
|
|
"Access-Control-Allow-Methods": "*",
|
|
"Access-Control-Allow-Headers": "*",
|
|
"Access-Control-Expose-Headers": "*"
|
|
}
|
|
)
|
|
|
|
elif method == "mcp/callTool" or method == "tools/call":
|
|
# Call tool - special handling
|
|
tool_name = params.get("name")
|
|
arguments = params.get("arguments", {})
|
|
|
|
if not tool_name:
|
|
error_response = {
|
|
"jsonrpc": "2.0",
|
|
"id": message_id,
|
|
"error": {
|
|
"code": -32602,
|
|
"message": "Invalid params: tool name is required"
|
|
}
|
|
}
|
|
await self.client_sessions[session_id]["queue"].put(error_response)
|
|
return JSONResponse(
|
|
{"status": "error", "message": "Tool name is required"},
|
|
status_code=400,
|
|
headers={
|
|
"Access-Control-Allow-Origin": "*",
|
|
"Access-Control-Allow-Credentials": "true",
|
|
"Access-Control-Allow-Methods": "*",
|
|
"Access-Control-Allow-Headers": "*",
|
|
"Access-Control-Expose-Headers": "*"
|
|
}
|
|
)
|
|
|
|
# Get MCP instance
|
|
mcp = self.mcp_server
|
|
|
|
# Check if it's a streaming tool call
|
|
stream_mode = "stream" in request.query_params or params.get("stream", False)
|
|
|
|
logger.info(f"Calling tool [Session ID: {session_id}, Tool: {tool_name}, Args: {arguments}, Stream mode: {stream_mode}]")
|
|
|
|
if stream_mode:
|
|
# Streaming tool call
|
|
logger.info(f"Using streaming response to handle tool call [Session ID: {session_id}, Tool: {tool_name}]")
|
|
|
|
try:
|
|
# Find tool
|
|
tool_instance = None
|
|
for tool in await mcp.list_tools():
|
|
if getattr(tool, 'name', '') == tool_name:
|
|
tool_instance = tool
|
|
break
|
|
|
|
if not tool_instance:
|
|
raise Exception(f"Tool {tool_name} does not exist")
|
|
|
|
# Define callback function
|
|
async def callback(content, metadata):
|
|
# Send partial result
|
|
partial_message = {
|
|
"jsonrpc": "2.0",
|
|
"id": message_id,
|
|
"partial": True,
|
|
"result": {
|
|
"content": content,
|
|
"metadata": metadata
|
|
}
|
|
}
|
|
# Put message into queue
|
|
await self.client_sessions[session_id]["queue"].put(partial_message)
|
|
|
|
# If visualization data is included, broadcast to all clients
|
|
if metadata and "visualization" in metadata:
|
|
await self.broadcast_visualization_data(metadata["visualization"])
|
|
|
|
# Build argument dictionary
|
|
kwargs = dict(arguments)
|
|
kwargs['callback'] = callback
|
|
|
|
# Execute tool call
|
|
# Fix: Do not call tool object directly, instead create async task to call call_tool method
|
|
# func = tool_instance.func if hasattr(tool_instance, 'func') else tool_instance
|
|
|
|
# Start async task to execute streaming tool
|
|
# asyncio.create_task(self._execute_stream_tool(func, kwargs, message_id, session_id))
|
|
# Modified to use call_tool method
|
|
asyncio.create_task(self._execute_stream_tool_wrapper(tool_name, kwargs, message_id, session_id, request))
|
|
|
|
# Return received confirmation
|
|
return JSONResponse(
|
|
{"status": "processing"},
|
|
headers={
|
|
"Access-Control-Allow-Origin": "*",
|
|
"Access-Control-Allow-Credentials": "true",
|
|
"Access-Control-Allow-Methods": "*",
|
|
"Access-Control-Allow-Headers": "*",
|
|
"Access-Control-Expose-Headers": "*"
|
|
}
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Streaming tool processing error: {str(e)}")
|
|
error_message = {
|
|
"jsonrpc": "2.0",
|
|
"id": message_id,
|
|
"success": False,
|
|
"error": str(e)
|
|
}
|
|
# Put error message into queue
|
|
await self.client_sessions[session_id]["queue"].put(error_message)
|
|
return JSONResponse(
|
|
{"status": "error", "message": str(e)},
|
|
status_code=500,
|
|
headers={
|
|
"Access-Control-Allow-Origin": "*",
|
|
"Access-Control-Allow-Credentials": "true",
|
|
"Access-Control-Allow-Methods": "*",
|
|
"Access-Control-Allow-Headers": "*",
|
|
"Access-Control-Expose-Headers": "*"
|
|
}
|
|
)
|
|
else:
|
|
# Non-streaming tool call
|
|
logger.info(f"Using standard response to handle tool call [Session ID: {session_id}, Tool: {tool_name}]")
|
|
|
|
try:
|
|
# Find tool
|
|
tool_instance = None
|
|
for tool in await mcp.list_tools():
|
|
if getattr(tool, 'name', '') == tool_name:
|
|
tool_instance = tool
|
|
break
|
|
|
|
if not tool_instance:
|
|
error_response = {
|
|
"jsonrpc": "2.0",
|
|
"id": message_id,
|
|
"error": {
|
|
"code": -32601,
|
|
"message": f"Tool '{tool_name}' not found"
|
|
}
|
|
}
|
|
await self.client_sessions[session_id]["queue"].put(error_response)
|
|
return JSONResponse(
|
|
{"status": "error", "message": f"Tool '{tool_name}' not found"},
|
|
status_code=404,
|
|
headers={
|
|
"Access-Control-Allow-Origin": "*",
|
|
"Access-Control-Allow-Credentials": "true",
|
|
"Access-Control-Allow-Methods": "*",
|
|
"Access-Control-Allow-Headers": "*",
|
|
"Access-Control-Expose-Headers": "*"
|
|
}
|
|
)
|
|
|
|
# Execute tool call
|
|
# Fix: Do not call tool object directly, instead call custom call_tool method
|
|
result = await self.call_tool(tool_name, arguments, request)
|
|
|
|
# Special formatting for result
|
|
# If result is already in the correct format, use directly
|
|
if isinstance(result, dict) and "content" in result and isinstance(result["content"], list):
|
|
formatted_result = result
|
|
else:
|
|
# Otherwise, format into standard format
|
|
formatted_result = {
|
|
"content": [
|
|
{
|
|
"type": "json",
|
|
"text": result if isinstance(result, str) else json.dumps(result, ensure_ascii=False)
|
|
}
|
|
]
|
|
}
|
|
|
|
# Build response
|
|
response = {
|
|
"jsonrpc": "2.0",
|
|
"id": message_id,
|
|
"result": formatted_result
|
|
}
|
|
|
|
# Put response into queue
|
|
await self.client_sessions[session_id]["queue"].put(response)
|
|
|
|
# Return received confirmation
|
|
return JSONResponse(
|
|
{"status": "success"},
|
|
headers={
|
|
"Access-Control-Allow-Origin": "*",
|
|
"Access-Control-Allow-Credentials": "true",
|
|
"Access-Control-Allow-Methods": "*",
|
|
"Access-Control-Allow-Headers": "*",
|
|
"Access-Control-Expose-Headers": "*"
|
|
}
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Tool call error: {str(e)}", exc_info=True)
|
|
|
|
# Build error response
|
|
if str(e).startswith('{"code":'):
|
|
# If it's a JSON formatted error, use directly
|
|
try:
|
|
error_obj = json.loads(str(e))
|
|
error_response = {
|
|
"jsonrpc": "2.0",
|
|
"id": message_id,
|
|
"error": error_obj
|
|
}
|
|
except:
|
|
# If parsing fails, use standard format
|
|
error_response = {
|
|
"jsonrpc": "2.0",
|
|
"id": message_id,
|
|
"error": {
|
|
"code": -32000,
|
|
"message": str(e)
|
|
}
|
|
}
|
|
else:
|
|
# Normal error string
|
|
error_response = {
|
|
"jsonrpc": "2.0",
|
|
"id": message_id,
|
|
"error": {
|
|
"code": -32000,
|
|
"message": str(e)
|
|
}
|
|
}
|
|
|
|
# Put error response into queue
|
|
await self.client_sessions[session_id]["queue"].put(error_response)
|
|
|
|
# Return error status
|
|
return JSONResponse(
|
|
{"status": "error", "message": str(e)},
|
|
status_code=500,
|
|
headers={
|
|
"Access-Control-Allow-Origin": "*",
|
|
"Access-Control-Allow-Credentials": "true",
|
|
"Access-Control-Allow-Methods": "*",
|
|
"Access-Control-Allow-Headers": "*",
|
|
"Access-Control-Expose-Headers": "*"
|
|
}
|
|
)
|
|
else:
|
|
# Other message types, forward directly to MCP for handling
|
|
logger.info(f"Processing general message [Session ID: {session_id}]")
|
|
|
|
try:
|
|
# Process message
|
|
# FastMCP object doesn't have process_message method, build response directly instead
|
|
# result = await mcp.process_message(body)
|
|
|
|
# Build response
|
|
response = {
|
|
"jsonrpc": "2.0",
|
|
"id": message_id,
|
|
"result": {
|
|
"status": "ok",
|
|
"message": "Message received, but unable to process unrecognized message type"
|
|
}
|
|
}
|
|
except Exception as e:
|
|
logger.error(f"Error processing message [Session ID: {session_id}]: {str(e)}")
|
|
response = {
|
|
"jsonrpc": "2.0",
|
|
"id": message_id,
|
|
"error": {
|
|
"code": -32000,
|
|
"message": f"Error processing message: {str(e)}"
|
|
}
|
|
}
|
|
|
|
# Put response into queue
|
|
if response:
|
|
await self.client_sessions[session_id]["queue"].put(response)
|
|
|
|
# Return received confirmation to HTTP request
|
|
return JSONResponse(
|
|
{"status": "received"},
|
|
headers={
|
|
"Access-Control-Allow-Origin": "*",
|
|
"Access-Control-Allow-Credentials": "true",
|
|
"Access-Control-Allow-Methods": "*",
|
|
"Access-Control-Allow-Headers": "*",
|
|
"Access-Control-Expose-Headers": "*"
|
|
}
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Error processing message: {str(e)}")
|
|
# Send error response
|
|
error_response = {
|
|
"jsonrpc": "2.0",
|
|
"id": message_id if 'message_id' in locals() else str(uuid.uuid4()),
|
|
"error": {
|
|
"code": -32000,
|
|
"message": str(e)
|
|
}
|
|
}
|
|
|
|
await self.client_sessions[session_id]["queue"].put(error_response)
|
|
return JSONResponse(
|
|
{"status": "error", "message": str(e)},
|
|
status_code=500,
|
|
headers={
|
|
"Access-Control-Allow-Origin": "*",
|
|
"Access-Control-Allow-Credentials": "true",
|
|
"Access-Control-Allow-Methods": "*",
|
|
"Access-Control-Allow-Headers": "*",
|
|
"Access-Control-Expose-Headers": "*"
|
|
}
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Error processing request: {str(e)}", exc_info=True)
|
|
return JSONResponse(
|
|
{
|
|
"jsonrpc": "2.0",
|
|
"id": "unknown",
|
|
"error": {
|
|
"code": -32000,
|
|
"message": f"Error processing request: {str(e)}"
|
|
}
|
|
},
|
|
status_code=500,
|
|
headers={
|
|
"Access-Control-Allow-Origin": "*",
|
|
"Access-Control-Allow-Credentials": "true",
|
|
"Access-Control-Allow-Methods": "*",
|
|
"Access-Control-Allow-Headers": "*",
|
|
"Access-Control-Expose-Headers": "*"
|
|
}
|
|
)
|
|
|
|
async def _execute_stream_tool_wrapper(self, tool_name, kwargs, message_id, session_id, request):
|
|
"""Wrapper for streaming tool calls
|
|
|
|
Args:
|
|
tool_name: Tool name
|
|
kwargs: Function parameters
|
|
message_id: Message ID
|
|
session_id: Session ID
|
|
request: Request object
|
|
"""
|
|
try:
|
|
# Execute by calling the standard tool method
|
|
result = await self.call_tool(tool_name, kwargs, request)
|
|
|
|
# Send completion message
|
|
final_message = {
|
|
"jsonrpc": "2.0",
|
|
"id": message_id,
|
|
"success": True,
|
|
"result": result
|
|
}
|
|
|
|
# Check if session still exists
|
|
if session_id in self.client_sessions:
|
|
await self.client_sessions[session_id]["queue"].put(final_message)
|
|
logger.info(f"Streaming tool execution completed [Session ID: {session_id}, Message ID: {message_id}]")
|
|
else:
|
|
logger.warning(f"Streaming tool execution completed but session closed [Session ID: {session_id}]")
|
|
except Exception as e:
|
|
logger.error(f"Streaming tool execution failed [Session ID: {session_id}]: {str(e)}")
|
|
|
|
# Send error message
|
|
error_message = {
|
|
"jsonrpc": "2.0",
|
|
"id": message_id,
|
|
"success": False,
|
|
"error": str(e)
|
|
}
|
|
|
|
# Check if session still exists
|
|
if session_id in self.client_sessions:
|
|
await self.client_sessions[session_id]["queue"].put(error_message)
|
|
else:
|
|
logger.warning(f"Streaming tool execution failed but session closed [Session ID: {session_id}]")
|
|
|
|
async def broadcast_tool_result(self, tool_name, result_data):
|
|
"""Broadcast tool call result to all clients
|
|
|
|
Args:
|
|
tool_name: Tool name
|
|
result_data: Result data
|
|
"""
|
|
logger.info(f"Broadcasting tool result: {tool_name}")
|
|
message = {
|
|
"jsonrpc": "2.0",
|
|
"method": "notifications/tool_result",
|
|
"params": {
|
|
"type": "tool_result",
|
|
"tool": tool_name,
|
|
"result": result_data
|
|
}
|
|
}
|
|
await self.broadcast_message(message)
|
|
|
|
async def call_tool(self, tool_name, arguments, request):
|
|
"""
|
|
Calls the specified tool and returns the result
|
|
|
|
Args:
|
|
tool_name: Tool name
|
|
arguments: Tool parameters
|
|
request: Original request object
|
|
|
|
Returns:
|
|
Tool call result
|
|
"""
|
|
logger.info(f"Calling tool: {tool_name}, Parameters: {json.dumps(arguments, ensure_ascii=False)}")
|
|
|
|
# Get recent query content, used to handle random_string parameter
|
|
recent_query = self._extract_recent_query(request)
|
|
|
|
# Handle tool name mapping - Add support for standard name tools
|
|
tool_mapping = {
|
|
# --- Retained/New Tools ---
|
|
"exec_query": "mcp_doris_exec_query",
|
|
"get_table_schema": "mcp_doris_get_table_schema",
|
|
"get_db_table_list": "mcp_doris_get_db_table_list",
|
|
"get_db_list": "mcp_doris_get_db_list",
|
|
"get_table_comment": "mcp_doris_get_table_comment",
|
|
"get_table_column_comments": "mcp_doris_get_table_column_comments",
|
|
"get_table_indexes": "mcp_doris_get_table_indexes",
|
|
"get_recent_audit_logs": "mcp_doris_get_recent_audit_logs",
|
|
"get_catalog_list": "mcp_doris_get_catalog_list"
|
|
}
|
|
|
|
# If it's a standard name, convert to MCP name
|
|
mapped_tool_name = tool_mapping.get(tool_name, tool_name)
|
|
|
|
# Import tool functions from mcp_doris_tools
|
|
try:
|
|
# Import tool module
|
|
import doris_mcp_server.tools.mcp_doris_tools as mcp_tools
|
|
|
|
# Get the corresponding tool function
|
|
tool_function = getattr(mcp_tools, mapped_tool_name, None)
|
|
|
|
if not tool_function:
|
|
# If it doesn't exist in mcp_tools, try getting the tool using the MCP server instance
|
|
mcp = self.mcp_server
|
|
# Find the corresponding tool
|
|
for tool in await mcp.list_tools():
|
|
if getattr(tool, 'name', '') == mapped_tool_name:
|
|
tool_function = tool
|
|
break
|
|
|
|
if not tool_function:
|
|
raise ValueError(f"Tool not found: {tool_name} / {mapped_tool_name}")
|
|
|
|
# Process common input parameter conversions
|
|
processed_args = self._process_tool_arguments(mapped_tool_name, arguments, recent_query)
|
|
|
|
# Call the tool function
|
|
try:
|
|
# Log tool type and attribute information to help debugging
|
|
logger.debug(f"Tool function type: {type(tool_function)}")
|
|
logger.debug(f"Tool function attributes: {dir(tool_function)}")
|
|
|
|
if callable(tool_function):
|
|
logger.debug("Tool function is callable, calling directly")
|
|
result = await tool_function(**processed_args)
|
|
elif hasattr(tool_function, 'run'):
|
|
logger.debug("Tool function has run method, calling run method")
|
|
result = await tool_function.run(**processed_args)
|
|
elif hasattr(tool_function, 'execute'):
|
|
logger.debug("Tool function has execute method, calling execute method")
|
|
result = await tool_function.execute(**processed_args)
|
|
elif hasattr(tool_function, 'call'):
|
|
logger.debug("Tool function has call method, calling call method")
|
|
result = await tool_function.call(**processed_args)
|
|
elif hasattr(tool_function, '__call__'):
|
|
logger.debug("Tool function has __call__ method, calling __call__ method")
|
|
result = await tool_function.__call__(**processed_args)
|
|
else:
|
|
# If it's a dict type, try getting the function from it
|
|
if isinstance(tool_function, dict) and 'function' in tool_function:
|
|
logger.debug("Tool is a dictionary type, trying to get 'function' key")
|
|
actual_func = tool_function['function']
|
|
if callable(actual_func):
|
|
result = await actual_func(**processed_args)
|
|
else:
|
|
raise ValueError(f"Function in dictionary is not callable: {type(actual_func)}")
|
|
else:
|
|
raise ValueError(f"Unsupported tool type: {type(tool_function)}, Attributes: {dir(tool_function)}")
|
|
except Exception as e:
|
|
logger.error(f"Failed to call tool function: {str(e)}", exc_info=True)
|
|
raise ValueError(f"Error calling tool: {str(e)}")
|
|
|
|
# Return tool execution result
|
|
return result
|
|
except AttributeError as e:
|
|
logger.error(f"Tool function does not exist: {mapped_tool_name}, Error: {str(e)}")
|
|
raise ValueError(f"Tool function does not exist: {mapped_tool_name}")
|
|
except Exception as e:
|
|
logger.error(f"Error calling tool: {str(e)}", exc_info=True)
|
|
raise ValueError(f"Error calling tool: {str(e)}")
|
|
|
|
def _process_tool_arguments(self, tool_name, arguments, recent_query):
|
|
"""
|
|
Process tool parameters, supporting special handling logic
|
|
|
|
Args:
|
|
tool_name: Tool name (MCP internal name, e.g., mcp_doris_...)
|
|
arguments: Original parameters
|
|
recent_query: Recent query content
|
|
|
|
Returns:
|
|
Processed parameter dictionary
|
|
"""
|
|
# Copy parameters to avoid modifying the original object
|
|
processed_args = dict(arguments)
|
|
|
|
# Handle potential random_string parameter as fallback
|
|
if "random_string" in processed_args and tool_name.startswith("mcp_doris_"):
|
|
random_string = processed_args.pop("random_string", "")
|
|
logger.debug(f"Processing random_string parameter for tool {tool_name}: '{random_string}'")
|
|
|
|
# 1. For exec_query
|
|
if tool_name == "mcp_doris_exec_query":
|
|
if not processed_args.get("sql"):
|
|
sql_fallback = random_string or recent_query
|
|
if sql_fallback:
|
|
if not random_string and recent_query:
|
|
import re
|
|
sql_matches = re.findall(r'```sql\s*([\s\S]+?)\s*```', recent_query)
|
|
if sql_matches:
|
|
sql_fallback = sql_matches[0].strip()
|
|
|
|
if sql_fallback:
|
|
logger.info(f"Using random_string/recent_query as SQL for exec_query: {sql_fallback[:100]}...")
|
|
processed_args["sql"] = sql_fallback
|
|
else:
|
|
logger.warning(f"exec_query missing sql parameter, and random_string/recent_query is empty or SQL cannot be extracted")
|
|
else:
|
|
logger.warning(f"exec_query missing sql parameter, and both random_string and recent_query are empty")
|
|
|
|
# 2. For tools requiring table_name
|
|
elif tool_name in [
|
|
"mcp_doris_get_table_schema",
|
|
"mcp_doris_get_table_comment",
|
|
"mcp_doris_get_table_column_comments",
|
|
"mcp_doris_get_table_indexes"
|
|
]:
|
|
if not processed_args.get("table_name"):
|
|
table_fallback = random_string
|
|
if table_fallback:
|
|
logger.info(f"Using random_string/recent_query as table_name for {tool_name}: {table_fallback}")
|
|
processed_args["table_name"] = table_fallback
|
|
else:
|
|
logger.warning(f"{tool_name} missing table_name parameter, and random_string/recent_query is empty or table name cannot be extracted")
|
|
|
|
# 3. Other tools
|
|
else:
|
|
logger.debug(f"Tool {tool_name} does not apply random_string fallback logic, or logic is undefined.")
|
|
|
|
# Ensure return is outside the main if block
|
|
return processed_args
|
|
|
|
def _extract_recent_query(self, request: Request) -> Optional[str]:
|
|
"""
|
|
Extract the most recent user query from the request
|
|
|
|
Args:
|
|
request: Request object
|
|
|
|
Returns:
|
|
Optional[str]: The most recent user query, or None if not found
|
|
"""
|
|
try:
|
|
# Try to extract message history from request body
|
|
body = None
|
|
body_bytes = getattr(request, "_body", None)
|
|
if body_bytes:
|
|
try:
|
|
body = json.loads(body_bytes)
|
|
except:
|
|
pass
|
|
|
|
if not body:
|
|
body = getattr(request, "_json", {})
|
|
|
|
# Find the most recent user message from message history
|
|
messages = body.get("params", {}).get("messages", [])
|
|
if messages:
|
|
# Iterate messages in reverse to find the most recent user message
|
|
for msg in reversed(messages):
|
|
if msg.get("role") == "user":
|
|
return msg.get("content", "")
|
|
|
|
# If not found in message history, try extracting from the original message
|
|
message = body.get("params", {}).get("message", {})
|
|
if message and message.get("role") == "user":
|
|
return message.get("content", "")
|
|
|
|
return None
|
|
except Exception as e:
|
|
logger.error(f"Error extracting recent query: {str(e)}")
|
|
return None
|
|
|
|
def format_tool_result(self, result):
|
|
"""
|
|
Format tool call result into a unified format
|
|
|
|
Args:
|
|
result: Original tool call result
|
|
|
|
Returns:
|
|
Formatted result
|
|
"""
|
|
try:
|
|
# If result is already a dictionary, return directly
|
|
if isinstance(result, dict):
|
|
return result
|
|
|
|
# If it's a string, try parsing as JSON
|
|
if isinstance(result, str):
|
|
try:
|
|
return json.loads(result)
|
|
except json.JSONDecodeError:
|
|
# Pure text result
|
|
return {"content": result}
|
|
|
|
# If it's another type, convert to string
|
|
return {"content": str(result)}
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error formatting tool result: {str(e)}")
|
|
return {"error": str(e)}
|
|
|
|
def _get_session_id(self, request: Request) -> str:
|
|
"""
|
|
Get session ID from the request
|
|
|
|
Tries to get session ID from the following locations (in priority order):
|
|
1. Query parameter session_id
|
|
2. session_id field in request body
|
|
3. X-Session-ID header
|
|
|
|
Args:
|
|
request: Request object
|
|
|
|
Returns:
|
|
str: Session ID, or None if not found
|
|
"""
|
|
# Get from query parameter
|
|
session_id = request.query_params.get("session_id")
|
|
if session_id:
|
|
return session_id
|
|
|
|
# Try getting from request body
|
|
try:
|
|
body = getattr(request, "_json", None)
|
|
if not body:
|
|
body_bytes = getattr(request, "_body", None)
|
|
if body_bytes:
|
|
try:
|
|
body = json.loads(body_bytes)
|
|
except:
|
|
pass
|
|
|
|
if body and isinstance(body, dict) and "session_id" in body:
|
|
return body["session_id"]
|
|
except:
|
|
pass
|
|
|
|
# Get from request header
|
|
session_id = request.headers.get("X-Session-ID")
|
|
if session_id:
|
|
return session_id
|
|
|
|
return None
|