vector_mcp 0.3.4 → 0.5.0

data/lib/vector_mcp/token_store.rb

@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+require "concurrent-ruby"
+require "securerandom"
+
+module VectorMCP
+  # Thread-safe bidirectional store mapping arbitrary string values to opaque
+  # tokens and back. The store has no knowledge of domain semantics: callers
+  # supply the prefix, and the store guarantees that the same (value, prefix)
+  # pair always yields the same token within its lifetime.
+  class TokenStore
+    # Regexp describing the token format emitted by {#tokenize}.
+    TOKEN_PATTERN = /\A[A-Z]+_[0-9A-F]{8}\z/
+
+    def initialize
+      @forward = Concurrent::Hash.new
+      @reverse = Concurrent::Hash.new
+      @mutex = Mutex.new
+    end
+
+    # Return an opaque token for +value+. Calling this repeatedly with the
+    # same +value+ and +prefix+ returns the same token.
+    #
+    # @param value [String] the value to tokenize.
+    # @param prefix [String] the token prefix (uppercase recommended).
+    # @return [String] a token of the form +"PREFIX_XXXXXXXX"+.
+    def tokenize(value, prefix:)
+      key = [prefix, value]
+      existing = @forward[key]
+      return existing if existing
+
+      @mutex.synchronize do
+        existing = @forward[key]
+        return existing if existing
+
+        token = generate_token(prefix)
+        # Populate the reverse map first so any thread that observes the
+        # token in @forward can always resolve it.
+        @reverse[token] = value
+        @forward[key] = token
+        token
+      end
+    end
+
+    # Resolve a token back to its original value.
+    #
+    # @param token [String] a token previously returned by {#tokenize}.
+    # @return [String, nil] the original value, or +nil+ if unknown.
+    def resolve(token)
+      @reverse[token]
+    end
+
+    # Predicate: does +string+ look like a token issued by this class?
+    # This check is purely structural and does not consult the store.
+    #
+    # @param string [Object] the value to test.
+    # @return [Boolean]
+    def token?(string)
+      string.is_a?(String) && TOKEN_PATTERN.match?(string)
+    end
+
+    # Remove all mappings. Intended for test teardown.
+    # @return [void]
+    def clear
+      @mutex.synchronize do
+        @forward.clear
+        @reverse.clear
+      end
+    end
+
+    private
+
+    def generate_token(prefix)
+      loop do
+        candidate = "#{prefix}_#{SecureRandom.hex(4).upcase}"
+        return candidate unless @reverse.key?(candidate)
+      end
+    end
+  end
+end

data/lib/vector_mcp/tool.rb

@@ -0,0 +1,221 @@
+# frozen_string_literal: true
+
+require "date"
+require "time"
+require_relative "errors"
+
+module VectorMCP
+  # Abstract base class for declarative tool definitions.
+  #
+  # Subclass this to define tools using a class-level DSL instead of
+  # the block-based +register_tool+ API. The two styles are fully
+  # interchangeable -- both produce the same +VectorMCP::Definitions::Tool+
+  # struct that the rest of the system consumes.
+  #
+  # @example
+  #   class ListProviders < VectorMCP::Tool
+  #     tool_name   "list_providers"
+  #     description "List providers filtered by category or status"
+  #
+  #     param :category, type: :string, desc: "Filter by category slug"
+  #     param :active,   type: :boolean, default: true
+  #
+  #     def call(args, session)
+  #       Provider.where(active: args.fetch("active", true))
+  #     end
+  #   end
+  #
+  #   server.register(ListProviders)
+  #
+  class Tool
+    # Maps Ruby symbol types to JSON Schema property fragments.
+    # Each value is merged into the generated property hash, so it may carry
+    # both +type+ and +format+ (or any other JSON Schema keyword).
+    TYPE_MAP = {
+      string: { "type" => "string" },
+      integer: { "type" => "integer" },
+      number: { "type" => "number" },
+      boolean: { "type" => "boolean" },
+      array: { "type" => "array" },
+      object: { "type" => "object" },
+      date: { "type" => "string", "format" => "date" },
+      datetime: { "type" => "string", "format" => "date-time" }
+    }.freeze
+
+    # Maps Ruby symbol types to a coercer lambda. Types not listed here
+    # pass their values through unchanged. Coercers receive the raw value
+    # (or nil) and return the coerced value. They must be total over the
+    # values JSON Schema validation would accept.
+    COERCERS = {
+      date: ->(v) { v.nil? || v.is_a?(Date) ? v : Date.parse(v.to_s) },
+      datetime: ->(v) { v.nil? || v.is_a?(Time) ? v : Time.parse(v.to_s) }
+    }.freeze
+
+    # Ensures each subclass gets its own +@params+ array so sibling
+    # classes do not share mutable state.
+    def self.inherited(subclass)
+      super
+      subclass.instance_variable_set(:@params, [])
+    end
+
+    # Sets or retrieves the tool name.
+    #
+    # When called with an argument, stores the name.
+    # When called without, returns the stored name or derives one from the class name.
+    #
+    # @param name [String, Symbol, nil] The tool name to set.
+    # @return [String] The tool name.
+    def self.tool_name(name = nil)
+      if name
+        @tool_name = name.to_s
+      else
+        @tool_name || derive_tool_name
+      end
+    end
+
+    # Sets or retrieves the tool description.
+    #
+    # @param text [String, nil] The description to set.
+    # @return [String, nil] The description.
+    def self.description(text = nil)
+      if text
+        @description = text
+      else
+        @description
+      end
+    end
+
+    # Declares a parameter for the tool's input schema.
+    #
+    # @param name [Symbol, String] The parameter name.
+    # @param type [Symbol] The parameter type (:string, :integer, :number, :boolean, :array, :object).
+    # @param desc [String, nil] A human-readable description.
+    # @param required [Boolean] Whether the parameter is required (default: false).
+    # @param options [Hash] Additional JSON Schema keywords (enum:, default:, format:, items:, etc.).
+    def self.param(name, type: :string, desc: nil, required: false, **options)
+      @params << {
+        name: name.to_s,
+        type: type,
+        desc: desc,
+        required: required,
+        options: options
+      }
+    end
+
+    # Builds a +VectorMCP::Definitions::Tool+ struct from the DSL metadata.
+    #
+    # @return [VectorMCP::Definitions::Tool]
+    # @raise [ArgumentError] If the subclass is missing a description or +#call+ method.
+    def self.to_definition
+      validate_tool_class!
+
+      VectorMCP::Definitions::Tool.new(
+        tool_name,
+        description,
+        build_input_schema,
+        build_handler
+      )
+    end
+
+    # The handler method that subclasses must implement.
+    #
+    # @param _args [Hash] The tool arguments (string keys).
+    # @param _session [VectorMCP::Session] The current session.
+    # @return [Object] The tool result.
+    def call(_args, _session)
+      raise NotImplementedError, "#{self.class.name} must implement #call(args, session)"
+    end
+
+    # Derives a snake_case tool name from the class name.
+    def self.derive_tool_name
+      base = name&.split("::")&.last
+      return "unnamed_tool" unless base
+
+      base.gsub(/([a-z\d])([A-Z])/, '\1_\2').downcase
+    end
+    private_class_method :derive_tool_name
+
+    # Validates that the subclass is properly configured.
+    def self.validate_tool_class!
+      raise ArgumentError, "#{name || self} must declare a description" unless description
+      return if method_defined?(:call) && instance_method(:call).owner != VectorMCP::Tool
+
+      raise ArgumentError, "#{name || self} must implement #call"
+    end
+    private_class_method :validate_tool_class!
+
+    # Builds a 2-arity handler lambda. A new instance is created per invocation.
+    # Arguments are coerced based on the declared param types before the
+    # handler sees them (e.g. :date param strings become Date objects).
+    def self.build_handler
+      klass = self
+      params = @params
+      ->(args, session) { klass.new.call(klass.coerce_args(args, params), session) }
+    end
+    private_class_method :build_handler
+
+    # Applies coercers to the raw argument hash. Returns a new hash; does
+    # not mutate the original. Keys without a coercible type pass through.
+    # Keys that are absent from +args+ stay absent — coercion only fires
+    # for keys actually present.
+    #
+    # A parse failure on a client-supplied value is translated into
+    # +VectorMCP::InvalidParamsError+ (JSON-RPC -32602) so the client sees
+    # a "bad request" response instead of a generic internal error.
+    # This is needed because the +json-schema+ gem does not enforce
+    # +format: date+ upstream (it does enforce +format: date-time+), so
+    # malformed +:date+ values would otherwise crash inside +Date.parse+.
+    def self.coerce_args(args, params)
+      coerced = args.dup
+      params.each do |param|
+        name = param[:name]
+        next unless coerced.key?(name)
+
+        coercer = COERCERS[param[:type]]
+        next unless coercer
+
+        begin
+          coerced[name] = coercer.call(coerced[name])
+        rescue ArgumentError, TypeError => e
+          # Date::Error < ArgumentError in Ruby 3.2+, so ArgumentError alone covers Date.parse failures.
+          raise VectorMCP::InvalidParamsError.new(
+            "Invalid #{param[:type]} value for param '#{name}': #{e.message}",
+            details: { param: name, type: param[:type], message: e.message }
+          )
+        end
+      end
+      coerced
+    end
+
+    # Builds a JSON Schema hash from the declared params.
+    def self.build_input_schema
+      properties = {}
+      required = []
+
+      @params.each do |param|
+        type_fragment = TYPE_MAP.fetch(param[:type]) do
+          raise ArgumentError, "Unknown param type :#{param[:type]} for param '#{param[:name]}' in #{name}. " \
+                               "Valid types: #{TYPE_MAP.keys.join(", ")}"
+        end
+
+        prop = type_fragment.dup
+        prop["description"] = param[:desc] if param[:desc]
+
+        param[:options].each do |key, value|
+          prop[key.to_s] = value
+        end
+
+        properties[param[:name]] = prop
+        required << param[:name] if param[:required]
+      end
+
+      schema = {
+        "type" => "object",
+        "properties" => properties
+      }
+      schema["required"] = required unless required.empty?
+      schema
+    end
+    private_class_method :build_input_schema
+  end
+end

data/lib/vector_mcp/transport/base_session_manager.rb

@@ -183,22 +183,6 @@ module VectorMCP
         end
       end
 
-      # Broadcasts a message to all sessions that support messaging.
-      #
-      # @param message [Hash] The message to broadcast
-      # @return [Integer] Number of sessions the message was sent to
-      def broadcast_message(message)
-        count = 0
-        @sessions.each_value do |session|
-          next unless can_send_message_to_session?(session)
-
-          count += 1 if message_sent_to_session?(session, message)
-        end
-
-        # Message broadcasted to recipients
-        count
-      end
-
       protected
 
       # Hook called when a session is created. Override in subclasses for transport-specific logic.
@@ -225,7 +209,7 @@ module VectorMCP
       end
 
       # Determines if this session manager should enable automatic cleanup.
-      # Override in subclasses that don't need automatic cleanup (e.g., stdio with single session).
+      # Override in subclasses that don't need automatic cleanup (e.g., transports with single persistent session).
       #
       # @return [Boolean] True if auto-cleanup should be enabled
       def auto_cleanup_enabled?

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

@@ -17,7 +17,7 @@ module VectorMCP
       # @api private
       class EventStore
         # Event data structure
-        Event = Struct.new(:id, :data, :type, :timestamp, :session_id) do
+        Event = Struct.new(:id, :data, :type, :timestamp, :session_id, :stream_id) do
           def to_sse_format
             lines = []
             lines << "id: #{id}"
@@ -36,7 +36,8 @@ module VectorMCP
         def initialize(max_events)
           @max_events = max_events
           @events = Concurrent::Array.new
-          @event_index = Concurrent::Hash.new # event_id -> index for fast lookup
+          @event_index = Concurrent::Hash.new # event_id -> logical position
+          @offset = 0 # number of events shifted off the front
           @current_sequence = Concurrent::AtomicFixnum.new(0)
         end
 
@@ -45,26 +46,23 @@ module VectorMCP
         # @param data [String] The event data
         # @param type [String] The event type (optional)
         # @param session_id [String, nil] The session ID to scope this event to
+        # @param stream_id [String, nil] The stream ID to scope this event to
         # @return [String] The generated event ID
-        def store_event(data, type = nil, session_id: nil)
+        def store_event(data, type = nil, session_id: nil, stream_id: nil)
           event_id = generate_event_id
           timestamp = Time.now
 
-          event = Event.new(event_id, data, type, timestamp, session_id)
+          event = Event.new(event_id, data, type, timestamp, session_id, stream_id)
 
-          # Add to events array
+          # Add to events array and record logical position
           @events.push(event)
-
-          # Update index
-          @event_index[event_id] = @events.length - 1
+          @event_index[event_id] = @offset + @events.length - 1
 
           # Maintain circular buffer
           if @events.length > @max_events
             removed_event = @events.shift
             @event_index.delete(removed_event.id)
-
-            # Update all indices after removal
-            @event_index.transform_values! { |index| index - 1 }
+            @offset += 1
           end
 
           event_id
@@ -74,24 +72,37 @@ module VectorMCP
         #
         # @param last_event_id [String] The last event ID received by client
         # @param session_id [String, nil] Filter events to this session only
+        # @param stream_id [String, nil] Filter events to this stream only
         # @return [Array<Event>] Array of events after the specified ID
-        def get_events_after(last_event_id, session_id: nil)
+        def get_events_after(last_event_id, session_id: nil, stream_id: nil)
           events = if last_event_id.nil?
                      @events.to_a
                    else
-                     last_index = @event_index[last_event_id]
-                     return [] if last_index.nil?
+                     logical = @event_index[last_event_id]
+                     return [] if logical.nil?
 
-                     start_index = last_index + 1
-                     return [] if start_index >= @events.length
+                     physical = logical - @offset + 1
+                     return [] if physical >= @events.length
 
-                     @events[start_index..]
+                     @events[physical..]
                    end
 
           events = events.select { |e| e.session_id == session_id } if session_id
+          events = events.select { |e| e.stream_id == stream_id } if stream_id
           events
         end
 
+        # Retrieves a specific event by ID.
+        #
+        # @param event_id [String] The event ID to look up
+        # @return [Event, nil] The stored event, or nil if it is no longer retained
+        def get_event(event_id)
+          logical = @event_index[event_id]
+          return nil if logical.nil?
+
+          @events[logical - @offset]
+        end
+
         # Gets the total number of stored events.
         #
         # @return [Integer] Number of events currently stored
@@ -127,6 +138,7 @@ module VectorMCP
         def clear
           @events.clear
           @event_index.clear
+          @offset = 0
         end
 
         # Gets statistics about the event store.

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

@@ -34,7 +34,7 @@ module VectorMCP
           end
 
           def streaming?
-            metadata[:streaming_connection] && !metadata[:streaming_connection].nil?
+            !streaming_connection.nil? || !streaming_connections.empty?
           end
 
           def streaming_connection
@@ -44,6 +44,25 @@ module VectorMCP
           def streaming_connection=(connection)
             metadata[:streaming_connection] = connection
           end
+
+          def streaming_connections
+            metadata[:streaming_connections] ||= Concurrent::Hash.new
+          end
+
+          def add_streaming_connection(connection)
+            streaming_connections[connection.stream_id] = connection
+            self.streaming_connection = connection
+          end
+
+          def remove_streaming_connection(connection = nil)
+            if connection
+              streaming_connections.delete(connection.stream_id)
+              self.streaming_connection = streaming_connections.values.first if streaming_connection == connection
+            else
+              streaming_connections.clear
+              self.streaming_connection = nil
+            end
+          end
         end
 
         # Initializes a new HTTP stream session manager.
@@ -51,24 +70,19 @@ module VectorMCP
         # @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
+        # Creates a new session. RequestContext.from_rack_env handles nil
+        # rack_env by falling back to a minimal context, so we don't need to
+        # branch here.
         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)
+          request_context = VectorMCP::RequestContext.from_rack_env(rack_env, "http_stream")
+          session_context = VectorMCP::Session.new(
+            @transport.server, @transport, id: session_id, request_context: request_context
+          )
 
+          session = Session.new(session_id, session_context, now, now, create_session_metadata)
           @sessions[session_id] = session
 
           logger.info { "Session created: #{session_id}" }
@@ -97,19 +111,6 @@ module VectorMCP
           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
@@ -131,17 +132,18 @@ module VectorMCP
         # @param connection [Object] The streaming connection object
         # @return [void]
         def set_streaming_connection(session, connection)
-          session.streaming_connection = connection
+          session.add_streaming_connection(connection)
           session.touch!
-          logger.debug { "Streaming connection associated: #{session.id}" }
+          logger.debug { "Streaming connection associated: #{session.id} (stream #{connection.stream_id})" }
         end
 
         # Removes streaming connection from a session.
         #
         # @param session [Session] The session to remove streaming from
+        # @param connection [Object, nil] The specific connection to remove, or nil to clear all
         # @return [void]
-        def remove_streaming_connection(session)
-          session.streaming_connection = nil
+        def remove_streaming_connection(session, connection = nil)
+          session.remove_streaming_connection(connection)
           session.touch!
           logger.debug { "Streaming connection removed: #{session.id}" }
         end
@@ -155,7 +157,7 @@ module VectorMCP
 
         # Override: Returns metadata for new HTTP stream sessions.
         def create_session_metadata
-          { streaming_connection: nil }
+          { streaming_connection: nil, streaming_connections: Concurrent::Hash.new }
         end
 
         # Override: Checks if a session can receive messages (has streaming connection).
@@ -175,15 +177,18 @@ module VectorMCP
         # @param session [Session] The session whose connection to close
         # @return [void]
         def close_streaming_connection(session)
-          return unless session&.streaming_connection
+          return unless session&.streaming?
+
+          connections = session.streaming_connections.values
+          connections = [session.streaming_connection] if connections.empty? && session.streaming_connection
 
-          begin
-            session.streaming_connection.close
+          connections.each do |connection|
+            connection.close
           rescue StandardError => e
             logger.warn { "Error closing streaming connection for #{session.id}: #{e.message}" }
           end
 
-          session.streaming_connection = nil
+          session.remove_streaming_connection
         end
       end
     end