## Background Since the Server has a map of ServersSessions->subscriptions, and takes care of routing ResourceUpdate messages to ServerSessions, I think the decision was made to leave *ServerSession out of the SubscribeHandler/UnsubscribeHandler signatures. However, resource subscriptions are not specific to the a particular client and this information is kept private in the Session today. ## Proposal Align the Subscribe/Unsubscribe handlers with the other two handlers that are client specific (ProgressNotificationHandler, and RootListChangeHanlder), and include *ServerSession in the signature. It comes in practice because because a gateway server that needs to forward subscriptions still needs this information. The routing logic in Server is convenient and means that most servers won't need to know which ServerSessions are subscribed. But I think there are still use cases where the user will need to know.
373 lines
11 KiB
Go
373 lines
11 KiB
Go
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
|
|
// Use of this source code is governed by an MIT-style
|
|
// license that can be found in the LICENSE file.
|
|
|
|
package mcp
|
|
|
|
import (
|
|
"context"
|
|
"log"
|
|
"slices"
|
|
"testing"
|
|
|
|
"github.com/google/go-cmp/cmp"
|
|
"github.com/modelcontextprotocol/go-sdk/jsonschema"
|
|
)
|
|
|
|
type testItem struct {
|
|
Name string
|
|
Value string
|
|
}
|
|
|
|
type testListParams struct {
|
|
Cursor string
|
|
}
|
|
|
|
func (p *testListParams) cursorPtr() *string {
|
|
return &p.Cursor
|
|
}
|
|
|
|
type testListResult struct {
|
|
Items []*testItem
|
|
NextCursor string
|
|
}
|
|
|
|
func (r *testListResult) nextCursorPtr() *string {
|
|
return &r.NextCursor
|
|
}
|
|
|
|
var allTestItems = []*testItem{
|
|
{"alpha", "val-A"},
|
|
{"bravo", "val-B"},
|
|
{"charlie", "val-C"},
|
|
{"delta", "val-D"},
|
|
{"echo", "val-E"},
|
|
{"foxtrot", "val-F"},
|
|
{"golf", "val-G"},
|
|
{"hotel", "val-H"},
|
|
{"india", "val-I"},
|
|
{"juliet", "val-J"},
|
|
{"kilo", "val-K"},
|
|
}
|
|
|
|
// getCursor encodes a string input into a URL-safe base64 cursor,
|
|
// fatally logging any encoding errors.
|
|
func getCursor(input string) string {
|
|
cursor, err := encodeCursor(input)
|
|
if err != nil {
|
|
log.Fatalf("encodeCursor(%s) error = %v", input, err)
|
|
}
|
|
return cursor
|
|
}
|
|
|
|
func TestServerPaginateBasic(t *testing.T) {
|
|
testCases := []struct {
|
|
name string
|
|
initialItems []*testItem
|
|
inputCursor string
|
|
inputPageSize int
|
|
wantFeatures []*testItem
|
|
wantNextCursor string
|
|
wantErr bool
|
|
}{
|
|
{
|
|
name: "FirstPage_DefaultSize_Full",
|
|
initialItems: allTestItems,
|
|
inputCursor: "",
|
|
inputPageSize: 5,
|
|
wantFeatures: allTestItems[0:5],
|
|
wantNextCursor: getCursor("echo"), // Based on last item of first page
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "SecondPage_DefaultSize_Full",
|
|
initialItems: allTestItems,
|
|
inputCursor: getCursor("echo"),
|
|
inputPageSize: 5,
|
|
wantFeatures: allTestItems[5:10],
|
|
wantNextCursor: getCursor("juliet"), // Based on last item of second page
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "SecondPage_DefaultSize_Full_OutOfOrder",
|
|
initialItems: append(allTestItems[5:], allTestItems[0:5]...),
|
|
inputCursor: getCursor("echo"),
|
|
inputPageSize: 5,
|
|
wantFeatures: allTestItems[5:10],
|
|
wantNextCursor: getCursor("juliet"), // Based on last item of second page
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "SecondPage_DefaultSize_Full_Duplicates",
|
|
initialItems: append(allTestItems, allTestItems[0:5]...),
|
|
inputCursor: getCursor("echo"),
|
|
inputPageSize: 5,
|
|
wantFeatures: allTestItems[5:10],
|
|
wantNextCursor: getCursor("juliet"), // Based on last item of second page
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "LastPage_Remaining",
|
|
initialItems: allTestItems,
|
|
inputCursor: getCursor("juliet"),
|
|
inputPageSize: 5,
|
|
wantFeatures: allTestItems[10:11], // Only 1 item left
|
|
wantNextCursor: "", // No more pages
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "PageSize_1",
|
|
initialItems: allTestItems,
|
|
inputCursor: "",
|
|
inputPageSize: 1,
|
|
wantFeatures: allTestItems[0:1],
|
|
wantNextCursor: getCursor("alpha"),
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "PageSize_All",
|
|
initialItems: allTestItems,
|
|
inputCursor: "",
|
|
inputPageSize: len(allTestItems), // Page size equals total
|
|
wantFeatures: allTestItems,
|
|
wantNextCursor: "", // No more pages
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "PageSize_LargerThanAll",
|
|
initialItems: allTestItems,
|
|
inputCursor: "",
|
|
inputPageSize: len(allTestItems) + 5, // Page size larger than total
|
|
wantFeatures: allTestItems,
|
|
wantNextCursor: "",
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "EmptySet",
|
|
initialItems: nil,
|
|
inputCursor: "",
|
|
inputPageSize: 5,
|
|
wantFeatures: nil,
|
|
wantNextCursor: "",
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "InvalidCursor",
|
|
initialItems: allTestItems,
|
|
inputCursor: "not-a-valid-gob-base64-cursor",
|
|
inputPageSize: 5,
|
|
wantFeatures: nil, // Should be nil for error cases
|
|
wantNextCursor: "",
|
|
wantErr: true,
|
|
},
|
|
{
|
|
name: "AboveNonExistentID",
|
|
initialItems: allTestItems,
|
|
inputCursor: getCursor("dne"), // A UID that doesn't exist
|
|
inputPageSize: 5,
|
|
wantFeatures: allTestItems[4:9], // Should return elements above UID.
|
|
wantNextCursor: getCursor("india"),
|
|
wantErr: false,
|
|
},
|
|
}
|
|
|
|
for _, tc := range testCases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
fs := newFeatureSet(func(t *testItem) string { return t.Name })
|
|
fs.add(tc.initialItems...)
|
|
params := &testListParams{Cursor: tc.inputCursor}
|
|
gotResult, err := paginateList(fs, tc.inputPageSize, params, &testListResult{}, func(res *testListResult, items []*testItem) {
|
|
res.Items = items
|
|
})
|
|
if (err != nil) != tc.wantErr {
|
|
t.Errorf("paginateList(%s) error, got %v, wantErr %v", tc.name, err, tc.wantErr)
|
|
}
|
|
if tc.wantErr {
|
|
return
|
|
}
|
|
if diff := cmp.Diff(tc.wantFeatures, gotResult.Items); diff != "" {
|
|
t.Errorf("paginateList(%s) mismatch (-want +got):\n%s", tc.name, diff)
|
|
}
|
|
if tc.wantNextCursor != gotResult.NextCursor {
|
|
t.Errorf("paginateList(%s) nextCursor, got %v, want %v", tc.name, gotResult.NextCursor, tc.wantNextCursor)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestServerPaginateVariousPageSizes(t *testing.T) {
|
|
fs := newFeatureSet(func(t *testItem) string { return t.Name })
|
|
fs.add(allTestItems...)
|
|
// Try all possible page sizes, ensuring we get the correct list of items.
|
|
for pageSize := 1; pageSize < len(allTestItems)+1; pageSize++ {
|
|
var gotItems []*testItem
|
|
var nextCursor string
|
|
wantChunks := slices.Collect(slices.Chunk(allTestItems, pageSize))
|
|
index := 0
|
|
// Iterate through all pages, comparing sub-slices to the paginated list.
|
|
for {
|
|
params := &testListParams{Cursor: nextCursor}
|
|
gotResult, err := paginateList(fs, pageSize, params, &testListResult{}, func(res *testListResult, items []*testItem) {
|
|
res.Items = items
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("paginateList() unexpected error for pageSize %d, cursor %q: %v", pageSize, nextCursor, err)
|
|
}
|
|
if diff := cmp.Diff(wantChunks[index], gotResult.Items); diff != "" {
|
|
t.Errorf("paginateList mismatch (-want +got):\n%s", diff)
|
|
}
|
|
gotItems = append(gotItems, gotResult.Items...)
|
|
nextCursor = gotResult.NextCursor
|
|
if nextCursor == "" {
|
|
break
|
|
}
|
|
index++
|
|
}
|
|
|
|
if len(gotItems) != len(allTestItems) {
|
|
t.Fatalf("paginateList() returned %d items, want %d", len(allTestItems), len(gotItems))
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestServerCapabilities(t *testing.T) {
|
|
tool := &Tool{Name: "t", InputSchema: &jsonschema.Schema{}}
|
|
testCases := []struct {
|
|
name string
|
|
configureServer func(s *Server)
|
|
serverOpts ServerOptions
|
|
wantCapabilities *serverCapabilities
|
|
}{
|
|
{
|
|
name: "No capabilities",
|
|
configureServer: func(s *Server) {},
|
|
wantCapabilities: &serverCapabilities{
|
|
Logging: &loggingCapabilities{},
|
|
},
|
|
},
|
|
{
|
|
name: "With prompts",
|
|
configureServer: func(s *Server) {
|
|
s.AddPrompt(&Prompt{Name: "p"}, nil)
|
|
},
|
|
wantCapabilities: &serverCapabilities{
|
|
Logging: &loggingCapabilities{},
|
|
Prompts: &promptCapabilities{ListChanged: true},
|
|
},
|
|
},
|
|
{
|
|
name: "With resources",
|
|
configureServer: func(s *Server) {
|
|
s.AddResource(&Resource{URI: "file:///r"}, nil)
|
|
},
|
|
wantCapabilities: &serverCapabilities{
|
|
Logging: &loggingCapabilities{},
|
|
Resources: &resourceCapabilities{ListChanged: true},
|
|
},
|
|
},
|
|
{
|
|
name: "With resource templates",
|
|
configureServer: func(s *Server) {
|
|
s.AddResourceTemplate(&ResourceTemplate{URITemplate: "file:///rt"}, nil)
|
|
},
|
|
wantCapabilities: &serverCapabilities{
|
|
Logging: &loggingCapabilities{},
|
|
Resources: &resourceCapabilities{ListChanged: true},
|
|
},
|
|
},
|
|
{
|
|
name: "With resource subscriptions",
|
|
configureServer: func(s *Server) {
|
|
s.AddResourceTemplate(&ResourceTemplate{URITemplate: "file:///rt"}, nil)
|
|
},
|
|
serverOpts: ServerOptions{
|
|
SubscribeHandler: func(ctx context.Context, _ *ServerSession, sp *SubscribeParams) error {
|
|
return nil
|
|
},
|
|
UnsubscribeHandler: func(ctx context.Context, _ *ServerSession, up *UnsubscribeParams) error {
|
|
return nil
|
|
},
|
|
},
|
|
wantCapabilities: &serverCapabilities{
|
|
Logging: &loggingCapabilities{},
|
|
Resources: &resourceCapabilities{ListChanged: true, Subscribe: true},
|
|
},
|
|
},
|
|
{
|
|
name: "With tools",
|
|
configureServer: func(s *Server) {
|
|
s.AddTool(tool, nil)
|
|
},
|
|
wantCapabilities: &serverCapabilities{
|
|
Logging: &loggingCapabilities{},
|
|
Tools: &toolCapabilities{ListChanged: true},
|
|
},
|
|
},
|
|
{
|
|
name: "With completions",
|
|
configureServer: func(s *Server) {},
|
|
serverOpts: ServerOptions{
|
|
CompletionHandler: func(ctx context.Context, ss *ServerSession, params *CompleteParams) (*CompleteResult, error) {
|
|
return nil, nil
|
|
},
|
|
},
|
|
wantCapabilities: &serverCapabilities{
|
|
Logging: &loggingCapabilities{},
|
|
Completions: &completionCapabilities{},
|
|
},
|
|
},
|
|
{
|
|
name: "With all capabilities",
|
|
configureServer: func(s *Server) {
|
|
s.AddPrompt(&Prompt{Name: "p"}, nil)
|
|
s.AddResource(&Resource{URI: "file:///r"}, nil)
|
|
s.AddResourceTemplate(&ResourceTemplate{URITemplate: "file:///rt"}, nil)
|
|
s.AddTool(tool, nil)
|
|
},
|
|
serverOpts: ServerOptions{
|
|
SubscribeHandler: func(ctx context.Context, _ *ServerSession, sp *SubscribeParams) error {
|
|
return nil
|
|
},
|
|
UnsubscribeHandler: func(ctx context.Context, _ *ServerSession, up *UnsubscribeParams) error {
|
|
return nil
|
|
},
|
|
CompletionHandler: func(ctx context.Context, ss *ServerSession, params *CompleteParams) (*CompleteResult, error) {
|
|
return nil, nil
|
|
},
|
|
},
|
|
wantCapabilities: &serverCapabilities{
|
|
Completions: &completionCapabilities{},
|
|
Logging: &loggingCapabilities{},
|
|
Prompts: &promptCapabilities{ListChanged: true},
|
|
Resources: &resourceCapabilities{ListChanged: true, Subscribe: true},
|
|
Tools: &toolCapabilities{ListChanged: true},
|
|
},
|
|
},
|
|
{
|
|
name: "With initial capabilities",
|
|
configureServer: func(s *Server) {},
|
|
serverOpts: ServerOptions{
|
|
HasPrompts: true,
|
|
HasResources: true,
|
|
HasTools: true,
|
|
},
|
|
wantCapabilities: &serverCapabilities{
|
|
Logging: &loggingCapabilities{},
|
|
Prompts: &promptCapabilities{ListChanged: true},
|
|
Resources: &resourceCapabilities{ListChanged: true},
|
|
Tools: &toolCapabilities{ListChanged: true},
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tc := range testCases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
server := NewServer(testImpl, &tc.serverOpts)
|
|
tc.configureServer(server)
|
|
gotCapabilities := server.capabilities()
|
|
if diff := cmp.Diff(tc.wantCapabilities, gotCapabilities); diff != "" {
|
|
t.Errorf("capabilities() mismatch (-want +got):\n%s", diff)
|
|
}
|
|
})
|
|
}
|
|
}
|