internal/mcp/protocol: featureSet[T]

What the spec calls "lists" of features (tools, prompts and so on)
are actually sets. Define a common data structure we can use to
implement all of them, and use it for server tools and prompts.

Change-Id: I257d0776b87a746ff23cd49ebcc7b317f4ea65e1
Reviewed-on: https://go-review.googlesource.com/c/tools/+/671359
Reviewed-by: Robert Findley <rfindley@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

mcp/features.go

@@ -0,0 +1,73 @@
+// 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 (
+	"iter"
+
+	"github.com/modelcontextprotocol/go-sdk/internal/util"
+)
+
+// This file contains implementations that are common to all features.
+// A feature is an item provided to a peer. In the 2025-03-26 spec,
+// the features are prompt, tool, resource and root.
+
+// A featureSet is a collection of features of type T.
+// Every feature has a unique ID, and the spec never mentions
+// an ordering for the List calls, so what it calls a "list" is actually a set.
+type featureSet[T any] struct {
+	uniqueID func(T) string
+	features map[string]T
+}
+
+// newFeatureSet creates a new featureSet for features of type T.
+// The argument function should return the unique ID for a single feature.
+func newFeatureSet[T any](uniqueIDFunc func(T) string) *featureSet[T] {
+	return &featureSet[T]{
+		uniqueID: uniqueIDFunc,
+		features: make(map[string]T),
+	}
+}
+
+// add adds each feature to the set if it is not present,
+// or replaces an existing feature.
+func (s *featureSet[T]) add(fs ...T) {
+	for _, f := range fs {
+		s.features[s.uniqueID(f)] = f
+	}
+}
+
+// remove removes all features with the given uids from the set if present,
+// and returns whether any were removed.
+// It is not an error to remove a nonexistent feature.
+func (s *featureSet[T]) remove(uids ...string) bool {
+	changed := false
+	for _, uid := range uids {
+		if _, ok := s.features[uid]; ok {
+			changed = true
+			delete(s.features, uid)
+		}
+	}
+	return changed
+}
+
+// get returns the feature with the given uid.
+// If there is none, it returns zero, false.
+func (s *featureSet[T]) get(uid string) (T, bool) {
+	t, ok := s.features[uid]
+	return t, ok
+}
+
+// all returns an iterator over of all the features in the set
+// sorted by unique ID.
+func (s *featureSet[T]) all() iter.Seq[T] {
+	return func(yield func(T) bool) {
+		for _, f := range util.Sorted(s.features) {
+			if !yield(f) {
+				return
+			}
+		}
+	}
+}

mcp/mcp_test.go

@@ -131,25 +131,28 @@ func TestEndToEnd(t *testing.T) {
 	if err != nil {
 		t.Errorf("tools/list failed: %v", err)
 	}
-	wantTools := []protocol.Tool{{
-		Name:        "greet",
-		Description: "say hi",
-		InputSchema: &jsonschema.Schema{
-			Type:     "object",
-			Required: []string{"Name"},
-			Properties: map[string]*jsonschema.Schema{
-				"Name": {Type: "string"},
+	wantTools := []protocol.Tool{
+		{
+			Name:        "fail",
+			Description: "just fail",
+			InputSchema: &jsonschema.Schema{
+				Type:                 "object",
+				AdditionalProperties: falseSchema,
 			},
-			AdditionalProperties: falseSchema,
 		},
-	}, {
-		Name:        "fail",
-		Description: "just fail",
-		InputSchema: &jsonschema.Schema{
-			Type:                 "object",
-			AdditionalProperties: falseSchema,
+		{
+			Name:        "greet",
+			Description: "say hi",
+			InputSchema: &jsonschema.Schema{
+				Type:     "object",
+				Required: []string{"Name"},
+				Properties: map[string]*jsonschema.Schema{
+					"Name": {Type: "string"},
+				},
+				AdditionalProperties: falseSchema,
+			},
 		},
-	}}
+	}
 	if diff := cmp.Diff(wantTools, gotTools, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
 		t.Fatalf("tools/list mismatch (-want +got):\n%s", diff)
 	}

mcp/server.go

@@ -21,13 +21,14 @@ import (
 // Servers expose server-side MCP features, which can serve one or more MCP
 // sessions by using [Server.Start] or [Server.Run].
 type Server struct {
+	// fixed at creation
 	name    string
 	version string
 	opts    ServerOptions
 
 	mu      sync.Mutex
-	prompts []*Prompt
-	tools   []*Tool
+	prompts *featureSet[*Prompt]
+	tools   *featureSet[*Tool]
 	conns   []*ServerConnection
 }
 
@@ -51,16 +52,29 @@ func NewServer(name, version string, opts *ServerOptions) *Server {
 		name:    name,
 		version: version,
 		opts:    *opts,
+		prompts: newFeatureSet(func(p *Prompt) string { return p.Definition.Name }),
+		tools:   newFeatureSet(func(t *Tool) string { return t.Definition.Name }),
 	}
 }
 
 // AddPrompts adds the given prompts to the server.
-//
-// TODO(rfindley): notify connected clients of any changes.
 func (s *Server) AddPrompts(prompts ...*Prompt) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
-	s.prompts = append(s.prompts, prompts...)
+	s.prompts.add(prompts...)
+	// Assume there was a change, since add replaces existing prompts.
+	// (It's possible a prompt was replaced with an identical one, but not worth checking.)
+	// TODO(rfindley): notify connected clients
+}
+
+// RemovePrompts removes if the prompts with the given names.
+// It is not an error to remove a nonexistent prompt.
+func (s *Server) RemovePrompts(names ...string) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.prompts.remove(names...) {
+		// TODO: notify
+	}
 }
 
 // AddTools adds the given tools to the server.
@@ -69,7 +83,20 @@ func (s *Server) AddPrompts(prompts ...*Prompt) {
 func (s *Server) AddTools(tools ...*Tool) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
-	s.tools = append(s.tools, tools...)
+	s.tools.add(tools...)
+	// Assume there was a change, since add replaces existing tools.
+	// (It's possible a tool was replaced with an identical one, but not worth checking.)
+	// TODO(rfindley): notify connected clients
+}
+
+// RemoveTools removes if the tools with the given names.
+// It is not an error to remove a nonexistent tool.
+func (s *Server) RemoveTools(names ...string) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.tools.remove(names...) {
+		// TODO: notify
+	}
 }
 
 // Clients returns an iterator that yields the current set of client
@@ -84,9 +111,8 @@ func (s *Server) Clients() iter.Seq[*ServerConnection] {
 func (s *Server) listPrompts(_ context.Context, _ *ServerConnection, params *protocol.ListPromptsParams) (*protocol.ListPromptsResult, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
-
 	res := new(protocol.ListPromptsResult)
-	for _, p := range s.prompts {
+	for p := range s.prompts.all() {
 		res.Prompts = append(res.Prompts, p.Definition)
 	}
 	return res, nil
@@ -94,15 +120,10 @@ func (s *Server) listPrompts(_ context.Context, _ *ServerConnection, params *pro
 
 func (s *Server) getPrompt(ctx context.Context, cc *ServerConnection, params *protocol.GetPromptParams) (*protocol.GetPromptResult, error) {
 	s.mu.Lock()
-	var prompt *Prompt
-	if i := slices.IndexFunc(s.prompts, func(t *Prompt) bool {
-		return t.Definition.Name == params.Name
-	}); i >= 0 {
-		prompt = s.prompts[i]
-	}
+	prompt, ok := s.prompts.get(params.Name)
 	s.mu.Unlock()
-
-	if prompt == nil {
+	if !ok {
+		// TODO: surface the error code over the wire, instead of flattening it into the string.
 		return nil, fmt.Errorf("%s: unknown prompt %q", jsonrpc2.ErrInvalidParams, params.Name)
 	}
 	return prompt.Handler(ctx, cc, params.Arguments)
@@ -111,9 +132,8 @@ func (s *Server) getPrompt(ctx context.Context, cc *ServerConnection, params *pr
 func (s *Server) listTools(_ context.Context, _ *ServerConnection, params *protocol.ListToolsParams) (*protocol.ListToolsResult, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
-
 	res := new(protocol.ListToolsResult)
-	for _, t := range s.tools {
+	for t := range s.tools.all() {
 		res.Tools = append(res.Tools, t.Definition)
 	}
 	return res, nil
@@ -121,15 +141,9 @@ func (s *Server) listTools(_ context.Context, _ *ServerConnection, params *proto
 
 func (s *Server) callTool(ctx context.Context, cc *ServerConnection, params *protocol.CallToolParams) (*protocol.CallToolResult, error) {
 	s.mu.Lock()
-	var tool *Tool
-	if i := slices.IndexFunc(s.tools, func(t *Tool) bool {
-		return t.Definition.Name == params.Name
-	}); i >= 0 {
-		tool = s.tools[i]
-	}
+	tool, ok := s.tools.get(params.Name)
 	s.mu.Unlock()
-
-	if tool == nil {
+	if !ok {
 		return nil, fmt.Errorf("%s: unknown tool %q", jsonrpc2.ErrInvalidParams, params.Name)
 	}
 	return tool.Handler(ctx, cc, params.Arguments)