app/routes/mcp.py

"""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())