clean up and oauth experiments

AUTH_README.md

@@ -0,0 +1,204 @@
+# W&B MCP Server Authentication
+
+The W&B MCP Server uses **Bearer token authentication** with W&B API keys for secure access to your Weights & Biases data.
+
+## How It Works
+
+### API Key Authentication
+
+The server uses standard HTTP Bearer token authentication where your W&B API key serves as the Bearer token:
+
+```http
+Authorization: Bearer YOUR_WANDB_API_KEY
+```
+
+**Key Features:**
+- Each client provides their own W&B API key
+- Server uses the client's key for all W&B operations  
+- Perfect isolation between users
+- No server-side API key needed for HTTP transport
+
+### Getting Your W&B API Key
+
+1. Go to [https://wandb.ai/authorize](https://wandb.ai/authorize)
+2. Log in with your W&B account
+3. Copy your API key (exactly 40 characters, no spaces)
+4. Use it as the Bearer token in your MCP client
+
+**Important**: W&B API keys must be exactly 40 alphanumeric characters. Make sure you copy the entire key without any extra spaces or line breaks.
+
+## Client Configuration
+
+### Mistral LeChat
+
+In Mistral LeChat, add a Custom MCP Connector:
+
+1. **Server URL**: `https://your-space.hf.space/mcp`
+2. **Authentication**: Choose "HTTP Bearer Token"
+3. **Token**: Enter your W&B API key
+
+### Claude Desktop / Cursor
+
+Configure in your MCP settings:
+
+```json
+{
+  "mcpServers": {
+    "wandb": {
+      "transport": "http",
+      "url": "http://localhost:8080/mcp",
+      "headers": {
+        "Authorization": "Bearer YOUR_WANDB_API_KEY",
+        "Accept": "application/json, text/event-stream"
+      }
+    }
+  }
+}
+```
+
+**Important Headers:**
+- `Authorization`: Your W&B API key as Bearer token
+- `Accept`: Must include both `application/json` and `text/event-stream` for MCP Streamable HTTP
+
+### Python Client Example
+
+```python
+import requests
+
+# Initialize MCP session
+response = requests.post(
+    "http://localhost:8080/mcp",
+    headers={
+        "Authorization": "Bearer YOUR_WANDB_API_KEY",
+        "Accept": "application/json, text/event-stream",
+        "Content-Type": "application/json"
+    },
+    json={
+        "jsonrpc": "2.0",
+        "method": "initialize",
+        "params": {},
+        "id": 1
+    }
+)
+```
+
+## Security Considerations
+
+### Non-Expiring API Keys
+
+W&B API keys don't expire by default, similar to GitHub Personal Access Tokens or OpenAI API keys. This is a design choice by W&B for developer convenience.
+
+**Best Practices:**
+- Rotate keys regularly (quarterly recommended)
+- Use separate keys for different services
+- Monitor usage at [wandb.ai/settings](https://wandb.ai/settings)
+- Revoke compromised keys immediately
+- Never commit keys to version control
+
+### Multi-User Deployment
+
+For HuggingFace Spaces or shared deployments:
+- Server requires no API key configuration
+- Each user provides their own key
+- Keys are used transiently per request
+- No keys are stored or logged
+
+## OAuth: Why We Can't Support Full OAuth 2.0
+
+### What We Tried
+
+We attempted to implement OAuth 2.0 support to provide a seamless authentication experience, especially for clients like Mistral LeChat that expect OAuth for custom connectors. This included:
+
+1. **OAuth Discovery Endpoints**: `/.well-known/oauth-authorization-server`
+2. **Authorization Flow**: Redirect to W&B's Auth0 login
+3. **Token Exchange**: Accept W&B API keys as "access tokens"
+4. **Device Flow**: Guide users to get their API key
+
+### Why It Doesn't Work
+
+**Fundamental Limitations:**
+
+1. **W&B Doesn't Provide OAuth for Third Parties**
+   - W&B uses Auth0 internally but doesn't allow third-party OAuth client registration
+   - No way to register our MCP server as an OAuth application
+   - Can't receive authorization codes or callbacks from W&B
+
+2. **API Keys Are Not OAuth Tokens**
+   - W&B provides permanent API keys, not temporary OAuth tokens
+   - No refresh mechanism, no expiration, no scopes
+   - Keys are managed through W&B's web interface, not OAuth flows
+
+3. **Cross-Domain Issues**
+   - OAuth requires the authorization server and resource server to cooperate
+   - W&B's Auth0 instance (`wandb.auth0.com`) doesn't know about our server
+   - Can't validate tokens or handle callbacks
+
+### What Would Be Needed for Full OAuth
+
+For proper OAuth 2.0 support, W&B would need to:
+
+1. **Allow OAuth Client Registration**
+   - Let developers register OAuth applications
+   - Provide client ID and secret
+   - Support redirect URIs for callbacks
+
+2. **Issue Real OAuth Tokens**
+   - Temporary access tokens (e.g., 1-hour expiry)
+   - Refresh tokens for obtaining new access tokens
+   - Scoped permissions (read-only, write, admin)
+
+3. **Provide Token Validation**
+   - Introspection endpoint for validating tokens
+   - Revocation endpoint for invalidating tokens
+   - JWKS endpoint for JWT validation
+
+### Current Solution
+
+Given these limitations, we use W&B API keys directly as Bearer tokens. This approach:
+- ✅ Works with all W&B functionality
+- ✅ Compatible with MCP specification
+- ✅ Simple and reliable
+- ✅ Follows industry patterns (GitHub, OpenAI)
+- ❌ Requires manual key management
+- ❌ No automatic token refresh
+
+## Troubleshooting
+
+### Common Issues
+
+#### "Authorization required" (401)
+- Ensure Bearer token is included in header
+- Check API key is valid at [wandb.ai/settings](https://wandb.ai/settings)
+
+#### "Not Acceptable" (406)
+- Include `Accept: application/json, text/event-stream` header
+- Required for MCP Streamable HTTP transport
+
+#### "Missing session ID" (400)
+- Call `initialize` method first
+- Include session ID from response in subsequent requests
+
+#### "Invalid API key format"
+- W&B API keys are ~40 alphanumeric characters
+- Get a valid key from [wandb.ai/authorize](https://wandb.ai/authorize)
+
+### Testing Authentication
+
+```bash
+# Test with curl
+curl -X POST http://localhost:8080/mcp \
+  -H "Authorization: Bearer YOUR_WANDB_API_KEY" \
+  -H "Accept: application/json, text/event-stream" \
+  -H "Content-Type: application/json" \
+  -d '{"jsonrpc":"2.0","method":"initialize","params":{},"id":1}'
+```
+
+## Development Mode
+
+For local development only, you can disable authentication:
+
+```bash
+export MCP_AUTH_DISABLED=true
+```
+
+⚠️ **Warning**: Never use this in production or on public servers!

HUGGINGFACE_DEPLOYMENT.md

@@ -37,18 +37,18 @@ The application runs as a FastAPI server on port 7860 (HF Spaces default) with:
 
 ## Environment Variables
 
-Required in HF Spaces settings:
-```
-WANDB_API_KEY=your_api_key_here
-```
+No environment variables are required! The server works without any configuration.
 
-Optional:
+Optional (for server-side logging only):
 ```
-WANDB_ENTITY=your_wandb_entity
-MCP_LOGS_WANDB_PROJECT=wandb-mcp-logs
+WANDB_API_KEY=your_api_key_here  # Only if you want server-side Weave tracing
+WANDB_ENTITY=your_wandb_entity  # For server-side Weave tracing
+MCP_LOGS_WANDB_PROJECT=wandb-mcp-logs  # Project for server logs
 WEAVE_DISABLED=false  # Set to enable Weave tracing
 ```
 
+**Note**: For normal operation, users provide their own W&B API keys as Bearer tokens. No server configuration needed.
+
 ## Deployment Steps
 
 1. **Create a new Space on Hugging Face**
@@ -151,6 +151,153 @@ For MCP clients that support streamable HTTP:
 }
 ```
 
+## MCP Architecture & Key Learnings
+
+### Understanding MCP and FastMCP
+
+The Model Context Protocol (MCP) is a protocol for communication between AI assistants and external tools/services. Through our experimentation, we discovered several important aspects:
+
+#### 1. FastMCP Framework
+- **FastMCP** is a Python framework that simplifies MCP server implementation
+- It provides decorators (`@mcp.tool()`) for easy tool registration
+- Internally uses Starlette for HTTP handling
+- Supports multiple transports: stdio, SSE, and streamable HTTP
+
+#### 2. Streamable HTTP Transport
+The streamable HTTP transport (introduced in [MCP PR #206](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/206)) is the modern approach for remote MCP:
+
+- **Single endpoint** (`/mcp`) handles all communication
+- **Dual mode operation**:
+  - Regular POST requests for stateless operations
+  - SSE (Server-Sent Events) upgrade for streaming responses
+- **Key advantages**:
+  - Stateless servers possible (no persistent connections required)
+  - Better infrastructure compatibility ("just HTTP")
+  - Supports both request-response and streaming patterns
+
+#### 3. Implementation Patterns
+
+##### The HuggingFace Pattern
+Based on the [reference implementation](https://huggingface.co/spaces/Jofthomas/Multiple_mcp_fastapi_template), the correct pattern is:
+
+```python
+# Create MCP server
+mcp = FastMCP("server-name")
+
+# Register tools
+@mcp.tool()
+def my_tool(): ...
+
+# Get streamable HTTP app (returns Starlette app)
+mcp_app = mcp.streamable_http_app()
+
+# Mount in FastAPI
+app.mount("/", mcp_app)  # Note: mount at root, not at /mcp
+```
+
+##### Why Mount at Root?
+- `streamable_http_app()` creates internal routes at `/mcp`
+- Mounting at `/mcp` would create `/mcp/mcp` (double path)
+- Mounting at root gives us the clean `/mcp` endpoint
+
+#### 4. Session Management
+- FastMCP includes a `session_manager` for handling stateful operations
+- Use lifespan context manager to properly initialize/cleanup:
+  ```python
+  async with mcp.session_manager.run():
+      yield
+  ```
+
+#### 5. Response Format
+- MCP uses **Server-Sent Events (SSE)** for responses
+- Responses are prefixed with `event: message` and `data: `
+- JSON-RPC format for the actual message content
+- Example response:
+  ```
+  event: message
+  data: {"jsonrpc":"2.0","id":1,"result":{...}}
+  ```
+
+### Critical Implementation Details
+
+#### 1. Required Headers
+Clients MUST send:
+- `Content-Type: application/json`
+- `Accept: application/json, text/event-stream`
+
+Without the correct Accept header, the server returns a "Not Acceptable" error.
+
+#### 2. Lazy Loading Pattern
+To avoid initialization issues (e.g., API keys required at import time):
+```python
+# Instead of this:
+_service = Service()  # Fails if no API key
+
+# Use lazy loading:
+_service = None
+def get_service():
+    global _service
+    if _service is None:
+        _service = Service()
+    return _service
+```
+
+#### 3. Environment Setup for HF Spaces
+Critical for avoiding permission errors:
+```python
+os.environ["WANDB_CACHE_DIR"] = "/tmp/.wandb_cache"
+os.environ["HOME"] = "/tmp"
+```
+
+### Common Pitfalls & Solutions
+
+| Issue | Symptom | Solution |
+|-------|---------|----------|
+| Double path (`/mcp/mcp`) | 404 errors on `/mcp` | Mount streamable_http_app() at root (`/`) |
+| Missing Accept header | "Not Acceptable" error | Include `Accept: application/json, text/event-stream` |
+| Import-time API key errors | Server fails to start | Use lazy loading pattern |
+| Permission errors in HF Spaces | `mkdir /.cache: permission denied` | Set cache dirs to `/tmp` |
+| Can't access MCP methods | Methods not exposed | Use FastMCP's built-in decorators and methods |
+
+### Testing Strategy
+
+1. **Local Testing**: Always test with correct headers
+2. **Check Routes**: Verify mounting creates `/mcp` endpoint
+3. **Test Initialize First**: This method doesn't require session state
+4. **SSE Response Parsing**: Remember responses are SSE formatted, not plain JSON
+
+### Evolution of Our Implementation
+
+Our journey to the correct implementation went through several iterations:
+
+#### Attempt 1: Direct Protocol Implementation
+- **Approach**: Implement MCP protocol directly in FastAPI
+- **Issue**: Reinventing the wheel, not using FastMCP's built-in capabilities
+- **Learning**: FastMCP already handles the protocol complexity
+
+#### Attempt 2: Trying to Extract FastMCP's Internal App
+- **Approach**: Access FastMCP's internal FastAPI app via attributes
+- **Issue**: FastMCP doesn't expose its app in an accessible way
+- **Learning**: Need to use FastMCP's intended methods
+
+#### Attempt 3: Using http_app() Method
+- **Approach**: Try various methods like `http_app()`, `asgi_app()`, etc.
+- **Issue**: These methods either don't exist or don't work as expected
+- **Learning**: Documentation and examples are crucial
+
+#### Attempt 4: The Correct Pattern
+- **Approach**: Use `streamable_http_app()` following HuggingFace example
+- **Success**: Works perfectly when mounted at root
+- **Key Insight**: The example pattern exists for a reason - follow it!
+
+### Key Takeaways
+
+1. **Follow Existing Examples**: The HuggingFace example was the key to success
+2. **Understand the Protocol**: MCP uses SSE for good reasons (streaming, stateless option)
+3. **Lazy Loading is Critical**: Avoid initialization-time dependencies
+4. **Environment Matters**: HF Spaces has specific constraints (ports, permissions)
+5. **Test Incrementally**: Start with basic endpoints before complex operations
+
 ## Differences from Standard Deployment
 
 | Feature | Standard | HF Spaces |

app.py

@@ -22,7 +22,7 @@ os.environ["HOME"] = "/tmp"
 os.environ["WANDB_SILENT"] = "True"
 os.environ["WEAVE_SILENT"] = "True"
 
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
 from fastapi.responses import HTMLResponse, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from mcp.server.fastmcp import FastMCP
@@ -38,11 +38,7 @@ from wandb_mcp_server.server import (
 )
 
 # Import authentication
-from wandb_mcp_server.auth import (
-    mcp_auth_middleware,
-    create_resource_metadata_response,
-    MCPAuthConfig
-)
+from wandb_mcp_server.auth import mcp_auth_middleware
 
 # Configure logging
 logging.basicConfig(
@@ -124,11 +120,8 @@ 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))
+# Removed OAuth endpoints - only API key authentication is supported
+# See AUTH_README.md for details on why full OAuth isn't feasible
 
 @app.get("/health")
 async def health():

auth.example.env

@@ -0,0 +1,34 @@
+# W&B MCP Server Authentication Configuration Example
+# Copy this file to .env and add your values
+
+# ================================
+# API KEY AUTHENTICATION
+# ================================
+
+# For STDIO transport (required):
+# Get your key from: https://wandb.ai/authorize
+WANDB_API_KEY=your_wandb_api_key_here
+
+# For HTTP transport (optional):
+# Users provide their own API keys as Bearer tokens
+# No server configuration needed
+
+# ================================
+# OPTIONAL SETTINGS
+# ================================
+
+# W&B Entity (team or username) for server-side logging
+# WANDB_ENTITY=your_team_or_username
+
+# Project for MCP operation traces (if Weave tracing enabled)
+# MCP_LOGS_WANDB_PROJECT=wandb-mcp-logs
+
+# Disable Weave tracing (set to true to disable)
+# WEAVE_DISABLED=true
+
+# ================================
+# DEVELOPMENT ONLY
+# ================================
+
+# Disable authentication entirely (NEVER use in production!)
+# MCP_AUTH_DISABLED=true

src/wandb_mcp_server/auth.py

@@ -29,30 +29,29 @@ class MCPAuthConfig:
     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"
-        )
+    pass  # Simple config, no OAuth metadata needed
 
 
 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.
+    W&B API keys are typically 40 characters but we'll be permissive.
     """
-    if not token or len(token) < 20 or len(token) > 100:
+    if not token:
         return False
+    
+    # Strip any whitespace that might have been included
+    token = token.strip()
+    
+    # Be permissive - accept keys between 20 and 100 characters
+    # The actual W&B API will validate the exact format
+    if 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
 
 
@@ -77,26 +76,24 @@ async def validate_bearer_token(
             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}"'
+                "WWW-Authenticate": 'Bearer realm="W&B MCP"'
             }
         )
     
-    token = credentials.credentials
+    token = credentials.credentials.strip()  # Strip any whitespace
     
     # 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",
+            detail=f"Invalid W&B API key format. Got {len(token)} characters. "
+                   f"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}"'
+                "WWW-Authenticate": 'Bearer realm="W&B MCP", error="invalid_token"'
             }
         )
     
-    logger.debug("Bearer token validated successfully")
+    logger.debug(f"Bearer token validated successfully (length: {len(token)})")
     return token
 
 
@@ -124,14 +121,19 @@ async def mcp_auth_middleware(request: Request, call_next):
         authorization = request.headers.get("Authorization", "")
         credentials = None
         if authorization.startswith("Bearer "):
+            # Remove "Bearer " prefix and strip any whitespace
+            token = authorization[7:].strip()
             credentials = HTTPAuthorizationCredentials(
                 scheme="Bearer",
-                credentials=authorization[7:]  # Remove "Bearer " prefix
+                credentials=token
             )
         
         # Validate and get the W&B API key
         wandb_api_key = await validate_bearer_token(credentials, config)
         
+        # Make sure the key is clean (no extra whitespace or encoding issues)
+        wandb_api_key = wandb_api_key.strip()
+        
         # 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
@@ -139,8 +141,16 @@ async def mcp_auth_middleware(request: Request, call_next):
         # 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")
+        
+        # Set the clean API key
         os.environ["WANDB_API_KEY"] = wandb_api_key
         
+        # Debug logging (remove in production)
+        logger.info(f"Auth middleware: Set WANDB_API_KEY with length={len(wandb_api_key)}, "
+                   f"first_6={wandb_api_key[:6] if len(wandb_api_key) >= 6 else 'N/A'}..., "
+                   f"last_4={wandb_api_key[-4:] if len(wandb_api_key) >= 4 else 'N/A'}, "
+                   f"is_40_chars={len(wandb_api_key) == 40}")
+        
         try:
             # Continue processing
             response = await call_next(request)
@@ -166,23 +176,9 @@ async def mcp_auth_middleware(request: Request, call_next):
             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}"'
+                "WWW-Authenticate": 'Bearer realm="W&B MCP"'
             }
         )
 
 
-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",
-    }
+# OAuth-related functions removed - see AUTH_README.md for details

src/wandb_mcp_server/mcp_tools/count_traces.py

@@ -180,6 +180,11 @@ def count_traces(
     if not api_key:
         logger.error("WANDB_API_KEY not found in environment variables.")
         raise ValueError("WANDB_API_KEY is required to query Weave traces count.")
+    
+    # Debug logging to diagnose API key issues
+    logger.debug(f"Using W&B API key: length={len(api_key)}, "
+                f"first_6={api_key[:6] if len(api_key) >= 6 else 'N/A'}..., "
+                f"last_4={api_key[-4:] if len(api_key) >= 4 else 'N/A'}")
 
     request_body: Dict[str, Any] = {"project_id": project_id}
     filter_payload: Dict[
@@ -296,6 +301,11 @@ def count_traces(
         if response.status_code != 200:
             error_msg = f"Error querying Weave trace count: {response.status_code} - {response.text}"
             logger.error(error_msg)
+            # Log API key info for debugging
+            logger.error(f"API key info: length={len(api_key)}, is_40_chars={len(api_key) == 40}")
+            if "40 characters" in response.text:
+                logger.error(f"W&B requires exactly 40 character API keys. Current key has {len(api_key)} characters.")
+                logger.error(f"Key preview: {api_key[:8]}...{api_key[-4:] if len(api_key) >= 12 else ''}")
             # Log request body for easier debugging on error
             logger.debug(f"Failed request body: {json.dumps(request_body)}")
             raise Exception(error_msg)