From d0bc3bc8906797d62354c6a1400225437cad46f2 Mon Sep 17 00:00:00 2001
From: Sam Thanawalla <samthanawalla@google.com>
Date: Tue, 27 May 2025 19:44:25 +0000
Subject: [PATCH] internal/mcp: add pagination for prompts

Change-Id: Ib231427299dbfe807fb80ad8da6ba16d4a8fde03
Reviewed-on: https://go-review.googlesource.com/c/tools/+/676615
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Jonathan Amsterdam <jba@google.com>
---
 mcp/server.go              | 11 ++++++++++-
 mcp/server_example_test.go | 22 +++++++++++++++++++++-
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/mcp/server.go b/mcp/server.go
index 3930a62..f5cde12 100644
--- a/mcp/server.go
+++ b/mcp/server.go
@@ -179,8 +179,17 @@ func (s *Server) Sessions() iter.Seq[*ServerSession] {
 func (s *Server) listPrompts(_ context.Context, _ *ServerSession, params *ListPromptsParams) (*ListPromptsResult, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
+	var cursor string
+	if params != nil {
+		cursor = params.Cursor
+	}
+	prompts, nextCursor, err := paginateList(s.prompts, cursor, s.opts.PageSize)
+	if err != nil {
+		return nil, err
+	}
 	res := new(ListPromptsResult)
-	for p := range s.prompts.all() {
+	res.NextCursor = nextCursor
+	for _, p := range prompts {
 		res.Prompts = append(res.Prompts, p.Prompt)
 	}
 	return res, nil
diff --git a/mcp/server_example_test.go b/mcp/server_example_test.go
index 0887d47..93679ac 100644
--- a/mcp/server_example_test.go
+++ b/mcp/server_example_test.go
@@ -72,7 +72,7 @@ func createSessions(ctx context.Context) (*mcp.ClientSession, *mcp.ServerSession
 	return clientSession, serverSession, server
 }
 
-func TestListTool(t *testing.T) {
+func TestListTools(t *testing.T) {
 	toolA := mcp.NewTool("apple", "apple tool", SayHi)
 	toolB := mcp.NewTool("banana", "banana tool", SayHi)
 	toolC := mcp.NewTool("cherry", "cherry tool", SayHi)
@@ -111,3 +111,23 @@ func TestListResources(t *testing.T) {
 		t.Fatalf("ListResources() mismatch (-want +got):\n%s", diff)
 	}
 }
+
+func TestListPrompts(t *testing.T) {
+	promptA := mcp.NewPrompt("apple", "apple prompt", testPromptHandler[struct{}])
+	promptB := mcp.NewPrompt("banana", "banana prompt", testPromptHandler[struct{}])
+	promptC := mcp.NewPrompt("cherry", "cherry prompt", testPromptHandler[struct{}])
+	prompts := []*mcp.ServerPrompt{promptA, promptB, promptC}
+	wantPrompts := []*mcp.Prompt{promptA.Prompt, promptB.Prompt, promptC.Prompt}
+	ctx := context.Background()
+	clientSession, serverSession, server := createSessions(ctx)
+	defer clientSession.Close()
+	defer serverSession.Close()
+	server.AddPrompts(prompts...)
+	res, err := clientSession.ListPrompts(ctx, nil)
+	if err != nil {
+		t.Fatal("ListPrompts() failed:", err)
+	}
+	if diff := cmp.Diff(wantPrompts, res.Prompts, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
+		t.Fatalf("ListPrompts() mismatch (-want +got):\n%s", diff)
+	}
+}