From 7f85c10e9019a23de9f186083d804edb0672aaa7 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 14 Aug 2026 18:33:08 +0000 Subject: [PATCH 1/4] Initial plan From 4b5e91bc05d4116b8dd0304d9f844a07f6184c92 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 14 Aug 2026 18:44:55 +0000 Subject: [PATCH 2/4] Upgrade mcp dependency to 2.0 and update Server API usage Co-authored-by: shreddd <143514+shreddd@users.noreply.github.com> --- pyproject.toml | 2 +- src/dsagt/mcp/server.py | 17 +++++++++-------- tests/mcp_helpers.py | 22 ++++++++------------- tests/test_dsagt_server.py | 17 ++++++----------- tests/test_kb_search_filters.py | 12 +++++------- tests/test_knowledge_server.py | 34 ++++++++++++++------------------- tests/test_memory_tools.py | 16 +++++++--------- 7 files changed, 50 insertions(+), 70 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 0f40469..7b705f0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,7 +16,7 @@ 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 "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..4fa3972 100644 --- a/src/dsagt/mcp/server.py +++ b/src/dsagt/mcp/server.py @@ -76,14 +76,15 @@ 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 + async def on_list_tools(ctx, params) -> types.ListToolsResult: + return types.ListToolsResult(tools=tools) - @server.call_tool() - async def call_tool(tool_name: str, arguments: dict) -> list[types.TextContent]: + async def on_call_tool( + ctx, params: types.CallToolRequestParams + ) -> types.CallToolResult: + tool_name = params.name + arguments = params.arguments or {} handler = handlers[tool_name] # KeyError = bug in list_tools schema with open_span(tool_name, source=tool_category.get(tool_name)) as span: try: @@ -104,9 +105,9 @@ async def call_tool(tool_name: str, arguments: dict) -> list[types.TextContent]: if isinstance(result, str) else json.dumps(result, ensure_ascii=False) ) - return [types.TextContent(type="text", text=text)] + 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..d2d72ee 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): 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") From 1c0ef3c3b16aefc53f17e74b10c220d6ecbcf656 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 14 Aug 2026 18:46:53 +0000 Subject: [PATCH 3/4] Address code review: avoid silently defaulting None arguments Co-authored-by: shreddd <143514+shreddd@users.noreply.github.com> --- src/dsagt/mcp/server.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/dsagt/mcp/server.py b/src/dsagt/mcp/server.py index 4fa3972..7d1cb63 100644 --- a/src/dsagt/mcp/server.py +++ b/src/dsagt/mcp/server.py @@ -84,7 +84,7 @@ async def on_call_tool( ctx, params: types.CallToolRequestParams ) -> types.CallToolResult: tool_name = params.name - arguments = params.arguments or {} + arguments = params.arguments handler = handlers[tool_name] # KeyError = bug in list_tools schema with open_span(tool_name, source=tool_category.get(tool_name)) as span: try: From 86d2a0f11488e775574ae151afc8a63d94556f32 Mon Sep 17 00:00:00 2001 From: aarontuor Date: Fri, 28 Aug 2026 07:21:42 -0700 Subject: [PATCH 4/4] Validate tool arguments in the dispatch shell The mcp 2.x server invokes on_call_tool without validating arguments against the tool's input_schema, and params.arguments is None when omitted. Malformed calls reached handlers raw and returned "Unexpected error: 'NoneType' object is not subscriptable" plus a server-log traceback for what is a client-side error. Coerce None arguments to {} and validate with jsonschema before dispatch, so a malformed call returns "Input validation error: 'query' is a required property" as under mcp 1.x. Validation runs before the categorization span, matching 1.x, where rejected calls never reached dsagt code. Declare jsonschema as a direct dependency (previously only transitive via mcp). Co-Authored-By: Claude Fable 5 --- pyproject.toml | 1 + src/dsagt/mcp/server.py | 14 +++++++++- tests/test_dsagt_server.py | 52 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 66 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 7b705f0..bf6f4d2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,7 @@ dependencies = [ "httpx>=0.27.0", # HTTP client; API embeddings "questionary>=2.0", # interactive `dsagt init` select/checkbox menus "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 7d1cb63..ec6d43c 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 @@ -76,6 +77,7 @@ def build_dispatch_server( in the single-concern test servers / one-shot tools. """ tool_category = tool_category or {} + schemas = {tool.name: tool.input_schema for tool in tools} async def on_list_tools(ctx, params) -> types.ListToolsResult: return types.ListToolsResult(tools=tools) @@ -84,8 +86,18 @@ async def on_call_tool( ctx, params: types.CallToolRequestParams ) -> types.CallToolResult: tool_name = params.name - arguments = params.arguments + # ``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 {} handler = handlers[tool_name] # KeyError = bug in list_tools schema + try: + jsonschema.validate(instance=arguments, schema=schemas[tool_name]) + except jsonschema.ValidationError as e: + error = {"status": "error", "error": f"Input validation error: {e.message}"} + return types.CallToolResult( + content=[types.TextContent(type="text", text=json.dumps(error))] + ) with open_span(tool_name, source=tool_category.get(tool_name)) as span: try: result = await handler(arguments) diff --git a/tests/test_dsagt_server.py b/tests/test_dsagt_server.py index d2d72ee..ba85189 100644 --- a/tests/test_dsagt_server.py +++ b/tests/test_dsagt_server.py @@ -130,6 +130,58 @@ 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"} + + class TestBuildKbFromConfig: """``_build_kb_from_config`` validates embedding config before building a KB.