mcp: implement sampling with tools (#699)
Summary Add support for tool use within sampling requests, as described in the MCP 2025-11-25 spec's sampling.tools capability. New types Content types: - ToolUseContent — represents a tool invocation request from the assistant - ToolResultContent — represents the result of a tool invocation, with nested []Content for the result blocks Capability types: - SamplingCapabilities gains Tools and Context sub-fields - SamplingToolsCapabilities, SamplingContextCapabilities Tool-enabled sampling (parallel tool calls): - CreateMessageWithToolsParams with SamplingMessageV2 (supports array content) - CreateMessageWithToolsResult (supports array content) - ToolChoice — controls tool invocation mode (auto, required, none) - ServerSession.CreateMessageWithTools - ClientOptions.CreateMessageWithToolsHandler Design Following the TypeScript SDK's pattern, tool-enabled sampling uses separate types from basic sampling to avoid breaking the existing API. The basic CreateMessage/CreateMessageResult path is unchanged. Both paths share the same wire method (sampling/createMessage) and go through the method info table. The table uses the broader CreateMessageWithToolsResult type internally; CreateMessage downconverts (erroring if multiple content blocks are returned). This is documented as a rough edge to unify in v2. Setting CreateMessageWithToolsHandler automatically infers the sampling.tools capability. It is a panic to set both CreateMessageHandler and CreateMessageWithToolsHandler. References - https://modelcontextprotocol.io/specification/2025-11-25/client/sampling - TypeScript SDK: CreateMessageResultWithToolsSchema / CreateMessageRequestParamsWithTools --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Robert Findley
GitHub
parent
14869a46ee
commit
3cf9f9908d
9 changed files with 1641 additions and 38 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
@ -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{
|
||||
|
|
@ -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)
|
||||
|
|
@ -198,10 +222,13 @@ 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 = &SamplingToolsCapabilities{}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -453,12 +480,27 @@ 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)
|
||||
}
|
||||
return c.opts.CreateMessageHandler(ctx, req)
|
||||
if c.opts.CreateMessageHandler != nil {
|
||||
// Downconvert the request for the basic handler.
|
||||
baseParams, err := req.Params.toBase()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
baseReq := &CreateMessageRequest{
|
||||
Session: req.Session,
|
||||
Params: baseParams,
|
||||
}
|
||||
res, err := c.opts.CreateMessageHandler(ctx, baseReq)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return res.toWithTools(), nil
|
||||
}
|
||||
return nil, &jsonrpc.Error{Code: codeUnsupportedMethod, Message: "client does not support CreateMessage"}
|
||||
}
|
||||
|
||||
// urlElicitationMiddleware returns middleware that automatically handles URL elicitation
|
||||
|
|
|
|||
185
mcp/content.go
185
mcp/content.go
|
|
@ -13,7 +13,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)
|
||||
|
|
@ -182,6 +185,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 {
|
||||
|
|
@ -194,28 +295,55 @@ type ResourceContents struct {
|
|||
|
||||
// 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"`
|
||||
MIMEType string `json:"mimeType,omitempty"`
|
||||
Data []byte `json:"data,omitempty"`
|
||||
Resource *ResourceContents `json:"resource,omitempty"`
|
||||
URI string `json:"uri,omitempty"`
|
||||
Name string `json:"name,omitempty"`
|
||||
Title string `json:"title,omitempty"`
|
||||
Description string `json:"description,omitempty"`
|
||||
Size *int64 `json:"size,omitempty"`
|
||||
Meta Meta `json:"_meta,omitempty"`
|
||||
Annotations *Annotations `json:"annotations,omitempty"`
|
||||
Icons []Icon `json:"icons,omitempty"`
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text,omitempty"` // TextContent
|
||||
MIMEType string `json:"mimeType,omitempty"` // ImageContent, AudioContent, ResourceLink
|
||||
Data []byte `json:"data,omitempty"` // ImageContent, AudioContent
|
||||
Resource *ResourceContents `json:"resource,omitempty"` // EmbeddedResource
|
||||
URI string `json:"uri,omitempty"` // ResourceLink
|
||||
Name string `json:"name,omitempty"` // ResourceLink, ToolUseContent
|
||||
Title string `json:"title,omitempty"` // ResourceLink
|
||||
Description string `json:"description,omitempty"` // ResourceLink
|
||||
Size *int64 `json:"size,omitempty"` // ResourceLink
|
||||
Meta Meta `json:"_meta,omitempty"` // all types
|
||||
Annotations *Annotations `json:"annotations,omitempty"` // all types except ToolUseContent, ToolResultContent
|
||||
Icons []Icon `json:"icons,omitempty"` // ResourceLink
|
||||
ID string `json:"id,omitempty"` // ToolUseContent
|
||||
Input map[string]any `json:"input,omitempty"` // ToolUseContent
|
||||
ToolUseID string `json:"toolUseId,omitempty"` // ToolResultContent
|
||||
NestedContent []*wireContent `json:"content,omitempty"` // ToolResultContent
|
||||
StructuredContent any `json:"structuredContent,omitempty"` // ToolResultContent
|
||||
IsError bool `json:"isError,omitempty"` // ToolResultContent
|
||||
}
|
||||
|
||||
// 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 {
|
||||
|
|
@ -254,6 +382,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)
|
||||
}
|
||||
|
|
|
|||
232
mcp/protocol.go
232
mcp/protocol.go
|
|
@ -188,7 +188,6 @@ type RootCapabilities struct {
|
|||
// this schema, but this is not a closed set: any client can define its own,
|
||||
// additional capabilities.
|
||||
type ClientCapabilities struct {
|
||||
|
||||
// NOTE: any addition to ClientCapabilities must also be reflected in
|
||||
// [ClientCapabilities.clone].
|
||||
|
||||
|
|
@ -216,7 +215,12 @@ type ClientCapabilities struct {
|
|||
func (c *ClientCapabilities) clone() *ClientCapabilities {
|
||||
cp := *c
|
||||
cp.RootsV2 = shallowClone(c.RootsV2)
|
||||
cp.Sampling = shallowClone(c.Sampling)
|
||||
if c.Sampling != nil {
|
||||
x := *c.Sampling
|
||||
x.Tools = shallowClone(c.Sampling.Tools)
|
||||
x.Context = shallowClone(c.Sampling.Context)
|
||||
cp.Sampling = &x
|
||||
}
|
||||
if c.Elicitation != nil {
|
||||
x := *c.Elicitation
|
||||
x.Form = shallowClone(c.Elicitation.Form)
|
||||
|
|
@ -357,6 +361,11 @@ type CreateMessageParams struct {
|
|||
Meta `json:"_meta,omitempty"`
|
||||
// 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.
|
||||
//
|
||||
// The default is "none". Values "thisServer" and
|
||||
// "allServers" are soft-deprecated. Servers SHOULD only use these values if
|
||||
// the client declares ClientCapabilities.sampling.context. These values may
|
||||
// be removed in future spec releases.
|
||||
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.
|
||||
|
|
@ -379,6 +388,106 @@ func (x *CreateMessageParams) isParams() {}
|
|||
func (x *CreateMessageParams) GetProgressToken() any { return getProgressToken(x) }
|
||||
func (x *CreateMessageParams) SetProgressToken(t any) { setProgressToken(x, t) }
|
||||
|
||||
// CreateMessageWithToolsParams is a sampling request that includes tools.
|
||||
// It extends the basic [CreateMessageParams] fields with tools, tool choice,
|
||||
// and messages that support array content (for parallel tool calls).
|
||||
//
|
||||
// Use with [ServerSession.CreateMessageWithTools].
|
||||
type CreateMessageWithToolsParams struct {
|
||||
Meta `json:"_meta,omitempty"`
|
||||
IncludeContext string `json:"includeContext,omitempty"`
|
||||
MaxTokens int64 `json:"maxTokens"`
|
||||
// Messages supports array content for tool_use and tool_result blocks.
|
||||
Messages []*SamplingMessageV2 `json:"messages"`
|
||||
Metadata any `json:"metadata,omitempty"`
|
||||
ModelPreferences *ModelPreferences `json:"modelPreferences,omitempty"`
|
||||
StopSequences []string `json:"stopSequences,omitempty"`
|
||||
SystemPrompt string `json:"systemPrompt,omitempty"`
|
||||
Temperature float64 `json:"temperature,omitempty"`
|
||||
// Tools is the list of tools available for the model to use.
|
||||
Tools []*Tool `json:"tools,omitempty"`
|
||||
// ToolChoice controls how the model should use tools.
|
||||
ToolChoice *ToolChoice `json:"toolChoice,omitempty"`
|
||||
}
|
||||
|
||||
func (x *CreateMessageWithToolsParams) isParams() {}
|
||||
func (x *CreateMessageWithToolsParams) GetProgressToken() any { return getProgressToken(x) }
|
||||
func (x *CreateMessageWithToolsParams) SetProgressToken(t any) { setProgressToken(x, t) }
|
||||
|
||||
// toBase converts to CreateMessageParams by taking the content block from each
|
||||
// message. Tools and ToolChoice are dropped. Returns an error if any message
|
||||
// has multiple content blocks, since SamplingMessage only supports one.
|
||||
func (p *CreateMessageWithToolsParams) toBase() (*CreateMessageParams, error) {
|
||||
var msgs []*SamplingMessage
|
||||
for _, m := range p.Messages {
|
||||
if len(m.Content) > 1 {
|
||||
return nil, fmt.Errorf("message has %d content blocks; use CreateMessageWithToolsHandler to support multiple content", len(m.Content))
|
||||
}
|
||||
var content Content
|
||||
if len(m.Content) > 0 {
|
||||
content = m.Content[0]
|
||||
}
|
||||
msgs = append(msgs, &SamplingMessage{Content: content, Role: m.Role})
|
||||
}
|
||||
return &CreateMessageParams{
|
||||
Meta: p.Meta,
|
||||
IncludeContext: p.IncludeContext,
|
||||
MaxTokens: p.MaxTokens,
|
||||
Messages: msgs,
|
||||
Metadata: p.Metadata,
|
||||
ModelPreferences: p.ModelPreferences,
|
||||
StopSequences: p.StopSequences,
|
||||
SystemPrompt: p.SystemPrompt,
|
||||
Temperature: p.Temperature,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// SamplingMessageV2 describes a message issued to or received from an
|
||||
// LLM API, supporting array content for parallel tool calls. The "V2" refers
|
||||
// to the 2025-11-25 spec, which changed content from a single block to
|
||||
// single-or-array. In v2 of the SDK, this will replace [SamplingMessage].
|
||||
//
|
||||
// When marshaling, a single-element Content slice is marshaled as a single
|
||||
// object for compatibility with pre-2025-11-25 implementations. When
|
||||
// unmarshaling, a single JSON content object is accepted and wrapped in a
|
||||
// one-element slice.
|
||||
type SamplingMessageV2 struct {
|
||||
Content []Content `json:"content"`
|
||||
Role Role `json:"role"`
|
||||
}
|
||||
|
||||
var samplingWithToolsAllow = map[string]bool{
|
||||
"text": true, "image": true, "audio": true,
|
||||
"tool_use": true, "tool_result": true,
|
||||
}
|
||||
|
||||
// MarshalJSON marshals the message. A single-element Content slice is marshaled
|
||||
// as a single object for backward compatibility.
|
||||
func (m *SamplingMessageV2) MarshalJSON() ([]byte, error) {
|
||||
if len(m.Content) == 1 {
|
||||
return json.Marshal(&SamplingMessage{Content: m.Content[0], Role: m.Role})
|
||||
}
|
||||
type msg SamplingMessageV2 // avoid recursion
|
||||
return json.Marshal((*msg)(m))
|
||||
}
|
||||
|
||||
func (m *SamplingMessageV2) UnmarshalJSON(data []byte) error {
|
||||
type msg SamplingMessageV2 // avoid recursion
|
||||
var wire struct {
|
||||
msg
|
||||
Content json.RawMessage `json:"content"`
|
||||
}
|
||||
if err := json.Unmarshal(data, &wire); err != nil {
|
||||
return err
|
||||
}
|
||||
var err error
|
||||
if wire.msg.Content, err = unmarshalContent(wire.Content, samplingWithToolsAllow); err != nil {
|
||||
return err
|
||||
}
|
||||
*m = SamplingMessageV2(wire.msg)
|
||||
return nil
|
||||
}
|
||||
|
||||
// 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
|
||||
|
|
@ -392,6 +501,12 @@ type CreateMessageResult struct {
|
|||
Model string `json:"model"`
|
||||
Role Role `json:"role"`
|
||||
// The reason why sampling stopped, if known.
|
||||
//
|
||||
// Standard values:
|
||||
// - "endTurn": natural end of the assistant's turn
|
||||
// - "stopSequence": a stop sequence was encountered
|
||||
// - "maxTokens": reached the maximum token limit
|
||||
// - "toolUse": the model wants to use one or more tools
|
||||
StopReason string `json:"stopReason,omitempty"`
|
||||
}
|
||||
|
||||
|
|
@ -413,6 +528,84 @@ func (r *CreateMessageResult) UnmarshalJSON(data []byte) error {
|
|||
return nil
|
||||
}
|
||||
|
||||
// CreateMessageWithToolsResult is the client's response to a
|
||||
// sampling/create_message request that included tools. Content is a slice to
|
||||
// support parallel tool calls (multiple tool_use blocks in one response).
|
||||
//
|
||||
// Use [ServerSession.CreateMessageWithTools] to send a sampling request with
|
||||
// tools and receive this result type.
|
||||
//
|
||||
// When unmarshaling, a single JSON content object is accepted and wrapped in a
|
||||
// one-element slice, for compatibility with clients that return a single block.
|
||||
type CreateMessageWithToolsResult struct {
|
||||
Meta `json:"_meta,omitempty"`
|
||||
Content []Content `json:"content"`
|
||||
Model string `json:"model"`
|
||||
Role Role `json:"role"`
|
||||
// The reason why sampling stopped.
|
||||
//
|
||||
// Standard values: "endTurn", "stopSequence", "maxTokens", "toolUse".
|
||||
StopReason string `json:"stopReason,omitempty"`
|
||||
}
|
||||
|
||||
// createMessageWithToolsResultAllow lists content types valid in assistant responses.
|
||||
// tool_result is excluded: it only appears in user messages.
|
||||
var createMessageWithToolsResultAllow = map[string]bool{
|
||||
"text": true, "image": true, "audio": true,
|
||||
"tool_use": true,
|
||||
}
|
||||
|
||||
func (*CreateMessageWithToolsResult) isResult() {}
|
||||
|
||||
// MarshalJSON marshals the result. When Content has a single element, it is
|
||||
// marshaled as a single object for compatibility with pre-2025-11-25
|
||||
// implementations that expect a single content block.
|
||||
func (r *CreateMessageWithToolsResult) MarshalJSON() ([]byte, error) {
|
||||
if len(r.Content) == 1 {
|
||||
return json.Marshal(&CreateMessageResult{
|
||||
Meta: r.Meta,
|
||||
Content: r.Content[0],
|
||||
Model: r.Model,
|
||||
Role: r.Role,
|
||||
StopReason: r.StopReason,
|
||||
})
|
||||
}
|
||||
type result CreateMessageWithToolsResult // avoid recursion
|
||||
return json.Marshal((*result)(r))
|
||||
}
|
||||
|
||||
func (r *CreateMessageWithToolsResult) UnmarshalJSON(data []byte) error {
|
||||
type result CreateMessageWithToolsResult // avoid recursion
|
||||
var wire struct {
|
||||
result
|
||||
Content json.RawMessage `json:"content"`
|
||||
}
|
||||
if err := json.Unmarshal(data, &wire); err != nil {
|
||||
return err
|
||||
}
|
||||
var err error
|
||||
if wire.result.Content, err = unmarshalContent(wire.Content, createMessageWithToolsResultAllow); err != nil {
|
||||
return err
|
||||
}
|
||||
*r = CreateMessageWithToolsResult(wire.result)
|
||||
return nil
|
||||
}
|
||||
|
||||
// toWithTools converts a CreateMessageResult to CreateMessageWithToolsResult.
|
||||
func (r *CreateMessageResult) toWithTools() *CreateMessageWithToolsResult {
|
||||
var content []Content
|
||||
if r.Content != nil {
|
||||
content = []Content{r.Content}
|
||||
}
|
||||
return &CreateMessageWithToolsResult{
|
||||
Meta: r.Meta,
|
||||
Content: content,
|
||||
Model: r.Model,
|
||||
Role: r.Role,
|
||||
StopReason: r.StopReason,
|
||||
}
|
||||
}
|
||||
|
||||
type GetPromptParams struct {
|
||||
// This property is reserved by the protocol to allow clients and servers to
|
||||
// attach additional metadata to their responses.
|
||||
|
|
@ -982,7 +1175,27 @@ func (x *RootsListChangedParams) SetProgressToken(t any) { setProgressToken(x, t
|
|||
// below directly above ClientCapabilities.
|
||||
|
||||
// SamplingCapabilities describes the client's support for sampling.
|
||||
type SamplingCapabilities struct{}
|
||||
type SamplingCapabilities struct {
|
||||
// Context indicates the client supports includeContext values other than "none".
|
||||
Context *SamplingContextCapabilities `json:"context,omitempty"`
|
||||
// Tools indicates the client supports tools and toolChoice in sampling requests.
|
||||
Tools *SamplingToolsCapabilities `json:"tools,omitempty"`
|
||||
}
|
||||
|
||||
// SamplingContextCapabilities indicates the client supports context inclusion.
|
||||
type SamplingContextCapabilities struct{}
|
||||
|
||||
// SamplingToolsCapabilities indicates the client supports tool use in sampling.
|
||||
type SamplingToolsCapabilities struct{}
|
||||
|
||||
// ToolChoice controls how the model uses tools during sampling.
|
||||
type ToolChoice struct {
|
||||
// Mode controls tool invocation behavior:
|
||||
// - "auto": Model decides whether to use tools (default)
|
||||
// - "required": Model must use at least one tool
|
||||
// - "none": Model must not use any tools
|
||||
Mode string `json:"mode,omitempty"`
|
||||
}
|
||||
|
||||
// ElicitationCapabilities describes the capabilities for elicitation.
|
||||
//
|
||||
|
|
@ -993,14 +1206,15 @@ type ElicitationCapabilities struct {
|
|||
}
|
||||
|
||||
// FormElicitationCapabilities describes capabilities for form elicitation.
|
||||
type FormElicitationCapabilities struct {
|
||||
}
|
||||
type FormElicitationCapabilities struct{}
|
||||
|
||||
// URLElicitationCapabilities describes capabilities for url elicitation.
|
||||
type URLElicitationCapabilities struct {
|
||||
}
|
||||
type URLElicitationCapabilities struct{}
|
||||
|
||||
// Describes a message issued to or received from an LLM API.
|
||||
//
|
||||
// For assistant messages, Content may be text, image, audio, or tool_use.
|
||||
// For user messages, Content may be text, image, audio, or tool_result.
|
||||
type SamplingMessage struct {
|
||||
Content Content `json:"content"`
|
||||
Role Role `json:"role"`
|
||||
|
|
@ -1017,8 +1231,9 @@ func (m *SamplingMessage) UnmarshalJSON(data []byte) error {
|
|||
if err := json.Unmarshal(data, &wire); err != nil {
|
||||
return err
|
||||
}
|
||||
// Allow text, image, audio, tool_use, and tool_result in sampling messages
|
||||
var err error
|
||||
if wire.msg.Content, err = contentFromWire(wire.Content, map[string]bool{"text": true, "image": true, "audio": true}); err != nil {
|
||||
if wire.msg.Content, err = contentFromWire(wire.Content, map[string]bool{"text": true, "image": true, "audio": true, "tool_use": true, "tool_result": true}); err != nil {
|
||||
return err
|
||||
}
|
||||
*m = SamplingMessage(wire.msg)
|
||||
|
|
@ -1297,7 +1512,6 @@ type ToolCapabilities struct {
|
|||
|
||||
// ServerCapabilities describes capabilities that a server supports.
|
||||
type ServerCapabilities struct {
|
||||
|
||||
// NOTE: any addition to ServerCapabilities must also be reflected in
|
||||
// [ServerCapabilities.clone].
|
||||
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ package mcp
|
|||
import (
|
||||
"encoding/json"
|
||||
"maps"
|
||||
"reflect"
|
||||
"testing"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
|
|
@ -488,6 +489,587 @@ func TestCompleteResult(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestToolUseContent_MarshalJSON(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
content *ToolUseContent
|
||||
want *ToolUseContent
|
||||
}{
|
||||
{
|
||||
name: "basic tool use",
|
||||
content: &ToolUseContent{
|
||||
ID: "tool_123",
|
||||
Name: "calculator",
|
||||
Input: map[string]any{
|
||||
"operation": "add",
|
||||
"x": 1.0,
|
||||
"y": 2.0,
|
||||
},
|
||||
},
|
||||
want: &ToolUseContent{
|
||||
ID: "tool_123",
|
||||
Name: "calculator",
|
||||
Input: map[string]any{
|
||||
"operation": "add",
|
||||
"x": 1.0,
|
||||
"y": 2.0,
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "nil input marshals as empty object",
|
||||
content: &ToolUseContent{
|
||||
ID: "tool_456",
|
||||
Name: "no_args_tool",
|
||||
Input: nil,
|
||||
},
|
||||
want: &ToolUseContent{
|
||||
ID: "tool_456",
|
||||
Name: "no_args_tool",
|
||||
Input: map[string]any{},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
data, err := tt.content.MarshalJSON()
|
||||
if err != nil {
|
||||
t.Fatalf("MarshalJSON() error = %v", err)
|
||||
}
|
||||
wire := &wireContent{}
|
||||
if err := json.Unmarshal(data, wire); err != nil {
|
||||
t.Fatalf("Unmarshal wire error = %v", err)
|
||||
}
|
||||
got, err := contentFromWire(wire, map[string]bool{"tool_use": true})
|
||||
if err != nil {
|
||||
t.Fatalf("contentFromWire() error = %v", err)
|
||||
}
|
||||
if diff := cmp.Diff(tt.want, got); diff != "" {
|
||||
t.Errorf("mismatch (-want +got):\n%s", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestToolUseContent_UnmarshalJSON(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
json string
|
||||
want *ToolUseContent
|
||||
}{
|
||||
{
|
||||
name: "basic tool use",
|
||||
json: `{"type":"tool_use","id":"tool_123","name":"calculator","input":{"x":1,"y":2}}`,
|
||||
want: &ToolUseContent{
|
||||
ID: "tool_123",
|
||||
Name: "calculator",
|
||||
Input: map[string]any{
|
||||
"x": 1.0,
|
||||
"y": 2.0,
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "with meta",
|
||||
json: `{"type":"tool_use","id":"t1","name":"calc","input":{"x":1},"_meta":{"requestId":"req-123"}}`,
|
||||
want: &ToolUseContent{
|
||||
ID: "t1",
|
||||
Name: "calc",
|
||||
Input: map[string]any{"x": 1.0},
|
||||
Meta: Meta{"requestId": "req-123"},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
wire := &wireContent{}
|
||||
if err := json.Unmarshal([]byte(tt.json), wire); err != nil {
|
||||
t.Fatalf("Unmarshal() error = %v", err)
|
||||
}
|
||||
got, err := contentFromWire(wire, map[string]bool{"tool_use": true})
|
||||
if err != nil {
|
||||
t.Fatalf("contentFromWire() error = %v", err)
|
||||
}
|
||||
if diff := cmp.Diff(tt.want, got); diff != "" {
|
||||
t.Errorf("mismatch (-want +got):\n%s", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestToolResultContent_MarshalJSON(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
content *ToolResultContent
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "basic tool result",
|
||||
content: &ToolResultContent{
|
||||
ToolUseID: "tool_123",
|
||||
Content: []Content{&TextContent{Text: "42"}},
|
||||
},
|
||||
want: `{"type":"tool_result","toolUseId":"tool_123","content":[{"type":"text","text":"42"}]}`,
|
||||
},
|
||||
{
|
||||
name: "tool result with error",
|
||||
content: &ToolResultContent{
|
||||
ToolUseID: "tool_456",
|
||||
Content: []Content{&TextContent{Text: "division by zero"}},
|
||||
IsError: true,
|
||||
},
|
||||
want: `{"type":"tool_result","toolUseId":"tool_456","content":[{"type":"text","text":"division by zero"}],"isError":true}`,
|
||||
},
|
||||
{
|
||||
name: "tool result with structured content",
|
||||
content: &ToolResultContent{
|
||||
ToolUseID: "tool_789",
|
||||
Content: []Content{&TextContent{Text: `{"result": 42}`}},
|
||||
StructuredContent: map[string]any{"result": 42.0},
|
||||
},
|
||||
want: `{"type":"tool_result","toolUseId":"tool_789","content":[{"type":"text","text":"{\"result\": 42}"}],"structuredContent":{"result":42}}`,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
data, err := tt.content.MarshalJSON()
|
||||
if err != nil {
|
||||
t.Fatalf("MarshalJSON() error = %v", err)
|
||||
}
|
||||
if diff := cmp.Diff(tt.want, string(data)); diff != "" {
|
||||
t.Errorf("mismatch (-want +got):\n%s", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestToolResultContent_UnmarshalJSON(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
json string
|
||||
wantID string
|
||||
wantErr bool
|
||||
checkFn func(t *testing.T, got *ToolResultContent)
|
||||
}{
|
||||
{
|
||||
name: "basic tool result",
|
||||
json: `{"type":"tool_result","toolUseId":"tool_123","content":[{"type":"text","text":"42"}],"isError":false}`,
|
||||
wantID: "tool_123",
|
||||
checkFn: func(t *testing.T, got *ToolResultContent) {
|
||||
if len(got.Content) != 1 {
|
||||
t.Fatalf("len(Content) = %d, want 1", len(got.Content))
|
||||
}
|
||||
tc, ok := got.Content[0].(*TextContent)
|
||||
if !ok {
|
||||
t.Fatalf("Content[0] type = %T, want *TextContent", got.Content[0])
|
||||
}
|
||||
if tc.Text != "42" {
|
||||
t.Errorf("Text = %v, want 42", tc.Text)
|
||||
}
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "image nested content",
|
||||
json: `{"type":"tool_result","toolUseId":"t1","content":[{"type":"image","mimeType":"image/png","data":"YWJj"}]}`,
|
||||
wantID: "t1",
|
||||
checkFn: func(t *testing.T, got *ToolResultContent) {
|
||||
if len(got.Content) != 1 {
|
||||
t.Fatalf("len(Content) = %d, want 1", len(got.Content))
|
||||
}
|
||||
img, ok := got.Content[0].(*ImageContent)
|
||||
if !ok {
|
||||
t.Fatalf("Content[0] type = %T, want *ImageContent", got.Content[0])
|
||||
}
|
||||
if img.MIMEType != "image/png" {
|
||||
t.Errorf("MIMEType = %v, want image/png", img.MIMEType)
|
||||
}
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
wire := &wireContent{}
|
||||
if err := json.Unmarshal([]byte(tt.json), wire); err != nil {
|
||||
t.Fatalf("Unmarshal() error = %v", err)
|
||||
}
|
||||
got, err := contentFromWire(wire, map[string]bool{"tool_result": true})
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Fatalf("contentFromWire() error = %v, wantErr %v", err, tt.wantErr)
|
||||
}
|
||||
if tt.wantErr {
|
||||
return
|
||||
}
|
||||
result, ok := got.(*ToolResultContent)
|
||||
if !ok {
|
||||
t.Fatalf("type = %T, want *ToolResultContent", got)
|
||||
}
|
||||
if result.ToolUseID != tt.wantID {
|
||||
t.Errorf("ToolUseID = %v, want %v", result.ToolUseID, tt.wantID)
|
||||
}
|
||||
if tt.checkFn != nil {
|
||||
tt.checkFn(t, result)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSamplingMessage_UnmarshalJSON(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
json string
|
||||
wantRole Role
|
||||
wantType string // expected Content type name
|
||||
}{
|
||||
{
|
||||
name: "tool_use content",
|
||||
json: `{"content":{"type":"tool_use","id":"tool_1","name":"calc","input":{}},"role":"assistant"}`,
|
||||
wantRole: "assistant",
|
||||
wantType: "*mcp.ToolUseContent",
|
||||
},
|
||||
{
|
||||
name: "tool_result content",
|
||||
json: `{"content":{"type":"tool_result","toolUseId":"tool_1","content":[{"type":"text","text":"42"}]},"role":"user"}`,
|
||||
wantRole: "user",
|
||||
wantType: "*mcp.ToolResultContent",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
var msg SamplingMessage
|
||||
if err := json.Unmarshal([]byte(tt.json), &msg); err != nil {
|
||||
t.Fatalf("Unmarshal() error = %v", err)
|
||||
}
|
||||
if msg.Role != tt.wantRole {
|
||||
t.Errorf("Role = %v, want %v", msg.Role, tt.wantRole)
|
||||
}
|
||||
gotType := reflect.TypeOf(msg.Content).String()
|
||||
if gotType != tt.wantType {
|
||||
t.Errorf("Content type = %v, want %v", gotType, tt.wantType)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSamplingCapabilities_MarshalJSON(t *testing.T) {
|
||||
caps := &SamplingCapabilities{
|
||||
Tools: &SamplingToolsCapabilities{},
|
||||
Context: &SamplingContextCapabilities{},
|
||||
}
|
||||
data, err := json.Marshal(caps)
|
||||
if err != nil {
|
||||
t.Fatalf("Marshal() error = %v", err)
|
||||
}
|
||||
var got SamplingCapabilities
|
||||
if err := json.Unmarshal(data, &got); err != nil {
|
||||
t.Fatalf("Unmarshal() error = %v", err)
|
||||
}
|
||||
if got.Tools == nil {
|
||||
t.Error("Tools capability should not be nil")
|
||||
}
|
||||
if got.Context == nil {
|
||||
t.Error("Context capability should not be nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSamplingCapabilities_UnmarshalJSON(t *testing.T) {
|
||||
// Empty struct should marshal/unmarshal correctly (backward compatibility).
|
||||
caps := &SamplingCapabilities{}
|
||||
data, err := json.Marshal(caps)
|
||||
if err != nil {
|
||||
t.Fatalf("Marshal() error = %v", err)
|
||||
}
|
||||
var got SamplingCapabilities
|
||||
if err := json.Unmarshal(data, &got); err != nil {
|
||||
t.Fatalf("Unmarshal() error = %v", err)
|
||||
}
|
||||
if got.Tools != nil {
|
||||
t.Error("Tools capability should be nil for empty capabilities")
|
||||
}
|
||||
if got.Context != nil {
|
||||
t.Error("Context capability should be nil for empty capabilities")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateMessageWithToolsParams_MarshalUnmarshalJSON(t *testing.T) {
|
||||
params := &CreateMessageWithToolsParams{
|
||||
MaxTokens: 1000,
|
||||
Messages: []*SamplingMessageV2{
|
||||
{
|
||||
Role: "user",
|
||||
Content: []Content{&TextContent{Text: "Calculate 1+1"}},
|
||||
},
|
||||
},
|
||||
Tools: []*Tool{
|
||||
{
|
||||
Name: "calculator",
|
||||
Description: "A calculator tool",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"x": map[string]any{"type": "number"},
|
||||
"y": map[string]any{"type": "number"},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
ToolChoice: &ToolChoice{Mode: "auto"},
|
||||
}
|
||||
|
||||
data, err := json.Marshal(params)
|
||||
if err != nil {
|
||||
t.Fatalf("Marshal() error = %v", err)
|
||||
}
|
||||
|
||||
var got CreateMessageWithToolsParams
|
||||
if err := json.Unmarshal(data, &got); err != nil {
|
||||
t.Fatalf("Unmarshal() error = %v", err)
|
||||
}
|
||||
|
||||
if len(got.Tools) != 1 {
|
||||
t.Fatalf("len(Tools) = %v, want 1", len(got.Tools))
|
||||
}
|
||||
if got.Tools[0].Name != "calculator" {
|
||||
t.Errorf("Tools[0].Name = %v, want calculator", got.Tools[0].Name)
|
||||
}
|
||||
if got.ToolChoice == nil || got.ToolChoice.Mode != "auto" {
|
||||
t.Errorf("ToolChoice = %v, want {Mode: auto}", got.ToolChoice)
|
||||
}
|
||||
}
|
||||
|
||||
func TestToolChoice_MarshalUnmarshalJSON(t *testing.T) {
|
||||
choices := []*ToolChoice{
|
||||
{Mode: "auto"},
|
||||
{Mode: "required"},
|
||||
{Mode: "none"},
|
||||
}
|
||||
|
||||
for _, tc := range choices {
|
||||
data, err := json.Marshal(tc)
|
||||
if err != nil {
|
||||
t.Fatalf("Marshal(%v) error = %v", tc.Mode, err)
|
||||
}
|
||||
var got ToolChoice
|
||||
if err := json.Unmarshal(data, &got); err != nil {
|
||||
t.Fatalf("Unmarshal() error = %v", err)
|
||||
}
|
||||
if diff := cmp.Diff(*tc, got); diff != "" {
|
||||
t.Errorf("mismatch (-want +got):\n%s", diff)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateMessageWithToolsResult_MarshalJSON(t *testing.T) {
|
||||
// Single-element Content marshals as object (not array) for backward compat.
|
||||
result := &CreateMessageWithToolsResult{
|
||||
Model: "test",
|
||||
Role: "assistant",
|
||||
Content: []Content{&TextContent{Text: "hello"}},
|
||||
}
|
||||
|
||||
data, err := json.Marshal(result)
|
||||
if err != nil {
|
||||
t.Fatalf("Marshal() error = %v", err)
|
||||
}
|
||||
|
||||
var raw map[string]json.RawMessage
|
||||
if err := json.Unmarshal(data, &raw); err != nil {
|
||||
t.Fatalf("Unmarshal raw error = %v", err)
|
||||
}
|
||||
|
||||
content := raw["content"]
|
||||
for _, b := range content {
|
||||
if b == ' ' || b == '\t' || b == '\n' || b == '\r' {
|
||||
continue
|
||||
}
|
||||
if b == '[' {
|
||||
t.Errorf("single-element Content marshaled as array, want object")
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateMessageWithToolsResult_UnmarshalJSON(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
json string
|
||||
wantLen int
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "single tool_use content",
|
||||
json: `{"content":{"type":"tool_use","id":"tool_1","name":"calculator","input":{"x":1}},"model":"test-model","role":"assistant","stopReason":"toolUse"}`,
|
||||
wantLen: 1,
|
||||
},
|
||||
{
|
||||
name: "array of tool_use content",
|
||||
json: `{"content":[{"type":"tool_use","id":"t1","name":"calc","input":{"x":1}},{"type":"tool_use","id":"t2","name":"search","input":{"q":"hi"}}],"model":"test","role":"assistant","stopReason":"toolUse"}`,
|
||||
wantLen: 2,
|
||||
},
|
||||
{
|
||||
name: "empty array",
|
||||
json: `{"content":[],"model":"m","role":"assistant"}`,
|
||||
wantLen: 0,
|
||||
},
|
||||
{
|
||||
name: "null content",
|
||||
json: `{"content":null,"model":"m","role":"assistant"}`,
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "rejects tool_result",
|
||||
json: `{"content":{"type":"tool_result","toolUseId":"t1","content":[]},"model":"m","role":"assistant"}`,
|
||||
wantErr: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
var got CreateMessageWithToolsResult
|
||||
err := json.Unmarshal([]byte(tt.json), &got)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Fatalf("Unmarshal() error = %v, wantErr %v", err, tt.wantErr)
|
||||
}
|
||||
if tt.wantErr {
|
||||
return
|
||||
}
|
||||
if len(got.Content) != tt.wantLen {
|
||||
t.Errorf("len(Content) = %d, want %d", len(got.Content), tt.wantLen)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSamplingMessageV2_MarshalJSON(t *testing.T) {
|
||||
msg := &SamplingMessageV2{
|
||||
Role: "user",
|
||||
Content: []Content{},
|
||||
}
|
||||
data, err := json.Marshal(msg)
|
||||
if err != nil {
|
||||
t.Fatalf("Marshal() error = %v", err)
|
||||
}
|
||||
var got SamplingMessageV2
|
||||
if err := json.Unmarshal(data, &got); err != nil {
|
||||
t.Fatalf("Unmarshal() error = %v", err)
|
||||
}
|
||||
if len(got.Content) != 0 {
|
||||
t.Errorf("len(Content) = %d, want 0", len(got.Content))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSamplingMessageV2_UnmarshalJSON(t *testing.T) {
|
||||
// Text + tool_use in the same message (valid per spec for assistant).
|
||||
msg := &SamplingMessageV2{
|
||||
Role: "assistant",
|
||||
Content: []Content{
|
||||
&TextContent{Text: "Let me check the weather."},
|
||||
&ToolUseContent{ID: "c1", Name: "weather", Input: map[string]any{"city": "SF"}},
|
||||
},
|
||||
}
|
||||
data, err := json.Marshal(msg)
|
||||
if err != nil {
|
||||
t.Fatalf("Marshal() error = %v", err)
|
||||
}
|
||||
var got SamplingMessageV2
|
||||
if err := json.Unmarshal(data, &got); err != nil {
|
||||
t.Fatalf("Unmarshal() error = %v", err)
|
||||
}
|
||||
if len(got.Content) != 2 {
|
||||
t.Fatalf("len(Content) = %d, want 2", len(got.Content))
|
||||
}
|
||||
if _, ok := got.Content[0].(*TextContent); !ok {
|
||||
t.Errorf("Content[0] type = %T, want *TextContent", got.Content[0])
|
||||
}
|
||||
if _, ok := got.Content[1].(*ToolUseContent); !ok {
|
||||
t.Errorf("Content[1] type = %T, want *ToolUseContent", got.Content[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestToBase_Conversion(t *testing.T) {
|
||||
// Single-content messages convert successfully.
|
||||
params := &CreateMessageWithToolsParams{
|
||||
MaxTokens: 1000,
|
||||
Messages: []*SamplingMessageV2{
|
||||
{Role: "user", Content: []Content{&TextContent{Text: "hello"}}},
|
||||
},
|
||||
Tools: []*Tool{{Name: "calc"}},
|
||||
ToolChoice: &ToolChoice{Mode: "auto"},
|
||||
}
|
||||
base, err := params.toBase()
|
||||
if err != nil {
|
||||
t.Fatalf("toBase() error = %v", err)
|
||||
}
|
||||
if base.MaxTokens != 1000 {
|
||||
t.Errorf("MaxTokens = %d, want 1000", base.MaxTokens)
|
||||
}
|
||||
if tc, ok := base.Messages[0].Content.(*TextContent); !ok || tc.Text != "hello" {
|
||||
t.Errorf("Messages[0].Content = %v, want TextContent{hello}", base.Messages[0].Content)
|
||||
}
|
||||
|
||||
// Multi-content messages return an error.
|
||||
params2 := &CreateMessageWithToolsParams{
|
||||
MaxTokens: 1000,
|
||||
Messages: []*SamplingMessageV2{
|
||||
{Role: "assistant", Content: []Content{
|
||||
&ToolUseContent{ID: "c1", Name: "calc", Input: map[string]any{}},
|
||||
&ToolUseContent{ID: "c2", Name: "search", Input: map[string]any{}},
|
||||
}},
|
||||
},
|
||||
}
|
||||
if _, err := params2.toBase(); err == nil {
|
||||
t.Error("toBase() should return error for multi-content message")
|
||||
}
|
||||
}
|
||||
|
||||
func TestToWithTools_Conversion(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
result *CreateMessageResult
|
||||
wantLen int
|
||||
}{
|
||||
{
|
||||
name: "with content",
|
||||
result: &CreateMessageResult{
|
||||
Model: "test",
|
||||
Role: "assistant",
|
||||
Content: &TextContent{Text: "hello"},
|
||||
StopReason: "endTurn",
|
||||
},
|
||||
wantLen: 1,
|
||||
},
|
||||
{
|
||||
name: "nil content",
|
||||
result: &CreateMessageResult{
|
||||
Model: "test",
|
||||
Role: "assistant",
|
||||
},
|
||||
wantLen: 0,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
wt := tt.result.toWithTools()
|
||||
if wt.Model != tt.result.Model {
|
||||
t.Errorf("Model = %v, want %v", wt.Model, tt.result.Model)
|
||||
}
|
||||
if len(wt.Content) != tt.wantLen {
|
||||
t.Fatalf("len(Content) = %d, want %d", len(wt.Content), tt.wantLen)
|
||||
}
|
||||
if tt.wantLen > 0 {
|
||||
if tc, ok := wt.Content[0].(*TextContent); !ok || tc.Text != "hello" {
|
||||
t.Errorf("Content[0] = %v, want TextContent{hello}", wt.Content[0])
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestContentUnmarshal(t *testing.T) {
|
||||
// Verify that types with a Content field round-trip properly.
|
||||
roundtrip := func(in, out any) {
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ type (
|
|||
|
||||
type (
|
||||
CreateMessageRequest = ClientRequest[*CreateMessageParams]
|
||||
CreateMessageWithToolsRequest = ClientRequest[*CreateMessageWithToolsParams]
|
||||
ElicitRequest = ClientRequest[*ElicitParams]
|
||||
initializedClientRequest = ClientRequest[*InitializedParams]
|
||||
InitializeRequest = ClientRequest[*InitializeParams]
|
||||
|
|
|
|||
546
mcp/sampling_test.go
Normal file
546
mcp/sampling_test.go
Normal file
|
|
@ -0,0 +1,546 @@
|
|||
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
|
||||
// Use of this source code is governed by an MIT-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package mcp
|
||||
|
||||
// TODO: move other sampling-related tests to this file.
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestSamplingWithTools_Integration(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
ct, st := NewInMemoryTransports()
|
||||
|
||||
// Track what the client received
|
||||
var receivedParams *CreateMessageWithToolsParams
|
||||
|
||||
// Client with tools capability, using CreateMessageWithToolsHandler
|
||||
client := NewClient(testImpl, &ClientOptions{
|
||||
CreateMessageWithToolsHandler: func(_ context.Context, req *CreateMessageWithToolsRequest) (*CreateMessageWithToolsResult, error) {
|
||||
receivedParams = req.Params
|
||||
// Return a tool use response
|
||||
return &CreateMessageWithToolsResult{
|
||||
Model: "test-model",
|
||||
Role: "assistant",
|
||||
Content: []Content{&ToolUseContent{
|
||||
ID: "tool_call_1",
|
||||
Name: "calculator",
|
||||
Input: map[string]any{"x": 1.0, "y": 2.0},
|
||||
}},
|
||||
StopReason: "toolUse",
|
||||
}, nil
|
||||
},
|
||||
Capabilities: &ClientCapabilities{
|
||||
Sampling: &SamplingCapabilities{Tools: &SamplingToolsCapabilities{}},
|
||||
},
|
||||
})
|
||||
|
||||
server := NewServer(testImpl, nil)
|
||||
ss, err := server.Connect(ctx, st, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer ss.Close()
|
||||
|
||||
cs, err := client.Connect(ctx, ct, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer cs.Close()
|
||||
|
||||
// Server sends CreateMessageWithTools
|
||||
result, err := ss.CreateMessageWithTools(ctx, &CreateMessageWithToolsParams{
|
||||
MaxTokens: 1000,
|
||||
Messages: []*SamplingMessageV2{
|
||||
{Role: "user", Content: []Content{&TextContent{Text: "Calculate 1+2"}}},
|
||||
},
|
||||
Tools: []*Tool{
|
||||
{
|
||||
Name: "calculator",
|
||||
Description: "A calculator",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"x": map[string]any{"type": "number"},
|
||||
"y": map[string]any{"type": "number"},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
ToolChoice: &ToolChoice{Mode: "auto"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("CreateMessageWithTools() error = %v", err)
|
||||
}
|
||||
|
||||
// Verify client received the tools
|
||||
if receivedParams == nil {
|
||||
t.Fatal("client did not receive params")
|
||||
}
|
||||
if len(receivedParams.Tools) != 1 {
|
||||
t.Errorf("client received %d tools, want 1", len(receivedParams.Tools))
|
||||
}
|
||||
if receivedParams.Tools[0].Name != "calculator" {
|
||||
t.Errorf("tool name = %v, want calculator", receivedParams.Tools[0].Name)
|
||||
}
|
||||
if receivedParams.ToolChoice == nil || receivedParams.ToolChoice.Mode != "auto" {
|
||||
t.Errorf("tool choice mode = %v, want auto", receivedParams.ToolChoice)
|
||||
}
|
||||
|
||||
// Verify server received the tool use response
|
||||
if result.StopReason != "toolUse" {
|
||||
t.Errorf("StopReason = %v, want toolUse", result.StopReason)
|
||||
}
|
||||
if len(result.Content) != 1 {
|
||||
t.Fatalf("len(Content) = %d, want 1", len(result.Content))
|
||||
}
|
||||
toolUse, ok := result.Content[0].(*ToolUseContent)
|
||||
if !ok {
|
||||
t.Fatalf("Content[0] type = %T, want *ToolUseContent", result.Content[0])
|
||||
}
|
||||
if toolUse.ID != "tool_call_1" {
|
||||
t.Errorf("ToolUse.ID = %v, want tool_call_1", toolUse.ID)
|
||||
}
|
||||
if toolUse.Name != "calculator" {
|
||||
t.Errorf("ToolUse.Name = %v, want calculator", toolUse.Name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSamplingWithToolResult_Integration(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
ct, st := NewInMemoryTransports()
|
||||
|
||||
// Track messages received by client
|
||||
var receivedMessages []*SamplingMessage
|
||||
|
||||
client := NewClient(testImpl, &ClientOptions{
|
||||
CreateMessageHandler: func(_ context.Context, req *CreateMessageRequest) (*CreateMessageResult, error) {
|
||||
receivedMessages = req.Params.Messages
|
||||
return &CreateMessageResult{
|
||||
Model: "test-model",
|
||||
Role: "assistant",
|
||||
Content: &TextContent{Text: "The result is 3"},
|
||||
}, nil
|
||||
},
|
||||
Capabilities: &ClientCapabilities{
|
||||
Sampling: &SamplingCapabilities{Tools: &SamplingToolsCapabilities{}},
|
||||
},
|
||||
})
|
||||
|
||||
server := NewServer(testImpl, nil)
|
||||
ss, err := server.Connect(ctx, st, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer ss.Close()
|
||||
|
||||
cs, err := client.Connect(ctx, ct, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer cs.Close()
|
||||
|
||||
// Server sends CreateMessage with tool result in messages
|
||||
_, err = ss.CreateMessage(ctx, &CreateMessageParams{
|
||||
MaxTokens: 1000,
|
||||
Messages: []*SamplingMessage{
|
||||
{Role: "user", Content: &TextContent{Text: "Calculate 1+2"}},
|
||||
{Role: "assistant", Content: &ToolUseContent{
|
||||
ID: "tool_1",
|
||||
Name: "calculator",
|
||||
Input: map[string]any{"x": 1.0, "y": 2.0},
|
||||
}},
|
||||
{Role: "user", Content: &ToolResultContent{
|
||||
ToolUseID: "tool_1",
|
||||
Content: []Content{&TextContent{Text: "3"}},
|
||||
}},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("CreateMessage() error = %v", err)
|
||||
}
|
||||
|
||||
// Verify client received all messages including tool content
|
||||
if len(receivedMessages) != 3 {
|
||||
t.Fatalf("received %d messages, want 3", len(receivedMessages))
|
||||
}
|
||||
|
||||
// Check first message is text
|
||||
if _, ok := receivedMessages[0].Content.(*TextContent); !ok {
|
||||
t.Errorf("message[0] content type = %T, want *TextContent", receivedMessages[0].Content)
|
||||
}
|
||||
|
||||
// Check second message is tool use
|
||||
toolUse, ok := receivedMessages[1].Content.(*ToolUseContent)
|
||||
if !ok {
|
||||
t.Fatalf("message[1] content type = %T, want *ToolUseContent", receivedMessages[1].Content)
|
||||
}
|
||||
if toolUse.ID != "tool_1" {
|
||||
t.Errorf("toolUse.ID = %v, want tool_1", toolUse.ID)
|
||||
}
|
||||
|
||||
// Check third message is tool result
|
||||
toolResult, ok := receivedMessages[2].Content.(*ToolResultContent)
|
||||
if !ok {
|
||||
t.Fatalf("message[2] content type = %T, want *ToolResultContent", receivedMessages[2].Content)
|
||||
}
|
||||
if toolResult.ToolUseID != "tool_1" {
|
||||
t.Errorf("toolResult.ToolUseID = %v, want tool_1", toolResult.ToolUseID)
|
||||
}
|
||||
if len(toolResult.Content) != 1 {
|
||||
t.Fatalf("toolResult.Content len = %d, want 1", len(toolResult.Content))
|
||||
}
|
||||
if tc, ok := toolResult.Content[0].(*TextContent); !ok || tc.Text != "3" {
|
||||
t.Errorf("toolResult.Content[0] = %v, want TextContent with '3'", toolResult.Content[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestSamplingToolsCapabilities(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
t.Run("client with explicit tools capability", func(t *testing.T) {
|
||||
ct, st := NewInMemoryTransports()
|
||||
|
||||
client := NewClient(testImpl, &ClientOptions{
|
||||
CreateMessageHandler: func(_ context.Context, _ *CreateMessageRequest) (*CreateMessageResult, error) {
|
||||
return &CreateMessageResult{Model: "m", Content: &TextContent{}}, nil
|
||||
},
|
||||
Capabilities: &ClientCapabilities{
|
||||
Sampling: &SamplingCapabilities{
|
||||
Tools: &SamplingToolsCapabilities{},
|
||||
Context: &SamplingContextCapabilities{},
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
server := NewServer(testImpl, nil)
|
||||
ss, err := server.Connect(ctx, st, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer ss.Close()
|
||||
|
||||
cs, err := client.Connect(ctx, ct, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer cs.Close()
|
||||
|
||||
// Check server sees client capabilities
|
||||
caps := ss.InitializeParams().Capabilities
|
||||
if caps.Sampling == nil {
|
||||
t.Fatal("client should advertise sampling capability")
|
||||
}
|
||||
if caps.Sampling.Tools == nil {
|
||||
t.Error("client should advertise sampling.tools capability")
|
||||
}
|
||||
if caps.Sampling.Context == nil {
|
||||
t.Error("client should advertise sampling.context capability")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("client without tools capability", func(t *testing.T) {
|
||||
ct, st := NewInMemoryTransports()
|
||||
|
||||
client := NewClient(testImpl, &ClientOptions{
|
||||
CreateMessageHandler: func(_ context.Context, _ *CreateMessageRequest) (*CreateMessageResult, error) {
|
||||
return &CreateMessageResult{Model: "m", Content: &TextContent{}}, nil
|
||||
},
|
||||
// No Capabilities.Sampling.Tools set
|
||||
})
|
||||
|
||||
server := NewServer(testImpl, nil)
|
||||
ss, err := server.Connect(ctx, st, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer ss.Close()
|
||||
|
||||
cs, err := client.Connect(ctx, ct, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer cs.Close()
|
||||
|
||||
// Check server sees client capabilities
|
||||
caps := ss.InitializeParams().Capabilities
|
||||
if caps.Sampling == nil {
|
||||
t.Fatal("client should advertise sampling capability")
|
||||
}
|
||||
if caps.Sampling.Tools != nil {
|
||||
t.Error("client should NOT advertise sampling.tools capability")
|
||||
}
|
||||
if caps.Sampling.Context != nil {
|
||||
t.Error("client should NOT advertise sampling.context capability")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("client with inferred tools capability", func(t *testing.T) {
|
||||
ct, st := NewInMemoryTransports()
|
||||
|
||||
client := NewClient(testImpl, &ClientOptions{
|
||||
CreateMessageWithToolsHandler: func(_ context.Context, _ *CreateMessageWithToolsRequest) (*CreateMessageWithToolsResult, error) {
|
||||
return &CreateMessageWithToolsResult{Model: "m", Content: []Content{&TextContent{}}}, nil
|
||||
},
|
||||
// No explicit Capabilities set — tools should be inferred.
|
||||
})
|
||||
|
||||
server := NewServer(testImpl, nil)
|
||||
ss, err := server.Connect(ctx, st, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer ss.Close()
|
||||
|
||||
cs, err := client.Connect(ctx, ct, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer cs.Close()
|
||||
|
||||
caps := ss.InitializeParams().Capabilities
|
||||
if caps.Sampling == nil {
|
||||
t.Fatal("client should advertise sampling capability")
|
||||
}
|
||||
if caps.Sampling.Tools == nil {
|
||||
t.Error("client should infer sampling.tools capability from CreateMessageWithToolsHandler")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestSamplingToolResultWithError_Integration(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
ct, st := NewInMemoryTransports()
|
||||
|
||||
var receivedMessages []*SamplingMessage
|
||||
|
||||
client := NewClient(testImpl, &ClientOptions{
|
||||
CreateMessageHandler: func(_ context.Context, req *CreateMessageRequest) (*CreateMessageResult, error) {
|
||||
receivedMessages = req.Params.Messages
|
||||
return &CreateMessageResult{
|
||||
Model: "test-model",
|
||||
Role: "assistant",
|
||||
Content: &TextContent{Text: "I see the tool failed"},
|
||||
}, nil
|
||||
},
|
||||
Capabilities: &ClientCapabilities{
|
||||
Sampling: &SamplingCapabilities{Tools: &SamplingToolsCapabilities{}},
|
||||
},
|
||||
})
|
||||
|
||||
server := NewServer(testImpl, nil)
|
||||
ss, err := server.Connect(ctx, st, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer ss.Close()
|
||||
|
||||
cs, err := client.Connect(ctx, ct, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer cs.Close()
|
||||
|
||||
// Server sends CreateMessage with error tool result, preceded by
|
||||
// the original request and tool use for a more realistic scenario.
|
||||
_, err = ss.CreateMessage(ctx, &CreateMessageParams{
|
||||
MaxTokens: 1000,
|
||||
Messages: []*SamplingMessage{
|
||||
{Role: "user", Content: &TextContent{Text: "Divide 1 by 0"}},
|
||||
{Role: "assistant", Content: &ToolUseContent{
|
||||
ID: "tool_1",
|
||||
Name: "calculator",
|
||||
Input: map[string]any{"op": "div", "x": 1.0, "y": 0.0},
|
||||
}},
|
||||
{Role: "user", Content: &ToolResultContent{
|
||||
ToolUseID: "tool_1",
|
||||
Content: []Content{&TextContent{Text: "division by zero"}},
|
||||
IsError: true,
|
||||
}},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("CreateMessage() error = %v", err)
|
||||
}
|
||||
|
||||
if len(receivedMessages) != 3 {
|
||||
t.Fatalf("received %d messages, want 3", len(receivedMessages))
|
||||
}
|
||||
|
||||
toolResult, ok := receivedMessages[2].Content.(*ToolResultContent)
|
||||
if !ok {
|
||||
t.Fatalf("content type = %T, want *ToolResultContent", receivedMessages[2].Content)
|
||||
}
|
||||
if !toolResult.IsError {
|
||||
t.Error("IsError should be true")
|
||||
}
|
||||
if toolResult.ToolUseID != "tool_1" {
|
||||
t.Errorf("ToolUseID = %v, want tool_1", toolResult.ToolUseID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParallelToolCalls_Integration(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
ct, st := NewInMemoryTransports()
|
||||
|
||||
// Client returns parallel tool use results
|
||||
client := NewClient(testImpl, &ClientOptions{
|
||||
CreateMessageWithToolsHandler: func(_ context.Context, req *CreateMessageWithToolsRequest) (*CreateMessageWithToolsResult, error) {
|
||||
return &CreateMessageWithToolsResult{
|
||||
Model: "test-model",
|
||||
Role: "assistant",
|
||||
Content: []Content{
|
||||
&ToolUseContent{ID: "call_1", Name: "weather", Input: map[string]any{"city": "SF"}},
|
||||
&ToolUseContent{ID: "call_2", Name: "weather", Input: map[string]any{"city": "NY"}},
|
||||
},
|
||||
StopReason: "toolUse",
|
||||
}, nil
|
||||
},
|
||||
Capabilities: &ClientCapabilities{
|
||||
Sampling: &SamplingCapabilities{Tools: &SamplingToolsCapabilities{}},
|
||||
},
|
||||
})
|
||||
|
||||
server := NewServer(testImpl, nil)
|
||||
ss, err := server.Connect(ctx, st, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer ss.Close()
|
||||
|
||||
cs, err := client.Connect(ctx, ct, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer cs.Close()
|
||||
|
||||
result, err := ss.CreateMessageWithTools(ctx, &CreateMessageWithToolsParams{
|
||||
MaxTokens: 1000,
|
||||
Messages: []*SamplingMessageV2{
|
||||
{Role: "user", Content: []Content{&TextContent{Text: "Weather in SF and NY"}}},
|
||||
},
|
||||
Tools: []*Tool{
|
||||
{Name: "weather", InputSchema: map[string]any{"type": "object"}},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("CreateMessageWithTools() error = %v", err)
|
||||
}
|
||||
|
||||
if len(result.Content) != 2 {
|
||||
t.Fatalf("len(Content) = %d, want 2", len(result.Content))
|
||||
}
|
||||
for i, c := range result.Content {
|
||||
tu, ok := c.(*ToolUseContent)
|
||||
if !ok {
|
||||
t.Fatalf("Content[%d] type = %T, want *ToolUseContent", i, c)
|
||||
}
|
||||
if tu.Name != "weather" {
|
||||
t.Errorf("Content[%d].Name = %v, want weather", i, tu.Name)
|
||||
}
|
||||
}
|
||||
if result.Content[0].(*ToolUseContent).ID != "call_1" {
|
||||
t.Errorf("Content[0].ID = %v, want call_1", result.Content[0].(*ToolUseContent).ID)
|
||||
}
|
||||
if result.Content[1].(*ToolUseContent).ID != "call_2" {
|
||||
t.Errorf("Content[1].ID = %v, want call_2", result.Content[1].(*ToolUseContent).ID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewClient_BothHandlersPanics(t *testing.T) {
|
||||
defer func() {
|
||||
r := recover()
|
||||
if r == nil {
|
||||
t.Fatal("expected panic when both handlers set")
|
||||
}
|
||||
msg, ok := r.(string)
|
||||
if !ok || !strings.Contains(msg, "CreateMessageHandler") {
|
||||
t.Errorf("unexpected panic: %v", r)
|
||||
}
|
||||
}()
|
||||
NewClient(testImpl, &ClientOptions{
|
||||
CreateMessageHandler: func(context.Context, *CreateMessageRequest) (*CreateMessageResult, error) {
|
||||
return nil, nil
|
||||
},
|
||||
CreateMessageWithToolsHandler: func(context.Context, *CreateMessageWithToolsRequest) (*CreateMessageWithToolsResult, error) {
|
||||
return nil, nil
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func TestCreateMessage_MultipleContentError(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
ct, st := NewInMemoryTransports()
|
||||
|
||||
// Client returns multiple content blocks via CreateMessageWithToolsHandler
|
||||
client := NewClient(testImpl, &ClientOptions{
|
||||
CreateMessageWithToolsHandler: func(_ context.Context, _ *CreateMessageWithToolsRequest) (*CreateMessageWithToolsResult, error) {
|
||||
return &CreateMessageWithToolsResult{
|
||||
Model: "test",
|
||||
Role: "assistant",
|
||||
Content: []Content{
|
||||
&TextContent{Text: "a"},
|
||||
&TextContent{Text: "b"},
|
||||
},
|
||||
}, nil
|
||||
},
|
||||
})
|
||||
|
||||
server := NewServer(testImpl, nil)
|
||||
ss, err := server.Connect(ctx, st, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer ss.Close()
|
||||
|
||||
cs, err := client.Connect(ctx, ct, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer cs.Close()
|
||||
|
||||
// Server calls CreateMessage (singular), should get error
|
||||
_, err = ss.CreateMessage(ctx, &CreateMessageParams{
|
||||
MaxTokens: 100,
|
||||
Messages: []*SamplingMessage{{Role: "user", Content: &TextContent{Text: "hi"}}},
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("expected error for multiple content blocks")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "CreateMessageWithTools") {
|
||||
t.Errorf("error should mention CreateMessageWithTools, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientCapabilities_CloneSampling(t *testing.T) {
|
||||
caps := &ClientCapabilities{
|
||||
Sampling: &SamplingCapabilities{
|
||||
Tools: &SamplingToolsCapabilities{},
|
||||
Context: &SamplingContextCapabilities{},
|
||||
},
|
||||
}
|
||||
cloned := caps.clone()
|
||||
|
||||
// Verify deep copy — Sampling pointer should differ.
|
||||
// (Tools and Context are empty structs, so Go may reuse the same address;
|
||||
// we just check they're non-nil and that mutating Sampling doesn't alias.)
|
||||
if cloned.Sampling == caps.Sampling {
|
||||
t.Error("Sampling pointer should differ after clone")
|
||||
}
|
||||
if cloned.Sampling.Tools == nil {
|
||||
t.Error("cloned Sampling.Tools should not be nil")
|
||||
}
|
||||
if cloned.Sampling.Context == nil {
|
||||
t.Error("cloned Sampling.Context should not be nil")
|
||||
}
|
||||
// Verify mutation doesn't affect original.
|
||||
cloned.Sampling.Tools = nil
|
||||
if caps.Sampling.Tools == nil {
|
||||
t.Error("modifying cloned Sampling.Tools should not affect original")
|
||||
}
|
||||
}
|
||||
|
|
@ -1163,6 +1163,10 @@ func (ss *ServerSession) ListRoots(ctx context.Context, params *ListRootsParams)
|
|||
}
|
||||
|
||||
// CreateMessage sends a sampling request to the client.
|
||||
//
|
||||
// If the client returns multiple content blocks (e.g. parallel tool calls),
|
||||
// CreateMessage returns an error. Use [ServerSession.CreateMessageWithTools]
|
||||
// for tool-enabled sampling.
|
||||
func (ss *ServerSession) CreateMessage(ctx context.Context, params *CreateMessageParams) (*CreateMessageResult, error) {
|
||||
if err := ss.checkInitialized(methodCreateMessage); err != nil {
|
||||
return nil, err
|
||||
|
|
@ -1175,7 +1179,44 @@ func (ss *ServerSession) CreateMessage(ctx context.Context, params *CreateMessag
|
|||
p2.Messages = []*SamplingMessage{} // avoid JSON "null"
|
||||
params = &p2
|
||||
}
|
||||
return handleSend[*CreateMessageResult](ctx, methodCreateMessage, newServerRequest(ss, orZero[Params](params)))
|
||||
res, err := handleSend[*CreateMessageWithToolsResult](ctx, methodCreateMessage, newServerRequest(ss, orZero[Params](params)))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Downconvert to singular content.
|
||||
if len(res.Content) > 1 {
|
||||
return nil, fmt.Errorf("CreateMessage result has %d content blocks; use CreateMessageWithTools for multiple content", len(res.Content))
|
||||
}
|
||||
var content Content
|
||||
if len(res.Content) > 0 {
|
||||
content = res.Content[0]
|
||||
}
|
||||
return &CreateMessageResult{
|
||||
Meta: res.Meta,
|
||||
Content: content,
|
||||
Model: res.Model,
|
||||
Role: res.Role,
|
||||
StopReason: res.StopReason,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// CreateMessageWithTools sends a sampling request with tools to the client,
|
||||
// returning a [CreateMessageWithToolsResult] that supports array content
|
||||
// (for parallel tool calls). Use this instead of [ServerSession.CreateMessage]
|
||||
// when the request includes tools.
|
||||
func (ss *ServerSession) CreateMessageWithTools(ctx context.Context, params *CreateMessageWithToolsParams) (*CreateMessageWithToolsResult, error) {
|
||||
if err := ss.checkInitialized(methodCreateMessage); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if params == nil {
|
||||
params = &CreateMessageWithToolsParams{Messages: []*SamplingMessageV2{}}
|
||||
}
|
||||
if params.Messages == nil {
|
||||
p2 := *params
|
||||
p2.Messages = []*SamplingMessageV2{} // avoid JSON "null"
|
||||
params = &p2
|
||||
}
|
||||
return handleSend[*CreateMessageWithToolsResult](ctx, methodCreateMessage, newServerRequest(ss, orZero[Params](params)))
|
||||
}
|
||||
|
||||
// Elicit sends an elicitation request to the client asking for user input.
|
||||
|
|
|
|||
Loading…
Reference in a new issue