fix(mcp): add async support and test_request_context to mcp_auth_hook

aminghadersohi · aminghadersohi · commit 15c9c58f2bdf · 2025-11-13T00:00:43.000-05:00
This commit extends PR apache#36013 to fully resolve the 'Working outside of application context' error for async MCP tools. Changes: - Add async/sync detection using inspect.iscoroutinefunction() - Create separate async_wrapper and sync_wrapper functions - Add test_request_context() nested inside app_context() for g object access - Support await for async tools, regular call for sync tools The original PR added app_context() which partially fixed the issue, but async tools like list_charts still failed because: 1. The decorator didn't await async functions 2. Only app_context() was used, but g object requires test_request_context() Both contexts are needed: - app_context() provides current_app, config, db - test_request_context() provides g, request Tested with get_instance_info and list_charts - both work without context errors.
diff --git a/superset/mcp_service/auth.py b/superset/mcp_service/auth.py
@@ -16,35 +16,15 @@
 # under the License.
 
 """
-Authentication and authorization hooks for MCP tools.
-
-This module provides a flexible authentication framework that can be customized
-by deployers through configuration and the FastMCP auth provider system.
-
-Default Behavior:
-- Uses configured user (MCP_DEV_USERNAME) for authentication
-- Leverages Superset's existing RBAC and security manager
-- Provides dataset-level access control via has_dataset_access()
-
-Customization Options for Deployers:
-1. JWT Authentication: Configure BearerAuthProvider in mcp_config.py with your
-   JWT issuer, audience, and public key/JWKS URI. See create_default_mcp_auth_factory()
-   for reference implementation.
-
-2. User Resolution: Override get_user_from_request() to extract user from JWT
-   claims, headers, or other authentication mechanisms.
-
-3. Permission Checks: Add custom permission logic in mcp_auth_hook before tool
-   execution. Use Superset's security_manager for role-based checks.
-
-4. Audit Logging: Extend mcp_auth_hook to log tool access, integrate with your
-   SIEM/logging infrastructure, or track usage metrics.
-
-5. Rate Limiting: Integrate rate limiting middleware in the mcp_auth_hook wrapper
-   based on user, IP, or tool-specific quotas.
-
-The hook pattern allows production deployments to layer additional security
-without modifying tool implementations.
+Minimal authentication hooks for MCP tools.
+This is a placeholder implementation that provides basic user context.
+
+Future enhancements (to be added in separate PRs):
+- JWT token authentication and validation
+- User impersonation support
+- Permission checking with scopes
+- Comprehensive audit logging
+- Field-level permissions
 """
 
 import logging
@@ -69,18 +49,6 @@ def get_user_from_request() -> User:
     Default implementation uses MCP_DEV_USERNAME from configuration for
     development and single-user deployments.
 
-    Deployers can customize this function to:
-    - Extract user from JWT claims (when using BearerAuthProvider)
-    - Implement user impersonation based on request headers
-    - Add fallback authentication for service accounts
-    - Integrate with external identity providers
-
-    Example JWT customization:
-        from flask import request
-        token = request.headers.get('Authorization', '').replace('Bearer ', '')
-        # Parse JWT and extract user from claims
-        # Then query user from database as shown below
-
     Returns:
         User object with roles and groups eagerly loaded
 
@@ -93,7 +61,6 @@ def get_user_from_request() -> User:
     from superset.extensions import db
 
     # Use configured username for authentication
-    # Deployers: Replace this with your authentication logic
     username = current_app.config.get("MCP_DEV_USERNAME")
 
     if not username:
@@ -154,97 +121,142 @@ def mcp_auth_hook(tool_func: F) -> F:
     """
     Authentication and authorization decorator for MCP tools.
 
-    This hook ensures proper Flask app context and user authentication for all
-    MCP tool calls. Deployers can extend this decorator to add:
-
-    - Permission checks: Validate user has specific roles or tool access rights
-    - JWT scope validation: Check OAuth scopes match required tool permissions
-    - Audit logging: Track tool usage, parameters, and results for compliance
-    - Rate limiting: Enforce per-user or per-tool rate limits
-    - Custom middleware: Inject organization-specific security policies
-
-    Default behavior:
-    1. Ensures Flask app context is active (critical for stateless_http mode)
-    2. Authenticates user via get_user_from_request()
-    3. Sets g.user for Superset's security manager integration
-    4. Manages database session lifecycle
-
-    Example extensions:
-        # Add permission check before tool execution:
-        if tool_func.__name__ in ['execute_sql'] and not user.is_admin():
-            raise PermissionError("SQL execution requires admin role")
-
-        # Add audit logging after successful execution:
-        audit_log.info({
-            'user': user.username,
-            'tool': tool_func.__name__,
-            'timestamp': datetime.now(),
-            'result': 'success'
-        })
+    This decorator assumes Flask application context and g.user
+    have already been set by WorkspaceContextMiddleware.
+
+    Supports both sync and async tool functions.
+
+    TODO (future PR): Add permission checking
+    TODO (future PR): Add JWT scope validation
+    TODO (future PR): Add comprehensive audit logging
     """
+    import asyncio
     import functools
+    import inspect
 
-    @functools.wraps(tool_func)
-    def wrapper(*args: Any, **kwargs: Any) -> Any:
-        from superset.extensions import db
-        from superset.mcp_service.flask_singleton import get_flask_app
-
-        # Get Flask app instance
-        app = get_flask_app()
-
-        # Ensure Flask app context is active
-        # This is critical for stateless_http mode where context isn't preserved
-        with app.app_context():
-            # Get user and set Flask context OUTSIDE try block
-            user = get_user_from_request()
-
-            # Force load relationships NOW while session is definitely active
-            # This prevents lazy-loading errors after session closure
-            user_roles = user.roles  # noqa: F841
-            if hasattr(user, "groups"):
-                user_groups = user.groups  # noqa: F841
-
-            g.user = user
-
-            try:
-                # Deployers: Add custom permission checks here
-                # Example: Check tool-specific permissions, validate JWT scopes,
-                # or enforce organizational policies
-
-                logger.debug(
-                    "MCP tool call: user=%s, tool=%s", user.username, tool_func.__name__
-                )
-
-                result = tool_func(*args, **kwargs)
-
-                # Deployers: Add audit logging or post-execution hooks here
-                # Example: Log successful tool execution, track metrics,
-                # or trigger downstream workflows
-
-                return result
-
-            except Exception:
-                # On error, rollback and cleanup session
-                # pylint: disable=consider-using-transaction
-                try:
-                    db.session.rollback()
-                    db.session.remove()
-                except Exception as cleanup_error:
-                    # Log cleanup failure but re-raise original exception
-                    # We want to preserve the original error, not the cleanup error
-                    logger.warning(
-                        "Error cleaning up session after exception: %s", cleanup_error
-                    )
-                raise
-
-            finally:
-                # Only rollback if session is still active (no exception occurred)
-                # Do NOT call remove() on success to avoid detaching user
-                try:
-                    if db.session.is_active:
-                        # pylint: disable=consider-using-transaction
-                        db.session.rollback()
-                except Exception as e:
-                    logger.warning("Error in finally block: %s", e)
+    # Check if the tool function is async
+    is_async = inspect.iscoroutinefunction(tool_func)
 
-    return wrapper  # type: ignore[return-value]
+    if is_async:
+
+        @functools.wraps(tool_func)
+        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+            from superset.extensions import db
+            from superset.mcp_service.flask_singleton import get_flask_app
+
+            # Get Flask app instance
+            app = get_flask_app()
+
+            # Ensure Flask app AND request context are active
+            # app_context() provides current_app, config, db
+            # test_request_context() provides g, request
+            with app.app_context():
+                with app.test_request_context():
+                    # Get user from config
+                    user = get_user_from_request()
+
+                    # Validate user has necessary relationships loaded
+                    # (Force access to ensure they're loaded if lazy)
+                    user_roles = user.roles  # noqa: F841
+                    if hasattr(user, "groups"):
+                        user_groups = user.groups  # noqa: F841
+
+                    g.user = user
+
+                    try:
+                        # TODO: Add permission checks here in future PR
+                        # TODO: Add audit logging here in future PR
+
+                        logger.debug(
+                            "MCP tool call: user=%s, tool=%s",
+                            user.username,
+                            tool_func.__name__,
+                        )
+
+                        result = await tool_func(*args, **kwargs)
+
+                        return result
+
+                    except Exception:
+                        # On error, rollback and cleanup session
+                        # pylint: disable=consider-using-transaction
+                        try:
+                            db.session.rollback()
+                            db.session.remove()
+                        except Exception as e:
+                            logger.warning("Error cleaning up session after exception: %s", e)
+                        raise
+
+                    finally:
+                        # Only rollback if session is still active (no exception occurred)
+                        # Do NOT call remove() on success to avoid detaching user
+                        try:
+                            if db.session.is_active:
+                                # pylint: disable=consider-using-transaction
+                                db.session.rollback()
+                        except Exception as e:
+                            logger.warning("Error in finally block: %s", e)
+
+        return async_wrapper  # type: ignore[return-value]
+
+    else:
+
+        @functools.wraps(tool_func)
+        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+            from superset.extensions import db
+            from superset.mcp_service.flask_singleton import get_flask_app
+
+            # Get Flask app instance
+            app = get_flask_app()
+
+            # Ensure Flask app AND request context are active
+            # app_context() provides current_app, config, db
+            # test_request_context() provides g, request
+            with app.app_context():
+                with app.test_request_context():
+                    # Get user from config
+                    user = get_user_from_request()
+
+                    # Validate user has necessary relationships loaded
+                    # (Force access to ensure they're loaded if lazy)
+                    user_roles = user.roles  # noqa: F841
+                    if hasattr(user, "groups"):
+                        user_groups = user.groups  # noqa: F841
+
+                    g.user = user
+
+                    try:
+                        # TODO: Add permission checks here in future PR
+                        # TODO: Add audit logging here in future PR
+
+                        logger.debug(
+                            "MCP tool call: user=%s, tool=%s",
+                            user.username,
+                            tool_func.__name__,
+                        )
+
+                        result = tool_func(*args, **kwargs)
+
+                        return result
+
+                    except Exception:
+                        # On error, rollback and cleanup session
+                        # pylint: disable=consider-using-transaction
+                        try:
+                            db.session.rollback()
+                            db.session.remove()
+                        except Exception as e:
+                            logger.warning("Error cleaning up session after exception: %s", e)
+                        raise
+
+                    finally:
+                        # Only rollback if session is still active (no exception occurred)
+                        # Do NOT call remove() on success to avoid detaching user
+                        try:
+                            if db.session.is_active:
+                                # pylint: disable=consider-using-transaction
+                                db.session.rollback()
+                        except Exception as e:
+                            logger.warning("Error in finally block: %s", e)
+
+        return sync_wrapper  # type: ignore[return-value]
