diff --git a/pyproject.toml b/pyproject.toml index c2ab1863..bd7611f3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,7 @@ dependencies = [ ] [project.optional-dependencies] -mcp = ["mcp>=1.28.1"] +mcp = ["mcp>=2.0"] dev = ["pre-commit", "pytest", "pytest-cov", "pytest-doctestplus", "pytest-asyncio", "ruff", "docutils"] doc = [ "jupyter-book>=2.1.6,<3", diff --git a/src/gdm/mcp/server.py b/src/gdm/mcp/server.py index 12a620cd..a31761be 100644 --- a/src/gdm/mcp/server.py +++ b/src/gdm/mcp/server.py @@ -13,9 +13,9 @@ import typer from gdm.distribution import DistributionSystem -from mcp.server import Server +from mcp.server import Server, ServerRequestContext from mcp.server.stdio import stdio_server -from mcp.types import TextContent, Tool +from mcp.types import CallToolResult, ListToolsResult, TextContent, Tool from gdm.mcp import __version__ from gdm.mcp.exceptions import GDMMCPException @@ -51,9 +51,6 @@ logging.basicConfig(level=logging.INFO) logger = logging.getLogger("gdm_mcp") -# Create MCP server instance -app = Server("grid-data-models-mcp") - # Runtime toggle for serving tool calls. Control tools remain available. _TOOL_CALLS_ENABLED = True _CONTROL_TOOLS = {"set_tool_calls_enabled", "get_tool_calls_enabled"} @@ -185,10 +182,9 @@ def _split_tool_input_schema() -> dict[str, Any]: } -@app.list_tools() -async def list_tools() -> list[Tool]: +async def list_tools(ctx: ServerRequestContext, params=None) -> ListToolsResult: """List all available MCP tools.""" - return [ + tools = [ # Validation tools Tool( name="diagnose_system", @@ -686,6 +682,7 @@ async def list_tools() -> list[Tool]: }, ), ] + return ListToolsResult(tools=tools) # Tool dispatch map @@ -717,27 +714,30 @@ async def list_tools() -> list[Tool]: } -@app.call_tool() -async def call_tool(name: str, arguments: Any) -> list[TextContent]: +async def call_tool(ctx: ServerRequestContext, params: Any) -> CallToolResult: """Handle tool calls from MCP clients.""" + name = params.name + arguments = params.arguments or {} try: logger.info(f"Tool called: {name} with arguments: {arguments}") if not _TOOL_CALLS_ENABLED and name not in _CONTROL_TOOLS: - return [ - TextContent( - type="text", - text=json.dumps( - { - "error": ( - "Tool calls are currently disabled. " - "Use set_tool_calls_enabled to re-enable." - ) - }, - indent=2, - ), - ) - ] + return CallToolResult( + content=[ + TextContent( + type="text", + text=json.dumps( + { + "error": ( + "Tool calls are currently disabled. " + "Use set_tool_calls_enabled to re-enable." + ) + }, + indent=2, + ), + ) + ] + ) handler = _TOOL_HANDLERS.get(name) if handler is not None: @@ -745,18 +745,27 @@ async def call_tool(name: str, arguments: Any) -> list[TextContent]: else: result = {"error": f"Unknown tool: {name}"} - return [TextContent(type="text", text=json.dumps(result, indent=2, default=str))] + return CallToolResult( + content=[TextContent(type="text", text=json.dumps(result, indent=2, default=str))] + ) except GDMMCPException as e: logger.error(f"GDM MCP error in {name}: {str(e)}") - return [TextContent(type="text", text=json.dumps({"error": str(e)}, indent=2))] + return CallToolResult( + content=[TextContent(type="text", text=json.dumps({"error": str(e)}, indent=2))], + is_error=True, + ) except Exception as e: logger.error(f"Unexpected error in {name}: {str(e)}", exc_info=True) - return [ - TextContent( - type="text", text=json.dumps({"error": f"Unexpected error: {str(e)}"}, indent=2) - ) - ] + return CallToolResult( + content=[ + TextContent( + type="text", + text=json.dumps({"error": f"Unexpected error: {str(e)}"}, indent=2), + ) + ], + is_error=True, + ) # Tool implementations @@ -1079,6 +1088,13 @@ def _run_server( # Run the server import asyncio + app = Server( + "grid-data-models-mcp", + version=__version__, + on_list_tools=list_tools, + on_call_tool=call_tool, + ) + async def run(): async with stdio_server() as (read_stream, write_stream): await app.run(read_stream, write_stream, app.create_initialization_options()) diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 97fbdd44..4ee93bea 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -4,6 +4,7 @@ import json import os import sqlite3 +from unittest.mock import MagicMock import gdm.mcp.server as mcp_server import pytest @@ -11,6 +12,22 @@ from gdm.distribution.components import DistributionBus, DistributionVoltageSource +def _make_call_tool_params(name: str, arguments: dict | None = None): + """Create a mock CallToolRequestParams for testing.""" + params = MagicMock() + params.name = name + params.arguments = arguments or {} + return params + + +def _call_tool(name: str, arguments: dict | None = None): + """Helper to call tool_handler with mock context.""" + ctx = MagicMock() + params = _make_call_tool_params(name, arguments) + result = asyncio.run(mcp_server.call_tool(ctx, params)) + return result + + def test_load_system_with_fallback_name_for_null_name(simple_system, tmp_path): """Falls back to file stem when serialized system name is null.""" system_path = tmp_path / "null_name_system.json" @@ -95,8 +112,8 @@ def test_get_tool_calls_enabled_reports_current_state(): """Control status tool should report current runtime toggle state.""" mcp_server._TOOL_CALLS_ENABLED = True - response = asyncio.run(mcp_server.call_tool("get_tool_calls_enabled", {})) - payload = json.loads(response[0].text) + response = _call_tool("get_tool_calls_enabled", {}) + payload = json.loads(response.content[0].text) assert payload["tool_calls_enabled"] is True @@ -105,19 +122,17 @@ def test_set_tool_calls_enabled_disables_non_control_calls(): """Disabling should block normal tools while allowing control tools.""" mcp_server._TOOL_CALLS_ENABLED = True - disable_response = asyncio.run( - mcp_server.call_tool("set_tool_calls_enabled", {"enabled": False}) - ) - disable_payload = json.loads(disable_response[0].text) + disable_response = _call_tool("set_tool_calls_enabled", {"enabled": False}) + disable_payload = json.loads(disable_response.content[0].text) assert disable_payload["tool_calls_enabled"] is False - blocked_response = asyncio.run(mcp_server.call_tool("unknown_normal_tool", {})) - blocked_payload = json.loads(blocked_response[0].text) + blocked_response = _call_tool("unknown_normal_tool", {}) + blocked_payload = json.loads(blocked_response.content[0].text) assert "disabled" in blocked_payload["error"].lower() # Control tools remain callable so clients can re-enable. - status_response = asyncio.run(mcp_server.call_tool("get_tool_calls_enabled", {})) - status_payload = json.loads(status_response[0].text) + status_response = _call_tool("get_tool_calls_enabled", {}) + status_payload = json.loads(status_response.content[0].text) assert status_payload["tool_calls_enabled"] is False @@ -125,22 +140,21 @@ def test_set_tool_calls_enabled_can_reenable(): """Re-enabling should restore normal call flow.""" mcp_server._TOOL_CALLS_ENABLED = False - enable_response = asyncio.run( - mcp_server.call_tool("set_tool_calls_enabled", {"enabled": True}) - ) - enable_payload = json.loads(enable_response[0].text) + enable_response = _call_tool("set_tool_calls_enabled", {"enabled": True}) + enable_payload = json.loads(enable_response.content[0].text) assert enable_payload["tool_calls_enabled"] is True - unknown_response = asyncio.run(mcp_server.call_tool("unknown_normal_tool", {})) - unknown_payload = json.loads(unknown_response[0].text) + unknown_response = _call_tool("unknown_normal_tool", {}) + unknown_payload = json.loads(unknown_response.content[0].text) assert "unknown tool" in unknown_payload["error"].lower() def test_list_tools_includes_reduce_system(): """Tool list should expose model-reduction capability.""" - tools = asyncio.run(mcp_server.list_tools()) - tool_names = {tool.name for tool in tools} + ctx = MagicMock() + result = asyncio.run(mcp_server.list_tools(ctx, None)) + tool_names = {tool.name for tool in result.tools} assert "reduce_system" in tool_names assert "save_system" in tool_names