vector_mcp 0.1.0 → 0.3.0

data/lib/vector_mcp/session.rb

@@ -1,5 +1,9 @@
 # frozen_string_literal: true
 
+require_relative "sampling/request"
+require_relative "sampling/result"
+require_relative "errors"
+
 module VectorMCP
   # Represents the state of a single client-server connection session in MCP.
   # It tracks initialization status, and negotiated capabilities between the client and server.
@@ -10,21 +14,23 @@ module VectorMCP
   # @attr_reader client_info [Hash, nil] Information about the client, received during initialization.
   # @attr_reader client_capabilities [Hash, nil] Capabilities supported by the client, received during initialization.
   class Session
-    attr_reader :server_info, :server_capabilities, :protocol_version, :client_info, :client_capabilities
+    attr_reader :server_info, :server_capabilities, :protocol_version, :client_info, :client_capabilities, :server, :transport, :id
+    attr_accessor :data # For user-defined session-specific storage
 
     # Initializes a new session.
     #
-    # @param server_info [Hash] Hash containing server information (e.g., name, version).
-    # @param server_capabilities [Hash] Hash describing server capabilities.
-    # @param protocol_version [String] The protocol version the server adheres to.
-    def initialize(server_info:, server_capabilities:, protocol_version:)
-      @server_info = server_info
-      @server_capabilities = server_capabilities
-      @protocol_version = protocol_version
-
-      @initialized = false
+    # @param server [VectorMCP::Server] The server instance managing this session.
+    # @param transport [VectorMCP::Transport::Base, nil] The transport handling this session. Required for sampling.
+    # @param id [String] A unique identifier for this session (e.g., from transport layer).
+    def initialize(server, transport = nil, id: SecureRandom.uuid)
+      @server = server
+      @transport = transport # Store the transport for sending requests
+      @id = id
+      @initialized_state = :pending # :pending, :succeeded, :failed
       @client_info = nil
       @client_capabilities = nil
+      @data = {} # Initialize user data hash
+      @logger = server.logger
     end
 
     # Marks the session as initialized using parameters from the client's `initialize` request.
@@ -33,35 +39,106 @@ module VectorMCP
     #   Expected keys include "protocolVersion", "clientInfo", and "capabilities".
     # @return [Hash] A hash suitable for the server's `initialize` response result.
     def initialize!(params)
-      client_protocol_version = params["protocolVersion"]
+      raise InitializationError, "Session already initialized or initialization attempt in progress." unless @initialized_state == :pending
 
-      if client_protocol_version != @protocol_version
-        # raise VectorMCP::ProtocolError.new("Unsupported protocol version: #{client_protocol_version}", code: -32603)
-        VectorMCP.logger.warn("Client requested protocol version '#{client_protocol_version}', server using '#{@protocol_version}'")
-      end
+      # TODO: More robust validation of params against MCP spec for initialize request
+      params["protocolVersion"]
+      client_capabilities_raw = params["capabilities"]
+      client_info_raw = params["clientInfo"]
+
+      # For now, we mostly care about clientInfo and capabilities for the session object.
+      # Protocol version matching is more of a server/transport concern at a lower level if strict checks are needed.
+      @client_info = client_info_raw.transform_keys(&:to_sym) if client_info_raw.is_a?(Hash)
+      @client_capabilities = client_capabilities_raw.transform_keys(&:to_sym) if client_capabilities_raw.is_a?(Hash)
 
-      @client_info = params["clientInfo"] || {}
-      @client_capabilities = params["capabilities"] || {}
-      @initialized = true
+      @initialized_state = :succeeded
+      @logger.info("[Session #{@id}] Initialized successfully. Client: #{@client_info&.dig(:name)}")
 
-      # Return the initialize result (will be sent by transport)
       {
-        protocolVersion: @protocol_version,
-        serverInfo: @server_info,
-        capabilities: @server_capabilities
+        protocolVersion: @server.protocol_version,
+        serverInfo: @server.server_info,
+        capabilities: @server.server_capabilities
       }
+    rescue StandardError => e
+      @initialized_state = :failed
+      @logger.error("[Session #{@id}] Initialization failed: #{e.message}")
+      # Re-raise as an InitializationError if it's not already one of our ProtocolErrors
+      raise e if e.is_a?(ProtocolError)
+
+      raise InitializationError, "Initialization processing error: #{e.message}", details: { original_error: e.to_s }
     end
 
     # Checks if the session has been successfully initialized.
     #
     # @return [Boolean] True if the session is initialized, false otherwise.
     def initialized?
-      @initialized
+      @initialized_state == :succeeded
     end
 
     # Helper to check client capabilities later if needed
     # def supports?(capability_key)
     #   @client_capabilities.key?(capability_key.to_s)
     # end
+
+    # --- MCP Sampling Method ---
+
+    # Initiates an MCP sampling request to the client associated with this session.
+    # This is a blocking call that waits for the client's response.
+    #
+    # @param request_params [Hash] Parameters for the `sampling/createMessage` request.
+    #   See `VectorMCP::Sampling::Request` for expected structure (e.g., :messages, :max_tokens).
+    # @param timeout [Numeric, nil] Optional timeout in seconds for this specific request.
+    #   Defaults to the transport's default request timeout.
+    # @return [VectorMCP::Sampling::Result] The result of the sampling operation.
+    # @raise [VectorMCP::SamplingError] if the sampling request fails, is rejected, or times out.
+    # @raise [StandardError] if the session's transport does not support `send_request`.
+    def sample(request_params, timeout: nil)
+      validate_sampling_preconditions
+
+      sampling_req_obj = VectorMCP::Sampling::Request.new(request_params)
+      @logger.info("[Session #{@id}] Sending sampling/createMessage request to client.")
+
+      send_sampling_request(sampling_req_obj, timeout)
+    end
+
+    private
+
+    # Validates that sampling can be performed on this session.
+    # @api private
+    # @raise [StandardError, InitializationError] if preconditions are not met.
+    def validate_sampling_preconditions
+      unless @transport.respond_to?(:send_request)
+        raise StandardError, "Session's transport does not support sending requests (required for sampling)."
+      end
+
+      return if initialized?
+
+      @logger.warn("[Session #{@id}] Attempted to send sampling request on a non-initialized session.")
+      raise InitializationError, "Cannot send sampling request: session not initialized."
+    end
+
+    # Sends the sampling request and handles the response.
+    # @api private
+    # @param sampling_req_obj [VectorMCP::Sampling::Request] The sampling request object.
+    # @param timeout [Numeric, nil] Optional timeout for the request.
+    # @return [VectorMCP::Sampling::Result] The sampling result.
+    # @raise [VectorMCP::SamplingError] if the request fails.
+    def send_sampling_request(sampling_req_obj, timeout)
+      send_request_args = ["sampling/createMessage", sampling_req_obj.to_h]
+      send_request_kwargs = {}
+      send_request_kwargs[:timeout] = timeout if timeout
+
+      raw_result = @transport.send_request(*send_request_args, **send_request_kwargs)
+      VectorMCP::Sampling::Result.new(raw_result)
+    rescue ArgumentError => e
+      @logger.error("[Session #{@id}] Invalid parameters for sampling request or result: #{e.message}")
+      raise VectorMCP::SamplingError, "Invalid sampling parameters or malformed client response: #{e.message}", details: { original_error: e.to_s }
+    rescue VectorMCP::SamplingError => e
+      @logger.warn("[Session #{@id}] Sampling request failed: #{e.message}")
+      raise e
+    rescue StandardError => e
+      @logger.error("[Session #{@id}] Unexpected error during sampling: #{e.class.name}: #{e.message}")
+      raise VectorMCP::SamplingError, "An unexpected error occurred during sampling: #{e.message}", details: { original_error: e.to_s }
+    end
   end
 end

data/lib/vector_mcp/transport/sse/client_connection.rb

@@ -0,0 +1,113 @@
+# frozen_string_literal: true
+
+module VectorMCP
+  module Transport
+    class SSE
+      # Manages individual client connection state for SSE transport.
+      # Each client connection has a unique session ID, message queue, and streaming thread.
+      class ClientConnection
+        attr_reader :session_id, :message_queue, :logger
+        attr_accessor :stream_thread, :stream_io
+
+        # Initializes a new client connection.
+        #
+        # @param session_id [String] Unique identifier for this client session
+        # @param logger [Logger] Logger instance for debugging and error reporting
+        def initialize(session_id, logger)
+          @session_id = session_id
+          @logger = logger
+          @message_queue = Queue.new
+          @stream_thread = nil
+          @stream_io = nil
+          @closed = false
+          @mutex = Mutex.new
+
+          logger.debug { "Client connection created: #{session_id}" }
+        end
+
+        # Checks if the connection is closed
+        #
+        # @return [Boolean] true if connection is closed
+        def closed?
+          @mutex.synchronize { @closed }
+        end
+
+        # Closes the client connection and cleans up resources.
+        # This method is thread-safe and can be called multiple times.
+        def close
+          @mutex.synchronize do
+            return if @closed
+
+            @closed = true
+            logger.debug { "Closing client connection: #{session_id}" }
+
+            # Close the message queue to signal streaming thread to stop
+            @message_queue.close if @message_queue.respond_to?(:close)
+
+            # Close the stream I/O if it exists
+            begin
+              @stream_io&.close
+            rescue StandardError => e
+              logger.warn { "Error closing stream I/O for #{session_id}: #{e.message}" }
+            end
+
+            # Stop the streaming thread
+            if @stream_thread&.alive?
+              @stream_thread.kill
+              @stream_thread.join(1) # Wait up to 1 second for clean shutdown
+            end
+
+            logger.debug { "Client connection closed: #{session_id}" }
+          end
+        end
+
+        # Enqueues a message to be sent to this client.
+        # This method is thread-safe.
+        #
+        # @param message [Hash] The JSON-RPC message to send
+        # @return [Boolean] true if message was enqueued successfully
+        def enqueue_message(message)
+          return false if closed?
+
+          begin
+            @message_queue.push(message)
+            logger.debug { "Message enqueued for client #{session_id}: #{message.inspect}" }
+            true
+          rescue ClosedQueueError
+            logger.warn { "Attempted to enqueue message to closed queue for client #{session_id}" }
+            false
+          rescue StandardError => e
+            logger.error { "Error enqueuing message for client #{session_id}: #{e.message}" }
+            false
+          end
+        end
+
+        # Dequeues the next message from the client's message queue.
+        # This method blocks until a message is available or the queue is closed.
+        #
+        # @return [Hash, nil] The next message, or nil if queue is closed
+        def dequeue_message
+          return nil if closed?
+
+          begin
+            @message_queue.pop
+          rescue ClosedQueueError
+            nil
+          rescue StandardError => e
+            logger.error { "Error dequeuing message for client #{session_id}: #{e.message}" }
+            nil
+          end
+        end
+
+        # Gets the current queue size
+        #
+        # @return [Integer] Number of messages waiting in the queue
+        def queue_size
+          @message_queue.size
+        rescue StandardError
+          0
+        end
+      end
+    end
+  end
+end

data/lib/vector_mcp/transport/sse/message_handler.rb

@@ -0,0 +1,166 @@
+# frozen_string_literal: true
+
+require "json"
+
+module VectorMCP
+  module Transport
+    class SSE
+      # Handles JSON-RPC message processing for POST requests.
+      # Processes incoming messages and sends responses via SSE streams.
+      class MessageHandler
+        # Initializes a new message handler.
+        #
+        # @param server [VectorMCP::Server] The MCP server instance
+        # @param session [VectorMCP::Session] The server session
+        # @param logger [Logger] Logger instance for debugging
+        def initialize(server, session, logger)
+          @server = server
+          @session = session
+          @logger = logger
+        end
+
+        # Handles a POST message request from a client.
+        #
+        # @param env [Hash] Rack environment hash
+        # @param client_conn [ClientConnection] The client connection
+        # @return [Array] Rack response triplet
+        def handle_post_message(env, client_conn)
+          request_body = read_request_body(env)
+          return error_response(nil, -32_600, "Request body is empty") if request_body.nil? || request_body.empty?
+
+          message = parse_json_message(request_body, client_conn)
+          return message if message.is_a?(Array) # Error response
+
+          process_message(message, client_conn)
+        rescue VectorMCP::ProtocolError => e
+          @logger.error { "Protocol error for client #{client_conn.session_id}: #{e.message}" }
+          request_id = e.request_id || message&.dig("id")
+          enqueue_error_response(client_conn, request_id, e.code, e.message, e.details)
+          error_response(request_id, e.code, e.message, e.details)
+        rescue StandardError => e
+          @logger.error { "Unexpected error for client #{client_conn.session_id}: #{e.message}\n#{e.backtrace.join("\n")}" }
+          request_id = message&.dig("id")
+          enqueue_error_response(client_conn, request_id, -32_603, "Internal server error")
+          error_response(request_id, -32_603, "Internal server error")
+        end
+
+        private
+
+        # Reads the request body from the Rack environment.
+        #
+        # @param env [Hash] Rack environment
+        # @return [String, nil] Request body as string
+        def read_request_body(env)
+          input = env["rack.input"]
+          return nil unless input
+
+          body = input.read
+          input.rewind if input.respond_to?(:rewind)
+          body
+        end
+
+        # Parses JSON message from request body.
+        #
+        # @param body_str [String] JSON string from request body
+        # @param client_conn [ClientConnection] Client connection for error handling
+        # @return [Hash, Array] Parsed message or error response triplet
+        def parse_json_message(body_str, client_conn)
+          JSON.parse(body_str)
+        rescue JSON::ParserError => e
+          @logger.error { "JSON parse error for client #{client_conn.session_id}: #{e.message}" }
+          malformed_id = VectorMCP::Util.extract_id_from_invalid_json(body_str)
+          enqueue_error_response(client_conn, malformed_id, -32_700, "Parse error")
+          error_response(malformed_id, -32_700, "Parse error")
+        end
+
+        # Processes a valid JSON-RPC message.
+        #
+        # @param message [Hash] Parsed JSON-RPC message
+        # @param client_conn [ClientConnection] Client connection
+        # @return [Array] Rack response triplet
+        def process_message(message, client_conn)
+          # Handle the message through the server
+          response_data = @server.handle_message(message, @session, client_conn.session_id)
+
+          # If it's a request (has id), send response via SSE
+          if message["id"]
+            enqueue_success_response(client_conn, message["id"], response_data)
+          else
+            @logger.debug { "Processed notification for client #{client_conn.session_id}" }
+          end
+
+          # Always return 202 Accepted for valid POST messages
+          success_response(message["id"])
+        end
+
+        # Enqueues a successful response to the client's SSE stream.
+        #
+        # @param client_conn [ClientConnection] Client connection
+        # @param request_id [String, Integer] Original request ID
+        # @param result [Object] Response result data
+        def enqueue_success_response(client_conn, request_id, result)
+          response = {
+            jsonrpc: "2.0",
+            id: request_id,
+            result: result
+          }
+          StreamManager.enqueue_message(client_conn, response)
+        end
+
+        # Enqueues an error response to the client's SSE stream.
+        #
+        # @param client_conn [ClientConnection] Client connection
+        # @param request_id [String, Integer, nil] Original request ID
+        # @param code [Integer] Error code
+        # @param message [String] Error message
+        # @param data [Object, nil] Additional error data
+        def enqueue_error_response(client_conn, request_id, code, message, data = nil)
+          error_payload = { code: code, message: message }
+          error_payload[:data] = data if data
+
+          error_response = {
+            jsonrpc: "2.0",
+            id: request_id,
+            error: error_payload
+          }
+          StreamManager.enqueue_message(client_conn, error_response)
+        end
+
+        # Creates a successful HTTP response for the POST request.
+        #
+        # @param request_id [String, Integer, nil] Request ID
+        # @return [Array] Rack response triplet
+        def success_response(request_id)
+          body = { status: "accepted", id: request_id }.to_json
+          [202, { "Content-Type" => "application/json" }, [body]]
+        end
+
+        # Creates an error HTTP response for the POST request.
+        #
+        # @param id [String, Integer, nil] Request ID
+        # @param code [Integer] Error code
+        # @param message [String] Error message
+        # @param data [Object, nil] Additional error data
+        # @return [Array] Rack response triplet
+        def error_response(id, code, message, data = nil)
+          status = case code
+                   when -32_700, -32_600, -32_602 then 400 # Parse, Invalid Request, Invalid Params
+                   when -32_601, -32_001 then 404 # Method Not Found, Not Found
+                   else 500 # Internal Error, Server Error
+                   end
+
+          error_payload = { code: code, message: message }
+          error_payload[:data] = data if data
+
+          body = {
+            jsonrpc: "2.0",
+            id: id,
+            error: error_payload
+          }.to_json
+
+          [status, { "Content-Type" => "application/json" }, [body]]
+        end
+      end
+    end
+  end
+end

data/lib/vector_mcp/transport/sse/puma_config.rb

@@ -0,0 +1,77 @@
+# frozen_string_literal: true
+
+module VectorMCP
+  module Transport
+    class SSE
+      # Configures Puma server for production-ready SSE transport.
+      # Handles server setup, threading, and resource management.
+      class PumaConfig
+        attr_reader :host, :port, :logger
+
+        # Initializes Puma configuration.
+        #
+        # @param host [String] Host to bind to
+        # @param port [Integer] Port to listen on
+        # @param logger [Logger] Logger instance
+        def initialize(host, port, logger)
+          @host = host
+          @port = port
+          @logger = logger
+        end
+
+        # Configures a Puma server instance.
+        #
+        # @param server [Puma::Server] The Puma server to configure
+        def configure(server)
+          server.add_tcp_listener(host, port)
+
+          # Configure threading for production use
+          configure_threading(server)
+
+          # Set up server options
+          configure_server_options(server)
+
+          logger.debug { "Puma server configured for #{host}:#{port}" }
+        end
+
+        private
+
+        # Configures threading parameters for optimal performance.
+        #
+        # @param server [Puma::Server] The Puma server
+        def configure_threading(server)
+          # Set thread pool size based on CPU cores and expected load
+          min_threads = 2
+          max_threads = [4, Etc.nprocessors * 2].max
+
+          # Puma 6.x does not expose min_threads= and max_threads= as public API.
+          # Thread pool sizing should be set via Puma DSL/config before server creation.
+          # For legacy compatibility, set if possible, otherwise log a warning.
+          if server.respond_to?(:min_threads=) && server.respond_to?(:max_threads=)
+            server.min_threads = min_threads
+            server.max_threads = max_threads
+            logger.debug { "Puma configured with #{min_threads}-#{max_threads} threads" }
+          else
+            logger.warn { "Puma::Server does not support direct thread pool sizing; set threads via Puma config DSL before server creation." }
+          end
+        end
+
+        # Configures server-specific options.
+        #
+        # @param server [Puma::Server] The Puma server
+        def configure_server_options(server)
+          # Set server-specific options for SSE handling
+          server.leak_stack_on_error = false if server.respond_to?(:leak_stack_on_error=)
+
+          # Configure timeouts appropriate for SSE connections
+          # SSE connections should be long-lived, so we set generous timeouts
+          if server.respond_to?(:first_data_timeout=)
+            server.first_data_timeout = 30 # 30 seconds to send first data
+          end
+
+          logger.debug { "Puma server options configured for SSE transport" }
+        end
+      end
+    end
+  end
+end

data/lib/vector_mcp/transport/sse/stream_manager.rb

@@ -0,0 +1,92 @@
+# frozen_string_literal: true
+
+module VectorMCP
+  module Transport
+    class SSE
+      # Manages Server-Sent Events streaming for client connections.
+      # Handles creation of streaming responses and message broadcasting.
+      class StreamManager
+        class << self
+          # Creates an SSE streaming response body for a client connection.
+          #
+          # @param client_conn [ClientConnection] The client connection to stream to
+          # @param endpoint_url [String] The URL for the client to POST messages to
+          # @param logger [Logger] Logger instance for debugging
+          # @return [Enumerator] Rack-compatible streaming response body
+          def create_sse_stream(client_conn, endpoint_url, logger)
+            Enumerator.new do |yielder|
+              # Send initial endpoint event
+              yielder << format_sse_event("endpoint", endpoint_url)
+              logger.debug { "Sent endpoint event to client #{client_conn.session_id}: #{endpoint_url}" }
+
+              # Start streaming thread for this client
+              client_conn.stream_thread = Thread.new do
+                stream_messages_to_client(client_conn, yielder, logger)
+              end
+
+              # Keep the connection alive by yielding from the streaming thread
+              client_conn.stream_thread.join
+            rescue StandardError => e
+              logger.error { "Error in SSE stream for client #{client_conn.session_id}: #{e.message}\n#{e.backtrace.join("\n")}" }
+            ensure
+              logger.debug { "SSE stream ended for client #{client_conn.session_id}" }
+              client_conn.close
+            end
+          end
+
+          # Enqueues a message to a specific client connection.
+          #
+          # @param client_conn [ClientConnection] The target client connection
+          # @param message [Hash] The JSON-RPC message to send
+          # @return [Boolean] true if message was enqueued successfully
+          def enqueue_message(client_conn, message)
+            return false unless client_conn && !client_conn.closed?
+
+            client_conn.enqueue_message(message)
+          end
+
+          private
+
+          # Streams messages from a client's queue to the SSE connection.
+          # This method runs in a dedicated thread per client.
+          #
+          # @param client_conn [ClientConnection] The client connection
+          # @param yielder [Enumerator::Yielder] The response yielder
+          # @param logger [Logger] Logger instance
+          def stream_messages_to_client(client_conn, yielder, logger)
+            logger.debug { "Starting message streaming thread for client #{client_conn.session_id}" }
+
+            loop do
+              message = client_conn.dequeue_message
+              break if message.nil? # Queue closed or connection closed
+
+              begin
+                json_message = message.to_json
+                sse_data = format_sse_event("message", json_message)
+                yielder << sse_data
+
+                logger.debug { "Streamed message to client #{client_conn.session_id}: #{json_message}" }
+              rescue StandardError => e
+                logger.error { "Error streaming message to client #{client_conn.session_id}: #{e.message}" }
+                break
+              end
+            end
+
+            logger.debug { "Message streaming thread ended for client #{client_conn.session_id}" }
+          rescue StandardError => e
+            logger.error { "Fatal error in streaming thread for client #{client_conn.session_id}: #{e.message}\n#{e.backtrace.join("\n")}" }
+          end
+
+          # Formats data as a Server-Sent Event.
+          #
+          # @param event [String] The event type
+          # @param data [String] The event data
+          # @return [String] Properly formatted SSE event
+          def format_sse_event(event, data)
+            "event: #{event}\ndata: #{data}\n\n"
+          end
+        end
+      end
+    end
+  end
+end