actionmcp 0.50.13 → 0.52.0

data/lib/action_mcp/client/session_store.rb

@@ -0,0 +1,231 @@
+# frozen_string_literal: true
+
+module ActionMCP
+  module Client
+    # Abstract interface for session storage
+    module SessionStore
+      # Load session data by ID
+      def load_session(session_id)
+        raise NotImplementedError, "#{self.class} must implement #load_session"
+      end
+
+      # Save session data
+      def save_session(session_id, session_data)
+        raise NotImplementedError, "#{self.class} must implement #save_session"
+      end
+
+      # Delete session
+      def delete_session(session_id)
+        raise NotImplementedError, "#{self.class} must implement #delete_session"
+      end
+
+      # Check if session exists
+      def session_exists?(session_id)
+        raise NotImplementedError, "#{self.class} must implement #session_exists?"
+      end
+
+      # Update specific session attributes
+      def update_session(session_id, attributes)
+        session_data = load_session(session_id)
+        return nil unless session_data
+
+        session_data.merge!(attributes)
+        save_session(session_id, session_data)
+        # Return the reloaded session to get the actual saved values
+        load_session(session_id)
+      end
+    end
+
+    # Volatile session store for development (data lost on restart)
+    class VolatileSessionStore
+      include SessionStore
+
+      def initialize
+        @sessions = Concurrent::Hash.new
+      end
+
+      def load_session(session_id)
+        @sessions[session_id]
+      end
+
+      def save_session(session_id, session_data)
+        @sessions[session_id] = session_data.dup
+      end
+
+      def delete_session(session_id)
+        @sessions.delete(session_id)
+      end
+
+      def session_exists?(session_id)
+        @sessions.key?(session_id)
+      end
+
+      def clear_all
+        @sessions.clear
+      end
+
+      def session_count
+        @sessions.size
+      end
+    end
+
+    # ActiveRecord-backed session store for production
+    class ActiveRecordSessionStore
+      include SessionStore
+
+      def load_session(session_id)
+        session = ActionMCP::Session.find_by(id: session_id)
+        return nil unless session
+
+        {
+          id: session.id,
+          protocol_version: session.protocol_version,
+          client_info: session.client_info,
+          client_capabilities: session.client_capabilities,
+          server_info: session.server_info,
+          server_capabilities: session.server_capabilities,
+          created_at: session.created_at,
+          updated_at: session.updated_at
+        }
+      end
+
+      def save_session(session_id, session_data)
+        session = ActionMCP::Session.find_or_initialize_by(id: session_id)
+
+        # Only assign attributes that exist in the database
+        attributes = {}
+        attributes[:protocol_version] = session_data[:protocol_version] if session_data.key?(:protocol_version)
+        attributes[:client_info] = session_data[:client_info] if session_data.key?(:client_info)
+        attributes[:client_capabilities] = session_data[:client_capabilities] if session_data.key?(:client_capabilities)
+        attributes[:server_info] = session_data[:server_info] if session_data.key?(:server_info)
+        attributes[:server_capabilities] = session_data[:server_capabilities] if session_data.key?(:server_capabilities)
+
+        # Store any extra data in a jsonb column if available
+        # For now, we'll skip last_event_id and session_data as they don't exist in the DB
+
+        session.assign_attributes(attributes)
+        session.save!
+        session_data
+      end
+
+      def delete_session(session_id)
+        ActionMCP::Session.find_by(id: session_id)&.destroy
+      end
+
+      def session_exists?(session_id)
+        ActionMCP::Session.exists?(id: session_id)
+      end
+
+      def cleanup_expired_sessions(older_than: 24.hours.ago)
+        ActionMCP::Session.where("updated_at < ?", older_than).delete_all
+      end
+    end
+
+    # Test session store that tracks all operations for assertions
+    class TestSessionStore < VolatileSessionStore
+      attr_reader :operations, :saved_sessions, :loaded_sessions,
+                  :deleted_sessions, :updated_sessions
+
+      def initialize
+        super
+        @operations = Concurrent::Array.new
+        @saved_sessions = Concurrent::Array.new
+        @loaded_sessions = Concurrent::Array.new
+        @deleted_sessions = Concurrent::Array.new
+        @updated_sessions = Concurrent::Array.new
+      end
+
+      def load_session(session_id)
+        session = super
+        @operations << { type: :load, session_id: session_id, found: !session.nil? }
+        @loaded_sessions << session_id if session
+        session
+      end
+
+      def save_session(session_id, session_data)
+        super
+        @operations << { type: :save, session_id: session_id, data: session_data }
+        @saved_sessions << session_id
+      end
+
+      def delete_session(session_id)
+        result = super
+        @operations << { type: :delete, session_id: session_id }
+        @deleted_sessions << session_id
+        result
+      end
+
+      def update_session(session_id, attributes)
+        result = super
+        @operations << { type: :update, session_id: session_id, attributes: attributes }
+        @updated_sessions << session_id if result
+        result
+      end
+
+      # Test helper methods
+      def session_saved?(session_id)
+        @saved_sessions.include?(session_id)
+      end
+
+      def session_loaded?(session_id)
+        @loaded_sessions.include?(session_id)
+      end
+
+      def session_deleted?(session_id)
+        @deleted_sessions.include?(session_id)
+      end
+
+      def session_updated?(session_id)
+        @updated_sessions.include?(session_id)
+      end
+
+      def operation_count(type = nil)
+        if type
+          @operations.count { |op| op[:type] == type }
+        else
+          @operations.size
+        end
+      end
+
+      def last_saved_data(session_id)
+        @operations.reverse.find { |op| op[:type] == :save && op[:session_id] == session_id }&.dig(:data)
+      end
+
+      def reset_tracking!
+        @operations.clear
+        @saved_sessions.clear
+        @loaded_sessions.clear
+        @deleted_sessions.clear
+        @updated_sessions.clear
+      end
+    end
+
+    # Factory for creating session stores
+    class SessionStoreFactory
+      def self.create(type = nil, **options)
+        type ||= default_type
+
+        case type.to_sym
+        when :volatile, :memory
+          VolatileSessionStore.new
+        when :active_record, :persistent
+          ActiveRecordSessionStore.new
+        when :test
+          TestSessionStore.new
+        else
+          raise ArgumentError, "Unknown session store type: #{type}"
+        end
+      end
+
+      def self.default_type
+        if Rails.env.test?
+          :volatile  # Use volatile for tests unless explicitly using :test
+        elsif Rails.env.production?
+          :active_record
+        else
+          :volatile
+        end
+      end
+    end
+  end
+end

data/lib/action_mcp/client/streamable_http_transport.rb

@@ -0,0 +1,291 @@
+# frozen_string_literal: true
+
+require_relative "transport"
+require_relative "session_store"
+
+module ActionMCP
+  module Client
+    # StreamableHTTP transport implementation following MCP specification
+    class StreamableHttpTransport < TransportBase
+      class ConnectionError < StandardError; end
+      class AuthenticationError < StandardError; end
+
+      SSE_TIMEOUT = 10
+      ENDPOINT_TIMEOUT = 5
+
+      attr_reader :session_id, :last_event_id
+
+      def initialize(url, session_store:, session_id: nil, **options)
+        super(url, session_store: session_store, **options)
+        @session_id = session_id
+        @last_event_id = nil
+        @buffer = +""
+        @current_event = nil
+        @reconnect_attempts = 0
+        @max_reconnect_attempts = options[:max_reconnect_attempts] || 3
+        @reconnect_delay = options[:reconnect_delay] || 1.0
+
+        setup_http_client
+      end
+
+      def connect
+        log_debug("Connecting via StreamableHTTP to #{@url}")
+
+        # Load session if session_id provided
+        load_session_state if @session_id
+
+        # Start SSE stream if server supports it
+        start_sse_stream
+
+        set_connected(true)
+        set_ready(true)
+        log_debug("StreamableHTTP connection established")
+        true
+      rescue StandardError => e
+        handle_error(e)
+        false
+      end
+
+      def disconnect
+        return true unless connected?
+
+        log_debug("Disconnecting StreamableHTTP")
+        stop_sse_stream
+        save_session_state if @session_id
+        set_connected(false)
+        set_ready(false)
+        true
+      rescue StandardError => e
+        handle_error(e)
+        false
+      end
+
+      def send_message(message)
+        unless ready?
+          raise ConnectionError, "Transport not ready"
+        end
+
+        headers = build_post_headers
+        json_data = message.is_a?(String) ? message : message.to_json
+
+        log_debug("Sending message via POST")
+        response = @http_client.post(@url, json_data, headers)
+
+        handle_post_response(response, message)
+        true
+      rescue StandardError => e
+        handle_error(e)
+        false
+      end
+
+      private
+
+      def setup_http_client
+        require "faraday"
+        @http_client = Faraday.new do |f|
+          f.headers["User-Agent"] = user_agent
+          f.options.timeout = nil # No read timeout for SSE
+          f.options.open_timeout = SSE_TIMEOUT
+          f.adapter :net_http
+        end
+      end
+
+      def build_get_headers
+        headers = {
+          "Accept" => "text/event-stream",
+          "Cache-Control" => "no-cache"
+        }
+        headers["mcp-session-id"] = @session_id if @session_id
+        headers["Last-Event-ID"] = @last_event_id if @last_event_id
+        headers
+      end
+
+      def build_post_headers
+        headers = {
+          "Content-Type" => "application/json",
+          "Accept" => "application/json, text/event-stream"
+        }
+        headers["mcp-session-id"] = @session_id if @session_id
+        headers
+      end
+
+      def start_sse_stream
+        log_debug("Starting SSE stream")
+        @sse_thread = Thread.new { run_sse_stream }
+      end
+
+      def stop_sse_stream
+        return unless @sse_thread
+
+        log_debug("Stopping SSE stream")
+        @stop_requested = true
+        @sse_thread.kill if @sse_thread.alive?
+        @sse_thread = nil
+        @stop_requested = false
+      end
+
+      def run_sse_stream
+        headers = build_get_headers
+
+        @http_client.get(@url, nil, headers) do |req|
+          req.options.on_data = proc do |chunk, _bytes|
+            break if @stop_requested
+            process_sse_chunk(chunk)
+          end
+        end
+      rescue StandardError => e
+        handle_sse_error(e)
+      end
+
+      def process_sse_chunk(chunk)
+        @buffer << chunk
+        process_complete_events while @buffer.include?("\n\n")
+      end
+
+      def process_complete_events
+        event_data, _separator, rest = @buffer.partition("\n\n")
+        @buffer = rest
+
+        return if event_data.strip.empty?
+
+        parse_sse_event(event_data)
+      end
+
+      def parse_sse_event(event_data)
+        lines = event_data.split("\n")
+        event_id = nil
+        data_lines = []
+
+        lines.each do |line|
+          if line.start_with?("id:")
+            event_id = line[3..-1].strip
+          elsif line.start_with?("data:")
+            data_lines << line[5..-1].strip
+          end
+        end
+
+        return if data_lines.empty?
+
+        @last_event_id = event_id if event_id
+
+        begin
+          message_data = data_lines.join("\n")
+          message = MultiJson.load(message_data)
+          handle_message(message)
+        rescue MultiJson::ParseError => e
+          log_error("Failed to parse SSE message: #{e}")
+        end
+      end
+
+      def handle_post_response(response, original_message)
+        # Extract session ID from response headers
+        if response.headers["mcp-session-id"]
+          @session_id = response.headers["mcp-session-id"]
+        end
+
+        case response.status
+        when 200
+          handle_success_response(response)
+        when 202
+          # Accepted - message received, no immediate response
+          log_debug("Message accepted (202)")
+        when 401
+          raise AuthenticationError, "Authentication required"
+        when 405
+          # Method not allowed - server doesn't support this operation
+          log_debug("Server returned 405 - operation not supported")
+        else
+          handle_error_response(response)
+        end
+      end
+
+      def handle_success_response(response)
+        content_type = response.headers["content-type"]
+
+        if content_type&.include?("application/json")
+          # Direct JSON response
+          handle_json_response(response)
+        elsif content_type&.include?("text/event-stream")
+          # SSE response stream
+          handle_sse_response_stream(response)
+        end
+      end
+
+      def handle_json_response(response)
+        begin
+          message = MultiJson.load(response.body)
+          handle_message(message)
+        rescue MultiJson::ParseError => e
+          log_error("Failed to parse JSON response: #{e}")
+        end
+      end
+
+      def handle_sse_response_stream(response)
+        # Handle SSE stream from POST response
+        response.body.each_line do |line|
+          process_sse_chunk(line)
+        end
+      end
+
+      def handle_error_response(response)
+        error_msg = "HTTP #{response.status}: #{response.reason_phrase}"
+        if response.body && !response.body.empty?
+          error_msg << " - #{response.body}"
+        end
+        raise ConnectionError, error_msg
+      end
+
+      def handle_sse_error(error)
+        log_error("SSE stream error: #{error.message}")
+
+        if should_reconnect?
+          schedule_reconnect
+        else
+          handle_error(error)
+        end
+      end
+
+      def should_reconnect?
+        connected? && @reconnect_attempts < @max_reconnect_attempts
+      end
+
+      def schedule_reconnect
+        @reconnect_attempts += 1
+        delay = @reconnect_delay * @reconnect_attempts
+
+        log_debug("Scheduling SSE reconnect in #{delay}s (attempt #{@reconnect_attempts})")
+
+        Thread.new do
+          sleep(delay)
+          start_sse_stream unless @stop_requested
+        end
+      end
+
+      def load_session_state
+        session_data = @session_store.load_session(@session_id)
+        return unless session_data
+
+        @last_event_id = session_data[:last_event_id]
+        log_debug("Loaded session state: last_event_id=#{@last_event_id}")
+      end
+
+      def save_session_state
+        return unless @session_id
+
+        session_data = {
+          id: @session_id,
+          last_event_id: @last_event_id,
+          session_data: {},
+          protocol_version: PROTOCOL_VERSION
+        }
+
+        @session_store.save_session(@session_id, session_data)
+        log_debug("Saved session state")
+      end
+
+      def user_agent
+        "ActionMCP-StreamableHTTP/#{ActionMCP.gem_version}"
+      end
+    end
+  end
+end

data/lib/action_mcp/client/transport.rb

@@ -0,0 +1,107 @@
+# frozen_string_literal: true
+
+module ActionMCP
+  module Client
+    # Base transport interface for MCP client connections
+    module Transport
+      # Called when transport should establish connection
+      def connect
+        raise NotImplementedError, "#{self.class} must implement #connect"
+      end
+
+      # Called when transport should close connection
+      def disconnect
+        raise NotImplementedError, "#{self.class} must implement #disconnect"
+      end
+
+      # Send a message through the transport
+      def send_message(message)
+        raise NotImplementedError, "#{self.class} must implement #send_message"
+      end
+
+      # Check if transport is ready to send/receive
+      def ready?
+        raise NotImplementedError, "#{self.class} must implement #ready?"
+      end
+
+      # Check if transport is connected
+      def connected?
+        raise NotImplementedError, "#{self.class} must implement #connected?"
+      end
+
+      # Set callback for received messages
+      def on_message(&block)
+        @message_callback = block
+      end
+
+      # Set callback for errors
+      def on_error(&block)
+        @error_callback = block
+      end
+
+      # Set callback for connection events
+      def on_connect(&block)
+        @connect_callback = block
+      end
+
+      # Set callback for disconnection events
+      def on_disconnect(&block)
+        @disconnect_callback = block
+      end
+
+      protected
+
+      def handle_message(message)
+        @message_callback&.call(message)
+      end
+
+      def handle_error(error)
+        @error_callback&.call(error)
+      end
+
+      def handle_connect
+        @connect_callback&.call
+      end
+
+      def handle_disconnect
+        @disconnect_callback&.call
+      end
+    end
+
+    # Base class for transport implementations
+    class TransportBase
+      include Transport
+      include Logging
+
+      attr_reader :url, :options, :session_store
+
+      def initialize(url, session_store:, logger: ActionMCP.logger, **options)
+        @url = url
+        @session_store = session_store
+        @logger = logger
+        @options = options
+        @connected = false
+        @ready = false
+      end
+
+      def connected?
+        @connected
+      end
+
+      def ready?
+        @ready
+      end
+
+      protected
+
+      def set_connected(state)
+        @connected = state
+        state ? handle_connect : handle_disconnect
+      end
+
+      def set_ready(state)
+        @ready = state
+      end
+    end
+  end
+end

data/lib/action_mcp/client.rb

@@ -1,24 +1,64 @@
 # frozen_string_literal: true
 
+require_relative "client/transport"
+require_relative "client/session_store"
+require_relative "client/streamable_http_transport"
+
 module ActionMCP
   # Creates a client appropriate for the given endpoint.
   #
   # @param endpoint [String] The endpoint to connect to (URL).
+  # @param transport [Symbol] The transport type to use (:streamable_http, :sse for legacy)
+  # @param session_store [Symbol] The session store type (:memory, :active_record)
+  # @param session_id [String] Optional session ID for resuming connections
   # @param logger [Logger] The logger to use. Default is Logger.new($stdout).
   # @param options [Hash] Additional options to pass to the client constructor.
   #
-  # @return [Client::SSEClient] An instance of SSEClient for HTTP(S) endpoints.
+  # @return [Client::Base] An instance of the appropriate client.
   #
-  # @example
+  # @example Basic usage
   #   client = ActionMCP.create_client("http://127.0.0.1:3001/action_mcp")
   #   client.connect
-  def self.create_client(endpoint, logger: Logger.new($stdout), **options)
+  #
+  # @example With specific transport and session store
+  #   client = ActionMCP.create_client(
+  #     "http://127.0.0.1:3001/action_mcp",
+  #     transport: :streamable_http,
+  #     session_store: :active_record,
+  #     session_id: "existing-session-123"
+  #   )
+  #
+  # @example Memory-based for development
+  #   client = ActionMCP.create_client(
+  #     "http://127.0.0.1:3001/action_mcp",
+  #     session_store: :memory
+  #   )
+  def self.create_client(endpoint, transport: :streamable_http, session_store: nil, session_id: nil, logger: Logger.new($stdout), **options)
     unless endpoint =~ %r{\Ahttps?://}
       raise ArgumentError, "Only HTTP(S) endpoints are supported. STDIO and other transports are not supported."
     end
 
-    logger.info("Creating SSE client for endpoint: #{endpoint}")
-    Client::SSEClient.new(endpoint, logger: logger, **options)
+    # Create session store
+    store = Client::SessionStoreFactory.create(session_store, **options)
+
+    # Create transport
+    transport_instance = create_transport(transport, endpoint, session_store: store, session_id: session_id, logger: logger, **options)
+
+    logger.info("Creating #{transport} client for endpoint: #{endpoint}")
+    # Pass session_id to the client
+    Client::Base.new(transport: transport_instance, logger: logger, session_id: session_id, **options)
+  end
+
+  private_class_method def self.create_transport(type, endpoint, **options)
+    case type.to_sym
+    when :streamable_http
+      Client::StreamableHttpTransport.new(endpoint, **options)
+    when :sse
+      # Legacy SSE transport (wrapped for compatibility)
+      Client::LegacySSETransport.new(endpoint, **options)
+    else
+      raise ArgumentError, "Unknown transport type: #{type}"
+    end
   end
 
   module Client