Upgrade mcp dependency from 1.x to 2.0

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

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,15 +77,27 @@ 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
-
-    @server.call_tool()
-    async def call_tool(tool_name: str, arguments: dict) -> list[types.TextContent]:
+    schemas = {tool.name: tool.input_schema for tool in tools}
+
+    async def on_list_tools(ctx, params) -> types.ListToolsResult:
+        return types.ListToolsResult(tools=tools)
+
+    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 {}
         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)
@@ -104,9 +117,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

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
 
 
 # ---------------------------------------------------------------------------

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,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.
 

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):

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"]

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,24 +163,23 @@ 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
 
     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")