diff --git a/mcp/client.go b/mcp/client.go index ab0b119..b9311dd 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -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 diff --git a/mcp/cmd_test.go b/mcp/cmd_test.go index 4d64a03..a42a7b0 100644 --- a/mcp/cmd_test.go +++ b/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) } } diff --git a/mcp/mcp.go b/mcp/mcp.go index 86d05ee..0516477 100644 --- a/mcp/mcp.go +++ b/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: // diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 8fbac1f..6dd69ad 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -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) diff --git a/mcp/server_example_test.go b/mcp/server_example_test.go index eabc812..7cafab2 100644 --- a/mcp/server_example_test.go +++ b/mcp/server_example_test.go @@ -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 diff --git a/mcp/sse_example_test.go b/mcp/sse_example_test.go index 2f53a6c..a261969 100644 --- a/mcp/sse_example_test.go +++ b/mcp/sse_example_test.go @@ -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) } diff --git a/mcp/sse_test.go b/mcp/sse_test.go index d259f4e..e1ec72f 100644 --- a/mcp/sse_test.go +++ b/mcp/sse_test.go @@ -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() } }) } diff --git a/mcp/transport.go b/mcp/transport.go index d06bcb1..7a629b4 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -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 }