From 5b56d08db4e1a60ebb54b0af225afae4c7eec56d Mon Sep 17 00:00:00 2001 From: basilalshukaili Date: Tue, 2 Jun 2026 21:02:19 +0000 Subject: [PATCH] mcptest: add SetSamplingHandler and SetElicitationHandler The mcptest.Server helper previously had no way to test tools that call server.RequestSampling or server.RequestElicitation, because the underlying client was created without sampling/elicitation options and the stdio transport request-handler was never wired up. This change: - Adds SetSamplingHandler(client.SamplingHandler) to mcptest.Server - Adds SetElicitationHandler(client.ElicitationHandler) to mcptest.Server - Passes the handlers as client.ClientOption to client.NewClient - Uses client.Start() instead of transport.Start() so the bidirectional request handler (needed for server-to-client sampling/elicitation calls over stdio) is registered before the Initialize handshake - Automatically calls mcpServer.EnableSampling() when a sampling handler is set, removing boilerplate from callers Two tests are added in mcptest_sampling_elicitation_test.go covering both paths end-to-end through the real stdio transport. --- mcptest/mcptest.go | 48 +++- mcptest/mcptest_sampling_elicitation_test.go | 220 +++++++++++++++++++ 2 files changed, 265 insertions(+), 3 deletions(-) create mode 100644 mcptest/mcptest_sampling_elicitation_test.go diff --git a/mcptest/mcptest.go b/mcptest/mcptest.go index 4ee27eb21..b0fdbb27e 100644 --- a/mcptest/mcptest.go +++ b/mcptest/mcptest.go @@ -27,6 +27,9 @@ type Server struct { serverOpts []server.ServerOption clientInfo mcp.Implementation + samplingHandler client.SamplingHandler + elicitationHandler client.ElicitationHandler + cancel func() serverReader *io.PipeReader @@ -135,6 +138,22 @@ func (s *Server) SetClientInfo(info mcp.Implementation) { s.clientInfo = info } +// SetSamplingHandler registers a handler that responds to sampling requests +// (server.RequestSampling) made by tools under test. Must be called before Start(). +// The test client will advertise the sampling capability during initialization so +// that the server is allowed to issue sampling requests. +func (s *Server) SetSamplingHandler(h client.SamplingHandler) { + s.samplingHandler = h +} + +// SetElicitationHandler registers a handler that responds to elicitation requests +// (server.RequestElicitation) made by tools under test. Must be called before Start(). +// The test client will advertise the elicitation capability during initialization so +// that the server is allowed to issue elicitation requests. +func (s *Server) SetElicitationHandler(h client.ElicitationHandler) { + s.elicitationHandler = h +} + // Start starts the server in a goroutine. Make sure to defer Close() after Start(). // When using NewServer(), the returned server is already started. func (s *Server) Start(ctx context.Context) error { @@ -142,6 +161,10 @@ func (s *Server) Start(ctx context.Context) error { ctx, s.cancel = context.WithCancel(ctx) + // Capture handler state for the goroutine. Start must be called after any + // SetSamplingHandler / SetElicitationHandler calls, so there is no data race. + samplingHandler := s.samplingHandler + // Start the MCP server in a goroutine go func() { defer s.wg.Done() @@ -153,6 +176,13 @@ func (s *Server) Start(ctx context.Context) error { mcpServer.AddResources(s.resources...) mcpServer.AddResourceTemplates(s.resourceTemplates...) + // Automatically enable sampling on the server when the test supplies a + // sampling handler, so tools can call server.RequestSampling without + // any additional server-side setup. + if samplingHandler != nil { + mcpServer.EnableSampling() + } + logger := log.New(&s.logBuffer, "", 0) stdioServer := server.NewStdioServer(mcpServer) @@ -164,11 +194,23 @@ func (s *Server) Start(ctx context.Context) error { }() s.transport = transport.NewIO(s.clientReader, s.clientWriter, io.NopCloser(&s.logBuffer)) - if err := s.transport.Start(ctx); err != nil { - return fmt.Errorf("transport.Start(): %w", err) + + // Build client options from registered handlers. + var clientOpts []client.ClientOption + if s.samplingHandler != nil { + clientOpts = append(clientOpts, client.WithSamplingHandler(s.samplingHandler)) } + if s.elicitationHandler != nil { + clientOpts = append(clientOpts, client.WithElicitationHandler(s.elicitationHandler)) + } + + s.client = client.NewClient(s.transport, clientOpts...) - s.client = client.NewClient(s.transport) + // Use client.Start instead of transport.Start so that bidirectional request + // handlers (sampling, elicitation) are registered before the Initialize handshake. + if err := s.client.Start(ctx); err != nil { + return fmt.Errorf("client.Start(): %w", err) + } var initReq mcp.InitializeRequest initReq.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION diff --git a/mcptest/mcptest_sampling_elicitation_test.go b/mcptest/mcptest_sampling_elicitation_test.go new file mode 100644 index 000000000..b1161667f --- /dev/null +++ b/mcptest/mcptest_sampling_elicitation_test.go @@ -0,0 +1,220 @@ +package mcptest_test + +import ( + "context" + "testing" + + "github.com/mark3labs/mcp-go/mcp" + "github.com/mark3labs/mcp-go/mcptest" + "github.com/mark3labs/mcp-go/server" +) + +// TestServerWithSamplingHandler verifies that a tool which calls server.RequestSampling +// can be tested end-to-end using mcptest, without any external LLM. +func TestServerWithSamplingHandler(t *testing.T) { + ctx := t.Context() + + const wantReply = "42" + + // A sampling handler that returns a fixed answer regardless of the question. + samplingHandler := &fixedSamplingHandler{reply: wantReply} + + srv := mcptest.NewUnstartedServer(t) + defer srv.Close() + + // A tool that delegates its response to the LLM via sampling. + srv.AddTool( + mcp.NewTool("ask_llm", + mcp.WithDescription("Ask the LLM a question and return its response."), + mcp.WithString("question", + mcp.Required(), + mcp.Description("The question to ask."), + ), + ), + func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + question, err := req.RequireString("question") + if err != nil { + return mcp.NewToolResultError(err.Error()), nil + } + + samplingReq := mcp.CreateMessageRequest{ + CreateMessageParams: mcp.CreateMessageParams{ + Messages: []mcp.SamplingMessage{ + { + Role: mcp.RoleUser, + Content: mcp.NewTextContent(question), + }, + }, + MaxTokens: 64, + }, + } + + mcpServer := server.ServerFromContext(ctx) + result, err := mcpServer.RequestSampling(ctx, samplingReq) + if err != nil { + return mcp.NewToolResultError("sampling failed: " + err.Error()), nil + } + + text, ok := result.Content.(mcp.TextContent) + if !ok { + return mcp.NewToolResultError("unexpected content type from sampling"), nil + } + return mcp.NewToolResultText(text.Text), nil + }, + ) + + srv.SetSamplingHandler(samplingHandler) + + if err := srv.Start(ctx); err != nil { + t.Fatal("Start:", err) + } + + var callReq mcp.CallToolRequest + callReq.Params.Name = "ask_llm" + callReq.Params.Arguments = map[string]any{"question": "What is 6*7?"} + + result, err := srv.Client().CallTool(ctx, callReq) + if err != nil { + t.Fatal("CallTool:", err) + } + + got, err := resultToString(result) + if err != nil { + t.Fatal(err) + } + + if got != wantReply { + t.Errorf("got %q, want %q", got, wantReply) + } + if samplingHandler.callCount != 1 { + t.Errorf("expected sampling handler called once, got %d", samplingHandler.callCount) + } +} + +// TestServerWithElicitationHandler verifies that a tool which calls +// server.RequestElicitation can be tested end-to-end using mcptest. +func TestServerWithElicitationHandler(t *testing.T) { + ctx := t.Context() + + // An elicitation handler that always accepts with a canned response. + elicitationHandler := &fixedElicitationHandler{ + response: map[string]any{ + "confirmed": true, + }, + } + + srv := mcptest.NewUnstartedServer(t) + defer srv.Close() + + // The server must declare elicitation capability so the spec allows it to + // issue elicitation/create requests. + srv.AddServerOptions(server.WithElicitation()) + + // A tool that asks the user to confirm an action before proceeding. + srv.AddTool( + mcp.NewTool("confirm_action", + mcp.WithDescription("Ask the user to confirm before proceeding."), + mcp.WithString("action", + mcp.Required(), + mcp.Description("Description of the action to confirm."), + ), + ), + func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + action, err := req.RequireString("action") + if err != nil { + return mcp.NewToolResultError(err.Error()), nil + } + + elicitReq := mcp.ElicitationRequest{ + Params: mcp.ElicitationParams{ + Message: "Please confirm: " + action, + RequestedSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "confirmed": map[string]any{ + "type": "boolean", + }, + }, + "required": []string{"confirmed"}, + }, + }, + } + + mcpServer := server.ServerFromContext(ctx) + elicitResult, err := mcpServer.RequestElicitation(ctx, elicitReq) + if err != nil { + return mcp.NewToolResultError("elicitation failed: " + err.Error()), nil + } + + switch elicitResult.Action { + case mcp.ElicitationResponseActionAccept: + return mcp.NewToolResultText("confirmed"), nil + case mcp.ElicitationResponseActionDecline: + return mcp.NewToolResultText("declined"), nil + default: + return mcp.NewToolResultText("cancelled"), nil + } + }, + ) + + srv.SetElicitationHandler(elicitationHandler) + + if err := srv.Start(ctx); err != nil { + t.Fatal("Start:", err) + } + + var callReq mcp.CallToolRequest + callReq.Params.Name = "confirm_action" + callReq.Params.Arguments = map[string]any{"action": "delete all records"} + + result, err := srv.Client().CallTool(ctx, callReq) + if err != nil { + t.Fatal("CallTool:", err) + } + + got, err := resultToString(result) + if err != nil { + t.Fatal(err) + } + + if got != "confirmed" { + t.Errorf("got %q, want %q", got, "confirmed") + } + if elicitationHandler.callCount != 1 { + t.Errorf("expected elicitation handler called once, got %d", elicitationHandler.callCount) + } +} + +// fixedSamplingHandler is a test double that always returns a preset text reply. +type fixedSamplingHandler struct { + reply string + callCount int +} + +func (h *fixedSamplingHandler) CreateMessage(_ context.Context, _ mcp.CreateMessageRequest) (*mcp.CreateMessageResult, error) { + h.callCount++ + return &mcp.CreateMessageResult{ + SamplingMessage: mcp.SamplingMessage{ + Role: mcp.RoleAssistant, + Content: mcp.NewTextContent(h.reply), + }, + Model: "test-model", + StopReason: "endTurn", + }, nil +} + +// fixedElicitationHandler is a test double that always accepts with a preset content map. +type fixedElicitationHandler struct { + response map[string]any + callCount int +} + +func (h *fixedElicitationHandler) Elicit(_ context.Context, _ mcp.ElicitationRequest) (*mcp.ElicitationResult, error) { + h.callCount++ + return &mcp.ElicitationResult{ + ElicitationResponse: mcp.ElicitationResponse{ + Action: mcp.ElicitationResponseActionAccept, + Content: h.response, + }, + }, nil +}