37ae90648b
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>
73 lines
2 KiB
Go
73 lines
2 KiB
Go
// 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
|
|
}
|
|
}
|
|
}
|
|
}
|