modelcontextprotocol/go-sdk
Watch
1
0
Fork
You've already forked go-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

internal/mcp: reinstate ClientSession

Browse source 
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 <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Sam Thanawalla <samthanawalla@google.com>
Reviewed-by: Jonathan Amsterdam <jba@google.com>
This commit is contained in:
Rob Findley 2025-05-14 17:53:48 +00:00 • committed by Robert Findley
commit a8b98da480
11 changed files with 227 additions and 196 deletions
Download patch file Download diff file Expand all files Collapse all files

25
design/design.md
View file

@ -34,7 +34,6 @@ compatible with mcp-go, translating between them should be straightforward in
most cases. most cases.
(Later, we will provide a detailed translation guide.) (Later, we will provide a detailed translation guide.)
# Requirements # Requirements
These may be obvious, but it's worthwhile to define goals for an official MCP 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) func (*StreamableClientTransport) Connect(context.Context) (Stream, error)
``` ```
Finally, we also provide an in-memory transport, for scenarios such as testing, Finally, we also provide a couple of transport implementations for special scenarios.
where the MCP client and server are in the same process. 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 ```go
// An InMemoryTransport is a [Transport] that communicates over an in-memory
// network connection, using newline-delimited JSON.
type InMemoryTransport struct { /* ... */ } 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 **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] func (*ClientSession) ResourceTemplates(context.Context, *ListResourceTemplatesParams) iter.Seq2[ResourceTemplate, error]
``` ```
### Middleware ### Middleware
We provide a mechanism to add MCP-level middleware on the server side, which runs after the 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, In our design, each resource and resource template is associated with a function that reads it,
with this signature: with this signature:
```go ```go
type ResourceHandler func(context.Context, *ServerSession, *ReadResourceParams) (*ReadResourceResult, error) type ResourceHandler func(context.Context, *ServerSession, *ReadResourceParams) (*ReadResourceResult, error)
``` ```
The arguments include the `ServerSession` so the handler can observe the client's roots. 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` 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 or `NewBlobResourceContents`. If the handler omits the URI or MIME type, the server will populate them from the
resource. resource.
The `ServerResource` and `ServerResourceTemplate` types hold the association between the resource and its handler: The `ServerResource` and `ServerResourceTemplate` types hold the association between the resource and its handler:
```go ```go
type ServerResource struct { type ServerResource struct {
Resource Resource Resource Resource
@ -939,12 +952,14 @@ The `ReadResource` method finds a resource or resource template matching the arg
its assocated handler. its assocated handler.
To read files from the local filesystem, we recommend using `FileResourceHandler` to construct a handler: To read files from the local filesystem, we recommend using `FileResourceHandler` to construct a handler:
```go ```go
// FileResourceHandler returns a ResourceHandler that reads paths using dir as a root directory. // FileResourceHandler returns a ResourceHandler that reads paths using dir as a root directory.
// It protects against path traversal attacks. // 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. // 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 func (*Server) FileResourceHandler(dir string) ResourceHandler
``` ```
Here is an example: Here is an example:
```go ```go

4
examples/hello/main.go
View file

@ -50,8 +50,8 @@ func main() {
}) })
http.ListenAndServe(*httpAddr, handler) http.ListenAndServe(*httpAddr, handler)
} else { } else {
opts := &mcp.ConnectionOptions{Logger: os.Stderr} t := mcp.NewLoggingTransport(mcp.NewStdIOTransport(), os.Stderr)
if err := server.Run(context.Background(), mcp.NewStdIOTransport(), opts); err != nil { if err := server.Run(context.Background(), t); err != nil {
fmt.Fprintf(os.Stderr, "Server failed: %v", err) fmt.Fprintf(os.Stderr, "Server failed: %v", err)
} }
} }

109
mcp/client.go
View file

@ -16,29 +16,26 @@ import (
) )
// A Client is an MCP client, which may be connected to an MCP server // 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 { type Client struct {
name string name string
version string version string
transport Transport opts ClientOptions
opts ClientOptions mu sync.Mutex
mu sync.Mutex roots *featureSet[*Root]
conn *jsonrpc2.Connection sessions []*ClientSession
roots *featureSet[*Root]
initializeResult *initializeResult
} }
// NewClient creates a new Client. // 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. // 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{ c := &Client{
name: name, name: name,
version: version, version: version,
transport: t, roots: newFeatureSet(func(r *Root) string { return r.URI }),
roots: newFeatureSet(func(r *Root) string { return r.URI }),
} }
if opts != nil { if opts != nil {
c.opts = *opts c.opts = *opts
@ -48,63 +45,79 @@ func NewClient(name, version string, t Transport, opts *ClientOptions) *Client {
// ClientOptions configures the behavior of the client. // ClientOptions configures the behavior of the client.
type ClientOptions struct { 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]. // 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() c.mu.Lock()
defer c.mu.Unlock() defer c.mu.Unlock()
c.conn = conn c.sessions = append(c.sessions, cs)
return c return cs
} }
// disconnect implements the binder[*Client] interface, so that // disconnect implements the binder[*Client] interface, so that
// Clients can be connected using [connect]. // Clients can be connected using [connect].
func (c *Client) disconnect(*Client) { func (c *Client) disconnect(cs *ClientSession) {
// Do nothing. In particular, do not set conn to nil: it needs to exist so it can c.mu.Lock()
// return an error. 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 // 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 // when it is no longer needed. However, if the connection is closed by the
// server, calls or notifications will return an error wrapping // server, calls or notifications will return an error wrapping
// [ErrConnectionClosed]. // [ErrConnectionClosed].
func (c *Client) Start(ctx context.Context) (err error) { func (c *Client) Connect(ctx context.Context, t Transport) (cs *ClientSession, err error) {
defer func() { cs, err = connect(ctx, t, c)
if err != nil {
_ = c.Close()
}
}()
_, err = connect(ctx, c.transport, &c.opts.ConnectionOptions, c)
if err != nil { if err != nil {
return err return nil, err
} }
params := &initializeParams{ params := &initializeParams{
ClientInfo: &implementation{Name: c.name, Version: c.version}, ClientInfo: &implementation{Name: c.name, Version: c.version},
} }
if err := call(ctx, c.conn, "initialize", params, &c.initializeResult); err != nil { if err := call(ctx, cs.conn, "initialize", params, &cs.initializeResult); err != nil {
return err _ = cs.Close()
return nil, err
} }
if err := c.conn.Notify(ctx, "notifications/initialized", &initializedParams{}); err != nil { if err := cs.conn.Notify(ctx, "notifications/initialized", &initializedParams{}); err != nil {
return err _ = 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 // Close performs a graceful close of the connection, preventing new requests
// from being handled, and waiting for ongoing requests to return. Close then // from being handled, and waiting for ongoing requests to return. Close then
// terminates the connection. // terminates the connection.
func (c *Client) Close() error { func (c *ClientSession) Close() error {
return c.conn.Close() return c.conn.Close()
} }
// Wait waits for the connection to be closed by the server. // Wait waits for the connection to be closed by the server.
// Generally, clients should be responsible for closing the connection. // Generally, clients should be responsible for closing the connection.
func (c *Client) Wait() error { func (c *ClientSession) Wait() error {
return c.conn.Wait() return c.conn.Wait()
} }
@ -136,7 +149,7 @@ func (c *Client) listRoots(_ context.Context, _ *ListRootsParams) (*ListRootsRes
}, nil }, 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 // TODO: when we switch to ClientSessions, use a copy of the server's dispatch function, or
// maybe just add another type parameter. // 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 return struct{}{}, nil
case "roots/list": case "roots/list":
// ListRootsParams happens to be unused. // ListRootsParams happens to be unused.
return c.listRoots(ctx, nil) return c.client.listRoots(ctx, nil)
} }
return nil, jsonrpc2.ErrNotHandled return nil, jsonrpc2.ErrNotHandled
} }
// Ping makes an MCP "ping" request to the server. // 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) return call(ctx, c.conn, "ping", params, nil)
} }
// ListPrompts lists prompts that are currently available on the server. // 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) return standardCall[ListPromptsResult](ctx, c.conn, "prompts/list", params)
} }
// GetPrompt gets a prompt from the server. // 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) return standardCall[GetPromptResult](ctx, c.conn, "prompts/get", params)
} }
// ListTools lists tools that are currently available on the server. // 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) return standardCall[ListToolsResult](ctx, c.conn, "tools/list", params)
} }
// CallTool calls the tool with the given name and arguments. // CallTool calls the tool with the given name and arguments.
// Pass a [CallToolOptions] to provide additional request fields. // 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() { defer func() {
if err != nil { if err != nil {
err = fmt.Errorf("calling tool %q: %w", name, err) 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. // 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) return standardCall[ListResourcesResult](ctx, c.conn, "resources/list", params)
} }
// ReadResource ask the server to read a resource and return its contents. // 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) return standardCall[ReadResourceResult](ctx, c.conn, "resources/read", params)
} }

11
mcp/cmd_test.go
View file

@ -32,7 +32,7 @@ func runServer() {
server := mcp.NewServer("greeter", "v0.0.1", nil) server := mcp.NewServer("greeter", "v0.0.1", nil)
server.AddTools(mcp.NewTool("greet", "say hi", SayHi)) 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) log.Fatal(err)
} }
} }
@ -48,11 +48,12 @@ func TestCmdTransport(t *testing.T) {
cmd := exec.Command(exe) cmd := exec.Command(exe)
cmd.Env = append(os.Environ(), runAsServer+"=true") cmd.Env = append(os.Environ(), runAsServer+"=true")
client := mcp.NewClient("client", "v0.0.1", mcp.NewCommandTransport(cmd), nil) client := mcp.NewClient("client", "v0.0.1", nil)
if err := client.Start(ctx); err != nil { session, err := client.Connect(ctx, mcp.NewCommandTransport(cmd))
if err != nil {
log.Fatal(err) 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 { if err != nil {
log.Fatal(err) log.Fatal(err)
} }
@ -62,7 +63,7 @@ func TestCmdTransport(t *testing.T) {
if diff := cmp.Diff(want, got); diff != "" { if diff := cmp.Diff(want, got); diff != "" {
t.Errorf("greet returned unexpected content (-want +got):\n%s", 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) t.Fatalf("closing server: %v", err)
} }
} }

55
mcp/mcp_test.go
View file

@ -33,7 +33,7 @@ func sayHi(ctx context.Context, cc *ServerSession, v hiParams) ([]*Content, erro
func TestEndToEnd(t *testing.T) { func TestEndToEnd(t *testing.T) {
ctx := context.Background() ctx := context.Background()
ct, st := NewInMemoryTransport() ct, st := NewInMemoryTransports()
s := NewServer("testServer", "v1.0.0", nil) s := NewServer("testServer", "v1.0.0", nil)
@ -63,7 +63,7 @@ func TestEndToEnd(t *testing.T) {
) )
// Connect the server. // Connect the server.
ss, err := s.Connect(ctx, st, nil) ss, err := s.Connect(ctx, st)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
@ -81,19 +81,20 @@ func TestEndToEnd(t *testing.T) {
clientWG.Done() clientWG.Done()
}() }()
c := NewClient("testClient", "v1.0.0", ct, nil) c := NewClient("testClient", "v1.0.0", nil)
c.AddRoots(&Root{URI: "file:///root"}) c.AddRoots(&Root{URI: "file:///root"})
// Connect the client. // Connect the client.
if err := c.Start(ctx); err != nil { cs, err := c.Connect(ctx, ct)
if err != nil {
t.Fatal(err) 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.Fatalf("ping failed: %v", err)
} }
t.Run("prompts", func(t *testing.T) { t.Run("prompts", func(t *testing.T) {
res, err := c.ListPrompts(ctx, nil) res, err := cs.ListPrompts(ctx, nil)
if err != nil { if err != nil {
t.Errorf("prompts/list failed: %v", err) 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) 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 { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
@ -124,13 +125,13 @@ func TestEndToEnd(t *testing.T) {
t.Errorf("prompts/get 'code_review' mismatch (-want +got):\n%s", diff) 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.Errorf("fail returned unexpected error: got %v, want containing %v", err, failure)
} }
}) })
t.Run("tools", func(t *testing.T) { t.Run("tools", func(t *testing.T) {
res, err := c.ListTools(ctx, nil) res, err := cs.ListTools(ctx, nil)
if err != nil { if err != nil {
t.Errorf("tools/list failed: %v", err) 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) 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 { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
@ -171,7 +172,7 @@ func TestEndToEnd(t *testing.T) {
t.Errorf("tools/call 'greet' mismatch (-want +got):\n%s", diff) 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 // Counter-intuitively, when a tool fails, we don't expect an RPC error for
// call tool: instead, the failure is embedded in the result. // call tool: instead, the failure is embedded in the result.
if err != nil { if err != nil {
@ -212,7 +213,7 @@ func TestEndToEnd(t *testing.T) {
&ServerResource{resource1, readHandler}, &ServerResource{resource1, readHandler},
&ServerResource{resource2, readHandler}) &ServerResource{resource2, readHandler})
lrres, err := c.ListResources(ctx, nil) lrres, err := cs.ListResources(ctx, nil)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
@ -228,7 +229,7 @@ func TestEndToEnd(t *testing.T) {
{"file:///nonexistent.txt", ""}, {"file:///nonexistent.txt", ""},
// TODO(jba): add resource template cases when we implement them // 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 { if err != nil {
var werr *jsonrpc2.WireError var werr *jsonrpc2.WireError
if errors.As(err, &werr) && werr.Code == codeResourceNotFound { if errors.As(err, &werr) && werr.Code == codeResourceNotFound {
@ -267,7 +268,7 @@ func TestEndToEnd(t *testing.T) {
}) })
// Disconnect. // Disconnect.
c.Close() cs.Close()
clientWG.Wait() clientWG.Wait()
// After disconnecting, neither client nor server should have any // 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 // The caller should cancel either the client connection or server connection
// when the connections are no longer needed. // 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() t.Helper()
ctx := context.Background() ctx := context.Background()
ct, st := NewInMemoryTransport() ct, st := NewInMemoryTransports()
s := NewServer("testServer", "v1.0.0", nil) s := NewServer("testServer", "v1.0.0", nil)
// The 'greet' tool says hi. // The 'greet' tool says hi.
s.AddTools(tools...) s.AddTools(tools...)
ss, err := s.Connect(ctx, st, nil) ss, err := s.Connect(ctx, st)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
c := NewClient("testClient", "v1.0.0", ct, nil) c := NewClient("testClient", "v1.0.0", nil)
if err := c.Start(ctx); err != nil { cs, err := c.Connect(ctx, ct)
if err != nil {
t.Fatal(err) t.Fatal(err)
} }
return ss, c return ss, cs
} }
func TestServerClosing(t *testing.T) { func TestServerClosing(t *testing.T) {
@ -329,28 +331,29 @@ func TestServerClosing(t *testing.T) {
func TestBatching(t *testing.T) { func TestBatching(t *testing.T) {
ctx := context.Background() ctx := context.Background()
ct, st := NewInMemoryTransport() ct, st := NewInMemoryTransports()
s := NewServer("testServer", "v1.0.0", nil) s := NewServer("testServer", "v1.0.0", nil)
_, err := s.Connect(ctx, st, nil) _, err := s.Connect(ctx, st)
if err != nil { if err != nil {
t.Fatal(err) 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 // 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. // 'initialize' to block. Therefore, we can only test with a size of 1.
const batchSize = 1 const batchSize = 1
BatchSize(ct, batchSize) BatchSize(ct, batchSize)
if err := c.Start(ctx); err != nil { cs, err := c.Connect(ctx, ct)
if err != nil {
t.Fatal(err) t.Fatal(err)
} }
defer c.Close() defer cs.Close()
errs := make(chan error, batchSize) errs := make(chan error, batchSize)
for i := range batchSize { for i := range batchSize {
go func() { go func() {
_, err := c.ListTools(ctx, nil) _, err := cs.ListTools(ctx, nil)
errs <- err errs <- err
}() }()
time.Sleep(2 * time.Millisecond) time.Sleep(2 * time.Millisecond)

74
mcp/server.go
View file

@ -30,7 +30,7 @@ type Server struct {
prompts *featureSet[*ServerPrompt] prompts *featureSet[*ServerPrompt]
tools *featureSet[*ServerTool] tools *featureSet[*ServerTool]
resources *featureSet[*ServerResource] resources *featureSet[*ServerResource]
conns []*ServerSession sessions []*ServerSession
} }
// ServerOptions is used to configure behavior of the server. // 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. // Sessions returns an iterator that yields the current set of server sessions.
func (s *Server) Sessions() iter.Seq[*ServerSession] { func (s *Server) Sessions() iter.Seq[*ServerSession] {
s.mu.Lock() s.mu.Lock()
clients := slices.Clone(s.conns) clients := slices.Clone(s.sessions)
s.mu.Unlock() s.mu.Unlock()
return slices.Values(clients) 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 runs the server over the given transport, which must be persistent.
// //
// Run blocks until the client terminates the connection. // Run blocks until the client terminates the connection.
func (s *Server) Run(ctx context.Context, t Transport, opts *ConnectionOptions) error { func (s *Server) Run(ctx context.Context, t Transport) error {
ss, err := s.Connect(ctx, t, opts) ss, err := s.Connect(ctx, t)
if err != nil { if err != nil {
return err 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 { func (s *Server) bind(conn *jsonrpc2.Connection) *ServerSession {
cc := &ServerSession{conn: conn, server: s} cc := &ServerSession{conn: conn, server: s}
s.mu.Lock() s.mu.Lock()
s.conns = append(s.conns, cc) s.sessions = append(s.sessions, cc)
s.mu.Unlock() s.mu.Unlock()
return cc return cc
} }
@ -271,7 +271,7 @@ func (s *Server) bind(conn *jsonrpc2.Connection) *ServerSession {
func (s *Server) disconnect(cc *ServerSession) { func (s *Server) disconnect(cc *ServerSession) {
s.mu.Lock() s.mu.Lock()
defer s.mu.Unlock() 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 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 // It returns a connection object that may be used to terminate the connection
// (with [Connection.Close]), or await client termination (with // (with [Connection.Close]), or await client termination (with
// [Connection.Wait]). // [Connection.Wait]).
func (s *Server) Connect(ctx context.Context, t Transport, opts *ConnectionOptions) (*ServerSession, error) { func (s *Server) Connect(ctx context.Context, t Transport) (*ServerSession, error) {
return connect(ctx, t, opts, s) return connect(ctx, t, s)
} }
// A ServerSession is a logical connection from a single MCP client. Its // 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. // Ping makes an MCP "ping" request to the client.
func (cc *ServerSession) Ping(ctx context.Context, _ *PingParams) error { func (ss *ServerSession) Ping(ctx context.Context, _ *PingParams) error {
return call(ctx, cc.conn, "ping", nil, nil) return call(ctx, ss.conn, "ping", nil, nil)
} }
func (cc *ServerSession) ListRoots(ctx context.Context, params *ListRootsParams) (*ListRootsResult, error) { func (ss *ServerSession) ListRoots(ctx context.Context, params *ListRootsParams) (*ListRootsResult, error) {
return standardCall[ListRootsResult](ctx, cc.conn, "roots/list", params) return standardCall[ListRootsResult](ctx, ss.conn, "roots/list", params)
} }
func (cc *ServerSession) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) { func (ss *ServerSession) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) {
cc.mu.Lock() ss.mu.Lock()
initialized := cc.initialized initialized := ss.initialized
cc.mu.Unlock() ss.mu.Unlock()
// From the spec: // From the spec:
// "The client SHOULD NOT send requests other than pings before the server // "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 { switch req.Method {
case "initialize": case "initialize":
return dispatch(ctx, cc, req, cc.initialize) return dispatch(ctx, ss, req, ss.initialize)
case "ping": case "ping":
// The spec says that 'ping' expects an empty object result. // The spec says that 'ping' expects an empty object result.
return struct{}{}, nil return struct{}{}, nil
case "prompts/list": case "prompts/list":
return dispatch(ctx, cc, req, cc.server.listPrompts) return dispatch(ctx, ss, req, ss.server.listPrompts)
case "prompts/get": case "prompts/get":
return dispatch(ctx, cc, req, cc.server.getPrompt) return dispatch(ctx, ss, req, ss.server.getPrompt)
case "tools/list": case "tools/list":
return dispatch(ctx, cc, req, cc.server.listTools) return dispatch(ctx, ss, req, ss.server.listTools)
case "tools/call": case "tools/call":
return dispatch(ctx, cc, req, cc.server.callTool) return dispatch(ctx, ss, req, ss.server.callTool)
case "resources/list": case "resources/list":
return dispatch(ctx, cc, req, cc.server.listResources) return dispatch(ctx, ss, req, ss.server.listResources)
case "resources/read": case "resources/read":
return dispatch(ctx, cc, req, cc.server.readResource) return dispatch(ctx, ss, req, ss.server.readResource)
case "notifications/initialized": case "notifications/initialized":
} }
return nil, jsonrpc2.ErrNotHandled return nil, jsonrpc2.ErrNotHandled
} }
func (cc *ServerSession) initialize(ctx context.Context, _ *ServerSession, params *initializeParams) (*initializeResult, error) { func (ss *ServerSession) initialize(ctx context.Context, _ *ServerSession, params *initializeParams) (*initializeResult, error) {
cc.mu.Lock() ss.mu.Lock()
cc.initializeParams = params ss.initializeParams = params
cc.mu.Unlock() ss.mu.Unlock()
// Mark the connection as initialized when this method exits. TODO: // Mark the connection as initialized when this method exits. TODO:
// Technically, the server should not be considered initialized until it has // 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 // connection to implement that easily. In any case, once we've initialized
// here, we can handle requests. // here, we can handle requests.
defer func() { defer func() {
cc.mu.Lock() ss.mu.Lock()
cc.initialized = true ss.initialized = true
cc.mu.Unlock() ss.mu.Unlock()
}() }()
return &initializeResult{ return &initializeResult{
@ -388,10 +388,10 @@ func (cc *ServerSession) initialize(ctx context.Context, _ *ServerSession, param
ListChanged: false, // not yet supported ListChanged: false, // not yet supported
}, },
}, },
Instructions: cc.server.opts.Instructions, Instructions: ss.server.opts.Instructions,
ServerInfo: &implementation{ ServerInfo: &implementation{
Name: cc.server.name, Name: ss.server.name,
Version: cc.server.version, Version: ss.server.version,
}, },
}, nil }, nil
} }
@ -399,13 +399,13 @@ func (cc *ServerSession) initialize(ctx context.Context, _ *ServerSession, param
// Close performs a graceful shutdown of the connection, preventing new // Close performs a graceful shutdown of the connection, preventing new
// requests from being handled, and waiting for ongoing requests to return. // requests from being handled, and waiting for ongoing requests to return.
// Close then terminates the connection. // Close then terminates the connection.
func (cc *ServerSession) Close() error { func (ss *ServerSession) Close() error {
return cc.conn.Close() return ss.conn.Close()
} }
// Wait waits for the connection to be closed by the client. // Wait waits for the connection to be closed by the client.
func (cc *ServerSession) Wait() error { func (ss *ServerSession) Wait() error {
return cc.conn.Wait() return ss.conn.Wait()
} }
// dispatch turns a strongly type request handler into a jsonrpc2 handler. // dispatch turns a strongly type request handler into a jsonrpc2 handler.

13
mcp/server_example_test.go
View file

@ -24,28 +24,29 @@ func SayHi(ctx context.Context, cc *mcp.ServerSession, params *SayHiParams) ([]*
func ExampleServer() { func ExampleServer() {
ctx := context.Background() ctx := context.Background()
clientTransport, serverTransport := mcp.NewInMemoryTransport() clientTransport, serverTransport := mcp.NewInMemoryTransports()
server := mcp.NewServer("greeter", "v0.0.1", nil) server := mcp.NewServer("greeter", "v0.0.1", nil)
server.AddTools(mcp.NewTool("greet", "say hi", SayHi)) server.AddTools(mcp.NewTool("greet", "say hi", SayHi))
serverSession, err := server.Connect(ctx, serverTransport, nil) serverSession, err := server.Connect(ctx, serverTransport)
if err != nil { if err != nil {
log.Fatal(err) log.Fatal(err)
} }
client := mcp.NewClient("client", "v0.0.1", clientTransport, nil) client := mcp.NewClient("client", "v0.0.1", nil)
if err := client.Start(ctx); err != nil { clientSession, err := client.Connect(ctx, clientTransport)
if err != nil {
log.Fatal(err) 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 { if err != nil {
log.Fatal(err) log.Fatal(err)
} }
fmt.Println(res.Content[0].Text) fmt.Println(res.Content[0].Text)
client.Close() clientSession.Close()
serverSession.Wait() serverSession.Wait()
// Output: Hi user // Output: Hi user

2
mcp/sse.go
View file

@ -229,7 +229,7 @@ func (h *SSEHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
// TODO(hxjiang): getServer returns nil will panic. // TODO(hxjiang): getServer returns nil will panic.
server := h.getServer(req) server := h.getServer(req)
ss, err := server.Connect(req.Context(), transport, nil) ss, err := server.Connect(req.Context(), transport)
if err != nil { if err != nil {
http.Error(w, "connection failed", http.StatusInternalServerError) http.Error(w, "connection failed", http.StatusInternalServerError)
return return

9
mcp/sse_example_test.go
View file

@ -34,13 +34,14 @@ func ExampleSSEHandler() {
ctx := context.Background() ctx := context.Background()
transport := mcp.NewSSEClientTransport(httpServer.URL) transport := mcp.NewSSEClientTransport(httpServer.URL)
client := mcp.NewClient("test", "v1.0.0", transport, nil) client := mcp.NewClient("test", "v1.0.0", nil)
if err := client.Start(ctx); err != nil { cs, err := client.Connect(ctx, transport)
if err != nil {
log.Fatal(err) 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 { if err != nil {
log.Fatal(err) log.Fatal(err)
} }

19
mcp/sse_test.go
View file

@ -35,15 +35,16 @@ func TestSSEServer(t *testing.T) {
clientTransport := NewSSEClientTransport(httpServer.URL) clientTransport := NewSSEClientTransport(httpServer.URL)
c := NewClient("testClient", "v1.0.0", clientTransport, nil) c := NewClient("testClient", "v1.0.0", nil)
if err := c.Start(ctx); err != nil { cs, err := c.Connect(ctx, clientTransport)
if err != nil {
t.Fatal(err) t.Fatal(err)
} }
if err := c.Ping(ctx, nil); err != nil { if err := cs.Ping(ctx, nil); err != nil {
t.Fatal(err) t.Fatal(err)
} }
cc := <-conns ss := <-conns
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 { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
@ -57,11 +58,11 @@ func TestSSEServer(t *testing.T) {
// Test that closing either end of the connection terminates the other // Test that closing either end of the connection terminates the other
// end. // end.
if closeServerFirst { if closeServerFirst {
c.Close() cs.Close()
cc.Wait() ss.Wait()
} else { } else {
cc.Close() ss.Close()
c.Wait() cs.Wait()
} }
}) })
} }

104
mcp/transport.go
View file

@ -41,13 +41,6 @@ type Stream interface {
io.Closer 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 // A StdIOTransport is a [Transport] that communicates over stdin/stdout using
// newline-delimited JSON. // newline-delimited JSON.
type StdIOTransport struct { type StdIOTransport struct {
@ -76,9 +69,9 @@ type InMemoryTransport struct {
ioTransport ioTransport
} }
// NewInMemoryTransport returns two InMemoryTransports that connect to each // NewInMemoryTransports returns two InMemoryTransports that connect to each
// other. // other.
func NewInMemoryTransport() (*InMemoryTransport, *InMemoryTransport) { func NewInMemoryTransports() (*InMemoryTransport, *InMemoryTransport) {
c1, c2 := net.Pipe() c1, c2 := net.Pipe()
return &InMemoryTransport{ioTransport{c1}}, &InMemoryTransport{ioTransport{c2}} return &InMemoryTransport{ioTransport{c1}}, &InMemoryTransport{ioTransport{c2}}
} }
@ -93,11 +86,7 @@ type binder[T handler] interface {
disconnect(T) disconnect(T)
} }
func connect[H handler](ctx context.Context, t Transport, opts *ConnectionOptions, b binder[H]) (H, error) { func connect[H handler](ctx context.Context, t Transport, b binder[H]) (H, error) {
if opts == nil {
opts = new(ConnectionOptions)
}
var zero H var zero H
stream, err := t.Connect(ctx) stream, err := t.Connect(ctx)
if err != nil { 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. // If logging is configured, write message logs.
reader, writer := jsonrpc2.Reader(stream), jsonrpc2.Writer(stream) reader, writer := jsonrpc2.Reader(stream), jsonrpc2.Writer(stream)
if opts.Logger != nil {
reader = loggingReader(opts.Logger, reader)
writer = loggingWriter(opts.Logger, writer)
}
var ( var (
h H h H
preempter canceller preempter canceller
@ -178,54 +162,66 @@ func call(ctx context.Context, conn *jsonrpc2.Connection, method string, params,
return nil return nil
} }
// The helpers below are used to bind transports to jsonrpc2. // A LoggingTransport is a [Transport] that delegates to another transport,
// writing RPC logs to an io.Writer.
// A readerFunc implements jsonrpc2.Reader.Read. type LoggingTransport struct {
type readerFunc func(context.Context) (jsonrpc2.Message, int64, error) delegate Transport
w io.Writer
func (f readerFunc) Read(ctx context.Context) (jsonrpc2.Message, int64, error) {
return f(ctx)
} }
// A writerFunc implements jsonrpc2.Writer.Write. // NewLoggingTransport creates a new LoggingTransport that delegates to the
type writerFunc func(context.Context, jsonrpc2.Message) (int64, error) // 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) { // Connect connects the underlying transport, returning a [Stream] that writes
return f(ctx, msg) // 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. // loggingReader is a stream middleware that logs incoming messages.
func loggingReader(w io.Writer, delegate jsonrpc2.Reader) jsonrpc2.Reader { func (s *loggingStream) Read(ctx context.Context) (jsonrpc2.Message, int64, error) {
return readerFunc(func(ctx context.Context) (jsonrpc2.Message, int64, error) { msg, n, err := s.delegate.Read(ctx)
msg, n, err := delegate.Read(ctx) if err != nil {
fmt.Fprintf(s.w, "read error: %v", err)
} else {
data, err := jsonrpc2.EncodeMessage(msg)
if err != nil { if err != nil {
fmt.Fprintf(w, "read error: %v", err) fmt.Fprintf(s.w, "LoggingTransport: failed to marshal: %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))
} }
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. // loggingWriter is a stream middleware that logs outgoing messages.
func loggingWriter(w io.Writer, delegate jsonrpc2.Writer) jsonrpc2.Writer { func (s *loggingStream) Write(ctx context.Context, msg jsonrpc2.Message) (int64, error) {
return writerFunc(func(ctx context.Context, msg jsonrpc2.Message) (int64, error) { n, err := s.delegate.Write(ctx, msg)
n, err := delegate.Write(ctx, msg) if err != nil {
fmt.Fprintf(s.w, "write error: %v", err)
} else {
data, err := jsonrpc2.EncodeMessage(msg)
if err != nil { if err != nil {
fmt.Fprintf(w, "write error: %v", err) fmt.Fprintf(s.w, "LoggingTransport: failed to marshal: %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))
} }
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 // A rwc binds an io.ReadCloser and io.WriteCloser together to create an