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:
parent
4e9d520c2e
commit
d0bc3bc890
2 changed files with 31 additions and 2 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue