Skip to content

Commit 1d83eb7

Browse files
committed
fix: address cubic AI review feedback
- Use accumulator with extend() instead of list concatenation to avoid quadratic CPU/memory churn on large paginated listings - Absorb final aggregated tools result as complete to prune stale session state (x-mcp-header maps and output schemas) Fixes cubic review comments on PR #3156 Signed-off-by: dongjiang <dongjiang1989@126.com>
1 parent c617269 commit 1d83eb7

1 file changed

Lines changed: 8 additions & 2 deletions

File tree

src/mcp/client/client.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -992,17 +992,21 @@ async def _drain_all_pages(
992992
pages = 1
993993
max_pages = self.list_max_pages
994994

995+
# Use an accumulator to avoid quadratic list copying
996+
accumulator = get_items(first)
997+
995998
while cursor is not None:
996999
if cursor in seen:
9971000
raise CursorCycleError(method, cursor)
9981001
seen.add(cursor)
9991002
if max_pages > 0 and pages >= max_pages:
10001003
raise PaginationExceededError(method, max_pages)
10011004
page = await fetch_page(cursor)
1002-
set_items(first, get_items(first) + get_items(page))
1005+
accumulator.extend(get_items(page))
10031006
cursor = page.next_cursor
10041007
pages += 1
10051008

1009+
set_items(first, accumulator)
10061010
# Strip the terminal cursor from the aggregated result.
10071011
first.next_cursor = None
10081012
return first
@@ -1014,12 +1018,14 @@ async def list_all_tools(self) -> ListToolsResult:
10141018
PaginationExceededError: The walk exceeded ``list_max_pages``.
10151019
CursorCycleError: The server returned a repeated cursor.
10161020
"""
1017-
return await self._drain_all_pages(
1021+
result = await self._drain_all_pages(
10181022
"tools/list",
10191023
fetch_page=lambda c: self.list_tools(cursor=c, cache_mode="bypass"),
10201024
get_items=lambda r: list(r.tools),
10211025
set_items=lambda r, items: setattr(r, "tools", items), # noqa: B010
10221026
)
1027+
# Absorb the final aggregated result as complete to prune stale session state
1028+
return self.session._absorb_tool_listing(result, complete=True) # pyright: ignore[reportPrivateUsage]
10231029

10241030
async def list_all_resources(self) -> ListResourcesResult:
10251031
"""Fetch all resources from the server, draining pagination automatically.

0 commit comments

Comments
 (0)