大規模履歴向けread_messages pagination対応

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)