From 7cb840a70509b9bbc1d1edb55253169db1a80168 Mon Sep 17 00:00:00 2001 From: Raghvendra Singh Date: Thu, 24 Sep 2026 16:08:41 +0530 Subject: [PATCH] fix(tools): derive the tool-schema guard limits from the result ceiling DomainDispatcher compiles every child's output schema with a bare ToolSchemaGuard(), so the guard's instance limits stay at the library defaults (1 MiB, 10,000 nodes, 256 KiB per string) whatever PerformanceConfig.max_result_bytes says. A result the query runtime had already bounded to a larger configured ceiling was then refused one step later as CHILD_RESULT_INVALID, with a message that did not say why. - SchemaLimits.for_result_bytes(max_result_bytes) derives limits from the ceiling: instance = ceiling + the base instance budget (the envelope around the rows), nodes = instance / 2 so the byte limit stays the binding one, one string may be as long as the result; nothing shrinks below the base. - DorisToolsManager resolves its limits from configured_result_limits() and hands the dispatcher a guard built from them; an explicit schema_limits argument wins. - create_doris_mcp_server follows the tools manager's limits when no schema_limits is passed (an explicit value still wins). - The CHILD_RESULT_INVALID envelope carries the validator's reason in details.reason ("exceeds output limits" vs "does not match output schema"). Tests cover the derivation, the default-vs-derived guard on a result inside the default ceiling, the manager and factory resolution, and the reason on the error envelope. --- CHANGELOG.md | 6 +++ doris_mcp_server/protocol.py | 14 +++++- doris_mcp_server/schema_validation.py | 32 ++++++++++++- doris_mcp_server/tools/domain_dispatcher.py | 4 +- doris_mcp_server/tools/tools_manager.py | 13 ++++++ test/protocol/test_mcp_v2_protocol.py | 17 +++++++ test/protocol/test_schema_validation.py | 50 +++++++++++++++++++++ test/tools/test_domain_dispatcher.py | 31 +++++++++++++ test/tools/test_tools_manager.py | 37 ++++++++++++++- 9 files changed, 198 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6b5363d..88ffd0a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -69,6 +69,12 @@ under **Unreleased** until a new version is selected and published. ### Fixed +- Derived the tool-schema guard's instance limits from the configured result + ceiling (`PerformanceConfig.max_result_bytes`) in `DorisToolsManager` and + `create_doris_mcp_server`, so a result the query runtime has already bounded + is no longer refused as `CHILD_RESULT_INVALID` by the guard's own defaults; + the refusal now carries its reason in `details.reason`, and + `DorisToolsManager` accepts an explicit `schema_limits`. - Preserved fail-closed distribution provenance across version parsing, brandless FE/BE component fallback, and certification evidence validation, preventing unknown or mismatched brands from inheriting Apache Doris diff --git a/doris_mcp_server/protocol.py b/doris_mcp_server/protocol.py index 5e79bc3..2a0d8fe 100644 --- a/doris_mcp_server/protocol.py +++ b/doris_mcp_server/protocol.py @@ -288,6 +288,12 @@ def _tools_for_protocol( ] +def _manager_schema_limits(tools_manager: object) -> SchemaLimits: + """The limits a tools manager resolved from its configuration, else defaults.""" + limits = getattr(tools_manager, "schema_limits", None) + return limits if isinstance(limits, SchemaLimits) else DEFAULT_SCHEMA_LIMITS + + def create_doris_mcp_server( *, resources_manager: ResourcesManager, @@ -299,7 +305,7 @@ def create_doris_mcp_server( list_page_size: int = DEFAULT_LIST_PAGE_SIZE, state_handle_secret: str | bytes | None = None, state_handle_ttl_seconds: int = DEFAULT_STATE_HANDLE_TTL_SECONDS, - schema_limits: SchemaLimits = DEFAULT_SCHEMA_LIMITS, + schema_limits: SchemaLimits | None = None, required_client_capabilities: Mapping[str, ClientCapabilities] | None = None, required_tool_capabilities: Mapping[str, ClientCapabilities] | None = None, ) -> Server: @@ -314,7 +320,11 @@ def create_doris_mcp_server( ), default_ttl_seconds=state_handle_ttl_seconds, ) - schema_guard = ToolSchemaGuard(schema_limits) + schema_guard = ToolSchemaGuard( + schema_limits + if schema_limits is not None + else _manager_schema_limits(tools_manager) + ) async def list_resources( ctx: ServerRequestContext, diff --git a/doris_mcp_server/schema_validation.py b/doris_mcp_server/schema_validation.py index 5359b70..cc3b1aa 100644 --- a/doris_mcp_server/schema_validation.py +++ b/doris_mcp_server/schema_validation.py @@ -22,7 +22,7 @@ import json from collections import OrderedDict from collections.abc import Mapping, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import Any from urllib.parse import unquote @@ -59,6 +59,36 @@ class SchemaLimits: max_validation_errors: int = 16 max_cached_schemas: int = 256 + @classmethod + def for_result_bytes( + cls, + max_result_bytes: int, + *, + base: SchemaLimits | None = None, + ) -> SchemaLimits: + """Return limits that admit a structured result of ``max_result_bytes``. + + Query children bound their rows by the configured result ceiling + (``PerformanceConfig.max_result_bytes``), while the output guard applies + instance limits of its own. When the two disagree, the guard refuses a + result the executor has already accepted and truncated. This derives the + guard from the ceiling instead: the instance budget grows to the ceiling + plus ``base.max_instance_bytes`` (the envelope around the rows: column + descriptors, metadata, the result wrapper), the node budget follows the + bytes (a JSON node is never shorter than two bytes, so the byte limit + stays the binding one), and one string may be as long as the result + itself. Every other limit is ``base``'s, and nothing shrinks below it. + """ + base = base if base is not None else DEFAULT_SCHEMA_LIMITS + ceiling = max(int(max_result_bytes), 0) + instance_bytes = max(base.max_instance_bytes, ceiling + base.max_instance_bytes) + return replace( + base, + max_instance_bytes=instance_bytes, + max_instance_nodes=max(base.max_instance_nodes, instance_bytes // 2), + max_instance_string_length=max(base.max_instance_string_length, ceiling), + ) + DEFAULT_SCHEMA_LIMITS = SchemaLimits() diff --git a/doris_mcp_server/tools/domain_dispatcher.py b/doris_mcp_server/tools/domain_dispatcher.py index cd0e26b..9fed8b1 100644 --- a/doris_mcp_server/tools/domain_dispatcher.py +++ b/doris_mcp_server/tools/domain_dispatcher.py @@ -791,7 +791,7 @@ async def _execute( "status_code": exc.status_code, }, ) - except ToolOutputValidationError: + except ToolOutputValidationError as exc: logger.exception("Formal child output validation failed for %s", feature_id) self._audit( feature_id, @@ -805,7 +805,7 @@ async def _execute( "Child result did not match the declared schema.", child_tool=child.name, manifest_version=manifest.manifest_version, - details={"rediscover": False}, + details={"rediscover": False, "reason": str(exc)}, ) except Exception as exc: logger.error( diff --git a/doris_mcp_server/tools/tools_manager.py b/doris_mcp_server/tools/tools_manager.py index 6a5dda1..2ddaad6 100644 --- a/doris_mcp_server/tools/tools_manager.py +++ b/doris_mcp_server/tools/tools_manager.py @@ -27,6 +27,8 @@ from mcp.types import Tool from ..auth.operation_policy import authorize_operation +from ..result_limits import configured_result_limits +from ..schema_validation import SchemaLimits, ToolSchemaGuard from ..state_handles import StateHandleCodec from ..utils.adbc_query_tools import DorisADBCQueryTools from ..utils.analysis_tools import SQLAnalyzer, TableAnalyzer @@ -96,9 +98,19 @@ def __init__( tool_providers: Iterable[CustomToolProvider] | None = None, domain_availability_provider: DomainAvailabilityProvider | None = None, tool_exposure_mode: str | ToolExposureMode | None = None, + schema_limits: SchemaLimits | None = None, ) -> None: self.connection_manager = connection_manager config = getattr(connection_manager, "config", None) + # The output guard follows the configured result ceiling: a result the + # executor was allowed to return must be one the guard accepts. + self.schema_limits = ( + schema_limits + if schema_limits is not None + else SchemaLimits.for_result_bytes( + configured_result_limits(config).max_bytes + ) + ) self._tool_provider_runtime = ToolProviderRuntime.create(config, tool_providers) self.query_executor = DorisQueryExecutor(connection_manager) self.table_analyzer = TableAnalyzer(connection_manager) @@ -216,6 +228,7 @@ def __init__( self._domain_dispatcher = DomainDispatcher( self, self._domain_manifest_service, + schema_guard=ToolSchemaGuard(self.schema_limits), state_handle_codec=StateHandleCodec( state_handle_secret, default_ttl_seconds=state_handle_ttl, diff --git a/test/protocol/test_mcp_v2_protocol.py b/test/protocol/test_mcp_v2_protocol.py index f46cd51..ba24f2e 100644 --- a/test/protocol/test_mcp_v2_protocol.py +++ b/test/protocol/test_mcp_v2_protocol.py @@ -200,6 +200,23 @@ async def get_prompt(self, name: str, arguments: dict) -> GetPromptResult: ) +def test_server_schema_limits_follow_the_tools_manager(): + from doris_mcp_server.protocol import _manager_schema_limits + from doris_mcp_server.schema_validation import DEFAULT_SCHEMA_LIMITS, SchemaLimits + + class WideToolsManager(StubToolsManager): + schema_limits = SchemaLimits.for_result_bytes(16 * 1024 * 1024) + + class OddToolsManager(StubToolsManager): + schema_limits = "not limits" + + assert _manager_schema_limits(WideToolsManager()) == SchemaLimits.for_result_bytes( + 16 * 1024 * 1024 + ) + assert _manager_schema_limits(StubToolsManager()) is DEFAULT_SCHEMA_LIMITS + assert _manager_schema_limits(OddToolsManager()) is DEFAULT_SCHEMA_LIMITS + + def create_test_server( required_client_capabilities: dict[str, ClientCapabilities] | None = None, required_tool_capabilities: dict[str, ClientCapabilities] | None = None, diff --git a/test/protocol/test_schema_validation.py b/test/protocol/test_schema_validation.py index f082f93..6a79187 100644 --- a/test/protocol/test_schema_validation.py +++ b/test/protocol/test_schema_validation.py @@ -23,7 +23,9 @@ import pytest from mcp.types import Tool +from doris_mcp_server.result_limits import DEFAULT_MAX_RESULT_BYTES from doris_mcp_server.schema_validation import ( + DEFAULT_SCHEMA_LIMITS, JSON_SCHEMA_2020_12, SchemaLimits, ToolArgumentsValidationError, @@ -238,3 +240,51 @@ def test_output_schema_requires_structured_content(): compiled = ToolSchemaGuard().compile_tool(_advanced_tool()) with pytest.raises(ToolOutputValidationError, match="no structured content"): compiled.validate_output(None) + + +def test_limits_for_result_bytes_follow_the_result_ceiling(): + base = DEFAULT_SCHEMA_LIMITS + ceiling = 16 * 1024 * 1024 + limits = SchemaLimits.for_result_bytes(ceiling) + + assert limits.max_instance_bytes == ceiling + base.max_instance_bytes + assert limits.max_instance_nodes == limits.max_instance_bytes // 2 + assert limits.max_instance_string_length == ceiling + assert limits.max_instance_depth == base.max_instance_depth + assert limits.max_schema_bytes == base.max_schema_bytes + + # Nothing shrinks below the base, whatever the ceiling. + small = SchemaLimits.for_result_bytes(0) + assert small.max_instance_bytes == base.max_instance_bytes + assert small.max_instance_nodes == base.max_instance_bytes // 2 + assert small.max_instance_string_length == base.max_instance_string_length + + custom = SchemaLimits.for_result_bytes( + 1024, base=SchemaLimits(max_instance_bytes=4096, max_instance_nodes=4) + ) + assert custom.max_instance_bytes == 4096 + 1024 + assert custom.max_instance_nodes == (4096 + 1024) // 2 + + +def test_default_guard_refuses_a_result_inside_the_default_result_ceiling(): + tool = Tool( + name="rows", + description="Rows under the default result ceiling.", + input_schema={"type": "object"}, + output_schema={"type": "object"}, + ) + # 6,000 small rows: far below DEFAULT_MAX_RESULT_BYTES, above the default + # node budget; one 300 KiB cell: below the ceiling, above the string budget. + payloads = [ + {"rows": [{"n": index} for index in range(6_000)]}, + {"cell": "x" * 300_000}, + ] + default_schema = ToolSchemaGuard().compile_tool(tool) + derived_schema = ToolSchemaGuard( + SchemaLimits.for_result_bytes(DEFAULT_MAX_RESULT_BYTES) + ).compile_tool(tool) + + for payload in payloads: + with pytest.raises(ToolOutputValidationError, match="exceeds output limits"): + default_schema.validate_output(payload) + derived_schema.validate_output(payload) diff --git a/test/tools/test_domain_dispatcher.py b/test/tools/test_domain_dispatcher.py index 3cbedef..2e12c88 100644 --- a/test/tools/test_domain_dispatcher.py +++ b/test/tools/test_domain_dispatcher.py @@ -1203,6 +1203,37 @@ async def test_output_schema_violation_becomes_stable_server_error() -> None: assert result.mode == "error" assert result.error.code is DomainErrorCode.CHILD_RESULT_INVALID + assert "output schema" in result.error.details["reason"] + + +@pytest.mark.asyncio +async def test_output_over_the_guard_limits_names_the_limit_in_the_error() -> None: + manager = DorisToolsManager( + _connection_manager(), + domain_availability_provider=SelectiveAvailabilityProvider( + {"doris_query.execute_query"} + ), + schema_limits=SchemaLimits(max_instance_nodes=8), + ) + manager._exec_query_tool = AsyncMock(return_value=_successful_query()) + + with patch.object( + manager.domain_dispatcher, + "_normalize_child_result", + return_value={"rows": [{"n": index} for index in range(64)]}, + ): + result = await manager.domain_dispatcher.call_domain( + "doris_query", + { + "child_tool": "execute_query", + "arguments": {"sql": "SELECT 1"}, + }, + None, + ) + + assert result.mode == "error" + assert result.error.code is DomainErrorCode.CHILD_RESULT_INVALID + assert "output limits" in result.error.details["reason"] @pytest.mark.asyncio diff --git a/test/tools/test_tools_manager.py b/test/tools/test_tools_manager.py index 119f38f..d3e3769 100644 --- a/test/tools/test_tools_manager.py +++ b/test/tools/test_tools_manager.py @@ -23,7 +23,8 @@ import pytest -from doris_mcp_server.schema_validation import ToolSchemaGuard +from doris_mcp_server.result_limits import DEFAULT_MAX_RESULT_BYTES +from doris_mcp_server.schema_validation import SchemaLimits, ToolSchemaGuard from doris_mcp_server.tools.domain_dispatcher import ToolNotFoundError from doris_mcp_server.tools.doris_feature_matrix import ( EXPECTED_DOMAIN_CHILDREN, @@ -394,3 +395,37 @@ async def test_all_production_tool_schemas_compile_as_bounded_2020_12( compiled = ToolSchemaGuard().compile_catalog(tools) assert set(compiled) == {tool.name for tool in tools} + + +def _connection_manager_with_result_bytes(max_result_bytes): + connection_manager = Mock() + connection_manager.get_connection = AsyncMock() + connection_manager.config.adbc.default_max_rows = 1000 + connection_manager.config.adbc.default_timeout = 30 + connection_manager.config.adbc.default_return_format = "dict" + connection_manager.config.performance.max_result_bytes = max_result_bytes + return connection_manager + + +def test_schema_limits_follow_the_configured_result_ceiling(): + manager = DorisToolsManager(_connection_manager_with_result_bytes(16 * 1024 * 1024)) + + assert manager.schema_limits == SchemaLimits.for_result_bytes(16 * 1024 * 1024) + + +def test_schema_limits_default_to_the_default_result_ceiling(): + manager = DorisToolsManager(_connection_manager_with_result_bytes(None)) + + assert manager.schema_limits == SchemaLimits.for_result_bytes( + DEFAULT_MAX_RESULT_BYTES + ) + + +def test_explicit_schema_limits_win_over_the_configured_ceiling(): + explicit = SchemaLimits(max_instance_nodes=8) + manager = DorisToolsManager( + _connection_manager_with_result_bytes(16 * 1024 * 1024), + schema_limits=explicit, + ) + + assert manager.schema_limits is explicit