From 347a1a0078a252d464b7b1d1f814436d7f3e228b Mon Sep 17 00:00:00 2001 From: Rob Findley Date: Wed, 7 May 2025 16:43:03 +0000 Subject: [PATCH] internal/mcp: rename ClientConnection to ServerConnection 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 Auto-Submit: Robert Findley Reviewed-by: Jonathan Amsterdam --- README.md | 3 +- examples/hello/main.go | 4 +-- examples/sse/main.go | 2 +- mcp/client.go | 4 +-- mcp/mcp.go | 4 +-- mcp/mcp_test.go | 12 ++++---- mcp/prompt.go | 6 ++-- mcp/prompt_test.go | 2 +- mcp/server.go | 60 ++++++++++++++++++++------------------ mcp/server_example_test.go | 2 +- mcp/sse.go | 8 ++--- mcp/sse_example_test.go | 2 +- mcp/sse_test.go | 8 ++--- mcp/tool.go | 6 ++-- mcp/tool_test.go | 2 +- mcp/transport.go | 3 +- 16 files changed, 65 insertions(+), 63 deletions(-) diff --git a/README.md b/README.md index 9fb0ddf..da0fec0 100644 --- a/README.md +++ b/README.md @@ -59,7 +59,7 @@ type HiParams struct { Name string `json:"name"` } -func SayHi(ctx context.Context, cc *mcp.ClientConnection, params *HiParams) ([]mcp.Content, error) { +func SayHi(ctx context.Context, cc *mcp.ServerConnection, params *HiParams) ([]mcp.Content, error) { return []mcp.Content{ mcp.TextContent{Text: "Hi " + params.Name}, }, nil @@ -110,6 +110,7 @@ server.AddPrompts(mcp.MakePrompt("greet", "", PromptHi)) Resources are a core primitive in the Model Context Protocol (MCP) that allow servers to expose data and content that can be read by clients and used as context for LLM interactions. + Resources are not supported yet. ## Testing diff --git a/examples/hello/main.go b/examples/hello/main.go index c86983b..c75960e 100644 --- a/examples/hello/main.go +++ b/examples/hello/main.go @@ -21,13 +21,13 @@ type HiParams struct { Name string `json:"name"` } -func SayHi(ctx context.Context, cc *mcp.ClientConnection, params *HiParams) ([]mcp.Content, error) { +func SayHi(ctx context.Context, cc *mcp.ServerConnection, params *HiParams) ([]mcp.Content, error) { return []mcp.Content{ mcp.TextContent{Text: "Hi " + params.Name}, }, nil } -func PromptHi(ctx context.Context, cc *mcp.ClientConnection, params *HiParams) (*protocol.GetPromptResult, error) { +func PromptHi(ctx context.Context, cc *mcp.ServerConnection, params *HiParams) (*protocol.GetPromptResult, error) { return &protocol.GetPromptResult{ Description: "Code review prompt", Messages: []protocol.PromptMessage{ diff --git a/examples/sse/main.go b/examples/sse/main.go index f87680e..1dcdd6e 100644 --- a/examples/sse/main.go +++ b/examples/sse/main.go @@ -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.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 diff --git a/mcp/client.go b/mcp/client.go index 92a1d18..76794c1 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -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 diff --git a/mcp/mcp.go b/mcp/mcp.go index c9a43de..734b96c 100644 --- a/mcp/mcp.go +++ b/mcp/mcp.go @@ -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 diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index ff4f270..0efd763 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -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(): diff --git a/mcp/prompt.go b/mcp/prompt.go index fc6c66f..bd57190 100644 --- a/mcp/prompt.go +++ b/mcp/prompt.go @@ -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) diff --git a/mcp/prompt_test.go b/mcp/prompt_test.go index 2cffd2c..35075a4 100644 --- a/mcp/prompt_test.go +++ b/mcp/prompt_test.go @@ -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") } diff --git a/mcp/server.go b/mcp/server.go index 89dfb7f..9c9351d 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -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, ¶ms); err != nil { return nil, err diff --git a/mcp/server_example_test.go b/mcp/server_example_test.go index bb4260a..614cc41 100644 --- a/mcp/server_example_test.go +++ b/mcp/server_example_test.go @@ -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 diff --git a/mcp/sse.go b/mcp/sse.go index e1891bc..2104755 100644 --- a/mcp/sse.go +++ b/mcp/sse.go @@ -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 diff --git a/mcp/sse_example_test.go b/mcp/sse_example_test.go index 125bcc0..11a680f 100644 --- a/mcp/sse_example_test.go +++ b/mcp/sse_example_test.go @@ -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 diff --git a/mcp/sse_test.go b/mcp/sse_test.go index 327eff6..b2e3d0e 100644 --- a/mcp/sse_test.go +++ b/mcp/sse_test.go @@ -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) diff --git a/mcp/tool.go b/mcp/tool.go index 9dfce8b..3c8bacc 100644 --- a/mcp/tool.go +++ b/mcp/tool.go @@ -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) diff --git a/mcp/tool_test.go b/mcp/tool_test.go index c7eeb94..f6230b2 100644 --- a/mcp/tool_test.go +++ b/mcp/tool_test.go @@ -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") } diff --git a/mcp/transport.go b/mcp/transport.go index 419f44d..15dc09e 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -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) }