diff --git a/mcp/client.go b/mcp/client.go index e1602d3..8de5050 100644 --- a/mcp/client.go +++ b/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)), - "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 } diff --git a/mcp/generate.go b/mcp/generate.go index feea820..e291f47 100644 --- a/mcp/generate.go +++ b/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 != "" { diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index a53a4f5..0165b98 100644 --- a/mcp/mcp_test.go +++ b/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 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. // diff --git a/mcp/protocol.go b/mcp/protocol.go index 0d085b4..12bfc35 100644 --- a/mcp/protocol.go +++ b/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 { - } `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. diff --git a/mcp/resource.go b/mcp/resource.go index 04f1a00..8265df3 100644 --- a/mcp/resource.go +++ b/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. diff --git a/mcp/server.go b/mcp/server.go index 9d35715..642b16d 100644 --- a/mcp/server.go +++ b/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. diff --git a/mcp/shared.go b/mcp/shared.go index 643fb13..632de61 100644 --- a/mcp/shared.go +++ b/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 +)