optimize capability and notification (mark3labs#184)

Author: dugenkui03

client/client.go

@@ -16,11 +16,21 @@ import (
 type Client struct {
 	transport transport.Interface
 
-	initialized   bool
-	notifications []func(mcp.JSONRPCNotification)
-	notifyMu      sync.RWMutex
-	requestID     atomic.Int64
-	capabilities  mcp.ServerCapabilities
+	initialized        bool
+	notifications      []func(mcp.JSONRPCNotification)
+	notifyMu           sync.RWMutex
+	requestID          atomic.Int64
+	clientCapabilities mcp.ClientCapabilities
+	serverCapabilities mcp.ServerCapabilities
+}
+
+type ClientOption func(*Client)
+
+// WithClientCapabilities sets the client capabilities for the client.
+func WithClientCapabilities(capabilities mcp.ClientCapabilities) ClientOption {
+	return func(c *Client) {
+		c.clientCapabilities = capabilities
+	}
 }
 
 // NewClient creates a new MCP client with the given transport.
@@ -31,10 +41,16 @@ type Client struct {
 //	if err != nil {
 //	    log.Fatalf("Failed to create client: %v", err)
 //	}
-func NewClient(transport transport.Interface) *Client {
-	return &Client{
+func NewClient(transport transport.Interface, options ...ClientOption) *Client {
+	client := &Client{
 		transport: transport,
 	}
+
+	for _, opt := range options {
+		opt(client)
+	}
+
+	return client
 }
 
 // Start initiates the connection to the server.
@@ -115,7 +131,7 @@ func (c *Client) Initialize(
 	params := struct {
 		ProtocolVersion string                 `json:"protocolVersion"`
 		ClientInfo      mcp.Implementation     `json:"clientInfo"`
-		Capabilities    mcp.ClientCapabilities `json:"capabilities"`
+		Capabilities    mcp.ClientCapabilities `json:"serverCapabilities"`
 	}{
 		ProtocolVersion: request.Params.ProtocolVersion,
 		ClientInfo:      request.Params.ClientInfo,
@@ -132,8 +148,8 @@ func (c *Client) Initialize(
 		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
 	}
 
-	// Store capabilities
-	c.capabilities = result.Capabilities
+	// Store serverCapabilities
+	c.serverCapabilities = result.Capabilities
 
 	// Send initialized notification
 	notification := mcp.JSONRPCNotification{
@@ -406,3 +422,13 @@ func listByPage[T any](
 func (c *Client) GetTransport() transport.Interface {
 	return c.transport
 }
+
+// GetServerCapabilities returns the server capabilities.
+func (c *Client) GetServerCapabilities() mcp.ServerCapabilities {
+	return c.serverCapabilities
+}
+
+// GetClientCapabilities returns the client capabilities.
+func (c *Client) GetClientCapabilities() mcp.ClientCapabilities {
+	return c.clientCapabilities
+}

mcp/types.go

@@ -46,6 +46,20 @@ const (
 	// Invokes a specific tool with provided parameters.
 	// https://modelcontextprotocol.io/specification/2024-11-05/server/tools/
 	MethodToolsCall MCPMethod = "tools/call"
+
+	// Notifies when the list of available resources changes.
+	// https://modelcontextprotocol.io/specification/2025-03-26/server/resources#list-changed-notification
+	MethodNotificationResourcesListChanged = "notifications/resources/list_changed"
+
+	MethodNotificationResourceUpdated = "notifications/resources/updated"
+
+	// Notifies when the list of available prompt templates changes.
+	// https://modelcontextprotocol.io/specification/2025-03-26/server/prompts#list-changed-notification
+	MethodNotificationPromptsListChanged = "notifications/prompts/list_changed"
+
+	// Notifies when the list of available tools changes.
+	// https://spec.modelcontextprotocol.io/specification/2024-11-05/server/tools/list_changed/
+	MethodNotificationToolsListChanged = "notifications/tools/list_changed"
 )
 
 type URITemplate struct {
@@ -226,6 +240,11 @@ const (
 	INTERNAL_ERROR   = -32603
 )
 
+// MCP error codes
+const (
+	RESOURCE_NOT_FOUND = -32002
+)
+
 /* Empty result */
 
 // EmptyResult represents a response that indicates success but carries no data.

server/server.go

@@ -419,6 +419,12 @@ func (s *MCPServer) AddResource(
 		resource: resource,
 		handler:  handler,
 	}
+
+	// When the list of available resources changes, servers that declared the listChanged capability SHOULD send a notification
+	if s.capabilities.resources.listChanged {
+		// Send notification to all initialized sessions
+		s.SendNotificationToAllClients(mcp.MethodNotificationResourcesListChanged, nil)
+	}
 }
 
 // RemoveResource removes a resource from the server
@@ -450,6 +456,12 @@ func (s *MCPServer) AddResourceTemplate(
 		template: template,
 		handler:  handler,
 	}
+
+	// When the list of available resources changes, servers that declared the listChanged capability SHOULD send a notification
+	if s.capabilities.resources.listChanged {
+		// Send notification to all initialized sessions
+		s.SendNotificationToAllClients(mcp.MethodNotificationResourcesListChanged, nil)
+	}
 }
 
 // AddPrompt registers a new prompt handler with the given name
@@ -464,6 +476,12 @@ func (s *MCPServer) AddPrompt(prompt mcp.Prompt, handler PromptHandlerFunc) {
 	defer s.promptsMu.Unlock()
 	s.prompts[prompt.Name] = prompt
 	s.promptHandlers[prompt.Name] = handler
+
+	// When the list of available resources changes, servers that declared the listChanged capability SHOULD send a notification.
+	if s.capabilities.prompts.listChanged {
+		// Send notification to all initialized sessions
+		s.SendNotificationToAllClients(mcp.MethodNotificationPromptsListChanged, nil)
+	}
 }
 
 // AddTool registers a new tool and its handler
@@ -485,8 +503,11 @@ func (s *MCPServer) AddTools(tools ...ServerTool) {
 	}
 	s.toolsMu.Unlock()
 
-	// Send notification to all initialized sessions
-	s.SendNotificationToAllClients("notifications/tools/list_changed", nil)
+	// When the list of available tools changes, servers that declared the listChanged capability SHOULD send a notification.
+	if s.capabilities.tools.listChanged {
+		// Send notification to all initialized sessions
+		s.SendNotificationToAllClients(mcp.MethodNotificationToolsListChanged, nil)
+	}
 }
 
 // SetTools replaces all existing tools with the provided list
@@ -505,8 +526,11 @@ func (s *MCPServer) DeleteTools(names ...string) {
 	}
 	s.toolsMu.Unlock()
 
-	// Send notification to all initialized sessions
-	s.SendNotificationToAllClients("notifications/tools/list_changed", nil)
+	// When the list of available tools changes, servers that declared the listChanged capability SHOULD send a notification.
+	if s.capabilities.tools.listChanged {
+		// Send notification to all initialized sessions
+		s.SendNotificationToAllClients(mcp.MethodNotificationToolsListChanged, nil)
+	}
 }
 
 // AddNotificationHandler registers a new handler for incoming notifications
@@ -737,7 +761,7 @@ func (s *MCPServer) handleReadResource(
 
 	return nil, &requestError{
 		id:   id,
-		code: mcp.INVALID_PARAMS,
+		code: mcp.RESOURCE_NOT_FOUND,
 		err:  fmt.Errorf("handler not found for resource URI '%s': %w", request.Params.URI, ErrResourceNotFound),
 	}
 }

server/server_test.go

@@ -199,7 +199,7 @@ func TestMCPServer_Tools(t *testing.T) {
 			},
 			expectedNotifications: 1,
 			validate: func(t *testing.T, notifications []mcp.JSONRPCNotification, toolsList mcp.JSONRPCMessage) {
-				assert.Equal(t, "notifications/tools/list_changed", notifications[0].Method)
+				assert.Equal(t, mcp.MethodNotificationToolsListChanged, notifications[0].Method)
 				tools := toolsList.(mcp.JSONRPCResponse).Result.(mcp.ListToolsResult).Tools
 				assert.Len(t, tools, 2)
 				assert.Equal(t, "test-tool-1", tools[0].Name)
@@ -241,7 +241,7 @@ func TestMCPServer_Tools(t *testing.T) {
 			expectedNotifications: 5,
 			validate: func(t *testing.T, notifications []mcp.JSONRPCNotification, toolsList mcp.JSONRPCMessage) {
 				for _, notification := range notifications {
-					assert.Equal(t, "notifications/tools/list_changed", notification.Method)
+					assert.Equal(t, mcp.MethodNotificationToolsListChanged, notification.Method)
 				}
 				tools := toolsList.(mcp.JSONRPCResponse).Result.(mcp.ListToolsResult).Tools
 				assert.Len(t, tools, 2)
@@ -269,8 +269,8 @@ func TestMCPServer_Tools(t *testing.T) {
 			},
 			expectedNotifications: 2,
 			validate: func(t *testing.T, notifications []mcp.JSONRPCNotification, toolsList mcp.JSONRPCMessage) {
-				assert.Equal(t, "notifications/tools/list_changed", notifications[0].Method)
-				assert.Equal(t, "notifications/tools/list_changed", notifications[1].Method)
+				assert.Equal(t, mcp.MethodNotificationToolsListChanged, notifications[0].Method)
+				assert.Equal(t, mcp.MethodNotificationToolsListChanged, notifications[1].Method)
 				tools := toolsList.(mcp.JSONRPCResponse).Result.(mcp.ListToolsResult).Tools
 				assert.Len(t, tools, 2)
 				assert.Equal(t, "test-tool-1", tools[0].Name)
@@ -294,9 +294,9 @@ func TestMCPServer_Tools(t *testing.T) {
 			expectedNotifications: 2,
 			validate: func(t *testing.T, notifications []mcp.JSONRPCNotification, toolsList mcp.JSONRPCMessage) {
 				// One for SetTools
-				assert.Equal(t, "notifications/tools/list_changed", notifications[0].Method)
+				assert.Equal(t, mcp.MethodNotificationToolsListChanged, notifications[0].Method)
 				// One for DeleteTools
-				assert.Equal(t, "notifications/tools/list_changed", notifications[1].Method)
+				assert.Equal(t, mcp.MethodNotificationToolsListChanged, notifications[1].Method)
 
 				// Expect a successful response with an empty list of tools
 				resp, ok := toolsList.(mcp.JSONRPCResponse)
@@ -312,7 +312,7 @@ func TestMCPServer_Tools(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			ctx := context.Background()
-			server := NewMCPServer("test-server", "1.0.0")
+			server := NewMCPServer("test-server", "1.0.0", WithToolCapabilities(true))
 			_ = server.HandleMessage(ctx, []byte(`{
 				"jsonrpc": "2.0",
 				"id": 1,
@@ -929,7 +929,7 @@ func TestMCPServer_HandleUndefinedHandlers(t *testing.T) {
                         "uri": "undefined-resource"
                     }
                 }`,
-			expectedErr: mcp.INVALID_PARAMS,
+			expectedErr: mcp.RESOURCE_NOT_FOUND,
 			validateCallbacks: func(t *testing.T, err error, beforeResults beforeResult) {
 				assert.Equal(t, mcp.MethodResourcesRead, beforeResults.method)
 				assert.True(t, errors.Is(err, ErrResourceNotFound))