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"},