Skip to content

Follow MCP tool-list pagination in discovery and remote tool lookups - #444

Open
anxkhn wants to merge 1 commit into
google:mainfrom
anxkhn:fix/mcp-tool-pagination
Open

anxkhn wants to merge 1 commit into
google:mainfrom
anxkhn:fix/mcp-tool-pagination

Conversation

@anxkhn

@anxkhn anxkhn commented Sep 19, 2026

Copy link
Copy Markdown
Contributor

What's wrong

Three places assume a single tools/list response contains the complete tool catalogue of an MCP backend:

  • MCPService.Tools in internal/node/mcp_service.go calls session.ListTools(ctx, nil) once and caches the names it finds.
  • SamNode.fetchToolsForRemoteService in internal/node/mcp_handlers.go builds the remote catalogue rows from that one page.
  • SamNode.fetchRemoteToolDescription in the same file searches only that page before returning tool not found on peer.

None of them consume NextCursor. A backend that paginates its tool list therefore loses every tool after the first page: the names never reach the announcement keys built in discovery_source.go, find_remote_tools under-reports, and describe_remote_tool claims a tool does not exist even when the caller passes the correct name. With a backend page size of one, two tools are enough to trigger it, independently of the intentional announcement key cap.

The existing tests could not catch this because their fake server used the default page size and returned everything in one response.

The fix

Add a small listAllTools helper in internal/node/mcp_service.go that drains the SDK's ClientSession.Tools iterator, which follows NextCursor for us, and propagates the first error it sees. All three call sites now use it.

Because the helper returns either a complete slice or an error, the partial-result behaviour stays safe: MCPService.Tools returns before writing its name cache on failure, and fetchToolsForRemoteService still emits its existing error row for the filtered service. Sorting, name prefixes, service filtering and session cleanup are unchanged, as is the announcement key cap. The now-unreachable listRes == nil guards are dropped. No new dependencies.

Testing

The fake MCP server helper in internal/node/mcp_handlers_test.go gained a newFakeMCPHandlerWithOptions variant so a test can pass &mcp.ServerOptions{PageSize: 1}; newFakeMCPHandler keeps its old behaviour. The discovery, remote-catalogue and describe-tool tests now run against a one-tool-per-page backend, and the describe test gained an alpha tool ahead of review_pr so the tool under test lives on a later page. Each of these fails on the current code and passes with the fix.

go test ./internal/node -run 'TestMCPService_ToolsPagination|TestRemoteToolCataloguePagination|TestRemoteToolDescriptionPagination' -count=1 -timeout=10s

Result: ok github.com/google/sam/internal/node 1.606s.

Signed-off-by: Anas Khan <anxkhn28@gmail.com>

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces pagination support for fetching MCP tools by replacing direct calls to session.ListTools with a new helper function listAllTools that iterates over all pages. Corresponding tests have been updated to verify pagination behavior. The review feedback suggests adding a nil check for the session parameter in the new listAllTools helper function to prevent potential nil pointer dereferences.

Comment on lines +197 to +206
func listAllTools(ctx context.Context, session *mcp.ClientSession) ([]*mcp.Tool, error) {
var tools []*mcp.Tool
for tool, err := range session.Tools(ctx, nil) {
if err != nil {
return nil, err
}
tools = append(tools, tool)
}
return tools, nil
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

To ensure defensive programming and prevent potential nil pointer dereference panics, we should add a guard check to verify that the session parameter is not nil before calling session.Tools(ctx, nil).

Suggested change
func listAllTools(ctx context.Context, session *mcp.ClientSession) ([]*mcp.Tool, error) {
var tools []*mcp.Tool
for tool, err := range session.Tools(ctx, nil) {
if err != nil {
return nil, err
}
tools = append(tools, tool)
}
return tools, nil
}
func listAllTools(ctx context.Context, session *mcp.ClientSession) ([]*mcp.Tool, error) {
if session == nil {
return nil, fmt.Errorf("mcp session is nil")
}
var tools []*mcp.Tool
for tool, err := range session.Tools(ctx, nil) {
if err != nil {
return nil, err
}
tools = append(tools, tool)
}
return tools, nil
}

@aojea

aojea commented Sep 20, 2026

Copy link
Copy Markdown
Collaborator

Thanks — the diagnosis is right and the tests prove it (I confirmed all three fail on main and pass here). One blocking change and two smaller ones before this can merge.

1. Bound the drain (blocking)

listAllTools turns one request into as many as the server wants. The SDK's paginate loop (mcp/client.go) has no page or item cap and does not check ctx itself; it only observes the deadline indirectly when a request goes over the wire. Two things follow:

  • A server that always returns a non-empty nextCursor makes one find_remote_tools call spin until the 5s peer deadline (I measured ~10k tools/list round-trips in 2s over loopback), then listAllTools returns nil, context deadline exceeded — the whole catalogue is lost, not just the extra pages.
  • If that server also negotiates protocol 2026-07-28 and sets ttlMs on the result, the SDK's per-cursor cache (mcp/cache.go) answers the repeated cursor from memory. After two wire requests the loop never touches the network again, so ctx is never checked: listAllTools spins forever and tools grows until OOM. I reproduced this: the call had not returned 4s after its context expired.

fetchToolsForRemoteService and fetchRemoteToolDescription read this from a remote peer's backend, which under AGENTS.md's Zero Trust rule is untrusted input. MCPService.Tools runs it against the local backend on every discovery tick.

What to do, in node:

  • Add a constant, e.g. maxToolsPerService = 256, and make the cap item-count based (a time bound does not help in the cached case).
  • In listAllTools, stop iterating once len(tools) == maxToolsPerService, log at debug that the list was truncated, and return what was collected. Truncating matches what capKeys in discovery_source.go already does for announcements; failing would make a legitimately large server disappear.
  • In fetchRemoteToolDescription, don't drain at all: range over session.Tools(ctx, nil) directly, return on the first tool.Name == actualToolName, and give up with tool not found on peer after maxToolsPerService items.
  • Add a test that proves the bound. A minimal hostile handler is enough — no SDK server needed, just answer initialize and return a fixed page with "nextCursor": "again" on every tools/list:
func TestListAllTools_BoundedAgainstEndlessCursor(t *testing.T) {
	var pages atomic.Int64
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, _ := io.ReadAll(r.Body)
		var req struct{ ID any `json:"id"`; Method string `json:"method"` }
		_ = json.Unmarshal(body, &req)
		if req.ID == nil { w.WriteHeader(http.StatusAccepted); return } // notification
		var result any = map[string]any{}
		switch req.Method {
		case "initialize":
			result = map[string]any{
				"protocolVersion": "2026-07-28",
				"capabilities":    map[string]any{"tools": map[string]any{}},
				"serverInfo":      map[string]any{"name": "hostile", "version": "0"},
			}
		case "tools/list":
			pages.Add(1)
			result = map[string]any{
				"tools":      []any{map[string]any{"name": "x", "inputSchema": map[string]any{"type": "object"}}},
				"nextCursor": "again",
				"ttlMs":      60000, "cacheScope": "public", // makes the SDK serve repeats from cache
			}
		}
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(map[string]any{"jsonrpc": "2.0", "id": req.ID, "result": result})
	}))
	defer srv.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	session, err := mcp.NewClient(&mcp.Implementation{Name: "t", Version: "0"}, nil).
		Connect(ctx, &mcp.StreamableClientTransport{Endpoint: srv.URL}, nil)
	if err != nil { t.Fatal(err) }
	defer session.Close()

	tools, err := listAllTools(ctx, session)
	if err != nil { t.Fatal(err) }
	if len(tools) != maxToolsPerService {
		t.Fatalf("got %d tools, want cap %d", len(tools), maxToolsPerService)
	}
}

Without the cap this test hangs past its deadline; with it, it returns immediately. Keep it well under the 10s integration budget.

2. Keep the original tests, add the paginated ones

TestMCPService_Tools, TestHandleFindRemoteTools_SinglePeer and TestHandleDescribeRemoteTool_RoundTrip were renamed rather than extended, so the default single-page path lost its coverage and the file's TestHandleX_Case naming is broken. Restore the originals and either add a PageSize: 1 sibling for each or turn each into a two-case subtest table (default page size, one tool per page).

3. Nits

  • newFakeMCPHandler lost its t.Helper() when it became a wrapper; put it back.
  • The bot's session == nil guard isn't needed — session cannot be nil after a successful Connect/ConnectMCPSession. Ignore that suggestion.

No new dependencies, everything stays in node.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants