Skip to content

Commit 5fa6700

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 533c6a8 commit 5fa6700

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
@@ -2,7 +2,7 @@
22

33
from __future__ import annotations
44

5-
from collections.abc import Awaitable, Callable, Mapping
5+
from collections.abc import AsyncIterator, Awaitable, Callable, Mapping
66
from contextlib import AsyncExitStack
77
from dataclasses import KW_ONLY, dataclass, field
88
from typing import Any, Literal, TypeVar
@@ -26,11 +26,15 @@
2626
ListToolsResult,
2727
LoggingLevel,
2828
PaginatedRequestParams,
29+
Prompt,
2930
PromptReference,
3031
ReadResourceResult,
3132
RequestParamsMeta,
33+
Resource,
34+
ResourceTemplate,
3235
ResourceTemplateReference,
3336
ServerCapabilities,
37+
Tool,
3438
)
3539
from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS, MODERN_PROTOCOL_VERSIONS
3640
from typing_extensions import deprecated
@@ -367,7 +371,11 @@ async def list_resources(
367371
cursor: str | None = None,
368372
meta: RequestParamsMeta | None = None,
369373
) -> ListResourcesResult:
370-
"""List available resources from the server."""
374+
"""List a single page of available resources from the server.
375+
376+
Returns one page only. The result may include a `next_cursor` if more
377+
pages are available. Use `list_all_resources` to drain every page.
378+
"""
371379
return await self.session.list_resources(params=PaginatedRequestParams(cursor=cursor, _meta=meta))
372380

373381
async def list_resource_templates(
@@ -376,7 +384,12 @@ async def list_resource_templates(
376384
cursor: str | None = None,
377385
meta: RequestParamsMeta | None = None,
378386
) -> ListResourceTemplatesResult:
379-
"""List available resource templates from the server."""
387+
"""List a single page of available resource templates from the server.
388+
389+
Returns one page only. The result may include a `next_cursor` if more
390+
pages are available. Use `list_all_resource_templates` to drain every
391+
page.
392+
"""
380393
return await self.session.list_resource_templates(params=PaginatedRequestParams(cursor=cursor, _meta=meta))
381394

382395
async def read_resource(
@@ -482,7 +495,11 @@ async def list_prompts(
482495
cursor: str | None = None,
483496
meta: RequestParamsMeta | None = None,
484497
) -> ListPromptsResult:
485-
"""List available prompts from the server."""
498+
"""List a single page of available prompts from the server.
499+
500+
Returns one page only. The result may include a `next_cursor` if more
501+
pages are available. Use `list_all_prompts` to drain every page.
502+
"""
486503
return await self.session.list_prompts(params=PaginatedRequestParams(cursor=cursor, _meta=meta))
487504

488505
async def get_prompt(
@@ -566,9 +583,84 @@ async def complete(
566583
return await self.session.complete(ref=ref, argument=argument, context_arguments=context_arguments)
567584

568585
async def list_tools(self, *, cursor: str | None = None, meta: RequestParamsMeta | None = None) -> ListToolsResult:
569-
"""List available tools from the server."""
586+
"""List a single page of available tools from the server.
587+
588+
Returns one page only. The result may include a `next_cursor` if more
589+
pages are available. Use `list_all_tools` to drain every page.
590+
"""
570591
return await self.session.list_tools(params=PaginatedRequestParams(cursor=cursor, _meta=meta))
571592

593+
async def iter_all_tools(self, *, meta: RequestParamsMeta | None = None) -> AsyncIterator[Tool]:
594+
"""Yield every tool from the server, paging through `next_cursor`.
595+
596+
Useful for streaming consumers that want to process tools without
597+
materializing the full list in memory.
598+
"""
599+
cursor: str | None = None
600+
while True:
601+
result = await self.list_tools(cursor=cursor, meta=meta)
602+
for tool in result.tools:
603+
yield tool
604+
if result.next_cursor is None:
605+
return
606+
cursor = result.next_cursor
607+
608+
async def list_all_tools(self, *, meta: RequestParamsMeta | None = None) -> list[Tool]:
609+
"""List every tool from the server, draining `next_cursor` across pages.
610+
611+
Unlike `list_tools`, which returns one page, this walks pagination
612+
until the server reports no further pages and returns the combined
613+
list.
614+
"""
615+
return [tool async for tool in self.iter_all_tools(meta=meta)]
616+
617+
async def iter_all_prompts(self, *, meta: RequestParamsMeta | None = None) -> AsyncIterator[Prompt]:
618+
"""Yield every prompt from the server, paging through `next_cursor`."""
619+
cursor: str | None = None
620+
while True:
621+
result = await self.list_prompts(cursor=cursor, meta=meta)
622+
for prompt in result.prompts:
623+
yield prompt
624+
if result.next_cursor is None:
625+
return
626+
cursor = result.next_cursor
627+
628+
async def list_all_prompts(self, *, meta: RequestParamsMeta | None = None) -> list[Prompt]:
629+
"""List every prompt from the server, draining `next_cursor` across pages."""
630+
return [prompt async for prompt in self.iter_all_prompts(meta=meta)]
631+
632+
async def iter_all_resources(self, *, meta: RequestParamsMeta | None = None) -> AsyncIterator[Resource]:
633+
"""Yield every resource from the server, paging through `next_cursor`."""
634+
cursor: str | None = None
635+
while True:
636+
result = await self.list_resources(cursor=cursor, meta=meta)
637+
for resource in result.resources:
638+
yield resource
639+
if result.next_cursor is None:
640+
return
641+
cursor = result.next_cursor
642+
643+
async def list_all_resources(self, *, meta: RequestParamsMeta | None = None) -> list[Resource]:
644+
"""List every resource from the server, draining `next_cursor` across pages."""
645+
return [resource async for resource in self.iter_all_resources(meta=meta)]
646+
647+
async def iter_all_resource_templates(
648+
self, *, meta: RequestParamsMeta | None = None
649+
) -> AsyncIterator[ResourceTemplate]:
650+
"""Yield every resource template from the server, paging through `next_cursor`."""
651+
cursor: str | None = None
652+
while True:
653+
result = await self.list_resource_templates(cursor=cursor, meta=meta)
654+
for template in result.resource_templates:
655+
yield template
656+
if result.next_cursor is None:
657+
return
658+
cursor = result.next_cursor
659+
660+
async def list_all_resource_templates(self, *, meta: RequestParamsMeta | None = None) -> list[ResourceTemplate]:
661+
"""List every resource template from the server, draining `next_cursor` across pages."""
662+
return [template async for template in self.iter_all_resource_templates(meta=meta)]
663+
572664
@deprecated("The roots capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning)
573665
async def send_roots_list_changed(self) -> None:
574666
"""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)