internal/mcp: keep an ordered list of feature keys

This change is necessary to support pagination which will be implemented
in subsequent CLs.

Change-Id: I8bc616a57463ea569dc7cadf1d4f06c69602ddd5
Reviewed-on: https://go-review.googlesource.com/c/tools/+/674135
Reviewed-by: Jonathan Amsterdam <jba@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

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) {

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
 		}
 	}
 }

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)
+
+		}
+	}
+}