MemRelay/scripts/validate_basic_memory.py

"""Validate the Basic Memory MCP contract used by MemRelay."""

from __future__ import annotations

import argparse
import asyncio
import json
from typing import Any

from fastmcp import Client


REQUIRED_TOOLS = {
    "delete_note",
    "list_directory",
    "move_note",
    "read_note",
    "search_notes",
    "write_note",
}


async def call(client: Client, name: str, arguments: dict[str, Any]) -> Any:
    result = await client.call_tool(name, arguments)
    if result.data is not None:
        return result.data
    return [getattr(item, "text", str(item)) for item in result.content]


async def delete_fixture(client: Client, identifier: str) -> None:
    await call(
        client,
        "delete_note",
        {"identifier": identifier, "output_format": "json"},
    )


async def validate(endpoint: str, semantic: bool) -> dict[str, Any]:
    async with Client(endpoint) as client:
        tools = {tool.name for tool in await client.list_tools()}
        missing = REQUIRED_TOOLS - tools
        if missing:
            raise RuntimeError(f"Basic Memory is missing required tools: {sorted(missing)}")

        await delete_fixture(client, "Stage Zero Note")
        for index in range(5):
            await delete_fixture(client, f"Concurrent Note {index}")
        await delete_fixture(client, "Credential Memory")
        await delete_fixture(client, "Weather Memory")

        written = await call(
            client,
            "write_note",
            {
                "title": "Stage Zero Note",
                "content": (
                    "# Stage Zero Note\n\n"
                    "中文语义与 English context.\n\n"
                    "- [decision] Markdown remains authoritative."
                ),
                "directory": "global",
                "tags": ["stage-zero", "multilingual"],
                "note_type": "decision",
                "metadata": {
                    "stable_id": "stage-zero-001",
                    "scope": "global",
                    "revision": 1,
                    "status": "active",
                },
                "overwrite": True,
                "output_format": "json",
            },
        )
        read = await call(
            client,
            "read_note",
            {
                "identifier": "Stage Zero Note",
                "output_format": "json",
                "include_frontmatter": True,
            },
        )
        search = await call(
            client,
            "search_notes",
            {
                "query": "Markdown authoritative",
                "search_type": "text",
                "output_format": "json",
            },
        )
        concurrent = await asyncio.gather(
            *(
                call(
                    client,
                    "write_note",
                    {
                        "title": f"Concurrent Note {index}",
                        "content": f"Concurrent content {index}",
                        "directory": "projects/stage-zero",
                        "metadata": {"request_id": f"stage0-{index}"},
                        "overwrite": True,
                        "output_format": "json",
                    },
                )
                for index in range(5)
            )
        )
        moved = await call(
            client,
            "move_note",
            {
                "identifier": "Stage Zero Note",
                "destination_path": "archive/stage-zero-note.md",
                "output_format": "json",
            },
        )
        archived = await call(
            client,
            "read_note",
            {
                "identifier": "archive/stage-zero-note",
                "output_format": "json",
                "include_frontmatter": True,
            },
        )

        semantic_result = None
        if semantic:
            await call(
                client,
                "write_note",
                {
                    "title": "Credential Memory",
                    "content": "团队账号密码和访问令牌应统一保存在密码库中。",
                    "directory": "projects/stage-zero",
                    "overwrite": True,
                    "output_format": "json",
                },
            )
            await call(
                client,
                "write_note",
                {
                    "title": "Weather Memory",
                    "content": "今天的天气预报是晴天，午后温度较高。",
                    "directory": "projects/stage-zero",
                    "overwrite": True,
                    "output_format": "json",
                },
            )
            semantic_result = await call(
                client,
                "search_notes",
                {
                    "query": "Where should team passwords and access tokens be stored?",
                    "search_type": "vector",
                    "min_similarity": 0,
                    "output_format": "json",
                },
            )
            titles = [item["title"] for item in semantic_result.get("results", [])]
            if "Credential Memory" not in titles:
                raise RuntimeError("Cross-language semantic search did not return the target note")

        await delete_fixture(client, "Stage Zero Note")
        for index in range(5):
            await delete_fixture(client, f"Concurrent Note {index}")
        await delete_fixture(client, "Credential Memory")
        await delete_fixture(client, "Weather Memory")

    return {
        "available_tool_count": len(tools),
        "write": written,
        "read": read,
        "search": search,
        "concurrent_write_count": len(concurrent),
        "move": moved,
        "archive_read": archived,
        "semantic_search": semantic_result,
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--endpoint", default="http://127.0.0.1:8000/mcp")
    parser.add_argument("--semantic", action="store_true")
    args = parser.parse_args()
    print(
        json.dumps(
            asyncio.run(validate(args.endpoint, args.semantic)),
            ensure_ascii=False,
            indent=2,
        )
    )


if __name__ == "__main__":
    main()