Expose Management API rate limit info via a client callback

lib/auth0.rb

@@ -12,6 +12,7 @@
 require_relative "auth0/internal/errors/type_error"
 require_relative "auth0/internal/http/base_request"
 require_relative "auth0/internal/json/request"
+require_relative "auth0/internal/http/rate_limit"
 require_relative "auth0/internal/http/raw_client"
 require_relative "auth0/internal/multipart/multipart_encoder"
 require_relative "auth0/internal/multipart/multipart_form_data_part"

lib/auth0/auth_client.rb

@@ -91,8 +91,23 @@ def management
         opts[:max_retries] = @management_max_retries if @management_max_retries
         opts[:headers] = @management_additional_headers if @management_additional_headers
         @_management = Auth0::Management.new(**opts)
+        attach_rate_limit_handler(@_management)
       end
       @_management
     end
+
+    private
+
+    # Attaches the configured rate limit handler to the management client's
+    # underlying raw client. Management is generated and builds its own raw
+    # client, so we set the handler on it after construction.
+    # @param management [Auth0::Management]
+    # @return [void]
+    def attach_rate_limit_handler(management)
+      return if @management_rate_limit_handler.nil?
+
+      raw_client = management.instance_variable_get(:@raw_client)
+      raw_client.rate_limit_handler = @management_rate_limit_handler if raw_client
+    end
   end
 end

lib/auth0/internal/http/rate_limit.rb

@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+module Auth0
+  module Internal
+    module Http
+      # Rate limit information parsed from the `x-ratelimit-*` headers Auth0
+      # returns on Management API responses.
+      #
+      # @see https://auth0.com/docs/troubleshoot/customer-support/operational-policies/rate-limit-policy
+      class RateLimit
+        # @return [Integer, nil] the maximum number of requests allowed in the current window
+        attr_reader :limit
+        # @return [Integer, nil] the number of requests remaining in the current window
+        attr_reader :remaining
+        # @return [Time, nil] the UTC time at which the current window resets
+        attr_reader :reset
+
+        # @param limit [Integer, nil]
+        # @param remaining [Integer, nil]
+        # @param reset [Time, nil]
+        def initialize(limit:, remaining:, reset:)
+          @limit = limit
+          @remaining = remaining
+          @reset = reset
+        end
+
+        # Build a RateLimit from an HTTP response. Header lookups are
+        # case-insensitive (delegated to the response), and missing or
+        # non-numeric values become nil rather than a misleading 0.
+        #
+        # @param response [Net::HTTPResponse] anything responding to `[]` with header access
+        # @return [Auth0::Internal::Http::RateLimit]
+        def self.from_response(response)
+          reset = to_integer(response["x-ratelimit-reset"])
+
+          new(
+            limit: to_integer(response["x-ratelimit-limit"]),
+            remaining: to_integer(response["x-ratelimit-remaining"]),
+            reset: reset.nil? ? nil : Time.at(reset).utc
+          )
+        end
+
+        def self.to_integer(value)
+          Integer(value.to_s.strip, exception: false)
+        end
+        private_class_method :to_integer
+      end
+    end
+  end
+end

lib/auth0/internal/http/raw_client.rb

@@ -20,14 +20,21 @@ class RawClient
         # @return [String] The base URL for requests
         attr_reader :base_url
 
+        # @return [#call, nil] Optional callback invoked with an
+        #   {Auth0::Internal::Http::RateLimit} after every response.
+        attr_accessor :rate_limit_handler
+
         # @param base_url [String] The base url for the request.
         # @param max_retries [Integer] The number of times to retry a failed request, defaults to 2.
         # @param timeout [Float] The timeout for the request, defaults to 60.0 seconds.
         # @param headers [Hash] The headers for the request.
-        def initialize(base_url:, max_retries: 2, timeout: 60.0, headers: {})
+        # @param rate_limit_handler [#call, nil] Optional callback invoked with the
+        #   parsed rate limit (from the `x-ratelimit-*` headers) after every response.
+        def initialize(base_url:, max_retries: 2, timeout: 60.0, headers: {}, rate_limit_handler: nil)
           @base_url = base_url
           @max_retries = max_retries
           @timeout = timeout
+          @rate_limit_handler = rate_limit_handler
 
           # Auth0 telemetry in standard format
           telemetry = {
@@ -45,7 +52,7 @@ def initialize(base_url:, max_retries: 2, timeout: 60.0, headers: {})
         end
 
         # @param request [Auth0::Internal::Http::BaseRequest] The HTTP request.
-        # @return [HTTP::Response] The HTTP response.
+        # @return [Net::HTTPResponse] The HTTP response.
         def send(request)
           url = build_url(request)
           attempt = 0
@@ -74,9 +81,23 @@ def send(request)
             attempt += 1
           end
 
+          notify_rate_limit(response)
           response
         end
 
+        # Invokes the rate limit handler with the rate limit parsed from the
+        # response headers. Runs after retries, on every response. A handler
+        # error must never break the request, so it is swallowed.
+        # @param response [Net::HTTPResponse] The HTTP response.
+        # @return [void]
+        def notify_rate_limit(response)
+          return if @rate_limit_handler.nil?
+
+          @rate_limit_handler.call(RateLimit.from_response(response))
+        rescue StandardError
+          nil
+        end
+
         # Determines if a request should be retried based on the response status code.
         # @param response [Net::HTTPResponse] The HTTP response.
         # @param attempt [Integer] The current retry attempt (0-indexed).

lib/auth0/mixins/initializer.rb

@@ -21,6 +21,7 @@ def initialize(config)
         @management_timeout = options[:management_timeout]
         @management_max_retries = options[:management_max_retries]
         @management_additional_headers = options[:management_additional_headers]
+        @management_rate_limit_handler = options[:management_rate_limit_handler]
         extend Auth0::Api::AuthenticationEndpoints
 
         @client_id = options[:client_id]

test/unit/internal/http/test_rate_limit.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+require "test_helper"
+
+describe Auth0::Internal::Http::RateLimit do
+  RateLimit = Auth0::Internal::Http::RateLimit
+
+  describe ".from_response" do
+    it "parses the x-ratelimit-* headers" do
+      response = {
+        "x-ratelimit-limit" => "100",
+        "x-ratelimit-remaining" => "42",
+        "x-ratelimit-reset" => "1724000000"
+      }
+
+      rate_limit = RateLimit.from_response(response)
+
+      _(rate_limit.limit).must_equal 100
+      _(rate_limit.remaining).must_equal 42
+      _(rate_limit.reset).must_equal Time.at(1_724_000_000).utc
+    end
+
+    it "reports a remaining of 0 as an integer, not nil" do
+      _(RateLimit.from_response("x-ratelimit-remaining" => "0").remaining).must_equal 0
+    end
+
+    it "returns nil for missing or non-numeric values instead of a misleading 0" do
+      rate_limit = RateLimit.from_response(
+        "x-ratelimit-limit" => "",
+        "x-ratelimit-remaining" => "not-a-number"
+      )
+
+      _(rate_limit.limit).must_be_nil
+      _(rate_limit.remaining).must_be_nil
+      _(rate_limit.reset).must_be_nil
+    end
+  end
+end

test/unit/internal/http/test_raw_client.rb

@@ -0,0 +1,102 @@
+# frozen_string_literal: true
+
+require "test_helper"
+
+describe Auth0::Internal::Http::RawClient do
+  module TestRawClient
+    # Minimal stand-in for a Net::HTTPResponse.
+    class FakeHttpResponse
+      def initialize(code:, body:, headers:)
+        @code = code
+        @body = body
+        @headers = headers
+      end
+
+      attr_reader :code, :body
+
+      def [](name)
+        @headers[name]
+      end
+    end
+
+    # Minimal stand-in for the Net::HTTP connection.
+    class FakeConnection
+      def initialize(response)
+        @response = response
+      end
+
+      def open_timeout=(_); end
+      def read_timeout=(_); end
+      def write_timeout=(_); end
+      def continue_timeout=(_); end
+
+      def request(_http_request)
+        @response
+      end
+    end
+
+    def self.build_request
+      Auth0::Internal::JSON::Request.new(
+        base_url: nil,
+        method: "GET",
+        path: "users",
+        query: {},
+        request_options: {}
+      )
+    end
+
+    def self.build_response
+      FakeHttpResponse.new(
+        code: "200",
+        body: "{}",
+        headers: {
+          "x-ratelimit-limit" => "100",
+          "x-ratelimit-remaining" => "12",
+          "x-ratelimit-reset" => "1724000000"
+        }
+      )
+    end
+  end
+
+  def send_with(client, response)
+    client.stub(:connect, TestRawClient::FakeConnection.new(response)) do
+      client.send(TestRawClient.build_request)
+    end
+  end
+
+  it "invokes the rate limit handler with the parsed rate limit and returns the response unchanged" do
+    captured = nil
+    client = Auth0::Internal::Http::RawClient.new(
+      base_url: "https://tenant.auth0.com",
+      max_retries: 0,
+      rate_limit_handler: ->(rate_limit) { captured = rate_limit }
+    )
+    response = TestRawClient.build_response
+
+    result = send_with(client, response)
+
+    _(result).must_be_same_as response
+    _(captured).must_be_instance_of Auth0::Internal::Http::RateLimit
+    _(captured.limit).must_equal 100
+    _(captured.remaining).must_equal 12
+    _(captured.reset).must_equal Time.at(1_724_000_000).utc
+  end
+
+  it "returns the response unchanged when no handler is configured" do
+    client = Auth0::Internal::Http::RawClient.new(base_url: "https://tenant.auth0.com", max_retries: 0)
+    response = TestRawClient.build_response
+
+    _(send_with(client, response)).must_be_same_as response
+  end
+
+  it "does not let a handler error break the request" do
+    client = Auth0::Internal::Http::RawClient.new(
+      base_url: "https://tenant.auth0.com",
+      max_retries: 0,
+      rate_limit_handler: ->(_rate_limit) { raise "boom" }
+    )
+    response = TestRawClient.build_response
+
+    _(send_with(client, response)).must_be_same_as response
+  end
+end

test/unit/test_auth_client_rate_limit.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+require "test_helper"
+
+describe Auth0::Client do
+  def build_client(**extra)
+    Auth0::Client.new(domain: "tenant.auth0.com", token: "test-token", **extra)
+  end
+
+  it "attaches the configured management_rate_limit_handler to the management raw client" do
+    handler = ->(_rate_limit) {}
+    client = build_client(management_rate_limit_handler: handler)
+
+    raw_client = client.management.instance_variable_get(:@raw_client)
+
+    _(raw_client.rate_limit_handler).must_be_same_as handler
+  end
+
+  it "leaves the handler unset when none is configured" do
+    raw_client = build_client.management.instance_variable_get(:@raw_client)
+
+    _(raw_client.rate_limit_handler).must_be_nil
+  end
+end