diff --git a/examples/hello/main.go b/examples/hello/main.go index c099e57..d661bc9 100644 --- a/examples/hello/main.go +++ b/examples/hello/main.go @@ -20,7 +20,7 @@ type SayHiParams struct { Name string `json:"name" mcp:"the name to say hi to"` } -func SayHi(ctx context.Context, params *SayHiParams) ([]mcp.Content, error) { +func SayHi(ctx context.Context, cc *mcp.ClientConnection, params *SayHiParams) ([]mcp.Content, error) { return []mcp.Content{ mcp.TextContent{Text: "Hi " + params.Name}, }, nil diff --git a/mcp/client.go b/mcp/client.go index 518b22e..1700561 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -81,9 +81,9 @@ func (c *Client) disconnect(sc *ServerConnection) { // Connect connects the MCP client over the given transport and initializes an // MCP session. // -// It returns a connection object that may be used to query the MCP server, -// terminate the connection (with [Connection.Close]), or await server -// termination (with [Connection.Wait]). +// It returns an initialized [ServerConnection] object that may be used to +// query the MCP server, terminate the connection (with [Connection.Close]), or +// await server termination (with [Connection.Wait]). // // Typically, it is the responsibility of the client to close the connection // when it is no longer needed. However, if the connection is closed by the @@ -105,7 +105,7 @@ func (c *Client) Connect(ctx context.Context, t Transport, opts *ConnectionOptio if err := call(ctx, sc.conn, "initialize", params, &sc.initializeResult); err != nil { return nil, err } - if err := sc.conn.Notify(ctx, "initialized", &protocol.InitializedParams{}); err != nil { + if err := sc.conn.Notify(ctx, "notifications/initialized", &protocol.InitializedParams{}); err != nil { return nil, err } return sc, nil @@ -135,11 +135,21 @@ func (cc *ServerConnection) Wait() error { } func (sc *ServerConnection) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) { + // No need to check that the connection is initialized, since we initialize + // it in Connect. switch req.Method { + case "ping": + // The spec says that 'ping' expects an empty object result. + return struct{}{}, nil } return nil, jsonrpc2.ErrNotHandled } +// Ping makes an MCP "ping" request to the server. +func (sc *ServerConnection) Ping(ctx context.Context) error { + return call(ctx, sc.conn, "ping", nil, nil) +} + // ListTools lists tools that are currently available on the server. func (sc *ServerConnection) ListTools(ctx context.Context) ([]protocol.Tool, error) { var ( diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 9e318b5..444bc7f 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -7,6 +7,7 @@ package mcp import ( "context" "errors" + "fmt" "slices" "strings" "sync" @@ -22,7 +23,10 @@ type hiParams struct { Name string } -func sayHi(_ context.Context, v hiParams) ([]Content, error) { +func sayHi(ctx context.Context, cc *ClientConnection, v hiParams) ([]Content, error) { + if err := cc.Ping(ctx); err != nil { + return nil, fmt.Errorf("ping failed: %v", err) + } return []Content{TextContent{Text: "hi " + v.Name}}, nil } @@ -37,7 +41,7 @@ func TestEndToEnd(t *testing.T) { // The 'fail' tool returns this error. failure := errors.New("mcp failure") - s.AddTools(MakeTool("fail", "just fail", func(context.Context, struct{}) ([]Content, error) { + s.AddTools(MakeTool("fail", "just fail", func(context.Context, *ClientConnection, struct{}) ([]Content, error) { return nil, failure })) @@ -67,10 +71,15 @@ func TestEndToEnd(t *testing.T) { if err != nil { t.Fatal(err) } + if got := slices.Collect(c.Servers()); len(got) != 1 { t.Errorf("after connection, Servers() has length %d, want 1", len(got)) } + if err := sc.Ping(ctx); err != nil { + t.Fatalf("ping failed: %v", err) + } + gotTools, err := sc.ListTools(ctx) if err != nil { t.Errorf("tools/list failed: %v", err) diff --git a/mcp/server.go b/mcp/server.go index 917d149..1cf0e90 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -82,7 +82,7 @@ func (s *Server) listTools(_ context.Context, _ *ClientConnection, params *proto return res, nil } -func (s *Server) callTool(ctx context.Context, _ *ClientConnection, params *protocol.CallToolParams) (*protocol.CallToolResult, error) { +func (s *Server) callTool(ctx context.Context, cc *ClientConnection, params *protocol.CallToolParams) (*protocol.CallToolResult, error) { s.mu.Lock() var tool *Tool if i := slices.IndexFunc(s.tools, func(t *Tool) bool { @@ -95,7 +95,7 @@ func (s *Server) callTool(ctx context.Context, _ *ClientConnection, params *prot if tool == nil { return nil, fmt.Errorf("%s: unknown tool %q", jsonrpc2.ErrInvalidParams, params.Name) } - return tool.Handler(ctx, params.Arguments) + return tool.Handler(ctx, cc, params.Arguments) } // Run runs the server over the given transport, which must be persistent. @@ -148,10 +148,31 @@ type ClientConnection struct { conn *jsonrpc2.Connection mu sync.Mutex - initializeParams *protocol.InitializeParams // set once initialize has been received + initializeParams *protocol.InitializeParams + initialized bool +} + +// Ping makes an MCP "ping" request to the client. +func (cc *ClientConnection) Ping(ctx context.Context) error { + return call(ctx, cc.conn, "ping", nil, nil) } func (cc *ClientConnection) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) { + cc.mu.Lock() + initialized := cc.initialized + cc.mu.Unlock() + + // From the spec: + // "The client SHOULD NOT send requests other than pings before the server + // has responded to the initialize request." + switch req.Method { + case "initialize", "ping": + default: + if !initialized { + return nil, fmt.Errorf("method %q is invalid during session ininitialization", req.Method) + } + } + // TODO: embed the incoming request ID in the ClientContext (or, more likely, // a wrapper around it), so that we can correlate responses and notifications // to the handler; this is required for the new session-based transport. @@ -160,6 +181,10 @@ func (cc *ClientConnection) handle(ctx context.Context, req *jsonrpc2.Request) ( case "initialize": return dispatch(ctx, cc, req, cc.initialize) + case "ping": + // The spec says that 'ping' expects an empty object result. + return struct{}{}, nil + case "tools/list": return dispatch(ctx, cc, req, cc.server.listTools) @@ -176,6 +201,17 @@ func (cc *ClientConnection) initialize(ctx context.Context, _ *ClientConnection, cc.initializeParams = params cc.mu.Unlock() + // Mark the connection as initialized when this method exits. TODO: + // Technically, the server should not be considered initialized until it has + // *responded*, but we don't have adequate visibility into the jsonrpc2 + // 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() + }() + return &protocol.InitializeResult{ // TODO(rfindley): support multiple protocol versions. ProtocolVersion: "2024-11-05", @@ -204,11 +240,14 @@ func (cc *ClientConnection) Wait() error { return cc.conn.Wait() } -func dispatch[TParams, TResult any](ctx context.Context, conn *ClientConnection, req *jsonrpc2.Request, f func(context.Context, *ClientConnection, TParams) (TResult, error)) (TResult, error) { +// dispatch turns a strongly type handler into a jsonrpc2 handler. +// +// Importantly, it returns nil if the handler returned an error, which is a +// requirement of the jsonrpc2 package. +func dispatch[TConn, TParams, TResult any](ctx context.Context, conn TConn, req *jsonrpc2.Request, f func(context.Context, TConn, TParams) (TResult, error)) (any, error) { var params TParams if err := json.Unmarshal(req.Params, ¶ms); err != nil { - var zero TResult - return zero, err + return nil, err } return f(ctx, conn, params) } diff --git a/mcp/sse_test.go b/mcp/sse_test.go index 80f6a64..13196ab 100644 --- a/mcp/sse_test.go +++ b/mcp/sse_test.go @@ -41,6 +41,9 @@ func TestSSEServer(t *testing.T) { if err != nil { t.Fatal(err) } + if err := sc.Ping(ctx); err != nil { + t.Fatal(err) + } cc := <-clients gotHi, err := sc.CallTool(ctx, "greet", hiParams{"user"}) if err != nil { diff --git a/mcp/tool.go b/mcp/tool.go index 0293326..452fbfd 100644 --- a/mcp/tool.go +++ b/mcp/tool.go @@ -13,7 +13,7 @@ import ( ) // A ToolHandler handles a call to tools/call. -type ToolHandler func(context.Context, json.RawMessage) (*protocol.CallToolResult, error) +type ToolHandler func(context.Context, *ClientConnection, json.RawMessage) (*protocol.CallToolResult, error) // A Tool is a tool definition that is bound to a tool handler. type Tool struct { @@ -29,17 +29,17 @@ type Tool struct { // It is the caller's responsibility that the handler request type can produce // a valid schema, as documented by [jsonschema.ForType]; otherwise, MakeTool // panics. -func MakeTool[TReq any](name, description string, handler func(context.Context, TReq) ([]Content, error)) *Tool { +func MakeTool[TReq any](name, description string, handler func(context.Context, *ClientConnection, TReq) ([]Content, error)) *Tool { schema, err := jsonschema.ForType[TReq]() if err != nil { panic(err) } - wrapped := func(ctx context.Context, args json.RawMessage) (*protocol.CallToolResult, error) { + wrapped := func(ctx context.Context, cc *ClientConnection, args json.RawMessage) (*protocol.CallToolResult, error) { var v TReq if err := unmarshalSchema(args, schema, &v); err != nil { return nil, err } - content, err := handler(ctx, v) + content, err := handler(ctx, cc, v) if err != nil { return &protocol.CallToolResult{ Content: marshalContent([]Content{TextContent{Text: err.Error()}}), diff --git a/mcp/util.go b/mcp/util.go index 94f983a..d929cf1 100644 --- a/mcp/util.go +++ b/mcp/util.go @@ -4,7 +4,9 @@ package mcp -import "crypto/rand" +import ( + "crypto/rand" +) func assert(cond bool, msg string) { if !cond {