diff --git a/README.md b/README.md index 439233de..bf21f707 100644 --- a/README.md +++ b/README.md @@ -2386,13 +2386,18 @@ If your application needs the complete collection regardless of how the server i `client.tools`, `client.resources`, `client.resource_templates`, and `client.prompts` auto-iterate through all pages and return a plain array of items, guaranteeing the full collection regardless of the server's `page_size` setting. When a server paginates, they issue multiple JSON-RPC round -trips per call and break out of the pagination loop if the server returns the same `nextCursor` -twice in a row as a safety measure. +trips per call. Two guards keep that loop finite: it stops when the server returns a `nextCursor` +it has already sent, and it stops after `max_pages` pages. ```ruby tools = client.tools # => Array of every tool on the server. ``` +`MCP::Client.new` accepts an optional `max_pages:` keyword that caps how many pages these methods +will walk. It defaults to `1_000`; a server that keeps offering a fresh `nextCursor` past that +point raises `MCP::Client::PaginationLimitError` rather than being followed indefinitely. Raise it +if you legitimately expect more pages than that. + Use these when you want the complete list; use `list_tools(cursor:)` etc. when you need fine-grained iteration (e.g. to stream-process pages without loading everything into memory). diff --git a/lib/mcp/client.rb b/lib/mcp/client.rb index e1de318e..7c9d8b93 100644 --- a/lib/mcp/client.rb +++ b/lib/mcp/client.rb @@ -12,6 +12,12 @@ module MCP class Client + # Upper bound on the number of pages the all-pages methods (`tools`, `resources`, + # `resource_templates`, `prompts`) will walk. The cursor guard in `fetch_all_pages` only + # stops a server that repeats or cycles cursors; one that returns a fresh `nextCursor` on + # every response would otherwise be followed indefinitely, growing the retained pages with it. + MAX_PAGES = 1_000 + class ServerError < StandardError attr_reader :code, :data @@ -52,6 +58,11 @@ def initialize(message, code:) # server-returned JSON-RPC error, which is raised as `ServerError`. class ValidationError < StandardError; end + # Raised when an all-pages method reaches `max_pages` while the server is still offering + # another cursor. Use the single-page `list_*` methods to walk such a collection with + # a policy of your own. + class PaginationLimitError < StandardError; end + # Raised when a server answers with a SEP-2322 Multi Round-Trip `input_required` result instead of # a final result. The result is not an error on the wire: it asks the client to fulfill the server's # `inputRequests` (a map of id => `{ "method" => ..., "params" => ... }` request objects with @@ -108,6 +119,8 @@ def initialize(message, request, original_error: nil) # @param transport [Object] The transport object to use for communication with the server. # The transport should be a duck type that responds to `send_request`. See the README for more details. # @param input_required_max_rounds [Integer] Cap on SEP-2322 driver rounds. + # @param max_pages [Integer] Maximum number of pages the all-pages methods ({#tools}, {#resources}, + # {#resource_templates}, {#prompts}) will walk before raising {MCP::Client::PaginationLimitError}. # # Once a handler is registered through `on_elicitation`, `on_sampling`, or `on_roots`, `call_tool`, # `get_prompt`, and `read_resource` resume `input_required` results automatically; without handlers @@ -117,8 +130,15 @@ def initialize(message, request, original_error: nil) # @example # transport = MCP::Client::HTTP.new(url: "http://localhost:3000") # client = MCP::Client.new(transport: transport) - def initialize(transport:, input_required_max_rounds: DEFAULT_INPUT_REQUIRED_MAX_ROUNDS) + def initialize(transport:, input_required_max_rounds: DEFAULT_INPUT_REQUIRED_MAX_ROUNDS, max_pages: MAX_PAGES) + # `nil` or a non-positive value would make the pagination unbounded and silently + # disable the protection, so reject it up front. + unless max_pages.is_a?(Integer) && max_pages > 0 + raise ArgumentError, "max_pages must be a positive Integer" + end + @transport = transport + @max_pages = max_pages # Populated by `on_elicitation`, `on_sampling`, and `on_roots`. The same handler answers both ways # the server can ask for input: a real server-to-client request, and an embedded request inside # a SEP-2322 `input_required` result. @@ -301,6 +321,7 @@ def list_tools(cursor: nil, meta: nil, cancellation: nil) # Cancelling it aborts whichever page is currently in flight; pages already returned are kept, # but the call raises `MCP::CancelledError` instead of returning the partial set. # @return [Array] An array of available tools. + # @raise [MCP::Client::PaginationLimitError] If the server offers more than `max_pages` pages. # # @example # tools = client.tools @@ -342,6 +363,7 @@ def list_resources(cursor: nil, meta: nil, cancellation: nil) # # @param cancellation [MCP::Cancellation, nil] Optional cancellation token (see {#tools}). # @return [Array] An array of available resources. + # @raise [MCP::Client::PaginationLimitError] See {#tools}. def resources(cancellation: nil) # TODO: consider renaming to `list_all_resources`. fetch_all_pages { |cursor| list_resources(cursor: cursor, cancellation: cancellation) }.flat_map(&:resources) @@ -377,6 +399,7 @@ def list_resource_templates(cursor: nil, meta: nil, cancellation: nil) # # @param cancellation [MCP::Cancellation, nil] Optional cancellation token (see {#tools}). # @return [Array] An array of available resource templates. + # @raise [MCP::Client::PaginationLimitError] See {#tools}. def resource_templates(cancellation: nil) # TODO: consider renaming to `list_all_resource_templates`. fetch_all_pages { |cursor| list_resource_templates(cursor: cursor, cancellation: cancellation) }.flat_map(&:resource_templates) @@ -412,6 +435,7 @@ def list_prompts(cursor: nil, meta: nil, cancellation: nil) # # @param cancellation [MCP::Cancellation, nil] Optional cancellation token (see {#tools}). # @return [Array] An array of available prompts. + # @raise [MCP::Client::PaginationLimitError] See {#tools}. def prompts(cancellation: nil) # TODO: consider renaming to `list_all_prompts`. fetch_all_pages { |cursor| list_prompts(cursor: cursor, cancellation: cancellation) }.flat_map(&:prompts) @@ -695,7 +719,8 @@ def transport_connect_accepts_mode? # Walks every page of a list endpoint, following `next_cursor`, and returns # the page results. The `seen` set guards against a server that repeats or - # cycles cursors, so the loop always terminates. + # cycles cursors, and `@max_pages` bounds one that returns a fresh cursor every + # time, so the loop always terminates. def fetch_all_pages pages = [] seen = Set.new @@ -707,6 +732,11 @@ def fetch_all_pages next_cursor = page.next_cursor break if next_cursor.nil? || seen.include?(next_cursor) + if pages.size >= @max_pages + raise PaginationLimitError, "Server returned more than #{@max_pages} pages; pass a larger `max_pages:` to " \ + "`MCP::Client.new` if this is expected." + end + seen << next_cursor cursor = next_cursor end diff --git a/test/mcp/client_test.rb b/test/mcp/client_test.rb index d944698f..bbd06520 100644 --- a/test/mcp/client_test.rb +++ b/test/mcp/client_test.rb @@ -1645,6 +1645,59 @@ def test_tools_breaks_when_server_returns_same_cursor_repeatedly assert_equal(2, tools.size) end + def test_tools_raises_when_the_server_offers_more_pages_than_max_pages + transport = mock + + pages = (1..3).map do |n| + { + "result" => { + "tools" => [{ "name" => "tool#{n}", "description" => "tool#{n}", "inputSchema" => {} }], + "nextCursor" => "cursor#{n}", + }, + } + end + + # A fresh cursor on every response defeats the repeat guard, so only `max_pages` stops the walk. + transport.expects(:send_request).times(3).returns(*pages) + + client = Client.new(transport: transport, max_pages: 3) + + error = assert_raises(Client::PaginationLimitError) { client.tools } + assert_equal( + "Server returned more than 3 pages; pass a larger `max_pages:` to `MCP::Client.new` if this is expected.", + error.message, + ) + end + + def test_tools_returns_every_page_when_the_server_stops_exactly_at_max_pages + transport = mock + + page1 = { + "result" => { + "tools" => [{ "name" => "tool1", "description" => "tool1", "inputSchema" => {} }], + "nextCursor" => "cursor1", + }, + } + page2 = { + "result" => { + "tools" => [{ "name" => "tool2", "description" => "tool2", "inputSchema" => {} }], + }, + } + + transport.expects(:send_request).twice.returns(page1, page2) + + client = Client.new(transport: transport, max_pages: 2) + tools = client.tools + + assert_equal(["tool1", "tool2"], tools.map(&:name)) + end + + def test_raises_argument_error_when_max_pages_is_not_positive + error = assert_raises(ArgumentError) { Client.new(transport: mock, max_pages: 0) } + + assert_equal("max_pages must be a positive Integer", error.message) + end + def test_tools_breaks_when_server_cycles_between_cursors transport = mock