diff options
Diffstat (limited to 'tests/mcp_protocol_test.py')
| -rw-r--r-- | tests/mcp_protocol_test.py | 236 |
1 files changed, 236 insertions, 0 deletions
diff --git a/tests/mcp_protocol_test.py b/tests/mcp_protocol_test.py new file mode 100644 index 0000000..118043f --- /dev/null +++ b/tests/mcp_protocol_test.py @@ -0,0 +1,236 @@ +#!/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.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:]}" + ) + 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 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", + "delete_message", + "rebuild_index", + "get_system_health", + } + assert tools["write_message"]["inputSchema"]["required"] == [ + "session_id", + "role", + "content", + ] + assert tools["read_messages"]["inputSchema"]["required"] == ["session_id"] + assert tools["search_messages"]["inputSchema"]["required"] == ["query"] + assert tools["search_messages"]["inputSchema"]["properties"]["limit"]["type"] == "integer" + assert tools["delete_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"]) + 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 + + 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 + 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() + finally: + shutil.rmtree(data_dir, ignore_errors=True) + + +if __name__ == "__main__": + main() |