update app routing

HUGGINGFACE_DEPLOYMENT.md

@@ -7,15 +7,18 @@ This repository is configured for deployment on Hugging Face Spaces as a Model C
 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 (supports both regular HTTP and SSE upgrade)
+- **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
 
 ## Key Changes for HF Spaces
 
 ### 1. app.py
 - Creates a FastAPI application that serves the landing page
-- Mounts the MCP server as a sub-application at `/mcp`
+- Implements MCP streamable HTTP protocol directly (no FastMCP mounting issues)
+- Handles both POST requests (JSON-RPC) and GET requests (SSE) at `/mcp`
 - Configured to run on `0.0.0.0:7860` (HF Spaces requirement)
-- No need for static/templates directories - index.html is served directly
+- Sets W&B cache directories to `/tmp` to avoid permission issues
 
 ### 2. server.py
 - Exports necessary functions for HF Spaces initialization

app.py

@@ -2,13 +2,17 @@
 """
 HuggingFace Spaces entry point for the Weights & Biases MCP Server.
 
-Simplified approach for HF Spaces deployment.
+This implements MCP streamable HTTP transport directly in FastAPI.
 """
 
 import os
 import sys
 import logging
+import json
+import asyncio
 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"))
@@ -21,8 +25,8 @@ os.environ["HOME"] = "/tmp"
 os.environ["WANDB_SILENT"] = "True"
 os.environ["WEAVE_SILENT"] = "True"
 
-from fastapi import FastAPI
-from fastapi.responses import HTMLResponse
+from fastapi import FastAPI, Request, Response, HTTPException
+from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 import uvicorn
 
@@ -31,7 +35,7 @@ logging.basicConfig(
     level=logging.INFO,
     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
 )
-logger = logging.getLogger("huggingface-spaces-app")
+logger = logging.getLogger("wandb-mcp-server")
 
 # Read the index.html file content
 INDEX_HTML_PATH = Path(__file__).parent / "index.html"
@@ -54,6 +58,115 @@ 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
+
 @app.get("/", response_class=HTMLResponse)
 async def index():
     """Serve the landing page."""
@@ -62,93 +175,213 @@ async def index():
 @app.get("/health")
 async def health():
     """Health check endpoint."""
-    wandb_configured = bool(os.environ.get("WANDB_API_KEY"))
     return {
         "status": "healthy",
         "service": "wandb-mcp-server",
-        "wandb_configured": wandb_configured
+        "wandb_configured": wandb_configured,
+        "mcp_initialized": mcp_initialized,
+        "tools_count": len(mcp_tools)
     }
 
-# Import and setup MCP functionality after basic app setup
-try:
-    from wandb_mcp_server.server import (
-        validate_and_get_api_key,
-        setup_wandb_login,
-        configure_wandb_logging,
-        initialize_weave_tracing,
-        create_mcp_server,
-        ServerMCPArgs
-    )
-    
-    # Setup W&B on import
-    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")
-    )
+# 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:
-        api_key = validate_and_get_api_key(args)
-        setup_wandb_login(api_key)
-        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 MCP server instance
-    mcp_server = create_mcp_server("http", "0.0.0.0", 7860)
-    
-    # Try to extract the FastAPI app from FastMCP and mount it
-    if hasattr(mcp_server, 'app'):
-        # Mount MCP app routes under /mcp
-        from fastapi import APIRouter
-        mcp_router = APIRouter(prefix="/mcp")
+        body = await request.json()
+        method = body.get("method", "")
+        params = body.get("params", {})
+        request_id = body.get("id")
         
-        # Copy routes from MCP app to our router
-        if hasattr(mcp_server.app, 'routes'):
-            for route in mcp_server.app.routes:
-                if hasattr(route, 'endpoint') and hasattr(route, 'path'):
-                    mcp_router.add_api_route(
-                        path=route.path,
-                        endpoint=route.endpoint,
-                        methods=route.methods if hasattr(route, 'methods') else ["POST", "GET"],
-                    )
+        # 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"
+                    }
+                }
+            }
         
-        app.include_router(mcp_router)
-        logger.info("MCP routes mounted at /mcp")
-    else:
-        # Fallback: Create simple MCP endpoint
-        @app.post("/mcp")
-        async def mcp_endpoint():
-            return {"message": "MCP server is running", "status": "ready"}
+        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}
+            }
         
-        @app.get("/mcp")
-        async def mcp_sse():
-            from fastapi.responses import StreamingResponse
-            import json
+        elif method == "tools/call":
+            tool_name = params.get("name")
+            tool_args = params.get("arguments", {})
             
-            async def event_stream():
-                yield f"data: {json.dumps({'status': 'connected'})}\n\n"
+            if tool_name not in mcp_tools:
+                return {
+                    "jsonrpc": "2.0",
+                    "id": request_id,
+                    "error": {
+                        "code": -32601,
+                        "message": f"Tool not found: {tool_name}"
+                    }
+                }
             
-            return StreamingResponse(event_stream(), media_type="text/event-stream")
+            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)}"
+                    }
+                }
         
-        logger.warning("Using fallback MCP endpoints")
-    
-except Exception as e:
-    logger.error(f"Error setting up MCP server: {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"}
+        )
     
-    # Provide fallback endpoints even if MCP setup fails
-    @app.post("/mcp")
-    async def mcp_fallback():
-        return {"error": "MCP server initialization failed", "details": str(e)}
+    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"
     
-    @app.get("/mcp")
-    async def mcp_sse_fallback():
-        return {"error": "MCP SSE not available", "details": str(e)}
+    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": "*"
+        }
+    )
+
+# 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"
+        }
+    )
 
 def main():
     """Main entry point for HuggingFace Spaces."""
@@ -158,7 +391,7 @@ def main():
     logger.info(f"Starting server on {host}:{port}")
     logger.info("Landing page: /")
     logger.info("Health check: /health")
-    logger.info("MCP endpoint: /mcp")
+    logger.info("MCP endpoint (Streamable HTTP): /mcp")
     
     uvicorn.run(
         app,