vector_mcp 0.3.1 → 0.3.3

data/lib/vector_mcp/transport/http_stream/session_manager.rb

@@ -0,0 +1,189 @@
+# frozen_string_literal: true
+
+require "securerandom"
+require "concurrent-ruby"
+require_relative "../base_session_manager"
+
+module VectorMCP
+  module Transport
+    class HttpStream
+      # Manages HTTP stream sessions with automatic cleanup and thread safety.
+      # Extends BaseSessionManager with HTTP streaming-specific functionality.
+      #
+      # Handles:
+      # - Session creation and lifecycle management
+      # - Thread-safe session storage using concurrent-ruby
+      # - Automatic session cleanup based on timeout
+      # - Session context integration with VectorMCP::Session
+      # - HTTP streaming connection management
+      #
+      # @api private
+      class SessionManager < BaseSessionManager
+        # HTTP stream session data structure extending base session
+        Session = Struct.new(:id, :context, :created_at, :last_accessed_at, :metadata) do
+          def touch!
+            self.last_accessed_at = Time.now
+          end
+
+          def expired?(timeout)
+            Time.now - last_accessed_at > timeout
+          end
+
+          def age
+            Time.now - created_at
+          end
+
+          def streaming?
+            metadata[:streaming_connection] && !metadata[:streaming_connection].nil?
+          end
+
+          def streaming_connection
+            metadata[:streaming_connection]
+          end
+
+          def streaming_connection=(connection)
+            metadata[:streaming_connection] = connection
+          end
+        end
+
+        # Initializes a new HTTP stream session manager.
+        #
+        # @param transport [HttpStream] The parent transport instance
+        # @param session_timeout [Integer] Session timeout in seconds
+
+        # Optimized session creation with reduced object allocation and faster context creation
+        def create_session(session_id = nil, rack_env = nil)
+          session_id ||= generate_session_id
+          now = Time.now
+
+          # Optimize session context creation - use cached minimal context when rack_env is nil
+          session_context = if rack_env
+                              create_session_with_context(session_id, rack_env)
+                            else
+                              create_minimal_session_context(session_id)
+                            end
+
+          # Pre-allocate metadata hash for better performance
+          metadata = { streaming_connection: nil }
+
+          # Create internal session record with streaming connection metadata
+          session = Session.new(session_id, session_context, now, now, metadata)
+
+          @sessions[session_id] = session
+
+          logger.info { "Session created: #{session_id}" }
+          session
+        end
+
+        # Override to add rack_env support
+        def get_or_create_session(session_id = nil, rack_env = nil)
+          if session_id
+            session = get_session(session_id)
+            if session
+              # Update existing session context if rack_env is provided
+              if rack_env
+                request_context = VectorMCP::RequestContext.from_rack_env(rack_env, "http_stream")
+                session.context.request_context = request_context
+              end
+              return session
+            end
+
+            # If session_id was provided but not found, create with that ID
+            return create_session(session_id, rack_env)
+          end
+
+          create_session(nil, rack_env)
+        end
+
+        # Creates a VectorMCP::Session with proper request context from Rack environment
+        def create_session_with_context(session_id, rack_env)
+          request_context = VectorMCP::RequestContext.from_rack_env(rack_env, "http_stream")
+          VectorMCP::Session.new(@transport.server, @transport, id: session_id, request_context: request_context)
+        end
+
+        # Creates a minimal session context for each session (no caching to prevent contamination)
+        def create_minimal_session_context(session_id)
+          # Create a new minimal context for each session to prevent cross-session contamination
+          minimal_context = VectorMCP::RequestContext.minimal("http_stream")
+          VectorMCP::Session.new(@transport.server, @transport, id: session_id, request_context: minimal_context)
+        end
+
+        # Terminates a session by ID.
+        #
+        # @param session_id [String] The session ID to terminate
+        # @return [Boolean] True if session was found and terminated
+        # rubocop:disable Naming/PredicateMethod
+        def terminate_session(session_id)
+          session = @sessions.delete(session_id)
+          return false unless session
+
+          on_session_terminated(session)
+          logger.info { "Session terminated: #{session_id}" }
+          true
+        end
+        # rubocop:enable Naming/PredicateMethod
+
+        # Associates a streaming connection with a session.
+        #
+        # @param session [Session] The session to associate with
+        # @param connection [Object] The streaming connection object
+        # @return [void]
+        def set_streaming_connection(session, connection)
+          session.streaming_connection = connection
+          session.touch!
+          logger.debug { "Streaming connection associated: #{session.id}" }
+        end
+
+        # Removes streaming connection from a session.
+        #
+        # @param session [Session] The session to remove streaming from
+        # @return [void]
+        def remove_streaming_connection(session)
+          session.streaming_connection = nil
+          session.touch!
+          logger.debug { "Streaming connection removed: #{session.id}" }
+        end
+
+        protected
+
+        # Override: Called when a session is terminated to clean up streaming connections.
+        def on_session_terminated(session)
+          close_streaming_connection(session)
+        end
+
+        # Override: Returns metadata for new HTTP stream sessions.
+        def create_session_metadata
+          { streaming_connection: nil }
+        end
+
+        # Override: Checks if a session can receive messages (has streaming connection).
+        def can_send_message_to_session?(session)
+          session.streaming?
+        end
+
+        # Override: Sends a message to a session via the stream handler.
+        def message_sent_to_session?(session, message)
+          @transport.stream_handler.send_message_to_session(session, message)
+        end
+
+        private
+
+        # Closes a session's streaming connection if it exists.
+        #
+        # @param session [Session] The session whose connection to close
+        # @return [void]
+        def close_streaming_connection(session)
+          return unless session&.streaming_connection
+
+          begin
+            session.streaming_connection.close
+          rescue StandardError => e
+            logger.warn { "Error closing streaming connection for #{session.id}: #{e.message}" }
+          end
+
+          session.streaming_connection = nil
+        end
+      end
+    end
+  end
+end

data/lib/vector_mcp/transport/http_stream/stream_handler.rb

@@ -0,0 +1,269 @@
+# frozen_string_literal: true
+
+require "concurrent-ruby"
+
+module VectorMCP
+  module Transport
+    class HttpStream
+      # Handles Server-Sent Events streaming for HTTP transport.
+      #
+      # Manages:
+      # - SSE connection lifecycle
+      # - Event streaming with resumability
+      # - Last-Event-ID header processing
+      # - Connection health monitoring
+      #
+      # @api private
+      class StreamHandler
+        attr_reader :transport, :logger
+
+        # Streaming connection data structure
+        StreamingConnection = Struct.new(:session, :yielder, :thread, :closed) do
+          def close
+            self.closed = true
+            thread&.kill
+          end
+
+          def closed?
+            closed
+          end
+        end
+
+        # Initializes a new stream handler.
+        #
+        # @param transport [HttpStream] The parent transport instance
+        def initialize(transport)
+          @transport = transport
+          @logger = transport.logger
+          @active_connections = Concurrent::Hash.new
+        end
+
+        # Handles a streaming request (GET request for SSE).
+        #
+        # @param env [Hash] The Rack environment
+        # @param session [SessionManager::Session] The session for this request
+        # @return [Array] Rack response triplet for SSE
+        def handle_streaming_request(env, session)
+          last_event_id = extract_last_event_id(env)
+
+          logger.info("Starting SSE stream for session #{session.id}")
+
+          headers = build_sse_headers
+          body = create_sse_stream(session, last_event_id)
+
+          [200, headers, body]
+        end
+
+        # Sends a message to a specific session.
+        #
+        # @param session [SessionManager::Session] The target session
+        # @param message [Hash] The message to send
+        # @return [Boolean] True if message was sent successfully
+        def send_message_to_session(session, message)
+          return false unless session.streaming?
+
+          connection = @active_connections[session.id]
+          return false unless connection && !connection.closed?
+
+          begin
+            # Store event for resumability
+            event_data = message.to_json
+            event_id = @transport.event_store.store_event(event_data, "message")
+
+            # Send via SSE
+            sse_event = format_sse_event(event_data, "message", event_id)
+            connection.yielder << sse_event
+
+            logger.debug("Message sent to session #{session.id}")
+
+            true
+          rescue StandardError => e
+            logger.error("Error sending message to session #{session.id}: #{e.message}")
+
+            # Mark connection as closed and clean up
+            cleanup_connection(session)
+            false
+          end
+        end
+
+        # Gets the number of active streaming connections.
+        #
+        # @return [Integer] Number of active connections
+        def active_connection_count
+          @active_connections.size
+        end
+
+        # Cleans up all active connections.
+        #
+        # @return [void]
+        def cleanup_all_connections
+          logger.info("Cleaning up all streaming connections: #{@active_connections.size}")
+
+          @active_connections.each_value(&:close)
+
+          @active_connections.clear
+        end
+
+        private
+
+        # Extracts Last-Event-ID from request headers.
+        #
+        # @param env [Hash] The Rack environment
+        # @return [String, nil] The last event ID or nil
+        def extract_last_event_id(env)
+          env["HTTP_LAST_EVENT_ID"]
+        end
+
+        # Builds SSE response headers.
+        #
+        # @return [Hash] SSE headers
+        def build_sse_headers
+          {
+            "Content-Type" => "text/event-stream",
+            "Cache-Control" => "no-cache",
+            "Connection" => "keep-alive",
+            "X-Accel-Buffering" => "no",
+            "Access-Control-Allow-Origin" => "*",
+            "Access-Control-Allow-Headers" => "Last-Event-ID"
+          }
+        end
+
+        # Creates an SSE stream for a session.
+        #
+        # @param session [SessionManager::Session] The session
+        # @param last_event_id [String, nil] The last event ID for resumability
+        # @return [Enumerator] SSE stream enumerator
+        def create_sse_stream(session, last_event_id)
+          Enumerator.new do |yielder|
+            connection = StreamingConnection.new(session, yielder, nil, false)
+
+            # Register connection
+            @active_connections[session.id] = connection
+            @transport.session_manager.set_streaming_connection(session, connection)
+
+            # Start streaming thread
+            connection.thread = Thread.new do
+              stream_to_client(session, yielder, last_event_id)
+            rescue StandardError => e
+              logger.error("Error in streaming thread for #{session.id}: #{e.message}")
+            ensure
+              cleanup_connection(session)
+            end
+
+            # Keep connection alive until thread completes
+            connection.thread.join
+          end
+        end
+
+        # Streams events to a client.
+        #
+        # @param session [SessionManager::Session] The session
+        # @param yielder [Enumerator::Yielder] The SSE yielder
+        # @param last_event_id [String, nil] The last event ID for resumability
+        # @return [void]
+        def stream_to_client(session, yielder, last_event_id)
+          # Send initial connection event
+          connection_event = {
+            jsonrpc: "2.0",
+            method: "connection/established",
+            params: {
+              session_id: session.id,
+              timestamp: Time.now.iso8601
+            }
+          }
+
+          event_id = @transport.event_store.store_event(connection_event.to_json, "connection")
+          yielder << format_sse_event(connection_event.to_json, "connection", event_id)
+
+          # Replay missed events if resuming
+          replay_events(yielder, last_event_id) if last_event_id
+
+          # Send periodic keep-alive events
+          keep_alive_loop(session, yielder)
+        end
+
+        # Replays events after a specific event ID.
+        #
+        # @param yielder [Enumerator::Yielder] The SSE yielder
+        # @param last_event_id [String] The last event ID received by client
+        # @return [void]
+        def replay_events(yielder, last_event_id)
+          missed_events = @transport.event_store.get_events_after(last_event_id)
+
+          logger.info("Replaying #{missed_events.length} missed events from #{last_event_id}")
+
+          missed_events.each do |event|
+            yielder << event.to_sse_format
+          end
+        end
+
+        # Keeps the connection alive with periodic heartbeat events.
+        #
+        # @param session [SessionManager::Session] The session
+        # @param yielder [Enumerator::Yielder] The SSE yielder
+        # @return [void]
+        def keep_alive_loop(session, yielder)
+          start_time = Time.now
+          max_duration = 300 # 5 minutes maximum connection time
+
+          loop do
+            sleep(30) # Send heartbeat every 30 seconds
+
+            connection = @active_connections[session.id]
+            break if connection.nil? || connection.closed?
+
+            # Check if connection has been alive too long
+            if Time.now - start_time > max_duration
+              logger.debug("Connection for #{session.id} reached maximum duration, closing")
+              break
+            end
+
+            # Send heartbeat
+            heartbeat_event = {
+              jsonrpc: "2.0",
+              method: "heartbeat",
+              params: { timestamp: Time.now.iso8601 }
+            }
+
+            begin
+              event_id = @transport.event_store.store_event(heartbeat_event.to_json, "heartbeat")
+              yielder << format_sse_event(heartbeat_event.to_json, "heartbeat", event_id)
+            rescue StandardError
+              logger.debug("Heartbeat failed for #{session.id}, connection likely closed")
+              break
+            end
+          end
+        end
+
+        # Formats data as an SSE event.
+        #
+        # @param data [String] The event data
+        # @param type [String] The event type
+        # @param event_id [String] The event ID
+        # @return [String] Formatted SSE event
+        def format_sse_event(data, type, event_id)
+          lines = []
+          lines << "id: #{event_id}"
+          lines << "event: #{type}" if type
+          lines << "data: #{data}"
+          lines << ""
+          "#{lines.join("\n")}\n"
+        end
+
+        # Cleans up a specific connection.
+        #
+        # @param session [SessionManager::Session] The session to clean up
+        # @return [void]
+        def cleanup_connection(session)
+          connection = @active_connections.delete(session.id)
+          return unless connection
+
+          connection.close
+          @transport.session_manager.remove_streaming_connection(session)
+
+          logger.debug("Streaming connection cleaned up for #{session.id}")
+        end
+      end
+    end
+  end
+end