heh

app.py

@@ -7,14 +7,13 @@ Hosted on Hugging Face Spaces with HTTP transport
 import json
 import asyncio
 import logging
+import uuid
 from typing import Any, Dict, List, Optional
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 import uvicorn
-import asyncio
-import json
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -36,6 +35,9 @@ app.add_middleware(
 chunks_data = None
 docs_data = None
 
+# Session management for MCP
+sessions = {}
+
 def load_data():
     """Load the documentation chunks and metadata"""
     global chunks_data, docs_data
@@ -321,109 +323,73 @@ async def list_docs():
     
     return {"documents": docs_data}
 
-@app.get("/mcp/stream")
-async def mcp_sse():
-    """MCP SSE endpoint for mcp-remote"""
-    async def event_generator():
-        # Send MCP initialization response (correct format from docs)
-        init_response = {
-            "jsonrpc": "2.0",
-            "id": 1,
-            "result": {
-                "protocolVersion": "2025-06-18",
-                "capabilities": {
-                    "tools": {
-                        "listChanged": True
-                    },
-                    "resources": {}
-                },
-                "serverInfo": {
-                    "name": "mcp-docs-server",
-                    "version": "1.0.0"
-                }
-            }
+def is_initialize_request(body: dict) -> bool:
+    """Check if request is an MCP initialize request"""
+    return (body.get("jsonrpc") == "2.0" and 
+            body.get("method") == "initialize" and 
+            "id" in body)
+
+def create_mcp_server():
+    """Create MCP server instance"""
+    return {
+        "name": "mcp-docs-server",
+        "version": "1.0.0",
+        "capabilities": {
+            "tools": {"listChanged": True},
+            "resources": {}
         }
-        yield f"data: {json.dumps(init_response)}\n\n"
+    }
+
+@app.post("/mcp")
+async def mcp_post_handler(request: Request):
+    """Handle MCP requests with proper session management"""
+    try:
+        body = await request.json()
+        session_id = request.headers.get("mcp-session-id")
         
-        # Send tools list response (correct format from docs)
-        tools_response = {
-            "jsonrpc": "2.0",
-            "id": 2,
-            "result": {
-                "tools": [
-                    {
-                        "name": "search_docs",
-                        "title": "Search Documentation",
-                        "description": "Search through MCP documentation chunks",
-                        "inputSchema": {
-                            "type": "object",
-                            "properties": {
-                                "query": {
-                                    "type": "string",
-                                    "description": "Search query for MCP documentation"
-                                },
-                                "limit": {
-                                    "type": "integer",
-                                    "description": "Maximum number of results",
-                                    "default": 5
-                                }
-                            },
-                            "required": ["query"]
-                        }
-                    },
-                    {
-                        "name": "get_chunk",
-                        "title": "Get Documentation Chunk",
-                        "description": "Get a specific documentation chunk by ID",
-                        "inputSchema": {
-                            "type": "object",
-                            "properties": {
-                                "chunk_id": {
-                                    "type": "string",
-                                    "description": "Chunk ID to retrieve"
-                                }
-                            },
-                            "required": ["chunk_id"]
-                        }
+        # Handle initialization request
+        if is_initialize_request(body):
+            # Create new session
+            new_session_id = str(uuid.uuid4())
+            sessions[new_session_id] = {
+                "server": create_mcp_server(),
+                "initialized": True
+            }
+            
+            # Return initialization response
+            return {
+                "jsonrpc": "2.0",
+                "id": body.get("id"),
+                "result": {
+                    "protocolVersion": "2025-06-18",
+                    "capabilities": {
+                        "tools": {"listChanged": True},
+                        "resources": {}
                     },
-                    {
-                        "name": "list_docs",
-                        "title": "List Documents",
-                        "description": "List all available documents",
-                        "inputSchema": {
-                            "type": "object",
-                            "properties": {}
-                        }
+                    "serverInfo": {
+                        "name": "mcp-docs-server",
+                        "version": "1.0.0"
                     }
-                ]
+                }
             }
-        }
-        yield f"data: {json.dumps(tools_response)}\n\n"
         
-        # Keep connection alive
-        while True:
-            await asyncio.sleep(30)
-            yield f"data: {json.dumps({'jsonrpc': '2.0', 'method': 'ping', 'params': {}})}\n\n"
-    
-    return StreamingResponse(
-        event_generator(),
-        media_type="text/event-stream",
-        headers={
-            "Cache-Control": "no-cache",
-            "Connection": "keep-alive",
-            "Access-Control-Allow-Origin": "*",
-            "Access-Control-Allow-Headers": "*",
-        }
-    )
-
-@app.post("/mcp/stream")
-async def mcp_post(request: dict):
-    """Handle MCP requests via POST (for mcp-remote bidirectional communication)"""
-    try:
-        method = request.get("method")
-        params = request.get("params", {})
-        request_id = request.get("id")
+        # Handle other requests with session
+        if not session_id or session_id not in sessions:
+            return {
+                "jsonrpc": "2.0",
+                "error": {
+                    "code": -32000,
+                    "message": "Bad Request: No valid session ID provided"
+                },
+                "id": body.get("id")
+            }
+        
+        session = sessions[session_id]
+        method = body.get("method")
+        params = body.get("params", {})
+        request_id = body.get("id")
         
+        # Handle tools/list
         if method == "tools/list":
             return {
                 "jsonrpc": "2.0",
@@ -478,6 +444,7 @@ async def mcp_post(request: dict):
                 }
             }
         
+        # Handle tools/call
         elif method == "tools/call":
             tool_name = params.get("name")
             arguments = params.get("arguments", {})
@@ -624,7 +591,7 @@ async def mcp_post(request: dict):
     except Exception as e:
         return {
             "jsonrpc": "2.0",
-            "id": request.get("id"),
+            "id": body.get("id") if 'body' in locals() else None,
             "error": {
                 "code": -32603,
                 "message": f"Internal error: {str(e)}"