Skip to content

Commit b4febf1

Browse files
authored
mcp: advertise a capability if and only if added (#93)
If tools are added to a Server when Run is called, it will send toolCapabilities to the client. Otherwise it won't. Ditto for prompts and resources. For #56.
1 parent bfa5e30 commit b4febf1

File tree

4 files changed

+118
-16
lines changed

4 files changed

+118
-16
lines changed

mcp/features.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,9 @@ func (s *featureSet[T]) get(uid string) (T, bool) {
6666
return t, ok
6767
}
6868

69+
// len returns the number of features in the set.
70+
func (s *featureSet[T]) len() int { return len(s.features) }
71+
6972
// all returns an iterator over of all the features in the set
7073
// sorted by unique ID.
7174
func (s *featureSet[T]) all() iter.Seq[T] {

mcp/server.go

Lines changed: 27 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -204,6 +204,26 @@ func (s *Server) RemoveResourceTemplates(uriTemplates ...string) {
204204
func() bool { return s.resourceTemplates.remove(uriTemplates...) })
205205
}
206206

207+
func (s *Server) capabilities() *serverCapabilities {
208+
s.mu.Lock()
209+
defer s.mu.Unlock()
210+
211+
caps := &serverCapabilities{
212+
Completions: &completionCapabilities{},
213+
Logging: &loggingCapabilities{},
214+
}
215+
if s.tools.len() > 0 {
216+
caps.Tools = &toolCapabilities{ListChanged: true}
217+
}
218+
if s.prompts.len() > 0 {
219+
caps.Prompts = &promptCapabilities{ListChanged: true}
220+
}
221+
if s.resources.len() > 0 || s.resourceTemplates.len() > 0 {
222+
caps.Resources = &resourceCapabilities{ListChanged: true}
223+
}
224+
return caps
225+
}
226+
207227
func (s *Server) complete(ctx context.Context, ss *ServerSession, params *CompleteParams) (Result, error) {
208228
if s.opts.CompletionHandler == nil {
209229
return nil, jsonrpc2.ErrMethodNotFound
@@ -407,6 +427,11 @@ func fileResourceHandler(dir string) ResourceHandler {
407427
//
408428
// Run blocks until the client terminates the connection or the provided
409429
// context is cancelled. If the context is cancelled, Run closes the connection.
430+
//
431+
// If tools have been added to the server before this call, then the server will
432+
// advertise the capability for tools, including the ability to send list-changed notifications.
433+
// If no tools have been added, the server will not have the tool capability.
434+
// The same goes for other features like prompts and resources.
410435
func (s *Server) Run(ctx context.Context, t Transport) error {
411436
ss, err := s.Connect(ctx, t)
412437
if err != nil {
@@ -659,20 +684,8 @@ func (ss *ServerSession) initialize(ctx context.Context, params *InitializeParam
659684
// TODO(rfindley): alter behavior when falling back to an older version:
660685
// reject unsupported features.
661686
ProtocolVersion: version,
662-
Capabilities: &serverCapabilities{
663-
Completions: &completionCapabilities{},
664-
Prompts: &promptCapabilities{
665-
ListChanged: true,
666-
},
667-
Tools: &toolCapabilities{
668-
ListChanged: true,
669-
},
670-
Resources: &resourceCapabilities{
671-
ListChanged: true,
672-
},
673-
Logging: &loggingCapabilities{},
674-
},
675-
Instructions: ss.server.opts.Instructions,
687+
Capabilities: ss.server.capabilities(),
688+
Instructions: ss.server.opts.Instructions,
676689
ServerInfo: &implementation{
677690
Name: ss.server.name,
678691
Version: ss.server.version,

mcp/server_test.go

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -227,3 +227,91 @@ func TestServerPaginateVariousPageSizes(t *testing.T) {
227227
}
228228
}
229229
}
230+
231+
func TestServerCapabilities(t *testing.T) {
232+
testCases := []struct {
233+
name string
234+
configureServer func(s *Server)
235+
wantCapabilities *serverCapabilities
236+
}{
237+
{
238+
name: "No capabilities",
239+
configureServer: func(s *Server) {},
240+
wantCapabilities: &serverCapabilities{
241+
Completions: &completionCapabilities{},
242+
Logging: &loggingCapabilities{},
243+
},
244+
},
245+
{
246+
name: "With prompts",
247+
configureServer: func(s *Server) {
248+
s.AddPrompt(&Prompt{Name: "p"}, nil)
249+
},
250+
wantCapabilities: &serverCapabilities{
251+
Completions: &completionCapabilities{},
252+
Logging: &loggingCapabilities{},
253+
Prompts: &promptCapabilities{ListChanged: true},
254+
},
255+
},
256+
{
257+
name: "With resources",
258+
configureServer: func(s *Server) {
259+
s.AddResource(&Resource{URI: "file:///r"}, nil)
260+
},
261+
wantCapabilities: &serverCapabilities{
262+
Completions: &completionCapabilities{},
263+
Logging: &loggingCapabilities{},
264+
Resources: &resourceCapabilities{ListChanged: true},
265+
},
266+
},
267+
{
268+
name: "With resource templates",
269+
configureServer: func(s *Server) {
270+
s.AddResourceTemplate(&ResourceTemplate{URITemplate: "file:///rt"}, nil)
271+
},
272+
wantCapabilities: &serverCapabilities{
273+
Completions: &completionCapabilities{},
274+
Logging: &loggingCapabilities{},
275+
Resources: &resourceCapabilities{ListChanged: true},
276+
},
277+
},
278+
{
279+
name: "With tools",
280+
configureServer: func(s *Server) {
281+
s.AddTool(&Tool{Name: "t"}, nil)
282+
},
283+
wantCapabilities: &serverCapabilities{
284+
Completions: &completionCapabilities{},
285+
Logging: &loggingCapabilities{},
286+
Tools: &toolCapabilities{ListChanged: true},
287+
},
288+
},
289+
{
290+
name: "With all capabilities",
291+
configureServer: func(s *Server) {
292+
s.AddPrompt(&Prompt{Name: "p"}, nil)
293+
s.AddResource(&Resource{URI: "file:///r"}, nil)
294+
s.AddResourceTemplate(&ResourceTemplate{URITemplate: "file:///rt"}, nil)
295+
s.AddTool(&Tool{Name: "t"}, nil)
296+
},
297+
wantCapabilities: &serverCapabilities{
298+
Completions: &completionCapabilities{},
299+
Logging: &loggingCapabilities{},
300+
Prompts: &promptCapabilities{ListChanged: true},
301+
Resources: &resourceCapabilities{ListChanged: true},
302+
Tools: &toolCapabilities{ListChanged: true},
303+
},
304+
},
305+
}
306+
307+
for _, tc := range testCases {
308+
t.Run(tc.name, func(t *testing.T) {
309+
server := NewServer("", "", nil)
310+
tc.configureServer(server)
311+
gotCapabilities := server.capabilities()
312+
if diff := cmp.Diff(tc.wantCapabilities, gotCapabilities); diff != "" {
313+
t.Errorf("capabilities() mismatch (-want +got):\n%s", diff)
314+
}
315+
})
316+
}
317+
}

mcp/streamable_test.go

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -159,8 +159,6 @@ func TestStreamableServerTransport(t *testing.T) {
159159
Capabilities: &serverCapabilities{
160160
Completions: &completionCapabilities{},
161161
Logging: &loggingCapabilities{},
162-
Prompts: &promptCapabilities{ListChanged: true},
163-
Resources: &resourceCapabilities{ListChanged: true},
164162
Tools: &toolCapabilities{ListChanged: true},
165163
},
166164
ProtocolVersion: latestProtocolVersion,

0 commit comments

Comments
 (0)