From 294e884f2656ee8ea21f7819f73af5c7ee9a224b Mon Sep 17 00:00:00 2001 From: Shashank Pachava Date: Sat, 2 Aug 2025 15:13:19 -0400 Subject: [PATCH] mcp: export server capabilities and add ClientSession InitializeResult accessor Export ServerCapabilities and related subtypes and update protocol, server, and tests to use exported types. Add InitializeResult method on ClientSession for accessing the initialize result from sessions. --- mcp/client.go | 2 ++ mcp/protocol.go | 24 +++++++------- mcp/server.go | 14 ++++---- mcp/server_test.go | 72 +++++++++++++++++++++--------------------- mcp/streamable_test.go | 16 +++++----- 5 files changed, 65 insertions(+), 63 deletions(-) diff --git a/mcp/client.go b/mcp/client.go index 9e3f3935..cb442a71 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -163,6 +163,8 @@ type ClientSession struct { mcpConn Connection } +func (cs *ClientSession) InitializeResult() *InitializeResult { return cs.initializeResult } + func (cs *ClientSession) setConn(c Connection) { cs.mcpConn = c } diff --git a/mcp/protocol.go b/mcp/protocol.go index 3ca6cb5e..a0bae2e2 100644 --- a/mcp/protocol.go +++ b/mcp/protocol.go @@ -324,7 +324,7 @@ type InitializeResult struct { // This property is reserved by the protocol to allow clients and servers to // attach additional metadata to their responses. Meta `json:"_meta,omitempty"` - Capabilities *serverCapabilities `json:"capabilities"` + Capabilities *ServerCapabilities `json:"capabilities"` // Instructions describing how to use the server and its features. // // This can be used by clients to improve the LLM's understanding of available @@ -925,19 +925,19 @@ type Implementation struct { } // Present if the server supports argument autocompletion suggestions. -type completionCapabilities struct{} +type CompletionCapabilities struct{} // Present if the server supports sending log messages to the client. -type loggingCapabilities struct{} +type LoggingCapabilities struct{} // Present if the server offers any prompt templates. -type promptCapabilities struct { +type PromptCapabilities struct { // Whether this server supports notifications for changes to the prompt list. ListChanged bool `json:"listChanged,omitempty"` } // Present if the server offers any resources to read. -type resourceCapabilities struct { +type ResourceCapabilities struct { // Whether this server supports notifications for changes to the resource list. ListChanged bool `json:"listChanged,omitempty"` // Whether this server supports subscribing to resource updates. @@ -947,23 +947,23 @@ type resourceCapabilities struct { // Capabilities that a server may support. Known capabilities are defined here, // in this schema, but this is not a closed set: any server can define its own, // additional capabilities. -type serverCapabilities struct { +type ServerCapabilities struct { // Present if the server supports argument autocompletion suggestions. - Completions *completionCapabilities `json:"completions,omitempty"` + Completions *CompletionCapabilities `json:"completions,omitempty"` // Experimental, non-standard capabilities that the server supports. Experimental map[string]struct{} `json:"experimental,omitempty"` // Present if the server supports sending log messages to the client. - Logging *loggingCapabilities `json:"logging,omitempty"` + Logging *LoggingCapabilities `json:"logging,omitempty"` // Present if the server offers any prompt templates. - Prompts *promptCapabilities `json:"prompts,omitempty"` + Prompts *PromptCapabilities `json:"prompts,omitempty"` // Present if the server offers any resources to read. - Resources *resourceCapabilities `json:"resources,omitempty"` + Resources *ResourceCapabilities `json:"resources,omitempty"` // Present if the server offers any tools to call. - Tools *toolCapabilities `json:"tools,omitempty"` + Tools *ToolCapabilities `json:"tools,omitempty"` } // Present if the server offers any tools to call. -type toolCapabilities struct { +type ToolCapabilities struct { // Whether this server supports notifications for changes to the tool list. ListChanged bool `json:"listChanged,omitempty"` } diff --git a/mcp/server.go b/mcp/server.go index b96edbe8..2b40bd63 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -233,23 +233,23 @@ func (s *Server) RemoveResourceTemplates(uriTemplates ...string) { func() bool { return s.resourceTemplates.remove(uriTemplates...) }) } -func (s *Server) capabilities() *serverCapabilities { +func (s *Server) capabilities() *ServerCapabilities { s.mu.Lock() defer s.mu.Unlock() - caps := &serverCapabilities{ + caps := &ServerCapabilities{ // TODO(samthanawalla): check for completionHandler before advertising capability. - Completions: &completionCapabilities{}, - Logging: &loggingCapabilities{}, + Completions: &CompletionCapabilities{}, + Logging: &LoggingCapabilities{}, } if s.opts.HasTools || s.tools.len() > 0 { - caps.Tools = &toolCapabilities{ListChanged: true} + caps.Tools = &ToolCapabilities{ListChanged: true} } if s.opts.HasPrompts || s.prompts.len() > 0 { - caps.Prompts = &promptCapabilities{ListChanged: true} + caps.Prompts = &PromptCapabilities{ListChanged: true} } if s.opts.HasResources || s.resources.len() > 0 || s.resourceTemplates.len() > 0 { - caps.Resources = &resourceCapabilities{ListChanged: true} + caps.Resources = &ResourceCapabilities{ListChanged: true} if s.opts.SubscribeHandler != nil { caps.Resources.Subscribe = true } diff --git a/mcp/server_test.go b/mcp/server_test.go index 79c607e9..60523321 100644 --- a/mcp/server_test.go +++ b/mcp/server_test.go @@ -236,14 +236,14 @@ func TestServerCapabilities(t *testing.T) { name string configureServer func(s *Server) serverOpts ServerOptions - wantCapabilities *serverCapabilities + wantCapabilities *ServerCapabilities }{ { name: "No capabilities", configureServer: func(s *Server) {}, - wantCapabilities: &serverCapabilities{ - Completions: &completionCapabilities{}, - Logging: &loggingCapabilities{}, + wantCapabilities: &ServerCapabilities{ + Completions: &CompletionCapabilities{}, + Logging: &LoggingCapabilities{}, }, }, { @@ -251,10 +251,10 @@ func TestServerCapabilities(t *testing.T) { configureServer: func(s *Server) { s.AddPrompt(&Prompt{Name: "p"}, nil) }, - wantCapabilities: &serverCapabilities{ - Completions: &completionCapabilities{}, - Logging: &loggingCapabilities{}, - Prompts: &promptCapabilities{ListChanged: true}, + wantCapabilities: &ServerCapabilities{ + Completions: &CompletionCapabilities{}, + Logging: &LoggingCapabilities{}, + Prompts: &PromptCapabilities{ListChanged: true}, }, }, { @@ -262,10 +262,10 @@ func TestServerCapabilities(t *testing.T) { configureServer: func(s *Server) { s.AddResource(&Resource{URI: "file:///r"}, nil) }, - wantCapabilities: &serverCapabilities{ - Completions: &completionCapabilities{}, - Logging: &loggingCapabilities{}, - Resources: &resourceCapabilities{ListChanged: true}, + wantCapabilities: &ServerCapabilities{ + Completions: &CompletionCapabilities{}, + Logging: &LoggingCapabilities{}, + Resources: &ResourceCapabilities{ListChanged: true}, }, }, { @@ -273,10 +273,10 @@ func TestServerCapabilities(t *testing.T) { configureServer: func(s *Server) { s.AddResourceTemplate(&ResourceTemplate{URITemplate: "file:///rt"}, nil) }, - wantCapabilities: &serverCapabilities{ - Completions: &completionCapabilities{}, - Logging: &loggingCapabilities{}, - Resources: &resourceCapabilities{ListChanged: true}, + wantCapabilities: &ServerCapabilities{ + Completions: &CompletionCapabilities{}, + Logging: &LoggingCapabilities{}, + Resources: &ResourceCapabilities{ListChanged: true}, }, }, { @@ -292,10 +292,10 @@ func TestServerCapabilities(t *testing.T) { return nil }, }, - wantCapabilities: &serverCapabilities{ - Completions: &completionCapabilities{}, - Logging: &loggingCapabilities{}, - Resources: &resourceCapabilities{ListChanged: true, Subscribe: true}, + wantCapabilities: &ServerCapabilities{ + Completions: &CompletionCapabilities{}, + Logging: &LoggingCapabilities{}, + Resources: &ResourceCapabilities{ListChanged: true, Subscribe: true}, }, }, { @@ -303,10 +303,10 @@ func TestServerCapabilities(t *testing.T) { configureServer: func(s *Server) { s.AddTool(tool, nil) }, - wantCapabilities: &serverCapabilities{ - Completions: &completionCapabilities{}, - Logging: &loggingCapabilities{}, - Tools: &toolCapabilities{ListChanged: true}, + wantCapabilities: &ServerCapabilities{ + Completions: &CompletionCapabilities{}, + Logging: &LoggingCapabilities{}, + Tools: &ToolCapabilities{ListChanged: true}, }, }, { @@ -325,12 +325,12 @@ func TestServerCapabilities(t *testing.T) { return nil }, }, - wantCapabilities: &serverCapabilities{ - Completions: &completionCapabilities{}, - Logging: &loggingCapabilities{}, - Prompts: &promptCapabilities{ListChanged: true}, - Resources: &resourceCapabilities{ListChanged: true, Subscribe: true}, - Tools: &toolCapabilities{ListChanged: true}, + wantCapabilities: &ServerCapabilities{ + Completions: &CompletionCapabilities{}, + Logging: &LoggingCapabilities{}, + Prompts: &PromptCapabilities{ListChanged: true}, + Resources: &ResourceCapabilities{ListChanged: true, Subscribe: true}, + Tools: &ToolCapabilities{ListChanged: true}, }, }, { @@ -341,12 +341,12 @@ func TestServerCapabilities(t *testing.T) { HasResources: true, HasTools: true, }, - wantCapabilities: &serverCapabilities{ - Completions: &completionCapabilities{}, - Logging: &loggingCapabilities{}, - Prompts: &promptCapabilities{ListChanged: true}, - Resources: &resourceCapabilities{ListChanged: true}, - Tools: &toolCapabilities{ListChanged: true}, + wantCapabilities: &ServerCapabilities{ + Completions: &CompletionCapabilities{}, + Logging: &LoggingCapabilities{}, + Prompts: &PromptCapabilities{ListChanged: true}, + Resources: &ResourceCapabilities{ListChanged: true}, + Tools: &ToolCapabilities{ListChanged: true}, }, }, } diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go index 81af2760..dab8f40f 100644 --- a/mcp/streamable_test.go +++ b/mcp/streamable_test.go @@ -313,10 +313,10 @@ func TestStreamableServerTransport(t *testing.T) { // Predefined steps, to avoid repetition below. initReq := req(1, methodInitialize, &InitializeParams{}) initResp := resp(1, &InitializeResult{ - Capabilities: &serverCapabilities{ - Completions: &completionCapabilities{}, - Logging: &loggingCapabilities{}, - Tools: &toolCapabilities{ListChanged: true}, + Capabilities: &ServerCapabilities{ + Completions: &CompletionCapabilities{}, + Logging: &LoggingCapabilities{}, + Tools: &ToolCapabilities{ListChanged: true}, }, ProtocolVersion: latestProtocolVersion, ServerInfo: &Implementation{Name: "testServer", Version: "v1.0.0"}, @@ -731,10 +731,10 @@ func TestStreamableClientTransportApplicationJSON(t *testing.T) { } } initResult := &InitializeResult{ - Capabilities: &serverCapabilities{ - Completions: &completionCapabilities{}, - Logging: &loggingCapabilities{}, - Tools: &toolCapabilities{ListChanged: true}, + Capabilities: &ServerCapabilities{ + Completions: &CompletionCapabilities{}, + Logging: &LoggingCapabilities{}, + Tools: &ToolCapabilities{ListChanged: true}, }, ProtocolVersion: latestProtocolVersion, ServerInfo: &Implementation{Name: "testServer", Version: "v1.0.0"},