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 ClientConnection to ServerConnection

Browse source 
It's too confusing that you send server requests through the Client, and
send client requests through the ClientConnection. Rename
ClientConnection to ServerConnection, so that the types on the
server side of the control diagram start with 'Server'.

Change-Id: Idfd27af464a9296a2341cd23b2075cd9d0eb6b45
Reviewed-on: https://go-review.googlesource.com/c/tools/+/670538
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Auto-Submit: Robert Findley <rfindley@google.com>
Reviewed-by: Jonathan Amsterdam <jba@google.com>
This commit is contained in:
Rob Findley 2025-05-07 16:43:03 +00:00 • committed by Gopher Robot
commit 347a1a0078
13 changed files with 65 additions and 63 deletions
Download patch file Download diff file Expand all files Collapse all files

4
mcp/client.go
View file

@ -49,7 +49,7 @@ type ClientOptions struct {
ConnectionOptions
}
// bind implements the binder[*ServerConnection] interface, so that Clients can
// bind implements the binder[*Client] interface, so that Clients can
// be connected using [connect].
func (c *Client) bind(conn *jsonrpc2.Connection) *Client {
c.mu.Lock()
@ -58,7 +58,7 @@ func (c *Client) bind(conn *jsonrpc2.Connection) *Client {
return c
}
// disconnect implements the binder[*ServerConnection] interface, so that
// disconnect implements the binder[*Client] interface, so that
// Clients can be connected using [connect].
func (c *Client) disconnect(*Client) {
// Do nothing. In particular, do not set conn to nil: it needs to exist so it can

4
mcp/mcp.go
View file

@ -10,7 +10,7 @@
//
// Client Server
// ⇅ (jsonrpc2) ⇅
// Client Transport ⇄ Server Transport ⇄ ClientConnection
// Client Transport ⇄ Server Transport ⇄ ServerConnection
//
// A [Client] is an MCP client, which can be configured with various client
// capabilities. Clients may be connected to a [Server] instance
@ -19,7 +19,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
// [ClientConnection].
// [ServerConnection].
//
// A [Transport] connects a bidirectional [Stream] of jsonrpc2 messages. In
// practice, transports in the MCP spec are are either client transports or

12
mcp/mcp_test.go
View file

@ -24,7 +24,7 @@ type hiParams struct {
Name string
}
func sayHi(ctx context.Context, cc *ClientConnection, v hiParams) ([]Content, error) {
func sayHi(ctx context.Context, cc *ServerConnection, v hiParams) ([]Content, error) {
if err := cc.Ping(ctx); 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, *ClientConnection, struct{}) ([]Content, error) {
NewTool("fail", "just fail", func(context.Context, *ServerConnection, struct{}) ([]Content, error) {
return nil, failure
}),
)
s.AddPrompts(
NewPrompt("code_review", "do a code review", func(_ context.Context, _ *ClientConnection, params struct{ Code string }) (*protocol.GetPromptResult, error) {
NewPrompt("code_review", "do a code review", func(_ context.Context, _ *ServerConnection, params struct{ Code string }) (*protocol.GetPromptResult, error) {
return &protocol.GetPromptResult{
Description: "Code review prompt",
Messages: []protocol.PromptMessage{
@ -57,7 +57,7 @@ func TestEndToEnd(t *testing.T) {
},
}, nil
}),
NewPrompt("fail", "", func(_ context.Context, _ *ClientConnection, params struct{}) (*protocol.GetPromptResult, error) {
NewPrompt("fail", "", func(_ context.Context, _ *ServerConnection, params struct{}) (*protocol.GetPromptResult, error) {
return nil, failure
}),
)
@ -195,7 +195,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 ...*Tool) (*ClientConnection, *Client) {
func basicConnection(t *testing.T, tools ...*Tool) (*ServerConnection, *Client) {
t.Helper()
ctx := context.Background()
@ -283,7 +283,7 @@ func TestCancellation(t *testing.T) {
cancelled = make(chan struct{}, 1) // don't block the request
)
slowRequest := func(ctx context.Context, cc *ClientConnection, v struct{}) ([]Content, error) {
slowRequest := func(ctx context.Context, cc *ServerConnection, v struct{}) ([]Content, error) {
start <- struct{}{}
select {
case <-ctx.Done():

6
mcp/prompt.go
View file

@ -17,7 +17,7 @@ import (
)
// A PromptHandler handles a call to prompts/get.
type PromptHandler func(context.Context, *ClientConnection, map[string]string) (*protocol.GetPromptResult, error)
type PromptHandler func(context.Context, *ServerConnection, map[string]string) (*protocol.GetPromptResult, error)
// A Prompt is a prompt definition bound to a prompt handler.
type Prompt struct {
@ -33,7 +33,7 @@ type Prompt 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, *ClientConnection, TReq) (*protocol.GetPromptResult, error), opts ...PromptOption) *Prompt {
func NewPrompt[TReq any](name, description string, handler func(context.Context, *ServerConnection, TReq) (*protocol.GetPromptResult, error), opts ...PromptOption) *Prompt {
schema, err := jsonschema.For[TReq]()
if err != nil {
panic(err)
@ -61,7 +61,7 @@ func NewPrompt[TReq any](name, description string, handler func(context.Context,
Required: required[name],
})
}
prompt.Handler = func(ctx context.Context, cc *ClientConnection, args map[string]string) (*protocol.GetPromptResult, error) {
prompt.Handler = func(ctx context.Context, cc *ServerConnection, args map[string]string) (*protocol.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

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

60
mcp/server.go
View file

@ -28,7 +28,7 @@ type Server struct {
mu sync.Mutex
prompts []*Prompt
tools []*Tool
clients []*ClientConnection
conns []*ServerConnection
}
// ServerOptions is used to configure behavior of the server.
@ -74,14 +74,14 @@ func (s *Server) AddTools(tools ...*Tool) {
// Clients returns an iterator that yields the current set of client
// connections.
func (s *Server) Clients() iter.Seq[*ClientConnection] {
func (s *Server) Clients() iter.Seq[*ServerConnection] {
s.mu.Lock()
clients := slices.Clone(s.clients)
clients := slices.Clone(s.conns)
s.mu.Unlock()
return slices.Values(clients)
}
func (s *Server) listPrompts(_ context.Context, _ *ClientConnection, params *protocol.ListPromptsParams) (*protocol.ListPromptsResult, error) {
func (s *Server) listPrompts(_ context.Context, _ *ServerConnection, params *protocol.ListPromptsParams) (*protocol.ListPromptsResult, error) {
s.mu.Lock()
defer s.mu.Unlock()
@ -92,7 +92,7 @@ func (s *Server) listPrompts(_ context.Context, _ *ClientConnection, params *pro
return res, nil
}
func (s *Server) getPrompt(ctx context.Context, cc *ClientConnection, params *protocol.GetPromptParams) (*protocol.GetPromptResult, error) {
func (s *Server) getPrompt(ctx context.Context, cc *ServerConnection, params *protocol.GetPromptParams) (*protocol.GetPromptResult, error) {
s.mu.Lock()
var prompt *Prompt
if i := slices.IndexFunc(s.prompts, func(t *Prompt) bool {
@ -108,7 +108,7 @@ func (s *Server) getPrompt(ctx context.Context, cc *ClientConnection, params *pr
return prompt.Handler(ctx, cc, params.Arguments)
}
func (s *Server) listTools(_ context.Context, _ *ClientConnection, params *protocol.ListToolsParams) (*protocol.ListToolsResult, error) {
func (s *Server) listTools(_ context.Context, _ *ServerConnection, params *protocol.ListToolsParams) (*protocol.ListToolsResult, error) {
s.mu.Lock()
defer s.mu.Unlock()
@ -119,7 +119,7 @@ func (s *Server) listTools(_ context.Context, _ *ClientConnection, params *proto
return res, nil
}
func (s *Server) callTool(ctx context.Context, cc *ClientConnection, params *protocol.CallToolParams) (*protocol.CallToolResult, error) {
func (s *Server) callTool(ctx context.Context, cc *ServerConnection, params *protocol.CallToolParams) (*protocol.CallToolResult, error) {
s.mu.Lock()
var tool *Tool
if i := slices.IndexFunc(s.tools, func(t *Tool) bool {
@ -146,22 +146,22 @@ func (s *Server) Run(ctx context.Context, t Transport, opts *ConnectionOptions)
return cc.Wait()
}
// bind implements the binder[*ClientConnection] interface, so that Servers can
// bind implements the binder[*ServerConnection] interface, so that Servers can
// be connected using [connect].
func (s *Server) bind(conn *jsonrpc2.Connection) *ClientConnection {
cc := &ClientConnection{conn: conn, server: s}
func (s *Server) bind(conn *jsonrpc2.Connection) *ServerConnection {
cc := &ServerConnection{conn: conn, server: s}
s.mu.Lock()
s.clients = append(s.clients, cc)
s.conns = append(s.conns, cc)
s.mu.Unlock()
return cc
}
// disconnect implements the binder[*ClientConnection] interface, so that
// disconnect implements the binder[*ServerConnection] interface, so that
// Servers can be connected using [connect].
func (s *Server) disconnect(cc *ClientConnection) {
func (s *Server) disconnect(cc *ServerConnection) {
s.mu.Lock()
defer s.mu.Unlock()
s.clients = slices.DeleteFunc(s.clients, func(cc2 *ClientConnection) bool {
s.conns = slices.DeleteFunc(s.conns, func(cc2 *ServerConnection) bool {
return cc2 == cc
})
}
@ -172,15 +172,17 @@ func (s *Server) disconnect(cc *ClientConnection) {
// 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) (*ClientConnection, error) {
func (s *Server) Connect(ctx context.Context, t Transport, opts *ConnectionOptions) (*ServerConnection, error) {
return connect(ctx, t, opts, s)
}
// A ClientConnection is a connection with an MCP client.
// 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].
//
// It handles messages from the client, and can be used to send messages to the
// client. Create a connection by calling [Server.Connect].
type ClientConnection struct {
// Call [ServerConnection.Close] to close the connection, or await client
// termination with [ServerConnection.Wait].
type ServerConnection struct {
server *Server
conn *jsonrpc2.Connection
@ -190,11 +192,11 @@ type ClientConnection struct {
}
// Ping makes an MCP "ping" request to the client.
func (cc *ClientConnection) Ping(ctx context.Context) error {
func (cc *ServerConnection) Ping(ctx context.Context) error {
return call(ctx, cc.conn, "ping", nil, nil)
}
func (cc *ClientConnection) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) {
func (cc *ServerConnection) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) {
cc.mu.Lock()
initialized := cc.initialized
cc.mu.Unlock()
@ -210,7 +212,7 @@ func (cc *ClientConnection) handle(ctx context.Context, req *jsonrpc2.Request) (
}
}
// TODO: embed the incoming request ID in the ClientContext (or, more likely,
// TODO: embed the incoming request ID in the client context (or, more likely,
// a wrapper around it), so that we can correlate responses and notifications
// to the handler; this is required for the new session-based transport.
@ -239,7 +241,7 @@ func (cc *ClientConnection) handle(ctx context.Context, req *jsonrpc2.Request) (
return nil, jsonrpc2.ErrNotHandled
}
func (cc *ClientConnection) initialize(ctx context.Context, _ *ClientConnection, params *protocol.InitializeParams) (*protocol.InitializeResult, error) {
func (cc *ServerConnection) initialize(ctx context.Context, _ *ServerConnection, params *protocol.InitializeParams) (*protocol.InitializeResult, error) {
cc.mu.Lock()
cc.initializeParams = params
cc.mu.Unlock()
@ -274,15 +276,15 @@ func (cc *ClientConnection) initialize(ctx context.Context, _ *ClientConnection,
}, nil
}
// Close performs a graceful close of the connection, preventing new requests
// from being handled, and waiting for ongoing requests to return. Close then
// terminates the connection.
func (cc *ClientConnection) Close() error {
// 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 {
return cc.conn.Close()
}
// Wait waits for the connection to be closed by the client.
func (cc *ClientConnection) Wait() error {
func (cc *ServerConnection) Wait() error {
return cc.conn.Wait()
}
@ -290,7 +292,7 @@ func (cc *ClientConnection) Wait() error {
//
// Importantly, it returns nil if the handler returned an error, which is a
// requirement of the jsonrpc2 package.
func dispatch[TConn, TParams, TResult any](ctx context.Context, conn TConn, req *jsonrpc2.Request, f func(context.Context, TConn, TParams) (TResult, error)) (any, error) {
func dispatch[TParams, TResult any](ctx context.Context, conn *ServerConnection, req *jsonrpc2.Request, f func(context.Context, *ServerConnection, TParams) (TResult, error)) (any, error) {
var params TParams
if err := json.Unmarshal(req.Params, &params); err != nil {
return nil, err

2
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.ClientConnection, params *SayHiParams) ([]mcp.Content, error) {
func SayHi(ctx context.Context, cc *mcp.ServerConnection, params *SayHiParams) ([]mcp.Content, error) {
return []mcp.Content{
mcp.TextContent{Text: "Hi " + params.Name},
}, nil

8
mcp/sse.go
View file

@ -63,8 +63,8 @@ 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
onClient func(*ClientConnection) // for testing; must not block
getServer func(request *http.Request) *Server
onConnection func(*ServerConnection) // for testing; must not block
mu sync.Mutex
sessions map[string]*sseSession
@ -177,8 +177,8 @@ func (h *SSEHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
http.Error(w, "connection failed", http.StatusInternalServerError)
return
}
if h.onClient != nil {
h.onClient(cc)
if h.onConnection != nil {
h.onConnection(cc)
}
defer cc.Close() // close the transport when the GET exits

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.ClientConnection, params *AddParams) ([]mcp.Content, error) {
func Add(ctx context.Context, cc *mcp.ServerConnection, params *AddParams) ([]mcp.Content, error) {
return []mcp.Content{
mcp.TextContent{Text: fmt.Sprintf("%d", params.X+params.Y)},
}, nil

8
mcp/sse_test.go
View file

@ -24,10 +24,10 @@ func TestSSEServer(t *testing.T) {
sseHandler := NewSSEHandler(func(*http.Request) *Server { return server })
clients := make(chan *ClientConnection, 1)
sseHandler.onClient = func(cc *ClientConnection) {
conns := make(chan *ServerConnection, 1)
sseHandler.onConnection = func(cc *ServerConnection) {
select {
case clients <- cc:
case conns <- cc:
default:
}
}
@ -43,7 +43,7 @@ func TestSSEServer(t *testing.T) {
if err := c.Ping(ctx); err != nil {
t.Fatal(err)
}
cc := <-clients
cc := <-conns
gotHi, err := c.CallTool(ctx, "greet", map[string]any{"name": "user"})
if err != nil {
t.Fatal(err)

6
mcp/tool.go
View file

@ -15,7 +15,7 @@ import (
)
// A ToolHandler handles a call to tools/call.
type ToolHandler func(context.Context, *ClientConnection, map[string]json.RawMessage) (*protocol.CallToolResult, error)
type ToolHandler func(context.Context, *ServerConnection, map[string]json.RawMessage) (*protocol.CallToolResult, error)
// A Tool is a tool definition that is bound to a tool handler.
type Tool struct {
@ -36,12 +36,12 @@ type Tool 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, *ClientConnection, TReq) ([]Content, error), opts ...ToolOption) *Tool {
func NewTool[TReq any](name, description string, handler func(context.Context, *ServerConnection, TReq) ([]Content, error), opts ...ToolOption) *Tool {
schema, err := jsonschema.For[TReq]()
if err != nil {
panic(err)
}
wrapped := func(ctx context.Context, cc *ClientConnection, args map[string]json.RawMessage) (*protocol.CallToolResult, error) {
wrapped := func(ctx context.Context, cc *ServerConnection, args map[string]json.RawMessage) (*protocol.CallToolResult, error) {
// For simplicity, just marshal and unmarshal the arguments.
// This could be avoided in the future.
rawArgs, err := json.Marshal(args)

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.ClientConnection, T) ([]mcp.Content, error) {
func testToolHandler[T any](context.Context, *mcp.ServerConnection, T) ([]mcp.Content, error) {
panic("not implemented")
}

3
mcp/transport.go
View file

@ -73,8 +73,7 @@ func NewLocalTransport() (*IOTransport, *IOTransport) {
return &IOTransport{c1}, &IOTransport{c2}
}
// handler is an unexported version of jsonrpc2.Handler, to be implemented by
// [ServerConnection] and [ClientConnection].
// handler is an unexported version of jsonrpc2.Handler.
type handler interface {
handle(ctx context.Context, req *jsonrpc2.Request) (result any, err error)
}