Update start_mcp.py

start_mcp.py

@@ -9,7 +9,7 @@ import os
 import sys
 import uuid
 
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 
@@ -20,7 +20,6 @@ logger = logging.getLogger("m3gnet-mcp-server")
 def _load_service():
     try:
         from mcp_service import mcp  # noqa: WPS433 (runtime import)
-
         return mcp
     except Exception as exc:  # noqa: BLE001
         logger.error("Failed to import MCP service: %s", exc, exc_info=True)
@@ -29,7 +28,6 @@ def _load_service():
 
 def run_stdio():
     """Run the FastMCP service over stdio."""
-
     service = _load_service()
     logger.info("Starting stdio server mode")
     service.run()
@@ -37,19 +35,22 @@ def run_stdio():
 
 def run_sse(host: str, port: int):
     """Run the FastMCP service with a lightweight SSE transport."""
-
     service = _load_service()
     logger.info("Starting SSE server on %s:%s", host, port)
 
-    import uvicorn  # Imported lazily to keep stdio mode lightweight
+    import uvicorn
     from starlette.applications import Starlette
+    from starlette.middleware import Middleware
+    from starlette.middleware.cors import CORSMiddleware
     from starlette.requests import Request
-    from starlette.responses import HTMLResponse, JSONResponse, StreamingResponse
+    from starlette.responses import HTMLResponse, JSONResponse, StreamingResponse, Response
     from starlette.routing import Route
 
+    # Store active SSE connections
     connections: Dict[str, asyncio.Queue] = {}
 
-    async def sse_endpoint(request: Request):  # noqa: WPS430
+    async def sse_endpoint(request: Request):
+        """SSE endpoint - establishes connection and sends events."""
         connection_id = str(uuid.uuid4())
         queue: asyncio.Queue = asyncio.Queue()
         connections[connection_id] = queue
@@ -57,13 +58,17 @@ def run_sse(host: str, port: int):
 
         async def event_stream():
             try:
+                # Send initial endpoint event telling client where to POST messages
                 yield f"event: endpoint\ndata: /messages?sessionId={connection_id}\n\n"
                 while True:
                     try:
-                        message = await asyncio.wait_for(queue.get(), timeout=15)
+                        message = await asyncio.wait_for(queue.get(), timeout=30)
                         yield f"event: message\ndata: {json.dumps(message)}\n\n"
                     except asyncio.TimeoutError:
+                        # Send keep-alive comment to prevent connection timeout
                         yield ": keep-alive\n\n"
+            except asyncio.CancelledError:
+                logger.info("SSE stream cancelled: %s", connection_id)
             finally:
                 connections.pop(connection_id, None)
                 logger.info("SSE connection closed: %s", connection_id)
@@ -73,47 +78,79 @@ def run_sse(host: str, port: int):
             "Connection": "keep-alive",
             "X-Accel-Buffering": "no",
             "Access-Control-Allow-Origin": "*",
+            "Access-Control-Allow-Methods": "GET, POST, OPTIONS",
+            "Access-Control-Allow-Headers": "*",
         }
         return StreamingResponse(event_stream(), media_type="text/event-stream", headers=headers)
 
-    async def messages_endpoint(request: Request):  # noqa: WPS430
+    async def messages_endpoint(request: Request):
+        """Messages endpoint - receives MCP protocol messages from client."""
+        # Handle CORS preflight
+        if request.method == "OPTIONS":
+            return Response(
+                status_code=200,
+                headers={
+                    "Access-Control-Allow-Origin": "*",
+                    "Access-Control-Allow-Methods": "GET, POST, OPTIONS",
+                    "Access-Control-Allow-Headers": "*",
+                }
+            )
+
         session_id = request.query_params.get("sessionId")
-        if session_id not in connections:
-            return JSONResponse({"error": "invalid session"}, status_code=400)
+        if not session_id or session_id not in connections:
+            logger.warning("Invalid session: %s", session_id)
+            return JSONResponse(
+                {"jsonrpc": "2.0", "error": {"code": -32600, "message": "Invalid session"}},
+                status_code=400,
+                headers={"Access-Control-Allow-Origin": "*"}
+            )
+
+        try:
+            body = await request.json()
+        except Exception as e:
+            logger.error("Failed to parse JSON body: %s", e)
+            return JSONResponse(
+                {"jsonrpc": "2.0", "error": {"code": -32700, "message": "Parse error"}},
+                status_code=400,
+                headers={"Access-Control-Allow-Origin": "*"}
+            )
 
-        body = await request.json()
         method = body.get("method")
         params = body.get("params", {})
-        message_id = body.get("id")
+        message_id = body.get("id")  # Can be None for notifications
 
         logger.info("Received method=%s id=%s", method, message_id)
 
-        # Notifications (methods starting with "notifications/") do not expect a response
-        is_notification = method and method.startswith("notifications/")
+        # Check if this is a notification (no id field means notification)
+        is_notification = message_id is None
 
         async def handle_request():
+            result: Optional[Dict[str, Any]] = None
+            error: Optional[Dict[str, Any]] = None
+
             try:
                 if method == "initialize":
-                    result: Dict[str, Any] = {
+                    result = {
                         "protocolVersion": "2024-11-05",
                         "serverInfo": {"name": "M3GNet-MCP", "version": "1.0.0"},
                         "capabilities": {"tools": {}},
                     }
                 elif method == "tools/list":
                     tools = []
-                    for tool in service._tool_manager._tools.values():  # noqa: SLF001 (internal attribute)
-                        tools.append(
-                            {
-                                "name": tool.name,
-                                "description": tool.description or tool.name,
-                                "inputSchema": getattr(tool, "parameters", {"type": "object", "properties": {}}),
-                            }
-                        )
+                    for tool in service._tool_manager._tools.values():
+                        tool_schema = getattr(tool, "parameters", None)
+                        if tool_schema is None:
+                            tool_schema = {"type": "object", "properties": {}}
+                        tools.append({
+                            "name": tool.name,
+                            "description": tool.description or tool.name,
+                            "inputSchema": tool_schema,
+                        })
                     result = {"tools": tools}
                 elif method == "tools/call":
                     tool_name = params.get("name")
                     arguments = params.get("arguments", {})
-                    tool = service._tool_manager._tools.get(tool_name)  # noqa: SLF001
+                    tool = service._tool_manager._tools.get(tool_name)
                     if tool is None:
                         raise ValueError(f"Unknown tool: {tool_name}")
                     logger.info("Invoking tool '%s'", tool_name)
@@ -121,48 +158,84 @@ def run_sse(host: str, port: int):
                         tool_response = await tool.fn(**arguments)
                     else:
                         tool_response = tool.fn(**arguments)
-                    result = {"content": [{"type": "text", "text": json.dumps(tool_response)}]}
-                elif is_notification:
+                    result = {"content": [{"type": "text", "text": json.dumps(tool_response, default=str)}]}
+                elif method and method.startswith("notifications/"):
                     # Notifications do not require a response per JSON-RPC spec
                     logger.info("Handled notification: %s", method)
-                    return
+                    return  # Exit without sending response
                 else:
-                    raise ValueError(f"Unsupported method: {method}")
-
-                # Only send response if message_id is present (i.e., it's a request, not a notification)
-                if message_id is not None:
-                    response = {"jsonrpc": "2.0", "id": message_id, "result": result}
-                    await connections[session_id].put(response)
+                    error = {"code": -32601, "message": f"Method not found: {method}"}
 
-            except Exception as exc:  # noqa: BLE001
+            except Exception as exc:
                 logger.exception("Error handling request")
-                # Only send error response if we have a valid message_id
-                if message_id is not None:
-                    response = {
-                        "jsonrpc": "2.0",
-                        "id": message_id,
-                        "error": {"code": -32000, "message": str(exc)},
-                    }
-                    await connections[session_id].put(response)
+                error = {"code": -32000, "message": str(exc)}
+
+            # Only send response for requests (not notifications)
+            if not is_notification and session_id in connections:
+                if error:
+                    response = {"jsonrpc": "2.0", "id": message_id, "error": error}
+                else:
+                    response = {"jsonrpc": "2.0", "id": message_id, "result": result}
+                await connections[session_id].put(response)
 
+        # Run handler in background
         asyncio.create_task(handle_request())
-        return JSONResponse({"ok": True})
+
+        # Return accepted response
+        return JSONResponse(
+            {"ok": True},
+            headers={"Access-Control-Allow-Origin": "*"}
+        )
 
     async def health(_: Request):
-        return JSONResponse({"status": "ok", "service": "M3GNet-MCP"})
+        """Health check endpoint."""
+        return JSONResponse(
+            {"status": "ok", "service": "M3GNet-MCP"},
+            headers={"Access-Control-Allow-Origin": "*"}
+        )
 
     async def home(_: Request):
+        """Homepage with service info."""
         return HTMLResponse("""<!DOCTYPE html>
-<html><head><title>M3GNet MCP</title></head>
-<body><h1>M3GNet MCP Service</h1><p>SSE endpoint: <code>/sse</code></p></body></html>""")
+<html>
+<head><title>M3GNet MCP</title></head>
+<body>
+<h1>M3GNet MCP Service</h1>
+<p>Status: <span style="color:green">Running</span></p>
+<p>SSE endpoint: <code>/sse</code></p>
+<p>Messages endpoint: <code>/messages</code></p>
+<h2>Cursor Configuration</h2>
+<pre>
+{
+  "mcpServers": {
+    "m3gnet": {
+      "url": "http://localhost:7860/sse"
+    }
+  }
+}
+</pre>
+</body>
+</html>""")
+
+    # Create app with CORS middleware
+    middleware = [
+        Middleware(
+            CORSMiddleware,
+            allow_origins=["*"],
+            allow_methods=["*"],
+            allow_headers=["*"],
+            allow_credentials=True,
+        )
+    ]
 
     app = Starlette(
         routes=[
-            Route("/", home),
-            Route("/health", health),
-            Route("/sse", sse_endpoint),
-            Route("/messages", messages_endpoint, methods=["POST"]),
-        ]
+            Route("/", home, methods=["GET"]),
+            Route("/health", health, methods=["GET"]),
+            Route("/sse", sse_endpoint, methods=["GET"]),
+            Route("/messages", messages_endpoint, methods=["GET", "POST", "OPTIONS"]),
+        ],
+        middleware=middleware,
     )
 
     uvicorn.run(app, host=host, port=port, log_level="info")