diff --git a/raven/agent/tools/mcp.py b/raven/agent/tools/mcp.py index 3c50d7c..6f72a80 100644 --- a/raven/agent/tools/mcp.py +++ b/raven/agent/tools/mcp.py @@ -13,6 +13,8 @@ from raven.sandbox import SandboxInitError if TYPE_CHECKING: + from mcp import ClientSession + from raven.sandbox import SandboxExecutor @@ -75,6 +77,26 @@ async def execute(self, **kwargs: Any) -> str | ToolResult: return text or "(no output)" +async def _collect_tools(session: "ClientSession") -> list: + """Page through tools/list until the server stops returning a nextCursor. + + A server may paginate tool discovery when it exposes more tools than its + page size. Track seen cursors so a misbehaving server that re-issues the + same cursor cannot trap discovery in an infinite loop. + """ + all_tools = [] + seen_cursors = {None} + cursor = None + while True: + page = await session.list_tools(cursor) + all_tools.extend(page.tools) + cursor = page.nextCursor + if not cursor or cursor in seen_cursors: + break + seen_cursors.add(cursor) + return all_tools + + async def connect_mcp_servers( mcp_servers: dict, registry: ToolRegistry, @@ -158,13 +180,14 @@ def httpx_client_factory( session = await stack.enter_async_context(ClientSession(read, write)) await session.initialize() - tools = await session.list_tools() - for tool_def in tools.tools: + all_tools = await _collect_tools(session) + + for tool_def in all_tools: wrapper = MCPToolWrapper(session, name, tool_def, tool_timeout=cfg.tool_timeout) registry.register(wrapper) logger.debug("MCP: registered tool '{}' from server '{}'", wrapper.name, name) - logger.info("MCP server '{}': connected, {} tools registered", name, len(tools.tools)) + logger.info("MCP server '{}': connected, {} tools registered", name, len(all_tools)) except (Exception, BaseExceptionGroup) as e: # BaseExceptionGroup is raised by anyio task groups (e.g. streamableHttp cancel # scope failures) and is not a subclass of Exception in Python 3.11+. diff --git a/tests/test_mcp_tools_pagination.py b/tests/test_mcp_tools_pagination.py new file mode 100644 index 0000000..88f3cd5 --- /dev/null +++ b/tests/test_mcp_tools_pagination.py @@ -0,0 +1,64 @@ +"""Pagination tests for MCP tools/list discovery (issue #301).""" + +from __future__ import annotations + +from types import SimpleNamespace + +from raven.agent.tools.mcp import _collect_tools + + +def _page(tool_names: list[str], next_cursor: str | None = None) -> SimpleNamespace: + return SimpleNamespace( + tools=[SimpleNamespace(name=n) for n in tool_names], + nextCursor=next_cursor, + ) + + +class _FakeSession: + """Replays a scripted list_tools response keyed by the cursor argument.""" + + def __init__(self, pages: dict): + self._pages = pages + self.calls: list = [] + + async def list_tools(self, cursor: str | None = None) -> SimpleNamespace: + self.calls.append(cursor) + return self._pages[cursor] + + +async def test_single_page_returns_tools_without_follow_up() -> None: + session = _FakeSession({None: _page(["a", "b"])}) + + tools = await _collect_tools(session) + + assert [t.name for t in tools] == ["a", "b"] + assert session.calls == [None] + + +async def test_multi_page_follows_cursor_chain() -> None: + session = _FakeSession( + { + None: _page(["a"], "c1"), + "c1": _page(["b"], "c2"), + "c2": _page(["c"]), + } + ) + + tools = await _collect_tools(session) + + assert [t.name for t in tools] == ["a", "b", "c"] + assert session.calls == [None, "c1", "c2"] + + +async def test_repeated_cursor_stops_instead_of_looping() -> None: + session = _FakeSession( + { + None: _page(["a"], "c1"), + "c1": _page(["b"], "c1"), + } + ) + + tools = await _collect_tools(session) + + assert [t.name for t in tools] == ["a", "b"] + assert session.calls == [None, "c1"]