From 0e4e368768ffa238b44bc02660fe788bd2feeb1f Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Wed, 5 Aug 2026 13:52:34 -0500 Subject: [PATCH] feat(mcp): add observability to MCP service (#41921) Co-authored-by: Claude --- superset/mcp_service/PRODUCTION.md | 48 +- .../mcp_service/chart/tool/get_chart_data.py | 78 +++ .../chart/tool/get_chart_type_schema.py | 4 +- .../mcp_service/dataset/tool/query_dataset.py | 1 + superset/mcp_service/mcp_config.py | 28 + superset/mcp_service/middleware.py | 286 ++++++++-- superset/mcp_service/tag/tool/get_tag_info.py | 3 + .../chart/tool/test_get_chart_data.py | 213 +++++++ .../unit_tests/mcp_service/test_middleware.py | 322 +++++++++++ .../mcp_service/test_middleware_logging.py | 539 ++++++++++++++++++ 10 files changed, 1486 insertions(+), 36 deletions(-) diff --git a/superset/mcp_service/PRODUCTION.md b/superset/mcp_service/PRODUCTION.md index 3ef3538b0e7..ad402fee266 100644 --- a/superset/mcp_service/PRODUCTION.md +++ b/superset/mcp_service/PRODUCTION.md @@ -217,19 +217,65 @@ All MCP tools return consistent error schemas: **Error Tracking (Sentry)**: +MCP tool execution runs on the FastMCP/Starlette asyncio stack, not Flask +request handling, so `FlaskIntegration` alone does **not** capture MCP tool +errors — it only sees the regular Superset web app's Flask requests. Use the +vendor-neutral `MCP_ERROR_HOOK` config to forward system-class MCP errors +(unexpected exceptions — not user errors like bad params or permission +denials) to Sentry instead: + ```python # superset_config.py import sentry_sdk from sentry_sdk.integrations.flask import FlaskIntegration +from sentry_sdk.scrubber import EventScrubber sentry_sdk.init( dsn="https://your-dsn@sentry.io/project-id", - integrations=[FlaskIntegration()], + integrations=[FlaskIntegration()], # covers the Flask web app only environment="production", traces_sample_rate=0.1, # 10% of transactions + # The hook forwards the RAW exception, which may embed connection + # strings or tokens in its message — scrub events before they leave + # the process. The default EventScrubber covers common secret keys; + # add a before_send for anything deployment-specific. + event_scrubber=EventScrubber(recursive=True), + send_default_pii=False, ) + + +def _mcp_error_hook(error: Exception, context: dict) -> None: + """Forward system-class MCP tool errors to Sentry. + + ``context`` always contains tool_name, mcp_call_id, user_id, + error_type, sanitized_message, and duration_ms. Values may be + unavailable depending on the capture path: user_id and duration_ms + are None on the last-resort path, mcp_call_id is None outside a tool + call, and tool_name falls back to "unknown" for non-tool messages. + ``error`` is the RAW exception — only ``sanitized_message`` has been + scrubbed — hence the event_scrubber/before_send above. + """ + with sentry_sdk.new_scope() as scope: + scope.set_tag("mcp.tool", context.get("tool_name")) + scope.set_tag("mcp.call_id", context.get("mcp_call_id")) + scope.set_user({"id": context.get("user_id")}) + sentry_sdk.capture_exception(error) + + +MCP_ERROR_HOOK = _mcp_error_hook ``` +`MCP_ERROR_HOOK` is invoked from `GlobalErrorHandlerMiddleware` (the primary +capture point, for every system-class error) and from +`StructuredContentStripperMiddleware`'s last-resort exception handler (for +errors that slip past the primary handler entirely). Hook failures are +caught and logged; they never affect the MCP response. + +The hook runs synchronously on the asyncio event loop, so a blocking hook +stalls all in-flight tool handling — hand events to a background transport +rather than doing network I/O inline. The Sentry SDK already does this: +`capture_exception` enqueues to a background worker. + **Metrics Export (Prometheus)**: ```python diff --git a/superset/mcp_service/chart/tool/get_chart_data.py b/superset/mcp_service/chart/tool/get_chart_data.py index 0ed596d0c87..ed1bb582aea 100644 --- a/superset/mcp_service/chart/tool/get_chart_data.py +++ b/superset/mcp_service/chart/tool/get_chart_data.py @@ -344,6 +344,10 @@ async def get_chart_data( # noqa: C901 ) cached_form_data = get_cached_form_data(request.form_data_key) if not cached_form_data: + logger.warning( + "get_chart_data: no cached form_data for form_data_key=%s", + request.form_data_key, + ) return ChartError( error="No cached chart data found for form_data_key. " "The cache may have expired.", @@ -352,11 +356,22 @@ async def get_chart_data( # noqa: C901 try: cached_form_data_dict = utils_json.loads(cached_form_data) except (TypeError, ValueError) as e: + logger.warning( + "get_chart_data: failed to parse cached form_data " + "for form_data_key=%s: %s", + request.form_data_key, + e, + ) return ChartError( error=f"Failed to parse cached form_data: {e}", error_type="ParseError", ) if not isinstance(cached_form_data_dict, dict): + logger.warning( + "get_chart_data: cached form_data is not a JSON object " + "for form_data_key=%s", + request.form_data_key, + ) return ChartError( error="Cached form_data is not a valid JSON object.", error_type="ParseError", @@ -383,6 +398,7 @@ async def get_chart_data( # noqa: C901 with event_logger.log_context(action="mcp.get_chart_data.chart_lookup"): await ctx.debug("Looking up chart: identifier=%s" % (request.identifier,)) if request.identifier is None: + logger.warning("get_chart_data: called without a chart identifier") return ChartError( error="Chart identifier is required", error_type="ValidationError", @@ -395,6 +411,9 @@ async def get_chart_data( # noqa: C901 if not chart: await ctx.warning("Chart not found: identifier=%s" % (request.identifier,)) + logger.warning( + "get_chart_data: chart not found: identifier=%s", request.identifier + ) safe_id = escape_llm_context_delimiters(str(request.identifier)[:200]) return ChartError( error=( @@ -424,6 +443,11 @@ async def get_chart_data( # noqa: C901 "Chart found but dataset is not accessible: %s" % (validation_result.error,) ) + logger.warning( + "get_chart_data: dataset not accessible for chart_id=%s: %s", + chart.id, + validation_result.error, + ) return ChartError( error=validation_result.error or "Chart's dataset is not accessible. " @@ -574,6 +598,12 @@ async def get_chart_data( # noqa: C901 "Re-save the chart to populate query_context." % (chart.id, viz_type) ) + logger.warning( + "get_chart_data: cannot construct fallback query for " + "chart_id=%s (viz_type=%s): no metrics/columns found", + chart.id, + viz_type, + ) return ChartError( error=( f"Chart {chart.id} (type: {viz_type}) has no " @@ -645,6 +675,12 @@ async def get_chart_data( # noqa: C901 "Empty query results: chart_id=%s, chart_type=%s" % (chart.id, chart.viz_type) ) + logger.warning( + "get_chart_data: empty query results for chart_id=%s, " + "chart_type=%s", + chart.id, + chart.viz_type, + ) return ChartError( error=f"No query results returned for chart {chart.id}. " f"This may occur with chart types like big_number.", @@ -669,6 +705,10 @@ async def get_chart_data( # noqa: C901 # Check if we have data to work with if not data: await ctx.warning("No data in query results: chart_id=%s" % (chart.id,)) + logger.warning( + "get_chart_data: no data in query results for chart_id=%s", + chart.id, + ) return ChartError( error=f"No data available for chart {chart.id}", error_type="NoData" ) @@ -855,6 +895,12 @@ async def get_chart_data( # noqa: C901 ) ) + except (OAuth2RedirectError, OAuth2Error): + # OAuth errors subclass SupersetException and would otherwise be + # swallowed by the generic handler below; re-raise so the + # dedicated outer handlers return the OAuth redirect message + # instead of a generic DataError. + raise except (CommandException, SupersetException, ValueError) as data_error: await ctx.error( "Data retrieval failed: chart_id=%s, error=%s, error_type=%s" @@ -875,6 +921,10 @@ async def get_chart_data( # noqa: C901 "Chart data requires OAuth authentication: identifier=%s" % request.identifier ) + logger.info( + "get_chart_data: OAuth authentication required for identifier=%s", + request.identifier, + ) return ChartError( error=build_oauth2_redirect_message(ex), error_type="OAUTH2_REDIRECT", @@ -883,6 +933,10 @@ async def get_chart_data( # noqa: C901 await ctx.error( "OAuth2 configuration error: identifier=%s" % request.identifier ) + logger.warning( + "get_chart_data: OAuth2 configuration error for identifier=%s", + request.identifier, + ) return ChartError( error=OAUTH2_CONFIG_ERROR_MESSAGE, error_type="OAUTH2_REDIRECT_ERROR", @@ -930,6 +984,11 @@ async def _query_from_form_data( datasource_id = parts[0] if not datasource_id: + logger.warning( + "get_chart_data: cached form_data has no datasource information " + "(form_data_key=%s)", + request.form_data_key, + ) return ChartError( error="Cached form_data does not contain datasource information.", error_type="InvalidFormData", @@ -956,6 +1015,11 @@ async def _query_from_form_data( result = command.run() if not result or "queries" not in result or len(result["queries"]) == 0: + logger.warning( + "get_chart_data: empty query results for unsaved chart " + "(form_data_key=%s)", + request.form_data_key, + ) return ChartError( error="No query results returned for unsaved chart.", error_type="EmptyQuery", @@ -966,6 +1030,10 @@ async def _query_from_form_data( raw_columns = query_result.get("colnames", []) if not data: + logger.warning( + "get_chart_data: no data for unsaved chart (form_data_key=%s)", + request.form_data_key, + ) return ChartError( error="No data available for unsaved chart.", error_type="NoData", @@ -1031,6 +1099,11 @@ async def _query_from_form_data( ) ) + except (OAuth2RedirectError, OAuth2Error): + # OAuth errors subclass SupersetException; re-raise so the caller's + # outer OAuth handlers return the redirect instead of a generic + # DataError. + raise except (CommandException, SupersetException, ValueError) as e: logger.error("Error querying unsaved chart data: %s", e) return ChartError( @@ -1174,6 +1247,11 @@ def _try_xlsxwriter_fallback( except ImportError: from superset.mcp_service.chart.schemas import ChartError + logger.warning( + "get_chart_data: Excel export failed for chart_id=%s — " + "neither openpyxl nor xlsxwriter is installed", + chart.id, + ) return ChartError( error="Excel export requires openpyxl or xlsxwriter package", error_type="ExportError", diff --git a/superset/mcp_service/chart/tool/get_chart_type_schema.py b/superset/mcp_service/chart/tool/get_chart_type_schema.py index af71848a27c..1823b2542e6 100644 --- a/superset/mcp_service/chart/tool/get_chart_type_schema.py +++ b/superset/mcp_service/chart/tool/get_chart_type_schema.py @@ -27,6 +27,7 @@ from typing import Any, Dict from pydantic import TypeAdapter from superset_core.mcp.decorators import tool, ToolAnnotations +from superset.extensions import event_logger from superset.mcp_service.chart.schemas import ( BigNumberChartConfig, BoxPlotChartConfig, @@ -245,4 +246,5 @@ def get_chart_type_schema( Returns the JSON Schema for the requested chart type, optionally with working examples. """ - return _get_chart_type_schema_impl(chart_type, include_examples) + with event_logger.log_context(action="mcp.get_chart_type_schema.lookup"): + return _get_chart_type_schema_impl(chart_type, include_examples) diff --git a/superset/mcp_service/dataset/tool/query_dataset.py b/superset/mcp_service/dataset/tool/query_dataset.py index 578d8ac14ab..6f8938ff55b 100644 --- a/superset/mcp_service/dataset/tool/query_dataset.py +++ b/superset/mcp_service/dataset/tool/query_dataset.py @@ -408,6 +408,7 @@ async def query_dataset( # noqa: C901 ) except SQLAlchemyError as exc: + logger.exception("Database error while querying dataset") await ctx.error("Database error: %s" % (str(exc),)) return DatasetError.create( error=f"Database error: {exc}", diff --git a/superset/mcp_service/mcp_config.py b/superset/mcp_service/mcp_config.py index 5094251e1c7..95190e921f7 100644 --- a/superset/mcp_service/mcp_config.py +++ b/superset/mcp_service/mcp_config.py @@ -83,6 +83,34 @@ MCP_RBAC_ENABLED = True # MCP_DISABLED_TOOLS = {"extensions.myorg.myext.some_tool"} MCP_DISABLED_TOOLS: set[str] = set() +# Pluggable error-capture hook, invoked for system-class MCP tool errors +# (unexpected exceptions — database down, bugs — not user errors like bad +# params or permission denials). Lets operators forward failures to an +# external error tracker (e.g. Sentry) without the OSS repo depending on any +# particular vendor SDK: FlaskIntegration does not see FastMCP tool +# execution, since it runs on the asyncio/Starlette stack, not a Flask +# request. See PRODUCTION.md "Error Tracking" for a Sentry wiring example. +# +# Signature: hook(error: Exception, context: dict[str, Any]) -> None +# ``context`` always contains the keys "tool_name", "mcp_call_id", +# "user_id", "error_type", "sanitized_message", and "duration_ms" — but +# values may be unavailable depending on the capture path: "user_id" and +# "duration_ms" are None on the last-resort path +# (StructuredContentStripperMiddleware), "mcp_call_id" is None outside a +# tool call, and "tool_name" falls back to "unknown" for non-tool +# messages. Only "sanitized_message" is scrubbed — the ``error`` argument +# is the RAW exception and may contain sensitive data (connection +# strings, tokens); sanitize it before exporting, or report +# "sanitized_message" instead. +# +# The hook runs SYNCHRONOUSLY on the asyncio event loop, so a blocking hook +# stalls all in-flight tool handling. Do not perform network I/O inline; +# hand the event to a background transport (the Sentry SDK's +# capture_exception already queues to a worker thread). Exceptions raised +# by the hook itself are caught and logged as a warning; they never affect +# the MCP response. +MCP_ERROR_HOOK: Callable[[Exception, dict[str, Any]], None] | None = None + # ============================================================================= # MCP Chart Plugin Filtering # ============================================================================= diff --git a/superset/mcp_service/middleware.py b/superset/mcp_service/middleware.py index 110412c2692..61ec04f1b87 100644 --- a/superset/mcp_service/middleware.py +++ b/superset/mcp_service/middleware.py @@ -16,6 +16,7 @@ # under the License. import logging +import re import secrets import time from contextvars import ContextVar @@ -26,7 +27,7 @@ from fastmcp.exceptions import ToolError, ValidationError as FastMCPValidationEr from fastmcp.server.middleware import Middleware, MiddlewareContext from fastmcp.server.middleware.middleware import CallNext from fastmcp.tools.tool import Tool, ToolResult -from flask import g, has_app_context +from flask import g from pydantic import ValidationError from sqlalchemy.exc import OperationalError, TimeoutError from starlette.exceptions import HTTPException @@ -37,7 +38,7 @@ from superset.commands.exceptions import ( ObjectNotFoundError, ) from superset.exceptions import SupersetException, SupersetSecurityException -from superset.extensions import event_logger +from superset.extensions import event_logger, stats_logger_manager from superset.mcp_service.auth import ( _get_app_context_manager, get_user_from_request, @@ -63,6 +64,11 @@ from superset.utils.core import get_user_id logger = logging.getLogger(__name__) _mcp_call_id_var: ContextVar[str | None] = ContextVar("mcp_call_id", default=None) +# Conservative shape for a tool-name segment embedded in a StatsD metric key. +# Matches registered tool names (snake_case, plus dots for extension-prefixed +# tools) while rejecting StatsD metadata characters and unbounded lengths. +_METRIC_TOOL_NAME_RE = re.compile(r"[A-Za-z0-9_][A-Za-z0-9_.\-]{0,127}") + def _sanitize_error_for_logging(error: Exception) -> str: """Sanitize error messages to prevent information disclosure in logs.""" @@ -134,6 +140,29 @@ def _sanitize_error_for_logging(error: Exception) -> str: return error_str +def _invoke_error_hook(error: Exception, hook_context: dict[str, Any]) -> None: + """Invoke the operator-configured ``MCP_ERROR_HOOK``, if any. + + Kept vendor-neutral (no ``sentry_sdk`` import here) so the OSS repo has + no hard dependency on any particular error tracker — operators wire + their own hook (e.g. calling ``sentry_sdk.capture_exception``) via + ``MCP_ERROR_HOOK`` in ``superset_config.py``. Hook failures are logged + and swallowed; they must never affect the MCP response. + """ + try: + from superset.mcp_service.flask_singleton import get_flask_app + + hook = get_flask_app().config.get("MCP_ERROR_HOOK") + except Exception: # noqa: BLE001 + return + if hook is None: + return + try: + hook(error, hook_context) + except Exception as hook_error: # noqa: BLE001 + logger.warning("MCP_ERROR_HOOK raised an exception: %s", hook_error) + + # Errors caused by the LLM/user — expected in normal MCP operation. # Agents send bad params, try tools they lack access to, request nonexistent # resources. These are 400-class errors and should be logged at WARNING. @@ -242,6 +271,31 @@ class LoggingMiddleware(Middleware): return False return bool(isinstance(payload, dict) and payload.get("error_type")) + @staticmethod + def _extract_error_type_from_response(result: ToolResult) -> str | None: + """Extract the ``error_type`` field from a serialized error response. + + Structured MCP error schemas (ChartError, DashboardError, etc.) embed + an ``error_type`` string. Parsing it here — instead of discarding it + after the substring sniff in ``_is_error_response`` — lets it flow + into the log line, curated payload, and metric tag. + """ + from superset.utils.json import loads as json_loads + + try: + text = result.content[0].text + except (AttributeError, IndexError): + return None + try: + payload = json_loads(text) + except (ValueError, TypeError): + return None + if isinstance(payload, dict): + error_type = payload.get("error_type") + if isinstance(error_type, str): + return error_type + return None + def _extract_context_info( self, context: MiddlewareContext ) -> tuple[ @@ -416,16 +470,21 @@ class LoggingMiddleware(Middleware): mcp_tool=mcp_tool, error_type=error_type, ) - if has_app_context(): - event_logger.log( - user_id=user_id, - action="mcp_tool_call", - dashboard_id=dashboard_id, - duration_ms=duration_ms, - slice_id=slice_id, - referrer=None, - curated_payload=payload, - ) + try: + with _get_app_context_manager(): + event_logger.log( + user_id=user_id, + action="mcp_tool_call", + dashboard_id=dashboard_id, + duration_ms=duration_ms, + slice_id=slice_id, + referrer=None, + curated_payload=payload, + ) + except Exception as log_error: # noqa: BLE001 + # A failing event logger or app-context setup must not mask the + # tool result or prevent metrics and structured logs below. + logger.warning("Failed to log mcp_tool_call event: %s", log_error) extra_parts = [] if mcp_tool is not None: extra_parts.append(f"mcp_tool={mcp_tool}") @@ -449,6 +508,76 @@ class LoggingMiddleware(Middleware): extra, ) + @staticmethod + async def _resolve_metric_tool_name( + context: MiddlewareContext, + tool_name: str | None, + mcp_tool: str | None, + ) -> str: + """Return a StatsD-safe tool segment for the per-tool metric keys. + + Both ``mcp_tool`` (the ``call_tool`` proxy's ``name`` argument) and + ``tool_name`` (the raw message name) are client-controlled input. + Using them verbatim in a metric key would let any authenticated + client mint unbounded metric series or inject StatsD metadata + characters (``\\n``/``:``/``|``) into the wire format. Only names + that resolve in the FastMCP tool registry are used; anything else + falls back to a constant. The raw name still reaches the curated + payload and log line, which are not StatsD keys. + """ + candidate = mcp_tool or tool_name + if not candidate: + return "unknown" + try: + registered = await context.fastmcp_context.fastmcp.get_tool(candidate) + except (AttributeError, TypeError): + # No registry reachable from this context (e.g. unit tests with + # mocked contexts) — accept only conservatively-shaped names. + if _METRIC_TOOL_NAME_RE.fullmatch(candidate): + return candidate + registered = None + except Exception: # noqa: BLE001 + # Registry reachable but the lookup failed (NotFoundError in + # FastMCP versions that raise instead of returning None) — + # treat as unregistered. + registered = None + if registered is not None: + return candidate + return "call_tool" if mcp_tool else "unknown" + + async def _emit_call_metrics( + self, + context: MiddlewareContext, + tool_name: str | None, + mcp_tool: str | None, + *, + success: bool, + raised_is_user_error: bool | None, + duration_ms: int, + ) -> None: + """Emit the per-tool outcome counter and timing for one call. + + Single emission point for the per-tool outcome counters — + GlobalErrorHandlerMiddleware (inner) re-raises every failure as + ToolError, so counting there as well would double-count raised + errors. Mirrors base_api.py's success/warning/error split: raised + user errors → warning, raised system errors → error. Structured + error responses (``raised_is_user_error`` is None) carry a + free-form error_type that cannot be reliably classified, so they + count as error (the parsed error_type is in the curated payload). + """ + metric_tool = await self._resolve_metric_tool_name(context, tool_name, mcp_tool) + if success: + outcome = "success" + elif raised_is_user_error: + outcome = "warning" + else: + outcome = "error" + stats_logger_manager.instance.incr(f"mcp.tool.{metric_tool}.{outcome}") + stats_logger_manager.instance.timing( + f"mcp.tool.{metric_tool}.time", duration_ms + ) + async def on_call_tool( self, context: MiddlewareContext, @@ -467,9 +596,12 @@ class LoggingMiddleware(Middleware): success = False error_type: str | None = None result: Any = None + raised_is_user_error: bool | None = None try: result = await call_next(context) success = not self._is_error_response(result) + if not success and isinstance(result, ToolResult): + error_type = self._extract_error_type_from_response(result) if isinstance(result, ToolResult): existing_meta = result.meta or {} result = ToolResult( @@ -479,10 +611,21 @@ class LoggingMiddleware(Middleware): ) return result except Exception as exc: - error_type = type(exc).__name__ + # GlobalErrorHandlerMiddleware (inner) wraps tool exceptions in + # ToolError with the original attached as __cause__; unwrap it so + # error_type and the user/system classification reflect the real + # failure rather than the ToolError wrapper. + original = ( + exc.__cause__ + if isinstance(exc, ToolError) and exc.__cause__ is not None + else exc + ) + error_type = type(original).__name__ + raised_is_user_error = _is_user_error(original) success = False raise finally: + duration_ms = int((time.time() - start_time) * 1000) self._log_call_tool_result( context=context, tool_name=tool_name, @@ -499,6 +642,19 @@ class LoggingMiddleware(Middleware): result=result, start_time=start_time, ) + try: + await self._emit_call_metrics( + context, + tool_name, + mcp_tool, + success=success, + raised_is_user_error=raised_is_user_error, + duration_ms=duration_ms, + ) + except Exception as metrics_error: # noqa: BLE001 + # A failing stats backend must never mask the tool's real + # result or exception — metrics are a side effect only. + logger.warning("Failed to emit MCP tool metrics: %s", metrics_error) async def on_message( self, @@ -509,24 +665,27 @@ class LoggingMiddleware(Middleware): agent_id, user_id, dashboard_id, slice_id, dataset_id, params = ( self._extract_context_info(context) ) - if has_app_context(): - event_logger.log( - user_id=user_id, - action="mcp_message", - dashboard_id=dashboard_id, - duration_ms=None, - slice_id=slice_id, - referrer=None, - curated_payload={ - "tool": getattr(context.message, "name", None), - "agent_id": agent_id, - "params": _sanitize_params(params), - "method": context.method, - "dashboard_id": dashboard_id, - "slice_id": slice_id, - "dataset_id": dataset_id, - }, - ) + try: + with _get_app_context_manager(): + event_logger.log( + user_id=user_id, + action="mcp_message", + dashboard_id=dashboard_id, + duration_ms=None, + slice_id=slice_id, + referrer=None, + curated_payload={ + "tool": getattr(context.message, "name", None), + "agent_id": agent_id, + "params": _sanitize_params(params), + "method": context.method, + "dashboard_id": dashboard_id, + "slice_id": slice_id, + "dataset_id": dataset_id, + }, + ) + except Exception as log_error: # noqa: BLE001 + logger.warning("Failed to log mcp_message event: %s", log_error) logger.info( "MCP message: tool=%s, agent_id=%s, user_id=%s, method=%s", getattr(context.message, "name", None), @@ -593,8 +752,42 @@ class StructuredContentStripperMiddleware(Middleware): # GlobalErrorHandlerMiddleware, ValueError, TypeError, etc. — # will cause encoding failures on the wire. mcp_call_id = _mcp_call_id_var.get(None) + # This is the documented "must never propagate" point. The + # client-facing text must be SANITIZED — an exception that + # bypasses GlobalErrorHandlerMiddleware could otherwise leak + # raw internals (SQL fragments, connection strings, tokens) to + # the caller; every other client-facing error path already + # runs through _sanitize_error_for_logging. That call (and + # str(e) inside it) can itself raise on a pathological + # __str__, so guard it and fall back to the exception class + # name, which never propagates. + try: + sanitized_message = _sanitize_error_for_logging(e) + except Exception: # noqa: BLE001 + sanitized_message = type(e).__name__ + error_text = f"Error: {sanitized_message}" + if not isinstance(e, ToolError): + # GlobalErrorHandlerMiddleware converts every exception it + # sees into ToolError (and already invokes MCP_ERROR_HOOK + # for system-class errors there). A non-ToolError reaching + # this final catch means it slipped past that handler + # entirely — invoke the hook here as the true last-resort + # capture point. All contract keys are populated so hooks + # can index them unconditionally; user_id and duration_ms + # are unknown at this layer and passed as None. + _invoke_error_hook( + e, + { + "tool_name": getattr(context.message, "name", "unknown"), + "mcp_call_id": mcp_call_id, + "user_id": None, + "error_type": type(e).__name__, + "sanitized_message": sanitized_message, + "duration_ms": None, + }, + ) return ToolResult( - content=[mt.TextContent(type="text", text=f"Error: {e}")], + content=[mt.TextContent(type="text", text=error_text)], meta={"mcp_call_id": mcp_call_id} if mcp_call_id else None, ) if isinstance(result, ToolResult) and result.structured_content is not None: @@ -731,6 +924,28 @@ class GlobalErrorHandlerMiddleware(Middleware): except Exception as log_error: logger.warning("Failed to log error event: %s", log_error) + # No stats emission here: this handler re-raises every failure as + # ToolError, which the outer LoggingMiddleware catches and counts + # (with the user/system classification recovered via __cause__). + # Emitting a counter here as well would double-count raised errors. + + mcp_call_id = _mcp_call_id_var.get(None) + if not is_user: + # System-class errors only — user errors (bad params, permission + # denials) are expected MCP traffic and would otherwise flood an + # error tracker. + _invoke_error_hook( + error, + { + "tool_name": tool_name, + "mcp_call_id": mcp_call_id, + "user_id": user_id, + "error_type": type(error).__name__, + "sanitized_message": sanitized_error, + "duration_ms": duration_ms, + }, + ) + # Handle specific error types with appropriate responses if isinstance(error, ToolError): # Tool errors are already formatted for MCP @@ -801,8 +1016,11 @@ class GlobalErrorHandlerMiddleware(Middleware): f"Connection error in {tool_name}: {_sanitize_error_for_logging(error)}" ) from error else: - # Generic internal errors — truly unexpected - error_id = f"err_{int(time.time())}" + # Generic internal errors — truly unexpected. Reuse the per-call + # mcp_call_id (set by LoggingMiddleware.on_call_tool) instead of a + # second-granularity timestamp, which collides under concurrent + # failures. + error_id = mcp_call_id or f"err_{secrets.token_hex(8)}" logger.error("Unexpected error [%s] in %s: %s", error_id, tool_name, error) raise ToolError( diff --git a/superset/mcp_service/tag/tool/get_tag_info.py b/superset/mcp_service/tag/tool/get_tag_info.py index f06f1c8332c..ab4cde6a929 100644 --- a/superset/mcp_service/tag/tool/get_tag_info.py +++ b/superset/mcp_service/tag/tool/get_tag_info.py @@ -97,6 +97,9 @@ async def get_tag_info(request: GetTagInfoRequest, ctx: Context) -> TagInfo | Ta return result except Exception as e: + logger.exception( + "Tag information retrieval failed: identifier=%s", request.identifier + ) await ctx.error( "Tag information retrieval failed: identifier=%s, error=%s, error_type=%s" % (request.identifier, str(e), type(e).__name__) diff --git a/tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py b/tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py index be386200253..8cb3b6ef714 100644 --- a/tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py +++ b/tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py @@ -1327,6 +1327,219 @@ class TestChartLookupEagerLoading: assert _extract_metrics_load_path(query_options[0]) == ["table", "metrics"] +class TestOAuthErrorRouting: + """Query-time OAuth errors must reach the dedicated OAuth handlers. + + OAuth2RedirectError/OAuth2Error subclass SupersetException, so without + the explicit re-raise ahead of the generic inner handler they would be + swallowed into a generic DataError and the client would never see the + OAuth redirect message. + """ + + def _make_chart(self) -> Any: + from unittest.mock import MagicMock + + from superset.utils import json as utils_json + + chart = MagicMock() + chart.id = 1 + chart.slice_name = "My Chart" + chart.viz_type = "table" + chart.query_context = None + chart.params = utils_json.dumps( + {"viz_type": "table", "metrics": ["count"], "groupby": ["gender"]} + ) + chart.datasource_id = 1 + chart.datasource_type = "table" + return chart + + def _patch_query_path( + self, + monkeypatch: pytest.MonkeyPatch, + chart: Any, + run_error: Exception, + ) -> None: + """Route the query-execution path into a command that raises.""" + from unittest.mock import MagicMock + + chart_data_module = importlib.import_module( + "superset.mcp_service.chart.tool.get_chart_data" + ) + query_context_factory_module = importlib.import_module( + "superset.common.query_context_factory" + ) + get_data_command_module = importlib.import_module( + "superset.commands.chart.data.get_data_command" + ) + + validation = MagicMock() + validation.is_valid = True + validation.warnings = [] + + class QueryContextFactory: + def create(self, **kwargs: Any) -> object: + return object() + + class RaisingChartDataCommand: + def __init__(self, query_context: object) -> None: + self.query_context = query_context + + def validate(self) -> None: + pass + + def run(self) -> dict[str, Any]: + raise run_error + + monkeypatch.setattr( + chart_data_module, + "find_chart_by_identifier", + lambda *args, **kwargs: chart, + ) + monkeypatch.setattr( + chart_data_module, + "validate_chart_dataset", + lambda *args, **kwargs: validation, + ) + monkeypatch.setattr( + query_context_factory_module, + "QueryContextFactory", + QueryContextFactory, + ) + monkeypatch.setattr( + get_data_command_module, + "ChartDataCommand", + RaisingChartDataCommand, + ) + + @pytest.mark.asyncio + async def test_oauth2_redirect_error_returns_oauth_redirect_type( + self, + mcp_server: Any, + mock_auth: Any, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + from fastmcp import Client + + from superset.exceptions import OAuth2RedirectError + from superset.utils import json as utils_json + + self._patch_query_path( + monkeypatch, + self._make_chart(), + OAuth2RedirectError( + "https://example.com/oauth", + "tab-1", + "https://example.com/redirect", + ), + ) + + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_chart_data", {"request": {"identifier": 1}} + ) + data = utils_json.loads(result.content[0].text) + + assert data["error_type"] == "OAUTH2_REDIRECT" + assert data["error_type"] != "DataError" + + @pytest.mark.asyncio + async def test_oauth2_error_returns_oauth_redirect_error_type( + self, + mcp_server: Any, + mock_auth: Any, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + from fastmcp import Client + + from superset.exceptions import OAuth2Error + from superset.utils import json as utils_json + + self._patch_query_path( + monkeypatch, + self._make_chart(), + OAuth2Error("token refresh failed"), + ) + + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_chart_data", {"request": {"identifier": 1}} + ) + data = utils_json.loads(result.content[0].text) + + assert data["error_type"] == "OAUTH2_REDIRECT_ERROR" + + @pytest.mark.asyncio + async def test_oauth2_redirect_on_form_data_key_path( + self, + mcp_server: Any, + mock_auth: Any, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """The unsaved-chart (form_data_key, no identifier) path runs its + own command via _query_from_form_data — an OAuth error there must + also surface as OAUTH2_REDIRECT, not a generic DataError.""" + from fastmcp import Client + + from superset.exceptions import OAuth2RedirectError + from superset.utils import json as utils_json + + chart_data_module = importlib.import_module( + "superset.mcp_service.chart.tool.get_chart_data" + ) + get_data_command_module = importlib.import_module( + "superset.commands.chart.data.get_data_command" + ) + + class RaisingChartDataCommand: + def __init__(self, query_context: object) -> None: + self.query_context = query_context + + def validate(self) -> None: + pass + + def run(self) -> dict[str, Any]: + raise OAuth2RedirectError( + "https://example.com/oauth", + "tab-1", + "https://example.com/redirect", + ) + + cached_form_data = utils_json.dumps( + { + "datasource_id": 1, + "datasource_type": "table", + "viz_type": "table", + "metrics": ["count"], + "groupby": ["gender"], + "row_limit": 10, + } + ) + monkeypatch.setattr( + chart_data_module, + "get_cached_form_data", + lambda *args, **kwargs: cached_form_data, + ) + monkeypatch.setattr( + chart_data_module, + "build_query_context_from_form_data", + lambda *args, **kwargs: object(), + ) + monkeypatch.setattr( + get_data_command_module, + "ChartDataCommand", + RaisingChartDataCommand, + ) + + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_chart_data", {"request": {"form_data_key": "cached-key"}} + ) + data = utils_json.loads(result.content[0].text) + + assert data["error_type"] == "OAUTH2_REDIRECT" + assert data["error_type"] != "DataError" + + # --------------------------------------------------------------------------- # Tests for _recommend_visualizations # --------------------------------------------------------------------------- diff --git a/tests/unit_tests/mcp_service/test_middleware.py b/tests/unit_tests/mcp_service/test_middleware.py index c3130a3a64e..d3a49866695 100644 --- a/tests/unit_tests/mcp_service/test_middleware.py +++ b/tests/unit_tests/mcp_service/test_middleware.py @@ -43,6 +43,7 @@ from superset.mcp_service.middleware import ( GlobalErrorHandlerMiddleware, RBACToolVisibilityMiddleware, ResponseSizeGuardMiddleware, + StructuredContentStripperMiddleware, ) @@ -1662,3 +1663,324 @@ class TestRBACToolVisibilityMiddleware: result = await middleware.on_list_tools(MagicMock(), call_next) assert result == tools + + +class TestGlobalErrorHandlerStatsMetrics: + """GlobalErrorHandlerMiddleware must NOT emit per-tool outcome + counters: it re-raises every failure as ToolError, which the outer + LoggingMiddleware catches and counts (classified via __cause__). + Emitting here as well would double-count raised errors.""" + + @pytest.mark.asyncio + async def test_no_stats_emitted_for_system_error(self) -> None: + middleware = GlobalErrorHandlerMiddleware() + context = MagicMock() + context.message.name = "execute_sql" + context.method = "tools/call" + call_next = AsyncMock(side_effect=OperationalError("db error", {}, Exception())) + + with ( + patch("superset.mcp_service.middleware.get_user_id", return_value=1), + patch("superset.mcp_service.middleware.event_logger"), + patch("superset.mcp_service.middleware.stats_logger_manager") as mock_stats, + pytest.raises(ToolError), + ): + await middleware.on_message(context, call_next) + + mock_stats.instance.incr.assert_not_called() + mock_stats.instance.timing.assert_not_called() + + @pytest.mark.asyncio + async def test_no_stats_emitted_for_user_error(self) -> None: + middleware = GlobalErrorHandlerMiddleware() + context = MagicMock() + context.message.name = "list_charts" + context.method = "tools/call" + call_next = AsyncMock(side_effect=ValueError("bad param")) + + with ( + patch("superset.mcp_service.middleware.get_user_id", return_value=1), + patch("superset.mcp_service.middleware.event_logger"), + patch("superset.mcp_service.middleware.stats_logger_manager") as mock_stats, + pytest.raises(ToolError), + ): + await middleware.on_message(context, call_next) + + mock_stats.instance.incr.assert_not_called() + mock_stats.instance.timing.assert_not_called() + + +class TestGlobalErrorHandlerErrorHook: + """Test that _handle_error invokes MCP_ERROR_HOOK for system-class + errors only, and never lets a raising hook affect the MCP response.""" + + @pytest.mark.asyncio + async def test_invokes_hook_for_system_error(self) -> None: + middleware = GlobalErrorHandlerMiddleware() + context = MagicMock() + context.message.name = "execute_sql" + context.method = "tools/call" + error = OperationalError("db error", {}, Exception()) + call_next = AsyncMock(side_effect=error) + mock_hook = MagicMock() + mock_flask_app = MagicMock() + mock_flask_app.config.get.return_value = mock_hook + + with ( + patch("superset.mcp_service.middleware.get_user_id", return_value=1), + patch("superset.mcp_service.middleware.event_logger"), + patch( + "superset.mcp_service.flask_singleton.get_flask_app", + return_value=mock_flask_app, + ), + pytest.raises(ToolError), + ): + await middleware.on_message(context, call_next) + + mock_hook.assert_called_once() + hook_error, hook_context = mock_hook.call_args[0] + assert hook_error is error + assert hook_context["tool_name"] == "execute_sql" + assert hook_context["error_type"] == "OperationalError" + + @pytest.mark.asyncio + async def test_does_not_invoke_hook_for_user_error(self) -> None: + middleware = GlobalErrorHandlerMiddleware() + context = MagicMock() + context.message.name = "list_charts" + context.method = "tools/call" + call_next = AsyncMock(side_effect=ValueError("bad param")) + mock_hook = MagicMock() + mock_flask_app = MagicMock() + mock_flask_app.config.get.return_value = mock_hook + + with ( + patch("superset.mcp_service.middleware.get_user_id", return_value=1), + patch("superset.mcp_service.middleware.event_logger"), + patch( + "superset.mcp_service.flask_singleton.get_flask_app", + return_value=mock_flask_app, + ), + pytest.raises(ToolError), + ): + await middleware.on_message(context, call_next) + + mock_hook.assert_not_called() + + @pytest.mark.asyncio + async def test_no_hook_configured_is_a_noop(self) -> None: + """Default config (MCP_ERROR_HOOK=None) must not raise.""" + middleware = GlobalErrorHandlerMiddleware() + context = MagicMock() + context.message.name = "execute_sql" + context.method = "tools/call" + call_next = AsyncMock(side_effect=OperationalError("db error", {}, Exception())) + mock_flask_app = MagicMock() + mock_flask_app.config.get.return_value = None + + with ( + patch("superset.mcp_service.middleware.get_user_id", return_value=1), + patch("superset.mcp_service.middleware.event_logger"), + patch( + "superset.mcp_service.flask_singleton.get_flask_app", + return_value=mock_flask_app, + ), + pytest.raises(ToolError), + ): + await middleware.on_message(context, call_next) + + @pytest.mark.asyncio + async def test_hook_exception_is_swallowed(self) -> None: + """A raising MCP_ERROR_HOOK must not affect the MCP response.""" + middleware = GlobalErrorHandlerMiddleware() + context = MagicMock() + context.message.name = "execute_sql" + context.method = "tools/call" + call_next = AsyncMock(side_effect=OperationalError("db error", {}, Exception())) + mock_hook = MagicMock(side_effect=RuntimeError("sentry unreachable")) + mock_flask_app = MagicMock() + mock_flask_app.config.get.return_value = mock_hook + + with ( + patch("superset.mcp_service.middleware.get_user_id", return_value=1), + patch("superset.mcp_service.middleware.event_logger"), + patch( + "superset.mcp_service.flask_singleton.get_flask_app", + return_value=mock_flask_app, + ), + pytest.raises(ToolError, match="Database error"), + ): + await middleware.on_message(context, call_next) + + mock_hook.assert_called_once() + + +class TestGlobalErrorHandlerErrorIdUsesCallId: + """Test that the generic 'Internal error' branch uses mcp_call_id + instead of a collision-prone f'err_{int(time.time())}' timestamp.""" + + @pytest.mark.asyncio + async def test_unexpected_error_uses_mcp_call_id(self) -> None: + from superset.mcp_service.middleware import _mcp_call_id_var + + middleware = GlobalErrorHandlerMiddleware() + context = MagicMock() + context.message.name = "list_charts" + context.method = "tools/call" + call_next = AsyncMock(side_effect=RuntimeError("boom")) + + token = _mcp_call_id_var.set("abc123deadbeef") + try: + with ( + patch("superset.mcp_service.middleware.get_user_id", return_value=1), + patch("superset.mcp_service.middleware.event_logger"), + pytest.raises(ToolError) as exc_info, + ): + await middleware.on_message(context, call_next) + finally: + _mcp_call_id_var.reset(token) + + assert "abc123deadbeef" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_falls_back_to_generated_id_when_no_call_id_set(self) -> None: + """When no mcp_call_id is in context (e.g. a non-tool-call message + path), fall back to a generated ID rather than raising.""" + from superset.mcp_service.middleware import _mcp_call_id_var + + middleware = GlobalErrorHandlerMiddleware() + context = MagicMock() + context.message.name = "list_charts" + context.method = "tools/call" + call_next = AsyncMock(side_effect=RuntimeError("boom")) + + token = _mcp_call_id_var.set(None) + try: + with ( + patch("superset.mcp_service.middleware.get_user_id", return_value=1), + patch("superset.mcp_service.middleware.event_logger"), + pytest.raises(ToolError, match="Error ID: err_"), + ): + await middleware.on_message(context, call_next) + finally: + _mcp_call_id_var.reset(token) + + +class TestStructuredContentStripperErrorHook: + """Test the last-resort MCP_ERROR_HOOK capture point in + StructuredContentStripperMiddleware.on_call_tool's except block.""" + + @pytest.mark.asyncio + async def test_invokes_hook_for_exception_bypassing_error_handler(self) -> None: + """A non-ToolError exception reaching this final catch means it + slipped past GlobalErrorHandlerMiddleware entirely — invoke the + hook here as the true last-resort capture point.""" + middleware = StructuredContentStripperMiddleware() + context = MagicMock() + context.message.name = "list_charts" + call_next = AsyncMock(side_effect=RuntimeError("boom")) + mock_hook = MagicMock() + mock_flask_app = MagicMock() + mock_flask_app.config.get.return_value = mock_hook + + with patch( + "superset.mcp_service.flask_singleton.get_flask_app", + return_value=mock_flask_app, + ): + result = await middleware.on_call_tool(context, call_next) + + assert result.content[0].text.startswith("Error:") + mock_hook.assert_called_once() + hook_error, hook_context = mock_hook.call_args[0] + assert isinstance(hook_error, RuntimeError) + assert hook_context["tool_name"] == "list_charts" + # The context contract: all keys always present, even on the + # last-resort path where user_id/duration_ms are unknown. + assert set(hook_context) == { + "tool_name", + "mcp_call_id", + "user_id", + "error_type", + "sanitized_message", + "duration_ms", + } + assert hook_context["user_id"] is None + assert hook_context["duration_ms"] is None + assert hook_context["error_type"] == "RuntimeError" + + @pytest.mark.asyncio + async def test_does_not_double_invoke_hook_for_tool_error(self) -> None: + """ToolError has already been classified and hooked by + GlobalErrorHandlerMiddleware — avoid double-reporting the same + failure to the error tracker.""" + middleware = StructuredContentStripperMiddleware() + context = MagicMock() + context.message.name = "list_charts" + call_next = AsyncMock(side_effect=ToolError("already handled")) + mock_hook = MagicMock() + mock_flask_app = MagicMock() + mock_flask_app.config.get.return_value = mock_hook + + with patch( + "superset.mcp_service.flask_singleton.get_flask_app", + return_value=mock_flask_app, + ): + result = await middleware.on_call_tool(context, call_next) + + assert result.content[0].text.startswith("Error:") + mock_hook.assert_not_called() + + @pytest.mark.asyncio + async def test_hostile_str_does_not_escape_last_resort_handler(self) -> None: + """The last-resort handler must never propagate — even when the + exception's own __str__ raises, it must fall back to the class + name rather than letting a formatting error escape to the MCP SDK + (the exact encoding failure this handler exists to prevent).""" + + class HostileStrError(Exception): + def __str__(self) -> str: + raise RuntimeError("hostile __str__") + + middleware = StructuredContentStripperMiddleware() + context = MagicMock() + context.message.name = "list_charts" + call_next = AsyncMock(side_effect=HostileStrError()) + mock_flask_app = MagicMock() + mock_flask_app.config.get.return_value = None + + with patch( + "superset.mcp_service.flask_singleton.get_flask_app", + return_value=mock_flask_app, + ): + result = await middleware.on_call_tool(context, call_next) + + assert result.content[0].text == "Error: HostileStrError" + + @pytest.mark.asyncio + async def test_client_facing_text_is_sanitized(self) -> None: + """An exception bypassing GlobalErrorHandlerMiddleware must not + leak raw internals to the client — the last-resort response text + goes through the same sanitizer as every other error path.""" + middleware = StructuredContentStripperMiddleware() + context = MagicMock() + context.message.name = "execute_sql" + # Connection string with embedded credentials — must be redacted. + call_next = AsyncMock( + side_effect=ValueError( + "connect failed: postgresql://user:s3cret@db.internal:5432/prod" + ) + ) + mock_flask_app = MagicMock() + mock_flask_app.config.get.return_value = None + + with patch( + "superset.mcp_service.flask_singleton.get_flask_app", + return_value=mock_flask_app, + ): + result = await middleware.on_call_tool(context, call_next) + + text = result.content[0].text + assert text.startswith("Error:") + assert "s3cret" not in text + assert "db.internal" not in text + assert "[REDACTED]" in text diff --git a/tests/unit_tests/mcp_service/test_middleware_logging.py b/tests/unit_tests/mcp_service/test_middleware_logging.py index 9e6893df856..ccab0f23e9c 100644 --- a/tests/unit_tests/mcp_service/test_middleware_logging.py +++ b/tests/unit_tests/mcp_service/test_middleware_logging.py @@ -34,6 +34,7 @@ import pytest from fastmcp.exceptions import ToolError from fastmcp.tools.tool import ToolResult from mcp import types as mt +from sqlalchemy.exc import OperationalError from superset.mcp_service.middleware import LoggingMiddleware @@ -753,3 +754,541 @@ class TestMiddlewareChainOrder: "on_list_tools must return [] on exception — " "ToolError cannot be encoded in a tools/list response." ) + + +class TestExtractErrorTypeFromResponse: + """Tests for LoggingMiddleware._extract_error_type_from_response().""" + + def test_extracts_error_type(self) -> None: + """Parses the error_type field out of a structured error response + instead of just sniffing for its presence.""" + result = ToolResult( + content=[ + mt.TextContent( + type="text", + text='{"error": "Chart 999 not found", "error_type": "not_found"}', + ) + ] + ) + assert ( + LoggingMiddleware._extract_error_type_from_response(result) == "not_found" + ) + + def test_returns_none_for_non_dict_payload(self) -> None: + result = ToolResult(content=[mt.TextContent(type="text", text="[1, 2, 3]")]) + assert LoggingMiddleware._extract_error_type_from_response(result) is None + + def test_returns_none_for_invalid_json(self) -> None: + result = ToolResult(content=[mt.TextContent(type="text", text="not json")]) + assert LoggingMiddleware._extract_error_type_from_response(result) is None + + def test_returns_none_for_missing_error_type(self) -> None: + result = ToolResult(content=[mt.TextContent(type="text", text='{"ok": true}')]) + assert LoggingMiddleware._extract_error_type_from_response(result) is None + + def test_returns_none_for_non_string_error_type(self) -> None: + result = ToolResult( + content=[mt.TextContent(type="text", text='{"error_type": 404}')] + ) + assert LoggingMiddleware._extract_error_type_from_response(result) is None + + def test_returns_none_for_empty_content(self) -> None: + assert ( + LoggingMiddleware._extract_error_type_from_response(ToolResult(content=[])) + is None + ) + + +class TestOnCallToolErrorTypeExtraction: + """Tests that on_call_tool surfaces the parsed error_type from + structured error responses (ChartError, DashboardError, etc.), rather + than discarding it after the substring sniff.""" + + @patch("superset.mcp_service.middleware.event_logger") + @patch("superset.mcp_service.middleware.get_user_id", return_value=42) + @pytest.mark.asyncio + async def test_logs_error_type_from_structured_response( + self, mock_get_user_id, mock_event_logger + ) -> None: + middleware = LoggingMiddleware() + ctx = _make_context(name="get_chart_info") + + error_json = ( + '{"error": "Chart 999999 not found",' + ' "error_type": "not_found",' + ' "timestamp": "2026-04-09T00:00:00Z"}' + ) + error_result = ToolResult( + content=[mt.TextContent(type="text", text=error_json)] + ) + call_next = AsyncMock(return_value=error_result) + + await middleware.on_call_tool(ctx, call_next) + + payload = mock_event_logger.log.call_args[1]["curated_payload"] + assert payload["error_type"] == "not_found" + + @patch("superset.mcp_service.middleware.event_logger") + @patch("superset.mcp_service.middleware.get_user_id", return_value=42) + @pytest.mark.asyncio + async def test_exception_error_type_takes_precedence( + self, mock_get_user_id, mock_event_logger + ) -> None: + """When call_next raises, error_type comes from the exception class + (there is no structured response to parse).""" + middleware = LoggingMiddleware() + ctx = _make_context(name="execute_sql") + call_next = AsyncMock(side_effect=ValueError("boom")) + + with pytest.raises(ValueError, match="boom"): + await middleware.on_call_tool(ctx, call_next) + + payload = mock_event_logger.log.call_args[1]["curated_payload"] + assert payload["error_type"] == "ValueError" + + +class TestOnCallToolStatsMetrics: + """Tests that on_call_tool emits per-tool StatsD counters and timing, + mirroring the success/warning/error split in base_api.py. This is the + single emission point — GlobalErrorHandlerMiddleware must not emit.""" + + @patch("superset.mcp_service.middleware.stats_logger_manager") + @patch("superset.mcp_service.middleware.event_logger") + @patch("superset.mcp_service.middleware.get_user_id", return_value=42) + @pytest.mark.asyncio + async def test_emits_success_counter_and_timing( + self, mock_get_user_id, mock_event_logger, mock_stats + ) -> None: + middleware = LoggingMiddleware() + ctx = _make_context(name="list_charts") + call_next = AsyncMock(return_value="ok") + + await middleware.on_call_tool(ctx, call_next) + + mock_stats.instance.incr.assert_called_once_with("mcp.tool.list_charts.success") + mock_stats.instance.timing.assert_called_once() + timing_key, timing_value = mock_stats.instance.timing.call_args[0] + assert timing_key == "mcp.tool.list_charts.time" + assert isinstance(timing_value, int) + assert timing_value >= 0 + + @patch("superset.mcp_service.middleware.stats_logger_manager") + @patch("superset.mcp_service.middleware.event_logger") + @patch("superset.mcp_service.middleware.get_user_id", return_value=42) + @pytest.mark.asyncio + async def test_emits_warning_counter_for_raised_user_error( + self, mock_get_user_id, mock_event_logger, mock_stats + ) -> None: + """A raised user-class error (ValueError) counts as warning, + matching base_api.py's 4xx handling.""" + middleware = LoggingMiddleware() + ctx = _make_context(name="execute_sql") + call_next = AsyncMock(side_effect=ValueError("boom")) + + with pytest.raises(ValueError, match="boom"): + await middleware.on_call_tool(ctx, call_next) + + mock_stats.instance.incr.assert_called_once_with("mcp.tool.execute_sql.warning") + + @patch("superset.mcp_service.middleware.stats_logger_manager") + @patch("superset.mcp_service.middleware.event_logger") + @patch("superset.mcp_service.middleware.get_user_id", return_value=42) + @pytest.mark.asyncio + async def test_emits_error_counter_for_raised_system_error( + self, mock_get_user_id, mock_event_logger, mock_stats + ) -> None: + """A raised system-class error (RuntimeError) counts as error.""" + middleware = LoggingMiddleware() + ctx = _make_context(name="execute_sql") + call_next = AsyncMock(side_effect=RuntimeError("db down")) + + with pytest.raises(RuntimeError, match="db down"): + await middleware.on_call_tool(ctx, call_next) + + mock_stats.instance.incr.assert_called_once_with("mcp.tool.execute_sql.error") + + @patch("superset.mcp_service.middleware.stats_logger_manager") + @patch("superset.mcp_service.middleware.event_logger") + @patch("superset.mcp_service.middleware.get_user_id", return_value=42) + @pytest.mark.asyncio + async def test_unwraps_tool_error_cause_for_classification( + self, mock_get_user_id, mock_event_logger, mock_stats + ) -> None: + """A ToolError wrapping a system error (as raised by + GlobalErrorHandlerMiddleware via ``raise ... from error``) must be + classified by its __cause__, not by the ToolError wrapper.""" + middleware = LoggingMiddleware() + ctx = _make_context(name="execute_sql") + wrapped = ToolError("Database error in execute_sql") + wrapped.__cause__ = OperationalError("db error", {}, Exception()) + call_next = AsyncMock(side_effect=wrapped) + + with pytest.raises(ToolError): + await middleware.on_call_tool(ctx, call_next) + + mock_stats.instance.incr.assert_called_once_with("mcp.tool.execute_sql.error") + payload = mock_event_logger.log.call_args[1]["curated_payload"] + assert payload["error_type"] == "OperationalError" + + @patch("superset.mcp_service.middleware.stats_logger_manager") + @patch("superset.mcp_service.middleware.event_logger") + @patch("superset.mcp_service.middleware.get_user_id", return_value=42) + @pytest.mark.asyncio + async def test_emits_error_counter_for_structured_error_response( + self, mock_get_user_id, mock_event_logger, mock_stats + ) -> None: + """A tool that returns an error schema (rather than raising) must + still increment the error counter, not the success counter.""" + middleware = LoggingMiddleware() + ctx = _make_context(name="get_chart_info") + error_result = ToolResult( + content=[ + mt.TextContent( + type="text", + text='{"error": "not found", "error_type": "not_found"}', + ) + ] + ) + call_next = AsyncMock(return_value=error_result) + + await middleware.on_call_tool(ctx, call_next) + + mock_stats.instance.incr.assert_called_once_with( + "mcp.tool.get_chart_info.error" + ) + + @patch("superset.mcp_service.middleware.stats_logger_manager") + @patch("superset.mcp_service.middleware.event_logger") + @patch("superset.mcp_service.middleware.get_user_id", return_value=42) + @pytest.mark.asyncio + async def test_metric_uses_resolved_tool_name_for_call_tool_proxy( + self, mock_get_user_id, mock_event_logger, mock_stats + ) -> None: + """When invoked via the call_tool search proxy, the metric key + must use the real tool name, not the literal 'call_tool'.""" + middleware = LoggingMiddleware() + ctx = _make_context( + name="call_tool", + params={"name": "list_datasets", "arguments": {}}, + ) + call_next = AsyncMock(return_value="ok") + + await middleware.on_call_tool(ctx, call_next) + + mock_stats.instance.incr.assert_called_once_with( + "mcp.tool.list_datasets.success" + ) + + @patch("superset.mcp_service.middleware.logger") + @patch("superset.mcp_service.middleware.stats_logger_manager") + @patch("superset.mcp_service.middleware.event_logger") + @patch("superset.mcp_service.middleware.get_user_id", return_value=42) + @pytest.mark.asyncio + async def test_metrics_failure_does_not_mask_tool_result( + self, mock_get_user_id, mock_event_logger, mock_stats, mock_logger + ) -> None: + """A failing stats backend must not turn a successful tool call + into an error — metrics are a side effect only. The swallowed + error must still be logged, not silently dropped.""" + mock_stats.instance.incr.side_effect = RuntimeError("metrics backend down") + middleware = LoggingMiddleware() + ctx = _make_context(name="list_charts") + call_next = AsyncMock(return_value="ok") + + result = await middleware.on_call_tool(ctx, call_next) + + assert result == "ok" + warning_messages = [c.args[0] for c in mock_logger.warning.call_args_list] + assert any("Failed to emit MCP tool metrics" in m for m in warning_messages) + + @patch("superset.mcp_service.middleware.logger") + @patch("superset.mcp_service.middleware.stats_logger_manager") + @patch("superset.mcp_service.middleware.event_logger") + @patch("superset.mcp_service.middleware.get_user_id", return_value=42) + @pytest.mark.asyncio + async def test_metrics_failure_does_not_mask_tool_exception( + self, mock_get_user_id, mock_event_logger, mock_stats, mock_logger + ) -> None: + """A failing stats backend must not replace the tool's real + exception with its own error in the finally block.""" + mock_stats.instance.incr.side_effect = RuntimeError("metrics backend down") + middleware = LoggingMiddleware() + ctx = _make_context(name="execute_sql") + call_next = AsyncMock(side_effect=ValueError("boom")) + + with pytest.raises(ValueError, match="boom"): + await middleware.on_call_tool(ctx, call_next) + + warning_messages = [c.args[0] for c in mock_logger.warning.call_args_list] + assert any("Failed to emit MCP tool metrics" in m for m in warning_messages) + + +class TestResolveMetricToolName: + """Tests for the StatsD metric-key validation in + _resolve_metric_tool_name — client-controlled tool names must never + reach a metric key unvalidated.""" + + @pytest.mark.asyncio + async def test_registered_tool_name_is_used(self) -> None: + """A name that resolves in the FastMCP tool registry is used.""" + ctx = _make_context() + ctx.fastmcp_context.fastmcp.get_tool = AsyncMock(return_value=MagicMock()) + + result = await LoggingMiddleware._resolve_metric_tool_name( + ctx, "call_tool", "list_datasets" + ) + + assert result == "list_datasets" + ctx.fastmcp_context.fastmcp.get_tool.assert_awaited_once_with("list_datasets") + + @pytest.mark.asyncio + async def test_unregistered_proxy_name_falls_back_to_call_tool(self) -> None: + """A bogus call_tool proxy name must not mint a new metric series.""" + ctx = _make_context() + ctx.fastmcp_context.fastmcp.get_tool = AsyncMock(return_value=None) + + result = await LoggingMiddleware._resolve_metric_tool_name( + ctx, "call_tool", "totally_fake_tool_9000" + ) + + assert result == "call_tool" + + @pytest.mark.asyncio + async def test_unregistered_direct_name_falls_back_to_unknown(self) -> None: + """A bogus direct tool name must not mint a new metric series.""" + ctx = _make_context() + ctx.fastmcp_context.fastmcp.get_tool = AsyncMock(return_value=None) + + result = await LoggingMiddleware._resolve_metric_tool_name( + ctx, "totally_fake_tool_9000", None + ) + + assert result == "unknown" + + @pytest.mark.asyncio + async def test_registry_raise_treated_as_unregistered(self) -> None: + """get_tool raising (NotFoundError-style) means unregistered.""" + ctx = _make_context() + ctx.fastmcp_context.fastmcp.get_tool = AsyncMock( + side_effect=KeyError("not found") + ) + + result = await LoggingMiddleware._resolve_metric_tool_name( + ctx, "call_tool", "totally_fake_tool_9000" + ) + + assert result == "call_tool" + + @pytest.mark.asyncio + async def test_injection_shaped_name_rejected_without_registry(self) -> None: + """With no registry reachable (mocked context), StatsD metadata + characters must still never reach the metric key.""" + ctx = _make_context() # plain MagicMock — get_tool is not awaitable + + result = await LoggingMiddleware._resolve_metric_tool_name( + ctx, "call_tool", "evil\nfake.metric:1|c" + ) + + assert result == "call_tool" + + @pytest.mark.asyncio + async def test_injection_shaped_name_rejected_via_registry(self) -> None: + """The production path: a reachable registry that does not know + the hostile name must fall back to the constant — the metadata + characters never reach the metric key.""" + ctx = _make_context() + ctx.fastmcp_context.fastmcp.get_tool = AsyncMock(return_value=None) + + result = await LoggingMiddleware._resolve_metric_tool_name( + ctx, "call_tool", "evil\nfake.metric:1|c" + ) + + assert result == "call_tool" + + @pytest.mark.asyncio + async def test_overlong_name_rejected_without_registry(self) -> None: + ctx = _make_context() + + result = await LoggingMiddleware._resolve_metric_tool_name(ctx, "x" * 500, None) + + assert result == "unknown" + + @pytest.mark.asyncio + async def test_missing_name_returns_unknown(self) -> None: + ctx = _make_context() + + assert ( + await LoggingMiddleware._resolve_metric_tool_name(ctx, None, None) + == "unknown" + ) + + +class TestChainLevelStatsMetrics: + """Drive failures through the REAL middleware chain from + build_middleware_list() and assert the exact set of stats calls — + pins that per-tool outcome counters are emitted exactly once per call + (no double-counting between LoggingMiddleware and + GlobalErrorHandlerMiddleware).""" + + @patch("superset.mcp_service.middleware.stats_logger_manager") + @patch("superset.mcp_service.middleware.event_logger") + @patch("superset.mcp_service.middleware.get_user_id", return_value=42) + @pytest.mark.asyncio + async def test_raised_user_error_counts_warning_exactly_once( + self, mock_get_user_id, mock_event_logger, mock_stats + ) -> None: + from superset.mcp_service.server import build_middleware_list + + middleware_list = build_middleware_list() + + async def failing_tool(context: Any) -> Any: + raise ValueError("chart not found") + + chain = failing_tool + for mw in reversed(middleware_list): + chain = partial(mw, call_next=chain) + + ctx = _make_context(name="get_chart_info") + result = await chain(ctx) + + # Stripper (outermost) converts the ToolError to a safe ToolResult + assert isinstance(result, ToolResult) + assert result.content[0].text.startswith("Error:") + + # Exactly ONE outcome counter for the whole chain: ValueError is a + # user error, classified via the ToolError __cause__ unwrap. + incr_keys = [c.args[0] for c in mock_stats.instance.incr.call_args_list] + assert incr_keys == ["mcp.tool.get_chart_info.warning"] + timing_keys = [c.args[0] for c in mock_stats.instance.timing.call_args_list] + assert timing_keys == ["mcp.tool.get_chart_info.time"] + + @patch("superset.mcp_service.middleware.stats_logger_manager") + @patch("superset.mcp_service.middleware.event_logger") + @patch("superset.mcp_service.middleware.get_user_id", return_value=42) + @pytest.mark.asyncio + async def test_raised_system_error_counts_error_exactly_once( + self, mock_get_user_id, mock_event_logger, mock_stats + ) -> None: + from superset.mcp_service.server import build_middleware_list + + middleware_list = build_middleware_list() + + async def failing_tool(context: Any) -> Any: + raise RuntimeError("infrastructure down") + + chain = failing_tool + for mw in reversed(middleware_list): + chain = partial(mw, call_next=chain) + + ctx = _make_context(name="execute_sql") + result = await chain(ctx) + + assert isinstance(result, ToolResult) + assert result.content[0].text.startswith("Error:") + + incr_keys = [c.args[0] for c in mock_stats.instance.incr.call_args_list] + assert incr_keys == ["mcp.tool.execute_sql.error"] + + @patch("superset.mcp_service.middleware.stats_logger_manager") + @patch("superset.mcp_service.middleware.event_logger") + @patch("superset.mcp_service.middleware.get_user_id", return_value=42) + @pytest.mark.asyncio + async def test_success_counts_success_exactly_once( + self, mock_get_user_id, mock_event_logger, mock_stats + ) -> None: + from superset.mcp_service.server import build_middleware_list + + middleware_list = build_middleware_list() + + async def ok_tool(context: Any) -> Any: + return ToolResult(content=[mt.TextContent(type="text", text="ok")]) + + chain = ok_tool + for mw in reversed(middleware_list): + chain = partial(mw, call_next=chain) + + ctx = _make_context(name="list_charts") + await chain(ctx) + + incr_keys = [c.args[0] for c in mock_stats.instance.incr.call_args_list] + assert incr_keys == ["mcp.tool.list_charts.success"] + + @patch("superset.mcp_service.middleware.stats_logger_manager") + @patch("superset.mcp_service.middleware.event_logger") + @patch("superset.mcp_service.middleware.get_user_id", return_value=42) + @pytest.mark.asyncio + async def test_structured_error_response_counts_error_exactly_once( + self, mock_get_user_id, mock_event_logger, mock_stats + ) -> None: + """A tool that returns an error schema (no raise) counts one + .error through the real chain — not .success, not doubled.""" + from superset.mcp_service.server import build_middleware_list + + middleware_list = build_middleware_list() + + async def error_returning_tool(context: Any) -> Any: + return ToolResult( + content=[ + mt.TextContent( + type="text", + text='{"error": "not found", "error_type": "not_found"}', + ) + ] + ) + + chain = error_returning_tool + for mw in reversed(middleware_list): + chain = partial(mw, call_next=chain) + + ctx = _make_context(name="get_chart_info") + await chain(ctx) + + incr_keys = [c.args[0] for c in mock_stats.instance.incr.call_args_list] + assert incr_keys == ["mcp.tool.get_chart_info.error"] + + +class TestAppContextFixForAuditRows: + """Regression tests for the has_app_context() skip bug: the DB audit + row (event_logger.log) must never be silently dropped just because + the per-tool app context has already exited by the time the + middleware finally block runs. Both on_call_tool and on_message must + wrap the log call in _get_app_context_manager() rather than + conditionally skipping it.""" + + @patch("superset.mcp_service.middleware.event_logger") + @patch("superset.mcp_service.middleware.get_user_id", return_value=42) + @patch("superset.mcp_service.middleware._get_app_context_manager") + @pytest.mark.asyncio + async def test_on_call_tool_uses_app_context_manager( + self, mock_get_app_context_manager, mock_get_user_id, mock_event_logger + ) -> None: + from contextlib import nullcontext + + mock_get_app_context_manager.return_value = nullcontext() + middleware = LoggingMiddleware() + ctx = _make_context(name="list_charts") + call_next = AsyncMock(return_value="ok") + + await middleware.on_call_tool(ctx, call_next) + + mock_get_app_context_manager.assert_called_once() + mock_event_logger.log.assert_called_once() + + @patch("superset.mcp_service.middleware.event_logger") + @patch("superset.mcp_service.middleware.get_user_id", return_value=42) + @patch("superset.mcp_service.middleware._get_app_context_manager") + @pytest.mark.asyncio + async def test_on_message_uses_app_context_manager( + self, mock_get_app_context_manager, mock_get_user_id, mock_event_logger + ) -> None: + from contextlib import nullcontext + + mock_get_app_context_manager.return_value = nullcontext() + middleware = LoggingMiddleware() + ctx = _make_context(method="resources/read", name="instance/metadata") + call_next = AsyncMock(return_value="resource_data") + + await middleware.on_message(ctx, call_next) + + mock_get_app_context_manager.assert_called_once() + mock_event_logger.log.assert_called_once()