From f8098af1c5095d8be5a12dc4b42bf09c0424c101 Mon Sep 17 00:00:00 2001 From: Jonathan Amsterdam Date: Sat, 3 May 2025 06:54:00 -0400 Subject: [PATCH] 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 Reviewed-by: Robert Findley --- mcp/client.go | 30 +++++++++++++++++++----------- mcp/cmd_test.go | 4 ++-- mcp/mcp.go | 2 +- mcp/mcp_test.go | 13 ++++++------- mcp/server_example_test.go | 4 ++-- mcp/sse_example_test.go | 4 ++-- mcp/sse_test.go | 4 ++-- mcp/transport.go | 2 +- 8 files changed, 35 insertions(+), 28 deletions(-) diff --git a/mcp/client.go b/mcp/client.go index 48c2c0c..92a1d18 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -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 } diff --git a/mcp/cmd_test.go b/mcp/cmd_test.go index 762754a..53cdb89 100644 --- a/mcp/cmd_test.go +++ b/mcp/cmd_test.go @@ -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"}) diff --git a/mcp/mcp.go b/mcp/mcp.go index 3737553..c9a43de 100644 --- a/mcp/mcp.go +++ b/mcp/mcp.go @@ -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] diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 69aa45a..be2497f 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -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() diff --git a/mcp/server_example_test.go b/mcp/server_example_test.go index 7cafab2..a52b46f 100644 --- a/mcp/server_example_test.go +++ b/mcp/server_example_test.go @@ -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) } diff --git a/mcp/sse_example_test.go b/mcp/sse_example_test.go index a261969..3007c1d 100644 --- a/mcp/sse_example_test.go +++ b/mcp/sse_example_test.go @@ -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() diff --git a/mcp/sse_test.go b/mcp/sse_test.go index c208ba4..d557425 100644 --- a/mcp/sse_test.go +++ b/mcp/sse_test.go @@ -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 { diff --git a/mcp/transport.go b/mcp/transport.go index 313f0b0..419f44d 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -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. //