diff options
Diffstat (limited to 'tests/mcp_protocol_test.py')
| -rw-r--r-- | tests/mcp_protocol_test.py | 114 |
1 files changed, 114 insertions, 0 deletions
diff --git a/tests/mcp_protocol_test.py b/tests/mcp_protocol_test.py index f739168..4d1deb2 100644 --- a/tests/mcp_protocol_test.py +++ b/tests/mcp_protocol_test.py @@ -15,6 +15,7 @@ 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, @@ -41,6 +42,7 @@ class Server: 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 @@ -132,6 +134,8 @@ def main() -> None: "content", ] assert tools["read_messages"]["inputSchema"]["required"] == ["session_id"] + 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"] @@ -165,6 +169,89 @@ def main() -> None: 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, @@ -265,6 +352,33 @@ def main() -> None: 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) |