Expose Management API rate limit info via a client callback

EXAMPLES.md

@@ -166,6 +166,26 @@ management.roles.create(name: 'admin', description: 'Administrator')
 management.organizations.list(take: 20)
 ```
 
+## Rate Limit Monitoring
+
+Pass a `rate_limit_handler` when building an `Auth0::Client` to observe Auth0's `x-ratelimit-*` headers, without changing the return value of your calls. It covers Management API calls and the Authentication API calls made through the client's HTTP layer (it is not available on `Auth0::Management` directly):
+
+```ruby
+client = Auth0::Client.new(
+  domain: ENV['AUTH0_RUBY_DOMAIN'],
+  client_id: ENV['AUTH0_RUBY_CLIENT_ID'],
+  client_secret: ENV['AUTH0_RUBY_CLIENT_SECRET'],
+  rate_limit_handler: lambda do |rate_limit|
+    # rate_limit.limit / rate_limit.remaining (Integers), rate_limit.reset (UTC Time)
+    Rails.logger.info("Auth0 rate limit: #{rate_limit.remaining}/#{rate_limit.limit}, resets at #{rate_limit.reset}")
+  end
+)
+
+client.users.get(id: 'auth0|123456') # handler fires with this response's rate limit
+```
+
+The handler is invoked on every response, including the `429`s that trigger an automatic retry. An error raised inside the handler is caught (and warned) so it never breaks the request.
+
 ## Organizations
 
 [Organizations](https://auth0.com/docs/organizations) is a set of features that provide better support for developers who build and maintain SaaS and Business-to-Business (B2B) applications.

README.md

@@ -104,6 +104,7 @@ client = Auth0::Client.new(
 | `management_timeout` | Float | Timeout in seconds for Management API calls. | `60` |
 | `management_max_retries` | Integer | Maximum retries for Management API calls. | `2` |
 | `management_additional_headers` | Hash | Additional HTTP headers for Management API calls. | `nil` |
+| `rate_limit_handler` | `#call` | Callback for [rate-limit monitoring](#rate-limit-monitoring). | `nil` |
 
 #### Accessing the Management Client Directly
 
@@ -283,6 +284,23 @@ client.users.list(
 )
 ```
 
+### Rate Limit Monitoring
+
+Auth0 returns `x-ratelimit-limit`, `x-ratelimit-remaining`, and `x-ratelimit-reset` headers on API responses. Pass a `rate_limit_handler` when building an `Auth0::Client` to be notified of this data, so you can monitor how close you are to the limit (for example, emit a metric and alert before you run out). It covers Management API calls and the Authentication API calls made through the client's HTTP layer (the JWKS fetch used by `validate_id_token` goes through a separate object and is not reported). It is not available on `Auth0::Management` directly.
+
+```ruby
+client = Auth0::Client.new(
+  domain: ENV['AUTH0_RUBY_DOMAIN'],
+  client_id: ENV['AUTH0_RUBY_CLIENT_ID'],
+  client_secret: ENV['AUTH0_RUBY_CLIENT_SECRET'],
+  rate_limit_handler: lambda do |rate_limit|
+    StatsD.gauge('auth0.rate_limit.remaining', rate_limit.remaining) if rate_limit.remaining
+  end
+)
+```
+
+The handler receives an object with `#limit`, `#remaining` (Integers), and `#reset` (a UTC `Time`). Each is `nil` when the corresponding header is absent or non-numeric. It is invoked on every response, including the `429`s that trigger an automatic retry, so you can observe the point at which the limit was reached. The return value of your API call is unchanged, and an exception raised inside the handler is caught (and warned) so it can never break a request.
+
 ### Errors
 
 Management API errors use the `Auth0::Errors` namespace:

lib/auth0/auth_client.rb

@@ -85,14 +85,39 @@ def respond_to_missing?(method_name, include_private = false)
     def management
       current_token = get_token
       if @_management.nil? || @_management_token != current_token
-        @_management_token = current_token
         opts = { token: current_token, base_url: "#{@base_uri}/api/v2" }
         opts[:timeout] = @management_timeout if @management_timeout
         opts[:max_retries] = @management_max_retries if @management_max_retries
         opts[:headers] = @management_additional_headers if @management_additional_headers
-        @_management = Auth0::Management.new(**opts)
+        management = Auth0::Management.new(**opts)
+        # Attach before caching so a failed attach doesn't leave a handler-less
+        # management cached (the next call would then skip the rebuild).
+        attach_rate_limit_handler(management)
+        @_management = management
+        @_management_token = current_token
       end
       @_management
     end
+
+    private
+
+    # Attaches the configured rate limit handler to the management client's raw
+    # client. Fails loudly rather than silently skipping, so a future change to
+    # the generated client's internals surfaces instead of quietly disabling it.
+    # @param management [Auth0::Management]
+    # @return [void]
+    # @raise [Auth0::Exception] if the management client exposes no compatible raw client
+    def attach_rate_limit_handler(management)
+      return if @rate_limit_handler.nil?
+
+      raw_client = management.instance_variable_get(:@raw_client)
+      unless raw_client.respond_to?(:rate_limit_handler=)
+        raise Auth0::Exception,
+              "Unable to attach rate_limit_handler: the management client does not expose a compatible raw client. " \
+              "This likely means the ruby-auth0 internals changed. Please report it."
+      end
+
+      raw_client.rate_limit_handler = @rate_limit_handler
+    end
   end
 end

lib/auth0/internal/http/rate_limit.rb

@@ -0,0 +1,103 @@
+# frozen_string_literal: true
+
+module Auth0
+  module Internal
+    module Http
+      # Rate limit information parsed from the `x-ratelimit-*` headers Auth0
+      # returns on 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 from an object that exposes headers via `#[]` using the HTTP
+        # header name (e.g. a `Net::HTTPResponse`, whose `#[]` is
+        # case-insensitive). Used by the Management API (RawClient) path.
+        #
+        # @param response [#[]] responds to `[]` with header access
+        # @return [Auth0::Internal::Http::RateLimit]
+        def self.from_http_response(response)
+          build(
+            response["x-ratelimit-limit"],
+            response["x-ratelimit-remaining"],
+            response["x-ratelimit-reset"]
+          )
+        end
+
+        # Build from a plain headers hash (e.g. RestClient's, whose keys are
+        # symbols like `:x_ratelimit_remaining`). Keys are matched
+        # case-insensitively and dash/underscore-agnostically. Used by the
+        # Authentication API (HTTPProxy) path.
+        #
+        # @param headers [Hash, nil]
+        # @return [Auth0::Internal::Http::RateLimit]
+        def self.from_headers(headers)
+          normalized = (headers || {}).each_with_object({}) do |(key, value), acc|
+            acc[key.to_s.downcase.tr("-", "_")] = value
+          end
+
+          build(
+            normalized["x_ratelimit_limit"],
+            normalized["x_ratelimit_remaining"],
+            normalized["x_ratelimit_reset"]
+          )
+        end
+
+        # Invokes the handler with the given rate limit. A handler error must
+        # never break the caller's request, so it is rescued and warned rather
+        # than raised. Parsing happens before this call, so a parsing bug is not
+        # misattributed to the handler.
+        #
+        # @param handler [#call, nil] the configured rate limit handler
+        # @param rate_limit [Auth0::Internal::Http::RateLimit]
+        # @return [void]
+        def self.deliver(handler, rate_limit)
+          return if handler.nil?
+
+          handler.call(rate_limit)
+        rescue StandardError => e
+          warn "[auth0] rate_limit_handler raised #{e.class}: #{e.message}"
+        end
+
+        # @param limit [String, nil] raw `x-ratelimit-limit` value
+        # @param remaining [String, nil] raw `x-ratelimit-remaining` value
+        # @param reset [String, nil] raw `x-ratelimit-reset` value
+        # @return [Auth0::Internal::Http::RateLimit]
+        def self.build(limit, remaining, reset)
+          reset_epoch = to_integer(reset)
+
+          new(
+            limit: to_integer(limit),
+            remaining: to_integer(remaining),
+            reset: reset_epoch.nil? ? nil : Time.at(reset_epoch).utc
+          )
+        end
+        private_class_method :build
+
+        # Parse an integer header value, returning nil for blank or non-numeric
+        # input (so a malformed header is never silently reported as 0).
+        #
+        # @param value [String, nil] the raw header value
+        # @return [Integer, nil]
+        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

@@ -3,6 +3,8 @@
 # This file was auto-generated by Fern from our API Definition.
 # Modified by Auth0 to use Auth0 telemetry format with dynamic versioning.
 
+require_relative "rate_limit"
+
 module Auth0
   module Internal
     module Http
@@ -20,14 +22,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 +54,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
@@ -67,6 +76,9 @@ def send(request)
 
             response = conn.request(http_request)
 
+            # Report rate limit on every response, including retried 429s.
+            RateLimit.deliver(@rate_limit_handler, RateLimit.from_http_response(response)) if @rate_limit_handler
+
             break unless should_retry?(response, attempt)
 
             delay = retry_delay(response, attempt)

lib/auth0/mixins/httpproxy.rb

@@ -1,13 +1,14 @@
 require "addressable/uri"
 require "retryable"
 require_relative "../exception.rb"
+require_relative "../internal/http/rate_limit"
 
 module Auth0
   module Mixins
     # here's the proxy for Rest calls based on rest-client, we're building all request on that gem
     # for now, if you want to feel free to use your own http client
     module HTTPProxy
-      attr_accessor :headers, :base_uri, :timeout, :retry_count
+      attr_accessor :headers, :base_uri, :timeout, :retry_count, :rate_limit_handler
       DEFAULT_RETRIES = 3
       MAX_ALLOWED_RETRIES = 10
       MAX_REQUEST_RETRY_JITTER = 250
@@ -95,6 +96,13 @@ def request(method, uri, body = {}, extra_headers = {})
           call(method, encode_uri(uri), timeout, headers, body.to_json)
         end
 
+        # Surface rate limit for each response, including 429 retries. `call`
+        # returns nil for some errors (e.g. a dropped connection), so skip it then.
+        if @rate_limit_handler && !result.nil?
+          rate_limit = Auth0::Internal::Http::RateLimit.from_headers(result.headers)
+          Auth0::Internal::Http::RateLimit.deliver(@rate_limit_handler, rate_limit)
+        end
+
         case result.code
         when 200...226 then safe_parse_json(result.body)
         when 400       then raise Auth0::BadRequest.new(result.body, code: result.code, headers: result.headers)

lib/auth0/mixins/initializer.rb

@@ -21,6 +21,11 @@ def initialize(config)
         @management_timeout = options[:management_timeout]
         @management_max_retries = options[:management_max_retries]
         @management_additional_headers = options[:management_additional_headers]
+        handler = options[:rate_limit_handler]
+        raise ArgumentError, "rate_limit_handler must respond to #call" if handler && !handler.respond_to?(:call)
+
+        @rate_limit_handler = handler
+
         extend Auth0::Api::AuthenticationEndpoints
 
         @client_id = options[:client_id]