ruby_llm-mcp 0.5.0 → 0.6.0

data/lib/ruby_llm/mcp/tool.rb

@@ -36,6 +36,10 @@ module RubyLLM
         @mcp_name = tool_response["name"]
         @description = tool_response["description"].to_s
         @parameters = create_parameters(tool_response["inputSchema"])
+
+        @input_schema = tool_response["inputSchema"]
+        @output_schema = tool_response["outputSchema"]
+
         @annotations = tool_response["annotations"] ? Annotation.new(tool_response["annotations"]) : nil
       end
 
@@ -59,6 +63,15 @@ module RubyLLM
           return { error: "Tool execution error: #{text_values}" }
         end
 
+        if result.value.key?("structuredContent") && !@output_schema.nil?
+          is_valid = JSON::Validator.validate(@output_schema, result.value["structuredContent"])
+          unless is_valid
+            return { error: "Structued outputs was not invalid: #{result.value['structuredContent']}" }
+          end
+
+          return text_values
+        end
+
         if text_values.empty?
           create_content_for_message(result.value.dig("content", 0))
         else
@@ -79,12 +92,12 @@ module RubyLLM
 
       private
 
-      def create_parameters(input_schema)
+      def create_parameters(schema)
         params = {}
-        return params if input_schema["properties"].nil?
+        return params if schema["properties"].nil?
 
-        input_schema["properties"].each_key do |key|
-          param_data = input_schema.dig("properties", key)
+        schema["properties"].each_key do |key|
+          param_data = schema.dig("properties", key)
 
           param = if param_data.key?("oneOf") || param_data.key?("anyOf") || param_data.key?("allOf")
                     process_union_parameter(key, param_data)
@@ -152,10 +165,25 @@ module RubyLLM
             "name" => name,
             "description" => description,
             "uri" => content.dig("resource", "uri"),
-            "content" => content["resource"]
+            "mimeType" => content.dig("resource", "mimeType"),
+            "content_response" => {
+              "text" => content.dig("resource", "text"),
+              "blob" => content.dig("resource", "blob")
+            }
+          }
+
+          resource = Resource.new(coordinator, resource_data)
+          resource.to_content
+        when "resource_link"
+          resource_data = {
+            "name" => content["name"],
+            "uri" => content["uri"],
+            "description" => content["description"],
+            "mimeType" => content["mimeType"]
           }
 
           resource = Resource.new(coordinator, resource_data)
+          @coordinator.register_resource(resource)
           resource.to_content
         end
       end

data/lib/ruby_llm/mcp/transports/sse.rb

@@ -10,15 +10,16 @@ module RubyLLM
   module MCP
     module Transports
       class SSE
-        include Timeout
+        include Support::Timeout
 
         attr_reader :headers, :id, :coordinator
 
-        def initialize(url:, coordinator:, request_timeout:, headers: {})
+        def initialize(url:, coordinator:, request_timeout:, version: :http2, headers: {})
           @event_url = url
           @messages_url = nil
           @coordinator = coordinator
           @request_timeout = request_timeout
+          @version = version
 
           uri = URI.parse(url)
           @root_url = "#{uri.scheme}://#{uri.host}"
@@ -44,7 +45,7 @@ module RubyLLM
           RubyLLM::MCP.logger.info "Initializing SSE transport to #{@event_url} with client ID #{@client_id}"
         end
 
-        def request(body, add_id: true, wait_for_response: true) # rubocop:disable Metrics/MethodLength
+        def request(body, add_id: true, wait_for_response: true)
           if add_id
             @id_mutex.synchronize { @id_counter += 1 }
             request_id = @id_counter
@@ -59,34 +60,20 @@ module RubyLLM
           end
 
           begin
-            http_client = HTTPClient.connection.with(timeout: { request_timeout: @request_timeout / 1000 },
-                                                     headers: @headers)
-            response = http_client.post(@messages_url, body: JSON.generate(body))
-
-            unless response.status == 200
-              @pending_mutex.synchronize { @pending_requests.delete(request_id.to_s) }
-              RubyLLM::MCP.logger.error "SSE request failed: #{response.status} - #{response.body}"
-              raise Errors::TransportError.new(
-                message: "Failed to request #{@messages_url}: #{response.status} - #{response.body}",
-                code: response.status
-              )
-            end
-          rescue StandardError => e
+            send_request(body, request_id)
+          rescue Errors::TransportError, Errors::TimeoutError => e
             @pending_mutex.synchronize { @pending_requests.delete(request_id.to_s) }
-            RubyLLM::MCP.logger.error "SSE request error (ID: #{request_id}): #{e.message}"
-            raise RubyLLM::MCP::Errors::TransportError.new(
-              message: e.message,
-              code: -1,
-              error: e
-            )
+            RubyLLM::MCP.logger.error "Request error (ID: #{request_id}): #{e.message}"
+            raise e
           end
+
           return unless wait_for_response
 
           begin
             with_timeout(@request_timeout / 1000, request_id: request_id) do
               response_queue.pop
             end
-          rescue RubyLLM::MCP::Errors::TimeoutError => e
+          rescue Errors::TimeoutError => e
             @pending_mutex.synchronize { @pending_requests.delete(request_id.to_s) }
             RubyLLM::MCP.logger.error "SSE request timeout (ID: #{request_id}) after #{@request_timeout / 1000} seconds"
             raise e
@@ -117,6 +104,23 @@ module RubyLLM
 
         private
 
+        def send_request(body, request_id)
+          http_client = Support::HTTPClient.connection.with(timeout: { request_timeout: @request_timeout / 1000 },
+                                                            headers: @headers)
+          response = http_client.post(@messages_url, body: JSON.generate(body))
+          handle_httpx_error_response!(response,
+                                       context: { location: "message endpoint request", request_id: request_id })
+
+          unless [200, 202].include?(response.status)
+            message = "Failed to have a successful request to #{@messages_url}: #{response.status} - #{response.body}"
+            RubyLLM::MCP.logger.error(message)
+            raise Errors::TransportError.new(
+              message: message,
+              code: response.status
+            )
+          end
+        end
+
         def start_sse_listener
           @connection_mutex.synchronize do
             return if sse_thread_running?
@@ -167,15 +171,36 @@ module RubyLLM
           sse_client = sse_client.with(
             headers: @headers
           )
+
+          if @version == :http1
+            sse_client = sse_client.with(
+              ssl: { alpn_protocols: ["http/1.1"] }
+            )
+          end
+
           response = sse_client.get(@event_url, stream: true)
+
+          event_buffer = []
           response.each_line do |event_line|
             unless @running
               response.body.close
               next
             end
 
-            event = parse_event(event_line)
-            process_event(event)
+            # Strip the line and check if it's empty (indicates end of event)
+            line = event_line.strip
+
+            if line.empty?
+              # End of event - process the accumulated buffer
+              if event_buffer.any?
+                event = parse_event(event_buffer.join("\n"))
+                process_event(event)
+                event_buffer.clear
+              end
+            else
+              # Accumulate the line for the current event
+              event_buffer << line
+            end
           end
         end
 
@@ -187,6 +212,25 @@ module RubyLLM
           sleep 1
         end
 
+        def handle_httpx_error_response!(response, context:)
+          return false unless response.is_a?(HTTPX::ErrorResponse)
+
+          error = response.error
+
+          if error.is_a?(HTTPX::ReadTimeoutError)
+            raise Errors::TimeoutError.new(
+              message: "Request timed out after #{@request_timeout / 1000} seconds"
+            )
+          end
+
+          error_message = response.error&.message || "Request failed"
+
+          raise Errors::TransportError.new(
+            code: nil,
+            message: "Request Error #{context}: #{error_message}"
+          )
+        end
+
         def process_event(raw_event)
           # Return if we believe that are getting a partial event
           return if raw_event[:data].nil?

data/lib/ruby_llm/mcp/transports/stdio.rb

@@ -9,11 +9,11 @@ module RubyLLM
   module MCP
     module Transports
       class Stdio
-        include Timeout
+        include Support::Timeout
 
         attr_reader :command, :stdin, :stdout, :stderr, :id, :coordinator
 
-        def initialize(command:, request_timeout:, coordinator:, args: [], env: {})
+        def initialize(command:, coordinator:, request_timeout:, args: [], env: {})
           @request_timeout = request_timeout
           @command = command
           @coordinator = coordinator

data/lib/ruby_llm/mcp/transports/streamable_http.rb

@@ -27,6 +27,21 @@ module RubyLLM
         end
       end
 
+      class OAuthOptions
+        attr_reader :issuer, :client_id, :client_secret, :scope
+
+        def initialize(issuer:, client_id:, client_secret:, scopes:)
+          @issuer = issuer
+          @client_id = client_id
+          @client_secret = client_secret
+          @scope = scopes
+        end
+
+        def enabled?
+          @issuer && @client_id && @client_secret && @scope
+        end
+      end
+
       # Options for starting SSE connections
       class StartSSEOptions
         attr_reader :resumption_token, :on_resumption_token, :replay_message_id
@@ -40,7 +55,7 @@ module RubyLLM
 
       # Main StreamableHTTP transport class
       class StreamableHTTP
-        include Timeout
+        include Support::Timeout
 
         attr_reader :session_id, :protocol_version, :coordinator
 
@@ -49,6 +64,10 @@ module RubyLLM
           request_timeout:,
           coordinator:,
           headers: {},
+          reconnection: {},
+          version: :http2,
+          oauth: nil,
+          rate_limit: nil,
           reconnection_options: nil,
           session_id: nil
         )
@@ -57,11 +76,19 @@ module RubyLLM
           @request_timeout = request_timeout
           @headers = headers || {}
           @session_id = session_id
+
+          @version = version
           @reconnection_options = reconnection_options || ReconnectionOptions.new
           @protocol_version = nil
+          @session_id = session_id
+
           @resource_metadata_url = nil
           @client_id = SecureRandom.uuid
 
+          @reconnection_options = ReconnectionOptions.new(**reconnection)
+          @oauth_options = OAuthOptions.new(**oauth) unless oauth.nil?
+          @rate_limiter = Support::RateLimiter.new(**rate_limit) if rate_limit
+
           @id_counter = 0
           @id_mutex = Mutex.new
           @pending_requests = {}
@@ -79,6 +106,11 @@ module RubyLLM
         end
 
         def request(body, add_id: true, wait_for_response: true)
+          if @rate_limiter&.exceeded?
+            sleep(1) while @rate_limiter&.exceeded?
+          end
+          @rate_limiter&.add
+
           # Generate a unique request ID for requests
           if add_id && body.is_a?(Hash) && !body.key?("id")
             @id_mutex.synchronize { @id_counter += 1 }
@@ -202,7 +234,7 @@ module RubyLLM
         end
 
         def create_connection
-          client = HTTPClient.connection.with(
+          client = Support::HTTPClient.connection.with(
             timeout: {
               connect_timeout: 10,
               read_timeout: @request_timeout / 1000,
@@ -210,6 +242,18 @@ module RubyLLM
               operation_timeout: @request_timeout / 1000
             }
           )
+
+          if @oauth_options&.enabled?
+            client = client.plugin(:oauth).oauth_auth(
+              issuer: @oauth_options.issuer,
+              client_id: @oauth_options.client_id,
+              client_secret: @oauth_options.client_secret,
+              scope: @oauth_options.scope
+            )
+
+            client.with_access_token
+          end
+
           register_client(client)
         end
 
@@ -219,6 +263,7 @@ module RubyLLM
           headers["mcp-session-id"] = @session_id if @session_id
           headers["mcp-protocol-version"] = @protocol_version if @protocol_version
           headers["X-CLIENT-ID"] = @client_id
+          headers["Origin"] = @uri.to_s
 
           headers
         end
@@ -259,7 +304,8 @@ module RubyLLM
         def create_connection_with_streaming_callbacks(request_id)
           buffer = +""
 
-          client = HTTPClient.connection.plugin(:callbacks).on_response_body_chunk do |request, _response, chunk|
+          client = Support::HTTPClient.connection.plugin(:callbacks)
+                                      .on_response_body_chunk do |request, _response, chunk|
             next unless @running && !@abort_controller
 
             RubyLLM::MCP.logger.debug "Received chunk: #{chunk.bytesize} bytes for #{request.uri}"
@@ -274,6 +320,17 @@ module RubyLLM
               operation_timeout: @request_timeout / 1000
             }
           )
+
+          if @oauth_options&.enabled?
+            client = client.plugin(:oauth).oauth_auth(
+              issuer: @oauth_options.issuer,
+              client_id: @oauth_options.client_id,
+              client_secret: @oauth_options.client_secret,
+              scope: @oauth_options.scope
+            )
+
+            client.with_access_token
+          end
           register_client(client)
         end
 
@@ -422,8 +479,8 @@ module RubyLLM
             end
 
             # Set up SSE streaming connection with callbacks
-            connection = create_connection_with_sse_callbacks(options)
-            response = connection.get(@url, headers: headers)
+            connection = create_connection_with_sse_callbacks(options, headers)
+            response = connection.get(@url)
 
             # Handle HTTPX error responses first
             error_result = handle_httpx_error_response!(response, context: { location: "SSE connection" },
@@ -463,12 +520,32 @@ module RubyLLM
           end
         end
 
-        def create_connection_with_sse_callbacks(options)
-          buffer = +""
+        def create_connection_with_sse_callbacks(options, headers)
+          client = HTTPX.plugin(:callbacks)
+          client = add_on_response_body_chunk_callback(client, options)
+
+          client = client.with(
+            timeout: {
+              connect_timeout: 10,
+              read_timeout: @request_timeout / 1000,
+              write_timeout: @request_timeout / 1000,
+              operation_timeout: @request_timeout / 1000
+            },
+            headers: headers
+          )
 
-          client = HTTPX
-                   .plugin(:callbacks)
-                   .on_response_body_chunk do |request, response, chunk|
+          if @version == :http1
+            client = client.with(
+              ssl: { alpn_protocols: ["http/1.1"] }
+            )
+          end
+
+          register_client(client)
+        end
+
+        def add_on_response_body_chunk_callback(client, options)
+          buffer = +""
+          client.on_response_body_chunk do |request, response, chunk|
             # Only process chunks for text/event-stream and if still running
             next unless @running && !@abort_controller
 
@@ -495,15 +572,6 @@ module RubyLLM
               end
             end
           end
-            .with(
-              timeout: {
-                connect_timeout: 10,
-                read_timeout: @request_timeout / 1000,
-                write_timeout: @request_timeout / 1000,
-                operation_timeout: @request_timeout / 1000
-              }
-            )
-          register_client(client)
         end
 
         def calculate_reconnection_delay(attempt)

data/lib/ruby_llm/mcp/transports/support/http_client.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+require "httpx"
+
+module RubyLLM
+  module MCP
+    module Transports
+      module Support
+        class HTTPClient
+          CONNECTION_KEY = :ruby_llm_mcp_client_connection
+
+          def self.connection
+            Thread.current[CONNECTION_KEY] ||= build_connection
+          end
+
+          def self.build_connection
+            HTTPX.with(
+              pool_options: {
+                max_connections: RubyLLM::MCP.config.max_connections,
+                pool_timeout: RubyLLM::MCP.config.pool_timeout
+              }
+            )
+          end
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/mcp/transports/support/rate_limit.rb

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module MCP
+    module Transports
+      module Support
+        class RateLimit
+          def initialize(limit: 10, interval: 1000)
+            @limit = limit
+            @interval = interval
+            @timestamps = []
+            @mutex = Mutex.new
+          end
+
+          def exceeded?
+            now = current_time
+
+            @mutex.synchronize do
+              purge_old(now)
+              @timestamps.size >= @limit
+            end
+          end
+
+          def add
+            now = current_time
+
+            @mutex.synchronize do
+              purge_old(now)
+              @timestamps << now
+            end
+          end
+
+          private
+
+          def current_time
+            Process.clock_gettime(Process::CLOCK_MONOTONIC)
+          end
+
+          def purge_old(now)
+            cutoff = now - @interval
+            @timestamps.reject! { |t| t < cutoff }
+          end
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/mcp/transports/support/timeout.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module MCP
+    module Transports
+      module Support
+        module Timeout
+          def with_timeout(seconds, request_id: nil)
+            result = nil
+            exception = nil
+
+            worker = Thread.new do
+              result = yield
+            rescue StandardError => e
+              exception = e
+            end
+
+            if worker.join(seconds)
+              raise exception if exception
+
+              result
+            else
+              worker.kill # stop the thread (can still have some risk if shared resources)
+              raise RubyLLM::MCP::Errors::TimeoutError.new(
+                message: "Request timed out after #{seconds} seconds",
+                request_id: request_id
+              )
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/mcp/version.rb

@@ -2,6 +2,6 @@
 
 module RubyLLM
   module MCP
-    VERSION = "0.5.0"
+    VERSION = "0.6.0"
   end
 end

data/lib/ruby_llm/mcp.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "ruby_llm"
+require "json-schema"
 require "zeitwerk"
 require_relative "chat"
 
@@ -9,20 +10,21 @@ module RubyLLM
     module_function
 
     def clients(config = RubyLLM::MCP.config.mcp_configuration)
-      @clients ||= {}
-      config.map do |options|
-        @clients[options[:name]] ||= Client.new(**options)
+      if @clients.nil?
+        @clients = {}
+        config.map do |options|
+          @clients[options[:name]] ||= Client.new(**options)
+        end
       end
+      @clients
     end
 
     def add_client(options)
-      @clients ||= {}
-      @clients[options[:name]] ||= Client.new(**options)
+      clients[options[:name]] ||= Client.new(**options)
     end
 
     def remove_client(name)
-      @clients ||= {}
-      client = @clients.delete(name)
+      client = clients.delete(name)
       client&.stop
       client
     end
@@ -33,8 +35,18 @@ module RubyLLM
 
     def establish_connection(&)
       clients.each(&:start)
-      yield clients
-    ensure
+      if block_given?
+        begin
+          yield clients
+        ensure
+          close_connection
+        end
+      else
+        clients
+      end
+    end
+
+    def close_connection
       clients.each do |client|
         client.stop if client.alive?
       end

data/lib/tasks/release.rake

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+namespace :release do
+  desc "Release a new version of the gem"
+  task :version do
+    # Load the current version from version.rb
+    require_relative "../../lib/ruby_llm/schema/version"
+    version = RubyLlm::Schema::VERSION
+
+    puts "Releasing version #{version}..."
+
+    # Make sure we are on the main branch
+    system "git checkout main"
+    system "git pull origin main"
+
+    # Create a new tag for the version
+    system "git tag -a v#{version} -m 'Release version #{version}'"
+    system "git push origin v#{version}"
+
+    system "gem build ruby_llm-mcp.gemspec"
+    system "gem push ruby_llm-mcp-#{version}.gem"
+  end
+end