internal/mcp: implement ping, and test request interleaving

Also validate that requests observe the lifecycle rules defined in the
spec.

Change-Id: I302282f81e053e76926e0bb295f7dbf8d526f02d
Reviewed-on: https://go-review.googlesource.com/c/tools/+/667575
Reviewed-by: Jonathan Amsterdam <[email protected]>
Commit-Queue: Robert Findley <[email protected]>
Auto-Submit: Robert Findley <[email protected]>
TryBot-Bypass: Robert Findley <[email protected]>

Author: findleyr

internal/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 (

internal/mcp/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

internal/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)

internal/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, &params); err != nil {
-		var zero TResult
-		return zero, err
+		return nil, err
 	}
 	return f(ctx, conn, params)
 }

internal/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 {

internal/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()}}),

internal/mcp/util.go

@@ -4,7 +4,9 @@
 
 package mcp
 
-import "crypto/rand"
+import (
+	"crypto/rand"
+)
 
 func assert(cond bool, msg string) {
 	if !cond {