ruby_llm_swarm-mcp 0.8.0 → 0.8.1

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

@@ -0,0 +1,435 @@
+# frozen_string_literal: true
+
+require "json"
+require "uri"
+require "httpx"
+require "timeout"
+require "securerandom"
+
+module RubyLLM
+  module MCP
+    module Transports
+      class SSE
+        include Support::Timeout
+
+        attr_reader :headers, :id, :coordinator, :oauth_provider
+
+        def initialize(url:, coordinator:, request_timeout:, options: {})
+          @event_url = url
+          @messages_url = nil
+          @coordinator = coordinator
+          @request_timeout = request_timeout
+          @version = options[:version] || options["version"] || :http2
+          @oauth_provider = options[:oauth_provider] || options["oauth_provider"]
+
+          uri = URI.parse(url)
+          @root_url = "#{uri.scheme}://#{uri.host}"
+          @root_url += ":#{uri.port}" if uri.port != uri.default_port
+
+          @client_id = SecureRandom.uuid
+          custom_headers = options[:headers] || options["headers"] || {}
+          @headers = custom_headers.merge({
+                                            "Accept" => "text/event-stream",
+                                            "Content-Type" => "application/json",
+                                            "Cache-Control" => "no-cache",
+                                            "X-CLIENT-ID" => @client_id
+                                          })
+
+          @id_counter = 0
+          @id_mutex = Mutex.new
+          @pending_requests = {}
+          @pending_mutex = Mutex.new
+          @connection_mutex = Mutex.new
+          @running = false
+          @sse_thread = nil
+
+          RubyLLM::MCP.logger.info "Initializing SSE transport to #{@event_url} with client ID #{@client_id}"
+          RubyLLM::MCP.logger.debug "OAuth provider: #{@oauth_provider ? 'present' : 'none'}" if @oauth_provider
+        end
+
+        def request(body, add_id: true, wait_for_response: true)
+          if add_id
+            @id_mutex.synchronize { @id_counter += 1 }
+            request_id = @id_counter
+            body["id"] = request_id
+          end
+
+          response_queue = Queue.new
+          if wait_for_response
+            @pending_mutex.synchronize do
+              @pending_requests[request_id.to_s] = response_queue
+            end
+          end
+
+          begin
+            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 "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 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
+          end
+        end
+
+        def alive?
+          @running
+        end
+
+        def start
+          return if @running
+
+          @running = true
+          start_sse_listener
+        end
+
+        def close
+          RubyLLM::MCP.logger.info "Closing SSE transport connection"
+          @running = false
+          @sse_thread&.join(1) # Give the thread a second to clean up
+          @sse_thread = nil
+        end
+
+        def set_protocol_version(version)
+          @protocol_version = version
+        end
+
+        private
+
+        def send_request(body, request_id)
+          request_headers = build_request_headers
+          http_client = Support::HTTPClient.connection.with(
+            timeout: { request_timeout: @request_timeout / 1000 },
+            headers: request_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)
+            handle_send_request_error(response)
+          end
+        end
+
+        def handle_send_request_error(response)
+          response_body = response.respond_to?(:body) ? response.body.to_s : "Unknown error"
+          status_code = response.respond_to?(:status) ? response.status : "Unknown"
+
+          # Try to parse JSON error
+          error_message = begin
+            error_body = JSON.parse(response_body)
+            if error_body.is_a?(Hash) && error_body["error"]
+              msg = error_body["error"]["message"] || error_body["error"]["code"] || error_body["error"].to_s
+              msg.to_s.strip.empty? ? "Empty error (full response: #{response_body})" : msg
+            else
+              response_body
+            end
+          rescue JSON::ParserError
+            response_body
+          end
+
+          full_message = "Failed to have a successful request to #{@messages_url}: #{status_code} - #{error_message}"
+          RubyLLM::MCP.logger.error(full_message)
+
+          # Special handling for 403 with OAuth
+          if status_code == 403 && @oauth_provider
+            raise Errors::TransportError.new(
+              message: "Authorization failed (403 Forbidden): #{error_message}. Check token scope and resource \
+                permissions at #{@oauth_provider.server_url}.",
+              code: status_code
+            )
+          end
+
+          raise Errors::TransportError.new(
+            message: full_message,
+            code: status_code
+          )
+        end
+
+        def start_sse_listener
+          @connection_mutex.synchronize do
+            return if sse_thread_running?
+
+            RubyLLM::MCP.logger.info "Starting SSE listener thread"
+
+            response_queue = Queue.new
+            @pending_mutex.synchronize do
+              @pending_requests["endpoint"] = response_queue
+            end
+
+            @sse_thread = Thread.new do
+              listen_for_events while @running
+            end
+            @sse_thread.abort_on_exception = true
+
+            with_timeout(@request_timeout / 1000) do
+              endpoint = response_queue.pop
+              set_message_endpoint(endpoint)
+            end
+          end
+        end
+
+        def set_message_endpoint(endpoint)
+          uri = URI.parse(endpoint)
+
+          @messages_url = if uri.host.nil?
+                            "#{@root_url}#{endpoint}"
+                          else
+                            endpoint
+                          end
+
+          RubyLLM::MCP.logger.info "SSE message endpoint set to: #{@messages_url}"
+        end
+
+        def sse_thread_running?
+          @sse_thread&.alive?
+        end
+
+        def listen_for_events
+          stream_events_from_server
+        rescue StandardError => e
+          handle_connection_error("SSE connection error", e)
+        end
+
+        def stream_events_from_server
+          sse_client = create_sse_client
+          response = sse_client.get(@event_url, stream: true)
+          validate_sse_response!(response)
+          process_event_stream(response)
+        end
+
+        def create_sse_client
+          stream_headers = build_request_headers
+          sse_client = HTTPX.plugin(:stream).with(headers: stream_headers)
+          return sse_client unless @version == :http1
+
+          sse_client.with(ssl: { alpn_protocols: ["http/1.1"] })
+        end
+
+        # Build request headers with OAuth authorization if available
+        def build_request_headers
+          headers = @headers.dup
+
+          if @oauth_provider
+            RubyLLM::MCP.logger.debug "OAuth provider present, attempting to get token..."
+            RubyLLM::MCP.logger.debug "  Server URL: #{@oauth_provider.server_url}"
+
+            token = @oauth_provider.access_token
+            if token
+              headers["Authorization"] = token.to_header
+              RubyLLM::MCP.logger.debug "✓ Applied OAuth authorization header: #{token.to_header[0..30]}..."
+            else
+              RubyLLM::MCP.logger.warn "✗ OAuth provider present but no valid token available!"
+              RubyLLM::MCP.logger.warn "  This means the token is not in storage or has expired"
+              RubyLLM::MCP.logger.warn "  Check that authentication completed successfully"
+            end
+          else
+            RubyLLM::MCP.logger.debug "No OAuth provider configured for this transport"
+          end
+
+          headers
+        end
+
+        def validate_sse_response!(response)
+          return unless response.status >= 400
+
+          error_body = read_error_body(response)
+
+          # Try to parse as JSON to get better error details
+          error_message = begin
+            error_data = JSON.parse(error_body)
+            if error_data.is_a?(Hash) && error_data["error"]
+              msg = error_data["error"]["message"] || error_data["error"]["code"] || error_data["error"].to_s
+              # If we still don't have a message, include the full error object
+              msg.to_s.strip.empty? ? "Empty error (full response: #{error_body})" : msg
+            else
+              error_body
+            end
+          rescue JSON::ParserError
+            error_body
+          end
+
+          full_error_message = "HTTP #{response.status} error from SSE endpoint: #{error_message}"
+          RubyLLM::MCP.logger.error full_error_message
+
+          handle_client_error!(full_error_message, response.status, error_message) if response.status < 500
+
+          raise StandardError, full_error_message
+        end
+
+        def handle_client_error!(full_error_message, status_code, error_message)
+          @running = false
+
+          # Special handling for 401 Unauthorized - OAuth authentication required
+          if status_code == 401
+            raise Errors::AuthenticationRequiredError.new(
+              message: "OAuth authentication required. Server returned 401 Unauthorized.",
+              code: 401
+            )
+          end
+
+          # Special handling for 403 Forbidden with OAuth
+          if status_code == 403 && @oauth_provider
+            raise Errors::TransportError.new(
+              message: "Authorization failed (403 Forbidden): #{error_message}. \
+                Check token scope and resource permissions at #{@oauth_provider.server_url}.",
+              code: status_code
+            )
+          end
+
+          raise Errors::TransportError.new(
+            message: full_error_message,
+            code: status_code
+          )
+        end
+
+        def process_event_stream(response)
+          event_buffer = []
+          response.each_line do |event_line|
+            break unless handle_event_line?(event_line, event_buffer, response)
+          end
+        end
+
+        def handle_event_line?(event_line, event_buffer, response)
+          unless @running
+            response.body.close
+            return false
+          end
+
+          line = event_line.strip
+
+          if line.empty?
+            process_buffered_event(event_buffer)
+          else
+            event_buffer << line
+          end
+
+          true
+        end
+
+        def process_buffered_event(event_buffer)
+          return unless event_buffer.any?
+
+          events = parse_event(event_buffer.join("\n"))
+          events.each { |event| process_event(event) }
+          event_buffer.clear
+        end
+
+        def read_error_body(response)
+          # Try to read the error body from the response
+          body = ""
+          begin
+            response.each do |chunk|
+              body << chunk
+            end
+          rescue StandardError
+            # If we can't read the body, just use what we have
+          end
+          body.strip.empty? ? "No error details provided" : body.strip
+        end
+
+        def handle_connection_error(message, error)
+          return unless @running
+
+          error_message = "#{message}: #{error.message}"
+          RubyLLM::MCP.logger.error "#{error_message}. Reconnecting in 1 seconds..."
+          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?
+
+          if raw_event[:event] == "endpoint"
+            request_id = "endpoint"
+            event = raw_event[:data]
+            return if event.nil?
+
+            RubyLLM::MCP.logger.debug "Received endpoint event: #{event}"
+            @pending_mutex.synchronize do
+              response_queue = @pending_requests.delete(request_id)
+              response_queue&.push(event)
+            end
+          else
+            event = begin
+              JSON.parse(raw_event[:data])
+            rescue JSON::ParserError => e
+              # We can sometimes get partial endpoint events, so we will ignore them
+              unless @endpoint.nil?
+                RubyLLM::MCP.logger.info "Failed to parse SSE event data: #{raw_event[:data]} - #{e.message}"
+              end
+
+              nil
+            end
+            return if event.nil?
+
+            request_id = event["id"]&.to_s
+            result = RubyLLM::MCP::Result.new(event)
+
+            result = @coordinator.process_result(result)
+            return if result.nil?
+
+            @pending_mutex.synchronize do
+              # You can receieve duplicate events for the same request id, and we will ignore thoses
+              if result.matching_id?(request_id) && @pending_requests.key?(request_id)
+                response_queue = @pending_requests.delete(request_id)
+                response_queue&.push(result)
+              end
+            end
+          end
+        end
+
+        def parse_event(raw)
+          event_blocks = raw.split(/\n\s*\n/)
+
+          events = event_blocks.map do |event_block|
+            event = {}
+            event_block.each_line do |line|
+              case line
+              when /^data:\s*(.*)/
+                (event[:data] ||= []) << ::Regexp.last_match(1)
+              when /^event:\s*(.*)/
+                event[:event] = ::Regexp.last_match(1)
+              when /^id:\s*(.*)/
+                event[:id] = ::Regexp.last_match(1)
+              end
+            end
+            event[:data] = event[:data]&.join("\n")
+            event
+          end
+
+          events.reject { |event| event.empty? || event[:data].nil? }
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,231 @@
+# frozen_string_literal: true
+
+require "open3"
+require "json"
+require "timeout"
+require "securerandom"
+
+module RubyLLM
+  module MCP
+    module Transports
+      class Stdio
+        include Support::Timeout
+
+        attr_reader :command, :stdin, :stdout, :stderr, :id, :coordinator
+
+        def initialize(command:, coordinator:, request_timeout:, options: {})
+          @request_timeout = request_timeout
+          @command = command
+          @coordinator = coordinator
+          @args = options[:args] || options["args"] || []
+          @env = options[:env] || options["env"] || {}
+          @client_id = SecureRandom.uuid
+          # NOTE: Stdio transport doesn't use OAuth (local process communication)
+
+          @id_counter = 0
+          @id_mutex = Mutex.new
+          @pending_requests = {}
+          @pending_mutex = Mutex.new
+          @running = false
+          @reader_thread = nil
+          @stderr_thread = nil
+        end
+
+        def request(body, add_id: true, wait_for_response: true)
+          if add_id
+            @id_mutex.synchronize { @id_counter += 1 }
+            request_id = @id_counter
+            body["id"] = request_id
+          end
+
+          response_queue = Queue.new
+          if wait_for_response
+            @pending_mutex.synchronize do
+              @pending_requests[request_id.to_s] = response_queue
+            end
+          end
+
+          begin
+            body = JSON.generate(body)
+            RubyLLM::MCP.logger.debug "Sending Request: #{body}"
+            @stdin.puts(body)
+            @stdin.flush
+          rescue IOError, Errno::EPIPE => e
+            @pending_mutex.synchronize { @pending_requests.delete(request_id.to_s) }
+            restart_process
+            raise RubyLLM::MCP::Errors::TransportError.new(message: e.message, error: 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
+            @pending_mutex.synchronize { @pending_requests.delete(request_id.to_s) }
+            log_message = "Stdio request timeout (ID: #{request_id}) after #{@request_timeout / 1000} seconds"
+            RubyLLM::MCP.logger.error(log_message)
+            raise e
+          end
+        end
+
+        def alive?
+          @running
+        end
+
+        def start
+          start_process unless @running
+          @running = true
+        end
+
+        def close
+          @running = false
+
+          [@stdin, @stdout, @stderr].each do |stream|
+            stream&.close
+          rescue IOError, Errno::EBADF
+            nil
+          end
+
+          [@wait_thread, @reader_thread, @stderr_thread].each do |thread|
+            thread&.join(1)
+          rescue StandardError
+            nil
+          end
+
+          @stdin = @stdout = @stderr = nil
+          @wait_thread = @reader_thread = @stderr_thread = nil
+        end
+
+        def set_protocol_version(version)
+          @protocol_version = version
+        end
+
+        private
+
+        def start_process
+          close if @stdin || @stdout || @stderr || @wait_thread
+
+          @stdin, @stdout, @stderr, @wait_thread = if @env.empty?
+                                                     Open3.popen3(@command, *@args)
+                                                   else
+                                                     Open3.popen3(@env, @command, *@args)
+                                                   end
+
+          start_reader_thread
+          start_stderr_thread
+        end
+
+        def restart_process
+          RubyLLM::MCP.logger.error "Process connection lost. Restarting..."
+          start_process
+        end
+
+        def start_reader_thread
+          @reader_thread = Thread.new do
+            read_stdout_loop
+          end
+
+          @reader_thread.abort_on_exception = true
+        end
+
+        def read_stdout_loop
+          while @running
+            begin
+              handle_stdout_read
+            rescue IOError, Errno::EPIPE => e
+              handle_stream_error(e, "Reader")
+              break unless @running
+            rescue StandardError => e
+              RubyLLM::MCP.logger.error "Error in reader thread: #{e.message}, #{e.backtrace.join("\n")}"
+              sleep 1
+            end
+          end
+        end
+
+        def handle_stdout_read
+          if @stdout.closed? || @wait_thread.nil? || !@wait_thread.alive?
+            if @running
+              sleep 1
+              restart_process
+            end
+            return
+          end
+
+          line = @stdout.gets
+          return unless line && !line.strip.empty?
+
+          process_response(line.strip)
+        end
+
+        def handle_stream_error(error, stream_name)
+          # Check @running to distinguish graceful shutdown from unexpected errors.
+          # During shutdown, streams are closed intentionally and shouldn't trigger restarts.
+          if @running
+            RubyLLM::MCP.logger.error "#{stream_name} error: #{error.message}. Restarting in 1 second..."
+            sleep 1
+            restart_process
+          else
+            # Graceful shutdown in progress
+            RubyLLM::MCP.logger.debug "#{stream_name} thread exiting during shutdown"
+          end
+        end
+
+        def start_stderr_thread
+          @stderr_thread = Thread.new do
+            read_stderr_loop
+          end
+
+          @stderr_thread.abort_on_exception = true
+        end
+
+        def read_stderr_loop
+          while @running
+            begin
+              handle_stderr_read
+            rescue IOError, Errno::EPIPE => e
+              handle_stream_error(e, "Stderr reader")
+              break unless @running
+            rescue StandardError => e
+              RubyLLM::MCP.logger.error "Error in stderr thread: #{e.message}"
+              sleep 1
+            end
+          end
+        end
+
+        def handle_stderr_read
+          if @stderr.closed? || @wait_thread.nil? || !@wait_thread.alive?
+            sleep 1
+            return
+          end
+
+          line = @stderr.gets
+          return unless line && !line.strip.empty?
+
+          RubyLLM::MCP.logger.info(line.strip)
+        end
+
+        def process_response(line)
+          response = JSON.parse(line)
+          request_id = response["id"]&.to_s
+          result = RubyLLM::MCP::Result.new(response)
+          RubyLLM::MCP.logger.debug "Result Received: #{result.inspect}"
+
+          result = @coordinator.process_result(result)
+          return if result.nil?
+
+          # Handle regular responses (tool calls, etc.)
+          @pending_mutex.synchronize do
+            if result.matching_id?(request_id) && @pending_requests.key?(request_id)
+              response_queue = @pending_requests.delete(request_id)
+              response_queue&.push(result)
+            end
+          end
+        rescue JSON::ParserError => e
+          RubyLLM::MCP.logger.error("Error parsing response as JSON: #{e.message}\nRaw response: #{line}")
+        end
+      end
+    end
+  end
+end