mcp: remove ServerConnection

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 <[email protected]>
Reviewed-by: Robert Findley <[email protected]>

Author: jba

internal/mcp/client.go

@@ -2,31 +2,27 @@
 // Use of this source code is governed by a BSD-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"
 
 	jsonrpc2 "golang.org/x/tools/internal/jsonrpc2_v2"
 	"golang.org/x/tools/internal/mcp/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,44 +110,44 @@ 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,
 			Arguments: args,
 		}
 		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

internal/mcp/cmd_test.go

@@ -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)
 	}
 }

internal/mcp/mcp.go

@@ -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:
 //