You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by ru...@apache.org on 2026/10/20 03:59:37 UTC
(superset) branch master updated: feat(mcp): add caller-scoped PNG chart previews (#44558)
This is an automated email from the ASF dual-hosted git repository.
rusackas pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new 1e18d01d76d feat(mcp): add caller-scoped PNG chart previews (#44558)
1e18d01d76d is described below
commit 1e18d01d76d6736dd3c22770f3840b8378776fd7
Author: Deepak Jain (DJ) <de...@users.noreply.github.com>
AuthorDate: Mon Oct 19 20:59:29 2026 -0700
feat(mcp): add caller-scoped PNG chart previews (#44558)
Signed-off-by: Deepak Jain <de...@gmail.com>
---
superset/mcp_service/README.md | 41 ++
superset/mcp_service/chart/schemas.py | 27 +-
.../mcp_service/chart/tool/get_chart_preview.py | 141 ++++++-
superset/mcp_service/constants.py | 1 +
superset/mcp_service/mcp_config.py | 3 +
superset/mcp_service/middleware.py | 27 ++
superset/utils/webdriver.py | 4 +-
.../chart/tool/test_get_chart_preview.py | 413 +++++++++++++++++++++
tests/unit_tests/mcp_service/test_middleware.py | 5 +-
9 files changed, 652 insertions(+), 10 deletions(-)
diff --git a/superset/mcp_service/README.md b/superset/mcp_service/README.md
index 2079ba94884..82254ae4472 100644
--- a/superset/mcp_service/README.md
+++ b/superset/mcp_service/README.md
@@ -803,3 +803,44 @@ curl http://localhost:5008/health
1. **Database Connection Errors**: Ensure the MCP service has the same database credentials as Superset
2. **Authentication Failures**: Verify `MCP_DEV_USERNAME` matches an existing Superset user
3. **Screenshot Generation Fails**: Check WebDriver configuration and ensure Chrome/Firefox is available in the container
+
+## Rendered PNG chart previews
+
+Call `get_chart_preview` with `{"identifier": 123, "format": "png"}` to
+render a saved chart with Superset's native renderer. The response content has
+`type: "png"`, `mime_type: "image/png"`, base64 `data`, and the actual image
+`width` and `height`. The default preview format remains ASCII.
+
+PNG requires a non-guest authenticated user and the normal chart/data access
+permissions. When granular export controls are enabled, image export permission
+is also required. Unsaved chart state (`form_data_key`) and extra dashboard
+filters (`extra_form_data`) are not supported by this format and return an
+explicit error. Use the other preview formats for those requests.
+
+The MCP host needs Playwright with Chromium installed and a reachable
+`WEBDRIVER_BASEURL`. Each request creates a separate browser, authenticates it as
+the calling user, and closes it after capture; images are not stored in the shared
+thumbnail cache. Requested viewport dimensions must be between 64 and 4096 pixels.
+Capture uses the configured screenshot timeouts. Cancelling an MCP request does
+not immediately interrupt its rendering worker.
+
+PNG responses have a separate serialized byte ceiling,
+`MCP_RESPONSE_SIZE_CONFIG["png_max_bytes"]` (default: 1,000,000 bytes),
+including base64 image data and preview metadata. The guard rejects larger
+responses with an error asking for smaller viewport dimensions or a URL preview.
+A viewport within the 64..4096 range can still exceed this byte ceiling,
+particularly with higher configured pixel density or detailed charts. Images are
+not truncated or silently rescaled. Other preview formats retain the normal
+`max_bytes` budget (default: 50,000 bytes). Operators can override `png_max_bytes`
+in `superset_config.py`; disabling the guard or explicitly excluding the tool
+also disables its image budget.
+
+PNG capture uses a dedicated two-worker executor per MCP process, separate from
+the pool used for transport authentication. Additional captures wait for a
+worker; cancellation does not release a running capture's worker until it exits.
+
+The executor lives for the MCP process lifetime. Start each MCP worker after any
+process fork, and do not fork a worker after it has begun rendering. Python waits
+for running captures when the process exits, so shutdown can wait for a browser
+operation to finish or reach its configured timeout. Cancellation does not
+terminate that browser operation or free its slot early.
diff --git a/superset/mcp_service/chart/schemas.py b/superset/mcp_service/chart/schemas.py
index 547c5ef4ea6..0b9ac7e58cc 100644
--- a/superset/mcp_service/chart/schemas.py
+++ b/superset/mcp_service/chart/schemas.py
@@ -6804,22 +6804,24 @@ class GetChartPreviewRequest(QueryCacheControl):
)
return self
- format: Literal["url", "ascii", "table", "vega_lite"] = Field(
+ format: Literal["url", "ascii", "table", "vega_lite", "png"] = Field(
default="ascii",
description=(
"Preview format: 'ascii' for text art (default), "
"'url' for explore link, "
"'table' for data table, "
- "'vega_lite' for interactive JSON specification"
+ "'vega_lite' for interactive JSON specification, "
+ "'png' for a rendered saved chart (authenticated users only; "
+ "no unsaved state or extra filters)"
),
)
width: int | None = Field(
default=800,
- description="Preview width in pixels (for url and vega_lite formats)",
+ description="Preview width in pixels (for url, vega_lite, and png formats)",
)
height: int | None = Field(
default=600,
- description="Preview height in pixels (for url and vega_lite formats)",
+ description="Preview height in pixels (for url, vega_lite, and png formats)",
)
ascii_width: int | None = Field(
default=80, description="ASCII chart width in characters (for ascii format)"
@@ -6838,6 +6840,16 @@ class GetChartPreviewRequest(QueryCacheControl):
)
+class PNGPreview(BaseModel):
+ """Rendered chart image, encoded without a shared thumbnail cache."""
+
+ type: Literal["png"] = "png"
+ mime_type: Literal["image/png"] = "image/png"
+ data: str = Field(..., description="Base64-encoded PNG image")
+ width: int = Field(..., description="Rendered PNG width in pixels")
+ height: int = Field(..., description="Rendered PNG height in pixels")
+
+
# Discriminated union preview formats for type safety
class URLPreview(BaseModel):
"""URL-based preview format."""
@@ -6901,7 +6913,12 @@ class TablePreview(BaseModel):
# Modern discriminated union using | syntax
ChartPreviewContent = Annotated[
- URLPreview | InteractivePreview | ASCIIPreview | VegaLitePreview | TablePreview,
+ URLPreview
+ | InteractivePreview
+ | ASCIIPreview
+ | VegaLitePreview
+ | TablePreview
+ | PNGPreview,
Field(discriminator="type"),
]
diff --git a/superset/mcp_service/chart/tool/get_chart_preview.py b/superset/mcp_service/chart/tool/get_chart_preview.py
index e2443e503e0..9f9aed4fd0c 100644
--- a/superset/mcp_service/chart/tool/get_chart_preview.py
+++ b/superset/mcp_service/chart/tool/get_chart_preview.py
@@ -19,15 +19,29 @@
MCP tool: get_chart_preview
"""
+import asyncio
+import base64
import logging
+from collections.abc import Callable
+from concurrent.futures import ThreadPoolExecutor
+from contextvars import copy_context
+from io import BytesIO
from typing import Any, Dict, List, Protocol
from fastmcp import Context
+from flask import g
+from PIL import Image
from sqlalchemy.exc import SQLAlchemyError
from superset_core.mcp.decorators import tool, ToolAnnotations
+from superset import security_manager
from superset.commands.exceptions import CommandException
-from superset.exceptions import OAuth2Error, OAuth2RedirectError, SupersetException
+from superset.exceptions import (
+ OAuth2Error,
+ OAuth2RedirectError,
+ SupersetException,
+ SupersetSecurityException,
+)
from superset.extensions import db, event_logger
from superset.mcp_service import guest_scope
from superset.mcp_service.chart.ascii_charts import (
@@ -70,6 +84,7 @@ from superset.mcp_service.chart.schemas import (
GetChartPreviewRequest,
InteractivePreview,
PerformanceMetadata,
+ PNGPreview,
TablePreview,
URLPreview,
VegaLitePreview,
@@ -84,6 +99,9 @@ from superset.utils.core import DatasourceType
logger = logging.getLogger(__name__)
+# Keep browser captures separate from the default pool used for transport auth.
+_PNG_EXECUTOR = ThreadPoolExecutor(max_workers=2, thread_name_prefix="mcp-png")
+
def _finalize_response(
response: ChartPreview | ChartError,
@@ -1188,6 +1206,116 @@ class PreviewFormatGenerator:
return strategy.generate()
+async def _run_png_render(
+ render: Callable[[], PNGPreview | ChartError],
+) -> PNGPreview | ChartError:
+ """Run the blocking PNG render in a worker thread with error mapping.
+
+ Chart-level access denials raised inside render() surface as Forbidden,
+ not as a rendering failure. Browser errors may contain URLs or page data,
+ so those stay server-side.
+ """
+ try:
+ context = copy_context()
+ return await asyncio.get_running_loop().run_in_executor(
+ _PNG_EXECUTOR, context.run, render
+ )
+ except SupersetSecurityException:
+ return ChartError(error="Chart access denied", error_type="Forbidden")
+ except Exception:
+ logger.exception("PNG chart rendering failed")
+ return ChartError(error="Chart rendering failed", error_type="RenderError")
+
+
+async def _generate_png_preview(
+ chart_id: int, request: GetChartPreviewRequest
+) -> PNGPreview | ChartError:
+ """Render saved charts as the caller in an isolated browser and app context."""
+ if (
+ not chart_id
+ or request.form_data_key
+ or request.extra_form_data
+ or guest_scope.is_guest_read()
+ ):
+ return ChartError(
+ error="PNG previews require a saved chart and a non-guest user, "
+ "without unsaved state or extra filters.",
+ error_type="UnsupportedFormat",
+ )
+ user_id = getattr(getattr(g, "user", None), "id", None)
+ if not isinstance(user_id, int):
+ return ChartError(error="Authentication required", error_type="Forbidden")
+ width = 800 if request.width is None else request.width
+ height = 600 if request.height is None else request.height
+ if not (64 <= width <= 4096 and 64 <= height <= 4096):
+ return ChartError(
+ error="PNG dimensions must be between 64 and 4096 pixels.",
+ error_type="ValidationError",
+ )
+
+ def render() -> PNGPreview | ChartError:
+ from superset import is_feature_enabled
+ from superset.mcp_service.auth import _mcp_tool_call_context
+ from superset.utils.core import override_user
+ from superset.utils.screenshots import validate_screenshot_image
+ from superset.utils.urls import get_url_path
+ from superset.utils.webdriver import (
+ _PlaywrightBrowserManager,
+ WebDriverPlaywright,
+ )
+
+ # Resolve ORM objects in a fresh context; neither the request session nor
+ # Flask's mutable g is shared with the rendering worker.
+ with _mcp_tool_call_context():
+ user = security_manager.find_user(id=user_id)
+ if user is None or not user.is_active:
+ return ChartError(
+ error="Authentication required", error_type="Forbidden"
+ )
+ with override_user(user):
+ chart = find_chart_by_identifier(chart_id)
+ if chart is None:
+ return ChartError(error="Chart not found", error_type="NotFound")
+ security_manager.raise_for_access(chart=chart)
+ if is_feature_enabled(
+ "GRANULAR_EXPORT_CONTROLS"
+ ) and not security_manager.can_access("can_export_image", "Superset"):
+ return ChartError(
+ error="Image export is not permitted", error_type="Forbidden"
+ )
+ url = get_url_path(
+ "Superset.slice", slice_id=chart_id, standalone="true"
+ )
+ manager = _PlaywrightBrowserManager()
+ try:
+ driver = WebDriverPlaywright(
+ "",
+ (width, height),
+ require_complete_capture=True,
+ browser_manager=manager,
+ )
+ image = driver.get_screenshot(url, "chart-container", user=user)
+ if (
+ image is None
+ or not image.startswith(b"\x89PNG\r\n\x1a\n")
+ or validate_screenshot_image(image)
+ ):
+ return ChartError(
+ error="Chart rendering failed", error_type="RenderError"
+ )
+ with Image.open(BytesIO(image)) as rendered:
+ image_width, image_height = rendered.size
+ return PNGPreview(
+ data=base64.b64encode(image).decode("ascii"),
+ width=image_width,
+ height=image_height,
+ )
+ finally:
+ manager._cleanup()
+
+ return await _run_png_render(render)
+
+
async def _get_chart_preview_internal( # noqa: C901
request: GetChartPreviewRequest,
ctx: Context,
@@ -1477,7 +1605,11 @@ async def _get_chart_preview_internal( # noqa: C901
action="mcp.get_chart_preview.preview_generation"
):
preview_generator = PreviewFormatGenerator(chart, request)
- content = preview_generator.generate()
+ content = (
+ await _generate_png_preview(chart.id, request)
+ if request.format == "png"
+ else preview_generator.generate()
+ )
if isinstance(content, ChartError):
await ctx.error(
@@ -1584,7 +1716,10 @@ async def get_chart_preview(
) -> ChartPreview | ChartError:
"""Get chart preview by ID or UUID.
- Returns preview URL or formatted content (ascii, table, vega_lite).
+ Returns preview URL or formatted content (ascii, table, vega_lite, png).
+
+ PNG renders saved charts as the calling non-guest user; unsaved state and
+ extra filters are unsupported.
Pass extra_form_data (e.g. a dashboard's active native filters) to render
the preview over the filtered data rather than the full dataset.
diff --git a/superset/mcp_service/constants.py b/superset/mcp_service/constants.py
index 0b435dba815..69f90a09c9b 100644
--- a/superset/mcp_service/constants.py
+++ b/superset/mcp_service/constants.py
@@ -29,6 +29,7 @@ MAX_PAGE_SIZE = 100 # Maximum allowed page_size to prevent oversized responses
# Response size guard defaults
DEFAULT_MAX_RESPONSE_BYTES = 50_000 # ~50KB preserves the former 25K-token guard
+DEFAULT_MAX_PNG_RESPONSE_BYTES = 1_000_000 # Includes base64 and preview metadata
DEFAULT_WARN_THRESHOLD_PCT = 80 # Log warnings above 80% of limit
# Phase 2 list-field truncation cap; matches MAX_PAGE_SIZE
DEFAULT_MAX_LIST_ITEMS = 100
diff --git a/superset/mcp_service/mcp_config.py b/superset/mcp_service/mcp_config.py
index 6be34dab5e5..466ab282ee2 100644
--- a/superset/mcp_service/mcp_config.py
+++ b/superset/mcp_service/mcp_config.py
@@ -28,6 +28,7 @@ from superset.constants import CHANGE_ME_GUEST_TOKEN_JWT_SECRET
from superset.mcp_service.composite_token_verifier import CompositeTokenVerifier
from superset.mcp_service.constants import (
DEFAULT_MAX_LIST_ITEMS,
+ DEFAULT_MAX_PNG_RESPONSE_BYTES,
DEFAULT_MAX_RESPONSE_BYTES,
DEFAULT_WARN_THRESHOLD_PCT,
)
@@ -475,6 +476,7 @@ MCP_CACHE_CONFIG: dict[str, Any] = {
# --------------
# - enabled: Toggle the guard on/off (default: True)
# - max_bytes: Maximum serialized response size in bytes (default: 50,000)
+# - png_max_bytes: Maximum serialized PNG preview response size (default: 1,000,000)
# - excluded_tools: Tools to skip checking (e.g., streaming tools)
# - warn_threshold_pct: Log warnings above this % of limit (default: 80%)
# - max_list_items: Cap applied to list fields (e.g. ``charts``,
@@ -494,6 +496,7 @@ MCP_CACHE_CONFIG: dict[str, Any] = {
MCP_RESPONSE_SIZE_CONFIG: dict[str, Any] = {
"enabled": True, # Enabled by default to protect LLM clients
"max_bytes": DEFAULT_MAX_RESPONSE_BYTES,
+ "png_max_bytes": DEFAULT_MAX_PNG_RESPONSE_BYTES,
"warn_threshold_pct": DEFAULT_WARN_THRESHOLD_PCT,
"max_list_items": DEFAULT_MAX_LIST_ITEMS,
"excluded_tools": [ # Tools to skip size checking
diff --git a/superset/mcp_service/middleware.py b/superset/mcp_service/middleware.py
index 08776ec69c0..fca8966d66b 100644
--- a/superset/mcp_service/middleware.py
+++ b/superset/mcp_service/middleware.py
@@ -59,6 +59,7 @@ from superset.mcp_service.auth import (
from superset.mcp_service.constants import (
CONNECTION_ERROR_TYPES,
DEFAULT_MAX_LIST_ITEMS,
+ DEFAULT_MAX_PNG_RESPONSE_BYTES,
DEFAULT_MAX_RESPONSE_BYTES,
DEFAULT_WARN_THRESHOLD_PCT,
)
@@ -1545,6 +1546,7 @@ class ResponseSizeGuardMiddleware(Middleware):
- max_bytes: Maximum serialized response size in bytes (default: 50,000)
- warn_threshold_pct: Log warnings above this % of limit (default: 80%)
- max_list_items: Cap for list fields during dynamic truncation (default: 100)
+ - png_max_bytes: Maximum PNG preview response size (default: 1,000,000)
- excluded_tools: Tools to skip checking
"""
@@ -1554,8 +1556,10 @@ class ResponseSizeGuardMiddleware(Middleware):
warn_threshold_pct: int = DEFAULT_WARN_THRESHOLD_PCT,
excluded_tools: list[str] | str | None = None,
max_list_items: int = DEFAULT_MAX_LIST_ITEMS,
+ png_max_bytes: int = DEFAULT_MAX_PNG_RESPONSE_BYTES,
) -> None:
self.max_bytes = max_bytes
+ self.png_max_bytes = max(1, png_max_bytes)
self.warn_threshold_pct = warn_threshold_pct
self.warn_threshold = int(max_bytes * warn_threshold_pct / 100)
if isinstance(excluded_tools, str):
@@ -2208,6 +2212,26 @@ class ResponseSizeGuardMiddleware(Middleware):
# the oversized path below rather than slipping through unmeasured.
actual_bytes = get_response_size_bytes(estimation_target)
+ # Choose the image budget from the returned payload, never request args.
+ # Other formats of this tool retain the normal text-response budget.
+ content = (
+ estimation_target.get("content")
+ if isinstance(estimation_target, dict)
+ else None
+ )
+ if (
+ tool_name == "get_chart_preview"
+ and isinstance(content, dict)
+ and content.get("type") == "png"
+ ):
+ if actual_bytes > self.png_max_bytes:
+ raise ToolError(
+ f"PNG preview is {actual_bytes} bytes, exceeding the "
+ f"{self.png_max_bytes}-byte limit. Request smaller width and "
+ "height, or use the URL preview format."
+ )
+ return response
+
# Log warning if approaching limit
if actual_bytes > self.warn_threshold:
logger.warning(
@@ -2332,6 +2356,9 @@ def create_response_size_guard_middleware() -> ResponseSizeGuardMiddleware | Non
),
excluded_tools=config.get("excluded_tools"),
max_list_items=max_list_items,
+ png_max_bytes=_safe_int_config(
+ config, "png_max_bytes", DEFAULT_MAX_PNG_RESPONSE_BYTES
+ ),
)
logger.info(
diff --git a/superset/utils/webdriver.py b/superset/utils/webdriver.py
index e32066a5617..a5b6b247b03 100644
--- a/superset/utils/webdriver.py
+++ b/superset/utils/webdriver.py
@@ -206,8 +206,10 @@ class WebDriverPlaywright(WebDriverProxy):
window: WindowSize | None = None,
*,
require_complete_capture: bool = False,
+ browser_manager: _PlaywrightBrowserManager | None = None,
) -> None:
super().__init__(driver_type, window)
+ self._browser_manager = browser_manager or _browser_manager
self._require_complete_capture = require_complete_capture
@staticmethod
@@ -950,7 +952,7 @@ class WebDriverPlaywright(WebDriverProxy):
browser_args = app.config["WEBDRIVER_OPTION_ARGS"]
try:
- browser = _browser_manager.get_browser(browser_args)
+ browser = self._browser_manager.get_browser(browser_args)
except Exception as ex:
logger.exception(
"Failed to launch the headless browser with args %s%s",
diff --git a/tests/unit_tests/mcp_service/chart/tool/test_get_chart_preview.py b/tests/unit_tests/mcp_service/chart/tool/test_get_chart_preview.py
index c762d3a57f8..0c8ed406ad9 100644
--- a/tests/unit_tests/mcp_service/chart/tool/test_get_chart_preview.py
+++ b/tests/unit_tests/mcp_service/chart/tool/test_get_chart_preview.py
@@ -27,6 +27,8 @@ from unittest.mock import MagicMock, patch
import pytest
from superset.common.db_query_status import QueryStatus
+from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
+from superset.exceptions import SupersetSecurityException
from superset.mcp_service.chart.schemas import (
AccessibilityMetadata,
ASCIIPreview,
@@ -2268,3 +2270,414 @@ def test_saved_preview_fallback_dispatch(
assert isinstance(result, VegaLitePreview)
if viz_type == "funnel":
assert result.specification["encoding"]["y"]["field"] == "stage"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("user_id", [11, 29])
+async def test_png_preview_renders_as_caller_in_isolated_context(app_context, user_id):
+ import base64
+ import threading
+ from io import BytesIO
+
+ from flask import g
+ from PIL import Image
+
+ module = importlib.import_module(
+ "superset.mcp_service.chart.tool.get_chart_preview"
+ )
+ caller_thread = threading.get_ident()
+ g.user = SimpleNamespace(id=user_id)
+ g.request_marker = "request-only"
+ user = SimpleNamespace(id=user_id, is_active=True)
+ chart = SimpleNamespace(id=104)
+ png = BytesIO()
+ Image.new("RGB", (80, 60), "white").save(png, format="PNG")
+
+ def screenshot(url, element, *, user):
+ assert threading.get_ident() != caller_thread
+ assert g.user.id == user_id
+ assert not hasattr(g, "request_marker")
+ assert user.id == user_id
+ assert url == "http://localhost/superset/slice/104/?standalone=true"
+ assert element == "chart-container"
+ return png.getvalue()
+
+ with (
+ patch.object(module, "security_manager", new=MagicMock()) as manager,
+ patch.object(module, "find_chart_by_identifier", return_value=chart),
+ patch.object(module.guest_scope, "is_guest_read", return_value=False),
+ patch("superset.is_feature_enabled", return_value=False),
+ patch(
+ "superset.utils.urls.get_url_path",
+ return_value="http://localhost/superset/slice/104/?standalone=true",
+ ),
+ patch("superset.utils.webdriver._PlaywrightBrowserManager") as browser_manager,
+ patch("superset.utils.webdriver.WebDriverPlaywright") as driver,
+ ):
+ manager.find_user.return_value = user
+ driver.return_value.get_screenshot.side_effect = screenshot
+ result = await module._generate_png_preview(
+ 104, GetChartPreviewRequest(identifier=104, format="png")
+ )
+ assert base64.b64decode(result.data) == png.getvalue()
+ assert (result.width, result.height) == (80, 60)
+ manager.raise_for_access.assert_called_once_with(chart=chart)
+ browser_manager.return_value._cleanup.assert_called_once()
+ assert (
+ driver.call_args.kwargs["browser_manager"] is browser_manager.return_value
+ )
+ assert g.request_marker == "request-only"
+ assert g.user.id == user_id
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "guest,kwargs",
+ [
+ (True, {}),
+ (False, {"form_data_key": "unsaved"}),
+ (False, {"extra_form_data": {"filters": []}}),
+ ],
+)
+async def test_png_preview_rejects_unpropagated_context(app_context, guest, kwargs):
+ module = importlib.import_module(
+ "superset.mcp_service.chart.tool.get_chart_preview"
+ )
+ with (
+ patch.object(module.guest_scope, "is_guest_read", return_value=guest),
+ patch("superset.utils.webdriver._PlaywrightBrowserManager") as manager,
+ ):
+ result = await module._generate_png_preview(
+ 104, GetChartPreviewRequest(identifier=104, format="png", **kwargs)
+ )
+ assert result.error_type == "UnsupportedFormat"
+ manager.assert_not_called()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("failure", ["denied", "missing", "inactive", "export"])
+async def test_png_preview_authorizes_before_browser(app_context, failure):
+ from flask import g
+
+ module = importlib.import_module(
+ "superset.mcp_service.chart.tool.get_chart_preview"
+ )
+ g.user = SimpleNamespace(id=11)
+ with (
+ patch.object(module, "security_manager", new=MagicMock()) as manager,
+ patch.object(
+ module,
+ "find_chart_by_identifier",
+ return_value=None if failure == "missing" else SimpleNamespace(id=104),
+ ),
+ patch.object(module.guest_scope, "is_guest_read", return_value=False),
+ patch("superset.is_feature_enabled", return_value=failure == "export"),
+ patch("superset.utils.webdriver._PlaywrightBrowserManager") as browser,
+ ):
+ manager.find_user.return_value = SimpleNamespace(
+ id=11, is_active=failure != "inactive"
+ )
+ manager.can_access.return_value = False
+ if failure == "denied":
+ # raise_for_access raises SupersetSecurityException on denial; the
+ # message must not leak into the client-facing error.
+ manager.raise_for_access.side_effect = SupersetSecurityException(
+ SupersetError(
+ message="secret denied URL",
+ error_type=SupersetErrorType.GENERIC_BACKEND_ERROR,
+ level=ErrorLevel.ERROR,
+ )
+ )
+ result = await module._generate_png_preview(
+ 104, GetChartPreviewRequest(identifier=104, format="png")
+ )
+ assert isinstance(result, ChartError)
+ assert "secret" not in result.error
+ if failure == "denied":
+ # Authorization denials must surface as Forbidden, not RenderError.
+ assert result.error_type == "Forbidden"
+ browser.assert_not_called()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("image", [None, b"not an image", b"\xff\xd8\xffjpeg"])
+async def test_png_preview_cleans_up_failed_capture(app_context, image):
+ from flask import g
+
+ module = importlib.import_module(
+ "superset.mcp_service.chart.tool.get_chart_preview"
+ )
+ g.user = SimpleNamespace(id=11)
+ with (
+ patch.object(module, "security_manager", new=MagicMock()) as manager,
+ patch.object(
+ module, "find_chart_by_identifier", return_value=SimpleNamespace(id=104)
+ ),
+ patch.object(module.guest_scope, "is_guest_read", return_value=False),
+ patch("superset.is_feature_enabled", return_value=False),
+ patch(
+ "superset.utils.urls.get_url_path",
+ return_value="http://localhost/chart/104",
+ ),
+ patch("superset.utils.webdriver._PlaywrightBrowserManager") as browser,
+ patch("superset.utils.webdriver.WebDriverPlaywright") as driver,
+ ):
+ manager.find_user.return_value = SimpleNamespace(id=11, is_active=True)
+ driver.return_value.get_screenshot.return_value = image
+ result = await module._generate_png_preview(
+ 104, GetChartPreviewRequest(identifier=104, format="png")
+ )
+ assert result.error_type == "RenderError"
+ browser.return_value._cleanup.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_png_preview_cleans_up_browser_exception(app_context):
+ from flask import g
+
+ module = importlib.import_module(
+ "superset.mcp_service.chart.tool.get_chart_preview"
+ )
+ g.user = SimpleNamespace(id=11)
+ with (
+ patch.object(module, "security_manager", new=MagicMock()) as manager,
+ patch.object(
+ module, "find_chart_by_identifier", return_value=SimpleNamespace(id=104)
+ ),
+ patch.object(module.guest_scope, "is_guest_read", return_value=False),
+ patch("superset.is_feature_enabled", return_value=False),
+ patch(
+ "superset.utils.urls.get_url_path",
+ return_value="http://localhost/chart/104",
+ ),
+ patch("superset.utils.webdriver._PlaywrightBrowserManager") as browser,
+ patch("superset.utils.webdriver.WebDriverPlaywright") as driver,
+ ):
+ manager.find_user.return_value = SimpleNamespace(id=11, is_active=True)
+ driver.return_value.get_screenshot.side_effect = RuntimeError(
+ "sensitive page data"
+ )
+ result = await module._generate_png_preview(
+ 104, GetChartPreviewRequest(identifier=104, format="png")
+ )
+ assert result.error_type == "RenderError"
+ assert "sensitive" not in result.error
+ browser.return_value._cleanup.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_png_worker_preserves_parent_mcp_orm_session(app_context):
+ from flask import g
+ from flask_appbuilder.security.sqla.models import User
+ from sqlalchemy import inspect
+
+ from superset.extensions import db
+ from superset.mcp_service.auth import _mcp_tool_call_context
+ from superset.mcp_service.session_scope import mcp_session_scopefunc
+
+ module = importlib.import_module(
+ "superset.mcp_service.chart.tool.get_chart_preview"
+ )
+ with patch.object(db.session.registry, "scopefunc", mcp_session_scopefunc):
+ with _mcp_tool_call_context():
+ g.user = SimpleNamespace(id=11)
+ parent_session = db.session()
+ parent_key = mcp_session_scopefunc()
+ parent_entity = User(username="parent-only")
+ parent_session.add(parent_entity)
+ worker_keys = []
+
+ def lookup(*, id):
+ worker_keys.append(mcp_session_scopefunc())
+ assert db.session() is not parent_session
+ return SimpleNamespace(id=id, is_active=False)
+
+ with (
+ patch.object(module, "security_manager", new=MagicMock()) as manager,
+ patch.object(module.guest_scope, "is_guest_read", return_value=False),
+ ):
+ manager.find_user.side_effect = lookup
+ result = await module._generate_png_preview(
+ 104, GetChartPreviewRequest(identifier=104, format="png")
+ )
+ assert worker_keys
+ assert worker_keys[0] != parent_key
+ assert result.error_type == "Forbidden"
+ assert db.session() is parent_session
+ assert inspect(parent_entity).session is parent_session
+ parent_session.expunge(parent_entity)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("dimension", [0, 63, 4097])
+@pytest.mark.parametrize("field", ["width", "height"])
+async def test_png_preview_rejects_invalid_viewport(app_context, dimension, field):
+ from flask import g
+
+ module = importlib.import_module(
+ "superset.mcp_service.chart.tool.get_chart_preview"
+ )
+ g.user = SimpleNamespace(id=11)
+ with (
+ patch.object(module.guest_scope, "is_guest_read", return_value=False),
+ patch("superset.utils.webdriver._PlaywrightBrowserManager") as browser,
+ ):
+ result = await module._generate_png_preview(
+ 104,
+ GetChartPreviewRequest(identifier=104, format="png", **{field: dimension}),
+ )
+ assert result.error_type == "ValidationError"
+ browser.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_png_preview_passes_default_response_size_guard():
+ import base64
+ import os
+ from io import BytesIO
+ from unittest.mock import AsyncMock
+
+ from PIL import Image
+
+ from superset.mcp_service.chart.schemas import PNGPreview
+ from superset.mcp_service.mcp_config import MCP_RESPONSE_SIZE_CONFIG
+ from superset.mcp_service.middleware import ResponseSizeGuardMiddleware
+ from superset.mcp_service.utils.response_size_utils import get_response_size_bytes
+
+ # Actual PNG bytes with enough detail to exceed the default text budget.
+ png = BytesIO()
+ Image.frombytes("L", (800, 600), os.urandom(800 * 600)).save(png, format="PNG")
+ preview = PNGPreview(
+ data=base64.b64encode(png.getvalue()).decode("ascii"), width=800, height=600
+ )
+ response = {"content": preview.model_dump()}
+ assert get_response_size_bytes(response) > MCP_RESPONSE_SIZE_CONFIG["max_bytes"]
+ context = MagicMock()
+ context.message.name = "get_chart_preview"
+ context.message.arguments = {"identifier": 104, "format": "png"}
+ call_next = AsyncMock(return_value=response)
+ middleware = ResponseSizeGuardMiddleware(
+ max_bytes=MCP_RESPONSE_SIZE_CONFIG["max_bytes"],
+ excluded_tools=MCP_RESPONSE_SIZE_CONFIG["excluded_tools"],
+ )
+ result = await middleware.on_call_tool(context, call_next)
+ assert PNGPreview.model_validate(result["content"]) == preview
+ call_next.assert_awaited_once_with(context)
+
+
+@pytest.mark.asyncio
+async def test_png_workers_do_not_starve_auth_or_release_on_cancellation():
+ import asyncio
+ import threading
+ from contextvars import ContextVar
+
+ module = importlib.import_module(
+ "superset.mcp_service.chart.tool.get_chart_preview"
+ )
+ started = [threading.Event(), threading.Event(), threading.Event()]
+ release = threading.Event()
+ marker = ContextVar("png_test_marker", default=None)
+ marker.set("caller")
+
+ def render(index):
+ assert marker.get() == "caller"
+ started[index].set()
+ assert release.wait(timeout=10)
+ return None
+
+ tasks = [
+ asyncio.create_task(module._run_png_render(lambda i=i: render(i)))
+ for i in range(3)
+ ]
+ try:
+ for event in started[:2]:
+ assert await asyncio.to_thread(event.wait, 5)
+ assert not started[2].is_set()
+ tasks[0].cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await tasks[0]
+ assert not started[2].is_set()
+ assert await asyncio.wait_for(asyncio.to_thread(lambda: "auth"), 1) == "auth"
+ finally:
+ release.set()
+ await asyncio.gather(*tasks, return_exceptions=True)
+ assert started[2].is_set()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("size", [800, 4096])
+async def test_png_response_budget_over_real_mcp_transport(size):
+ import base64
+ import os
+ from io import BytesIO
+
+ from fastmcp import Client, FastMCP
+ from fastmcp.exceptions import ToolError
+ from PIL import Image
+
+ from superset.mcp_service.middleware import ResponseSizeGuardMiddleware
+
+ png = BytesIO()
+ Image.frombytes("L", (size, size), os.urandom(size * size)).save(png, format="PNG")
+ encoded = base64.b64encode(png.getvalue()).decode("ascii")
+ server = FastMCP("PNG budget test")
+ server.add_middleware(ResponseSizeGuardMiddleware())
+
+ @server.tool
+ def get_chart_preview() -> dict:
+ return {
+ "content": {"type": "png", "data": encoded, "width": size, "height": size}
+ }
+
+ async with Client(server) as client:
+ if size == 4096:
+ with pytest.raises(ToolError, match="PNG preview.*1000000.*smaller"):
+ await client.call_tool("get_chart_preview")
+ else:
+ result = await client.call_tool("get_chart_preview")
+ assert result.data["content"]["data"] == encoded
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("preview_type", ["ascii", "table", "vega_lite"])
+async def test_non_png_previews_keep_text_response_budget(preview_type):
+ from unittest.mock import AsyncMock
+
+ from fastmcp.exceptions import ToolError
+
+ from superset.mcp_service.mcp_config import MCP_RESPONSE_SIZE_CONFIG
+ from superset.mcp_service.middleware import ResponseSizeGuardMiddleware
+
+ middleware = ResponseSizeGuardMiddleware(
+ excluded_tools=MCP_RESPONSE_SIZE_CONFIG["excluded_tools"],
+ )
+ context = MagicMock()
+ context.message.name = "get_chart_preview"
+ # The returned format, not caller-supplied arguments, chooses the limit.
+ context.message.arguments = {"format": "png"}
+ response = {"content": {"type": preview_type, "data": "x" * 60_000}}
+ with pytest.raises(ToolError):
+ await middleware.on_call_tool(context, AsyncMock(return_value=response))
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("headroom", [0, -1])
+async def test_png_budget_counts_metadata_at_exact_boundary(headroom):
+ from unittest.mock import AsyncMock
+
+ from fastmcp.exceptions import ToolError
+
+ from superset.mcp_service.middleware import ResponseSizeGuardMiddleware
+ from superset.mcp_service.utils.response_size_utils import get_response_size_bytes
+
+ response = {"chart_name": "A chart", "content": {"type": "png", "data": "AAAA"}}
+ budget = get_response_size_bytes(response) + headroom
+ middleware = ResponseSizeGuardMiddleware(png_max_bytes=budget)
+ context = MagicMock()
+ context.message.name = "get_chart_preview"
+ call_next = AsyncMock(return_value=response)
+ if headroom < 0:
+ with pytest.raises(ToolError, match="PNG preview"):
+ await middleware.on_call_tool(context, call_next)
+ else:
+ assert await middleware.on_call_tool(context, call_next) is response
diff --git a/tests/unit_tests/mcp_service/test_middleware.py b/tests/unit_tests/mcp_service/test_middleware.py
index 8112de0afa7..040d21ef662 100644
--- a/tests/unit_tests/mcp_service/test_middleware.py
+++ b/tests/unit_tests/mcp_service/test_middleware.py
@@ -1388,7 +1388,7 @@ class TestCreateResponseSizeGuardMiddleware:
"""Test create_response_size_guard_middleware factory function."""
def test_default_config_checks_chart_preview(self) -> None:
- """Should size-check chart preview responses by default."""
+ """Only PNG responses may use the separate image budget."""
mock_flask_app = MagicMock()
mock_flask_app.config.get.return_value = MCP_RESPONSE_SIZE_CONFIG
@@ -1399,6 +1399,7 @@ class TestCreateResponseSizeGuardMiddleware:
middleware = create_response_size_guard_middleware()
assert middleware is not None
+ assert middleware.png_max_bytes == 1_000_000
assert "get_chart_preview" not in middleware.excluded_tools
assert "health_check" in middleware.excluded_tools
@@ -1407,6 +1408,7 @@ class TestCreateResponseSizeGuardMiddleware:
mock_config = {
"enabled": True,
"max_bytes": 30000,
+ "png_max_bytes": 250000,
"warn_threshold_pct": 75,
"excluded_tools": ["health_check"],
}
@@ -1423,6 +1425,7 @@ class TestCreateResponseSizeGuardMiddleware:
assert middleware is not None
assert isinstance(middleware, ResponseSizeGuardMiddleware)
assert middleware.max_bytes == 30000
+ assert middleware.png_max_bytes == 250000
assert middleware.warn_threshold_pct == 75
assert "health_check" in middleware.excluded_tools