Skip to content

Commit af97501

Browse files
vasayxtxvasayxtx
committed
feat: add missing SetPrompts, DeleteResources, and SetResources methods
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 baa7153 commit af97501

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
@@ -350,6 +350,32 @@ func (s *MCPServer) AddResource(
350350
s.AddResources(ServerResource{Resource: resource, Handler: handler})
351351
}
352352

353+
// DeleteResources removes resources from the server
354+
func (s *MCPServer) DeleteResources(uris ...string) {
355+
s.resourcesMu.Lock()
356+
var exists bool
357+
for _, uri := range uris {
358+
if _, ok := s.resources[uri]; ok {
359+
delete(s.resources, uri)
360+
exists = true
361+
}
362+
}
363+
s.resourcesMu.Unlock()
364+
365+
// Send notification to all initialized sessions if listChanged capability is enabled and we actually remove a resource
366+
if exists && s.capabilities.resources != nil && s.capabilities.resources.listChanged {
367+
s.SendNotificationToAllClients(mcp.MethodNotificationResourcesListChanged, nil)
368+
}
369+
}
370+
371+
// SetResources replaces all existing resources with the provided list
372+
func (s *MCPServer) SetResources(resources ...ServerResource) {
373+
s.resourcesMu.Lock()
374+
s.resources = make(map[string]resourceEntry, len(resources))
375+
s.resourcesMu.Unlock()
376+
s.AddResources(resources...)
377+
}
378+
353379
// RemoveResource removes a resource from the server
354380
func (s *MCPServer) RemoveResource(uri string) {
355381
s.resourcesMu.Lock()
@@ -429,6 +455,15 @@ func (s *MCPServer) DeletePrompts(names ...string) {
429455
}
430456
}
431457

458+
// SetPrompts replaces all existing prompts with the provided list
459+
func (s *MCPServer) SetPrompts(prompts ...ServerPrompt) {
460+
s.promptsMu.Lock()
461+
s.prompts = make(map[string]mcp.Prompt, len(prompts))
462+
s.promptHandlers = make(map[string]PromptHandlerFunc, len(prompts))
463+
s.promptsMu.Unlock()
464+
s.AddPrompts(prompts...)
465+
}
466+
432467
// AddTool registers a new tool and its handler
433468
func (s *MCPServer) AddTool(tool mcp.Tool, handler ToolHandlerFunc) {
434469
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
@@ -962,6 +962,50 @@ func TestMCPServer_Prompts(t *testing.T) {
962962
assert.Equal(t, "test-prompt-2", prompts[1].Name)
963963
},
964964
},
965+
{
966+
name: "SetPrompts sends single notifications/prompts/list_changed with one active session",
967+
action: func(t *testing.T, server *MCPServer, notificationChannel chan mcp.JSONRPCNotification) {
968+
err := server.RegisterSession(context.TODO(), &fakeSession{
969+
sessionID: "test",
970+
notificationChannel: notificationChannel,
971+
initialized: true,
972+
})
973+
require.NoError(t, err)
974+
server.SetPrompts(ServerPrompt{
975+
Prompt: mcp.Prompt{
976+
Name: "test-prompt-1",
977+
Description: "A test prompt",
978+
Arguments: []mcp.PromptArgument{
979+
{
980+
Name: "arg1",
981+
Description: "First argument",
982+
},
983+
},
984+
},
985+
Handler: nil,
986+
}, ServerPrompt{
987+
Prompt: mcp.Prompt{
988+
Name: "test-prompt-2",
989+
Description: "Another test prompt",
990+
Arguments: []mcp.PromptArgument{
991+
{
992+
Name: "arg2",
993+
Description: "Second argument",
994+
},
995+
},
996+
},
997+
Handler: nil,
998+
})
999+
},
1000+
expectedNotifications: 1,
1001+
validate: func(t *testing.T, notifications []mcp.JSONRPCNotification, promptsList mcp.JSONRPCMessage) {
1002+
assert.Equal(t, mcp.MethodNotificationPromptsListChanged, notifications[0].Method)
1003+
prompts := promptsList.(mcp.JSONRPCResponse).Result.(mcp.ListPromptsResult).Prompts
1004+
assert.Len(t, prompts, 2)
1005+
assert.Equal(t, "test-prompt-1", prompts[0].Name)
1006+
assert.Equal(t, "test-prompt-2", prompts[1].Name)
1007+
},
1008+
},
9651009
}
9661010
for _, tt := range tests {
9671011
t.Run(tt.name, func(t *testing.T) {
@@ -997,6 +1041,211 @@ func TestMCPServer_Prompts(t *testing.T) {
9971041
}
9981042
}
9991043

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

0 commit comments

Comments
 (0)