This is an automated email from the ASF dual-hosted git repository. rusackas pushed a commit to branch fix/mcp-worker-metadata-sizing in repository https://gitbox.apache.org/repos/asf/superset.git
commit 58f5686b1acc741b3924a3f71c8997101b9d35b9 Author: Evan Rusackas <[email protected]> AuthorDate: Wed Sep 30 19:13:06 2026 -0700 fix(mcp): fail open consistently, fix a stale cancel id, fix 2 leaks Re-review of #44581 after the metadata-worker-sizing fix landed (admission_counts) and the coverage gap was independently closed. This covers what's left from the earlier review pass that those commits didn't touch: - on_list_tools and both search-transform _get_visible_tools overrides called run_in_metadata_thread with no try/except. That helper reloads the caller's ORM user itself before the passed-in filter ever runs, so a failure there (e.g. the metadata pool exhausted) propagated uncaught -- the filter's own fail-open handling never got a chance to run. All three now fail open consistently; call-time RBAC still enforces permissions. - QueryCancellation.capture() never reset cancel_id to None before a re-capture attempt. A multi-statement call whose second capture failed kept register() republishing the first (already-finished) statement's handle instead of going unavailable for the second. - WorkerPool.cancel() only released its permit on RuntimeError; any other exception from submitting cancellation I/O leaked a permit for the life of the process. - run_in_worker's pool.submit() call was outside the try/finally that resets _mcp_user_id_var, so a busy-pool rejection left a stale user id from a previous call in that context for audit logging. - run_in_worker and TransportContext.forward both relabel any TimeoutError raised while awaiting as "our deadline fired," including one genuinely raised by the tool/driver itself (e.g. a socket-level read timeout, a TimeoutError subclass since Python 3.10). Both now check whether the awaited future is already done (and, for the unshielded forward() case, not cancelled) before relabeling it, so a real driver timeout propagates as itself instead of a misleading generic deadline message with the original exception discarded. Deliberately left alone, flagged rather than guessed at: - cancellation.py's before_cursor_execute listener can raise (WorkerDeadlineExceeded/ToolError via admission) from inside SQLAlchemy's own event dispatch, outside the try/except that routes genuine DBAPI errors through handle_error's cleanup. Fixing this safely needs a clearer picture of how SQLAlchemy's connection pool actually treats a non-DBAPI exception raised from this hook than I have confidence in. - METADATA_ONLY_TOOLS is a hand-maintained list disconnected from where tools declare their own metadata; restructuring that is more invasive than this PR's scope. - run_in_worker's execute() reloads g.user via a metadata round trip that JWT/API-key auth immediately overwrites; skipping it safely needs tracing all four auth priority paths in get_user_from_request, not just the common ones. mcp_service and sql/execution suites pass (6454 passed; the 3 pre-existing failures in test_time_range_validation.py and test_query_dataset.py::test_query_dataset_reexecutes_across_rollover reproduce identically without this commit). superset/sql/'s 100% coverage gate still passes untouched. All touched files pass pre-commit (mypy, ruff, pylint). Co-Authored-By: Evan Rusackas <[email protected]> Co-Authored-By: Claude Sonnet 5 <[email protected]> --- superset/mcp_service/middleware.py | 12 ++- superset/mcp_service/server.py | 30 ++++++-- superset/mcp_service/worker.py | 98 +++++++++++++++++++------ tests/unit_tests/mcp_service/test_middleware.py | 21 ++++++ tests/unit_tests/mcp_service/test_worker.py | 55 +++++++++++++- 5 files changed, 185 insertions(+), 31 deletions(-) diff --git a/superset/mcp_service/middleware.py b/superset/mcp_service/middleware.py index dcdb733192f..5617225f425 100644 --- a/superset/mcp_service/middleware.py +++ b/superset/mcp_service/middleware.py @@ -1200,7 +1200,17 @@ class RBACToolVisibilityMiddleware(Middleware): # User and permission lookups need a metadata connection. Tool workers # can hold every pooled connection while they wait on this loop, so the # loop must never wait for one itself. - return await run_in_metadata_thread(self._visible_tools, tools) + try: + return await run_in_metadata_thread(self._visible_tools, tools) + except Exception: # noqa: BLE001 + # A failure setting up the metadata thread itself (e.g. the + # metadata pool is exhausted) happens before _visible_tools's own + # try/except ever runs. Fail open the same way that branch does; + # call-time RBAC still enforces permissions. + logger.warning( + "MCP tool list: failed to filter tools, showing all", exc_info=True + ) + return tools @staticmethod def _visible_tools(tools: list[Tool]) -> list[Tool]: diff --git a/superset/mcp_service/server.py b/superset/mcp_service/server.py index dc43ac5095c..51153ae3972 100644 --- a/superset/mcp_service/server.py +++ b/superset/mcp_service/server.py @@ -430,6 +430,28 @@ def _filter_tools_by_current_user_permission(tools: Sequence[Any]) -> list[Any]: return [tool for tool in tools if _tool_allowed_for_current_user(tool)] +async def _filter_visible_tools_fail_open(tools: Sequence[Any]) -> Sequence[Any]: + """Run the permission filter in the metadata thread, failing open on error. + + ``run_in_metadata_thread`` reloads the caller's ORM user itself before the + filter ever runs (e.g. a metadata-pool-exhaustion failure), so a bare + ``await run_in_metadata_thread(...)`` here would raise before any fail-open + handling inside the filter gets a chance to run. Call-time RBAC still + enforces permissions, so an unexpected failure here shows every tool + rather than breaking search, matching + ``RBACToolVisibilityMiddleware.on_list_tools``. + """ + try: + return await run_in_metadata_thread( + _filter_tools_by_current_user_permission, tools + ) + except Exception: # noqa: BLE001 + logger.warning( + "MCP tool search: failed to filter tools, showing all", exc_info=True + ) + return tools + + def _create_search_result_serializer( config: dict[str, Any], ) -> Any: @@ -671,9 +693,7 @@ def _create_search_transform( # noqa: C901 async def _get_visible_tools(self, ctx: Context) -> Sequence[Any]: """Return only tools visible to the current authenticated user.""" tools = await super()._get_visible_tools(ctx) - return await run_in_metadata_thread( - _filter_tools_by_current_user_permission, tools - ) + return await _filter_visible_tools_fail_open(tools) def _make_call_tool(self) -> Any: """Build the normalized ``call_tool`` proxy for regex search.""" @@ -695,9 +715,7 @@ def _create_search_transform( # noqa: C901 tools = await super()._get_visible_tools(ctx) # Permission lookups need a metadata connection; see # RBACToolVisibilityMiddleware.on_list_tools. - return await run_in_metadata_thread( - _filter_tools_by_current_user_permission, tools - ) + return await _filter_visible_tools_fail_open(tools) def _make_call_tool(self) -> Any: """Build the normalized ``call_tool`` proxy for BM25 search.""" diff --git a/superset/mcp_service/worker.py b/superset/mcp_service/worker.py index 57b35a738ea..eee53dbd9b5 100644 --- a/superset/mcp_service/worker.py +++ b/superset/mcp_service/worker.py @@ -178,15 +178,26 @@ class WorkerPool: try: future = self.cancellations.submit(fn) except RuntimeError: + # The cancellations executor is shut down; report failure to + # dispatch rather than raise. self.cancel_slots.release() return False + except BaseException: + self.cancel_slots.release() + raise def completed(_: Future[None]) -> None: """Release cancellation capacity before admitting another tool call.""" self.cancel_slots.release() finished() - future.add_done_callback(completed) + try: + future.add_done_callback(completed) + except BaseException: + # The permit is otherwise never released: completed() runs only if + # the callback was registered, and a raise here means it was not. + self.cancel_slots.release() + raise return True @@ -420,12 +431,20 @@ class TransportContext: async def forward(*args: Any, **kwargs: Any) -> Any: self.call.check() future = asyncio.run_coroutine_threadsafe(value(*args, **kwargs), self.loop) + wrapped = asyncio.wrap_future(future) try: result = await asyncio.wait_for( - asyncio.wrap_future(future), + wrapped, max(0, self.call.deadline - time.monotonic()), ) except TimeoutError: + if wrapped.done() and not wrapped.cancelled(): + # wait_for cancels the awaited future on its own timeout + # (no shield here), so a TimeoutError that leaves it + # already done rather than cancelled was genuinely + # raised by the forwarded call itself. Propagate that + # one as itself, not relabeled as our deadline. + raise raise WorkerDeadlineExceeded() from None self.call.check() return result @@ -554,6 +573,41 @@ async def run_in_transport_thread( ) +def _handle_worker_timeout_or_cancel( + exc: BaseException, + wrapped: asyncio.Future[Any], + call: WorkerCall, + seconds: float, +) -> None: + """Decide how a run_in_worker deadline or cancellation should surface. + + Always raises; the caller's except block delegates here instead of + inlining the decision, to keep run_in_worker's own branching simple. + """ + call.abandon() + # Retrieve late exceptions, including worker CancelledError, without + # retaining a task on the transport loop after the request has ended. + wrapped.add_done_callback( + lambda done: None if done.cancelled() else done.exception() + ) + if isinstance(exc, asyncio.CancelledError): + raise exc + if isinstance(exc, TimeoutError) and wrapped.done(): + # wait_for raises this same exception type both on its own deadline + # and when the shielded future already finished with a TimeoutError + # of its own (e.g. a driver-level socket timeout). Shielding means + # our deadline firing leaves the future still pending; only a + # TimeoutError genuinely raised by the tool/driver leaves it already + # done. Propagate that one as itself rather than relabeling it as + # our deadline. + raise exc + raise ToolError( + f"MCP tool timed out after {seconds:g} seconds. " + "Warehouse cancellation was requested where supported. " + f"Call id: {call.call_id}" + ) from None + + async def run_in_worker( fn: Callable[..., Coroutine[Any, Any, Any]], args: tuple[Any, ...], @@ -636,27 +690,18 @@ async def run_in_worker( call.check() return result - future = pool.submit(lambda: context.run(execute), call.finished, call.admitted) - wrapped = asyncio.wrap_future(future) try: - # Shield the future: cancellation must not release its pool slot early. - return await asyncio.wait_for( - asyncio.shield(wrapped), timeout=max(0, call.deadline - time.monotonic()) - ) - except (TimeoutError, WorkerDeadlineExceeded, asyncio.CancelledError) as exc: - call.abandon() - # Retrieve late exceptions, including worker CancelledError, without - # retaining a task on the transport loop after the request has ended. - wrapped.add_done_callback( - lambda done: None if done.cancelled() else done.exception() - ) - if isinstance(exc, asyncio.CancelledError): - raise - raise ToolError( - f"MCP tool timed out after {seconds:g} seconds. " - "Warehouse cancellation was requested where supported. " - f"Call id: {call.call_id}" - ) from None + future = pool.submit(lambda: context.run(execute), call.finished, call.admitted) + wrapped = asyncio.wrap_future(future) + try: + # Shield the future: cancellation must not release its pool slot + # early. + return await asyncio.wait_for( + asyncio.shield(wrapped), + timeout=max(0, call.deadline - time.monotonic()), + ) + except (TimeoutError, WorkerDeadlineExceeded, asyncio.CancelledError) as exc: + _handle_worker_timeout_or_cancel(exc, wrapped, call, seconds) finally: from superset.mcp_service.auth import _mcp_user_id_var @@ -685,9 +730,16 @@ class QueryCancellation: self.context = copy_context() def capture(self) -> None: - """Capture an engine handle without making unsupported drivers fail.""" + """Capture an engine handle without making unsupported drivers fail. + + Cleared up front: a multi-statement call re-captures per statement + (see ``refresh``), and a failed re-capture must not leave register() + republishing a stale handle for a statement that has already + finished. + """ from superset.tasks.query_cancel import capture_cancel_query_id + self.cancel_id = None try: self.cancel_id = capture_cancel_query_id(self.database, self.cursor) except Exception: diff --git a/tests/unit_tests/mcp_service/test_middleware.py b/tests/unit_tests/mcp_service/test_middleware.py index 47c3f054757..d6decb56e53 100644 --- a/tests/unit_tests/mcp_service/test_middleware.py +++ b/tests/unit_tests/mcp_service/test_middleware.py @@ -2426,6 +2426,27 @@ class TestRBACToolVisibilityMiddleware: assert result == tools + @pytest.mark.asyncio + async def test_fails_open_when_metadata_thread_setup_raises(self) -> None: + """Returns all tools when run_in_metadata_thread's own setup fails. + + Unlike test_fails_open_on_exception, this failure happens before + _visible_tools ever runs -- e.g. the metadata pool is exhausted while + run_in_metadata_thread reloads the caller's ORM user -- so + _visible_tools's own try/except never gets a chance to fail open. + """ + tools = [self._make_tool("list_charts"), self._make_tool("generate_chart")] + call_next = AsyncMock(return_value=tools) + middleware = RBACToolVisibilityMiddleware() + + with patch( + "superset.mcp_service.middleware.run_in_metadata_thread", + side_effect=RuntimeError("metadata pool exhausted"), + ): + result = await middleware.on_list_tools(MagicMock(), call_next) + + assert result == tools + @pytest.mark.asyncio async def test_fails_open_when_user_is_none(self, app) -> None: """Returns all tools when get_user_from_request returns None.""" diff --git a/tests/unit_tests/mcp_service/test_worker.py b/tests/unit_tests/mcp_service/test_worker.py index d6268283d1c..71e6d061bd7 100644 --- a/tests/unit_tests/mcp_service/test_worker.py +++ b/tests/unit_tests/mcp_service/test_worker.py @@ -24,7 +24,7 @@ from collections.abc import Iterator from contextlib import contextmanager, ExitStack from contextvars import ContextVar from typing import Any -from unittest.mock import Mock, patch +from unittest.mock import MagicMock, Mock, patch import pytest from fastmcp import Client, FastMCP @@ -380,6 +380,59 @@ async def test_failed_warehouse_discards_metadata_session(app: Any) -> None: pool.cancellations.shutdown() +def test_query_cancellation_capture_resets_stale_id_on_failed_recapture() -> None: + """A failed re-capture must not leave register() republishing a stale id + from an earlier, already-finished statement in a multi-statement call. + """ + from superset.mcp_service.worker import QueryCancellation + + call = MagicMock() + call.expired.is_set.return_value = False + database = MagicMock() + database.db_engine_spec.has_implicit_cancel.return_value = False + database.db_engine_spec.has_query_id_during_execute = False + cancellation = QueryCancellation(call, database, Mock(), None, None) + + with patch( + "superset.tasks.query_cancel.capture_cancel_query_id", + return_value="statement-1-id", + ): + cancellation.capture() + assert cancellation.cancel_id == "statement-1-id" + cancellation.register() + assert call.cancel_query is not None + + with patch( + "superset.tasks.query_cancel.capture_cancel_query_id", + side_effect=RuntimeError("driver refused a second handle"), + ): + cancellation.capture() + assert cancellation.cancel_id is None + cancellation.register() + # No usable handle for this statement: cancellation must go unavailable, + # not silently keep dispatching against the finished first statement. + assert call.cancel_query is None + + +def test_worker_pool_cancel_releases_slot_on_unexpected_submit_error() -> None: + """Any failure to dispatch cancellation I/O must release its permit, not + just the already-handled RuntimeError-on-shutdown case.""" + from superset.mcp_service.worker import WorkerPool + + pool = WorkerPool(1) + try: + with patch.object(pool.cancellations, "submit", side_effect=ValueError("boom")): + with pytest.raises(ValueError, match="boom"): + pool.cancel(Mock(), Mock()) + # The permit was released despite the unexpected exception type, so a + # later cancellation can still be admitted. + assert pool.cancel_slots.acquire(blocking=False) + pool.cancel_slots.release() + finally: + pool.executor.shutdown() + pool.cancellations.shutdown() + + @pytest.mark.asyncio async def test_sqlalchemy_warehouse_queries_register_cancellation(app: Any) -> None: """Metadata discovery also uses the engine hooks, not just chart/SQL Lab SQL."""
