internal/mcp: sampling

Implement sampling (the CreateMessage method).

Change-Id: Ie44a671dd62831122b25cd7386b5bfcdece7ff60
Reviewed-on: https://go-review.googlesource.com/c/tools/+/673176
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Robert Findley <rfindley@google.com>

mcp/client.go

@@ -46,7 +46,11 @@ func NewClient(name, version string, opts *ClientOptions) *Client {
 }
 
 // ClientOptions configures the behavior of the client.
-type ClientOptions struct{}
+type ClientOptions struct {
+	// Handler for sampling.
+	// Called when a server calls CreateMessage.
+	CreateMessageHandler func(context.Context, *ClientSession, *CreateMessageParams) (*CreateMessageResult, error)
+}
 
 // bind implements the binder[*ClientSession] interface, so that Clients can
 // be connected using [connect].
@@ -83,8 +87,13 @@ func (c *Client) Connect(ctx context.Context, t Transport) (cs *ClientSession, e
 	if err != nil {
 		return nil, err
 	}
+	caps := &ClientCapabilities{}
+	if c.opts.CreateMessageHandler != nil {
+		caps.Sampling = &SamplingCapabilities{}
+	}
 	params := &InitializeParams{
-		ClientInfo: &implementation{Name: c.name, Version: c.version},
+		ClientInfo:   &implementation{Name: c.name, Version: c.version},
+		Capabilities: caps,
 	}
 	if err := call(ctx, cs.conn, "initialize", params, &cs.initializeResult); err != nil {
 		_ = cs.Close()
@@ -150,6 +159,14 @@ func (c *Client) listRoots(_ context.Context, _ *ClientSession, _ *ListRootsPara
 	}, nil
 }
 
+func (c *Client) createMessage(ctx context.Context, cs *ClientSession, params *CreateMessageParams) (*CreateMessageResult, error) {
+	if c.opts.CreateMessageHandler == nil {
+		// TODO: wrap or annotate this error? Pick a standard code?
+		return nil, &jsonrpc2.WireError{Code: CodeUnsupportedMethod, Message: "client does not support CreateMessage"}
+	}
+	return c.opts.CreateMessageHandler(ctx, cs, params)
+}
+
 // AddMiddleware wraps the client's current method handler using the provided
 // middleware. Middleware is applied from right to left, so that the first one
 // is executed first.
@@ -164,8 +181,9 @@ func (c *Client) AddMiddleware(middleware ...Middleware[ClientSession]) {
 
 // clientMethodInfos maps from the RPC method name to serverMethodInfos.
 var clientMethodInfos = map[string]methodInfo[ClientSession]{
-	"ping":       newMethodInfo(sessionMethod((*ClientSession).ping)),
+	"ping":                   newMethodInfo(sessionMethod((*ClientSession).ping)),
-	"roots/list": newMethodInfo(clientMethod((*Client).listRoots)),
+	"roots/list":             newMethodInfo(clientMethod((*Client).listRoots)),
+	"sampling/createMessage": newMethodInfo(clientMethod((*Client).createMessage)),
 	// TODO: notifications
 }
 

mcp/generate.go

@@ -68,7 +68,14 @@ var declarations = config{
 		Name:   "-",
 		Fields: config{"Params": {Name: "CancelledParams"}},
 	},
-	"ClientCapabilities": {},
+	"ClientCapabilities": {
+		Fields: config{"Sampling": {Name: "SamplingCapabilities"}},
+	},
+	"CreateMessageRequest": {
+		Name:   "-",
+		Fields: config{"Params": {Name: "CreateMessageParams"}},
+	},
+	"CreateMessageResult": {},
 	"GetPromptRequest": {
 		Name:   "-",
 		Fields: config{"Params": {Name: "GetPromptParams"}},
@@ -103,7 +110,9 @@ var declarations = config{
 		Name:   "-",
 		Fields: config{"Params": {Name: "ListToolsParams"}},
 	},
-	"ListToolsResult": {},
+	"ListToolsResult":  {},
+	"ModelHint":        {},
+	"ModelPreferences": {},
 	"PingRequest": {
 		Name:   "-",
 		Fields: config{"Params": {Name: "PingParams"}},
@@ -124,6 +133,8 @@ var declarations = config{
 	"Role":     {},
 	"Root":     {},
 
+	"SamplingCapabilities": {Substitute: "struct{}"},
+	"SamplingMessage":      {},
 	"ServerCapabilities": {
 		Name: "serverCapabilities",
 		Fields: config{
@@ -350,6 +361,11 @@ func writeType(w io.Writer, config *typeConfig, def *jsonschema.Schema, named ma
 					fieldTypeSchema = rs
 				}
 				needPointer := isStruct(fieldTypeSchema)
+				// Special case: there are no sampling capabilities defined, but
+				// we want it to be a struct for future expansion.
+				if !needPointer && name == "sampling" {
+					needPointer = true
+				}
 				if config != nil && config.Fields[export] != nil {
 					r := config.Fields[export]
 					if r.Substitute != "" {

mcp/mcp_test.go

@@ -85,7 +85,12 @@ func TestEndToEnd(t *testing.T) {
 		clientWG.Done()
 	}()
 
-	c := NewClient("testClient", "v1.0.0", nil)
+	opts := &ClientOptions{
+		CreateMessageHandler: func(context.Context, *ClientSession, *CreateMessageParams) (*CreateMessageResult, error) {
+			return &CreateMessageResult{Model: "aModel"}, nil
+		},
+	}
+	c := NewClient("testClient", "v1.0.0", opts)
 	rootAbs, err := filepath.Abs(filepath.FromSlash("testdata/files"))
 	if err != nil {
 		t.Fatal(err)
@@ -233,8 +238,7 @@ func TestEndToEnd(t *testing.T) {
 		} {
 			rres, err := cs.ReadResource(ctx, &ReadResourceParams{URI: tt.uri})
 			if err != nil {
-				var werr *jsonrpc2.WireError
+				if code := errorCode(err); code == CodeResourceNotFound {
-				if errors.As(err, &werr) && werr.Code == codeResourceNotFound {
 					if tt.mimeType != "" {
 						t.Errorf("%s: not found but expected it to be", tt.uri)
 					}
@@ -252,13 +256,7 @@ func TestEndToEnd(t *testing.T) {
 		}
 	})
 	t.Run("roots", func(t *testing.T) {
-		// Take the server's first ServerSession.
+		rootRes, err := ss.ListRoots(ctx, &ListRootsParams{})
-		var sc *ServerSession
-		for sc = range s.Sessions() {
-			break
-		}
-
-		rootRes, err := sc.ListRoots(ctx, &ListRootsParams{})
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -268,6 +266,16 @@ func TestEndToEnd(t *testing.T) {
 			t.Errorf("roots/list mismatch (-want +got):\n%s", diff)
 		}
 	})
+	t.Run("sampling", func(t *testing.T) {
+		// TODO: test that a client that doesn't have the handler returns CodeUnsupportedMethod.
+		res, err := ss.CreateMessage(ctx, &CreateMessageParams{})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if g, w := res.Model, "aModel"; g != w {
+			t.Errorf("got %q, want %q", g, w)
+		}
+	})
 
 	// Disconnect.
 	cs.Close()
@@ -280,6 +288,20 @@ func TestEndToEnd(t *testing.T) {
 	}
 }
 
+// errorCode returns the code associated with err.
+// If err is nil, it returns 0.
+// If there is no code, it returns -1.
+func errorCode(err error) int64 {
+	if err == nil {
+		return 0
+	}
+	var werr *jsonrpc2.WireError
+	if errors.As(err, &werr) {
+		return werr.Code
+	}
+	return -1
+}
+
 // basicConnection returns a new basic client-server connection configured with
 // the provided tools.
 //

mcp/protocol.go

@@ -78,8 +78,45 @@ type ClientCapabilities struct {
 		ListChanged bool `json:"listChanged,omitempty"`
 	} `json:"roots,omitempty"`
 	// Present if the client supports sampling from an LLM.
-	Sampling struct {
+	Sampling *SamplingCapabilities `json:"sampling,omitempty"`
-	} `json:"sampling,omitempty"`
+}
+
+type CreateMessageParams struct {
+	// A request to include context from one or more MCP servers (including the
+	// caller), to be attached to the prompt. The client MAY ignore this request.
+	IncludeContext string `json:"includeContext,omitempty"`
+	// The maximum number of tokens to sample, as requested by the server. The
+	// client MAY choose to sample fewer tokens than requested.
+	MaxTokens int64              `json:"maxTokens"`
+	Messages  []*SamplingMessage `json:"messages"`
+	// Optional metadata to pass through to the LLM provider. The format of this
+	// metadata is provider-specific.
+	Metadata struct {
+	} `json:"metadata,omitempty"`
+	// The server's preferences for which model to select. The client MAY ignore
+	// these preferences.
+	ModelPreferences *ModelPreferences `json:"modelPreferences,omitempty"`
+	StopSequences    []string          `json:"stopSequences,omitempty"`
+	// An optional system prompt the server wants to use for sampling. The client
+	// MAY modify or omit this prompt.
+	SystemPrompt string  `json:"systemPrompt,omitempty"`
+	Temperature  float64 `json:"temperature,omitempty"`
+}
+
+// The client's response to a sampling/create_message request from the server.
+// The client should inform the user before returning the sampled message, to
+// allow them to inspect the response (human in the loop) and decide whether to
+// allow the server to see it.
+type CreateMessageResult struct {
+	// This result property is reserved by the protocol to allow clients and servers
+	// to attach additional metadata to their responses.
+	Meta    map[string]json.RawMessage `json:"_meta,omitempty"`
+	Content *Content                   `json:"content"`
+	// The name of the model that generated the message.
+	Model string `json:"model"`
+	Role  Role   `json:"role"`
+	// The reason why sampling stopped, if known.
+	StopReason string `json:"stopReason,omitempty"`
 }
 
 type GetPromptParams struct {
@@ -204,6 +241,58 @@ type ListToolsResult struct {
 	Tools      []*Tool `json:"tools"`
 }
 
+// Hints to use for model selection.
+//
+// Keys not declared here are currently left unspecified by the spec and are up
+// to the client to interpret.
+type ModelHint struct {
+	// A hint for a model name.
+	//
+	// The client SHOULD treat this as a substring of a model name; for example: -
+	// `claude-3-5-sonnet` should match `claude-3-5-sonnet-20241022` - `sonnet`
+	// should match `claude-3-5-sonnet-20241022`, `claude-3-sonnet-20240229`, etc. -
+	// `claude` should match any Claude model
+	//
+	// The client MAY also map the string to a different provider's model name or a
+	// different model family, as long as it fills a similar niche; for example: -
+	// `gemini-1.5-flash` could match `claude-3-haiku-20240307`
+	Name string `json:"name,omitempty"`
+}
+
+// The server's preferences for model selection, requested of the client during
+// sampling.
+//
+// Because LLMs can vary along multiple dimensions, choosing the "best" model is
+// rarely straightforward. Different models excel in different areas—some are
+// faster but less capable, others are more capable but more expensive, and so
+// on. This interface allows servers to express their priorities across multiple
+// dimensions to help clients make an appropriate selection for their use case.
+//
+// These preferences are always advisory. The client MAY ignore them. It is also
+// up to the client to decide how to interpret these preferences and how to
+// balance them against other considerations.
+type ModelPreferences struct {
+	// How much to prioritize cost when selecting a model. A value of 0 means cost
+	// is not important, while a value of 1 means cost is the most important factor.
+	CostPriority float64 `json:"costPriority,omitempty"`
+	// Optional hints to use for model selection.
+	//
+	// If multiple hints are specified, the client MUST evaluate them in order (such
+	// that the first match is taken).
+	//
+	// The client SHOULD prioritize these hints over the numeric priorities, but MAY
+	// still use the priorities to select from ambiguous matches.
+	Hints []*ModelHint `json:"hints,omitempty"`
+	// How much to prioritize intelligence and capabilities when selecting a model.
+	// A value of 0 means intelligence is not important, while a value of 1 means
+	// intelligence is the most important factor.
+	IntelligencePriority float64 `json:"intelligencePriority,omitempty"`
+	// How much to prioritize sampling speed (latency) when selecting a model. A
+	// value of 0 means speed is not important, while a value of 1 means speed is
+	// the most important factor.
+	SpeedPriority float64 `json:"speedPriority,omitempty"`
+}
+
 type PingParams struct {
 	Meta struct {
 		// If specified, the caller is requesting out-of-band progress notifications for
@@ -297,6 +386,16 @@ type Root struct {
 	URI string `json:"uri"`
 }
 
+// Present if the client supports sampling from an LLM.
+type SamplingCapabilities struct {
+}
+
+// Describes a message issued to or received from an LLM API.
+type SamplingMessage struct {
+	Content *Content `json:"content"`
+	Role    Role     `json:"role"`
+}
+
 // Definition for a tool the client can call.
 type Tool struct {
 	// Optional additional tool information.

mcp/resource.go

@@ -32,20 +32,12 @@ type ResourceHandler func(context.Context, *ServerSession, *ReadResourceParams)
 // not be found.
 func ResourceNotFoundError(uri string) error {
 	return &jsonrpc2.WireError{
-		Code:    codeResourceNotFound,
+		Code:    CodeResourceNotFound,
 		Message: "Resource not found",
 		Data:    json.RawMessage(fmt.Sprintf(`{"uri":%q}`, uri)),
 	}
 }
 
-// The error code to return when a resource isn't found.
-// See https://modelcontextprotocol.io/specification/2025-03-26/server/resources#error-handling
-// However, the code they chose in in the wrong space
-// (see https://github.com/modelcontextprotocol/modelcontextprotocol/issues/509).
-// so we pick a different one, arbirarily for now (until they fix it).
-// The immediate problem is that jsonprc2 defines -32002 as "server closing".
-const codeResourceNotFound = -31002
-
 // readFileResource reads from the filesystem at a URI relative to dirFilepath, respecting
 // the roots.
 // dirFilepath and rootFilepaths are absolute filesystem paths.

mcp/server.go

@@ -328,6 +328,11 @@ func (ss *ServerSession) ListRoots(ctx context.Context, params *ListRootsParams)
 	return standardCall[ListRootsResult](ctx, ss.conn, "roots/list", params)
 }
 
+// CreateMessage sends a sampling request to the client.
+func (ss *ServerSession) CreateMessage(ctx context.Context, params *CreateMessageParams) (*CreateMessageResult, error) {
+	return standardCall[CreateMessageResult](ctx, ss.conn, "sampling/createMessage", params)
+}
+
 // AddMiddleware wraps the server's current method handler using the provided
 // middleware. Middleware is applied from right to left, so that the first one
 // is executed first.

mcp/shared.go

@@ -134,3 +134,16 @@ func sessionMethod[S ClientSession | ServerSession, P, R any](f func(*S, context
 		return f(sess, ctx, p)
 	}
 }
+
+// Error codes
+const (
+	// The error code to return when a resource isn't found.
+	// See https://modelcontextprotocol.io/specification/2025-03-26/server/resources#error-handling
+	// However, the code they chose in in the wrong space
+	// (see https://github.com/modelcontextprotocol/modelcontextprotocol/issues/509).
+	// so we pick a different one, arbirarily for now (until they fix it).
+	// The immediate problem is that jsonprc2 defines -32002 as "server closing".
+	CodeResourceNotFound = -31002
+	// The error code if the method exists and was called properly, but the peer does not support it.
+	CodeUnsupportedMethod = -31001
+)