From 37ae90648b596ffd1dbd75854484c98142bf12b2 Mon Sep 17 00:00:00 2001 From: Jonathan Amsterdam Date: Fri, 9 May 2025 20:35:53 -0400 Subject: [PATCH] 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 LUCI-TryBot-Result: Go LUCI --- mcp/features.go | 73 +++++++++++++++++++++++++++++++++++++++++++++++++ mcp/mcp_test.go | 35 +++++++++++++----------- mcp/server.go | 66 ++++++++++++++++++++++++++------------------ 3 files changed, 132 insertions(+), 42 deletions(-) create mode 100644 mcp/features.go diff --git a/mcp/features.go b/mcp/features.go new file mode 100644 index 0000000..ea66be5 --- /dev/null +++ b/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 + } + } + } +} diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 0efd763..60ac84c 100644 --- a/mcp/mcp_test.go +++ b/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) } diff --git a/mcp/server.go b/mcp/server.go index 9c9351d..8060941 100644 --- a/mcp/server.go +++ b/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)