From a8b98da4807052ecc8d61b17aba5ffcef6d5d922 Mon Sep 17 00:00:00 2001 From: Rob Findley Date: Wed, 14 May 2025 17:53:48 +0000 Subject: [PATCH] internal/mcp: reinstate ClientSession Add a logical client session type back to the SDK, to bring it closer to the design. Also, remove 'ConnectionOptions', which is better implemented using middleware, and reimplement logging as a LoggingTransport middleware, which is also added to the design. Change-Id: I908b4f837d513ffc60ae5b6171a31519cffafebc Reviewed-on: https://go-review.googlesource.com/c/tools/+/672795 LUCI-TryBot-Result: Go LUCI Reviewed-by: Sam Thanawalla Reviewed-by: Jonathan Amsterdam --- design/design.md | 25 +++++++-- examples/hello/main.go | 4 +- mcp/client.go | 109 +++++++++++++++++++++---------------- mcp/cmd_test.go | 11 ++-- mcp/mcp_test.go | 55 ++++++++++--------- mcp/server.go | 74 ++++++++++++------------- mcp/server_example_test.go | 13 +++-- mcp/sse.go | 2 +- mcp/sse_example_test.go | 9 +-- mcp/sse_test.go | 19 ++++--- mcp/transport.go | 104 +++++++++++++++++------------------ 11 files changed, 228 insertions(+), 197 deletions(-) diff --git a/design/design.md b/design/design.md index 120ce53..21b43ac 100644 --- a/design/design.md +++ b/design/design.md @@ -34,7 +34,6 @@ compatible with mcp-go, translating between them should be straightforward in most cases. (Later, we will provide a detailed translation guide.) - # Requirements These may be obvious, but it's worthwhile to define goals for an official MCP @@ -283,12 +282,24 @@ func NewStreamableClientTransport(url string) *StreamableClientTransport { func (*StreamableClientTransport) Connect(context.Context) (Stream, error) ``` -Finally, we also provide an in-memory transport, for scenarios such as testing, -where the MCP client and server are in the same process. +Finally, we also provide a couple of transport implementations for special scenarios. +An InMemoryTransport can be used when the client and server reside in the same +process. A LoggingTransport is a middleware layer that logs RPC logs to a desired +location, specified as an io.Writer. ```go +// An InMemoryTransport is a [Transport] that communicates over an in-memory +// network connection, using newline-delimited JSON. type InMemoryTransport struct { /* ... */ } -func NewInMemoryTransport() (*InMemoryTransport, *InMemoryTransport) + +// NewInMemoryTransports returns two InMemoryTransports that connect to each +// other. +func NewInMemoryTransports() (*InMemoryTransport, *InMemoryTransport) + +// A LoggingTransport is a [Transport] that delegates to another transport, +// writing RPC logs to an io.Writer. +type LoggingTransport struct { /* ... */ } +func NewLoggingTransport(delegate Transport, w io.Writer) *LoggingTransport ``` **Differences from mcp-go**: The Go team has a battle-tested JSON-RPC @@ -538,7 +549,6 @@ func (*ClientSession) Resources(context.Context, *ListResourceParams) iter.Seq2[ func (*ClientSession) ResourceTemplates(context.Context, *ListResourceTemplatesParams) iter.Seq2[ResourceTemplate, error] ``` - ### Middleware We provide a mechanism to add MCP-level middleware on the server side, which runs after the @@ -902,15 +912,18 @@ handler to a Go function using reflection to derive its arguments. We provide In our design, each resource and resource template is associated with a function that reads it, with this signature: + ```go type ResourceHandler func(context.Context, *ServerSession, *ReadResourceParams) (*ReadResourceResult, error) ``` + The arguments include the `ServerSession` so the handler can observe the client's roots. The handler should return the resource contents in a `ReadResourceResult`, calling either `NewTextResourceContents` or `NewBlobResourceContents`. If the handler omits the URI or MIME type, the server will populate them from the resource. The `ServerResource` and `ServerResourceTemplate` types hold the association between the resource and its handler: + ```go type ServerResource struct { Resource Resource @@ -939,12 +952,14 @@ The `ReadResource` method finds a resource or resource template matching the arg its assocated handler. To read files from the local filesystem, we recommend using `FileResourceHandler` to construct a handler: + ```go // FileResourceHandler returns a ResourceHandler that reads paths using dir as a root directory. // It protects against path traversal attacks. // It will not read any file that is not in the root set of the client session requesting the resource. func (*Server) FileResourceHandler(dir string) ResourceHandler ``` + Here is an example: ```go diff --git a/examples/hello/main.go b/examples/hello/main.go index 52a8132..40b4b22 100644 --- a/examples/hello/main.go +++ b/examples/hello/main.go @@ -50,8 +50,8 @@ func main() { }) http.ListenAndServe(*httpAddr, handler) } else { - opts := &mcp.ConnectionOptions{Logger: os.Stderr} - if err := server.Run(context.Background(), mcp.NewStdIOTransport(), opts); err != nil { + t := mcp.NewLoggingTransport(mcp.NewStdIOTransport(), os.Stderr) + if err := server.Run(context.Background(), t); err != nil { fmt.Fprintf(os.Stderr, "Server failed: %v", err) } } diff --git a/mcp/client.go b/mcp/client.go index d1ff374..4a7694a 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -16,29 +16,26 @@ import ( ) // A Client is an MCP client, which may be connected to an MCP server -// using the [Client.Start] method. +// using the [Client.Connect] method. type Client struct { - name string - version string - transport Transport - opts ClientOptions - mu sync.Mutex - conn *jsonrpc2.Connection - roots *featureSet[*Root] - initializeResult *initializeResult + name string + version string + opts ClientOptions + mu sync.Mutex + roots *featureSet[*Root] + sessions []*ClientSession } // NewClient creates a new Client. // -// Use [Client.Start] to connect it to an MCP server. +// Use [Client.Connect] to connect it to an MCP server. // // If non-nil, the provided options configure the Client. -func NewClient(name, version string, t Transport, opts *ClientOptions) *Client { +func NewClient(name, version string, opts *ClientOptions) *Client { c := &Client{ - name: name, - version: version, - transport: t, - roots: newFeatureSet(func(r *Root) string { return r.URI }), + name: name, + version: version, + roots: newFeatureSet(func(r *Root) string { return r.URI }), } if opts != nil { c.opts = *opts @@ -48,63 +45,79 @@ func NewClient(name, version string, t Transport, opts *ClientOptions) *Client { // ClientOptions configures the behavior of the client. type ClientOptions struct { - ConnectionOptions } -// bind implements the binder[*Client] interface, so that Clients can +// bind implements the binder[*ClientSession] interface, so that Clients can // be connected using [connect]. -func (c *Client) bind(conn *jsonrpc2.Connection) *Client { +func (c *Client) bind(conn *jsonrpc2.Connection) *ClientSession { + cs := &ClientSession{ + conn: conn, + client: c, + } c.mu.Lock() defer c.mu.Unlock() - c.conn = conn - return c + c.sessions = append(c.sessions, cs) + return cs } // disconnect implements the binder[*Client] interface, so that // Clients can be connected using [connect]. -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. +func (c *Client) disconnect(cs *ClientSession) { + c.mu.Lock() + defer c.mu.Unlock() + c.sessions = slices.DeleteFunc(c.sessions, func(cs2 *ClientSession) bool { + return cs2 == cs + }) } -// Start begins an MCP session by connecting the MCP client over its transport. +// Connect begins an MCP session by connecting to a server over the given +// transport, and initializing the session. // // 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) Start(ctx context.Context) (err error) { - defer func() { - if err != nil { - _ = c.Close() - } - }() - _, err = connect(ctx, c.transport, &c.opts.ConnectionOptions, c) +func (c *Client) Connect(ctx context.Context, t Transport) (cs *ClientSession, err error) { + cs, err = connect(ctx, t, c) if err != nil { - return err + return nil, err } params := &initializeParams{ ClientInfo: &implementation{Name: c.name, Version: c.version}, } - if err := call(ctx, c.conn, "initialize", params, &c.initializeResult); err != nil { - return err + if err := call(ctx, cs.conn, "initialize", params, &cs.initializeResult); err != nil { + _ = cs.Close() + return nil, err } - if err := c.conn.Notify(ctx, "notifications/initialized", &initializedParams{}); err != nil { - return err + if err := cs.conn.Notify(ctx, "notifications/initialized", &initializedParams{}); err != nil { + _ = cs.Close() + return nil, err } - return nil + return cs, nil +} + +// A ClientSession is a logical connection with an MCP server. Its +// methods can be used to send requests or notifications to the server. Create +// a session by calling [Client.Connect]. +// +// Call [ClientSession.Close] to close the connection, or await client +// termination with [ServerSession.Wait]. +type ClientSession struct { + conn *jsonrpc2.Connection + client *Client + initializeResult *initializeResult } // 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 (c *Client) Close() error { +func (c *ClientSession) 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 (c *Client) Wait() error { +func (c *ClientSession) Wait() error { return c.conn.Wait() } @@ -136,7 +149,7 @@ func (c *Client) listRoots(_ context.Context, _ *ListRootsParams) (*ListRootsRes }, nil } -func (c *Client) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) { +func (c *ClientSession) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) { // TODO: when we switch to ClientSessions, use a copy of the server's dispatch function, or // maybe just add another type parameter. // @@ -148,34 +161,34 @@ func (c *Client) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) return struct{}{}, nil case "roots/list": // ListRootsParams happens to be unused. - return c.listRoots(ctx, nil) + return c.client.listRoots(ctx, nil) } return nil, jsonrpc2.ErrNotHandled } // Ping makes an MCP "ping" request to the server. -func (c *Client) Ping(ctx context.Context, params *PingParams) error { +func (c *ClientSession) Ping(ctx context.Context, params *PingParams) error { return call(ctx, c.conn, "ping", params, nil) } // ListPrompts lists prompts that are currently available on the server. -func (c *Client) ListPrompts(ctx context.Context, params *ListPromptsParams) (*ListPromptsResult, error) { +func (c *ClientSession) ListPrompts(ctx context.Context, params *ListPromptsParams) (*ListPromptsResult, error) { return standardCall[ListPromptsResult](ctx, c.conn, "prompts/list", params) } // GetPrompt gets a prompt from the server. -func (c *Client) GetPrompt(ctx context.Context, params *GetPromptParams) (*GetPromptResult, error) { +func (c *ClientSession) GetPrompt(ctx context.Context, params *GetPromptParams) (*GetPromptResult, error) { return standardCall[GetPromptResult](ctx, c.conn, "prompts/get", params) } // ListTools lists tools that are currently available on the server. -func (c *Client) ListTools(ctx context.Context, params *ListToolsParams) (*ListToolsResult, error) { +func (c *ClientSession) ListTools(ctx context.Context, params *ListToolsParams) (*ListToolsResult, error) { return standardCall[ListToolsResult](ctx, c.conn, "tools/list", params) } // CallTool calls the tool with the given name and arguments. // Pass a [CallToolOptions] to provide additional request fields. -func (c *Client) CallTool(ctx context.Context, name string, args map[string]any, opts *CallToolOptions) (_ *CallToolResult, err error) { +func (c *ClientSession) CallTool(ctx context.Context, name string, args map[string]any, opts *CallToolOptions) (_ *CallToolResult, err error) { defer func() { if err != nil { err = fmt.Errorf("calling tool %q: %w", name, err) @@ -201,12 +214,12 @@ type CallToolOptions struct { } // ListResources lists the resources that are currently available on the server. -func (c *Client) ListResources(ctx context.Context, params *ListResourcesParams) (*ListResourcesResult, error) { +func (c *ClientSession) ListResources(ctx context.Context, params *ListResourcesParams) (*ListResourcesResult, error) { return standardCall[ListResourcesResult](ctx, c.conn, "resources/list", params) } // ReadResource ask the server to read a resource and return its contents. -func (c *Client) ReadResource(ctx context.Context, params *ReadResourceParams) (*ReadResourceResult, error) { +func (c *ClientSession) ReadResource(ctx context.Context, params *ReadResourceParams) (*ReadResourceResult, error) { return standardCall[ReadResourceResult](ctx, c.conn, "resources/read", params) } diff --git a/mcp/cmd_test.go b/mcp/cmd_test.go index fd048f4..922dbd2 100644 --- a/mcp/cmd_test.go +++ b/mcp/cmd_test.go @@ -32,7 +32,7 @@ func runServer() { server := mcp.NewServer("greeter", "v0.0.1", nil) server.AddTools(mcp.NewTool("greet", "say hi", SayHi)) - if err := server.Run(ctx, mcp.NewStdIOTransport(), nil); err != nil { + if err := server.Run(ctx, mcp.NewStdIOTransport()); err != nil { log.Fatal(err) } } @@ -48,11 +48,12 @@ func TestCmdTransport(t *testing.T) { cmd := exec.Command(exe) cmd.Env = append(os.Environ(), runAsServer+"=true") - client := mcp.NewClient("client", "v0.0.1", mcp.NewCommandTransport(cmd), nil) - if err := client.Start(ctx); err != nil { + client := mcp.NewClient("client", "v0.0.1", nil) + session, err := client.Connect(ctx, mcp.NewCommandTransport(cmd)) + if err != nil { log.Fatal(err) } - got, err := client.CallTool(ctx, "greet", map[string]any{"name": "user"}, nil) + got, err := session.CallTool(ctx, "greet", map[string]any{"name": "user"}, nil) if err != nil { log.Fatal(err) } @@ -62,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 := client.Close(); err != nil { + if err := session.Close(); err != nil { t.Fatalf("closing server: %v", err) } } diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 0a87f93..1e8447c 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -33,7 +33,7 @@ func sayHi(ctx context.Context, cc *ServerSession, v hiParams) ([]*Content, erro func TestEndToEnd(t *testing.T) { ctx := context.Background() - ct, st := NewInMemoryTransport() + ct, st := NewInMemoryTransports() s := NewServer("testServer", "v1.0.0", nil) @@ -63,7 +63,7 @@ func TestEndToEnd(t *testing.T) { ) // Connect the server. - ss, err := s.Connect(ctx, st, nil) + ss, err := s.Connect(ctx, st) if err != nil { t.Fatal(err) } @@ -81,19 +81,20 @@ func TestEndToEnd(t *testing.T) { clientWG.Done() }() - c := NewClient("testClient", "v1.0.0", ct, nil) + c := NewClient("testClient", "v1.0.0", nil) c.AddRoots(&Root{URI: "file:///root"}) // Connect the client. - if err := c.Start(ctx); err != nil { + cs, err := c.Connect(ctx, ct) + if err != nil { t.Fatal(err) } - if err := c.Ping(ctx, nil); err != nil { + if err := cs.Ping(ctx, nil); err != nil { t.Fatalf("ping failed: %v", err) } t.Run("prompts", func(t *testing.T) { - res, err := c.ListPrompts(ctx, nil) + res, err := cs.ListPrompts(ctx, nil) if err != nil { t.Errorf("prompts/list failed: %v", err) } @@ -109,7 +110,7 @@ func TestEndToEnd(t *testing.T) { t.Fatalf("prompts/list mismatch (-want +got):\n%s", diff) } - gotReview, err := c.GetPrompt(ctx, &GetPromptParams{Name: "code_review", Arguments: map[string]string{"Code": "1+1"}}) + gotReview, err := cs.GetPrompt(ctx, &GetPromptParams{Name: "code_review", Arguments: map[string]string{"Code": "1+1"}}) if err != nil { t.Fatal(err) } @@ -124,13 +125,13 @@ func TestEndToEnd(t *testing.T) { t.Errorf("prompts/get 'code_review' mismatch (-want +got):\n%s", diff) } - if _, err := c.GetPrompt(ctx, &GetPromptParams{Name: "fail"}); err == nil || !strings.Contains(err.Error(), failure.Error()) { + if _, err := cs.GetPrompt(ctx, &GetPromptParams{Name: "fail"}); err == nil || !strings.Contains(err.Error(), failure.Error()) { t.Errorf("fail returned unexpected error: got %v, want containing %v", err, failure) } }) t.Run("tools", func(t *testing.T) { - res, err := c.ListTools(ctx, nil) + res, err := cs.ListTools(ctx, nil) if err != nil { t.Errorf("tools/list failed: %v", err) } @@ -160,7 +161,7 @@ func TestEndToEnd(t *testing.T) { t.Fatalf("tools/list mismatch (-want +got):\n%s", diff) } - gotHi, err := c.CallTool(ctx, "greet", map[string]any{"name": "user"}, nil) + gotHi, err := cs.CallTool(ctx, "greet", map[string]any{"name": "user"}, nil) if err != nil { t.Fatal(err) } @@ -171,7 +172,7 @@ func TestEndToEnd(t *testing.T) { t.Errorf("tools/call 'greet' mismatch (-want +got):\n%s", diff) } - gotFail, err := c.CallTool(ctx, "fail", map[string]any{}, nil) + gotFail, err := cs.CallTool(ctx, "fail", map[string]any{}, nil) // 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 { @@ -212,7 +213,7 @@ func TestEndToEnd(t *testing.T) { &ServerResource{resource1, readHandler}, &ServerResource{resource2, readHandler}) - lrres, err := c.ListResources(ctx, nil) + lrres, err := cs.ListResources(ctx, nil) if err != nil { t.Fatal(err) } @@ -228,7 +229,7 @@ func TestEndToEnd(t *testing.T) { {"file:///nonexistent.txt", ""}, // TODO(jba): add resource template cases when we implement them } { - rres, err := c.ReadResource(ctx, &ReadResourceParams{URI: tt.uri}) + rres, err := cs.ReadResource(ctx, &ReadResourceParams{URI: tt.uri}) if err != nil { var werr *jsonrpc2.WireError if errors.As(err, &werr) && werr.Code == codeResourceNotFound { @@ -267,7 +268,7 @@ func TestEndToEnd(t *testing.T) { }) // Disconnect. - c.Close() + cs.Close() clientWG.Wait() // After disconnecting, neither client nor server should have any @@ -282,26 +283,27 @@ 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 ...*ServerTool) (*ServerSession, *Client) { +func basicConnection(t *testing.T, tools ...*ServerTool) (*ServerSession, *ClientSession) { t.Helper() ctx := context.Background() - ct, st := NewInMemoryTransport() + ct, st := NewInMemoryTransports() s := NewServer("testServer", "v1.0.0", nil) // The 'greet' tool says hi. s.AddTools(tools...) - ss, err := s.Connect(ctx, st, nil) + ss, err := s.Connect(ctx, st) if err != nil { t.Fatal(err) } - c := NewClient("testClient", "v1.0.0", ct, nil) - if err := c.Start(ctx); err != nil { + c := NewClient("testClient", "v1.0.0", nil) + cs, err := c.Connect(ctx, ct) + if err != nil { t.Fatal(err) } - return ss, c + return ss, cs } func TestServerClosing(t *testing.T) { @@ -329,28 +331,29 @@ func TestServerClosing(t *testing.T) { func TestBatching(t *testing.T) { ctx := context.Background() - ct, st := NewInMemoryTransport() + ct, st := NewInMemoryTransports() s := NewServer("testServer", "v1.0.0", nil) - _, err := s.Connect(ctx, st, nil) + _, err := s.Connect(ctx, st) if err != nil { t.Fatal(err) } - c := NewClient("testClient", "v1.0.0", ct, nil) + c := NewClient("testClient", "v1.0.0", 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.Start(ctx); err != nil { + cs, err := c.Connect(ctx, ct) + if err != nil { t.Fatal(err) } - defer c.Close() + defer cs.Close() errs := make(chan error, batchSize) for i := range batchSize { go func() { - _, err := c.ListTools(ctx, nil) + _, err := cs.ListTools(ctx, nil) errs <- err }() time.Sleep(2 * time.Millisecond) diff --git a/mcp/server.go b/mcp/server.go index 9c27100..3b5bac9 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -30,7 +30,7 @@ type Server struct { prompts *featureSet[*ServerPrompt] tools *featureSet[*ServerTool] resources *featureSet[*ServerResource] - conns []*ServerSession + sessions []*ServerSession } // ServerOptions is used to configure behavior of the server. @@ -160,7 +160,7 @@ func (s *Server) RemoveResources(uris ...string) { // Sessions returns an iterator that yields the current set of server sessions. func (s *Server) Sessions() iter.Seq[*ServerSession] { s.mu.Lock() - clients := slices.Clone(s.conns) + clients := slices.Clone(s.sessions) s.mu.Unlock() return slices.Values(clients) } @@ -248,8 +248,8 @@ func (s *Server) readResource(ctx context.Context, ss *ServerSession, params *Re // Run runs the server over the given transport, which must be persistent. // // Run blocks until the client terminates the connection. -func (s *Server) Run(ctx context.Context, t Transport, opts *ConnectionOptions) error { - ss, err := s.Connect(ctx, t, opts) +func (s *Server) Run(ctx context.Context, t Transport) error { + ss, err := s.Connect(ctx, t) if err != nil { return err } @@ -261,7 +261,7 @@ func (s *Server) Run(ctx context.Context, t Transport, opts *ConnectionOptions) func (s *Server) bind(conn *jsonrpc2.Connection) *ServerSession { cc := &ServerSession{conn: conn, server: s} s.mu.Lock() - s.conns = append(s.conns, cc) + s.sessions = append(s.sessions, cc) s.mu.Unlock() return cc } @@ -271,7 +271,7 @@ func (s *Server) bind(conn *jsonrpc2.Connection) *ServerSession { func (s *Server) disconnect(cc *ServerSession) { s.mu.Lock() defer s.mu.Unlock() - s.conns = slices.DeleteFunc(s.conns, func(cc2 *ServerSession) bool { + s.sessions = slices.DeleteFunc(s.sessions, func(cc2 *ServerSession) bool { return cc2 == cc }) } @@ -282,8 +282,8 @@ func (s *Server) disconnect(cc *ServerSession) { // It returns a connection object that may be used to terminate the connection // (with [Connection.Close]), or await client termination (with // [Connection.Wait]). -func (s *Server) Connect(ctx context.Context, t Transport, opts *ConnectionOptions) (*ServerSession, error) { - return connect(ctx, t, opts, s) +func (s *Server) Connect(ctx context.Context, t Transport) (*ServerSession, error) { + return connect(ctx, t, s) } // A ServerSession is a logical connection from a single MCP client. Its @@ -302,18 +302,18 @@ type ServerSession struct { } // Ping makes an MCP "ping" request to the client. -func (cc *ServerSession) Ping(ctx context.Context, _ *PingParams) error { - return call(ctx, cc.conn, "ping", nil, nil) +func (ss *ServerSession) Ping(ctx context.Context, _ *PingParams) error { + return call(ctx, ss.conn, "ping", nil, nil) } -func (cc *ServerSession) ListRoots(ctx context.Context, params *ListRootsParams) (*ListRootsResult, error) { - return standardCall[ListRootsResult](ctx, cc.conn, "roots/list", params) +func (ss *ServerSession) ListRoots(ctx context.Context, params *ListRootsParams) (*ListRootsResult, error) { + return standardCall[ListRootsResult](ctx, ss.conn, "roots/list", params) } -func (cc *ServerSession) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) { - cc.mu.Lock() - initialized := cc.initialized - cc.mu.Unlock() +func (ss *ServerSession) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) { + ss.mu.Lock() + initialized := ss.initialized + ss.mu.Unlock() // From the spec: // "The client SHOULD NOT send requests other than pings before the server @@ -332,39 +332,39 @@ func (cc *ServerSession) handle(ctx context.Context, req *jsonrpc2.Request) (any switch req.Method { case "initialize": - return dispatch(ctx, cc, req, cc.initialize) + return dispatch(ctx, ss, req, ss.initialize) case "ping": // The spec says that 'ping' expects an empty object result. return struct{}{}, nil case "prompts/list": - return dispatch(ctx, cc, req, cc.server.listPrompts) + return dispatch(ctx, ss, req, ss.server.listPrompts) case "prompts/get": - return dispatch(ctx, cc, req, cc.server.getPrompt) + return dispatch(ctx, ss, req, ss.server.getPrompt) case "tools/list": - return dispatch(ctx, cc, req, cc.server.listTools) + return dispatch(ctx, ss, req, ss.server.listTools) case "tools/call": - return dispatch(ctx, cc, req, cc.server.callTool) + return dispatch(ctx, ss, req, ss.server.callTool) case "resources/list": - return dispatch(ctx, cc, req, cc.server.listResources) + return dispatch(ctx, ss, req, ss.server.listResources) case "resources/read": - return dispatch(ctx, cc, req, cc.server.readResource) + return dispatch(ctx, ss, req, ss.server.readResource) case "notifications/initialized": } return nil, jsonrpc2.ErrNotHandled } -func (cc *ServerSession) initialize(ctx context.Context, _ *ServerSession, params *initializeParams) (*initializeResult, error) { - cc.mu.Lock() - cc.initializeParams = params - cc.mu.Unlock() +func (ss *ServerSession) initialize(ctx context.Context, _ *ServerSession, params *initializeParams) (*initializeResult, error) { + ss.mu.Lock() + ss.initializeParams = params + ss.mu.Unlock() // Mark the connection as initialized when this method exits. TODO: // Technically, the server should not be considered initialized until it has @@ -372,9 +372,9 @@ func (cc *ServerSession) initialize(ctx context.Context, _ *ServerSession, param // connection to implement that easily. In any case, once we've initialized // here, we can handle requests. defer func() { - cc.mu.Lock() - cc.initialized = true - cc.mu.Unlock() + ss.mu.Lock() + ss.initialized = true + ss.mu.Unlock() }() return &initializeResult{ @@ -388,10 +388,10 @@ func (cc *ServerSession) initialize(ctx context.Context, _ *ServerSession, param ListChanged: false, // not yet supported }, }, - Instructions: cc.server.opts.Instructions, + Instructions: ss.server.opts.Instructions, ServerInfo: &implementation{ - Name: cc.server.name, - Version: cc.server.version, + Name: ss.server.name, + Version: ss.server.version, }, }, nil } @@ -399,13 +399,13 @@ func (cc *ServerSession) initialize(ctx context.Context, _ *ServerSession, param // Close performs a graceful shutdown of the connection, preventing new // requests from being handled, and waiting for ongoing requests to return. // Close then terminates the connection. -func (cc *ServerSession) Close() error { - return cc.conn.Close() +func (ss *ServerSession) Close() error { + return ss.conn.Close() } // Wait waits for the connection to be closed by the client. -func (cc *ServerSession) Wait() error { - return cc.conn.Wait() +func (ss *ServerSession) Wait() error { + return ss.conn.Wait() } // dispatch turns a strongly type request handler into a jsonrpc2 handler. diff --git a/mcp/server_example_test.go b/mcp/server_example_test.go index 85fccd1..2be9035 100644 --- a/mcp/server_example_test.go +++ b/mcp/server_example_test.go @@ -24,28 +24,29 @@ func SayHi(ctx context.Context, cc *mcp.ServerSession, params *SayHiParams) ([]* func ExampleServer() { ctx := context.Background() - clientTransport, serverTransport := mcp.NewInMemoryTransport() + clientTransport, serverTransport := mcp.NewInMemoryTransports() server := mcp.NewServer("greeter", "v0.0.1", nil) server.AddTools(mcp.NewTool("greet", "say hi", SayHi)) - serverSession, err := server.Connect(ctx, serverTransport, nil) + serverSession, err := server.Connect(ctx, serverTransport) if err != nil { log.Fatal(err) } - client := mcp.NewClient("client", "v0.0.1", clientTransport, nil) - if err := client.Start(ctx); err != nil { + client := mcp.NewClient("client", "v0.0.1", nil) + clientSession, err := client.Connect(ctx, clientTransport) + if err != nil { log.Fatal(err) } - res, err := client.CallTool(ctx, "greet", map[string]any{"name": "user"}, nil) + res, err := clientSession.CallTool(ctx, "greet", map[string]any{"name": "user"}, nil) if err != nil { log.Fatal(err) } fmt.Println(res.Content[0].Text) - client.Close() + clientSession.Close() serverSession.Wait() // Output: Hi user diff --git a/mcp/sse.go b/mcp/sse.go index c1f923c..81a5130 100644 --- a/mcp/sse.go +++ b/mcp/sse.go @@ -229,7 +229,7 @@ func (h *SSEHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) { // TODO(hxjiang): getServer returns nil will panic. server := h.getServer(req) - ss, err := server.Connect(req.Context(), transport, nil) + ss, err := server.Connect(req.Context(), transport) if err != nil { http.Error(w, "connection failed", http.StatusInternalServerError) return diff --git a/mcp/sse_example_test.go b/mcp/sse_example_test.go index ca66cd8..102fce5 100644 --- a/mcp/sse_example_test.go +++ b/mcp/sse_example_test.go @@ -34,13 +34,14 @@ func ExampleSSEHandler() { ctx := context.Background() transport := mcp.NewSSEClientTransport(httpServer.URL) - client := mcp.NewClient("test", "v1.0.0", transport, nil) - if err := client.Start(ctx); err != nil { + client := mcp.NewClient("test", "v1.0.0", nil) + cs, err := client.Connect(ctx, transport) + if err != nil { log.Fatal(err) } - defer client.Close() + defer cs.Close() - res, err := client.CallTool(ctx, "add", map[string]any{"x": 1, "y": 2}, nil) + res, err := cs.CallTool(ctx, "add", map[string]any{"x": 1, "y": 2}, nil) if err != nil { log.Fatal(err) } diff --git a/mcp/sse_test.go b/mcp/sse_test.go index 8cd71a1..8b90958 100644 --- a/mcp/sse_test.go +++ b/mcp/sse_test.go @@ -35,15 +35,16 @@ func TestSSEServer(t *testing.T) { clientTransport := NewSSEClientTransport(httpServer.URL) - c := NewClient("testClient", "v1.0.0", clientTransport, nil) - if err := c.Start(ctx); err != nil { + c := NewClient("testClient", "v1.0.0", nil) + cs, err := c.Connect(ctx, clientTransport) + if err != nil { t.Fatal(err) } - if err := c.Ping(ctx, nil); err != nil { + if err := cs.Ping(ctx, nil); err != nil { t.Fatal(err) } - cc := <-conns - gotHi, err := c.CallTool(ctx, "greet", map[string]any{"name": "user"}, nil) + ss := <-conns + gotHi, err := cs.CallTool(ctx, "greet", map[string]any{"name": "user"}, nil) if err != nil { t.Fatal(err) } @@ -57,11 +58,11 @@ func TestSSEServer(t *testing.T) { // Test that closing either end of the connection terminates the other // end. if closeServerFirst { - c.Close() - cc.Wait() + cs.Close() + ss.Wait() } else { - cc.Close() - c.Wait() + ss.Close() + cs.Wait() } }) } diff --git a/mcp/transport.go b/mcp/transport.go index 1cbc1d2..f69299d 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -41,13 +41,6 @@ type Stream interface { io.Closer } -// ConnectionOptions configures the behavior of an individual client<->server -// connection. -type ConnectionOptions struct { - SessionID string // if set, the session ID - Logger io.Writer // if set, write RPC logs -} - // A StdIOTransport is a [Transport] that communicates over stdin/stdout using // newline-delimited JSON. type StdIOTransport struct { @@ -76,9 +69,9 @@ type InMemoryTransport struct { ioTransport } -// NewInMemoryTransport returns two InMemoryTransports that connect to each +// NewInMemoryTransports returns two InMemoryTransports that connect to each // other. -func NewInMemoryTransport() (*InMemoryTransport, *InMemoryTransport) { +func NewInMemoryTransports() (*InMemoryTransport, *InMemoryTransport) { c1, c2 := net.Pipe() return &InMemoryTransport{ioTransport{c1}}, &InMemoryTransport{ioTransport{c2}} } @@ -93,11 +86,7 @@ type binder[T handler] interface { disconnect(T) } -func connect[H handler](ctx context.Context, t Transport, opts *ConnectionOptions, b binder[H]) (H, error) { - if opts == nil { - opts = new(ConnectionOptions) - } - +func connect[H handler](ctx context.Context, t Transport, b binder[H]) (H, error) { var zero H stream, err := t.Connect(ctx) if err != nil { @@ -105,11 +94,6 @@ func connect[H handler](ctx context.Context, t Transport, opts *ConnectionOption } // If logging is configured, write message logs. reader, writer := jsonrpc2.Reader(stream), jsonrpc2.Writer(stream) - if opts.Logger != nil { - reader = loggingReader(opts.Logger, reader) - writer = loggingWriter(opts.Logger, writer) - } - var ( h H preempter canceller @@ -178,54 +162,66 @@ func call(ctx context.Context, conn *jsonrpc2.Connection, method string, params, return nil } -// The helpers below are used to bind transports to jsonrpc2. - -// A readerFunc implements jsonrpc2.Reader.Read. -type readerFunc func(context.Context) (jsonrpc2.Message, int64, error) - -func (f readerFunc) Read(ctx context.Context) (jsonrpc2.Message, int64, error) { - return f(ctx) +// A LoggingTransport is a [Transport] that delegates to another transport, +// writing RPC logs to an io.Writer. +type LoggingTransport struct { + delegate Transport + w io.Writer } -// A writerFunc implements jsonrpc2.Writer.Write. -type writerFunc func(context.Context, jsonrpc2.Message) (int64, error) +// NewLoggingTransport creates a new LoggingTransport that delegates to the +// provided transport, writing RPC logs to the provided io.Writer. +func NewLoggingTransport(delegate Transport, w io.Writer) *LoggingTransport { + return &LoggingTransport{delegate, w} +} -func (f writerFunc) Write(ctx context.Context, msg jsonrpc2.Message) (int64, error) { - return f(ctx, msg) +// Connect connects the underlying transport, returning a [Stream] that writes +// logs to the configured destination. +func (t *LoggingTransport) Connect(ctx context.Context) (Stream, error) { + delegate, err := t.delegate.Connect(ctx) + if err != nil { + return nil, err + } + return &loggingStream{delegate, t.w}, nil +} + +type loggingStream struct { + delegate Stream + w io.Writer } // loggingReader is a stream middleware that logs incoming messages. -func loggingReader(w io.Writer, delegate jsonrpc2.Reader) jsonrpc2.Reader { - return readerFunc(func(ctx context.Context) (jsonrpc2.Message, int64, error) { - msg, n, err := delegate.Read(ctx) +func (s *loggingStream) Read(ctx context.Context) (jsonrpc2.Message, int64, error) { + msg, n, err := s.delegate.Read(ctx) + if err != nil { + fmt.Fprintf(s.w, "read error: %v", err) + } else { + data, err := jsonrpc2.EncodeMessage(msg) if err != nil { - fmt.Fprintf(w, "read error: %v", err) - } else { - data, err := jsonrpc2.EncodeMessage(msg) - if err != nil { - fmt.Fprintf(w, "LoggingFramer: failed to marshal: %v", err) - } - fmt.Fprintf(w, "read: %s", string(data)) + fmt.Fprintf(s.w, "LoggingTransport: failed to marshal: %v", err) } - return msg, n, err - }) + fmt.Fprintf(s.w, "read: %s", string(data)) + } + return msg, n, err } // loggingWriter is a stream middleware that logs outgoing messages. -func loggingWriter(w io.Writer, delegate jsonrpc2.Writer) jsonrpc2.Writer { - return writerFunc(func(ctx context.Context, msg jsonrpc2.Message) (int64, error) { - n, err := delegate.Write(ctx, msg) +func (s *loggingStream) Write(ctx context.Context, msg jsonrpc2.Message) (int64, error) { + n, err := s.delegate.Write(ctx, msg) + if err != nil { + fmt.Fprintf(s.w, "write error: %v", err) + } else { + data, err := jsonrpc2.EncodeMessage(msg) if err != nil { - fmt.Fprintf(w, "write error: %v", err) - } else { - data, err := jsonrpc2.EncodeMessage(msg) - if err != nil { - fmt.Fprintf(w, "LoggingFramer: failed to marshal: %v", err) - } - fmt.Fprintf(w, "write: %s", string(data)) + fmt.Fprintf(s.w, "LoggingTransport: failed to marshal: %v", err) } - return n, err - }) + fmt.Fprintf(s.w, "write: %s", string(data)) + } + return n, err +} + +func (s *loggingStream) Close() error { + return s.delegate.Close() } // A rwc binds an io.ReadCloser and io.WriteCloser together to create an