path: root/tests/mcp_protocol_test.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import shutil
import subprocess
import tempfile
import threading
from collections.abc import Mapping
from pathlib import Path
from typing import Any


class Server:
    def __init__(self, executable: Path, data_dir: Path) -> None:
        self.stderr: list[str] = []
        self.last_response_size = 0
        self.process = subprocess.Popen(
            [str(executable), "--data-dir", str(data_dir)],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            encoding="utf-8",
            bufsize=1,
        )
        assert self.process.stdin and self.process.stdout and self.process.stderr
        self._stderr_thread = threading.Thread(target=self._drain_stderr, daemon=True)
        self._stderr_thread.start()

    def _drain_stderr(self) -> None:
        assert self.process.stderr
        self.stderr.extend(self.process.stderr)

    def send(self, message: Mapping[str, Any] | str) -> dict[str, Any]:
        assert self.process.stdin and self.process.stdout
        line = message if isinstance(message, str) else json.dumps(message, separators=(",", ":"))
        self.process.stdin.write(line + "\n")
        self.process.stdin.flush()
        response_line = self.process.stdout.readline()
        if not response_line:
            raise AssertionError(
                f"server closed: returncode={self.process.poll()} stderr={''.join(self.stderr)[-4000:]}"
            )
        self.last_response_size = len(response_line.encode("utf-8"))
        response = json.loads(response_line)
        assert response.get("jsonrpc") == "2.0", response
        return response

    def notify(self, message: Mapping[str, Any]) -> None:
        assert self.process.stdin
        self.process.stdin.write(json.dumps(message, separators=(",", ":")) + "\n")
        self.process.stdin.flush()

    def close(self) -> None:
        if self.process.stdin:
            self.process.stdin.close()
        returncode = self.process.wait(timeout=15)
        self._stderr_thread.join(timeout=2)
        assert returncode == 0, f"returncode={returncode} stderr={''.join(self.stderr)}"
        stderr = "".join(self.stderr)
        assert "AddressSanitizer" not in stderr and "ERROR:" not in stderr, stderr


def request(request_id: int, method: str, params: Mapping[str, Any] | None = None) -> dict[str, Any]:
    message: dict[str, Any] = {"jsonrpc": "2.0", "id": request_id, "method": method}
    if params is not None:
        message["params"] = params
    return message


def tool_call(server: Server, request_id: int, name: str, arguments: Mapping[str, Any]) -> dict[str, Any]:
    response = server.send(
        request(request_id, "tools/call", {"name": name, "arguments": arguments})
    )
    assert response.get("id") == request_id, response
    return response


def tool_payload(response: Mapping[str, Any], expect_error: bool = False) -> dict[str, Any]:
    result = response["result"]
    assert bool(result.get("isError", False)) is expect_error, response
    return json.loads(result["content"][0]["text"])


def assert_stable_recency_order(messages: list[dict[str, Any]]) -> None:
    for left, right in zip(messages, messages[1:]):
        left_time = int(left.get("updated_at", left["timestamp"]))
        right_time = int(right.get("updated_at", right["timestamp"]))
        left_created = int(left["timestamp"])
        right_created = int(right["timestamp"])
        assert left_time > right_time or (
            left_time == right_time
            and (
                left_created > right_created
                or (left_created == right_created and left["message_id"] <= right["message_id"])
            )
        ), (left, right)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--server", type=Path, required=True)
    parser.add_argument("--iterations", type=int, default=1000)
    args = parser.parse_args()
    data_dir = Path(tempfile.mkdtemp(prefix="tiny-convmem-p0-"))
    server = Server(args.server.resolve(), data_dir)
    next_id = 1
    try:
        response = server.send(
            request(
                next_id,
                "initialize",
                {
                    "protocolVersion": "2024-11-05",
                    "capabilities": {},
                    "clientInfo": {"name": "p0-regression", "version": "1"},
                },
            )
        )
        assert response["result"]["protocolVersion"] == "2024-11-05"
        next_id += 1
        server.notify({"jsonrpc": "2.0", "method": "notifications/initialized", "params": {}})

        listed = server.send(request(next_id, "tools/list", {}))
        assert listed["id"] == next_id  # also proves the notification emitted no response
        next_id += 1
        tools = {item["name"]: item for item in listed["result"]["tools"]}
        assert set(tools) == {
            "write_message",
            "read_messages",
            "search_messages",
            "update_message",
            "delete_message",
            "rebuild_index",
            "get_system_health",
        }
        expected_descriptions = {
            "write_message": "Append and durably persist a conversation message. Use it to store information that should remain available after server restarts.",
            "read_messages": "Read persisted messages from one known session in chronological order. Returns a paginated history for reconstructing that conversation.",
            "search_messages": "Search persisted message content and tags by text relevance. Use it to recall matching information globally or within one session.",
            "update_message": "Update message content or metadata while preserving identity, session, role, and creation time.",
            "delete_message": "Delete a persisted message so it is immediately excluded from reads and searches. Physical removal occurs after rebuild_index compacts the store.",
            "rebuild_index": "Compact the store and rebuild its snapshot and in-memory indexes. Use it for maintenance or to physically remove logically deleted records; it may block other requests while running.",
            "get_system_health": "Inspect storage health, recovery state, message counts, and file sizes. Use it to diagnose persistence or recovery problems.",
        }
        assert {name: tool["description"] for name, tool in tools.items()} == expected_descriptions
        expected_property_descriptions = {
            "write_message": {
                "session_id": "Conversation session that owns the message.",
                "role": "Message author role: user, assistant, or system.",
                "content": "UTF-8 message text to persist.",
                "metadata_json": "Optional JSON object string containing tool_call_id, tags, memory_tier, memory_kind, importance, summary, relations, and/or source_message_ids; defaults to {}.",
            },
            "read_messages": {
                "view": "Read view: history (default), context, sessions, or export.",
                "session_id": "Conversation session whose chronological history to read.",
                "message_id": "Message ULID required by the context view.",
                "before": "Context messages before the target, from 0 to 10.",
                "after": "Context messages after the target, from 0 to 10.",
                "limit": "Maximum messages to return in this page, from 1 to 1000; defaults to 100.",
                "cursor": "Opaque nextCursor returned by a previous read_messages call; omit for the first page.",
            },
            "search_messages": {
                "query": "UTF-8 text to match against persisted message content and tags.",
                "session_id": "Conversation session to search; omit to search across all sessions.",
                "memory_kind": "Optional exact memory kind filter.",
                "memory_tier": "Optional exact memory tier filter.",
                "tag": "Optional exact tag filter.",
                "min_importance": "Minimum importance from 0 to 100; defaults to 0.",
                "limit": "Maximum ranked matches to return, from 1 to 100; defaults to 20.",
            },
            "delete_message": {
                "message_id": "Canonical 26-character ULID of the message to delete.",
                "session_id": "Optional session guard; when provided, delete only if the message belongs to this session.",
            },
            "update_message": {
                "message_id": "ULID of the live message to update.",
                "content": "Replacement UTF-8 content; omit to retain content.",
                "metadata_json": "Replacement metadata JSON; omit to retain metadata.",
                "expected_revision": "Optional positive revision for optimistic concurrency.",
            },
        }
        for tool_name, descriptions in expected_property_descriptions.items():
            properties = tools[tool_name]["inputSchema"]["properties"]
            assert {name: value["description"] for name, value in properties.items()} == descriptions
        assert tools["write_message"]["inputSchema"]["required"] == [
            "session_id",
            "role",
            "content",
        ]
        assert tools["read_messages"]["inputSchema"].get("required", []) == []
        assert tools["read_messages"]["inputSchema"]["properties"]["limit"]["type"] == "integer"
        assert tools["read_messages"]["inputSchema"]["properties"]["cursor"]["type"] == "string"
        assert tools["search_messages"]["inputSchema"]["required"] == ["query"]
        assert tools["search_messages"]["inputSchema"]["properties"]["limit"]["type"] == "integer"
        assert tools["delete_message"]["inputSchema"]["required"] == ["message_id"]
        assert tools["update_message"]["inputSchema"]["required"] == ["message_id"]

        message_ids: list[str] = []
        for index in range(20):
            payload = tool_payload(
                tool_call(
                    server,
                    next_id,
                    "write_message",
                    {"session_id": "target", "role": "user", "content": f"needle target {index}"},
                )
            )
            next_id += 1
            message_ids.append(payload["message"]["message_id"])
        updated = tool_payload(
            tool_call(
                server,
                next_id,
                "update_message",
                {
                    "message_id": message_ids[0],
                    "content": "needle target revised",
                    "expected_revision": 1,
                },
            )
        )
        next_id += 1
        assert updated["changed"] is True
        assert updated["message"]["revision"] == 2
        assert updated["message"]["content"] == "needle target revised"
        unchanged = tool_payload(
            tool_call(
                server,
                next_id,
                "update_message",
                {
                    "message_id": message_ids[0],
                    "content": "needle target revised",
                    "expected_revision": 2,
                },
            )
        )
        next_id += 1
        assert unchanged["changed"] is False
        assert unchanged["message"]["revision"] == 2
        io_before = tool_payload(tool_call(server, next_id, "get_system_health", {}))["io"]
        next_id += 1
        context = tool_payload(
            tool_call(
                server,
                next_id,
                "read_messages",
                {"view": "context", "message_id": message_ids[1], "before": 1, "after": 1},
            )
        )
        next_id += 1
        assert context["message"]["message_id"] == message_ids[1]
        assert [item["message_id"] for item in context["chronological_context"]["before"]] == [
            message_ids[0]
        ]
        assert [item["message_id"] for item in context["chronological_context"]["after"]] == [
            message_ids[2]
        ]
        io_after = tool_payload(tool_call(server, next_id, "get_system_health", {}))["io"]
        next_id += 1
        assert io_after == io_before
        sessions_view = tool_payload(
            tool_call(server, next_id, "read_messages", {"view": "sessions", "limit": 10})
        )
        next_id += 1
        assert any(item["session_id"] == "target" for item in sessions_view["sessions"])
        exported = tool_payload(
            tool_call(
                server,
                next_id,
                "read_messages",
                {"view": "export", "session_id": "target", "limit": 100},
            )
        )
        next_id += 1
        export_lines = exported["jsonl"].splitlines()
        assert json.loads(export_lines[0])["format_version"] == 2
        assert len(export_lines) == 21
        assert all(json.loads(line)["session_id"] == "target" for line in export_lines[1:])
        io_after_views = tool_payload(tool_call(server, next_id, "get_system_health", {}))["io"]
        next_id += 1
        assert io_after_views == io_before
        relation_message = tool_payload(
            tool_call(
                server,
                next_id,
                "write_message",
                {
                    "session_id": "relations",
                    "role": "assistant",
                    "content": "summary of the revised record",
                    "metadata_json": json.dumps(
                        {
                            "memory_kind": "summary",
                            "importance": 90,
                            "summary": "record summary",
                            "relations": [
                                {"type": "derived_from", "message_id": message_ids[0]}
                            ],
                            "source_message_ids": [message_ids[0]],
                        }
                    ),
                },
            )
        )["message"]
        next_id += 1
        relation_context = tool_payload(
            tool_call(
                server,
                next_id,
                "read_messages",
                {"view": "context", "message_id": relation_message["message_id"]},
            )
        )
        next_id += 1
        assert relation_context["relations"][0]["type"] == "derived_from"
        assert relation_context["relations"][0]["message"]["message_id"] == message_ids[0]
        incoming_context = tool_payload(
            tool_call(
                server,
                next_id,
                "read_messages",
                {"view": "context", "message_id": message_ids[0]},
            )
        )
        next_id += 1
        assert any(
            relation["direction"] == "incoming"
            and relation["message"]["message_id"] == relation_message["message_id"]
            for relation in incoming_context["relations"]
        )
        explained = tool_payload(
            tool_call(
                server,
                next_id,
                "search_messages",
                {
                    "query": "record summary",
                    "memory_kind": "summary",
                    "min_importance": 80,
                    "limit": 5,
                },
            )
        )["results"]
        next_id += 1
        assert explained[0]["message_id"] == relation_message["message_id"]
        assert "summary" in explained[0]["matched_fields"]
        assert explained[0]["score_breakdown"]["summary"] > 0
        for index in range(10):
            tool_payload(
                tool_call(
                    server,
                    next_id,
                    "write_message",
                    {"session_id": "other", "role": "user", "content": f"needle other {index}"},
                )
            )
            next_id += 1

        history = tool_payload(
            tool_call(server, next_id, "read_messages", {"session_id": "target"})
        )["history"]
        next_id += 1
        assert len(history) == 20

        paged_ids: list[str] = []
        for index in range(105):
            payload = tool_payload(
                tool_call(
                    server,
                    next_id,
                    "write_message",
                    {"session_id": "paged", "role": "user", "content": f"page item {index}"},
                )
            )
            next_id += 1
            paged_ids.append(payload["message"]["message_id"])
        first_page = tool_payload(
            tool_call(server, next_id, "read_messages", {"session_id": "paged"})
        )
        next_id += 1
        assert len(first_page["history"]) == 100 and first_page.get("nextCursor")
        second_page = tool_payload(
            tool_call(
                server,
                next_id,
                "read_messages",
                {"session_id": "paged", "cursor": first_page["nextCursor"]},
            )
        )
        next_id += 1
        assert len(second_page["history"]) == 5 and "nextCursor" not in second_page
        assert [item["message_id"] for item in first_page["history"] + second_page["history"]] == paged_ids
        for invalid_limit in (0, 1001):
            assert tool_payload(
                tool_call(
                    server,
                    next_id,
                    "read_messages",
                    {"session_id": "paged", "limit": invalid_limit},
                ),
                True,
            )["status"] == "ERROR"
            next_id += 1
        assert tool_payload(
            tool_call(
                server,
                next_id,
                "read_messages",
                {"session_id": "other", "cursor": first_page["nextCursor"]},
            ),
            True,
        )["status"] == "ERROR"
        next_id += 1

        large_content = "\\" * (512 * 1024)
        for index in range(6):
            tool_payload(
                tool_call(
                    server,
                    next_id,
                    "write_message",
                    {"session_id": "large-page", "role": "user", "content": large_content},
                )
            )
            next_id += 1
        large_page = tool_payload(
            tool_call(
                server,
                next_id,
                "read_messages",
                {"session_id": "large-page", "limit": 100},
            )
        )
        next_id += 1
        assert 0 < len(large_page["history"]) < 6 and large_page.get("nextCursor")
        assert server.last_response_size <= 8 * 1024 * 1024
        assert tool_payload(
            tool_call(
                server,
                next_id,
                "read_messages",
                {"session_id": "paged", "cursor": "broken"},
            ),
            True,
        )["status"] == "ERROR"
        next_id += 1

        ranked = tool_payload(
            tool_call(
                server,
                next_id,
                "search_messages",
                {"query": "needle target", "limit": 100},
            )
        )["results"]
        next_id += 1
        assert len(ranked) == 30
        assert all(item["session_id"] == "target" for item in ranked[:20])
        assert all(item["session_id"] == "other" for item in ranked[20:])
        assert_stable_recency_order(ranked[:20])
        assert_stable_recency_order(ranked[20:])
        ranked_ids_before = [item["message_id"] for item in ranked]

        for index in range(args.iterations):
            arguments: dict[str, Any] = {"query": "needle"}
            if index % 2:
                arguments["session_id"] = "target"
            results = tool_payload(
                tool_call(server, next_id, "search_messages", arguments)
            )["results"]
            next_id += 1
            assert len(results) == 20
            if index % 2:
                assert all(item["session_id"] == "target" for item in results)
            if (index + 1) % 100 == 0:
                print(f"completed {index + 1}/{args.iterations} alternating searches", flush=True)

        assert tool_payload(
            tool_call(server, next_id, "search_messages", {"query": "needle", "limit": 0}), True
        )["status"] == "ERROR"
        next_id += 1
        assert tool_payload(
            tool_call(server, next_id, "search_messages", {"query": "needle", "limit": 101}), True
        )["status"] == "ERROR"
        next_id += 1
        fractional = tool_call(
            server, next_id, "search_messages", {"query": "needle", "limit": 1.5}
        )
        assert fractional["error"]["code"] == -32602
        next_id += 1
        assert tool_payload(tool_call(server, next_id, "search_messages", {}), True)["status"] == "ERROR"
        next_id += 1

        assert tool_payload(
            tool_call(server, next_id, "delete_message", {"message_id": message_ids[0]})
        )["status"] == "OK"
        next_id += 1
        tool_payload(tool_call(server, next_id, "rebuild_index", {}))
        next_id += 1
        ranked_after_rebuild = tool_payload(
            tool_call(
                server,
                next_id,
                "search_messages",
                {"query": "needle target", "limit": 100},
            )
        )["results"]
        next_id += 1
        assert [item["message_id"] for item in ranked_after_rebuild] == [
            message_id for message_id in ranked_ids_before if message_id != message_ids[0]
        ]
        assert tool_payload(tool_call(server, next_id, "get_system_health", {}))["status"] == "OK"
        next_id += 1

        assert server.send("{bad}")["error"]["code"] == -32700
        invalid = server.send({"id": next_id, "method": "tools/list"})
        assert invalid["error"]["code"] == -32600
        next_id += 1
        missing_method = server.send({"jsonrpc": "2.0", "id": next_id})
        assert missing_method["error"]["code"] == -32600
        next_id += 1
        null_id = server.send(
            {"jsonrpc": "2.0", "id": None, "method": "tools/list", "params": {}}
        )
        assert null_id["error"]["code"] == -32600 and null_id["id"] is None
        invalid_params = server.send(
            {"jsonrpc": "2.0", "id": next_id, "method": "tools/list", "params": []}
        )
        assert invalid_params["error"]["code"] == -32602
        next_id += 1
        unknown_method = server.send(request(next_id, "unknown/method", {}))
        assert unknown_method["error"]["code"] == -32601
        next_id += 1
        unknown_tool = tool_call(server, next_id, "unknown_tool", {})
        assert unknown_tool["error"]["code"] == -32602
        next_id += 1
        missing_arguments = server.send(
            request(next_id, "tools/call", {"name": "search_messages"})
        )
        assert missing_arguments["error"]["code"] == -32602
    except Exception:
        if server.process.poll() is None:
            server.process.kill()
        server.process.wait(timeout=5)
        raise
    else:
        server.close()
        restarted = Server(args.server.resolve(), data_dir)
        try:
            restarted.send(
                request(
                    1,
                    "initialize",
                    {
                        "protocolVersion": "2024-11-05",
                        "capabilities": {},
                        "clientInfo": {"name": "pagination-restart", "version": "1"},
                    },
                )
            )
            restarted.notify(
                {"jsonrpc": "2.0", "method": "notifications/initialized", "params": {}}
            )
            resumed = tool_payload(
                tool_call(
                    restarted,
                    2,
                    "read_messages",
                    {"session_id": "paged", "cursor": first_page["nextCursor"]},
                )
            )
            assert [item["message_id"] for item in resumed["history"]] == paged_ids[100:]
        finally:
            restarted.close()
    finally:
        shutil.rmtree(data_dir, ignore_errors=True)


if __name__ == "__main__":
    main()