jason/google-mcp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #513 from fmgs31/drive_files_pagination_fix

Browse Source 
Bugfix: Adding missing nextPageToken for pagination in drive tools
This commit is contained in:
Taylor Wilsdon
2026-02-28 17:45:48 -04:00
committed by GitHub
parent bbdd61d08d 8b69a49d40
commit 328619412d
3 changed files with 194 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
gdrive/drive_helpers.py
View File

@@ -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,
detailed: bool = True,
) -> Dict[str, Any]:
"""
@@ -192,6 +193,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)
detailed: Whether to request size, modifiedTime, and webViewLink fields.
Defaults to True to preserve existing behavior.
@@ -210,6 +212,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:

24
gdrive/drive_tools.py
View File

@@ -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,
@@ -69,6 +70,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').
@@ -78,6 +80,7 @@ async def search_drive_files(
Returns:
str: A formatted list of found files/folders with their details (ID, name, type, and optionally 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}'"
@@ -106,6 +109,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,
detailed=detailed,
)
@@ -114,9 +118,9 @@ async def search_drive_files(
if not files:
return f"No files found for '{query}'."
formatted_files_text_parts = [
f"Found {len(files)} files for {user_google_email} matching '{query}':"
]
next_token = results.get("nextPageToken")
header = f"Found {len(files)} files for {user_google_email} matching '{query}':"
formatted_files_text_parts = [header]
for item in files:
if detailed:
size_str = f", Size: {item.get('size', 'N/A')}" if "size" in item else ""
@@ -127,6 +131,8 @@ async def search_drive_files(
formatted_files_text_parts.append(
f'- Name: "{item["name"]}" (ID: {item["id"]}, Type: {item["mimeType"]})'
)
if next_token:
formatted_files_text_parts.append(f"nextPageToken: {next_token}")
text_output = "\n".join(formatted_files_text_parts)
return text_output
@@ -419,6 +425,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,
@@ -433,6 +440,7 @@ 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.
@@ -440,6 +448,7 @@ async def list_drive_items(
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}'"
@@ -454,6 +463,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,
detailed=detailed,
)
@@ -462,9 +472,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")
header = (
f"Found {len(files)} items in folder '{folder_id}' for {user_google_email}:"
]
)
formatted_items_text_parts = [header]
for item in files:
if detailed:
size_str = f", Size: {item.get('size', 'N/A')}" if "size" in item else ""
@@ -475,6 +487,8 @@ async def list_drive_items(
formatted_items_text_parts.append(
f'- Name: "{item["name"]}" (ID: {item["id"]}, Type: {item["mimeType"]})'
)
if next_token:
formatted_items_text_parts.append(f"nextPageToken: {next_token}")
text_output = "\n".join(formatted_items_text_parts)
return text_output

170
tests/gdrive/test_drive_tools.py
View File

@@ -30,6 +30,176 @@ def _unwrap(tool):
# ---------------------------------------------------------------------------
# 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 at the end of 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 result.endswith("nextPageToken: next_tok_xyz")
@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 at the end of 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 result.endswith("nextPageToken: next_list_tok")
@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
# Helpers
# ---------------------------------------------------------------------------