(superset) 01/01: fix(mcp): clean up nits from PNG chart preview review (#44558)

[ru...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

rusackas pushed a commit to branch fix/mcp-png-preview-nits
in repository https://gitbox.apache.org/repos/asf/superset.git

commit 18eb97edb55535f343eb8f395b4a82ad7841734a
Author: Evan Rusackas <ev...@rusackas.com>
AuthorDate: Mon Oct 19 21:04:57 2026 -0700

    fix(mcp): clean up nits from PNG chart preview review (#44558)
    
    Small follow-ups from #44558's review that weren't worth holding the
    merge for: reuse the existing PNG_MAGIC_BYTES constant instead of a
    second hardcoded magic-byte check, avoid constructing
    PreviewFormatGenerator on the png branch where it's never used, log
    mcp_response_size_exceeded on PNG rejections the same way the text
    path does, and name the PNG-guarded tool set as a constant instead of
    a bare string literal.
    
    Co-Authored-By: Evan Rusackas <ev...@preset.io>
    Co-Authored-By: Claude Sonnet 5 <no...@anthropic.com>
---
 .../mcp_service/chart/tool/get_chart_preview.py    | 10 +++++----
 superset/mcp_service/middleware.py                 | 25 +++++++++++++++++++++-
 superset/mcp_service/utils/response_size_utils.py  |  4 ++++
 3 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/superset/mcp_service/chart/tool/get_chart_preview.py b/superset/mcp_service/chart/tool/get_chart_preview.py
index 9f9aed4fd0c..300928deba2 100644
--- a/superset/mcp_service/chart/tool/get_chart_preview.py
+++ b/superset/mcp_service/chart/tool/get_chart_preview.py
@@ -1257,7 +1257,10 @@ async def _generate_png_preview(
         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.screenshots import (
+            PNG_MAGIC_BYTES,
+            validate_screenshot_image,
+        )
         from superset.utils.urls import get_url_path
         from superset.utils.webdriver import (
             _PlaywrightBrowserManager,
@@ -1297,7 +1300,7 @@ async def _generate_png_preview(
                     image = driver.get_screenshot(url, "chart-container", user=user)
                     if (
                         image is None
-                        or not image.startswith(b"\x89PNG\r\n\x1a\n")
+                        or not image.startswith(PNG_MAGIC_BYTES)
                         or validate_screenshot_image(image)
                     ):
                         return ChartError(
@@ -1604,11 +1607,10 @@ async def _get_chart_preview_internal(  # noqa: C901
         with event_logger.log_context(
             action="mcp.get_chart_preview.preview_generation"
         ):
-            preview_generator = PreviewFormatGenerator(chart, request)
             content = (
                 await _generate_png_preview(chart.id, request)
                 if request.format == "png"
-                else preview_generator.generate()
+                else PreviewFormatGenerator(chart, request).generate()
             )
 
         if isinstance(content, ChartError):
diff --git a/superset/mcp_service/middleware.py b/superset/mcp_service/middleware.py
index fca8966d66b..199cb693f95 100644
--- a/superset/mcp_service/middleware.py
+++ b/superset/mcp_service/middleware.py
@@ -71,6 +71,7 @@ from superset.mcp_service.utils.response_size_utils import (
     format_size_limit_error,
     get_response_size_bytes,
     INFO_TOOLS,
+    PNG_PREVIEW_TOOLS,
     RESPONSE_AWARE_SUGGESTION_TOOLS,
     string_clip_chars,
     STRING_FIELD_TRUNCATION_TOOLS,
@@ -2220,11 +2221,33 @@ class ResponseSizeGuardMiddleware(Middleware):
             else None
         )
         if (
-            tool_name == "get_chart_preview"
+            tool_name in PNG_PREVIEW_TOOLS
             and isinstance(content, dict)
             and content.get("type") == "png"
         ):
             if actual_bytes > self.png_max_bytes:
+                logger.warning(
+                    "Response blocked for %s: %d bytes exceeds png_max_bytes of %d",
+                    tool_name,
+                    actual_bytes,
+                    self.png_max_bytes,
+                )
+                try:
+                    event_logger.log(
+                        user_id=get_user_id(),
+                        action="mcp_response_size_exceeded",
+                        dashboard_id=None,
+                        duration_ms=None,
+                        slice_id=None,
+                        referrer=None,
+                        curated_payload={
+                            "tool": tool_name,
+                            "actual_bytes": actual_bytes,
+                            "max_bytes": self.png_max_bytes,
+                        },
+                    )
+                except Exception as log_error:  # noqa: BLE001
+                    logger.warning("Failed to log size exceeded event: %s", log_error)
                 raise ToolError(
                     f"PNG preview is {actual_bytes} bytes, exceeding the "
                     f"{self.png_max_bytes}-byte limit. Request smaller width and "
diff --git a/superset/mcp_service/utils/response_size_utils.py b/superset/mcp_service/utils/response_size_utils.py
index d217fd65a74..38b9614fc1c 100644
--- a/superset/mcp_service/utils/response_size_utils.py
+++ b/superset/mcp_service/utils/response_size_utils.py
@@ -295,6 +295,10 @@ def _identify_large_fields(response: ToolResponse) -> List[str]:
 # the size guard passes the parsed payload to format_size_limit_error.
 RESPONSE_AWARE_SUGGESTION_TOOLS = frozenset({"get_dashboard_layout"})
 
+# Tools whose PNG-formatted responses are checked against png_max_bytes
+# instead of the normal text-response budget.
+PNG_PREVIEW_TOOLS = frozenset({"get_chart_preview"})
+
 _MAX_SUGGESTED_TAB_IDS = 5
 _MAX_SUGGESTED_TAB_ID_CHARS = 64