internal/mcp: NewClient takes Transport

Pass the client Transport to NewClient instead of Connect.
Combine the ConnectionOptions with the ClientOptions.

Change-Id: I3acd6cf6d5bf46666742f2a73e2acaef2057ea77
Reviewed-on: https://go-review.googlesource.com/c/tools/+/670215
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-03 06:54:00 -04:00
commit f8098af1c5
8 changed files with 35 additions and 28 deletions

View file

@ -16,10 +16,12 @@ import (
)
// A Client is an MCP client, which may be connected to an MCP server
// using the [Client.Connect] method.
// using the [Client.Start] method.
type Client struct {
name string
version string
transport Transport
opts ClientOptions
mu sync.Mutex
conn *jsonrpc2.Connection
initializeResult *protocol.InitializeResult
@ -27,18 +29,25 @@ type Client struct {
// NewClient creates a new Client.
//
// Use [Client.Connect] to connect it to an MCP server.
// Use [Client.Start] to connect it to an MCP server.
//
// If non-nil, the provided options configure the Client.
func NewClient(name, version string, opts *ClientOptions) *Client {
return &Client{
name: name,
version: version,
func NewClient(name, version string, t Transport, opts *ClientOptions) *Client {
c := &Client{
name: name,
version: version,
transport: t,
}
if opts != nil {
c.opts = *opts
}
return c
}
// ClientOptions configures the behavior of the client.
type ClientOptions struct{}
type ClientOptions struct {
ConnectionOptions
}
// bind implements the binder[*ServerConnection] interface, so that Clients can
// be connected using [connect].
@ -56,20 +65,19 @@ func (c *Client) disconnect(*Client) {
// return an error.
}
// Connect connects the MCP client over the given transport and initializes an
// MCP session.
// Start begins an MCP session by connecting the MCP client over its transport.
//
// 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) (err error) {
func (c *Client) Start(ctx context.Context) (err error) {
defer func() {
if err != nil {
_ = c.Close()
}
}()
_, err = connect(ctx, t, opts, c)
_, err = connect(ctx, c.transport, &c.opts.ConnectionOptions, c)
if err != nil {
return err
}

View file

@ -49,8 +49,8 @@ func TestCmdTransport(t *testing.T) {
cmd := exec.Command(exe)
cmd.Env = append(os.Environ(), runAsServer+"=true")
client := mcp.NewClient("client", "v0.0.1", nil)
if err := client.Connect(ctx, mcp.NewCommandTransport(cmd), nil); err != nil {
client := mcp.NewClient("client", "v0.0.1", mcp.NewCommandTransport(cmd), nil)
if err := client.Start(ctx); err != nil {
log.Fatal(err)
}
got, err := client.CallTool(ctx, "greet", map[string]any{"name": "user"})

View file

@ -14,7 +14,7 @@
//
// A [Client] is an MCP client, which can be configured with various client
// capabilities. Clients may be connected to a [Server] instance
// using the [Client.Connect] method.
// using the [Client.Start] 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]

View file

@ -81,10 +81,10 @@ func TestEndToEnd(t *testing.T) {
clientWG.Done()
}()
c := NewClient("testClient", "v1.0.0", nil)
c := NewClient("testClient", "v1.0.0", ct, nil)
// Connect the client.
if err := c.Connect(ctx, ct, nil); err != nil {
if err := c.Start(ctx); err != nil {
t.Fatal(err)
}
@ -210,8 +210,8 @@ func basicConnection(t *testing.T, tools ...*Tool) (*ClientConnection, *Client)
t.Fatal(err)
}
c := NewClient("testClient", "v1.0.0", nil)
if err := c.Connect(ctx, ct, nil); err != nil {
c := NewClient("testClient", "v1.0.0", ct, nil)
if err := c.Start(ctx); err != nil {
t.Fatal(err)
}
return cc, c
@ -250,13 +250,12 @@ func TestBatching(t *testing.T) {
t.Fatal(err)
}
c := NewClient("testClient", "v1.0.0", nil)
opts := new(ConnectionOptions)
c := NewClient("testClient", "v1.0.0", ct, nil)
// TODO: this test is broken, because increasing the batch size here causes
// 'initialize' to block. Therefore, we can only test with a size of 1.
const batchSize = 1
BatchSize(ct, batchSize)
if err := c.Connect(ctx, ct, opts); err != nil {
if err := c.Start(ctx); err != nil {
t.Fatal(err)
}
defer c.Close()

View file

@ -34,8 +34,8 @@ func ExampleServer() {
log.Fatal(err)
}
client := mcp.NewClient("client", "v0.0.1", nil)
if err := client.Connect(ctx, clientTransport, nil); err != nil {
client := mcp.NewClient("client", "v0.0.1", clientTransport, nil)
if err := client.Start(ctx); err != nil {
log.Fatal(err)
}

View file

@ -34,8 +34,8 @@ func ExampleSSEHandler() {
ctx := context.Background()
transport := mcp.NewSSEClientTransport(httpServer.URL)
client := mcp.NewClient("test", "v1.0.0", nil)
if err := client.Connect(ctx, transport, nil); err != nil {
client := mcp.NewClient("test", "v1.0.0", transport, nil)
if err := client.Start(ctx); err != nil {
log.Fatal(err)
}
defer client.Close()

View file

@ -36,8 +36,8 @@ func TestSSEServer(t *testing.T) {
clientTransport := NewSSEClientTransport(httpServer.URL)
c := NewClient("testClient", "v1.0.0", nil)
if err := c.Connect(ctx, clientTransport, nil); err != nil {
c := NewClient("testClient", "v1.0.0", clientTransport, nil)
if err := c.Start(ctx); err != nil {
t.Fatal(err)
}
if err := c.Ping(ctx); err != nil {

View file

@ -27,7 +27,7 @@ var ErrConnectionClosed = errors.New("connection closed")
// and server.
//
// Transports should be used for at most one call to [Server.Connect] or
// [Client.Connect].
// [Client.Start].
type Transport interface {
// Connect returns the logical stream.
//