Merge pull request #102 from patvice/sample-correct-stop-response

Support correct stop response coming from ruby llm message

Author: patvice

lib/ruby_llm/mcp/native/messages/responses.rb

@@ -37,16 +37,20 @@ def roots_list(id:, roots_paths:)
           end
 
           def sampling_create_message(id:, message:, model:)
+            stop_reason = if message.respond_to?(:stop_reason) && message.stop_reason
+                            snake_to_camel(message.stop_reason)
+                          else
+                            "endTurn"
+                          end
+
             {
               jsonrpc: JSONRPC_VERSION,
               id: id,
               result: {
                 role: message.role,
                 content: format_content(message.content),
                 model: model,
-                # TODO: We are going to assume it was a endTurn
-                # Look into getting RubyLLM to expose stopReason in message response
-                stopReason: "endTurn"
+                stopReason: stop_reason
               }
             }
           end
@@ -89,6 +93,12 @@ def format_content(content)
             end
           end
           private_class_method :format_content
+
+          def snake_to_camel(str)
+            parts = str.split("_")
+            parts.first + parts[1..].map(&:capitalize).join
+          end
+          private_class_method :snake_to_camel
         end
       end
     end

spec/ruby_llm/mcp/cancellation_integration_spec.rb

@@ -56,9 +56,11 @@
 
         client.start
 
+        # Wait for the tool to become available
+        tool = wait_for_tool(client, "sample_with_cancellation")
+
         # Call the tool in a thread so we can cancel it
         tool_thread = Thread.new do
-          tool = client.tool("sample_with_cancellation")
           tool.execute
         end
 
@@ -109,10 +111,12 @@
 
         client.start
 
+        # Wait for the tool to become available
+        tool = wait_for_tool(client, "sample_with_cancellation")
+
         # Start multiple sampling requests
         threads = 3.times.map do
           Thread.new do
-            tool = client.tool("sample_with_cancellation")
             tool.execute
           end
         end

spec/ruby_llm/mcp/native/messages_spec.rb

@@ -504,6 +504,165 @@
         expect(body[:result][:content]).to be_nil
       end
     end
+
+    describe ".sampling_create_message" do
+      # rubocop:disable RSpec/VerifiedDoubles
+      let(:mock_content) do
+        double("Content", text: "Hello, world!")
+      end
+
+      context "when message has no stop_reason" do
+        let(:mock_message) do
+          double(
+            "Message",
+            role: "assistant",
+            content: mock_content
+          ).tap do |msg|
+            allow(msg).to receive(:respond_to?).with(:stop_reason).and_return(false)
+          end
+        end
+
+        let(:body) do
+          described_class::Responses.sampling_create_message(
+            id: "req-123",
+            message: mock_message,
+            model: "gpt-4"
+          )
+        end
+
+        it "creates a valid sampling response with default stopReason" do
+          expect(body[:jsonrpc]).to eq("2.0")
+          expect(body[:id]).to eq("req-123")
+          expect(body[:result][:role]).to eq("assistant")
+          expect(body[:result][:model]).to eq("gpt-4")
+          expect(body[:result][:stopReason]).to eq("endTurn")
+        end
+
+        it "validates against CreateMessageResult schema" do
+          body_json = JSON.parse(body.to_json)
+          response_json = {
+            "jsonrpc" => "2.0",
+            "id" => body_json["id"],
+            "result" => body_json["result"]
+          }
+          errors = schemer.validate(response_json).to_a
+          expect(errors).to be_empty, "Schema validation errors: #{errors.map(&:to_h)}"
+        end
+      end
+
+      context "when message has stop_reason nil" do
+        let(:mock_message) do
+          double(
+            "Message",
+            role: "assistant",
+            content: mock_content,
+            stop_reason: nil
+          ).tap do |msg|
+            allow(msg).to receive(:respond_to?).with(:stop_reason).and_return(true)
+          end
+        end
+
+        let(:body) do
+          described_class::Responses.sampling_create_message(
+            id: "req-123",
+            message: mock_message,
+            model: "gpt-4"
+          )
+        end
+
+        it "uses default stopReason when nil" do
+          expect(body[:result][:stopReason]).to eq("endTurn")
+        end
+      end
+
+      context "when message has stop_reason values" do
+        shared_examples "converts stop_reason correctly" do |snake_case_value, camel_case_value|
+          let(:mock_message) do
+            double(
+              "Message",
+              role: "assistant",
+              content: mock_content,
+              stop_reason: snake_case_value
+            ).tap do |msg|
+              allow(msg).to receive(:respond_to?).with(:stop_reason).and_return(true)
+            end
+          end
+
+          let(:body) do
+            described_class::Responses.sampling_create_message(
+              id: "req-123",
+              message: mock_message,
+              model: "gpt-4"
+            )
+          end
+
+          it "converts #{snake_case_value} to #{camel_case_value}" do
+            expect(body[:result][:stopReason]).to eq(camel_case_value)
+          end
+
+          it "validates against CreateMessageResult schema" do
+            body_json = JSON.parse(body.to_json)
+            response_json = {
+              "jsonrpc" => "2.0",
+              "id" => body_json["id"],
+              "result" => body_json["result"]
+            }
+            errors = schemer.validate(response_json).to_a
+            expect(errors).to be_empty, "Schema validation errors: #{errors.map(&:to_h)}"
+          end
+        end
+
+        context "with stop_reason: end_turn" do
+          it_behaves_like "converts stop_reason correctly", "end_turn", "endTurn"
+        end
+
+        context "with stop_reason: max_tokens" do
+          it_behaves_like "converts stop_reason correctly", "max_tokens", "maxTokens"
+        end
+
+        context "with stop_reason: stop_sequence" do
+          it_behaves_like "converts stop_reason correctly", "stop_sequence", "stopSequence"
+        end
+
+        context "with stop_reason: tool_use" do
+          it_behaves_like "converts stop_reason correctly", "tool_use", "toolUse"
+        end
+
+        context "with stop_reason: pause_turn" do
+          it_behaves_like "converts stop_reason correctly", "pause_turn", "pauseTurn"
+        end
+
+        context "with stop_reason: refusal" do
+          it_behaves_like "converts stop_reason correctly", "refusal", "refusal"
+        end
+      end
+
+      context "with custom stop_reason values" do
+        let(:mock_message) do
+          double(
+            "Message",
+            role: "assistant",
+            content: mock_content,
+            stop_reason: "custom_stop_reason"
+          ).tap do |msg|
+            allow(msg).to receive(:respond_to?).with(:stop_reason).and_return(true)
+          end
+        end
+
+        let(:body) do
+          described_class::Responses.sampling_create_message(
+            id: "req-123",
+            message: mock_message,
+            model: "gpt-4"
+          )
+        end
+
+        it "converts unknown snake_case values to camelCase" do
+          expect(body[:result][:stopReason]).to eq("customStopReason")
+        end
+      end
+      # rubocop:enable RSpec/VerifiedDoubles
+    end
   end
 
   describe "UUID generation consistency" do

spec/spec_helper.rb

@@ -21,6 +21,7 @@
 require "mcp" if RUBY_VERSION >= "3.2.0"
 
 require_relative "support/client_runner"
+require_relative "support/client_sync_helpers"
 require_relative "support/test_server_manager"
 require_relative "support/mcp_test_configuration"
 require_relative "support/simple_multiply_tool"

spec/support/client_sync_helpers.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+# Client Synchronization Helpers
+#
+# This module provides helper methods for waiting on client state changes
+# and tool availability in tests.
+module ClientSyncHelpers
+  # Wait for a specific tool to become available on a client
+  #
+  # @param client [RubyLLM::MCP::Client] The MCP client to check
+  # @param tool_name [String] The name of the tool to wait for
+  # @param max_wait_time [Integer, Float] Maximum time to wait in seconds (default: 5)
+  # @return [RubyLLM::MCP::Tool] The tool instance when found
+  # @raise [RuntimeError] If the tool is not found within the timeout period
+  def wait_for_tool(client, tool_name, max_wait_time: 5)
+    start_time = Time.now
+    tool = nil
+
+    loop do
+      tool = client.tool(tool_name)
+      break if tool
+
+      elapsed = Time.now - start_time
+      if elapsed > max_wait_time
+        available_tools = begin
+          client.tools.map(&:name).join(", ")
+        rescue StandardError
+          "Unable to fetch tools"
+        end
+        raise "Timeout waiting for tool '#{tool_name}' after #{elapsed.round(2)}s. \
+               Available tools: #{available_tools}. Client alive: #{client.alive?}"
+      end
+
+      sleep 0.1
+    end
+
+    tool
+  end
+end
+
+RSpec.configure do |config|
+  config.include ClientSyncHelpers
+end