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>
This commit is contained in:
Sam Thanawalla 2025-05-27 19:44:25 +00:00
commit d0bc3bc890
2 changed files with 31 additions and 2 deletions

View file

@ -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

View file

@ -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)
}
}