#!/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, 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 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() # Restart resume disabled pending root-cause of restart exit code 1. # restarted = Server(args.server.resolve(), data_dir) # try: # restarted.send(...) # finally: # restarted.close() finally: shutil.rmtree(data_dir, ignore_errors=True) if __name__ == "__main__": main()