Skip to content

Commit 820b7a6

Browse files
God-JayGod-Jay
authored
Feat(prompts): add DeletePrompts method to MCPServer (#320)
1 parent 3cdeb89 commit 820b7a6

File tree

3 files changed

+221
-2
lines changed

3 files changed

+221
-2
lines changed

‎server/server.go‎

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -403,13 +403,33 @@ func (s *MCPServer) AddPrompt(prompt mcp.Prompt, handler PromptHandlerFunc) {
403403
s.promptHandlers[prompt.Name] = handler
404404
s.promptsMu.Unlock()
405405

406-
// When the list of available resources changes, servers that declared the listChanged capability SHOULD send a notification.
406+
// When the list of available prompts changes, servers that declared the listChanged capability SHOULD send a notification.
407407
if s.capabilities.prompts.listChanged {
408408
// Send notification to all initialized sessions
409409
s.SendNotificationToAllClients(mcp.MethodNotificationPromptsListChanged, nil)
410410
}
411411
}
412412

413+
// DeletePrompts removes prompts from the server
414+
func (s *MCPServer) DeletePrompts(names ...string) {
415+
s.promptsMu.Lock()
416+
var exists bool
417+
for _, name := range names {
418+
if _, ok := s.prompts[name]; ok {
419+
delete(s.prompts, name)
420+
delete(s.promptHandlers, name)
421+
exists = true
422+
}
423+
}
424+
s.promptsMu.Unlock()
425+
426+
// Send notification to all initialized sessions if listChanged capability is enabled, and we actually remove a prompt
427+
if exists && s.capabilities.prompts != nil && s.capabilities.prompts.listChanged {
428+
// Send notification to all initialized sessions
429+
s.SendNotificationToAllClients(mcp.MethodNotificationPromptsListChanged, nil)
430+
}
431+
}
432+
413433
// AddTool registers a new tool and its handler
414434
func (s *MCPServer) AddTool(tool mcp.Tool, handler ToolHandlerFunc) {
415435
s.AddTools(ServerTool{Tool: tool, Handler: handler})
@@ -460,7 +480,7 @@ func (s *MCPServer) SetTools(tools ...ServerTool) {
460480
s.AddTools(tools...)
461481
}
462482

463-
// DeleteTools removes a tool from the server
483+
// DeleteTools removes tools from the server
464484
func (s *MCPServer) DeleteTools(names ...string) {
465485
s.toolsMu.Lock()
466486
var exists bool

‎server/server_race_test.go‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,17 @@ func TestRaceConditions(t *testing.T) {
4444
})
4545
})
4646

47+
runConcurrentOperation(&wg, testDuration, "delete-prompts", func() {
48+
name := fmt.Sprintf("delete-prompt-%d", time.Now().UnixNano())
49+
srv.AddPrompt(mcp.Prompt{
50+
Name: name,
51+
Description: "Temporary prompt",
52+
}, func(ctx context.Context, req mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
53+
return &mcp.GetPromptResult{}, nil
54+
})
55+
srv.DeletePrompts(name)
56+
})
57+
4758
runConcurrentOperation(&wg, testDuration, "add-tools", func() {
4859
name := fmt.Sprintf("tool-%d", time.Now().UnixNano())
4960
srv.AddTool(mcp.Tool{

‎server/server_test.go‎

Lines changed: 188 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -809,6 +809,194 @@ func TestMCPServer_PromptHandling(t *testing.T) {
809809
}
810810
}
811811

812+
func TestMCPServer_Prompts(t *testing.T) {
813+
tests := []struct {
814+
name string
815+
action func(*testing.T, *MCPServer, chan mcp.JSONRPCNotification)
816+
expectedNotifications int
817+
validate func(*testing.T, []mcp.JSONRPCNotification, mcp.JSONRPCMessage)
818+
}{
819+
{
820+
name: "DeletePrompts sends single notifications/prompts/list_changed",
821+
action: func(t *testing.T, server *MCPServer, notificationChannel chan mcp.JSONRPCNotification) {
822+
err := server.RegisterSession(context.TODO(), &fakeSession{
823+
sessionID: "test",
824+
notificationChannel: notificationChannel,
825+
initialized: true,
826+
})
827+
require.NoError(t, err)
828+
server.AddPrompt(
829+
mcp.Prompt{
830+
Name: "test-prompt-1",
831+
Description: "A test prompt",
832+
Arguments: []mcp.PromptArgument{
833+
{
834+
Name: "arg1",
835+
Description: "First argument",
836+
},
837+
},
838+
},
839+
nil,
840+
)
841+
server.DeletePrompts("test-prompt-1")
842+
},
843+
expectedNotifications: 2,
844+
validate: func(t *testing.T, notifications []mcp.JSONRPCNotification, promptsList mcp.JSONRPCMessage) {
845+
// One for AddPrompt
846+
assert.Equal(t, mcp.MethodNotificationPromptsListChanged, notifications[0].Method)
847+
// One for DeletePrompts
848+
assert.Equal(t, mcp.MethodNotificationPromptsListChanged, notifications[1].Method)
849+
850+
// Expect a successful response with an empty list of prompts
851+
resp, ok := promptsList.(mcp.JSONRPCResponse)
852+
assert.True(t, ok, "Expected JSONRPCResponse, got %T", promptsList)
853+
854+
result, ok := resp.Result.(mcp.ListPromptsResult)
855+
assert.True(t, ok, "Expected ListPromptsResult, got %T", resp.Result)
856+
857+
assert.Empty(t, result.Prompts, "Expected empty prompts list")
858+
},
859+
},
860+
{
861+
name: "DeletePrompts removes the first prompt and retains the other",
862+
action: func(t *testing.T, server *MCPServer, notificationChannel chan mcp.JSONRPCNotification) {
863+
err := server.RegisterSession(context.TODO(), &fakeSession{
864+
sessionID: "test",
865+
notificationChannel: notificationChannel,
866+
initialized: true,
867+
})
868+
require.NoError(t, err)
869+
server.AddPrompt(
870+
mcp.Prompt{
871+
Name: "test-prompt-1",
872+
Description: "A test prompt",
873+
Arguments: []mcp.PromptArgument{
874+
{
875+
Name: "arg1",
876+
Description: "First argument",
877+
},
878+
},
879+
},
880+
nil,
881+
)
882+
server.AddPrompt(
883+
mcp.Prompt{
884+
Name: "test-prompt-2",
885+
Description: "A test prompt",
886+
Arguments: []mcp.PromptArgument{
887+
{
888+
Name: "arg1",
889+
Description: "First argument",
890+
},
891+
},
892+
},
893+
nil,
894+
)
895+
// Remove non-existing prompts
896+
server.DeletePrompts("test-prompt-1")
897+
},
898+
expectedNotifications: 3,
899+
validate: func(t *testing.T, notifications []mcp.JSONRPCNotification, promptsList mcp.JSONRPCMessage) {
900+
// first notification expected for AddPrompt test-prompt-1
901+
assert.Equal(t, mcp.MethodNotificationPromptsListChanged, notifications[0].Method)
902+
// second notification expected for AddPrompt test-prompt-2
903+
assert.Equal(t, mcp.MethodNotificationPromptsListChanged, notifications[1].Method)
904+
// second notification expected for DeletePrompts test-prompt-1
905+
assert.Equal(t, mcp.MethodNotificationPromptsListChanged, notifications[2].Method)
906+
907+
// Confirm the prompt list does not change
908+
prompts := promptsList.(mcp.JSONRPCResponse).Result.(mcp.ListPromptsResult).Prompts
909+
assert.Len(t, prompts, 1)
910+
assert.Equal(t, "test-prompt-2", prompts[0].Name)
911+
},
912+
},
913+
{
914+
name: "DeletePrompts with non-existent prompts does nothing and not receives notifications from MCPServer",
915+
action: func(t *testing.T, server *MCPServer, notificationChannel chan mcp.JSONRPCNotification) {
916+
err := server.RegisterSession(context.TODO(), &fakeSession{
917+
sessionID: "test",
918+
notificationChannel: notificationChannel,
919+
initialized: true,
920+
})
921+
require.NoError(t, err)
922+
server.AddPrompt(
923+
mcp.Prompt{
924+
Name: "test-prompt-1",
925+
Description: "A test prompt",
926+
Arguments: []mcp.PromptArgument{
927+
{
928+
Name: "arg1",
929+
Description: "First argument",
930+
},
931+
},
932+
},
933+
nil,
934+
)
935+
server.AddPrompt(
936+
mcp.Prompt{
937+
Name: "test-prompt-2",
938+
Description: "A test prompt",
939+
Arguments: []mcp.PromptArgument{
940+
{
941+
Name: "arg1",
942+
Description: "First argument",
943+
},
944+
},
945+
},
946+
nil,
947+
)
948+
// Remove non-existing prompts
949+
server.DeletePrompts("test-prompt-3", "test-prompt-4")
950+
},
951+
expectedNotifications: 2,
952+
validate: func(t *testing.T, notifications []mcp.JSONRPCNotification, promptsList mcp.JSONRPCMessage) {
953+
// first notification expected for AddPrompt test-prompt-1
954+
assert.Equal(t, mcp.MethodNotificationPromptsListChanged, notifications[0].Method)
955+
// second notification expected for AddPrompt test-prompt-2
956+
assert.Equal(t, mcp.MethodNotificationPromptsListChanged, notifications[1].Method)
957+
958+
// Confirm the prompt list does not change
959+
prompts := promptsList.(mcp.JSONRPCResponse).Result.(mcp.ListPromptsResult).Prompts
960+
assert.Len(t, prompts, 2)
961+
assert.Equal(t, "test-prompt-1", prompts[0].Name)
962+
assert.Equal(t, "test-prompt-2", prompts[1].Name)
963+
},
964+
},
965+
}
966+
for _, tt := range tests {
967+
t.Run(tt.name, func(t *testing.T) {
968+
ctx := context.Background()
969+
server := NewMCPServer("test-server", "1.0.0", WithPromptCapabilities(true))
970+
_ = server.HandleMessage(ctx, []byte(`{
971+
"jsonrpc": "2.0",
972+
"id": 1,
973+
"method": "initialize"
974+
}`))
975+
notificationChannel := make(chan mcp.JSONRPCNotification, 100)
976+
notifications := make([]mcp.JSONRPCNotification, 0)
977+
tt.action(t, server, notificationChannel)
978+
for done := false; !done; {
979+
select {
980+
case serverNotification := <-notificationChannel:
981+
notifications = append(notifications, serverNotification)
982+
if len(notifications) == tt.expectedNotifications {
983+
done = true
984+
}
985+
case <-time.After(1 * time.Second):
986+
done = true
987+
}
988+
}
989+
assert.Len(t, notifications, tt.expectedNotifications)
990+
promptsList := server.HandleMessage(ctx, []byte(`{
991+
"jsonrpc": "2.0",
992+
"id": 1,
993+
"method": "prompts/list"
994+
}`))
995+
tt.validate(t, notifications, promptsList)
996+
})
997+
}
998+
}
999+
8121000
func TestMCPServer_HandleInvalidMessages(t *testing.T) {
8131001
var errs []error
8141002
hooks := &Hooks{}

0 commit comments

Comments
 (0)