Update app/routes/mcp.py

app/routes/mcp.py

@@ -1,9 +1,10 @@
-"""MCP protocol endpoints - SSE and JSON-RPC style."""
+"""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
@@ -17,16 +18,244 @@ 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."""
+    """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."""
+    """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
@@ -54,55 +283,4 @@ async def execute_tool_stream(tool_input: MCPToolInput) -> EventSourceResponse:
             "data": json.dumps({"status": "completed"}),
         }
 
-    return EventSourceResponse(event_generator())
-
-
-# ─── JSON-RPC 2.0 Compatible Endpoint ─────────────────────────────
-
-@router.post("/jsonrpc")
-async def jsonrpc_handler(request: Request) -> JSONResponse:
-    """Handle JSON-RPC 2.0 style MCP requests."""
-    try:
-        body = await request.json()
-    except Exception:
-        return JSONResponse(
-            status_code=400,
-            content={
-                "jsonrpc": "2.0",
-                "error": {"code": -32700, "message": "Parse error"},
-                "id": None,
-            },
-        )
-
-    request_id = body.get("id")
-    method = body.get("method", "")
-    params = body.get("params", {})
-
-    if method == "tools/list":
-        return JSONResponse(
-            content={
-                "jsonrpc": "2.0",
-                "result": {"tools": TOOL_DEFINITIONS},
-                "id": request_id,
-            }
-        )
-    elif method == "tools/call":
-        tool_name = params.get("name", "")
-        tool_args = params.get("arguments", {})
-        result = await handle_tool_call(tool_name, tool_args)
-        return JSONResponse(
-            content={
-                "jsonrpc": "2.0",
-                "result": result.model_dump(),
-                "id": request_id,
-            }
-        )
-    else:
-        return JSONResponse(
-            status_code=400,
-            content={
-                "jsonrpc": "2.0",
-                "error": {"code": -32601, "message": f"Method not found: {method}"},
-                "id": request_id,
-            },
-        )
+    return EventSourceResponse(event_generator())