go-sdk/mcp/server_test.go
Jim Clark cb17593a5d
update Subscribe/Unsubscribe handlers in ServerOpts (#231)
## 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.
2025-08-04 14:46:49 -04:00

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