You are viewing a plain text version of this content. The canonical link for it is here.
Posted to notifications@superset.apache.org by "aminghadersohi (via GitHub)" <gi...@apache.org> on 2026/08/06 06:46:50 UTC

[PR] fix(mcp): return all chart query results [superset]

aminghadersohi opened a new pull request, #42824:
URL: https://github.com/apache/superset/pull/42824

   ### SUMMARY
   
   `get_chart_data` previously selected only the first entry returned by `ChartDataCommand`, which dropped secondary results from charts whose query context contains multiple queries.
   
   This change preserves the existing top-level `columns`, `data`, and row metadata as the first-query view for backward compatibility, and adds an optional `query_results` field containing every query result for multi-query charts. Single-query charts retain their existing behavior. Response-size enforcement applies one shared row cap across the complete multi-query response so additional results cannot bypass the configured token limit.
   
   Mixed Timeseries is the chart type that explicitly builds two queries in the MCP chart query helper. The response handling also covers any saved query context containing multiple queries rather than coupling serialization to that visualization type.
   
   ### BEFORE/AFTER SCREENSHOTS OR ANIMATED GIF
   
   Not applicable; this changes an MCP tool response without UI changes.
   
   ### TESTING INSTRUCTIONS
   
   - `pytest -q tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py tests/unit_tests/mcp_service/test_middleware.py` (182 passed)
   - `pytest -q tests/unit_tests/mcp_service` (3319 passed, 1 unrelated health-check failure because the local test configuration returned a null application version)
   - `uvx pre-commit run --files superset/mcp_service/chart/schemas.py superset/mcp_service/chart/tool/get_chart_data.py superset/mcp_service/utils/token_utils.py tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py tests/unit_tests/mcp_service/test_middleware.py`
   - `pre-commit run --all-files` was also run. Python checks passed; the repository-wide frontend hooks could not complete because frontend dependencies such as `glob` and `postcss-styled-syntax` are not installed, and repository-wide Ruff reported pre-existing violations outside this change.
   
   The regression coverage executes the unsaved Mixed Timeseries response path with two command results and asserts that both are returned. Additional coverage verifies the additive single-query behavior and shared response-size limiting.
   
   ### ADDITIONAL INFORMATION
   
   - [ ] Has associated issue:
   - [ ] Required feature flags:
   - [ ] Changes UI
   - [ ] Includes DB Migration (follow approval process in [SIP-59](https://github.com/apache/superset/issues/13351))
     - [ ] Migration is atomic, supports rollback & is backwards-compatible
     - [ ] Confirm DB migration upgrade and downgrade tested
     - [ ] Runtime estimates and downtime expectations provided
   - [x] Introduces new feature or API
   - [ ] Removes existing feature or API
   
   ### BACKWARD COMPATIBILITY
   
   This is additive. Existing top-level fields continue to represent the first query, while multi-query responses additionally expose `query_results`. Clients that ignore unknown fields are unaffected.
   
   ### EVAL EVIDENCE
   
   Not applicable to this deterministic serialization change; no prompt, model routing, or model-generated output changed. The MCP regression suite results are listed above.
   
   ### COST & LATENCY DELTA
   
   No model or parameter changes. Multi-query responses include more data by design, bounded by the existing response token limit using response-wide truncation. No external model-cost or latency benchmark is available in the OSS worktree.
   
   ### PROMPT / NON-DETERMINISM
   
   No prompt changes and no non-deterministic behavior.
   
   ### REVIEW GUIDANCE
   
   Review the additive schema and response construction first, then the shared-budget truncation logic. Concurrent work touching `get_chart_data` may need to rebase carefully around the response construction paths.
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org


Re: [PR] fix(mcp): return all chart query results [superset]

Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on PR #42824:
URL: https://github.com/apache/superset/pull/42824#issuecomment-5262502497

   <!-- Bito Reply -->
   The flagged issue is correct. In `_build_query_results`, the `row_count` is currently set to `len(query_result.get("data", []))`, which represents the total rows before slicing, rather than the length of the sliced `data` list. To resolve this, `row_count` should be set to the length of the sliced data list, while `total_rows` retains the original count.
   
   Here is the corrected implementation for `superset/mcp_service/chart/tool/get_chart_data.py`:
   
   ```python
               data=(
                   query_result.get("data", [])[:limit]
                   if limit
                   else query_result.get("data", [])
               ),
               row_count=len(query_result.get("data", [])[:limit] if limit else query_result.get("data", [])),
               total_rows=query_result.get("rowcount"),
   ```
   
   There are no other comments in the PR to address.
   
   **superset/mcp_service/chart/tool/get_chart_data.py**
   ```
   data=(
                   query_result.get("data", [])[:limit]
                   if limit
                   else query_result.get("data", [])
               ),
               row_count=len(query_result.get("data", [])[:limit] if limit else query_result.get("data", [])),
               total_rows=query_result.get("rowcount"),
   ```


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org


Re: [PR] fix(mcp): return all chart query results [superset]

Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on code in PR #42824:
URL: https://github.com/apache/superset/pull/42824#discussion_r3769932200


##########
tests/unit_tests/mcp_service/test_middleware.py:
##########
@@ -471,6 +471,36 @@ async def test_truncates_get_chart_data_rows(self) -> None:
         assert result["_response_truncated"] is True
         assert len(result["data"]) < 200
 
+    @pytest.mark.asyncio
+    async def test_truncates_multi_query_chart_rows_across_whole_response(self) -> None:
+        """All query results share the response's token budget."""
+        middleware = ResponseSizeGuardMiddleware(token_limit=500)
+        context = MagicMock()
+        context.message.name = "get_chart_data"
+        context.message.arguments = {}
+        row = {f"col_{i}": f"value_{i}" for i in range(10)}
+        large_response = {
+            "chart_id": 1,
+            "data": [row] * 200,
+            "row_count": 200,
+            "query_results": [
+                {"query_index": 0, "data": [row] * 200, "row_count": 200},
+                {"query_index": 1, "data": [row] * 200, "row_count": 200},
+            ],
+        }
+        call_next = AsyncMock(return_value=large_response)
+
+        with (
+            patch("superset.mcp_service.middleware.get_user_id", return_value=1),
+            patch("superset.mcp_service.middleware.event_logger"),
+        ):
+            result = await middleware.on_call_tool(context, call_next)
+
+        assert isinstance(result, dict)
+        assert result["_response_truncated"] is True
+        assert all(query["data"] for query in result["query_results"])
+        assert all(len(query["data"]) < 200 for query in result["query_results"])
+

Review Comment:
   Addressed in 16b9da44b3. The regression now asserts that every nested query result is truncated to the same returned row count and that each nested `row_count` exactly matches its truncated `data` length.
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org


Re: [PR] fix(mcp): return all chart query results [superset]

Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on PR #42824:
URL: https://github.com/apache/superset/pull/42824#issuecomment-5296501924

   Fixed the unit-test failure in 285eeaa6ee. Reserving the longest final truncation note correctly changes the 700-token boundary from two rows per query to one, so the regression now asserts the invariant that the note reports 400 distinct source rows without coupling the test to the exact kept-row count. Changed-file pre-commit passes.
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org


Re: [PR] fix(mcp): return all chart query results [superset]

Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on PR #42824:
URL: https://github.com/apache/superset/pull/42824#issuecomment-5201321250

   AI Code Review is in progress (usually takes 3 to 15 minutes unless it's a very large PR).


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org


Re: [PR] fix(mcp): return all chart query results [superset]

Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on PR #42824:
URL: https://github.com/apache/superset/pull/42824#issuecomment-5272204354

   The two experimental Playwright failures are unrelated to this backend MCP diff. Run 31615757119 timed out in existing embedded-dashboard tests (`dashboard renders in embedded iframe` / `charts render inside embedded iframe`) and the existing pivot-collapse scroll test, all while waiting 30 seconds for embedded iframe UI. The PR changes only MCP Python response shaping and unit tests; no embedded frontend or Playwright code is touched. I retried the failed run once as requested. Separately, I addressed the latest review suggestion in 16b9da44b3 by strengthening the shared-budget truncation assertions.
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org


Re: [PR] fix(mcp): return all chart query results [superset]

Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on PR #42824:
URL: https://github.com/apache/superset/pull/42824#issuecomment-5273832665

   AI Code Review is in progress (usually takes 3 to 15 minutes unless it's a very large PR).
   
   <details>
   <summary>
   Bito Usage Guide
   </summary>
   
   **Commands**
         
   Type the following command in the pull request comment and save the comment.
         
   - `/review` - Manually triggers a full AI review.
   
   - `/pause` - Pauses automatic reviews on this pull request.
   - `/resume` - Resumes automatic reviews.
   - `/resolve` - Marks all Bito-posted review comments as resolved.
   - `/abort` - Cancels all in-progress reviews.
   
         
   Refer to the <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/available-commands" target="_blank" rel="noopener noreferrer">documentation</a> for additional commands.
         
   **Configuration**
   
   This repository uses `Superset` You can customize the agent settings <a href="https://alpha.bito.ai/home/ai-agents/code-review-agent" target="_blank" rel="noopener noreferrer">here</a> or contact your Bito workspace admin at evan@preset.io.
         
   **Documentation & Help**
   - <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/getting-started/install-run-using-bito-cloud/create-or-customize-an-agent-instance" target="_blank" rel="noopener noreferrer">Customize agent settings</a>
   - <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/implementing-custom-code-review-rules" target="_blank" rel="noopener noreferrer">Review rules</a>
   - <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent" target="_blank" rel="noopener noreferrer">General documentation</a>
   - <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/faqs" target="_blank" rel="noopener noreferrer">FAQ</a>
   
   </details>
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org


Re: [PR] fix(mcp): return all chart query results [superset]

Posted by "codeant-ai-for-open-source[bot] (via GitHub)" <gi...@apache.org>.
codeant-ai-for-open-source[bot] commented on code in PR #42824:
URL: https://github.com/apache/superset/pull/42824#discussion_r3763728793


##########
superset/mcp_service/chart/tool/get_chart_data.py:
##########
@@ -279,6 +286,29 @@ def _sanitize_chart_data_for_llm_context(chart_data: ChartData) -> ChartData:
     return ChartData.model_validate(payload)
 
 
+def _build_query_results(
+    query_results: list[dict[str, Any]], limit: int | None
+) -> list[ChartQueryResult] | None:
+    """Preserve every result when a chart executes more than one query."""
+    if len(query_results) <= 1:
+        return None
+
+    return [
+        ChartQueryResult(
+            query_index=index,
+            columns=query_result.get("colnames", []),
+            data=(
+                query_result.get("data", [])[:limit]
+                if limit
+                else query_result.get("data", [])
+            ),
+            row_count=len(query_result.get("data", [])),
+            total_rows=query_result.get("rowcount"),
+        )

Review Comment:
   **Suggestion:** The nested result slices `data` to `limit`, but `row_count` is calculated from the unsliced source list. Consequently, a response containing 10 returned rows can report `row_count=100`, while `total_rows` reports the same source count, contradicting `ChartQueryResult`'s “Rows returned” contract. Set `row_count` to the number of rows actually placed in the nested `data` field, while retaining the source count in `total_rows`. [api mismatch]
   
   <details>
   <summary><b>Severity Level:</b> Major ⚠️</summary>
   
   ```mdx
   - ⚠️ Multi-query chart responses expose incorrect nested row counts.
   - ⚠️ MCP consumers may misinterpret returned versus available rows.
   - ⚠️ Mixed Timeseries results are affected when a limit is configured.
   ```
   </details>
   
   [![Fix in Cursor](https://new-codeant-butcket.s3.us-west-1.amazonaws.com/badges/fix-in-cursor-flat.svg)](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=9ba1d1c1365b456da427d505b36efe16&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset) [![Fix in VSCode Claude](https://new-codeant-butcket.s3.us-west-1.amazonaws.com/badges/fix-in-vscode-claude-flat.svg)](https://app.codeant.ai/fix-in-ide?tool=vscode-claude&prompt_id=9ba1d1c1365b456da427d505b36efe16&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset)
   <details>
   <summary><b>Prompt for AI Agent 🤖 </b></summary>
   
   ```mdx
   This is a comment left during a code review.
   
   **Path:** superset/mcp_service/chart/tool/get_chart_data.py
   **Line:** 297:307
   **Comment:**
   	*Api Mismatch: The nested result slices `data` to `limit`, but `row_count` is calculated from the unsliced source list. Consequently, a response containing 10 returned rows can report `row_count=100`, while `total_rows` reports the same source count, contradicting `ChartQueryResult`'s “Rows returned” contract. Set `row_count` to the number of rows actually placed in the nested `data` field, while retaining the source count in `total_rows`.
   
   Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
   Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
   ```
   </details>
   <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F42824&comment_hash=5648c2238655a5957f6c92a430e3c7358fa6b1f4015a1dfa9460a67d4abfc9a1&reaction=like'>👍</a> | <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F42824&comment_hash=5648c2238655a5957f6c92a430e3c7358fa6b1f4015a1dfa9460a67d4abfc9a1&reaction=dislike'>👎</a>



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org


Re: [PR] fix(mcp): return all chart query results [superset]

Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on code in PR #42824:
URL: https://github.com/apache/superset/pull/42824#discussion_r3764276851


##########
superset/mcp_service/chart/tool/get_chart_data.py:
##########
@@ -279,6 +286,29 @@ def _sanitize_chart_data_for_llm_context(chart_data: ChartData) -> ChartData:
     return ChartData.model_validate(payload)
 
 
+def _build_query_results(
+    query_results: list[dict[str, Any]], limit: int | None
+) -> list[ChartQueryResult] | None:
+    """Preserve every result when a chart executes more than one query."""
+    if len(query_results) <= 1:
+        return None
+
+    return [
+        ChartQueryResult(
+            query_index=index,
+            columns=query_result.get("colnames", []),
+            data=(
+                query_result.get("data", [])[:limit]
+                if limit
+                else query_result.get("data", [])
+            ),
+            row_count=len(query_result.get("data", [])),
+            total_rows=query_result.get("rowcount"),
+        )

Review Comment:
   Fixed in 9edfa951c1. `_build_query_results` now computes `row_count` from the sliced `returned_data`, while `total_rows` continues to carry the source result count. I also added a regression test covering `limit=1` across both query results.
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org


Re: [PR] fix(mcp): return all chart query results [superset]

Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on PR #42824:
URL: https://github.com/apache/superset/pull/42824#issuecomment-5262311919

   Thanks for flagging the initial patch coverage. The multi-query response construction is covered by regression tests for a genuine Mixed Timeseries query context, the full unsaved-chart tool response path, legacy single-query behavior, and response-wide truncation across multiple query results. The focused suite passed before push; the rebased tree’s changed-file pre-commit suite also passes. The latest push will refresh the coverage report against the current base.
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org


Re: [PR] fix(mcp): return all chart query results [superset]

Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on code in PR #42824:
URL: https://github.com/apache/superset/pull/42824#discussion_r3785474621


##########
superset/mcp_service/utils/token_utils.py:
##########
@@ -802,6 +802,63 @@ def _truncate_rows_field(
     return None
 
 
+def _truncate_chart_query_results(
+    data: Dict[str, Any], token_limit: int, advice: str
+) -> list[str] | None:
+    """Apply one response-wide row cap to every result of a multi-query chart."""
+    from superset.utils import json as utils_json
+
+    query_results = data.get("query_results")
+    if not isinstance(query_results, list) or not query_results:
+        return None
+
+    row_lists = [data.get("data", [])]
+    row_lists.extend(
+        result.get("data", [])
+        for result in query_results
+        if isinstance(result, dict) and isinstance(result.get("data"), list)
+    )
+    originals = [list(rows) for rows in row_lists]
+    if not any(originals):
+        return None
+
+    original_count = sum(len(rows) for rows in originals)
+    data["_response_truncated"] = True
+    data["_truncation_notes"] = [
+        f"Result truncated from {original_count} rows across multiple queries "
+        f"(limit ~{token_limit:,} tokens). {advice}"
+    ]

Review Comment:
   Fixed in 3d176a838e. The bisection placeholder now uses the longest possible final-note shape, and a token-limit=700 regression asserts the fully materialized response remains within budget.
   



##########
superset/mcp_service/utils/token_utils.py:
##########
@@ -802,6 +802,63 @@ def _truncate_rows_field(
     return None
 
 
+def _truncate_chart_query_results(
+    data: Dict[str, Any], token_limit: int, advice: str
+) -> list[str] | None:
+    """Apply one response-wide row cap to every result of a multi-query chart."""
+    from superset.utils import json as utils_json
+
+    query_results = data.get("query_results")
+    if not isinstance(query_results, list) or not query_results:
+        return None
+
+    row_lists = [data.get("data", [])]
+    row_lists.extend(
+        result.get("data", [])
+        for result in query_results
+        if isinstance(result, dict) and isinstance(result.get("data"), list)
+    )
+    originals = [list(rows) for rows in row_lists]
+    if not any(originals):
+        return None
+
+    original_count = sum(len(rows) for rows in originals)

Review Comment:
   Fixed in 3d176a838e. Truncation notes now count only the distinct rows in `query_results`; the aliased top-level first-query rows are still capped in sync but are excluded from original/kept totals. The regression verifies the note reports 4 of 400 rows rather than the duplicated 6 of 600.
   



##########
superset/mcp_service/chart/tool/get_chart_data.py:
##########
@@ -883,6 +913,9 @@ async def get_chart_data(  # noqa: C901
                     chart_type=chart.viz_type or "unknown",
                     columns=columns,
                     data=data[: request.limit] if request.limit else data,
+                    query_results=_build_query_results(

Review Comment:
   Handled in 3d176a838e. Both saved and cached-form-data paths now return a multi-query response when any query has rows, even if query 0 is empty. They return `NoData` only when every query is empty. The unsaved Mixed Timeseries regression now exercises an empty primary result with a populated secondary result.
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org


Re: [PR] fix(mcp): return all chart query results [superset]

Posted by "codecov[bot] (via GitHub)" <gi...@apache.org>.
codecov[bot] commented on PR #42824:
URL: https://github.com/apache/superset/pull/42824#issuecomment-5201385287

   ## [Codecov](https://app.codecov.io/gh/apache/superset/pull/42824?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
   :x: Patch coverage is `15.68627%` with `43 lines` in your changes missing coverage. Please review.
   :white_check_mark: Project coverage is 56.15%. Comparing base ([`1de35d1`](https://app.codecov.io/gh/apache/superset/commit/1de35d136117e171c92acaf6cf898ab8ab395a42?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)) to head ([`4087586`](https://app.codecov.io/gh/apache/superset/commit/4087586adc3b0704729db8fe3c7980743284a728?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)).
   :warning: Report is 4 commits behind head on master.
   
   | [Files with missing lines](https://app.codecov.io/gh/apache/superset/pull/42824?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Patch % | Lines |
   |---|---|---|
   | [superset/mcp\_service/utils/token\_utils.py](https://app.codecov.io/gh/apache/superset/pull/42824?src=pr&el=tree&filepath=superset%2Fmcp_service%2Futils%2Ftoken_utils.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c3VwZXJzZXQvbWNwX3NlcnZpY2UvdXRpbHMvdG9rZW5fdXRpbHMucHk=) | 0.00% | [38 Missing :warning: ](https://app.codecov.io/gh/apache/superset/pull/42824?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
   | [superset/mcp\_service/chart/tool/get\_chart\_data.py](https://app.codecov.io/gh/apache/superset/pull/42824?src=pr&el=tree&filepath=superset%2Fmcp_service%2Fchart%2Ftool%2Fget_chart_data.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c3VwZXJzZXQvbWNwX3NlcnZpY2UvY2hhcnQvdG9vbC9nZXRfY2hhcnRfZGF0YS5weQ==) | 16.66% | [5 Missing :warning: ](https://app.codecov.io/gh/apache/superset/pull/42824?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
   > :exclamation:  There is a different number of reports uploaded between BASE (1de35d1) and HEAD (4087586). Click for more details.
   > 
   > <details><summary>HEAD has 63 uploads less than BASE</summary>
   >
   >| Flag | BASE (1de35d1) | HEAD (4087586) |
   >|------|------|------|
   >|python|55|1|
   >|presto|10|1|
   ></details>
   
   <details><summary>Additional details and impacted files</summary>
   
   
   
   ```diff
   @@            Coverage Diff             @@
   ##           master   #42824      +/-   ##
   ==========================================
   - Coverage   65.73%   56.15%   -9.59%     
   ==========================================
     Files        2843     2843              
     Lines      162672   162727      +55     
     Branches    37255    37267      +12     
   ==========================================
   - Hits       106935    91372   -15563     
   - Misses      53645    70512   +16867     
   + Partials     2092      843    -1249     
   ```
   
   | [Flag](https://app.codecov.io/gh/apache/superset/pull/42824/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Coverage Δ | |
   |---|---|---|
   | [hive](https://app.codecov.io/gh/apache/superset/pull/42824/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
   | [mysql](https://app.codecov.io/gh/apache/superset/pull/42824/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
   | [postgres](https://app.codecov.io/gh/apache/superset/pull/42824/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
   | [presto](https://app.codecov.io/gh/apache/superset/pull/42824/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `39.85% <15.68%> (-0.02%)` | :arrow_down: |
   | [python](https://app.codecov.io/gh/apache/superset/pull/42824/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `39.85% <15.68%> (-19.36%)` | :arrow_down: |
   | [sqlite](https://app.codecov.io/gh/apache/superset/pull/42824/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
   | [unit](https://app.codecov.io/gh/apache/superset/pull/42824/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
   
   Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#carryforward-flags-in-the-pull-request-comment) to find out more.
   </details>
   
   [:umbrella: View full report in Codecov by Harness](https://app.codecov.io/gh/apache/superset/pull/42824?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).   
   :loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
   <details><summary> :rocket: New features to boost your workflow: </summary>
   
   - :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
   - :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
   </details>


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org


Re: [PR] fix(mcp): return all chart query results [superset]

Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on PR #42824:
URL: https://github.com/apache/superset/pull/42824#issuecomment-5271494433

   AI Code Review is in progress (usually takes 3 to 15 minutes unless it's a very large PR).
   
   <details>
   <summary>
   Bito Usage Guide
   </summary>
   
   **Commands**
         
   Type the following command in the pull request comment and save the comment.
         
   - `/review` - Manually triggers a full AI review.
   
   - `/pause` - Pauses automatic reviews on this pull request.
   - `/resume` - Resumes automatic reviews.
   - `/resolve` - Marks all Bito-posted review comments as resolved.
   - `/abort` - Cancels all in-progress reviews.
   
         
   Refer to the <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/available-commands" target="_blank" rel="noopener noreferrer">documentation</a> for additional commands.
         
   **Configuration**
   
   This repository uses `Superset` You can customize the agent settings <a href="https://alpha.bito.ai/home/ai-agents/code-review-agent" target="_blank" rel="noopener noreferrer">here</a> or contact your Bito workspace admin at evan@preset.io.
         
   **Documentation & Help**
   - <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/getting-started/install-run-using-bito-cloud/create-or-customize-an-agent-instance" target="_blank" rel="noopener noreferrer">Customize agent settings</a>
   - <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/implementing-custom-code-review-rules" target="_blank" rel="noopener noreferrer">Review rules</a>
   - <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent" target="_blank" rel="noopener noreferrer">General documentation</a>
   - <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/faqs" target="_blank" rel="noopener noreferrer">FAQ</a>
   
   </details>
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org


Re: [PR] fix(mcp): return all chart query results [superset]

Posted by "gabotorresruiz (via GitHub)" <gi...@apache.org>.
gabotorresruiz commented on code in PR #42824:
URL: https://github.com/apache/superset/pull/42824#discussion_r3784737942


##########
superset/mcp_service/utils/token_utils.py:
##########
@@ -802,6 +802,63 @@ def _truncate_rows_field(
     return None
 
 
+def _truncate_chart_query_results(
+    data: Dict[str, Any], token_limit: int, advice: str
+) -> list[str] | None:
+    """Apply one response-wide row cap to every result of a multi-query chart."""
+    from superset.utils import json as utils_json
+
+    query_results = data.get("query_results")
+    if not isinstance(query_results, list) or not query_results:
+        return None
+
+    row_lists = [data.get("data", [])]
+    row_lists.extend(
+        result.get("data", [])
+        for result in query_results
+        if isinstance(result, dict) and isinstance(result.get("data"), list)
+    )
+    originals = [list(rows) for rows in row_lists]
+    if not any(originals):
+        return None
+
+    original_count = sum(len(rows) for rows in originals)
+    data["_response_truncated"] = True
+    data["_truncation_notes"] = [
+        f"Result truncated from {original_count} rows across multiple queries "
+        f"(limit ~{token_limit:,} tokens). {advice}"
+    ]

Review Comment:
   This block worries me a bit. The placeholder note reserved during the bisect is shorter than the final note swapped in afterwards ("Result truncated from N rows across" vs "Result truncated: K of N rows returned across"), so the search can land a cap that fits with the placeholder but overflows once the final note replaces it. `_truncate_rows_field` reserves the longest-possible note for exactly this reason (see the comment at its placeholder). The middleware re-checks the final size and discards the truncation entirely when it is over, so the caller gets the hard "Response too large" error instead of truncated rows. I verified it empirically on this branch: scanning token limits 300 to 7999 against a 2-query 200-row payload, 101 limits produce a final payload 1 to 4 tokens over the limit with kept > 1, and running the real ResponseSizeGuardMiddleware at token_limit=700 raises the hard ToolError where kept=2 truncation nearly fit. To be fair, this is not a regression: master hard-
 blocks all of these responses anyway. But it is cheap insurance to make the new path succeed everywhere it can:
   
   ```python
       placeholder_note = (
           f"Result truncated: {original_count} of {original_count} rows returned "
           f"across multiple queries (limit ~{token_limit:,} tokens). {advice}"
       )
       data["_truncation_notes"] = [placeholder_note]
   ```
   
   That is provably the longest the real note can be since kept <= original_count. A test like `test_multi_query_truncation_result_fits_budget` asserting `estimate_token_count(dumps(result)) <= token_limit` at a step-boundary limit (700 with the existing middleware-test fixture reproduces it) would lock this in; the existing test misses it because 500 does not land in an overshoot window.



##########
superset/mcp_service/utils/token_utils.py:
##########
@@ -802,6 +802,63 @@ def _truncate_rows_field(
     return None
 
 
+def _truncate_chart_query_results(
+    data: Dict[str, Any], token_limit: int, advice: str
+) -> list[str] | None:
+    """Apply one response-wide row cap to every result of a multi-query chart."""
+    from superset.utils import json as utils_json
+
+    query_results = data.get("query_results")
+    if not isinstance(query_results, list) or not query_results:
+        return None
+
+    row_lists = [data.get("data", [])]
+    row_lists.extend(
+        result.get("data", [])
+        for result in query_results
+        if isinstance(result, dict) and isinstance(result.get("data"), list)
+    )
+    originals = [list(rows) for rows in row_lists]
+    if not any(originals):
+        return None
+
+    original_count = sum(len(rows) for rows in originals)

Review Comment:
   Just a small NIT: `original_count` and `kept_count` sum across `row_lists`, which includes both the top-level `data` list and `query_results[0]["data"]`, and those hold the same rows (the top-level fields are the first-query alias). So the note a 2-query chart with 200 rows per query produces reads "truncated from 600 rows" when there are 400 distinct rows, and the "K of N returned" numbers are inflated the same way. I verified this on the branch: at token_limit=1000 the note reads "Result truncated: 6 of 600 rows returned" while each query actually returns 2 rows. Since the note is written for the LLM to relay to users, I would compute both counts from the `query_results` entries only (while still capping the aliased top-level list in sync, as the code already does).



##########
superset/mcp_service/chart/tool/get_chart_data.py:
##########
@@ -883,6 +913,9 @@ async def get_chart_data(  # noqa: C901
                     chart_type=chart.viz_type or "unknown",
                     columns=columns,
                     data=data[: request.limit] if request.limit else data,
+                    query_results=_build_query_results(

Review Comment:
   Not a blocker, and pre-existing, but very much in the spirit of this PR: the `if not data` guard above (line 736, and its sibling in `_query_from_form_data`) only inspects `result["queries"][0]`. I verified on this branch that a mixed_timeseries chart whose first query returns zero rows (grouped and filtered to empty) while the second query has data returns a NoData ChartError, and the second query's rows are dropped entirely, even though this PR's machinery could now carry them. Are we certain we want a multi-query chart to error when only its primary query is empty? If handling the empty-first-query response shape (top-level columns and data come from query 0) is too invasive here, a named follow-up would be a fine answer.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org


Re: [PR] fix(mcp): return all chart query results [superset]

Posted by "netlify[bot] (via GitHub)" <gi...@apache.org>.
netlify[bot] commented on PR #42824:
URL: https://github.com/apache/superset/pull/42824#issuecomment-5295643556

   ### <span aria-hidden="true">✅</span> Deploy Preview for *superset-docs-preview* ready!
   
   
   |  Name | Link |
   |:-:|------------------------|
   |<span aria-hidden="true">🔨</span> Latest commit | 3d176a838e058c866683f463cfa125c79851d20a |
   |<span aria-hidden="true">🔍</span> Latest deploy log | https://app.netlify.com/projects/superset-docs-preview/deploys/6a7f405c3272520008303565 |
   |<span aria-hidden="true">😎</span> Deploy Preview | [https://deploy-preview-42824--superset-docs-preview.netlify.app](https://deploy-preview-42824--superset-docs-preview.netlify.app) |
   |<span aria-hidden="true">📱</span> Preview on mobile | <details><summary> Toggle QR Code... </summary><br /><br />![QR Code](https://app.netlify.com/qr-code/eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJ1cmwiOiJodHRwczovL2RlcGxveS1wcmV2aWV3LTQyODI0LS1zdXBlcnNldC1kb2NzLXByZXZpZXcubmV0bGlmeS5hcHAifQ.erz2gjcNjRUwi8nWH5rwWtKQbqkVP1GL5yy-qsDN4i4)<br /><br />_Use your smartphone camera to open QR code link._</details> |
   |<span aria-hidden="true">🤖</span> Make changes | [Run an agent on this branch](https://app.netlify.com/projects/superset-docs-preview/agent-runs#DctBDoMgEADAv-y5pBsQbfoDn7HCtpgg4AIHY_r3cpvL3FAkH6XBG-ABm1ByYZiOPX0DeZaaw_4Ubl2SohiVCyRNnZ3lUsK1x1ZH9FxivlY_6kzLZ0LrjF601Yj4MmjsbOH3Bw) |
   ---
   <!-- [superset-docs-preview Preview](https://deploy-preview-42824--superset-docs-preview.netlify.app) -->
   _To edit notification comments on pull requests, go to your [Netlify project configuration](https://app.netlify.com/projects/superset-docs-preview/configuration/notifications#deploy-notifications)._


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org


Re: [PR] fix(mcp): return all chart query results [superset]

Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on PR #42824:
URL: https://github.com/apache/superset/pull/42824#issuecomment-5297273227

   AI Code Review is in progress (usually takes 3 to 15 minutes unless it's a very large PR).
   
   <details>
   <summary>
   Bito Usage Guide
   </summary>
   
   **Commands**
         
   Type the following command in the pull request comment and save the comment.
         
   - `/review` - Manually triggers a full AI review.
   
   - `/pause` - Pauses automatic reviews on this pull request.
   - `/resume` - Resumes automatic reviews.
   - `/resolve` - Marks all Bito-posted review comments as resolved.
   - `/abort` - Cancels all in-progress reviews.
   
         
   Refer to the <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/available-commands" target="_blank" rel="noopener noreferrer">documentation</a> for additional commands.
         
   **Configuration**
   
   This repository uses `Superset` You can customize the agent settings <a href="https://alpha.bito.ai/home/ai-agents/code-review-agent" target="_blank" rel="noopener noreferrer">here</a> or contact your Bito workspace admin at evan@preset.io.
         
   **Documentation & Help**
   - <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/getting-started/install-run-using-bito-cloud/create-or-customize-an-agent-instance" target="_blank" rel="noopener noreferrer">Customize agent settings</a>
   - <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/implementing-custom-code-review-rules" target="_blank" rel="noopener noreferrer">Review rules</a>
   - <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent" target="_blank" rel="noopener noreferrer">General documentation</a>
   - <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/faqs" target="_blank" rel="noopener noreferrer">FAQ</a>
   
   </details>
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org


Re: [PR] fix(mcp): return all chart query results [superset]

Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi merged PR #42824:
URL: https://github.com/apache/superset/pull/42824


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org


Re: [PR] fix(mcp): return all chart query results [superset]

Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on code in PR #42824:
URL: https://github.com/apache/superset/pull/42824#discussion_r3769933627


##########
tests/unit_tests/mcp_service/test_middleware.py:
##########
@@ -471,6 +471,36 @@ async def test_truncates_get_chart_data_rows(self) -> None:
         assert result["_response_truncated"] is True
         assert len(result["data"]) < 200
 
+    @pytest.mark.asyncio
+    async def test_truncates_multi_query_chart_rows_across_whole_response(self) -> None:
+        """All query results share the response's token budget."""
+        middleware = ResponseSizeGuardMiddleware(token_limit=500)
+        context = MagicMock()
+        context.message.name = "get_chart_data"
+        context.message.arguments = {}
+        row = {f"col_{i}": f"value_{i}" for i in range(10)}
+        large_response = {
+            "chart_id": 1,
+            "data": [row] * 200,
+            "row_count": 200,
+            "query_results": [
+                {"query_index": 0, "data": [row] * 200, "row_count": 200},
+                {"query_index": 1, "data": [row] * 200, "row_count": 200},
+            ],
+        }
+        call_next = AsyncMock(return_value=large_response)
+
+        with (
+            patch("superset.mcp_service.middleware.get_user_id", return_value=1),
+            patch("superset.mcp_service.middleware.event_logger"),
+        ):
+            result = await middleware.on_call_tool(context, call_next)
+
+        assert isinstance(result, dict)
+        assert result["_response_truncated"] is True
+        assert all(query["data"] for query in result["query_results"])
+        assert all(len(query["data"]) < 200 for query in result["query_results"])
+

Review Comment:
   <!-- Bito Reply -->
   The regression test added in the latest commit correctly verifies the shared-token-budget logic by asserting that all nested query results are truncated and that their `row_count` values accurately reflect the truncated data length. This addresses the reviewer's concern regarding the potential for silent partial truncation or stale row counts in multi-query chart responses.
   
   **tests/unit_tests/mcp_service/test_middleware.py**
   ```
   assert isinstance(result, dict)
           assert result["_response_truncated"] is True
           assert all(query["data"] for query in result["query_results"])
           assert all(len(query["data"]) < 200 for query in result["query_results"])
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org


Re: [PR] fix(mcp): return all chart query results [superset]

Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on code in PR #42824:
URL: https://github.com/apache/superset/pull/42824#discussion_r3769560877


##########
tests/unit_tests/mcp_service/test_middleware.py:
##########
@@ -471,6 +471,36 @@ async def test_truncates_get_chart_data_rows(self) -> None:
         assert result["_response_truncated"] is True
         assert len(result["data"]) < 200
 
+    @pytest.mark.asyncio
+    async def test_truncates_multi_query_chart_rows_across_whole_response(self) -> None:
+        """All query results share the response's token budget."""
+        middleware = ResponseSizeGuardMiddleware(token_limit=500)
+        context = MagicMock()
+        context.message.name = "get_chart_data"
+        context.message.arguments = {}
+        row = {f"col_{i}": f"value_{i}" for i in range(10)}
+        large_response = {
+            "chart_id": 1,
+            "data": [row] * 200,
+            "row_count": 200,
+            "query_results": [
+                {"query_index": 0, "data": [row] * 200, "row_count": 200},
+                {"query_index": 1, "data": [row] * 200, "row_count": 200},
+            ],
+        }
+        call_next = AsyncMock(return_value=large_response)
+
+        with (
+            patch("superset.mcp_service.middleware.get_user_id", return_value=1),
+            patch("superset.mcp_service.middleware.event_logger"),
+        ):
+            result = await middleware.on_call_tool(context, call_next)
+
+        assert isinstance(result, dict)
+        assert result["_response_truncated"] is True
+        assert all(query["data"] for query in result["query_results"])
+        assert all(len(query["data"]) < 200 for query in result["query_results"])
+

Review Comment:
   <div>
   
   
   <div id="suggestion">
   <div id="issue"><b>Test incomplete on shared budget invariant</b></div>
   <div id="fix">
   
   Test asserts rows are truncated but does not verify the two core properties of the shared-token-budget logic: (1) all query_results are truncated to the same row count, confirming they share a budget; (2) each nested row_count reflects the truncated length, not the original 200. Without these assertions a regression could silently truncate only some results or leave row_count stale.
   </div>
   
   
   </div>
   
   
   
   
   <small><i>Code Review Run #68608e</i></small>
   </div>
   
   ---
   Should Bito avoid suggestions like this for future reviews? (<a href=https://alpha.bito.ai/home/ai-agents/review-rules>Manage Rules</a>)
   - [ ] Yes, avoid them



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org


Re: [PR] fix(mcp): return all chart query results [superset]

Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on PR #42824:
URL: https://github.com/apache/superset/pull/42824#issuecomment-5263459651

   Addressed in 9edfa951c1: nested multi-query `row_count` is derived from the limited data actually returned, and a regression test verifies the behavior for both results.
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org