add api key based http auth

app.py

@@ -23,7 +23,7 @@ os.environ["WANDB_SILENT"] = "True"
 os.environ["WEAVE_SILENT"] = "True"
 
 from fastapi import FastAPI
-from fastapi.responses import HTMLResponse
+from fastapi.responses import HTMLResponse, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from mcp.server.fastmcp import FastMCP
 
@@ -37,6 +37,13 @@ from wandb_mcp_server.server import (
     ServerMCPArgs
 )
 
+# Import authentication
+from wandb_mcp_server.auth import (
+    mcp_auth_middleware,
+    create_resource_metadata_response,
+    MCPAuthConfig
+)
+
 # Configure logging
 logging.basicConfig(
     level=logging.INFO,
@@ -61,15 +68,17 @@ args = ServerMCPArgs(
 )
 
 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")
+api_key = validate_and_get_api_key(args)
+if api_key:
+    try:
+        setup_wandb_login(api_key)
+        initialize_weave_tracing()
+        wandb_configured = True
+        logger.info("Server W&B API key configured successfully")
+    except Exception as e:
+        logger.warning(f"Failed to configure server W&B API key: {e}")
+else:
+    logger.info("No server W&B API key configured - clients will provide their own")
 
 # Create the MCP server
 logger.info("Creating W&B MCP server...")
@@ -103,12 +112,24 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+# Add authentication middleware for MCP endpoints
+@app.middleware("http")
+async def auth_middleware(request, call_next):
+    """Add OAuth 2.1 Bearer token authentication for MCP endpoints."""
+    return await mcp_auth_middleware(request, call_next)
+
 # Add custom routes
 @app.get("/", response_class=HTMLResponse)
 async def index():
     """Serve the landing page."""
     return INDEX_HTML_CONTENT
 
+@app.get("/.well-known/oauth-protected-resource")
+async def resource_metadata():
+    """OAuth 2.0 Protected Resource Metadata endpoint (RFC 9728)."""
+    config = MCPAuthConfig()
+    return JSONResponse(create_resource_metadata_response(config))
+
 @app.get("/health")
 async def health():
     """Health check endpoint."""
@@ -119,11 +140,14 @@ async def health():
     except:
         tool_count = 0
     
+    auth_status = "disabled" if os.environ.get("MCP_AUTH_DISABLED", "false").lower() == "true" else "enabled"
+    
     return {
         "status": "healthy",
         "service": "wandb-mcp-server",
         "wandb_configured": wandb_configured,
-        "tools_registered": tool_count
+        "tools_registered": tool_count,
+        "authentication": auth_status
     }
 
 # Mount the MCP streamable HTTP app

src/wandb_mcp_server/auth.py

@@ -0,0 +1,188 @@
+"""
+Authentication middleware for W&B MCP Server.
+
+Implements Bearer token validation for HTTP transport as per 
+MCP specification: https://modelcontextprotocol.io/specification/draft/basic/authorization
+
+Clients send their W&B API keys as Bearer tokens, which the server
+then uses for all W&B operations on behalf of that client.
+"""
+
+import os
+import logging
+import re
+from typing import Optional, Dict, Any
+from fastapi import HTTPException, Request, status
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from fastapi.responses import JSONResponse
+
+logger = logging.getLogger(__name__)
+
+# Bearer token security scheme
+bearer_scheme = HTTPBearer(auto_error=False)
+
+
+class MCPAuthConfig:
+    """
+    Configuration for MCP authentication.
+    
+    For HTTP transport: Accepts any W&B API key as a Bearer token.
+    The server uses the client's token for all W&B operations.
+    """
+    
+    def __init__(self):
+        self.resource_metadata_url = os.environ.get(
+            "MCP_RESOURCE_METADATA_URL", 
+            "/.well-known/oauth-protected-resource"
+        )
+        # Point to W&B's Auth0 instance for reference
+        self.authorization_server = os.environ.get(
+            "MCP_AUTH_SERVER", 
+            "https://wandb.auth0.com"
+        )
+
+
+def is_valid_wandb_api_key(token: str) -> bool:
+    """
+    Check if a token looks like a valid W&B API key format.
+    W&B API keys are typically 40 characters of alphanumeric + some special chars.
+    """
+    if not token or len(token) < 20 or len(token) > 100:
+        return False
+    # Basic validation - W&B keys contain alphanumeric and some special characters
+    # This is a permissive check since W&B key format may vary
+    if re.match(r'^[a-zA-Z0-9_\-\.]+$', token):
+        return True
+    return False
+
+
+async def validate_bearer_token(
+    credentials: Optional[HTTPAuthorizationCredentials],
+    config: MCPAuthConfig
+) -> str:
+    """
+    Validate Bearer token (W&B API key) for MCP access.
+    
+    Accepts any valid-looking W&B API key. The actual validation
+    happens when the key is used to call W&B APIs.
+    
+    Returns:
+        The W&B API key to use for operations
+        
+    Raises:
+        HTTPException: 401 Unauthorized with WWW-Authenticate header
+    """
+    if not credentials or not credentials.credentials:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Authorization required - please provide your W&B API key as a Bearer token",
+            headers={
+                "WWW-Authenticate": f'Bearer realm="W&B MCP", '
+                                   f'resource_metadata="{config.resource_metadata_url}"'
+            }
+        )
+    
+    token = credentials.credentials
+    
+    # Basic format validation
+    if not is_valid_wandb_api_key(token):
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid W&B API key format. Get your key at: https://wandb.ai/authorize",
+            headers={
+                "WWW-Authenticate": f'Bearer realm="W&B MCP", '
+                                   f'error="invalid_token", '
+                                   f'resource_metadata="{config.resource_metadata_url}"'
+            }
+        )
+    
+    logger.debug("Bearer token validated successfully")
+    return token
+
+
+async def mcp_auth_middleware(request: Request, call_next):
+    """
+    FastAPI middleware for MCP authentication on HTTP transport.
+    
+    Only applies to MCP endpoints (/mcp/*).
+    Extracts the client's W&B API key from the Bearer token and stores it
+    for use in W&B operations.
+    """
+    # Only apply auth to MCP endpoints
+    if not request.url.path.startswith("/mcp"):
+        return await call_next(request)
+    
+    # Skip auth if explicitly disabled (development only)
+    if os.environ.get("MCP_AUTH_DISABLED", "false").lower() == "true":
+        logger.warning("MCP authentication is disabled - endpoints are publicly accessible")
+        return await call_next(request)
+    
+    config = MCPAuthConfig()
+    
+    try:
+        # Extract bearer token from Authorization header
+        authorization = request.headers.get("Authorization", "")
+        credentials = None
+        if authorization.startswith("Bearer "):
+            credentials = HTTPAuthorizationCredentials(
+                scheme="Bearer",
+                credentials=authorization[7:]  # Remove "Bearer " prefix
+            )
+        
+        # Validate and get the W&B API key
+        wandb_api_key = await validate_bearer_token(credentials, config)
+        
+        # Store the API key in request state for W&B operations
+        # The MCP tools should access this from the request context
+        request.state.wandb_api_key = wandb_api_key
+        
+        # For now, we'll set it in environment (in production, use contextvars)
+        # Save the original value to restore later
+        original_api_key = os.environ.get("WANDB_API_KEY")
+        os.environ["WANDB_API_KEY"] = wandb_api_key
+        
+        try:
+            # Continue processing
+            response = await call_next(request)
+        finally:
+            # Restore original environment
+            if original_api_key:
+                os.environ["WANDB_API_KEY"] = original_api_key
+            elif "WANDB_API_KEY" in os.environ:
+                del os.environ["WANDB_API_KEY"]
+            
+        return response
+        
+    except HTTPException as e:
+        # Return proper error response
+        return JSONResponse(
+            status_code=e.status_code,
+            content={"error": e.detail},
+            headers=e.headers
+        )
+    except Exception as e:
+        logger.error(f"Authentication error: {e}")
+        return JSONResponse(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            content={"error": "Authentication failed"},
+            headers={
+                "WWW-Authenticate": f'Bearer realm="W&B MCP", '
+                                   f'resource_metadata="{config.resource_metadata_url}"'
+            }
+        )
+
+
+def create_resource_metadata_response(config: MCPAuthConfig) -> Dict[str, Any]:
+    """
+    Create OAuth 2.0 Protected Resource Metadata response (RFC 9728).
+    
+    This tells MCP clients that we use W&B API keys as Bearer tokens.
+    Points to W&B's Auth0 instance where users can get their API keys.
+    """
+    return {
+        "resource": os.environ.get("MCP_SERVER_URL", "https://wandb-mcp-server.hf.space"),
+        "authorization_servers": [config.authorization_server],
+        "bearer_methods_supported": ["header"],
+        "resource_documentation": "https://github.com/wandb/wandb-mcp-server",
+        "authentication_note": "Use your W&B API key as a Bearer token. Get your key at https://wandb.ai/authorize",
+    }

src/wandb_mcp_server/mcp_tools/query_weave.py

@@ -6,7 +6,14 @@ from wandb_mcp_server.weave_api.models import QueryResult
 
 logger = get_rich_logger(__name__)
 
-_trace_service = TraceService()
+# Lazy load the trace service to avoid requiring API key at import time
+_trace_service = None
+
+def get_trace_service():
+    global _trace_service
+    if _trace_service is None:
+        _trace_service = TraceService()
+    return _trace_service
 
 QUERY_WEAVE_TRACES_TOOL_DESCRIPTION = """
 Query Weave traces, trace metadata, and trace costs with filtering and sorting options.
@@ -301,7 +308,7 @@ def query_traces(
     but delegates to our new implementation.
     """
     # If api_key was provided, create a new service with that key
-    service = _trace_service
+    service = get_trace_service()
     if api_key:
         service = TraceService(
             api_key=api_key,
@@ -410,7 +417,7 @@ async def query_paginated_weave_traces(
         QueryResult: A Pydantic model containing the query results
     """
     # If api_key was provided, create a new service with that key
-    service = _trace_service
+    service = get_trace_service()
     if api_key:
         service = TraceService(
             api_key=api_key,

src/wandb_mcp_server/server.py

@@ -119,10 +119,13 @@ def setup_wandb_login(api_key: str) -> None:
         sys.stderr = original_stderr
 
 
-def validate_and_get_api_key(args: ServerMCPArgs) -> str:
+def validate_and_get_api_key(args: ServerMCPArgs) -> Optional[str]:
     """
     Validate and retrieve the W&B API key from various sources.
     
+    For HTTP transport: API key is optional (clients provide their own)
+    For STDIO transport: API key is required from environment
+    
     Priority order:
     1. Command-line argument (--wandb-api-key)
     2. Environment variable (WANDB_API_KEY)
@@ -133,16 +136,25 @@ def validate_and_get_api_key(args: ServerMCPArgs) -> str:
         args: Parsed command-line arguments
         
     Returns:
-        The W&B API key
+        The W&B API key if found, None otherwise
         
     Raises:
-        ValueError: If no API key is found
+        ValueError: If no API key is found for STDIO transport
     """
     api_key = args.wandb_api_key or get_server_args().wandb_api_key
     
+    # For HTTP transport, API key is optional (clients provide their own)
+    if args.transport == "http":
+        if api_key:
+            logger.info("Server W&B API key configured (for server operations)")
+        else:
+            logger.info("No server W&B API key configured (clients will provide their own)")
+        return api_key
+    
+    # For STDIO transport, API key is required
     if not api_key:
         raise ValueError(
-            "WANDB_API_KEY must be set. Options:\n"
+            "WANDB_API_KEY must be set for STDIO transport. Options:\n"
             "1. Command-line: --wandb-api-key YOUR_KEY\n"
             "2. Environment: export WANDB_API_KEY=YOUR_KEY\n"
             "3. .env file: WANDB_API_KEY=YOUR_KEY\n"
@@ -387,14 +399,27 @@ def create_mcp_server(transport: str, host: str = "localhost", port: Optional[in
         
     Raises:
         ValueError: If transport type is invalid
+    
+    Authentication:
+        - STDIO transport: Uses environment variables (WANDB_API_KEY required)
+        - HTTP transport: Clients provide W&B API key as Bearer token
+          Set MCP_AUTH_DISABLED=true to disable auth (development only)
     """
     if transport == "http":
         port = port if port is not None else 8080
         logger.info(f"Configuring HTTP server on {host}:{port}")
         mcp = FastMCP("weave-mcp-server", host=host, port=port, stateless_http=True)
+        
+        # Log authentication status for HTTP
+        if os.environ.get("MCP_AUTH_DISABLED", "false").lower() == "true":
+            logger.warning("⚠️  MCP authentication is DISABLED - server is publicly accessible")
+        else:
+            logger.info("🔒 MCP authentication enabled - clients must provide W&B API key as Bearer token")
+            
     elif transport == "stdio":
         logger.info("Configuring stdio server")
         mcp = FastMCP("weave-mcp-server")
+        logger.info("STDIO transport uses environment variable authentication")
     else:
         raise ValueError(f"Invalid transport type: {transport}. Must be 'stdio' or 'http'")
     
@@ -422,11 +447,12 @@ def cli():
         --wandb-api-key KEY         W&B API key (can also use env var)
         
     Environment Variables:
-        WANDB_API_KEY               W&B API key for authentication
+        WANDB_API_KEY               W&B API key (required for STDIO, optional for HTTP)
         MCP_SERVER_LOG_LEVEL        Server log level (DEBUG, INFO, WARNING, ERROR)
         WANDB_SILENT                Set to "False" to enable W&B output (default: True)
         WEAVE_SILENT                Set to "False" to enable Weave output (default: True)
         WANDB_DEBUG                 Set to "true" to enable W&B debug logging
+        MCP_AUTH_DISABLED           Set to "true" to disable HTTP auth (dev only)
     """
     # Parse command line arguments
     import simple_parsing
@@ -438,8 +464,9 @@ def cli():
     # Validate and get API key
     api_key = validate_and_get_api_key(args)
     
-    # Perform W&B login
-    setup_wandb_login(api_key)
+    # Perform W&B login only if we have an API key
+    if api_key:
+        setup_wandb_login(api_key)
     
     # Initialize Weave tracing for MCP tool calls
     weave_initialized = initialize_weave_tracing()