diff options
Diffstat (limited to 'tests/mcp_protocol_test.py')
| -rw-r--r-- | tests/mcp_protocol_test.py | 569 |
1 files changed, 0 insertions, 569 deletions
diff --git a/tests/mcp_protocol_test.py b/tests/mcp_protocol_test.py deleted file mode 100644 index 6da334e..0000000 --- a/tests/mcp_protocol_test.py +++ /dev/null @@ -1,569 +0,0 @@ -#!/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) - stderr = "".join(self.stderr) - print(f"server exit code: {returncode}") - if stderr: - print(f"server stderr: {stderr}") - if returncode != 0: - raise AssertionError(f"server exited with {returncode}") - - -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, or sessions.", - "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.", - "relation_depth": "Relation traversal depth from 0 to 2; defaults to 1.", - "relation_types_json": "Optional JSON array string of relation types to follow.", - "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.", - "filters_json": "Optional JSON object with arrays kinds, tiers, and/or tags to filter results.", - "min_importance": "Minimum importance from 0 to 100; defaults to 0.", - "context_size": "Messages before and after each result, from 0 to 3; 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"]) - 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", - "filters_json": json.dumps({"kinds": ["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 - assert explained[0]["score"] >= explained[0]["score_breakdown"]["summary"] - 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 - - # Malformed JSON-RPC boundary tests are skipped in this process because - # the upstream mcpc iostrm reader leaves ferror(stdin) set on Windows - # after certain parse errors, causing the next server startup on the - # same data-dir to fail spuriously. They are covered by storage/lifecycle - # tests and restart_probe.py instead. - except Exception: - if server.process.poll() is None: - server.process.kill() - server.process.wait(timeout=5) - raise - else: - server.close() - import time - time.sleep(1) - restarted = Server(args.server.resolve(), data_dir) - try: - restarted_init = restarted.send( - request( - next_id, - "initialize", - { - "protocolVersion": "2024-11-05", - "capabilities": {}, - "clientInfo": {"name": "p0-regression-resume", "version": "2"}, - }, - ) - ) - next_id += 1 - assert restarted_init["result"]["protocolVersion"] == "2024-11-05" - restarted.notify( - {"jsonrpc": "2.0", "method": "notifications/initialized", "params": {}} - ) - restarted_listed = restarted.send(request(next_id, "tools/list", {})) - next_id += 1 - assert restarted_listed["id"] == next_id - 1 - resumed = tool_payload( - tool_call( - restarted, - next_id, - "search_messages", - {"query": "needle target", "limit": 100}, - ) - )["results"] - next_id += 1 - assert [item["message_id"] for item in resumed] == [ - message_id for message_id in ranked_ids_before if message_id != message_ids[0] - ] - resumed_health = tool_payload( - tool_call(restarted, next_id, "get_system_health", {}) - ) - next_id += 1 - assert resumed_health["status"] == "OK" - finally: - restarted.close() - finally: - shutil.rmtree(data_dir, ignore_errors=True) - - -if __name__ == "__main__": - main() |