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

mcp: remove ServerConnection

Browse source 
Remove ServerConnection.
- It's not in typescript or mcpgo APIs.
- Amortizing a client over multiple servers doesn't have any value.

Change-Id: I725b42569733d6836864d96b16898a10ae91e570
Reviewed-on: https://go-review.googlesource.com/c/tools/+/669575
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Robert Findley <rfindley@google.com>
This commit is contained in:
Jonathan Amsterdam 2025-05-02 12:40:29 -04:00 • committed by Robert Findley
commit 7a76f3502a
8 changed files with 83 additions and 134 deletions
Download patch file Download diff file Expand all files Collapse all files

115
mcp/client.go
View file

@ -2,31 +2,27 @@
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file.
// TODO: consider passing Transport to NewClient and merging {Connection,Client}Options
package mcp
import (
"context"
"encoding/json"
"fmt"
"iter"
"slices"
"sync"
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
"github.com/modelcontextprotocol/go-sdk/internal/protocol"
)
// A Client is an MCP client, which may be connected to one or more MCP servers
// A Client is an MCP client, which may be connected to an MCP server
// using the [Client.Connect] method.
//
// TODO(rfindley): revisit the many-to-one relationship of clients and servers.
// It is a bit odd.
type Client struct {
name string
version string
mu sync.Mutex
servers []*ServerConnection
name string
version string
mu sync.Mutex
conn *jsonrpc2.Connection
initializeResult *protocol.InitializeResult
}
// NewClient creates a new Client.
@ -41,99 +37,68 @@ func NewClient(name, version string, opts *ClientOptions) *Client {
}
}
// Servers returns an iterator that yields the current set of server
// connections.
func (c *Client) Servers() iter.Seq[*ServerConnection] {
c.mu.Lock()
clients := slices.Clone(c.servers)
c.mu.Unlock()
return slices.Values(clients)
}
// ClientOptions configures the behavior of the client, and apply to every
// client-server connection created using [Client.Connect].
// ClientOptions configures the behavior of the client.
type ClientOptions struct{}
// bind implements the binder[*ServerConnection] interface, so that Clients can
// be connected using [connect].
func (c *Client) bind(conn *jsonrpc2.Connection) *ServerConnection {
sc := &ServerConnection{
conn: conn,
client: c,
}
func (c *Client) bind(conn *jsonrpc2.Connection) *Client {
c.mu.Lock()
c.servers = append(c.servers, sc)
c.mu.Unlock()
return sc
defer c.mu.Unlock()
c.conn = conn
return c
}
// disconnect implements the binder[*ServerConnection] interface, so that
// Clients can be connected using [connect].
func (c *Client) disconnect(sc *ServerConnection) {
c.mu.Lock()
defer c.mu.Unlock()
c.servers = slices.DeleteFunc(c.servers, func(sc2 *ServerConnection) bool {
return sc2 == sc
})
func (c *Client) disconnect(*Client) {
// Do nothing. In particular, do not set conn to nil: it needs to exist so it can
// return an error.
}
// Connect connects the MCP client over the given transport and initializes an
// MCP session.
//
// It returns an initialized [ServerConnection] object that may be used to
// query the MCP server, terminate the connection (with [Connection.Close]), or
// await server termination (with [Connection.Wait]).
//
// Typically, it is the responsibility of the client to close the connection
// when it is no longer needed. However, if the connection is closed by the
// server, calls or notifications will return an error wrapping
// [ErrConnectionClosed].
func (c *Client) Connect(ctx context.Context, t Transport, opts *ConnectionOptions) (sc *ServerConnection, err error) {
func (c *Client) Connect(ctx context.Context, t Transport, opts *ConnectionOptions) (err error) {
defer func() {
if sc != nil && err != nil {
_ = sc.Close()
if err != nil {
_ = c.Close()
}
}()
sc, err = connect(ctx, t, opts, c)
_, err = connect(ctx, t, opts, c)
if err != nil {
return nil, err
return err
}
params := &protocol.InitializeParams{
ClientInfo: protocol.Implementation{Name: c.name, Version: c.version},
}
if err := call(ctx, sc.conn, "initialize", params, &sc.initializeResult); err != nil {
return nil, err
if err := call(ctx, c.conn, "initialize", params, &c.initializeResult); err != nil {
return err
}
if err := sc.conn.Notify(ctx, "notifications/initialized", &protocol.InitializedParams{}); err != nil {
return nil, err
if err := c.conn.Notify(ctx, "notifications/initialized", &protocol.InitializedParams{}); err != nil {
return err
}
return sc, nil
}
// A ServerConnection is a connection with an MCP server.
//
// It handles messages from the client, and can be used to send messages to the
// client. Create a connection by calling [Server.Connect].
type ServerConnection struct {
conn *jsonrpc2.Connection
client *Client
initializeResult *protocol.InitializeResult
return 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 *ServerConnection) Close() error {
return cc.conn.Close()
func (c *Client) Close() error {
return c.conn.Close()
}
// Wait waits for the connection to be closed by the server.
// Generally, clients should be responsible for closing the connection.
func (cc *ServerConnection) Wait() error {
return cc.conn.Wait()
func (c *Client) Wait() error {
return c.conn.Wait()
}
func (sc *ServerConnection) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) {
func (*Client) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) {
// No need to check that the connection is initialized, since we initialize
// it in Connect.
switch req.Method {
@ -145,24 +110,24 @@ func (sc *ServerConnection) handle(ctx context.Context, req *jsonrpc2.Request) (
}
// Ping makes an MCP "ping" request to the server.
func (sc *ServerConnection) Ping(ctx context.Context) error {
return call(ctx, sc.conn, "ping", nil, nil)
func (c *Client) Ping(ctx context.Context) error {
return call(ctx, c.conn, "ping", nil, nil)
}
// ListPrompts lists prompts that are currently available on the server.
func (sc *ServerConnection) ListPrompts(ctx context.Context) ([]protocol.Prompt, error) {
func (c *Client) ListPrompts(ctx context.Context) ([]protocol.Prompt, error) {
var (
params = &protocol.ListPromptsParams{}
result protocol.ListPromptsResult
)
if err := call(ctx, sc.conn, "prompts/list", params, &result); err != nil {
if err := call(ctx, c.conn, "prompts/list", params, &result); err != nil {
return nil, err
}
return result.Prompts, nil
}
// GetPrompt gets a prompt from the server.
func (sc *ServerConnection) GetPrompt(ctx context.Context, name string, args map[string]string) (*protocol.GetPromptResult, error) {
func (c *Client) GetPrompt(ctx context.Context, name string, args map[string]string) (*protocol.GetPromptResult, error) {
var (
params = &protocol.GetPromptParams{
Name: name,
@ -170,19 +135,19 @@ func (sc *ServerConnection) GetPrompt(ctx context.Context, name string, args map
}
result = &protocol.GetPromptResult{}
)
if err := call(ctx, sc.conn, "prompts/get", params, result); err != nil {
if err := call(ctx, c.conn, "prompts/get", params, result); err != nil {
return nil, err
}
return result, nil
}
// ListTools lists tools that are currently available on the server.
func (sc *ServerConnection) ListTools(ctx context.Context) ([]protocol.Tool, error) {
func (c *Client) ListTools(ctx context.Context) ([]protocol.Tool, error) {
var (
params = &protocol.ListToolsParams{}
result protocol.ListToolsResult
)
if err := call(ctx, sc.conn, "tools/list", params, &result); err != nil {
if err := call(ctx, c.conn, "tools/list", params, &result); err != nil {
return nil, err
}
return result.Tools, nil
@ -193,7 +158,7 @@ func (sc *ServerConnection) ListTools(ctx context.Context) ([]protocol.Tool, err
// TODO(jba): make the following true:
// If the provided arguments do not conform to the schema for the given tool,
// the call fails.
func (sc *ServerConnection) CallTool(ctx context.Context, name string, args map[string]any) (_ *protocol.CallToolResult, err error) {
func (c *Client) CallTool(ctx context.Context, name string, args map[string]any) (_ *protocol.CallToolResult, err error) {
defer func() {
if err != nil {
err = fmt.Errorf("calling tool %q: %w", name, err)
@ -214,7 +179,7 @@ func (sc *ServerConnection) CallTool(ctx context.Context, name string, args map[
}
result protocol.CallToolResult
)
if err := call(ctx, sc.conn, "tools/call", params, &result); err != nil {
if err := call(ctx, c.conn, "tools/call", params, &result); err != nil {
return nil, err
}
return &result, nil

7
mcp/cmd_test.go
View file

@ -50,11 +50,10 @@ func TestCmdTransport(t *testing.T) {
cmd.Env = append(os.Environ(), runAsServer+"=true")
client := mcp.NewClient("client", "v0.0.1", nil)
serverConn, err := client.Connect(ctx, mcp.NewCommandTransport(cmd), nil)
if err != nil {
if err := client.Connect(ctx, mcp.NewCommandTransport(cmd), nil); err != nil {
log.Fatal(err)
}
got, err := serverConn.CallTool(ctx, "greet", map[string]any{"name": "user"})
got, err := client.CallTool(ctx, "greet", map[string]any{"name": "user"})
if err != nil {
log.Fatal(err)
}
@ -64,7 +63,7 @@ func TestCmdTransport(t *testing.T) {
if diff := cmp.Diff(want, got); diff != "" {
t.Errorf("greet returned unexpected content (-want +got):\n%s", diff)
}
if err := serverConn.Close(); err != nil {
if err := client.Close(); err != nil {
t.Fatalf("closing server: %v", err)
}
}

15
mcp/mcp.go
View file

@ -9,13 +9,13 @@
// 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) ⇅
// ServerConnection ⇄ Client Transport ⇄ Server Transport ⇄ ClientConnection
// Client Server
// ⇅ (jsonrpc2) ⇅
// Client Transport ⇄ Server Transport ⇄ ClientConnection
//
// A [Client] is an MCP client, which can be configured with various client
// capabilities. Clients may be connected to one or more [Server] instances
// using the [Client.Connect] method, which creates a [ServerConnection].
// capabilities. Clients may be connected to a [Server] instance
// using the [Client.Connect] method.
//
// Similarly, a [Server] is an MCP server, which can be configured with various
// server capabilities. Servers may be connected to one or more [Client]
@ -44,12 +44,11 @@
// client := mcp.NewClient("mcp-client", "v1.0.0", nil)
// // Connect to a server over stdin/stdout
// transport := mcp.NewCommandTransport(exec.Command("myserver"))
// serverConn, err := client.Connect(ctx, transport, nil)
// if err != nil {
// if err := client.Connect(ctx, transport, nil); err != nil {
// log.Fatal(err)
// }
// // Call a tool on the server.
// content, err := serverConn.CallTool(ctx, "greet", map[string]any{"name": "you"})
// content, err := client.CallTool(ctx, "greet", map[string]any{"name": "you"})
//
// Here is an example of the corresponding server, connected over stdin/stdout:
//

50
mcp/mcp_test.go
View file

@ -83,20 +83,15 @@ func TestEndToEnd(t *testing.T) {
c := NewClient("testClient", "v1.0.0", nil)
// Connect the client.
sc, err := c.Connect(ctx, ct, nil)
if err != nil {
if err := c.Connect(ctx, ct, nil); err != nil {
t.Fatal(err)
}
if got := slices.Collect(c.Servers()); len(got) != 1 {
t.Errorf("after connection, Servers() has length %d, want 1", len(got))
}
if err := sc.Ping(ctx); err != nil {
if err := c.Ping(ctx); err != nil {
t.Fatalf("ping failed: %v", err)
}
gotPrompts, err := sc.ListPrompts(ctx)
gotPrompts, err := c.ListPrompts(ctx)
if err != nil {
t.Errorf("prompts/list failed: %v", err)
}
@ -112,7 +107,7 @@ func TestEndToEnd(t *testing.T) {
t.Fatalf("prompts/list mismatch (-want +got):\n%s", diff)
}
gotReview, err := sc.GetPrompt(ctx, "code_review", map[string]string{"Code": "1+1"})
gotReview, err := c.GetPrompt(ctx, "code_review", map[string]string{"Code": "1+1"})
if err != nil {
t.Fatal(err)
}
@ -127,11 +122,11 @@ func TestEndToEnd(t *testing.T) {
t.Errorf("prompts/get 'code_review' mismatch (-want +got):\n%s", diff)
}
if _, err := sc.GetPrompt(ctx, "fail", map[string]string{}); err == nil || !strings.Contains(err.Error(), failure.Error()) {
if _, err := c.GetPrompt(ctx, "fail", map[string]string{}); err == nil || !strings.Contains(err.Error(), failure.Error()) {
t.Errorf("fail returned unexpected error: got %v, want containing %v", err, failure)
}
gotTools, err := sc.ListTools(ctx)
gotTools, err := c.ListTools(ctx)
if err != nil {
t.Errorf("tools/list failed: %v", err)
}
@ -158,7 +153,7 @@ func TestEndToEnd(t *testing.T) {
t.Fatalf("tools/list mismatch (-want +got):\n%s", diff)
}
gotHi, err := sc.CallTool(ctx, "greet", map[string]any{"name": "user"})
gotHi, err := c.CallTool(ctx, "greet", map[string]any{"name": "user"})
if err != nil {
t.Fatal(err)
}
@ -169,7 +164,7 @@ func TestEndToEnd(t *testing.T) {
t.Errorf("tools/call 'greet' mismatch (-want +got):\n%s", diff)
}
gotFail, err := sc.CallTool(ctx, "fail", map[string]any{})
gotFail, err := c.CallTool(ctx, "fail", map[string]any{})
// Counter-intuitively, when a tool fails, we don't expect an RPC error for
// call tool: instead, the failure is embedded in the result.
if err != nil {
@ -184,7 +179,7 @@ func TestEndToEnd(t *testing.T) {
}
// Disconnect.
sc.Close()
c.Close()
clientWG.Wait()
// After disconnecting, neither client nor server should have any
@ -192,9 +187,6 @@ func TestEndToEnd(t *testing.T) {
for range s.Clients() {
t.Errorf("unexpected client after disconnection")
}
for range c.Servers() {
t.Errorf("unexpected server after disconnection")
}
}
// basicConnection returns a new basic client-server connection configured with
@ -202,7 +194,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, *ServerConnection) {
func basicConnection(t *testing.T, tools ...*Tool) (*ClientConnection, *Client) {
t.Helper()
ctx := context.Background()
@ -218,32 +210,31 @@ func basicConnection(t *testing.T, tools ...*Tool) (*ClientConnection, *ServerCo
}
c := NewClient("testClient", "v1.0.0", nil)
sc, err := c.Connect(ctx, ct, nil)
if err != nil {
if err := c.Connect(ctx, ct, nil); err != nil {
t.Fatal(err)
}
return cc, sc
return cc, c
}
func TestServerClosing(t *testing.T) {
cc, sc := basicConnection(t, MakeTool("greet", "say hi", sayHi))
defer sc.Close()
cc, c := basicConnection(t, MakeTool("greet", "say hi", sayHi))
defer c.Close()
ctx := context.Background()
var wg sync.WaitGroup
wg.Add(1)
go func() {
if err := sc.Wait(); err != nil {
if err := c.Wait(); err != nil {
t.Errorf("server connection failed: %v", err)
}
wg.Done()
}()
if _, err := sc.CallTool(ctx, "greet", map[string]any{"name": "user"}); err != nil {
if _, err := c.CallTool(ctx, "greet", map[string]any{"name": "user"}); err != nil {
t.Fatalf("after connecting: %v", err)
}
cc.Close()
wg.Wait()
if _, err := sc.CallTool(ctx, "greet", map[string]any{"name": "user"}); !errors.Is(err, ErrConnectionClosed) {
if _, err := c.CallTool(ctx, "greet", map[string]any{"name": "user"}); !errors.Is(err, ErrConnectionClosed) {
t.Errorf("after disconnection, got error %v, want EOF", err)
}
}
@ -264,16 +255,15 @@ func TestBatching(t *testing.T) {
// 'initialize' to block. Therefore, we can only test with a size of 1.
const batchSize = 1
BatchSize(ct, batchSize)
sc, err := c.Connect(ctx, ct, opts)
if err != nil {
if err := c.Connect(ctx, ct, opts); err != nil {
t.Fatal(err)
}
defer sc.Close()
defer c.Close()
errs := make(chan error, batchSize)
for i := range batchSize {
go func() {
_, err := sc.ListTools(ctx)
_, err := c.ListTools(ctx)
errs <- err
}()
time.Sleep(2 * time.Millisecond)

7
mcp/server_example_test.go
View file

@ -35,18 +35,17 @@ func ExampleServer() {
}
client := mcp.NewClient("client", "v0.0.1", nil)
serverConnection, err := client.Connect(ctx, clientTransport, nil)
if err != nil {
if err := client.Connect(ctx, clientTransport, nil); err != nil {
log.Fatal(err)
}
res, err := serverConnection.CallTool(ctx, "greet", map[string]any{"name": "user"})
res, err := client.CallTool(ctx, "greet", map[string]any{"name": "user"})
if err != nil {
log.Fatal(err)
}
fmt.Println(res.Content[0].Text)
serverConnection.Close()
client.Close()
clientConnection.Wait()
// Output: Hi user

8
mcp/sse_example_test.go
View file

@ -34,13 +34,13 @@ func ExampleSSEHandler() {
ctx := context.Background()
transport := mcp.NewSSEClientTransport(httpServer.URL)
serverConn, err := mcp.NewClient("test", "v1.0.0", nil).Connect(ctx, transport, nil)
if err != nil {
client := mcp.NewClient("test", "v1.0.0", nil)
if err := client.Connect(ctx, transport, nil); err != nil {
log.Fatal(err)
}
defer serverConn.Close()
defer client.Close()
res, err := serverConn.CallTool(ctx, "add", map[string]any{"x": 1, "y": 2})
res, err := client.CallTool(ctx, "add", map[string]any{"x": 1, "y": 2})
if err != nil {
log.Fatal(err)
}

13
mcp/sse_test.go
View file

@ -36,16 +36,15 @@ func TestSSEServer(t *testing.T) {
clientTransport := NewSSEClientTransport(httpServer.URL)
client := NewClient("testClient", "v1.0.0", nil)
sc, err := client.Connect(ctx, clientTransport, nil)
if err != nil {
c := NewClient("testClient", "v1.0.0", nil)
if err := c.Connect(ctx, clientTransport, nil); err != nil {
t.Fatal(err)
}
if err := sc.Ping(ctx); err != nil {
if err := c.Ping(ctx); err != nil {
t.Fatal(err)
}
cc := <-clients
gotHi, err := sc.CallTool(ctx, "greet", map[string]any{"name": "user"})
gotHi, err := c.CallTool(ctx, "greet", map[string]any{"name": "user"})
if err != nil {
t.Fatal(err)
}
@ -59,11 +58,11 @@ func TestSSEServer(t *testing.T) {
// Test that closing either end of the connection terminates the other
// end.
if closeServerFirst {
sc.Close()
c.Close()
cc.Wait()
} else {
cc.Close()
sc.Wait()
c.Wait()
}
})
}

2
mcp/transport.go
View file

@ -77,7 +77,6 @@ func NewLocalTransport() (*IOTransport, *IOTransport) {
// [ServerConnection] and [ClientConnection].
type handler interface {
handle(ctx context.Context, req *jsonrpc2.Request) (result any, err error)
comparable
}
type binder[T handler] interface {
@ -122,7 +121,6 @@ func connect[H handler](ctx context.Context, t Transport, opts *ConnectionOption
},
})
assert(preempter.conn != nil, "unbound preempter")
assert(h != zero, "unbound connection")
return h, nil
}