mcp: route multi round-trip input requests through client receiving middleware

mcp/mrtr.go

@@ -317,16 +317,30 @@ func fulfillInputRequests(ctx context.Context, cs *ClientSession, requests Input
 }
 
 func fulfillInputRequest(ctx context.Context, cs *ClientSession, ir InputRequest) (InputResponse, error) {
+	var (
+		method string
+		params Params
+	)
 	switch p := ir.(type) {
 	case *ElicitParams:
-		return cs.client.elicit(ctx, newClientRequest(cs, p))
+		method, params = methodElicit, p
 	case *CreateMessageParams:
-		return cs.client.createMessage(ctx, &CreateMessageWithToolsRequest{Session: cs, Params: createMessageParamsToWithTools(p)})
+		method, params = methodCreateMessage, createMessageParamsToWithTools(p)
 	case *CreateMessageWithToolsParams:
-		return cs.client.createMessage(ctx, &CreateMessageWithToolsRequest{Session: cs, Params: p})
+		method, params = methodCreateMessage, p
 	case *ListRootsParams:
-		return cs.client.listRoots(ctx, newClientRequest(cs, p))
+		method, params = methodListRoots, p
 	default:
 		return nil, fmt.Errorf("unknown input request type: %T", ir)
 	}
+	req := cs.receivingMethodInfos()[method].newRequest(cs, params, nil, jsonrpc.ID{})
+	res, err := cs.receivingMethodHandler()(ctx, method, req)
+	if err != nil {
+		return nil, err
+	}
+	resp, ok := res.(InputResponse)
+	if !ok {
+		return nil, fmt.Errorf("%s: result of type %T is not an input response", method, res)
+	}
+	return resp, nil
 }

mcp/mrtr_test.go

@@ -12,6 +12,7 @@ import (
 	"fmt"
 	"slices"
 	"strings"
+	"sync"
 	"sync/atomic"
 	"testing"
 
@@ -200,6 +201,81 @@ func TestMultiRoundTrip_AutoRetry(t *testing.T) {
 	}
 }
 
+// inputRequestServer returns a server with a tool "ask" that asks the client
+// for an elicitation, a sampling message and its roots before it completes.
+func inputRequestServer() *Server {
+	srv := NewServer(testImpl, nil)
+	AddTool(srv, &Tool{Name: "ask"}, func(ctx context.Context, req *CallToolRequest, _ struct{}) (*CallToolResult, any, error) {
+		if len(req.Params.InputResponses) == 0 {
+			return &CallToolResult{InputRequests: InputRequestMap{
+				"confirm": &ElicitParams{Message: "OK?"},
+				"draft": &CreateMessageParams{
+					Messages:  []*SamplingMessage{{Role: "user", Content: &TextContent{Text: "write"}}},
+					MaxTokens: 50,
+				},
+				"roots": &ListRootsParams{},
+			}}, nil, nil
+		}
+		return &CallToolResult{Content: []Content{&TextContent{Text: "done"}}}, nil, nil
+	})
+	return srv
+}
+
+var inputRequestClientOptions = &ClientOptions{
+	ElicitationHandler: func(context.Context, *ElicitRequest) (*ElicitResult, error) {
+		return &ElicitResult{Action: "accept"}, nil
+	},
+	CreateMessageHandler: func(context.Context, *CreateMessageRequest) (*CreateMessageResult, error) {
+		return &CreateMessageResult{Model: "test-model", Role: "assistant", Content: &TextContent{Text: "response"}}, nil
+	},
+}
+
+// TestMultiRoundTrip_InputRequestsReachReceivingMiddleware checks that client
+// receiving middleware sees the requests a server makes for input in the same
+// way whether they arrive as server-to-client requests (2025-11-25) or as
+// input requests in a multi round-trip result (2026-07-28).
+func TestMultiRoundTrip_InputRequestsReachReceivingMiddleware(t *testing.T) {
+	for _, tc := range []struct {
+		name    string
+		connect func(*testing.T, *Server, *ClientOptions) *ClientSession
+	}{
+		{"2025-11-25", mustConnectOldProtocol},
+		{"2026-07-28", mustConnect},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			cs := tc.connect(t, inputRequestServer(), inputRequestClientOptions)
+			cs.client.AddRoots(&Root{URI: "file:///workspace", Name: "workspace"})
+			var (
+				mu   sync.Mutex
+				seen = map[string]string{} // method -> params type
+			)
+			cs.client.AddReceivingMiddleware(func(next MethodHandler) MethodHandler {
+				return func(ctx context.Context, method string, req Request) (Result, error) {
+					switch method {
+					case methodElicit, methodCreateMessage, methodListRoots:
+						mu.Lock()
+						seen[method] = fmt.Sprintf("%T", req.GetParams())
+						mu.Unlock()
+					}
+					return next(ctx, method, req)
+				}
+			})
+
+			if _, err := cs.CallTool(context.Background(), &CallToolParams{Name: "ask"}); err != nil {
+				t.Fatalf("CallTool() error = %v", err)
+			}
+			want := map[string]string{
+				methodElicit:        "*mcp.ElicitParams",
+				methodCreateMessage: "*mcp.CreateMessageWithToolsParams",
+				methodListRoots:     "*mcp.ListRootsParams",
+			}
+			if diff := cmp.Diff(want, seen); diff != "" {
+				t.Errorf("receiving middleware saw (-want +got):\n%s", diff)
+			}
+		})
+	}
+}
+
 // TestMultiRoundTrip_AutoRetryDoesNotMutateCallerParams verifies that
 // the client middleware carries inputResponses and requestState on a
 // copy of the caller's params: after CallTool returns, the caller's