separate client and server api key for http and weave tracing

HUGGINGFACE_DEPLOYMENT.md

@@ -39,15 +39,7 @@ The application runs as a FastAPI server on port 7860 (HF Spaces default) with:
 
 No environment variables are required! The server works without any configuration.
 
-Optional (for server-side logging only):
-```
-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.
+**Note**: Users provide their own W&B API keys as Bearer tokens. No server configuration needed.
 
 ## Deployment Steps
 

auth.example.env

@@ -1,30 +1,20 @@
-# W&B MCP Server Authentication Configuration Example
+# W&B MCP Server Configuration Example
 # Copy this file to .env and add your values
 
 # ================================
-# API KEY AUTHENTICATION
+# FOR STDIO TRANSPORT
 # ================================
 
-# For STDIO transport (required):
+# Required for STDIO transport only
 # 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
+# FOR HTTP TRANSPORT
 # ================================
 
-# 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
+# No configuration needed!
+# Users provide their own API keys as Bearer tokens
 
 # ================================
 # DEVELOPMENT ONLY

src/wandb_mcp_server/auth.py

@@ -135,32 +135,21 @@ async def mcp_auth_middleware(request: Request, call_next):
         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
         
-        # 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
+        # Set the API key for this request
+        # Note: We don't restore the original value because with streaming responses,
+        # the tool execution happens after call_next returns. Each request sets its own 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}")
+        # Debug logging
+        logger.debug(f"Auth middleware: Set WANDB_API_KEY with length={len(wandb_api_key)}, "
+                    f"is_40_chars={len(wandb_api_key) == 40}")
+        
+        # Continue processing without restoring the env var
+        # Each request will set its own API key
+        response = await call_next(request)
         
-        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:

src/wandb_mcp_server/mcp_tools/count_traces.py

@@ -176,6 +176,7 @@ def count_traces(
     """
     project_id = f"{entity_name}/{project_name}"
 
+    # Get API key from environment (set by auth middleware for HTTP, or by user for STDIO)
     api_key = os.environ.get("WANDB_API_KEY")
     if not api_key:
         logger.error("WANDB_API_KEY not found in environment variables.")

src/wandb_mcp_server/mcp_tools/create_report.py

@@ -130,6 +130,7 @@ def create_report(
         processed_plots_html = None
 
     try:
+        # W&B will use WANDB_API_KEY from environment
         wandb.init(
             entity=entity_name, project=project_name, job_type="mcp_report_creation"
         )

src/wandb_mcp_server/mcp_tools/list_wandb_entities_projects.py

@@ -70,6 +70,7 @@ def list_entity_projects(entity: str | None = None) -> dict[str, list[dict[str,
             - tags: List of project tags
     """
     # Initialize wandb API
+    # Will use WANDB_API_KEY from environment (set by auth middleware or user)
     api = wandb.Api()
 
     # Merge entity and teams into a single list

src/wandb_mcp_server/mcp_tools/query_wandb_gql.py

@@ -591,7 +591,8 @@ def query_paginated_wandb_gql(
     api = None
     limit_key = None
     try:
-        api = wandb.Api()
+        # Use API key from environment (set by auth middleware for HTTP, or by user for STDIO)
+        api = wandb.Api()  # Will use WANDB_API_KEY from environment
         logger.info(
             "--- Inside query_paginated_wandb_gql: Step 0: Execute Initial Query ---"
         )

src/wandb_mcp_server/weave_api/service.py

@@ -62,6 +62,7 @@ class TraceService:
 
     def __init__(
         self,
+        api_key: Optional[str] = None,
         server_url: Optional[str] = None,
         retries: int = 3,
         timeout: int = 10,
@@ -69,18 +70,26 @@ class TraceService:
         """Initialize the TraceService.
 
         Args:
+            api_key: W&B API key. If not provided, uses WANDB_API_KEY env var.
             server_url: Weave API server URL. Defaults to 'https://trace.wandb.ai'.
             retries: Number of retries for failed requests.
             timeout: Request timeout in seconds.
         """
-        # Call get_server_args() to ensure API key is loaded from .netrc or env var
-        # and a warning is logged by get_server_args if no key is found.
-        server_config = get_server_args()
+        # If no API key provided, try to get from environment
+        if api_key is None:
+            import os
+            # Try to get from environment (set by auth middleware for HTTP or user for STDIO)
+            api_key = os.environ.get("WANDB_API_KEY")
+            
+            # If still no key, try get_server_args as fallback
+            if not api_key:
+                server_config = get_server_args()
+                api_key = server_config.wandb_api_key
 
         # Pass the resolved API key to WeaveApiClient.
-        # If server_config.wandb_api_key is None or "", WeaveApiClient will raise its ValueError.
+        # If api_key is None or "", WeaveApiClient will raise its ValueError.
         self.client = WeaveApiClient(
-            api_key=server_config.wandb_api_key,
+            api_key=api_key,
             server_url=server_url,
             retries=retries,
             timeout=timeout,