diff --git a/examples/sse/main.go b/examples/sse/main.go index f72de6e..4c199f8 100644 --- a/examples/sse/main.go +++ b/examples/sse/main.go @@ -16,7 +16,7 @@ import ( var httpAddr = flag.String("http", "", "use SSE HTTP at this address") type SayHiParams struct { - Name string `json:"name" mcp:"the name to say hi to"` + Name string `json:"name"` } func SayHi(ctx context.Context, cc *mcp.ServerSession, params *SayHiParams) ([]*mcp.Content, error) { diff --git a/mcp/features.go b/mcp/features.go index ea66be5..1777b33 100644 --- a/mcp/features.go +++ b/mcp/features.go @@ -6,8 +6,8 @@ package mcp import ( "iter" - - "github.com/modelcontextprotocol/go-sdk/internal/util" + "maps" + "slices" ) // This file contains implementations that are common to all features. @@ -17,9 +17,11 @@ import ( // 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. +// TODO: switch to an ordered map type featureSet[T any] struct { - uniqueID func(T) string - features map[string]T + uniqueID func(T) string + features map[string]T + sortedKeys []string // lazily computed; nil after add or remove } // newFeatureSet creates a new featureSet for features of type T. @@ -37,6 +39,7 @@ func (s *featureSet[T]) add(fs ...T) { for _, f := range fs { s.features[s.uniqueID(f)] = f } + s.sortedKeys = nil } // remove removes all features with the given uids from the set if present, @@ -50,6 +53,9 @@ func (s *featureSet[T]) remove(uids ...string) bool { delete(s.features, uid) } } + if changed { + s.sortedKeys = nil + } return changed } @@ -63,11 +69,41 @@ func (s *featureSet[T]) get(uid string) (T, bool) { // all returns an iterator over of all the features in the set // sorted by unique ID. func (s *featureSet[T]) all() iter.Seq[T] { + s.sortKeys() return func(yield func(T) bool) { - for _, f := range util.Sorted(s.features) { - if !yield(f) { - return - } + s.yieldFrom(0, yield) + } +} + +// above returns an iterator over features in the set whose unique IDs are +// greater than `uid`, in ascending ID order. +func (s *featureSet[T]) above(uid string) iter.Seq[T] { + s.sortKeys() + index, found := slices.BinarySearch(s.sortedKeys, uid) + if found { + index++ + } + return func(yield func(T) bool) { + s.yieldFrom(index, yield) + } +} + +// sortKeys is a helper that maintains a sorted list of feature IDs. It +// computes this list lazily upon its first call after a modification, or +// if it's nil. +func (s *featureSet[T]) sortKeys() { + if s.sortedKeys != nil { + return + } + s.sortedKeys = slices.Sorted(maps.Keys(s.features)) +} + +// yieldFrom is a helper that iterates over the features in the set, +// starting at the given index, and calls the yield function for each one. +func (s *featureSet[T]) yieldFrom(index int, yield func(T) bool) { + for i := index; i < len(s.sortedKeys); i++ { + if !yield(s.features[s.sortedKeys[i]]) { + return } } } diff --git a/mcp/features_test.go b/mcp/features_test.go new file mode 100644 index 0000000..478457d --- /dev/null +++ b/mcp/features_test.go @@ -0,0 +1,75 @@ +// 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 ( + "context" + "slices" + "testing" + + "github.com/google/go-cmp/cmp" + "github.com/google/go-cmp/cmp/cmpopts" + "github.com/modelcontextprotocol/go-sdk/jsonschema" +) + +type SayHiParams struct { + Name string `json:"name"` +} + +func SayHi(ctx context.Context, cc *ServerSession, params *SayHiParams) ([]*Content, error) { + return []*Content{ + NewTextContent("Hi " + params.Name), + }, nil +} + +func TestFeatureSetOrder(t *testing.T) { + toolA := NewTool("apple", "apple tool", SayHi).Tool + toolB := NewTool("banana", "banana tool", SayHi).Tool + toolC := NewTool("cherry", "cherry tool", SayHi).Tool + + testCases := []struct { + tools []*Tool + want []*Tool + }{ + {[]*Tool{toolA, toolB, toolC}, []*Tool{toolA, toolB, toolC}}, + {[]*Tool{toolB, toolC, toolA}, []*Tool{toolA, toolB, toolC}}, + {[]*Tool{toolA, toolC}, []*Tool{toolA, toolC}}, + {[]*Tool{toolA, toolA, toolA}, []*Tool{toolA}}, + {[]*Tool{}, nil}, + } + for _, tc := range testCases { + fs := newFeatureSet(func(t *Tool) string { return t.Name }) + fs.add(tc.tools...) + got := slices.Collect(fs.all()) + if diff := cmp.Diff(got, tc.want, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { + t.Errorf("expected %v, got %v, (-want +got):\n%s", tc.want, got, diff) + } + } +} + +func TestFeatureSetAbove(t *testing.T) { + toolA := NewTool("apple", "apple tool", SayHi).Tool + toolB := NewTool("banana", "banana tool", SayHi).Tool + toolC := NewTool("cherry", "cherry tool", SayHi).Tool + + testCases := []struct { + tools []*Tool + above string + want []*Tool + }{ + {[]*Tool{toolA, toolB, toolC}, "apple", []*Tool{toolB, toolC}}, + {[]*Tool{toolA, toolB, toolC}, "banana", []*Tool{toolC}}, + {[]*Tool{toolA, toolB, toolC}, "cherry", nil}, + } + for _, tc := range testCases { + fs := newFeatureSet(func(t *Tool) string { return t.Name }) + fs.add(tc.tools...) + got := slices.Collect(fs.above(tc.above)) + if diff := cmp.Diff(got, tc.want, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { + t.Errorf("expected %v, got %v, (-want +got):\n%s", tc.want, got, diff) + + } + } +}