fixed mcp routing issue

HUGGINGFACE_DEPLOYMENT.md

@@ -8,15 +8,16 @@ The application runs as a FastAPI server on port 7860 (HF Spaces default) with:
 - **Main landing page**: `/` - Serves the index.html with setup instructions
 - **Health check**: `/health` - Returns server status and W&B configuration
 - **MCP endpoint**: `/mcp` - Streamable HTTP transport endpoint for MCP
-  - POST `/mcp` - Handles JSON-RPC requests (initialize, tools/list, tools/call)
-  - GET `/mcp` - SSE endpoint for server-initiated messages and long-lived connections
+  - Uses Server-Sent Events (SSE) for responses
+  - Requires `Accept: application/json, text/event-stream` header
+  - Supports initialize, tools/list, tools/call methods
 
 ## Key Changes for HF Spaces
 
 ### 1. app.py
 - Creates a FastAPI application that serves the landing page
-- Implements MCP streamable HTTP protocol directly (no FastMCP mounting issues)
-- Handles both POST requests (JSON-RPC) and GET requests (SSE) at `/mcp`
+- Mounts FastMCP server using `mcp.streamable_http_app()` pattern (following HuggingFace example)
+- Uses lifespan context manager for session management
 - Configured to run on `0.0.0.0:7860` (HF Spaces requirement)
 - Sets W&B cache directories to `/tmp` to avoid permission issues
 
@@ -29,6 +30,11 @@ The application runs as a FastAPI server on port 7860 (HF Spaces default) with:
 - FastAPI and uvicorn moved to main dependencies (not optional)
 - All dependencies listed in requirements.txt for HF Spaces
 
+### 4. Lazy Loading Fix
+- Fixed `TraceService` initialization in `query_weave.py` to use lazy loading
+- This allows the server to start even without a W&B API key
+- The service is only initialized when first needed
+
 ## Environment Variables
 
 Required in HF Spaces settings:
@@ -98,7 +104,41 @@ The server will start on http://localhost:7860
 
 ## MCP Client Configuration
 
-Example configuration for Claude Desktop or other MCP clients:
+### Important Notes
+
+The MCP server uses the Streamable HTTP transport which:
+- Returns responses in Server-Sent Events (SSE) format
+- Requires the client to send `Accept: application/json, text/event-stream` header
+- Uses session management for stateful operations
+
+### Testing with curl
+
+```bash
+# Initialize the server
+curl -X POST https://[your-username]-[space-name].hf.space/mcp \
+  -H "Content-Type: application/json" \
+  -H "Accept: application/json, text/event-stream" \
+  -d '{
+    "jsonrpc": "2.0",
+    "method": "initialize",
+    "params": {
+      "protocolVersion": "0.1.0",
+      "capabilities": {},
+      "clientInfo": {"name": "test-client", "version": "1.0"}
+    },
+    "id": 1
+  }'
+
+# List available tools
+curl -X POST https://[your-username]-[space-name].hf.space/mcp \
+  -H "Content-Type: application/json" \
+  -H "Accept: application/json, text/event-stream" \
+  -d '{"jsonrpc":"2.0","method":"tools/list","params":{},"id":2}'
+```
+
+### MCP Client Configuration
+
+For MCP clients that support streamable HTTP:
 
 ```json
 {

app.py

@@ -2,17 +2,14 @@
 """
 HuggingFace Spaces entry point for the Weights & Biases MCP Server.
 
-This implements MCP streamable HTTP transport directly in FastAPI.
+Using the correct FastMCP mounting pattern with streamable_http_app().
 """
 
 import os
 import sys
 import logging
-import json
-import asyncio
+import contextlib
 from pathlib import Path
-from typing import Dict, Any, Optional, List
-from datetime import datetime
 
 # Add the src directory to Python path
 sys.path.insert(0, str(Path(__file__).parent / "src"))
@@ -25,10 +22,20 @@ os.environ["HOME"] = "/tmp"
 os.environ["WANDB_SILENT"] = "True"
 os.environ["WEAVE_SILENT"] = "True"
 
-from fastapi import FastAPI, Request, Response, HTTPException
-from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse
+from fastapi import FastAPI
+from fastapi.responses import HTMLResponse
 from fastapi.middleware.cors import CORSMiddleware
-import uvicorn
+from mcp.server.fastmcp import FastMCP
+
+# Import W&B setup functions
+from wandb_mcp_server.server import (
+    validate_and_get_api_key,
+    setup_wandb_login,
+    configure_wandb_logging,
+    initialize_weave_tracing,
+    register_tools,
+    ServerMCPArgs
+)
 
 # Configure logging
 logging.basicConfig(
@@ -42,11 +49,49 @@ INDEX_HTML_PATH = Path(__file__).parent / "index.html"
 with open(INDEX_HTML_PATH, "r") as f:
     INDEX_HTML_CONTENT = f.read()
 
-# Create FastAPI app
+# Initialize W&B
+logger.info("Initializing W&B configuration...")
+configure_wandb_logging()
+
+args = ServerMCPArgs(
+    transport="http",
+    host="0.0.0.0",
+    port=7860,
+    wandb_api_key=os.environ.get("WANDB_API_KEY")
+)
+
+wandb_configured = False
+try:
+    api_key = validate_and_get_api_key(args)
+    setup_wandb_login(api_key)
+    initialize_weave_tracing()
+    wandb_configured = True
+    logger.info("W&B API configured successfully")
+except ValueError as e:
+    logger.warning(f"W&B API key not configured: {e}")
+    logger.warning("Server will start but W&B operations will fail")
+
+# Create the MCP server
+logger.info("Creating W&B MCP server...")
+mcp = FastMCP("wandb-mcp-server")
+
+# Register all W&B tools
+register_tools(mcp)
+
+# Create lifespan context manager for session management
+@contextlib.asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Manage MCP session lifecycle."""
+    async with mcp.session_manager.run():
+        logger.info("MCP session manager started")
+        yield
+        logger.info("MCP session manager stopped")
+
+# Create the main FastAPI app with lifespan
 app = FastAPI(
     title="Weights & Biases MCP Server",
-    description="Model Context Protocol server for querying W&B data",
-    version="0.1.0"
+    description="Model Context Protocol server for W&B",
+    lifespan=lifespan
 )
 
 # Add CORS middleware
@@ -58,115 +103,7 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# MCP server state
-mcp_initialized = False
-mcp_tools = {}
-wandb_configured = False
-
-@app.on_event("startup")
-async def startup_event():
-    """Initialize W&B and MCP tools on startup."""
-    global mcp_initialized, mcp_tools, wandb_configured
-    
-    logger.info("Starting Weights & Biases MCP Server on HuggingFace Spaces")
-    
-    try:
-        # Import W&B components
-        from wandb_mcp_server.server import (
-            validate_and_get_api_key,
-            setup_wandb_login,
-            configure_wandb_logging,
-            initialize_weave_tracing,
-            ServerMCPArgs
-        )
-        
-        # Configure W&B
-        configure_wandb_logging()
-        
-        args = ServerMCPArgs(
-            transport="http",
-            host="0.0.0.0",
-            port=7860,
-            wandb_api_key=os.environ.get("WANDB_API_KEY")
-        )
-        
-        try:
-            api_key = validate_and_get_api_key(args)
-            setup_wandb_login(api_key)
-            initialize_weave_tracing()
-            wandb_configured = True
-            logger.info("W&B API configured successfully")
-        except ValueError as e:
-            logger.warning(f"W&B API key not configured: {e}")
-            logger.warning("Server will start but W&B operations will fail")
-        
-        # Import MCP tools
-        from wandb_mcp_server.mcp_tools.query_weave import (
-            QUERY_WEAVE_TRACES_TOOL_DESCRIPTION,
-            query_paginated_weave_traces
-        )
-        from wandb_mcp_server.mcp_tools.count_traces import (
-            COUNT_WEAVE_TRACES_TOOL_DESCRIPTION,
-            count_traces
-        )
-        from wandb_mcp_server.mcp_tools.query_wandb_gql import (
-            QUERY_WANDB_GQL_TOOL_DESCRIPTION,
-            query_paginated_wandb_gql
-        )
-        from wandb_mcp_server.mcp_tools.create_report import (
-            CREATE_WANDB_REPORT_TOOL_DESCRIPTION,
-            create_report
-        )
-        from wandb_mcp_server.mcp_tools.list_wandb_entities_projects import (
-            LIST_ENTITY_PROJECTS_TOOL_DESCRIPTION,
-            list_entity_projects
-        )
-        from wandb_mcp_server.mcp_tools.query_wandbot import (
-            WANDBOT_TOOL_DESCRIPTION,
-            query_wandbot_api
-        )
-        
-        # Register tools with their descriptions
-        mcp_tools = {
-            "query_weave_traces_tool": {
-                "function": query_paginated_weave_traces,
-                "description": QUERY_WEAVE_TRACES_TOOL_DESCRIPTION,
-                "async": True
-            },
-            "count_weave_traces_tool": {
-                "function": count_traces,
-                "description": COUNT_WEAVE_TRACES_TOOL_DESCRIPTION,
-                "async": False
-            },
-            "query_wandb_tool": {
-                "function": query_paginated_wandb_gql,
-                "description": QUERY_WANDB_GQL_TOOL_DESCRIPTION,
-                "async": False
-            },
-            "create_wandb_report_tool": {
-                "function": create_report,
-                "description": CREATE_WANDB_REPORT_TOOL_DESCRIPTION,
-                "async": False
-            },
-            "query_wandb_entity_projects": {
-                "function": list_entity_projects,
-                "description": LIST_ENTITY_PROJECTS_TOOL_DESCRIPTION,
-                "async": False
-            },
-            "query_wandb_support_bot": {
-                "function": query_wandbot_api,
-                "description": WANDBOT_TOOL_DESCRIPTION,
-                "async": False
-            }
-        }
-        
-        mcp_initialized = True
-        logger.info(f"MCP server initialized with {len(mcp_tools)} tools")
-        
-    except Exception as e:
-        logger.error(f"Error initializing MCP server: {e}")
-        mcp_initialized = False
-
+# Add custom routes
 @app.get("/", response_class=HTMLResponse)
 async def index():
     """Serve the landing page."""
@@ -175,231 +112,35 @@ async def index():
 @app.get("/health")
 async def health():
     """Health check endpoint."""
+    # list_tools is async, so we need to handle it properly
+    try:
+        tools = await mcp.list_tools()
+        tool_count = len(tools)
+    except:
+        tool_count = 0
+    
     return {
         "status": "healthy",
         "service": "wandb-mcp-server",
         "wandb_configured": wandb_configured,
-        "mcp_initialized": mcp_initialized,
-        "tools_count": len(mcp_tools)
+        "tools_registered": tool_count
     }
 
-# MCP Streamable HTTP Implementation
-@app.post("/mcp")
-async def handle_mcp_post(request: Request):
-    """
-    Handle MCP POST requests following the streamable HTTP transport protocol.
-    """
-    if not mcp_initialized:
-        return JSONResponse(
-            status_code=503,
-            content={
-                "jsonrpc": "2.0",
-                "error": {
-                    "code": -32603,
-                    "message": "MCP server not initialized"
-                }
-            }
-        )
-    
-    try:
-        body = await request.json()
-        method = body.get("method", "")
-        params = body.get("params", {})
-        request_id = body.get("id")
-        
-        # Handle different MCP methods
-        if method == "initialize":
-            return {
-                "jsonrpc": "2.0",
-                "id": request_id,
-                "result": {
-                    "protocolVersion": "0.1.0",
-                    "capabilities": {
-                        "tools": {"listChanged": True},
-                        "prompts": {"listChanged": False},
-                        "resources": {"listChanged": False}
-                    },
-                    "serverInfo": {
-                        "name": "wandb-mcp-server",
-                        "version": "0.1.0"
-                    }
-                }
-            }
-        
-        elif method == "tools/list":
-            tools_list = []
-            for tool_name, tool_info in mcp_tools.items():
-                # Extract a shorter description (first line)
-                desc_lines = tool_info["description"].split('\n')
-                short_desc = desc_lines[0] if desc_lines else f"W&B tool: {tool_name}"
-                
-                tools_list.append({
-                    "name": tool_name,
-                    "description": short_desc,
-                    "inputSchema": {
-                        "type": "object",
-                        "properties": {},
-                        "required": []
-                    }
-                })
-            
-            return {
-                "jsonrpc": "2.0",
-                "id": request_id,
-                "result": {"tools": tools_list}
-            }
-        
-        elif method == "tools/call":
-            tool_name = params.get("name")
-            tool_args = params.get("arguments", {})
-            
-            if tool_name not in mcp_tools:
-                return {
-                    "jsonrpc": "2.0",
-                    "id": request_id,
-                    "error": {
-                        "code": -32601,
-                        "message": f"Tool not found: {tool_name}"
-                    }
-                }
-            
-            try:
-                tool_info = mcp_tools[tool_name]
-                tool_function = tool_info["function"]
-                
-                # Execute the tool
-                if tool_info["async"]:
-                    result = await tool_function(**tool_args)
-                else:
-                    result = tool_function(**tool_args)
-                
-                # Format the result
-                if isinstance(result, str):
-                    content_text = result
-                elif isinstance(result, dict):
-                    content_text = json.dumps(result, indent=2)
-                else:
-                    content_text = str(result)
-                
-                return {
-                    "jsonrpc": "2.0",
-                    "id": request_id,
-                    "result": {
-                        "content": [
-                            {
-                                "type": "text",
-                                "text": content_text
-                            }
-                        ]
-                    }
-                }
-                
-            except Exception as e:
-                logger.error(f"Error executing tool {tool_name}: {e}")
-                return {
-                    "jsonrpc": "2.0",
-                    "id": request_id,
-                    "error": {
-                        "code": -32603,
-                        "message": f"Tool execution error: {str(e)}"
-                    }
-                }
-        
-        else:
-            return {
-                "jsonrpc": "2.0",
-                "id": request_id,
-                "error": {
-                    "code": -32601,
-                    "message": f"Method not found: {method}"
-                }
-            }
-            
-    except Exception as e:
-        logger.error(f"Error handling MCP request: {e}")
-        return JSONResponse(
-            status_code=500,
-            content={
-                "jsonrpc": "2.0",
-                "error": {
-                    "code": -32603,
-                    "message": f"Internal error: {str(e)}"
-                }
-            }
-        )
-
-@app.get("/mcp")
-async def handle_mcp_sse(request: Request):
-    """
-    Handle MCP GET requests for SSE (Server-Sent Events) streaming.
-    This enables server-initiated messages and long-lived connections.
-    """
-    if not mcp_initialized:
-        return JSONResponse(
-            status_code=503,
-            content={"error": "MCP server not initialized"}
-        )
-    
-    async def event_stream():
-        """Generate server-sent events for MCP."""
-        try:
-            # Send initial connection confirmation
-            yield f"data: {json.dumps({'jsonrpc': '2.0', 'method': 'connection/ready', 'params': {'status': 'connected', 'timestamp': datetime.utcnow().isoformat()}})}\n\n"
-            
-            # Keep connection alive with periodic heartbeats
-            while True:
-                await asyncio.sleep(30)
-                yield f"data: {json.dumps({'jsonrpc': '2.0', 'method': 'ping', 'params': {'timestamp': datetime.utcnow().isoformat()}})}\n\n"
-                
-        except asyncio.CancelledError:
-            logger.info("SSE connection closed")
-            raise
-        except Exception as e:
-            logger.error(f"Error in SSE stream: {e}")
-            yield f"data: {json.dumps({'jsonrpc': '2.0', 'method': 'error', 'params': {'message': str(e)}})}\n\n"
-    
-    return StreamingResponse(
-        event_stream(),
-        media_type="text/event-stream",
-        headers={
-            "Cache-Control": "no-cache",
-            "Connection": "keep-alive",
-            "X-Accel-Buffering": "no",
-            "Access-Control-Allow-Origin": "*"
-        }
-    )
+# Mount the MCP streamable HTTP app
+# Note: streamable_http_app() creates internal routes at /mcp
+# So we mount at root to avoid /mcp/mcp double path
+mcp_app = mcp.streamable_http_app()
+logger.info("Mounting MCP streamable HTTP app")
+# Mount at root, so MCP endpoint will be at /mcp (not /mcp/mcp)
+app.mount("/", mcp_app)
 
-# Additional MCP endpoints for better compatibility
-@app.options("/mcp")
-async def handle_mcp_options():
-    """Handle OPTIONS requests for CORS preflight."""
-    return Response(
-        status_code=200,
-        headers={
-            "Access-Control-Allow-Origin": "*",
-            "Access-Control-Allow-Methods": "GET, POST, OPTIONS",
-            "Access-Control-Allow-Headers": "Content-Type, Accept",
-            "Access-Control-Max-Age": "3600"
-        }
-    )
+# Port for HF Spaces
+PORT = int(os.environ.get("PORT", "7860"))
 
-def main():
-    """Main entry point for HuggingFace Spaces."""
-    port = 7860
-    host = "0.0.0.0"
-    
-    logger.info(f"Starting server on {host}:{port}")
+if __name__ == "__main__":
+    import uvicorn
+    logger.info(f"Starting server on 0.0.0.0:{PORT}")
     logger.info("Landing page: /")
     logger.info("Health check: /health")
-    logger.info("MCP endpoint (Streamable HTTP): /mcp")
-    
-    uvicorn.run(
-        app,
-        host=host,
-        port=port,
-        log_level="info",
-        reload=False
-    )
-
-if __name__ == "__main__":
-    main()
+    logger.info("MCP endpoint: /mcp")
+    uvicorn.run(app, host="0.0.0.0", port=PORT)

uv.lock

@@ -2092,6 +2092,7 @@ name = "wandb-mcp-server"
 version = "0.1.0"
 source = { editable = "." }
 dependencies = [
+    { name = "fastapi" },
     { name = "httpx" },
     { name = "mcp", extra = ["cli"] },
     { name = "networkx" },
@@ -2100,16 +2101,13 @@ dependencies = [
     { name = "requests" },
     { name = "simple-parsing" },
     { name = "tiktoken" },
+    { name = "uvicorn" },
     { name = "wandb" },
     { name = "wandb-workspaces" },
     { name = "weave" },
 ]
 
 [package.optional-dependencies]
-http = [
-    { name = "fastapi" },
-    { name = "uvicorn" },
-]
 test = [
     { name = "anthropic" },
     { name = "litellm" },
@@ -2125,7 +2123,7 @@ dev = [
 [package.metadata]
 requires-dist = [
     { name = "anthropic", marker = "extra == 'test'", specifier = ">=0.50.0" },
-    { name = "fastapi", marker = "extra == 'http'", specifier = ">=0.104.0" },
+    { name = "fastapi", specifier = ">=0.104.0" },
     { name = "httpx", specifier = ">=0.28.1" },
     { name = "litellm", marker = "extra == 'test'", specifier = ">=1.67.2" },
     { name = "mcp", extras = ["cli"], specifier = ">=1.0.0" },
@@ -2137,12 +2135,12 @@ requires-dist = [
     { name = "requests", specifier = ">=2.31.0" },
     { name = "simple-parsing", specifier = ">=0.1.7" },
     { name = "tiktoken", specifier = ">=0.9.0" },
-    { name = "uvicorn", marker = "extra == 'http'", specifier = ">=0.24.0" },
+    { name = "uvicorn", specifier = ">=0.24.0" },
     { name = "wandb", specifier = ">=0.19.8" },
     { name = "wandb-workspaces", specifier = ">=0.1.12" },
     { name = "weave", specifier = ">=0.51.56" },
 ]
-provides-extras = ["test", "http"]
+provides-extras = ["test"]
 
 [package.metadata.requires-dev]
 dev = [{ name = "ruff", specifier = ">=0.11.12" }]