feat: multi-round-trip request implementation (SEP-2322) (#950)

### Context

Design (`design/mrtr.md`) and implementation proposal for Multi
Round-Trip Requests (MRTR) per
[SEP-2322](de6d76fba3/seps/2322-MRTR.md).

### Changes

- Introduce `InputRequest` and `InputResponse` sealed interfaces with
corresponding `InputRequestMap` and `InputResponseMap` types for custom
map JSON codec implementation.
- Extend `CallToolParams`, `GetPromptParams`, and `ReadResourceParams`
with `InputResponses`, `RequestState` and an unexported `resultType`
("complete" or "input_required") for retry round-trips. The type is set
by the SDK based on input requests presence.
- Add client-side MRTR middleware enabled by default that automatically
fulfills input requests and retries.
- Added MRTROptions on ClientOptions to allow disabling the middleware
and configuring the number of retries.
- Add server-side MRTR middleware for backward compatibility:
transparently bridges input requests to direct
Elicit/CreateMessage/ListRoots calls for older clients
This commit is contained in:
Yaroslav 2026-05-29 09:29:34 +02:00 • committed by GitHub
commit 189a85ad92
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 1587 additions and 23 deletions

View file

@ -58,13 +58,17 @@ func NewClient(impl *Implementation, options *ClientOptions) *Client {
opts.Logger = ensureLogger(nil)
}
return &Client{
c := &Client{
impl: impl,
opts: opts,
roots: newFeatureSet(func(r *Root) string { return r.URI }),
sendingMethodHandler_: defaultSendingMethodHandler,
receivingMethodHandler_: defaultReceivingMethodHandler[*ClientSession],
}
if opts.MultiRoundTrip == nil || !opts.MultiRoundTrip.Disabled {
c.AddSendingMiddleware(clientMultiRoundTripMiddleware())
}
return c
}
// ClientOptions configures the behavior of the client.
@ -154,6 +158,10 @@ type ClientOptions struct {
ResourceUpdatedHandler func(context.Context, *ResourceUpdatedNotificationRequest)
LoggingMessageHandler func(context.Context, *LoggingMessageRequest)
ProgressNotificationHandler func(context.Context, *ProgressNotificationClientRequest)
// MultiRoundTrip configures the automatic MultiRoundTrip (Multi Round-Trip Requests) middleware.
// By default (nil), the middleware is enabled with default settings.
// Set Disabled to true to opt out of automatic MultiRoundTrip handling.
MultiRoundTrip *MultiRoundTripOptions
// If non-zero, defines an interval for regular "ping" requests.
// If the peer fails to respond to pings originating from the keepalive check,
// the session is automatically closed.

View file

@ -15,6 +15,7 @@ import (
"testing"
"github.com/google/go-cmp/cmp"
"github.com/google/go-cmp/cmp/cmpopts"
"github.com/modelcontextprotocol/go-sdk/mcp"
)
@ -223,4 +224,4 @@ func TestContentUnmarshalNilWithInvalidContent(t *testing.T) {
}
}
var ctrCmpOpts = []cmp.Option{cmp.AllowUnexported(mcp.CallToolResult{})}
var ctrCmpOpts = []cmp.Option{cmpopts.IgnoreUnexported(mcp.CallToolResult{}, mcp.GetPromptResult{}, mcp.ReadResourceResult{})}

View file

@ -24,6 +24,7 @@ import (
"time"
"github.com/google/go-cmp/cmp"
"github.com/google/go-cmp/cmp/cmpopts"
"github.com/google/jsonschema-go/jsonschema"
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
"github.com/modelcontextprotocol/go-sdk/jsonrpc"
@ -206,7 +207,7 @@ func TestEndToEnd(t *testing.T) {
Role: "user",
}},
}
if diff := cmp.Diff(wantReview, gotReview); diff != "" {
if diff := cmp.Diff(wantReview, gotReview, ctrCmpOpts...); diff != "" {
t.Errorf("prompts/get 'code_review' mismatch (-want +got):\n%s", diff)
}
@ -2371,4 +2372,4 @@ func TestSetErrorPreservesContent(t *testing.T) {
}
}
var ctrCmpOpts = []cmp.Option{cmp.AllowUnexported(CallToolResult{})}
var ctrCmpOpts = []cmp.Option{cmpopts.IgnoreUnexported(CallToolResult{}, GetPromptResult{}, ReadResourceResult{})}

268
mcp/mrtr.go Normal file
View file

@ -0,0 +1,268 @@
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
// Use of this source code is governed by the license
// that can be found in the LICENSE file.
package mcp
import (
"context"
"fmt"
"log/slog"
"sync"
"github.com/modelcontextprotocol/go-sdk/jsonrpc"
"golang.org/x/sync/errgroup"
)
const maxMultiRoundTripRetries = 10
const maxLoadSheddingMultiRoundTripRetries = 3
// MultiRoundTripOptions configures the client-side multi round-trip request (SEP-2322)
// middleware. The middleware is enabled by default and automatically fulfills input
// requests from the server by invoking the appropriate client handlers and
// retrying the original call.
type MultiRoundTripOptions struct {
// Disabled prevents the automatic multi-round-tirp middleware from being installed.
// When true, the client returns input-required results directly and callers must
// handle the retry loop themselves using [CallToolResult.NeedsInput],
// [GetPromptResult.NeedsInput], or [ReadResourceResult.NeedsInput].
Disabled bool
}
type multiRoundTripResponse interface {
setResultType(resultType)
inputRequests() map[string]InputRequest
requestState() string
hasContent() bool
}
func handleMultiRoundTripResult(ss *ServerSession, logger *slog.Logger, res multiRoundTripResponse) error {
if res == nil {
return nil
}
hasInputRequests := res.inputRequests() != nil
if hasInputRequests && res.hasContent() {
logger.Warn("handler returned both content and inputRequests")
return &jsonrpc.Error{
Code: jsonrpc.CodeInternalError,
Message: "server bug: result has both content and inputRequests",
}
}
if clientSupportsMultiRoundTrip(ss) {
// For older clients the resultType is left unset. Input requests will be handled
// by serverMultiRoundTripMiddleware client calls and handler reinvocation.
if hasInputRequests {
res.setResultType(resultTypeInputRequired)
} else {
res.setResultType(resultTypeComplete)
}
}
return nil
}
func clientSupportsMultiRoundTrip(ss *ServerSession) bool {
protocolVersion := latestProtocolVersion
if iparams := ss.InitializeParams(); iparams != nil {
protocolVersion = iparams.ProtocolVersion
}
return protocolVersion >= protocolVersion20260630
}
func clientMultiRoundTripMiddleware() Middleware {
return func(next MethodHandler) MethodHandler {
return func(ctx context.Context, method string, req Request) (Result, error) {
if method != methodCallTool && method != methodGetPrompt && method != methodReadResource {
return next(ctx, method, req)
}
loadSheddingFailures := 0
for retries := 1; ; retries++ {
res, err := next(ctx, method, req)
if err != nil {
return res, err
}
mrtrResult, ok := res.(multiRoundTripResponse)
if !ok {
return res, nil
}
reqMap := mrtrResult.inputRequests()
if reqMap == nil {
return res, nil
}
if len(reqMap) == 0 {
loadSheddingFailures++
}
if loadSheddingFailures >= maxLoadSheddingMultiRoundTripRetries {
return nil, fmt.Errorf("multi-round-trip: exceeded maximum load-shedding retries (%d)", maxLoadSheddingMultiRoundTripRetries)
}
if retries >= maxMultiRoundTripRetries {
return nil, fmt.Errorf("multi-round-trip: exceeded maximum retries (%d)", maxMultiRoundTripRetries)
}
cs, ok := req.GetSession().(*ClientSession)
if !ok {
return res, nil
}
responses, err := fulfillInputRequests(ctx, cs, reqMap)
if err != nil {
return nil, err
}
setMultiRoundTripRetryParams(req, responses, mrtrResult.requestState())
}
}
}
}
// serverMultiRoundTripMiddleware is a receiving middleware for servers that transparently
// handles multi-round-trip for clients on older protocol versions. When a handler returns
// InputRequests and the client does not support multi-round-trip, the middleware fulfills
// the requests by calling the client directly and reinvokes the handler once with the responses.
func serverMultiRoundTripMiddleware() Middleware {
return func(next MethodHandler) MethodHandler {
return func(ctx context.Context, method string, req Request) (Result, error) {
if method != methodCallTool && method != methodGetPrompt && method != methodReadResource {
return next(ctx, method, req)
}
ss, ok := req.GetSession().(*ServerSession)
if !ok {
return next(ctx, method, req)
}
if clientSupportsMultiRoundTrip(ss) {
return next(ctx, method, req)
}
res, err := next(ctx, method, req)
if err != nil {
return res, err
}
mrtrResult, ok := res.(multiRoundTripResponse)
if !ok {
return res, nil
}
reqMap := mrtrResult.inputRequests()
if reqMap == nil {
return res, nil
}
if len(reqMap) == 0 {
return nil, fmt.Errorf("the server is busy, retry later")
}
responses, err := fulfillServerInputRequests(ctx, ss, reqMap)
if err != nil {
return nil, err
}
setMultiRoundTripRetryParams(req, responses, mrtrResult.requestState())
return next(ctx, method, req)
}
}
}
func fulfillServerInputRequests(ctx context.Context, ss *ServerSession, requests InputRequestMap) (InputResponseMap, error) {
g, ctx := errgroup.WithContext(ctx)
var mu sync.Mutex
responses := make(InputResponseMap, len(requests))
for id, ir := range requests {
g.Go(func() error {
resp, err := fulfillServerInputRequest(ctx, ss, ir)
if err != nil {
return fmt.Errorf("fulfilling input request %q: %w", id, err)
}
mu.Lock()
responses[id] = resp
mu.Unlock()
return nil
})
}
if err := g.Wait(); err != nil {
return nil, fmt.Errorf("multi-round-trip: %w", err)
}
return responses, nil
}
func fulfillServerInputRequest(ctx context.Context, ss *ServerSession, ir InputRequest) (InputResponse, error) {
switch p := ir.(type) {
case *ElicitParams:
return ss.Elicit(ctx, p)
case *CreateMessageParams:
return ss.CreateMessageWithTools(ctx, createMessageParamsToWithTools(p))
case *CreateMessageWithToolsParams:
return ss.CreateMessageWithTools(ctx, p)
case *ListRootsParams:
return ss.ListRoots(ctx, p)
default:
return nil, fmt.Errorf("unknown input request type: %T", ir)
}
}
func createMessageParamsToWithTools(p *CreateMessageParams) *CreateMessageWithToolsParams {
var msgs []*SamplingMessageV2
for _, m := range p.Messages {
msgs = append(msgs, &SamplingMessageV2{Content: []Content{m.Content}, Role: m.Role})
}
return &CreateMessageWithToolsParams{
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,
}
}
func setMultiRoundTripRetryParams(req Request, responses InputResponseMap, state string) {
switch p := req.GetParams().(type) {
case *CallToolParams:
p.InputResponses = responses
p.RequestState = state
case *CallToolParamsRaw:
p.InputResponses = responses
p.RequestState = state
case *GetPromptParams:
p.InputResponses = responses
p.RequestState = state
case *ReadResourceParams:
p.InputResponses = responses
p.RequestState = state
}
}
func fulfillInputRequests(ctx context.Context, cs *ClientSession, requests InputRequestMap) (InputResponseMap, error) {
g, ctx := errgroup.WithContext(ctx)
var mu sync.Mutex
responses := make(InputResponseMap, len(requests))
for id, ir := range requests {
g.Go(func() error {
resp, err := fulfillInputRequest(ctx, cs, ir)
if err != nil {
return fmt.Errorf("fulfilling input request %q: %w", id, err)
}
mu.Lock()
responses[id] = resp
mu.Unlock()
return nil
})
}
if err := g.Wait(); err != nil {
return nil, fmt.Errorf("multi round-trip: %w", err)
}
return responses, nil
}
func fulfillInputRequest(ctx context.Context, cs *ClientSession, ir InputRequest) (InputResponse, error) {
switch p := ir.(type) {
case *ElicitParams:
return cs.client.elicit(ctx, newClientRequest(cs, p))
case *CreateMessageParams:
return cs.client.createMessage(ctx, &CreateMessageWithToolsRequest{Session: cs, Params: createMessageParamsToWithTools(p)})
case *CreateMessageWithToolsParams:
return cs.client.createMessage(ctx, &CreateMessageWithToolsRequest{Session: cs, Params: p})
case *ListRootsParams:
return cs.client.listRoots(ctx, newClientRequest(cs, p))
default:
return nil, fmt.Errorf("unknown input request type: %T", ir)
}
}

546
mcp/mrtr_test.go Normal file
View file

@ -0,0 +1,546 @@
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
// Use of this source code is governed by the license
// that can be found in the LICENSE file.
package mcp
import (
"context"
"fmt"
"slices"
"sync/atomic"
"testing"
"github.com/google/go-cmp/cmp"
"github.com/google/jsonschema-go/jsonschema"
)
func TestMultiRoundTrip_ManualRetry(t *testing.T) {
type deployResult struct {
Deployed bool `json:"deployed"`
Reason string `json:"reason,omitempty"`
}
orig := supportedProtocolVersions
supportedProtocolVersions = append(slices.Clone(orig), protocolVersion20260630)
t.Cleanup(func() { supportedProtocolVersions = orig })
ctx := context.Background()
srv := NewServer(testImpl, nil)
AddTool(srv, &Tool{Name: "deploy"}, func(ctx context.Context, req *CallToolRequest, input struct{}) (*CallToolResult, *deployResult, error) {
if len(req.Params.InputResponses) == 0 {
return &CallToolResult{
InputRequests: InputRequestMap{"confirm": &ElicitParams{Message: "Deploy to production?"}},
RequestState: "deployment-123",
}, nil, nil
}
resp, ok := req.Params.InputResponses["confirm"]
if !ok {
return &CallToolResult{
InputRequests: InputRequestMap{"confirm": &ElicitParams{Message: "Please confirm (retry)"}},
}, nil, nil
}
if req.Params.RequestState == "" {
return &CallToolResult{}, &deployResult{Deployed: false, Reason: "no_state"}, nil
}
if elicitResult := resp.(*ElicitResult); elicitResult != nil && elicitResult.Action != "accept" {
return &CallToolResult{}, &deployResult{Deployed: false, Reason: "cancelled"}, nil
}
return &CallToolResult{}, &deployResult{Deployed: true}, nil
})
conn := mustConnect(t, srv, &ClientOptions{
MultiRoundTrip: &MultiRoundTripOptions{Disabled: true},
})
// Round 1: initiate deployment
res, err := conn.CallTool(ctx, &CallToolParams{Name: "deploy"})
if err != nil {
t.Fatalf("CallTool() error = %v", err)
}
if !res.NeedsInput() {
t.Fatal("NeedsInput() = false, want true")
}
if got := len(res.InputRequests); got != 1 {
t.Fatalf("len(res.InputRequests) = %d, want 1", got)
}
if _, ok := res.InputRequests["confirm"].(*ElicitParams); !ok {
t.Fatalf("res.InputRequests[confirm] type = %T, want *ElicitParams", res.InputRequests["confirm"])
}
// Round 2: retry with confirmation
res, err = conn.CallTool(ctx, &CallToolParams{
Name: "deploy",
InputResponses: InputResponseMap{
"confirm": &ElicitResult{Action: "accept", Content: map[string]any{"ok": true}},
},
RequestState: res.RequestState,
})
if err != nil {
t.Fatalf("CallTool() follow-up error = %v", err)
}
if res.NeedsInput() {
t.Fatal("NeedsInput() = true after follow-up, want false")
}
if diff := cmp.Diff(map[string]any{"deployed": true}, res.StructuredContent, ctrCmpOpts...); diff != "" {
t.Errorf("result mismatch (-want +got):\n%s", diff)
}
}
func TestMultiRoundTrip_AutoRetry(t *testing.T) {
orig := supportedProtocolVersions
supportedProtocolVersions = append(slices.Clone(orig), protocolVersion20260630)
t.Cleanup(func() { supportedProtocolVersions = orig })
tests := []struct {
name string
inputRequests InputRequestMap
wantResult map[string]any
}{
{
name: "elicit",
inputRequests: InputRequestMap{
"confirm": &ElicitParams{Message: "Deploy?"},
},
wantResult: map[string]any{"ids": []any{"confirm"}},
},
{
name: "createMessage",
inputRequests: InputRequestMap{
"summarize": &CreateMessageParams{
Messages: []*SamplingMessage{{Role: "user", Content: &TextContent{Text: "summarize"}}},
MaxTokens: 100,
},
},
wantResult: map[string]any{"ids": []any{"summarize"}},
},
{
name: "listRoots",
inputRequests: InputRequestMap{
"roots": &ListRootsParams{},
},
wantResult: map[string]any{"ids": []any{"roots"}},
},
{
name: "all three",
inputRequests: InputRequestMap{
"confirm": &ElicitParams{Message: "OK?"},
"draft": &CreateMessageParams{
Messages: []*SamplingMessage{{Role: "user", Content: &TextContent{Text: "write"}}},
MaxTokens: 50,
},
"roots": &ListRootsParams{},
},
wantResult: map[string]any{"ids": []any{"confirm", "draft", "roots"}},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
ctx := context.Background()
srv := NewServer(testImpl, nil)
inputRequests := tt.inputRequests
AddTool(srv, &Tool{Name: "act"}, func(ctx context.Context, req *CallToolRequest, input struct{}) (*CallToolResult, any, error) {
if len(req.Params.InputResponses) == 0 {
return &CallToolResult{
InputRequests: inputRequests,
RequestState: "state-1",
}, nil, nil
}
// Collect the IDs of fulfilled responses.
var ids []string
for id := range req.Params.InputResponses {
ids = append(ids, id)
}
slices.Sort(ids)
return &CallToolResult{}, map[string]any{"ids": ids}, nil
})
conn := mustConnect(t, srv, &ClientOptions{
ElicitationHandler: func(_ context.Context, req *ElicitRequest) (*ElicitResult, error) {
return &ElicitResult{Action: "accept"}, nil
},
CreateMessageHandler: func(_ context.Context, req *CreateMessageRequest) (*CreateMessageResult, error) {
return &CreateMessageResult{
Model: "test-model",
Role: "assistant",
Content: &TextContent{Text: "response"},
}, nil
},
})
conn.client.AddRoots(&Root{URI: "file:///workspace", Name: "workspace"})
res, err := conn.CallTool(ctx, &CallToolParams{Name: "act"})
if err != nil {
t.Fatalf("CallTool() error = %v", err)
}
if res.NeedsInput() {
t.Fatal("NeedsInput() = true after auto-retry, want false")
}
// Sort the expected IDs for stable comparison.
if wantIDs, ok := tt.wantResult["ids"].([]any); ok {
slices.SortFunc(wantIDs, func(a, b any) int {
if a.(string) < b.(string) {
return -1
}
return 1
})
}
if diff := cmp.Diff(tt.wantResult, res.StructuredContent, ctrCmpOpts...); diff != "" {
t.Errorf("result mismatch (-want +got):\n%s", diff)
}
})
}
}
func TestMultiRoundTrip_MaxRetries(t *testing.T) {
testCases := []struct {
name string
requests InputRequestMap
wantRetries int
}{
{
name: "load shedding",
requests: InputRequestMap{},
wantRetries: maxLoadSheddingMultiRoundTripRetries,
},
{
name: "input request",
requests: InputRequestMap{"confirm": &ElicitParams{Message: "Again?"}},
wantRetries: maxMultiRoundTripRetries,
},
}
orig := supportedProtocolVersions
supportedProtocolVersions = append(slices.Clone(orig), protocolVersion20260630)
t.Cleanup(func() { supportedProtocolVersions = orig })
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
ctx := context.Background()
var serverCalls atomic.Int32
srv := NewServer(testImpl, nil)
AddTool(srv, &Tool{Name: "loop"}, func(ctx context.Context, req *CallToolRequest, input struct{}) (*CallToolResult, any, error) {
serverCalls.Add(1)
return &CallToolResult{InputRequests: tc.requests, RequestState: "loop-state"}, nil, nil
})
conn := mustConnect(t, srv, &ClientOptions{
ElicitationHandler: func(_ context.Context, req *ElicitRequest) (*ElicitResult, error) {
return &ElicitResult{Action: "accept"}, nil
},
})
_, err := conn.CallTool(ctx, &CallToolParams{Name: "loop"})
if err == nil {
t.Fatal("CallTool() err = nil, want error for exceeded max retries")
}
if serverCalls.Load() != int32(tc.wantRetries) {
t.Errorf("serverCalls = %d, want %d", serverCalls.Load(), tc.wantRetries)
}
})
}
}
func TestMultiRoundTrip_ServerMiddleware(t *testing.T) {
// multiRoundTripToolHandler returns a ToolHandler (plain, non-generic) that requests
// the given inputRequests on the first call and returns the fulfilled
// response IDs on the second.
multiRoundTripToolHandler := func(inputRequests InputRequestMap) ToolHandler {
return func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) {
if len(req.Params.InputResponses) == 0 {
return &CallToolResult{
InputRequests: inputRequests,
RequestState: "state-1",
}, nil
}
var ids []string
for id := range req.Params.InputResponses {
ids = append(ids, id)
}
slices.Sort(ids)
content := &TextContent{Text: fmt.Sprintf("%v", ids)}
return &CallToolResult{Content: []Content{content}}, nil
}
}
tests := []struct {
name string
inputRequests InputRequestMap
wantText string
}{
{
name: "elicit via ToolHandler",
inputRequests: InputRequestMap{
"confirm": &ElicitParams{Message: "Sure?"},
},
wantText: "[confirm]",
},
{
name: "elicit and listRoots via ToolHandler",
inputRequests: InputRequestMap{
"confirm": &ElicitParams{Message: "OK?"},
"roots": &ListRootsParams{},
},
wantText: "[confirm roots]",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
ctx := context.Background()
srv := NewServer(testImpl, nil)
srv.AddTool(
&Tool{Name: "act", InputSchema: &jsonschema.Schema{Type: "object"}},
multiRoundTripToolHandler(tt.inputRequests),
)
// Connect with an OLD protocol version where multi-round-trip is not supported.
// The server middleware should handle it transparently.
st, ct := NewInMemoryTransports()
ss, err := srv.Connect(t.Context(), st, nil)
if err != nil {
t.Fatalf("server.Connect() error = %v", err)
}
t.Cleanup(func() { _ = ss.Close() })
c := NewClient(testImpl, &ClientOptions{
MultiRoundTrip: &MultiRoundTripOptions{Disabled: true},
ElicitationHandler: func(_ context.Context, req *ElicitRequest) (*ElicitResult, error) {
return &ElicitResult{Action: "accept"}, nil
},
})
c.AddRoots(&Root{URI: "file:///workspace", Name: "workspace"})
cs, err := c.Connect(t.Context(), ct, &ClientSessionOptions{})
if err != nil {
t.Fatalf("client.Connect() error = %v", err)
}
t.Cleanup(func() { _ = cs.Close() })
res, err := cs.CallTool(ctx, &CallToolParams{Name: "act"})
if err != nil {
t.Fatalf("CallTool() error = %v", err)
}
if got := res.Content[0].(*TextContent).Text; got != tt.wantText {
t.Errorf("result text = %q, want %q", got, tt.wantText)
}
})
}
}
func TestMultiRoundTrip_GetPrompt_AutoRetry(t *testing.T) {
orig := supportedProtocolVersions
supportedProtocolVersions = append(slices.Clone(orig), protocolVersion20260630)
t.Cleanup(func() { supportedProtocolVersions = orig })
ctx := context.Background()
srv := NewServer(testImpl, nil)
srv.AddPrompt(&Prompt{Name: "review"}, func(_ context.Context, req *GetPromptRequest) (*GetPromptResult, error) {
if len(req.Params.InputResponses) == 0 {
return &GetPromptResult{
InputRequests: InputRequestMap{"confirm": &ElicitParams{Message: "Include sensitive data?"}},
RequestState: "prompt-state",
}, nil
}
return &GetPromptResult{
Description: "Code review prompt",
Messages: []*PromptMessage{{Role: "user", Content: &TextContent{Text: "review this code"}}},
}, nil
})
conn := mustConnect(t, srv, &ClientOptions{
ElicitationHandler: func(_ context.Context, _ *ElicitRequest) (*ElicitResult, error) {
return &ElicitResult{Action: "accept"}, nil
},
})
res, err := conn.GetPrompt(ctx, &GetPromptParams{Name: "review"})
if err != nil {
t.Fatalf("GetPrompt() error = %v", err)
}
if res.NeedsInput() {
t.Fatal("NeedsInput() = true after auto-retry, want false")
}
if len(res.Messages) != 1 {
t.Fatalf("len(res.Messages) = %d, want 1", len(res.Messages))
}
if got := res.Messages[0].Content.(*TextContent).Text; got != "review this code" {
t.Errorf("message text = %q, want %q", got, "review this code")
}
}
func TestMultiRoundTrip_GetPrompt_ManualRetry(t *testing.T) {
orig := supportedProtocolVersions
supportedProtocolVersions = append(slices.Clone(orig), protocolVersion20260630)
t.Cleanup(func() { supportedProtocolVersions = orig })
ctx := context.Background()
srv := NewServer(testImpl, nil)
srv.AddPrompt(&Prompt{Name: "review"}, func(_ context.Context, req *GetPromptRequest) (*GetPromptResult, error) {
if len(req.Params.InputResponses) == 0 {
return &GetPromptResult{
InputRequests: InputRequestMap{"confirm": &ElicitParams{Message: "Include sensitive data?"}},
RequestState: "prompt-state",
}, nil
}
return &GetPromptResult{
Description: "Code review prompt",
Messages: []*PromptMessage{{Role: "user", Content: &TextContent{Text: "review this code"}}},
}, nil
})
conn := mustConnect(t, srv, &ClientOptions{
MultiRoundTrip: &MultiRoundTripOptions{Disabled: true},
})
res, err := conn.GetPrompt(ctx, &GetPromptParams{Name: "review"})
if err != nil {
t.Fatalf("GetPrompt() error = %v", err)
}
if !res.NeedsInput() {
t.Fatal("NeedsInput() = false, want true")
}
if _, ok := res.InputRequests["confirm"].(*ElicitParams); !ok {
t.Fatalf("InputRequests[confirm] type = %T, want *ElicitParams", res.InputRequests["confirm"])
}
res, err = conn.GetPrompt(ctx, &GetPromptParams{
Name: "review",
InputResponses: InputResponseMap{"confirm": &ElicitResult{Action: "accept"}},
RequestState: res.RequestState,
})
if err != nil {
t.Fatalf("GetPrompt() follow-up error = %v", err)
}
if res.NeedsInput() {
t.Fatal("NeedsInput() = true after follow-up, want false")
}
if len(res.Messages) != 1 {
t.Fatalf("len(res.Messages) = %d, want 1", len(res.Messages))
}
}
func TestMultiRoundTrip_ReadResource_AutoRetry(t *testing.T) {
orig := supportedProtocolVersions
supportedProtocolVersions = append(slices.Clone(orig), protocolVersion20260630)
t.Cleanup(func() { supportedProtocolVersions = orig })
ctx := context.Background()
srv := NewServer(testImpl, nil)
srv.AddResource(&Resource{URI: "test://data", Name: "data"}, func(_ context.Context, req *ReadResourceRequest) (*ReadResourceResult, error) {
if len(req.Params.InputResponses) == 0 {
return &ReadResourceResult{
InputRequests: InputRequestMap{"auth": &ElicitParams{Message: "Authenticate?"}},
RequestState: "resource-state",
}, nil
}
return &ReadResourceResult{
Contents: []*ResourceContents{{URI: "test://data", Text: "resource data"}},
}, nil
})
conn := mustConnect(t, srv, &ClientOptions{
ElicitationHandler: func(_ context.Context, _ *ElicitRequest) (*ElicitResult, error) {
return &ElicitResult{Action: "accept"}, nil
},
})
res, err := conn.ReadResource(ctx, &ReadResourceParams{URI: "test://data"})
if err != nil {
t.Fatalf("ReadResource() error = %v", err)
}
if res.NeedsInput() {
t.Fatal("NeedsInput() = true after auto-retry, want false")
}
if len(res.Contents) != 1 {
t.Fatalf("len(res.Contents) = %d, want 1", len(res.Contents))
}
if got := res.Contents[0].Text; got != "resource data" {
t.Errorf("resource text = %q, want %q", got, "resource data")
}
}
func TestMultiRoundTrip_ReadResource_ManualRetry(t *testing.T) {
orig := supportedProtocolVersions
supportedProtocolVersions = append(slices.Clone(orig), protocolVersion20260630)
t.Cleanup(func() { supportedProtocolVersions = orig })
ctx := context.Background()
srv := NewServer(testImpl, nil)
srv.AddResource(&Resource{URI: "test://data", Name: "data"}, func(_ context.Context, req *ReadResourceRequest) (*ReadResourceResult, error) {
if len(req.Params.InputResponses) == 0 {
return &ReadResourceResult{
InputRequests: InputRequestMap{"auth": &ElicitParams{Message: "Authenticate?"}},
RequestState: "resource-state",
}, nil
}
return &ReadResourceResult{
Contents: []*ResourceContents{{URI: "test://data", Text: "resource data"}},
}, nil
})
conn := mustConnect(t, srv, &ClientOptions{
MultiRoundTrip: &MultiRoundTripOptions{Disabled: true},
})
res, err := conn.ReadResource(ctx, &ReadResourceParams{URI: "test://data"})
if err != nil {
t.Fatalf("ReadResource() error = %v", err)
}
if !res.NeedsInput() {
t.Fatal("NeedsInput() = false, want true")
}
if _, ok := res.InputRequests["auth"].(*ElicitParams); !ok {
t.Fatalf("InputRequests[auth] type = %T, want *ElicitParams", res.InputRequests["auth"])
}
res, err = conn.ReadResource(ctx, &ReadResourceParams{
URI: "test://data",
InputResponses: InputResponseMap{"auth": &ElicitResult{Action: "accept"}},
RequestState: res.RequestState,
})
if err != nil {
t.Fatalf("ReadResource() follow-up error = %v", err)
}
if res.NeedsInput() {
t.Fatal("NeedsInput() = true after follow-up, want false")
}
if len(res.Contents) != 1 {
t.Fatalf("len(res.Contents) = %d, want 1", len(res.Contents))
}
}
func mustConnect(t *testing.T, s *Server, clientOpts *ClientOptions) *ClientSession {
t.Helper()
st, ct := NewInMemoryTransports()
ss, err := s.Connect(t.Context(), st, nil)
if err != nil {
t.Fatalf("server.Connect() error = %v", err)
}
t.Cleanup(func() {
_ = ss.Close()
})
c := NewClient(testImpl, clientOpts)
cs, err := c.Connect(t.Context(), ct, &ClientSessionOptions{protocolVersion: protocolVersion20260630})
if err != nil {
t.Fatalf("client.Connect() error = %v", err)
}
t.Cleanup(func() {
_ = cs.Close()
})
return cs
}

View file

@ -13,6 +13,178 @@ import (
"github.com/modelcontextprotocol/go-sdk/internal/mcpgodebug"
)
// resultType indicates whether a result is complete or requires further input
// from the client via the multi round-trip request protocol.
type resultType string
const (
// resultTypeComplete indicates the result is final.
// This is the default when ResultType is empty.
resultTypeComplete resultType = "complete"
// resultTypeInputRequired indicates the server needs additional client
// input before it can complete the request. The client should fulfill the
// InputRequests and retry the call with the responses.
resultTypeInputRequired resultType = "input_required"
)
// InputRequest is a type for parameters that a server can include in the response
// to request input from client (SEP-2322). Implementations are [*ElicitParams],
// [*CreateMessageParams], and [*ListRootsParams].
type InputRequest interface{ isInputRequest() }
// InputRequestMap maps server-assigned request IDs to [InputRequest] values.
// It is used in result types to tell the client what input the server needs.
type InputRequestMap map[string]InputRequest
func (m InputRequestMap) MarshalJSON() ([]byte, error) {
if m == nil {
return json.Marshal(map[string]any(nil))
}
type wire struct {
Method string `json:"method"`
Params InputRequest `json:"params,omitempty"`
}
typeToMethod := func(v InputRequest) (string, error) {
switch v.(type) {
case *ElicitParams:
return methodElicit, nil
case *CreateMessageParams, *CreateMessageWithToolsParams:
return methodCreateMessage, nil
case *ListRootsParams:
return methodListRoots, nil
default:
return "", fmt.Errorf("unsupported type: %T", v)
}
}
converted := map[string]*wire{}
for k, v := range m {
method, err := typeToMethod(v)
if err != nil {
return nil, err
}
converted[k] = &wire{Method: method, Params: v}
}
return json.Marshal(converted)
}
func (m *InputRequestMap) UnmarshalJSON(data []byte) error {
type raw struct {
Method string `json:"method"`
Params json.RawMessage `json:"params"`
}
var rawMap map[string]*raw
if err := json.Unmarshal(data, &rawMap); err != nil {
return err
}
if rawMap == nil {
return nil
}
result := make(InputRequestMap, len(rawMap))
for k, raw := range rawMap {
switch raw.Method {
case methodElicit:
var p ElicitParams
if err := json.Unmarshal(raw.Params, &p); err != nil {
return err
}
result[k] = &p
case methodCreateMessage:
var p CreateMessageWithToolsParams
if err := json.Unmarshal(raw.Params, &p); err != nil {
return err
}
result[k] = &p
case methodListRoots:
var p ListRootsParams
if err := json.Unmarshal(raw.Params, &p); err != nil {
return err
}
result[k] = &p
default:
return fmt.Errorf("unsupported InputRequest method: %q", raw.Method)
}
}
*m = result
return nil
}
// InputResponse is a type for results that a client sends back when fulfilling
// a server input request (SEP-2322). Implementations are [*ElicitResult],
// [*CreateMessageResult], and [*ListRootsResult].
type InputResponse interface{ isInputResponse() }
// InputResponseMap maps request IDs (from [InputRequestMap]) to [InputResponse]
// values. It is used in params types when retrying a call after an
// input-required result.
type InputResponseMap map[string]InputResponse
func (m InputResponseMap) MarshalJSON() ([]byte, error) {
type wire struct {
Method string `json:"method"`
Result InputResponse `json:"result,omitempty"`
}
typeToMethod := func(v InputResponse) (string, error) {
switch v.(type) {
case *ElicitResult:
return methodElicit, nil
case *CreateMessageResult, *CreateMessageWithToolsResult:
return methodCreateMessage, nil
case *ListRootsResult:
return methodListRoots, nil
default:
return "", fmt.Errorf("unsupported type: %T", v)
}
}
converted := map[string]*wire{}
for k, v := range m {
method, err := typeToMethod(v)
if err != nil {
return nil, err
}
converted[k] = &wire{Method: method, Result: v}
}
return json.Marshal(converted)
}
func (m *InputResponseMap) UnmarshalJSON(data []byte) error {
type raw struct {
Method string `json:"method"`
Result json.RawMessage `json:"result"`
}
var rawMap map[string]*raw
if err := json.Unmarshal(data, &rawMap); err != nil {
return err
}
result := make(InputResponseMap, len(rawMap))
for k, raw := range rawMap {
switch raw.Method {
case methodElicit:
var p ElicitResult
if err := json.Unmarshal(raw.Result, &p); err != nil {
return err
}
result[k] = &p
case methodCreateMessage:
var p CreateMessageWithToolsResult
if err := json.Unmarshal(raw.Result, &p); err != nil {
return err
}
result[k] = &p
case methodListRoots:
var p ListRootsResult
if err := json.Unmarshal(raw.Result, &p); err != nil {
return err
}
result[k] = &p
default:
return fmt.Errorf("unsupported InputResponse method: %q", raw.Method)
}
}
*m = result
return nil
}
// Optional annotations for the client. The client can use annotations to inform
// how objects are used or displayed.
type Annotations struct {
@ -46,6 +218,14 @@ type CallToolParams struct {
// Arguments holds the tool arguments. It can hold any value that can be
// marshaled to JSON.
Arguments any `json:"arguments,omitempty"`
// InputResponses maps input request IDs to responses, provided when
// retrying a call after receiving a result with ResultType
// ResultTypeInputRequired.
InputResponses InputResponseMap `json:"inputResponses,omitempty"`
// RequestState is the opaque state from the previous input-required result.
// The client must echo this back when retrying.
RequestState string `json:"requestState,omitempty"`
}
// CallToolParamsRaw is passed to tool handlers on the server. Its arguments
@ -61,6 +241,14 @@ type CallToolParamsRaw struct {
// is the responsibility of the tool handler to unmarshal and validate the
// Arguments (see [AddTool]).
Arguments json.RawMessage `json:"arguments,omitempty"`
// InputResponses maps input request IDs to responses, provided when
// retrying a call after receiving a result with ResultType
// ResultTypeInputRequired.
InputResponses InputResponseMap `json:"inputResponses,omitempty"`
// RequestState is the opaque state from the previous input-required result.
// The client must echo this back when retrying.
RequestState string `json:"requestState,omitempty"`
}
// A CallToolResult is the server's response to a tool call.
@ -107,6 +295,24 @@ type CallToolResult struct {
// the Content field.
IsError bool `json:"isError,omitempty"`
// InputRequests is a map of server-assigned IDs to input requests.
// Populated only when ResultType is ResultTypeInputRequired.
// The client must fulfill these and echo the IDs back in InputResponses
// when retrying the call.
InputRequests InputRequestMap `json:"inputRequests,omitempty"`
// RequestState is an opaque string the client must echo back when
// retrying after an input-required result. Servers use this to carry
// context between independent requests.
//
// Unauthenticated servers must encrypt, sign and verify this value.
RequestState string `json:"requestState,omitempty"`
// ResultType indicates whether this result is complete or requires further
// client input. Empty or ResultTypeComplete means the call succeeded
// normally. ResultTypeInputRequired means the client should fulfill the
// InputRequests and retry the call.
resultType resultType
// The error passed to setError, if any.
// It is not marshaled, and therefore it is only visible on the server.
// Its only use is in server sending middleware, where it can be accessed
@ -145,13 +351,49 @@ func (r *CallToolResult) GetError() error {
func (*CallToolResult) isResult() {}
// UnmarshalJSON handles the unmarshalling of content into the Content
// interface.
func (r *CallToolResult) setResultType(rt resultType) { r.resultType = rt }
func (r *CallToolResult) requestState() string { return r.RequestState }
func (r *CallToolResult) inputRequests() map[string]InputRequest {
if r == nil {
return nil
}
return r.InputRequests
}
func (r *CallToolResult) hasContent() bool {
return len(r.Content) > 0 || r.StructuredContent != nil
}
// NeedsInput reports whether this result requires further client input.
// This is true when the server returned ResultType "input_required".
// When NeedsInput returns true, check InputRequests for the set of
// requests the server needs fulfilled before retrying the call.
// An empty InputRequests with NeedsInput true indicates load-shedding.
func (r *CallToolResult) NeedsInput() bool { return r.resultType == resultTypeInputRequired }
func (x *CallToolResult) MarshalJSON() ([]byte, error) {
type res CallToolResult // avoid recursion
type wire struct {
res
ResultType resultType `json:"resultType,omitempty"`
InputRequests json.RawMessage `json:"inputRequests,omitempty"` // shadows res.InputRequests
}
w := wire{res: res(*x), ResultType: x.resultType}
if x.InputRequests != nil {
ir, err := json.Marshal(x.InputRequests)
if err != nil {
return nil, err
}
w.InputRequests = ir
}
return json.Marshal(w)
}
func (x *CallToolResult) UnmarshalJSON(data []byte) error {
type res CallToolResult // avoid recursion
var wire struct {
res
Content []*wireContent `json:"content"`
Content []*wireContent `json:"content"`
ResultType resultType `json:"resultType"`
}
if err := internaljson.Unmarshal(data, &wire); err != nil {
return err
@ -160,6 +402,7 @@ func (x *CallToolResult) UnmarshalJSON(data []byte) error {
if wire.res.Content, err = contentsFromWire(wire.Content, nil); err != nil {
return err
}
wire.res.resultType = wire.ResultType
*x = CallToolResult(wire.res)
return nil
}
@ -426,6 +669,7 @@ type CreateMessageParams struct {
}
func (x *CreateMessageParams) isParams() {}
func (x *CreateMessageParams) isInputRequest() {}
func (x *CreateMessageParams) isNil() bool { return x == nil }
func (x *CreateMessageParams) GetProgressToken() any { return getProgressToken(x) }
func (x *CreateMessageParams) SetProgressToken(t any) { setProgressToken(x, t) }
@ -453,6 +697,7 @@ type CreateMessageWithToolsParams struct {
}
func (x *CreateMessageWithToolsParams) isParams() {}
func (x *CreateMessageWithToolsParams) isInputRequest() {}
func (x *CreateMessageWithToolsParams) isNil() bool { return x == nil }
func (x *CreateMessageWithToolsParams) GetProgressToken() any { return getProgressToken(x) }
func (x *CreateMessageWithToolsParams) SetProgressToken(t any) { setProgressToken(x, t) }
@ -553,7 +798,8 @@ type CreateMessageResult struct {
StopReason string `json:"stopReason,omitempty"`
}
func (*CreateMessageResult) isResult() {}
func (*CreateMessageResult) isResult() {}
func (*CreateMessageResult) isInputResponse() {}
func (r *CreateMessageResult) UnmarshalJSON(data []byte) error {
type result CreateMessageResult // avoid recursion
var wire struct {
@ -598,7 +844,8 @@ var createMessageWithToolsResultAllow = map[string]bool{
"tool_use": true,
}
func (*CreateMessageWithToolsResult) isResult() {}
func (*CreateMessageWithToolsResult) isResult() {}
func (*CreateMessageWithToolsResult) isInputResponse() {}
// MarshalJSON marshals the result. When Content has a single element, it is
// marshaled as a single object for compatibility with pre-2025-11-25
@ -657,6 +904,13 @@ type GetPromptParams struct {
Arguments map[string]string `json:"arguments,omitempty"`
// The name of the prompt or prompt template.
Name string `json:"name"`
// InputResponses maps input request IDs to responses, provided when
// retrying a call after receiving a result with ResultType
// ResultTypeInputRequired.
InputResponses InputResponseMap `json:"inputResponses,omitempty"`
// RequestState is the opaque state from the previous input-required result.
RequestState string `json:"requestState,omitempty"`
}
func (x *GetPromptParams) isParams() {}
@ -672,10 +926,67 @@ type GetPromptResult struct {
// An optional description for the prompt.
Description string `json:"description,omitempty"`
Messages []*PromptMessage `json:"messages"`
// InputRequests is populated when ResultType is ResultTypeInputRequired.
// See [CallToolResult.InputRequests].
InputRequests InputRequestMap `json:"inputRequests,omitempty"`
// RequestState is the opaque state for multi-round-trip retries.
// See [CallToolResult.RequestState].
RequestState string `json:"requestState,omitempty"`
// ResultType indicates whether this result is complete or requires further
// client input. See [CallToolResult.ResultType] for details.
resultType resultType
}
func (*GetPromptResult) isResult() {}
func (r *GetPromptResult) setResultType(rt resultType) { r.resultType = rt }
func (r *GetPromptResult) requestState() string { return r.RequestState }
func (r *GetPromptResult) inputRequests() map[string]InputRequest {
if r == nil {
return nil
}
return r.InputRequests
}
func (r *GetPromptResult) hasContent() bool { return len(r.Messages) > 0 }
// NeedsInput reports whether this result requires further client input.
// See [CallToolResult.NeedsInput] for details.
func (r *GetPromptResult) NeedsInput() bool { return r.resultType == resultTypeInputRequired }
func (x *GetPromptResult) MarshalJSON() ([]byte, error) {
type res GetPromptResult
type wire struct {
res
ResultType resultType `json:"resultType,omitempty"`
InputRequests json.RawMessage `json:"inputRequests,omitempty"` // shadows res.InputRequests
}
w := wire{res: res(*x), ResultType: x.resultType}
if x.InputRequests != nil {
ir, err := json.Marshal(x.InputRequests)
if err != nil {
return nil, err
}
w.InputRequests = ir
}
return json.Marshal(w)
}
func (x *GetPromptResult) UnmarshalJSON(data []byte) error {
type res GetPromptResult
var wire struct {
res
ResultType resultType `json:"resultType"`
}
if err := internaljson.Unmarshal(data, &wire); err != nil {
return err
}
wire.res.resultType = wire.ResultType
*x = GetPromptResult(wire.res)
return nil
}
// InitializeParams is sent by the client to initialize the session.
type InitializeParams struct {
// This property is reserved by the protocol to allow clients and servers to
@ -845,6 +1156,7 @@ type ListRootsParams struct {
}
func (x *ListRootsParams) isParams() {}
func (x *ListRootsParams) isInputRequest() {}
func (x *ListRootsParams) isNil() bool { return x == nil }
func (x *ListRootsParams) GetProgressToken() any { return getProgressToken(x) }
func (x *ListRootsParams) SetProgressToken(t any) { setProgressToken(x, t) }
@ -859,7 +1171,8 @@ type ListRootsResult struct {
Roots []*Root `json:"roots"`
}
func (*ListRootsResult) isResult() {}
func (*ListRootsResult) isResult() {}
func (*ListRootsResult) isInputResponse() {}
type ListToolsParams struct {
// This property is reserved by the protocol to allow clients and servers to
@ -1104,6 +1417,13 @@ type ReadResourceParams struct {
// The URI of the resource to read. The URI can use any protocol; it is up to
// the server how to interpret it.
URI string `json:"uri"`
// InputResponses maps input request IDs to responses, provided when
// retrying a call after receiving a result with ResultType
// ResultTypeInputRequired.
InputResponses InputResponseMap `json:"inputResponses,omitempty"`
// RequestState is the opaque state from the previous input-required result.
RequestState string `json:"requestState,omitempty"`
}
func (x *ReadResourceParams) isParams() {}
@ -1117,10 +1437,67 @@ type ReadResourceResult struct {
// attach additional metadata to their responses.
Meta `json:"_meta,omitempty"`
Contents []*ResourceContents `json:"contents"`
// InputRequests is populated when ResultType is ResultTypeInputRequired.
// See [CallToolResult.InputRequests].
InputRequests InputRequestMap `json:"inputRequests,omitempty"`
// RequestState is the opaque state for multi-round-trip retries.
// See [CallToolResult.RequestState].
RequestState string `json:"requestState,omitempty"`
// ResultType indicates whether this result is complete or requires further
// client input. See [CallToolResult.ResultType] for details.
resultType resultType
}
func (*ReadResourceResult) isResult() {}
func (r *ReadResourceResult) setResultType(rt resultType) { r.resultType = rt }
func (r *ReadResourceResult) requestState() string { return r.RequestState }
func (r *ReadResourceResult) inputRequests() map[string]InputRequest {
if r == nil {
return nil
}
return r.InputRequests
}
func (r *ReadResourceResult) hasContent() bool { return len(r.Contents) > 0 }
// NeedsInput reports whether this result requires further client input.
// See [CallToolResult.NeedsInput] for details.
func (r *ReadResourceResult) NeedsInput() bool { return r.resultType == resultTypeInputRequired }
func (x *ReadResourceResult) MarshalJSON() ([]byte, error) {
type res ReadResourceResult
type wire struct {
res
ResultType resultType `json:"resultType,omitempty"`
InputRequests json.RawMessage `json:"inputRequests,omitempty"` // shadows res.InputRequests
}
w := wire{res: res(*x), ResultType: x.resultType}
if x.InputRequests != nil {
ir, err := json.Marshal(x.InputRequests)
if err != nil {
return nil, err
}
w.InputRequests = ir
}
return json.Marshal(w)
}
func (x *ReadResourceResult) UnmarshalJSON(data []byte) error {
type res ReadResourceResult
var wire struct {
res
ResultType resultType `json:"resultType"`
}
if err := internaljson.Unmarshal(data, &wire); err != nil {
return err
}
wire.res.resultType = wire.ResultType
*x = ReadResourceResult(wire.res)
return nil
}
// A known resource that the server is capable of reading.
type Resource struct {
// See [specification/2025-06-18/basic/index#general-fields] for notes on _meta
@ -1520,8 +1897,9 @@ type ElicitParams struct {
ElicitationID string `json:"elicitationId,omitempty"`
}
func (x *ElicitParams) isParams() {}
func (x *ElicitParams) isNil() bool { return x == nil }
func (x *ElicitParams) isParams() {}
func (x *ElicitParams) isInputRequest() {}
func (x *ElicitParams) isNil() bool { return x == nil }
func (x *ElicitParams) GetProgressToken() any { return getProgressToken(x) }
func (x *ElicitParams) SetProgressToken(t any) { setProgressToken(x, t) }
@ -1541,7 +1919,8 @@ type ElicitResult struct {
Content map[string]any `json:"content,omitempty"`
}
func (*ElicitResult) isResult() {}
func (*ElicitResult) isResult() {}
func (*ElicitResult) isInputResponse() {}
// ElicitationCompleteParams is sent from the server to the client, informing it that an out-of-band elicitation interaction has completed.
type ElicitationCompleteParams struct {

View file

@ -1151,6 +1151,71 @@ func TestToWithTools_Conversion(t *testing.T) {
}
}
func TestInputRequestMapJSON(t *testing.T) {
t.Run("nil is omitted from JSON", func(t *testing.T) {
result := CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}
data, err := json.Marshal(&result)
if err != nil {
t.Fatal(err)
}
var raw map[string]json.RawMessage
if err := json.Unmarshal(data, &raw); err != nil {
t.Fatal(err)
}
if _, ok := raw["inputRequests"]; ok {
t.Errorf("nil InputRequests should be omitted, got %s", raw["inputRequests"])
}
})
t.Run("non-nil empty round-trips", func(t *testing.T) {
result := CallToolResult{
Content: []Content{&TextContent{Text: "ok"}},
InputRequests: InputRequestMap{},
}
data, err := json.Marshal(&result)
if err != nil {
t.Fatal(err)
}
var raw map[string]json.RawMessage
if err := json.Unmarshal(data, &raw); err != nil {
t.Fatal(err)
}
if string(raw["inputRequests"]) != "{}" {
t.Errorf("empty InputRequests should marshal to {}, got %s", raw["inputRequests"])
}
var got CallToolResult
if err := json.Unmarshal(data, &got); err != nil {
t.Fatal(err)
}
if got.InputRequests == nil {
t.Error("empty InputRequests should round-trip as non-nil")
}
})
t.Run("populated round-trips", func(t *testing.T) {
result := CallToolResult{
Content: []Content{&TextContent{Text: "ok"}},
InputRequests: InputRequestMap{
"r1": &ElicitParams{Message: "confirm?"},
},
}
data, err := json.Marshal(&result)
if err != nil {
t.Fatal(err)
}
var got CallToolResult
if err := json.Unmarshal(data, &got); err != nil {
t.Fatal(err)
}
if got.InputRequests == nil {
t.Fatal("InputRequests should not be nil after round-trip")
}
if _, ok := got.InputRequests["r1"]; !ok {
t.Error("expected key r1 in InputRequests")
}
})
}
func TestContentUnmarshal(t *testing.T) {
// Verify that types with a Content field round-trip properly.
roundtrip := func(in, out any) {

View file

@ -187,7 +187,7 @@ func NewServer(impl *Implementation, options *ServerOptions) *Server {
opts.Logger = ensureLogger(nil)
}
return &Server{
s := &Server{
impl: impl,
opts: opts,
prompts: newFeatureSet(func(p *serverPrompt) string { return p.prompt.Name }),
@ -199,6 +199,8 @@ func NewServer(impl *Implementation, options *ServerOptions) *Server {
resourceSubscriptions: make(map[string]map[*ServerSession]bool),
pendingNotifications: make(map[string]*time.Timer),
}
s.AddReceivingMiddleware(serverMultiRoundTripMiddleware())
return s
}
// AddPrompt adds a [Prompt] to the server, or replaces one with the same name.
@ -370,8 +372,12 @@ func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out], cache *SchemaCa
}
// Marshal the output and put the RawMessage in the StructuredContent field.
// Skip when the handler returned input requests (multi round-trip): content and
// inputRequests are mutually exclusive on the wire.
var outval any = out
if elemZero != nil {
if res.InputRequests != nil {
outval = nil
} else if elemZero != nil {
// Avoid typed nil, which will serialize as JSON null.
// Instead, use the zero value of the unpointered type.
var z Out
@ -742,7 +748,13 @@ func (s *Server) getPrompt(ctx context.Context, req *GetPromptRequest) (*GetProm
Message: fmt.Sprintf("unknown prompt %q", req.Params.Name),
}
}
return prompt.handler(ctx, req)
res, err := prompt.handler(ctx, req)
if err == nil && res != nil {
if err := handleMultiRoundTripResult(req.Session, s.opts.Logger, res); err != nil {
return nil, err
}
}
return res, err
}
func (s *Server) listTools(_ context.Context, req *ListToolsRequest) (*ListToolsResult, error) {
@ -775,10 +787,15 @@ func (s *Server) callTool(ctx context.Context, req *CallToolRequest) (*CallToolR
}
}
res, err := st.handler(ctx, req)
if err == nil && res != nil && res.Content == nil {
res2 := *res
res2.Content = []Content{} // avoid "null"
res = &res2
if err == nil && res != nil {
if err := handleMultiRoundTripResult(req.Session, s.opts.Logger, res); err != nil {
return nil, err
}
if res.Content == nil && res.resultType != resultTypeInputRequired {
res2 := *res
res2.Content = []Content{} // avoid "null"
res = &res2
}
}
return res, err
}
@ -826,6 +843,12 @@ func (s *Server) readResource(ctx context.Context, req *ReadResourceRequest) (*R
if err != nil {
return nil, err
}
if err := handleMultiRoundTripResult(req.Session, s.opts.Logger, res); err != nil {
return nil, err
}
if res.resultType == resultTypeInputRequired {
return res, nil
}
if res == nil || res.Contents == nil {
return nil, fmt.Errorf("reading resource %s: read handler returned nil information", uri)
}

View file

@ -769,6 +769,12 @@ func req(id int64, method string, params any) *jsonrpc.Request {
return r
}
func completeCallToolResult() *CallToolResult {
r := &CallToolResult{Content: []Content{}}
r.resultType = resultTypeComplete
return r
}
func resp(id int64, result any, err error) *jsonrpc.Response {
return &jsonrpc.Response{
ID: jsonrpc2.Int64ID(id),
@ -1947,7 +1953,7 @@ func TestStreamableMcpHeaderValidation(t *testing.T) {
},
messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Meta: testMeta, Name: "my-tool"})},
wantStatusCode: http.StatusOK,
wantMessages: []jsonrpc.Message{resp(2, &CallToolResult{Content: []Content{}}, nil)},
wantMessages: []jsonrpc.Message{resp(2, completeCallToolResult(), nil)},
},
{
method: "POST",
@ -1991,7 +1997,7 @@ func TestStreamableMcpHeaderValidation(t *testing.T) {
},
messages: []jsonrpc.Message{req(6, "tools/call", &CallToolParams{Meta: testMeta, Name: "my-tool"})},
wantStatusCode: http.StatusOK,
wantMessages: []jsonrpc.Message{resp(6, &CallToolResult{Content: []Content{}}, nil)},
wantMessages: []jsonrpc.Message{resp(6, completeCallToolResult(), nil)},
},
{
method: "POST",
@ -2007,7 +2013,7 @@ func TestStreamableMcpHeaderValidation(t *testing.T) {
Arguments: map[string]any{"region": "us-west1", "query": "SELECT 1"},
})},
wantStatusCode: http.StatusOK,
wantMessages: []jsonrpc.Message{resp(7, &CallToolResult{Content: []Content{}}, nil)},
wantMessages: []jsonrpc.Message{resp(7, completeCallToolResult(), nil)},
},
{
method: "POST",