model-context-protocol-rb 0.3.4 → 0.5.0

data/lib/model_context_protocol/server/registry.rb

@@ -39,8 +39,8 @@ module ModelContextProtocol
     end
 
     def register(klass)
-      metadata = klass.metadata
-      entry = {klass: klass}.merge(metadata)
+      definition = klass.definition
+      entry = {klass: klass}.merge(definition)
 
       case klass.ancestors
       when ->(ancestors) { ancestors.include?(ModelContextProtocol::Server::Prompt) }
@@ -74,45 +74,121 @@ module ModelContextProtocol
       find_by_name(@tools, name)
     end
 
-    def prompts_data
-      PromptsData[prompts: @prompts.map { |entry| entry.except(:klass) }]
+    def prompts_data(cursor: nil, page_size: nil, cursor_ttl: nil)
+      items = @prompts.map { |entry| entry.except(:klass) }
+
+      if cursor || page_size
+        paginated = Server::Pagination.paginate(
+          items,
+          cursor: cursor,
+          page_size: page_size || 100,
+          cursor_ttl: cursor_ttl
+        )
+
+        PromptsData[prompts: paginated.items, next_cursor: paginated.next_cursor]
+      else
+        PromptsData[prompts: items]
+      end
     end
 
-    def resources_data
-      ResourcesData[resources: @resources.map { |entry| entry.except(:klass) }]
+    def resources_data(cursor: nil, page_size: nil, cursor_ttl: nil)
+      items = @resources.map { |entry| entry.except(:klass) }
+
+      if cursor || page_size
+        paginated = Server::Pagination.paginate(
+          items,
+          cursor: cursor,
+          page_size: page_size || 100,
+          cursor_ttl: cursor_ttl
+        )
+
+        ResourcesData[resources: paginated.items, next_cursor: paginated.next_cursor]
+      else
+        ResourcesData[resources: items]
+      end
     end
 
-    def resource_templates_data
-      ResourceTemplatesData[resource_templates: @resource_templates.map { |entry| entry.except(:klass, :completions) }]
+    def resource_templates_data(cursor: nil, page_size: nil, cursor_ttl: nil)
+      items = @resource_templates.map { |entry| entry.except(:klass, :completions) }
+
+      if cursor || page_size
+        paginated = Server::Pagination.paginate(
+          items,
+          cursor: cursor,
+          page_size: page_size || 100,
+          cursor_ttl: cursor_ttl
+        )
+
+        ResourceTemplatesData[resource_templates: paginated.items, next_cursor: paginated.next_cursor]
+      else
+        ResourceTemplatesData[resource_templates: items]
+      end
     end
 
-    def tools_data
-      ToolsData[tools: @tools.map { |entry| entry.except(:klass) }]
+    def tools_data(cursor: nil, page_size: nil, cursor_ttl: nil)
+      items = @tools.map { |entry| entry.except(:klass) }
+
+      if cursor || page_size
+        paginated = Server::Pagination.paginate(
+          items,
+          cursor: cursor,
+          page_size: page_size || 100,
+          cursor_ttl: cursor_ttl
+        )
+
+        ToolsData[tools: paginated.items, next_cursor: paginated.next_cursor]
+      else
+        ToolsData[tools: items]
+      end
     end
 
     private
 
-    PromptsData = Data.define(:prompts) do
+    PromptsData = Data.define(:prompts, :next_cursor) do
+      def initialize(prompts:, next_cursor: nil)
+        super
+      end
+
       def serialized
-        {prompts:}
+        result = {prompts:}
+        result[:nextCursor] = next_cursor if next_cursor
+        result
       end
     end
 
-    ResourcesData = Data.define(:resources) do
+    ResourcesData = Data.define(:resources, :next_cursor) do
+      def initialize(resources:, next_cursor: nil)
+        super
+      end
+
       def serialized
-        {resources:}
+        result = {resources:}
+        result[:nextCursor] = next_cursor if next_cursor
+        result
       end
     end
 
-    ResourceTemplatesData = Data.define(:resource_templates) do
+    ResourceTemplatesData = Data.define(:resource_templates, :next_cursor) do
+      def initialize(resource_templates:, next_cursor: nil)
+        super
+      end
+
       def serialized
-        {resourceTemplates: resource_templates}
+        result = {resourceTemplates: resource_templates}
+        result[:nextCursor] = next_cursor if next_cursor
+        result
       end
     end
 
-    ToolsData = Data.define(:tools) do
+    ToolsData = Data.define(:tools, :next_cursor) do
+      def initialize(tools:, next_cursor: nil)
+        super
+      end
+
       def serialized
-        {tools:}
+        result = {tools:}
+        result[:nextCursor] = next_cursor if next_cursor
+        result
       end
     end
 

data/lib/model_context_protocol/server/resource.rb

@@ -1,12 +1,13 @@
 module ModelContextProtocol
   class Server::Resource
-    attr_reader :mime_type, :uri, :context, :logger
+    include ModelContextProtocol::Server::Cancellable
+    include ModelContextProtocol::Server::Progressable
 
-    def initialize(logger, context = {})
+    attr_reader :mime_type, :uri
+
+    def initialize
       @mime_type = self.class.mime_type
       @uri = self.class.uri
-      @context = context
-      @logger = logger
     end
 
     def call
@@ -15,59 +16,76 @@ module ModelContextProtocol
 
     TextResponse = Data.define(:resource, :text) do
       def serialized
-        {contents: [{mimeType: resource.mime_type, text:, uri: resource.uri}]}
+        content = {mimeType: resource.mime_type, text:, uri: resource.uri}
+        content[:title] = resource.class.title if resource.class.title
+        annotations = resource.class.annotations&.serialized
+        content[:annotations] = annotations if annotations
+        {contents: [content]}
       end
     end
     private_constant :TextResponse
 
     BinaryResponse = Data.define(:blob, :resource) do
       def serialized
-        {contents: [{blob:, mimeType: resource.mime_type, uri: resource.uri}]}
+        content = {blob:, mimeType: resource.mime_type, uri: resource.uri}
+        content[:title] = resource.class.title if resource.class.title
+        annotations = resource.class.annotations&.serialized
+        content[:annotations] = annotations if annotations
+        {contents: [content]}
       end
     end
     private_constant :BinaryResponse
 
-    private def respond_with(type, **options)
-      case [type, options]
-      in [:text, {text:}]
+    private def respond_with(**kwargs)
+      case [kwargs]
+      in [{text:}]
         TextResponse[resource: self, text:]
-      in [:binary, {blob:}]
-        BinaryResponse[blob:, resource: self]
+      in [{binary:}]
+        BinaryResponse[blob: binary, resource: self]
       else
-        raise ModelContextProtocol::Server::ResponseArgumentsError, "Invalid arguments: #{type}, #{options}"
+        raise ModelContextProtocol::Server::ResponseArgumentsError, "Invalid arguments: #{options}"
       end
     end
 
     class << self
-      attr_reader :name, :description, :mime_type, :uri
+      attr_reader :name, :description, :title, :mime_type, :uri, :annotations
 
-      def with_metadata(&block)
-        metadata_dsl = MetadataDSL.new
-        metadata_dsl.instance_eval(&block)
+      def define(&block)
+        definition_dsl = DefinitionDSL.new
+        definition_dsl.instance_eval(&block)
 
-        @name = metadata_dsl.name
-        @description = metadata_dsl.description
-        @mime_type = metadata_dsl.mime_type
-        @uri = metadata_dsl.uri
+        @name = definition_dsl.name
+        @description = definition_dsl.description
+        @title = definition_dsl.title
+        @mime_type = definition_dsl.mime_type
+        @uri = definition_dsl.uri
+        @annotations = definition_dsl.defined_annotations
       end
 
       def inherited(subclass)
         subclass.instance_variable_set(:@name, @name)
         subclass.instance_variable_set(:@description, @description)
+        subclass.instance_variable_set(:@title, @title)
         subclass.instance_variable_set(:@mime_type, @mime_type)
         subclass.instance_variable_set(:@uri, @uri)
+        subclass.instance_variable_set(:@annotations, @annotations&.dup)
       end
 
-      def call(logger, context = {})
-        new(logger, context).call
+      def call
+        new.call
       end
 
-      def metadata
-        {name: @name, description: @description, mimeType: @mime_type, uri: @uri}
+      def definition
+        result = {name: @name, description: @description, mimeType: @mime_type, uri: @uri}
+        result[:title] = @title if @title
+        result[:annotations] = @annotations.serialized if @annotations
+        result
       end
     end
 
-    class MetadataDSL
+    class DefinitionDSL
+      attr_reader :defined_annotations
+
       def name(value = nil)
         @name = value if value
         @name
@@ -78,6 +96,11 @@ module ModelContextProtocol
         @description
       end
 
+      def title(value = nil)
+        @title = value if value
+        @title
+      end
+
       def mime_type(value = nil)
         @mime_type = value if value
         @mime_type
@@ -87,6 +110,56 @@ module ModelContextProtocol
         @uri = value if value
         @uri
       end
+
+      def annotations(&block)
+        @defined_annotations = AnnotationsDSL.new
+        @defined_annotations.instance_eval(&block)
+        @defined_annotations
+      end
+    end
+
+    class AnnotationsDSL
+      VALID_AUDIENCE_VALUES = [:user, :assistant].freeze
+
+      def initialize
+        @audience = nil
+        @priority = nil
+        @last_modified = nil
+      end
+
+      def audience(value)
+        normalized_value = Array(value).map(&:to_sym)
+        invalid_values = normalized_value - VALID_AUDIENCE_VALUES
+        unless invalid_values.empty?
+          raise ArgumentError, "Invalid audience values: #{invalid_values.join(", ")}. Valid values are: #{VALID_AUDIENCE_VALUES.join(", ")}"
+        end
+        @audience = normalized_value.map(&:to_s)
+      end
+
+      def priority(value)
+        unless value.is_a?(Numeric) && value >= 0.0 && value <= 1.0
+          raise ArgumentError, "Priority must be a number between 0.0 and 1.0, got: #{value}"
+        end
+        @priority = value.to_f
+      end
+
+      def last_modified(value)
+        # Validate ISO 8601 format
+        begin
+          Time.iso8601(value)
+        rescue ArgumentError
+          raise ArgumentError, "lastModified must be in ISO 8601 format (e.g., '2025-01-12T15:00:58Z'), got: #{value}"
+        end
+        @last_modified = value
+      end
+
+      def serialized
+        result = {}
+        result[:audience] = @audience if @audience
+        result[:priority] = @priority if @priority
+        result[:lastModified] = @last_modified if @last_modified
+        result.empty? ? nil : result
+      end
     end
   end
 end

data/lib/model_context_protocol/server/resource_template.rb

@@ -3,15 +3,15 @@ module ModelContextProtocol
     class << self
       attr_reader :name, :description, :mime_type, :uri_template, :completions
 
-      def with_metadata(&block)
-        metadata_dsl = MetadataDSL.new
-        metadata_dsl.instance_eval(&block)
-
-        @name = metadata_dsl.name
-        @description = metadata_dsl.description
-        @mime_type = metadata_dsl.mime_type
-        @uri_template = metadata_dsl.uri_template
-        @completions = metadata_dsl.completions
+      def define(&block)
+        definition_dsl = DefinitionDSL.new
+        definition_dsl.instance_eval(&block)
+
+        @name = definition_dsl.name
+        @description = definition_dsl.description
+        @mime_type = definition_dsl.mime_type
+        @uri_template = definition_dsl.uri_template
+        @completions = definition_dsl.completions
       end
 
       def inherited(subclass)
@@ -32,7 +32,7 @@ module ModelContextProtocol
         completion.call(param_name.to_s, value)
       end
 
-      def metadata
+      def definition
         {
           name: @name,
           description: @description,
@@ -43,7 +43,7 @@ module ModelContextProtocol
       end
     end
 
-    class MetadataDSL
+    class DefinitionDSL
       attr_reader :completions
 
       def initialize
@@ -85,8 +85,21 @@ module ModelContextProtocol
         @completions = {}
       end
 
-      def completion(param_name, completion_class)
-        @completions[param_name.to_s] = completion_class
+      def completion(param_name, completion_class_or_values)
+        @completions[param_name.to_s] = if completion_class_or_values.is_a?(Array)
+          create_array_completion(completion_class_or_values)
+        else
+          completion_class_or_values
+        end
+      end
+
+      private
+
+      def create_array_completion(values)
+        ModelContextProtocol::Server::Completion.define do
+          filtered_values = values.grep(/#{argument_value}/)
+          respond_with values: filtered_values
+        end
       end
     end
   end

data/lib/model_context_protocol/server/router.rb

@@ -1,3 +1,5 @@
+require_relative "cancellable"
+
 module ModelContextProtocol
   class Server::Router
     # Raised when an invalid method is provided.
@@ -12,14 +14,45 @@ module ModelContextProtocol
       @handlers[method] = handler
     end
 
-    def route(message)
+    # Route a message to its handler with request tracking support
+    #
+    # @param message [Hash] the JSON-RPC message
+    # @param request_store [Object] the request store for tracking cancellation
+    # @param session_id [String, nil] the session ID for HTTP transport
+    # @param transport [Object, nil] the transport for sending notifications
+    # @return [Object] the handler result, or nil if cancelled
+    def route(message, request_store: nil, session_id: nil, transport: nil)
       method = message["method"]
       handler = @handlers[method]
       raise MethodNotFoundError, "Method not found: #{method}" unless handler
 
-      with_environment(@configuration&.environment_variables) do
-        handler.call(message)
+      request_id = message["id"]
+      progress_token = message.dig("params", "_meta", "progressToken")
+
+      if request_id && request_store
+        request_store.register_request(request_id, session_id)
+      end
+
+      result = nil
+      begin
+        with_environment(@configuration&.environment_variables) do
+          context = {request_id:, request_store:, session_id:, progress_token:, transport:}
+
+          Thread.current[:mcp_context] = context
+
+          result = handler.call(message)
+        end
+      rescue Server::Cancellable::CancellationError
+        return nil
+      ensure
+        if request_id && request_store
+          request_store.unregister_request(request_id)
+        end
+
+        Thread.current[:mcp_context] = nil
       end
+
+      result
     end
 
     private

data/lib/model_context_protocol/server/stdio_transport/request_store.rb

@@ -0,0 +1,102 @@
+module ModelContextProtocol
+  class Server::StdioTransport
+    # Thread-safe in-memory storage for tracking active requests and their cancellation status.
+    # This store is used by StdioTransport to manage request lifecycle and handle cancellation.
+    class RequestStore
+      def initialize
+        @mutex = Mutex.new
+        @requests = {}
+      end
+
+      # Register a new request with its associated thread
+      #
+      # @param request_id [String] the unique request identifier
+      # @param thread [Thread] the thread processing this request (defaults to current thread)
+      # @return [void]
+      def register_request(request_id, thread = Thread.current)
+        @mutex.synchronize do
+          @requests[request_id] = {
+            thread:,
+            cancelled: false,
+            started_at: Time.now
+          }
+        end
+      end
+
+      # Mark a request as cancelled
+      #
+      # @param request_id [String] the unique request identifier
+      # @return [Boolean] true if request was found and marked cancelled, false otherwise
+      def mark_cancelled(request_id)
+        @mutex.synchronize do
+          if (request = @requests[request_id])
+            request[:cancelled] = true
+            return true
+          end
+          false
+        end
+      end
+
+      # Check if a request has been cancelled
+      #
+      # @param request_id [String] the unique request identifier
+      # @return [Boolean] true if the request is cancelled, false otherwise
+      def cancelled?(request_id)
+        @mutex.synchronize do
+          @requests[request_id]&.fetch(:cancelled, false) || false
+        end
+      end
+
+      # Unregister a request (typically called when request completes)
+      #
+      # @param request_id [String] the unique request identifier
+      # @return [Hash, nil] the removed request data, or nil if not found
+      def unregister_request(request_id)
+        @mutex.synchronize do
+          @requests.delete(request_id)
+        end
+      end
+
+      # Get information about a specific request
+      #
+      # @param request_id [String] the unique request identifier
+      # @return [Hash, nil] request information or nil if not found
+      def get_request(request_id)
+        @mutex.synchronize do
+          @requests[request_id]&.dup
+        end
+      end
+
+      # Get all active request IDs
+      #
+      # @return [Array<String>] list of active request IDs
+      def active_requests
+        @mutex.synchronize do
+          @requests.keys.dup
+        end
+      end
+
+      # Clean up old requests (useful for preventing memory leaks)
+      #
+      # @param max_age_seconds [Integer] maximum age of requests to keep
+      # @return [Array<String>] list of cleaned up request IDs
+      def cleanup_old_requests(max_age_seconds = 300)
+        cutoff_time = Time.now - max_age_seconds
+        removed_ids = []
+
+        @mutex.synchronize do
+          @requests.delete_if do |request_id, data|
+            if data[:started_at] < cutoff_time
+              removed_ids << request_id
+              true
+            else
+              false
+            end
+          end
+        end
+
+        removed_ids
+      end
+    end
+  end
+end

data/lib/model_context_protocol/server/stdio_transport.rb

@@ -1,3 +1,5 @@
+require_relative "stdio_transport/request_store"
+
 module ModelContextProtocol
   class Server::StdioTransport
     Response = Data.define(:id, :result) do
@@ -12,15 +14,15 @@ module ModelContextProtocol
       end
     end
 
-    attr_reader :router, :configuration
+    attr_reader :router, :configuration, :request_store
 
     def initialize(router:, configuration:)
       @router = router
       @configuration = configuration
+      @request_store = RequestStore.new
     end
 
     def handle
-      # Connect logger to transport
       @configuration.logger.connect_transport(self)
 
       loop do
@@ -29,10 +31,19 @@ module ModelContextProtocol
 
         begin
           message = JSON.parse(line.chomp)
-          next if message["method"].start_with?("notifications")
 
-          result = router.route(message)
-          send_message(Response[id: message["id"], result: result.serialized])
+          if message["method"] == "notifications/cancelled"
+            handle_cancellation(message)
+            next
+          end
+
+          next if message["method"]&.start_with?("notifications/")
+
+          result = router.route(message, request_store: @request_store, transport: self)
+
+          if result
+            send_message(Response[id: message["id"], result: result.serialized])
+          end
         rescue ModelContextProtocol::Server::ParameterValidationError => validation_error
           @configuration.logger.error("Validation error", error: validation_error.message)
           send_message(
@@ -61,12 +72,26 @@ module ModelContextProtocol
       $stdout.puts(JSON.generate(notification))
       $stdout.flush
     rescue IOError => e
-      # Handle broken pipe gracefully
       @configuration.logger.debug("Failed to send notification", error: e.message) if @configuration.logging_enabled?
     end
 
     private
 
+    # Handle a cancellation notification from the client
+    #
+    # @param message [Hash] the cancellation notification message
+    def handle_cancellation(message)
+      params = message["params"]
+      return unless params
+
+      request_id = params["requestId"]
+      return unless request_id
+
+      @request_store.mark_cancelled(request_id)
+    rescue
+      nil
+    end
+
     def receive_message
       $stdin.gets
     end

data/lib/model_context_protocol/server/streamable_http_transport/event_counter.rb

@@ -0,0 +1,35 @@
+module ModelContextProtocol
+  class Server::StreamableHttpTransport
+    class EventCounter
+      COUNTER_KEY_PREFIX = "event_counter:"
+
+      def initialize(redis_client, server_instance)
+        @redis = redis_client
+        @server_instance = server_instance
+        @counter_key = "#{COUNTER_KEY_PREFIX}#{server_instance}"
+
+        if @redis.exists(@counter_key) == 0
+          @redis.set(@counter_key, 0)
+        end
+      end
+
+      def next_event_id
+        count = @redis.incr(@counter_key)
+        "#{@server_instance}-#{count}"
+      end
+
+      def current_count
+        count = @redis.get(@counter_key)
+        count ? count.to_i : 0
+      end
+
+      def reset
+        @redis.set(@counter_key, 0)
+      end
+
+      def set_count(value)
+        @redis.set(@counter_key, value.to_i)
+      end
+    end
+  end
+end