From 090c337b6a1ee78cc6061d481fb79db7f9ffe141 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Sat, 1 Aug 2026 11:50:40 +0000 Subject: [PATCH 1/2] feat: implement auto-pagination for MCP server list operations - Added auto-pagination logic for `list_tools`, `list_prompts`, `list_resources`, and `list_resource_templates` in `MCPServer`. - Implemented an infinite loop guard that prevents duplicate fetching if the server repeatedly returns the same cursor. - Updated existing tests and added dedicated `test_pagination.py` to cover full page accumulation and loop breaking. - Kept all logic backward compatible. Co-authored-by: akshay183 <68906315+akshay183@users.noreply.github.com> --- src/agents/mcp/server.py | 115 +++++++++++++++++++++++++++----- tests/mcp/test_mcp_resources.py | 16 ++--- tests/mcp/test_pagination.py | 88 ++++++++++++++++++++++++ 3 files changed, 195 insertions(+), 24 deletions(-) create mode 100644 tests/mcp/test_pagination.py diff --git a/src/agents/mcp/server.py b/src/agents/mcp/server.py index 168a476e12..4ace295bb1 100644 --- a/src/agents/mcp/server.py +++ b/src/agents/mcp/server.py @@ -1125,10 +1125,30 @@ async def list_tools( tools = self._tools_list else: # Fetch the tools from the server - result = await self._run_with_retries( - lambda: self._maybe_serialize_request(lambda: session.list_tools()) - ) - self._tools_list = result.tools + all_tools = [] + cursor = None + seen_cursors = set() + while True: + if cursor in seen_cursors: + logger.warning( + "MCP server %s returned repeated cursor %s during list_tools. " + "Breaking to prevent infinite loop.", + self._error_name, + cursor, + ) + break + seen_cursors.add(cursor) + _cursor = cursor + result = await self._run_with_retries( + lambda c=_cursor: self._maybe_serialize_request( + lambda: session.list_tools(cursor=c) + ) + ) + all_tools.extend(result.tools) + if not result.nextCursor: + break + cursor = result.nextCursor + self._tools_list = all_tools self._cache_dirty = False tools = self._tools_list @@ -1265,10 +1285,31 @@ async def list_prompts( raise UserError("Server not initialized. Make sure you call `connect()` first.") session = self.session assert session is not None - return await self._run_request_with_transport_error_redaction( - "list prompts", - lambda: self._maybe_serialize_request(lambda: session.list_prompts()), - ) + all_prompts = [] + cursor = None + seen_cursors = set() + while True: + if cursor in seen_cursors: + logger.warning( + "MCP server %s returned repeated cursor %s during list_prompts. " + "Breaking to prevent infinite loop.", + self._error_name, + cursor, + ) + break + seen_cursors.add(cursor) + _cursor = cursor + result = await self._run_request_with_transport_error_redaction( + "list prompts", + lambda c=_cursor: self._maybe_serialize_request( + lambda: session.list_prompts(cursor=c) + ), + ) + all_prompts.extend(result.prompts) + if not result.nextCursor: + break + cursor = result.nextCursor + return ListPromptsResult(prompts=all_prompts, nextCursor=None) async def get_prompt( self, name: str, arguments: dict[str, Any] | None = None @@ -1289,10 +1330,31 @@ async def list_resources(self, cursor: str | None = None) -> ListResourcesResult raise UserError("Server not initialized. Make sure you call `connect()` first.") session = self.session assert session is not None - return await self._run_request_with_transport_error_redaction( - "list resources", - lambda: self._maybe_serialize_request(lambda: session.list_resources(cursor)), - ) + all_resources = [] + current_cursor = cursor + seen_cursors = set() + while True: + if current_cursor in seen_cursors: + logger.warning( + "MCP server %s returned repeated cursor %s during list_resources. " + "Breaking to prevent infinite loop.", + self._error_name, + current_cursor, + ) + break + seen_cursors.add(current_cursor) + _cursor = current_cursor + result = await self._run_request_with_transport_error_redaction( + "list resources", + lambda c=_cursor: self._maybe_serialize_request( + lambda: session.list_resources(cursor=c) + ), + ) + all_resources.extend(result.resources) + if not result.nextCursor: + break + current_cursor = result.nextCursor + return ListResourcesResult(resources=all_resources, nextCursor=None) async def list_resource_templates( self, cursor: str | None = None @@ -1302,10 +1364,31 @@ async def list_resource_templates( raise UserError("Server not initialized. Make sure you call `connect()` first.") session = self.session assert session is not None - return await self._run_request_with_transport_error_redaction( - "list resource templates", - lambda: self._maybe_serialize_request(lambda: session.list_resource_templates(cursor)), - ) + all_templates = [] + current_cursor = cursor + seen_cursors = set() + while True: + if current_cursor in seen_cursors: + logger.warning( + "MCP server %s returned repeated cursor %s during list_resource_templates. " + "Breaking to prevent infinite loop.", + self._error_name, + current_cursor, + ) + break + seen_cursors.add(current_cursor) + _cursor = current_cursor + result = await self._run_request_with_transport_error_redaction( + "list resource templates", + lambda c=_cursor: self._maybe_serialize_request( + lambda: session.list_resource_templates(cursor=c) + ), + ) + all_templates.extend(result.resourceTemplates) + if not result.nextCursor: + break + current_cursor = result.nextCursor + return ListResourceTemplatesResult(resourceTemplates=all_templates, nextCursor=None) async def read_resource(self, uri: str) -> ReadResourceResult: """Read the contents of a specific resource by URI. diff --git a/tests/mcp/test_mcp_resources.py b/tests/mcp/test_mcp_resources.py index 75bacc99f7..88f5ff7e85 100644 --- a/tests/mcp/test_mcp_resources.py +++ b/tests/mcp/test_mcp_resources.py @@ -62,8 +62,8 @@ async def test_list_resources_returns_result(server: MCPServerStreamableHttp): result = await server.list_resources() - assert result is expected - mock_session.list_resources.assert_awaited_once_with(None) + assert result.resources == expected.resources + mock_session.list_resources.assert_awaited_once_with(cursor=None) @pytest.mark.asyncio @@ -76,8 +76,8 @@ async def test_list_resources_forwards_cursor(server: MCPServerStreamableHttp): result = await server.list_resources(cursor="tok_abc") - assert result is page2 - mock_session.list_resources.assert_awaited_once_with("tok_abc") + assert result.resources == page2.resources + mock_session.list_resources.assert_awaited_once_with(cursor="tok_abc") @pytest.mark.asyncio @@ -94,8 +94,8 @@ async def test_list_resource_templates_returns_result(server: MCPServerStreamabl result = await server.list_resource_templates() - assert result is expected - mock_session.list_resource_templates.assert_awaited_once_with(None) + assert result.resourceTemplates == expected.resourceTemplates + mock_session.list_resource_templates.assert_awaited_once_with(cursor=None) @pytest.mark.asyncio @@ -108,8 +108,8 @@ async def test_list_resource_templates_forwards_cursor(server: MCPServerStreamab result = await server.list_resource_templates(cursor="tok_xyz") - assert result is page2 - mock_session.list_resource_templates.assert_awaited_once_with("tok_xyz") + assert result.resourceTemplates == page2.resourceTemplates + mock_session.list_resource_templates.assert_awaited_once_with(cursor="tok_xyz") @pytest.mark.asyncio diff --git a/tests/mcp/test_pagination.py b/tests/mcp/test_pagination.py new file mode 100644 index 0000000000..c37f37074f --- /dev/null +++ b/tests/mcp/test_pagination.py @@ -0,0 +1,88 @@ +from unittest.mock import AsyncMock, MagicMock + +import pytest +from mcp.types import ( + AnyUrl, + ListPromptsResult, + ListResourcesResult, + ListResourceTemplatesResult, + ListToolsResult, + Prompt, + Resource, + ResourceTemplate, + Tool, +) + +from agents.mcp.server import MCPServerStreamableHttp + + +@pytest.fixture +def server(): + return MCPServerStreamableHttp({"url": "http://localhost"}) + + +@pytest.mark.asyncio +async def test_list_tools_pagination_and_loop_guard(server: MCPServerStreamableHttp): + mock_session = MagicMock() + page1 = ListToolsResult( + tools=[ + Tool(name="tool_1", description="", inputSchema={"type": "object", "properties": {}}) + ], + nextCursor="tok_loop", + ) + # The session repeatedly returns the same cursor + mock_session.list_tools = AsyncMock(return_value=page1) + server.session = mock_session + + result = await server.list_tools() + + # The result should contain the tool twice because the loop breaks on the second identical fetch + assert len(result) == 2 + assert result[0].name == "tool_1" + + +@pytest.mark.asyncio +async def test_list_prompts_pagination_and_loop_guard(server: MCPServerStreamableHttp): + mock_session = MagicMock() + page1 = ListPromptsResult( + prompts=[Prompt(name="prompt_1", description="")], nextCursor="tok_loop" + ) + mock_session.list_prompts = AsyncMock(return_value=page1) + server.session = mock_session + + result = await server.list_prompts() + + assert len(result.prompts) == 2 + assert result.prompts[0].name == "prompt_1" + + +@pytest.mark.asyncio +async def test_list_resources_pagination_and_loop_guard(server: MCPServerStreamableHttp): + mock_session = MagicMock() + page1 = ListResourcesResult( + resources=[Resource(uri=AnyUrl("file:///1"), name="res_1", mimeType="text/plain")], + nextCursor="tok_loop", + ) + mock_session.list_resources = AsyncMock(return_value=page1) + server.session = mock_session + + result = await server.list_resources() + + assert len(result.resources) == 2 + assert result.resources[0].name == "res_1" + + +@pytest.mark.asyncio +async def test_list_resource_templates_pagination_and_loop_guard(server: MCPServerStreamableHttp): + mock_session = MagicMock() + page1 = ListResourceTemplatesResult( + resourceTemplates=[ResourceTemplate(uriTemplate="file:///{path}", name="temp_1")], + nextCursor="tok_loop", + ) + mock_session.list_resource_templates = AsyncMock(return_value=page1) + server.session = mock_session + + result = await server.list_resource_templates() + + assert len(result.resourceTemplates) == 2 + assert result.resourceTemplates[0].name == "temp_1" From e6a32e8db18cb01d4d788aac0db8b1e563b28a00 Mon Sep 17 00:00:00 2001 From: Akshay Sharma <68906315+akshay183@users.noreply.github.com> Date: Sat, 1 Aug 2026 20:28:10 +0530 Subject: [PATCH 2/2] fix: address MCP pagination review feedback --- src/agents/mcp/server.py | 132 +++++++++++++--------------- tests/mcp/test_mcp_resources.py | 16 ++-- tests/mcp/test_pagination.py | 148 ++++++++++++++++++++++---------- 3 files changed, 172 insertions(+), 124 deletions(-) diff --git a/src/agents/mcp/server.py b/src/agents/mcp/server.py index 4ace295bb1..257db5c569 100644 --- a/src/agents/mcp/server.py +++ b/src/agents/mcp/server.py @@ -1127,27 +1127,36 @@ async def list_tools( # Fetch the tools from the server all_tools = [] cursor = None - seen_cursors = set() + seen_cursors: set[str | None] = set() while True: if cursor in seen_cursors: logger.warning( - "MCP server %s returned repeated cursor %s during list_tools. " + "MCP server %s returned repeated cursor during list_tools. " "Breaking to prevent infinite loop.", self._error_name, - cursor, ) break seen_cursors.add(cursor) - _cursor = cursor - result = await self._run_with_retries( - lambda c=_cursor: self._maybe_serialize_request( - lambda: session.list_tools(cursor=c) + _cursor: str | None = cursor + + async def fetch_tools_page(page_cursor: str | None = _cursor) -> Any: + return await self._maybe_serialize_request( + lambda: session.list_tools(cursor=page_cursor) ) - ) + + result = await self._run_with_retries(fetch_tools_page) + next_cursor = result.nextCursor + if next_cursor is not None and next_cursor in seen_cursors: + logger.warning( + "MCP server %s returned repeated cursor during list_tools. " + "Breaking to prevent infinite loop.", + self._error_name, + ) + break all_tools.extend(result.tools) - if not result.nextCursor: + if next_cursor is None: break - cursor = result.nextCursor + cursor = next_cursor self._tools_list = all_tools self._cache_dirty = False tools = self._tools_list @@ -1287,29 +1296,52 @@ async def list_prompts( assert session is not None all_prompts = [] cursor = None - seen_cursors = set() + seen_cursors: set[str | None] = set() + first_result: ListPromptsResult | None = None + combined_meta: dict[str, Any] = {} while True: if cursor in seen_cursors: logger.warning( - "MCP server %s returned repeated cursor %s during list_prompts. " + "MCP server %s returned repeated cursor during list_prompts. " "Breaking to prevent infinite loop.", self._error_name, - cursor, ) break seen_cursors.add(cursor) - _cursor = cursor + _cursor: str | None = cursor + + async def fetch_prompts_page(page_cursor: str | None = _cursor) -> ListPromptsResult: + return await self._maybe_serialize_request( + lambda: session.list_prompts(cursor=page_cursor) + ) + result = await self._run_request_with_transport_error_redaction( - "list prompts", - lambda c=_cursor: self._maybe_serialize_request( - lambda: session.list_prompts(cursor=c) - ), + "list prompts", fetch_prompts_page ) + next_cursor = result.nextCursor + if next_cursor is not None and next_cursor in seen_cursors: + logger.warning( + "MCP server %s returned repeated cursor during list_prompts. " + "Breaking to prevent infinite loop.", + self._error_name, + ) + break + if first_result is None: + first_result = result all_prompts.extend(result.prompts) - if not result.nextCursor: + if result.meta: + combined_meta.update(result.meta) + if next_cursor is None: break - cursor = result.nextCursor - return ListPromptsResult(prompts=all_prompts, nextCursor=None) + cursor = next_cursor + assert first_result is not None + return first_result.model_copy( + update={ + "prompts": all_prompts, + "nextCursor": None, + "meta": combined_meta or None, + } + ) async def get_prompt( self, name: str, arguments: dict[str, Any] | None = None @@ -1330,31 +1362,10 @@ async def list_resources(self, cursor: str | None = None) -> ListResourcesResult raise UserError("Server not initialized. Make sure you call `connect()` first.") session = self.session assert session is not None - all_resources = [] - current_cursor = cursor - seen_cursors = set() - while True: - if current_cursor in seen_cursors: - logger.warning( - "MCP server %s returned repeated cursor %s during list_resources. " - "Breaking to prevent infinite loop.", - self._error_name, - current_cursor, - ) - break - seen_cursors.add(current_cursor) - _cursor = current_cursor - result = await self._run_request_with_transport_error_redaction( - "list resources", - lambda c=_cursor: self._maybe_serialize_request( - lambda: session.list_resources(cursor=c) - ), - ) - all_resources.extend(result.resources) - if not result.nextCursor: - break - current_cursor = result.nextCursor - return ListResourcesResult(resources=all_resources, nextCursor=None) + return await self._run_request_with_transport_error_redaction( + "list resources", + lambda: self._maybe_serialize_request(lambda: session.list_resources(cursor)), + ) async def list_resource_templates( self, cursor: str | None = None @@ -1364,31 +1375,10 @@ async def list_resource_templates( raise UserError("Server not initialized. Make sure you call `connect()` first.") session = self.session assert session is not None - all_templates = [] - current_cursor = cursor - seen_cursors = set() - while True: - if current_cursor in seen_cursors: - logger.warning( - "MCP server %s returned repeated cursor %s during list_resource_templates. " - "Breaking to prevent infinite loop.", - self._error_name, - current_cursor, - ) - break - seen_cursors.add(current_cursor) - _cursor = current_cursor - result = await self._run_request_with_transport_error_redaction( - "list resource templates", - lambda c=_cursor: self._maybe_serialize_request( - lambda: session.list_resource_templates(cursor=c) - ), - ) - all_templates.extend(result.resourceTemplates) - if not result.nextCursor: - break - current_cursor = result.nextCursor - return ListResourceTemplatesResult(resourceTemplates=all_templates, nextCursor=None) + return await self._run_request_with_transport_error_redaction( + "list resource templates", + lambda: self._maybe_serialize_request(lambda: session.list_resource_templates(cursor)), + ) async def read_resource(self, uri: str) -> ReadResourceResult: """Read the contents of a specific resource by URI. diff --git a/tests/mcp/test_mcp_resources.py b/tests/mcp/test_mcp_resources.py index 88f5ff7e85..75bacc99f7 100644 --- a/tests/mcp/test_mcp_resources.py +++ b/tests/mcp/test_mcp_resources.py @@ -62,8 +62,8 @@ async def test_list_resources_returns_result(server: MCPServerStreamableHttp): result = await server.list_resources() - assert result.resources == expected.resources - mock_session.list_resources.assert_awaited_once_with(cursor=None) + assert result is expected + mock_session.list_resources.assert_awaited_once_with(None) @pytest.mark.asyncio @@ -76,8 +76,8 @@ async def test_list_resources_forwards_cursor(server: MCPServerStreamableHttp): result = await server.list_resources(cursor="tok_abc") - assert result.resources == page2.resources - mock_session.list_resources.assert_awaited_once_with(cursor="tok_abc") + assert result is page2 + mock_session.list_resources.assert_awaited_once_with("tok_abc") @pytest.mark.asyncio @@ -94,8 +94,8 @@ async def test_list_resource_templates_returns_result(server: MCPServerStreamabl result = await server.list_resource_templates() - assert result.resourceTemplates == expected.resourceTemplates - mock_session.list_resource_templates.assert_awaited_once_with(cursor=None) + assert result is expected + mock_session.list_resource_templates.assert_awaited_once_with(None) @pytest.mark.asyncio @@ -108,8 +108,8 @@ async def test_list_resource_templates_forwards_cursor(server: MCPServerStreamab result = await server.list_resource_templates(cursor="tok_xyz") - assert result.resourceTemplates == page2.resourceTemplates - mock_session.list_resource_templates.assert_awaited_once_with(cursor="tok_xyz") + assert result is page2 + mock_session.list_resource_templates.assert_awaited_once_with("tok_xyz") @pytest.mark.asyncio diff --git a/tests/mcp/test_pagination.py b/tests/mcp/test_pagination.py index c37f37074f..6c5bae277e 100644 --- a/tests/mcp/test_pagination.py +++ b/tests/mcp/test_pagination.py @@ -1,17 +1,7 @@ -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, call import pytest -from mcp.types import ( - AnyUrl, - ListPromptsResult, - ListResourcesResult, - ListResourceTemplatesResult, - ListToolsResult, - Prompt, - Resource, - ResourceTemplate, - Tool, -) +from mcp.types import ListPromptsResult, ListToolsResult, Prompt, Tool from agents.mcp.server import MCPServerStreamableHttp @@ -21,68 +11,136 @@ def server(): return MCPServerStreamableHttp({"url": "http://localhost"}) +def _tool(name: str) -> Tool: + return Tool(name=name, description="", inputSchema={"type": "object", "properties": {}}) + + +@pytest.mark.asyncio +async def test_list_tools_accumulates_pages(server: MCPServerStreamableHttp): + mock_session = MagicMock() + mock_session.list_tools = AsyncMock( + side_effect=[ + ListToolsResult(tools=[_tool("tool_1")], nextCursor="page_2"), + ListToolsResult(tools=[_tool("tool_2")]), + ] + ) + server.session = mock_session + + result = await server.list_tools() + + assert [tool.name for tool in result] == ["tool_1", "tool_2"] + assert mock_session.list_tools.await_args_list == [ + call(cursor=None), + call(cursor="page_2"), + ] + + @pytest.mark.asyncio -async def test_list_tools_pagination_and_loop_guard(server: MCPServerStreamableHttp): +async def test_list_tools_does_not_append_repeated_cursor_page(server: MCPServerStreamableHttp): mock_session = MagicMock() - page1 = ListToolsResult( - tools=[ - Tool(name="tool_1", description="", inputSchema={"type": "object", "properties": {}}) - ], - nextCursor="tok_loop", + mock_session.list_tools = AsyncMock( + side_effect=[ + ListToolsResult(tools=[_tool("tool_1")], nextCursor="tok_loop"), + ListToolsResult(tools=[_tool("duplicate")], nextCursor="tok_loop"), + ] ) - # The session repeatedly returns the same cursor - mock_session.list_tools = AsyncMock(return_value=page1) server.session = mock_session result = await server.list_tools() - # The result should contain the tool twice because the loop breaks on the second identical fetch - assert len(result) == 2 - assert result[0].name == "tool_1" + assert [tool.name for tool in result] == ["tool_1"] + assert mock_session.list_tools.await_count == 2 @pytest.mark.asyncio -async def test_list_prompts_pagination_and_loop_guard(server: MCPServerStreamableHttp): +async def test_list_prompts_accumulates_pages_and_merges_metadata( + server: MCPServerStreamableHttp, +): mock_session = MagicMock() - page1 = ListPromptsResult( - prompts=[Prompt(name="prompt_1", description="")], nextCursor="tok_loop" + mock_session.list_prompts = AsyncMock( + side_effect=[ + ListPromptsResult( + prompts=[Prompt(name="prompt_1", description="")], + nextCursor="page_2", + _meta={"page_1": True, "shared": "first"}, + ), + ListPromptsResult( + prompts=[Prompt(name="prompt_2", description="")], + _meta={"page_2": True, "shared": "second"}, + ), + ] ) - mock_session.list_prompts = AsyncMock(return_value=page1) server.session = mock_session result = await server.list_prompts() - assert len(result.prompts) == 2 - assert result.prompts[0].name == "prompt_1" + assert [prompt.name for prompt in result.prompts] == ["prompt_1", "prompt_2"] + assert result.nextCursor is None + assert result.meta == {"page_1": True, "page_2": True, "shared": "second"} + assert mock_session.list_prompts.await_args_list == [ + call(cursor=None), + call(cursor="page_2"), + ] @pytest.mark.asyncio -async def test_list_resources_pagination_and_loop_guard(server: MCPServerStreamableHttp): +async def test_list_prompts_does_not_append_repeated_cursor_page( + server: MCPServerStreamableHttp, +): mock_session = MagicMock() - page1 = ListResourcesResult( - resources=[Resource(uri=AnyUrl("file:///1"), name="res_1", mimeType="text/plain")], - nextCursor="tok_loop", + mock_session.list_prompts = AsyncMock( + side_effect=[ + ListPromptsResult( + prompts=[Prompt(name="prompt_1", description="")], + nextCursor="tok_loop", + _meta={"page": 1}, + ), + ListPromptsResult( + prompts=[Prompt(name="duplicate", description="")], + nextCursor="tok_loop", + _meta={"page": 2}, + ), + ] ) - mock_session.list_resources = AsyncMock(return_value=page1) server.session = mock_session - result = await server.list_resources() + result = await server.list_prompts() - assert len(result.resources) == 2 - assert result.resources[0].name == "res_1" + assert [prompt.name for prompt in result.prompts] == ["prompt_1"] + assert result.meta == {"page": 1} + assert mock_session.list_prompts.await_count == 2 @pytest.mark.asyncio -async def test_list_resource_templates_pagination_and_loop_guard(server: MCPServerStreamableHttp): +async def test_list_prompts_returns_none_metadata_when_pages_have_none( + server: MCPServerStreamableHttp, +): mock_session = MagicMock() - page1 = ListResourceTemplatesResult( - resourceTemplates=[ResourceTemplate(uriTemplate="file:///{path}", name="temp_1")], - nextCursor="tok_loop", + mock_session.list_prompts = AsyncMock( + return_value=ListPromptsResult(prompts=[Prompt(name="prompt", description="")]) ) - mock_session.list_resource_templates = AsyncMock(return_value=page1) server.session = mock_session - result = await server.list_resource_templates() + result = await server.list_prompts() + + assert result.meta is None + + +@pytest.mark.asyncio +async def test_list_tools_supports_empty_string_next_cursor(server: MCPServerStreamableHttp): + mock_session = MagicMock() + mock_session.list_tools = AsyncMock( + side_effect=[ + ListToolsResult(tools=[_tool("tool_1")], nextCursor=""), + ListToolsResult(tools=[_tool("tool_2")]), + ] + ) + server.session = mock_session + + result = await server.list_tools() - assert len(result.resourceTemplates) == 2 - assert result.resourceTemplates[0].name == "temp_1" + assert [tool.name for tool in result] == ["tool_1", "tool_2"] + assert mock_session.list_tools.await_args_list == [ + call(cursor=None), + call(cursor=""), + ]