From 73c35f786429b04542f5162915bf166a2e3cbc72 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. Co-authored-by: Matt Liegey --- mcp/client.go | 6 ++++ mcp/protocol.go | 24 ++++++++-------- mcp/server.go | 14 ++++----- mcp/server_test.go | 64 +++++++++++++++++++++--------------------- mcp/streamable_test.go | 13 +++++---- 5 files changed, 64 insertions(+), 57 deletions(-) diff --git a/mcp/client.go b/mcp/client.go index 88eea7da..65a7a954 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -177,6 +177,12 @@ type clientSessionState struct { InitializeResult *InitializeResult } +func (cs *ClientSession) InitializeResult() *InitializeResult { return cs.state.InitializeResult } + +func (cs *ClientSession) setConn(c Connection) { + cs.mcpConn = c +} + func (cs *ClientSession) ID() string { if c, ok := cs.mcpConn.(hasSessionID); ok { return c.SessionID() diff --git a/mcp/protocol.go b/mcp/protocol.go index d2d343b8..7fbfccf0 100644 --- a/mcp/protocol.go +++ b/mcp/protocol.go @@ -338,7 +338,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 @@ -971,19 +971,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. @@ -993,23 +993,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 ed4ec720..5bc626b3 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -242,27 +242,27 @@ 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{ - Logging: &loggingCapabilities{}, + caps := &ServerCapabilities{ + 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 } } if s.opts.CompletionHandler != nil { - caps.Completions = &completionCapabilities{} + caps.Completions = &CompletionCapabilities{} } return caps } diff --git a/mcp/server_test.go b/mcp/server_test.go index 202ab5d9..5482d51f 100644 --- a/mcp/server_test.go +++ b/mcp/server_test.go @@ -236,13 +236,13 @@ 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{ - Logging: &loggingCapabilities{}, + wantCapabilities: &ServerCapabilities{ + Logging: &LoggingCapabilities{}, }, }, { @@ -250,9 +250,9 @@ func TestServerCapabilities(t *testing.T) { configureServer: func(s *Server) { s.AddPrompt(&Prompt{Name: "p"}, nil) }, - wantCapabilities: &serverCapabilities{ - Logging: &loggingCapabilities{}, - Prompts: &promptCapabilities{ListChanged: true}, + wantCapabilities: &ServerCapabilities{ + Logging: &LoggingCapabilities{}, + Prompts: &PromptCapabilities{ListChanged: true}, }, }, { @@ -260,9 +260,9 @@ func TestServerCapabilities(t *testing.T) { configureServer: func(s *Server) { s.AddResource(&Resource{URI: "file:///r"}, nil) }, - wantCapabilities: &serverCapabilities{ - Logging: &loggingCapabilities{}, - Resources: &resourceCapabilities{ListChanged: true}, + wantCapabilities: &ServerCapabilities{ + Logging: &LoggingCapabilities{}, + Resources: &ResourceCapabilities{ListChanged: true}, }, }, { @@ -270,9 +270,9 @@ func TestServerCapabilities(t *testing.T) { configureServer: func(s *Server) { s.AddResourceTemplate(&ResourceTemplate{URITemplate: "file:///rt"}, nil) }, - wantCapabilities: &serverCapabilities{ - Logging: &loggingCapabilities{}, - Resources: &resourceCapabilities{ListChanged: true}, + wantCapabilities: &ServerCapabilities{ + Logging: &LoggingCapabilities{}, + Resources: &ResourceCapabilities{ListChanged: true}, }, }, { @@ -288,9 +288,9 @@ func TestServerCapabilities(t *testing.T) { return nil }, }, - wantCapabilities: &serverCapabilities{ - Logging: &loggingCapabilities{}, - Resources: &resourceCapabilities{ListChanged: true, Subscribe: true}, + wantCapabilities: &ServerCapabilities{ + Logging: &LoggingCapabilities{}, + Resources: &ResourceCapabilities{ListChanged: true, Subscribe: true}, }, }, { @@ -298,9 +298,9 @@ func TestServerCapabilities(t *testing.T) { configureServer: func(s *Server) { s.AddTool(tool, nil) }, - wantCapabilities: &serverCapabilities{ - Logging: &loggingCapabilities{}, - Tools: &toolCapabilities{ListChanged: true}, + wantCapabilities: &ServerCapabilities{ + Logging: &LoggingCapabilities{}, + Tools: &ToolCapabilities{ListChanged: true}, }, }, { @@ -311,9 +311,9 @@ func TestServerCapabilities(t *testing.T) { return nil, nil }, }, - wantCapabilities: &serverCapabilities{ - Logging: &loggingCapabilities{}, - Completions: &completionCapabilities{}, + wantCapabilities: &ServerCapabilities{ + Logging: &LoggingCapabilities{}, + Completions: &CompletionCapabilities{}, }, }, { @@ -335,12 +335,12 @@ func TestServerCapabilities(t *testing.T) { return nil, 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}, }, }, { @@ -351,11 +351,11 @@ func TestServerCapabilities(t *testing.T) { HasResources: true, HasTools: true, }, - wantCapabilities: &serverCapabilities{ - Logging: &loggingCapabilities{}, - Prompts: &promptCapabilities{ListChanged: true}, - Resources: &resourceCapabilities{ListChanged: true}, - Tools: &toolCapabilities{ListChanged: true}, + wantCapabilities: &ServerCapabilities{ + 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 11600fbc..3db1a060 100644 --- a/mcp/streamable_test.go +++ b/mcp/streamable_test.go @@ -409,9 +409,9 @@ func TestStreamableServerTransport(t *testing.T) { // Predefined steps, to avoid repetition below. initReq := req(1, methodInitialize, &InitializeParams{}) initResp := resp(1, &InitializeResult{ - Capabilities: &serverCapabilities{ - Logging: &loggingCapabilities{}, - Tools: &toolCapabilities{ListChanged: true}, + Capabilities: &ServerCapabilities{ + Logging: &LoggingCapabilities{}, + Tools: &ToolCapabilities{ListChanged: true}, }, ProtocolVersion: latestProtocolVersion, ServerInfo: &Implementation{Name: "testServer", Version: "v1.0.0"}, @@ -891,9 +891,10 @@ func TestStreamableClientTransportApplicationJSON(t *testing.T) { } } initResult := &InitializeResult{ - Capabilities: &serverCapabilities{ - 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"},