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 <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
Robert Findley
parent
0c7d8180a6
commit
a8b98da480
11 changed files with 227 additions and 196 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
109
mcp/client.go
109
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)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
}
|
||||
})
|
||||
}
|
||||
|
|
|
|||
104
mcp/transport.go
104
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
|
||||
|
|
|
|||
Loading…
Reference in a new issue