fix authetication

app.py

@@ -99,11 +99,54 @@ if __name__ == "__main__":
         from starlette.requests import Request
         from database.api_keys import generate_api_key as db_generate_api_key
 
-        # NOTE: Custom middleware for FastMCP has known limitations (GitHub issue #817)
-        # Headers/state set in middleware are NOT accessible inside FastMCP tools.
-        # For local development, use SKIP_AUTH=true with ENV=development in .env
-        # For production, API key validation happens via extract_api_key_from_request() in server.py
-        logger.info("[Auth] Using FastMCP built-in request handling for authentication")
+        # Import session auth store from server.py
+        from server import store_session_api_key, _session_auth_store
+
+        # =====================================================================
+        # API KEY CAPTURE MIDDLEWARE (Starlette)
+        # This simple middleware captures the API key from SSE connections
+        # and stores it for later lookup by session_id during tool calls.
+        #
+        # Flow:
+        # 1. SSE connects: /sse?api_key=xxx → Store api_key temporarily
+        # 2. Messages come: /messages/?session_id=yyy → Link session to api_key
+        # 3. FastMCP middleware retrieves api_key from session store
+        # =====================================================================
+        from starlette.middleware.base import BaseHTTPMiddleware
+
+        class ApiKeyCaptureMiddleware(BaseHTTPMiddleware):
+            """Captures API key from SSE and links it to session_id"""
+
+            async def dispatch(self, request: Request, call_next):
+                api_key = request.query_params.get('api_key')
+                session_id = request.query_params.get('session_id')
+                path = request.url.path
+
+                # SSE connection with api_key - store for later linking
+                if api_key and path == '/sse':
+                    _session_auth_store['_pending_api_key'] = api_key
+                    logger.info(f"[ApiKeyCapture] SSE with api_key {api_key[:15]}...")
+
+                # Messages request - link session to api_key
+                if session_id and path.startswith('/messages'):
+                    # Check if session already has api_key
+                    if session_id not in _session_auth_store:
+                        # Get from URL if present
+                        if api_key:
+                            store_session_api_key(session_id, api_key)
+                            logger.info(f"[ApiKeyCapture] Session {session_id[:12]}... linked via URL")
+                        # Otherwise use pending key from SSE connection
+                        elif '_pending_api_key' in _session_auth_store:
+                            pending = _session_auth_store.pop('_pending_api_key')
+                            store_session_api_key(session_id, pending)
+                            logger.info(f"[ApiKeyCapture] Session {session_id[:12]}... linked via pending")
+
+                return await call_next(request)
+
+        # Add the capture middleware - must be added before mcp.run()
+        # We'll add it via mcp.settings or directly to the ASGI app
+        logger.info("[Auth] ApiKeyCaptureMiddleware configured")
+        logger.info("[Auth] FastMCP ApiKeyAuthMiddleware handles authentication in tools")
 
         @mcp.custom_route("/", methods=["GET"])
         async def landing_page(request):
@@ -430,12 +473,38 @@ if __name__ == "__main__":
         logger.info("[OK] API key generation page added at /generate-key")
         logger.info("[OK] MCP SSE endpoint available at /sse")
 
-        # Run MCP server with SSE transport (includes both /sse and custom routes)
-        mcp.run(
-            transport="sse",
-            host=HF_SPACE_HOST,
-            port=HF_SPACE_PORT
-        )
+        # Try to use SSE app with custom middleware, fallback to standard mcp.run()
+        try:
+            from starlette.middleware import Middleware
+            import uvicorn
+
+            # Create middleware list
+            custom_middleware = [
+                Middleware(ApiKeyCaptureMiddleware)
+            ]
+
+            # Get SSE app with custom middleware
+            sse_app = mcp.sse_app(custom_middleware=custom_middleware)
+            logger.info("[Auth] Running with ApiKeyCaptureMiddleware via sse_app()")
+
+            # Run with uvicorn
+            uvicorn.run(
+                sse_app,
+                host=HF_SPACE_HOST,
+                port=HF_SPACE_PORT,
+                log_level="info"
+            )
+        except (AttributeError, TypeError) as e:
+            # Fallback: FastMCP version doesn't support custom_middleware
+            logger.warning(f"[Auth] custom_middleware not supported: {e}")
+            logger.warning("[Auth] Falling back to standard mcp.run()")
+
+            # Run MCP server with SSE transport (includes both /sse and custom routes)
+            mcp.run(
+                transport="sse",
+                host=HF_SPACE_HOST,
+                port=HF_SPACE_PORT
+            )
 
     except Exception as e:
         logger.error(f"Failed to start server: {e}")

server.py

@@ -115,6 +115,81 @@ mcp = FastMCP(
     version="1.0.0"
 )
 
+# ============================================================================
+# FASTMCP AUTHENTICATION MIDDLEWARE
+# Uses FastMCP's native middleware system to properly pass user context to tools
+# Reference: https://gelembjuk.com/blog/post/authentication-remote-mcp-server-python/
+# ============================================================================
+try:
+    from fastmcp.server.middleware import Middleware, MiddlewareContext
+    from fastmcp.exceptions import ToolError
+
+    class ApiKeyAuthMiddleware(Middleware):
+        """
+        FastMCP middleware for API key authentication.
+        Validates API key from request query params and injects user info into context.
+        """
+
+        async def on_call_tool(self, context: MiddlewareContext, call_next):
+            """Intercept tool calls to validate authentication"""
+            try:
+                from fastmcp.server.dependencies import get_http_request
+                request = get_http_request()
+
+                # Get API key from query params
+                api_key = request.query_params.get('api_key')
+
+                # Also check session store if no direct api_key
+                if not api_key:
+                    session_id = request.query_params.get('session_id')
+                    if session_id:
+                        api_key = get_api_key_from_session(session_id)
+
+                if api_key:
+                    # Validate API key and get user info
+                    from database.api_keys import verify_api_key
+                    user_info = verify_api_key(api_key)
+
+                    if user_info:
+                        # Store user info in context for tools to access
+                        context.fastmcp_context.set_state("user_id", user_info['user_id'])
+                        context.fastmcp_context.set_state("user_email", user_info['email'])
+                        context.fastmcp_context.set_state("user_scopes", user_info.get('scopes', []))
+                        context.fastmcp_context.set_state("user_name", user_info.get('name', ''))
+                        logger.info(f"✅ Auth middleware: User {user_info['email']} authenticated")
+                        return await call_next(context)
+                    else:
+                        logger.warning(f"❌ Auth middleware: Invalid API key {api_key[:10]}...")
+
+                # Check SKIP_AUTH for development
+                env = os.getenv("ENV") or os.getenv("ENVIRONMENT", "production")
+                skip_auth = os.getenv("SKIP_AUTH", "false").lower() == "true"
+
+                if skip_auth and env.lower() != "production":
+                    # Development mode - use dev user
+                    context.fastmcp_context.set_state("user_id", "dev-user")
+                    context.fastmcp_context.set_state("user_email", "[email protected]")
+                    context.fastmcp_context.set_state("user_scopes", ["admin"])
+                    context.fastmcp_context.set_state("user_name", "Development User")
+                    logger.warning(f"⚠️ Auth middleware: SKIP_AUTH enabled, using dev user")
+                    return await call_next(context)
+
+                # No valid authentication
+                raise ToolError("Authentication required. Please provide a valid API key.")
+
+            except ImportError:
+                # get_http_request not available (stdio transport)
+                logger.debug("Auth middleware: No HTTP request available")
+                return await call_next(context)
+
+    # Register the middleware with FastMCP
+    mcp.add_middleware(ApiKeyAuthMiddleware())
+    logger.info("[Auth] FastMCP ApiKeyAuthMiddleware registered")
+
+except ImportError as e:
+    logger.warning(f"[Auth] Could not import FastMCP middleware: {e}")
+    logger.warning("[Auth] Falling back to per-tool authentication")
+
 # Initialize shared services
 logger.info("Initializing FleetMind MCP Server...")
 geocoding_service = GeocodingService()
@@ -131,18 +206,40 @@ except Exception as e:
 # AUTHENTICATION - API KEY SYSTEM
 # ============================================================================
 
-def get_authenticated_user():
+def get_authenticated_user(ctx=None):
     """
-    Get authenticated user by validating API key from request context.
+    Get authenticated user from multiple sources:
+    1. FastMCP Context state (set by middleware) - PREFERRED
+    2. HTTP request API key extraction (fallback)
+    3. Environment variable (fallback for testing)
+    4. SKIP_AUTH bypass (development only)
+
+    Args:
+        ctx: Optional FastMCP Context object from tool function
 
     Returns:
         User info dict with user_id, email, scopes, name or None if not authenticated
     """
     try:
-        # METHOD 1: Extract API key from current HTTP request
+        # METHOD 1: Get user from FastMCP Context state (set by middleware)
+        # This is the preferred method when using FastMCP middleware
+        if ctx is not None:
+            try:
+                user_id = ctx.get_state("user_id")
+                if user_id:
+                    return {
+                        'user_id': user_id,
+                        'email': ctx.get_state("user_email") or "",
+                        'scopes': ctx.get_state("user_scopes") or [],
+                        'name': ctx.get_state("user_name") or ""
+                    }
+            except Exception:
+                pass  # Context state not available, try other methods
+
+        # METHOD 2: Extract API key from current HTTP request
         api_key = extract_api_key_from_request()
 
-        # METHOD 2: Fallback to environment variable for testing
+        # METHOD 3: Fallback to environment variable for testing
         if not api_key:
             api_key = os.getenv("FLEETMIND_API_KEY")
             if api_key:
@@ -159,7 +256,7 @@ def get_authenticated_user():
                 logger.warning(f"❌ Invalid API key: {api_key[:10]}...")
                 return None
 
-        # METHOD 3: Development bypass mode (local testing only)
+        # METHOD 4: Development bypass mode (local testing only)
         # SECURITY: Only allow SKIP_AUTH in development environments
         # Check both ENV and ENVIRONMENT variables for compatibility
         env = os.getenv("ENV") or os.getenv("ENVIRONMENT", "production")