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: [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]
