Skip to content

Commit c7390fe

Browse files
JinlkjJinlkj
authored
Feature/pagination functionality (#107)
* add pagination functionality * update test * solve race problem * solve race pro
1 parent 62cdf71 commit c7390fe

File tree

6 files changed

+386
-101
lines changed

6 files changed

+386
-101
lines changed

‎client/client.go

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@ package client
33

44
import (
55
"context"
6+
"encoding/json"
7+
"fmt"
68

79
"github.com/mark3labs/mcp-go/mcp"
810
)
@@ -18,12 +20,25 @@ type MCPClient interface {
1820
// Ping checks if the server is alive
1921
Ping(ctx context.Context) error
2022

23+
// ListResourcesByPage manually list resources by page.
24+
ListResourcesByPage(
25+
ctx context.Context,
26+
request mcp.ListResourcesRequest,
27+
) (*mcp.ListResourcesResult, error)
28+
2129
// ListResources requests a list of available resources from the server
2230
ListResources(
2331
ctx context.Context,
2432
request mcp.ListResourcesRequest,
2533
) (*mcp.ListResourcesResult, error)
2634

35+
// ListResourceTemplatesByPage manually list resource templates by page.
36+
ListResourceTemplatesByPage(
37+
ctx context.Context,
38+
request mcp.ListResourceTemplatesRequest,
39+
) (*mcp.ListResourceTemplatesResult,
40+
error)
41+
2742
// ListResourceTemplates requests a list of available resource templates from the server
2843
ListResourceTemplates(
2944
ctx context.Context,
@@ -43,6 +58,12 @@ type MCPClient interface {
4358
// Unsubscribe cancels notifications for a specific resource
4459
Unsubscribe(ctx context.Context, request mcp.UnsubscribeRequest) error
4560

61+
// ListPromptsByPage manually list prompts by page.
62+
ListPromptsByPage(
63+
ctx context.Context,
64+
request mcp.ListPromptsRequest,
65+
) (*mcp.ListPromptsResult, error)
66+
4667
// ListPrompts requests a list of available prompts from the server
4768
ListPrompts(
4869
ctx context.Context,
@@ -55,6 +76,12 @@ type MCPClient interface {
5576
request mcp.GetPromptRequest,
5677
) (*mcp.GetPromptResult, error)
5778

79+
// ListToolsByPage manually list tools by page.
80+
ListToolsByPage(
81+
ctx context.Context,
82+
request mcp.ListToolsRequest,
83+
) (*mcp.ListToolsResult, error)
84+
5885
// ListTools requests a list of available tools from the server
5986
ListTools(
6087
ctx context.Context,
@@ -81,4 +108,23 @@ type MCPClient interface {
81108

82109
// OnNotification registers a handler for notifications
83110
OnNotification(handler func(notification mcp.JSONRPCNotification))
111+
112+
sendRequest(ctx context.Context, method string, params interface{}) (*json.RawMessage, error)
113+
}
114+
115+
func listByPage[T any](
116+
ctx context.Context,
117+
client MCPClient,
118+
request mcp.PaginatedRequest,
119+
method string,
120+
) (*T, error) {
121+
response, err := client.sendRequest(ctx, method, request.Params)
122+
if err != nil {
123+
return nil, err
124+
}
125+
var result T
126+
if err := json.Unmarshal(*response, &result); err != nil {
127+
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
128+
}
129+
return &result, nil
84130
}

‎client/sse.go

Lines changed: 104 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -402,42 +402,77 @@ func (c *SSEMCPClient) Ping(ctx context.Context) error {
402402
return err
403403
}
404404

405-
func (c *SSEMCPClient) ListResources(
405+
// ListResourcesByPage manually list resources by page.
406+
func (c *SSEMCPClient) ListResourcesByPage(
406407
ctx context.Context,
407408
request mcp.ListResourcesRequest,
408409
) (*mcp.ListResourcesResult, error) {
409-
response, err := c.sendRequest(ctx, "resources/list", request.Params)
410+
result, err := listByPage[mcp.ListResourcesResult](ctx, c, request.PaginatedRequest, "resources/list")
410411
if err != nil {
411412
return nil, err
412413
}
414+
return result, nil
415+
}
413416

414-
var result mcp.ListResourcesResult
415-
if err := json.Unmarshal(*response, &result); err != nil {
416-
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
417+
func (c *SSEMCPClient) ListResources(
418+
ctx context.Context,
419+
request mcp.ListResourcesRequest,
420+
) (*mcp.ListResourcesResult, error) {
421+
result, err := c.ListResourcesByPage(ctx, request)
422+
if err != nil {
423+
return nil, err
417424
}
425+
for result.NextCursor != "" {
426+
select {
427+
case <-ctx.Done():
428+
return nil, ctx.Err()
429+
default:
430+
request.Params.Cursor = result.NextCursor
431+
newPageRes, err := c.ListResourcesByPage(ctx, request)
432+
if err != nil {
433+
return nil, err
434+
}
435+
result.Resources = append(result.Resources, newPageRes.Resources...)
436+
result.NextCursor = newPageRes.NextCursor
437+
}
438+
}
439+
return result, nil
440+
}
418441

419-
return &result, nil
442+
func (c *SSEMCPClient) ListResourceTemplatesByPage(
443+
ctx context.Context,
444+
request mcp.ListResourceTemplatesRequest,
445+
) (*mcp.ListResourceTemplatesResult, error) {
446+
result, err := listByPage[mcp.ListResourceTemplatesResult](ctx, c, request.PaginatedRequest, "resources/templates/list")
447+
if err != nil {
448+
return nil, err
449+
}
450+
return result, nil
420451
}
421452

422453
func (c *SSEMCPClient) ListResourceTemplates(
423454
ctx context.Context,
424455
request mcp.ListResourceTemplatesRequest,
425456
) (*mcp.ListResourceTemplatesResult, error) {
426-
response, err := c.sendRequest(
427-
ctx,
428-
"resources/templates/list",
429-
request.Params,
430-
)
457+
result, err := c.ListResourceTemplatesByPage(ctx, request)
431458
if err != nil {
432459
return nil, err
433460
}
434-
435-
var result mcp.ListResourceTemplatesResult
436-
if err := json.Unmarshal(*response, &result); err != nil {
437-
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
461+
for result.NextCursor != "" {
462+
select {
463+
case <-ctx.Done():
464+
return nil, ctx.Err()
465+
default:
466+
request.Params.Cursor = result.NextCursor
467+
newPageRes, err := c.ListResourceTemplatesByPage(ctx, request)
468+
if err != nil {
469+
return nil, err
470+
}
471+
result.ResourceTemplates = append(result.ResourceTemplates, newPageRes.ResourceTemplates...)
472+
result.NextCursor = newPageRes.NextCursor
473+
}
438474
}
439-
440-
return &result, nil
475+
return result, nil
441476
}
442477

443478
func (c *SSEMCPClient) ReadResource(
@@ -468,21 +503,40 @@ func (c *SSEMCPClient) Unsubscribe(
468503
return err
469504
}
470505

471-
func (c *SSEMCPClient) ListPrompts(
506+
func (c *SSEMCPClient) ListPromptsByPage(
472507
ctx context.Context,
473508
request mcp.ListPromptsRequest,
474509
) (*mcp.ListPromptsResult, error) {
475-
response, err := c.sendRequest(ctx, "prompts/list", request.Params)
510+
result, err := listByPage[mcp.ListPromptsResult](ctx, c, request.PaginatedRequest, "prompts/list")
476511
if err != nil {
477512
return nil, err
478513
}
514+
return result, nil
515+
}
479516

480-
var result mcp.ListPromptsResult
481-
if err := json.Unmarshal(*response, &result); err != nil {
482-
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
517+
func (c *SSEMCPClient) ListPrompts(
518+
ctx context.Context,
519+
request mcp.ListPromptsRequest,
520+
) (*mcp.ListPromptsResult, error) {
521+
result, err := c.ListPromptsByPage(ctx, request)
522+
if err != nil {
523+
return nil, err
483524
}
484-
485-
return &result, nil
525+
for result.NextCursor != "" {
526+
select {
527+
case <-ctx.Done():
528+
return nil, ctx.Err()
529+
default:
530+
request.Params.Cursor = result.NextCursor
531+
newPageRes, err := c.ListPromptsByPage(ctx, request)
532+
if err != nil {
533+
return nil, err
534+
}
535+
result.Prompts = append(result.Prompts, newPageRes.Prompts...)
536+
result.NextCursor = newPageRes.NextCursor
537+
}
538+
}
539+
return result, nil
486540
}
487541

488542
func (c *SSEMCPClient) GetPrompt(
@@ -497,21 +551,40 @@ func (c *SSEMCPClient) GetPrompt(
497551
return mcp.ParseGetPromptResult(response)
498552
}
499553

500-
func (c *SSEMCPClient) ListTools(
554+
func (c *SSEMCPClient) ListToolsByPage(
501555
ctx context.Context,
502556
request mcp.ListToolsRequest,
503557
) (*mcp.ListToolsResult, error) {
504-
response, err := c.sendRequest(ctx, "tools/list", request.Params)
558+
result, err := listByPage[mcp.ListToolsResult](ctx, c, request.PaginatedRequest, "tools/list")
505559
if err != nil {
506560
return nil, err
507561
}
562+
return result, nil
563+
}
508564

509-
var result mcp.ListToolsResult
510-
if err := json.Unmarshal(*response, &result); err != nil {
511-
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
565+
func (c *SSEMCPClient) ListTools(
566+
ctx context.Context,
567+
request mcp.ListToolsRequest,
568+
) (*mcp.ListToolsResult, error) {
569+
result, err := c.ListToolsByPage(ctx, request)
570+
if err != nil {
571+
return nil, err
512572
}
513-
514-
return &result, nil
573+
for result.NextCursor != "" {
574+
select {
575+
case <-ctx.Done():
576+
return nil, ctx.Err()
577+
default:
578+
request.Params.Cursor = result.NextCursor
579+
newPageRes, err := c.ListToolsByPage(ctx, request)
580+
if err != nil {
581+
return nil, err
582+
}
583+
result.Tools = append(result.Tools, newPageRes.Tools...)
584+
result.NextCursor = newPageRes.NextCursor
585+
}
586+
}
587+
return result, nil
515588
}
516589

517590
func (c *SSEMCPClient) CallTool(

0 commit comments

Comments
 (0)