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: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
