Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions superset/mcp_service/middleware.py
Original file line number Diff line number Diff line change
Expand Up @@ -246,7 +246,7 @@ def _extract_context_info(
dashboard_id = None
slice_id = None
dataset_id = None
params = getattr(context.message, "params", {}) or {}
params = getattr(context.message, "arguments", {}) or {}
if hasattr(context, "metadata") and context.metadata:
agent_id = context.metadata.get("agent_id")
if not agent_id and hasattr(context, "session") and context.session:
Expand Down Expand Up @@ -1152,7 +1152,7 @@ async def on_call_tool(
)

if estimated_tokens > self.token_limit:
params = getattr(context.message, "params", {}) or {}
params = getattr(context.message, "arguments", {}) or {}
return self._handle_oversized_response(
tool_name, response, estimated_tokens, params
)
Expand Down
56 changes: 28 additions & 28 deletions tests/unit_tests/mcp_service/test_middleware.py
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,7 @@ async def test_allows_small_response(self) -> None:
# Create mock context
context = MagicMock()
context.message.name = "list_charts"
context.message.params = {}
context.message.arguments = {}

# Create mock call_next that returns small response
small_response = {"charts": [{"id": 1, "name": "test"}]}
Expand All @@ -118,7 +118,7 @@ async def test_blocks_large_response(self) -> None:
# Create mock context
context = MagicMock()
context.message.name = "list_charts"
context.message.params = {"page_size": 100}
context.message.arguments = {"page_size": 100}

# Create large response
large_response = {
Expand Down Expand Up @@ -148,7 +148,7 @@ async def test_skips_excluded_tools(self) -> None:
# Create mock context for excluded tool
context = MagicMock()
context.message.name = "health_check"
context.message.params = {}
context.message.arguments = {}

# Create response that would exceed limit
large_response = {"data": "x" * 10000}
Expand All @@ -173,7 +173,7 @@ async def test_logs_warning_at_threshold(self) -> None:

context = MagicMock()
context.message.name = "list_charts"
context.message.params = {}
context.message.arguments = {}

response = {"data": "approaching the limit"}
call_next = AsyncMock(return_value=response)
Expand Down Expand Up @@ -201,7 +201,7 @@ async def test_error_includes_suggestions(self) -> None:

context = MagicMock()
context.message.name = "list_charts"
context.message.params = {"page_size": 100}
context.message.arguments = {"page_size": 100}

large_response = {"charts": [{"id": i} for i in range(1000)]}
call_next = AsyncMock(return_value=large_response)
Expand All @@ -226,7 +226,7 @@ async def test_logs_size_exceeded_event(self) -> None:

context = MagicMock()
context.message.name = "list_charts"
context.message.params = {}
context.message.arguments = {}

large_response = {"data": "x" * 10000}
call_next = AsyncMock(return_value=large_response)
Expand All @@ -250,7 +250,7 @@ async def test_truncates_info_tool_instead_of_blocking(self) -> None:

context = MagicMock()
context.message.name = "get_dataset_info"
context.message.params = {}
context.message.arguments = {}

# Large info tool response with a big description
large_response = {
Expand Down Expand Up @@ -279,7 +279,7 @@ async def test_truncates_chart_info_with_large_form_data(self) -> None:

context = MagicMock()
context.message.name = "get_chart_info"
context.message.params = {}
context.message.arguments = {}

large_response = {
"id": 1,
Expand All @@ -305,7 +305,7 @@ async def test_still_blocks_non_info_tools(self) -> None:

context = MagicMock()
context.message.name = "list_charts" # Not an info tool
context.message.params = {}
context.message.arguments = {}

large_response = {"data": "x" * 10000}
call_next = AsyncMock(return_value=large_response)
Expand All @@ -324,7 +324,7 @@ async def test_logs_truncation_event(self) -> None:

context = MagicMock()
context.message.name = "get_dashboard_info"
context.message.params = {}
context.message.arguments = {}

large_response = {
"id": 1,
Expand Down Expand Up @@ -357,7 +357,7 @@ async def test_truncates_dashboard_info_with_custom_max_list_items(self) -> None

context = MagicMock()
context.message.name = "get_dashboard_info"
context.message.params = {}
context.message.arguments = {}

large_response = {
"id": 1,
Expand Down Expand Up @@ -387,7 +387,7 @@ async def test_truncates_execute_sql_rows_instead_of_blocking(self) -> None:

context = MagicMock()
context.message.name = "execute_sql"
context.message.params = {}
context.message.arguments = {}

row = {f"col_{i}": f"value_{i}" for i in range(10)}
large_response = {
Expand Down Expand Up @@ -417,7 +417,7 @@ async def test_truncates_query_dataset_data_field(self) -> None:

context = MagicMock()
context.message.name = "query_dataset"
context.message.params = {}
context.message.arguments = {}

row = {f"col_{i}": f"value_{i}" for i in range(10)}
large_response = {
Expand Down Expand Up @@ -446,7 +446,7 @@ async def test_truncates_get_chart_data_rows(self) -> None:

context = MagicMock()
context.message.name = "get_chart_data"
context.message.params = {}
context.message.arguments = {}

row = {f"col_{i}": f"value_{i}" for i in range(10)}
large_response = {
Expand Down Expand Up @@ -476,7 +476,7 @@ async def test_data_query_truncation_updates_row_count(self) -> None:

context = MagicMock()
context.message.name = "execute_sql"
context.message.params = {}
context.message.arguments = {}

row = {f"col_{i}": f"value_{i}" for i in range(10)}
large_response = {
Expand All @@ -502,7 +502,7 @@ async def test_data_query_truncation_note_mentions_limit_clause(self) -> None:

context = MagicMock()
context.message.name = "execute_sql"
context.message.params = {}
context.message.arguments = {}

row = {f"col_{i}": f"value_{i}" for i in range(10)}
large_response = {
Expand All @@ -529,7 +529,7 @@ async def test_data_query_truncation_logs_truncation_event(self) -> None:

context = MagicMock()
context.message.name = "execute_sql"
context.message.params = {}
context.message.arguments = {}

row = {f"col_{i}": f"value_{i}" for i in range(10)}
large_response = {
Expand Down Expand Up @@ -557,7 +557,7 @@ async def test_truncates_get_chart_data_csv_export(self) -> None:

context = MagicMock()
context.message.name = "get_chart_data"
context.message.params = {}
context.message.arguments = {}

large_response: dict[str, Any] = {
"chart_id": 1,
Expand Down Expand Up @@ -595,7 +595,7 @@ async def test_data_query_blocks_when_single_row_still_exceeds_limit(self) -> No

context = MagicMock()
context.message.name = "execute_sql"
context.message.params = {}
context.message.arguments = {}

huge_row = {"col": "x" * 5000}
large_response = {
Expand All @@ -619,7 +619,7 @@ async def test_data_query_under_limit_passes_through(self) -> None:

context = MagicMock()
context.message.name = "execute_sql"
context.message.params = {}
context.message.arguments = {}

small_response = {
"status": "success",
Expand Down Expand Up @@ -949,7 +949,7 @@ async def test_info_tool_result_is_truncated_and_rewrapped(self) -> None:
middleware = ResponseSizeGuardMiddleware(token_limit=500)
context = MagicMock()
context.message.name = "get_dataset_info"
context.message.params = {}
context.message.arguments = {}

large_payload = {"id": 1, "table_name": "test", "description": "x" * 50000}
tool_result = self._make_tool_result(large_payload)
Expand All @@ -975,7 +975,7 @@ async def test_small_tool_result_passes_through_unchanged(self) -> None:
middleware = ResponseSizeGuardMiddleware(token_limit=25000)
context = MagicMock()
context.message.name = "get_chart_info"
context.message.params = {}
context.message.arguments = {}

small_payload = {"id": 1, "name": "My Chart"}
tool_result = self._make_tool_result(small_payload)
Expand All @@ -995,7 +995,7 @@ async def test_large_non_info_tool_result_is_blocked(self) -> None:
middleware = ResponseSizeGuardMiddleware(token_limit=100)
context = MagicMock()
context.message.name = "list_charts"
context.message.params = {}
context.message.arguments = {}

large_payload = {
"charts": [{"id": i, "name": f"chart_{i}"} for i in range(500)]
Expand Down Expand Up @@ -1026,7 +1026,7 @@ async def test_data_query_tool_result_is_truncated_and_rewrapped(self) -> None:
middleware = ResponseSizeGuardMiddleware(token_limit=500)
context = MagicMock()
context.message.name = "execute_sql"
context.message.params = {}
context.message.arguments = {}

row = {f"col_{i}": f"value_{i}" for i in range(10)}
large_payload = {
Expand Down Expand Up @@ -1058,7 +1058,7 @@ async def test_meta_preserved_after_truncation(self) -> None:
middleware = ResponseSizeGuardMiddleware(token_limit=500)
context = MagicMock()
context.message.name = "get_dashboard_info"
context.message.params = {}
context.message.arguments = {}

meta = {"request_id": "abc-123"}
large_payload = {"id": 1, "title": "My Dashboard", "description": "x" * 50000}
Expand Down Expand Up @@ -1093,7 +1093,7 @@ class ChartInfo(BaseModel):

context = MagicMock()
context.message.name = "get_chart_info"
context.message.params = {}
context.message.arguments = {}

response = ChartInfo(id=1, name="Test Chart")
call_next = AsyncMock(return_value=response)
Expand All @@ -1113,7 +1113,7 @@ async def test_list_response(self) -> None:

context = MagicMock()
context.message.name = "list_charts"
context.message.params = {}
context.message.arguments = {}

response = [{"id": 1}, {"id": 2}, {"id": 3}]
call_next = AsyncMock(return_value=response)
Expand All @@ -1133,7 +1133,7 @@ async def test_string_response(self) -> None:

context = MagicMock()
context.message.name = "health_check"
context.message.params = {}
context.message.arguments = {}

response = "OK"
call_next = AsyncMock(return_value=response)
Expand Down
29 changes: 28 additions & 1 deletion tests/unit_tests/mcp_service/test_middleware_logging.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ def _make_context(
ctx.method = method
message = MagicMock()
message.name = name
message.params = params or {}
message.arguments = params or {}
ctx.message = message
if metadata is not None:
ctx.metadata = metadata
Expand Down Expand Up @@ -503,6 +503,33 @@ def test_extract_slice_id_from_slice_id(self, mock_get_user_id) -> None:

assert slice_id == 66

@patch("superset.mcp_service.middleware.get_user_id", return_value=1)
def test_extract_reads_arguments_on_real_call_tool_request_params(
self, mock_get_user_id
) -> None:
"""Regression test: the real MCP ``CallToolRequestParams`` object
exposes tool arguments as ``.arguments``, not ``.params`` -- a
``MagicMock``-based context would auto-vivify a ``.params``
attribute and hide a mismatch. Using the real SDK type here
ensures params/dashboard_id/etc. are actually populated instead
of silently logging as empty."""
middleware = LoggingMiddleware()
message = mt.CallToolRequestParams(
name="get_dashboard_info",
arguments={"dashboard_id": 7},
)
ctx = MagicMock()
ctx.message = message
ctx.metadata = None
ctx.session = None

agent_id, user_id, dashboard_id, slice_id, dataset_id, params = (
middleware._extract_context_info(ctx)
)

assert params == {"dashboard_id": 7}
assert dashboard_id == 7


class TestIsErrorResponse:
"""Tests for LoggingMiddleware._is_error_response()."""
Expand Down
Loading