#447: make off_quota? aware of GitHub Search API quota

.rubocop.yml

@@ -34,6 +34,7 @@ Elegant/GoodMethodName:
     - delete_one
     - dump_headers
     - extract_remaining_count
+    - extract_search_remaining_count
     - faraday_value
     - github_graph
     - handle_rate_limit_request
@@ -53,6 +54,7 @@ Elegant/GoodMethodName:
     - off_quota?
     - organization_memberships
     - organization_repositories
+    - patched_body
     - print_trace!
     - publish_to
     - pull_request
@@ -85,7 +87,6 @@ Elegant/GoodMethodName:
     - track_request
     - unmask_repos
     - update_organization_membership
-    - update_remaining_count
     - user_name_by_id
     - user_repository_invitations
     - with_disable_auto_paginate

lib/fbe/middleware/rate_limit.rb

@@ -24,13 +24,16 @@
 # Copyright:: Copyright (c) 2024-2026 Zerocracy
 # License:: MIT
 class Fbe::Middleware::RateLimit < Faraday::Middleware
+  # NOT thread-safe: assumes single-threaded use (judges run sequentially in judges-action).
+  #
   # Initializes the rate limit middleware.
   #
   # @param [Object] app The next middleware in the stack
   def initialize(app)
     super
     @cached = nil
     @remaining = nil
+    @searchleft = nil
     @counter = 0
   end
 
@@ -42,7 +45,7 @@ def call(env)
     if env.url.path == '/rate_limit'
       handle_rate_limit_request(env)
     else
-      track_request
+      track_request(env.url.path)
       @app.call(env)
     end
   end
@@ -56,22 +59,24 @@ def call(env)
   def handle_rate_limit_request(env)
     if @cached.nil? || @counter >= 100
       response = @app.call(env)
-      @cached = response.dup
+      @cached = response
       @remaining = extract_remaining_count(response)
+      @searchleft = extract_search_remaining_count(response)
       @counter = 0
       response
     else
-      response = @cached.dup
-      update_remaining_count(response)
-      Faraday::Response.new(response_env(env, response))
+      Faraday::Response.new(response_env(env, @cached))
     end
   end
 
   # Tracks non-rate_limit requests and decrements counter.
-  def track_request
-    return if @remaining.nil?
-    @remaining -= 1 if @remaining.positive?
+  def track_request(path = nil)
     @counter += 1
+    if path&.start_with?('/search/')
+      @searchleft -= 1 if @searchleft&.positive?
+    elsif @remaining
+      @remaining -= 1 if @remaining.positive?
+    end
   end
 
   # Extracts the remaining count from the response body.
@@ -80,52 +85,56 @@ def track_request
   # @return [Integer] The remaining requests count
   def extract_remaining_count(response)
     body = response.body
-    if body.is_a?(String)
-      begin
-        body = JSON.parse(body)
-      rescue JSON::ParserError
-        return 0
-      end
-    end
+    body = JSON.parse(body) if body.is_a?(String)
     return 0 unless body.is_a?(Hash)
     body.dig('rate', 'remaining') || 0
   end
 
-  # Updates the remaining count in the response body.
+  # Extracts the search-resource remaining count from the response body.
   #
-  # @param [Faraday::Response] response The cached response to update
-  def update_remaining_count(response)
+  # @param [Faraday::Response] response The API response
+  # @return [Integer, nil] Remaining search-API requests or nil if absent
+  def extract_search_remaining_count(response)
     body = response.body
-    stringed = body.is_a?(String)
-    if stringed
-      begin
-        body = JSON.parse(body)
-      rescue JSON::ParserError
-        return
+    body = JSON.parse(body) if body.is_a?(String)
+    return nil unless body.is_a?(Hash)
+    body.dig('resources', 'search', 'remaining')
+  end
+
+  # Builds a fresh body with the current remaining counts written in,
+  # without mutating the cached response. Uses a JSON round-trip for
+  # the deep copy so we only handle JSON-shaped data.
+  #
+  # @param [Object] original The cached response body (Hash or JSON String)
+  # @return [Object] A new body of the same type with remaining counts updated
+  def patched_body(original)
+    stringed = original.is_a?(String)
+    body =
+      if stringed
+        JSON.parse(original)
+      elsif original.is_a?(Hash)
+        JSON.parse(original.to_json)
+      else
+        return original
       end
-    end
-    return unless body.is_a?(Hash) && body['rate']
-    body['rate']['remaining'] = @remaining
-    return unless stringed
-    response.instance_variable_set(:@body, body.to_json)
+    body['rate']['remaining'] = @remaining if body['rate']
+    body.dig('resources', 'search')&.[]=('remaining', @searchleft) if @searchleft
+    stringed ? body.to_json : body
   end
 
-  # Creates a response environment for the cached response.
+  # Builds a response environment that mirrors the cached response,
+  # preserving Faraday::Env invariants by dup-ing the original env
+  # and only overriding body and rate-limit headers.
   #
   # @param [Faraday::Env] env The original request environment
   # @param [Faraday::Response] response The cached response
-  # @return [Hash] Response environment hash
+  # @return [Faraday::Env] Response env ready to wrap in Faraday::Response
   def response_env(env, response)
-    headers = response.headers.dup
-    headers['x-ratelimit-remaining'] = @remaining.to_s if @remaining
-    {
-      method: env.method,
-      url: env.url,
-      request_headers: env.request_headers,
-      request_body: env.request_body,
-      status: response.status,
-      response_headers: headers,
-      body: response.body
-    }
+    served = response.env.dup
+    served.request_headers = env.request_headers
+    served.body = patched_body(response.body)
+    served.response_headers = response.headers.dup
+    served.response_headers['x-ratelimit-remaining'] = @remaining.to_s if @remaining
+    served
   end
 end

lib/fbe/octo.rb

@@ -28,6 +28,10 @@
 # When we are off quota.
 class Fbe::OffQuota < StandardError; end
 
+Fbe::SEARCH_METHODS = %i[
+  search_issues search_commits search_repositories search_users search_code search_topics
+].freeze
+
 # Makes a call to the GitHub API.
 #
 # It is supposed to be used instead of +Octokit::Client+, because it
@@ -154,13 +158,34 @@ def print_trace!(all: false, max: 5)
               @trace.clear
             end
           end
-          def off_quota?(threshold: 50) # rubocop:disable Layout/EmptyLineBetweenDefs
+          def off_quota?(threshold: nil, resource: :core) # rubocop:disable Layout/EmptyLineBetweenDefs, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
+            threshold ||= resource == :search ? 5 : 50
+            label = resource == :search ? 'GitHub Search API' : 'GitHub API'
             left = @origin.rate_limit!.remaining
+            got = false
+            if resource == :search && @origin.respond_to?(:last_response)
+              body = @origin.last_response&.body
+              body = JSON.parse(body) if body.is_a?(String)
+              if body.is_a?(Hash)
+                fresh = body.dig('resources', 'search', 'remaining') || body.dig(:resources, :search, :remaining)
+                if fresh
+                  left = Integer(fresh)
+                  got = true
+                end
+              end
+            end
+            if resource == :search && !got
+              klass = @origin.respond_to?(:last_response) ? @origin.last_response&.body&.class : nil
+              @loog.warn(
+                "Search-quota check fell back to core remaining (#{left}); " \
+                "search count unavailable (last_response body class: #{klass.inspect})"
+              )
+            end
             if left < threshold
-              @loog.info("Too much GitHub API quota consumed already (#{left} < #{threshold})")
+              @loog.info("Too much #{label} quota consumed already (#{left} < #{threshold})")
               true
             else
-              @loog.debug("Still #{left} GitHub API quota left (>#{threshold})")
+              @loog.debug("Still #{left} #{label} quota left (>#{threshold})")
               false
             end
           end
@@ -208,7 +233,12 @@ def with_disable_auto_paginate # rubocop:disable Layout/EmptyLineBetweenDefs
         end
       o =
         intercepted(o) do |e, m, _args, _r|
-          if e == :before && m != :off_quota? && m != :print_trace! && m != :rate_limit && o.off_quota?
+          next unless e == :before
+          next if %i[off_quota? print_trace! rate_limit].include?(m)
+          if Fbe::SEARCH_METHODS.include?(m)
+            raise(Fbe::OffQuota, "We are off-quota on the search resource, can't do #{m}()") if
+              o.off_quota?(resource: :search)
+          elsif o.off_quota?
             raise(Fbe::OffQuota, "We are off-quota (remaining: #{o.rate_limit.remaining}), can't do #{m}()")
           end
         end