Merge pull request #458 from koic/feature_sampling_client

Support client-side sampling via `MCP::Client#on_sampling`

Author: koic

README.md

@@ -2346,6 +2346,44 @@ since servers deliver requests that are not tied to a client request on that str
 a JSON-RPC `-32601` (method not found) error. To handle methods other than `elicitation/create`, register directly on the transport with
 `http_transport.on_server_request("method/name") { |params| ... }`.
 
+#### Server-to-Client Requests (Sampling)
+
+Servers can also request an LLM completion from the client with [`sampling/createMessage`](https://modelcontextprotocol.io/specification/2025-11-25/client/sampling),
+letting a server leverage the client's model access without its own API keys.
+
+> MCP Sampling is deprecated as of protocol version `2026-07-28` (SEP-2577), while remaining fully supported under `2025-11-25`.
+> Register this handler to interoperate with servers that still send sampling requests during the deprecation window;
+> new servers should call LLM provider APIs directly.
+
+Register a handler and advertise the capability on `connect`:
+
+```ruby
+client.connect(capabilities: { sampling: {} })
+
+client.on_sampling do |params|
+  completion = my_llm.complete(params["messages"], max_tokens: params["maxTokens"])
+  {
+    role: "assistant",
+    content: { type: "text", text: completion.text },
+    model: completion.model,
+    stopReason: "endTurn",
+  }
+end
+```
+
+For trust and safety, the spec recommends a human in the loop able to review, edit, or reject the request and the generated response.
+To reject a request, raise `MCP::Client::ServerRequestError` with the spec's user-rejection code `-1`:
+
+```ruby
+client.on_sampling do |params|
+  raise MCP::Client::ServerRequestError.new("User rejected sampling request", code: -1) unless approved?(params)
+
+  generate_completion(params)
+end
+```
+
+Use `capabilities: { sampling: { tools: {} } }` to receive tool-enabled sampling requests. Like elicitation, this uses the same standalone GET SSE listening stream.
+
 #### HTTP Authorization
 
 By default, the HTTP transport layer provides no authentication to the server, but you can provide custom headers if you need authentication. For example, to use Bearer token authentication:

lib/mcp/client.rb

@@ -423,6 +423,44 @@ def on_elicitation(&handler)
       transport.on_server_request(Methods::ELICITATION_CREATE, &handler)
     end
 
+    # Registers a handler for `sampling/createMessage` requests the server sends while one of this client's requests is in flight.
+    # The handler receives the request `params` (`messages`, `maxTokens`, optionally `systemPrompt`, `modelPreferences`, `tools`,
+    # `toolChoice`, ...; string keys) and must return a `CreateMessageResult`-shaped Hash:
+    # `{ role: "assistant", content: { type: "text", text: "..." }, model: "...", stopReason: "..." }`.
+    #
+    # For trust and safety, the spec recommends a human in the loop able to review, edit, or reject the request and the generated response
+    # before it is returned to the server. To reject, raise `ServerRequestError` with the spec's user-rejection code `-1`.
+    #
+    # Requires a transport that supports server-to-client requests (e.g. `MCP::Client::HTTP`); pass `capabilities: { sampling: {} }` to
+    # `connect` (or `{ sampling: { tools: {} } }` to receive tool-enabled sampling requests) so the server knows it may send them.
+    #
+    # @example Forward the request to an LLM and return its completion
+    #
+    #   client.on_sampling do |params|
+    #     raise MCP::Client::ServerRequestError.new("User rejected sampling request", code: -1) unless approved?(params)
+    #
+    #     completion = my_llm.complete(params["messages"], max_tokens: params["maxTokens"])
+    #     {
+    #       role: "assistant",
+    #       content: { type: "text", text: completion.text },
+    #       model: completion.model,
+    #       stopReason: "endTurn",
+    #     }
+    #   end
+    #
+    # https://modelcontextprotocol.io/specification/2025-11-25/client/sampling
+    #
+    # @deprecated MCP Sampling (`sampling/createMessage`) is deprecated as of MCP protocol version 2026-07-28 (SEP-2577).
+    #   Register this handler only to interoperate with servers that still send sampling requests during the deprecation window;
+    #   new servers should call LLM provider APIs directly.
+    def on_sampling(&handler)
+      unless transport.respond_to?(:on_server_request)
+        raise ArgumentError, "The transport does not support server-to-client requests"
+      end
+
+      transport.on_server_request(Methods::SAMPLING_CREATE_MESSAGE, &handler)
+    end
+
     # Sends a `ping` request to the server to verify the connection is alive.
     # Per the MCP spec, the server responds with an empty result.
     #

test/mcp/client/http_test.rb

@@ -1020,6 +1020,67 @@ def test_send_request_dispatches_server_request_to_registered_handler
         assert_equal("Please accept with defaults", received_params["message"])
       end
 
+      def test_send_request_dispatches_sampling_request_to_registered_handler
+        request = {
+          jsonrpc: "2.0",
+          id: "test_id",
+          method: "tools/call",
+          params: { name: "ask_llm", arguments: {} },
+        }
+
+        sampling_request = {
+          jsonrpc: "2.0",
+          id: 0,
+          method: "sampling/createMessage",
+          params: {
+            messages: [{ role: "user", content: { type: "text", text: "Hi" } }],
+            maxTokens: 100,
+          },
+        }
+        tool_result = { jsonrpc: "2.0", id: "test_id", result: { content: [] } }
+        sse_body = "event: message\ndata: #{sampling_request.to_json}\n\n" \
+          "event: message\ndata: #{tool_result.to_json}\n\n"
+
+        stub_request(:post, url).with(
+          body: request.to_json,
+        ).to_return(
+          status: 200,
+          headers: { "Content-Type" => "text/event-stream" },
+          body: sse_body,
+        )
+
+        expected_response = {
+          jsonrpc: "2.0",
+          id: 0,
+          result: {
+            role: "assistant",
+            content: { type: "text", text: "Hello there" },
+            model: "test-model",
+            stopReason: "endTurn",
+          },
+        }
+        response_stub = stub_request(:post, url).with(
+          body: expected_response.to_json,
+        ).to_return(status: 202, body: "")
+
+        received_params = nil
+        client.on_server_request("sampling/createMessage") do |params|
+          received_params = params
+          {
+            role: "assistant",
+            content: { type: "text", text: "Hello there" },
+            model: "test-model",
+            stopReason: "endTurn",
+          }
+        end
+
+        response = client.send_request(request: request)
+
+        assert_equal({ "content" => [] }, response["result"])
+        assert_requested(response_stub)
+        assert_equal(100, received_params["maxTokens"])
+      end
+
       def test_send_request_answers_unregistered_server_request_with_method_not_found
         request = {
           jsonrpc: "2.0",

test/mcp/client_test.rb

@@ -66,6 +66,25 @@ def test_on_elicitation_raises_when_transport_does_not_support_server_requests
       assert_includes(error.message, "does not support server-to-client requests")
     end
 
+    def test_on_sampling_registers_handler_on_transport
+      transport = mock
+      handler = proc { { role: "assistant", content: { type: "text", text: "hi" } } }
+      transport.expects(:on_server_request).with("sampling/createMessage")
+
+      Client.new(transport: transport).on_sampling(&handler)
+    end
+
+    def test_on_sampling_raises_when_transport_does_not_support_server_requests
+      transport = mock
+      transport.stubs(:respond_to?).with(:on_server_request).returns(false)
+
+      error = assert_raises(ArgumentError) do
+        Client.new(transport: transport).on_sampling { { role: "assistant", content: { type: "text", text: "hi" } } }
+      end
+
+      assert_includes(error.message, "does not support server-to-client requests")
+    end
+
     def test_connected_delegates_to_transport_when_supported
       transport = mock
       transport.expects(:connected?).returns(true)