"""MCP protocol endpoints - Full MCP 2025-03-26 Streamable HTTP + legacy SSE.""" from __future__ import annotations import json import logging import uuid from typing import Any from fastapi import APIRouter, Request from fastapi.responses import JSONResponse from sse_starlette.sse import EventSourceResponse from app.mcp_handler import TOOL_DEFINITIONS, handle_tool_call from app.models import MCPListToolsResponse, MCPToolInput, MCPToolResult logger = logging.getLogger(__name__) router = APIRouter(prefix="/mcp", tags=["mcp"]) # ─── MCP Server Info ────────────────────────────────────────── MCP_SERVER_INFO = { "name": "mcp-code-executor", "version": "1.0.0", } MCP_CAPABILITIES = { "tools": { "listChanged": False, }, } SUPPORTED_PROTOCOL_VERSIONS = [ "2025-03-26", "2024-11-05", ] # ─── Helper: Build JSON-RPC Response ────────────────────────── def jsonrpc_response(result: Any, request_id: Any) -> dict: return { "jsonrpc": "2.0", "result": result, "id": request_id, } def jsonrpc_error(code: int, message: str, request_id: Any, data: Any = None) -> dict: err = {"code": code, "message": message} if data is not None: err["data"] = data return { "jsonrpc": "2.0", "error": err, "id": request_id, } # ─── Streamable HTTP Endpoint (MCP 2025-03-26) ──────────────── # Groq and other MCP clients POST JSON-RPC messages here. # This single endpoint handles initialize, tools/list, tools/call, # ping, and notifications. @router.post("/jsonrpc") async def mcp_jsonrpc_endpoint(request: Request) -> JSONResponse: """ Full MCP Streamable HTTP transport endpoint. Handles the complete MCP JSON-RPC protocol including: - initialize / initialized - tools/list - tools/call - ping - notifications (no response needed) """ try: body = await request.json() except Exception: return JSONResponse( status_code=400, content=jsonrpc_error(-32700, "Parse error", None), ) # Handle batch requests if isinstance(body, list): responses = [] for item in body: resp = await _handle_single_jsonrpc(item) if resp is not None: # Notifications don't get responses responses.append(resp) if not responses: return JSONResponse(status_code=204, content=None) return JSONResponse(content=responses) # Single request result = await _handle_single_jsonrpc(body) if result is None: # Notification - no response return JSONResponse(status_code=204, content=None) return JSONResponse(content=result) @router.get("/jsonrpc") async def mcp_jsonrpc_sse_get(request: Request) -> EventSourceResponse: """ SSE GET endpoint for MCP Streamable HTTP transport. Some clients open a GET to receive server-initiated messages. We keep it open as a heartbeat/keepalive. """ async def event_generator(): # Send an initial endpoint event (for SSE transport compatibility) yield { "event": "endpoint", "data": "/mcp/jsonrpc", } # Keep connection alive import asyncio try: while True: await asyncio.sleep(30) yield { "event": "ping", "data": "{}", } except asyncio.CancelledError: pass return EventSourceResponse(event_generator()) @router.delete("/jsonrpc") async def mcp_jsonrpc_delete(): """Handle session termination.""" return JSONResponse(status_code=200, content={"status": "session_terminated"}) async def _handle_single_jsonrpc(body: dict) -> dict | None: """Process a single JSON-RPC message.""" if not isinstance(body, dict): return jsonrpc_error(-32600, "Invalid Request", None) method = body.get("method", "") request_id = body.get("id") # None for notifications params = body.get("params", {}) logger.info("MCP JSON-RPC method: %s (id=%s)", method, request_id) # ── Notifications (no id = no response) ── if request_id is None: if method == "notifications/initialized": logger.info("Client sent initialized notification") elif method == "notifications/cancelled": logger.info("Client cancelled request: %s", params) else: logger.info("Received notification: %s", method) return None # ── Requests (have id = need response) ── if method == "initialize": return _handle_initialize(params, request_id) elif method == "ping": return jsonrpc_response({}, request_id) elif method == "tools/list": return _handle_tools_list(params, request_id) elif method == "tools/call": return await _handle_tools_call(params, request_id) elif method == "resources/list": return jsonrpc_response({"resources": []}, request_id) elif method == "resources/templates/list": return jsonrpc_response({"resourceTemplates": []}, request_id) elif method == "prompts/list": return jsonrpc_response({"prompts": []}, request_id) elif method == "completion/complete": return jsonrpc_response({"completion": {"values": []}}, request_id) else: return jsonrpc_error(-32601, f"Method not found: {method}", request_id) def _handle_initialize(params: dict, request_id: Any) -> dict: """Handle MCP initialize handshake.""" client_protocol = params.get("protocolVersion", "2024-11-05") client_info = params.get("clientInfo", {}) logger.info( "MCP Initialize: client=%s, protocol=%s", client_info.get("name", "unknown"), client_protocol, ) # Negotiate protocol version if client_protocol in SUPPORTED_PROTOCOL_VERSIONS: negotiated_version = client_protocol else: # Fall back to our latest supported negotiated_version = SUPPORTED_PROTOCOL_VERSIONS[0] return jsonrpc_response( { "protocolVersion": negotiated_version, "capabilities": MCP_CAPABILITIES, "serverInfo": MCP_SERVER_INFO, }, request_id, ) def _handle_tools_list(params: dict, request_id: Any) -> dict: """Handle tools/list request.""" # MCP supports pagination via cursor, but we return all tools at once return jsonrpc_response( {"tools": TOOL_DEFINITIONS}, request_id, ) async def _handle_tools_call(params: dict, request_id: Any) -> dict: """Handle tools/call request.""" tool_name = params.get("name", "") tool_args = params.get("arguments", {}) if not tool_name: return jsonrpc_error( -32602, "Invalid params: 'name' is required", request_id ) logger.info("MCP tools/call: %s", tool_name) result = await handle_tool_call(tool_name, tool_args) return jsonrpc_response( { "content": result.content, "isError": result.isError, }, request_id, ) # ─── Legacy REST Endpoints (non-MCP clients) ────────────────── @router.get("/tools") async def list_tools() -> MCPListToolsResponse: """List all available MCP tools (simple REST).""" return MCPListToolsResponse(tools=TOOL_DEFINITIONS) @router.post("/execute", response_model=MCPToolResult) async def execute_tool(tool_input: MCPToolInput) -> MCPToolResult: """Execute an MCP tool call (simple REST).""" logger.info("MCP tool call: %s", tool_input.name) result = await handle_tool_call(tool_input.name, tool_input.arguments) return result @router.post("/execute/stream") async def execute_tool_stream(tool_input: MCPToolInput) -> EventSourceResponse: """Execute an MCP tool call with SSE streaming response.""" async def event_generator(): yield { "event": "start", "data": json.dumps({"tool": tool_input.name, "status": "running"}), } result = await handle_tool_call(tool_input.name, tool_input.arguments) yield { "event": "result", "data": json.dumps(result.model_dump()), } yield { "event": "done", "data": json.dumps({"status": "completed"}), } return EventSourceResponse(event_generator())