You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by ju...@apache.org on 2026/08/14 20:28:07 UTC
(superset) branch master updated: fix(mcp): populate user_id in MCP audit logs (#42767)
This is an automated email from the ASF dual-hosted git repository.
justinpark pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new f4587218dd1 fix(mcp): populate user_id in MCP audit logs (#42767)
f4587218dd1 is described below
commit f4587218dd19d046c3e4d00063e7d27f8a2ed354
Author: JUST.in DO IT <ju...@airbnb.com>
AuthorDate: Fri Aug 14 13:28:01 2026 -0700
fix(mcp): populate user_id in MCP audit logs (#42767)
---
superset/mcp_service/auth.py | 54 ++++++++---
superset/mcp_service/middleware.py | 85 +++++++++++------
.../mcp_service/test_auth_user_resolution.py | 63 +++++++++++++
.../mcp_service/test_middleware_logging.py | 103 +++++++++++++++++++++
4 files changed, 266 insertions(+), 39 deletions(-)
diff --git a/superset/mcp_service/auth.py b/superset/mcp_service/auth.py
index 95617d03dd2..2a663d2faa2 100644
--- a/superset/mcp_service/auth.py
+++ b/superset/mcp_service/auth.py
@@ -46,6 +46,7 @@ Configuration:
import logging
from contextlib import AbstractContextManager, contextmanager
+from contextvars import ContextVar
from typing import Any, Callable, cast, Generator, TYPE_CHECKING, TypeAlias, TypeVar
from flask import current_app, g, has_app_context, has_request_context
@@ -79,6 +80,15 @@ F = TypeVar("F", bound=Callable[..., Any])
logger = logging.getLogger(__name__)
+# The resolved user's id for the current tool call, set by
+# _setup_user_context(). g.user is only valid for the lifetime of the
+# per-call app context pushed by _get_app_context_manager() (see its
+# docstring) and is gone by the time LoggingMiddleware's on_call_tool/
+# on_message finally-blocks run after call_next() returns. This ContextVar
+# survives that context pop (mirrors _mcp_call_id_var in middleware.py) so
+# audit logging can attribute the call to the right user.
+_mcp_user_id_var: ContextVar[int | None] = ContextVar("mcp_user_id", default=None)
+
# An MCP request resolves to a real DB ``User`` or, for embedded guests, a
# ``GuestUser`` (an AnonymousUserMixin, not a ``User`` subclass). Both are valid
# authenticated principals for tool execution.
@@ -879,9 +889,9 @@ def _assert_user_active(user: MCPUser | None) -> None:
)
-def _setup_user_context() -> MCPUser | None:
+def _resolve_user_with_retry() -> MCPUser | None:
"""
- Set up user context for MCP tool execution.
+ Resolve the current user, retrying once on a stale DB connection.
Includes retry logic for stale database connections (e.g., SSL dropped
by proxy/load balancer after idle periods). On OperationalError, the
@@ -890,14 +900,6 @@ def _setup_user_context() -> MCPUser | None:
Returns:
User object with roles and groups loaded, or None if no Flask context
"""
- # Clear stale g.user to prevent user impersonation across
- # tool calls when no per-request middleware refreshes it.
- # Only clear in app-context-only mode; preserve g.user when
- # a request context is active (external middleware set it).
-
- if not has_request_context():
- g.pop("user", None)
-
from sqlalchemy.exc import OperationalError
user = None # Ensure defined before loop in case of unexpected exit
@@ -914,7 +916,7 @@ def _setup_user_context() -> MCPUser | None:
if hasattr(user, "groups"):
user_groups = user.groups # noqa: F841
- break
+ return user
except RuntimeError as e:
# No Flask application context (e.g., prompts before middleware runs)
if "application context" in str(e):
@@ -948,8 +950,38 @@ def _setup_user_context() -> MCPUser | None:
g.pop("user", None)
raise
+ return user
+
+
+def _setup_user_context() -> MCPUser | None:
+ """
+ Set up user context for MCP tool execution.
+
+ Returns:
+ User object with roles and groups loaded, or None if no Flask context
+ """
+ # Clear stale g.user to prevent user impersonation across
+ # tool calls when no per-request middleware refreshes it.
+ # Only clear in app-context-only mode; preserve g.user when
+ # a request context is active (external middleware set it).
+
+ if not has_request_context():
+ g.pop("user", None)
+ # Clear any user_id left over from a previous call in this context
+ # (e.g. sequential calls sharing one asyncio task) so a failed/
+ # unauthenticated lookup below doesn't inherit a stale value.
+ _mcp_user_id_var.set(None)
+
+ user = _resolve_user_with_retry()
+ if user is None:
+ return None
+
_assert_user_active(user)
g.user = user
+ # GuestUser (embedded auth) has no numeric id; leave the ContextVar
+ # cleared (already reset above) rather than raise.
+ if (user_id := getattr(user, "id", None)) is not None:
+ _mcp_user_id_var.set(user_id)
return user
diff --git a/superset/mcp_service/middleware.py b/superset/mcp_service/middleware.py
index 9fbb5f3f9a2..de47a7c0dc7 100644
--- a/superset/mcp_service/middleware.py
+++ b/superset/mcp_service/middleware.py
@@ -41,6 +41,7 @@ from superset.exceptions import SupersetException, SupersetSecurityException
from superset.extensions import event_logger, stats_logger_manager
from superset.mcp_service.auth import (
_get_app_context_manager,
+ _mcp_user_id_var,
get_user_from_request,
is_tool_visible_to_current_user,
MCPNoAuthSourceError,
@@ -625,6 +626,24 @@ class LoggingMiddleware(Middleware):
success = False
raise
finally:
+ # user_id was captured before call_next() ran the tool, i.e.
+ # before the @tool auth decorator (superset/mcp_service/auth.py)
+ # resolves the user. It sets g.user on a per-call app context
+ # that _get_app_context_manager() pushes and pops around the
+ # tool's execution (see its docstring), so g.user/get_user_id()
+ # are back to their pre-call state by the time we get here —
+ # re-reading get_user_id() would still yield the stale value.
+ # _mcp_user_id_var is a plain ContextVar (not tied to that Flask
+ # app-context lifecycle) that _setup_user_context() sets before
+ # the context pops, so it survives to this point.
+ resolved_user_id = _mcp_user_id_var.get(None)
+ if resolved_user_id is not None:
+ user_id = resolved_user_id
+ # Reset so a later on_call_tool/on_message in the same asyncio
+ # task (e.g. an unprotected tool or resource/prompt read that
+ # never calls _setup_user_context()) doesn't inherit this call's
+ # resolved user id.
+ _mcp_user_id_var.set(None)
duration_ms = int((time.time() - start_time) * 1000)
self._log_call_tool_result(
context=context,
@@ -666,34 +685,44 @@ class LoggingMiddleware(Middleware):
self._extract_context_info(context)
)
try:
- with _get_app_context_manager():
- event_logger.log(
- user_id=user_id,
- action="mcp_message",
- dashboard_id=dashboard_id,
- duration_ms=None,
- slice_id=slice_id,
- referrer=None,
- curated_payload={
- "tool": getattr(context.message, "name", None),
- "agent_id": agent_id,
- "params": _sanitize_params(params),
- "method": context.method,
- "dashboard_id": dashboard_id,
- "slice_id": slice_id,
- "dataset_id": dataset_id,
- },
- )
- except Exception as log_error: # noqa: BLE001
- logger.warning("Failed to log mcp_message event: %s", log_error)
- logger.info(
- "MCP message: tool=%s, agent_id=%s, user_id=%s, method=%s",
- getattr(context.message, "name", None),
- agent_id,
- user_id,
- context.method,
- )
- return await call_next(context)
+ return await call_next(context)
+ finally:
+ # See the matching comment in on_call_tool: g.user/get_user_id()
+ # are stale here because the per-call app context has already
+ # been popped. _mcp_user_id_var survives it.
+ resolved_user_id = _mcp_user_id_var.get(None)
+ if resolved_user_id is not None:
+ user_id = resolved_user_id
+ # See the matching reset in on_call_tool.
+ _mcp_user_id_var.set(None)
+ try:
+ with _get_app_context_manager():
+ event_logger.log(
+ user_id=user_id,
+ action="mcp_message",
+ dashboard_id=dashboard_id,
+ duration_ms=None,
+ slice_id=slice_id,
+ referrer=None,
+ curated_payload={
+ "tool": getattr(context.message, "name", None),
+ "agent_id": agent_id,
+ "params": _sanitize_params(params),
+ "method": context.method,
+ "dashboard_id": dashboard_id,
+ "slice_id": slice_id,
+ "dataset_id": dataset_id,
+ },
+ )
+ except Exception as log_error: # noqa: BLE001
+ logger.warning("Failed to log mcp_message event: %s", log_error)
+ logger.info(
+ "MCP message: tool=%s, agent_id=%s, user_id=%s, method=%s",
+ getattr(context.message, "name", None),
+ agent_id,
+ user_id,
+ context.method,
+ )
class StructuredContentStripperMiddleware(Middleware):
diff --git a/tests/unit_tests/mcp_service/test_auth_user_resolution.py b/tests/unit_tests/mcp_service/test_auth_user_resolution.py
index a98dd9a7ddc..a9ff3ab5263 100644
--- a/tests/unit_tests/mcp_service/test_auth_user_resolution.py
+++ b/tests/unit_tests/mcp_service/test_auth_user_resolution.py
@@ -606,6 +606,69 @@ def test_setup_user_context_allows_active_user(app) -> None:
assert g.user is active_user
+# -- _mcp_user_id_var (ContextVar surviving the per-call app context pop) --
+#
+# g.user is only valid for the lifetime of the per-call app context that
+# _get_app_context_manager() pushes around tool execution; it's popped
+# before LoggingMiddleware's finally block runs, so get_user_id() there
+# always sees a stale/cleared g. _mcp_user_id_var is a plain ContextVar,
+# not tied to that app-context lifecycle, set here so it survives to be
+# read later for audit logging.
+
+
+def test_setup_user_context_sets_contextvar_for_active_user(app) -> None:
+ """_mcp_user_id_var carries the resolved user's id past this call."""
+ from superset.mcp_service.auth import _mcp_user_id_var, _setup_user_context
+
+ active_user = _make_mock_user("active_user")
+ active_user.is_active = True
+ active_user.active = True
+ active_user.id = 321
+
+ with app.test_request_context():
+ with patch(
+ "superset.mcp_service.auth.get_user_from_request",
+ return_value=active_user,
+ ):
+ _setup_user_context()
+ assert _mcp_user_id_var.get() == 321
+
+
+def test_setup_user_context_clears_stale_contextvar_on_failure(app) -> None:
+ """A previous call's user_id must not leak into a call that fails to
+ resolve a user (e.g. sequential calls sharing one asyncio task)."""
+ from superset.mcp_service.auth import _mcp_user_id_var, _setup_user_context
+
+ with app.test_request_context():
+ _mcp_user_id_var.set(999)
+ with patch(
+ "superset.mcp_service.auth.get_user_from_request",
+ side_effect=ValueError("no user"),
+ ):
+ with pytest.raises(ValueError, match="no user"):
+ _setup_user_context()
+ assert _mcp_user_id_var.get() is None
+
+
+def test_setup_user_context_leaves_contextvar_unset_for_guest_user(app) -> None:
+ """GuestUser (embedded auth) has no numeric id -- the ContextVar must
+ stay cleared rather than store a bogus value."""
+ from superset.mcp_service.auth import _mcp_user_id_var, _setup_user_context
+
+ guest_user = _make_mock_user("guest_user")
+ guest_user.is_active = True
+ guest_user.active = True
+ guest_user.id = None
+
+ with app.test_request_context():
+ with patch(
+ "superset.mcp_service.auth.get_user_from_request",
+ return_value=guest_user,
+ ):
+ _setup_user_context()
+ assert _mcp_user_id_var.get() is None
+
+
# -- Multi-issuer binding guard --
diff --git a/tests/unit_tests/mcp_service/test_middleware_logging.py b/tests/unit_tests/mcp_service/test_middleware_logging.py
index ccab0f23e9c..5b8af9a481f 100644
--- a/tests/unit_tests/mcp_service/test_middleware_logging.py
+++ b/tests/unit_tests/mcp_service/test_middleware_logging.py
@@ -36,6 +36,7 @@ from fastmcp.tools.tool import ToolResult
from mcp import types as mt
from sqlalchemy.exc import OperationalError
+from superset.mcp_service.auth import _mcp_user_id_var
from superset.mcp_service.middleware import LoggingMiddleware
@@ -417,6 +418,108 @@ class TestLoggingMiddlewareOnMessage:
assert payload["error_type"] == "PermissionError"
+class TestLoggingMiddlewareUserIdContextVar:
+ """Tests for _mcp_user_id_var overriding the pre-call user_id snapshot.
+
+ _extract_context_info() runs *before* call_next(), i.e. before the
+ tool's auth decorator (superset/mcp_service/auth.py) has resolved a
+ user, so get_user_id() there is always pre-auth (typically None).
+ _setup_user_context() sets _mcp_user_id_var while running inside
+ call_next(), right before the per-call app context (which g.user
+ depends on) is popped. These tests simulate that by setting the
+ ContextVar from within call_next(), mirroring the real auth decorator.
+ """
+
+ @patch("superset.mcp_service.middleware.event_logger")
+ @patch("superset.mcp_service.middleware.get_user_id", return_value=None)
+ @pytest.mark.asyncio
+ async def test_on_call_tool_uses_contextvar_set_during_call_next(
+ self, mock_get_user_id, mock_event_logger
+ ) -> None:
+ """Logged user_id reflects auth resolved during call_next(), not
+ the pre-call snapshot."""
+ middleware = LoggingMiddleware()
+ ctx = _make_context(name="list_charts")
+
+ async def fake_call_next(_ctx):
+ _mcp_user_id_var.set(42)
+ return "tool_result"
+
+ token = _mcp_user_id_var.set(None)
+ try:
+ result = await middleware.on_call_tool(ctx, fake_call_next)
+ finally:
+ _mcp_user_id_var.reset(token)
+
+ assert result == "tool_result"
+ call_kwargs = mock_event_logger.log.call_args[1]
+ assert call_kwargs["user_id"] == 42
+
+ @patch("superset.mcp_service.middleware.event_logger")
+ @patch("superset.mcp_service.middleware.get_user_id", return_value=7)
+ @pytest.mark.asyncio
+ async def test_on_call_tool_falls_back_when_contextvar_unset(
+ self, mock_get_user_id, mock_event_logger
+ ) -> None:
+ """When the call never runs _setup_user_context() (ContextVar
+ stays at its default), the pre-call get_user_id() value is used
+ as-is -- no regression for paths that don't set the ContextVar."""
+ middleware = LoggingMiddleware()
+ ctx = _make_context(name="health_check")
+ call_next = AsyncMock(return_value="ok")
+
+ token = _mcp_user_id_var.set(None)
+ try:
+ await middleware.on_call_tool(ctx, call_next)
+ finally:
+ _mcp_user_id_var.reset(token)
+
+ call_kwargs = mock_event_logger.log.call_args[1]
+ assert call_kwargs["user_id"] == 7
+
+ @patch("superset.mcp_service.middleware.event_logger")
+ @patch("superset.mcp_service.middleware.get_user_id", return_value=None)
+ @pytest.mark.asyncio
+ async def test_on_message_uses_contextvar_set_during_call_next(
+ self, mock_get_user_id, mock_event_logger
+ ) -> None:
+ middleware = LoggingMiddleware()
+ ctx = _make_context(method="resources/read", name="instance/metadata")
+
+ async def fake_call_next(_ctx):
+ _mcp_user_id_var.set(99)
+ return "resource_data"
+
+ token = _mcp_user_id_var.set(None)
+ try:
+ result = await middleware.on_message(ctx, fake_call_next)
+ finally:
+ _mcp_user_id_var.reset(token)
+
+ assert result == "resource_data"
+ call_kwargs = mock_event_logger.log.call_args[1]
+ assert call_kwargs["user_id"] == 99
+
+ @patch("superset.mcp_service.middleware.event_logger")
+ @patch("superset.mcp_service.middleware.get_user_id", return_value=13)
+ @pytest.mark.asyncio
+ async def test_on_message_falls_back_when_contextvar_unset(
+ self, mock_get_user_id, mock_event_logger
+ ) -> None:
+ middleware = LoggingMiddleware()
+ ctx = _make_context(method="resources/read", name="instance/metadata")
+ call_next = AsyncMock(return_value="resource_data")
+
+ token = _mcp_user_id_var.set(None)
+ try:
+ await middleware.on_message(ctx, call_next)
+ finally:
+ _mcp_user_id_var.reset(token)
+
+ call_kwargs = mock_event_logger.log.call_args[1]
+ assert call_kwargs["user_id"] == 13
+
+
class TestResolveToolName:
"""Tests for LoggingMiddleware._resolve_tool_name()."""