Skip to content

Commit c5fa4ea

Browse files
dugenkui03dugenkui03
committed
optimize capability and notification
1 parent 37ac814 commit c5fa4ea

File tree

5 files changed

+86
-23
lines changed

5 files changed

+86
-23
lines changed

‎client/client.go

Lines changed: 29 additions & 9 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,7 +41,7 @@ type Client struct {
3141
// if err != nil {
3242
// log.Fatalf("Failed to create client: %v", err)
3343
// }
34-
func NewClient(transport transport.Interface) *Client {
44+
func NewClient(transport transport.Interface, options ...ClientOption) *Client {
3545
return &Client{
3646
transport: transport,
3747
}
@@ -115,7 +125,7 @@ func (c *Client) Initialize(
115125
params := struct {
116126
ProtocolVersion string `json:"protocolVersion"`
117127
ClientInfo mcp.Implementation `json:"clientInfo"`
118-
Capabilities mcp.ClientCapabilities `json:"capabilities"`
128+
Capabilities mcp.ClientCapabilities `json:"serverCapabilities"`
119129
}{
120130
ProtocolVersion: request.Params.ProtocolVersion,
121131
ClientInfo: request.Params.ClientInfo,
@@ -132,8 +142,8 @@ func (c *Client) Initialize(
132142
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
133143
}
134144

135-
// Store capabilities
136-
c.capabilities = result.Capabilities
145+
// Store serverCapabilities
146+
c.serverCapabilities = result.Capabilities
137147

138148
// Send initialized notification
139149
notification := mcp.JSONRPCNotification{
@@ -406,3 +416,13 @@ func listByPage[T any](
406416
func (c *Client) GetTransport() transport.Interface {
407417
return c.transport
408418
}
419+
420+
// GetServerCapabilities returns the server capabilities.
421+
func (c *Client) GetServerCapabilities() mcp.ServerCapabilities {
422+
return c.serverCapabilities
423+
}
424+
425+
// GetClientCapabilities returns the client capabilities.
426+
func (c *Client) GetClientCapabilities() mcp.ClientCapabilities {
427+
return c.clientCapabilities
428+
}

‎client/sse_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ import (
1111
)
1212

1313
func TestSSEMCPClient(t *testing.T) {
14-
// Create MCP server with capabilities
14+
// Create MCP server with serverCapabilities
1515
mcpServer := server.NewMCPServer(
1616
"test-server",
1717
"1.0.0",

‎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://spec.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
@@ -417,6 +417,12 @@ func (s *MCPServer) AddResource(
417417
resource: resource,
418418
handler: handler,
419419
}
420+
421+
// When the list of available resources changes, servers that declared the listChanged capability SHOULD send a notification
422+
if s.capabilities.resources.listChanged {
423+
// Send notification to all initialized sessions
424+
s.sendNotificationToAllClients(mcp.MethodNotificationResourcesListChanged, nil)
425+
}
420426
}
421427

422428
// AddResourceTemplate registers a new resource template and its handler
@@ -436,6 +442,12 @@ func (s *MCPServer) AddResourceTemplate(
436442
template: template,
437443
handler: handler,
438444
}
445+
446+
// When the list of available resources changes, servers that declared the listChanged capability SHOULD send a notification
447+
if s.capabilities.resources.listChanged {
448+
// Send notification to all initialized sessions
449+
s.sendNotificationToAllClients(mcp.MethodNotificationResourcesListChanged, nil)
450+
}
439451
}
440452

441453
// AddPrompt registers a new prompt handler with the given name
@@ -450,6 +462,12 @@ func (s *MCPServer) AddPrompt(prompt mcp.Prompt, handler PromptHandlerFunc) {
450462
defer s.promptsMu.Unlock()
451463
s.prompts[prompt.Name] = prompt
452464
s.promptHandlers[prompt.Name] = handler
465+
466+
// When the list of available resources changes, servers that declared the listChanged capability SHOULD send a notification.
467+
if s.capabilities.prompts.listChanged {
468+
// Send notification to all initialized sessions
469+
s.sendNotificationToAllClients(mcp.MethodNotificationPromptsListChanged, nil)
470+
}
453471
}
454472

455473
// AddTool registers a new tool and its handler
@@ -471,8 +489,11 @@ func (s *MCPServer) AddTools(tools ...ServerTool) {
471489
}
472490
s.toolsMu.Unlock()
473491

474-
// Send notification to all initialized sessions
475-
s.sendNotificationToAllClients("notifications/tools/list_changed", nil)
492+
// When the list of available tools changes, servers that declared the listChanged capability SHOULD send a notification.
493+
if s.capabilities.tools.listChanged {
494+
// Send notification to all initialized sessions
495+
s.sendNotificationToAllClients(mcp.MethodNotificationToolsListChanged, nil)
496+
}
476497
}
477498

478499
// SetTools replaces all existing tools with the provided list
@@ -491,8 +512,11 @@ func (s *MCPServer) DeleteTools(names ...string) {
491512
}
492513
s.toolsMu.Unlock()
493514

494-
// Send notification to all initialized sessions
495-
s.sendNotificationToAllClients("notifications/tools/list_changed", nil)
515+
// When the list of available tools changes, servers that declared the listChanged capability SHOULD send a notification.
516+
if s.capabilities.tools.listChanged {
517+
// Send notification to all initialized sessions
518+
s.sendNotificationToAllClients(mcp.MethodNotificationToolsListChanged, nil)
519+
}
496520
}
497521

498522
// AddNotificationHandler registers a new handler for incoming notifications
@@ -723,7 +747,7 @@ func (s *MCPServer) handleReadResource(
723747

724748
return nil, &requestError{
725749
id: id,
726-
code: mcp.INVALID_PARAMS,
750+
code: mcp.RESOURCE_NOT_FOUND,
727751
err: fmt.Errorf("handler not found for resource URI '%s': %w", request.Params.URI, ErrResourceNotFound),
728752
}
729753
}

‎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,
@@ -861,7 +861,7 @@ func TestMCPServer_HandleUndefinedHandlers(t *testing.T) {
861861
"uri": "undefined-resource"
862862
}
863863
}`,
864-
expectedErr: mcp.INVALID_PARAMS,
864+
expectedErr: mcp.RESOURCE_NOT_FOUND,
865865
validateCallbacks: func(t *testing.T, err error, beforeResults beforeResult) {
866866
assert.Equal(t, mcp.MethodResourcesRead, beforeResults.method)
867867
assert.True(t, errors.Is(err, ErrResourceNotFound))

0 commit comments

Comments
 (0)