modelcontextprotocol/go-sdk
Watch
1
0
Fork
You've already forked go-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

internal/mcp: rename ServerConnection->ServerSession

Browse source 
Bring the Server implementation up to date with the spec by renaming
ServerConnection to ServerSession. Also update variable names, which
were in various states of disarray from the multiple renamings of this
type.

Change-Id: I4aec883db6d0500fb91e6daa51b702614877cf8a
Reviewed-on: https://go-review.googlesource.com/c/tools/+/672655
Auto-Submit: Robert Findley <rfindley@google.com>
TryBot-Bypass: Robert Findley <rfindley@google.com>
Commit-Queue: Robert Findley <rfindley@google.com>
Reviewed-by: Jonathan Amsterdam <jba@google.com>
This commit is contained in:
Rob Findley 2025-05-14 14:00:02 +00:00 • committed by Gopher Robot
commit ac828d5a82
14 changed files with 74 additions and 75 deletions
Download patch file Download diff file Expand all files Collapse all files

2
README.md
View file

@ -59,7 +59,7 @@ type HiParams struct {
Name string `json:"name"`
}
func SayHi(ctx context.Context, cc *mcp.ServerConnection, params *HiParams) ([]mcp.Content, error) {
func SayHi(ctx context.Context, cc *mcp.ServerSession, params *HiParams) ([]mcp.Content, error) {
return []mcp.Content{
mcp.TextContent{Text: "Hi " + params.Name},
}, nil

4
examples/hello/main.go
View file

@ -20,13 +20,13 @@ type HiParams struct {
Name string `json:"name"`
}
func SayHi(ctx context.Context, cc *mcp.ServerConnection, params *HiParams) ([]*mcp.Content, error) {
func SayHi(ctx context.Context, cc *mcp.ServerSession, params *HiParams) ([]*mcp.Content, error) {
return []*mcp.Content{
mcp.NewTextContent("Hi " + params.Name),
}, nil
}
func PromptHi(ctx context.Context, cc *mcp.ServerConnection, params *HiParams) (*mcp.GetPromptResult, error) {
func PromptHi(ctx context.Context, cc *mcp.ServerSession, params *HiParams) (*mcp.GetPromptResult, error) {
return &mcp.GetPromptResult{
Description: "Code review prompt",
Messages: []*mcp.PromptMessage{

2
examples/sse/main.go
View file

@ -19,7 +19,7 @@ type SayHiParams struct {
Name string `json:"name" mcp:"the name to say hi to"`
}
func SayHi(ctx context.Context, cc *mcp.ServerConnection, params *SayHiParams) ([]*mcp.Content, error) {
func SayHi(ctx context.Context, cc *mcp.ServerSession, params *SayHiParams) ([]*mcp.Content, error) {
return []*mcp.Content{
mcp.NewTextContent("Hi " + params.Name),
}, nil

8
mcp/mcp.go
View file

@ -10,9 +10,9 @@
// To get started, create either a [Client] or [Server], and connect it to a
// peer using a [Transport]. The diagram below illustrates how this works:
//
// Client Server
// ⇅ (jsonrpc2) ⇅
// Client Transport ⇄ Server Transport ⇄ ServerConnection
// Client Server
// ⇅ (jsonrpc2) ⇅
// Client Transport ⇄ Server Transport ⇄ ServerSession
//
// A [Client] is an MCP client, which can be configured with various client
// capabilities. Clients may be connected to a [Server] instance
@ -21,7 +21,7 @@
// Similarly, a [Server] is an MCP server, which can be configured with various
// server capabilities. Servers may be connected to one or more [Client]
// instances using the [Server.Connect] method, which creates a
// [ServerConnection].
// [ServerSession].
//
// A [Transport] connects a bidirectional [Stream] of jsonrpc2 messages. In
// practice, transports in the MCP spec are are either client transports or

32
mcp/mcp_test.go
View file

@ -24,7 +24,7 @@ type hiParams struct {
Name string
}
func sayHi(ctx context.Context, cc *ServerConnection, v hiParams) ([]*Content, error) {
func sayHi(ctx context.Context, cc *ServerSession, v hiParams) ([]*Content, error) {
if err := cc.Ping(ctx, nil); err != nil {
return nil, fmt.Errorf("ping failed: %v", err)
}
@ -43,13 +43,13 @@ func TestEndToEnd(t *testing.T) {
// The 'fail' tool returns this error.
failure := errors.New("mcp failure")
s.AddTools(
NewTool("fail", "just fail", func(context.Context, *ServerConnection, struct{}) ([]*Content, error) {
NewTool("fail", "just fail", func(context.Context, *ServerSession, struct{}) ([]*Content, error) {
return nil, failure
}),
)
s.AddPrompts(
NewPrompt("code_review", "do a code review", func(_ context.Context, _ *ServerConnection, params struct{ Code string }) (*GetPromptResult, error) {
NewPrompt("code_review", "do a code review", func(_ context.Context, _ *ServerSession, params struct{ Code string }) (*GetPromptResult, error) {
return &GetPromptResult{
Description: "Code review prompt",
Messages: []*PromptMessage{
@ -57,17 +57,17 @@ func TestEndToEnd(t *testing.T) {
},
}, nil
}),
NewPrompt("fail", "", func(_ context.Context, _ *ServerConnection, params struct{}) (*GetPromptResult, error) {
NewPrompt("fail", "", func(_ context.Context, _ *ServerSession, params struct{}) (*GetPromptResult, error) {
return nil, failure
}),
)
// Connect the server.
sc, err := s.Connect(ctx, st, nil)
ss, err := s.Connect(ctx, st, nil)
if err != nil {
t.Fatal(err)
}
if got := slices.Collect(s.Clients()); len(got) != 1 {
if got := slices.Collect(s.Sessions()); len(got) != 1 {
t.Errorf("after connection, Clients() has length %d, want 1", len(got))
}
@ -75,7 +75,7 @@ func TestEndToEnd(t *testing.T) {
var clientWG sync.WaitGroup
clientWG.Add(1)
go func() {
if err := sc.Wait(); err != nil {
if err := ss.Wait(); err != nil {
t.Errorf("server failed: %v", err)
}
clientWG.Done()
@ -198,7 +198,7 @@ func TestEndToEnd(t *testing.T) {
URI: "file:///nonexistent.txt",
}
readHandler := func(_ context.Context, _ *ServerConnection, p *ReadResourceParams) (*ReadResourceResult, error) {
readHandler := func(_ context.Context, _ *ServerSession, p *ReadResourceParams) (*ReadResourceResult, error) {
if p.URI == "file:///file1.txt" {
return &ReadResourceResult{
Contents: &ResourceContents{
@ -249,9 +249,9 @@ func TestEndToEnd(t *testing.T) {
}
})
t.Run("roots", func(t *testing.T) {
// Take the server's first ServerConnection.
var sc *ServerConnection
for sc = range s.Clients() {
// Take the server's first ServerSession.
var sc *ServerSession
for sc = range s.Sessions() {
break
}
@ -272,7 +272,7 @@ func TestEndToEnd(t *testing.T) {
// After disconnecting, neither client nor server should have any
// connections.
for range s.Clients() {
for range s.Sessions() {
t.Errorf("unexpected client after disconnection")
}
}
@ -282,7 +282,7 @@ func TestEndToEnd(t *testing.T) {
//
// The caller should cancel either the client connection or server connection
// when the connections are no longer needed.
func basicConnection(t *testing.T, tools ...*ServerTool) (*ServerConnection, *Client) {
func basicConnection(t *testing.T, tools ...*ServerTool) (*ServerSession, *Client) {
t.Helper()
ctx := context.Background()
@ -292,7 +292,7 @@ func basicConnection(t *testing.T, tools ...*ServerTool) (*ServerConnection, *Cl
// The 'greet' tool says hi.
s.AddTools(tools...)
cc, err := s.Connect(ctx, st, nil)
ss, err := s.Connect(ctx, st, nil)
if err != nil {
t.Fatal(err)
}
@ -301,7 +301,7 @@ func basicConnection(t *testing.T, tools ...*ServerTool) (*ServerConnection, *Cl
if err := c.Start(ctx); err != nil {
t.Fatal(err)
}
return cc, c
return ss, c
}
func TestServerClosing(t *testing.T) {
@ -370,7 +370,7 @@ func TestCancellation(t *testing.T) {
cancelled = make(chan struct{}, 1) // don't block the request
)
slowRequest := func(ctx context.Context, cc *ServerConnection, v struct{}) ([]*Content, error) {
slowRequest := func(ctx context.Context, cc *ServerSession, v struct{}) ([]*Content, error) {
start <- struct{}{}
select {
case <-ctx.Done():

6
mcp/prompt.go
View file

@ -16,7 +16,7 @@ import (
)
// A PromptHandler handles a call to prompts/get.
type PromptHandler func(context.Context, *ServerConnection, map[string]string) (*GetPromptResult, error)
type PromptHandler func(context.Context, *ServerSession, map[string]string) (*GetPromptResult, error)
// A Prompt is a prompt definition bound to a prompt handler.
type ServerPrompt struct {
@ -32,7 +32,7 @@ type ServerPrompt struct {
// of type string or *string. The argument names for the resulting prompt
// definition correspond to the JSON names of the request fields, and any
// fields that are not marked "omitempty" are considered required.
func NewPrompt[TReq any](name, description string, handler func(context.Context, *ServerConnection, TReq) (*GetPromptResult, error), opts ...PromptOption) *ServerPrompt {
func NewPrompt[TReq any](name, description string, handler func(context.Context, *ServerSession, TReq) (*GetPromptResult, error), opts ...PromptOption) *ServerPrompt {
schema, err := jsonschema.For[TReq]()
if err != nil {
panic(err)
@ -60,7 +60,7 @@ func NewPrompt[TReq any](name, description string, handler func(context.Context,
Required: required[name],
})
}
prompt.Handler = func(ctx context.Context, cc *ServerConnection, args map[string]string) (*GetPromptResult, error) {
prompt.Handler = func(ctx context.Context, cc *ServerSession, args map[string]string) (*GetPromptResult, error) {
// For simplicity, just marshal and unmarshal the arguments.
// This could be avoided in the future.
rawArgs, err := json.Marshal(args)

2
mcp/prompt_test.go
View file

@ -13,7 +13,7 @@ import (
)
// testPromptHandler is used for type inference in TestNewPrompt.
func testPromptHandler[T any](context.Context, *mcp.ServerConnection, T) (*mcp.GetPromptResult, error) {
func testPromptHandler[T any](context.Context, *mcp.ServerSession, T) (*mcp.GetPromptResult, error) {
panic("not implemented")
}

65
mcp/server.go
View file

@ -30,7 +30,7 @@ type Server struct {
prompts *featureSet[*ServerPrompt]
tools *featureSet[*ServerTool]
resources *featureSet[*ServerResource]
conns []*ServerConnection
conns []*ServerSession
}
// ServerOptions is used to configure behavior of the server.
@ -121,7 +121,7 @@ const codeResourceNotFound = -31002
// A ResourceHandler is a function that reads a resource.
// If it cannot find the resource, it should return the result of calling [ResourceNotFoundError].
type ResourceHandler func(context.Context, *ServerConnection, *ReadResourceParams) (*ReadResourceResult, error)
type ResourceHandler func(context.Context, *ServerSession, *ReadResourceParams) (*ReadResourceResult, error)
// A ServerResource associates a Resource with its handler.
type ServerResource struct {
@ -157,16 +157,15 @@ func (s *Server) RemoveResources(uris ...string) {
s.resources.remove(uris...)
}
// Clients returns an iterator that yields the current set of client
// connections.
func (s *Server) Clients() iter.Seq[*ServerConnection] {
// Sessions returns an iterator that yields the current set of server sessions.
func (s *Server) Sessions() iter.Seq[*ServerSession] {
s.mu.Lock()
clients := slices.Clone(s.conns)
s.mu.Unlock()
return slices.Values(clients)
}
func (s *Server) listPrompts(_ context.Context, _ *ServerConnection, params *ListPromptsParams) (*ListPromptsResult, error) {
func (s *Server) listPrompts(_ context.Context, _ *ServerSession, params *ListPromptsParams) (*ListPromptsResult, error) {
s.mu.Lock()
defer s.mu.Unlock()
res := new(ListPromptsResult)
@ -176,7 +175,7 @@ func (s *Server) listPrompts(_ context.Context, _ *ServerConnection, params *Lis
return res, nil
}
func (s *Server) getPrompt(ctx context.Context, cc *ServerConnection, params *GetPromptParams) (*GetPromptResult, error) {
func (s *Server) getPrompt(ctx context.Context, cc *ServerSession, params *GetPromptParams) (*GetPromptResult, error) {
s.mu.Lock()
prompt, ok := s.prompts.get(params.Name)
s.mu.Unlock()
@ -187,7 +186,7 @@ func (s *Server) getPrompt(ctx context.Context, cc *ServerConnection, params *Ge
return prompt.Handler(ctx, cc, params.Arguments)
}
func (s *Server) listTools(_ context.Context, _ *ServerConnection, params *ListToolsParams) (*ListToolsResult, error) {
func (s *Server) listTools(_ context.Context, _ *ServerSession, params *ListToolsParams) (*ListToolsResult, error) {
s.mu.Lock()
defer s.mu.Unlock()
res := new(ListToolsResult)
@ -197,7 +196,7 @@ func (s *Server) listTools(_ context.Context, _ *ServerConnection, params *ListT
return res, nil
}
func (s *Server) callTool(ctx context.Context, cc *ServerConnection, params *CallToolParams) (*CallToolResult, error) {
func (s *Server) callTool(ctx context.Context, cc *ServerSession, params *CallToolParams) (*CallToolResult, error) {
s.mu.Lock()
tool, ok := s.tools.get(params.Name)
s.mu.Unlock()
@ -207,7 +206,7 @@ func (s *Server) callTool(ctx context.Context, cc *ServerConnection, params *Cal
return tool.Handler(ctx, cc, params)
}
func (s *Server) listResources(_ context.Context, _ *ServerConnection, params *ListResourcesParams) (*ListResourcesResult, error) {
func (s *Server) listResources(_ context.Context, _ *ServerSession, params *ListResourcesParams) (*ListResourcesResult, error) {
s.mu.Lock()
defer s.mu.Unlock()
res := new(ListResourcesResult)
@ -217,7 +216,7 @@ func (s *Server) listResources(_ context.Context, _ *ServerConnection, params *L
return res, nil
}
func (s *Server) readResource(ctx context.Context, ss *ServerConnection, params *ReadResourceParams) (*ReadResourceResult, error) {
func (s *Server) readResource(ctx context.Context, ss *ServerSession, params *ReadResourceParams) (*ReadResourceResult, error) {
uri := params.URI
// Look up the resource URI in the list we have.
// This is a security check as well as an information lookup.
@ -250,29 +249,29 @@ func (s *Server) readResource(ctx context.Context, ss *ServerConnection, params
//
// Run blocks until the client terminates the connection.
func (s *Server) Run(ctx context.Context, t Transport, opts *ConnectionOptions) error {
cc, err := s.Connect(ctx, t, opts)
ss, err := s.Connect(ctx, t, opts)
if err != nil {
return err
}
return cc.Wait()
return ss.Wait()
}
// bind implements the binder[*ServerConnection] interface, so that Servers can
// bind implements the binder[*ServerSession] interface, so that Servers can
// be connected using [connect].
func (s *Server) bind(conn *jsonrpc2.Connection) *ServerConnection {
cc := &ServerConnection{conn: conn, server: s}
func (s *Server) bind(conn *jsonrpc2.Connection) *ServerSession {
cc := &ServerSession{conn: conn, server: s}
s.mu.Lock()
s.conns = append(s.conns, cc)
s.mu.Unlock()
return cc
}
// disconnect implements the binder[*ServerConnection] interface, so that
// disconnect implements the binder[*ServerSession] interface, so that
// Servers can be connected using [connect].
func (s *Server) disconnect(cc *ServerConnection) {
func (s *Server) disconnect(cc *ServerSession) {
s.mu.Lock()
defer s.mu.Unlock()
s.conns = slices.DeleteFunc(s.conns, func(cc2 *ServerConnection) bool {
s.conns = slices.DeleteFunc(s.conns, func(cc2 *ServerSession) bool {
return cc2 == cc
})
}
@ -283,17 +282,17 @@ func (s *Server) disconnect(cc *ServerConnection) {
// It returns a connection object that may be used to terminate the connection
// (with [Connection.Close]), or await client termination (with
// [Connection.Wait]).
func (s *Server) Connect(ctx context.Context, t Transport, opts *ConnectionOptions) (*ServerConnection, error) {
func (s *Server) Connect(ctx context.Context, t Transport, opts *ConnectionOptions) (*ServerSession, error) {
return connect(ctx, t, opts, s)
}
// A ServerConnection is a connection from a single MCP client. Its methods can
// be used to send requests or notifications to the client. Create a connection
// by calling [Server.Connect].
// A ServerSession is a logical connection from a single MCP client. Its
// methods can be used to send requests or notifications to the client. Create
// a session by calling [Server.Connect].
//
// Call [ServerConnection.Close] to close the connection, or await client
// termination with [ServerConnection.Wait].
type ServerConnection struct {
// Call [ServerSession.Close] to close the connection, or await client
// termination with [ServerSession.Wait].
type ServerSession struct {
server *Server
conn *jsonrpc2.Connection
@ -303,15 +302,15 @@ type ServerConnection struct {
}
// Ping makes an MCP "ping" request to the client.
func (cc *ServerConnection) Ping(ctx context.Context, _ *PingParams) error {
func (cc *ServerSession) Ping(ctx context.Context, _ *PingParams) error {
return call(ctx, cc.conn, "ping", nil, nil)
}
func (cc *ServerConnection) ListRoots(ctx context.Context, params *ListRootsParams) (*ListRootsResult, error) {
func (cc *ServerSession) ListRoots(ctx context.Context, params *ListRootsParams) (*ListRootsResult, error) {
return standardCall[ListRootsResult](ctx, cc.conn, "roots/list", params)
}
func (cc *ServerConnection) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) {
func (cc *ServerSession) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) {
cc.mu.Lock()
initialized := cc.initialized
cc.mu.Unlock()
@ -362,7 +361,7 @@ func (cc *ServerConnection) handle(ctx context.Context, req *jsonrpc2.Request) (
return nil, jsonrpc2.ErrNotHandled
}
func (cc *ServerConnection) initialize(ctx context.Context, _ *ServerConnection, params *initializeParams) (*initializeResult, error) {
func (cc *ServerSession) initialize(ctx context.Context, _ *ServerSession, params *initializeParams) (*initializeResult, error) {
cc.mu.Lock()
cc.initializeParams = params
cc.mu.Unlock()
@ -400,12 +399,12 @@ func (cc *ServerConnection) initialize(ctx context.Context, _ *ServerConnection,
// Close performs a graceful shutdown of the connection, preventing new
// requests from being handled, and waiting for ongoing requests to return.
// Close then terminates the connection.
func (cc *ServerConnection) Close() error {
func (cc *ServerSession) Close() error {
return cc.conn.Close()
}
// Wait waits for the connection to be closed by the client.
func (cc *ServerConnection) Wait() error {
func (cc *ServerSession) Wait() error {
return cc.conn.Wait()
}
@ -413,7 +412,7 @@ func (cc *ServerConnection) Wait() error {
//
// Importantly, it returns nil if the handler returned an error, which is a
// requirement of the jsonrpc2 package.
func dispatch[TParams, TResult any](ctx context.Context, conn *ServerConnection, req *jsonrpc2.Request, f func(context.Context, *ServerConnection, TParams) (TResult, error)) (any, error) {
func dispatch[TParams, TResult any](ctx context.Context, conn *ServerSession, req *jsonrpc2.Request, f func(context.Context, *ServerSession, TParams) (TResult, error)) (any, error) {
var params TParams
if err := json.Unmarshal(req.Params, &params); err != nil {
return nil, err

6
mcp/server_example_test.go
View file

@ -16,7 +16,7 @@ type SayHiParams struct {
Name string `json:"name" mcp:"the name to say hi to"`
}
func SayHi(ctx context.Context, cc *mcp.ServerConnection, params *SayHiParams) ([]*mcp.Content, error) {
func SayHi(ctx context.Context, cc *mcp.ServerSession, params *SayHiParams) ([]*mcp.Content, error) {
return []*mcp.Content{
mcp.NewTextContent("Hi " + params.Name),
}, nil
@ -29,7 +29,7 @@ func ExampleServer() {
server := mcp.NewServer("greeter", "v0.0.1", nil)
server.AddTools(mcp.NewTool("greet", "say hi", SayHi))
clientConnection, err := server.Connect(ctx, serverTransport, nil)
serverSession, err := server.Connect(ctx, serverTransport, nil)
if err != nil {
log.Fatal(err)
}
@ -46,7 +46,7 @@ func ExampleServer() {
fmt.Println(res.Content[0].Text)
client.Close()
clientConnection.Wait()
serverSession.Wait()
// Output: Hi user
}

8
mcp/sse.go
View file

@ -64,7 +64,7 @@ func writeEvent(w io.Writer, evt event) (int, error) {
// https://modelcontextprotocol.io/specification/2024-11-05/basic/transports
type SSEHandler struct {
getServer func(request *http.Request) *Server
onConnection func(*ServerConnection) // for testing; must not block
onConnection func(*ServerSession) // for testing; must not block
mu sync.Mutex
sessions map[string]*SSEServerTransport
@ -229,15 +229,15 @@ func (h *SSEHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
// TODO(hxjiang): getServer returns nil will panic.
server := h.getServer(req)
cc, err := server.Connect(req.Context(), transport, nil)
ss, err := server.Connect(req.Context(), transport, nil)
if err != nil {
http.Error(w, "connection failed", http.StatusInternalServerError)
return
}
if h.onConnection != nil {
h.onConnection(cc)
h.onConnection(ss)
}
defer cc.Close() // close the transport when the GET exits
defer ss.Close() // close the transport when the GET exits
select {
case <-req.Context().Done():

2
mcp/sse_example_test.go
View file

@ -18,7 +18,7 @@ type AddParams struct {
X, Y int
}
func Add(ctx context.Context, cc *mcp.ServerConnection, params *AddParams) ([]*mcp.Content, error) {
func Add(ctx context.Context, cc *mcp.ServerSession, params *AddParams) ([]*mcp.Content, error) {
return []*mcp.Content{
mcp.NewTextContent(fmt.Sprintf("%d", params.X+params.Y)),
}, nil

4
mcp/sse_test.go
View file

@ -23,8 +23,8 @@ func TestSSEServer(t *testing.T) {
sseHandler := NewSSEHandler(func(*http.Request) *Server { return server })
conns := make(chan *ServerConnection, 1)
sseHandler.onConnection = func(cc *ServerConnection) {
conns := make(chan *ServerSession, 1)
sseHandler.onConnection = func(cc *ServerSession) {
select {
case conns <- cc:
default:

6
mcp/tool.go
View file

@ -13,7 +13,7 @@ import (
)
// A ToolHandler handles a call to tools/call.
type ToolHandler func(context.Context, *ServerConnection, *CallToolParams) (*CallToolResult, error)
type ToolHandler func(context.Context, *ServerSession, *CallToolParams) (*CallToolResult, error)
// A Tool is a tool definition that is bound to a tool handler.
type ServerTool struct {
@ -34,12 +34,12 @@ type ServerTool struct {
//
// TODO: just have the handler return a CallToolResult: returning []Content is
// going to be inconsistent with other server features.
func NewTool[TReq any](name, description string, handler func(context.Context, *ServerConnection, TReq) ([]*Content, error), opts ...ToolOption) *ServerTool {
func NewTool[TReq any](name, description string, handler func(context.Context, *ServerSession, TReq) ([]*Content, error), opts ...ToolOption) *ServerTool {
schema, err := jsonschema.For[TReq]()
if err != nil {
panic(err)
}
wrapped := func(ctx context.Context, cc *ServerConnection, params *CallToolParams) (*CallToolResult, error) {
wrapped := func(ctx context.Context, cc *ServerSession, params *CallToolParams) (*CallToolResult, error) {
var v TReq
if err := unmarshalSchema(params.Arguments, schema, &v); err != nil {
return nil, err

2
mcp/tool_test.go
View file

@ -15,7 +15,7 @@ import (
)
// testToolHandler is used for type inference in TestNewTool.
func testToolHandler[T any](context.Context, *mcp.ServerConnection, T) ([]*mcp.Content, error) {
func testToolHandler[T any](context.Context, *mcp.ServerSession, T) ([]*mcp.Content, error) {
panic("not implemented")
}