diff --git a/pyproject.toml b/pyproject.toml index 3fa8d5c..7bc1425 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 5bc94a8..e98f22c 100644 --- a/src/dsagt/mcp/server.py +++ b/src/dsagt/mcp/server.py @@ -83,14 +83,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 handler = handlers[tool_name] # KeyError = bug in list_tools schema with open_span(tool_name, source=tool_category.get(tool_name)) as span: try: @@ -111,9 +112,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")