diff --git a/mcp/client.go b/mcp/client.go index 6f1073ad..22df9967 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -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{} @@ -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 }) } @@ -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 }) } @@ -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 }) } @@ -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 @@ -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 } } } diff --git a/mcp/client_test.go b/mcp/client_test.go index 0fe39c08..e1bb0910 100644 --- a/mcp/client_test.go +++ b/mcp/client_test.go @@ -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 @@ -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) @@ -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