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: sampling

Browse source 
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>
This commit is contained in:
Jonathan Amsterdam 2025-05-16 08:53:48 -04:00
commit 9bcc6c758f
7 changed files with 192 additions and 27 deletions
Download patch file Download diff file Expand all files Collapse all files

26
mcp/client.go
View file

@ -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)),
"roots/list": newMethodInfo(clientMethod((*Client).listRoots)),
"ping": newMethodInfo(sessionMethod((*ClientSession).ping)),
"roots/list": newMethodInfo(clientMethod((*Client).listRoots)),
"sampling/createMessage": newMethodInfo(clientMethod((*Client).createMessage)),
// TODO: notifications
}

20
mcp/generate.go
View file

@ -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 != "" {

42
mcp/mcp_test.go
View file

@ -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 errors.As(err, &werr) && werr.Code == codeResourceNotFound {
if code := errorCode(err); 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.
var sc *ServerSession
for sc = range s.Sessions() {
break
}
rootRes, err := sc.ListRoots(ctx, &ListRootsParams{})
rootRes, err := ss.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.
//

103
mcp/protocol.go
View file

@ -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 {
} `json:"sampling,omitempty"`
Sampling *SamplingCapabilities `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.

10
mcp/resource.go
View file

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

5
mcp/server.go
View file

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

13
mcp/shared.go
View file

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