vector_mcp 0.3.1 → 0.3.3

data/lib/vector_mcp/transport/sse.rb

@@ -9,6 +9,7 @@ require "concurrent-ruby"
 require_relative "../errors"
 require_relative "../util"
 require_relative "../session"
+require_relative "sse_session_manager"
 require_relative "sse/client_connection"
 require_relative "sse/stream_manager"
 require_relative "sse/message_handler"
@@ -41,7 +42,7 @@ module VectorMCP
     # @attr_reader port [Integer] The port number the server will listen on.
     # @attr_reader path_prefix [String] The base URL path for MCP endpoints (e.g., "/mcp").
     class SSE
-      attr_reader :logger, :server, :host, :port, :path_prefix
+      attr_reader :logger, :server, :host, :port, :path_prefix, :session_manager
 
       # Initializes a new SSE transport.
       #
@@ -50,6 +51,8 @@ module VectorMCP
       # @option options [String] :host ("localhost") The hostname or IP to bind to.
       # @option options [Integer] :port (8000) The port to listen on.
       # @option options [String] :path_prefix ("/mcp") The base path for HTTP endpoints.
+      # @option options [Boolean] :disable_session_manager (false) **DEPRECATED**: Whether to disable secure session isolation.
+      #   When false (default), each client gets isolated sessions. When true, all clients share a global session (security risk).
       def initialize(server, options = {})
         @server = server
         @logger = server.logger
@@ -61,9 +64,19 @@ module VectorMCP
         @sse_path = "#{@path_prefix}/sse"
         @message_path = "#{@path_prefix}/message"
 
-        # Thread-safe client storage using concurrent-ruby
+        # Thread-safe client storage using concurrent-ruby (legacy approach)
         @clients = Concurrent::Hash.new
         @session = nil # Global session for this transport instance, initialized in run
+
+        # Initialize session manager for secure multi-client session isolation (default behavior)
+        # Legacy shared session behavior can be enabled with disable_session_manager: true (deprecated)
+        if options[:disable_session_manager]
+          logger.warn("[DEPRECATED] SSE shared session mode is deprecated and poses security risks in multi-client scenarios. " \
+                      "Consider removing disable_session_manager: true to use secure per-client sessions.")
+          @session_manager = nil
+        else
+          @session_manager = SseSessionManager.new(self)
+        end
         @puma_server = nil
         @running = false
 
@@ -77,7 +90,8 @@ module VectorMCP
       # @raise [StandardError] if there's a fatal error during server startup.
       def run
         logger.info("Starting server with Puma SSE transport on #{@host}:#{@port}")
-        create_session
+        # Only create shared session if explicitly using legacy mode (deprecated)
+        create_session unless @session_manager
         start_puma_server
       rescue StandardError => e
         handle_fatal_error(e)
@@ -106,13 +120,32 @@ module VectorMCP
 
       # --- Public methods for Server to send notifications ---
 
+      # Sends a JSON-RPC notification to the first available client session.
+      # If no clients are connected, returns false.
+      #
+      # @param method [String] The method name of the notification.
+      # @param params [Hash, Array, nil] The parameters for the notification (optional).
+      # @return [Boolean] True if the message was sent successfully, false otherwise.
+      def send_notification(method, params = nil)
+        return false if @clients.empty?
+
+        # Send to first available client
+        first_client = @clients.values.first
+        return false unless first_client
+
+        message = { jsonrpc: "2.0", method: method }
+        message[:params] = params if params
+
+        StreamManager.enqueue_message(first_client, message)
+      end
+
       # Sends a JSON-RPC notification to a specific client session via its SSE stream.
       #
       # @param session_id [String] The ID of the client session to send the notification to.
       # @param method [String] The method name of the notification.
       # @param params [Hash, Array, nil] The parameters for the notification (optional).
       # @return [Boolean] True if the message was successfully enqueued, false otherwise (e.g., client not found).
-      def send_notification(session_id, method, params = nil)
+      def send_notification_to_session(session_id, method, params = nil)
         message = { jsonrpc: "2.0", method: method }
         message[:params] = params if params
 
@@ -128,7 +161,7 @@ module VectorMCP
       # @param params [Hash, Array, nil] The parameters for the notification (optional).
       # @return [void]
       def broadcast_notification(method, params = nil)
-        logger.debug { "Broadcasting notification '#{method}' to #{@clients.size} client(s)" }
+        # Broadcasting notification to clients
         message = { jsonrpc: "2.0", method: method }
         message[:params] = params if params
 
@@ -150,11 +183,26 @@ module VectorMCP
       # Stops the transport and cleans up resources
       def stop
         @running = false
-        cleanup_clients
+        if @session_manager
+          @session_manager.cleanup_all_sessions
+        else
+          cleanup_clients
+        end
         @puma_server&.stop
         logger.info("SSE transport stopped")
       end
 
+      # Cleans up all client connections (legacy mode)
+      def cleanup_clients
+        logger.info("Cleaning up #{@clients.size} client connection(s)")
+        @clients.each_value do |client_conn|
+          client_conn.close if client_conn.respond_to?(:close)
+        rescue StandardError => e
+          logger.warn("Error closing client connection: #{e.message}")
+        end
+        @clients.clear
+      end
+
       # --- Private methods ---
       private
 
@@ -177,8 +225,8 @@ module VectorMCP
         @puma_server.run.join # This blocks until server stops
         logger.info("Puma server stopped.")
       ensure
-        cleanup_clients
-        @session = nil
+        # Only cleanup if session manager is enabled
+        @session_manager&.cleanup_all_sessions
         logger.info("SSE transport and resources shut down.")
       end
 
@@ -194,13 +242,6 @@ module VectorMCP
         end
       end
 
-      # Cleans up resources for all connected clients on server shutdown.
-      def cleanup_clients
-        logger.info("Cleaning up #{@clients.size} client connection(s)...")
-        @clients.each_value(&:close)
-        @clients.clear
-      end
-
       # Handles fatal errors during server startup or main run loop.
       def handle_fatal_error(error)
         logger.fatal("Fatal error in SSE transport: #{error.message}\n#{error.backtrace.join("\n")}")
@@ -245,11 +286,17 @@ module VectorMCP
 
         # Create client connection
         client_conn = ClientConnection.new(session_id, logger)
-        @clients[session_id] = client_conn
+
+        # Store client connection
+        if @session_manager
+          @session_manager.register_client(session_id, client_conn)
+        else
+          @clients[session_id] = client_conn
+        end
 
         # Build message POST URL for this client
         message_post_url = build_post_url(session_id)
-        logger.debug("Client #{session_id} should POST messages to: #{message_post_url}")
+        # Client message POST URL configured
 
         # Set up SSE stream
         headers = sse_headers
@@ -268,10 +315,18 @@ module VectorMCP
                                 "Missing session_id parameter")
         end
 
-        client_conn = @clients[session_id]
+        # Get client connection and session
+        if @session_manager
+          client_conn = @session_manager.clients[session_id]
+          shared_session = @session_manager.shared_session
+        else
+          client_conn = @clients[session_id]
+          shared_session = @session
+        end
+
         return error_response(nil, VectorMCP::NotFoundError.new("Invalid session_id").code, "Invalid session_id") unless client_conn
 
-        MessageHandler.new(@server, @session, logger).handle_post_message(env, client_conn)
+        MessageHandler.new(@server, shared_session, logger).handle_post_message(env, client_conn)
       end
 
       # Helper methods

data/lib/vector_mcp/transport/sse_session_manager.rb

@@ -0,0 +1,188 @@
+# frozen_string_literal: true
+
+require_relative "base_session_manager"
+
+module VectorMCP
+  module Transport
+    # Session manager for SSE transport with single shared session and client connection management.
+    # Extends BaseSessionManager with SSE-specific functionality.
+    #
+    # The SSE transport uses a single shared session for all client connections,
+    # but manages multiple client connections separately.
+    class SseSessionManager < BaseSessionManager
+      attr_reader :clients
+
+      # Initializes a new SSE session manager.
+      #
+      # @param transport [SSE] The parent transport instance
+      # @param session_timeout [Integer] Session timeout in seconds
+      def initialize(transport, session_timeout = 300)
+        @clients = Concurrent::Hash.new
+        super
+
+        # Create the single shared session for SSE transport
+        @shared_session = create_shared_session
+      end
+
+      # Gets the shared session for SSE transport.
+      # SSE uses a single session shared across all client connections.
+      #
+      # @return [Session] The shared session
+      def shared_session
+        @shared_session.touch!
+        @shared_session
+      end
+
+      # Registers a client connection with the session manager.
+      #
+      # @param client_id [String] The client connection ID
+      # @param client_connection [Object] The client connection object
+      # @return [void]
+      def register_client(client_id, client_connection)
+        @clients[client_id] = client_connection
+        session_metadata_updated?(@shared_session.id, clients_count: @clients.size)
+        logger.debug { "Client registered: #{client_id}" }
+      end
+
+      # Unregisters a client connection from the session manager.
+      #
+      # @param client_id [String] The client connection ID
+      # @return [Boolean] True if client was found and removed
+      def client_unregistered?(client_id)
+        client = @clients.delete(client_id)
+        return false unless client
+
+        session_metadata_updated?(@shared_session.id, clients_count: @clients.size)
+        logger.debug { "Client unregistered: #{client_id}" }
+        true
+      end
+
+      # Gets all client connections.
+      #
+      # @return [Hash] Hash of client_id => client_connection
+      def all_clients
+        @clients.dup
+      end
+
+      # Gets the number of connected clients.
+      #
+      # @return [Integer] Number of connected clients
+      def client_count
+        @clients.size
+      end
+
+      # Cleans up all clients and the shared session.
+      #
+      # @return [void]
+      def cleanup_all_sessions
+        logger.info { "Cleaning up #{@clients.size} client connection(s)" }
+
+        @clients.each_value do |client_conn|
+          close_client_connection(client_conn)
+        end
+        @clients.clear
+
+        super
+      end
+
+      protected
+
+      # Override: SSE doesn't need automatic cleanup since it has a single shared session.
+      def auto_cleanup_enabled?
+        false
+      end
+
+      # Override: Called when the shared session is terminated.
+      def on_session_terminated(_session)
+        # Clean up all client connections when session is terminated
+        @clients.each_value do |client_conn|
+          close_client_connection(client_conn)
+        end
+        @clients.clear
+      end
+
+      # Override: Returns metadata for SSE sessions.
+      def create_session_metadata
+        { clients_count: 0, session_type: :sse_shared }
+      end
+
+      # Override: Checks if any clients are connected to receive messages.
+      def can_send_message_to_session?(_session)
+        !@clients.empty?
+      end
+
+      # Override: Sends a message to the first available client.
+      def send_message_to_session(_session, message)
+        return false if @clients.empty?
+
+        first_client = @clients.values.first
+        return false unless first_client
+
+        @transport.class::StreamManager.enqueue_message(first_client, message)
+      end
+
+      # Override: Broadcasts messages to all connected clients.
+      def broadcast_message(message)
+        count = 0
+        @clients.each_value do |client_conn|
+          count += 1 if @transport.class::StreamManager.enqueue_message(client_conn, message)
+        end
+
+        logger.debug { "Message broadcasted to #{count} client(s)" }
+        count
+      end
+
+      private
+
+      # Creates the single shared session for SSE transport.
+      #
+      # @return [BaseSessionManager::Session] The shared session
+      def create_shared_session(rack_env = nil)
+        session_id = "sse_shared_session_#{SecureRandom.uuid}"
+        now = Time.now
+
+        # Create VectorMCP session context with request context
+        session_context = create_session_with_context(session_id, rack_env)
+
+        # Create internal session record using base session manager struct
+        session = BaseSessionManager::Session.new(
+          session_id,
+          session_context,
+          now,
+          now,
+          create_session_metadata
+        )
+
+        @sessions[session_id] = session
+        logger.info { "Shared SSE session created: #{session_id}" }
+        session
+      end
+
+      # Creates a VectorMCP::Session with proper request context from Rack environment
+      def create_session_with_context(session_id, rack_env)
+        request_context = if rack_env
+                            # Create request context from Rack environment
+                            VectorMCP::RequestContext.from_rack_env(rack_env, "sse")
+                          else
+                            # Fallback to minimal context for cases where rack_env is not available
+                            VectorMCP::RequestContext.minimal("sse")
+                          end
+        VectorMCP::Session.new(@transport.server, @transport, id: session_id, request_context: request_context)
+      end
+
+      # Closes a client connection safely.
+      #
+      # @param client_conn [Object] The client connection to close
+      # @return [void]
+      def close_client_connection(client_conn)
+        return unless client_conn
+
+        begin
+          client_conn.close if client_conn.respond_to?(:close)
+        rescue StandardError => e
+          logger.warn { "Error closing client connection: #{e.message}" }
+        end
+      end
+    end
+  end
+end

data/lib/vector_mcp/transport/stdio.rb

@@ -4,6 +4,7 @@
 require "json"
 require_relative "../errors"
 require_relative "../util"
+require_relative "stdio_session_manager"
 require "securerandom" # For generating unique request IDs
 require "timeout"      # For request timeouts
 
@@ -20,6 +21,8 @@ module VectorMCP
       attr_reader :server
       # @return [Logger] The logger instance, shared with the server.
       attr_reader :logger
+      # @return [StdioSessionManager] The session manager for this transport.
+      attr_reader :session_manager
 
       # Timeout for waiting for a response to a server-initiated request (in seconds)
       DEFAULT_REQUEST_TIMEOUT = 30 # Configurable if needed
@@ -27,9 +30,12 @@ module VectorMCP
       # Initializes a new Stdio transport.
       #
       # @param server [VectorMCP::Server] The server instance that will handle messages.
-      def initialize(server)
+      # @param options [Hash] Optional configuration options.
+      # @option options [Boolean] :enable_session_manager (false) Whether to enable the unified session manager.
+      def initialize(server, options = {})
         @server = server
         @logger = server.logger
+        @session_manager = options[:enable_session_manager] ? StdioSessionManager.new(self) : nil
         @input_mutex = Mutex.new
         @output_mutex = Mutex.new
         @running = false
@@ -109,6 +115,43 @@ module VectorMCP
         write_message(notification)
       end
 
+      # Sends a JSON-RPC notification message to a specific session.
+      # For stdio transport, this behaves the same as send_notification since there's only one session.
+      #
+      # @param _session_id [String] The session ID (ignored for stdio transport).
+      # @param method [String] The method name of the notification.
+      # @param params [Hash, Array, nil] The parameters for the notification (optional).
+      # @return [Boolean] True if the notification was sent successfully.
+      # rubocop:disable Naming/PredicateMethod
+      def send_notification_to_session(_session_id, method, params = nil)
+        send_notification(method, params)
+        true
+      end
+      # rubocop:enable Naming/PredicateMethod
+
+      # Sends a JSON-RPC notification message to a specific session.
+      # For stdio transport, this behaves the same as send_notification since there's only one session.
+      #
+      # @param _session_id [String] The session ID (ignored for stdio transport).
+      # @param method [String] The method name of the notification.
+      # @param params [Hash, Array, nil] The parameters for the notification (optional).
+      # @return [Boolean] True if the notification was sent successfully.
+      def notification_sent_to_session?(_session_id, method, params = nil)
+        send_notification(method, params)
+        true
+      end
+
+      # Broadcasts a JSON-RPC notification message to all sessions.
+      # For stdio transport, this behaves the same as send_notification since there's only one session.
+      #
+      # @param method [String] The method name of the notification.
+      # @param params [Hash, Array, nil] The parameters for the notification (optional).
+      # @return [Integer] Number of sessions the notification was sent to (always 1 for stdio).
+      def broadcast_notification(method, params = nil)
+        send_notification(method, params)
+        1
+      end
+
       # Sends a server-initiated JSON-RPC request to the client and waits for a response.
       # This is a blocking call.
       #
@@ -126,7 +169,7 @@ module VectorMCP
         request_payload[:params] = params if params
 
         setup_request_tracking(request_id)
-        logger.debug "[Stdio Transport] Sending request ID #{request_id}: #{method}"
+        # Sending request to client
         write_message(request_payload)
 
         response = wait_for_response(request_id, method, timeout)
@@ -184,7 +227,6 @@ module VectorMCP
 
         return handle_outgoing_response(message) if outgoing_response?(message)
 
-        ensure_session_exists
         handle_server_message(message)
       end
 
@@ -196,7 +238,21 @@ module VectorMCP
         message["id"] && !message["method"] && (message.key?("result") || message.key?("error"))
       end
 
-      # Ensures a global session exists for this stdio transport.
+      # Gets the global session for this stdio transport.
+      # @api private
+      # @return [VectorMCP::Session] The current session.
+      def session
+        # Try session manager first, fallback to old method for backward compatibility
+        if @session_manager
+          session_wrapper = @session_manager.global_session
+          return session_wrapper.context if session_wrapper
+        end
+
+        # Fallback to old session creation for backward compatibility
+        ensure_session_exists
+      end
+
+      # Ensures a global session exists for this stdio transport (legacy method).
       # @api private
       # @return [VectorMCP::Session] The current session.
       def ensure_session_exists
@@ -208,11 +264,11 @@ module VectorMCP
       # @param message [Hash] The parsed message.
       # @return [void]
       def handle_server_message(message)
-        session = ensure_session_exists
-        session_id = session.id
+        current_session = session
+        session_id = current_session.id
 
         begin
-          result = @server.handle_message(message, session, session_id)
+          result = @server.handle_message(message, current_session, session_id)
           send_response(message["id"], result) if message["id"] && result
         rescue VectorMCP::ProtocolError => e
           handle_protocol_error(e, message)
@@ -223,11 +279,11 @@ module VectorMCP
 
       # --- Run helpers (private) ---
 
-      # Creates a new session for the stdio connection.
+      # Gets the session for the stdio connection.
       # @api private
-      # @return [VectorMCP::Session] The newly created session.
+      # @return [VectorMCP::Session] The session.
       def create_session
-        VectorMCP::Session.new(@server, self)
+        session
       end
 
       # Launches the input reading loop in a new thread.
@@ -250,6 +306,7 @@ module VectorMCP
       def shutdown_transport
         @running = false
         @input_thread&.kill if @input_thread&.alive?
+        @session_manager&.cleanup_all_sessions
         logger.info("Stdio transport shut down")
       end
 
@@ -261,7 +318,7 @@ module VectorMCP
       # @return [void]
       def handle_outgoing_response(message)
         request_id = message["id"]
-        logger.debug "[Stdio Transport] Received response for outgoing request ID #{request_id}"
+        # Received response for outgoing request
 
         @mutex.synchronize do
           # Store the response (convert keys to symbols for consistency)
@@ -272,7 +329,7 @@ module VectorMCP
           condition = @outgoing_request_conditions[request_id]
           if condition
             condition.signal
-            logger.debug "[Stdio Transport] Signaled condition for request ID #{request_id}"
+            # Signaled condition for request
           else
             logger.warn "[Stdio Transport] Received response for request ID #{request_id} but no thread is waiting"
           end
@@ -342,7 +399,7 @@ module VectorMCP
       # @return [void]
       def write_message(message)
         json_msg = message.to_json
-        logger.debug { "Sending stdio message: #{json_msg}" }
+        # Sending stdio message
 
         begin
           @output_mutex.synchronize do