Skip to content

Commit 781b732

Browse files
dugenkui03dugenkui03
authored
optimize capability and notification (#184)
1 parent 6760d87 commit 781b732

File tree

4 files changed

+92
-23
lines changed

4 files changed

+92
-23
lines changed

‎client/client.go

Lines changed: 36 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,21 @@ import (
1616
type Client struct {
1717
transport transport.Interface
1818

19-
initialized bool
20-
notifications []func(mcp.JSONRPCNotification)
21-
notifyMu sync.RWMutex
22-
requestID atomic.Int64
23-
capabilities mcp.ServerCapabilities
19+
initialized bool
20+
notifications []func(mcp.JSONRPCNotification)
21+
notifyMu sync.RWMutex
22+
requestID atomic.Int64
23+
clientCapabilities mcp.ClientCapabilities
24+
serverCapabilities mcp.ServerCapabilities
25+
}
26+
27+
type ClientOption func(*Client)
28+
29+
// WithClientCapabilities sets the client capabilities for the client.
30+
func WithClientCapabilities(capabilities mcp.ClientCapabilities) ClientOption {
31+
return func(c *Client) {
32+
c.clientCapabilities = capabilities
33+
}
2434
}
2535

2636
// NewClient creates a new MCP client with the given transport.
@@ -31,10 +41,16 @@ type Client struct {
3141
// if err != nil {
3242
// log.Fatalf("Failed to create client: %v", err)
3343
// }
34-
func NewClient(transport transport.Interface) *Client {
35-
return &Client{
44+
func NewClient(transport transport.Interface, options ...ClientOption) *Client {
45+
client := &Client{
3646
transport: transport,
3747
}
48+
49+
for _, opt := range options {
50+
opt(client)
51+
}
52+
53+
return client
3854
}
3955

4056
// Start initiates the connection to the server.
@@ -115,7 +131,7 @@ func (c *Client) Initialize(
115131
params := struct {
116132
ProtocolVersion string `json:"protocolVersion"`
117133
ClientInfo mcp.Implementation `json:"clientInfo"`
118-
Capabilities mcp.ClientCapabilities `json:"capabilities"`
134+
Capabilities mcp.ClientCapabilities `json:"serverCapabilities"`
119135
}{
120136
ProtocolVersion: request.Params.ProtocolVersion,
121137
ClientInfo: request.Params.ClientInfo,
@@ -132,8 +148,8 @@ func (c *Client) Initialize(
132148
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
133149
}
134150

135-
// Store capabilities
136-
c.capabilities = result.Capabilities
151+
// Store serverCapabilities
152+
c.serverCapabilities = result.Capabilities
137153

138154
// Send initialized notification
139155
notification := mcp.JSONRPCNotification{
@@ -406,3 +422,13 @@ func listByPage[T any](
406422
func (c *Client) GetTransport() transport.Interface {
407423
return c.transport
408424
}
425+
426+
// GetServerCapabilities returns the server capabilities.
427+
func (c *Client) GetServerCapabilities() mcp.ServerCapabilities {
428+
return c.serverCapabilities
429+
}
430+
431+
// GetClientCapabilities returns the client capabilities.
432+
func (c *Client) GetClientCapabilities() mcp.ClientCapabilities {
433+
return c.clientCapabilities
434+
}

‎mcp/types.go

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,20 @@ const (
4646
// Invokes a specific tool with provided parameters.
4747
// https://modelcontextprotocol.io/specification/2024-11-05/server/tools/
4848
MethodToolsCall MCPMethod = "tools/call"
49+
50+
// Notifies when the list of available resources changes.
51+
// https://modelcontextprotocol.io/specification/2025-03-26/server/resources#list-changed-notification
52+
MethodNotificationResourcesListChanged = "notifications/resources/list_changed"
53+
54+
MethodNotificationResourceUpdated = "notifications/resources/updated"
55+
56+
// Notifies when the list of available prompt templates changes.
57+
// https://modelcontextprotocol.io/specification/2025-03-26/server/prompts#list-changed-notification
58+
MethodNotificationPromptsListChanged = "notifications/prompts/list_changed"
59+
60+
// Notifies when the list of available tools changes.
61+
// https://spec.modelcontextprotocol.io/specification/2024-11-05/server/tools/list_changed/
62+
MethodNotificationToolsListChanged = "notifications/tools/list_changed"
4963
)
5064

5165
type URITemplate struct {
@@ -226,6 +240,11 @@ const (
226240
INTERNAL_ERROR = -32603
227241
)
228242

243+
// MCP error codes
244+
const (
245+
RESOURCE_NOT_FOUND = -32002
246+
)
247+
229248
/* Empty result */
230249

231250
// EmptyResult represents a response that indicates success but carries no data.

‎server/server.go

Lines changed: 29 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -419,6 +419,12 @@ func (s *MCPServer) AddResource(
419419
resource: resource,
420420
handler: handler,
421421
}
422+
423+
// When the list of available resources changes, servers that declared the listChanged capability SHOULD send a notification
424+
if s.capabilities.resources.listChanged {
425+
// Send notification to all initialized sessions
426+
s.SendNotificationToAllClients(mcp.MethodNotificationResourcesListChanged, nil)
427+
}
422428
}
423429

424430
// RemoveResource removes a resource from the server
@@ -450,6 +456,12 @@ func (s *MCPServer) AddResourceTemplate(
450456
template: template,
451457
handler: handler,
452458
}
459+
460+
// When the list of available resources changes, servers that declared the listChanged capability SHOULD send a notification
461+
if s.capabilities.resources.listChanged {
462+
// Send notification to all initialized sessions
463+
s.SendNotificationToAllClients(mcp.MethodNotificationResourcesListChanged, nil)
464+
}
453465
}
454466

455467
// AddPrompt registers a new prompt handler with the given name
@@ -464,6 +476,12 @@ func (s *MCPServer) AddPrompt(prompt mcp.Prompt, handler PromptHandlerFunc) {
464476
defer s.promptsMu.Unlock()
465477
s.prompts[prompt.Name] = prompt
466478
s.promptHandlers[prompt.Name] = handler
479+
480+
// When the list of available resources changes, servers that declared the listChanged capability SHOULD send a notification.
481+
if s.capabilities.prompts.listChanged {
482+
// Send notification to all initialized sessions
483+
s.SendNotificationToAllClients(mcp.MethodNotificationPromptsListChanged, nil)
484+
}
467485
}
468486

469487
// AddTool registers a new tool and its handler
@@ -485,8 +503,11 @@ func (s *MCPServer) AddTools(tools ...ServerTool) {
485503
}
486504
s.toolsMu.Unlock()
487505

488-
// Send notification to all initialized sessions
489-
s.SendNotificationToAllClients("notifications/tools/list_changed", nil)
506+
// When the list of available tools changes, servers that declared the listChanged capability SHOULD send a notification.
507+
if s.capabilities.tools.listChanged {
508+
// Send notification to all initialized sessions
509+
s.SendNotificationToAllClients(mcp.MethodNotificationToolsListChanged, nil)
510+
}
490511
}
491512

492513
// SetTools replaces all existing tools with the provided list
@@ -505,8 +526,11 @@ func (s *MCPServer) DeleteTools(names ...string) {
505526
}
506527
s.toolsMu.Unlock()
507528

508-
// Send notification to all initialized sessions
509-
s.SendNotificationToAllClients("notifications/tools/list_changed", nil)
529+
// When the list of available tools changes, servers that declared the listChanged capability SHOULD send a notification.
530+
if s.capabilities.tools.listChanged {
531+
// Send notification to all initialized sessions
532+
s.SendNotificationToAllClients(mcp.MethodNotificationToolsListChanged, nil)
533+
}
510534
}
511535

512536
// AddNotificationHandler registers a new handler for incoming notifications
@@ -737,7 +761,7 @@ func (s *MCPServer) handleReadResource(
737761

738762
return nil, &requestError{
739763
id: id,
740-
code: mcp.INVALID_PARAMS,
764+
code: mcp.RESOURCE_NOT_FOUND,
741765
err: fmt.Errorf("handler not found for resource URI '%s': %w", request.Params.URI, ErrResourceNotFound),
742766
}
743767
}

‎server/server_test.go

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -199,7 +199,7 @@ func TestMCPServer_Tools(t *testing.T) {
199199
},
200200
expectedNotifications: 1,
201201
validate: func(t *testing.T, notifications []mcp.JSONRPCNotification, toolsList mcp.JSONRPCMessage) {
202-
assert.Equal(t, "notifications/tools/list_changed", notifications[0].Method)
202+
assert.Equal(t, mcp.MethodNotificationToolsListChanged, notifications[0].Method)
203203
tools := toolsList.(mcp.JSONRPCResponse).Result.(mcp.ListToolsResult).Tools
204204
assert.Len(t, tools, 2)
205205
assert.Equal(t, "test-tool-1", tools[0].Name)
@@ -241,7 +241,7 @@ func TestMCPServer_Tools(t *testing.T) {
241241
expectedNotifications: 5,
242242
validate: func(t *testing.T, notifications []mcp.JSONRPCNotification, toolsList mcp.JSONRPCMessage) {
243243
for _, notification := range notifications {
244-
assert.Equal(t, "notifications/tools/list_changed", notification.Method)
244+
assert.Equal(t, mcp.MethodNotificationToolsListChanged, notification.Method)
245245
}
246246
tools := toolsList.(mcp.JSONRPCResponse).Result.(mcp.ListToolsResult).Tools
247247
assert.Len(t, tools, 2)
@@ -269,8 +269,8 @@ func TestMCPServer_Tools(t *testing.T) {
269269
},
270270
expectedNotifications: 2,
271271
validate: func(t *testing.T, notifications []mcp.JSONRPCNotification, toolsList mcp.JSONRPCMessage) {
272-
assert.Equal(t, "notifications/tools/list_changed", notifications[0].Method)
273-
assert.Equal(t, "notifications/tools/list_changed", notifications[1].Method)
272+
assert.Equal(t, mcp.MethodNotificationToolsListChanged, notifications[0].Method)
273+
assert.Equal(t, mcp.MethodNotificationToolsListChanged, notifications[1].Method)
274274
tools := toolsList.(mcp.JSONRPCResponse).Result.(mcp.ListToolsResult).Tools
275275
assert.Len(t, tools, 2)
276276
assert.Equal(t, "test-tool-1", tools[0].Name)
@@ -294,9 +294,9 @@ func TestMCPServer_Tools(t *testing.T) {
294294
expectedNotifications: 2,
295295
validate: func(t *testing.T, notifications []mcp.JSONRPCNotification, toolsList mcp.JSONRPCMessage) {
296296
// One for SetTools
297-
assert.Equal(t, "notifications/tools/list_changed", notifications[0].Method)
297+
assert.Equal(t, mcp.MethodNotificationToolsListChanged, notifications[0].Method)
298298
// One for DeleteTools
299-
assert.Equal(t, "notifications/tools/list_changed", notifications[1].Method)
299+
assert.Equal(t, mcp.MethodNotificationToolsListChanged, notifications[1].Method)
300300

301301
// Expect a successful response with an empty list of tools
302302
resp, ok := toolsList.(mcp.JSONRPCResponse)
@@ -312,7 +312,7 @@ func TestMCPServer_Tools(t *testing.T) {
312312
for _, tt := range tests {
313313
t.Run(tt.name, func(t *testing.T) {
314314
ctx := context.Background()
315-
server := NewMCPServer("test-server", "1.0.0")
315+
server := NewMCPServer("test-server", "1.0.0", WithToolCapabilities(true))
316316
_ = server.HandleMessage(ctx, []byte(`{
317317
"jsonrpc": "2.0",
318318
"id": 1,
@@ -929,7 +929,7 @@ func TestMCPServer_HandleUndefinedHandlers(t *testing.T) {
929929
"uri": "undefined-resource"
930930
}
931931
}`,
932-
expectedErr: mcp.INVALID_PARAMS,
932+
expectedErr: mcp.RESOURCE_NOT_FOUND,
933933
validateCallbacks: func(t *testing.T, err error, beforeResults beforeResult) {
934934
assert.Equal(t, mcp.MethodResourcesRead, beforeResults.method)
935935
assert.True(t, errors.Is(err, ErrResourceNotFound))

0 commit comments

Comments
 (0)