slice 2.6.0: MCPServer class + handle_message dispatch (12/12 tests; AC 2.6.3, 2.6.4, 2.6.7, 2.6.8)

lore_engine_poc/mcp_server.py

@@ -0,0 +1,189 @@
+"""Lore Engine POC — MCP server (slice 2.6+).
+
+This is the in-process JSON-RPC 2.0 dispatcher used by the stdio
+entry point in ``scripts/05_mcp_server.py``. The class is split
+out of the stdio script so the dispatcher logic can be unit-tested
+without spawning a subprocess — see ``tests/test_mcp/test_server.py``.
+
+The MCP protocol (https://modelcontextprotocol.io) we implement:
+
+* ``initialize`` — returns ``protocolVersion`` (we hard-code
+  ``2024-11-05``), ``serverInfo``, and an empty ``capabilities``
+  object.
+* ``tools/list`` — returns the schema for every registered tool
+  (name + description + inputSchema).
+* ``tools/call`` — dispatches to the registered tool function
+  by name, with ``arguments`` unpacked as kwargs. Tool failures
+  come back as ``isError: True`` in the ``result``; transport
+  errors (unknown tool, missing ``id``) come back as JSON-RPC
+  ``error`` envelopes with negative codes from the spec.
+
+Notifications (no ``id``) get no reply, per JSON-RPC 2.0.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from typing import Callable, Optional
+
+
+# JSON-RPC 2.0 standard error codes we emit.
+ERR_PARSE = -32700
+ERR_INVALID_REQUEST = -32600
+ERR_METHOD_NOT_FOUND = -32601
+ERR_INVALID_PARAMS = -32602
+ERR_INTERNAL = -32603
+
+
+# MCP protocolVersion we announce. Pinned at the slice-2.6+
+# version; bump if/when the spec adds a major capability.
+MCP_PROTOCOL_VERSION = "2024-11-05"
+
+
+@dataclass
+class ToolEntry:
+    """One MCP tool: name + description + JSON Schema + Python fn.
+
+    The ``fn`` is expected to take ``(graph, **kwargs)``; the
+    dispatcher calls it with ``tool.fn(graph, **arguments)``.
+    Type coercion (e.g. ``at_time`` string → time atom) is the
+    tool's job, not the dispatcher's.
+    """
+
+    name: str
+    description: str
+    input_schema: dict
+    fn: Callable
+
+
+class MCPServer:
+    """JSON-RPC 2.0 dispatcher over an in-memory graph + tool registry.
+
+    The class is **stateless across requests** (every call to
+    ``handle_message`` is independent) except for the tool
+    registry itself, which is fixed at construction time.
+    """
+
+    def __init__(self, graph, tool_registry):
+        self.graph = graph
+        # Build a name → entry lookup once.
+        self.tools = {t.name: t for t in tool_registry}
+
+    # ------------------------------------------------------------------
+    # Public dispatch
+    # ------------------------------------------------------------------
+
+    def handle_message(self, msg: dict) -> Optional[dict]:
+        """Dispatch one JSON-RPC 2.0 message.
+
+        Returns the response dict, or ``None`` for notifications
+        (methods named ``notifications/*`` with no ``id``).
+        """
+        msg_id = msg.get("id")
+        method = msg.get("method")
+
+        # Notifications — must not produce a reply.
+        if method and method.startswith("notifications/") and msg_id is None:
+            return None
+
+        if method == "initialize":
+            return self._initialize(msg)
+        if method == "tools/list":
+            return self._tools_list(msg)
+        if method == "tools/call":
+            return self._tools_call(msg)
+
+        return self._error(msg_id, ERR_METHOD_NOT_FOUND, f"Method not found: {method!r}")
+
+    # ------------------------------------------------------------------
+    # initialize
+    # ------------------------------------------------------------------
+
+    def _initialize(self, msg: dict) -> dict:
+        return {
+            "jsonrpc": "2.0",
+            "id": msg.get("id"),
+            "result": {
+                "protocolVersion": MCP_PROTOCOL_VERSION,
+                "serverInfo": {
+                    "name": "lore-engine-poc",
+                    "version": "0.1.0",
+                },
+                "capabilities": {},
+            },
+        }
+
+    # ------------------------------------------------------------------
+    # tools/list
+    # ------------------------------------------------------------------
+
+    def _tools_list(self, msg: dict) -> dict:
+        tools = [
+            {
+                "name": t.name,
+                "description": t.description,
+                "inputSchema": t.input_schema,
+            }
+            for t in self.tools.values()
+        ]
+        return {
+            "jsonrpc": "2.0",
+            "id": msg.get("id"),
+            "result": {"tools": tools},
+        }
+
+    # ------------------------------------------------------------------
+    # tools/call
+    # ------------------------------------------------------------------
+
+    def _tools_call(self, msg: dict) -> dict:
+        msg_id = msg.get("id")
+        params = msg.get("params", {}) or {}
+        name = params.get("name")
+        arguments = params.get("arguments", {}) or {}
+
+        tool = self.tools.get(name)
+        if tool is None:
+            return self._error(msg_id, ERR_INVALID_PARAMS, f"Unknown tool: {name!r}")
+
+        try:
+            result = tool.fn(self.graph, **arguments)
+        except Exception as exc:
+            # Tool-body errors come back as isError=True so callers
+            # can distinguish "the tool ran and failed" from
+            # "the dispatcher couldn't reach it".
+            return {
+                "jsonrpc": "2.0",
+                "id": msg_id,
+                "result": {
+                    "content": [{"type": "text", "text": str(exc)}],
+                    "isError": True,
+                },
+            }
+
+        return {
+            "jsonrpc": "2.0",
+            "id": msg_id,
+            "result": {
+                "content": [
+                    {"type": "text", "text": json.dumps(result, default=str)}
+                ],
+                "isError": False,
+            },
+        }
+
+    # ------------------------------------------------------------------
+    # Error envelope
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _error(msg_id, code: int, message: str) -> dict:
+        return {
+            "jsonrpc": "2.0",
+            "id": msg_id,
+            "error": {"code": code, "message": message},
+        }
+
+
+__all__ = ["MCPServer", "ToolEntry", "MCP_PROTOCOL_VERSION"]

tests/test_mcp/test_server.py

@@ -0,0 +1,270 @@
+"""Tests for MCPServer.handle_message (slice 2.6.0, in-process dispatch).
+
+These tests exercise the JSON-RPC 2.0 dispatcher without touching
+stdio. They build a :class:`MCPServer` with a fixture registry and
+poke ``handle_message({...})`` directly, then assert on the
+response dict shape.
+
+The registry used here is the test double ``_trivial_registry``,
+not the real ``TOOL_REGISTRY`` — those get exercised in
+``test_tool_registry.py`` and ``test_protocol.py``.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from typing import Callable
+
+import pytest
+
+from lore_engine_poc.mcp_server import MCPServer
+from lore_engine_poc.tools import Graph
+
+
+# ---------------------------------------------------------------------------
+# Test fixture: a tiny registry so we exercise dispatcher logic
+# without depending on the 12-tool real registry (covered elsewhere).
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class _Tool:
+    name: str
+    description: str
+    input_schema: dict
+    fn: Callable
+
+
+def _echo(graph, msg):
+    return {"echo": msg}
+
+
+def _failing(graph, **kwargs):
+    raise RuntimeError("boom from the tool body")
+
+
+_TRIVIAL_REGISTRY = [
+    _Tool(
+        name="echo",
+        description="Echo arguments back.",
+        input_schema={
+            "type": "object",
+            "properties": {"msg": {"type": "string"}},
+            "required": ["msg"],
+        },
+        fn=_echo,
+    ),
+    _Tool(
+        name="failing",
+        description="Always raises.",
+        input_schema={"type": "object", "properties": {}},
+        fn=_failing,
+    ),
+]
+
+
+@pytest.fixture
+def server():
+    return MCPServer(Graph(), _TRIVIAL_REGISTRY)
+
+
+# ---------------------------------------------------------------------------
+# JSON-RPC envelope — initialize
+# ---------------------------------------------------------------------------
+
+
+def test_initialize_returns_protocol_version_2024_11_05(server):
+    resp = server.handle_message(
+        {
+            "jsonrpc": "2.0",
+            "id": 1,
+            "method": "initialize",
+            "params": {
+                "protocolVersion": "2024-11-05",
+                "capabilities": {},
+                "clientInfo": {"name": "test", "version": "0.1"},
+            },
+        }
+    )
+    assert resp["jsonrpc"] == "2.0"
+    assert resp["id"] == 1
+    assert "result" in resp
+    result = resp["result"]
+    assert result["protocolVersion"] == "2024-11-05"
+    assert "serverInfo" in result
+    assert "name" in result["serverInfo"]
+
+
+# ---------------------------------------------------------------------------
+# tools/list
+# ---------------------------------------------------------------------------
+
+
+def test_tools_list_returns_every_registered_tool(server):
+    resp = server.handle_message(
+        {"jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {}}
+    )
+    assert resp["id"] == 2
+    tools = resp["result"]["tools"]
+    names = {t["name"] for t in tools}
+    assert names == {"echo", "failing"}
+
+
+def test_tools_list_shape_per_tool(server):
+    resp = server.handle_message(
+        {"jsonrpc": "2.0", "id": 3, "method": "tools/list", "params": {}}
+    )
+    for tool in resp["result"]["tools"]:
+        assert "name" in tool
+        assert "description" in tool
+        assert "inputSchema" in tool
+        assert tool["inputSchema"]["type"] == "object"
+
+
+# ---------------------------------------------------------------------------
+# tools/call — happy path
+# ---------------------------------------------------------------------------
+
+
+def test_tools_call_round_trip(server):
+    resp = server.handle_message(
+        {
+            "jsonrpc": "2.0",
+            "id": 4,
+            "method": "tools/call",
+            "params": {
+                "name": "echo",
+                "arguments": {"msg": "hello"},
+            },
+        }
+    )
+    assert resp["id"] == 4
+    result = resp["result"]
+    assert result["isError"] is False
+    assert len(result["content"]) == 1
+    content = result["content"][0]
+    assert content["type"] == "text"
+    # The text payload is JSON-serialised tool output.
+    payload = json.loads(content["text"])
+    assert payload == {"echo": "hello"}
+
+
+# ---------------------------------------------------------------------------
+# tools/call — error envelopes
+# ---------------------------------------------------------------------------
+
+
+def test_tools_call_unknown_tool_returns_error_envelope(server):
+    resp = server.handle_message(
+        {
+            "jsonrpc": "2.0",
+            "id": 5,
+            "method": "tools/call",
+            "params": {"name": "does_not_exist", "arguments": {}},
+        }
+    )
+    assert resp["id"] == 5
+    assert "error" in resp
+    assert resp["error"]["code"] == -32602
+    assert "does_not_exist" in resp["error"]["message"]
+
+
+def test_tools_call_missing_required_arg_returns_error(server):
+    # echo requires 'msg'. Omit it.
+    resp = server.handle_message(
+        {
+            "jsonrpc": "2.0",
+            "id": 6,
+            "method": "tools/call",
+            "params": {"name": "echo", "arguments": {}},
+        }
+    )
+    # The tool raises TypeError (missing positional 'msg'), which
+    # the dispatcher converts to isError=True, NOT a transport
+    # error.
+    assert resp["id"] == 6
+    result = resp["result"]
+    assert result["isError"] is True
+    assert "msg" in result["content"][0]["text"]
+
+
+def test_tools_call_tool_body_exception_is_tool_error_not_transport(server):
+    resp = server.handle_message(
+        {
+            "jsonrpc": "2.0",
+            "id": 7,
+            "method": "tools/call",
+            "params": {"name": "failing", "arguments": {}},
+        }
+    )
+    # Tool body exception → result envelope with isError=True.
+    # This is distinct from transport-level errors.
+    assert resp["id"] == 7
+    assert resp.get("error") is None
+    assert resp["result"]["isError"] is True
+    assert "boom from the tool body" in resp["result"]["content"][0]["text"]
+
+
+# ---------------------------------------------------------------------------
+# Unknown method / notifications
+# ---------------------------------------------------------------------------
+
+
+def test_unknown_method_returns_method_not_found(server):
+    resp = server.handle_message(
+        {"jsonrpc": "2.0", "id": 8, "method": "tools/banana", "params": {}}
+    )
+    assert resp["id"] == 8
+    assert resp["error"]["code"] == -32601
+
+
+def test_notification_returns_none(server):
+    """Notifications (no ``id``) get no reply at all."""
+    result = server.handle_message(
+        {"jsonrpc": "2.0", "method": "notifications/initialized", "params": {}}
+    )
+    assert result is None
+
+
+def test_notification_with_id_still_returns_response(server):
+    """A method named ``notifications/*`` but carrying an ``id``
+    is a malformed notification — for safety we still reply."""
+    resp = server.handle_message(
+        {
+            "jsonrpc": "2.0",
+            "id": 9,
+            "method": "notifications/something",
+            "params": {},
+        }
+    )
+    # id present → still returns a response. (Not strictly spec-
+    # compliant, but defensive.)
+    assert resp is not None
+
+
+# ---------------------------------------------------------------------------
+# Edge cases — bad message shape
+# ---------------------------------------------------------------------------
+
+
+def test_empty_message_returns_error(server):
+    resp = server.handle_message({})
+    assert "error" in resp
+    # id was missing → None
+    assert resp.get("id") is None
+
+
+def test_tools_call_with_no_arguments_key_works(server):
+    """Some clients omit ``arguments`` entirely; the dispatcher
+    must treat that as an empty kwargs dict, not an error."""
+    resp = server.handle_message(
+        {
+            "jsonrpc": "2.0",
+            "id": 10,
+            "method": "tools/call",
+            "params": {"name": "failing"},
+        }
+    )
+    # failing has no required args, so it runs and raises → isError
+    assert resp["result"]["isError"] is True