Skip to content

mcp: harden paginate against infinite loops and param mutation #1110

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
dongjiang1989 wants to merge 1 commit into
base: main
Choose a base branch
from
+238 −9
Open

mcp: harden paginate against infinite loops and param mutation #1110

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 54 additions & 7 deletions mcp/client.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -204,8 +204,19 @@ type ClientOptions struct {
// reset" guidance, letting a transient miss pass without tearing down an
// otherwise live session. Has no effect unless KeepAlive is non-zero.
KeepAliveFailureThreshold int
// ListMaxPages is the maximum number of pages to fetch during automatic
// pagination of list operations (Tools, Resources, ResourceTemplates,
// Prompts). A value of 0 uses the default of [DefaultListMaxPages] (64).
// A negative value means unlimited. A positive value caps the number of
// pages. This prevents runaway pagination loops caused by server-side
// cursor cycles.
ListMaxPages int
}

// DefaultListMaxPages is the default value for [ClientOptions.ListMaxPages],
// matching the TypeScript SDK's DEFAULT_LIST_MAX_PAGES.
const DefaultListMaxPages = 64

// toolContextKeyType is the context key type for passing tool definitions
// from CallTool to the transport layer.
type toolContextKeyType struct{}
Expand Down Expand Up @@ -1550,7 +1561,7 @@ func (cs *ClientSession) Tools(ctx context.Context, params *ListToolsParams) ite
if params == nil {
params = &ListToolsParams{}
}
return paginate(ctx, params, cs.ListTools, func(res *ListToolsResult) []*Tool {
return paginate(ctx, params, cs.client.opts.ListMaxPages, cs.ListTools, func(res *ListToolsResult) []*Tool {
return res.Tools
})
}
Expand All @@ -1563,7 +1574,7 @@ func (cs *ClientSession) Resources(ctx context.Context, params *ListResourcesPar
if params == nil {
params = &ListResourcesParams{}
}
return paginate(ctx, params, cs.ListResources, func(res *ListResourcesResult) []*Resource {
return paginate(ctx, params, cs.client.opts.ListMaxPages, cs.ListResources, func(res *ListResourcesResult) []*Resource {
return res.Resources
})
}
Expand All @@ -1576,7 +1587,7 @@ func (cs *ClientSession) ResourceTemplates(ctx context.Context, params *ListReso
if params == nil {
params = &ListResourceTemplatesParams{}
}
return paginate(ctx, params, cs.ListResourceTemplates, func(res *ListResourceTemplatesResult) []*ResourceTemplate {
return paginate(ctx, params, cs.client.opts.ListMaxPages, cs.ListResourceTemplates, func(res *ListResourceTemplatesResult) []*ResourceTemplate {
return res.ResourceTemplates
})
}
Expand All @@ -1589,16 +1600,38 @@ func (cs *ClientSession) Prompts(ctx context.Context, params *ListPromptsParams)
if params == nil {
params = &ListPromptsParams{}
}
return paginate(ctx, params, cs.ListPrompts, func(res *ListPromptsResult) []*Prompt {
return paginate(ctx, params, cs.client.opts.ListMaxPages, cs.ListPrompts, func(res *ListPromptsResult) []*Prompt {
return res.Prompts
})
}

// paginate is a generic helper function to provide a paginated iterator.
func paginate[P listParams, R listResult[T], T any](ctx context.Context, params P, listFunc func(context.Context, P) (R, error), items func(R) []*T) iter.Seq2[*T, error] {
//
// It fetches pages by calling listFunc until the result has no NextCursor,
// maxPages is exceeded (if non-zero), or a cursor cycle is detected.
// The caller's params struct is not mutated; a local copy is used instead.
func paginate[P listParams, R listResult[T], T any](ctx context.Context, params P, maxPages int, listFunc func(context.Context, P) (R, error), items func(R) []*T) iter.Seq2[*T, error] {
return func(yield func(*T, error) bool) {
// Copy the underlying struct so we don't mutate the caller's params.
// P is always a pointer to a struct (e.g. *ListToolsParams).
// We use reflect to create a shallow copy of the pointed-to struct.
localParams := params
if v := reflect.ValueOf(params); v.Kind() == reflect.Pointer {
cp := reflect.New(v.Type().Elem())
cp.Elem().Set(v.Elem())
localParams = cp.Interface().(P)
}
var seen map[string]bool
pages := 0
// 0 means use default (64); negative means unlimited.
effectiveMax := maxPages
if effectiveMax == 0 {
effectiveMax = DefaultListMaxPages
}

for {
res, err := listFunc(ctx, params)
pages++
res, err := listFunc(ctx, localParams)
if err != nil {
yield(nil, err)
return
Expand All @@ -1612,7 +1645,21 @@ func paginate[P listParams, R listResult[T], T any](ctx context.Context, params
if nextCursorVal == nil || *nextCursorVal == "" {
return
}
*params.cursorPtr() = *nextCursorVal
// Check max pages limit.
if effectiveMax > 0 && pages >= effectiveMax {
yield(nil, fmt.Errorf("mcp: pagination exceeded maximum page limit of %d", effectiveMax))
return
}
// Detect cursor cycles to prevent infinite loops.
if seen == nil {
seen = make(map[string]bool)
}
if seen[*nextCursorVal] {
yield(nil, fmt.Errorf("mcp: pagination detected cursor cycle: %q", *nextCursorVal))
return
}
seen[*nextCursorVal] = true
*localParams.cursorPtr() = *nextCursorVal
}
}
}
Expand Down
186 changes: 184 additions & 2 deletions mcp/client_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -147,7 +147,7 @@ func TestClientPaginateBasic(t *testing.T) {

var gotItems []*Item
var iterationErr error
seq := paginate(ctx, params, listFunc, func(r *ListTestResult) []*Item { return r.Items })
seq := paginate(ctx, params, 0, listFunc, func(r *ListTestResult) []*Item { return r.Items })
for item, err := range seq {
if err != nil {
iterationErr = err
Expand Down Expand Up @@ -200,7 +200,7 @@ func TestClientPaginateVariousPageSizes(t *testing.T) {
return res, nil
}
var gotItems []*Item
seq := paginate(ctx, &ListTestParams{}, listFunc, func(r *ListTestResult) []*Item { return r.Items })
seq := paginate(ctx, &ListTestParams{}, 0, listFunc, func(r *ListTestResult) []*Item { return r.Items })
for item, err := range seq {
if err != nil {
t.Fatalf("paginate() unexpected error during iteration: %v", err)
Expand All @@ -214,6 +214,188 @@ func TestClientPaginateVariousPageSizes(t *testing.T) {
}
}

func TestPaginateMaxPages(t *testing.T) {
ctx := context.Background()
// Create 10 pages of results (1 item each, all with next cursor).
all := allItems[:10]
results := generatePaginatedResults(all, 1)

listFunc := func(ctx context.Context, params *ListTestParams) (*ListTestResult, error) {
if len(results) == 0 {
t.Fatal("listFunc called more times than expected")
}
res := results[0]
results = results[1:]
return res, nil
}

t.Run("DefaultLimit64", func(t *testing.T) {
// With maxPages=0 (default 64), 10 pages should succeed.
results = generatePaginatedResults(all, 1)
var got []*Item
var iterErr error
seq := paginate(ctx, &ListTestParams{}, 0, listFunc, func(r *ListTestResult) []*Item { return r.Items })
for item, err := range seq {
if err != nil {
iterErr = err
break
}
got = append(got, item)
}
if iterErr != nil {
t.Fatalf("unexpected error: %v", iterErr)
}
if len(got) != len(all) {
t.Fatalf("got %d items, want %d", len(got), len(all))
}
})

t.Run("MaxPagesExceeded", func(t *testing.T) {
results = generatePaginatedResults(all, 1)
var got []*Item
var iterErr error
// maxPages=3: should stop after 3 pages with an error.
seq := paginate(ctx, &ListTestParams{}, 3, listFunc, func(r *ListTestResult) []*Item { return r.Items })
for item, err := range seq {
if err != nil {
iterErr = err
break
}
got = append(got, item)
}
if iterErr == nil {
t.Fatal("expected pagination error, got nil")
}
if len(got) != 3 {
t.Fatalf("got %d items, want 3", len(got))
}
})

t.Run("UnlimitedWithNegative", func(t *testing.T) {
results = generatePaginatedResults(all, 1)
var got []*Item
var iterErr error
// maxPages=-1: unlimited, should get all items.
seq := paginate(ctx, &ListTestParams{}, -1, listFunc, func(r *ListTestResult) []*Item { return r.Items })
for item, err := range seq {
if err != nil {
iterErr = err
break
}
got = append(got, item)
}
if iterErr != nil {
t.Fatalf("unexpected error: %v", iterErr)
}
if len(got) != len(all) {
t.Fatalf("got %d items, want %d", len(got), len(all))
}
})

t.Run("DefaultCapAt64", func(t *testing.T) {
// maxPages=0 defaults to DefaultListMaxPages (64). Create 100 pages; should stop at 64.
bigAll := make([]*Item, 100)
for i := range bigAll {
bigAll[i] = &Item{Name: fmt.Sprintf("item-%d", i)}
}
results := generatePaginatedResults(bigAll, 1)
listFunc := func(ctx context.Context, params *ListTestParams) (*ListTestResult, error) {
if len(results) == 0 {
t.Fatal("listFunc called more times than expected")
}
res := results[0]
results = results[1:]
return res, nil
}
var got []*Item
var iterErr error
seq := paginate(ctx, &ListTestParams{}, 0, listFunc, func(r *ListTestResult) []*Item { return r.Items })
for item, err := range seq {
if err != nil {
iterErr = err
break
}
got = append(got, item)
}
if iterErr == nil {
t.Fatal("expected pagination error at default cap, got nil")
}
if len(got) != DefaultListMaxPages {
t.Fatalf("got %d items, want %d (DefaultListMaxPages)", len(got), DefaultListMaxPages)
}
})
}

func TestPaginateCursorCycle(t *testing.T) {
ctx := context.Background()
callCount := 0
// Server returns a cursor cycle: page1 → cursorA, page2 → cursorB, page3 → cursorA again.
listFunc := func(ctx context.Context, params *ListTestParams) (*ListTestResult, error) {
callCount++
switch callCount {
case 1:
return &ListTestResult{Items: []*Item{{Name: "a"}}, NextCursor: "cursorA"}, nil
case 2:
return &ListTestResult{Items: []*Item{{Name: "b"}}, NextCursor: "cursorB"}, nil
case 3:
// Cycle: cursorA was already seen.
return &ListTestResult{Items: []*Item{{Name: "c"}}, NextCursor: "cursorA"}, nil
default:
t.Fatal("listFunc called after cycle should have been detected")
return nil, fmt.Errorf("unreachable")
}
}

var got []*Item
var iterErr error
seq := paginate(ctx, &ListTestParams{}, -1, listFunc, func(r *ListTestResult) []*Item { return r.Items })
for item, err := range seq {
if err != nil {
iterErr = err
break
}
got = append(got, item)
}
if iterErr == nil {
t.Fatal("expected cursor cycle error, got nil")
}
if len(got) != 3 {
t.Fatalf("got %d items, want 3 (items on cycling page are yielded before detection)", len(got))
}
}

func TestPaginateNoMutation(t *testing.T) {
ctx := context.Background()
params := &ListTestParams{Cursor: "initial"}
results := generatePaginatedResults(allItems[:3], 1)

listFunc := func(ctx context.Context, p *ListTestParams) (*ListTestResult, error) {
if len(results) == 0 {
return &ListTestResult{Items: nil, NextCursor: ""}, nil
}
res := results[0]
results = results[1:]
return res, nil
}

var got []*Item
seq := paginate(ctx, params, -1, listFunc, func(r *ListTestResult) []*Item { return r.Items })
for item, err := range seq {
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
got = append(got, item)
}

// The original params should not have been mutated.
if params.Cursor != "initial" {
t.Errorf("params.Cursor was mutated to %q, want %q", params.Cursor, "initial")
}
if len(got) != 3 {
t.Fatalf("got %d items, want 3", len(got))
}
}

func TestClientCapabilities(t *testing.T) {
testCases := []struct {
name string
Expand Down