Adding nextPageToken for pagination in drive tools

gdrive/drive_helpers.py

@@ -181,6 +181,7 @@ def build_drive_list_params(
     drive_id: Optional[str] = None,
     include_items_from_all_drives: bool = True,
     corpora: Optional[str] = None,
+    page_token: Optional[str] = None,
 ) -> Dict[str, Any]:
     """
     Helper function to build common list parameters for Drive API calls.
@@ -191,6 +192,7 @@ def build_drive_list_params(
         drive_id: Optional shared drive ID
         include_items_from_all_drives: Whether to include items from all drives
         corpora: Optional corpus specification
+        page_token: Optional page token for pagination (from a previous nextPageToken)
 
     Returns:
         Dictionary of parameters for Drive API list calls
@@ -203,6 +205,9 @@ def build_drive_list_params(
         "includeItemsFromAllDrives": include_items_from_all_drives,
     }
 
+    if page_token:
+        list_params["pageToken"] = page_token
+
     if drive_id:
         list_params["driveId"] = drive_id
         if corpora:

gdrive/drive_tools.py

@@ -57,6 +57,7 @@ async def search_drive_files(
     user_google_email: str,
     query: str,
     page_size: int = 10,
+    page_token: Optional[str] = None,
     drive_id: Optional[str] = None,
     include_items_from_all_drives: bool = True,
     corpora: Optional[str] = None,
@@ -68,6 +69,7 @@ async def search_drive_files(
         user_google_email (str): The user's Google email address. Required.
         query (str): The search query string. Supports Google Drive search operators.
         page_size (int): The maximum number of files to return. Defaults to 10.
+        page_token (Optional[str]): Page token from a previous response's nextPageToken to retrieve the next page of results.
         drive_id (Optional[str]): ID of the shared drive to search. If None, behavior depends on `corpora` and `include_items_from_all_drives`.
         include_items_from_all_drives (bool): Whether shared drive items should be included in results. Defaults to True. This is effective when not specifying a `drive_id`.
         corpora (Optional[str]): Bodies of items to query (e.g., 'user', 'domain', 'drive', 'allDrives').
@@ -76,6 +78,7 @@ async def search_drive_files(
 
     Returns:
         str: A formatted list of found files/folders with their details (ID, name, type, size, modified time, link).
+             Includes a nextPageToken line when more results are available.
     """
     logger.info(
         f"[search_drive_files] Invoked. Email: '{user_google_email}', Query: '{query}'"
@@ -104,6 +107,7 @@ async def search_drive_files(
         drive_id=drive_id,
         include_items_from_all_drives=include_items_from_all_drives,
         corpora=corpora,
+        page_token=page_token,
     )
 
     results = await asyncio.to_thread(service.files().list(**list_params).execute)
@@ -111,9 +115,11 @@ async def search_drive_files(
     if not files:
         return f"No files found for '{query}'."
 
-    formatted_files_text_parts = [
+    next_token = results.get("nextPageToken")
-        f"Found {len(files)} files for {user_google_email} matching '{query}':"
+    header = f"Found {len(files)} files for {user_google_email} matching '{query}':"
-    ]
+    if next_token:
+        header += f"\nnextPageToken: {next_token}"
+    formatted_files_text_parts = [header]
     for item in files:
         size_str = f", Size: {item.get('size', 'N/A')}" if "size" in item else ""
         formatted_files_text_parts.append(
@@ -411,6 +417,7 @@ async def list_drive_items(
     user_google_email: str,
     folder_id: str = "root",
     page_size: int = 100,
+    page_token: Optional[str] = None,
     drive_id: Optional[str] = None,
     include_items_from_all_drives: bool = True,
     corpora: Optional[str] = None,
@@ -424,12 +431,14 @@ async def list_drive_items(
         user_google_email (str): The user's Google email address. Required.
         folder_id (str): The ID of the Google Drive folder. Defaults to 'root'. For a shared drive, this can be the shared drive's ID to list its root, or a folder ID within that shared drive.
         page_size (int): The maximum number of items to return. Defaults to 100.
+        page_token (Optional[str]): Page token from a previous response's nextPageToken to retrieve the next page of results.
         drive_id (Optional[str]): ID of the shared drive. If provided, the listing is scoped to this drive.
         include_items_from_all_drives (bool): Whether items from all accessible shared drives should be included if `drive_id` is not set. Defaults to True.
         corpora (Optional[str]): Corpus to query ('user', 'drive', 'allDrives'). If `drive_id` is set and `corpora` is None, 'drive' is used. If None and no `drive_id`, API defaults apply.
 
     Returns:
         str: A formatted list of files/folders in the specified folder.
+             Includes a nextPageToken line when more results are available.
     """
     logger.info(
         f"[list_drive_items] Invoked. Email: '{user_google_email}', Folder ID: '{folder_id}'"
@@ -444,6 +453,7 @@ async def list_drive_items(
         drive_id=drive_id,
         include_items_from_all_drives=include_items_from_all_drives,
         corpora=corpora,
+        page_token=page_token,
     )
 
     results = await asyncio.to_thread(service.files().list(**list_params).execute)
@@ -451,9 +461,11 @@ async def list_drive_items(
     if not files:
         return f"No items found in folder '{folder_id}'."
 
-    formatted_items_text_parts = [
+    next_token = results.get("nextPageToken")
-        f"Found {len(files)} items in folder '{folder_id}' for {user_google_email}:"
+    header = f"Found {len(files)} items in folder '{folder_id}' for {user_google_email}:"
-    ]
+    if next_token:
+        header += f"\nnextPageToken: {next_token}"
+    formatted_items_text_parts = [header]
     for item in files:
         size_str = f", Size: {item.get('size', 'N/A')}" if "size" in item else ""
         formatted_items_text_parts.append(

tests/gdrive/test_drive_tools.py

@@ -1,7 +1,7 @@
 """
 Unit tests for Google Drive MCP tools.
 
-Tests create_drive_folder with mocked API responses.
+Tests create_drive_folder, search_drive_files, and list_drive_items with mocked API responses.
 """
 
 import pytest
@@ -11,6 +11,193 @@ import os
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
 
+from gdrive.drive_tools import list_drive_items, search_drive_files
+
+
+def _unwrap(fn):
+    """Unwrap a decorator chain to the original async function."""
+    if hasattr(fn, "fn"):
+        fn = fn.fn  # FunctionTool wrapper (other server versions)
+    while hasattr(fn, "__wrapped__"):
+        fn = fn.__wrapped__
+    return fn
+
+
+# ---------------------------------------------------------------------------
+# search_drive_files — page_token
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_search_drive_files_page_token_passed_to_api():
+    """page_token is forwarded to the Drive API as pageToken."""
+    mock_service = Mock()
+    mock_service.files().list().execute.return_value = {
+        "files": [
+            {
+                "id": "f1",
+                "name": "Report.pdf",
+                "mimeType": "application/pdf",
+                "webViewLink": "https://drive.google.com/file/f1",
+                "modifiedTime": "2024-01-01T00:00:00Z",
+            }
+        ]
+    }
+
+    await _unwrap(search_drive_files)(
+        service=mock_service,
+        user_google_email="user@example.com",
+        query="budget",
+        page_token="tok_abc123",
+    )
+
+    call_kwargs = mock_service.files.return_value.list.call_args.kwargs
+    assert call_kwargs.get("pageToken") == "tok_abc123"
+
+
+@pytest.mark.asyncio
+async def test_search_drive_files_next_page_token_in_output():
+    """nextPageToken from the API response is appended to the output."""
+    mock_service = Mock()
+    mock_service.files().list().execute.return_value = {
+        "files": [
+            {
+                "id": "f2",
+                "name": "Notes.docx",
+                "mimeType": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+                "webViewLink": "https://drive.google.com/file/f2",
+                "modifiedTime": "2024-02-01T00:00:00Z",
+            }
+        ],
+        "nextPageToken": "next_tok_xyz",
+    }
+
+    result = await _unwrap(search_drive_files)(
+        service=mock_service,
+        user_google_email="user@example.com",
+        query="notes",
+    )
+
+    assert "nextPageToken: next_tok_xyz" in result
+
+
+@pytest.mark.asyncio
+async def test_search_drive_files_no_next_page_token_when_absent():
+    """nextPageToken does not appear in output when the API has no more pages."""
+    mock_service = Mock()
+    mock_service.files().list().execute.return_value = {
+        "files": [
+            {
+                "id": "f3",
+                "name": "Summary.txt",
+                "mimeType": "text/plain",
+                "webViewLink": "https://drive.google.com/file/f3",
+                "modifiedTime": "2024-03-01T00:00:00Z",
+            }
+        ]
+        # no nextPageToken key
+    }
+
+    result = await _unwrap(search_drive_files)(
+        service=mock_service,
+        user_google_email="user@example.com",
+        query="summary",
+    )
+
+    assert "nextPageToken" not in result
+
+
+# ---------------------------------------------------------------------------
+# list_drive_items — page_token
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+@patch("gdrive.drive_tools.resolve_folder_id", new_callable=AsyncMock)
+async def test_list_drive_items_page_token_passed_to_api(mock_resolve_folder):
+    """page_token is forwarded to the Drive API as pageToken."""
+    mock_resolve_folder.return_value = "root"
+    mock_service = Mock()
+    mock_service.files().list().execute.return_value = {
+        "files": [
+            {
+                "id": "folder1",
+                "name": "Archive",
+                "mimeType": "application/vnd.google-apps.folder",
+                "webViewLink": "https://drive.google.com/drive/folders/folder1",
+                "modifiedTime": "2024-01-15T00:00:00Z",
+            }
+        ]
+    }
+
+    await _unwrap(list_drive_items)(
+        service=mock_service,
+        user_google_email="user@example.com",
+        page_token="tok_page2",
+    )
+
+    call_kwargs = mock_service.files.return_value.list.call_args.kwargs
+    assert call_kwargs.get("pageToken") == "tok_page2"
+
+
+@pytest.mark.asyncio
+@patch("gdrive.drive_tools.resolve_folder_id", new_callable=AsyncMock)
+async def test_list_drive_items_next_page_token_in_output(mock_resolve_folder):
+    """nextPageToken from the API response is appended to the output."""
+    mock_resolve_folder.return_value = "root"
+    mock_service = Mock()
+    mock_service.files().list().execute.return_value = {
+        "files": [
+            {
+                "id": "file99",
+                "name": "data.csv",
+                "mimeType": "text/csv",
+                "webViewLink": "https://drive.google.com/file/file99",
+                "modifiedTime": "2024-04-01T00:00:00Z",
+            }
+        ],
+        "nextPageToken": "next_list_tok",
+    }
+
+    result = await _unwrap(list_drive_items)(
+        service=mock_service,
+        user_google_email="user@example.com",
+    )
+
+    assert "nextPageToken: next_list_tok" in result
+
+
+@pytest.mark.asyncio
+@patch("gdrive.drive_tools.resolve_folder_id", new_callable=AsyncMock)
+async def test_list_drive_items_no_next_page_token_when_absent(mock_resolve_folder):
+    """nextPageToken does not appear in output when the API has no more pages."""
+    mock_resolve_folder.return_value = "root"
+    mock_service = Mock()
+    mock_service.files().list().execute.return_value = {
+        "files": [
+            {
+                "id": "file100",
+                "name": "readme.txt",
+                "mimeType": "text/plain",
+                "webViewLink": "https://drive.google.com/file/file100",
+                "modifiedTime": "2024-05-01T00:00:00Z",
+            }
+        ]
+        # no nextPageToken key
+    }
+
+    result = await _unwrap(list_drive_items)(
+        service=mock_service,
+        user_google_email="user@example.com",
+    )
+
+    assert "nextPageToken" not in result
+
+
+# ---------------------------------------------------------------------------
+# create_drive_folder
+# ---------------------------------------------------------------------------
+
 
 @pytest.mark.asyncio
 async def test_create_drive_folder():