Skip to content

mcp: implement sampling with tools #699

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
findleyr wants to merge 2 commits into
base: main
Choose a base branch
from
+1,750 −25
Open

mcp: implement sampling with tools #699

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
41908dd
mcp: implement sampling with tools
findleyr Feb 2, 2026
9154420
mcp: address minor observations from self-review
findleyr Feb 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions docs/rough_edges.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,3 +45,17 @@ v2.
**Workaround**: to advertise no capabilities, set
`ServerOptions.Capabilities` or `ClientOptions.Capabilities` to an empty
`&ServerCapabilities{}` or `&ClientCapabilities{}` respectively.

- `CreateMessageResult.Content` is singular `Content`, but the 2025-11-25 spec
allows `content` to be a single block or an array (for parallel tool calls).
We added `CreateMessageResultWithTools` (with `Content []Content`) as a
workaround, matching the TypeScript SDK's approach. In v2,
`CreateMessageResult` should use `[]Content` directly. Similarly,
`SamplingMessage.Content` should become `[]Content` to support sending
multiple tool_result blocks in a single user message.

- We didn't actually need CallToolParams and CallToolParamsRaw, since even when
we're unmarshalling into a custom Go type (for the mcp.AddTool convenience
wrapper) we need to first unmarshal into a `map[string]any` in order to do
server-side validation of required fields. CallToolParams could have just had
a map[string]any.
14 changes: 14 additions & 0 deletions internal/docs/rough_edges.src.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,3 +44,17 @@ v2.
**Workaround**: to advertise no capabilities, set
`ServerOptions.Capabilities` or `ClientOptions.Capabilities` to an empty
`&ServerCapabilities{}` or `&ClientCapabilities{}` respectively.

- `CreateMessageResult.Content` is singular `Content`, but the 2025-11-25 spec
allows `content` to be a single block or an array (for parallel tool calls).
We added `CreateMessageResultWithTools` (with `Content []Content`) as a
workaround, matching the TypeScript SDK's approach. In v2,
`CreateMessageResult` should use `[]Content` directly. Similarly,
`SamplingMessage.Content` should become `[]Content` to support sending
multiple tool_result blocks in a single user message.

- We didn't actually need CallToolParams and CallToolParamsRaw, since even when
we're unmarshalling into a custom Go type (for the mcp.AddTool convenience
wrapper) we need to first unmarshal into a `map[string]any` in order to do
server-side validation of required fields. CallToolParams could have just had
a map[string]any.
58 changes: 48 additions & 10 deletions mcp/client.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,9 @@ func NewClient(impl *Implementation, options *ClientOptions) *Client {
}
options = nil // prevent reuse

if opts.CreateMessageHandler != nil && opts.CreateMessageWithToolsHandler != nil {
panic("cannot set both CreateMessageHandler and CreateMessageWithToolsHandler; use CreateMessageWithToolsHandler for tool support, or CreateMessageHandler for basic sampling")
}
if opts.Logger == nil { // ensure we have a logger
opts.Logger = ensureLogger(nil)
}
Expand All @@ -77,6 +80,19 @@ type ClientOptions struct {
// non nil value for [ClientCapabilities.Sampling], that value overrides the
// inferred capability.
CreateMessageHandler func(context.Context, *CreateMessageRequest) (*CreateMessageResult, error)
// CreateMessageWithToolsHandler handles incoming sampling/createMessage
// requests that may involve tool use. It returns
// [CreateMessageWithToolsResult], which supports array content for parallel
// tool calls.
//
// Setting this handler causes the client to advertise the sampling
// capability with tools support (sampling.tools). As with
// [CreateMessageHandler], [ClientOptions.Capabilities].Sampling overrides
// the inferred capability.
//
// It is a panic to set both CreateMessageHandler and
// CreateMessageWithToolsHandler.
CreateMessageWithToolsHandler func(context.Context, *CreateMessageWithToolsRequest) (*CreateMessageWithToolsResult, error)
// ElicitationHandler handles incoming requests for elicitation/create.
//
// Setting ElicitationHandler to a non-nil value automatically causes the
Expand Down Expand Up @@ -109,7 +125,16 @@ type ClientOptions struct {
// are set in the Capabilities field, their values override the inferred
// value.
//
// For example, to to configure elicitation modes:
// For example, to advertise sampling with tools and context support:
//
// Capabilities: &ClientCapabilities{
// Sampling: &SamplingCapabilities{
// Tools: &SamplingToolsCapabilities{},
// Context: &SamplingContextCapabilities{},
// },
// }
//
// Or to configure elicitation modes:
//
// Capabilities: &ClientCapabilities{
// Elicitation: &ElicitationCapabilities{
Expand All @@ -119,8 +144,7 @@ type ClientOptions struct {
// }
//
// Conversely, if Capabilities does not set a field (for example, if the
// Elicitation field is nil), the inferred elicitation capability will be
// used.
// Elicitation field is nil), the inferred capability will be used.
Capabilities *ClientCapabilities
// ElicitationCompleteHandler handles incoming notifications for notifications/elicitation/complete.
ElicitationCompleteHandler func(context.Context, *ElicitationCompleteNotificationRequest)
Expand Down Expand Up @@ -198,11 +222,14 @@ func (c *Client) capabilities(protocolVersion string) *ClientCapabilities {
caps.Roots = *caps.RootsV2
}

// Augment with sampling capability if handler is set.
if c.opts.CreateMessageHandler != nil {
// Augment with sampling capability if a handler is set.
if c.opts.CreateMessageHandler != nil || c.opts.CreateMessageWithToolsHandler != nil {
if caps.Sampling == nil {
caps.Sampling = &SamplingCapabilities{}
}
if c.opts.CreateMessageWithToolsHandler != nil && caps.Sampling.Tools == nil {
caps.Sampling.Tools = &SamplingToolsCapabilities{}
}
}

// Augment with elicitation capability if handler is set.
Expand Down Expand Up @@ -453,12 +480,23 @@ func (c *Client) listRoots(_ context.Context, req *ListRootsRequest) (*ListRoots
}, nil
}

func (c *Client) createMessage(ctx context.Context, req *CreateMessageRequest) (*CreateMessageResult, error) {
if c.opts.CreateMessageHandler == nil {
// TODO: wrap or annotate this error? Pick a standard code?
return nil, &jsonrpc.Error{Code: codeUnsupportedMethod, Message: "client does not support CreateMessage"}
func (c *Client) createMessage(ctx context.Context, req *CreateMessageWithToolsRequest) (*CreateMessageWithToolsResult, error) {
if c.opts.CreateMessageWithToolsHandler != nil {
return c.opts.CreateMessageWithToolsHandler(ctx, req)
}
if c.opts.CreateMessageHandler != nil {
// Downconvert the request for the basic handler.
baseReq := &CreateMessageRequest{
Session: req.Session,
Params: req.Params.toBase(),
}
res, err := c.opts.CreateMessageHandler(ctx, baseReq)
if err != nil {
return nil, err
}
return res.toWithTools(), nil
}
return c.opts.CreateMessageHandler(ctx, req)
return nil, &jsonrpc.Error{Code: codeUnsupportedMethod, Message: "client does not support CreateMessage"}
}

// urlElicitationMiddleware returns middleware that automatically handles URL elicitation
Expand Down
161 changes: 156 additions & 5 deletions mcp/content.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,10 @@ import (
)

// A Content is a [TextContent], [ImageContent], [AudioContent],
// [ResourceLink], or [EmbeddedResource].
// [ResourceLink], [EmbeddedResource], [ToolUseContent], or [ToolResultContent].
//
// Note: [ToolUseContent] and [ToolResultContent] are only valid in sampling
// message contexts (CreateMessageParams/CreateMessageResult).
type Content interface {
MarshalJSON() ([]byte, error)
fromWire(*wireContent)
Expand Down Expand Up @@ -183,6 +186,104 @@ func (c *EmbeddedResource) fromWire(wire *wireContent) {
c.Annotations = wire.Annotations
}

// ToolUseContent represents a request from the assistant to invoke a tool.
// This content type is only valid in sampling messages.
type ToolUseContent struct {
// ID is a unique identifier for this tool use, used to match with ToolResultContent.
ID string
// Name is the name of the tool to invoke.
Name string
// Input contains the tool arguments as a JSON object.
Input map[string]any
Meta Meta
}

func (c *ToolUseContent) MarshalJSON() ([]byte, error) {
input := c.Input
if input == nil {
input = map[string]any{}
}
wire := struct {
Type string `json:"type"`
ID string `json:"id"`
Name string `json:"name"`
Input map[string]any `json:"input"`
Meta Meta `json:"_meta,omitempty"`
}{
Type: "tool_use",
ID: c.ID,
Name: c.Name,
Input: input,
Meta: c.Meta,
}
return json.Marshal(wire)
}

func (c *ToolUseContent) fromWire(wire *wireContent) {
c.ID = wire.ID
c.Name = wire.Name
c.Input = wire.Input
c.Meta = wire.Meta
}

// ToolResultContent represents the result of a tool invocation.
// This content type is only valid in sampling messages with role "user".
type ToolResultContent struct {
// ToolUseID references the ID from the corresponding ToolUseContent.
ToolUseID string
// Content holds the unstructured result of the tool call.
Content []Content
// StructuredContent holds an optional structured result as a JSON object.
StructuredContent any
// IsError indicates whether the tool call ended in an error.
IsError bool
Meta Meta
}

func (c *ToolResultContent) MarshalJSON() ([]byte, error) {
// Marshal nested content
var contentWire []*wireContent
for _, content := range c.Content {
data, err := content.MarshalJSON()
if err != nil {
return nil, err
}
var w wireContent
if err := json.Unmarshal(data, &w); err != nil {
return nil, err
}
contentWire = append(contentWire, &w)
}
if contentWire == nil {
contentWire = []*wireContent{} // avoid JSON null
}

wire := struct {
Type string `json:"type"`
ToolUseID string `json:"toolUseId"`
Content []*wireContent `json:"content"`
StructuredContent any `json:"structuredContent,omitempty"`
IsError bool `json:"isError,omitempty"`
Meta Meta `json:"_meta,omitempty"`
}{
Type: "tool_result",
ToolUseID: c.ToolUseID,
Content: contentWire,
StructuredContent: c.StructuredContent,
IsError: c.IsError,
Meta: c.Meta,
}
return json.Marshal(wire)
}

func (c *ToolResultContent) fromWire(wire *wireContent) {
c.ToolUseID = wire.ToolUseID
c.StructuredContent = wire.StructuredContent
c.IsError = wire.IsError
c.Meta = wire.Meta
// Content is handled separately in contentFromWire due to nested content
}

// ResourceContents contains the contents of a specific resource or
// sub-resource.
type ResourceContents struct {
Expand Down Expand Up @@ -224,10 +325,9 @@ func (r *ResourceContents) MarshalJSON() ([]byte, error) {

// wireContent is the wire format for content.
// It represents the protocol types TextContent, ImageContent, AudioContent,
// ResourceLink, and EmbeddedResource.
// ResourceLink, EmbeddedResource, ToolUseContent, and ToolResultContent.
// The Type field distinguishes them. In the protocol, each type has a constant
// value for the field.
// At most one of Text, Data, Resource, and URI is non-zero.
type wireContent struct {
Type string `json:"type"`
Text string `json:"text,omitempty"`
Expand All @@ -242,10 +342,40 @@ type wireContent struct {
Meta Meta `json:"_meta,omitempty"`
Annotations *Annotations `json:"annotations,omitempty"`
Icons []Icon `json:"icons,omitempty"`
// Fields for ToolUseContent (type: "tool_use")
ID string `json:"id,omitempty"`
Input map[string]any `json:"input,omitempty"`
// Fields for ToolResultContent (type: "tool_result")
ToolUseID string `json:"toolUseId,omitempty"`
NestedContent []*wireContent `json:"content,omitempty"` // nested content for tool_result
StructuredContent any `json:"structuredContent,omitempty"`
IsError bool `json:"isError,omitempty"`
}

// unmarshalContent unmarshals JSON that is either a single content object or
// an array of content objects. A single object is wrapped in a one-element slice.
func unmarshalContent(raw json.RawMessage, allow map[string]bool) ([]Content, error) {
if len(raw) == 0 || string(raw) == "null" {
return nil, fmt.Errorf("nil content")
}
// Try array first, then fall back to single object.
var wires []*wireContent
if err := json.Unmarshal(raw, &wires); err == nil {
return contentsFromWire(wires, allow)
}
var wire wireContent
if err := json.Unmarshal(raw, &wire); err != nil {
return nil, err
}
c, err := contentFromWire(&wire, allow)
if err != nil {
return nil, err
}
return []Content{c}, nil
}

func contentsFromWire(wires []*wireContent, allow map[string]bool) ([]Content, error) {
var blocks []Content
blocks := make([]Content, 0, len(wires))
for _, wire := range wires {
block, err := contentFromWire(wire, allow)
if err != nil {
Expand Down Expand Up @@ -284,6 +414,27 @@ func contentFromWire(wire *wireContent, allow map[string]bool) (Content, error)
v := new(EmbeddedResource)
v.fromWire(wire)
return v, nil
case "tool_use":
v := new(ToolUseContent)
v.fromWire(wire)
return v, nil
case "tool_result":
v := new(ToolResultContent)
v.fromWire(wire)
// Handle nested content - tool_result content can contain text, image, audio,
// resource_link, and resource (same as CallToolResult.content)
if wire.NestedContent != nil {
toolResultContentAllow := map[string]bool{
"text": true, "image": true, "audio": true,
"resource_link": true, "resource": true,
}
nestedContent, err := contentsFromWire(wire.NestedContent, toolResultContentAllow)
if err != nil {
return nil, fmt.Errorf("tool_result nested content: %w", err)
}
v.Content = nestedContent
}
return v, nil
}
return nil, fmt.Errorf("internal error: unrecognized content type %s", wire.Type)
return nil, fmt.Errorf("unrecognized content type %q", wire.Type)
}
Loading
Loading