mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-21 21:14:17 +02:00
Bound client auto-pagination loops to prevent unbounded list fetches (#3411)
* Cap client auto-pagination pages
🤖 Generated with GPT-5.2-Codex
* Raise on pagination limit instead of returning partial data
Add max_pages kwarg (default 250) to list_tools/list_resources/
list_resource_templates/list_prompts so users can control the bound.
This commit is contained in:
parent
799c4f1673
commit
85c71fa834
4 changed files with 160 additions and 12 deletions
|
|
@ -20,6 +20,8 @@ from fastmcp.utilities.logging import get_logger
|
|||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
AUTO_PAGINATION_MAX_PAGES = 250
|
||||
|
||||
# Type alias for task response union (SEP-1686 graceful degradation)
|
||||
PromptTaskResponseUnion = RootModel[
|
||||
mcp.types.CreateTaskResult | mcp.types.GetPromptResult
|
||||
|
|
@ -54,25 +56,31 @@ class ClientPromptsMixin:
|
|||
)
|
||||
return result
|
||||
|
||||
async def list_prompts(self: Client) -> list[mcp.types.Prompt]:
|
||||
async def list_prompts(
|
||||
self: Client,
|
||||
max_pages: int = AUTO_PAGINATION_MAX_PAGES,
|
||||
) -> list[mcp.types.Prompt]:
|
||||
"""Retrieve all prompts available on the server.
|
||||
|
||||
This method automatically fetches all pages if the server paginates results,
|
||||
returning the complete list. For manual pagination control (e.g., to handle
|
||||
large result sets incrementally), use list_prompts_mcp() with the cursor parameter.
|
||||
|
||||
Args:
|
||||
max_pages: Maximum number of pages to fetch before raising. Defaults to 250.
|
||||
|
||||
Returns:
|
||||
list[mcp.types.Prompt]: A list of all Prompt objects.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If called while the client is not connected.
|
||||
RuntimeError: If the page limit is reached before pagination completes.
|
||||
McpError: If the request results in a TimeoutError | JSONRPCError
|
||||
"""
|
||||
all_prompts: list[mcp.types.Prompt] = []
|
||||
cursor: str | None = None
|
||||
seen_cursors: set[str] = set()
|
||||
|
||||
while True:
|
||||
for _ in range(max_pages):
|
||||
result = await self.list_prompts_mcp(cursor=cursor)
|
||||
all_prompts.extend(result.prompts)
|
||||
if not result.nextCursor:
|
||||
|
|
@ -85,6 +93,13 @@ class ClientPromptsMixin:
|
|||
break
|
||||
seen_cursors.add(result.nextCursor)
|
||||
cursor = result.nextCursor
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"[{self.name}] Reached auto-pagination limit"
|
||||
f" ({max_pages} pages) for list_prompts."
|
||||
" Use list_prompts_mcp() with cursor for manual pagination,"
|
||||
" or increase max_pages."
|
||||
)
|
||||
|
||||
return all_prompts
|
||||
|
||||
|
|
|
|||
|
|
@ -19,6 +19,8 @@ from fastmcp.utilities.logging import get_logger
|
|||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
AUTO_PAGINATION_MAX_PAGES = 250
|
||||
|
||||
# Type alias for task response union (SEP-1686 graceful degradation)
|
||||
ResourceTaskResponseUnion = RootModel[
|
||||
mcp.types.CreateTaskResult | mcp.types.ReadResourceResult
|
||||
|
|
@ -53,25 +55,31 @@ class ClientResourcesMixin:
|
|||
)
|
||||
return result
|
||||
|
||||
async def list_resources(self: Client) -> list[mcp.types.Resource]:
|
||||
async def list_resources(
|
||||
self: Client,
|
||||
max_pages: int = AUTO_PAGINATION_MAX_PAGES,
|
||||
) -> list[mcp.types.Resource]:
|
||||
"""Retrieve all resources available on the server.
|
||||
|
||||
This method automatically fetches all pages if the server paginates results,
|
||||
returning the complete list. For manual pagination control (e.g., to handle
|
||||
large result sets incrementally), use list_resources_mcp() with the cursor parameter.
|
||||
|
||||
Args:
|
||||
max_pages: Maximum number of pages to fetch before raising. Defaults to 250.
|
||||
|
||||
Returns:
|
||||
list[mcp.types.Resource]: A list of all Resource objects.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If called while the client is not connected.
|
||||
RuntimeError: If the page limit is reached before pagination completes.
|
||||
McpError: If the request results in a TimeoutError | JSONRPCError
|
||||
"""
|
||||
all_resources: list[mcp.types.Resource] = []
|
||||
cursor: str | None = None
|
||||
seen_cursors: set[str] = set()
|
||||
|
||||
while True:
|
||||
for _ in range(max_pages):
|
||||
result = await self.list_resources_mcp(cursor=cursor)
|
||||
all_resources.extend(result.resources)
|
||||
if not result.nextCursor:
|
||||
|
|
@ -84,6 +92,13 @@ class ClientResourcesMixin:
|
|||
break
|
||||
seen_cursors.add(result.nextCursor)
|
||||
cursor = result.nextCursor
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"[{self.name}] Reached auto-pagination limit"
|
||||
f" ({max_pages} pages) for list_resources."
|
||||
" Use list_resources_mcp() with cursor for manual pagination,"
|
||||
" or increase max_pages."
|
||||
)
|
||||
|
||||
return all_resources
|
||||
|
||||
|
|
@ -110,7 +125,10 @@ class ClientResourcesMixin:
|
|||
)
|
||||
return result
|
||||
|
||||
async def list_resource_templates(self: Client) -> list[mcp.types.ResourceTemplate]:
|
||||
async def list_resource_templates(
|
||||
self: Client,
|
||||
max_pages: int = AUTO_PAGINATION_MAX_PAGES,
|
||||
) -> list[mcp.types.ResourceTemplate]:
|
||||
"""Retrieve all resource templates available on the server.
|
||||
|
||||
This method automatically fetches all pages if the server paginates results,
|
||||
|
|
@ -118,18 +136,21 @@ class ClientResourcesMixin:
|
|||
large result sets incrementally), use list_resource_templates_mcp() with the
|
||||
cursor parameter.
|
||||
|
||||
Args:
|
||||
max_pages: Maximum number of pages to fetch before raising. Defaults to 250.
|
||||
|
||||
Returns:
|
||||
list[mcp.types.ResourceTemplate]: A list of all ResourceTemplate objects.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If called while the client is not connected.
|
||||
RuntimeError: If the page limit is reached before pagination completes.
|
||||
McpError: If the request results in a TimeoutError | JSONRPCError
|
||||
"""
|
||||
all_templates: list[mcp.types.ResourceTemplate] = []
|
||||
cursor: str | None = None
|
||||
seen_cursors: set[str] = set()
|
||||
|
||||
while True:
|
||||
for _ in range(max_pages):
|
||||
result = await self.list_resource_templates_mcp(cursor=cursor)
|
||||
all_templates.extend(result.resourceTemplates)
|
||||
if not result.nextCursor:
|
||||
|
|
@ -143,6 +164,13 @@ class ClientResourcesMixin:
|
|||
break
|
||||
seen_cursors.add(result.nextCursor)
|
||||
cursor = result.nextCursor
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"[{self.name}] Reached auto-pagination limit"
|
||||
f" ({max_pages} pages) for list_resource_templates."
|
||||
" Use list_resource_templates_mcp() with cursor for manual pagination,"
|
||||
" or increase max_pages."
|
||||
)
|
||||
|
||||
return all_templates
|
||||
|
||||
|
|
|
|||
|
|
@ -25,6 +25,8 @@ from fastmcp.utilities.types import get_cached_typeadapter
|
|||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
AUTO_PAGINATION_MAX_PAGES = 250
|
||||
|
||||
# Type alias for task response union (SEP-1686 graceful degradation)
|
||||
ToolTaskResponseUnion = RootModel[mcp.types.CreateTaskResult | mcp.types.CallToolResult]
|
||||
|
||||
|
|
@ -57,25 +59,31 @@ class ClientToolsMixin:
|
|||
)
|
||||
return result
|
||||
|
||||
async def list_tools(self: Client) -> list[mcp.types.Tool]:
|
||||
async def list_tools(
|
||||
self: Client,
|
||||
max_pages: int = AUTO_PAGINATION_MAX_PAGES,
|
||||
) -> list[mcp.types.Tool]:
|
||||
"""Retrieve all tools available on the server.
|
||||
|
||||
This method automatically fetches all pages if the server paginates results,
|
||||
returning the complete list. For manual pagination control (e.g., to handle
|
||||
large result sets incrementally), use list_tools_mcp() with the cursor parameter.
|
||||
|
||||
Args:
|
||||
max_pages: Maximum number of pages to fetch before raising. Defaults to 250.
|
||||
|
||||
Returns:
|
||||
list[mcp.types.Tool]: A list of all Tool objects.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If called while the client is not connected.
|
||||
RuntimeError: If the page limit is reached before pagination completes.
|
||||
McpError: If the request results in a TimeoutError | JSONRPCError
|
||||
"""
|
||||
all_tools: list[mcp.types.Tool] = []
|
||||
cursor: str | None = None
|
||||
seen_cursors: set[str] = set()
|
||||
|
||||
while True:
|
||||
for _ in range(max_pages):
|
||||
result = await self.list_tools_mcp(cursor=cursor)
|
||||
all_tools.extend(result.tools)
|
||||
if not result.nextCursor:
|
||||
|
|
@ -88,6 +96,13 @@ class ClientToolsMixin:
|
|||
break
|
||||
seen_cursors.add(result.nextCursor)
|
||||
cursor = result.nextCursor
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"[{self.name}] Reached auto-pagination limit"
|
||||
f" ({max_pages} pages) for list_tools."
|
||||
" Use list_tools_mcp() with cursor for manual pagination,"
|
||||
" or increase max_pages."
|
||||
)
|
||||
|
||||
return all_tools
|
||||
|
||||
|
|
|
|||
|
|
@ -420,6 +420,96 @@ class TestPaginationCycleDetection:
|
|||
assert len(tools) == 1
|
||||
assert tools[0].name == "my_tool"
|
||||
|
||||
async def test_tools_raises_on_auto_pagination_limit(self) -> None:
|
||||
"""list_tools should raise RuntimeError after exceeding max_pages."""
|
||||
server = FastMCP()
|
||||
|
||||
@server.tool
|
||||
def my_tool() -> str:
|
||||
return "ok"
|
||||
|
||||
async with Client(server) as client:
|
||||
original = client.list_tools_mcp
|
||||
call_count = 0
|
||||
|
||||
async def returning_unique_cursor(
|
||||
*,
|
||||
cursor: str | None = None,
|
||||
) -> mcp.types.ListToolsResult:
|
||||
nonlocal call_count
|
||||
result = await original(cursor=cursor)
|
||||
call_count += 1
|
||||
result.nextCursor = f"cursor-{call_count}"
|
||||
return result
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
client, "list_tools_mcp", side_effect=returning_unique_cursor
|
||||
),
|
||||
pytest.raises(RuntimeError, match="auto-pagination limit"),
|
||||
):
|
||||
await client.list_tools(max_pages=5)
|
||||
|
||||
async def test_resources_raises_on_auto_pagination_limit(self) -> None:
|
||||
"""list_resources should raise RuntimeError after exceeding max_pages."""
|
||||
server = FastMCP()
|
||||
|
||||
@server.resource("test://r")
|
||||
def my_resource() -> str:
|
||||
return "data"
|
||||
|
||||
async with Client(server) as client:
|
||||
original = client.list_resources_mcp
|
||||
call_count = 0
|
||||
|
||||
async def returning_unique_cursor(
|
||||
*,
|
||||
cursor: str | None = None,
|
||||
) -> mcp.types.ListResourcesResult:
|
||||
nonlocal call_count
|
||||
result = await original(cursor=cursor)
|
||||
call_count += 1
|
||||
result.nextCursor = f"cursor-{call_count}"
|
||||
return result
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
client, "list_resources_mcp", side_effect=returning_unique_cursor
|
||||
),
|
||||
pytest.raises(RuntimeError, match="auto-pagination limit"),
|
||||
):
|
||||
await client.list_resources(max_pages=5)
|
||||
|
||||
async def test_prompts_raises_on_auto_pagination_limit(self) -> None:
|
||||
"""list_prompts should raise RuntimeError after exceeding max_pages."""
|
||||
server = FastMCP()
|
||||
|
||||
@server.prompt
|
||||
def my_prompt() -> str:
|
||||
return "text"
|
||||
|
||||
async with Client(server) as client:
|
||||
original = client.list_prompts_mcp
|
||||
call_count = 0
|
||||
|
||||
async def returning_unique_cursor(
|
||||
*,
|
||||
cursor: str | None = None,
|
||||
) -> mcp.types.ListPromptsResult:
|
||||
nonlocal call_count
|
||||
result = await original(cursor=cursor)
|
||||
call_count += 1
|
||||
result.nextCursor = f"cursor-{call_count}"
|
||||
return result
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
client, "list_prompts_mcp", side_effect=returning_unique_cursor
|
||||
),
|
||||
pytest.raises(RuntimeError, match="auto-pagination limit"),
|
||||
):
|
||||
await client.list_prompts(max_pages=5)
|
||||
|
||||
async def test_normal_pagination_unaffected(self) -> None:
|
||||
"""Cycle detection should not interfere with normal pagination."""
|
||||
server = FastMCP(list_page_size=10)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue