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: implement ping, and test request interleaving

Browse source 
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 <jba@google.com>
Commit-Queue: Robert Findley <rfindley@google.com>
Auto-Submit: Robert Findley <rfindley@google.com>
TryBot-Bypass: Robert Findley <rfindley@google.com>
This commit is contained in:
Rob Findley 2025-04-23 16:03:54 +00:00 • committed by Robert Findley
commit fe938f6eb7
6 changed files with 81 additions and 18 deletions
Download patch file Download diff file Expand all files Collapse all files

18
mcp/client.go
View file

@ -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 (

13
mcp/mcp_test.go
View file

@ -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)

51
mcp/server.go
View file

@ -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)
}

3
mcp/sse_test.go
View file

@ -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 {

8
mcp/tool.go
View file

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

4
mcp/util.go
View file

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