feat(CE): spending limit alerts schema, models and the refusal seam (#2267)

server/app/controllers/api/v1/connectors_controller.rb

@@ -131,7 +131,9 @@ def execute_model
         if @connector.source?
           result = ExecuteModel.call(
             connector: @connector,
-            payload: params[:payload]
+            payload: params[:payload],
+            workspace: current_workspace,
+            user: current_user
           )
 
           if result.success?
@@ -140,7 +142,8 @@ def execute_model
           else
             render_error(
               message: result["error"],
-              status: :unprocessable_content
+              status: spend_limit_status(result),
+              metadata: result.metadata
             )
           end
         else
@@ -153,6 +156,10 @@ def execute_model
 
       private
 
+      def spend_limit_status(result)
+        SpendingLimits::SpendLimitExceeded.blocked?(result.metadata) ? :too_many_requests : :unprocessable_content
+      end
+
       def set_connector
         @connector = current_workspace.connectors.find(params[:id])
       rescue ActiveRecord::RecordNotFound

server/app/mailers/spending_limit_mailer.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+class SpendingLimitMailer < ApplicationMailer
+  default from: Rails.configuration.x.mail_from
+
+  def threshold_alert_email
+    @alert_attrs = params
+
+    mail(
+      to: params[:recipients],
+      subject: "Spend limit \"#{params[:limit_name]}\" reached #{params[:percentage]}%"
+    )
+  end
+end

server/app/models/spending_limit.rb

@@ -0,0 +1,134 @@
+# frozen_string_literal: true
+
+class SpendingLimit < ApplicationRecord
+  MIN_ALERT_THRESHOLD = 1
+  MAX_ALERT_THRESHOLD = 99
+  DEFAULT_ALERT_THRESHOLDS = [80, 90].freeze
+  ALLOCATION_NAME = "Workspace allocation"
+  ALWAYS_ALERT_AT = 100
+  # An empty list means "every one of them", so these carry no null state.
+  TARGET_FILTERS = %i[scope_ids providers models].freeze
+  MAX_TARGET_FILTER_ENTRIES = 100
+
+  belongs_to :organization
+  belongs_to :workspace
+  belongs_to :created_by, class_name: "User", optional: true
+
+  has_many :spending_limit_counters, dependent: :destroy
+
+  # Explicit values: index_spending_limits_one_allocation_per_workspace is partial on level = 1.
+  enum :level, { limit: 0, allocation: 1 }, scopes: false
+  enum :scope_type, %i[workspace workflow data_app user role]
+  enum :limit_type, %i[cost tokens both]
+  enum :period, %i[daily weekly monthly]
+  enum :action_on_exhaust, %i[warn block]
+
+  before_validation :start_counting, on: :create
+  before_validation :normalize_alert_thresholds
+  before_validation :normalize_target_filters
+
+  validates :name, presence: true, uniqueness: { scope: :workspace_id }
+  # The partial unique index is the real guarantee; this only turns the ordinary case into a 422.
+  validates :workspace_id, uniqueness: { conditions: -> { where(level: :allocation) },
+                                         message: "already has an allocation" }, if: :allocation?
+  validates :cost_limit, presence: true, if: -> { cost? || both? }
+  validates :token_limit, presence: true, if: -> { tokens? || both? }
+  validates :cost_limit, numericality: { greater_than: 0 }, allow_nil: true
+  validates :token_limit, numericality: { only_integer: true, greater_than: 0 }, allow_nil: true
+  validates(*TARGET_FILTERS, length: { maximum: MAX_TARGET_FILTER_ENTRIES,
+                                       too_long: "has too many entries (maximum is #{MAX_TARGET_FILTER_ENTRIES})" })
+  validate :alert_thresholds_in_range
+  validate :allocation_name_is_reserved
+
+  def meters_cost?
+    cost? || both?
+  end
+
+  def meters_tokens?
+    tokens? || both?
+  end
+
+  def percent_used(snapshot)
+    percentages = []
+    percentages << ratio(snapshot&.spent_cost || 0, cost_limit) if meters_cost?
+    percentages << ratio(snapshot&.spent_tokens || 0, token_limit) if meters_tokens?
+    percentages.compact.max || 0
+  end
+
+  # LedgerScope clamps its sum to this.
+  def restart_counting!
+    update_columns(counting_started_at: Time.current) # rubocop:disable Rails/SkipsModelValidations
+  end
+
+  # Recovery reads this; settlement cannot, because a crossing is only visible in one upsert.
+  def thresholds_reached(counter)
+    used = percent_used(counter)
+    alerting_thresholds.select { |threshold| used >= threshold }
+  end
+
+  def thresholds_crossed(before, after)
+    wanted = alerting_thresholds
+    from = percent_used(before)
+    to = percent_used(after)
+    wanted.select { |threshold| from < threshold && to >= threshold }
+  end
+
+  scope :limits, -> { where(level: :limit) }
+  scope :allocations, -> { where(level: :allocation) }
+  scope :enabled, -> { where(enabled: true) }
+  scope :for_workspace, ->(workspace_id) { where(workspace_id:) }
+
+  # The array type turns a lone string into [] before any callback sees it, so wrap in the writer.
+  TARGET_FILTERS.each do |field|
+    define_method(:"#{field}=") { |value| super(Array.wrap(value)) }
+  end
+
+  private
+
+  def alerting_thresholds
+    (Array(alert_thresholds).map(&:to_i) + [ALWAYS_ALERT_AT]).uniq.sort
+  end
+
+  def start_counting
+    self.counting_started_at ||= Time.current
+  end
+
+  def ratio(spent, cap)
+    return nil if cap.blank? || cap.to_d <= 0
+
+    (spent.to_d / cap.to_d) * 100
+  end
+
+  # [] means no early warnings and is kept; only nil takes the default.
+  def normalize_alert_thresholds
+    return self.alert_thresholds = DEFAULT_ALERT_THRESHOLDS.dup if alert_thresholds.nil?
+    return unless alert_thresholds.is_a?(Array) && alert_thresholds.all?(Numeric)
+
+    self.alert_thresholds = alert_thresholds.map(&:to_i).uniq.sort
+  end
+
+  def normalize_target_filters
+    TARGET_FILTERS.each do |field|
+      entries = Array.wrap(public_send(field)).map { |entry| entry.to_s.strip }.reject(&:empty?).uniq
+      public_send(:"#{field}=", entries)
+    end
+  end
+
+  # The (workspace_id, name) index is level agnostic, so a limit may not take the allocation's name.
+  def allocation_name_is_reserved
+    return unless limit?
+    return unless name.to_s.strip.casecmp?(ALLOCATION_NAME)
+
+    errors.add(:name, "is reserved for the workspace allocation")
+  end
+
+  # The integer[] cast leaves a nested array alone, and that answers neither to_i nor between?.
+  def alert_thresholds_in_range
+    return errors.add(:alert_thresholds, "must be a list of numbers") unless alert_thresholds.is_a?(Array)
+    return if alert_thresholds.blank?
+    return errors.add(:alert_thresholds, "must be a list of numbers") unless alert_thresholds.all?(Integer)
+    return if alert_thresholds.all? { |threshold| threshold.between?(MIN_ALERT_THRESHOLD, MAX_ALERT_THRESHOLD) }
+
+    errors.add(:alert_thresholds, "must be between #{MIN_ALERT_THRESHOLD} and #{MAX_ALERT_THRESHOLD}")
+  end
+end

server/app/models/spending_limit_alert.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+class SpendingLimitAlert < ApplicationRecord
+  # Configurable thresholds stop at 99; reaching the cap always alerts.
+  MAX_THRESHOLD = 100
+
+  belongs_to :spending_limit_counter
+
+  validates :threshold, presence: true,
+                        numericality: { only_integer: true,
+                                        greater_than_or_equal_to: SpendingLimit::MIN_ALERT_THRESHOLD,
+                                        less_than_or_equal_to: MAX_THRESHOLD },
+                        uniqueness: { scope: :spending_limit_counter_id }
+
+  CLAIM_TTL = 15.minutes
+
+  scope :pending, lambda {
+    where(delivered_at: nil, abandoned_at: nil)
+      .where("claimed_at IS NULL OR claimed_at < ?", CLAIM_TTL.ago)
+  }
+
+  # ON CONFLICT, not a rescued create!: a uniqueness violation would abort the enclosing
+  # transaction and roll back the counter movement it shares.
+  def self.record_missing(counter, thresholds)
+    rows = Array(thresholds).map { |threshold| { spending_limit_counter_id: counter.id, threshold: threshold.to_i } }
+    return [] if rows.empty?
+
+    result = insert_all( # rubocop:disable Rails/SkipsModelValidations
+      rows, unique_by: :index_spend_alerts_on_counter_and_threshold, returning: :id, record_timestamps: true
+    )
+    result.rows.flatten
+  end
+
+  # Only for a crossing that can never mean anything again; every other skip state can come back.
+  def abandon!
+    update!(abandoned_at: Time.current)
+  end
+
+  def delivered?
+    delivered_at.present?
+  end
+end

server/app/models/spending_limit_counter.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+class SpendingLimitCounter < ApplicationRecord
+  belongs_to :spending_limit
+
+  # The FK cascade is what a window restart relies on; this covers only counter.destroy.
+  has_many :spending_limit_alerts, dependent: :delete_all
+
+  validates :window_start, presence: true, uniqueness: { scope: :spending_limit_id }
+  validates :spent_cost, numericality: { greater_than_or_equal_to: 0 }
+  validates :spent_tokens, numericality: { only_integer: true, greater_than_or_equal_to: 0 }
+  validates :request_count, numericality: { only_integer: true, greater_than_or_equal_to: 0 }
+  validates :unpriced_requests, numericality: { only_integer: true, greater_than_or_equal_to: 0 }
+
+  scope :recent, -> { order(window_start: :desc) }
+  scope :for_window, ->(window_start) { where(window_start:) }
+
+  def alerted?(threshold)
+    alerted_thresholds.include?(threshold)
+  end
+
+  # One statement, because a worker for another threshold writes the same column concurrently.
+  def record_threshold!(threshold)
+    value = threshold.to_i
+    append = ["alerted_thresholds = (SELECT array_agg(DISTINCT t ORDER BY t) " \
+              "FROM unnest(alerted_thresholds || ?::integer) AS t)", value]
+
+    self.class.where(id:)
+        .where.not("alerted_thresholds @> ARRAY[?]::integer[]", value)
+        .update_all(append) # rubocop:disable Rails/SkipsModelValidations
+  end
+end

server/app/services/spending_limits/spend_limit_exceeded.rb

@@ -0,0 +1,78 @@
+# frozen_string_literal: true
+
+module SpendingLimits
+  class SpendLimitExceeded < StandardError
+    CODE = "spend_limit_exceeded"
+    PUBLIC_MESSAGE = "This service is temporarily unavailable. Please try again later."
+
+    SCOPE_LABELS = {
+      "workspace" => "Workspace",
+      "workflow" => "AI Workflows",
+      "data_app" => "Data Apps",
+      "user" => "Users",
+      "role" => "Roles"
+    }.freeze
+
+    attr_reader :decision
+
+    def self.blocked?(metadata)
+      return false unless metadata.is_a?(Hash)
+
+      (metadata[:code] || metadata["code"]) == CODE
+    end
+
+    def initialize(decision = nil)
+      @decision = decision
+      super(build_message)
+    end
+
+    def limit
+      decision&.limit
+    end
+
+    def resets_at
+      decision&.resets_at
+    end
+
+    def scope
+      limit&.scope_type.to_s.presence
+    end
+
+    def exhausted
+      decision&.exhausted.to_s.presence
+    end
+
+    # No row id: this body reaches callers who may hold no spending_limit permission.
+    def metadata
+      {
+        code: CODE,
+        scope:,
+        exhausted:,
+        resets_at: resets_at&.utc&.iso8601
+      }.compact
+    end
+
+    def public_message
+      PUBLIC_MESSAGE
+    end
+
+    private
+
+    def build_message
+      return PUBLIC_MESSAGE if limit.nil?
+
+      "Spend limit '#{limit.name}' (#{qualifiers}) is used up for the current period."
+    end
+
+    def qualifiers
+      [SCOPE_LABELS[limit.scope_type.to_s] || limit.scope_type.to_s, provider_label].compact.join(", ")
+    end
+
+    def provider_label
+      providers = Array(limit.providers).map(&:to_s).reject(&:empty?)
+      return nil if providers.empty?
+
+      providers.to_sentence
+    end
+  end
+end