You are viewing a plain text version of this content. The canonical link for it is here.
Posted to notifications@superset.apache.org by "aminghadersohi (via GitHub)" <gi...@apache.org> on 2026/09/23 18:18:26 UTC
[PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
aminghadersohi opened a new pull request, #44581:
URL: https://github.com/apache/superset/pull/44581
### SUMMARY
Prevent synchronous warehouse calls from freezing the MCP transport loop. Run protected async tools in isolated thread workers, including indirect chart, dataset, preview, and metadata-discovery query paths.
Design choices:
- `MCP_TOOL_WORKERS=16` bounds admission; saturation returns โserver busyโ immediately.
- Use the request timeout, or `SQLLAB_TIMEOUT` when absent, for the whole call. Composed calls share that deadline.
- Reuse engine cancellation IDs/hooks and implicit cursor cancellation on a separate bounded pool. Uncancellable queries retain their slot; stuck cancellation also retains admission. No shared engine-spec statement-timeout API exists, so MCP does not rely on process-wide SIGALRM.
- Each worker owns its app context/session. Invalidate failed-query sessions and safely remove them before teardown; preserve acting-user and transport context.
### BEFORE/AFTER SCREENSHOTS OR ANIMATED GIF
Not applicable. Sleeping queries no longer prevent another MCP request from completing, and callers receive a clear timeout error at their deadline.
### TESTING INSTRUCTIONS
- Four timing regressions using real sleeping SQLite queries through chart/SQL Lab execution fail on `origin/master` (`c0c688d3e4`) and pass here.
- Worker tests cover overload, cancellation (including late handles), cleanup, and context isolation.
- MCP, SQL executor, model, and cancellation unit suites; pre-commit, Ruff, and mypy.
### ADDITIONAL INFORMATION
No feature flag, migration, or UI change. Driver cancellation is best-effort; a timeout does not guarantee warehouse writes were rolled back.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on PR #44581:
URL: https://github.com/apache/superset/pull/44581#issuecomment-5924737515
AI Code Review is in progress (usually takes 3 to 15 minutes unless it's a very large PR).
<details>
<summary>
Bito Usage Guide
</summary>
**Commands**
Type the following command in the pull request comment and save the comment.
- `/review` - Manually triggers an incremental AI Review.
- `/review full` - Manually triggers a full AI Review.
- `/pause` - Pauses automatic reviews on this pull request.
- `/resume` - Resumes automatic reviews.
- `/resolve` - Marks all Bito-posted review comments as resolved.
- `/abort` - Cancels all in-progress reviews.
Refer to the <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/available-commands" target="_blank" rel="noopener noreferrer">documentation</a> for additional commands.
**Configuration**
This repository uses `Superset` You can customize the agent settings <a href="https://alpha.bito.ai/home/ai-agents/code-review-agent" target="_blank" rel="noopener noreferrer">here</a> or contact your Bito workspace admin at evan@preset.io.
**Documentation & Help**
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/getting-started/install-run-using-bito-cloud/create-or-customize-an-agent-instance" target="_blank" rel="noopener noreferrer">Customize agent settings</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/implementing-custom-code-review-rules" target="_blank" rel="noopener noreferrer">Review rules</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent" target="_blank" rel="noopener noreferrer">General documentation</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/faqs" target="_blank" rel="noopener noreferrer">FAQ</a>
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "gabotorresruiz (via GitHub)" <gi...@apache.org>.
gabotorresruiz commented on code in PR #44581:
URL: https://github.com/apache/superset/pull/44581#discussion_r4157865812
##########
superset/mcp_service/worker.py:
##########
@@ -0,0 +1,910 @@
+# 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.
+"""Bounded, deadline-aware execution of MCP tools outside the transport loop.
+
+The tool's coroutine runs on a worker-owned loop because its database APIs are
+synchronous. Only transport notifications are marshalled back to the server
+loop. Flask and SQLAlchemy lifetimes belong to the worker, not the waiting
+request: a timed-out DBAPI call cannot outlive and reuse a torn-down session.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import logging
+import threading
+import time
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import contextmanager
+from contextvars import ContextVar, copy_context
+from typing import Any, Callable, Coroutine, Iterator, ParamSpec, TYPE_CHECKING, TypeVar
+from weakref import WeakKeyDictionary
+
+from fastmcp.exceptions import ToolError
+from flask import current_app, g, has_app_context, has_request_context
+from sqlalchemy import inspect as sa_inspect
+from sqlalchemy.orm.state import InstanceState
+from sqlalchemy.pool import QueuePool
+
+from superset.mcp_service.session_scope import _mcp_session_token
+
+if TYPE_CHECKING:
+ from flask import Flask
+
+ from superset.models.core import Database
+
+logger = logging.getLogger(__name__)
+_active_call: ContextVar[WorkerCall | None] = ContextVar(
+ "mcp_worker_call", default=None
+)
+_metadata_context_owned: ContextVar[bool] = ContextVar(
+ "mcp_metadata_context_owned", default=False
+)
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
+_pools_lock = threading.Lock()
+_pools: WeakKeyDictionary[Flask, WorkerPool] = WeakKeyDictionary()
+
+
+class WorkerDeadlineExceeded(BaseException):
+ """Stop abandoned work without being swallowed by tool error handlers."""
+
+
+BUSY_MESSAGE = "MCP server busy: all tool workers are occupied. Retry later."
+
+# Tools that only read Superset's metadata database. They are admitted under a
+# separate, larger bound so they keep answering while slow warehouse queries
+# hold every warehouse slot. Anything that can reach a warehouse, a semantic
+# layer or a screenshot service stays in the warehouse bound. A listed tool
+# that does reach a warehouse must still take a warehouse slot first (see
+# WorkerCall.admit_warehouse), so a wrong entry cannot overdraw the pool.
+METADATA_ONLY_TOOLS = frozenset(
+ {
+ "find_users",
+ "get_annotation_layer_info",
+ "get_chart_info",
+ "get_chart_type_schema",
+ "get_dashboard_info",
+ "get_database_info",
+ "get_dataset_info",
+ "get_instance_info",
+ "get_layer_annotation_info",
+ "get_query_info",
+ "get_report_info",
+ "get_rls_filter_info",
+ "get_role_info",
+ "get_saved_query_info",
+ "get_schema",
+ "get_tag_info",
+ "get_task_info",
+ "get_theme_info",
+ "get_user_info",
+ "health_check",
+ "list_annotation_layers",
+ "list_charts",
+ "list_dashboards",
+ "list_databases",
+ "list_datasets",
+ "list_layer_annotations",
+ "list_queries",
+ "list_reports",
+ "list_rls_filters",
+ "list_roles",
+ "list_saved_queries",
+ "list_tags",
+ "list_tasks",
+ "list_themes",
+ "list_users",
+ }
+)
+
+
+class WorkerPool:
+ """Bound submissions, including abandoned work; never queue behind a query."""
+
+ def __init__(
+ self,
+ size: int,
+ metadata_size: int = 0,
+ transport_size: int = 1,
+ auth_size: int = 1,
+ auth_pending: int = 32,
+ ) -> None:
+ if size < 1:
+ raise ValueError("MCP_TOOL_WORKERS must be positive")
+ if metadata_size < 0:
+ raise ValueError("MCP_METADATA_TOOL_WORKERS must not be negative")
+ self.slots = threading.BoundedSemaphore(size)
+ self.metadata_slots = (
+ threading.BoundedSemaphore(metadata_size) if metadata_size else None
+ )
+ self.executor = ThreadPoolExecutor(
+ size + metadata_size, thread_name_prefix="mcp-tool"
+ )
+ # Cancellation must not wait behind the warehouse work it is cancelling.
+ # Only calls holding a warehouse slot register cancellation.
+ self.cancel_slots = threading.BoundedSemaphore(size)
+ self.cancellations = ThreadPoolExecutor(size, thread_name_prefix="mcp-cancel")
+ # Transport-side metadata I/O (tools/list filtering, audit writes,
+ # error hooks) is bounded too, so it cannot overdraw the connections
+ # budgeted for tool workers. Callers wait for it on the event loop
+ # without blocking it.
+ self.transport = ThreadPoolExecutor(
+ transport_size, thread_name_prefix="mcp-metadata"
+ )
+ # API-key lookups run before a caller is authenticated, so they get
+ # their own thread: a flood of bad keys or a slow lookup must not
+ # delay audit writes and RBAC filtering. Pending lookups are capped;
+ # past the cap a key is rejected rather than queued.
+ self.auth = ThreadPoolExecutor(auth_size, thread_name_prefix="mcp-auth")
+ self.auth_slots = threading.BoundedSemaphore(auth_pending)
+
+ def admission(self, metadata_only: bool) -> threading.BoundedSemaphore:
+ """Choose the bound a call is admitted under."""
+ if metadata_only and self.metadata_slots is not None:
+ return self.metadata_slots
+ return self.slots
+
+ def submit(
+ self,
+ fn: Callable[[], Any],
+ finished: Callable[[], None],
+ slots: threading.BoundedSemaphore | None = None,
+ ) -> Future[Any]:
+ """Admit immediately or report overload without retaining a queued call."""
+ slots = self.slots if slots is None else slots
+ if not slots.acquire(blocking=False):
+ raise ToolError(BUSY_MESSAGE)
+ try:
+ future = self.executor.submit(fn)
+ except BaseException:
+ slots.release()
+ raise
+ future.add_done_callback(lambda _: finished())
+ return future
+
+ def cancel(self, fn: Callable[[], None], finished: Callable[[], None]) -> bool:
+ """Bound cancellation I/O separately so it cannot delay the caller."""
+ if not self.cancel_slots.acquire(blocking=False):
+ return False
+ 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()
+
+ 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
+
+
+DEFAULT_TOOL_WORKERS = 16
+DEFAULT_METADATA_TOOL_WORKERS = 16
+# Threads, and so metadata connections, for transport-side metadata I/O.
+TRANSPORT_METADATA_THREADS = 1
+# Threads, and so metadata connections, for transport-side API-key lookups.
+API_KEY_AUTH_THREADS = 1
Review Comment:
`CodeQL` is the only red check left, and it is 3 new high alerts rather than infrastructure. I looked at all three and I do not believe any of them is real, but all three land on lines this PR adds, so they will need either a small change here or a dismissal before this can merge.
- This line, logged at `worker.py:349`: the value reaching the log is the integer `1`. The sensitive-data heuristic is keying on the identifier reading like a credential, not on anything that actually flows there.
- `worker.py:301`: the value is `2 * configured + RESERVED_METADATA_THREADS`, an int derived from `app.config.get("MCP_TOOL_WORKERS")`. Values read out of `app.config` are treated as tainted because that same mapping holds secrets elsewhere.
- `composite_token_verifier.py:206`: the value is `username`, the first element of `_validate_api_key_sync`'s documented `(username, scopes)` return. The token itself is never logged.
Renaming `API_KEY_AUTH_THREADS` and `RESERVED_METADATA_THREADS` so neither identifier reads as a credential is the cheapest thing to try for the two in this file, and logging the user id rather than `username` would likely clear the third. I say "try" because I have not re-run the analysis against a rename. Dismissing all three is equally defensible given what they point at, but a red `CodeQL` on added lines tends to stall a merge either way, so it is worth deciding which.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on code in PR #44581:
URL: https://github.com/apache/superset/pull/44581#discussion_r4149559042
##########
superset/mcp_service/worker.py:
##########
@@ -0,0 +1,588 @@
+# 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.
+"""Bounded, deadline-aware execution of MCP tools outside the transport loop.
+
+The tool's coroutine runs on a worker-owned loop because its database APIs are
+synchronous. Only transport notifications are marshalled back to the server
+loop. Flask and SQLAlchemy lifetimes belong to the worker, not the waiting
+request: a timed-out DBAPI call cannot outlive and reuse a torn-down session.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import logging
+import threading
+import time
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import contextmanager
+from contextvars import ContextVar, copy_context
+from typing import Any, Callable, Coroutine, Iterator, ParamSpec, TYPE_CHECKING, TypeVar
+from weakref import WeakKeyDictionary
+
+from fastmcp.exceptions import ToolError
+from flask import current_app, g, has_app_context, has_request_context
+from sqlalchemy import inspect as sa_inspect
+from sqlalchemy.orm.state import InstanceState
+from sqlalchemy.pool import QueuePool
+
+from superset.mcp_service.session_scope import _mcp_session_token
+
+if TYPE_CHECKING:
+ from flask import Flask
+
+ from superset.models.core import Database
+
+logger = logging.getLogger(__name__)
+_active_call: ContextVar[WorkerCall | None] = ContextVar(
+ "mcp_worker_call", default=None
+)
+_metadata_context_owned: ContextVar[bool] = ContextVar(
+ "mcp_metadata_context_owned", default=False
+)
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
+_pools_lock = threading.Lock()
+_pools: WeakKeyDictionary[Flask, WorkerPool] = WeakKeyDictionary()
+
+
+class WorkerDeadlineExceeded(BaseException):
+ """Stop abandoned work without being swallowed by tool error handlers."""
+
+
+class WorkerPool:
+ """Bound submissions, including abandoned work; never queue behind a query."""
+
+ def __init__(self, size: int) -> None:
+ if size < 1:
+ raise ValueError("MCP_TOOL_WORKERS must be positive")
+ self.slots = threading.BoundedSemaphore(size)
+ self.executor = ThreadPoolExecutor(size, thread_name_prefix="mcp-tool")
+ # Cancellation must not wait behind the warehouse work it is cancelling.
+ self.cancel_slots = threading.BoundedSemaphore(size)
+ self.cancellations = ThreadPoolExecutor(size, thread_name_prefix="mcp-cancel")
+
+ def submit(
+ self, fn: Callable[[], Any], finished: Callable[[], None]
+ ) -> Future[Any]:
+ """Admit immediately or report overload without retaining a queued call."""
+ if not self.slots.acquire(blocking=False):
+ raise ToolError(
+ "MCP server busy: all tool workers are occupied. Retry later."
+ )
+ try:
+ future = self.executor.submit(fn)
+ except BaseException:
+ self.slots.release()
+ raise
+ future.add_done_callback(lambda _: finished())
+ return future
+
+ def cancel(self, fn: Callable[[], None], finished: Callable[[], None]) -> bool:
+ """Bound cancellation I/O separately so it cannot delay the caller."""
+ if not self.cancel_slots.acquire(blocking=False):
+ return False
+ try:
+ future = self.cancellations.submit(fn)
+ except RuntimeError:
+ self.cancel_slots.release()
+ return False
+
+ def completed(_: Future[None]) -> None:
+ """Release cancellation capacity before admitting another tool call."""
+ self.cancel_slots.release()
+ finished()
+
+ future.add_done_callback(completed)
+ return True
+
+
+DEFAULT_TOOL_WORKERS = 16
+
+
+def _metadata_pool_capacity(app: Flask) -> int | None:
+ """Return how many metadata connections can be checked out at once.
+
+ ``None`` means a checkout never waits for another holder to return one
+ (e.g. ``NullPool``, per-thread pools, or unlimited overflow).
+ """
+ from superset import db
+
+ with app.app_context():
+ pool = db.engine.pool
+ if not isinstance(pool, QueuePool):
+ return None
+ # SQLAlchemy has no public accessor for the configured overflow limit.
+ max_overflow = pool._max_overflow # pylint: disable=protected-access
+ return None if max_overflow < 0 else pool.size() + max_overflow
+
+
+def tool_worker_count(app: Flask) -> int:
+ """Admit only as many calls as the metadata pool can always serve.
+
+ An admitted call can hold one metadata connection for the whole of its
+ warehouse I/O, and its cancellation needs another. ``2 * workers + 1``
+ connections therefore always leave one that is only held by short metadata
+ lookups, so cancellation and transport-side lookups (tools/list filtering,
+ audit logging) never wait for a warehouse query to end on its own.
+ """
+ configured = app.config.get("MCP_TOOL_WORKERS")
+ capacity = _metadata_pool_capacity(app)
+ if capacity is None:
+ return DEFAULT_TOOL_WORKERS if configured is None else configured
+ limit = (capacity - 1) // 2
+ if limit < 1:
+ raise ValueError(
+ f"The metadata database pool allows {capacity} connections; "
+ "MCP tool execution needs at least 3"
+ )
+ if configured is None:
+ return min(DEFAULT_TOOL_WORKERS, limit)
+ if configured > limit:
+ logger.warning(
+ "MCP_TOOL_WORKERS=%s needs %s metadata database connections, but the "
+ "pool allows %s; admitting %s concurrent tool calls. Raise the pool's "
+ "pool_size/max_overflow in SQLALCHEMY_ENGINE_OPTIONS to admit more.",
+ configured,
+ 2 * configured + 1,
+ capacity,
+ limit,
+ )
+ return limit
+ return configured
+
+
+def _get_pool(app: Flask) -> WorkerPool:
+ """Lazily create a pool for this application, without import-time threads."""
+ with _pools_lock:
+ if app not in _pools:
+ size = tool_worker_count(app)
+ logger.info("MCP tool calls admitted concurrently: %s", size)
+ _pools[app] = WorkerPool(size)
+ return _pools[app]
+
+
+class WorkerCall:
+ """Thread-safe deadline and cancellation registration for one tool call."""
+
+ def __init__(self, app: Flask, pool: WorkerPool, seconds: float) -> None:
+ self.app = app
+ self.pool = pool
+ self.loop = asyncio.get_running_loop()
+ self.seconds = seconds
+ self.deadline = time.monotonic() + seconds
+ from superset.mcp_service.middleware import _mcp_call_id_var
+
+ self.call_id = _mcp_call_id_var.get() or uuid.uuid4().hex
+ self.expired = threading.Event()
+ self.lock = threading.RLock()
+ self.pending = 1
+ self.cancel_query: Callable[[], None] | None = None
+ self.cancel_dispatched = False
+ self.user_id: int | None = None
+
+ def check(self) -> None:
+ """Prevent an abandoned tool from starting more work or mutations."""
+ if self.expired.is_set() or time.monotonic() >= self.deadline:
+ # BaseException deliberately bypasses tools' broad Exception handlers.
+ raise WorkerDeadlineExceeded()
+
+ def dispatch_cancel(self) -> bool:
+ """Dispatch at most once for the active cursor, including late handles."""
+ with self.lock:
+ if self.cancel_query is None or self.cancel_dispatched:
+ return self.cancel_dispatched
+ self.pending += 1
+ self.cancel_dispatched = self.pool.cancel(self.cancel_query, self.finished)
+ if not self.cancel_dispatched:
+ self.pending -= 1
+ return self.cancel_dispatched
+
+ def finished(self) -> None:
+ """Retain admission until both query and cancellation I/O have ended.
+
+ A stuck cancellation therefore cannot consume the cancellation capacity
+ needed by newly admitted queries: it retains its original tool slot.
+ """
+ with self.lock:
+ self.pending -= 1
+ if self.pending == 0:
+ self.pool.slots.release()
+
+ def abandon(self) -> None:
+ """Signal abandonment and dispatch cancellation without blocking asyncio."""
+ self.expired.set()
+ dispatched = self.dispatch_cancel()
+ logger.warning(
+ "MCP call %s exceeded its deadline or disconnected; "
+ "cancellation dispatched=%s; worker slot retained until completion",
+ self.call_id,
+ dispatched,
+ )
+
+
+class TransportContext:
+ """Keep FastMCP's transport-bound async methods on their owning event loop."""
+
+ def __init__(
+ self, context: Any, loop: asyncio.AbstractEventLoop, call: WorkerCall
+ ) -> None:
+ self.context = context
+ self.loop = loop
+ self.call = call
+
+ def __getattr__(self, name: str) -> Any:
+ value = getattr(self.context, name)
+ if not asyncio.iscoroutinefunction(value):
+ return value
+
+ @functools.wraps(value)
+ async def forward(*args: Any, **kwargs: Any) -> Any:
+ self.call.check()
+ future = asyncio.run_coroutine_threadsafe(value(*args, **kwargs), self.loop)
+ try:
+ result = await asyncio.wait_for(
+ asyncio.wrap_future(future),
+ max(0, self.call.deadline - time.monotonic()),
+ )
+ except TimeoutError:
+ raise WorkerDeadlineExceeded() from None
+ self.call.check()
+ return result
+
+ return forward
+
+
+@contextmanager
+def _worker_context(app: Flask) -> Iterator[None]:
+ """Give a worker an independent app context and scoped metadata session."""
+ from superset.mcp_service.auth import _remove_session_safe
+
+ token = _mcp_session_token.set(object())
+ try:
+ with app.app_context():
+ try:
+ # Clean this scope before loading any ORM user into it.
+ _remove_session_safe()
+ yield
+ finally:
+ _remove_session_safe()
+ finally:
+ _mcp_session_token.reset(token)
+
+
+def get_context_user_id() -> int | None:
+ """Read the caller's identity without refreshing an expired ORM instance."""
+ if not has_app_context():
+ return None
+ user = getattr(g, "user", None)
+ state = sa_inspect(user, raiseerr=False)
+ if isinstance(state, InstanceState):
+ return state.identity[0] if state.identity else None
+ return getattr(user, "id", None)
+
+
+async def run_in_metadata_thread(
+ fn: Callable[_P, _T], *args: _P.args, **kwargs: _P.kwargs
+) -> _T:
+ """Run transport metadata I/O with thread-owned Flask and session lifetimes.
+
+ ``to_thread`` copies contextvars, including Flask contexts and MCP session
+ tokens. Replace those owners, reload ORM users, and retain only request and
+ routing data. Cleanup belongs to the thread even if its awaiter disconnects.
+ This uses the metadata executor rather than admitted tool workers, which
+ may themselves be waiting for transport notifications.
+ """
+ from contextlib import nullcontext
+
+ from flask.globals import _cv_request
+
+ if has_app_context():
+ app = current_app._get_current_object()
+ snapshot = dict(vars(g._get_current_object()))
+ else:
+ from superset.mcp_service.flask_singleton import get_flask_app
+
+ app = get_flask_app()
+ snapshot = {}
+ user = snapshot.pop("user", None)
+ state = sa_inspect(user, raiseerr=False)
+ is_orm_user = isinstance(state, InstanceState)
+ user_id = get_context_user_id() if is_orm_user else None
+ request_context = _cv_request.get(None)
+ request_copy = request_context.copy() if request_context is not None else None
+
+ def execute() -> _T:
+ """Own teardown rather than handing a live session back to asyncio."""
+ from superset import db, security_manager
+ from superset.sql.execution.cancellation import without_execution_hooks
+
+ active_token = _active_call.set(None)
+ owner_token = _metadata_context_owned.set(True)
+ try:
+ with without_execution_hooks(), _worker_context(app):
+ vars(g._get_current_object()).update(snapshot)
+ with request_copy if request_copy is not None else nullcontext():
+ if is_orm_user:
+ g.user = (
+ db.session.get(security_manager.user_model, user_id)
+ if user_id is not None
+ else None
+ )
+ elif user is not None:
+ g.user = user
+ return fn(*args, **kwargs)
+ finally:
+ _metadata_context_owned.reset(owner_token)
+ _active_call.reset(active_token)
+
+ return await asyncio.to_thread(execute)
+
+
+async def run_in_worker(
+ fn: Callable[..., Coroutine[Any, Any, Any]],
+ args: tuple[Any, ...],
+ kwargs: dict[str, Any],
+ seconds: float,
+) -> Any:
+ """Run a complete tool lifecycle in a bounded, independently owned worker."""
+ if active := _active_call.get():
+ # Composed tools share the outer deadline and worker/session ownership.
+ # A nested auth hook can inject the original transport Context again.
+ nested_kwargs = dict(kwargs)
+ if "ctx" in nested_kwargs and not isinstance(
+ nested_kwargs["ctx"], TransportContext
+ ):
+ nested_kwargs["ctx"] = TransportContext(
+ nested_kwargs["ctx"], active.loop, active
+ )
+ active.check()
+ return await fn(*args, **nested_kwargs)
+
+ if has_app_context():
+ app = current_app._get_current_object()
+ else:
+ from superset.mcp_service.flask_singleton import get_flask_app
+
+ app = get_flask_app()
+ pool = _get_pool(app)
+ call = WorkerCall(app, pool, seconds)
+ loop = asyncio.get_running_loop()
+ worker_kwargs = dict(kwargs)
+ if "ctx" in worker_kwargs:
+ worker_kwargs["ctx"] = TransportContext(worker_kwargs["ctx"], loop, call)
+ # Copy contextvars (token, tenant routing), but never share Flask g or a
+ # metadata Session. Request-backed middleware's user is reloaded by id.
+ from flask.globals import _cv_request
+
+ request_context = _cv_request.get(None)
+ request_copy = request_context.copy() if request_context is not None else None
+ globals_snapshot = dict(vars(g._get_current_object())) if has_app_context() else {}
+ user = globals_snapshot.pop("user", None) if has_request_context() else None
+ globals_snapshot.pop("user", None)
+ # An expired ORM user's .id can issue metadata I/O on the transport loop.
+ # The identity key is available without loading any attributes.
+ user_state = sa_inspect(user, raiseerr=False)
+ user_id = (
+ user_state.identity[0]
+ if isinstance(user_state, InstanceState) and user_state.identity
+ else getattr(user, "id", None)
+ )
+ guest = user if getattr(user, "is_guest_user", False) else None
+ context = copy_context()
+
+ def execute() -> Any:
+ """Own context teardown even when the caller stops waiting."""
+ from contextlib import nullcontext
+
+ from superset import db, security_manager
+ from superset.sql.execution.cancellation import cursor_scope
+
+ _active_call.set(call)
+ cursor_scope.set(warehouse_cursor)
+ with _worker_context(app):
+ vars(g._get_current_object()).update(globals_snapshot)
+ with request_copy if request_copy is not None else nullcontext():
+ if user_id is not None:
+ g.user = db.session.get(security_manager.user_model, user_id)
+ elif guest is not None:
+ g.user = guest
+ call.check()
+ result = asyncio.run(fn(*args, **worker_kwargs))
+ call.check()
+ return result
+
+ future = pool.submit(lambda: context.run(execute), call.finished)
+ 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
+ finally:
+ from superset.mcp_service.auth import _mcp_user_id_var
+
+ _mcp_user_id_var.set(call.user_id)
+
+
+class QueryCancellation:
+ """Bridge engine cancellation to a separate, freshly scoped metadata session."""
+
+ def __init__(
+ self,
+ call: WorkerCall,
+ database: Database,
+ cursor: Any,
+ catalog: str | None,
+ schema: str | None,
+ ) -> None:
+ self.call = call
+ self.database = database
+ self.database_id = database.id
+ self.spec = database.db_engine_spec
+ self.cursor = cursor
+ self.catalog = catalog
+ self.schema = schema
+ self.cancel_id: str | None = None
+ self.context = copy_context()
+
+ def capture(self) -> None:
+ """Capture an engine handle without making unsupported drivers fail."""
+ from superset.tasks.query_cancel import capture_cancel_query_id
+
+ try:
+ self.cancel_id = capture_cancel_query_id(self.database, self.cursor)
+ except Exception:
+ logger.warning(
+ "MCP call %s: cancel-id capture failed",
+ self.call.call_id,
+ exc_info=True,
+ )
+
+ def register(self) -> None:
+ """Publish a usable handle, including one obtained after the deadline."""
+ with self.call.lock:
+ if self.spec.has_implicit_cancel():
+ self.call.cancel_query = self.cancel_live_cursor
+ elif self.cancel_id is not None:
+ self.call.cancel_query = lambda: self.context.run(self.cancel)
+ else:
+ self.call.cancel_query = None
+ if self.call.expired.is_set():
+ self.call.dispatch_cancel()
+
+ def refresh(self) -> None:
+ """Capture post-execute handles before result fetching can block."""
+ if not self.spec.has_query_id_before_execute:
+ self.capture()
+ self.register()
+ self.call.check()
+
+ def cancel_live_cursor(self) -> None:
+ """Use the same implicit driver cancellation as engine handle_cursor."""
+ try:
+ self.cursor.cancel()
+ except Exception:
+ logger.warning(
+ "MCP call %s: cursor cancellation failed",
+ self.call.call_id,
+ exc_info=True,
+ )
+
+ def cancel(self) -> None:
+ """Resolve a fresh Database and acting user on the cancellation worker."""
+ from superset import db, security_manager
+ from superset.models.core import Database
+ from superset.sql.execution.cancellation import without_execution_hooks
+ from superset.tasks.query_cancel import cancel_chart_query
+
+ active_token = _active_call.set(None)
+ try:
+ with without_execution_hooks(), _worker_context(self.call.app):
+ if self.call.user_id is not None:
+ g.user = db.session.get(
+ security_manager.user_model, self.call.user_id
+ )
+ target = db.session.get(Database, self.database_id)
+ if target is not None and self.cancel_id is not None:
+ cancel_chart_query(
+ target, self.cancel_id, catalog=self.catalog, schema=self.schema
+ )
+ except Exception:
+ logger.warning(
+ "MCP call %s: cancellation failed", self.call.call_id, exc_info=True
+ )
+ finally:
+ _active_call.reset(active_token)
+
+
+@contextmanager
+def warehouse_cursor(
+ database: Database, cursor: Any, catalog: str | None, schema: str | None
+) -> Iterator[None]:
+ """Register Superset's engine cancellation hook for an MCP warehouse call."""
+ call = _active_call.get()
+ if call is None:
+ yield
+ return
+ from superset.sql.execution.cancellation import after_execute, check_deadline
+
+ call.check()
+ cancellation = QueryCancellation(call, database, cursor, catalog, schema)
+ with call.lock:
+ call.cancel_dispatched = False
+ cancellation.capture()
+ cancellation.register()
+ deadline_token = check_deadline.set(call.check)
+ execute_token = after_execute.set(cancellation.refresh)
+ try:
+ call.check()
+ yield
+ call.check()
+ except BaseException:
+ # Do not let error conversion, audit logging or cleanup reuse a metadata
+ # connection held across failed/abandoned warehouse I/O.
+ from superset import db
+
+ # Rollback listeners (including version history) may inspect the live
+ # connection. Finish the transaction before invalidation makes that
+ # inspection fail with PendingRollbackError.
+ db.session.rollback() # pylint: disable=consider-using-transaction
Review Comment:
Fixed in c6b379fa4f, using your second option. The teardown is gone from `warehouse_cursor`. The tool now handles a warehouse error with its session intact, and `_worker_context` rolls back and removes the session once the tool returns.
The old block never invalidated anything anyway. `rollback()` had already returned the connection to the pool, so `invalidate()` had nothing left to act on (checked on SQLAlchemy 2.0.52). The block amounted to a rollback plus a fresh session, and `_worker_context` still provides that.
No deadline or abandoned-call test needed a change. Only `test_failed_warehouse_discards_metadata_session` did, because it asserted the in-flight `invalidate()`. It now asserts that error handling keeps the same usable session and that the next call on the same thread gets a new one, with Continuum's listener present at teardown.
New tests in `test_worker_metadata_pool.py`, both failing on e09b3df:
- `test_warehouse_error_reaches_caller_unchanged` runs a real `Database.get_df` on a missing table inside a worker. The caller gets `sqlite3.OperationalError: no such table: table_that_does_not_exist`, and `database_name`, `backend` and `db_engine_spec` on the row loaded before the failure are still readable.
- `test_execute_sql_reports_missing_table_as_query_error` runs the real `execute_sql` through `build_middleware_list()`. It returns `success: false` with `error: "sqlite error: no such table: no_such_table"`. Nothing is raised, so `_is_user_error` never runs and `MCP_ERROR_HOOK` is not called. On e09b3df the same test gets your `DetachedInstanceError` message, logged as a system error.
##########
superset/mcp_service/worker.py:
##########
@@ -0,0 +1,588 @@
+# 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.
+"""Bounded, deadline-aware execution of MCP tools outside the transport loop.
+
+The tool's coroutine runs on a worker-owned loop because its database APIs are
+synchronous. Only transport notifications are marshalled back to the server
+loop. Flask and SQLAlchemy lifetimes belong to the worker, not the waiting
+request: a timed-out DBAPI call cannot outlive and reuse a torn-down session.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import logging
+import threading
+import time
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import contextmanager
+from contextvars import ContextVar, copy_context
+from typing import Any, Callable, Coroutine, Iterator, ParamSpec, TYPE_CHECKING, TypeVar
+from weakref import WeakKeyDictionary
+
+from fastmcp.exceptions import ToolError
+from flask import current_app, g, has_app_context, has_request_context
+from sqlalchemy import inspect as sa_inspect
+from sqlalchemy.orm.state import InstanceState
+from sqlalchemy.pool import QueuePool
+
+from superset.mcp_service.session_scope import _mcp_session_token
+
+if TYPE_CHECKING:
+ from flask import Flask
+
+ from superset.models.core import Database
+
+logger = logging.getLogger(__name__)
+_active_call: ContextVar[WorkerCall | None] = ContextVar(
+ "mcp_worker_call", default=None
+)
+_metadata_context_owned: ContextVar[bool] = ContextVar(
+ "mcp_metadata_context_owned", default=False
+)
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
+_pools_lock = threading.Lock()
+_pools: WeakKeyDictionary[Flask, WorkerPool] = WeakKeyDictionary()
+
+
+class WorkerDeadlineExceeded(BaseException):
+ """Stop abandoned work without being swallowed by tool error handlers."""
+
+
+class WorkerPool:
+ """Bound submissions, including abandoned work; never queue behind a query."""
+
+ def __init__(self, size: int) -> None:
+ if size < 1:
+ raise ValueError("MCP_TOOL_WORKERS must be positive")
+ self.slots = threading.BoundedSemaphore(size)
+ self.executor = ThreadPoolExecutor(size, thread_name_prefix="mcp-tool")
+ # Cancellation must not wait behind the warehouse work it is cancelling.
+ self.cancel_slots = threading.BoundedSemaphore(size)
+ self.cancellations = ThreadPoolExecutor(size, thread_name_prefix="mcp-cancel")
+
+ def submit(
+ self, fn: Callable[[], Any], finished: Callable[[], None]
+ ) -> Future[Any]:
+ """Admit immediately or report overload without retaining a queued call."""
+ if not self.slots.acquire(blocking=False):
+ raise ToolError(
+ "MCP server busy: all tool workers are occupied. Retry later."
+ )
+ try:
+ future = self.executor.submit(fn)
+ except BaseException:
+ self.slots.release()
+ raise
+ future.add_done_callback(lambda _: finished())
+ return future
+
+ def cancel(self, fn: Callable[[], None], finished: Callable[[], None]) -> bool:
+ """Bound cancellation I/O separately so it cannot delay the caller."""
+ if not self.cancel_slots.acquire(blocking=False):
+ return False
+ try:
+ future = self.cancellations.submit(fn)
+ except RuntimeError:
+ self.cancel_slots.release()
+ return False
+
+ def completed(_: Future[None]) -> None:
+ """Release cancellation capacity before admitting another tool call."""
+ self.cancel_slots.release()
+ finished()
+
+ future.add_done_callback(completed)
+ return True
+
+
+DEFAULT_TOOL_WORKERS = 16
+
+
+def _metadata_pool_capacity(app: Flask) -> int | None:
+ """Return how many metadata connections can be checked out at once.
+
+ ``None`` means a checkout never waits for another holder to return one
+ (e.g. ``NullPool``, per-thread pools, or unlimited overflow).
+ """
+ from superset import db
+
+ with app.app_context():
+ pool = db.engine.pool
+ if not isinstance(pool, QueuePool):
+ return None
+ # SQLAlchemy has no public accessor for the configured overflow limit.
+ max_overflow = pool._max_overflow # pylint: disable=protected-access
+ return None if max_overflow < 0 else pool.size() + max_overflow
+
+
+def tool_worker_count(app: Flask) -> int:
+ """Admit only as many calls as the metadata pool can always serve.
+
+ An admitted call can hold one metadata connection for the whole of its
+ warehouse I/O, and its cancellation needs another. ``2 * workers + 1``
+ connections therefore always leave one that is only held by short metadata
+ lookups, so cancellation and transport-side lookups (tools/list filtering,
+ audit logging) never wait for a warehouse query to end on its own.
+ """
+ configured = app.config.get("MCP_TOOL_WORKERS")
+ capacity = _metadata_pool_capacity(app)
+ if capacity is None:
+ return DEFAULT_TOOL_WORKERS if configured is None else configured
+ limit = (capacity - 1) // 2
Review Comment:
It wasn't deliberate. Fixed in 6eed3e4d10.
- The `(pool - 1) // 2` bound (`MCP_TOOL_WORKERS`) still covers anything that can reach a warehouse. Semantic-layer, chart data/preview, SQL and write tools stay under it.
- An explicit `METADATA_ONLY_TOOLS` list gets its own bound, `MCP_METADATA_TOOL_WORKERS` (default 16; `0` puts them back under the strict bound). The list is the read-only `list_*` and `get_*_info` tools plus `get_schema`, `get_instance_info`, `health_check`, `find_users` and `get_chart_type_schema`. These tools hold a metadata connection only for their own short queries, so the reservation math doesn't change. They wait for a connection on a worker thread, never on the loop.
- Backstop: `get_sqla_engine` enters `cancellable_engine` before any warehouse connection opens, prequeries included. There, a call admitted as metadata-only has to take a warehouse slot or fail fast as busy, so a misclassified tool can't overdraw the pool.
Tests:
- `test_metadata_only_tools_answer_while_warehouse_bound_is_full` fills the warehouse bound with calls holding a metadata connection, checks that one more is busy, then runs 12 concurrent `list_charts` calls, which all succeed.
- `test_metadata_only_call_takes_warehouse_slot_before_warehouse_io` covers the backstop.
- `test_metadata_only_tools_are_registered_tools` checks that every entry is a registered read-only tool.
- `test_metadata_only_tools_have_their_own_bound` covers the new setting.
The docs and `mcp_config.py` are updated.
##########
superset/mcp_service/worker.py:
##########
@@ -0,0 +1,588 @@
+# 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.
+"""Bounded, deadline-aware execution of MCP tools outside the transport loop.
+
+The tool's coroutine runs on a worker-owned loop because its database APIs are
+synchronous. Only transport notifications are marshalled back to the server
+loop. Flask and SQLAlchemy lifetimes belong to the worker, not the waiting
+request: a timed-out DBAPI call cannot outlive and reuse a torn-down session.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import logging
+import threading
+import time
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import contextmanager
+from contextvars import ContextVar, copy_context
+from typing import Any, Callable, Coroutine, Iterator, ParamSpec, TYPE_CHECKING, TypeVar
+from weakref import WeakKeyDictionary
+
+from fastmcp.exceptions import ToolError
+from flask import current_app, g, has_app_context, has_request_context
+from sqlalchemy import inspect as sa_inspect
+from sqlalchemy.orm.state import InstanceState
+from sqlalchemy.pool import QueuePool
+
+from superset.mcp_service.session_scope import _mcp_session_token
+
+if TYPE_CHECKING:
+ from flask import Flask
+
+ from superset.models.core import Database
+
+logger = logging.getLogger(__name__)
+_active_call: ContextVar[WorkerCall | None] = ContextVar(
+ "mcp_worker_call", default=None
+)
+_metadata_context_owned: ContextVar[bool] = ContextVar(
+ "mcp_metadata_context_owned", default=False
+)
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
+_pools_lock = threading.Lock()
+_pools: WeakKeyDictionary[Flask, WorkerPool] = WeakKeyDictionary()
+
+
+class WorkerDeadlineExceeded(BaseException):
+ """Stop abandoned work without being swallowed by tool error handlers."""
+
+
+class WorkerPool:
+ """Bound submissions, including abandoned work; never queue behind a query."""
+
+ def __init__(self, size: int) -> None:
+ if size < 1:
+ raise ValueError("MCP_TOOL_WORKERS must be positive")
+ self.slots = threading.BoundedSemaphore(size)
+ self.executor = ThreadPoolExecutor(size, thread_name_prefix="mcp-tool")
+ # Cancellation must not wait behind the warehouse work it is cancelling.
+ self.cancel_slots = threading.BoundedSemaphore(size)
+ self.cancellations = ThreadPoolExecutor(size, thread_name_prefix="mcp-cancel")
+
+ def submit(
+ self, fn: Callable[[], Any], finished: Callable[[], None]
+ ) -> Future[Any]:
+ """Admit immediately or report overload without retaining a queued call."""
+ if not self.slots.acquire(blocking=False):
+ raise ToolError(
+ "MCP server busy: all tool workers are occupied. Retry later."
+ )
+ try:
+ future = self.executor.submit(fn)
+ except BaseException:
+ self.slots.release()
+ raise
+ future.add_done_callback(lambda _: finished())
+ return future
+
+ def cancel(self, fn: Callable[[], None], finished: Callable[[], None]) -> bool:
+ """Bound cancellation I/O separately so it cannot delay the caller."""
+ if not self.cancel_slots.acquire(blocking=False):
+ return False
+ try:
+ future = self.cancellations.submit(fn)
+ except RuntimeError:
+ self.cancel_slots.release()
+ return False
+
+ def completed(_: Future[None]) -> None:
+ """Release cancellation capacity before admitting another tool call."""
+ self.cancel_slots.release()
+ finished()
+
+ future.add_done_callback(completed)
+ return True
+
+
+DEFAULT_TOOL_WORKERS = 16
+
+
+def _metadata_pool_capacity(app: Flask) -> int | None:
+ """Return how many metadata connections can be checked out at once.
+
+ ``None`` means a checkout never waits for another holder to return one
+ (e.g. ``NullPool``, per-thread pools, or unlimited overflow).
+ """
+ from superset import db
+
+ with app.app_context():
+ pool = db.engine.pool
+ if not isinstance(pool, QueuePool):
+ return None
+ # SQLAlchemy has no public accessor for the configured overflow limit.
+ max_overflow = pool._max_overflow # pylint: disable=protected-access
+ return None if max_overflow < 0 else pool.size() + max_overflow
+
+
+def tool_worker_count(app: Flask) -> int:
+ """Admit only as many calls as the metadata pool can always serve.
+
+ An admitted call can hold one metadata connection for the whole of its
+ warehouse I/O, and its cancellation needs another. ``2 * workers + 1``
+ connections therefore always leave one that is only held by short metadata
+ lookups, so cancellation and transport-side lookups (tools/list filtering,
+ audit logging) never wait for a warehouse query to end on its own.
+ """
+ configured = app.config.get("MCP_TOOL_WORKERS")
+ capacity = _metadata_pool_capacity(app)
+ if capacity is None:
+ return DEFAULT_TOOL_WORKERS if configured is None else configured
+ limit = (capacity - 1) // 2
+ if limit < 1:
+ raise ValueError(
+ f"The metadata database pool allows {capacity} connections; "
+ "MCP tool execution needs at least 3"
+ )
+ if configured is None:
+ return min(DEFAULT_TOOL_WORKERS, limit)
+ if configured > limit:
+ logger.warning(
+ "MCP_TOOL_WORKERS=%s needs %s metadata database connections, but the "
+ "pool allows %s; admitting %s concurrent tool calls. Raise the pool's "
+ "pool_size/max_overflow in SQLALCHEMY_ENGINE_OPTIONS to admit more.",
+ configured,
+ 2 * configured + 1,
+ capacity,
+ limit,
+ )
+ return limit
+ return configured
+
+
+def _get_pool(app: Flask) -> WorkerPool:
+ """Lazily create a pool for this application, without import-time threads."""
+ with _pools_lock:
+ if app not in _pools:
+ size = tool_worker_count(app)
+ logger.info("MCP tool calls admitted concurrently: %s", size)
+ _pools[app] = WorkerPool(size)
+ return _pools[app]
+
+
+class WorkerCall:
+ """Thread-safe deadline and cancellation registration for one tool call."""
+
+ def __init__(self, app: Flask, pool: WorkerPool, seconds: float) -> None:
+ self.app = app
+ self.pool = pool
+ self.loop = asyncio.get_running_loop()
+ self.seconds = seconds
+ self.deadline = time.monotonic() + seconds
+ from superset.mcp_service.middleware import _mcp_call_id_var
+
+ self.call_id = _mcp_call_id_var.get() or uuid.uuid4().hex
+ self.expired = threading.Event()
+ self.lock = threading.RLock()
+ self.pending = 1
+ self.cancel_query: Callable[[], None] | None = None
+ self.cancel_dispatched = False
+ self.user_id: int | None = None
+
+ def check(self) -> None:
+ """Prevent an abandoned tool from starting more work or mutations."""
+ if self.expired.is_set() or time.monotonic() >= self.deadline:
+ # BaseException deliberately bypasses tools' broad Exception handlers.
+ raise WorkerDeadlineExceeded()
+
+ def dispatch_cancel(self) -> bool:
+ """Dispatch at most once for the active cursor, including late handles."""
+ with self.lock:
+ if self.cancel_query is None or self.cancel_dispatched:
+ return self.cancel_dispatched
+ self.pending += 1
+ self.cancel_dispatched = self.pool.cancel(self.cancel_query, self.finished)
+ if not self.cancel_dispatched:
+ self.pending -= 1
+ return self.cancel_dispatched
+
+ def finished(self) -> None:
+ """Retain admission until both query and cancellation I/O have ended.
+
+ A stuck cancellation therefore cannot consume the cancellation capacity
+ needed by newly admitted queries: it retains its original tool slot.
+ """
+ with self.lock:
+ self.pending -= 1
+ if self.pending == 0:
+ self.pool.slots.release()
+
+ def abandon(self) -> None:
+ """Signal abandonment and dispatch cancellation without blocking asyncio."""
+ self.expired.set()
+ dispatched = self.dispatch_cancel()
+ logger.warning(
+ "MCP call %s exceeded its deadline or disconnected; "
+ "cancellation dispatched=%s; worker slot retained until completion",
+ self.call_id,
+ dispatched,
+ )
+
+
+class TransportContext:
+ """Keep FastMCP's transport-bound async methods on their owning event loop."""
+
+ def __init__(
+ self, context: Any, loop: asyncio.AbstractEventLoop, call: WorkerCall
+ ) -> None:
+ self.context = context
+ self.loop = loop
+ self.call = call
+
+ def __getattr__(self, name: str) -> Any:
+ value = getattr(self.context, name)
+ if not asyncio.iscoroutinefunction(value):
+ return value
+
+ @functools.wraps(value)
+ async def forward(*args: Any, **kwargs: Any) -> Any:
+ self.call.check()
+ future = asyncio.run_coroutine_threadsafe(value(*args, **kwargs), self.loop)
+ try:
+ result = await asyncio.wait_for(
+ asyncio.wrap_future(future),
+ max(0, self.call.deadline - time.monotonic()),
+ )
+ except TimeoutError:
+ raise WorkerDeadlineExceeded() from None
+ self.call.check()
+ return result
+
+ return forward
+
+
+@contextmanager
+def _worker_context(app: Flask) -> Iterator[None]:
+ """Give a worker an independent app context and scoped metadata session."""
+ from superset.mcp_service.auth import _remove_session_safe
+
+ token = _mcp_session_token.set(object())
+ try:
+ with app.app_context():
+ try:
+ # Clean this scope before loading any ORM user into it.
+ _remove_session_safe()
+ yield
+ finally:
+ _remove_session_safe()
+ finally:
+ _mcp_session_token.reset(token)
+
+
+def get_context_user_id() -> int | None:
+ """Read the caller's identity without refreshing an expired ORM instance."""
+ if not has_app_context():
+ return None
+ user = getattr(g, "user", None)
+ state = sa_inspect(user, raiseerr=False)
+ if isinstance(state, InstanceState):
+ return state.identity[0] if state.identity else None
+ return getattr(user, "id", None)
+
+
+async def run_in_metadata_thread(
+ fn: Callable[_P, _T], *args: _P.args, **kwargs: _P.kwargs
+) -> _T:
+ """Run transport metadata I/O with thread-owned Flask and session lifetimes.
+
+ ``to_thread`` copies contextvars, including Flask contexts and MCP session
+ tokens. Replace those owners, reload ORM users, and retain only request and
+ routing data. Cleanup belongs to the thread even if its awaiter disconnects.
+ This uses the metadata executor rather than admitted tool workers, which
+ may themselves be waiting for transport notifications.
+ """
+ from contextlib import nullcontext
+
+ from flask.globals import _cv_request
+
+ if has_app_context():
+ app = current_app._get_current_object()
+ snapshot = dict(vars(g._get_current_object()))
+ else:
+ from superset.mcp_service.flask_singleton import get_flask_app
+
+ app = get_flask_app()
+ snapshot = {}
+ user = snapshot.pop("user", None)
+ state = sa_inspect(user, raiseerr=False)
+ is_orm_user = isinstance(state, InstanceState)
+ user_id = get_context_user_id() if is_orm_user else None
+ request_context = _cv_request.get(None)
+ request_copy = request_context.copy() if request_context is not None else None
+
+ def execute() -> _T:
+ """Own teardown rather than handing a live session back to asyncio."""
+ from superset import db, security_manager
+ from superset.sql.execution.cancellation import without_execution_hooks
+
+ active_token = _active_call.set(None)
+ owner_token = _metadata_context_owned.set(True)
+ try:
+ with without_execution_hooks(), _worker_context(app):
+ vars(g._get_current_object()).update(snapshot)
+ with request_copy if request_copy is not None else nullcontext():
+ if is_orm_user:
+ g.user = (
+ db.session.get(security_manager.user_model, user_id)
+ if user_id is not None
+ else None
+ )
+ elif user is not None:
+ g.user = user
+ return fn(*args, **kwargs)
+ finally:
+ _metadata_context_owned.reset(owner_token)
+ _active_call.reset(active_token)
+
+ return await asyncio.to_thread(execute)
+
+
+async def run_in_worker(
+ fn: Callable[..., Coroutine[Any, Any, Any]],
+ args: tuple[Any, ...],
+ kwargs: dict[str, Any],
+ seconds: float,
+) -> Any:
+ """Run a complete tool lifecycle in a bounded, independently owned worker."""
+ if active := _active_call.get():
+ # Composed tools share the outer deadline and worker/session ownership.
+ # A nested auth hook can inject the original transport Context again.
+ nested_kwargs = dict(kwargs)
+ if "ctx" in nested_kwargs and not isinstance(
+ nested_kwargs["ctx"], TransportContext
+ ):
+ nested_kwargs["ctx"] = TransportContext(
+ nested_kwargs["ctx"], active.loop, active
+ )
+ active.check()
+ return await fn(*args, **nested_kwargs)
+
+ if has_app_context():
+ app = current_app._get_current_object()
+ else:
+ from superset.mcp_service.flask_singleton import get_flask_app
+
+ app = get_flask_app()
+ pool = _get_pool(app)
+ call = WorkerCall(app, pool, seconds)
+ loop = asyncio.get_running_loop()
+ worker_kwargs = dict(kwargs)
+ if "ctx" in worker_kwargs:
+ worker_kwargs["ctx"] = TransportContext(worker_kwargs["ctx"], loop, call)
+ # Copy contextvars (token, tenant routing), but never share Flask g or a
+ # metadata Session. Request-backed middleware's user is reloaded by id.
+ from flask.globals import _cv_request
+
+ request_context = _cv_request.get(None)
+ request_copy = request_context.copy() if request_context is not None else None
+ globals_snapshot = dict(vars(g._get_current_object())) if has_app_context() else {}
+ user = globals_snapshot.pop("user", None) if has_request_context() else None
+ globals_snapshot.pop("user", None)
+ # An expired ORM user's .id can issue metadata I/O on the transport loop.
+ # The identity key is available without loading any attributes.
+ user_state = sa_inspect(user, raiseerr=False)
+ user_id = (
+ user_state.identity[0]
+ if isinstance(user_state, InstanceState) and user_state.identity
+ else getattr(user, "id", None)
+ )
+ guest = user if getattr(user, "is_guest_user", False) else None
+ context = copy_context()
+
+ def execute() -> Any:
+ """Own context teardown even when the caller stops waiting."""
+ from contextlib import nullcontext
+
+ from superset import db, security_manager
+ from superset.sql.execution.cancellation import cursor_scope
+
+ _active_call.set(call)
+ cursor_scope.set(warehouse_cursor)
+ with _worker_context(app):
+ vars(g._get_current_object()).update(globals_snapshot)
+ with request_copy if request_copy is not None else nullcontext():
+ if user_id is not None:
+ g.user = db.session.get(security_manager.user_model, user_id)
+ elif guest is not None:
+ g.user = guest
+ call.check()
+ result = asyncio.run(fn(*args, **worker_kwargs))
+ call.check()
+ return result
+
+ future = pool.submit(lambda: context.run(execute), call.finished)
+ 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
+ finally:
+ from superset.mcp_service.auth import _mcp_user_id_var
+
+ _mcp_user_id_var.set(call.user_id)
+
+
+class QueryCancellation:
+ """Bridge engine cancellation to a separate, freshly scoped metadata session."""
+
+ def __init__(
+ self,
+ call: WorkerCall,
+ database: Database,
+ cursor: Any,
+ catalog: str | None,
+ schema: str | None,
+ ) -> None:
+ self.call = call
+ self.database = database
+ self.database_id = database.id
+ self.spec = database.db_engine_spec
+ self.cursor = cursor
+ self.catalog = catalog
+ self.schema = schema
+ self.cancel_id: str | None = None
+ self.context = copy_context()
+
+ def capture(self) -> None:
+ """Capture an engine handle without making unsupported drivers fail."""
+ from superset.tasks.query_cancel import capture_cancel_query_id
+
+ try:
+ self.cancel_id = capture_cancel_query_id(self.database, self.cursor)
+ except Exception:
+ logger.warning(
+ "MCP call %s: cancel-id capture failed",
+ self.call.call_id,
+ exc_info=True,
+ )
+
+ def register(self) -> None:
+ """Publish a usable handle, including one obtained after the deadline."""
+ with self.call.lock:
+ if self.spec.has_implicit_cancel():
Review Comment:
Understood and accepted, and it's documented. I didn't reuse the SQL Lab path because it would need new Trino-specific code:
- `TrinoEngineSpec.handle_cursor` writes `QUERY_CANCEL_KEY` onto a persisted SQL Lab `Query` row and commits it, then polls the cursor while `execute_with_cursor` runs `execute` on another thread. MCP's chart-data (`get_df`) and `SQLExecutor` paths have no `Query` row and call `db_engine_spec.execute` directly. That call blocks inside the Trino client until the query ends, so capturing the id after execute comes too late.
- The `query_cancel` seam goes through `get_cancel_query_id`, and Trino inherits `None` from `BaseEngineSpec`, as do BigQuery, Databricks and Athena. Supporting it would take a Trino override that reads `cursor.query_id` from the cancellation thread while `execute` is still blocked. That is new per-driver behaviour, better handled in its own change that also covers the other engines.
So for those engines, abandonment logs `dispatched=False` and the call keeps its slot until the driver returns. The docs say so. What keeps the process responsive meanwhile is the separate metadata-only bound from 6eed3e4d10 (see the thread above): list/get tools, `tools/list` and audit logging keep answering, and only additional warehouse calls get `MCP server busy`.
##########
tests/unit_tests/conftest.py:
##########
@@ -43,7 +44,11 @@ def get_session(mocker: MockerFixture) -> Callable[[], Session]:
"""
Create an in-memory SQLite db.session.to test models.
"""
- engine = create_engine("sqlite://")
+ # MCP unit tests hand this injected Session from setup to a tool worker
+ # sequentially. Concurrency/session ownership is tested with scoped sessions.
+ engine = create_engine(
+ "sqlite://", connect_args={"check_same_thread": False}, poolclass=StaticPool
Review Comment:
Done in 95fe5d0aa4. The shared `session` fixture is back to master's plain `create_engine("sqlite://")`, now supplied by an overridable `session_engine` fixture. Only the two modules that hand their session to a tool worker override it with `check_same_thread=False` and `StaticPool`:
- `test_get_dashboard_datasets.py` (`test_semantic_layers_are_loaded_before_serialization`)
- `test_update_dataset_metric.py` (`test_metric_read_modify_write_persists_clean_value_through_real_command`)
With the shared fixture reverted, those tests accounted for the only 3 failures in `tests/unit_tests/mcp_service`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on PR #44581:
URL: https://github.com/apache/superset/pull/44581#issuecomment-5842370750
AI Code Review is in progress (usually takes 3 to 15 minutes unless it's a very large PR).
<details>
<summary>
Bito Usage Guide
</summary>
**Commands**
Type the following command in the pull request comment and save the comment.
- `/review` - Manually triggers an incremental AI Review.
- `/review full` - Manually triggers a full AI Review.
- `/pause` - Pauses automatic reviews on this pull request.
- `/resume` - Resumes automatic reviews.
- `/resolve` - Marks all Bito-posted review comments as resolved.
- `/abort` - Cancels all in-progress reviews.
Refer to the <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/available-commands" target="_blank" rel="noopener noreferrer">documentation</a> for additional commands.
**Configuration**
This repository uses `Superset` You can customize the agent settings <a href="https://alpha.bito.ai/home/ai-agents/code-review-agent" target="_blank" rel="noopener noreferrer">here</a> or contact your Bito workspace admin at evan@preset.io.
**Documentation & Help**
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/getting-started/install-run-using-bito-cloud/create-or-customize-an-agent-instance" target="_blank" rel="noopener noreferrer">Customize agent settings</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/implementing-custom-code-review-rules" target="_blank" rel="noopener noreferrer">Review rules</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent" target="_blank" rel="noopener noreferrer">General documentation</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/faqs" target="_blank" rel="noopener noreferrer">FAQ</a>
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on PR #44581:
URL: https://github.com/apache/superset/pull/44581#issuecomment-5943202768
**Bito Automatic Review Skipped โ PR Already Merged**<br/><br/>Bito scheduled an automatic review for this pull request, but the review was skipped because this PR was merged before the review could be run.<br/>No action is needed if you didn't intend to review it. To get a review, you can type `/review` in a comment and save it
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on code in PR #44581:
URL: https://github.com/apache/superset/pull/44581#discussion_r4160690119
##########
superset/mcp_service/worker.py:
##########
@@ -0,0 +1,910 @@
+# 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.
+"""Bounded, deadline-aware execution of MCP tools outside the transport loop.
+
+The tool's coroutine runs on a worker-owned loop because its database APIs are
+synchronous. Only transport notifications are marshalled back to the server
+loop. Flask and SQLAlchemy lifetimes belong to the worker, not the waiting
+request: a timed-out DBAPI call cannot outlive and reuse a torn-down session.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import logging
+import threading
+import time
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import contextmanager
+from contextvars import ContextVar, copy_context
+from typing import Any, Callable, Coroutine, Iterator, ParamSpec, TYPE_CHECKING, TypeVar
+from weakref import WeakKeyDictionary
+
+from fastmcp.exceptions import ToolError
+from flask import current_app, g, has_app_context, has_request_context
+from sqlalchemy import inspect as sa_inspect
+from sqlalchemy.orm.state import InstanceState
+from sqlalchemy.pool import QueuePool
+
+from superset.mcp_service.session_scope import _mcp_session_token
+
+if TYPE_CHECKING:
+ from flask import Flask
+
+ from superset.models.core import Database
+
+logger = logging.getLogger(__name__)
+_active_call: ContextVar[WorkerCall | None] = ContextVar(
+ "mcp_worker_call", default=None
+)
+_metadata_context_owned: ContextVar[bool] = ContextVar(
+ "mcp_metadata_context_owned", default=False
+)
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
+_pools_lock = threading.Lock()
+_pools: WeakKeyDictionary[Flask, WorkerPool] = WeakKeyDictionary()
+
+
+class WorkerDeadlineExceeded(BaseException):
+ """Stop abandoned work without being swallowed by tool error handlers."""
+
+
+BUSY_MESSAGE = "MCP server busy: all tool workers are occupied. Retry later."
+
+# Tools that only read Superset's metadata database. They are admitted under a
+# separate, larger bound so they keep answering while slow warehouse queries
+# hold every warehouse slot. Anything that can reach a warehouse, a semantic
+# layer or a screenshot service stays in the warehouse bound. A listed tool
+# that does reach a warehouse must still take a warehouse slot first (see
+# WorkerCall.admit_warehouse), so a wrong entry cannot overdraw the pool.
+METADATA_ONLY_TOOLS = frozenset(
+ {
+ "find_users",
+ "get_annotation_layer_info",
+ "get_chart_info",
+ "get_chart_type_schema",
+ "get_dashboard_info",
+ "get_database_info",
+ "get_dataset_info",
+ "get_instance_info",
+ "get_layer_annotation_info",
+ "get_query_info",
+ "get_report_info",
+ "get_rls_filter_info",
+ "get_role_info",
+ "get_saved_query_info",
+ "get_schema",
+ "get_tag_info",
+ "get_task_info",
+ "get_theme_info",
+ "get_user_info",
+ "health_check",
+ "list_annotation_layers",
+ "list_charts",
+ "list_dashboards",
+ "list_databases",
+ "list_datasets",
+ "list_layer_annotations",
+ "list_queries",
+ "list_reports",
+ "list_rls_filters",
+ "list_roles",
+ "list_saved_queries",
+ "list_tags",
+ "list_tasks",
+ "list_themes",
+ "list_users",
+ }
+)
+
+
+class WorkerPool:
+ """Bound submissions, including abandoned work; never queue behind a query."""
+
+ def __init__(
+ self,
+ size: int,
+ metadata_size: int = 0,
+ transport_size: int = 1,
+ auth_size: int = 1,
+ auth_pending: int = 32,
+ ) -> None:
+ if size < 1:
+ raise ValueError("MCP_TOOL_WORKERS must be positive")
+ if metadata_size < 0:
+ raise ValueError("MCP_METADATA_TOOL_WORKERS must not be negative")
+ self.slots = threading.BoundedSemaphore(size)
+ self.metadata_slots = (
+ threading.BoundedSemaphore(metadata_size) if metadata_size else None
+ )
+ self.executor = ThreadPoolExecutor(
+ size + metadata_size, thread_name_prefix="mcp-tool"
+ )
+ # Cancellation must not wait behind the warehouse work it is cancelling.
+ # Only calls holding a warehouse slot register cancellation.
+ self.cancel_slots = threading.BoundedSemaphore(size)
+ self.cancellations = ThreadPoolExecutor(size, thread_name_prefix="mcp-cancel")
+ # Transport-side metadata I/O (tools/list filtering, audit writes,
+ # error hooks) is bounded too, so it cannot overdraw the connections
+ # budgeted for tool workers. Callers wait for it on the event loop
+ # without blocking it.
+ self.transport = ThreadPoolExecutor(
+ transport_size, thread_name_prefix="mcp-metadata"
+ )
+ # API-key lookups run before a caller is authenticated, so they get
+ # their own thread: a flood of bad keys or a slow lookup must not
+ # delay audit writes and RBAC filtering. Pending lookups are capped;
+ # past the cap a key is rejected rather than queued.
+ self.auth = ThreadPoolExecutor(auth_size, thread_name_prefix="mcp-auth")
+ self.auth_slots = threading.BoundedSemaphore(auth_pending)
+
+ def admission(self, metadata_only: bool) -> threading.BoundedSemaphore:
+ """Choose the bound a call is admitted under."""
+ if metadata_only and self.metadata_slots is not None:
+ return self.metadata_slots
+ return self.slots
+
+ def submit(
+ self,
+ fn: Callable[[], Any],
+ finished: Callable[[], None],
+ slots: threading.BoundedSemaphore | None = None,
+ ) -> Future[Any]:
+ """Admit immediately or report overload without retaining a queued call."""
+ slots = self.slots if slots is None else slots
+ if not slots.acquire(blocking=False):
+ raise ToolError(BUSY_MESSAGE)
+ try:
+ future = self.executor.submit(fn)
+ except BaseException:
+ slots.release()
+ raise
+ future.add_done_callback(lambda _: finished())
+ return future
+
+ def cancel(self, fn: Callable[[], None], finished: Callable[[], None]) -> bool:
+ """Bound cancellation I/O separately so it cannot delay the caller."""
+ if not self.cancel_slots.acquire(blocking=False):
+ return False
+ 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()
+
+ 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
+
+
+DEFAULT_TOOL_WORKERS = 16
+DEFAULT_METADATA_TOOL_WORKERS = 16
+# Threads, and so metadata connections, for transport-side metadata I/O.
+TRANSPORT_METADATA_THREADS = 1
+# Threads, and so metadata connections, for transport-side API-key lookups.
+API_KEY_CHECK_POOL_SIZE = 1
+# API-key lookups admitted at once, running or waiting; more are rejected.
+API_KEY_AUTH_PENDING = 32
+# Metadata connections held by transport-side threads rather than tool calls.
+METADATA_POOL_HEADROOM = TRANSPORT_METADATA_THREADS + API_KEY_CHECK_POOL_SIZE
+
+
+def _metadata_pool_capacity(app: Flask) -> int | None:
+ """Return how many metadata connections can be checked out at once.
+
+ ``None`` means a checkout never waits for another holder to return one
+ (e.g. ``NullPool``, per-thread pools, or unlimited overflow), including
+ for an application that never registered the metadata database: it has
+ no metadata connections to budget, and resolving ``db.engine`` for it
+ would raise rather than report a capacity.
+ """
+ from contextlib import nullcontext
+
+ from superset import db
+
+ if app.extensions.get("sqlalchemy") is not db:
+ return None
+
+ # Popping a pushed context tears down the caller's scoped session, and
+ # pools are created lazily from transport-side code that holds one.
+ in_app = has_app_context() and current_app._get_current_object() is app
+ with nullcontext() if in_app else app.app_context():
+ pool = db.engine.pool
+ if not isinstance(pool, QueuePool):
+ return None
+ # SQLAlchemy has no public accessor for the configured overflow limit.
+ max_overflow = pool._max_overflow # pylint: disable=protected-access
+ return None if max_overflow < 0 else pool.size() + max_overflow
+
+
+def admission_counts(app: Flask) -> tuple[int, int]:
+ """Split the metadata pool between warehouse and metadata-only calls.
+
+ Every holder of a metadata connection is bounded and budgeted:
+
+ - each warehouse-capable call can hold one for the whole of its warehouse
+ I/O, and its cancellation needs another (``2 * workers``);
+ - each metadata-only call holds at most one (``metadata_workers``);
+ - transport-side metadata I/O runs on ``TRANSPORT_METADATA_THREADS`` and
+ API-key lookups on ``API_KEY_CHECK_POOL_SIZE``.
+
+ ``2 * workers + metadata_workers + METADATA_POOL_HEADROOM`` never
+ exceeds the pool's capacity, so with every slot admitted no checkout
+ waits for another holder. Unless configured, warehouse calls get about
+ two thirds of the remaining connections and metadata-only calls the
+ rest, each up to 16. Explicit settings above the budget are reduced
+ with a warning. ``0`` metadata-only workers admits those tools under
+ the warehouse bound.
+
+ Returns ``(workers, metadata_workers)``.
+ """
+ configured = app.config.get("MCP_TOOL_WORKERS")
+ configured_metadata = app.config.get("MCP_METADATA_TOOL_WORKERS")
+ if configured_metadata is not None and configured_metadata < 0:
+ raise ValueError("MCP_METADATA_TOOL_WORKERS must not be negative")
+ capacity = _metadata_pool_capacity(app)
+ if capacity is None:
+ return (
+ DEFAULT_TOOL_WORKERS if configured is None else configured,
+ DEFAULT_METADATA_TOOL_WORKERS
+ if configured_metadata is None
+ else configured_metadata,
+ )
+ available = capacity - METADATA_POOL_HEADROOM
+ limit = available // 2
+ if limit < 1:
+ raise ValueError(
+ f"The metadata database pool allows {capacity} connections; "
+ "MCP tool execution needs at least "
+ f"{2 + METADATA_POOL_HEADROOM}"
+ )
+ if configured is None:
+ workers = min(DEFAULT_TOOL_WORKERS, max(1, available // 3))
+ elif configured > limit:
+ logger.warning(
+ "MCP_TOOL_WORKERS=%s needs %s metadata database connections, but the "
+ "pool allows %s; admitting %s concurrent tool calls. Raise the pool's "
+ "pool_size/max_overflow in SQLALCHEMY_ENGINE_OPTIONS to admit more.",
+ configured,
+ 2 * configured + METADATA_POOL_HEADROOM,
Review Comment:
Fixed in 310379c4ea: pool initialization and clamping logs use static messages, with no configuration-, engine-, or pool-count-derived arguments; admission behavior is unchanged. CodeQL has re-run successfully with no new alerts.
##########
superset/mcp_service/worker.py:
##########
@@ -0,0 +1,910 @@
+# 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.
+"""Bounded, deadline-aware execution of MCP tools outside the transport loop.
+
+The tool's coroutine runs on a worker-owned loop because its database APIs are
+synchronous. Only transport notifications are marshalled back to the server
+loop. Flask and SQLAlchemy lifetimes belong to the worker, not the waiting
+request: a timed-out DBAPI call cannot outlive and reuse a torn-down session.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import logging
+import threading
+import time
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import contextmanager
+from contextvars import ContextVar, copy_context
+from typing import Any, Callable, Coroutine, Iterator, ParamSpec, TYPE_CHECKING, TypeVar
+from weakref import WeakKeyDictionary
+
+from fastmcp.exceptions import ToolError
+from flask import current_app, g, has_app_context, has_request_context
+from sqlalchemy import inspect as sa_inspect
+from sqlalchemy.orm.state import InstanceState
+from sqlalchemy.pool import QueuePool
+
+from superset.mcp_service.session_scope import _mcp_session_token
+
+if TYPE_CHECKING:
+ from flask import Flask
+
+ from superset.models.core import Database
+
+logger = logging.getLogger(__name__)
+_active_call: ContextVar[WorkerCall | None] = ContextVar(
+ "mcp_worker_call", default=None
+)
+_metadata_context_owned: ContextVar[bool] = ContextVar(
+ "mcp_metadata_context_owned", default=False
+)
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
+_pools_lock = threading.Lock()
+_pools: WeakKeyDictionary[Flask, WorkerPool] = WeakKeyDictionary()
+
+
+class WorkerDeadlineExceeded(BaseException):
+ """Stop abandoned work without being swallowed by tool error handlers."""
+
+
+BUSY_MESSAGE = "MCP server busy: all tool workers are occupied. Retry later."
+
+# Tools that only read Superset's metadata database. They are admitted under a
+# separate, larger bound so they keep answering while slow warehouse queries
+# hold every warehouse slot. Anything that can reach a warehouse, a semantic
+# layer or a screenshot service stays in the warehouse bound. A listed tool
+# that does reach a warehouse must still take a warehouse slot first (see
+# WorkerCall.admit_warehouse), so a wrong entry cannot overdraw the pool.
+METADATA_ONLY_TOOLS = frozenset(
+ {
+ "find_users",
+ "get_annotation_layer_info",
+ "get_chart_info",
+ "get_chart_type_schema",
+ "get_dashboard_info",
+ "get_database_info",
+ "get_dataset_info",
+ "get_instance_info",
+ "get_layer_annotation_info",
+ "get_query_info",
+ "get_report_info",
+ "get_rls_filter_info",
+ "get_role_info",
+ "get_saved_query_info",
+ "get_schema",
+ "get_tag_info",
+ "get_task_info",
+ "get_theme_info",
+ "get_user_info",
+ "health_check",
+ "list_annotation_layers",
+ "list_charts",
+ "list_dashboards",
+ "list_databases",
+ "list_datasets",
+ "list_layer_annotations",
+ "list_queries",
+ "list_reports",
+ "list_rls_filters",
+ "list_roles",
+ "list_saved_queries",
+ "list_tags",
+ "list_tasks",
+ "list_themes",
+ "list_users",
+ }
+)
+
+
+class WorkerPool:
+ """Bound submissions, including abandoned work; never queue behind a query."""
+
+ def __init__(
+ self,
+ size: int,
+ metadata_size: int = 0,
+ transport_size: int = 1,
+ auth_size: int = 1,
+ auth_pending: int = 32,
+ ) -> None:
+ if size < 1:
+ raise ValueError("MCP_TOOL_WORKERS must be positive")
+ if metadata_size < 0:
+ raise ValueError("MCP_METADATA_TOOL_WORKERS must not be negative")
+ self.slots = threading.BoundedSemaphore(size)
+ self.metadata_slots = (
+ threading.BoundedSemaphore(metadata_size) if metadata_size else None
+ )
+ self.executor = ThreadPoolExecutor(
+ size + metadata_size, thread_name_prefix="mcp-tool"
+ )
+ # Cancellation must not wait behind the warehouse work it is cancelling.
+ # Only calls holding a warehouse slot register cancellation.
+ self.cancel_slots = threading.BoundedSemaphore(size)
+ self.cancellations = ThreadPoolExecutor(size, thread_name_prefix="mcp-cancel")
+ # Transport-side metadata I/O (tools/list filtering, audit writes,
+ # error hooks) is bounded too, so it cannot overdraw the connections
+ # budgeted for tool workers. Callers wait for it on the event loop
+ # without blocking it.
+ self.transport = ThreadPoolExecutor(
+ transport_size, thread_name_prefix="mcp-metadata"
+ )
+ # API-key lookups run before a caller is authenticated, so they get
+ # their own thread: a flood of bad keys or a slow lookup must not
+ # delay audit writes and RBAC filtering. Pending lookups are capped;
+ # past the cap a key is rejected rather than queued.
+ self.auth = ThreadPoolExecutor(auth_size, thread_name_prefix="mcp-auth")
+ self.auth_slots = threading.BoundedSemaphore(auth_pending)
+
+ def admission(self, metadata_only: bool) -> threading.BoundedSemaphore:
+ """Choose the bound a call is admitted under."""
+ if metadata_only and self.metadata_slots is not None:
+ return self.metadata_slots
+ return self.slots
+
+ def submit(
+ self,
+ fn: Callable[[], Any],
+ finished: Callable[[], None],
+ slots: threading.BoundedSemaphore | None = None,
+ ) -> Future[Any]:
+ """Admit immediately or report overload without retaining a queued call."""
+ slots = self.slots if slots is None else slots
+ if not slots.acquire(blocking=False):
+ raise ToolError(BUSY_MESSAGE)
+ try:
+ future = self.executor.submit(fn)
+ except BaseException:
+ slots.release()
+ raise
+ future.add_done_callback(lambda _: finished())
+ return future
+
+ def cancel(self, fn: Callable[[], None], finished: Callable[[], None]) -> bool:
+ """Bound cancellation I/O separately so it cannot delay the caller."""
+ if not self.cancel_slots.acquire(blocking=False):
+ return False
+ 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()
+
+ 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
+
+
+DEFAULT_TOOL_WORKERS = 16
+DEFAULT_METADATA_TOOL_WORKERS = 16
+# Threads, and so metadata connections, for transport-side metadata I/O.
+TRANSPORT_METADATA_THREADS = 1
+# Threads, and so metadata connections, for transport-side API-key lookups.
+API_KEY_CHECK_POOL_SIZE = 1
+# API-key lookups admitted at once, running or waiting; more are rejected.
+API_KEY_AUTH_PENDING = 32
+# Metadata connections held by transport-side threads rather than tool calls.
+METADATA_POOL_HEADROOM = TRANSPORT_METADATA_THREADS + API_KEY_CHECK_POOL_SIZE
+
+
+def _metadata_pool_capacity(app: Flask) -> int | None:
+ """Return how many metadata connections can be checked out at once.
+
+ ``None`` means a checkout never waits for another holder to return one
+ (e.g. ``NullPool``, per-thread pools, or unlimited overflow), including
+ for an application that never registered the metadata database: it has
+ no metadata connections to budget, and resolving ``db.engine`` for it
+ would raise rather than report a capacity.
+ """
+ from contextlib import nullcontext
+
+ from superset import db
+
+ if app.extensions.get("sqlalchemy") is not db:
+ return None
+
+ # Popping a pushed context tears down the caller's scoped session, and
+ # pools are created lazily from transport-side code that holds one.
+ in_app = has_app_context() and current_app._get_current_object() is app
+ with nullcontext() if in_app else app.app_context():
+ pool = db.engine.pool
+ if not isinstance(pool, QueuePool):
+ return None
+ # SQLAlchemy has no public accessor for the configured overflow limit.
+ max_overflow = pool._max_overflow # pylint: disable=protected-access
+ return None if max_overflow < 0 else pool.size() + max_overflow
+
+
+def admission_counts(app: Flask) -> tuple[int, int]:
+ """Split the metadata pool between warehouse and metadata-only calls.
+
+ Every holder of a metadata connection is bounded and budgeted:
+
+ - each warehouse-capable call can hold one for the whole of its warehouse
+ I/O, and its cancellation needs another (``2 * workers``);
+ - each metadata-only call holds at most one (``metadata_workers``);
+ - transport-side metadata I/O runs on ``TRANSPORT_METADATA_THREADS`` and
+ API-key lookups on ``API_KEY_CHECK_POOL_SIZE``.
+
+ ``2 * workers + metadata_workers + METADATA_POOL_HEADROOM`` never
+ exceeds the pool's capacity, so with every slot admitted no checkout
+ waits for another holder. Unless configured, warehouse calls get about
+ two thirds of the remaining connections and metadata-only calls the
+ rest, each up to 16. Explicit settings above the budget are reduced
+ with a warning. ``0`` metadata-only workers admits those tools under
+ the warehouse bound.
+
+ Returns ``(workers, metadata_workers)``.
+ """
+ configured = app.config.get("MCP_TOOL_WORKERS")
+ configured_metadata = app.config.get("MCP_METADATA_TOOL_WORKERS")
+ if configured_metadata is not None and configured_metadata < 0:
+ raise ValueError("MCP_METADATA_TOOL_WORKERS must not be negative")
+ capacity = _metadata_pool_capacity(app)
+ if capacity is None:
+ return (
+ DEFAULT_TOOL_WORKERS if configured is None else configured,
+ DEFAULT_METADATA_TOOL_WORKERS
+ if configured_metadata is None
+ else configured_metadata,
+ )
+ available = capacity - METADATA_POOL_HEADROOM
+ limit = available // 2
+ if limit < 1:
+ raise ValueError(
+ f"The metadata database pool allows {capacity} connections; "
+ "MCP tool execution needs at least "
+ f"{2 + METADATA_POOL_HEADROOM}"
+ )
+ if configured is None:
+ workers = min(DEFAULT_TOOL_WORKERS, max(1, available // 3))
+ elif configured > limit:
+ logger.warning(
+ "MCP_TOOL_WORKERS=%s needs %s metadata database connections, but the "
+ "pool allows %s; admitting %s concurrent tool calls. Raise the pool's "
+ "pool_size/max_overflow in SQLALCHEMY_ENGINE_OPTIONS to admit more.",
+ configured,
+ 2 * configured + METADATA_POOL_HEADROOM,
+ capacity,
+ limit,
+ )
+ workers = limit
+ else:
+ workers = configured
+ remaining = available - 2 * workers
+ if configured_metadata is None:
+ metadata_workers = min(DEFAULT_METADATA_TOOL_WORKERS, remaining)
+ elif configured_metadata > remaining:
+ logger.warning(
+ "MCP_METADATA_TOOL_WORKERS=%s exceeds the %s metadata database "
+ "connections left after %s warehouse-capable tool calls; admitting "
+ "%s metadata-only tool calls.",
+ configured_metadata,
+ remaining,
+ workers,
+ remaining,
+ )
+ metadata_workers = remaining
+ else:
+ metadata_workers = configured_metadata
+ return workers, metadata_workers
+
+
+def tool_worker_count(app: Flask) -> int:
+ """Admit only as many warehouse-capable calls as the pool can always serve."""
+ return admission_counts(app)[0]
+
+
+def metadata_tool_worker_count(app: Flask) -> int:
+ """Admit metadata-only calls from what warehouse calls leave of the pool."""
+ return admission_counts(app)[1]
+
+
+def _get_pool(app: Flask) -> WorkerPool:
+ """Lazily create a pool for this application, without import-time threads."""
+ with _pools_lock:
+ if app not in _pools:
+ size, metadata_size = admission_counts(app)
+ logger.info(
+ "MCP tool calls admitted concurrently: %s warehouse-capable, "
+ "%s metadata-only; %s transport metadata thread(s), "
+ "%s API-key lookup thread(s)",
+ size,
+ metadata_size,
+ TRANSPORT_METADATA_THREADS,
+ API_KEY_CHECK_POOL_SIZE,
Review Comment:
Fixed in 310379c4ea: pool initialization and clamping logs use static messages, with no configuration-, engine-, or pool-count-derived arguments; admission behavior is unchanged. CodeQL has re-run successfully with no new alerts.
##########
superset/mcp_service/worker.py:
##########
@@ -0,0 +1,910 @@
+# 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.
+"""Bounded, deadline-aware execution of MCP tools outside the transport loop.
+
+The tool's coroutine runs on a worker-owned loop because its database APIs are
+synchronous. Only transport notifications are marshalled back to the server
+loop. Flask and SQLAlchemy lifetimes belong to the worker, not the waiting
+request: a timed-out DBAPI call cannot outlive and reuse a torn-down session.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import logging
+import threading
+import time
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import contextmanager
+from contextvars import ContextVar, copy_context
+from typing import Any, Callable, Coroutine, Iterator, ParamSpec, TYPE_CHECKING, TypeVar
+from weakref import WeakKeyDictionary
+
+from fastmcp.exceptions import ToolError
+from flask import current_app, g, has_app_context, has_request_context
+from sqlalchemy import inspect as sa_inspect
+from sqlalchemy.orm.state import InstanceState
+from sqlalchemy.pool import QueuePool
+
+from superset.mcp_service.session_scope import _mcp_session_token
+
+if TYPE_CHECKING:
+ from flask import Flask
+
+ from superset.models.core import Database
+
+logger = logging.getLogger(__name__)
+_active_call: ContextVar[WorkerCall | None] = ContextVar(
+ "mcp_worker_call", default=None
+)
+_metadata_context_owned: ContextVar[bool] = ContextVar(
+ "mcp_metadata_context_owned", default=False
+)
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
+_pools_lock = threading.Lock()
+_pools: WeakKeyDictionary[Flask, WorkerPool] = WeakKeyDictionary()
+
+
+class WorkerDeadlineExceeded(BaseException):
+ """Stop abandoned work without being swallowed by tool error handlers."""
+
+
+BUSY_MESSAGE = "MCP server busy: all tool workers are occupied. Retry later."
+
+# Tools that only read Superset's metadata database. They are admitted under a
+# separate, larger bound so they keep answering while slow warehouse queries
+# hold every warehouse slot. Anything that can reach a warehouse, a semantic
+# layer or a screenshot service stays in the warehouse bound. A listed tool
+# that does reach a warehouse must still take a warehouse slot first (see
+# WorkerCall.admit_warehouse), so a wrong entry cannot overdraw the pool.
+METADATA_ONLY_TOOLS = frozenset(
+ {
+ "find_users",
+ "get_annotation_layer_info",
+ "get_chart_info",
+ "get_chart_type_schema",
+ "get_dashboard_info",
+ "get_database_info",
+ "get_dataset_info",
+ "get_instance_info",
+ "get_layer_annotation_info",
+ "get_query_info",
+ "get_report_info",
+ "get_rls_filter_info",
+ "get_role_info",
+ "get_saved_query_info",
+ "get_schema",
+ "get_tag_info",
+ "get_task_info",
+ "get_theme_info",
+ "get_user_info",
+ "health_check",
+ "list_annotation_layers",
+ "list_charts",
+ "list_dashboards",
+ "list_databases",
+ "list_datasets",
+ "list_layer_annotations",
+ "list_queries",
+ "list_reports",
+ "list_rls_filters",
+ "list_roles",
+ "list_saved_queries",
+ "list_tags",
+ "list_tasks",
+ "list_themes",
+ "list_users",
+ }
+)
+
+
+class WorkerPool:
+ """Bound submissions, including abandoned work; never queue behind a query."""
+
+ def __init__(
+ self,
+ size: int,
+ metadata_size: int = 0,
+ transport_size: int = 1,
+ auth_size: int = 1,
+ auth_pending: int = 32,
+ ) -> None:
+ if size < 1:
+ raise ValueError("MCP_TOOL_WORKERS must be positive")
+ if metadata_size < 0:
+ raise ValueError("MCP_METADATA_TOOL_WORKERS must not be negative")
+ self.slots = threading.BoundedSemaphore(size)
+ self.metadata_slots = (
+ threading.BoundedSemaphore(metadata_size) if metadata_size else None
+ )
+ self.executor = ThreadPoolExecutor(
+ size + metadata_size, thread_name_prefix="mcp-tool"
+ )
+ # Cancellation must not wait behind the warehouse work it is cancelling.
+ # Only calls holding a warehouse slot register cancellation.
+ self.cancel_slots = threading.BoundedSemaphore(size)
+ self.cancellations = ThreadPoolExecutor(size, thread_name_prefix="mcp-cancel")
+ # Transport-side metadata I/O (tools/list filtering, audit writes,
+ # error hooks) is bounded too, so it cannot overdraw the connections
+ # budgeted for tool workers. Callers wait for it on the event loop
+ # without blocking it.
+ self.transport = ThreadPoolExecutor(
+ transport_size, thread_name_prefix="mcp-metadata"
+ )
+ # API-key lookups run before a caller is authenticated, so they get
+ # their own thread: a flood of bad keys or a slow lookup must not
+ # delay audit writes and RBAC filtering. Pending lookups are capped;
+ # past the cap a key is rejected rather than queued.
+ self.auth = ThreadPoolExecutor(auth_size, thread_name_prefix="mcp-auth")
+ self.auth_slots = threading.BoundedSemaphore(auth_pending)
+
+ def admission(self, metadata_only: bool) -> threading.BoundedSemaphore:
+ """Choose the bound a call is admitted under."""
+ if metadata_only and self.metadata_slots is not None:
+ return self.metadata_slots
+ return self.slots
+
+ def submit(
+ self,
+ fn: Callable[[], Any],
+ finished: Callable[[], None],
+ slots: threading.BoundedSemaphore | None = None,
+ ) -> Future[Any]:
+ """Admit immediately or report overload without retaining a queued call."""
+ slots = self.slots if slots is None else slots
+ if not slots.acquire(blocking=False):
+ raise ToolError(BUSY_MESSAGE)
+ try:
+ future = self.executor.submit(fn)
+ except BaseException:
+ slots.release()
+ raise
+ future.add_done_callback(lambda _: finished())
+ return future
+
+ def cancel(self, fn: Callable[[], None], finished: Callable[[], None]) -> bool:
+ """Bound cancellation I/O separately so it cannot delay the caller."""
+ if not self.cancel_slots.acquire(blocking=False):
+ return False
+ 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()
+
+ 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
+
+
+DEFAULT_TOOL_WORKERS = 16
+DEFAULT_METADATA_TOOL_WORKERS = 16
+# Threads, and so metadata connections, for transport-side metadata I/O.
+TRANSPORT_METADATA_THREADS = 1
+# Threads, and so metadata connections, for transport-side API-key lookups.
+API_KEY_AUTH_THREADS = 1
Review Comment:
CodeQL is now clean on 310379c4ea: the check passed with no new alerts, including both remaining pool-log findings.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on code in PR #44581:
URL: https://github.com/apache/superset/pull/44581#discussion_r4152367639
##########
superset/mcp_service/composite_token_verifier.py:
##########
@@ -180,9 +180,13 @@ async def verify_token(self, token: str) -> AccessToken | None:
"""
if any(token.startswith(prefix) for prefix in self._api_key_prefixes):
if self._app is not None:
+ from superset.mcp_service.worker import transport_executor
+
loop = asyncio.get_running_loop()
+ # Metadata I/O on the transport side shares a bounded,
+ # pool-budgeted executor with audit writes and RBAC filtering.
result = await loop.run_in_executor(
- None, self._validate_api_key_sync, token
+ transport_executor(self._app), self._validate_api_key_sync, token
Review Comment:
Valid. Fixed in 8690ac8600: API-key lookups run on their own thread with a budgeted metadata connection, with at most 32 admitted at once (more are rejected, not queued). Audit writes and RBAC filtering keep the transport thread to themselves. Covered by `test_stuck_lookups_do_not_delay_transport_metadata_io` and `test_transport_validation_runs_on_its_own_thread`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on PR #44581:
URL: https://github.com/apache/superset/pull/44581#issuecomment-5812507557
AI Code Review is in progress (usually takes 3 to 15 minutes unless it's a very large PR).
<details>
<summary>
Bito Usage Guide
</summary>
**Commands**
Type the following command in the pull request comment and save the comment.
- `/review` - Manually triggers an incremental AI Review.
- `/review full` - Manually triggers a full AI Review.
- `/pause` - Pauses automatic reviews on this pull request.
- `/resume` - Resumes automatic reviews.
- `/resolve` - Marks all Bito-posted review comments as resolved.
- `/abort` - Cancels all in-progress reviews.
Refer to the <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/available-commands" target="_blank" rel="noopener noreferrer">documentation</a> for additional commands.
**Configuration**
This repository uses `Superset` You can customize the agent settings <a href="https://alpha.bito.ai/home/ai-agents/code-review-agent" target="_blank" rel="noopener noreferrer">here</a> or contact your Bito workspace admin at evan@preset.io.
**Documentation & Help**
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/getting-started/install-run-using-bito-cloud/create-or-customize-an-agent-instance" target="_blank" rel="noopener noreferrer">Customize agent settings</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/implementing-custom-code-review-rules" target="_blank" rel="noopener noreferrer">Review rules</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent" target="_blank" rel="noopener noreferrer">General documentation</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/faqs" target="_blank" rel="noopener noreferrer">FAQ</a>
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on code in PR #44581:
URL: https://github.com/apache/superset/pull/44581#discussion_r4160570897
##########
superset/mcp_service/worker.py:
##########
@@ -0,0 +1,910 @@
+# 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.
+"""Bounded, deadline-aware execution of MCP tools outside the transport loop.
+
+The tool's coroutine runs on a worker-owned loop because its database APIs are
+synchronous. Only transport notifications are marshalled back to the server
+loop. Flask and SQLAlchemy lifetimes belong to the worker, not the waiting
+request: a timed-out DBAPI call cannot outlive and reuse a torn-down session.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import logging
+import threading
+import time
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import contextmanager
+from contextvars import ContextVar, copy_context
+from typing import Any, Callable, Coroutine, Iterator, ParamSpec, TYPE_CHECKING, TypeVar
+from weakref import WeakKeyDictionary
+
+from fastmcp.exceptions import ToolError
+from flask import current_app, g, has_app_context, has_request_context
+from sqlalchemy import inspect as sa_inspect
+from sqlalchemy.orm.state import InstanceState
+from sqlalchemy.pool import QueuePool
+
+from superset.mcp_service.session_scope import _mcp_session_token
+
+if TYPE_CHECKING:
+ from flask import Flask
+
+ from superset.models.core import Database
+
+logger = logging.getLogger(__name__)
+_active_call: ContextVar[WorkerCall | None] = ContextVar(
+ "mcp_worker_call", default=None
+)
+_metadata_context_owned: ContextVar[bool] = ContextVar(
+ "mcp_metadata_context_owned", default=False
+)
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
+_pools_lock = threading.Lock()
+_pools: WeakKeyDictionary[Flask, WorkerPool] = WeakKeyDictionary()
+
+
+class WorkerDeadlineExceeded(BaseException):
+ """Stop abandoned work without being swallowed by tool error handlers."""
+
+
+BUSY_MESSAGE = "MCP server busy: all tool workers are occupied. Retry later."
+
+# Tools that only read Superset's metadata database. They are admitted under a
+# separate, larger bound so they keep answering while slow warehouse queries
+# hold every warehouse slot. Anything that can reach a warehouse, a semantic
+# layer or a screenshot service stays in the warehouse bound. A listed tool
+# that does reach a warehouse must still take a warehouse slot first (see
+# WorkerCall.admit_warehouse), so a wrong entry cannot overdraw the pool.
+METADATA_ONLY_TOOLS = frozenset(
+ {
+ "find_users",
+ "get_annotation_layer_info",
+ "get_chart_info",
+ "get_chart_type_schema",
+ "get_dashboard_info",
+ "get_database_info",
+ "get_dataset_info",
+ "get_instance_info",
+ "get_layer_annotation_info",
+ "get_query_info",
+ "get_report_info",
+ "get_rls_filter_info",
+ "get_role_info",
+ "get_saved_query_info",
+ "get_schema",
+ "get_tag_info",
+ "get_task_info",
+ "get_theme_info",
+ "get_user_info",
+ "health_check",
+ "list_annotation_layers",
+ "list_charts",
+ "list_dashboards",
+ "list_databases",
+ "list_datasets",
+ "list_layer_annotations",
+ "list_queries",
+ "list_reports",
+ "list_rls_filters",
+ "list_roles",
+ "list_saved_queries",
+ "list_tags",
+ "list_tasks",
+ "list_themes",
+ "list_users",
+ }
+)
+
+
+class WorkerPool:
+ """Bound submissions, including abandoned work; never queue behind a query."""
+
+ def __init__(
+ self,
+ size: int,
+ metadata_size: int = 0,
+ transport_size: int = 1,
+ auth_size: int = 1,
+ auth_pending: int = 32,
+ ) -> None:
+ if size < 1:
+ raise ValueError("MCP_TOOL_WORKERS must be positive")
+ if metadata_size < 0:
+ raise ValueError("MCP_METADATA_TOOL_WORKERS must not be negative")
+ self.slots = threading.BoundedSemaphore(size)
+ self.metadata_slots = (
+ threading.BoundedSemaphore(metadata_size) if metadata_size else None
+ )
+ self.executor = ThreadPoolExecutor(
+ size + metadata_size, thread_name_prefix="mcp-tool"
+ )
+ # Cancellation must not wait behind the warehouse work it is cancelling.
+ # Only calls holding a warehouse slot register cancellation.
+ self.cancel_slots = threading.BoundedSemaphore(size)
+ self.cancellations = ThreadPoolExecutor(size, thread_name_prefix="mcp-cancel")
+ # Transport-side metadata I/O (tools/list filtering, audit writes,
+ # error hooks) is bounded too, so it cannot overdraw the connections
+ # budgeted for tool workers. Callers wait for it on the event loop
+ # without blocking it.
+ self.transport = ThreadPoolExecutor(
+ transport_size, thread_name_prefix="mcp-metadata"
+ )
+ # API-key lookups run before a caller is authenticated, so they get
+ # their own thread: a flood of bad keys or a slow lookup must not
+ # delay audit writes and RBAC filtering. Pending lookups are capped;
+ # past the cap a key is rejected rather than queued.
+ self.auth = ThreadPoolExecutor(auth_size, thread_name_prefix="mcp-auth")
+ self.auth_slots = threading.BoundedSemaphore(auth_pending)
+
+ def admission(self, metadata_only: bool) -> threading.BoundedSemaphore:
+ """Choose the bound a call is admitted under."""
+ if metadata_only and self.metadata_slots is not None:
+ return self.metadata_slots
+ return self.slots
+
+ def submit(
+ self,
+ fn: Callable[[], Any],
+ finished: Callable[[], None],
+ slots: threading.BoundedSemaphore | None = None,
+ ) -> Future[Any]:
+ """Admit immediately or report overload without retaining a queued call."""
+ slots = self.slots if slots is None else slots
+ if not slots.acquire(blocking=False):
+ raise ToolError(BUSY_MESSAGE)
+ try:
+ future = self.executor.submit(fn)
+ except BaseException:
+ slots.release()
+ raise
+ future.add_done_callback(lambda _: finished())
+ return future
+
+ def cancel(self, fn: Callable[[], None], finished: Callable[[], None]) -> bool:
+ """Bound cancellation I/O separately so it cannot delay the caller."""
+ if not self.cancel_slots.acquire(blocking=False):
+ return False
+ 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()
+
+ 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
+
+
+DEFAULT_TOOL_WORKERS = 16
+DEFAULT_METADATA_TOOL_WORKERS = 16
+# Threads, and so metadata connections, for transport-side metadata I/O.
+TRANSPORT_METADATA_THREADS = 1
+# Threads, and so metadata connections, for transport-side API-key lookups.
+API_KEY_AUTH_THREADS = 1
Review Comment:
Fixed in 3a56a5ada0: renamed the pool-count constants and changed successful API-key validation logging to the user id, preserving pool sizing, scopes, and token claims. CodeQL will re-run on this commit; no alerts were dismissed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on PR #44581:
URL: https://github.com/apache/superset/pull/44581#issuecomment-5923708413
> `admission_counts` looks right, nice unifying the transport-thread bound in there too. Lifting the block.
>
> Put together a PR against your branch with a few smaller things from a deeper pass, rebased onto your latest since most of it had landed independently: [aminghadersohi#1](https://github.com/aminghadersohi/superset/pull/1). What's left:
>
> * `on_list_tools` and both search-transform `_get_visible_tools` overrides call `run_in_metadata_thread` with no try/except, so a setup failure there (pool exhausted, same scenario as above) propagates instead of failing open like the filter itself does.
> * `QueryCancellation.capture()` doesn't reset `cancel_id` before a re-capture, so a failed second capture re-dispatches cancellation against the first, already-finished statement.
> * `WorkerPool.cancel()` only releases its permit on `RuntimeError`; anything else leaks it.
> * `run_in_worker`'s `pool.submit()` is outside the `try/finally` resetting `_mcp_user_id_var`, so a busy-pool rejection leaves a stale user id for audit logging.
> * `run_in_worker`/`TransportContext.forward` both relabel any `TimeoutError` as our deadline, including one actually raised by the tool/driver.
>
> Take whatever's useful, happy to let the rest drop.
i mergdd it in
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on PR #44581:
URL: https://github.com/apache/superset/pull/44581#issuecomment-5925793717
Pushed 8690ac8600: API-key lookups get their own bounded thread, so they no longer share the transport metadata thread with audit writes and RBAC filtering.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on code in PR #44581:
URL: https://github.com/apache/superset/pull/44581#discussion_r4114433479
##########
tests/unit_tests/mcp_service/test_worker_metadata_pool.py:
##########
@@ -0,0 +1,530 @@
+# 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.
+"""Tool workers and the transport loop share the metadata connection pool."""
+
+import asyncio
+import threading
+import time
+from collections.abc import Iterator
+from pathlib import Path
+from typing import Any
+from unittest.mock import patch, PropertyMock
+
+import pytest
+from fastmcp import Client, Context, FastMCP
+from sqlalchemy import create_engine, text
+from sqlalchemy.engine import Engine
+from sqlalchemy.orm import Session
+from sqlalchemy.pool import NullPool
+
+from superset.extensions import db
+from superset.mcp_service.auth import mcp_auth_hook
+from superset.mcp_service.server import build_middleware_list
+from superset.mcp_service.session_scope import install_mcp_session_scoping
+
+WORKERS = 2
+# The smallest pool that admits WORKERS tool calls: one connection per tool
+# worker, one per cancellation worker, and one for transport-side lookups.
+POOL_SIZE = 2 * WORKERS + 1
+
+
+@pytest.fixture
+def metadata_engine(app: Any, tmp_path: Path) -> Iterator[Engine]:
+ """Bind the metadata session to a bounded pool, as in production."""
+ engine = create_engine(
+ f"sqlite:///{tmp_path / 'metadata.db'}",
+ pool_size=POOL_SIZE,
+ max_overflow=0,
+ pool_timeout=4,
+ connect_args={"check_same_thread": False},
+ )
+ previous = {
+ key: app.config.get(key) for key in ("MCP_TOOL_WORKERS", "SQLLAB_TIMEOUT")
+ }
+ app.config["MCP_TOOL_WORKERS"] = WORKERS
+ app.config["SQLLAB_TIMEOUT"] = 5
+ install_mcp_session_scoping()
+ db.session.remove()
+ try:
+ with (
+ patch.object(
+ type(db),
+ "engines",
+ new_callable=PropertyMock,
+ return_value={None: engine},
+ ),
+ patch.object(
+ type(db), "engine", new_callable=PropertyMock, return_value=engine
+ ),
+ ):
+ yield engine
+ finally:
+ db.session.remove()
+ engine.dispose()
+ app.config.update(previous)
+
+
+@pytest.mark.asyncio
+async def test_saturated_workers_holding_metadata_pool_do_not_freeze_loop(
+ metadata_engine: Engine,
+) -> None:
+ """Every tool worker waits on the loop while every connection is checked out.
+
+ If the transport loop itself waits for a metadata connection (tools/list
+ RBAC filtering, audit logging), neither side can progress until a pool or
+ call timeout expires, and the whole server stops responding meanwhile.
+ """
+ holding = threading.Barrier(WORKERS + 1)
+ notify = threading.Event()
+
+ async def hold_and_notify(ctx: Context) -> str:
+ """Keep a metadata transaction open while reporting progress."""
+ db.session.execute(text("SELECT 1"))
+ holding.wait(timeout=5)
+ notify.wait(timeout=5)
+ await ctx.info("progress")
+ return "done"
+
+ def user_lookup() -> None:
+ """Resolve the caller for tools/list with a metadata query."""
+ db.session.execute(text("SELECT 1"))
+
+ mcp = FastMCP("metadata pool regression", middleware=build_middleware_list())
+ with (
+ patch("superset.mcp_service.auth._setup_user_context", return_value=None),
+ patch(
+ "superset.mcp_service.middleware.get_user_from_request",
+ side_effect=user_lookup,
+ ),
+ ):
+ mcp.tool(mcp_auth_hook(hold_and_notify), name="hold_and_notify")
+ async with Client(mcp) as client:
+ calls = [
+ asyncio.create_task(client.call_tool("hold_and_notify", {}))
+ for _ in range(WORKERS)
+ ]
+ await asyncio.to_thread(holding.wait, 5)
+ # Other holders, such as cancellations, own the rest of the pool.
+ held = [metadata_engine.connect() for _ in range(POOL_SIZE - WORKERS)]
+ gaps = []
+
+ async def heartbeat() -> None:
+ """Measure how long the transport loop goes unscheduled."""
+ last = time.monotonic()
+ while True:
+ await asyncio.sleep(0.01)
+ now = time.monotonic()
+ gaps.append(now - last)
+ last = now
+
+ beat = asyncio.create_task(heartbeat())
+ started = time.monotonic()
+ # tools/list needs metadata connections for its audit record and
+ # RBAC filtering. Release the workers once it is waiting for one,
+ # so their progress reports reach the loop while it waits.
+ listing = asyncio.create_task(client.list_tools())
+ threading.Timer(0.2, notify.set).start()
+ try:
+ await asyncio.sleep(0.05)
+ pinged = time.monotonic()
+ await client.ping()
+ ping_latency = time.monotonic() - pinged
+ results = await asyncio.gather(*calls)
+ elapsed = time.monotonic() - started
+ finally:
+ for connection in held:
+ connection.close()
+ beat.cancel()
+ tools = await listing
+
+ assert [result.content[0].text for result in results] == ["done"] * WORKERS
+ assert [tool.name for tool in tools] == ["hold_and_notify"]
+ assert ping_latency < 0.5
+ assert elapsed < 1.5
+ assert max(gaps) < 0.5
+
+
+@pytest.mark.parametrize(
+ ("pool_options", "configured", "expected"),
+ [
+ # SQLAlchemy's default QueuePool lends 5 + 10 connections.
+ ({}, None, 7),
+ ({}, 16, 7),
+ ({}, 4, 4),
+ ({"pool_size": 20, "max_overflow": 20}, None, 16),
+ ({"pool_size": 20, "max_overflow": 20}, 19, 19),
+ ({"max_overflow": -1}, 32, 32),
+ ({"poolclass": NullPool}, None, 16),
+ ],
+)
+def test_tool_workers_leave_metadata_connections_for_cancellation(
+ app: Any,
+ tmp_path: Path,
+ pool_options: dict[str, Any],
+ configured: int | None,
+ expected: int,
+) -> None:
+ """Each call and its cancellation get a connection, with one to spare."""
+ from superset.mcp_service.worker import tool_worker_count
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'metadata.db'}", **pool_options)
+ previous = app.config.get("MCP_TOOL_WORKERS")
+ app.config["MCP_TOOL_WORKERS"] = configured
+ try:
+ with patch.object(
+ type(db), "engine", new_callable=PropertyMock, return_value=engine
+ ):
+ assert tool_worker_count(app) == expected
+ finally:
+ app.config["MCP_TOOL_WORKERS"] = previous
+ engine.dispose()
+
+
+def test_metadata_pool_too_small_for_cancellation_is_rejected(
+ app: Any, tmp_path: Path
+) -> None:
+ """Refuse a pool where a call and its cancellation could take every slot."""
+ from superset.mcp_service.worker import tool_worker_count
+
+ engine = create_engine(
+ f"sqlite:///{tmp_path / 'metadata.db'}", pool_size=2, max_overflow=0
+ )
+ with (
+ patch.object(
+ type(db), "engine", new_callable=PropertyMock, return_value=engine
+ ),
+ pytest.raises(ValueError, match="needs at least 3"),
+ ):
+ tool_worker_count(app)
+ engine.dispose()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("strategy", ["bm25", "regex"])
+@pytest.mark.parametrize("request_backed", [False, True])
+async def test_threaded_search_returns_metadata_connection(
+ app: Any, metadata_engine: Engine, strategy: str, request_backed: bool
+) -> None:
+ """An inherited Flask context must not bypass the search thread's cleanup."""
+ from unittest.mock import AsyncMock, Mock
+
+ from fastmcp.server.transforms.search.base import BaseSearchTransform
+ from flask import g
+
+ from superset.mcp_service.server import _create_search_transform
+ from superset.mcp_service.session_scope import mcp_session_scopefunc
+
+ sessions: dict[Any, Session] = {}
+
+ def permission(tool: Any) -> bool:
+ """Perform the metadata lookup used by database-backed RBAC."""
+ sessions[mcp_session_scopefunc()] = db.session()
+ db.session.execute(text("SELECT 1"))
+ return True
+
+ transform = _create_search_transform(
+ strategy=strategy, kwargs={}, make_normalizing_call_tool=lambda _: None
+ )
+ tool = Mock()
+ try:
+ with app.test_request_context("/mcp/") if request_backed else app.app_context():
+ g.user = Mock()
+ with (
+ patch.object(
+ BaseSearchTransform,
+ "_get_visible_tools",
+ new_callable=AsyncMock,
+ return_value=[tool],
+ ),
+ patch(
+ "superset.mcp_service.auth.is_tool_visible_to_current_user",
+ side_effect=permission,
+ ),
+ ):
+ assert await transform._get_visible_tools(Mock()) == [tool]
+ assert metadata_engine.pool.checkedout() == 0
+ assert metadata_engine.pool.checkedout() == 0
+ assert all(key not in db.session.registry.registry for key in sessions)
+ finally:
+ # Also clean up when run against the leaking implementation.
+ for key, session in sessions.items():
+ session.close()
+ db.session.registry.registry.pop(key, None)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("audit_path", ["error", "response_size", "error_hook"])
+async def test_error_audit_waiting_for_metadata_does_not_block_loop(
+ app: Any, metadata_engine: Engine, audit_path: str
+) -> None:
+ """A timeout/busy error must not synchronously wait for an audit connection."""
+ from unittest.mock import AsyncMock, Mock
+
+ from fastmcp.exceptions import ToolError
+
+ from superset.mcp_service.middleware import (
+ GlobalErrorHandlerMiddleware,
+ ResponseSizeGuardMiddleware,
+ ToolResultCompatibilityMiddleware,
+ )
+ from superset.models.core import Log
+ from superset.utils.log import DBEventLogger
+
+ Log.__table__.create(metadata_engine)
+ held = [metadata_engine.connect() for _ in range(POOL_SIZE)]
+ heartbeat = asyncio.Event()
+
+ async def beat() -> None:
+ """Prove the transport can run while the audit checkout is blocked."""
+ await asyncio.sleep(0.01)
+ heartbeat.set()
+
+ task = asyncio.create_task(beat())
+ release = threading.Timer(0.3, held[-1].close)
+ release.start()
+ try:
+ context = Mock(method="tools/call")
+ context.message.name = "unknown_tool"
+ context.message.arguments = {}
+ event_logger = DBEventLogger()
+
+ def error_hook(error: Exception, info: dict[str, Any]) -> None:
+ """An operator hook may itself perform metadata I/O."""
+ event_logger.log(
+ user_id=None,
+ action="test_error_hook",
+ dashboard_id=None,
+ duration_ms=None,
+ slice_id=None,
+ referrer=None,
+ curated_payload={"tool": info["tool_name"]},
+ )
+
+ with (
+ app.test_request_context("/mcp/"),
+ patch("superset.mcp_service.middleware.event_logger", event_logger),
+ patch(
+ "superset.mcp_service.flask_singleton.get_flask_app", return_value=app
+ ),
+ patch.dict(app.config, {"MCP_ERROR_HOOK": error_hook}),
+ ):
+ if audit_path == "error":
+ with pytest.raises(ToolError, match="server busy"):
+ await GlobalErrorHandlerMiddleware()._handle_error(
+ ToolError("MCP server busy"), context, "execute_sql", 0
+ )
+ elif audit_path == "response_size":
+ with pytest.raises(ToolError):
+ await ResponseSizeGuardMiddleware(max_bytes=10).on_call_tool(
+ context, AsyncMock(return_value={"data": "x" * 100})
+ )
+ else:
+ result = await ToolResultCompatibilityMiddleware().on_call_tool(
+ context, AsyncMock(side_effect=RuntimeError("failed"))
+ )
+ assert result.is_error
+ assert heartbeat.is_set()
+ assert metadata_engine.pool.checkedout() == POOL_SIZE - 1
+ finally:
+ release.join()
+ for connection in held:
+ connection.close()
+ await task
+ assert metadata_engine.pool.checkedout() == 0
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("fail", [False, True])
+async def test_metadata_thread_owns_user_context_and_session(
+ app: Any, metadata_engine: Engine, fail: bool
+) -> None:
+ """Do not share even an expired request user or an inherited session token."""
+ from contextlib import nullcontext
+
+ from flask import g, request
+ from flask_appbuilder.security.sqla.models import User
+ from sqlalchemy import event
+
+ from superset.mcp_service.session_scope import _mcp_session_token
+ from superset.mcp_service.worker import (
+ get_context_user_id,
+ run_in_metadata_thread,
+ )
+
+ User.__table__.create(metadata_engine)
+ with metadata_engine.begin() as connection:
+ connection.execute(
+ User.__table__.insert(),
+ {
+ "id": 42,
+ "username": "metadata-user",
+ "email": "metadata@example.com",
+ "first_name": "Meta",
+ "last_name": "Data",
+ },
+ )
+ token = _mcp_session_token.set(object())
+ threads: list[int] = []
+ sessions: list[Session] = []
+ main_thread = threading.get_ident()
+
+ def record_thread(*args: Any) -> None:
+ """Track all SQL, including implicit ORM refreshes."""
+ threads.append(threading.get_ident())
+
+ def lookup() -> None:
+ """Read a reloaded user while owning fresh request globals and session."""
+ assert g.user is not user
+ assert g.user.username == "metadata-user"
+ assert request.path == "/mcp/"
+ assert g.routing_marker == "caller"
+ g.routing_marker = "metadata"
+ session = db.session()
+ sessions.append(session)
+ assert session is not parent_session
+ session.execute(text("SELECT 1"))
+ if fail:
+ raise ValueError("metadata failed")
+
+ try:
+ with app.test_request_context("/mcp/"):
+ parent_session = db.session()
+ user = parent_session.get(User, 42)
+ parent_session.commit() # expire attributes and return the connection
+ g.user = user
+ g.routing_marker = "caller"
+ event.listen(metadata_engine, "before_cursor_execute", record_thread)
+ try:
+ assert get_context_user_id() == 42
+ with (
+ pytest.raises(ValueError, match="metadata failed")
+ if fail
+ else nullcontext()
+ ):
+ await run_in_metadata_thread(lookup)
+ assert g.user is user
+ assert g.routing_marker == "caller"
+ assert db.session() is parent_session
+ assert threads
+ assert main_thread not in threads
+ assert metadata_engine.pool.checkedout() == 0
+ assert all(
+ session not in db.session.registry.registry.values()
+ for session in sessions
+ )
+ finally:
+ event.remove(metadata_engine, "before_cursor_execute", record_thread)
+ finally:
+ db.session.remove()
+ _mcp_session_token.reset(token)
+
+
+@pytest.mark.asyncio
+async def test_cancelled_metadata_awaiter_leaves_cleanup_to_thread(
+ metadata_engine: Engine,
+) -> None:
+ """A disconnect must neither tear down live I/O nor leak its connection."""
+ from sqlalchemy import event
+
+ from superset.mcp_service.worker import run_in_metadata_thread
+
+ entered = threading.Event()
+ release = threading.Event()
+ returned = threading.Event()
+
+ def checkin(*args: Any) -> None:
+ """Signal teardown after the thread's metadata I/O has completed."""
+ returned.set()
+
+ def lookup() -> None:
+ """Keep using the same live session after its awaiter disconnects."""
+ db.session.execute(text("SELECT 1"))
+ entered.set()
+ release.wait(5)
+ db.session.execute(text("SELECT 2"))
+
+ event.listen(metadata_engine, "checkin", checkin)
+ task = asyncio.create_task(run_in_metadata_thread(lookup))
+ try:
+ assert await asyncio.to_thread(entered.wait, 2)
+ task.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await task
+ assert metadata_engine.pool.checkedout() == 1
+ assert not returned.is_set()
+ finally:
+ release.set()
+ assert await asyncio.to_thread(returned.wait, 2)
+ event.remove(metadata_engine, "checkin", checkin)
+ assert metadata_engine.pool.checkedout() == 0
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("configured", [False, True])
+async def test_error_hook_does_not_require_metadata_user_reload(
+ app: Any, metadata_engine: Engine, configured: bool
+) -> None:
+ """Error reporting survives metadata outages without sharing caller state."""
+ from unittest.mock import Mock
+
+ from flask import g
+ from flask_appbuilder.security.sqla.models import User
+ from sqlalchemy.orm import make_transient_to_detached
+
+ from superset.mcp_service.middleware import _invoke_error_hook_off_loop
+
+ main_thread = threading.get_ident()
+ error = RuntimeError("tool failed")
+ info = {"tool_name": "execute_sql", "user_id": 42}
+ sessions: list[Session] = []
+
+ with app.test_request_context("/mcp/"):
+ user = User(id=42)
+ make_transient_to_detached(user)
+ g.user = user
+ parent_session = db.session()
+
+ def capture(exc: Exception, context: dict[str, Any]) -> None:
+ """Capture errors without loading a user from the unavailable DB."""
+ assert threading.get_ident() != main_thread
+ assert exc is error
+ assert context == info
+ assert getattr(g, "user", None) is None
+ assert db.session() is not parent_session
+ sessions.append(db.session())
Review Comment:
Fixed in e09b3df7745c044ca757f52c5542839f1b237ff7. Confirmed that _invoke_error_hook catches AssertionError (and _invoke_error_hook_off_loop also catches exceptions). The hook now records observations; all assertions execute after awaiting it, including a non-vacuous observation count and session cleanup check.
Reviewed the other tests in this file: the assertions in the metadata lookup callback propagate through run_in_metadata_thread, so no other swallowed assertions required changes. The two additional review suggestions concern annotations and request context, not additional swallowed assertions.
Validation: the complete file passes (21 tests), and pre-commit run --files tests/unit_tests/mcp_service/test_worker_metadata_pool.py passes. Break-it proof: temporarily exposing the caller ORM user inside capture produced a failure at assert hook_user is None (1 failed, 1 passed); reverted the mutation and reran the complete file successfully.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on code in PR #44581:
URL: https://github.com/apache/superset/pull/44581#discussion_r4152043716
##########
superset/mcp_service/composite_token_verifier.py:
##########
@@ -180,9 +180,13 @@ async def verify_token(self, token: str) -> AccessToken | None:
"""
if any(token.startswith(prefix) for prefix in self._api_key_prefixes):
if self._app is not None:
+ from superset.mcp_service.worker import transport_executor
+
loop = asyncio.get_running_loop()
+ # Metadata I/O on the transport side shares a bounded,
+ # pool-budgeted executor with audit writes and RBAC filtering.
result = await loop.run_in_executor(
- None, self._validate_api_key_sync, token
+ transport_executor(self._app), self._validate_api_key_sync, token
Review Comment:
<div>
<div id="suggestion">
<div id="issue"><b>Single-thread executor contention</b></div>
<div id="fix">
This routes FAB validation onto `transport_executor(self._app)`, whose pool has a single thread (`TRANSPORT_METADATA_THREADS = 1`, worker.py:207). Every API-key auth now shares that one thread with audit writes and RBAC filtering (`run_in_metadata_thread` callers in middleware.py); a slow or timing-out FAB DB call serializes all of them, and unauthenticated `sst_` requests can hold the thread. Consider a larger or dedicated bound for auth validation.
</div>
</div>
<small><i>Code Review Run #b98e53</i></small>
</div>
---
Should Bito avoid suggestions like this for future reviews? (<a href=https://alpha.bito.ai/home/ai-agents/review-rules>Manage Rules</a>)
- [ ] Yes, avoid them
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "codeant-ai-for-open-source[bot] (via GitHub)" <gi...@apache.org>.
codeant-ai-for-open-source[bot] commented on code in PR #44581:
URL: https://github.com/apache/superset/pull/44581#discussion_r4085843716
##########
superset/mcp_service/auth.py:
##########
@@ -1259,7 +1277,27 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
_cleanup_session_on_error()
raise
- wrapper = async_wrapper
+ @functools.wraps(tool_func)
+ async def worker_wrapper(*args: Any, **kwargs: Any) -> Any:
+ from superset.mcp_service.worker import run_in_worker
+
+ bound = _tool_sig.bind_partial(*args, **kwargs)
+ request = bound.arguments.get("request")
+ seconds = getattr(request, "timeout", None)
+ if seconds is None:
+ if has_app_context():
+ seconds = current_app.config.get("SQLLAB_TIMEOUT", 30)
+ else:
+ from superset.mcp_service.flask_singleton import get_flask_app
+
+ seconds = get_flask_app().config.get("SQLLAB_TIMEOUT", 30)
+ # Bind ctx once on the transport loop, including positional callers.
+ bound.arguments.update(_inject_ctx(dict(bound.arguments)))
+ return await run_in_worker(
+ async_wrapper, (), dict(bound.arguments), seconds
+ )
+
+ wrapper = worker_wrapper
Review Comment:
**Suggestion:** Only async protected functions use `run_in_worker`; protected synchronous tools still execute directly on the transport thread and can freeze MCP while warehouse calls block.
**Assessment:** ๐ด `Critical` ยท ๐ `Occurrence: Sometimes` ยท ๐ท๏ธ `Incomplete implementation`
[](https://docs.codeant.ai/cli/resolve-pr-comments-skill) [](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=e36f7a05504e43d0b026dfdd1eb341d3&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset) [](https://app.codeant.ai/fix-in-ide?tool=vscode-claude&prompt_id=e36f7a05504e43d0b026dfdd1eb341d3&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset)
<details>
<summary><b>Prompt for AI Agent ๐ค </b></summary>
```mdx
This is a comment left during a code review.
**Path:** superset/mcp_service/auth.py
**Line:** 1300:1300
**Comment:**
*Incomplete Implementation: Only async protected functions use `run_in_worker`; protected synchronous tools still execute directly on the transport thread and can freeze MCP while warehouse calls block.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
```
</details>
<a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F44581&comment_hash=65ddad2d3b1ef5babaed6ed147ba3cdbf91a484595541c8c6949d3ff9771eedd&reaction=like'>๐</a> | <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F44581&comment_hash=65ddad2d3b1ef5babaed6ed147ba3cdbf91a484595541c8c6949d3ff9771eedd&reaction=dislike'>๐</a>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "codeant-ai-for-open-source[bot] (via GitHub)" <gi...@apache.org>.
codeant-ai-for-open-source[bot] commented on code in PR #44581:
URL: https://github.com/apache/superset/pull/44581#discussion_r4085849128
##########
superset/models/core.py:
##########
@@ -602,35 +602,38 @@ def get_sqla_engine( # pylint: disable=too-many-arguments
sqlalchemy_uri=sqlalchemy_uri,
cacheable=not prequeries,
)
- if prequeries:
- # SQLAlchemy connect event: runs prequeries on every new
- # DBAPI connection (e.g. SET search_path for PostgreSQL).
- def run_prequeries(
- dbapi_connection: Any,
- connection_record: Any, # pylint: disable=unused-argument
- ) -> None:
- cursor = dbapi_connection.cursor()
+ from superset.sql.execution.cancellation import cancellable_engine
+
+ with cancellable_engine(self, engine, catalog, schema):
+ if prequeries:
+ # SQLAlchemy connect event: runs prequeries on every new
+ # DBAPI connection (e.g. SET search_path for PostgreSQL).
+ def run_prequeries(
+ dbapi_connection: Any,
+ connection_record: Any, # pylint: disable=unused-argument
+ ) -> None:
+ cursor = dbapi_connection.cursor()
+ try:
+ for prequery in prequeries:
+ cursor.execute(prequery)
Review Comment:
**Suggestion:** Prequeries execute through a raw DBAPI cursor before cancellation registration, so a blocking prequery can exceed the MCP deadline without cancellation.
**Assessment:** ๐ `Major` ยท ๐ `Occurrence: Sometimes` ยท ๐ท๏ธ `Incomplete implementation`
[](https://docs.codeant.ai/cli/resolve-pr-comments-skill) [](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=691f8c603af04edf8fc2d4e6ce0412e0&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset) [](https://app.codeant.ai/fix-in-ide?tool=vscode-claude&prompt_id=691f8c603af04edf8fc2d4e6ce0412e0&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset)
<details>
<summary><b>Prompt for AI Agent ๐ค </b></summary>
```mdx
This is a comment left during a code review.
**Path:** superset/models/core.py
**Line:** 617:618
**Comment:**
*Incomplete Implementation: Prequeries execute through a raw DBAPI cursor before cancellation registration, so a blocking prequery can exceed the MCP deadline without cancellation.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
```
</details>
<a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F44581&comment_hash=830896e61f12bca216429e7a9296c3c6df9c9204b8c0f2635b601587c1b81e3d&reaction=like'>๐</a> | <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F44581&comment_hash=830896e61f12bca216429e7a9296c3c6df9c9204b8c0f2635b601587c1b81e3d&reaction=dislike'>๐</a>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on PR #44581:
URL: https://github.com/apache/superset/pull/44581#issuecomment-5850021896
AI Code Review is in progress (usually takes 3 to 15 minutes unless it's a very large PR).
<details>
<summary>
Bito Usage Guide
</summary>
**Commands**
Type the following command in the pull request comment and save the comment.
- `/review` - Manually triggers an incremental AI Review.
- `/review full` - Manually triggers a full AI Review.
- `/pause` - Pauses automatic reviews on this pull request.
- `/resume` - Resumes automatic reviews.
- `/resolve` - Marks all Bito-posted review comments as resolved.
- `/abort` - Cancels all in-progress reviews.
Refer to the <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/available-commands" target="_blank" rel="noopener noreferrer">documentation</a> for additional commands.
**Configuration**
This repository uses `Superset` You can customize the agent settings <a href="https://alpha.bito.ai/home/ai-agents/code-review-agent" target="_blank" rel="noopener noreferrer">here</a> or contact your Bito workspace admin at evan@preset.io.
**Documentation & Help**
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/getting-started/install-run-using-bito-cloud/create-or-customize-an-agent-instance" target="_blank" rel="noopener noreferrer">Customize agent settings</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/implementing-custom-code-review-rules" target="_blank" rel="noopener noreferrer">Review rules</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent" target="_blank" rel="noopener noreferrer">General documentation</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/faqs" target="_blank" rel="noopener noreferrer">FAQ</a>
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on code in PR #44581:
URL: https://github.com/apache/superset/pull/44581#discussion_r4114367478
##########
tests/unit_tests/mcp_service/test_worker_metadata_pool.py:
##########
@@ -0,0 +1,530 @@
+# 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.
+"""Tool workers and the transport loop share the metadata connection pool."""
+
+import asyncio
+import threading
+import time
+from collections.abc import Iterator
+from pathlib import Path
+from typing import Any
+from unittest.mock import patch, PropertyMock
+
+import pytest
+from fastmcp import Client, Context, FastMCP
+from sqlalchemy import create_engine, text
+from sqlalchemy.engine import Engine
+from sqlalchemy.orm import Session
+from sqlalchemy.pool import NullPool
+
+from superset.extensions import db
+from superset.mcp_service.auth import mcp_auth_hook
+from superset.mcp_service.server import build_middleware_list
+from superset.mcp_service.session_scope import install_mcp_session_scoping
+
+WORKERS = 2
+# The smallest pool that admits WORKERS tool calls: one connection per tool
+# worker, one per cancellation worker, and one for transport-side lookups.
+POOL_SIZE = 2 * WORKERS + 1
+
+
+@pytest.fixture
+def metadata_engine(app: Any, tmp_path: Path) -> Iterator[Engine]:
+ """Bind the metadata session to a bounded pool, as in production."""
+ engine = create_engine(
+ f"sqlite:///{tmp_path / 'metadata.db'}",
+ pool_size=POOL_SIZE,
+ max_overflow=0,
+ pool_timeout=4,
+ connect_args={"check_same_thread": False},
+ )
+ previous = {
+ key: app.config.get(key) for key in ("MCP_TOOL_WORKERS", "SQLLAB_TIMEOUT")
+ }
+ app.config["MCP_TOOL_WORKERS"] = WORKERS
+ app.config["SQLLAB_TIMEOUT"] = 5
+ install_mcp_session_scoping()
+ db.session.remove()
+ try:
+ with (
+ patch.object(
+ type(db),
+ "engines",
+ new_callable=PropertyMock,
+ return_value={None: engine},
+ ),
+ patch.object(
+ type(db), "engine", new_callable=PropertyMock, return_value=engine
+ ),
+ ):
+ yield engine
+ finally:
+ db.session.remove()
+ engine.dispose()
+ app.config.update(previous)
+
+
+@pytest.mark.asyncio
+async def test_saturated_workers_holding_metadata_pool_do_not_freeze_loop(
+ metadata_engine: Engine,
+) -> None:
+ """Every tool worker waits on the loop while every connection is checked out.
+
+ If the transport loop itself waits for a metadata connection (tools/list
+ RBAC filtering, audit logging), neither side can progress until a pool or
+ call timeout expires, and the whole server stops responding meanwhile.
+ """
+ holding = threading.Barrier(WORKERS + 1)
+ notify = threading.Event()
+
+ async def hold_and_notify(ctx: Context) -> str:
+ """Keep a metadata transaction open while reporting progress."""
+ db.session.execute(text("SELECT 1"))
+ holding.wait(timeout=5)
+ notify.wait(timeout=5)
+ await ctx.info("progress")
+ return "done"
+
+ def user_lookup() -> None:
+ """Resolve the caller for tools/list with a metadata query."""
+ db.session.execute(text("SELECT 1"))
+
+ mcp = FastMCP("metadata pool regression", middleware=build_middleware_list())
+ with (
+ patch("superset.mcp_service.auth._setup_user_context", return_value=None),
+ patch(
+ "superset.mcp_service.middleware.get_user_from_request",
+ side_effect=user_lookup,
+ ),
+ ):
+ mcp.tool(mcp_auth_hook(hold_and_notify), name="hold_and_notify")
+ async with Client(mcp) as client:
+ calls = [
+ asyncio.create_task(client.call_tool("hold_and_notify", {}))
+ for _ in range(WORKERS)
+ ]
+ await asyncio.to_thread(holding.wait, 5)
+ # Other holders, such as cancellations, own the rest of the pool.
+ held = [metadata_engine.connect() for _ in range(POOL_SIZE - WORKERS)]
+ gaps = []
+
+ async def heartbeat() -> None:
+ """Measure how long the transport loop goes unscheduled."""
+ last = time.monotonic()
+ while True:
+ await asyncio.sleep(0.01)
+ now = time.monotonic()
+ gaps.append(now - last)
+ last = now
+
+ beat = asyncio.create_task(heartbeat())
+ started = time.monotonic()
+ # tools/list needs metadata connections for its audit record and
+ # RBAC filtering. Release the workers once it is waiting for one,
+ # so their progress reports reach the loop while it waits.
+ listing = asyncio.create_task(client.list_tools())
+ threading.Timer(0.2, notify.set).start()
+ try:
+ await asyncio.sleep(0.05)
+ pinged = time.monotonic()
+ await client.ping()
+ ping_latency = time.monotonic() - pinged
+ results = await asyncio.gather(*calls)
+ elapsed = time.monotonic() - started
+ finally:
+ for connection in held:
+ connection.close()
+ beat.cancel()
+ tools = await listing
+
+ assert [result.content[0].text for result in results] == ["done"] * WORKERS
+ assert [tool.name for tool in tools] == ["hold_and_notify"]
+ assert ping_latency < 0.5
+ assert elapsed < 1.5
+ assert max(gaps) < 0.5
+
+
+@pytest.mark.parametrize(
+ ("pool_options", "configured", "expected"),
+ [
+ # SQLAlchemy's default QueuePool lends 5 + 10 connections.
+ ({}, None, 7),
+ ({}, 16, 7),
+ ({}, 4, 4),
+ ({"pool_size": 20, "max_overflow": 20}, None, 16),
+ ({"pool_size": 20, "max_overflow": 20}, 19, 19),
+ ({"max_overflow": -1}, 32, 32),
+ ({"poolclass": NullPool}, None, 16),
+ ],
+)
+def test_tool_workers_leave_metadata_connections_for_cancellation(
+ app: Any,
+ tmp_path: Path,
+ pool_options: dict[str, Any],
+ configured: int | None,
+ expected: int,
+) -> None:
+ """Each call and its cancellation get a connection, with one to spare."""
+ from superset.mcp_service.worker import tool_worker_count
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'metadata.db'}", **pool_options)
+ previous = app.config.get("MCP_TOOL_WORKERS")
+ app.config["MCP_TOOL_WORKERS"] = configured
+ try:
+ with patch.object(
+ type(db), "engine", new_callable=PropertyMock, return_value=engine
+ ):
+ assert tool_worker_count(app) == expected
+ finally:
+ app.config["MCP_TOOL_WORKERS"] = previous
+ engine.dispose()
+
+
+def test_metadata_pool_too_small_for_cancellation_is_rejected(
+ app: Any, tmp_path: Path
+) -> None:
+ """Refuse a pool where a call and its cancellation could take every slot."""
+ from superset.mcp_service.worker import tool_worker_count
+
+ engine = create_engine(
+ f"sqlite:///{tmp_path / 'metadata.db'}", pool_size=2, max_overflow=0
+ )
+ with (
+ patch.object(
+ type(db), "engine", new_callable=PropertyMock, return_value=engine
+ ),
+ pytest.raises(ValueError, match="needs at least 3"),
+ ):
+ tool_worker_count(app)
+ engine.dispose()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("strategy", ["bm25", "regex"])
+@pytest.mark.parametrize("request_backed", [False, True])
+async def test_threaded_search_returns_metadata_connection(
+ app: Any, metadata_engine: Engine, strategy: str, request_backed: bool
+) -> None:
+ """An inherited Flask context must not bypass the search thread's cleanup."""
+ from unittest.mock import AsyncMock, Mock
+
+ from fastmcp.server.transforms.search.base import BaseSearchTransform
+ from flask import g
+
+ from superset.mcp_service.server import _create_search_transform
+ from superset.mcp_service.session_scope import mcp_session_scopefunc
+
+ sessions: dict[Any, Session] = {}
+
+ def permission(tool: Any) -> bool:
+ """Perform the metadata lookup used by database-backed RBAC."""
+ sessions[mcp_session_scopefunc()] = db.session()
+ db.session.execute(text("SELECT 1"))
+ return True
+
+ transform = _create_search_transform(
+ strategy=strategy, kwargs={}, make_normalizing_call_tool=lambda _: None
+ )
+ tool = Mock()
+ try:
+ with app.test_request_context("/mcp/") if request_backed else app.app_context():
+ g.user = Mock()
+ with (
+ patch.object(
+ BaseSearchTransform,
+ "_get_visible_tools",
+ new_callable=AsyncMock,
+ return_value=[tool],
+ ),
+ patch(
+ "superset.mcp_service.auth.is_tool_visible_to_current_user",
+ side_effect=permission,
+ ),
+ ):
+ assert await transform._get_visible_tools(Mock()) == [tool]
+ assert metadata_engine.pool.checkedout() == 0
+ assert metadata_engine.pool.checkedout() == 0
+ assert all(key not in db.session.registry.registry for key in sessions)
+ finally:
+ # Also clean up when run against the leaking implementation.
+ for key, session in sessions.items():
+ session.close()
+ db.session.registry.registry.pop(key, None)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("audit_path", ["error", "response_size", "error_hook"])
+async def test_error_audit_waiting_for_metadata_does_not_block_loop(
+ app: Any, metadata_engine: Engine, audit_path: str
+) -> None:
+ """A timeout/busy error must not synchronously wait for an audit connection."""
+ from unittest.mock import AsyncMock, Mock
+
+ from fastmcp.exceptions import ToolError
+
+ from superset.mcp_service.middleware import (
+ GlobalErrorHandlerMiddleware,
+ ResponseSizeGuardMiddleware,
+ ToolResultCompatibilityMiddleware,
+ )
+ from superset.models.core import Log
+ from superset.utils.log import DBEventLogger
+
+ Log.__table__.create(metadata_engine)
+ held = [metadata_engine.connect() for _ in range(POOL_SIZE)]
+ heartbeat = asyncio.Event()
+
+ async def beat() -> None:
+ """Prove the transport can run while the audit checkout is blocked."""
+ await asyncio.sleep(0.01)
+ heartbeat.set()
+
+ task = asyncio.create_task(beat())
+ release = threading.Timer(0.3, held[-1].close)
+ release.start()
+ try:
+ context = Mock(method="tools/call")
+ context.message.name = "unknown_tool"
+ context.message.arguments = {}
+ event_logger = DBEventLogger()
+
+ def error_hook(error: Exception, info: dict[str, Any]) -> None:
+ """An operator hook may itself perform metadata I/O."""
+ event_logger.log(
+ user_id=None,
+ action="test_error_hook",
+ dashboard_id=None,
+ duration_ms=None,
+ slice_id=None,
+ referrer=None,
+ curated_payload={"tool": info["tool_name"]},
+ )
+
+ with (
+ app.test_request_context("/mcp/"),
+ patch("superset.mcp_service.middleware.event_logger", event_logger),
+ patch(
+ "superset.mcp_service.flask_singleton.get_flask_app", return_value=app
+ ),
+ patch.dict(app.config, {"MCP_ERROR_HOOK": error_hook}),
+ ):
+ if audit_path == "error":
+ with pytest.raises(ToolError, match="server busy"):
+ await GlobalErrorHandlerMiddleware()._handle_error(
+ ToolError("MCP server busy"), context, "execute_sql", 0
+ )
+ elif audit_path == "response_size":
+ with pytest.raises(ToolError):
+ await ResponseSizeGuardMiddleware(max_bytes=10).on_call_tool(
+ context, AsyncMock(return_value={"data": "x" * 100})
+ )
+ else:
+ result = await ToolResultCompatibilityMiddleware().on_call_tool(
+ context, AsyncMock(side_effect=RuntimeError("failed"))
+ )
+ assert result.is_error
+ assert heartbeat.is_set()
+ assert metadata_engine.pool.checkedout() == POOL_SIZE - 1
+ finally:
+ release.join()
+ for connection in held:
+ connection.close()
+ await task
+ assert metadata_engine.pool.checkedout() == 0
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("fail", [False, True])
+async def test_metadata_thread_owns_user_context_and_session(
+ app: Any, metadata_engine: Engine, fail: bool
+) -> None:
+ """Do not share even an expired request user or an inherited session token."""
+ from contextlib import nullcontext
+
+ from flask import g, request
+ from flask_appbuilder.security.sqla.models import User
+ from sqlalchemy import event
+
+ from superset.mcp_service.session_scope import _mcp_session_token
+ from superset.mcp_service.worker import (
+ get_context_user_id,
+ run_in_metadata_thread,
+ )
+
+ User.__table__.create(metadata_engine)
+ with metadata_engine.begin() as connection:
+ connection.execute(
+ User.__table__.insert(),
+ {
+ "id": 42,
+ "username": "metadata-user",
+ "email": "metadata@example.com",
+ "first_name": "Meta",
+ "last_name": "Data",
+ },
+ )
+ token = _mcp_session_token.set(object())
+ threads: list[int] = []
+ sessions: list[Session] = []
+ main_thread = threading.get_ident()
+
+ def record_thread(*args: Any) -> None:
+ """Track all SQL, including implicit ORM refreshes."""
+ threads.append(threading.get_ident())
+
+ def lookup() -> None:
+ """Read a reloaded user while owning fresh request globals and session."""
+ assert g.user is not user
+ assert g.user.username == "metadata-user"
+ assert request.path == "/mcp/"
+ assert g.routing_marker == "caller"
+ g.routing_marker = "metadata"
+ session = db.session()
+ sessions.append(session)
+ assert session is not parent_session
+ session.execute(text("SELECT 1"))
+ if fail:
+ raise ValueError("metadata failed")
+
+ try:
+ with app.test_request_context("/mcp/"):
+ parent_session = db.session()
+ user = parent_session.get(User, 42)
+ parent_session.commit() # expire attributes and return the connection
+ g.user = user
+ g.routing_marker = "caller"
+ event.listen(metadata_engine, "before_cursor_execute", record_thread)
+ try:
+ assert get_context_user_id() == 42
+ with (
+ pytest.raises(ValueError, match="metadata failed")
+ if fail
+ else nullcontext()
+ ):
+ await run_in_metadata_thread(lookup)
+ assert g.user is user
+ assert g.routing_marker == "caller"
+ assert db.session() is parent_session
+ assert threads
+ assert main_thread not in threads
+ assert metadata_engine.pool.checkedout() == 0
+ assert all(
+ session not in db.session.registry.registry.values()
+ for session in sessions
+ )
+ finally:
+ event.remove(metadata_engine, "before_cursor_execute", record_thread)
+ finally:
+ db.session.remove()
+ _mcp_session_token.reset(token)
+
+
+@pytest.mark.asyncio
+async def test_cancelled_metadata_awaiter_leaves_cleanup_to_thread(
+ metadata_engine: Engine,
+) -> None:
+ """A disconnect must neither tear down live I/O nor leak its connection."""
+ from sqlalchemy import event
+
+ from superset.mcp_service.worker import run_in_metadata_thread
+
+ entered = threading.Event()
+ release = threading.Event()
+ returned = threading.Event()
+
+ def checkin(*args: Any) -> None:
+ """Signal teardown after the thread's metadata I/O has completed."""
+ returned.set()
+
+ def lookup() -> None:
+ """Keep using the same live session after its awaiter disconnects."""
+ db.session.execute(text("SELECT 1"))
+ entered.set()
+ release.wait(5)
+ db.session.execute(text("SELECT 2"))
+
+ event.listen(metadata_engine, "checkin", checkin)
+ task = asyncio.create_task(run_in_metadata_thread(lookup))
+ try:
+ assert await asyncio.to_thread(entered.wait, 2)
+ task.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await task
+ assert metadata_engine.pool.checkedout() == 1
+ assert not returned.is_set()
+ finally:
+ release.set()
+ assert await asyncio.to_thread(returned.wait, 2)
+ event.remove(metadata_engine, "checkin", checkin)
+ assert metadata_engine.pool.checkedout() == 0
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("configured", [False, True])
+async def test_error_hook_does_not_require_metadata_user_reload(
+ app: Any, metadata_engine: Engine, configured: bool
+) -> None:
+ """Error reporting survives metadata outages without sharing caller state."""
+ from unittest.mock import Mock
+
+ from flask import g
+ from flask_appbuilder.security.sqla.models import User
+ from sqlalchemy.orm import make_transient_to_detached
+
+ from superset.mcp_service.middleware import _invoke_error_hook_off_loop
+
+ main_thread = threading.get_ident()
+ error = RuntimeError("tool failed")
+ info = {"tool_name": "execute_sql", "user_id": 42}
+ sessions: list[Session] = []
+
+ with app.test_request_context("/mcp/"):
+ user = User(id=42)
+ make_transient_to_detached(user)
+ g.user = user
+ parent_session = db.session()
+
+ def capture(exc: Exception, context: dict[str, Any]) -> None:
+ """Capture errors without loading a user from the unavailable DB."""
+ assert threading.get_ident() != main_thread
+ assert exc is error
+ assert context == info
+ assert getattr(g, "user", None) is None
+ assert db.session() is not parent_session
+ sessions.append(db.session())
Review Comment:
<div>
<div id="suggestion">
<div id="issue"><b>Swallowed test assertions</b></div>
<div id="fix">
`capture`'s asserts (503-508) run inside the `MCP_ERROR_HOOK` mock, but `_invoke_error_hook` catches every hook exception and only logs it (`middleware.py:227`), so these asserts can never fail the test โ a regression to shared caller state or main-thread hook execution would pass silently. Record observations in `capture` and assert them after `await _invoke_error_hook_off_loop(...)`.
</div>
</div>
<small><i>Code Review Run #56a20f</i></small>
</div>
---
Should Bito avoid suggestions like this for future reviews? (<a href=https://alpha.bito.ai/home/ai-agents/review-rules>Manage Rules</a>)
- [ ] Yes, avoid them
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on PR #44581:
URL: https://github.com/apache/superset/pull/44581#issuecomment-5800421437
AI Code Review is in progress (usually takes 3 to 15 minutes unless it's a very large PR).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "github-advanced-security[bot] (via GitHub)" <gi...@apache.org>.
github-advanced-security[bot] commented on code in PR #44581:
URL: https://github.com/apache/superset/pull/44581#discussion_r4152394193
##########
superset/mcp_service/worker.py:
##########
@@ -0,0 +1,904 @@
+# 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.
+"""Bounded, deadline-aware execution of MCP tools outside the transport loop.
+
+The tool's coroutine runs on a worker-owned loop because its database APIs are
+synchronous. Only transport notifications are marshalled back to the server
+loop. Flask and SQLAlchemy lifetimes belong to the worker, not the waiting
+request: a timed-out DBAPI call cannot outlive and reuse a torn-down session.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import logging
+import threading
+import time
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import contextmanager
+from contextvars import ContextVar, copy_context
+from typing import Any, Callable, Coroutine, Iterator, ParamSpec, TYPE_CHECKING, TypeVar
+from weakref import WeakKeyDictionary
+
+from fastmcp.exceptions import ToolError
+from flask import current_app, g, has_app_context, has_request_context
+from sqlalchemy import inspect as sa_inspect
+from sqlalchemy.orm.state import InstanceState
+from sqlalchemy.pool import QueuePool
+
+from superset.mcp_service.session_scope import _mcp_session_token
+
+if TYPE_CHECKING:
+ from flask import Flask
+
+ from superset.models.core import Database
+
+logger = logging.getLogger(__name__)
+_active_call: ContextVar[WorkerCall | None] = ContextVar(
+ "mcp_worker_call", default=None
+)
+_metadata_context_owned: ContextVar[bool] = ContextVar(
+ "mcp_metadata_context_owned", default=False
+)
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
+_pools_lock = threading.Lock()
+_pools: WeakKeyDictionary[Flask, WorkerPool] = WeakKeyDictionary()
+
+
+class WorkerDeadlineExceeded(BaseException):
+ """Stop abandoned work without being swallowed by tool error handlers."""
+
+
+BUSY_MESSAGE = "MCP server busy: all tool workers are occupied. Retry later."
+
+# Tools that only read Superset's metadata database. They are admitted under a
+# separate, larger bound so they keep answering while slow warehouse queries
+# hold every warehouse slot. Anything that can reach a warehouse, a semantic
+# layer or a screenshot service stays in the warehouse bound. A listed tool
+# that does reach a warehouse must still take a warehouse slot first (see
+# WorkerCall.admit_warehouse), so a wrong entry cannot overdraw the pool.
+METADATA_ONLY_TOOLS = frozenset(
+ {
+ "find_users",
+ "get_annotation_layer_info",
+ "get_chart_info",
+ "get_chart_type_schema",
+ "get_dashboard_info",
+ "get_database_info",
+ "get_dataset_info",
+ "get_instance_info",
+ "get_layer_annotation_info",
+ "get_query_info",
+ "get_report_info",
+ "get_rls_filter_info",
+ "get_role_info",
+ "get_saved_query_info",
+ "get_schema",
+ "get_tag_info",
+ "get_task_info",
+ "get_theme_info",
+ "get_user_info",
+ "health_check",
+ "list_annotation_layers",
+ "list_charts",
+ "list_dashboards",
+ "list_databases",
+ "list_datasets",
+ "list_layer_annotations",
+ "list_queries",
+ "list_reports",
+ "list_rls_filters",
+ "list_roles",
+ "list_saved_queries",
+ "list_tags",
+ "list_tasks",
+ "list_themes",
+ "list_users",
+ }
+)
+
+
+class WorkerPool:
+ """Bound submissions, including abandoned work; never queue behind a query."""
+
+ def __init__(
+ self,
+ size: int,
+ metadata_size: int = 0,
+ transport_size: int = 1,
+ auth_size: int = 1,
+ auth_pending: int = 32,
+ ) -> None:
+ if size < 1:
+ raise ValueError("MCP_TOOL_WORKERS must be positive")
+ if metadata_size < 0:
+ raise ValueError("MCP_METADATA_TOOL_WORKERS must not be negative")
+ self.slots = threading.BoundedSemaphore(size)
+ self.metadata_slots = (
+ threading.BoundedSemaphore(metadata_size) if metadata_size else None
+ )
+ self.executor = ThreadPoolExecutor(
+ size + metadata_size, thread_name_prefix="mcp-tool"
+ )
+ # Cancellation must not wait behind the warehouse work it is cancelling.
+ # Only calls holding a warehouse slot register cancellation.
+ self.cancel_slots = threading.BoundedSemaphore(size)
+ self.cancellations = ThreadPoolExecutor(size, thread_name_prefix="mcp-cancel")
+ # Transport-side metadata I/O (tools/list filtering, audit writes,
+ # error hooks) is bounded too, so it cannot overdraw the connections
+ # budgeted for tool workers. Callers wait for it on the event loop
+ # without blocking it.
+ self.transport = ThreadPoolExecutor(
+ transport_size, thread_name_prefix="mcp-metadata"
+ )
+ # API-key lookups run before a caller is authenticated, so they get
+ # their own thread: a flood of bad keys or a slow lookup must not
+ # delay audit writes and RBAC filtering. Pending lookups are capped;
+ # past the cap a key is rejected rather than queued.
+ self.auth = ThreadPoolExecutor(auth_size, thread_name_prefix="mcp-auth")
+ self.auth_slots = threading.BoundedSemaphore(auth_pending)
+
+ def admission(self, metadata_only: bool) -> threading.BoundedSemaphore:
+ """Choose the bound a call is admitted under."""
+ if metadata_only and self.metadata_slots is not None:
+ return self.metadata_slots
+ return self.slots
+
+ def submit(
+ self,
+ fn: Callable[[], Any],
+ finished: Callable[[], None],
+ slots: threading.BoundedSemaphore | None = None,
+ ) -> Future[Any]:
+ """Admit immediately or report overload without retaining a queued call."""
+ slots = self.slots if slots is None else slots
+ if not slots.acquire(blocking=False):
+ raise ToolError(BUSY_MESSAGE)
+ try:
+ future = self.executor.submit(fn)
+ except BaseException:
+ slots.release()
+ raise
+ future.add_done_callback(lambda _: finished())
+ return future
+
+ def cancel(self, fn: Callable[[], None], finished: Callable[[], None]) -> bool:
+ """Bound cancellation I/O separately so it cannot delay the caller."""
+ if not self.cancel_slots.acquire(blocking=False):
+ return False
+ 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()
+
+ 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
+
+
+DEFAULT_TOOL_WORKERS = 16
+DEFAULT_METADATA_TOOL_WORKERS = 16
+# Threads, and so metadata connections, for transport-side metadata I/O.
+TRANSPORT_METADATA_THREADS = 1
+# Threads, and so metadata connections, for transport-side API-key lookups.
+API_KEY_AUTH_THREADS = 1
+# API-key lookups admitted at once, running or waiting; more are rejected.
+API_KEY_AUTH_PENDING = 32
+# Metadata connections held by transport-side threads rather than tool calls.
+RESERVED_METADATA_THREADS = TRANSPORT_METADATA_THREADS + API_KEY_AUTH_THREADS
+
+
+def _metadata_pool_capacity(app: Flask) -> int | None:
+ """Return how many metadata connections can be checked out at once.
+
+ ``None`` means a checkout never waits for another holder to return one
+ (e.g. ``NullPool``, per-thread pools, or unlimited overflow).
+ """
+ from contextlib import nullcontext
+
+ from superset import db
+
+ # Popping a pushed context tears down the caller's scoped session, and
+ # pools are created lazily from transport-side code that holds one.
+ in_app = has_app_context() and current_app._get_current_object() is app
+ with nullcontext() if in_app else app.app_context():
+ pool = db.engine.pool
+ if not isinstance(pool, QueuePool):
+ return None
+ # SQLAlchemy has no public accessor for the configured overflow limit.
+ max_overflow = pool._max_overflow # pylint: disable=protected-access
+ return None if max_overflow < 0 else pool.size() + max_overflow
+
+
+def admission_counts(app: Flask) -> tuple[int, int]:
+ """Split the metadata pool between warehouse and metadata-only calls.
+
+ Every holder of a metadata connection is bounded and budgeted:
+
+ - each warehouse-capable call can hold one for the whole of its warehouse
+ I/O, and its cancellation needs another (``2 * workers``);
+ - each metadata-only call holds at most one (``metadata_workers``);
+ - transport-side metadata I/O runs on ``TRANSPORT_METADATA_THREADS`` and
+ API-key lookups on ``API_KEY_AUTH_THREADS``.
+
+ ``2 * workers + metadata_workers + RESERVED_METADATA_THREADS`` never
+ exceeds the pool's capacity, so with every slot admitted no checkout
+ waits for another holder. Unless configured, warehouse calls get about
+ two thirds of the remaining connections and metadata-only calls the
+ rest, each up to 16. Explicit settings above the budget are reduced
+ with a warning. ``0`` metadata-only workers admits those tools under
+ the warehouse bound.
+
+ Returns ``(workers, metadata_workers)``.
+ """
+ configured = app.config.get("MCP_TOOL_WORKERS")
+ configured_metadata = app.config.get("MCP_METADATA_TOOL_WORKERS")
+ if configured_metadata is not None and configured_metadata < 0:
+ raise ValueError("MCP_METADATA_TOOL_WORKERS must not be negative")
+ capacity = _metadata_pool_capacity(app)
+ if capacity is None:
+ return (
+ DEFAULT_TOOL_WORKERS if configured is None else configured,
+ DEFAULT_METADATA_TOOL_WORKERS
+ if configured_metadata is None
+ else configured_metadata,
+ )
+ available = capacity - RESERVED_METADATA_THREADS
+ limit = available // 2
+ if limit < 1:
+ raise ValueError(
+ f"The metadata database pool allows {capacity} connections; "
+ "MCP tool execution needs at least "
+ f"{2 + RESERVED_METADATA_THREADS}"
+ )
+ if configured is None:
+ workers = min(DEFAULT_TOOL_WORKERS, max(1, available // 3))
+ elif configured > limit:
+ logger.warning(
+ "MCP_TOOL_WORKERS=%s needs %s metadata database connections, but the "
+ "pool allows %s; admitting %s concurrent tool calls. Raise the pool's "
+ "pool_size/max_overflow in SQLALCHEMY_ENGINE_OPTIONS to admit more.",
+ configured,
+ 2 * configured + RESERVED_METADATA_THREADS,
+ capacity,
+ limit,
+ )
+ workers = limit
+ else:
+ workers = configured
+ remaining = available - 2 * workers
+ if configured_metadata is None:
+ metadata_workers = min(DEFAULT_METADATA_TOOL_WORKERS, remaining)
+ elif configured_metadata > remaining:
+ logger.warning(
+ "MCP_METADATA_TOOL_WORKERS=%s exceeds the %s metadata database "
+ "connections left after %s warehouse-capable tool calls; admitting "
+ "%s metadata-only tool calls.",
+ configured_metadata,
+ remaining,
+ workers,
+ remaining,
+ )
+ metadata_workers = remaining
+ else:
+ metadata_workers = configured_metadata
+ return workers, metadata_workers
+
+
+def tool_worker_count(app: Flask) -> int:
+ """Admit only as many warehouse-capable calls as the pool can always serve."""
+ return admission_counts(app)[0]
+
+
+def metadata_tool_worker_count(app: Flask) -> int:
+ """Admit metadata-only calls from what warehouse calls leave of the pool."""
+ return admission_counts(app)[1]
+
+
+def _get_pool(app: Flask) -> WorkerPool:
+ """Lazily create a pool for this application, without import-time threads."""
+ with _pools_lock:
+ if app not in _pools:
+ size, metadata_size = admission_counts(app)
+ logger.info(
+ "MCP tool calls admitted concurrently: %s warehouse-capable, "
+ "%s metadata-only; %s transport metadata thread(s), "
+ "%s API-key lookup thread(s)",
+ size,
+ metadata_size,
+ TRANSPORT_METADATA_THREADS,
+ API_KEY_AUTH_THREADS,
Review Comment:
## CodeQL / Clear-text logging of sensitive information
This expression logs [sensitive data (password)](1) as clear text.
[Show more details](https://github.com/apache/superset/security/code-scanning/3909)
##########
superset/mcp_service/worker.py:
##########
@@ -0,0 +1,904 @@
+# 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.
+"""Bounded, deadline-aware execution of MCP tools outside the transport loop.
+
+The tool's coroutine runs on a worker-owned loop because its database APIs are
+synchronous. Only transport notifications are marshalled back to the server
+loop. Flask and SQLAlchemy lifetimes belong to the worker, not the waiting
+request: a timed-out DBAPI call cannot outlive and reuse a torn-down session.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import logging
+import threading
+import time
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import contextmanager
+from contextvars import ContextVar, copy_context
+from typing import Any, Callable, Coroutine, Iterator, ParamSpec, TYPE_CHECKING, TypeVar
+from weakref import WeakKeyDictionary
+
+from fastmcp.exceptions import ToolError
+from flask import current_app, g, has_app_context, has_request_context
+from sqlalchemy import inspect as sa_inspect
+from sqlalchemy.orm.state import InstanceState
+from sqlalchemy.pool import QueuePool
+
+from superset.mcp_service.session_scope import _mcp_session_token
+
+if TYPE_CHECKING:
+ from flask import Flask
+
+ from superset.models.core import Database
+
+logger = logging.getLogger(__name__)
+_active_call: ContextVar[WorkerCall | None] = ContextVar(
+ "mcp_worker_call", default=None
+)
+_metadata_context_owned: ContextVar[bool] = ContextVar(
+ "mcp_metadata_context_owned", default=False
+)
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
+_pools_lock = threading.Lock()
+_pools: WeakKeyDictionary[Flask, WorkerPool] = WeakKeyDictionary()
+
+
+class WorkerDeadlineExceeded(BaseException):
+ """Stop abandoned work without being swallowed by tool error handlers."""
+
+
+BUSY_MESSAGE = "MCP server busy: all tool workers are occupied. Retry later."
+
+# Tools that only read Superset's metadata database. They are admitted under a
+# separate, larger bound so they keep answering while slow warehouse queries
+# hold every warehouse slot. Anything that can reach a warehouse, a semantic
+# layer or a screenshot service stays in the warehouse bound. A listed tool
+# that does reach a warehouse must still take a warehouse slot first (see
+# WorkerCall.admit_warehouse), so a wrong entry cannot overdraw the pool.
+METADATA_ONLY_TOOLS = frozenset(
+ {
+ "find_users",
+ "get_annotation_layer_info",
+ "get_chart_info",
+ "get_chart_type_schema",
+ "get_dashboard_info",
+ "get_database_info",
+ "get_dataset_info",
+ "get_instance_info",
+ "get_layer_annotation_info",
+ "get_query_info",
+ "get_report_info",
+ "get_rls_filter_info",
+ "get_role_info",
+ "get_saved_query_info",
+ "get_schema",
+ "get_tag_info",
+ "get_task_info",
+ "get_theme_info",
+ "get_user_info",
+ "health_check",
+ "list_annotation_layers",
+ "list_charts",
+ "list_dashboards",
+ "list_databases",
+ "list_datasets",
+ "list_layer_annotations",
+ "list_queries",
+ "list_reports",
+ "list_rls_filters",
+ "list_roles",
+ "list_saved_queries",
+ "list_tags",
+ "list_tasks",
+ "list_themes",
+ "list_users",
+ }
+)
+
+
+class WorkerPool:
+ """Bound submissions, including abandoned work; never queue behind a query."""
+
+ def __init__(
+ self,
+ size: int,
+ metadata_size: int = 0,
+ transport_size: int = 1,
+ auth_size: int = 1,
+ auth_pending: int = 32,
+ ) -> None:
+ if size < 1:
+ raise ValueError("MCP_TOOL_WORKERS must be positive")
+ if metadata_size < 0:
+ raise ValueError("MCP_METADATA_TOOL_WORKERS must not be negative")
+ self.slots = threading.BoundedSemaphore(size)
+ self.metadata_slots = (
+ threading.BoundedSemaphore(metadata_size) if metadata_size else None
+ )
+ self.executor = ThreadPoolExecutor(
+ size + metadata_size, thread_name_prefix="mcp-tool"
+ )
+ # Cancellation must not wait behind the warehouse work it is cancelling.
+ # Only calls holding a warehouse slot register cancellation.
+ self.cancel_slots = threading.BoundedSemaphore(size)
+ self.cancellations = ThreadPoolExecutor(size, thread_name_prefix="mcp-cancel")
+ # Transport-side metadata I/O (tools/list filtering, audit writes,
+ # error hooks) is bounded too, so it cannot overdraw the connections
+ # budgeted for tool workers. Callers wait for it on the event loop
+ # without blocking it.
+ self.transport = ThreadPoolExecutor(
+ transport_size, thread_name_prefix="mcp-metadata"
+ )
+ # API-key lookups run before a caller is authenticated, so they get
+ # their own thread: a flood of bad keys or a slow lookup must not
+ # delay audit writes and RBAC filtering. Pending lookups are capped;
+ # past the cap a key is rejected rather than queued.
+ self.auth = ThreadPoolExecutor(auth_size, thread_name_prefix="mcp-auth")
+ self.auth_slots = threading.BoundedSemaphore(auth_pending)
+
+ def admission(self, metadata_only: bool) -> threading.BoundedSemaphore:
+ """Choose the bound a call is admitted under."""
+ if metadata_only and self.metadata_slots is not None:
+ return self.metadata_slots
+ return self.slots
+
+ def submit(
+ self,
+ fn: Callable[[], Any],
+ finished: Callable[[], None],
+ slots: threading.BoundedSemaphore | None = None,
+ ) -> Future[Any]:
+ """Admit immediately or report overload without retaining a queued call."""
+ slots = self.slots if slots is None else slots
+ if not slots.acquire(blocking=False):
+ raise ToolError(BUSY_MESSAGE)
+ try:
+ future = self.executor.submit(fn)
+ except BaseException:
+ slots.release()
+ raise
+ future.add_done_callback(lambda _: finished())
+ return future
+
+ def cancel(self, fn: Callable[[], None], finished: Callable[[], None]) -> bool:
+ """Bound cancellation I/O separately so it cannot delay the caller."""
+ if not self.cancel_slots.acquire(blocking=False):
+ return False
+ 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()
+
+ 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
+
+
+DEFAULT_TOOL_WORKERS = 16
+DEFAULT_METADATA_TOOL_WORKERS = 16
+# Threads, and so metadata connections, for transport-side metadata I/O.
+TRANSPORT_METADATA_THREADS = 1
+# Threads, and so metadata connections, for transport-side API-key lookups.
+API_KEY_AUTH_THREADS = 1
+# API-key lookups admitted at once, running or waiting; more are rejected.
+API_KEY_AUTH_PENDING = 32
+# Metadata connections held by transport-side threads rather than tool calls.
+RESERVED_METADATA_THREADS = TRANSPORT_METADATA_THREADS + API_KEY_AUTH_THREADS
+
+
+def _metadata_pool_capacity(app: Flask) -> int | None:
+ """Return how many metadata connections can be checked out at once.
+
+ ``None`` means a checkout never waits for another holder to return one
+ (e.g. ``NullPool``, per-thread pools, or unlimited overflow).
+ """
+ from contextlib import nullcontext
+
+ from superset import db
+
+ # Popping a pushed context tears down the caller's scoped session, and
+ # pools are created lazily from transport-side code that holds one.
+ in_app = has_app_context() and current_app._get_current_object() is app
+ with nullcontext() if in_app else app.app_context():
+ pool = db.engine.pool
+ if not isinstance(pool, QueuePool):
+ return None
+ # SQLAlchemy has no public accessor for the configured overflow limit.
+ max_overflow = pool._max_overflow # pylint: disable=protected-access
+ return None if max_overflow < 0 else pool.size() + max_overflow
+
+
+def admission_counts(app: Flask) -> tuple[int, int]:
+ """Split the metadata pool between warehouse and metadata-only calls.
+
+ Every holder of a metadata connection is bounded and budgeted:
+
+ - each warehouse-capable call can hold one for the whole of its warehouse
+ I/O, and its cancellation needs another (``2 * workers``);
+ - each metadata-only call holds at most one (``metadata_workers``);
+ - transport-side metadata I/O runs on ``TRANSPORT_METADATA_THREADS`` and
+ API-key lookups on ``API_KEY_AUTH_THREADS``.
+
+ ``2 * workers + metadata_workers + RESERVED_METADATA_THREADS`` never
+ exceeds the pool's capacity, so with every slot admitted no checkout
+ waits for another holder. Unless configured, warehouse calls get about
+ two thirds of the remaining connections and metadata-only calls the
+ rest, each up to 16. Explicit settings above the budget are reduced
+ with a warning. ``0`` metadata-only workers admits those tools under
+ the warehouse bound.
+
+ Returns ``(workers, metadata_workers)``.
+ """
+ configured = app.config.get("MCP_TOOL_WORKERS")
+ configured_metadata = app.config.get("MCP_METADATA_TOOL_WORKERS")
+ if configured_metadata is not None and configured_metadata < 0:
+ raise ValueError("MCP_METADATA_TOOL_WORKERS must not be negative")
+ capacity = _metadata_pool_capacity(app)
+ if capacity is None:
+ return (
+ DEFAULT_TOOL_WORKERS if configured is None else configured,
+ DEFAULT_METADATA_TOOL_WORKERS
+ if configured_metadata is None
+ else configured_metadata,
+ )
+ available = capacity - RESERVED_METADATA_THREADS
+ limit = available // 2
+ if limit < 1:
+ raise ValueError(
+ f"The metadata database pool allows {capacity} connections; "
+ "MCP tool execution needs at least "
+ f"{2 + RESERVED_METADATA_THREADS}"
+ )
+ if configured is None:
+ workers = min(DEFAULT_TOOL_WORKERS, max(1, available // 3))
+ elif configured > limit:
+ logger.warning(
+ "MCP_TOOL_WORKERS=%s needs %s metadata database connections, but the "
+ "pool allows %s; admitting %s concurrent tool calls. Raise the pool's "
+ "pool_size/max_overflow in SQLALCHEMY_ENGINE_OPTIONS to admit more.",
+ configured,
+ 2 * configured + RESERVED_METADATA_THREADS,
Review Comment:
## CodeQL / Clear-text logging of sensitive information
This expression logs [sensitive data (password)](1) as clear text.
This expression logs [sensitive data (password)](1) as clear text.
[Show more details](https://github.com/apache/superset/security/code-scanning/3908)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on code in PR #44581:
URL: https://github.com/apache/superset/pull/44581#discussion_r4113937151
##########
superset/mcp_service/middleware.py:
##########
@@ -218,6 +221,16 @@ def _invoke_error_hook(error: Exception, hook_context: dict[str, Any]) -> None:
logger.warning("MCP_ERROR_HOOK raised an exception: %s", hook_error)
+async def _invoke_error_hook_off_loop(
+ error: Exception, hook_context: dict[str, Any]
+) -> None:
+ """Keep hook I/O and context-setup failures outside the error boundary."""
+ try:
+ await run_in_metadata_thread(_invoke_error_hook, error, hook_context)
+ except Exception as hook_error: # noqa: BLE001
+ logger.warning("Could not run MCP_ERROR_HOOK: %s", hook_error)
Review Comment:
<!-- Bito Reply -->
The suggestion to route the error hook through a metadata thread is appropriate and addresses the identified database dependency issue. By moving the hook execution outside the primary transport loop, the implementation ensures that hook I/O and context-setup failures do not interfere with the main error boundary, while also preventing unnecessary user-reload queries when no hook is configured.
**superset/mcp_service/middleware.py**
```
+async def _invoke_error_hook_off_loop(
+ error: Exception, hook_context: dict[str, Any]
+) -> None:
+ """Keep hook I/O and context-setup failures outside the error boundary."""
+ try:
+ await run_in_metadata_thread(_invoke_error_hook, error, hook_context)
+ except Exception as hook_error: # noqa: BLE001
+ logger.warning("Could not run MCP_ERROR_HOOK: %s", hook_error)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on PR #44581:
URL: https://github.com/apache/superset/pull/44581#issuecomment-5854910364
AI Code Review is in progress (usually takes 3 to 15 minutes unless it's a very large PR).
<details>
<summary>
Bito Usage Guide
</summary>
**Commands**
Type the following command in the pull request comment and save the comment.
- `/review` - Manually triggers an incremental AI Review.
- `/review full` - Manually triggers a full AI Review.
- `/pause` - Pauses automatic reviews on this pull request.
- `/resume` - Resumes automatic reviews.
- `/resolve` - Marks all Bito-posted review comments as resolved.
- `/abort` - Cancels all in-progress reviews.
Refer to the <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/available-commands" target="_blank" rel="noopener noreferrer">documentation</a> for additional commands.
**Configuration**
This repository uses `Superset` You can customize the agent settings <a href="https://alpha.bito.ai/home/ai-agents/code-review-agent" target="_blank" rel="noopener noreferrer">here</a> or contact your Bito workspace admin at evan@preset.io.
**Documentation & Help**
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/getting-started/install-run-using-bito-cloud/create-or-customize-an-agent-instance" target="_blank" rel="noopener noreferrer">Customize agent settings</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/implementing-custom-code-review-rules" target="_blank" rel="noopener noreferrer">Review rules</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent" target="_blank" rel="noopener noreferrer">General documentation</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/faqs" target="_blank" rel="noopener noreferrer">FAQ</a>
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on code in PR #44581:
URL: https://github.com/apache/superset/pull/44581#discussion_r4152368622
##########
superset/mcp_service/composite_token_verifier.py:
##########
@@ -180,9 +180,13 @@ async def verify_token(self, token: str) -> AccessToken | None:
"""
if any(token.startswith(prefix) for prefix in self._api_key_prefixes):
if self._app is not None:
+ from superset.mcp_service.worker import transport_executor
+
loop = asyncio.get_running_loop()
+ # Metadata I/O on the transport side shares a bounded,
+ # pool-budgeted executor with audit writes and RBAC filtering.
result = await loop.run_in_executor(
- None, self._validate_api_key_sync, token
+ transport_executor(self._app), self._validate_api_key_sync, token
Review Comment:
<!-- Bito Reply -->
The suggestion provided by the reviewer is valid and addresses a potential performance bottleneck by highlighting the risk of thread contention when multiple API-key lookups share a single-threaded executor. The implementation described in the user's reply, which uses a budgeted metadata connection pool to manage these lookups, effectively mitigates the risk of serialization and blocking identified in the review.
**superset/mcp_service/composite_token_verifier.py**
```
from superset.mcp_service.worker import transport_executor
loop = asyncio.get_running_loop()
# Metadata I/O on the transport side shares a bounded,
# pool-budgeted executor with audit writes and RBAC filtering.
result = await loop.run_in_executor(
transport_executor(self._app), self._validate_api_key_sync, token
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on code in PR #44581:
URL: https://github.com/apache/superset/pull/44581#discussion_r4090574075
##########
superset/mcp_service/auth.py:
##########
@@ -1259,7 +1277,27 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
_cleanup_session_on_error()
raise
- wrapper = async_wrapper
+ @functools.wraps(tool_func)
+ async def worker_wrapper(*args: Any, **kwargs: Any) -> Any:
+ from superset.mcp_service.worker import run_in_worker
+
+ bound = _tool_sig.bind_partial(*args, **kwargs)
+ request = bound.arguments.get("request")
+ seconds = getattr(request, "timeout", None)
+ if seconds is None:
+ if has_app_context():
+ seconds = current_app.config.get("SQLLAB_TIMEOUT", 30)
+ else:
+ from superset.mcp_service.flask_singleton import get_flask_app
+
+ seconds = get_flask_app().config.get("SQLLAB_TIMEOUT", 30)
+ # Bind ctx once on the transport loop, including positional callers.
+ bound.arguments.update(_inject_ctx(dict(bound.arguments)))
+ return await run_in_worker(
+ async_wrapper, (), dict(bound.arguments), seconds
+ )
+
+ wrapper = worker_wrapper
Review Comment:
Fixed in 19bb267df3 (direct-call test updates in f6d2f2d6b7). There are eleven registered synchronous tools; update_chart_preview reaches warehouse execution through generate_preview_from_form_data -> ChartDataCommand.run. FastMCP can offload sync calls itself, but that bypassed this PRโs bounded workers, deadline, and cancellation context. Both protected function kinds now use run_in_worker. The new sync chart/SQL Lab regressions failed before the fix and pass afterward; they check deadline responses and concurrent request progress. The complete MCP suite plus related modules passes (5,360 tests; four existing skips), and pre-commit passes.
##########
superset/models/core.py:
##########
@@ -602,35 +602,38 @@ def get_sqla_engine( # pylint: disable=too-many-arguments
sqlalchemy_uri=sqlalchemy_uri,
cacheable=not prequeries,
)
- if prequeries:
- # SQLAlchemy connect event: runs prequeries on every new
- # DBAPI connection (e.g. SET search_path for PostgreSQL).
- def run_prequeries(
- dbapi_connection: Any,
- connection_record: Any, # pylint: disable=unused-argument
- ) -> None:
- cursor = dbapi_connection.cursor()
+ from superset.sql.execution.cancellation import cancellable_engine
+
+ with cancellable_engine(self, engine, catalog, schema):
+ if prequeries:
+ # SQLAlchemy connect event: runs prequeries on every new
+ # DBAPI connection (e.g. SET search_path for PostgreSQL).
+ def run_prequeries(
+ dbapi_connection: Any,
+ connection_record: Any, # pylint: disable=unused-argument
+ ) -> None:
+ cursor = dbapi_connection.cursor()
+ try:
+ for prequery in prequeries:
+ cursor.execute(prequery)
Review Comment:
Fixed in 19bb267df3. The raw connect-event prequery cursor now enters cancellable_cursor before execution, checks the deadline before each statement, and refreshes cancellation handles afterward. The regression uses a blocking DBAPI prequery: it failed before the fix because cancellation was never dispatched; it now verifies cancellation, cursor cleanup, and that the next prequery never executes after timeout. The SQL suite also passes its 100% coverage gate (1,975 tests).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "github-advanced-security[bot] (via GitHub)" <gi...@apache.org>.
github-advanced-security[bot] commented on code in PR #44581:
URL: https://github.com/apache/superset/pull/44581#discussion_r4160611073
##########
superset/mcp_service/worker.py:
##########
@@ -0,0 +1,910 @@
+# 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.
+"""Bounded, deadline-aware execution of MCP tools outside the transport loop.
+
+The tool's coroutine runs on a worker-owned loop because its database APIs are
+synchronous. Only transport notifications are marshalled back to the server
+loop. Flask and SQLAlchemy lifetimes belong to the worker, not the waiting
+request: a timed-out DBAPI call cannot outlive and reuse a torn-down session.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import logging
+import threading
+import time
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import contextmanager
+from contextvars import ContextVar, copy_context
+from typing import Any, Callable, Coroutine, Iterator, ParamSpec, TYPE_CHECKING, TypeVar
+from weakref import WeakKeyDictionary
+
+from fastmcp.exceptions import ToolError
+from flask import current_app, g, has_app_context, has_request_context
+from sqlalchemy import inspect as sa_inspect
+from sqlalchemy.orm.state import InstanceState
+from sqlalchemy.pool import QueuePool
+
+from superset.mcp_service.session_scope import _mcp_session_token
+
+if TYPE_CHECKING:
+ from flask import Flask
+
+ from superset.models.core import Database
+
+logger = logging.getLogger(__name__)
+_active_call: ContextVar[WorkerCall | None] = ContextVar(
+ "mcp_worker_call", default=None
+)
+_metadata_context_owned: ContextVar[bool] = ContextVar(
+ "mcp_metadata_context_owned", default=False
+)
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
+_pools_lock = threading.Lock()
+_pools: WeakKeyDictionary[Flask, WorkerPool] = WeakKeyDictionary()
+
+
+class WorkerDeadlineExceeded(BaseException):
+ """Stop abandoned work without being swallowed by tool error handlers."""
+
+
+BUSY_MESSAGE = "MCP server busy: all tool workers are occupied. Retry later."
+
+# Tools that only read Superset's metadata database. They are admitted under a
+# separate, larger bound so they keep answering while slow warehouse queries
+# hold every warehouse slot. Anything that can reach a warehouse, a semantic
+# layer or a screenshot service stays in the warehouse bound. A listed tool
+# that does reach a warehouse must still take a warehouse slot first (see
+# WorkerCall.admit_warehouse), so a wrong entry cannot overdraw the pool.
+METADATA_ONLY_TOOLS = frozenset(
+ {
+ "find_users",
+ "get_annotation_layer_info",
+ "get_chart_info",
+ "get_chart_type_schema",
+ "get_dashboard_info",
+ "get_database_info",
+ "get_dataset_info",
+ "get_instance_info",
+ "get_layer_annotation_info",
+ "get_query_info",
+ "get_report_info",
+ "get_rls_filter_info",
+ "get_role_info",
+ "get_saved_query_info",
+ "get_schema",
+ "get_tag_info",
+ "get_task_info",
+ "get_theme_info",
+ "get_user_info",
+ "health_check",
+ "list_annotation_layers",
+ "list_charts",
+ "list_dashboards",
+ "list_databases",
+ "list_datasets",
+ "list_layer_annotations",
+ "list_queries",
+ "list_reports",
+ "list_rls_filters",
+ "list_roles",
+ "list_saved_queries",
+ "list_tags",
+ "list_tasks",
+ "list_themes",
+ "list_users",
+ }
+)
+
+
+class WorkerPool:
+ """Bound submissions, including abandoned work; never queue behind a query."""
+
+ def __init__(
+ self,
+ size: int,
+ metadata_size: int = 0,
+ transport_size: int = 1,
+ auth_size: int = 1,
+ auth_pending: int = 32,
+ ) -> None:
+ if size < 1:
+ raise ValueError("MCP_TOOL_WORKERS must be positive")
+ if metadata_size < 0:
+ raise ValueError("MCP_METADATA_TOOL_WORKERS must not be negative")
+ self.slots = threading.BoundedSemaphore(size)
+ self.metadata_slots = (
+ threading.BoundedSemaphore(metadata_size) if metadata_size else None
+ )
+ self.executor = ThreadPoolExecutor(
+ size + metadata_size, thread_name_prefix="mcp-tool"
+ )
+ # Cancellation must not wait behind the warehouse work it is cancelling.
+ # Only calls holding a warehouse slot register cancellation.
+ self.cancel_slots = threading.BoundedSemaphore(size)
+ self.cancellations = ThreadPoolExecutor(size, thread_name_prefix="mcp-cancel")
+ # Transport-side metadata I/O (tools/list filtering, audit writes,
+ # error hooks) is bounded too, so it cannot overdraw the connections
+ # budgeted for tool workers. Callers wait for it on the event loop
+ # without blocking it.
+ self.transport = ThreadPoolExecutor(
+ transport_size, thread_name_prefix="mcp-metadata"
+ )
+ # API-key lookups run before a caller is authenticated, so they get
+ # their own thread: a flood of bad keys or a slow lookup must not
+ # delay audit writes and RBAC filtering. Pending lookups are capped;
+ # past the cap a key is rejected rather than queued.
+ self.auth = ThreadPoolExecutor(auth_size, thread_name_prefix="mcp-auth")
+ self.auth_slots = threading.BoundedSemaphore(auth_pending)
+
+ def admission(self, metadata_only: bool) -> threading.BoundedSemaphore:
+ """Choose the bound a call is admitted under."""
+ if metadata_only and self.metadata_slots is not None:
+ return self.metadata_slots
+ return self.slots
+
+ def submit(
+ self,
+ fn: Callable[[], Any],
+ finished: Callable[[], None],
+ slots: threading.BoundedSemaphore | None = None,
+ ) -> Future[Any]:
+ """Admit immediately or report overload without retaining a queued call."""
+ slots = self.slots if slots is None else slots
+ if not slots.acquire(blocking=False):
+ raise ToolError(BUSY_MESSAGE)
+ try:
+ future = self.executor.submit(fn)
+ except BaseException:
+ slots.release()
+ raise
+ future.add_done_callback(lambda _: finished())
+ return future
+
+ def cancel(self, fn: Callable[[], None], finished: Callable[[], None]) -> bool:
+ """Bound cancellation I/O separately so it cannot delay the caller."""
+ if not self.cancel_slots.acquire(blocking=False):
+ return False
+ 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()
+
+ 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
+
+
+DEFAULT_TOOL_WORKERS = 16
+DEFAULT_METADATA_TOOL_WORKERS = 16
+# Threads, and so metadata connections, for transport-side metadata I/O.
+TRANSPORT_METADATA_THREADS = 1
+# Threads, and so metadata connections, for transport-side API-key lookups.
+API_KEY_CHECK_POOL_SIZE = 1
+# API-key lookups admitted at once, running or waiting; more are rejected.
+API_KEY_AUTH_PENDING = 32
+# Metadata connections held by transport-side threads rather than tool calls.
+METADATA_POOL_HEADROOM = TRANSPORT_METADATA_THREADS + API_KEY_CHECK_POOL_SIZE
+
+
+def _metadata_pool_capacity(app: Flask) -> int | None:
+ """Return how many metadata connections can be checked out at once.
+
+ ``None`` means a checkout never waits for another holder to return one
+ (e.g. ``NullPool``, per-thread pools, or unlimited overflow), including
+ for an application that never registered the metadata database: it has
+ no metadata connections to budget, and resolving ``db.engine`` for it
+ would raise rather than report a capacity.
+ """
+ from contextlib import nullcontext
+
+ from superset import db
+
+ if app.extensions.get("sqlalchemy") is not db:
+ return None
+
+ # Popping a pushed context tears down the caller's scoped session, and
+ # pools are created lazily from transport-side code that holds one.
+ in_app = has_app_context() and current_app._get_current_object() is app
+ with nullcontext() if in_app else app.app_context():
+ pool = db.engine.pool
+ if not isinstance(pool, QueuePool):
+ return None
+ # SQLAlchemy has no public accessor for the configured overflow limit.
+ max_overflow = pool._max_overflow # pylint: disable=protected-access
+ return None if max_overflow < 0 else pool.size() + max_overflow
+
+
+def admission_counts(app: Flask) -> tuple[int, int]:
+ """Split the metadata pool between warehouse and metadata-only calls.
+
+ Every holder of a metadata connection is bounded and budgeted:
+
+ - each warehouse-capable call can hold one for the whole of its warehouse
+ I/O, and its cancellation needs another (``2 * workers``);
+ - each metadata-only call holds at most one (``metadata_workers``);
+ - transport-side metadata I/O runs on ``TRANSPORT_METADATA_THREADS`` and
+ API-key lookups on ``API_KEY_CHECK_POOL_SIZE``.
+
+ ``2 * workers + metadata_workers + METADATA_POOL_HEADROOM`` never
+ exceeds the pool's capacity, so with every slot admitted no checkout
+ waits for another holder. Unless configured, warehouse calls get about
+ two thirds of the remaining connections and metadata-only calls the
+ rest, each up to 16. Explicit settings above the budget are reduced
+ with a warning. ``0`` metadata-only workers admits those tools under
+ the warehouse bound.
+
+ Returns ``(workers, metadata_workers)``.
+ """
+ configured = app.config.get("MCP_TOOL_WORKERS")
+ configured_metadata = app.config.get("MCP_METADATA_TOOL_WORKERS")
+ if configured_metadata is not None and configured_metadata < 0:
+ raise ValueError("MCP_METADATA_TOOL_WORKERS must not be negative")
+ capacity = _metadata_pool_capacity(app)
+ if capacity is None:
+ return (
+ DEFAULT_TOOL_WORKERS if configured is None else configured,
+ DEFAULT_METADATA_TOOL_WORKERS
+ if configured_metadata is None
+ else configured_metadata,
+ )
+ available = capacity - METADATA_POOL_HEADROOM
+ limit = available // 2
+ if limit < 1:
+ raise ValueError(
+ f"The metadata database pool allows {capacity} connections; "
+ "MCP tool execution needs at least "
+ f"{2 + METADATA_POOL_HEADROOM}"
+ )
+ if configured is None:
+ workers = min(DEFAULT_TOOL_WORKERS, max(1, available // 3))
+ elif configured > limit:
+ logger.warning(
+ "MCP_TOOL_WORKERS=%s needs %s metadata database connections, but the "
+ "pool allows %s; admitting %s concurrent tool calls. Raise the pool's "
+ "pool_size/max_overflow in SQLALCHEMY_ENGINE_OPTIONS to admit more.",
+ configured,
+ 2 * configured + METADATA_POOL_HEADROOM,
Review Comment:
## CodeQL / Clear-text logging of sensitive information
This expression logs [sensitive data (password)](1) as clear text.
This expression logs [sensitive data (password)](1) as clear text.
[Show more details](https://github.com/apache/superset/security/code-scanning/3914)
##########
superset/mcp_service/worker.py:
##########
@@ -0,0 +1,910 @@
+# 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.
+"""Bounded, deadline-aware execution of MCP tools outside the transport loop.
+
+The tool's coroutine runs on a worker-owned loop because its database APIs are
+synchronous. Only transport notifications are marshalled back to the server
+loop. Flask and SQLAlchemy lifetimes belong to the worker, not the waiting
+request: a timed-out DBAPI call cannot outlive and reuse a torn-down session.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import logging
+import threading
+import time
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import contextmanager
+from contextvars import ContextVar, copy_context
+from typing import Any, Callable, Coroutine, Iterator, ParamSpec, TYPE_CHECKING, TypeVar
+from weakref import WeakKeyDictionary
+
+from fastmcp.exceptions import ToolError
+from flask import current_app, g, has_app_context, has_request_context
+from sqlalchemy import inspect as sa_inspect
+from sqlalchemy.orm.state import InstanceState
+from sqlalchemy.pool import QueuePool
+
+from superset.mcp_service.session_scope import _mcp_session_token
+
+if TYPE_CHECKING:
+ from flask import Flask
+
+ from superset.models.core import Database
+
+logger = logging.getLogger(__name__)
+_active_call: ContextVar[WorkerCall | None] = ContextVar(
+ "mcp_worker_call", default=None
+)
+_metadata_context_owned: ContextVar[bool] = ContextVar(
+ "mcp_metadata_context_owned", default=False
+)
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
+_pools_lock = threading.Lock()
+_pools: WeakKeyDictionary[Flask, WorkerPool] = WeakKeyDictionary()
+
+
+class WorkerDeadlineExceeded(BaseException):
+ """Stop abandoned work without being swallowed by tool error handlers."""
+
+
+BUSY_MESSAGE = "MCP server busy: all tool workers are occupied. Retry later."
+
+# Tools that only read Superset's metadata database. They are admitted under a
+# separate, larger bound so they keep answering while slow warehouse queries
+# hold every warehouse slot. Anything that can reach a warehouse, a semantic
+# layer or a screenshot service stays in the warehouse bound. A listed tool
+# that does reach a warehouse must still take a warehouse slot first (see
+# WorkerCall.admit_warehouse), so a wrong entry cannot overdraw the pool.
+METADATA_ONLY_TOOLS = frozenset(
+ {
+ "find_users",
+ "get_annotation_layer_info",
+ "get_chart_info",
+ "get_chart_type_schema",
+ "get_dashboard_info",
+ "get_database_info",
+ "get_dataset_info",
+ "get_instance_info",
+ "get_layer_annotation_info",
+ "get_query_info",
+ "get_report_info",
+ "get_rls_filter_info",
+ "get_role_info",
+ "get_saved_query_info",
+ "get_schema",
+ "get_tag_info",
+ "get_task_info",
+ "get_theme_info",
+ "get_user_info",
+ "health_check",
+ "list_annotation_layers",
+ "list_charts",
+ "list_dashboards",
+ "list_databases",
+ "list_datasets",
+ "list_layer_annotations",
+ "list_queries",
+ "list_reports",
+ "list_rls_filters",
+ "list_roles",
+ "list_saved_queries",
+ "list_tags",
+ "list_tasks",
+ "list_themes",
+ "list_users",
+ }
+)
+
+
+class WorkerPool:
+ """Bound submissions, including abandoned work; never queue behind a query."""
+
+ def __init__(
+ self,
+ size: int,
+ metadata_size: int = 0,
+ transport_size: int = 1,
+ auth_size: int = 1,
+ auth_pending: int = 32,
+ ) -> None:
+ if size < 1:
+ raise ValueError("MCP_TOOL_WORKERS must be positive")
+ if metadata_size < 0:
+ raise ValueError("MCP_METADATA_TOOL_WORKERS must not be negative")
+ self.slots = threading.BoundedSemaphore(size)
+ self.metadata_slots = (
+ threading.BoundedSemaphore(metadata_size) if metadata_size else None
+ )
+ self.executor = ThreadPoolExecutor(
+ size + metadata_size, thread_name_prefix="mcp-tool"
+ )
+ # Cancellation must not wait behind the warehouse work it is cancelling.
+ # Only calls holding a warehouse slot register cancellation.
+ self.cancel_slots = threading.BoundedSemaphore(size)
+ self.cancellations = ThreadPoolExecutor(size, thread_name_prefix="mcp-cancel")
+ # Transport-side metadata I/O (tools/list filtering, audit writes,
+ # error hooks) is bounded too, so it cannot overdraw the connections
+ # budgeted for tool workers. Callers wait for it on the event loop
+ # without blocking it.
+ self.transport = ThreadPoolExecutor(
+ transport_size, thread_name_prefix="mcp-metadata"
+ )
+ # API-key lookups run before a caller is authenticated, so they get
+ # their own thread: a flood of bad keys or a slow lookup must not
+ # delay audit writes and RBAC filtering. Pending lookups are capped;
+ # past the cap a key is rejected rather than queued.
+ self.auth = ThreadPoolExecutor(auth_size, thread_name_prefix="mcp-auth")
+ self.auth_slots = threading.BoundedSemaphore(auth_pending)
+
+ def admission(self, metadata_only: bool) -> threading.BoundedSemaphore:
+ """Choose the bound a call is admitted under."""
+ if metadata_only and self.metadata_slots is not None:
+ return self.metadata_slots
+ return self.slots
+
+ def submit(
+ self,
+ fn: Callable[[], Any],
+ finished: Callable[[], None],
+ slots: threading.BoundedSemaphore | None = None,
+ ) -> Future[Any]:
+ """Admit immediately or report overload without retaining a queued call."""
+ slots = self.slots if slots is None else slots
+ if not slots.acquire(blocking=False):
+ raise ToolError(BUSY_MESSAGE)
+ try:
+ future = self.executor.submit(fn)
+ except BaseException:
+ slots.release()
+ raise
+ future.add_done_callback(lambda _: finished())
+ return future
+
+ def cancel(self, fn: Callable[[], None], finished: Callable[[], None]) -> bool:
+ """Bound cancellation I/O separately so it cannot delay the caller."""
+ if not self.cancel_slots.acquire(blocking=False):
+ return False
+ 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()
+
+ 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
+
+
+DEFAULT_TOOL_WORKERS = 16
+DEFAULT_METADATA_TOOL_WORKERS = 16
+# Threads, and so metadata connections, for transport-side metadata I/O.
+TRANSPORT_METADATA_THREADS = 1
+# Threads, and so metadata connections, for transport-side API-key lookups.
+API_KEY_CHECK_POOL_SIZE = 1
+# API-key lookups admitted at once, running or waiting; more are rejected.
+API_KEY_AUTH_PENDING = 32
+# Metadata connections held by transport-side threads rather than tool calls.
+METADATA_POOL_HEADROOM = TRANSPORT_METADATA_THREADS + API_KEY_CHECK_POOL_SIZE
+
+
+def _metadata_pool_capacity(app: Flask) -> int | None:
+ """Return how many metadata connections can be checked out at once.
+
+ ``None`` means a checkout never waits for another holder to return one
+ (e.g. ``NullPool``, per-thread pools, or unlimited overflow), including
+ for an application that never registered the metadata database: it has
+ no metadata connections to budget, and resolving ``db.engine`` for it
+ would raise rather than report a capacity.
+ """
+ from contextlib import nullcontext
+
+ from superset import db
+
+ if app.extensions.get("sqlalchemy") is not db:
+ return None
+
+ # Popping a pushed context tears down the caller's scoped session, and
+ # pools are created lazily from transport-side code that holds one.
+ in_app = has_app_context() and current_app._get_current_object() is app
+ with nullcontext() if in_app else app.app_context():
+ pool = db.engine.pool
+ if not isinstance(pool, QueuePool):
+ return None
+ # SQLAlchemy has no public accessor for the configured overflow limit.
+ max_overflow = pool._max_overflow # pylint: disable=protected-access
+ return None if max_overflow < 0 else pool.size() + max_overflow
+
+
+def admission_counts(app: Flask) -> tuple[int, int]:
+ """Split the metadata pool between warehouse and metadata-only calls.
+
+ Every holder of a metadata connection is bounded and budgeted:
+
+ - each warehouse-capable call can hold one for the whole of its warehouse
+ I/O, and its cancellation needs another (``2 * workers``);
+ - each metadata-only call holds at most one (``metadata_workers``);
+ - transport-side metadata I/O runs on ``TRANSPORT_METADATA_THREADS`` and
+ API-key lookups on ``API_KEY_CHECK_POOL_SIZE``.
+
+ ``2 * workers + metadata_workers + METADATA_POOL_HEADROOM`` never
+ exceeds the pool's capacity, so with every slot admitted no checkout
+ waits for another holder. Unless configured, warehouse calls get about
+ two thirds of the remaining connections and metadata-only calls the
+ rest, each up to 16. Explicit settings above the budget are reduced
+ with a warning. ``0`` metadata-only workers admits those tools under
+ the warehouse bound.
+
+ Returns ``(workers, metadata_workers)``.
+ """
+ configured = app.config.get("MCP_TOOL_WORKERS")
+ configured_metadata = app.config.get("MCP_METADATA_TOOL_WORKERS")
+ if configured_metadata is not None and configured_metadata < 0:
+ raise ValueError("MCP_METADATA_TOOL_WORKERS must not be negative")
+ capacity = _metadata_pool_capacity(app)
+ if capacity is None:
+ return (
+ DEFAULT_TOOL_WORKERS if configured is None else configured,
+ DEFAULT_METADATA_TOOL_WORKERS
+ if configured_metadata is None
+ else configured_metadata,
+ )
+ available = capacity - METADATA_POOL_HEADROOM
+ limit = available // 2
+ if limit < 1:
+ raise ValueError(
+ f"The metadata database pool allows {capacity} connections; "
+ "MCP tool execution needs at least "
+ f"{2 + METADATA_POOL_HEADROOM}"
+ )
+ if configured is None:
+ workers = min(DEFAULT_TOOL_WORKERS, max(1, available // 3))
+ elif configured > limit:
+ logger.warning(
+ "MCP_TOOL_WORKERS=%s needs %s metadata database connections, but the "
+ "pool allows %s; admitting %s concurrent tool calls. Raise the pool's "
+ "pool_size/max_overflow in SQLALCHEMY_ENGINE_OPTIONS to admit more.",
+ configured,
+ 2 * configured + METADATA_POOL_HEADROOM,
+ capacity,
+ limit,
+ )
+ workers = limit
+ else:
+ workers = configured
+ remaining = available - 2 * workers
+ if configured_metadata is None:
+ metadata_workers = min(DEFAULT_METADATA_TOOL_WORKERS, remaining)
+ elif configured_metadata > remaining:
+ logger.warning(
+ "MCP_METADATA_TOOL_WORKERS=%s exceeds the %s metadata database "
+ "connections left after %s warehouse-capable tool calls; admitting "
+ "%s metadata-only tool calls.",
+ configured_metadata,
+ remaining,
+ workers,
+ remaining,
+ )
+ metadata_workers = remaining
+ else:
+ metadata_workers = configured_metadata
+ return workers, metadata_workers
+
+
+def tool_worker_count(app: Flask) -> int:
+ """Admit only as many warehouse-capable calls as the pool can always serve."""
+ return admission_counts(app)[0]
+
+
+def metadata_tool_worker_count(app: Flask) -> int:
+ """Admit metadata-only calls from what warehouse calls leave of the pool."""
+ return admission_counts(app)[1]
+
+
+def _get_pool(app: Flask) -> WorkerPool:
+ """Lazily create a pool for this application, without import-time threads."""
+ with _pools_lock:
+ if app not in _pools:
+ size, metadata_size = admission_counts(app)
+ logger.info(
+ "MCP tool calls admitted concurrently: %s warehouse-capable, "
+ "%s metadata-only; %s transport metadata thread(s), "
+ "%s API-key lookup thread(s)",
+ size,
+ metadata_size,
+ TRANSPORT_METADATA_THREADS,
+ API_KEY_CHECK_POOL_SIZE,
Review Comment:
## CodeQL / Clear-text logging of sensitive information
This expression logs [sensitive data (password)](1) as clear text.
[Show more details](https://github.com/apache/superset/security/code-scanning/3915)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi merged PR #44581:
URL: https://github.com/apache/superset/pull/44581
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "codecov[bot] (via GitHub)" <gi...@apache.org>.
codecov[bot] commented on PR #44581:
URL: https://github.com/apache/superset/pull/44581#issuecomment-5800497843
## [Codecov](https://app.codecov.io/gh/apache/superset/pull/44581?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
:x: Patch coverage is `12.29947%` with `328 lines` in your changes missing coverage. Please review.
:white_check_mark: Project coverage is 57.02%. Comparing base ([`c0c688d`](https://app.codecov.io/gh/apache/superset/commit/c0c688d3e4b2696bf481757f5364a5dcc4184fbe?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)) to head ([`99e9c17`](https://app.codecov.io/gh/apache/superset/commit/99e9c1722b5803708ede42130277ca5987c27881?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)).
:warning: Report is 1 commits behind head on master.
| [Files with missing lines](https://app.codecov.io/gh/apache/superset/pull/44581?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Patch % | Lines |
|---|---|---|
| [superset/mcp\_service/worker.py](https://app.codecov.io/gh/apache/superset/pull/44581?src=pr&el=tree&filepath=superset%2Fmcp_service%2Fworker.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c3VwZXJzZXQvbWNwX3NlcnZpY2Uvd29ya2VyLnB5) | 0.00% | [256 Missing :warning: ](https://app.codecov.io/gh/apache/superset/pull/44581?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [superset/mcp\_service/auth.py](https://app.codecov.io/gh/apache/superset/pull/44581?src=pr&el=tree&filepath=superset%2Fmcp_service%2Fauth.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c3VwZXJzZXQvbWNwX3NlcnZpY2UvYXV0aC5weQ==) | 10.71% | [25 Missing :warning: ](https://app.codecov.io/gh/apache/superset/pull/44581?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [superset/sql/execution/cancellation.py](https://app.codecov.io/gh/apache/superset/pull/44581?src=pr&el=tree&filepath=superset%2Fsql%2Fexecution%2Fcancellation.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c3VwZXJzZXQvc3FsL2V4ZWN1dGlvbi9jYW5jZWxsYXRpb24ucHk=) | 53.84% | [17 Missing and 7 partials :warning: ](https://app.codecov.io/gh/apache/superset/pull/44581?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [superset/models/core.py](https://app.codecov.io/gh/apache/superset/pull/44581?src=pr&el=tree&filepath=superset%2Fmodels%2Fcore.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c3VwZXJzZXQvbW9kZWxzL2NvcmUucHk=) | 51.85% | [11 Missing and 2 partials :warning: ](https://app.codecov.io/gh/apache/superset/pull/44581?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [superset/sql/execution/executor.py](https://app.codecov.io/gh/apache/superset/pull/44581?src=pr&el=tree&filepath=superset%2Fsql%2Fexecution%2Fexecutor.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c3VwZXJzZXQvc3FsL2V4ZWN1dGlvbi9leGVjdXRvci5weQ==) | 0.00% | [9 Missing :warning: ](https://app.codecov.io/gh/apache/superset/pull/44581?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [superset/tasks/query\_cancel.py](https://app.codecov.io/gh/apache/superset/pull/44581?src=pr&el=tree&filepath=superset%2Ftasks%2Fquery_cancel.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c3VwZXJzZXQvdGFza3MvcXVlcnlfY2FuY2VsLnB5) | 0.00% | [1 Missing :warning: ](https://app.codecov.io/gh/apache/superset/pull/44581?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
> :exclamation: There is a different number of reports uploaded between BASE (c0c688d) and HEAD (99e9c17). Click for more details.
>
> <details><summary>HEAD has 9 uploads less than BASE</summary>
>
>| Flag | BASE (c0c688d) | HEAD (99e9c17) |
>|------|------|------|
>|python|9|2|
>|presto|2|1|
>|hive|2|1|
></details>
<details><summary>Additional details and impacted files</summary>
```diff
@@ Coverage Diff @@
## master #44581 +/- ##
===========================================
- Coverage 81.03% 57.02% -24.01%
===========================================
Files 2952 2957 +5
Lines 177996 178564 +568
Branches 41259 41317 +58
===========================================
- Hits 144231 101833 -42398
- Misses 31048 75748 +44700
+ Partials 2717 983 -1734
```
| [Flag](https://app.codecov.io/gh/apache/superset/pull/44581/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Coverage ฮ | |
|---|---|---|
| [hive](https://app.codecov.io/gh/apache/superset/pull/44581/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `36.69% <8.55%> (-0.14%)` | :arrow_down: |
| [mysql](https://app.codecov.io/gh/apache/superset/pull/44581/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
| [postgres](https://app.codecov.io/gh/apache/superset/pull/44581/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
| [presto](https://app.codecov.io/gh/apache/superset/pull/44581/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `38.62% <12.29%> (-0.14%)` | :arrow_down: |
| [python](https://app.codecov.io/gh/apache/superset/pull/44581/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `38.66% <12.29%> (-46.66%)` | :arrow_down: |
| [sqlite](https://app.codecov.io/gh/apache/superset/pull/44581/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
| [unit](https://app.codecov.io/gh/apache/superset/pull/44581/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#carryforward-flags-in-the-pull-request-comment) to find out more.
</details>
[:umbrella: View full report in Codecov by Harness](https://app.codecov.io/gh/apache/superset/pull/44581?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
<details><summary> :rocket: New features to boost your workflow: </summary>
- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
- :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on PR #44581:
URL: https://github.com/apache/superset/pull/44581#issuecomment-5840484364
AI Code Review is in progress (usually takes 3 to 15 minutes unless it's a very large PR).
<details>
<summary>
Bito Usage Guide
</summary>
**Commands**
Type the following command in the pull request comment and save the comment.
- `/review` - Manually triggers an incremental AI Review.
- `/review full` - Manually triggers a full AI Review.
- `/pause` - Pauses automatic reviews on this pull request.
- `/resume` - Resumes automatic reviews.
- `/resolve` - Marks all Bito-posted review comments as resolved.
- `/abort` - Cancels all in-progress reviews.
Refer to the <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/available-commands" target="_blank" rel="noopener noreferrer">documentation</a> for additional commands.
**Configuration**
This repository uses `Superset` You can customize the agent settings <a href="https://alpha.bito.ai/home/ai-agents/code-review-agent" target="_blank" rel="noopener noreferrer">here</a> or contact your Bito workspace admin at evan@preset.io.
**Documentation & Help**
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/getting-started/install-run-using-bito-cloud/create-or-customize-an-agent-instance" target="_blank" rel="noopener noreferrer">Customize agent settings</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/implementing-custom-code-review-rules" target="_blank" rel="noopener noreferrer">Review rules</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent" target="_blank" rel="noopener noreferrer">General documentation</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/faqs" target="_blank" rel="noopener noreferrer">FAQ</a>
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on code in PR #44581:
URL: https://github.com/apache/superset/pull/44581#discussion_r4112948712
##########
superset/mcp_service/middleware.py:
##########
@@ -218,6 +221,16 @@ def _invoke_error_hook(error: Exception, hook_context: dict[str, Any]) -> None:
logger.warning("MCP_ERROR_HOOK raised an exception: %s", hook_error)
+async def _invoke_error_hook_off_loop(
+ error: Exception, hook_context: dict[str, Any]
+) -> None:
+ """Keep hook I/O and context-setup failures outside the error boundary."""
+ try:
+ await run_in_metadata_thread(_invoke_error_hook, error, hook_context)
+ except Exception as hook_error: # noqa: BLE001
+ logger.warning("Could not run MCP_ERROR_HOOK: %s", hook_error)
Review Comment:
<div>
<div id="suggestion">
<div id="issue"><b>Error hook DB dependency</b></div>
<div id="fix">
`_invoke_error_hook` only reads `MCP_ERROR_HOOK` config via `get_flask_app()`; it never touches the metadata session. Routing it through `run_in_metadata_thread` adds a DB dependency โ `execute()` reloads `g.user` via `db.session.get` (worker.py) โ so during a metadata outage the hook is silently skipped (caught at line 230), and every tool error pays a needless user-reload query even when no hook is configured.
</div>
</div>
<small><i>Code Review Run #8985f4</i></small>
</div>
---
Should Bito avoid suggestions like this for future reviews? (<a href=https://alpha.bito.ai/home/ai-agents/review-rules>Manage Rules</a>)
- [ ] Yes, avoid them
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on code in PR #44581:
URL: https://github.com/apache/superset/pull/44581#discussion_r4113936769
##########
superset/mcp_service/middleware.py:
##########
@@ -218,6 +221,16 @@ def _invoke_error_hook(error: Exception, hook_context: dict[str, Any]) -> None:
logger.warning("MCP_ERROR_HOOK raised an exception: %s", hook_error)
+async def _invoke_error_hook_off_loop(
+ error: Exception, hook_context: dict[str, Any]
+) -> None:
+ """Keep hook I/O and context-setup failures outside the error boundary."""
+ try:
+ await run_in_metadata_thread(_invoke_error_hook, error, hook_context)
+ except Exception as hook_error: # noqa: BLE001
+ logger.warning("Could not run MCP_ERROR_HOOK: %s", hook_error)
Review Comment:
Valid finding; fixed in 480d27adca0a8039f3bc3e040ed8b2d62c5a4719. Error hooks run off the transport loop with their own application context and session cleanup, but without reloading g.user. An absent hook also performs no user lookup. Added regression coverage with an ORM request user and a failing metadata lookup, covering both configured and absent hooks and verifying worker/session isolation. Existing database-writing hook coverage still passes. MCP unit suite: 5,487 passed, 3 skipped; pre-commit passed on all branch-changed files, including mypy.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on code in PR #44581:
URL: https://github.com/apache/superset/pull/44581#discussion_r4114434011
##########
tests/unit_tests/mcp_service/test_worker_metadata_pool.py:
##########
@@ -0,0 +1,530 @@
+# 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.
+"""Tool workers and the transport loop share the metadata connection pool."""
+
+import asyncio
+import threading
+import time
+from collections.abc import Iterator
+from pathlib import Path
+from typing import Any
+from unittest.mock import patch, PropertyMock
+
+import pytest
+from fastmcp import Client, Context, FastMCP
+from sqlalchemy import create_engine, text
+from sqlalchemy.engine import Engine
+from sqlalchemy.orm import Session
+from sqlalchemy.pool import NullPool
+
+from superset.extensions import db
+from superset.mcp_service.auth import mcp_auth_hook
+from superset.mcp_service.server import build_middleware_list
+from superset.mcp_service.session_scope import install_mcp_session_scoping
+
+WORKERS = 2
+# The smallest pool that admits WORKERS tool calls: one connection per tool
+# worker, one per cancellation worker, and one for transport-side lookups.
+POOL_SIZE = 2 * WORKERS + 1
+
+
+@pytest.fixture
+def metadata_engine(app: Any, tmp_path: Path) -> Iterator[Engine]:
+ """Bind the metadata session to a bounded pool, as in production."""
+ engine = create_engine(
+ f"sqlite:///{tmp_path / 'metadata.db'}",
+ pool_size=POOL_SIZE,
+ max_overflow=0,
+ pool_timeout=4,
+ connect_args={"check_same_thread": False},
+ )
+ previous = {
+ key: app.config.get(key) for key in ("MCP_TOOL_WORKERS", "SQLLAB_TIMEOUT")
+ }
+ app.config["MCP_TOOL_WORKERS"] = WORKERS
+ app.config["SQLLAB_TIMEOUT"] = 5
+ install_mcp_session_scoping()
+ db.session.remove()
+ try:
+ with (
+ patch.object(
+ type(db),
+ "engines",
+ new_callable=PropertyMock,
+ return_value={None: engine},
+ ),
+ patch.object(
+ type(db), "engine", new_callable=PropertyMock, return_value=engine
+ ),
+ ):
+ yield engine
+ finally:
+ db.session.remove()
+ engine.dispose()
+ app.config.update(previous)
+
+
+@pytest.mark.asyncio
+async def test_saturated_workers_holding_metadata_pool_do_not_freeze_loop(
+ metadata_engine: Engine,
+) -> None:
+ """Every tool worker waits on the loop while every connection is checked out.
+
+ If the transport loop itself waits for a metadata connection (tools/list
+ RBAC filtering, audit logging), neither side can progress until a pool or
+ call timeout expires, and the whole server stops responding meanwhile.
+ """
+ holding = threading.Barrier(WORKERS + 1)
+ notify = threading.Event()
+
+ async def hold_and_notify(ctx: Context) -> str:
+ """Keep a metadata transaction open while reporting progress."""
+ db.session.execute(text("SELECT 1"))
+ holding.wait(timeout=5)
+ notify.wait(timeout=5)
+ await ctx.info("progress")
+ return "done"
+
+ def user_lookup() -> None:
+ """Resolve the caller for tools/list with a metadata query."""
+ db.session.execute(text("SELECT 1"))
+
+ mcp = FastMCP("metadata pool regression", middleware=build_middleware_list())
+ with (
+ patch("superset.mcp_service.auth._setup_user_context", return_value=None),
+ patch(
+ "superset.mcp_service.middleware.get_user_from_request",
+ side_effect=user_lookup,
+ ),
+ ):
+ mcp.tool(mcp_auth_hook(hold_and_notify), name="hold_and_notify")
+ async with Client(mcp) as client:
+ calls = [
+ asyncio.create_task(client.call_tool("hold_and_notify", {}))
+ for _ in range(WORKERS)
+ ]
+ await asyncio.to_thread(holding.wait, 5)
+ # Other holders, such as cancellations, own the rest of the pool.
+ held = [metadata_engine.connect() for _ in range(POOL_SIZE - WORKERS)]
+ gaps = []
+
+ async def heartbeat() -> None:
+ """Measure how long the transport loop goes unscheduled."""
+ last = time.monotonic()
+ while True:
+ await asyncio.sleep(0.01)
+ now = time.monotonic()
+ gaps.append(now - last)
+ last = now
+
+ beat = asyncio.create_task(heartbeat())
+ started = time.monotonic()
+ # tools/list needs metadata connections for its audit record and
+ # RBAC filtering. Release the workers once it is waiting for one,
+ # so their progress reports reach the loop while it waits.
+ listing = asyncio.create_task(client.list_tools())
+ threading.Timer(0.2, notify.set).start()
+ try:
+ await asyncio.sleep(0.05)
+ pinged = time.monotonic()
+ await client.ping()
+ ping_latency = time.monotonic() - pinged
+ results = await asyncio.gather(*calls)
+ elapsed = time.monotonic() - started
+ finally:
+ for connection in held:
+ connection.close()
+ beat.cancel()
+ tools = await listing
+
+ assert [result.content[0].text for result in results] == ["done"] * WORKERS
+ assert [tool.name for tool in tools] == ["hold_and_notify"]
+ assert ping_latency < 0.5
+ assert elapsed < 1.5
+ assert max(gaps) < 0.5
+
+
+@pytest.mark.parametrize(
+ ("pool_options", "configured", "expected"),
+ [
+ # SQLAlchemy's default QueuePool lends 5 + 10 connections.
+ ({}, None, 7),
+ ({}, 16, 7),
+ ({}, 4, 4),
+ ({"pool_size": 20, "max_overflow": 20}, None, 16),
+ ({"pool_size": 20, "max_overflow": 20}, 19, 19),
+ ({"max_overflow": -1}, 32, 32),
+ ({"poolclass": NullPool}, None, 16),
+ ],
+)
+def test_tool_workers_leave_metadata_connections_for_cancellation(
+ app: Any,
+ tmp_path: Path,
+ pool_options: dict[str, Any],
+ configured: int | None,
+ expected: int,
+) -> None:
+ """Each call and its cancellation get a connection, with one to spare."""
+ from superset.mcp_service.worker import tool_worker_count
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'metadata.db'}", **pool_options)
+ previous = app.config.get("MCP_TOOL_WORKERS")
+ app.config["MCP_TOOL_WORKERS"] = configured
+ try:
+ with patch.object(
+ type(db), "engine", new_callable=PropertyMock, return_value=engine
+ ):
+ assert tool_worker_count(app) == expected
+ finally:
+ app.config["MCP_TOOL_WORKERS"] = previous
+ engine.dispose()
+
+
+def test_metadata_pool_too_small_for_cancellation_is_rejected(
+ app: Any, tmp_path: Path
+) -> None:
+ """Refuse a pool where a call and its cancellation could take every slot."""
+ from superset.mcp_service.worker import tool_worker_count
+
+ engine = create_engine(
+ f"sqlite:///{tmp_path / 'metadata.db'}", pool_size=2, max_overflow=0
+ )
+ with (
+ patch.object(
+ type(db), "engine", new_callable=PropertyMock, return_value=engine
+ ),
+ pytest.raises(ValueError, match="needs at least 3"),
+ ):
+ tool_worker_count(app)
+ engine.dispose()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("strategy", ["bm25", "regex"])
+@pytest.mark.parametrize("request_backed", [False, True])
+async def test_threaded_search_returns_metadata_connection(
+ app: Any, metadata_engine: Engine, strategy: str, request_backed: bool
+) -> None:
+ """An inherited Flask context must not bypass the search thread's cleanup."""
+ from unittest.mock import AsyncMock, Mock
+
+ from fastmcp.server.transforms.search.base import BaseSearchTransform
+ from flask import g
+
+ from superset.mcp_service.server import _create_search_transform
+ from superset.mcp_service.session_scope import mcp_session_scopefunc
+
+ sessions: dict[Any, Session] = {}
+
+ def permission(tool: Any) -> bool:
+ """Perform the metadata lookup used by database-backed RBAC."""
+ sessions[mcp_session_scopefunc()] = db.session()
+ db.session.execute(text("SELECT 1"))
+ return True
+
+ transform = _create_search_transform(
+ strategy=strategy, kwargs={}, make_normalizing_call_tool=lambda _: None
+ )
+ tool = Mock()
+ try:
+ with app.test_request_context("/mcp/") if request_backed else app.app_context():
+ g.user = Mock()
+ with (
+ patch.object(
+ BaseSearchTransform,
+ "_get_visible_tools",
+ new_callable=AsyncMock,
+ return_value=[tool],
+ ),
+ patch(
+ "superset.mcp_service.auth.is_tool_visible_to_current_user",
+ side_effect=permission,
+ ),
+ ):
+ assert await transform._get_visible_tools(Mock()) == [tool]
+ assert metadata_engine.pool.checkedout() == 0
+ assert metadata_engine.pool.checkedout() == 0
+ assert all(key not in db.session.registry.registry for key in sessions)
+ finally:
+ # Also clean up when run against the leaking implementation.
+ for key, session in sessions.items():
+ session.close()
+ db.session.registry.registry.pop(key, None)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("audit_path", ["error", "response_size", "error_hook"])
+async def test_error_audit_waiting_for_metadata_does_not_block_loop(
+ app: Any, metadata_engine: Engine, audit_path: str
+) -> None:
+ """A timeout/busy error must not synchronously wait for an audit connection."""
+ from unittest.mock import AsyncMock, Mock
+
+ from fastmcp.exceptions import ToolError
+
+ from superset.mcp_service.middleware import (
+ GlobalErrorHandlerMiddleware,
+ ResponseSizeGuardMiddleware,
+ ToolResultCompatibilityMiddleware,
+ )
+ from superset.models.core import Log
+ from superset.utils.log import DBEventLogger
+
+ Log.__table__.create(metadata_engine)
+ held = [metadata_engine.connect() for _ in range(POOL_SIZE)]
+ heartbeat = asyncio.Event()
+
+ async def beat() -> None:
+ """Prove the transport can run while the audit checkout is blocked."""
+ await asyncio.sleep(0.01)
+ heartbeat.set()
+
+ task = asyncio.create_task(beat())
+ release = threading.Timer(0.3, held[-1].close)
+ release.start()
+ try:
+ context = Mock(method="tools/call")
+ context.message.name = "unknown_tool"
+ context.message.arguments = {}
+ event_logger = DBEventLogger()
+
+ def error_hook(error: Exception, info: dict[str, Any]) -> None:
+ """An operator hook may itself perform metadata I/O."""
+ event_logger.log(
+ user_id=None,
+ action="test_error_hook",
+ dashboard_id=None,
+ duration_ms=None,
+ slice_id=None,
+ referrer=None,
+ curated_payload={"tool": info["tool_name"]},
+ )
+
+ with (
+ app.test_request_context("/mcp/"),
+ patch("superset.mcp_service.middleware.event_logger", event_logger),
+ patch(
+ "superset.mcp_service.flask_singleton.get_flask_app", return_value=app
+ ),
+ patch.dict(app.config, {"MCP_ERROR_HOOK": error_hook}),
+ ):
+ if audit_path == "error":
+ with pytest.raises(ToolError, match="server busy"):
+ await GlobalErrorHandlerMiddleware()._handle_error(
+ ToolError("MCP server busy"), context, "execute_sql", 0
+ )
+ elif audit_path == "response_size":
+ with pytest.raises(ToolError):
+ await ResponseSizeGuardMiddleware(max_bytes=10).on_call_tool(
+ context, AsyncMock(return_value={"data": "x" * 100})
+ )
+ else:
+ result = await ToolResultCompatibilityMiddleware().on_call_tool(
+ context, AsyncMock(side_effect=RuntimeError("failed"))
+ )
+ assert result.is_error
+ assert heartbeat.is_set()
+ assert metadata_engine.pool.checkedout() == POOL_SIZE - 1
+ finally:
+ release.join()
+ for connection in held:
+ connection.close()
+ await task
+ assert metadata_engine.pool.checkedout() == 0
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("fail", [False, True])
+async def test_metadata_thread_owns_user_context_and_session(
+ app: Any, metadata_engine: Engine, fail: bool
+) -> None:
+ """Do not share even an expired request user or an inherited session token."""
+ from contextlib import nullcontext
+
+ from flask import g, request
+ from flask_appbuilder.security.sqla.models import User
+ from sqlalchemy import event
+
+ from superset.mcp_service.session_scope import _mcp_session_token
+ from superset.mcp_service.worker import (
+ get_context_user_id,
+ run_in_metadata_thread,
+ )
+
+ User.__table__.create(metadata_engine)
+ with metadata_engine.begin() as connection:
+ connection.execute(
+ User.__table__.insert(),
+ {
+ "id": 42,
+ "username": "metadata-user",
+ "email": "metadata@example.com",
+ "first_name": "Meta",
+ "last_name": "Data",
+ },
+ )
+ token = _mcp_session_token.set(object())
+ threads: list[int] = []
+ sessions: list[Session] = []
+ main_thread = threading.get_ident()
+
+ def record_thread(*args: Any) -> None:
+ """Track all SQL, including implicit ORM refreshes."""
+ threads.append(threading.get_ident())
+
+ def lookup() -> None:
+ """Read a reloaded user while owning fresh request globals and session."""
+ assert g.user is not user
+ assert g.user.username == "metadata-user"
+ assert request.path == "/mcp/"
+ assert g.routing_marker == "caller"
+ g.routing_marker = "metadata"
+ session = db.session()
+ sessions.append(session)
+ assert session is not parent_session
+ session.execute(text("SELECT 1"))
+ if fail:
+ raise ValueError("metadata failed")
+
+ try:
+ with app.test_request_context("/mcp/"):
+ parent_session = db.session()
+ user = parent_session.get(User, 42)
+ parent_session.commit() # expire attributes and return the connection
+ g.user = user
+ g.routing_marker = "caller"
+ event.listen(metadata_engine, "before_cursor_execute", record_thread)
+ try:
+ assert get_context_user_id() == 42
+ with (
+ pytest.raises(ValueError, match="metadata failed")
+ if fail
+ else nullcontext()
+ ):
+ await run_in_metadata_thread(lookup)
+ assert g.user is user
+ assert g.routing_marker == "caller"
+ assert db.session() is parent_session
+ assert threads
+ assert main_thread not in threads
+ assert metadata_engine.pool.checkedout() == 0
+ assert all(
+ session not in db.session.registry.registry.values()
+ for session in sessions
+ )
+ finally:
+ event.remove(metadata_engine, "before_cursor_execute", record_thread)
+ finally:
+ db.session.remove()
+ _mcp_session_token.reset(token)
+
+
+@pytest.mark.asyncio
+async def test_cancelled_metadata_awaiter_leaves_cleanup_to_thread(
+ metadata_engine: Engine,
+) -> None:
+ """A disconnect must neither tear down live I/O nor leak its connection."""
+ from sqlalchemy import event
+
+ from superset.mcp_service.worker import run_in_metadata_thread
+
+ entered = threading.Event()
+ release = threading.Event()
+ returned = threading.Event()
+
+ def checkin(*args: Any) -> None:
+ """Signal teardown after the thread's metadata I/O has completed."""
+ returned.set()
+
+ def lookup() -> None:
+ """Keep using the same live session after its awaiter disconnects."""
+ db.session.execute(text("SELECT 1"))
+ entered.set()
+ release.wait(5)
+ db.session.execute(text("SELECT 2"))
+
+ event.listen(metadata_engine, "checkin", checkin)
+ task = asyncio.create_task(run_in_metadata_thread(lookup))
+ try:
+ assert await asyncio.to_thread(entered.wait, 2)
+ task.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await task
+ assert metadata_engine.pool.checkedout() == 1
+ assert not returned.is_set()
+ finally:
+ release.set()
+ assert await asyncio.to_thread(returned.wait, 2)
+ event.remove(metadata_engine, "checkin", checkin)
+ assert metadata_engine.pool.checkedout() == 0
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("configured", [False, True])
+async def test_error_hook_does_not_require_metadata_user_reload(
+ app: Any, metadata_engine: Engine, configured: bool
+) -> None:
+ """Error reporting survives metadata outages without sharing caller state."""
+ from unittest.mock import Mock
+
+ from flask import g
+ from flask_appbuilder.security.sqla.models import User
+ from sqlalchemy.orm import make_transient_to_detached
+
+ from superset.mcp_service.middleware import _invoke_error_hook_off_loop
+
+ main_thread = threading.get_ident()
+ error = RuntimeError("tool failed")
+ info = {"tool_name": "execute_sql", "user_id": 42}
+ sessions: list[Session] = []
+
+ with app.test_request_context("/mcp/"):
+ user = User(id=42)
+ make_transient_to_detached(user)
+ g.user = user
+ parent_session = db.session()
+
+ def capture(exc: Exception, context: dict[str, Any]) -> None:
+ """Capture errors without loading a user from the unavailable DB."""
+ assert threading.get_ident() != main_thread
+ assert exc is error
+ assert context == info
+ assert getattr(g, "user", None) is None
+ assert db.session() is not parent_session
+ sessions.append(db.session())
Review Comment:
<!-- Bito Reply -->
The suggestion to record observations in the `capture` function and assert them after awaiting `_invoke_error_hook_off_loop` is correct and necessary. As noted, `_invoke_error_hook` catches all exceptions within the hook, meaning any assertions placed directly inside the mock function would be swallowed and fail to trigger a test failure, rendering the test ineffective at detecting regressions. Moving the assertions outside ensures they are properly evaluated by the test runner.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on code in PR #44581:
URL: https://github.com/apache/superset/pull/44581#discussion_r4150911741
##########
superset/sql/execution/cancellation.py:
##########
@@ -0,0 +1,170 @@
+# 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.
+"""Optional execution-owner hooks around a live warehouse cursor.
+
+Ordinary web and Celery execution are unchanged. An in-process execution owner
+can register cancellation while keeping its async runtime out of model code.
+"""
+
+from __future__ import annotations
+
+from contextlib import AbstractContextManager, contextmanager
+from contextvars import ContextVar
+from typing import Any, Callable, Iterator, TYPE_CHECKING
+
+from sqlalchemy import event
+from sqlalchemy.engine import Engine
+
+if TYPE_CHECKING:
+ from superset.models.core import Database
+
+cursor_scope: ContextVar[
+ Callable[[Database, Any, str | None, str | None], AbstractContextManager[None]]
+ | None
+] = ContextVar("warehouse_cursor_scope", default=None)
+
+
+@contextmanager
+def cancellable_cursor(
+ database: Database,
+ cursor: Any,
+ catalog: str | None = None,
+ schema: str | None = None,
+) -> Iterator[None]:
+ """Let the execution owner capture cancellation before blocking DBAPI I/O."""
+ scope = cursor_scope.get()
+ if scope is None:
+ yield
+ else:
+ with scope(database, cursor, catalog, schema):
+ yield
+
+
+before_warehouse_access: ContextVar[Callable[[], None] | None] = ContextVar(
+ "warehouse_before_access", default=None
+)
+check_deadline: ContextVar[Callable[[], None] | None] = ContextVar(
+ "warehouse_check_deadline", default=None
+)
+after_execute: ContextVar[Callable[[], None] | None] = ContextVar(
+ "warehouse_after_execute", default=None
+)
+
+
+def check_query_deadline() -> None:
+ """Stop before metadata access or another statement after warehouse I/O."""
+ if check := check_deadline.get():
+ check()
+
+
+def query_executed() -> None:
+ """Refresh cancellation handles exposed only after driver execution."""
+ if refresh := after_execute.get():
+ refresh()
+
+
+_engine_scope: ContextVar[tuple[Database, Engine, str | None, str | None] | None] = (
+ ContextVar("warehouse_engine_scope", default=None)
+)
+
+
+@contextmanager
+def without_execution_hooks() -> Iterator[None]:
+ """Keep cancellation I/O independent of the query it is cancelling.
+
+ Preserve unrelated context, such as authentication and tenant routing, while
+ preventing prequeries from checking the abandoned deadline or registering
+ cancellation recursively.
+ """
+ cursor_token = cursor_scope.set(None)
+ access_token = before_warehouse_access.set(None)
+ deadline_token = check_deadline.set(None)
+ execute_token = after_execute.set(None)
+ engine_token = _engine_scope.set(None)
+ try:
+ yield
+ finally:
+ _engine_scope.reset(engine_token)
+ after_execute.reset(execute_token)
+ check_deadline.reset(deadline_token)
+ before_warehouse_access.reset(access_token)
+ cursor_scope.reset(cursor_token)
+
+
+@contextmanager
+def cancellable_engine(
+ database: Database, engine: Engine, catalog: str | None, schema: str | None
+) -> Iterator[None]:
+ """Cover SQLAlchemy warehouse statements, including metadata discovery.
+
+ Listeners are installed once, never added/removed on shared cached engines.
+ Matching the engine excludes Superset metadata queries in the same context.
+ """
+ if cursor_scope.get() is None:
+ yield
+ return
+ if admit := before_warehouse_access.get():
+ # Before any connection is opened, including prequeries on connect.
+ admit()
Review Comment:
Added your test as `test_engine_admits_before_opening_a_connection` in `tests/unit_tests/sql/execution/test_cancellation.py` (91a3a66b0e). Locally, `pytest --cov=superset/sql/ ./tests/unit_tests/sql/ --cov-fail-under=100` now gives `cancellation.py` 100% and `TOTAL ... 100%` ("Required test coverage of 100% reached. Total coverage: 100.00%"), with 2,012 passed.
##########
superset/mcp_service/worker.py:
##########
@@ -0,0 +1,700 @@
+# 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.
+"""Bounded, deadline-aware execution of MCP tools outside the transport loop.
+
+The tool's coroutine runs on a worker-owned loop because its database APIs are
+synchronous. Only transport notifications are marshalled back to the server
+loop. Flask and SQLAlchemy lifetimes belong to the worker, not the waiting
+request: a timed-out DBAPI call cannot outlive and reuse a torn-down session.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import logging
+import threading
+import time
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import contextmanager
+from contextvars import ContextVar, copy_context
+from typing import Any, Callable, Coroutine, Iterator, ParamSpec, TYPE_CHECKING, TypeVar
+from weakref import WeakKeyDictionary
+
+from fastmcp.exceptions import ToolError
+from flask import current_app, g, has_app_context, has_request_context
+from sqlalchemy import inspect as sa_inspect
+from sqlalchemy.orm.state import InstanceState
+from sqlalchemy.pool import QueuePool
+
+from superset.mcp_service.session_scope import _mcp_session_token
+
+if TYPE_CHECKING:
+ from flask import Flask
+
+ from superset.models.core import Database
+
+logger = logging.getLogger(__name__)
+_active_call: ContextVar[WorkerCall | None] = ContextVar(
+ "mcp_worker_call", default=None
+)
+_metadata_context_owned: ContextVar[bool] = ContextVar(
+ "mcp_metadata_context_owned", default=False
+)
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
+_pools_lock = threading.Lock()
+_pools: WeakKeyDictionary[Flask, WorkerPool] = WeakKeyDictionary()
+
+
+class WorkerDeadlineExceeded(BaseException):
+ """Stop abandoned work without being swallowed by tool error handlers."""
+
+
+BUSY_MESSAGE = "MCP server busy: all tool workers are occupied. Retry later."
+
+# Tools that only read Superset's metadata database. They are admitted under a
+# separate, larger bound so they keep answering while slow warehouse queries
+# hold every warehouse slot. Anything that can reach a warehouse, a semantic
+# layer or a screenshot service stays in the warehouse bound. A listed tool
+# that does reach a warehouse must still take a warehouse slot first (see
+# WorkerCall.admit_warehouse), so a wrong entry cannot overdraw the pool.
+METADATA_ONLY_TOOLS = frozenset(
+ {
+ "find_users",
+ "get_annotation_layer_info",
+ "get_chart_info",
+ "get_chart_type_schema",
+ "get_dashboard_info",
+ "get_database_info",
+ "get_dataset_info",
+ "get_instance_info",
+ "get_layer_annotation_info",
+ "get_query_info",
+ "get_report_info",
+ "get_rls_filter_info",
+ "get_role_info",
+ "get_saved_query_info",
+ "get_schema",
+ "get_tag_info",
+ "get_task_info",
+ "get_theme_info",
+ "get_user_info",
+ "health_check",
+ "list_annotation_layers",
+ "list_charts",
+ "list_dashboards",
+ "list_databases",
+ "list_datasets",
+ "list_layer_annotations",
+ "list_queries",
+ "list_reports",
+ "list_rls_filters",
+ "list_roles",
+ "list_saved_queries",
+ "list_tags",
+ "list_tasks",
+ "list_themes",
+ "list_users",
+ }
+)
+
+
+class WorkerPool:
+ """Bound submissions, including abandoned work; never queue behind a query."""
+
+ def __init__(self, size: int, metadata_size: int = 0) -> None:
+ if size < 1:
+ raise ValueError("MCP_TOOL_WORKERS must be positive")
+ if metadata_size < 0:
+ raise ValueError("MCP_METADATA_TOOL_WORKERS must not be negative")
+ self.slots = threading.BoundedSemaphore(size)
+ self.metadata_slots = (
+ threading.BoundedSemaphore(metadata_size) if metadata_size else None
+ )
+ self.executor = ThreadPoolExecutor(
+ size + metadata_size, thread_name_prefix="mcp-tool"
+ )
+ # Cancellation must not wait behind the warehouse work it is cancelling.
+ # Only calls holding a warehouse slot register cancellation.
+ self.cancel_slots = threading.BoundedSemaphore(size)
+ self.cancellations = ThreadPoolExecutor(size, thread_name_prefix="mcp-cancel")
+
+ def admission(self, metadata_only: bool) -> threading.BoundedSemaphore:
+ """Choose the bound a call is admitted under."""
+ if metadata_only and self.metadata_slots is not None:
+ return self.metadata_slots
+ return self.slots
+
+ def submit(
+ self,
+ fn: Callable[[], Any],
+ finished: Callable[[], None],
+ slots: threading.BoundedSemaphore | None = None,
+ ) -> Future[Any]:
+ """Admit immediately or report overload without retaining a queued call."""
+ slots = self.slots if slots is None else slots
+ if not slots.acquire(blocking=False):
+ raise ToolError(BUSY_MESSAGE)
+ try:
+ future = self.executor.submit(fn)
+ except BaseException:
+ slots.release()
+ raise
+ future.add_done_callback(lambda _: finished())
+ return future
+
+ def cancel(self, fn: Callable[[], None], finished: Callable[[], None]) -> bool:
+ """Bound cancellation I/O separately so it cannot delay the caller."""
+ if not self.cancel_slots.acquire(blocking=False):
+ return False
+ try:
+ future = self.cancellations.submit(fn)
+ except RuntimeError:
+ self.cancel_slots.release()
+ return False
+
+ def completed(_: Future[None]) -> None:
+ """Release cancellation capacity before admitting another tool call."""
+ self.cancel_slots.release()
+ finished()
+
+ future.add_done_callback(completed)
+ return True
+
+
+DEFAULT_TOOL_WORKERS = 16
+DEFAULT_METADATA_TOOL_WORKERS = 16
+
+
+def _metadata_pool_capacity(app: Flask) -> int | None:
+ """Return how many metadata connections can be checked out at once.
+
+ ``None`` means a checkout never waits for another holder to return one
+ (e.g. ``NullPool``, per-thread pools, or unlimited overflow).
+ """
+ from superset import db
+
+ with app.app_context():
+ pool = db.engine.pool
+ if not isinstance(pool, QueuePool):
+ return None
+ # SQLAlchemy has no public accessor for the configured overflow limit.
+ max_overflow = pool._max_overflow # pylint: disable=protected-access
+ return None if max_overflow < 0 else pool.size() + max_overflow
+
+
+def tool_worker_count(app: Flask) -> int:
+ """Admit only as many calls as the metadata pool can always serve.
+
+ An admitted call can hold one metadata connection for the whole of its
+ warehouse I/O, and its cancellation needs another. ``2 * workers + 1``
+ connections therefore always leave one that is only held by short metadata
+ lookups, so cancellation and transport-side lookups (tools/list filtering,
+ audit logging) never wait for a warehouse query to end on its own.
+ """
+ configured = app.config.get("MCP_TOOL_WORKERS")
+ capacity = _metadata_pool_capacity(app)
+ if capacity is None:
+ return DEFAULT_TOOL_WORKERS if configured is None else configured
+ limit = (capacity - 1) // 2
+ if limit < 1:
+ raise ValueError(
+ f"The metadata database pool allows {capacity} connections; "
+ "MCP tool execution needs at least 3"
+ )
+ if configured is None:
+ return min(DEFAULT_TOOL_WORKERS, limit)
+ if configured > limit:
+ logger.warning(
+ "MCP_TOOL_WORKERS=%s needs %s metadata database connections, but the "
+ "pool allows %s; admitting %s concurrent tool calls. Raise the pool's "
+ "pool_size/max_overflow in SQLALCHEMY_ENGINE_OPTIONS to admit more.",
+ configured,
+ 2 * configured + 1,
+ capacity,
+ limit,
+ )
+ return limit
+ return configured
+
+
+def metadata_tool_worker_count(app: Flask) -> int:
+ """Admit metadata-only calls independently of warehouse capacity.
+
+ They hold a metadata connection only for their own short metadata queries,
+ never across warehouse I/O, so they cannot keep the connections reserved
+ for warehouse calls and their cancellation from cycling. They only wait
+ briefly for a connection, on a worker thread rather than the event loop.
+ """
+ configured = app.config.get("MCP_METADATA_TOOL_WORKERS")
+ count = DEFAULT_METADATA_TOOL_WORKERS if configured is None else configured
Review Comment:
Clamped in c9bc7895ae. Both bounds now come out of one budget, checked when the pool is created at startup (`admission_counts`):
`2 * MCP_TOOL_WORKERS + MCP_METADATA_TOOL_WORKERS + 1 <= pool_size + max_overflow`
- The `+ 1` is now an actual bound. Transport-side metadata I/O used the unbounded default executor; it now runs on one dedicated thread. That covers tools/list filtering, audit writes, error hooks, and `CompositeTokenVerifier`'s API-key lookup.
- With the default 5 + 10 pool the defaults are 4 warehouse-capable and 6 metadata-only calls, down from 7 and 16, so the default warehouse concurrency drops from 7 to 4. Explicit values above the budget are reduced with a startup warning. Docs and `mcp_config.py` describe the budget.
- `test_saturated_default_pool_never_waits_for_a_connection` runs on a real 5 + 10 `QueuePool` with `pool_timeout=1`. It holds all 4 warehouse slots, 4 cancellation connections, all 6 metadata-only slots and the transport thread at once (15 checked out). Every holder gets its connection and both bounds refuse the next call. Re-run with one metadata slot over the budget, it fails with `BrokenBarrierError`.
- `test_tool_workers_leave_metadata_connections_for_cancellation` checks the split and the inequality across 11 pool and config combinations.
- d8f971db50 fixes a side effect I hit while doing this. Creating the pool lazily from transport-side code popped an app context, and that removed the caller's scoped session. Covered by `test_lazy_pool_creation_keeps_callers_session`, which fails without the fix.
##########
superset/mcp_service/worker.py:
##########
@@ -0,0 +1,588 @@
+# 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.
+"""Bounded, deadline-aware execution of MCP tools outside the transport loop.
+
+The tool's coroutine runs on a worker-owned loop because its database APIs are
+synchronous. Only transport notifications are marshalled back to the server
+loop. Flask and SQLAlchemy lifetimes belong to the worker, not the waiting
+request: a timed-out DBAPI call cannot outlive and reuse a torn-down session.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import logging
+import threading
+import time
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import contextmanager
+from contextvars import ContextVar, copy_context
+from typing import Any, Callable, Coroutine, Iterator, ParamSpec, TYPE_CHECKING, TypeVar
+from weakref import WeakKeyDictionary
+
+from fastmcp.exceptions import ToolError
+from flask import current_app, g, has_app_context, has_request_context
+from sqlalchemy import inspect as sa_inspect
+from sqlalchemy.orm.state import InstanceState
+from sqlalchemy.pool import QueuePool
+
+from superset.mcp_service.session_scope import _mcp_session_token
+
+if TYPE_CHECKING:
+ from flask import Flask
+
+ from superset.models.core import Database
+
+logger = logging.getLogger(__name__)
+_active_call: ContextVar[WorkerCall | None] = ContextVar(
+ "mcp_worker_call", default=None
+)
+_metadata_context_owned: ContextVar[bool] = ContextVar(
+ "mcp_metadata_context_owned", default=False
+)
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
+_pools_lock = threading.Lock()
+_pools: WeakKeyDictionary[Flask, WorkerPool] = WeakKeyDictionary()
+
+
+class WorkerDeadlineExceeded(BaseException):
+ """Stop abandoned work without being swallowed by tool error handlers."""
+
+
+class WorkerPool:
+ """Bound submissions, including abandoned work; never queue behind a query."""
+
+ def __init__(self, size: int) -> None:
+ if size < 1:
+ raise ValueError("MCP_TOOL_WORKERS must be positive")
+ self.slots = threading.BoundedSemaphore(size)
+ self.executor = ThreadPoolExecutor(size, thread_name_prefix="mcp-tool")
+ # Cancellation must not wait behind the warehouse work it is cancelling.
+ self.cancel_slots = threading.BoundedSemaphore(size)
+ self.cancellations = ThreadPoolExecutor(size, thread_name_prefix="mcp-cancel")
+
+ def submit(
+ self, fn: Callable[[], Any], finished: Callable[[], None]
+ ) -> Future[Any]:
+ """Admit immediately or report overload without retaining a queued call."""
+ if not self.slots.acquire(blocking=False):
+ raise ToolError(
+ "MCP server busy: all tool workers are occupied. Retry later."
+ )
+ try:
+ future = self.executor.submit(fn)
+ except BaseException:
+ self.slots.release()
+ raise
+ future.add_done_callback(lambda _: finished())
+ return future
+
+ def cancel(self, fn: Callable[[], None], finished: Callable[[], None]) -> bool:
+ """Bound cancellation I/O separately so it cannot delay the caller."""
+ if not self.cancel_slots.acquire(blocking=False):
+ return False
+ try:
+ future = self.cancellations.submit(fn)
+ except RuntimeError:
+ self.cancel_slots.release()
+ return False
+
+ def completed(_: Future[None]) -> None:
+ """Release cancellation capacity before admitting another tool call."""
+ self.cancel_slots.release()
+ finished()
+
+ future.add_done_callback(completed)
+ return True
+
+
+DEFAULT_TOOL_WORKERS = 16
+
+
+def _metadata_pool_capacity(app: Flask) -> int | None:
+ """Return how many metadata connections can be checked out at once.
+
+ ``None`` means a checkout never waits for another holder to return one
+ (e.g. ``NullPool``, per-thread pools, or unlimited overflow).
+ """
+ from superset import db
+
+ with app.app_context():
+ pool = db.engine.pool
+ if not isinstance(pool, QueuePool):
+ return None
+ # SQLAlchemy has no public accessor for the configured overflow limit.
+ max_overflow = pool._max_overflow # pylint: disable=protected-access
+ return None if max_overflow < 0 else pool.size() + max_overflow
+
+
+def tool_worker_count(app: Flask) -> int:
+ """Admit only as many calls as the metadata pool can always serve.
+
+ An admitted call can hold one metadata connection for the whole of its
+ warehouse I/O, and its cancellation needs another. ``2 * workers + 1``
+ connections therefore always leave one that is only held by short metadata
+ lookups, so cancellation and transport-side lookups (tools/list filtering,
+ audit logging) never wait for a warehouse query to end on its own.
+ """
+ configured = app.config.get("MCP_TOOL_WORKERS")
+ capacity = _metadata_pool_capacity(app)
+ if capacity is None:
+ return DEFAULT_TOOL_WORKERS if configured is None else configured
+ limit = (capacity - 1) // 2
+ if limit < 1:
+ raise ValueError(
+ f"The metadata database pool allows {capacity} connections; "
+ "MCP tool execution needs at least 3"
+ )
+ if configured is None:
+ return min(DEFAULT_TOOL_WORKERS, limit)
+ if configured > limit:
+ logger.warning(
+ "MCP_TOOL_WORKERS=%s needs %s metadata database connections, but the "
+ "pool allows %s; admitting %s concurrent tool calls. Raise the pool's "
+ "pool_size/max_overflow in SQLALCHEMY_ENGINE_OPTIONS to admit more.",
+ configured,
+ 2 * configured + 1,
+ capacity,
+ limit,
+ )
+ return limit
+ return configured
+
+
+def _get_pool(app: Flask) -> WorkerPool:
+ """Lazily create a pool for this application, without import-time threads."""
+ with _pools_lock:
+ if app not in _pools:
+ size = tool_worker_count(app)
+ logger.info("MCP tool calls admitted concurrently: %s", size)
+ _pools[app] = WorkerPool(size)
+ return _pools[app]
+
+
+class WorkerCall:
+ """Thread-safe deadline and cancellation registration for one tool call."""
+
+ def __init__(self, app: Flask, pool: WorkerPool, seconds: float) -> None:
+ self.app = app
+ self.pool = pool
+ self.loop = asyncio.get_running_loop()
+ self.seconds = seconds
+ self.deadline = time.monotonic() + seconds
+ from superset.mcp_service.middleware import _mcp_call_id_var
+
+ self.call_id = _mcp_call_id_var.get() or uuid.uuid4().hex
+ self.expired = threading.Event()
+ self.lock = threading.RLock()
+ self.pending = 1
+ self.cancel_query: Callable[[], None] | None = None
+ self.cancel_dispatched = False
+ self.user_id: int | None = None
+
+ def check(self) -> None:
+ """Prevent an abandoned tool from starting more work or mutations."""
+ if self.expired.is_set() or time.monotonic() >= self.deadline:
+ # BaseException deliberately bypasses tools' broad Exception handlers.
+ raise WorkerDeadlineExceeded()
+
+ def dispatch_cancel(self) -> bool:
+ """Dispatch at most once for the active cursor, including late handles."""
+ with self.lock:
+ if self.cancel_query is None or self.cancel_dispatched:
+ return self.cancel_dispatched
+ self.pending += 1
+ self.cancel_dispatched = self.pool.cancel(self.cancel_query, self.finished)
+ if not self.cancel_dispatched:
+ self.pending -= 1
+ return self.cancel_dispatched
+
+ def finished(self) -> None:
+ """Retain admission until both query and cancellation I/O have ended.
+
+ A stuck cancellation therefore cannot consume the cancellation capacity
+ needed by newly admitted queries: it retains its original tool slot.
+ """
+ with self.lock:
+ self.pending -= 1
+ if self.pending == 0:
+ self.pool.slots.release()
+
+ def abandon(self) -> None:
+ """Signal abandonment and dispatch cancellation without blocking asyncio."""
+ self.expired.set()
+ dispatched = self.dispatch_cancel()
+ logger.warning(
+ "MCP call %s exceeded its deadline or disconnected; "
+ "cancellation dispatched=%s; worker slot retained until completion",
+ self.call_id,
+ dispatched,
+ )
+
+
+class TransportContext:
+ """Keep FastMCP's transport-bound async methods on their owning event loop."""
+
+ def __init__(
+ self, context: Any, loop: asyncio.AbstractEventLoop, call: WorkerCall
+ ) -> None:
+ self.context = context
+ self.loop = loop
+ self.call = call
+
+ def __getattr__(self, name: str) -> Any:
+ value = getattr(self.context, name)
+ if not asyncio.iscoroutinefunction(value):
+ return value
+
+ @functools.wraps(value)
+ async def forward(*args: Any, **kwargs: Any) -> Any:
+ self.call.check()
+ future = asyncio.run_coroutine_threadsafe(value(*args, **kwargs), self.loop)
+ try:
+ result = await asyncio.wait_for(
+ asyncio.wrap_future(future),
+ max(0, self.call.deadline - time.monotonic()),
+ )
+ except TimeoutError:
+ raise WorkerDeadlineExceeded() from None
+ self.call.check()
+ return result
+
+ return forward
+
+
+@contextmanager
+def _worker_context(app: Flask) -> Iterator[None]:
+ """Give a worker an independent app context and scoped metadata session."""
+ from superset.mcp_service.auth import _remove_session_safe
+
+ token = _mcp_session_token.set(object())
+ try:
+ with app.app_context():
+ try:
+ # Clean this scope before loading any ORM user into it.
+ _remove_session_safe()
+ yield
+ finally:
+ _remove_session_safe()
+ finally:
+ _mcp_session_token.reset(token)
+
+
+def get_context_user_id() -> int | None:
+ """Read the caller's identity without refreshing an expired ORM instance."""
+ if not has_app_context():
+ return None
+ user = getattr(g, "user", None)
+ state = sa_inspect(user, raiseerr=False)
+ if isinstance(state, InstanceState):
+ return state.identity[0] if state.identity else None
+ return getattr(user, "id", None)
+
+
+async def run_in_metadata_thread(
+ fn: Callable[_P, _T], *args: _P.args, **kwargs: _P.kwargs
+) -> _T:
+ """Run transport metadata I/O with thread-owned Flask and session lifetimes.
+
+ ``to_thread`` copies contextvars, including Flask contexts and MCP session
+ tokens. Replace those owners, reload ORM users, and retain only request and
+ routing data. Cleanup belongs to the thread even if its awaiter disconnects.
+ This uses the metadata executor rather than admitted tool workers, which
+ may themselves be waiting for transport notifications.
+ """
+ from contextlib import nullcontext
+
+ from flask.globals import _cv_request
+
+ if has_app_context():
+ app = current_app._get_current_object()
+ snapshot = dict(vars(g._get_current_object()))
+ else:
+ from superset.mcp_service.flask_singleton import get_flask_app
+
+ app = get_flask_app()
+ snapshot = {}
+ user = snapshot.pop("user", None)
+ state = sa_inspect(user, raiseerr=False)
+ is_orm_user = isinstance(state, InstanceState)
+ user_id = get_context_user_id() if is_orm_user else None
+ request_context = _cv_request.get(None)
+ request_copy = request_context.copy() if request_context is not None else None
+
+ def execute() -> _T:
+ """Own teardown rather than handing a live session back to asyncio."""
+ from superset import db, security_manager
+ from superset.sql.execution.cancellation import without_execution_hooks
+
+ active_token = _active_call.set(None)
+ owner_token = _metadata_context_owned.set(True)
+ try:
+ with without_execution_hooks(), _worker_context(app):
+ vars(g._get_current_object()).update(snapshot)
+ with request_copy if request_copy is not None else nullcontext():
+ if is_orm_user:
+ g.user = (
+ db.session.get(security_manager.user_model, user_id)
+ if user_id is not None
+ else None
+ )
+ elif user is not None:
+ g.user = user
+ return fn(*args, **kwargs)
+ finally:
+ _metadata_context_owned.reset(owner_token)
+ _active_call.reset(active_token)
+
+ return await asyncio.to_thread(execute)
+
+
+async def run_in_worker(
+ fn: Callable[..., Coroutine[Any, Any, Any]],
+ args: tuple[Any, ...],
+ kwargs: dict[str, Any],
+ seconds: float,
+) -> Any:
+ """Run a complete tool lifecycle in a bounded, independently owned worker."""
+ if active := _active_call.get():
+ # Composed tools share the outer deadline and worker/session ownership.
+ # A nested auth hook can inject the original transport Context again.
+ nested_kwargs = dict(kwargs)
+ if "ctx" in nested_kwargs and not isinstance(
+ nested_kwargs["ctx"], TransportContext
+ ):
+ nested_kwargs["ctx"] = TransportContext(
+ nested_kwargs["ctx"], active.loop, active
+ )
+ active.check()
+ return await fn(*args, **nested_kwargs)
+
+ if has_app_context():
+ app = current_app._get_current_object()
+ else:
+ from superset.mcp_service.flask_singleton import get_flask_app
+
+ app = get_flask_app()
+ pool = _get_pool(app)
+ call = WorkerCall(app, pool, seconds)
+ loop = asyncio.get_running_loop()
+ worker_kwargs = dict(kwargs)
+ if "ctx" in worker_kwargs:
+ worker_kwargs["ctx"] = TransportContext(worker_kwargs["ctx"], loop, call)
+ # Copy contextvars (token, tenant routing), but never share Flask g or a
+ # metadata Session. Request-backed middleware's user is reloaded by id.
+ from flask.globals import _cv_request
+
+ request_context = _cv_request.get(None)
+ request_copy = request_context.copy() if request_context is not None else None
+ globals_snapshot = dict(vars(g._get_current_object())) if has_app_context() else {}
+ user = globals_snapshot.pop("user", None) if has_request_context() else None
+ globals_snapshot.pop("user", None)
+ # An expired ORM user's .id can issue metadata I/O on the transport loop.
+ # The identity key is available without loading any attributes.
+ user_state = sa_inspect(user, raiseerr=False)
+ user_id = (
+ user_state.identity[0]
+ if isinstance(user_state, InstanceState) and user_state.identity
+ else getattr(user, "id", None)
+ )
+ guest = user if getattr(user, "is_guest_user", False) else None
+ context = copy_context()
+
+ def execute() -> Any:
+ """Own context teardown even when the caller stops waiting."""
+ from contextlib import nullcontext
+
+ from superset import db, security_manager
+ from superset.sql.execution.cancellation import cursor_scope
+
+ _active_call.set(call)
+ cursor_scope.set(warehouse_cursor)
+ with _worker_context(app):
+ vars(g._get_current_object()).update(globals_snapshot)
+ with request_copy if request_copy is not None else nullcontext():
+ if user_id is not None:
+ g.user = db.session.get(security_manager.user_model, user_id)
+ elif guest is not None:
+ g.user = guest
+ call.check()
+ result = asyncio.run(fn(*args, **worker_kwargs))
+ call.check()
+ return result
+
+ future = pool.submit(lambda: context.run(execute), call.finished)
+ 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
+ finally:
+ from superset.mcp_service.auth import _mcp_user_id_var
+
+ _mcp_user_id_var.set(call.user_id)
+
+
+class QueryCancellation:
+ """Bridge engine cancellation to a separate, freshly scoped metadata session."""
+
+ def __init__(
+ self,
+ call: WorkerCall,
+ database: Database,
+ cursor: Any,
+ catalog: str | None,
+ schema: str | None,
+ ) -> None:
+ self.call = call
+ self.database = database
+ self.database_id = database.id
+ self.spec = database.db_engine_spec
+ self.cursor = cursor
+ self.catalog = catalog
+ self.schema = schema
+ self.cancel_id: str | None = None
+ self.context = copy_context()
+
+ def capture(self) -> None:
+ """Capture an engine handle without making unsupported drivers fail."""
+ from superset.tasks.query_cancel import capture_cancel_query_id
+
+ try:
+ self.cancel_id = capture_cancel_query_id(self.database, self.cursor)
+ except Exception:
+ logger.warning(
+ "MCP call %s: cancel-id capture failed",
+ self.call.call_id,
+ exc_info=True,
+ )
+
+ def register(self) -> None:
+ """Publish a usable handle, including one obtained after the deadline."""
+ with self.call.lock:
+ if self.spec.has_implicit_cancel():
Review Comment:
A correction to my earlier answer: this is now implemented, in 2c1e351025.
- `TrinoEngineSpec.get_cancel_query_id` returns `cursor.query_id`, which is `None` until execution starts. SQL Lab (`sql_lab.py`), the Celery SQL task and the chart-data task seam all ask before executing, so they record nothing new, and SQL Lab's `handle_cursor` / `QUERY_CANCEL_KEY` path is unchanged.
- A new spec flag, `has_query_id_during_execute`, is True only for `TrinoEngineSpec`. With it, MCP registers cancellation without a handle. On abandonment the cancellation thread reads the queryId from the still-blocked cursor and calls the existing `cancel_chart_query` โ `TrinoEngineSpec.cancel_query` (`system.runtime.kill_query`).
- `test_abandoned_trino_query_is_killed_by_its_query_id` runs a real trino 0.339.0 DB-API cursor through `TrinoEngineSpec.execute` against a fake coordinator that keeps the query RUNNING. After the 0.3s deadline, the kill receives the coordinator's queryId on a thread other than the worker's. The worker's `execute` then ends with `ADMINISTRATIVELY_KILLED` for that queryId. Without the change the test fails because the worker stays blocked.
- Limits: the kill call is mocked in that test, so I have not run it against a real Trino. BigQuery, Databricks native, Athena and SQLite specs still provide no cancel id; for them abandonment still logs `dispatched=False`, and the call keeps its slot until the driver returns.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on PR #44581:
URL: https://github.com/apache/superset/pull/44581#issuecomment-5923233393
@rusackas the metadata sizing is fixed in c9bc7895ae: both admission bounds now come out of one budget, checked when the pool is created at startup: `2 * MCP_TOOL_WORKERS + MCP_METADATA_TOOL_WORKERS + 1 <= pool_size + max_overflow`. The `+ 1` is a real bound now (transport-side metadata I/O runs on one dedicated thread). Details in the `worker.py` thread.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "gabotorresruiz (via GitHub)" <gi...@apache.org>.
gabotorresruiz commented on code in PR #44581:
URL: https://github.com/apache/superset/pull/44581#discussion_r4150385770
##########
superset/mcp_service/worker.py:
##########
@@ -0,0 +1,700 @@
+# 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.
+"""Bounded, deadline-aware execution of MCP tools outside the transport loop.
+
+The tool's coroutine runs on a worker-owned loop because its database APIs are
+synchronous. Only transport notifications are marshalled back to the server
+loop. Flask and SQLAlchemy lifetimes belong to the worker, not the waiting
+request: a timed-out DBAPI call cannot outlive and reuse a torn-down session.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import logging
+import threading
+import time
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import contextmanager
+from contextvars import ContextVar, copy_context
+from typing import Any, Callable, Coroutine, Iterator, ParamSpec, TYPE_CHECKING, TypeVar
+from weakref import WeakKeyDictionary
+
+from fastmcp.exceptions import ToolError
+from flask import current_app, g, has_app_context, has_request_context
+from sqlalchemy import inspect as sa_inspect
+from sqlalchemy.orm.state import InstanceState
+from sqlalchemy.pool import QueuePool
+
+from superset.mcp_service.session_scope import _mcp_session_token
+
+if TYPE_CHECKING:
+ from flask import Flask
+
+ from superset.models.core import Database
+
+logger = logging.getLogger(__name__)
+_active_call: ContextVar[WorkerCall | None] = ContextVar(
+ "mcp_worker_call", default=None
+)
+_metadata_context_owned: ContextVar[bool] = ContextVar(
+ "mcp_metadata_context_owned", default=False
+)
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
+_pools_lock = threading.Lock()
+_pools: WeakKeyDictionary[Flask, WorkerPool] = WeakKeyDictionary()
+
+
+class WorkerDeadlineExceeded(BaseException):
+ """Stop abandoned work without being swallowed by tool error handlers."""
+
+
+BUSY_MESSAGE = "MCP server busy: all tool workers are occupied. Retry later."
+
+# Tools that only read Superset's metadata database. They are admitted under a
+# separate, larger bound so they keep answering while slow warehouse queries
+# hold every warehouse slot. Anything that can reach a warehouse, a semantic
+# layer or a screenshot service stays in the warehouse bound. A listed tool
+# that does reach a warehouse must still take a warehouse slot first (see
+# WorkerCall.admit_warehouse), so a wrong entry cannot overdraw the pool.
+METADATA_ONLY_TOOLS = frozenset(
+ {
+ "find_users",
+ "get_annotation_layer_info",
+ "get_chart_info",
+ "get_chart_type_schema",
+ "get_dashboard_info",
+ "get_database_info",
+ "get_dataset_info",
+ "get_instance_info",
+ "get_layer_annotation_info",
+ "get_query_info",
+ "get_report_info",
+ "get_rls_filter_info",
+ "get_role_info",
+ "get_saved_query_info",
+ "get_schema",
+ "get_tag_info",
+ "get_task_info",
+ "get_theme_info",
+ "get_user_info",
+ "health_check",
+ "list_annotation_layers",
+ "list_charts",
+ "list_dashboards",
+ "list_databases",
+ "list_datasets",
+ "list_layer_annotations",
+ "list_queries",
+ "list_reports",
+ "list_rls_filters",
+ "list_roles",
+ "list_saved_queries",
+ "list_tags",
+ "list_tasks",
+ "list_themes",
+ "list_users",
+ }
+)
+
+
+class WorkerPool:
+ """Bound submissions, including abandoned work; never queue behind a query."""
+
+ def __init__(self, size: int, metadata_size: int = 0) -> None:
+ if size < 1:
+ raise ValueError("MCP_TOOL_WORKERS must be positive")
+ if metadata_size < 0:
+ raise ValueError("MCP_METADATA_TOOL_WORKERS must not be negative")
+ self.slots = threading.BoundedSemaphore(size)
+ self.metadata_slots = (
+ threading.BoundedSemaphore(metadata_size) if metadata_size else None
+ )
+ self.executor = ThreadPoolExecutor(
+ size + metadata_size, thread_name_prefix="mcp-tool"
+ )
+ # Cancellation must not wait behind the warehouse work it is cancelling.
+ # Only calls holding a warehouse slot register cancellation.
+ self.cancel_slots = threading.BoundedSemaphore(size)
+ self.cancellations = ThreadPoolExecutor(size, thread_name_prefix="mcp-cancel")
+
+ def admission(self, metadata_only: bool) -> threading.BoundedSemaphore:
+ """Choose the bound a call is admitted under."""
+ if metadata_only and self.metadata_slots is not None:
+ return self.metadata_slots
+ return self.slots
+
+ def submit(
+ self,
+ fn: Callable[[], Any],
+ finished: Callable[[], None],
+ slots: threading.BoundedSemaphore | None = None,
+ ) -> Future[Any]:
+ """Admit immediately or report overload without retaining a queued call."""
+ slots = self.slots if slots is None else slots
+ if not slots.acquire(blocking=False):
+ raise ToolError(BUSY_MESSAGE)
+ try:
+ future = self.executor.submit(fn)
+ except BaseException:
+ slots.release()
+ raise
+ future.add_done_callback(lambda _: finished())
+ return future
+
+ def cancel(self, fn: Callable[[], None], finished: Callable[[], None]) -> bool:
+ """Bound cancellation I/O separately so it cannot delay the caller."""
+ if not self.cancel_slots.acquire(blocking=False):
+ return False
+ try:
+ future = self.cancellations.submit(fn)
+ except RuntimeError:
+ self.cancel_slots.release()
+ return False
+
+ def completed(_: Future[None]) -> None:
+ """Release cancellation capacity before admitting another tool call."""
+ self.cancel_slots.release()
+ finished()
+
+ future.add_done_callback(completed)
+ return True
+
+
+DEFAULT_TOOL_WORKERS = 16
+DEFAULT_METADATA_TOOL_WORKERS = 16
+
+
+def _metadata_pool_capacity(app: Flask) -> int | None:
+ """Return how many metadata connections can be checked out at once.
+
+ ``None`` means a checkout never waits for another holder to return one
+ (e.g. ``NullPool``, per-thread pools, or unlimited overflow).
+ """
+ from superset import db
+
+ with app.app_context():
+ pool = db.engine.pool
+ if not isinstance(pool, QueuePool):
+ return None
+ # SQLAlchemy has no public accessor for the configured overflow limit.
+ max_overflow = pool._max_overflow # pylint: disable=protected-access
+ return None if max_overflow < 0 else pool.size() + max_overflow
+
+
+def tool_worker_count(app: Flask) -> int:
+ """Admit only as many calls as the metadata pool can always serve.
+
+ An admitted call can hold one metadata connection for the whole of its
+ warehouse I/O, and its cancellation needs another. ``2 * workers + 1``
+ connections therefore always leave one that is only held by short metadata
+ lookups, so cancellation and transport-side lookups (tools/list filtering,
+ audit logging) never wait for a warehouse query to end on its own.
+ """
+ configured = app.config.get("MCP_TOOL_WORKERS")
+ capacity = _metadata_pool_capacity(app)
+ if capacity is None:
+ return DEFAULT_TOOL_WORKERS if configured is None else configured
+ limit = (capacity - 1) // 2
+ if limit < 1:
+ raise ValueError(
+ f"The metadata database pool allows {capacity} connections; "
+ "MCP tool execution needs at least 3"
+ )
+ if configured is None:
+ return min(DEFAULT_TOOL_WORKERS, limit)
+ if configured > limit:
+ logger.warning(
+ "MCP_TOOL_WORKERS=%s needs %s metadata database connections, but the "
+ "pool allows %s; admitting %s concurrent tool calls. Raise the pool's "
+ "pool_size/max_overflow in SQLALCHEMY_ENGINE_OPTIONS to admit more.",
+ configured,
+ 2 * configured + 1,
+ capacity,
+ limit,
+ )
+ return limit
+ return configured
+
+
+def metadata_tool_worker_count(app: Flask) -> int:
+ """Admit metadata-only calls independently of warehouse capacity.
+
+ They hold a metadata connection only for their own short metadata queries,
+ never across warehouse I/O, so they cannot keep the connections reserved
+ for warehouse calls and their cancellation from cycling. They only wait
+ briefly for a connection, on a worker thread rather than the event loop.
+ """
+ configured = app.config.get("MCP_METADATA_TOOL_WORKERS")
+ count = DEFAULT_METADATA_TOOL_WORKERS if configured is None else configured
Review Comment:
Not a blocker, and the docstring matches what I measured for the case this is designed for: with the real 5 plus 10 metadata pool and 7 warehouse calls each holding a connection across blocked warehouse I/O, 16 concurrent metadata-only calls finished in 0.07s.
The one thing I would raise is that this bound is never clamped against the pool, unlike its sibling: the default 16 is larger than the default capacity of 15. When all 15 are checked out (the 14 the admission math allows, plus one held by a transport-side audit write) all 16 metadata-only calls block inside `QueuePool` for the full `pool_timeout` and come back as `MCP tool timed out after 30 seconds`, each holding a worker thread for that long. Typed and graceful rather than a correctness problem, but `briefly` becomes the whole deadline in that window. Is clamping the default against capacity worth it here, or is the narrowness of the window the reason to leave it alone?
##########
superset/sql/execution/cancellation.py:
##########
@@ -0,0 +1,170 @@
+# 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.
+"""Optional execution-owner hooks around a live warehouse cursor.
+
+Ordinary web and Celery execution are unchanged. An in-process execution owner
+can register cancellation while keeping its async runtime out of model code.
+"""
+
+from __future__ import annotations
+
+from contextlib import AbstractContextManager, contextmanager
+from contextvars import ContextVar
+from typing import Any, Callable, Iterator, TYPE_CHECKING
+
+from sqlalchemy import event
+from sqlalchemy.engine import Engine
+
+if TYPE_CHECKING:
+ from superset.models.core import Database
+
+cursor_scope: ContextVar[
+ Callable[[Database, Any, str | None, str | None], AbstractContextManager[None]]
+ | None
+] = ContextVar("warehouse_cursor_scope", default=None)
+
+
+@contextmanager
+def cancellable_cursor(
+ database: Database,
+ cursor: Any,
+ catalog: str | None = None,
+ schema: str | None = None,
+) -> Iterator[None]:
+ """Let the execution owner capture cancellation before blocking DBAPI I/O."""
+ scope = cursor_scope.get()
+ if scope is None:
+ yield
+ else:
+ with scope(database, cursor, catalog, schema):
+ yield
+
+
+before_warehouse_access: ContextVar[Callable[[], None] | None] = ContextVar(
+ "warehouse_before_access", default=None
+)
+check_deadline: ContextVar[Callable[[], None] | None] = ContextVar(
+ "warehouse_check_deadline", default=None
+)
+after_execute: ContextVar[Callable[[], None] | None] = ContextVar(
+ "warehouse_after_execute", default=None
+)
+
+
+def check_query_deadline() -> None:
+ """Stop before metadata access or another statement after warehouse I/O."""
+ if check := check_deadline.get():
+ check()
+
+
+def query_executed() -> None:
+ """Refresh cancellation handles exposed only after driver execution."""
+ if refresh := after_execute.get():
+ refresh()
+
+
+_engine_scope: ContextVar[tuple[Database, Engine, str | None, str | None] | None] = (
+ ContextVar("warehouse_engine_scope", default=None)
+)
+
+
+@contextmanager
+def without_execution_hooks() -> Iterator[None]:
+ """Keep cancellation I/O independent of the query it is cancelling.
+
+ Preserve unrelated context, such as authentication and tenant routing, while
+ preventing prequeries from checking the abandoned deadline or registering
+ cancellation recursively.
+ """
+ cursor_token = cursor_scope.set(None)
+ access_token = before_warehouse_access.set(None)
+ deadline_token = check_deadline.set(None)
+ execute_token = after_execute.set(None)
+ engine_token = _engine_scope.set(None)
+ try:
+ yield
+ finally:
+ _engine_scope.reset(engine_token)
+ after_execute.reset(execute_token)
+ check_deadline.reset(deadline_token)
+ before_warehouse_access.reset(access_token)
+ cursor_scope.reset(cursor_token)
+
+
+@contextmanager
+def cancellable_engine(
+ database: Database, engine: Engine, catalog: str | None, schema: str | None
+) -> Iterator[None]:
+ """Cover SQLAlchemy warehouse statements, including metadata discovery.
+
+ Listeners are installed once, never added/removed on shared cached engines.
+ Matching the engine excludes Superset metadata queries in the same context.
+ """
+ if cursor_scope.get() is None:
+ yield
+ return
+ if admit := before_warehouse_access.get():
+ # Before any connection is opened, including prequeries on connect.
+ admit()
Review Comment:
This line is what has `unit-tests (current)` red. That job runs `pytest --cov=superset/sql/ ./tests/unit_tests/sql/ --cov-fail-under=100`, and `admit()` plus its branch are the only things missing:
```
superset/sql/execution/cancellation.py 69 1 16 1 98% 122
TOTAL 2151 1 710 1 99%
FAIL Required test coverage of 100% not reached. Total coverage: 99.93%
```
I reproduced the same command locally at this head and got the same numbers, and at the merge base `e90b9fb751` it is `100.00%`.
The hook is covered, just from a suite the gate does not measure: `test_metadata_only_call_takes_warehouse_slot_before_warehouse_io` exercises it from `tests/unit_tests/mcp_service`. `test_without_execution_hooks_restores_context` does enter `cancellable_engine` with `cursor_scope` set, but it never sets `before_warehouse_access`, so only the false branch is taken.
This closed it for me, back to `100.00%` with 2,012 passed:
```python
def test_engine_admits_before_opening_a_connection() -> None:
"""An execution owner is asked to admit warehouse access before connecting."""
admit = Mock()
engine = create_engine("sqlite://")
with ExitStack() as stack:
stack.callback(
cancellation.cursor_scope.reset, cancellation.cursor_scope.set(Mock())
)
stack.callback(
cancellation.before_warehouse_access.reset,
cancellation.before_warehouse_access.set(admit),
)
stack.callback(engine.dispose)
with cancellation.cancellable_engine(Mock(), engine, None, None):
admit.assert_called_once_with()
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "gabotorresruiz (via GitHub)" <gi...@apache.org>.
gabotorresruiz commented on code in PR #44581:
URL: https://github.com/apache/superset/pull/44581#discussion_r4149169235
##########
superset/mcp_service/worker.py:
##########
@@ -0,0 +1,588 @@
+# 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.
+"""Bounded, deadline-aware execution of MCP tools outside the transport loop.
+
+The tool's coroutine runs on a worker-owned loop because its database APIs are
+synchronous. Only transport notifications are marshalled back to the server
+loop. Flask and SQLAlchemy lifetimes belong to the worker, not the waiting
+request: a timed-out DBAPI call cannot outlive and reuse a torn-down session.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import logging
+import threading
+import time
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import contextmanager
+from contextvars import ContextVar, copy_context
+from typing import Any, Callable, Coroutine, Iterator, ParamSpec, TYPE_CHECKING, TypeVar
+from weakref import WeakKeyDictionary
+
+from fastmcp.exceptions import ToolError
+from flask import current_app, g, has_app_context, has_request_context
+from sqlalchemy import inspect as sa_inspect
+from sqlalchemy.orm.state import InstanceState
+from sqlalchemy.pool import QueuePool
+
+from superset.mcp_service.session_scope import _mcp_session_token
+
+if TYPE_CHECKING:
+ from flask import Flask
+
+ from superset.models.core import Database
+
+logger = logging.getLogger(__name__)
+_active_call: ContextVar[WorkerCall | None] = ContextVar(
+ "mcp_worker_call", default=None
+)
+_metadata_context_owned: ContextVar[bool] = ContextVar(
+ "mcp_metadata_context_owned", default=False
+)
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
+_pools_lock = threading.Lock()
+_pools: WeakKeyDictionary[Flask, WorkerPool] = WeakKeyDictionary()
+
+
+class WorkerDeadlineExceeded(BaseException):
+ """Stop abandoned work without being swallowed by tool error handlers."""
+
+
+class WorkerPool:
+ """Bound submissions, including abandoned work; never queue behind a query."""
+
+ def __init__(self, size: int) -> None:
+ if size < 1:
+ raise ValueError("MCP_TOOL_WORKERS must be positive")
+ self.slots = threading.BoundedSemaphore(size)
+ self.executor = ThreadPoolExecutor(size, thread_name_prefix="mcp-tool")
+ # Cancellation must not wait behind the warehouse work it is cancelling.
+ self.cancel_slots = threading.BoundedSemaphore(size)
+ self.cancellations = ThreadPoolExecutor(size, thread_name_prefix="mcp-cancel")
+
+ def submit(
+ self, fn: Callable[[], Any], finished: Callable[[], None]
+ ) -> Future[Any]:
+ """Admit immediately or report overload without retaining a queued call."""
+ if not self.slots.acquire(blocking=False):
+ raise ToolError(
+ "MCP server busy: all tool workers are occupied. Retry later."
+ )
+ try:
+ future = self.executor.submit(fn)
+ except BaseException:
+ self.slots.release()
+ raise
+ future.add_done_callback(lambda _: finished())
+ return future
+
+ def cancel(self, fn: Callable[[], None], finished: Callable[[], None]) -> bool:
+ """Bound cancellation I/O separately so it cannot delay the caller."""
+ if not self.cancel_slots.acquire(blocking=False):
+ return False
+ try:
+ future = self.cancellations.submit(fn)
+ except RuntimeError:
+ self.cancel_slots.release()
+ return False
+
+ def completed(_: Future[None]) -> None:
+ """Release cancellation capacity before admitting another tool call."""
+ self.cancel_slots.release()
+ finished()
+
+ future.add_done_callback(completed)
+ return True
+
+
+DEFAULT_TOOL_WORKERS = 16
+
+
+def _metadata_pool_capacity(app: Flask) -> int | None:
+ """Return how many metadata connections can be checked out at once.
+
+ ``None`` means a checkout never waits for another holder to return one
+ (e.g. ``NullPool``, per-thread pools, or unlimited overflow).
+ """
+ from superset import db
+
+ with app.app_context():
+ pool = db.engine.pool
+ if not isinstance(pool, QueuePool):
+ return None
+ # SQLAlchemy has no public accessor for the configured overflow limit.
+ max_overflow = pool._max_overflow # pylint: disable=protected-access
+ return None if max_overflow < 0 else pool.size() + max_overflow
+
+
+def tool_worker_count(app: Flask) -> int:
+ """Admit only as many calls as the metadata pool can always serve.
+
+ An admitted call can hold one metadata connection for the whole of its
+ warehouse I/O, and its cancellation needs another. ``2 * workers + 1``
+ connections therefore always leave one that is only held by short metadata
+ lookups, so cancellation and transport-side lookups (tools/list filtering,
+ audit logging) never wait for a warehouse query to end on its own.
+ """
+ configured = app.config.get("MCP_TOOL_WORKERS")
+ capacity = _metadata_pool_capacity(app)
+ if capacity is None:
+ return DEFAULT_TOOL_WORKERS if configured is None else configured
+ limit = (capacity - 1) // 2
Review Comment:
Not a blocker, but I would like to check the intent. With Superset's default metadata pool (`pool_size=5`, `max_overflow=10`) this resolves to 7, which I confirmed against a real `QueuePool`: capacity 15, `tool_worker_count` 7. Since every protected handler now goes through `run_in_worker`, that 7 is the ceiling for the whole process and not only for the tools that touch a warehouse. I ran 12 concurrent calls of a tool that does no warehouse I/O at all: on this branch 7 succeeded and 5 came back `MCP server busy`, while at the merge base all 12 succeeded. `tools/list` still answered in 8ms while saturated, so discovery is unaffected.
Was it a deliberate call to put metadata only tools such as `list_charts` behind the same admission bound as `execute_sql`? The connection reservation math only has to hold for handlers that can occupy a metadata connection across warehouse I/O, so a second and larger bound for the rest would keep the cheap tools answering while a few slow queries are in flight.
##########
superset/mcp_service/worker.py:
##########
@@ -0,0 +1,588 @@
+# 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.
+"""Bounded, deadline-aware execution of MCP tools outside the transport loop.
+
+The tool's coroutine runs on a worker-owned loop because its database APIs are
+synchronous. Only transport notifications are marshalled back to the server
+loop. Flask and SQLAlchemy lifetimes belong to the worker, not the waiting
+request: a timed-out DBAPI call cannot outlive and reuse a torn-down session.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import logging
+import threading
+import time
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import contextmanager
+from contextvars import ContextVar, copy_context
+from typing import Any, Callable, Coroutine, Iterator, ParamSpec, TYPE_CHECKING, TypeVar
+from weakref import WeakKeyDictionary
+
+from fastmcp.exceptions import ToolError
+from flask import current_app, g, has_app_context, has_request_context
+from sqlalchemy import inspect as sa_inspect
+from sqlalchemy.orm.state import InstanceState
+from sqlalchemy.pool import QueuePool
+
+from superset.mcp_service.session_scope import _mcp_session_token
+
+if TYPE_CHECKING:
+ from flask import Flask
+
+ from superset.models.core import Database
+
+logger = logging.getLogger(__name__)
+_active_call: ContextVar[WorkerCall | None] = ContextVar(
+ "mcp_worker_call", default=None
+)
+_metadata_context_owned: ContextVar[bool] = ContextVar(
+ "mcp_metadata_context_owned", default=False
+)
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
+_pools_lock = threading.Lock()
+_pools: WeakKeyDictionary[Flask, WorkerPool] = WeakKeyDictionary()
+
+
+class WorkerDeadlineExceeded(BaseException):
+ """Stop abandoned work without being swallowed by tool error handlers."""
+
+
+class WorkerPool:
+ """Bound submissions, including abandoned work; never queue behind a query."""
+
+ def __init__(self, size: int) -> None:
+ if size < 1:
+ raise ValueError("MCP_TOOL_WORKERS must be positive")
+ self.slots = threading.BoundedSemaphore(size)
+ self.executor = ThreadPoolExecutor(size, thread_name_prefix="mcp-tool")
+ # Cancellation must not wait behind the warehouse work it is cancelling.
+ self.cancel_slots = threading.BoundedSemaphore(size)
+ self.cancellations = ThreadPoolExecutor(size, thread_name_prefix="mcp-cancel")
+
+ def submit(
+ self, fn: Callable[[], Any], finished: Callable[[], None]
+ ) -> Future[Any]:
+ """Admit immediately or report overload without retaining a queued call."""
+ if not self.slots.acquire(blocking=False):
+ raise ToolError(
+ "MCP server busy: all tool workers are occupied. Retry later."
+ )
+ try:
+ future = self.executor.submit(fn)
+ except BaseException:
+ self.slots.release()
+ raise
+ future.add_done_callback(lambda _: finished())
+ return future
+
+ def cancel(self, fn: Callable[[], None], finished: Callable[[], None]) -> bool:
+ """Bound cancellation I/O separately so it cannot delay the caller."""
+ if not self.cancel_slots.acquire(blocking=False):
+ return False
+ try:
+ future = self.cancellations.submit(fn)
+ except RuntimeError:
+ self.cancel_slots.release()
+ return False
+
+ def completed(_: Future[None]) -> None:
+ """Release cancellation capacity before admitting another tool call."""
+ self.cancel_slots.release()
+ finished()
+
+ future.add_done_callback(completed)
+ return True
+
+
+DEFAULT_TOOL_WORKERS = 16
+
+
+def _metadata_pool_capacity(app: Flask) -> int | None:
+ """Return how many metadata connections can be checked out at once.
+
+ ``None`` means a checkout never waits for another holder to return one
+ (e.g. ``NullPool``, per-thread pools, or unlimited overflow).
+ """
+ from superset import db
+
+ with app.app_context():
+ pool = db.engine.pool
+ if not isinstance(pool, QueuePool):
+ return None
+ # SQLAlchemy has no public accessor for the configured overflow limit.
+ max_overflow = pool._max_overflow # pylint: disable=protected-access
+ return None if max_overflow < 0 else pool.size() + max_overflow
+
+
+def tool_worker_count(app: Flask) -> int:
+ """Admit only as many calls as the metadata pool can always serve.
+
+ An admitted call can hold one metadata connection for the whole of its
+ warehouse I/O, and its cancellation needs another. ``2 * workers + 1``
+ connections therefore always leave one that is only held by short metadata
+ lookups, so cancellation and transport-side lookups (tools/list filtering,
+ audit logging) never wait for a warehouse query to end on its own.
+ """
+ configured = app.config.get("MCP_TOOL_WORKERS")
+ capacity = _metadata_pool_capacity(app)
+ if capacity is None:
+ return DEFAULT_TOOL_WORKERS if configured is None else configured
+ limit = (capacity - 1) // 2
+ if limit < 1:
+ raise ValueError(
+ f"The metadata database pool allows {capacity} connections; "
+ "MCP tool execution needs at least 3"
+ )
+ if configured is None:
+ return min(DEFAULT_TOOL_WORKERS, limit)
+ if configured > limit:
+ logger.warning(
+ "MCP_TOOL_WORKERS=%s needs %s metadata database connections, but the "
+ "pool allows %s; admitting %s concurrent tool calls. Raise the pool's "
+ "pool_size/max_overflow in SQLALCHEMY_ENGINE_OPTIONS to admit more.",
+ configured,
+ 2 * configured + 1,
+ capacity,
+ limit,
+ )
+ return limit
+ return configured
+
+
+def _get_pool(app: Flask) -> WorkerPool:
+ """Lazily create a pool for this application, without import-time threads."""
+ with _pools_lock:
+ if app not in _pools:
+ size = tool_worker_count(app)
+ logger.info("MCP tool calls admitted concurrently: %s", size)
+ _pools[app] = WorkerPool(size)
+ return _pools[app]
+
+
+class WorkerCall:
+ """Thread-safe deadline and cancellation registration for one tool call."""
+
+ def __init__(self, app: Flask, pool: WorkerPool, seconds: float) -> None:
+ self.app = app
+ self.pool = pool
+ self.loop = asyncio.get_running_loop()
+ self.seconds = seconds
+ self.deadline = time.monotonic() + seconds
+ from superset.mcp_service.middleware import _mcp_call_id_var
+
+ self.call_id = _mcp_call_id_var.get() or uuid.uuid4().hex
+ self.expired = threading.Event()
+ self.lock = threading.RLock()
+ self.pending = 1
+ self.cancel_query: Callable[[], None] | None = None
+ self.cancel_dispatched = False
+ self.user_id: int | None = None
+
+ def check(self) -> None:
+ """Prevent an abandoned tool from starting more work or mutations."""
+ if self.expired.is_set() or time.monotonic() >= self.deadline:
+ # BaseException deliberately bypasses tools' broad Exception handlers.
+ raise WorkerDeadlineExceeded()
+
+ def dispatch_cancel(self) -> bool:
+ """Dispatch at most once for the active cursor, including late handles."""
+ with self.lock:
+ if self.cancel_query is None or self.cancel_dispatched:
+ return self.cancel_dispatched
+ self.pending += 1
+ self.cancel_dispatched = self.pool.cancel(self.cancel_query, self.finished)
+ if not self.cancel_dispatched:
+ self.pending -= 1
+ return self.cancel_dispatched
+
+ def finished(self) -> None:
+ """Retain admission until both query and cancellation I/O have ended.
+
+ A stuck cancellation therefore cannot consume the cancellation capacity
+ needed by newly admitted queries: it retains its original tool slot.
+ """
+ with self.lock:
+ self.pending -= 1
+ if self.pending == 0:
+ self.pool.slots.release()
+
+ def abandon(self) -> None:
+ """Signal abandonment and dispatch cancellation without blocking asyncio."""
+ self.expired.set()
+ dispatched = self.dispatch_cancel()
+ logger.warning(
+ "MCP call %s exceeded its deadline or disconnected; "
+ "cancellation dispatched=%s; worker slot retained until completion",
+ self.call_id,
+ dispatched,
+ )
+
+
+class TransportContext:
+ """Keep FastMCP's transport-bound async methods on their owning event loop."""
+
+ def __init__(
+ self, context: Any, loop: asyncio.AbstractEventLoop, call: WorkerCall
+ ) -> None:
+ self.context = context
+ self.loop = loop
+ self.call = call
+
+ def __getattr__(self, name: str) -> Any:
+ value = getattr(self.context, name)
+ if not asyncio.iscoroutinefunction(value):
+ return value
+
+ @functools.wraps(value)
+ async def forward(*args: Any, **kwargs: Any) -> Any:
+ self.call.check()
+ future = asyncio.run_coroutine_threadsafe(value(*args, **kwargs), self.loop)
+ try:
+ result = await asyncio.wait_for(
+ asyncio.wrap_future(future),
+ max(0, self.call.deadline - time.monotonic()),
+ )
+ except TimeoutError:
+ raise WorkerDeadlineExceeded() from None
+ self.call.check()
+ return result
+
+ return forward
+
+
+@contextmanager
+def _worker_context(app: Flask) -> Iterator[None]:
+ """Give a worker an independent app context and scoped metadata session."""
+ from superset.mcp_service.auth import _remove_session_safe
+
+ token = _mcp_session_token.set(object())
+ try:
+ with app.app_context():
+ try:
+ # Clean this scope before loading any ORM user into it.
+ _remove_session_safe()
+ yield
+ finally:
+ _remove_session_safe()
+ finally:
+ _mcp_session_token.reset(token)
+
+
+def get_context_user_id() -> int | None:
+ """Read the caller's identity without refreshing an expired ORM instance."""
+ if not has_app_context():
+ return None
+ user = getattr(g, "user", None)
+ state = sa_inspect(user, raiseerr=False)
+ if isinstance(state, InstanceState):
+ return state.identity[0] if state.identity else None
+ return getattr(user, "id", None)
+
+
+async def run_in_metadata_thread(
+ fn: Callable[_P, _T], *args: _P.args, **kwargs: _P.kwargs
+) -> _T:
+ """Run transport metadata I/O with thread-owned Flask and session lifetimes.
+
+ ``to_thread`` copies contextvars, including Flask contexts and MCP session
+ tokens. Replace those owners, reload ORM users, and retain only request and
+ routing data. Cleanup belongs to the thread even if its awaiter disconnects.
+ This uses the metadata executor rather than admitted tool workers, which
+ may themselves be waiting for transport notifications.
+ """
+ from contextlib import nullcontext
+
+ from flask.globals import _cv_request
+
+ if has_app_context():
+ app = current_app._get_current_object()
+ snapshot = dict(vars(g._get_current_object()))
+ else:
+ from superset.mcp_service.flask_singleton import get_flask_app
+
+ app = get_flask_app()
+ snapshot = {}
+ user = snapshot.pop("user", None)
+ state = sa_inspect(user, raiseerr=False)
+ is_orm_user = isinstance(state, InstanceState)
+ user_id = get_context_user_id() if is_orm_user else None
+ request_context = _cv_request.get(None)
+ request_copy = request_context.copy() if request_context is not None else None
+
+ def execute() -> _T:
+ """Own teardown rather than handing a live session back to asyncio."""
+ from superset import db, security_manager
+ from superset.sql.execution.cancellation import without_execution_hooks
+
+ active_token = _active_call.set(None)
+ owner_token = _metadata_context_owned.set(True)
+ try:
+ with without_execution_hooks(), _worker_context(app):
+ vars(g._get_current_object()).update(snapshot)
+ with request_copy if request_copy is not None else nullcontext():
+ if is_orm_user:
+ g.user = (
+ db.session.get(security_manager.user_model, user_id)
+ if user_id is not None
+ else None
+ )
+ elif user is not None:
+ g.user = user
+ return fn(*args, **kwargs)
+ finally:
+ _metadata_context_owned.reset(owner_token)
+ _active_call.reset(active_token)
+
+ return await asyncio.to_thread(execute)
+
+
+async def run_in_worker(
+ fn: Callable[..., Coroutine[Any, Any, Any]],
+ args: tuple[Any, ...],
+ kwargs: dict[str, Any],
+ seconds: float,
+) -> Any:
+ """Run a complete tool lifecycle in a bounded, independently owned worker."""
+ if active := _active_call.get():
+ # Composed tools share the outer deadline and worker/session ownership.
+ # A nested auth hook can inject the original transport Context again.
+ nested_kwargs = dict(kwargs)
+ if "ctx" in nested_kwargs and not isinstance(
+ nested_kwargs["ctx"], TransportContext
+ ):
+ nested_kwargs["ctx"] = TransportContext(
+ nested_kwargs["ctx"], active.loop, active
+ )
+ active.check()
+ return await fn(*args, **nested_kwargs)
+
+ if has_app_context():
+ app = current_app._get_current_object()
+ else:
+ from superset.mcp_service.flask_singleton import get_flask_app
+
+ app = get_flask_app()
+ pool = _get_pool(app)
+ call = WorkerCall(app, pool, seconds)
+ loop = asyncio.get_running_loop()
+ worker_kwargs = dict(kwargs)
+ if "ctx" in worker_kwargs:
+ worker_kwargs["ctx"] = TransportContext(worker_kwargs["ctx"], loop, call)
+ # Copy contextvars (token, tenant routing), but never share Flask g or a
+ # metadata Session. Request-backed middleware's user is reloaded by id.
+ from flask.globals import _cv_request
+
+ request_context = _cv_request.get(None)
+ request_copy = request_context.copy() if request_context is not None else None
+ globals_snapshot = dict(vars(g._get_current_object())) if has_app_context() else {}
+ user = globals_snapshot.pop("user", None) if has_request_context() else None
+ globals_snapshot.pop("user", None)
+ # An expired ORM user's .id can issue metadata I/O on the transport loop.
+ # The identity key is available without loading any attributes.
+ user_state = sa_inspect(user, raiseerr=False)
+ user_id = (
+ user_state.identity[0]
+ if isinstance(user_state, InstanceState) and user_state.identity
+ else getattr(user, "id", None)
+ )
+ guest = user if getattr(user, "is_guest_user", False) else None
+ context = copy_context()
+
+ def execute() -> Any:
+ """Own context teardown even when the caller stops waiting."""
+ from contextlib import nullcontext
+
+ from superset import db, security_manager
+ from superset.sql.execution.cancellation import cursor_scope
+
+ _active_call.set(call)
+ cursor_scope.set(warehouse_cursor)
+ with _worker_context(app):
+ vars(g._get_current_object()).update(globals_snapshot)
+ with request_copy if request_copy is not None else nullcontext():
+ if user_id is not None:
+ g.user = db.session.get(security_manager.user_model, user_id)
+ elif guest is not None:
+ g.user = guest
+ call.check()
+ result = asyncio.run(fn(*args, **worker_kwargs))
+ call.check()
+ return result
+
+ future = pool.submit(lambda: context.run(execute), call.finished)
+ 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
+ finally:
+ from superset.mcp_service.auth import _mcp_user_id_var
+
+ _mcp_user_id_var.set(call.user_id)
+
+
+class QueryCancellation:
+ """Bridge engine cancellation to a separate, freshly scoped metadata session."""
+
+ def __init__(
+ self,
+ call: WorkerCall,
+ database: Database,
+ cursor: Any,
+ catalog: str | None,
+ schema: str | None,
+ ) -> None:
+ self.call = call
+ self.database = database
+ self.database_id = database.id
+ self.spec = database.db_engine_spec
+ self.cursor = cursor
+ self.catalog = catalog
+ self.schema = schema
+ self.cancel_id: str | None = None
+ self.context = copy_context()
+
+ def capture(self) -> None:
+ """Capture an engine handle without making unsupported drivers fail."""
+ from superset.tasks.query_cancel import capture_cancel_query_id
+
+ try:
+ self.cancel_id = capture_cancel_query_id(self.database, self.cursor)
+ except Exception:
+ logger.warning(
+ "MCP call %s: cancel-id capture failed",
+ self.call.call_id,
+ exc_info=True,
+ )
+
+ def register(self) -> None:
+ """Publish a usable handle, including one obtained after the deadline."""
+ with self.call.lock:
+ if self.spec.has_implicit_cancel():
Review Comment:
Worth a sanity check on which engines actually reach this branch. Resolving the specs on this branch inside an app context, `TrinoEngineSpec`, `BigQueryEngineSpec`, `DatabricksNativeEngineSpec`, `AthenaEngineSpec` and `SqliteEngineSpec` all report `has_implicit_cancel() == False` and inherit `get_cancel_query_id` from `BaseEngineSpec`, which returns `None`; only seven modules under `superset/db_engine_specs/` override it at all. For those engines `register()` leaves `cancel_query` as `None`, `abandon()` logs `dispatched=False`, and the call keeps its admission slot for the full driver duration. Postgres does work end to end, I confirmed a timed out `pg_sleep(30)` leaves nothing in `pg_stat_activity`.
Trino is the interesting one, because it does support cancellation in SQL Lab, just through `handle_cursor` writing `QUERY_CANCEL_KEY` rather than through `get_cancel_query_id`. Is consulting that path out of scope here, or is the intent to leave it to the existing `query_cancel` seam? Together with the admission bound above, a handful of uncancellable slow queries takes the whole process to `MCP server busy` until the drivers return, so it would be good to know that is understood and accepted.
##########
tests/unit_tests/conftest.py:
##########
@@ -43,7 +44,11 @@ def get_session(mocker: MockerFixture) -> Callable[[], Session]:
"""
Create an in-memory SQLite db.session.to test models.
"""
- engine = create_engine("sqlite://")
+ # MCP unit tests hand this injected Session from setup to a tool worker
+ # sequentially. Concurrency/session ownership is tested with scoped sessions.
+ engine = create_engine(
+ "sqlite://", connect_args={"check_same_thread": False}, poolclass=StaticPool
Review Comment:
Just a small NIT, purely about the blast radius of this fixture. `check_same_thread=False` with a `StaticPool` means every unit test that takes `session` can now share one `Session` across threads without SQLAlchemy objecting, which is the exact thing the production code here works hard to avoid. Would scoping this to the MCP tests that need it, rather than the shared fixture, keep the rest of the suite able to catch that?
##########
superset/mcp_service/worker.py:
##########
@@ -0,0 +1,588 @@
+# 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.
+"""Bounded, deadline-aware execution of MCP tools outside the transport loop.
+
+The tool's coroutine runs on a worker-owned loop because its database APIs are
+synchronous. Only transport notifications are marshalled back to the server
+loop. Flask and SQLAlchemy lifetimes belong to the worker, not the waiting
+request: a timed-out DBAPI call cannot outlive and reuse a torn-down session.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import logging
+import threading
+import time
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import contextmanager
+from contextvars import ContextVar, copy_context
+from typing import Any, Callable, Coroutine, Iterator, ParamSpec, TYPE_CHECKING, TypeVar
+from weakref import WeakKeyDictionary
+
+from fastmcp.exceptions import ToolError
+from flask import current_app, g, has_app_context, has_request_context
+from sqlalchemy import inspect as sa_inspect
+from sqlalchemy.orm.state import InstanceState
+from sqlalchemy.pool import QueuePool
+
+from superset.mcp_service.session_scope import _mcp_session_token
+
+if TYPE_CHECKING:
+ from flask import Flask
+
+ from superset.models.core import Database
+
+logger = logging.getLogger(__name__)
+_active_call: ContextVar[WorkerCall | None] = ContextVar(
+ "mcp_worker_call", default=None
+)
+_metadata_context_owned: ContextVar[bool] = ContextVar(
+ "mcp_metadata_context_owned", default=False
+)
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
+_pools_lock = threading.Lock()
+_pools: WeakKeyDictionary[Flask, WorkerPool] = WeakKeyDictionary()
+
+
+class WorkerDeadlineExceeded(BaseException):
+ """Stop abandoned work without being swallowed by tool error handlers."""
+
+
+class WorkerPool:
+ """Bound submissions, including abandoned work; never queue behind a query."""
+
+ def __init__(self, size: int) -> None:
+ if size < 1:
+ raise ValueError("MCP_TOOL_WORKERS must be positive")
+ self.slots = threading.BoundedSemaphore(size)
+ self.executor = ThreadPoolExecutor(size, thread_name_prefix="mcp-tool")
+ # Cancellation must not wait behind the warehouse work it is cancelling.
+ self.cancel_slots = threading.BoundedSemaphore(size)
+ self.cancellations = ThreadPoolExecutor(size, thread_name_prefix="mcp-cancel")
+
+ def submit(
+ self, fn: Callable[[], Any], finished: Callable[[], None]
+ ) -> Future[Any]:
+ """Admit immediately or report overload without retaining a queued call."""
+ if not self.slots.acquire(blocking=False):
+ raise ToolError(
+ "MCP server busy: all tool workers are occupied. Retry later."
+ )
+ try:
+ future = self.executor.submit(fn)
+ except BaseException:
+ self.slots.release()
+ raise
+ future.add_done_callback(lambda _: finished())
+ return future
+
+ def cancel(self, fn: Callable[[], None], finished: Callable[[], None]) -> bool:
+ """Bound cancellation I/O separately so it cannot delay the caller."""
+ if not self.cancel_slots.acquire(blocking=False):
+ return False
+ try:
+ future = self.cancellations.submit(fn)
+ except RuntimeError:
+ self.cancel_slots.release()
+ return False
+
+ def completed(_: Future[None]) -> None:
+ """Release cancellation capacity before admitting another tool call."""
+ self.cancel_slots.release()
+ finished()
+
+ future.add_done_callback(completed)
+ return True
+
+
+DEFAULT_TOOL_WORKERS = 16
+
+
+def _metadata_pool_capacity(app: Flask) -> int | None:
+ """Return how many metadata connections can be checked out at once.
+
+ ``None`` means a checkout never waits for another holder to return one
+ (e.g. ``NullPool``, per-thread pools, or unlimited overflow).
+ """
+ from superset import db
+
+ with app.app_context():
+ pool = db.engine.pool
+ if not isinstance(pool, QueuePool):
+ return None
+ # SQLAlchemy has no public accessor for the configured overflow limit.
+ max_overflow = pool._max_overflow # pylint: disable=protected-access
+ return None if max_overflow < 0 else pool.size() + max_overflow
+
+
+def tool_worker_count(app: Flask) -> int:
+ """Admit only as many calls as the metadata pool can always serve.
+
+ An admitted call can hold one metadata connection for the whole of its
+ warehouse I/O, and its cancellation needs another. ``2 * workers + 1``
+ connections therefore always leave one that is only held by short metadata
+ lookups, so cancellation and transport-side lookups (tools/list filtering,
+ audit logging) never wait for a warehouse query to end on its own.
+ """
+ configured = app.config.get("MCP_TOOL_WORKERS")
+ capacity = _metadata_pool_capacity(app)
+ if capacity is None:
+ return DEFAULT_TOOL_WORKERS if configured is None else configured
+ limit = (capacity - 1) // 2
+ if limit < 1:
+ raise ValueError(
+ f"The metadata database pool allows {capacity} connections; "
+ "MCP tool execution needs at least 3"
+ )
+ if configured is None:
+ return min(DEFAULT_TOOL_WORKERS, limit)
+ if configured > limit:
+ logger.warning(
+ "MCP_TOOL_WORKERS=%s needs %s metadata database connections, but the "
+ "pool allows %s; admitting %s concurrent tool calls. Raise the pool's "
+ "pool_size/max_overflow in SQLALCHEMY_ENGINE_OPTIONS to admit more.",
+ configured,
+ 2 * configured + 1,
+ capacity,
+ limit,
+ )
+ return limit
+ return configured
+
+
+def _get_pool(app: Flask) -> WorkerPool:
+ """Lazily create a pool for this application, without import-time threads."""
+ with _pools_lock:
+ if app not in _pools:
+ size = tool_worker_count(app)
+ logger.info("MCP tool calls admitted concurrently: %s", size)
+ _pools[app] = WorkerPool(size)
+ return _pools[app]
+
+
+class WorkerCall:
+ """Thread-safe deadline and cancellation registration for one tool call."""
+
+ def __init__(self, app: Flask, pool: WorkerPool, seconds: float) -> None:
+ self.app = app
+ self.pool = pool
+ self.loop = asyncio.get_running_loop()
+ self.seconds = seconds
+ self.deadline = time.monotonic() + seconds
+ from superset.mcp_service.middleware import _mcp_call_id_var
+
+ self.call_id = _mcp_call_id_var.get() or uuid.uuid4().hex
+ self.expired = threading.Event()
+ self.lock = threading.RLock()
+ self.pending = 1
+ self.cancel_query: Callable[[], None] | None = None
+ self.cancel_dispatched = False
+ self.user_id: int | None = None
+
+ def check(self) -> None:
+ """Prevent an abandoned tool from starting more work or mutations."""
+ if self.expired.is_set() or time.monotonic() >= self.deadline:
+ # BaseException deliberately bypasses tools' broad Exception handlers.
+ raise WorkerDeadlineExceeded()
+
+ def dispatch_cancel(self) -> bool:
+ """Dispatch at most once for the active cursor, including late handles."""
+ with self.lock:
+ if self.cancel_query is None or self.cancel_dispatched:
+ return self.cancel_dispatched
+ self.pending += 1
+ self.cancel_dispatched = self.pool.cancel(self.cancel_query, self.finished)
+ if not self.cancel_dispatched:
+ self.pending -= 1
+ return self.cancel_dispatched
+
+ def finished(self) -> None:
+ """Retain admission until both query and cancellation I/O have ended.
+
+ A stuck cancellation therefore cannot consume the cancellation capacity
+ needed by newly admitted queries: it retains its original tool slot.
+ """
+ with self.lock:
+ self.pending -= 1
+ if self.pending == 0:
+ self.pool.slots.release()
+
+ def abandon(self) -> None:
+ """Signal abandonment and dispatch cancellation without blocking asyncio."""
+ self.expired.set()
+ dispatched = self.dispatch_cancel()
+ logger.warning(
+ "MCP call %s exceeded its deadline or disconnected; "
+ "cancellation dispatched=%s; worker slot retained until completion",
+ self.call_id,
+ dispatched,
+ )
+
+
+class TransportContext:
+ """Keep FastMCP's transport-bound async methods on their owning event loop."""
+
+ def __init__(
+ self, context: Any, loop: asyncio.AbstractEventLoop, call: WorkerCall
+ ) -> None:
+ self.context = context
+ self.loop = loop
+ self.call = call
+
+ def __getattr__(self, name: str) -> Any:
+ value = getattr(self.context, name)
+ if not asyncio.iscoroutinefunction(value):
+ return value
+
+ @functools.wraps(value)
+ async def forward(*args: Any, **kwargs: Any) -> Any:
+ self.call.check()
+ future = asyncio.run_coroutine_threadsafe(value(*args, **kwargs), self.loop)
+ try:
+ result = await asyncio.wait_for(
+ asyncio.wrap_future(future),
+ max(0, self.call.deadline - time.monotonic()),
+ )
+ except TimeoutError:
+ raise WorkerDeadlineExceeded() from None
+ self.call.check()
+ return result
+
+ return forward
+
+
+@contextmanager
+def _worker_context(app: Flask) -> Iterator[None]:
+ """Give a worker an independent app context and scoped metadata session."""
+ from superset.mcp_service.auth import _remove_session_safe
+
+ token = _mcp_session_token.set(object())
+ try:
+ with app.app_context():
+ try:
+ # Clean this scope before loading any ORM user into it.
+ _remove_session_safe()
+ yield
+ finally:
+ _remove_session_safe()
+ finally:
+ _mcp_session_token.reset(token)
+
+
+def get_context_user_id() -> int | None:
+ """Read the caller's identity without refreshing an expired ORM instance."""
+ if not has_app_context():
+ return None
+ user = getattr(g, "user", None)
+ state = sa_inspect(user, raiseerr=False)
+ if isinstance(state, InstanceState):
+ return state.identity[0] if state.identity else None
+ return getattr(user, "id", None)
+
+
+async def run_in_metadata_thread(
+ fn: Callable[_P, _T], *args: _P.args, **kwargs: _P.kwargs
+) -> _T:
+ """Run transport metadata I/O with thread-owned Flask and session lifetimes.
+
+ ``to_thread`` copies contextvars, including Flask contexts and MCP session
+ tokens. Replace those owners, reload ORM users, and retain only request and
+ routing data. Cleanup belongs to the thread even if its awaiter disconnects.
+ This uses the metadata executor rather than admitted tool workers, which
+ may themselves be waiting for transport notifications.
+ """
+ from contextlib import nullcontext
+
+ from flask.globals import _cv_request
+
+ if has_app_context():
+ app = current_app._get_current_object()
+ snapshot = dict(vars(g._get_current_object()))
+ else:
+ from superset.mcp_service.flask_singleton import get_flask_app
+
+ app = get_flask_app()
+ snapshot = {}
+ user = snapshot.pop("user", None)
+ state = sa_inspect(user, raiseerr=False)
+ is_orm_user = isinstance(state, InstanceState)
+ user_id = get_context_user_id() if is_orm_user else None
+ request_context = _cv_request.get(None)
+ request_copy = request_context.copy() if request_context is not None else None
+
+ def execute() -> _T:
+ """Own teardown rather than handing a live session back to asyncio."""
+ from superset import db, security_manager
+ from superset.sql.execution.cancellation import without_execution_hooks
+
+ active_token = _active_call.set(None)
+ owner_token = _metadata_context_owned.set(True)
+ try:
+ with without_execution_hooks(), _worker_context(app):
+ vars(g._get_current_object()).update(snapshot)
+ with request_copy if request_copy is not None else nullcontext():
+ if is_orm_user:
+ g.user = (
+ db.session.get(security_manager.user_model, user_id)
+ if user_id is not None
+ else None
+ )
+ elif user is not None:
+ g.user = user
+ return fn(*args, **kwargs)
+ finally:
+ _metadata_context_owned.reset(owner_token)
+ _active_call.reset(active_token)
+
+ return await asyncio.to_thread(execute)
+
+
+async def run_in_worker(
+ fn: Callable[..., Coroutine[Any, Any, Any]],
+ args: tuple[Any, ...],
+ kwargs: dict[str, Any],
+ seconds: float,
+) -> Any:
+ """Run a complete tool lifecycle in a bounded, independently owned worker."""
+ if active := _active_call.get():
+ # Composed tools share the outer deadline and worker/session ownership.
+ # A nested auth hook can inject the original transport Context again.
+ nested_kwargs = dict(kwargs)
+ if "ctx" in nested_kwargs and not isinstance(
+ nested_kwargs["ctx"], TransportContext
+ ):
+ nested_kwargs["ctx"] = TransportContext(
+ nested_kwargs["ctx"], active.loop, active
+ )
+ active.check()
+ return await fn(*args, **nested_kwargs)
+
+ if has_app_context():
+ app = current_app._get_current_object()
+ else:
+ from superset.mcp_service.flask_singleton import get_flask_app
+
+ app = get_flask_app()
+ pool = _get_pool(app)
+ call = WorkerCall(app, pool, seconds)
+ loop = asyncio.get_running_loop()
+ worker_kwargs = dict(kwargs)
+ if "ctx" in worker_kwargs:
+ worker_kwargs["ctx"] = TransportContext(worker_kwargs["ctx"], loop, call)
+ # Copy contextvars (token, tenant routing), but never share Flask g or a
+ # metadata Session. Request-backed middleware's user is reloaded by id.
+ from flask.globals import _cv_request
+
+ request_context = _cv_request.get(None)
+ request_copy = request_context.copy() if request_context is not None else None
+ globals_snapshot = dict(vars(g._get_current_object())) if has_app_context() else {}
+ user = globals_snapshot.pop("user", None) if has_request_context() else None
+ globals_snapshot.pop("user", None)
+ # An expired ORM user's .id can issue metadata I/O on the transport loop.
+ # The identity key is available without loading any attributes.
+ user_state = sa_inspect(user, raiseerr=False)
+ user_id = (
+ user_state.identity[0]
+ if isinstance(user_state, InstanceState) and user_state.identity
+ else getattr(user, "id", None)
+ )
+ guest = user if getattr(user, "is_guest_user", False) else None
+ context = copy_context()
+
+ def execute() -> Any:
+ """Own context teardown even when the caller stops waiting."""
+ from contextlib import nullcontext
+
+ from superset import db, security_manager
+ from superset.sql.execution.cancellation import cursor_scope
+
+ _active_call.set(call)
+ cursor_scope.set(warehouse_cursor)
+ with _worker_context(app):
+ vars(g._get_current_object()).update(globals_snapshot)
+ with request_copy if request_copy is not None else nullcontext():
+ if user_id is not None:
+ g.user = db.session.get(security_manager.user_model, user_id)
+ elif guest is not None:
+ g.user = guest
+ call.check()
+ result = asyncio.run(fn(*args, **worker_kwargs))
+ call.check()
+ return result
+
+ future = pool.submit(lambda: context.run(execute), call.finished)
+ 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
+ finally:
+ from superset.mcp_service.auth import _mcp_user_id_var
+
+ _mcp_user_id_var.set(call.user_id)
+
+
+class QueryCancellation:
+ """Bridge engine cancellation to a separate, freshly scoped metadata session."""
+
+ def __init__(
+ self,
+ call: WorkerCall,
+ database: Database,
+ cursor: Any,
+ catalog: str | None,
+ schema: str | None,
+ ) -> None:
+ self.call = call
+ self.database = database
+ self.database_id = database.id
+ self.spec = database.db_engine_spec
+ self.cursor = cursor
+ self.catalog = catalog
+ self.schema = schema
+ self.cancel_id: str | None = None
+ self.context = copy_context()
+
+ def capture(self) -> None:
+ """Capture an engine handle without making unsupported drivers fail."""
+ from superset.tasks.query_cancel import capture_cancel_query_id
+
+ try:
+ self.cancel_id = capture_cancel_query_id(self.database, self.cursor)
+ except Exception:
+ logger.warning(
+ "MCP call %s: cancel-id capture failed",
+ self.call.call_id,
+ exc_info=True,
+ )
+
+ def register(self) -> None:
+ """Publish a usable handle, including one obtained after the deadline."""
+ with self.call.lock:
+ if self.spec.has_implicit_cancel():
+ self.call.cancel_query = self.cancel_live_cursor
+ elif self.cancel_id is not None:
+ self.call.cancel_query = lambda: self.context.run(self.cancel)
+ else:
+ self.call.cancel_query = None
+ if self.call.expired.is_set():
+ self.call.dispatch_cancel()
+
+ def refresh(self) -> None:
+ """Capture post-execute handles before result fetching can block."""
+ if not self.spec.has_query_id_before_execute:
+ self.capture()
+ self.register()
+ self.call.check()
+
+ def cancel_live_cursor(self) -> None:
+ """Use the same implicit driver cancellation as engine handle_cursor."""
+ try:
+ self.cursor.cancel()
+ except Exception:
+ logger.warning(
+ "MCP call %s: cursor cancellation failed",
+ self.call.call_id,
+ exc_info=True,
+ )
+
+ def cancel(self) -> None:
+ """Resolve a fresh Database and acting user on the cancellation worker."""
+ from superset import db, security_manager
+ from superset.models.core import Database
+ from superset.sql.execution.cancellation import without_execution_hooks
+ from superset.tasks.query_cancel import cancel_chart_query
+
+ active_token = _active_call.set(None)
+ try:
+ with without_execution_hooks(), _worker_context(self.call.app):
+ if self.call.user_id is not None:
+ g.user = db.session.get(
+ security_manager.user_model, self.call.user_id
+ )
+ target = db.session.get(Database, self.database_id)
+ if target is not None and self.cancel_id is not None:
+ cancel_chart_query(
+ target, self.cancel_id, catalog=self.catalog, schema=self.schema
+ )
+ except Exception:
+ logger.warning(
+ "MCP call %s: cancellation failed", self.call.call_id, exc_info=True
+ )
+ finally:
+ _active_call.reset(active_token)
+
+
+@contextmanager
+def warehouse_cursor(
+ database: Database, cursor: Any, catalog: str | None, schema: str | None
+) -> Iterator[None]:
+ """Register Superset's engine cancellation hook for an MCP warehouse call."""
+ call = _active_call.get()
+ if call is None:
+ yield
+ return
+ from superset.sql.execution.cancellation import after_execute, check_deadline
+
+ call.check()
+ cancellation = QueryCancellation(call, database, cursor, catalog, schema)
+ with call.lock:
+ call.cancel_dispatched = False
+ cancellation.capture()
+ cancellation.register()
+ deadline_token = check_deadline.set(call.check)
+ execute_token = after_execute.set(cancellation.refresh)
+ try:
+ call.check()
+ yield
+ call.check()
+ except BaseException:
+ # Do not let error conversion, audit logging or cleanup reuse a metadata
+ # connection held across failed/abandoned warehouse I/O.
+ from superset import db
+
+ # Rollback listeners (including version history) may inspect the live
+ # connection. Finish the transaction before invalidation makes that
+ # inspection fail with PendingRollbackError.
+ db.session.rollback() # pylint: disable=consider-using-transaction
Review Comment:
This block worries me quite a bit: as written it turns every ordinary warehouse error inside an MCP tool into a `DetachedInstanceError`, and the real database error never reaches the caller.
`db.session.rollback()` expires every instance in the worker's session, and `db.session().invalidate()` plus `db.session.remove()` then detach them, all while the warehouse exception is still unwinding. The next frame out is `check_for_oauth2`, which `Database.get_sqla_engine` wraps around the entire engine scope. Its handler evaluates `database.is_oauth2_enabled()` (`superset/utils/oauth2.py:482`), which reaches `self.encrypted_extra` (`superset/models/core.py:1537`) on the instance that was just detached, so a `DetachedInstanceError` is raised out of that handler and replaces the original exception.
I verified it on this branch. Inside `run_in_worker`, `database.get_df("SELECT * FROM table_that_does_not_exist")` raises `DetachedInstanceError`, and `database.database_name`, `database.backend` and `database.db_engine_spec` all raise it afterwards as well. The same three lines at the merge base `5b9733fd`, and on this branch outside a worker, raise `OperationalError` and leave the row fully usable. Through the real `execute_sql` tool against a SQLite warehouse the client receives:
```
SQL execution failed: error=Instance <Database at 0x...> is not bound to a Session;
attribute refresh operation cannot proceed, database_id=3
```
instead of `no such table: no_such_table`. It is also misclassified on the way out, since `DetachedInstanceError` is a `SQLAlchemyError` and therefore not in `_USER_ERROR_TYPES`: `_is_user_error` returns `False`, so an agent typo is logged as a system error and fires `MCP_ERROR_HOOK`.
The deadline path itself is clean, because `WorkerDeadlineExceeded` is a `BaseException` and `check_for_oauth2` only catches `Exception`. What breaks is the everyday SQL mistake, which is the most common thing an agent does.
Two ways out that I think both work. Either gate this teardown on the case the docstring actually motivates, an abandoned or deadline expired call plus connection level failures, and let an ordinary warehouse `Exception` propagate with the session intact. Or move it out to `execute()` so it runs after the tool has finished handling its own error; `_worker_context` already calls `_remove_session_safe()` at the end of every worker call, so the next call gets a clean session either way.
On why the tests did not catch it: `test_failed_warehouse_discards_metadata_session` raises a bare `RuntimeError` straight into `cancellable_cursor`, so `check_for_oauth2` is never on the stack, and its assertions only cover that a *new* session works. A sibling such as `test_warehouse_error_reaches_caller_unchanged`, calling a real `Database.get_df` against a missing table and asserting both that the caller sees `OperationalError` and that a `Database` row loaded before the failure is still readable, would pin both halves.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on PR #44581:
URL: https://github.com/apache/superset/pull/44581#issuecomment-5853348725
AI Code Review is in progress (usually takes 3 to 15 minutes unless it's a very large PR).
<details>
<summary>
Bito Usage Guide
</summary>
**Commands**
Type the following command in the pull request comment and save the comment.
- `/review` - Manually triggers an incremental AI Review.
- `/review full` - Manually triggers a full AI Review.
- `/pause` - Pauses automatic reviews on this pull request.
- `/resume` - Resumes automatic reviews.
- `/resolve` - Marks all Bito-posted review comments as resolved.
- `/abort` - Cancels all in-progress reviews.
Refer to the <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/available-commands" target="_blank" rel="noopener noreferrer">documentation</a> for additional commands.
**Configuration**
This repository uses `Superset` You can customize the agent settings <a href="https://alpha.bito.ai/home/ai-agents/code-review-agent" target="_blank" rel="noopener noreferrer">here</a> or contact your Bito workspace admin at evan@preset.io.
**Documentation & Help**
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/getting-started/install-run-using-bito-cloud/create-or-customize-an-agent-instance" target="_blank" rel="noopener noreferrer">Customize agent settings</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/implementing-custom-code-review-rules" target="_blank" rel="noopener noreferrer">Review rules</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent" target="_blank" rel="noopener noreferrer">General documentation</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/faqs" target="_blank" rel="noopener noreferrer">FAQ</a>
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "netlify[bot] (via GitHub)" <gi...@apache.org>.
netlify[bot] commented on PR #44581:
URL: https://github.com/apache/superset/pull/44581#issuecomment-5800447646
### <span aria-hidden="true">โ
</span> Deploy Preview for *superset-docs-preview* ready!
| Name | Link |
|:-:|------------------------|
|<span aria-hidden="true">๐จ</span> Latest commit | 99e9c1722b5803708ede42130277ca5987c27881 |
|<span aria-hidden="true">๐</span> Latest deploy log | https://app.netlify.com/projects/superset-docs-preview/deploys/6ab417f4a8832400087cb54b |
|<span aria-hidden="true">๐</span> Deploy Preview | [https://deploy-preview-44581--superset-docs-preview.netlify.app](https://deploy-preview-44581--superset-docs-preview.netlify.app) |
|<span aria-hidden="true">๐ฑ</span> Preview on mobile | <details><summary> Toggle QR Code... </summary><br /><br /><br /><br />_Use your smartphone camera to open QR code link._</details> |
|<span aria-hidden="true">๐ค</span> Make changes | [Run an agent on this branch](https://app.netlify.com/projects/superset-docs-preview/agent-runs#Dcq7DYAgEADQXa72EkAU4gaOweeIJiKIWKhxd-le8V7IJcVcYQLowBazu6X5dMgFV0JjdBmPi8qNdY2UroqhED3Utqe8pXv27Y_GSq6CNFr3QjLGtHJ2kBa-Hw) |
---
<!-- [superset-docs-preview Preview](https://deploy-preview-44581--superset-docs-preview.netlify.app) -->
_To edit notification comments on pull requests, go to your [Netlify project configuration](https://app.netlify.com/projects/superset-docs-preview/configuration/notifications#deploy-notifications)._
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on PR #44581:
URL: https://github.com/apache/superset/pull/44581#issuecomment-5800500640
<!-- Bito Reply -->
The flagged issue is correct. Synchronous tools currently execute on the transport thread, which can block the MCP server if they perform long-running warehouse operations. To resolve this, you should wrap synchronous tool calls in `run_in_worker` to offload them to a worker thread, similar to how asynchronous tools are handled.
Would you like me to fetch all other comments on this PR to validate and implement fixes for them as well?
**superset/mcp_service/auth.py**
```
from superset.mcp_service.worker import run_in_worker
# Wrap synchronous tool execution:
# result = await run_in_worker(sync_tool_func, args, kwargs, timeout)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on PR #44581:
URL: https://github.com/apache/superset/pull/44581#issuecomment-5922328394
AI Code Review is in progress (usually takes 3 to 15 minutes unless it's a very large PR).
<details>
<summary>
Bito Usage Guide
</summary>
**Commands**
Type the following command in the pull request comment and save the comment.
- `/review` - Manually triggers an incremental AI Review.
- `/review full` - Manually triggers a full AI Review.
- `/pause` - Pauses automatic reviews on this pull request.
- `/resume` - Resumes automatic reviews.
- `/resolve` - Marks all Bito-posted review comments as resolved.
- `/abort` - Cancels all in-progress reviews.
Refer to the <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/available-commands" target="_blank" rel="noopener noreferrer">documentation</a> for additional commands.
**Configuration**
This repository uses `Superset` You can customize the agent settings <a href="https://alpha.bito.ai/home/ai-agents/code-review-agent" target="_blank" rel="noopener noreferrer">here</a> or contact your Bito workspace admin at evan@preset.io.
**Documentation & Help**
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/getting-started/install-run-using-bito-cloud/create-or-customize-an-agent-instance" target="_blank" rel="noopener noreferrer">Customize agent settings</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/implementing-custom-code-review-rules" target="_blank" rel="noopener noreferrer">Review rules</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent" target="_blank" rel="noopener noreferrer">General documentation</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/faqs" target="_blank" rel="noopener noreferrer">FAQ</a>
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] fix(mcp): enforce tool deadlines without blocking the server [superset]
Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on code in PR #44581:
URL: https://github.com/apache/superset/pull/44581#discussion_r4152808421
##########
superset/mcp_service/worker.py:
##########
@@ -0,0 +1,904 @@
+# 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.
+"""Bounded, deadline-aware execution of MCP tools outside the transport loop.
+
+The tool's coroutine runs on a worker-owned loop because its database APIs are
+synchronous. Only transport notifications are marshalled back to the server
+loop. Flask and SQLAlchemy lifetimes belong to the worker, not the waiting
+request: a timed-out DBAPI call cannot outlive and reuse a torn-down session.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import logging
+import threading
+import time
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import contextmanager
+from contextvars import ContextVar, copy_context
+from typing import Any, Callable, Coroutine, Iterator, ParamSpec, TYPE_CHECKING, TypeVar
+from weakref import WeakKeyDictionary
+
+from fastmcp.exceptions import ToolError
+from flask import current_app, g, has_app_context, has_request_context
+from sqlalchemy import inspect as sa_inspect
+from sqlalchemy.orm.state import InstanceState
+from sqlalchemy.pool import QueuePool
+
+from superset.mcp_service.session_scope import _mcp_session_token
+
+if TYPE_CHECKING:
+ from flask import Flask
+
+ from superset.models.core import Database
+
+logger = logging.getLogger(__name__)
+_active_call: ContextVar[WorkerCall | None] = ContextVar(
+ "mcp_worker_call", default=None
+)
+_metadata_context_owned: ContextVar[bool] = ContextVar(
+ "mcp_metadata_context_owned", default=False
+)
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
+_pools_lock = threading.Lock()
+_pools: WeakKeyDictionary[Flask, WorkerPool] = WeakKeyDictionary()
+
+
+class WorkerDeadlineExceeded(BaseException):
+ """Stop abandoned work without being swallowed by tool error handlers."""
+
+
+BUSY_MESSAGE = "MCP server busy: all tool workers are occupied. Retry later."
+
+# Tools that only read Superset's metadata database. They are admitted under a
+# separate, larger bound so they keep answering while slow warehouse queries
+# hold every warehouse slot. Anything that can reach a warehouse, a semantic
+# layer or a screenshot service stays in the warehouse bound. A listed tool
+# that does reach a warehouse must still take a warehouse slot first (see
+# WorkerCall.admit_warehouse), so a wrong entry cannot overdraw the pool.
+METADATA_ONLY_TOOLS = frozenset(
+ {
+ "find_users",
+ "get_annotation_layer_info",
+ "get_chart_info",
+ "get_chart_type_schema",
+ "get_dashboard_info",
+ "get_database_info",
+ "get_dataset_info",
+ "get_instance_info",
+ "get_layer_annotation_info",
+ "get_query_info",
+ "get_report_info",
+ "get_rls_filter_info",
+ "get_role_info",
+ "get_saved_query_info",
+ "get_schema",
+ "get_tag_info",
+ "get_task_info",
+ "get_theme_info",
+ "get_user_info",
+ "health_check",
+ "list_annotation_layers",
+ "list_charts",
+ "list_dashboards",
+ "list_databases",
+ "list_datasets",
+ "list_layer_annotations",
+ "list_queries",
+ "list_reports",
+ "list_rls_filters",
+ "list_roles",
+ "list_saved_queries",
+ "list_tags",
+ "list_tasks",
+ "list_themes",
+ "list_users",
+ }
+)
+
+
+class WorkerPool:
+ """Bound submissions, including abandoned work; never queue behind a query."""
+
+ def __init__(
+ self,
+ size: int,
+ metadata_size: int = 0,
+ transport_size: int = 1,
+ auth_size: int = 1,
+ auth_pending: int = 32,
+ ) -> None:
+ if size < 1:
+ raise ValueError("MCP_TOOL_WORKERS must be positive")
+ if metadata_size < 0:
+ raise ValueError("MCP_METADATA_TOOL_WORKERS must not be negative")
+ self.slots = threading.BoundedSemaphore(size)
+ self.metadata_slots = (
+ threading.BoundedSemaphore(metadata_size) if metadata_size else None
+ )
+ self.executor = ThreadPoolExecutor(
+ size + metadata_size, thread_name_prefix="mcp-tool"
+ )
+ # Cancellation must not wait behind the warehouse work it is cancelling.
+ # Only calls holding a warehouse slot register cancellation.
+ self.cancel_slots = threading.BoundedSemaphore(size)
+ self.cancellations = ThreadPoolExecutor(size, thread_name_prefix="mcp-cancel")
+ # Transport-side metadata I/O (tools/list filtering, audit writes,
+ # error hooks) is bounded too, so it cannot overdraw the connections
+ # budgeted for tool workers. Callers wait for it on the event loop
+ # without blocking it.
+ self.transport = ThreadPoolExecutor(
+ transport_size, thread_name_prefix="mcp-metadata"
+ )
+ # API-key lookups run before a caller is authenticated, so they get
+ # their own thread: a flood of bad keys or a slow lookup must not
+ # delay audit writes and RBAC filtering. Pending lookups are capped;
+ # past the cap a key is rejected rather than queued.
+ self.auth = ThreadPoolExecutor(auth_size, thread_name_prefix="mcp-auth")
+ self.auth_slots = threading.BoundedSemaphore(auth_pending)
+
+ def admission(self, metadata_only: bool) -> threading.BoundedSemaphore:
+ """Choose the bound a call is admitted under."""
+ if metadata_only and self.metadata_slots is not None:
+ return self.metadata_slots
+ return self.slots
+
+ def submit(
+ self,
+ fn: Callable[[], Any],
+ finished: Callable[[], None],
+ slots: threading.BoundedSemaphore | None = None,
+ ) -> Future[Any]:
+ """Admit immediately or report overload without retaining a queued call."""
+ slots = self.slots if slots is None else slots
+ if not slots.acquire(blocking=False):
+ raise ToolError(BUSY_MESSAGE)
+ try:
+ future = self.executor.submit(fn)
+ except BaseException:
+ slots.release()
+ raise
+ future.add_done_callback(lambda _: finished())
+ return future
+
+ def cancel(self, fn: Callable[[], None], finished: Callable[[], None]) -> bool:
+ """Bound cancellation I/O separately so it cannot delay the caller."""
+ if not self.cancel_slots.acquire(blocking=False):
+ return False
+ 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()
+
+ 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
+
+
+DEFAULT_TOOL_WORKERS = 16
+DEFAULT_METADATA_TOOL_WORKERS = 16
+# Threads, and so metadata connections, for transport-side metadata I/O.
+TRANSPORT_METADATA_THREADS = 1
+# Threads, and so metadata connections, for transport-side API-key lookups.
+API_KEY_AUTH_THREADS = 1
+# API-key lookups admitted at once, running or waiting; more are rejected.
+API_KEY_AUTH_PENDING = 32
+# Metadata connections held by transport-side threads rather than tool calls.
+RESERVED_METADATA_THREADS = TRANSPORT_METADATA_THREADS + API_KEY_AUTH_THREADS
+
+
+def _metadata_pool_capacity(app: Flask) -> int | None:
+ """Return how many metadata connections can be checked out at once.
+
+ ``None`` means a checkout never waits for another holder to return one
+ (e.g. ``NullPool``, per-thread pools, or unlimited overflow).
+ """
+ from contextlib import nullcontext
+
+ from superset import db
+
+ # Popping a pushed context tears down the caller's scoped session, and
+ # pools are created lazily from transport-side code that holds one.
+ in_app = has_app_context() and current_app._get_current_object() is app
+ with nullcontext() if in_app else app.app_context():
+ pool = db.engine.pool
+ if not isinstance(pool, QueuePool):
+ return None
+ # SQLAlchemy has no public accessor for the configured overflow limit.
+ max_overflow = pool._max_overflow # pylint: disable=protected-access
+ return None if max_overflow < 0 else pool.size() + max_overflow
+
+
+def admission_counts(app: Flask) -> tuple[int, int]:
+ """Split the metadata pool between warehouse and metadata-only calls.
+
+ Every holder of a metadata connection is bounded and budgeted:
+
+ - each warehouse-capable call can hold one for the whole of its warehouse
+ I/O, and its cancellation needs another (``2 * workers``);
+ - each metadata-only call holds at most one (``metadata_workers``);
+ - transport-side metadata I/O runs on ``TRANSPORT_METADATA_THREADS`` and
+ API-key lookups on ``API_KEY_AUTH_THREADS``.
+
+ ``2 * workers + metadata_workers + RESERVED_METADATA_THREADS`` never
+ exceeds the pool's capacity, so with every slot admitted no checkout
+ waits for another holder. Unless configured, warehouse calls get about
+ two thirds of the remaining connections and metadata-only calls the
+ rest, each up to 16. Explicit settings above the budget are reduced
+ with a warning. ``0`` metadata-only workers admits those tools under
+ the warehouse bound.
+
+ Returns ``(workers, metadata_workers)``.
+ """
+ configured = app.config.get("MCP_TOOL_WORKERS")
+ configured_metadata = app.config.get("MCP_METADATA_TOOL_WORKERS")
+ if configured_metadata is not None and configured_metadata < 0:
+ raise ValueError("MCP_METADATA_TOOL_WORKERS must not be negative")
+ capacity = _metadata_pool_capacity(app)
+ if capacity is None:
+ return (
+ DEFAULT_TOOL_WORKERS if configured is None else configured,
+ DEFAULT_METADATA_TOOL_WORKERS
+ if configured_metadata is None
+ else configured_metadata,
+ )
+ available = capacity - RESERVED_METADATA_THREADS
+ limit = available // 2
+ if limit < 1:
+ raise ValueError(
+ f"The metadata database pool allows {capacity} connections; "
+ "MCP tool execution needs at least "
+ f"{2 + RESERVED_METADATA_THREADS}"
+ )
+ if configured is None:
+ workers = min(DEFAULT_TOOL_WORKERS, max(1, available // 3))
+ elif configured > limit:
+ logger.warning(
+ "MCP_TOOL_WORKERS=%s needs %s metadata database connections, but the "
+ "pool allows %s; admitting %s concurrent tool calls. Raise the pool's "
+ "pool_size/max_overflow in SQLALCHEMY_ENGINE_OPTIONS to admit more.",
+ configured,
+ 2 * configured + RESERVED_METADATA_THREADS,
Review Comment:
False positive, so no code change; checked at 790e6b8ae1. The two values logged here are `2 * configured + RESERVED_METADATA_THREADS` and `capacity`, and both are ints. `configured` is the `MCP_TOOL_WORKERS` config int. `RESERVED_METADATA_THREADS` is the module constant `TRANSPORT_METADATA_THREADS + API_KEY_AUTH_THREADS` (1 + 1). `capacity` is `pool.size() + max_overflow`. CodeQL flags this because the constant name contains `API_KEY_AUTH`, but it is a thread count. No key, token, password or URI reaches this log line.
##########
superset/mcp_service/worker.py:
##########
@@ -0,0 +1,904 @@
+# 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.
+"""Bounded, deadline-aware execution of MCP tools outside the transport loop.
+
+The tool's coroutine runs on a worker-owned loop because its database APIs are
+synchronous. Only transport notifications are marshalled back to the server
+loop. Flask and SQLAlchemy lifetimes belong to the worker, not the waiting
+request: a timed-out DBAPI call cannot outlive and reuse a torn-down session.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import logging
+import threading
+import time
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import contextmanager
+from contextvars import ContextVar, copy_context
+from typing import Any, Callable, Coroutine, Iterator, ParamSpec, TYPE_CHECKING, TypeVar
+from weakref import WeakKeyDictionary
+
+from fastmcp.exceptions import ToolError
+from flask import current_app, g, has_app_context, has_request_context
+from sqlalchemy import inspect as sa_inspect
+from sqlalchemy.orm.state import InstanceState
+from sqlalchemy.pool import QueuePool
+
+from superset.mcp_service.session_scope import _mcp_session_token
+
+if TYPE_CHECKING:
+ from flask import Flask
+
+ from superset.models.core import Database
+
+logger = logging.getLogger(__name__)
+_active_call: ContextVar[WorkerCall | None] = ContextVar(
+ "mcp_worker_call", default=None
+)
+_metadata_context_owned: ContextVar[bool] = ContextVar(
+ "mcp_metadata_context_owned", default=False
+)
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
+_pools_lock = threading.Lock()
+_pools: WeakKeyDictionary[Flask, WorkerPool] = WeakKeyDictionary()
+
+
+class WorkerDeadlineExceeded(BaseException):
+ """Stop abandoned work without being swallowed by tool error handlers."""
+
+
+BUSY_MESSAGE = "MCP server busy: all tool workers are occupied. Retry later."
+
+# Tools that only read Superset's metadata database. They are admitted under a
+# separate, larger bound so they keep answering while slow warehouse queries
+# hold every warehouse slot. Anything that can reach a warehouse, a semantic
+# layer or a screenshot service stays in the warehouse bound. A listed tool
+# that does reach a warehouse must still take a warehouse slot first (see
+# WorkerCall.admit_warehouse), so a wrong entry cannot overdraw the pool.
+METADATA_ONLY_TOOLS = frozenset(
+ {
+ "find_users",
+ "get_annotation_layer_info",
+ "get_chart_info",
+ "get_chart_type_schema",
+ "get_dashboard_info",
+ "get_database_info",
+ "get_dataset_info",
+ "get_instance_info",
+ "get_layer_annotation_info",
+ "get_query_info",
+ "get_report_info",
+ "get_rls_filter_info",
+ "get_role_info",
+ "get_saved_query_info",
+ "get_schema",
+ "get_tag_info",
+ "get_task_info",
+ "get_theme_info",
+ "get_user_info",
+ "health_check",
+ "list_annotation_layers",
+ "list_charts",
+ "list_dashboards",
+ "list_databases",
+ "list_datasets",
+ "list_layer_annotations",
+ "list_queries",
+ "list_reports",
+ "list_rls_filters",
+ "list_roles",
+ "list_saved_queries",
+ "list_tags",
+ "list_tasks",
+ "list_themes",
+ "list_users",
+ }
+)
+
+
+class WorkerPool:
+ """Bound submissions, including abandoned work; never queue behind a query."""
+
+ def __init__(
+ self,
+ size: int,
+ metadata_size: int = 0,
+ transport_size: int = 1,
+ auth_size: int = 1,
+ auth_pending: int = 32,
+ ) -> None:
+ if size < 1:
+ raise ValueError("MCP_TOOL_WORKERS must be positive")
+ if metadata_size < 0:
+ raise ValueError("MCP_METADATA_TOOL_WORKERS must not be negative")
+ self.slots = threading.BoundedSemaphore(size)
+ self.metadata_slots = (
+ threading.BoundedSemaphore(metadata_size) if metadata_size else None
+ )
+ self.executor = ThreadPoolExecutor(
+ size + metadata_size, thread_name_prefix="mcp-tool"
+ )
+ # Cancellation must not wait behind the warehouse work it is cancelling.
+ # Only calls holding a warehouse slot register cancellation.
+ self.cancel_slots = threading.BoundedSemaphore(size)
+ self.cancellations = ThreadPoolExecutor(size, thread_name_prefix="mcp-cancel")
+ # Transport-side metadata I/O (tools/list filtering, audit writes,
+ # error hooks) is bounded too, so it cannot overdraw the connections
+ # budgeted for tool workers. Callers wait for it on the event loop
+ # without blocking it.
+ self.transport = ThreadPoolExecutor(
+ transport_size, thread_name_prefix="mcp-metadata"
+ )
+ # API-key lookups run before a caller is authenticated, so they get
+ # their own thread: a flood of bad keys or a slow lookup must not
+ # delay audit writes and RBAC filtering. Pending lookups are capped;
+ # past the cap a key is rejected rather than queued.
+ self.auth = ThreadPoolExecutor(auth_size, thread_name_prefix="mcp-auth")
+ self.auth_slots = threading.BoundedSemaphore(auth_pending)
+
+ def admission(self, metadata_only: bool) -> threading.BoundedSemaphore:
+ """Choose the bound a call is admitted under."""
+ if metadata_only and self.metadata_slots is not None:
+ return self.metadata_slots
+ return self.slots
+
+ def submit(
+ self,
+ fn: Callable[[], Any],
+ finished: Callable[[], None],
+ slots: threading.BoundedSemaphore | None = None,
+ ) -> Future[Any]:
+ """Admit immediately or report overload without retaining a queued call."""
+ slots = self.slots if slots is None else slots
+ if not slots.acquire(blocking=False):
+ raise ToolError(BUSY_MESSAGE)
+ try:
+ future = self.executor.submit(fn)
+ except BaseException:
+ slots.release()
+ raise
+ future.add_done_callback(lambda _: finished())
+ return future
+
+ def cancel(self, fn: Callable[[], None], finished: Callable[[], None]) -> bool:
+ """Bound cancellation I/O separately so it cannot delay the caller."""
+ if not self.cancel_slots.acquire(blocking=False):
+ return False
+ 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()
+
+ 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
+
+
+DEFAULT_TOOL_WORKERS = 16
+DEFAULT_METADATA_TOOL_WORKERS = 16
+# Threads, and so metadata connections, for transport-side metadata I/O.
+TRANSPORT_METADATA_THREADS = 1
+# Threads, and so metadata connections, for transport-side API-key lookups.
+API_KEY_AUTH_THREADS = 1
+# API-key lookups admitted at once, running or waiting; more are rejected.
+API_KEY_AUTH_PENDING = 32
+# Metadata connections held by transport-side threads rather than tool calls.
+RESERVED_METADATA_THREADS = TRANSPORT_METADATA_THREADS + API_KEY_AUTH_THREADS
+
+
+def _metadata_pool_capacity(app: Flask) -> int | None:
+ """Return how many metadata connections can be checked out at once.
+
+ ``None`` means a checkout never waits for another holder to return one
+ (e.g. ``NullPool``, per-thread pools, or unlimited overflow).
+ """
+ from contextlib import nullcontext
+
+ from superset import db
+
+ # Popping a pushed context tears down the caller's scoped session, and
+ # pools are created lazily from transport-side code that holds one.
+ in_app = has_app_context() and current_app._get_current_object() is app
+ with nullcontext() if in_app else app.app_context():
+ pool = db.engine.pool
+ if not isinstance(pool, QueuePool):
+ return None
+ # SQLAlchemy has no public accessor for the configured overflow limit.
+ max_overflow = pool._max_overflow # pylint: disable=protected-access
+ return None if max_overflow < 0 else pool.size() + max_overflow
+
+
+def admission_counts(app: Flask) -> tuple[int, int]:
+ """Split the metadata pool between warehouse and metadata-only calls.
+
+ Every holder of a metadata connection is bounded and budgeted:
+
+ - each warehouse-capable call can hold one for the whole of its warehouse
+ I/O, and its cancellation needs another (``2 * workers``);
+ - each metadata-only call holds at most one (``metadata_workers``);
+ - transport-side metadata I/O runs on ``TRANSPORT_METADATA_THREADS`` and
+ API-key lookups on ``API_KEY_AUTH_THREADS``.
+
+ ``2 * workers + metadata_workers + RESERVED_METADATA_THREADS`` never
+ exceeds the pool's capacity, so with every slot admitted no checkout
+ waits for another holder. Unless configured, warehouse calls get about
+ two thirds of the remaining connections and metadata-only calls the
+ rest, each up to 16. Explicit settings above the budget are reduced
+ with a warning. ``0`` metadata-only workers admits those tools under
+ the warehouse bound.
+
+ Returns ``(workers, metadata_workers)``.
+ """
+ configured = app.config.get("MCP_TOOL_WORKERS")
+ configured_metadata = app.config.get("MCP_METADATA_TOOL_WORKERS")
+ if configured_metadata is not None and configured_metadata < 0:
+ raise ValueError("MCP_METADATA_TOOL_WORKERS must not be negative")
+ capacity = _metadata_pool_capacity(app)
+ if capacity is None:
+ return (
+ DEFAULT_TOOL_WORKERS if configured is None else configured,
+ DEFAULT_METADATA_TOOL_WORKERS
+ if configured_metadata is None
+ else configured_metadata,
+ )
+ available = capacity - RESERVED_METADATA_THREADS
+ limit = available // 2
+ if limit < 1:
+ raise ValueError(
+ f"The metadata database pool allows {capacity} connections; "
+ "MCP tool execution needs at least "
+ f"{2 + RESERVED_METADATA_THREADS}"
+ )
+ if configured is None:
+ workers = min(DEFAULT_TOOL_WORKERS, max(1, available // 3))
+ elif configured > limit:
+ logger.warning(
+ "MCP_TOOL_WORKERS=%s needs %s metadata database connections, but the "
+ "pool allows %s; admitting %s concurrent tool calls. Raise the pool's "
+ "pool_size/max_overflow in SQLALCHEMY_ENGINE_OPTIONS to admit more.",
+ configured,
+ 2 * configured + RESERVED_METADATA_THREADS,
+ capacity,
+ limit,
+ )
+ workers = limit
+ else:
+ workers = configured
+ remaining = available - 2 * workers
+ if configured_metadata is None:
+ metadata_workers = min(DEFAULT_METADATA_TOOL_WORKERS, remaining)
+ elif configured_metadata > remaining:
+ logger.warning(
+ "MCP_METADATA_TOOL_WORKERS=%s exceeds the %s metadata database "
+ "connections left after %s warehouse-capable tool calls; admitting "
+ "%s metadata-only tool calls.",
+ configured_metadata,
+ remaining,
+ workers,
+ remaining,
+ )
+ metadata_workers = remaining
+ else:
+ metadata_workers = configured_metadata
+ return workers, metadata_workers
+
+
+def tool_worker_count(app: Flask) -> int:
+ """Admit only as many warehouse-capable calls as the pool can always serve."""
+ return admission_counts(app)[0]
+
+
+def metadata_tool_worker_count(app: Flask) -> int:
+ """Admit metadata-only calls from what warehouse calls leave of the pool."""
+ return admission_counts(app)[1]
+
+
+def _get_pool(app: Flask) -> WorkerPool:
+ """Lazily create a pool for this application, without import-time threads."""
+ with _pools_lock:
+ if app not in _pools:
+ size, metadata_size = admission_counts(app)
+ logger.info(
+ "MCP tool calls admitted concurrently: %s warehouse-capable, "
+ "%s metadata-only; %s transport metadata thread(s), "
+ "%s API-key lookup thread(s)",
+ size,
+ metadata_size,
+ TRANSPORT_METADATA_THREADS,
+ API_KEY_AUTH_THREADS,
Review Comment:
False positive, so no code change; checked at 790e6b8ae1. `API_KEY_AUTH_THREADS` is the module constant `1`: the number of threads that run API-key lookups. The other arguments are also int counts (`size`, `metadata_size`, `TRANSPORT_METADATA_THREADS`). CodeQL flags the name because it contains `API_KEY_AUTH`, but no key, token or credential value reaches this log line.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org