You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by jl...@apache.org on 2026/08/13 18:05:48 UTC
(superset) branch master updated: fix(mcp): give concurrent tool calls isolated db sessions (#42629)
This is an automated email from the ASF dual-hosted git repository.
sadpandajoe 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 d0658bacc89 fix(mcp): give concurrent tool calls isolated db sessions (#42629)
d0658bacc89 is described below
commit d0658bacc89daf7edac0006132f3d128deeae66d
Author: goingforstudying-ctrl <go...@gmail.com>
AuthorDate: Thu Aug 13 14:05:40 2026 -0400
fix(mcp): give concurrent tool calls isolated db sessions (#42629)
Co-authored-by: goingforstudying-ctrl <go...@users.noreply.github.com>
Co-authored-by: Evan Rusackas <ev...@preset.io>
Co-authored-by: Joe Li <jo...@preset.io>
---
superset/mcp_service/app.py | 28 +++
superset/mcp_service/auth.py | 71 +++++--
superset/mcp_service/server.py | 10 +
superset/mcp_service/session_scope.py | 69 ++++++
.../mcp_service/test_auth_user_resolution.py | 7 +-
tests/unit_tests/mcp_service/test_session_scope.py | 233 +++++++++++++++++++++
6 files changed, 402 insertions(+), 16 deletions(-)
diff --git a/superset/mcp_service/app.py b/superset/mcp_service/app.py
index b666ad9f291..7d233ee1d09 100644
--- a/superset/mcp_service/app.py
+++ b/superset/mcp_service/app.py
@@ -704,6 +704,27 @@ def create_mcp_app(
# Log instance creation
_log_instance_creation(name, auth, include_tags, exclude_tags)
+ # Install per-tool-call session scoping for callers that serve the app
+ # straight from the factory. init_fastmcp_server() and run_server()
+ # install it as well, but a deployment that only calls create_mcp_app()
+ # would otherwise keep the greenlet-scoped default and re-introduce the
+ # shared-session race between concurrent tool calls (#42622).
+ # The import stays lazy because this function also runs at module
+ # import time (``mcp = create_mcp_app()`` below), where
+ # superset.extensions can still be mid-initialization; in that case the
+ # server entry points install the scopefunc deterministically later.
+ try:
+ from superset.mcp_service.session_scope import ( # noqa: PLC0415
+ install_mcp_session_scoping,
+ )
+
+ install_mcp_session_scoping()
+ except ImportError:
+ logger.debug(
+ "superset.extensions not fully initialized during create_mcp_app(); "
+ "MCP session scoping deferred to the server entry point"
+ )
+
return mcp_instance
@@ -1007,6 +1028,13 @@ def init_fastmcp_server(
# circular import: flask_singleton imports from superset.extensions which
# re-enters mcp_service during startup; must stay lazy inside the function.
from superset.mcp_service.flask_singleton import app as flask_app # noqa: PLC0415
+ from superset.mcp_service.session_scope import ( # noqa: PLC0415
+ install_mcp_session_scoping,
+ )
+
+ # Give each tool call its own SQLAlchemy session; the greenlet-scoped
+ # default lets concurrent async calls share (and tear down) one session.
+ install_mcp_session_scoping()
# Derive branding from Superset's APP_NAME config (defaults to "Superset")
app_name = flask_app.config.get("APP_NAME", "Superset")
diff --git a/superset/mcp_service/auth.py b/superset/mcp_service/auth.py
index d4383ad489f..d4686addaf0 100644
--- a/superset/mcp_service/auth.py
+++ b/superset/mcp_service/auth.py
@@ -45,8 +45,8 @@ Configuration:
"""
import logging
-from contextlib import AbstractContextManager, nullcontext
-from typing import Any, Callable, cast, TYPE_CHECKING, TypeAlias, TypeVar
+from contextlib import AbstractContextManager, contextmanager
+from typing import Any, Callable, cast, Generator, TYPE_CHECKING, TypeAlias, TypeVar
from flask import current_app, g, has_app_context, has_request_context
from flask_appbuilder.security.sqla.models import User
@@ -62,6 +62,7 @@ from superset.mcp_service.mcp_config import (
get_mcp_api_key_enabled,
validate_multi_issuer_user_resolver,
)
+from superset.mcp_service.session_scope import _mcp_session_token
from superset.mcp_service.utils.error_sanitization import (
sanitize_for_log as _sanitize_for_log,
)
@@ -1021,7 +1022,8 @@ def _get_app_context_manager() -> AbstractContextManager[None]:
When a request context is present, external middleware (e.g.
Preset's WorkspaceContextMiddleware) has already set ``g.user``
- on a per-request app context — reuse it via ``nullcontext()``.
+ on a per-request app context — reuse it, but still tag the call
+ with a per-call session token (see ``_request_tool_call_context``).
When only a bare app context exists (no request context), push a
**new** app context so concurrent tool calls do not share one ``g``
@@ -1030,21 +1032,64 @@ def _get_app_context_manager() -> AbstractContextManager[None]:
When no context exists at all, push a fresh app context from the
Flask singleton.
+ Every path tags the call with a per-call session token (see
+ session_scope.py) so each tool call gets its own SQLAlchemy
+ session — otherwise all in-flight async calls on one greenlet
+ share the greenlet-scoped session, and one call's teardown
+ detaches the others' ORM instances.
+
This is the single source of truth for context selection — called
from both ``mcp_auth_hook`` (tool execution) and
``RBACToolVisibilityMiddleware`` (tools/list filtering).
"""
if has_request_context():
- return nullcontext()
- if has_app_context():
- # Push a new context for the CURRENT app (not get_flask_app()
- # which may return a different instance in test environments).
- return current_app._get_current_object().app_context()
- # Deferred: importing at module level would trigger create_app() before
- # Superset is fully initialised (e.g. during unit-test collection).
- from superset.mcp_service.flask_singleton import get_flask_app
-
- return get_flask_app().app_context()
+ return _request_tool_call_context()
+ return _mcp_tool_call_context()
+
+
+@contextmanager
+def _request_tool_call_context() -> Generator[None, None, None]:
+ """Tag a request-backed tool call with a per-call session token.
+
+ The request context is reused as-is (middleware already populated
+ ``g.user``), so no new app context is pushed — but without the
+ token, concurrent tool calls on the request's greenlet would share
+ the greenlet-scoped SQLAlchemy session, the same race that
+ ``_mcp_tool_call_context()`` prevents on the no-request path.
+ """
+ token = _mcp_session_token.set(object())
+ try:
+ yield
+ finally:
+ # Deregister this call's session while the token still resolves
+ # the registry to it; the request's own greenlet-scoped session
+ # is left for the request lifecycle to tear down.
+ _remove_session_safe()
+ _mcp_session_token.reset(token)
+
+
+@contextmanager
+def _mcp_tool_call_context() -> Generator[None, None, None]:
+ """Push a fresh app context tagged with a per-call session token."""
+ token = _mcp_session_token.set(object())
+ try:
+ if has_app_context():
+ # Push a new context for the CURRENT app (not get_flask_app()
+ # which may return a different instance in test environments).
+ with current_app._get_current_object().app_context():
+ yield
+ else:
+ # Deferred: importing at module level would trigger create_app()
+ # before Superset is fully initialised (e.g. during unit-test
+ # collection).
+ from superset.mcp_service.flask_singleton import get_flask_app
+
+ with get_flask_app().app_context():
+ yield
+ finally:
+ # Reset only after the app context popped, so teardown's
+ # db.session.remove() still resolves to this call's session.
+ _mcp_session_token.reset(token)
def mcp_auth_hook(tool_func: F) -> F: # noqa: C901
diff --git a/superset/mcp_service/server.py b/superset/mcp_service/server.py
index 81540c6380a..064b6305407 100644
--- a/superset/mcp_service/server.py
+++ b/superset/mcp_service/server.py
@@ -970,6 +970,16 @@ def run_server(
logging.info("Creating MCP app from factory configuration...")
factory_config = get_mcp_factory_config()
mcp_instance = create_mcp_app(**factory_config)
+ # The factory path bypasses init_fastmcp_server(), so install the
+ # per-tool-call session scoping here as well; without it concurrent
+ # tool calls share the greenlet-scoped db.session.
+ # Lazy import mirrors init_fastmcp_server() to avoid the circular
+ # import through superset.extensions during startup.
+ from superset.mcp_service.session_scope import ( # noqa: PLC0415
+ install_mcp_session_scoping,
+ )
+
+ install_mcp_session_scoping()
# Capture the actual auth object so the hello page reflects real auth state
auth_provider = factory_config.get("auth")
flask_app = None
diff --git a/superset/mcp_service/session_scope.py b/superset/mcp_service/session_scope.py
new file mode 100644
index 00000000000..9aca5dd55fb
--- /dev/null
+++ b/superset/mcp_service/session_scope.py
@@ -0,0 +1,69 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Per-tool-call SQLAlchemy session scoping for the MCP service.
+
+``db.session`` is a flask-sqlalchemy 2.5.1 scoped session whose registry is
+keyed by greenlet ident. Async MCP tool calls are asyncio tasks on one
+event-loop greenlet, so concurrent calls resolve to the SAME ``Session``
+even though each call pushes its own app context; the first call to finish
+then removes that shared session at app-context teardown and detaches the
+other calls' instances (DetachedInstanceError on the next attribute read).
+
+Keying the registry on a per-call ContextVar token instead gives every tool
+call its own ``Session`` while the app-context teardown removes exactly that
+call's session. Outside MCP tool calls the scope falls back to the greenlet
+ident, so web, CLI, and Celery behavior is unchanged.
+"""
+
+from __future__ import annotations
+
+import logging
+from contextvars import ContextVar
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+try:
+ from greenlet import getcurrent as _ident_func
+except ImportError: # pragma: no cover
+ from threading import get_ident as _ident_func
+
+# Set by _get_app_context_manager() for the duration of one MCP tool call.
+_mcp_session_token: ContextVar[Any] = ContextVar(
+ "superset_mcp_session_token", default=None
+)
+
+
+def mcp_session_scopefunc() -> Any:
+ """Registry key: per-call token inside MCP tool calls, greenlet otherwise."""
+ if (token := _mcp_session_token.get()) is not None:
+ return ("mcp_tool_call", id(token))
+ return _ident_func()
+
+
+def install_mcp_session_scoping() -> None:
+ """Point ``db.session``'s registry at the MCP-aware scopefunc (idempotent)."""
+ from superset.extensions import db
+
+ registry = getattr(db.session, "registry", None)
+ if registry is not None and hasattr(registry, "scopefunc"):
+ registry.scopefunc = mcp_session_scopefunc
+ else: # pragma: no cover
+ logger.warning(
+ "db.session has no scopefunc-based registry; "
+ "MCP per-call session scoping not installed"
+ )
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 33dc1354841..a98dd9a7ddc 100644
--- a/tests/unit_tests/mcp_service/test_auth_user_resolution.py
+++ b/tests/unit_tests/mcp_service/test_auth_user_resolution.py
@@ -465,7 +465,8 @@ def test_sync_wrapper_handles_ssl_error_on_pre_call_remove(app) -> None:
SAOperationalError(
"SSL connection has been closed unexpectedly", None, None
),
- None, # second call succeeds
+ None, # retry succeeds
+ None, # exit-path cleanup in _request_tool_call_context
]
with patch(
@@ -476,8 +477,8 @@ def test_sync_wrapper_handles_ssl_error_on_pre_call_remove(app) -> None:
assert result == "fresh"
assert mock_db.session.invalidate.called, "invalidate() must be called on SSL error"
- assert mock_db.session.remove.call_count == 2, (
- "remove() must be retried after SSL error"
+ assert mock_db.session.remove.call_count == 3, (
+ "remove() must be retried after SSL error, plus once more on exit"
)
diff --git a/tests/unit_tests/mcp_service/test_session_scope.py b/tests/unit_tests/mcp_service/test_session_scope.py
new file mode 100644
index 00000000000..6a7c62dae19
--- /dev/null
+++ b/tests/unit_tests/mcp_service/test_session_scope.py
@@ -0,0 +1,233 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tests for per-tool-call SQLAlchemy session scoping in the MCP service.
+
+db.session is keyed by greenlet ident, so concurrent async tool calls on
+one event-loop greenlet share a single Session; the first call's app-context
+teardown removes it and detaches the other calls' instances (issue #42622).
+The MCP scopefunc keys the registry on a per-call ContextVar token instead.
+"""
+
+import asyncio
+from collections.abc import Iterator
+from typing import Any
+
+import pytest
+from sqlalchemy import text
+
+from superset.extensions import db
+from superset.mcp_service.auth import _mcp_tool_call_context
+from superset.mcp_service.session_scope import (
+ install_mcp_session_scoping,
+ mcp_session_scopefunc,
+)
+
+# NOTE: most tests enter tool-call contexts via _mcp_tool_call_context()
+# directly rather than _get_app_context_manager(). The latter consults
+# Flask state (request/app context) to pick a context manager, while these
+# tests exercise the token mechanics in isolation. Request-backed tool
+# calls go through _request_tool_call_context(), covered by
+# test_request_context_tool_calls_get_isolated_sessions below.
+
+
+@pytest.fixture
+def mcp_scoping() -> Iterator[Any]:
+ """Install the MCP scopefunc and restore the original one afterwards."""
+ registry = db.session.registry
+ original_scopefunc = registry.scopefunc
+ install_mcp_session_scoping()
+ yield registry
+ registry.scopefunc = original_scopefunc
+
+
+def test_install_swaps_scopefunc_and_is_idempotent(mcp_scoping: Any) -> None:
+ assert mcp_scoping.scopefunc is mcp_session_scopefunc
+ install_mcp_session_scoping()
+ assert mcp_scoping.scopefunc is mcp_session_scopefunc
+
+
+def test_scopefunc_falls_back_to_greenlet_outside_tool_calls(
+ mcp_scoping: Any,
+) -> None:
+ # No per-call token set: the key must be the greenlet ident, exactly
+ # what flask-sqlalchemy uses by default, so web/CLI/Celery paths are
+ # unaffected.
+ import greenlet
+
+ assert mcp_session_scopefunc() is greenlet.getcurrent()
+
+
+@pytest.mark.asyncio
+async def test_concurrent_tool_calls_get_isolated_sessions(
+ mcp_scoping: Any,
+) -> None:
+ """Each async tool call gets its own Session; one call's app-context
+ teardown removes only that call's session, leaving the other usable."""
+ a_ready = asyncio.Event()
+ b_done = asyncio.Event()
+ outcome: dict[str, Any] = {}
+
+ async def call_a() -> None:
+ with _mcp_tool_call_context():
+ session_a = db.session()
+ outcome["a"] = session_a
+ a_ready.set()
+ await b_done.wait()
+ # B has fully torn down; A's session must still resolve and work.
+ assert db.session() is session_a
+ session_a.execute(text("SELECT 1"))
+ outcome["a_alive"] = True
+
+ async def call_b() -> None:
+ await a_ready.wait()
+ with _mcp_tool_call_context():
+ outcome["b"] = db.session()
+ # Leaving the context pops it, so teardown removes B's session here.
+ b_done.set()
+
+ await asyncio.gather(call_a(), call_b())
+
+ assert outcome["a"] is not outcome["b"]
+ assert outcome.get("a_alive") is True
+
+
+@pytest.mark.asyncio
+async def test_shared_session_teardown_breaks_other_call_without_fix(
+ mcp_scoping: Any,
+) -> None:
+ """Counterfactual: with the default greenlet scope restored, concurrent
+ calls resolve to the SAME Session and the first teardown breaks the
+ survivor — the exact failure this change prevents."""
+ registry = db.session.registry
+ original = registry.scopefunc
+ try:
+ import greenlet
+
+ registry.scopefunc = greenlet.getcurrent
+
+ a_ready = asyncio.Event()
+ b_done = asyncio.Event()
+ outcome: dict[str, Any] = {}
+
+ async def call_a() -> None:
+ with _mcp_tool_call_context():
+ outcome["a"] = db.session()
+ a_ready.set()
+ await b_done.wait()
+ outcome["a_has_session"] = db.session.registry.has()
+
+ async def call_b() -> None:
+ await a_ready.wait()
+ with _mcp_tool_call_context():
+ outcome["b"] = db.session()
+ b_done.set()
+
+ await asyncio.gather(call_a(), call_b())
+
+ assert outcome["a"] is outcome["b"]
+ # B's teardown removed the shared session out from under A.
+ assert outcome["a_has_session"] is False
+ finally:
+ registry.scopefunc = original
+
+
+def test_create_mcp_app_installs_session_scoping() -> None:
+ """Direct factory callers get per-call scoping without a server entry point.
+
+ Deployments that serve the app straight from ``create_mcp_app()`` never
+ pass through ``init_fastmcp_server()`` or ``run_server()``; the factory
+ itself must swap the scopefunc or concurrent tool calls share one
+ greenlet-scoped session again.
+ """
+ import greenlet
+
+ from superset.mcp_service.app import create_mcp_app
+
+ registry = db.session.registry
+ original_scopefunc = registry.scopefunc
+ try:
+ registry.scopefunc = greenlet.getcurrent
+ create_mcp_app(name="scoping-test")
+ assert registry.scopefunc is mcp_session_scopefunc
+ finally:
+ registry.scopefunc = original_scopefunc
+
+
+@pytest.mark.asyncio
+async def test_nested_tool_calls_keep_separate_sessions(mcp_scoping: Any) -> None:
+ """A nested tool-call context gets its own session and its teardown does
+ not disturb the enclosing call's session."""
+ with _mcp_tool_call_context():
+ outer = db.session()
+ with _mcp_tool_call_context():
+ inner = db.session()
+ assert inner is not outer
+ # inner context popped: outer session untouched
+ assert db.session() is outer
+ outer.execute(text("SELECT 1"))
+
+
+@pytest.mark.asyncio
+async def test_request_context_tool_calls_get_isolated_sessions(
+ mcp_scoping: Any,
+) -> None:
+ """WorkspaceContextMiddleware path: with a request context active,
+ _get_app_context_manager() reuses it but must still give each tool
+ call its own session token — request-backed calls would otherwise
+ share the greenlet-scoped session, the same race fixed elsewhere.
+ Exiting a call removes only that call's session; the request's own
+ session belongs to the request lifecycle.
+ """
+ from superset.mcp_service.auth import _get_app_context_manager
+ from superset.mcp_service.flask_singleton import get_flask_app
+
+ app = get_flask_app()
+ with app.test_request_context(path="/mcp"):
+ request_session = db.session()
+ a_ready = asyncio.Event()
+ b_done = asyncio.Event()
+ outcome: dict[str, Any] = {}
+
+ async def call_a() -> None:
+ with _get_app_context_manager():
+ session_a = db.session()
+ outcome["a"] = session_a
+ a_ready.set()
+ await b_done.wait()
+ # B has fully torn down; A's session must still resolve.
+ assert db.session() is session_a
+ session_a.execute(text("SELECT 1"))
+ outcome["a_alive"] = True
+
+ async def call_b() -> None:
+ await a_ready.wait()
+ with _get_app_context_manager():
+ outcome["b"] = db.session()
+ outcome["b_key"] = mcp_session_scopefunc()
+ # Exiting deregistered B's per-call session from the registry.
+ outcome["b_removed"] = outcome["b_key"] not in db.session.registry.registry
+ b_done.set()
+
+ await asyncio.gather(call_a(), call_b())
+
+ assert outcome["a"] is not outcome["b"]
+ assert outcome["a"] is not request_session
+ assert outcome["b"] is not request_session
+ assert outcome.get("a_alive") is True
+ assert outcome["b_removed"] is True
+ # The request's own greenlet-scoped session is untouched.
+ assert db.session() is request_session