Skip to content

Commit a7adecd

Browse files
adityasingh2400CJGjr
authored andcommitted
Add list_all_* helpers that drain pagination on the client
Currently Client.list_tools / list_prompts / list_resources / list_resource_templates return a single page and the caller has to loop on next_cursor manually. Add list_all_tools / list_all_prompts / list_all_resources / list_all_resource_templates that walk next_cursor until exhausted, plus iter_all_* async iterators for streaming consumers. The single-page methods get a docstring update pointing at the new drains. ClientSessionGroup switches its tool/prompt/resource aggregation to the drain helper so its consumers always see the full collection across multi-page servers. Implements the helper maxisbey endorsed in #2556. Rebased onto the v2 rework: types import from mcp_types, the stream-spy tests run in legacy mode, and the test Tool carries a valid input_schema.
1 parent 3a6f299 commit a7adecd

4 files changed

Lines changed: 384 additions & 19 deletions

File tree

src/mcp/client/client.py

Lines changed: 97 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import hashlib
66
import logging
77
import uuid
8-
from collections.abc import Awaitable, Callable, Mapping, Sequence
8+
from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence
99
from contextlib import AbstractAsyncContextManager, AsyncExitStack
1010
from dataclasses import KW_ONLY, dataclass, field
1111
from typing import Any, Literal, TypeVar, cast
@@ -32,12 +32,16 @@
3232
ListToolsResult,
3333
LoggingLevel,
3434
PaginatedRequestParams,
35+
Prompt,
3536
PromptReference,
3637
ReadResourceResult,
3738
RequestParamsMeta,
39+
Resource,
40+
ResourceTemplate,
3841
ResourceTemplateReference,
3942
Result,
4043
ServerCapabilities,
44+
Tool,
4145
)
4246
from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS, MODERN_PROTOCOL_VERSIONS
4347
from typing_extensions import deprecated
@@ -580,7 +584,11 @@ async def list_resources(
580584
meta: RequestParamsMeta | None = None,
581585
cache_mode: CacheMode = "use",
582586
) -> ListResourcesResult:
583-
"""List available resources from the server."""
587+
"""List a single page of available resources from the server.
588+
589+
Returns one page only. The result may include a `next_cursor` if more
590+
pages are available. Use `list_all_resources` to drain every page.
591+
"""
584592
return await self._cached_fetch(
585593
"resources/list",
586594
cursor=cursor,
@@ -596,7 +604,12 @@ async def list_resource_templates(
596604
meta: RequestParamsMeta | None = None,
597605
cache_mode: CacheMode = "use",
598606
) -> ListResourceTemplatesResult:
599-
"""List available resource templates from the server."""
607+
"""List a single page of available resource templates from the server.
608+
609+
Returns one page only. The result may include a `next_cursor` if more
610+
pages are available. Use `list_all_resource_templates` to drain every
611+
page.
612+
"""
600613
return await self._cached_fetch(
601614
"resources/templates/list",
602615
cursor=cursor,
@@ -814,7 +827,11 @@ async def list_prompts(
814827
meta: RequestParamsMeta | None = None,
815828
cache_mode: CacheMode = "use",
816829
) -> ListPromptsResult:
817-
"""List available prompts from the server."""
830+
"""List a single page of available prompts from the server.
831+
832+
Returns one page only. The result may include a `next_cursor` if more
833+
pages are available. Use `list_all_prompts` to drain every page.
834+
"""
818835
return await self._cached_fetch(
819836
"prompts/list",
820837
cursor=cursor,
@@ -912,7 +929,11 @@ async def list_tools(
912929
meta: RequestParamsMeta | None = None,
913930
cache_mode: CacheMode = "use",
914931
) -> ListToolsResult:
915-
"""List available tools from the server."""
932+
"""List a single page of available tools from the server.
933+
934+
Returns one page only. The result may include a `next_cursor` if more
935+
pages are available. Use `list_all_tools` to drain every page.
936+
"""
916937
return await self._cached_fetch(
917938
"tools/list",
918939
cursor=cursor,
@@ -927,6 +948,77 @@ async def list_tools(
927948
),
928949
)
929950

951+
async def iter_all_tools(self, *, meta: RequestParamsMeta | None = None) -> AsyncIterator[Tool]:
952+
"""Yield every tool from the server, paging through `next_cursor`.
953+
954+
Useful for streaming consumers that want to process tools without
955+
materializing the full list in memory.
956+
"""
957+
cursor: str | None = None
958+
while True:
959+
result = await self.list_tools(cursor=cursor, meta=meta)
960+
for tool in result.tools:
961+
yield tool
962+
if result.next_cursor is None:
963+
return
964+
cursor = result.next_cursor
965+
966+
async def list_all_tools(self, *, meta: RequestParamsMeta | None = None) -> list[Tool]:
967+
"""List every tool from the server, draining `next_cursor` across pages.
968+
969+
Unlike `list_tools`, which returns one page, this walks pagination
970+
until the server reports no further pages and returns the combined
971+
list.
972+
"""
973+
return [tool async for tool in self.iter_all_tools(meta=meta)]
974+
975+
async def iter_all_prompts(self, *, meta: RequestParamsMeta | None = None) -> AsyncIterator[Prompt]:
976+
"""Yield every prompt from the server, paging through `next_cursor`."""
977+
cursor: str | None = None
978+
while True:
979+
result = await self.list_prompts(cursor=cursor, meta=meta)
980+
for prompt in result.prompts:
981+
yield prompt
982+
if result.next_cursor is None:
983+
return
984+
cursor = result.next_cursor
985+
986+
async def list_all_prompts(self, *, meta: RequestParamsMeta | None = None) -> list[Prompt]:
987+
"""List every prompt from the server, draining `next_cursor` across pages."""
988+
return [prompt async for prompt in self.iter_all_prompts(meta=meta)]
989+
990+
async def iter_all_resources(self, *, meta: RequestParamsMeta | None = None) -> AsyncIterator[Resource]:
991+
"""Yield every resource from the server, paging through `next_cursor`."""
992+
cursor: str | None = None
993+
while True:
994+
result = await self.list_resources(cursor=cursor, meta=meta)
995+
for resource in result.resources:
996+
yield resource
997+
if result.next_cursor is None:
998+
return
999+
cursor = result.next_cursor
1000+
1001+
async def list_all_resources(self, *, meta: RequestParamsMeta | None = None) -> list[Resource]:
1002+
"""List every resource from the server, draining `next_cursor` across pages."""
1003+
return [resource async for resource in self.iter_all_resources(meta=meta)]
1004+
1005+
async def iter_all_resource_templates(
1006+
self, *, meta: RequestParamsMeta | None = None
1007+
) -> AsyncIterator[ResourceTemplate]:
1008+
"""Yield every resource template from the server, paging through `next_cursor`."""
1009+
cursor: str | None = None
1010+
while True:
1011+
result = await self.list_resource_templates(cursor=cursor, meta=meta)
1012+
for template in result.resource_templates:
1013+
yield template
1014+
if result.next_cursor is None:
1015+
return
1016+
cursor = result.next_cursor
1017+
1018+
async def list_all_resource_templates(self, *, meta: RequestParamsMeta | None = None) -> list[ResourceTemplate]:
1019+
"""List every resource template from the server, draining `next_cursor` across pages."""
1020+
return [template async for template in self.iter_all_resource_templates(meta=meta)]
1021+
9301022
@deprecated("The roots capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning)
9311023
async def send_roots_list_changed(self) -> None:
9321024
"""Send a notification that the roots list has changed."""

src/mcp/client/session_group.py

Lines changed: 29 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88

99
import contextlib
1010
import logging
11-
from collections.abc import Callable
11+
from collections.abc import Awaitable, Callable
1212
from dataclasses import dataclass
1313
from types import TracebackType
1414
from typing import Any, Literal, TypeAlias, overload
@@ -67,6 +67,28 @@ class StreamableHttpParameters(BaseModel):
6767
ServerParameters: TypeAlias = StdioServerParameters | SseServerParameters | StreamableHttpParameters
6868

6969

70+
async def _drain_paginated(
71+
fetch_page: Callable[..., Awaitable[Any]],
72+
attribute: str,
73+
) -> list[Any]:
74+
"""Drain a paginated `session.list_*` call across `next_cursor` pages.
75+
76+
`fetch_page` is one of the ClientSession `list_*` methods that takes a
77+
`params=PaginatedRequestParams(...)` keyword. `attribute` is the name of
78+
the list attribute on the result (e.g. `"tools"`, `"prompts"`).
79+
"""
80+
items: list[Any] = []
81+
cursor: str | None = None
82+
while True:
83+
params = types.PaginatedRequestParams(cursor=cursor) if cursor is not None else None
84+
result = await fetch_page(params=params)
85+
items.extend(getattr(result, attribute))
86+
next_cursor = getattr(result, "next_cursor", None)
87+
if next_cursor is None:
88+
return items
89+
cursor = next_cursor
90+
91+
7092
# Use dataclass instead of Pydantic BaseModel
7193
# because Pydantic BaseModel cannot handle Protocol fields.
7294
@dataclass
@@ -383,9 +405,11 @@ async def _aggregate_components(self, server_info: types.Implementation, session
383405
tools_temp: dict[str, types.Tool] = {}
384406
tool_to_session_temp: dict[str, mcp.ClientSession] = {}
385407

386-
# Query the server for its prompts and aggregate to list.
408+
# Query the server for its prompts and aggregate to list. Drain
409+
# pagination so we don't drop later pages on servers that split
410+
# results across multiple `next_cursor` responses.
387411
try:
388-
prompts = (await session.list_prompts()).prompts
412+
prompts = await _drain_paginated(session.list_prompts, "prompts")
389413
for prompt in prompts:
390414
name = self._component_name(prompt.name, server_info)
391415
prompts_temp[name] = prompt
@@ -395,7 +419,7 @@ async def _aggregate_components(self, server_info: types.Implementation, session
395419

396420
# Query the server for its resources and aggregate to list.
397421
try:
398-
resources = (await session.list_resources()).resources
422+
resources = await _drain_paginated(session.list_resources, "resources")
399423
for resource in resources:
400424
name = self._component_name(resource.name, server_info)
401425
resources_temp[name] = resource
@@ -405,7 +429,7 @@ async def _aggregate_components(self, server_info: types.Implementation, session
405429

406430
# Query the server for its tools and aggregate to list.
407431
try:
408-
tools = (await session.list_tools()).tools
432+
tools = await _drain_paginated(session.list_tools, "tools")
409433
for tool in tools:
410434
name = self._component_name(tool.name, server_info)
411435
tools_temp[name] = tool

0 commit comments

Comments
 (0)