diff --git a/pyproject.toml b/pyproject.toml index 0f40469..bf6f4d2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,7 +16,8 @@ dependencies = [ "pyyaml>=6.0", # config & tool-spec parsing "httpx>=0.27.0", # HTTP client; API embeddings "questionary>=2.0", # interactive `dsagt init` select/checkbox menus - "mcp>=1.0.0,<2.0.0", # MCP server protocol + "mcp>=2.0.0,<3.0.0", # MCP server protocol + "jsonschema>=4.0", # tool-argument validation in the dispatch shell "mlflow==3.11.1", # trace store & observability # Knowledge base # torch 2.2.2 (latest available for Intel Mac) was compiled against NumPy 1.x diff --git a/src/dsagt/mcp/server.py b/src/dsagt/mcp/server.py index 7010476..ecc22fe 100644 --- a/src/dsagt/mcp/server.py +++ b/src/dsagt/mcp/server.py @@ -35,6 +35,7 @@ import threading # noqa: E402 from pathlib import Path # noqa: E402 +import jsonschema # noqa: E402 import yaml # noqa: E402 import mcp.server.stdio # noqa: E402 @@ -60,12 +61,23 @@ def build_dispatch_server( ) -> Server: """Wrap a ``(tools, handlers)`` pair in a configured MCP ``Server``. - One dispatch contract for every concern module: catch + wrap errors, then - format by return type — a handler that returns ``str`` passes through, one - that returns ``dict`` is JSON-encoded. This is a superset of the old - per-server behavior (registry handlers returned ``str`` and never raised; - knowledge handlers returned ``dict`` and raised ``ValueError`` on bad - input), so it is behavior-preserving for both. + One dispatch contract for every concern module: reject what the tool's own + ``input_schema`` does not admit, run the handler, catch + wrap what it + raises, then format by return type — a handler that returns ``str`` passes + through, one that returns ``dict`` is JSON-encoded. Registry handlers + return ``str`` and never raise; knowledge handlers return ``dict`` and raise + ``ValueError`` on bad input; both are covered. + + Argument validation and the outer error boundary live here because the SDK + stopped providing them: through mcp 1.x the ``@server.call_tool()`` decorator + validated against ``inputSchema`` and turned any escaping exception into an + error result, and the v2 lowlevel server does neither. Two consequences + shape the code below. Nothing may escape this function — the v2 runner + converts an exception into a JSON-RPC protocol error, which tears down the + request instead of handing the agent something it can read and retry — and + every rejection carries ``is_error``, the only signal on the wire that a + call failed. The tool name is client-controlled, so an unknown one is a + rejection, not a bug. ``tool_category`` maps tool name → concern (``memory`` / ``skill`` / ``knowledge`` / ``registry``). Each call opens one categorization-root span @@ -76,15 +88,47 @@ def build_dispatch_server( in the single-concern test servers / one-shot tools. """ tool_category = tool_category or {} - server = Server(name) - - @server.list_tools() - async def list_tools() -> list[types.Tool]: - return tools + schemas = {tool.name: tool.input_schema for tool in tools} + + async def on_list_tools(ctx, params) -> types.ListToolsResult: + return types.ListToolsResult(tools=tools) + + def rejected(message: str) -> types.CallToolResult: + """A rejection the agent can read and retry from. + + ``is_error`` is what marks a result as failed on the wire; without it a + client renders a rejection as a successful call and the agent has no + signal to correct itself. + """ + error = {"status": "error", "error": message} + return types.CallToolResult( + content=[ + types.TextContent( + type="text", text=json.dumps(error, ensure_ascii=False) + ) + ], + is_error=True, + ) - @server.call_tool() - async def call_tool(tool_name: str, arguments: dict) -> list[types.TextContent]: - handler = handlers[tool_name] # KeyError = bug in list_tools schema + async def on_call_tool( + ctx, params: types.CallToolRequestParams + ) -> types.CallToolResult: + tool_name = params.name + # ``arguments`` is optional in the protocol (None when omitted), and the + # mcp server does not validate against input_schema before dispatch — + # reject malformed calls here so handlers can assume valid input. + arguments = params.arguments or {} + # The tool name is client-controlled — an agent inventing one, or holding + # a stale name across a restart, must get a rejection back rather than an + # escaping KeyError, which the runner turns into a JSON-RPC protocol + # error that tears down the request instead of informing the agent. + if tool_name not in handlers: + return rejected(f"Unknown tool: {tool_name}") + handler = handlers[tool_name] + try: + jsonschema.validate(instance=arguments, schema=schemas[tool_name]) + except jsonschema.ValidationError as e: + return rejected(f"Input validation error: {e.message}") with open_span(tool_name, source=tool_category.get(tool_name)) as span: try: result = await handler(arguments) @@ -99,14 +143,18 @@ async def call_tool(tool_name: str, arguments: dict) -> list[types.TextContent]: # the MLflow UI shows them instead of a null request. span.set_inputs(arguments) span.set_outputs(result) - text = ( - result - if isinstance(result, str) - else json.dumps(result, ensure_ascii=False) - ) - return [types.TextContent(type="text", text=text)] + try: + text = ( + result + if isinstance(result, str) + else json.dumps(result, ensure_ascii=False) + ) + except (TypeError, ValueError) as e: + logger.exception("Tool '%s' returned an unserializable result", tool_name) + return rejected(f"Unserializable result from {tool_name}: {e}") + return types.CallToolResult(content=[types.TextContent(type="text", text=text)]) - return server + return Server(name, on_list_tools=on_list_tools, on_call_tool=on_call_tool) HEARTBEAT_INTERVAL_S = 45.0 diff --git a/tests/mcp_helpers.py b/tests/mcp_helpers.py index 007e371..5748cf5 100644 --- a/tests/mcp_helpers.py +++ b/tests/mcp_helpers.py @@ -21,13 +21,10 @@ def call_tool_sync(server, name: str, arguments: dict) -> str: """Invoke a tool handler on an MCP server and return the response text.""" - req = types.CallToolRequest( - method="tools/call", - params=types.CallToolRequestParams(name=name, arguments=arguments), - ) - handler = server.request_handlers[types.CallToolRequest] - result = asyncio.run(handler(req)) - return result.root.content[0].text + params = types.CallToolRequestParams(name=name, arguments=arguments) + handler = server.get_request_handler("tools/call").handler + result = asyncio.run(handler(None, params)) + return result.content[0].text def call_tool_json(server, name: str, arguments: dict) -> dict: @@ -37,13 +34,10 @@ def call_tool_json(server, name: str, arguments: dict) -> dict: async def call_tool_async(server, name: str, arguments: dict) -> str: """Invoke a tool handler inside a running event loop.""" - req = types.CallToolRequest( - method="tools/call", - params=types.CallToolRequestParams(name=name, arguments=arguments), - ) - handler = server.request_handlers[types.CallToolRequest] - result = await handler(req) - return result.root.content[0].text + params = types.CallToolRequestParams(name=name, arguments=arguments) + handler = server.get_request_handler("tools/call").handler + result = await handler(None, params) + return result.content[0].text # --------------------------------------------------------------------------- diff --git a/tests/test_dsagt_server.py b/tests/test_dsagt_server.py index 87c7ffe..62f627e 100644 --- a/tests/test_dsagt_server.py +++ b/tests/test_dsagt_server.py @@ -16,9 +16,10 @@ import mcp.types as types import pytest +from mcp_helpers import call_tool_sync from dsagt.mcp.server import _build_kb_from_config, create_dsagt_server -from dsagt.registry import SkillRegistry, CodeRegistry +from dsagt.registry import CodeRegistry, SkillRegistry def _make_merged_server(tmp_path: Path): @@ -35,19 +36,13 @@ def _make_merged_server(tmp_path: Path): def _list_tools(server) -> list[str]: - handler = server.request_handlers[types.ListToolsRequest] - res = asyncio.run(handler(types.ListToolsRequest(method="tools/list"))) - return sorted(t.name for t in res.root.tools) + handler = server.get_request_handler("tools/list").handler + res = asyncio.run(handler(None, None)) + return sorted(t.name for t in res.tools) def _call(server, name: str, arguments: dict) -> str: - handler = server.request_handlers[types.CallToolRequest] - req = types.CallToolRequest( - method="tools/call", - params=types.CallToolRequestParams(name=name, arguments=arguments), - ) - res = asyncio.run(handler(req)) - return res.root.content[0].text + return call_tool_sync(server, name, arguments) def test_merged_server_exposes_all_tools(tmp_path): @@ -135,6 +130,103 @@ def test_dict_returning_handler_is_json_encoded(tmp_path): assert "sources" in parsed +class TestInputValidation: + """The dispatch shell validates arguments against the tool's input schema. + + The mcp 2.x server invokes ``on_call_tool`` without validating arguments + (and ``params.arguments`` is None when omitted), so the shell must reject + malformed calls before they reach a handler. + """ + + def _server(self): + from dsagt.mcp.server import build_dispatch_server + + async def echo(args): + return {"echoed": args["q"]} + + tools = [ + types.Tool( + name="demo", + description="d", + inputSchema={ + "type": "object", + "properties": {"q": {"type": "string"}}, + "required": ["q"], + }, + ) + ] + return build_dispatch_server("test", tools, {"demo": echo}) + + def test_missing_required_argument_rejected(self): + out = json.loads(_call(self._server(), "demo", {})) + assert out["status"] == "error" + assert "Input validation error" in out["error"] + assert "'q' is a required property" in out["error"] + + def test_omitted_arguments_rejected(self): + server = self._server() + handler = server.get_request_handler("tools/call").handler + params = types.CallToolRequestParams(name="demo") # arguments is None + res = asyncio.run(handler(None, params)) + out = json.loads(res.content[0].text) + assert out["status"] == "error" + assert "'q' is a required property" in out["error"] + + def test_wrong_type_rejected(self): + out = json.loads(_call(self._server(), "demo", {"q": 7})) + assert out["status"] == "error" + assert "Input validation error" in out["error"] + + def test_valid_arguments_dispatch(self): + out = json.loads(_call(self._server(), "demo", {"q": "hello"})) + assert out == {"echoed": "hello"} + + def _call_raw(self, server, name, arguments): + handler = server.get_request_handler("tools/call").handler + params = types.CallToolRequestParams(name=name, arguments=arguments) + return asyncio.run(handler(None, params)) + + def test_rejection_is_flagged_is_error(self): + """``is_error`` is the only signal a client has that a call failed. + + Without it a rejected call renders as a successful tool result and the + agent has nothing to correct itself from. + """ + res = self._call_raw(self._server(), "demo", {}) + assert res.is_error is True + + def test_successful_call_is_not_flagged(self): + res = self._call_raw(self._server(), "demo", {"q": "hello"}) + assert res.is_error is False + + def test_unknown_tool_is_rejected_not_raised(self): + """The tool name is client-controlled, so an unknown one must come back + as a readable rejection — an escaping KeyError becomes a JSON-RPC + protocol error that tears down the request instead.""" + res = self._call_raw(self._server(), "no_such_tool", {}) + assert res.is_error is True + out = json.loads(res.content[0].text) + assert out["status"] == "error" + assert "Unknown tool: no_such_tool" in out["error"] + + def test_unserializable_result_is_rejected_not_raised(self): + """A handler returning non-JSON data must not escape as a protocol + error either — ``json.dumps`` runs after the handler's own guard.""" + from dsagt.mcp.server import build_dispatch_server + + tools = [ + types.Tool(name="bad", description="d", inputSchema={"type": "object"}) + ] + + async def bad(args): + return {"obj": object()} + + server = build_dispatch_server("test", tools, {"bad": bad}) + res = self._call_raw(server, "bad", {}) + assert res.is_error is True + assert "Unserializable result" in json.loads(res.content[0].text)["error"] + + class TestBuildKbFromConfig: """``_build_kb_from_config`` validates embedding config before building a KB. diff --git a/tests/test_kb_search_filters.py b/tests/test_kb_search_filters.py index c9e7b4c..91df32f 100644 --- a/tests/test_kb_search_filters.py +++ b/tests/test_kb_search_filters.py @@ -9,10 +9,9 @@ from unittest.mock import MagicMock import pytest -import mcp.types as types +from mcp_helpers import call_tool_json as call_tool from dsagt.mcp.knowledge_tools import create_knowledge_server -from mcp_helpers import call_tool_json as call_tool def make_search_result(text, source_file, chunk_index=0, score=0.9, extra_meta=None): @@ -285,12 +284,11 @@ def test_empty_metadata_for_reference_collections(self, mock_kb): class TestSearchSchemaFilters: def _get_kb_search_schema(self, server): - req = types.ListToolsRequest(method="tools/list") - handler = server.request_handlers[types.ListToolsRequest] - result = asyncio.run(handler(req)) - for tool in result.root.tools: + handler = server.get_request_handler("tools/list").handler + result = asyncio.run(handler(None, None)) + for tool in result.tools: if tool.name == "kb_search": - return tool.inputSchema + return tool.input_schema raise AssertionError("kb_search not found") def test_filter_params_in_schema(self, server): diff --git a/tests/test_knowledge_server.py b/tests/test_knowledge_server.py index 0f35356..d054691 100644 --- a/tests/test_knowledge_server.py +++ b/tests/test_knowledge_server.py @@ -16,21 +16,16 @@ from unittest.mock import MagicMock import pytest -import mcp.types as types +from mcp_helpers import call_tool_async +from mcp_helpers import call_tool_json as call_tool from dsagt.mcp.knowledge_tools import create_knowledge_server, setup_runtime_kb -from mcp_helpers import call_tool_json as call_tool async def _call_tool_async(server, name: str, arguments: dict) -> dict: """Invoke a tool handler inside a running event loop.""" - req = types.CallToolRequest( - method="tools/call", - params=types.CallToolRequestParams(name=name, arguments=arguments), - ) - handler = server.request_handlers[types.CallToolRequest] - result = await handler(req) - return json.loads(result.root.content[0].text) + result = await call_tool_async(server, name, arguments) + return json.loads(result) async def call_tool_and_await_job( @@ -757,6 +752,7 @@ class TestOpenMPWorkaround: def test_kmp_duplicate_lib_ok_is_set(self): """KMP_DUPLICATE_LIB_OK is set after importing dsagt.mcp.knowledge_tools.""" import os + import dsagt.mcp.knowledge_tools # noqa: F401 assert os.environ.get("KMP_DUPLICATE_LIB_OK") == "TRUE" @@ -774,12 +770,11 @@ class TestRerankSchemaDefault: def _get_rerank_default(self, server): """Extract the rerank default from the kb_search tool schema.""" - req = types.ListToolsRequest(method="tools/list") - handler = server.request_handlers[types.ListToolsRequest] - result = asyncio.run(handler(req)) - for tool in result.root.tools: + handler = server.get_request_handler("tools/list").handler + result = asyncio.run(handler(None, None)) + for tool in result.tools: if tool.name == "kb_search": - return tool.inputSchema["properties"]["rerank"]["default"] + return tool.input_schema["properties"]["rerank"]["default"] raise AssertionError("kb_search tool not found") def test_rerank_default_from_kb(self, mock_kb): @@ -867,18 +862,17 @@ def test_multi_collection_merges_results(self, server, mock_kb): class TestKbSearchSchema: def _get_tool(self, server, name): - req = types.ListToolsRequest(method="tools/list") - handler = server.request_handlers[types.ListToolsRequest] - result = asyncio.run(handler(req)) - for tool in result.root.tools: + handler = server.get_request_handler("tools/list").handler + result = asyncio.run(handler(None, None)) + for tool in result.tools: if tool.name == name: return tool return None def test_kb_search_has_collections_param(self, server): tool = self._get_tool(server, "kb_search") - assert "collections" in tool.inputSchema["properties"] + assert "collections" in tool.input_schema["properties"] def test_kb_search_query_is_only_required(self, server): tool = self._get_tool(server, "kb_search") - assert tool.inputSchema["required"] == ["query"] + assert tool.input_schema["required"] == ["query"] diff --git a/tests/test_memory_tools.py b/tests/test_memory_tools.py index 4b05111..59dea72 100644 --- a/tests/test_memory_tools.py +++ b/tests/test_memory_tools.py @@ -11,11 +11,10 @@ from unittest.mock import MagicMock import pytest -import mcp.types as types +from mcp_helpers import call_tool_json as call_tool from dsagt.mcp.memory_tools import create_memory_server from dsagt.memory import ExplicitMemory -from mcp_helpers import call_tool_json as call_tool # --------------------------------------------------------------------------- # Fixtures @@ -164,10 +163,9 @@ def test_excludes_superseded(self, server): class TestToolSchemas: def _get_tool(self, server, name): - req = types.ListToolsRequest(method="tools/list") - handler = server.request_handlers[types.ListToolsRequest] - result = asyncio.run(handler(req)) - for tool in result.root.tools: + handler = server.get_request_handler("tools/list").handler + result = asyncio.run(handler(None, None)) + for tool in result.tools: if tool.name == name: return tool return None @@ -175,13 +173,13 @@ def _get_tool(self, server, name): def test_kb_remember_exists(self, server): tool = self._get_tool(server, "kb_remember") assert tool is not None - assert "text" in tool.inputSchema["properties"] - assert tool.inputSchema["required"] == ["text"] + assert "text" in tool.input_schema["properties"] + assert tool.input_schema["required"] == ["text"] def test_kb_remember_has_optional_params(self, server): tool = self._get_tool(server, "kb_remember") for param in ("category", "session_id", "supersedes"): - assert param in tool.inputSchema["properties"] + assert param in tool.input_schema["properties"] def test_kb_get_memories_exists(self, server): tool = self._get_tool(server, "kb_get_memories")