diff --git a/mcp/client_test.go b/mcp/client_test.go index 2f7abc0..fc37c3e 100644 --- a/mcp/client_test.go +++ b/mcp/client_test.go @@ -402,6 +402,26 @@ func TestClientCapabilities(t *testing.T) { RootsV2: &RootCapabilities{ListChanged: true}, }, }, + { + name: "extensions preserved", + configureClient: func(s *Client) {}, + clientOpts: func() ClientOptions { + caps := &ClientCapabilities{ + RootsV2: &RootCapabilities{ListChanged: true}, + } + caps.AddExtension("io.example/ext1", map[string]any{"key": "value"}) + caps.AddExtension("io.example/ext2", nil) + return ClientOptions{Capabilities: caps} + }(), + wantCapabilities: &ClientCapabilities{ + Extensions: map[string]any{ + "io.example/ext1": map[string]any{"key": "value"}, + "io.example/ext2": map[string]any{}, + }, + Roots: RootCapabilities{ListChanged: true}, + RootsV2: &RootCapabilities{ListChanged: true}, + }, + }, } for _, tc := range testCases { @@ -469,6 +489,23 @@ func TestClientCapabilitiesOverWire(t *testing.T) { }, wantCapabilities: &ClientCapabilities{}, }, + { + name: "Extensions over wire", + clientOpts: func() *ClientOptions { + caps := &ClientCapabilities{ + RootsV2: &RootCapabilities{ListChanged: true}, + } + caps.AddExtension("io.example/ext", map[string]any{"key": "value"}) + return &ClientOptions{Capabilities: caps} + }(), + wantCapabilities: &ClientCapabilities{ + Extensions: map[string]any{ + "io.example/ext": map[string]any{"key": "value"}, + }, + Roots: RootCapabilities{ListChanged: true}, + RootsV2: &RootCapabilities{ListChanged: true}, + }, + }, } for _, tc := range testCases { diff --git a/mcp/protocol.go b/mcp/protocol.go index 6e4a190..9bbaa0b 100644 --- a/mcp/protocol.go +++ b/mcp/protocol.go @@ -13,6 +13,7 @@ package mcp import ( "encoding/json" "fmt" + "maps" ) // Optional annotations for the client. The client can use annotations to inform @@ -192,7 +193,14 @@ type ClientCapabilities struct { // [ClientCapabilities.clone]. // Experimental reports non-standard capabilities that the client supports. + // The caller should not modify the map after assigning it. Experimental map[string]any `json:"experimental,omitempty"` + // Extensions reports extensions that the client supports. + // Keys are extension identifiers in "{vendor-prefix}/{extension-name}" format. + // Values are per-extension settings objects; use [ClientCapabilities.AddExtension] + // to ensure nil settings are normalized to empty objects. + // The caller should not modify the map or its values after assigning it. + Extensions map[string]any `json:"extensions,omitempty"` // Roots describes the client's support for roots. // // Deprecated: use RootsV2. As described in #607, Roots should have been a @@ -211,9 +219,26 @@ type ClientCapabilities struct { Elicitation *ElicitationCapabilities `json:"elicitation,omitempty"` } -// clone returns a deep copy of the ClientCapabilities. +// AddExtension adds an extension with the given name and settings. +// If settings is nil, an empty map is used to ensure valid JSON serialization +// (the spec requires an object, not null). +// The settings map should not be modified after the call. +func (c *ClientCapabilities) AddExtension(name string, settings map[string]any) { + if c.Extensions == nil { + c.Extensions = make(map[string]any) + } + if settings == nil { + settings = map[string]any{} + } + c.Extensions[name] = settings +} + +// clone returns a copy of the ClientCapabilities. +// Values in the Extensions and Experimental maps are shallow-copied. func (c *ClientCapabilities) clone() *ClientCapabilities { cp := *c + cp.Experimental = maps.Clone(c.Experimental) + cp.Extensions = maps.Clone(c.Extensions) cp.RootsV2 = shallowClone(c.RootsV2) if c.Sampling != nil { x := *c.Sampling @@ -1516,7 +1541,14 @@ type ServerCapabilities struct { // [ServerCapabilities.clone]. // Experimental reports non-standard capabilities that the server supports. + // The caller should not modify the map after assigning it. Experimental map[string]any `json:"experimental,omitempty"` + // Extensions reports extensions that the server supports. + // Keys are extension identifiers in "{vendor-prefix}/{extension-name}" format. + // Values are per-extension settings objects; use [ServerCapabilities.AddExtension] + // to ensure nil settings are normalized to empty objects. + // The caller should not modify the map or its values after assigning it. + Extensions map[string]any `json:"extensions,omitempty"` // Completions is present if the server supports argument autocompletion // suggestions. Completions *CompletionCapabilities `json:"completions,omitempty"` @@ -1530,9 +1562,26 @@ type ServerCapabilities struct { Tools *ToolCapabilities `json:"tools,omitempty"` } -// clone returns a deep copy of the ServerCapabilities. +// AddExtension adds an extension with the given name and settings. +// If settings is nil, an empty map is used to ensure valid JSON serialization +// (the spec requires an object, not null). +// The settings map should not be modified after the call. +func (c *ServerCapabilities) AddExtension(name string, settings map[string]any) { + if c.Extensions == nil { + c.Extensions = make(map[string]any) + } + if settings == nil { + settings = map[string]any{} + } + c.Extensions[name] = settings +} + +// clone returns a copy of the ServerCapabilities. +// Values in the Extensions and Experimental maps are shallow-copied. func (c *ServerCapabilities) clone() *ServerCapabilities { cp := *c + cp.Experimental = maps.Clone(c.Experimental) + cp.Extensions = maps.Clone(c.Extensions) cp.Completions = shallowClone(c.Completions) cp.Logging = shallowClone(c.Logging) cp.Prompts = shallowClone(c.Prompts) diff --git a/mcp/server_test.go b/mcp/server_test.go index 0f86f76..227e7be 100644 --- a/mcp/server_test.go +++ b/mcp/server_test.go @@ -441,6 +441,25 @@ func TestServerCapabilities(t *testing.T) { Logging: &LoggingCapabilities{}, }, }, + { + name: "extensions preserved", + configureServer: func(s *Server) {}, + serverOpts: func() ServerOptions { + caps := &ServerCapabilities{ + Logging: &LoggingCapabilities{}, + } + caps.AddExtension("io.example/ext1", map[string]any{"key": "value"}) + caps.AddExtension("io.example/ext2", nil) + return ServerOptions{Capabilities: caps} + }(), + wantCapabilities: &ServerCapabilities{ + Extensions: map[string]any{ + "io.example/ext1": map[string]any{"key": "value"}, + "io.example/ext2": map[string]any{}, + }, + Logging: &LoggingCapabilities{}, + }, + }, } for _, tc := range testCases { @@ -927,6 +946,23 @@ func TestServerCapabilitiesOverWire(t *testing.T) { Tools: &ToolCapabilities{ListChanged: true}, }, }, + { + name: "Extensions over wire", + serverOpts: func() *ServerOptions { + caps := &ServerCapabilities{ + Logging: &LoggingCapabilities{}, + } + caps.AddExtension("io.example/ext", map[string]any{"key": "value"}) + return &ServerOptions{Capabilities: caps} + }(), + configureServer: func(s *Server) {}, + wantCapabilities: &ServerCapabilities{ + Extensions: map[string]any{ + "io.example/ext": map[string]any{"key": "value"}, + }, + Logging: &LoggingCapabilities{}, + }, + }, } for _, tc := range testCases {