Skip to content

Commit a3d34d9

Browse files
vasayxtxvasayxtx
authored
feat: add missing SetPrompts, DeleteResources, and SetResources methods (#445)
Complete method parity across Tools, Prompts, and Resources: - SetPrompts/SetResources atomically replace all existing items - DeleteResources removes multiple resources by URI - Add comprehensive tests following existing patterns
1 parent 8a88d01 commit a3d34d9

File tree

2 files changed

+284
-0
lines changed

2 files changed

+284
-0
lines changed

‎server/server.go

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -365,6 +365,32 @@ func (s *MCPServer) AddResource(
365365
s.AddResources(ServerResource{Resource: resource, Handler: handler})
366366
}
367367

368+
// DeleteResources removes resources from the server
369+
func (s *MCPServer) DeleteResources(uris ...string) {
370+
s.resourcesMu.Lock()
371+
var exists bool
372+
for _, uri := range uris {
373+
if _, ok := s.resources[uri]; ok {
374+
delete(s.resources, uri)
375+
exists = true
376+
}
377+
}
378+
s.resourcesMu.Unlock()
379+
380+
// Send notification to all initialized sessions if listChanged capability is enabled and we actually remove a resource
381+
if exists && s.capabilities.resources != nil && s.capabilities.resources.listChanged {
382+
s.SendNotificationToAllClients(mcp.MethodNotificationResourcesListChanged, nil)
383+
}
384+
}
385+
386+
// SetResources replaces all existing resources with the provided list
387+
func (s *MCPServer) SetResources(resources ...ServerResource) {
388+
s.resourcesMu.Lock()
389+
s.resources = make(map[string]resourceEntry, len(resources))
390+
s.resourcesMu.Unlock()
391+
s.AddResources(resources...)
392+
}
393+
368394
// RemoveResource removes a resource from the server
369395
func (s *MCPServer) RemoveResource(uri string) {
370396
s.resourcesMu.Lock()
@@ -468,6 +494,15 @@ func (s *MCPServer) DeletePrompts(names ...string) {
468494
}
469495
}
470496

497+
// SetPrompts replaces all existing prompts with the provided list
498+
func (s *MCPServer) SetPrompts(prompts ...ServerPrompt) {
499+
s.promptsMu.Lock()
500+
s.prompts = make(map[string]mcp.Prompt, len(prompts))
501+
s.promptHandlers = make(map[string]PromptHandlerFunc, len(prompts))
502+
s.promptsMu.Unlock()
503+
s.AddPrompts(prompts...)
504+
}
505+
471506
// AddTool registers a new tool and its handler
472507
func (s *MCPServer) AddTool(tool mcp.Tool, handler ToolHandlerFunc) {
473508
s.AddTools(ServerTool{Tool: tool, Handler: handler})

‎server/server_test.go

Lines changed: 249 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -963,6 +963,50 @@ func TestMCPServer_Prompts(t *testing.T) {
963963
assert.Equal(t, "test-prompt-2", prompts[1].Name)
964964
},
965965
},
966+
{
967+
name: "SetPrompts sends single notifications/prompts/list_changed with one active session",
968+
action: func(t *testing.T, server *MCPServer, notificationChannel chan mcp.JSONRPCNotification) {
969+
err := server.RegisterSession(context.TODO(), &fakeSession{
970+
sessionID: "test",
971+
notificationChannel: notificationChannel,
972+
initialized: true,
973+
})
974+
require.NoError(t, err)
975+
server.SetPrompts(ServerPrompt{
976+
Prompt: mcp.Prompt{
977+
Name: "test-prompt-1",
978+
Description: "A test prompt",
979+
Arguments: []mcp.PromptArgument{
980+
{
981+
Name: "arg1",
982+
Description: "First argument",
983+
},
984+
},
985+
},
986+
Handler: nil,
987+
}, ServerPrompt{
988+
Prompt: mcp.Prompt{
989+
Name: "test-prompt-2",
990+
Description: "Another test prompt",
991+
Arguments: []mcp.PromptArgument{
992+
{
993+
Name: "arg2",
994+
Description: "Second argument",
995+
},
996+
},
997+
},
998+
Handler: nil,
999+
})
1000+
},
1001+
expectedNotifications: 1,
1002+
validate: func(t *testing.T, notifications []mcp.JSONRPCNotification, promptsList mcp.JSONRPCMessage) {
1003+
assert.Equal(t, mcp.MethodNotificationPromptsListChanged, notifications[0].Method)
1004+
prompts := promptsList.(mcp.JSONRPCResponse).Result.(mcp.ListPromptsResult).Prompts
1005+
assert.Len(t, prompts, 2)
1006+
assert.Equal(t, "test-prompt-1", prompts[0].Name)
1007+
assert.Equal(t, "test-prompt-2", prompts[1].Name)
1008+
},
1009+
},
9661010
}
9671011
for _, tt := range tests {
9681012
t.Run(tt.name, func(t *testing.T) {
@@ -998,6 +1042,211 @@ func TestMCPServer_Prompts(t *testing.T) {
9981042
}
9991043
}
10001044

1045+
func TestMCPServer_Resources(t *testing.T) {
1046+
tests := []struct {
1047+
name string
1048+
action func(*testing.T, *MCPServer, chan mcp.JSONRPCNotification)
1049+
expectedNotifications int
1050+
validate func(*testing.T, []mcp.JSONRPCNotification, mcp.JSONRPCMessage)
1051+
}{
1052+
{
1053+
name: "DeleteResources sends single notifications/resources/list_changed",
1054+
action: func(t *testing.T, server *MCPServer, notificationChannel chan mcp.JSONRPCNotification) {
1055+
err := server.RegisterSession(context.TODO(), &fakeSession{
1056+
sessionID: "test",
1057+
notificationChannel: notificationChannel,
1058+
initialized: true,
1059+
})
1060+
require.NoError(t, err)
1061+
server.AddResource(
1062+
mcp.Resource{
1063+
URI: "test://test-resource-1",
1064+
Name: "Test Resource 1",
1065+
},
1066+
func(ctx context.Context, request mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) {
1067+
return []mcp.ResourceContents{}, nil
1068+
},
1069+
)
1070+
server.DeleteResources("test://test-resource-1")
1071+
},
1072+
expectedNotifications: 2,
1073+
validate: func(t *testing.T, notifications []mcp.JSONRPCNotification, resourcesList mcp.JSONRPCMessage) {
1074+
// One for AddResource
1075+
assert.Equal(t, mcp.MethodNotificationResourcesListChanged, notifications[0].Method)
1076+
// One for DeleteResources
1077+
assert.Equal(t, mcp.MethodNotificationResourcesListChanged, notifications[1].Method)
1078+
1079+
// Expect a successful response with an empty list of resources
1080+
resp, ok := resourcesList.(mcp.JSONRPCResponse)
1081+
assert.True(t, ok, "Expected JSONRPCResponse, got %T", resourcesList)
1082+
1083+
result, ok := resp.Result.(mcp.ListResourcesResult)
1084+
assert.True(t, ok, "Expected ListResourcesResult, got %T", resp.Result)
1085+
1086+
assert.Empty(t, result.Resources, "Expected empty resources list")
1087+
},
1088+
},
1089+
{
1090+
name: "DeleteResources removes the first resource and retains the other",
1091+
action: func(t *testing.T, server *MCPServer, notificationChannel chan mcp.JSONRPCNotification) {
1092+
err := server.RegisterSession(context.TODO(), &fakeSession{
1093+
sessionID: "test",
1094+
notificationChannel: notificationChannel,
1095+
initialized: true,
1096+
})
1097+
require.NoError(t, err)
1098+
server.AddResource(
1099+
mcp.Resource{
1100+
URI: "test://test-resource-1",
1101+
Name: "Test Resource 1",
1102+
},
1103+
func(ctx context.Context, request mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) {
1104+
return []mcp.ResourceContents{}, nil
1105+
},
1106+
)
1107+
server.AddResource(
1108+
mcp.Resource{
1109+
URI: "test://test-resource-2",
1110+
Name: "Test Resource 2",
1111+
},
1112+
func(ctx context.Context, request mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) {
1113+
return []mcp.ResourceContents{}, nil
1114+
},
1115+
)
1116+
server.DeleteResources("test://test-resource-1")
1117+
},
1118+
expectedNotifications: 3,
1119+
validate: func(t *testing.T, notifications []mcp.JSONRPCNotification, resourcesList mcp.JSONRPCMessage) {
1120+
// first notification expected for AddResource test-resource-1
1121+
assert.Equal(t, mcp.MethodNotificationResourcesListChanged, notifications[0].Method)
1122+
// second notification expected for AddResource test-resource-2
1123+
assert.Equal(t, mcp.MethodNotificationResourcesListChanged, notifications[1].Method)
1124+
// third notification expected for DeleteResources test-resource-1
1125+
assert.Equal(t, mcp.MethodNotificationResourcesListChanged, notifications[2].Method)
1126+
1127+
// Confirm the resource list contains only test-resource-2
1128+
resources := resourcesList.(mcp.JSONRPCResponse).Result.(mcp.ListResourcesResult).Resources
1129+
assert.Len(t, resources, 1)
1130+
assert.Equal(t, "test://test-resource-2", resources[0].URI)
1131+
},
1132+
},
1133+
{
1134+
name: "DeleteResources with non-existent resources does nothing and not receives notifications from MCPServer",
1135+
action: func(t *testing.T, server *MCPServer, notificationChannel chan mcp.JSONRPCNotification) {
1136+
err := server.RegisterSession(context.TODO(), &fakeSession{
1137+
sessionID: "test",
1138+
notificationChannel: notificationChannel,
1139+
initialized: true,
1140+
})
1141+
require.NoError(t, err)
1142+
server.AddResource(
1143+
mcp.Resource{
1144+
URI: "test://test-resource-1",
1145+
Name: "Test Resource 1",
1146+
},
1147+
func(ctx context.Context, request mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) {
1148+
return []mcp.ResourceContents{}, nil
1149+
},
1150+
)
1151+
server.AddResource(
1152+
mcp.Resource{
1153+
URI: "test://test-resource-2",
1154+
Name: "Test Resource 2",
1155+
},
1156+
func(ctx context.Context, request mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) {
1157+
return []mcp.ResourceContents{}, nil
1158+
},
1159+
)
1160+
// Remove non-existing resources
1161+
server.DeleteResources("test://test-resource-3", "test://test-resource-4")
1162+
},
1163+
expectedNotifications: 2,
1164+
validate: func(t *testing.T, notifications []mcp.JSONRPCNotification, resourcesList mcp.JSONRPCMessage) {
1165+
// first notification expected for AddResource test-resource-1
1166+
assert.Equal(t, mcp.MethodNotificationResourcesListChanged, notifications[0].Method)
1167+
// second notification expected for AddResource test-resource-2
1168+
assert.Equal(t, mcp.MethodNotificationResourcesListChanged, notifications[1].Method)
1169+
1170+
// Confirm the resource list does not change
1171+
resources := resourcesList.(mcp.JSONRPCResponse).Result.(mcp.ListResourcesResult).Resources
1172+
assert.Len(t, resources, 2)
1173+
// Resources are sorted by name
1174+
assert.Equal(t, "test://test-resource-1", resources[0].URI)
1175+
assert.Equal(t, "test://test-resource-2", resources[1].URI)
1176+
},
1177+
},
1178+
{
1179+
name: "SetResources sends single notifications/resources/list_changed with one active session",
1180+
action: func(t *testing.T, server *MCPServer, notificationChannel chan mcp.JSONRPCNotification) {
1181+
err := server.RegisterSession(context.TODO(), &fakeSession{
1182+
sessionID: "test",
1183+
notificationChannel: notificationChannel,
1184+
initialized: true,
1185+
})
1186+
require.NoError(t, err)
1187+
server.SetResources(ServerResource{
1188+
Resource: mcp.Resource{
1189+
URI: "test://test-resource-1",
1190+
Name: "Test Resource 1",
1191+
},
1192+
Handler: func(ctx context.Context, request mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) {
1193+
return []mcp.ResourceContents{}, nil
1194+
},
1195+
}, ServerResource{
1196+
Resource: mcp.Resource{
1197+
URI: "test://test-resource-2",
1198+
Name: "Test Resource 2",
1199+
},
1200+
Handler: func(ctx context.Context, request mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) {
1201+
return []mcp.ResourceContents{}, nil
1202+
},
1203+
})
1204+
},
1205+
expectedNotifications: 1,
1206+
validate: func(t *testing.T, notifications []mcp.JSONRPCNotification, resourcesList mcp.JSONRPCMessage) {
1207+
assert.Equal(t, mcp.MethodNotificationResourcesListChanged, notifications[0].Method)
1208+
resources := resourcesList.(mcp.JSONRPCResponse).Result.(mcp.ListResourcesResult).Resources
1209+
assert.Len(t, resources, 2)
1210+
// Resources are sorted by name
1211+
assert.Equal(t, "test://test-resource-1", resources[0].URI)
1212+
assert.Equal(t, "test://test-resource-2", resources[1].URI)
1213+
},
1214+
},
1215+
}
1216+
for _, tt := range tests {
1217+
t.Run(tt.name, func(t *testing.T) {
1218+
ctx := context.Background()
1219+
server := NewMCPServer("test-server", "1.0.0", WithResourceCapabilities(true, true))
1220+
_ = server.HandleMessage(ctx, []byte(`{
1221+
"jsonrpc": "2.0",
1222+
"id": 1,
1223+
"method": "initialize"
1224+
}`))
1225+
notificationChannel := make(chan mcp.JSONRPCNotification, 100)
1226+
notifications := make([]mcp.JSONRPCNotification, 0)
1227+
tt.action(t, server, notificationChannel)
1228+
for done := false; !done; {
1229+
select {
1230+
case serverNotification := <-notificationChannel:
1231+
notifications = append(notifications, serverNotification)
1232+
if len(notifications) == tt.expectedNotifications {
1233+
done = true
1234+
}
1235+
case <-time.After(1 * time.Second):
1236+
done = true
1237+
}
1238+
}
1239+
assert.Len(t, notifications, tt.expectedNotifications)
1240+
resourcesList := server.HandleMessage(ctx, []byte(`{
1241+
"jsonrpc": "2.0",
1242+
"id": 1,
1243+
"method": "resources/list"
1244+
}`))
1245+
tt.validate(t, notifications, resourcesList)
1246+
})
1247+
}
1248+
}
1249+
10011250
func TestMCPServer_HandleInvalidMessages(t *testing.T) {
10021251
var errs []error
10031252
hooks := &Hooks{}

0 commit comments

Comments
 (0)