大規模履歴向けread_messages pagination対応

1 files changed, 21 insertions, 3 deletions

diff --git a/scripts/python/benchmark.py b/scripts/python/benchmark.py
index 5a146ce..0580287 100644
--- a/scripts/python/benchmark.py
+++ b/scripts/python/benchmark.py
@@ -274,13 +274,31 @@ async def write_message(
     return msg
 
 
-async def read_messages(session: ClientSession, session_id: str) -> list[dict]:
-    result = await session.call_tool("read_messages", {"session_id": session_id})
+async def read_messages_page(
+    session: ClientSession, session_id: str, limit: int = 100, cursor: str | None = None
+) -> tuple[list[dict], str | None]:
+    arguments: dict[str, Any] = {"session_id": session_id, "limit": limit}
+    if cursor is not None:
+        arguments["cursor"] = cursor
+    result = await session.call_tool("read_messages", arguments)
     payload = json_result(result)
     history = payload.get("history", [])
     if not isinstance(history, list):
         raise RuntimeError("history is not a list")
-    return history
+    next_cursor = payload.get("nextCursor")
+    if next_cursor is not None and not isinstance(next_cursor, str):
+        raise RuntimeError("nextCursor is not a string")
+    return history, next_cursor
+
+
+async def read_messages(session: ClientSession, session_id: str) -> list[dict]:
+    history: list[dict] = []
+    cursor: str | None = None
+    while True:
+        page, cursor = await read_messages_page(session, session_id, cursor=cursor)
+        history.extend(page)
+        if cursor is None:
+            return history
 
 
 async def search_messages(