t6_mem0_v2/mcp_server/http_server.py

"""
T6 Mem0 v2 MCP Server - HTTP/SSE Transport
Exposes MCP server via HTTP for n8n MCP Client Tool
"""

import logging
import asyncio
from typing import AsyncIterator
from fastapi import FastAPI, Request, Response
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from mcp.server import Server
from mcp.types import (
    JSONRPCRequest,
    JSONRPCResponse,
    JSONRPCError,
    Tool,
    TextContent,
    ImageContent,
    EmbeddedResource
)
from mem0 import Memory
import json

from config import mem0_config, settings
from mcp_server.tools import MemoryTools

# Configure logging
logging.basicConfig(
    level=getattr(logging, settings.log_level.upper()),
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)

logger = logging.getLogger(__name__)

# Initialize FastAPI
app = FastAPI(
    title="T6 Mem0 v2 MCP Server",
    description="Model Context Protocol server for memory operations via HTTP/SSE",
    version="2.0.0"
)

# Enable CORS for n8n
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


class MCPHTTPServer:
    """MCP Server with HTTP/SSE transport"""

    def __init__(self):
        self.server = Server("t6-mem0-v2")
        self.memory: Memory | None = None
        self.tools: MemoryTools | None = None
        self.setup_handlers()

    def setup_handlers(self):
        """Setup MCP server handlers"""

        @self.server.list_resources()
        async def list_resources():
            return []

        @self.server.read_resource()
        async def read_resource(uri: str) -> str:
            logger.warning(f"Resource read not implemented: {uri}")
            return ""

        @self.server.list_tools()
        async def list_tools() -> list[Tool]:
            logger.info("Listing tools")
            if not self.tools:
                raise RuntimeError("Tools not initialized")
            return self.tools.get_tool_definitions()

        @self.server.call_tool()
        async def call_tool(name: str, arguments: dict) -> list[TextContent | ImageContent | EmbeddedResource]:
            logger.info(f"Tool called: {name}")
            logger.debug(f"Arguments: {arguments}")

            if not self.tools:
                raise RuntimeError("Tools not initialized")

            handlers = {
                "add_memory": self.tools.handle_add_memory,
                "search_memories": self.tools.handle_search_memories,
                "get_memory": self.tools.handle_get_memory,
                "get_all_memories": self.tools.handle_get_all_memories,
                "update_memory": self.tools.handle_update_memory,
                "delete_memory": self.tools.handle_delete_memory,
                "delete_all_memories": self.tools.handle_delete_all_memories,
            }

            handler = handlers.get(name)
            if not handler:
                logger.error(f"Unknown tool: {name}")
                return [TextContent(type="text", text=f"Error: Unknown tool '{name}'")]

            try:
                return await handler(arguments)
            except Exception as e:
                logger.error(f"Tool execution failed: {e}", exc_info=True)
                return [TextContent(type="text", text=f"Error executing tool: {str(e)}")]

    async def initialize(self):
        """Initialize memory service"""
        logger.info("Initializing T6 Mem0 v2 MCP HTTP Server")
        logger.info(f"Environment: {settings.environment}")

        try:
            logger.info("Initializing Mem0...")
            self.memory = Memory.from_config(config_dict=mem0_config)
            logger.info("Mem0 initialized successfully")

            self.tools = MemoryTools(self.memory)
            logger.info("Tools initialized successfully")

            logger.info("T6 Mem0 v2 MCP HTTP Server ready")

        except Exception as e:
            logger.error(f"Failed to initialize server: {e}", exc_info=True)
            raise


# Global server instance
mcp_server = MCPHTTPServer()


@app.on_event("startup")
async def startup():
    """Initialize MCP server on startup"""
    await mcp_server.initialize()


@app.get("/health")
async def health():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "service": "t6-mem0-v2-mcp-http",
        "transport": "http-streamable"
    }


@app.post("/mcp")
async def mcp_endpoint(request: Request):
    """
    MCP HTTP Streamable endpoint

    This endpoint handles MCP JSON-RPC requests and returns responses
    compatible with n8n's MCP Client Tool.
    """
    try:
        # Parse JSON-RPC request
        body = await request.json()
        logger.info(f"Received MCP request: {body.get('method', 'unknown')}")

        # Handle different MCP methods
        method = body.get("method")
        params = body.get("params", {})
        request_id = body.get("id")

        if method == "tools/list":
            # List available tools
            tools = mcp_server.tools.get_tool_definitions()
            response = {
                "jsonrpc": "2.0",
                "id": request_id,
                "result": {
                    "tools": [
                        {
                            "name": tool.name,
                            "description": tool.description,
                            "inputSchema": tool.inputSchema
                        }
                        for tool in tools
                    ]
                }
            }

        elif method == "tools/call":
            # Call a tool
            tool_name = params.get("name")
            arguments = params.get("arguments", {})

            # Route to appropriate tool handler
            handlers = {
                "add_memory": mcp_server.tools.handle_add_memory,
                "search_memories": mcp_server.tools.handle_search_memories,
                "get_memory": mcp_server.tools.handle_get_memory,
                "get_all_memories": mcp_server.tools.handle_get_all_memories,
                "update_memory": mcp_server.tools.handle_update_memory,
                "delete_memory": mcp_server.tools.handle_delete_memory,
                "delete_all_memories": mcp_server.tools.handle_delete_all_memories,
            }

            handler = handlers.get(tool_name)
            if not handler:
                response = {
                    "jsonrpc": "2.0",
                    "id": request_id,
                    "error": {
                        "code": -32602,
                        "message": f"Unknown tool: {tool_name}"
                    }
                }
            else:
                result = await handler(arguments)

                # Convert TextContent to dict
                content = []
                for item in result:
                    if isinstance(item, TextContent):
                        content.append({
                            "type": "text",
                            "text": item.text
                        })

                response = {
                    "jsonrpc": "2.0",
                    "id": request_id,
                    "result": {
                        "content": content
                    }
                }

        elif method == "initialize":
            # Handle initialization
            response = {
                "jsonrpc": "2.0",
                "id": request_id,
                "result": {
                    "protocolVersion": "2024-11-05",
                    "capabilities": {
                        "tools": {}
                    },
                    "serverInfo": {
                        "name": "t6-mem0-v2",
                        "version": "2.0.0"
                    }
                }
            }

        else:
            # Unknown method
            response = {
                "jsonrpc": "2.0",
                "id": request_id,
                "error": {
                    "code": -32601,
                    "message": f"Method not found: {method}"
                }
            }

        logger.info(f"Sending response for {method}")
        return response

    except Exception as e:
        logger.error(f"Error processing MCP request: {e}", exc_info=True)
        return {
            "jsonrpc": "2.0",
            "id": body.get("id") if "body" in locals() else None,
            "error": {
                "code": -32603,
                "message": f"Internal error: {str(e)}"
            }
        }


if __name__ == "__main__":
    import uvicorn

    port = settings.mcp_port

    logger.info(f"Starting MCP HTTP Server on port {port}")
    uvicorn.run(
        "mcp_server.http_server:app",
        host="0.0.0.0",
        port=port,
        log_level=settings.log_level.lower()
    )