ruby_llm-agents 0.4.0 → 1.0.0.beta.1

data/lib/ruby_llm/agents/infrastructure/budget/forecaster.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Agents
+    module Budget
+      # Budget forecasting based on current spending trends
+      #
+      # @api private
+      module Forecaster
+        class << self
+          # Calculates budget forecasts based on current spending trends
+          #
+          # @param tenant_id [String, nil] The tenant identifier
+          # @param budget_config [Hash] Budget configuration
+          # @return [Hash, nil] Forecast information
+          def calculate_forecast(tenant_id: nil, budget_config:)
+            return nil unless budget_config[:enabled]
+            return nil unless budget_config[:global_daily] || budget_config[:global_monthly]
+
+            daily_current = BudgetQuery.current_spend(:global, :daily, tenant_id: tenant_id)
+            monthly_current = BudgetQuery.current_spend(:global, :monthly, tenant_id: tenant_id)
+
+            # Calculate hours elapsed today and days elapsed this month
+            hours_elapsed = Time.current.hour + (Time.current.min / 60.0)
+            hours_elapsed = [hours_elapsed, 1].max # Avoid division by zero
+            days_in_month = Time.current.end_of_month.day
+            day_of_month = Time.current.day
+            days_elapsed = day_of_month - 1 + (hours_elapsed / 24.0)
+            days_elapsed = [days_elapsed, 1].max
+
+            forecast = {}
+
+            # Daily forecast
+            if budget_config[:global_daily]
+              daily_rate = daily_current / hours_elapsed
+              projected_daily = daily_rate * 24
+              forecast[:daily] = {
+                current: daily_current.round(4),
+                projected: projected_daily.round(4),
+                limit: budget_config[:global_daily],
+                on_track: projected_daily <= budget_config[:global_daily],
+                hours_remaining: (24 - hours_elapsed).round(1),
+                rate_per_hour: daily_rate.round(6)
+              }
+            end
+
+            # Monthly forecast
+            if budget_config[:global_monthly]
+              monthly_rate = monthly_current / days_elapsed
+              projected_monthly = monthly_rate * days_in_month
+              days_remaining = days_in_month - day_of_month
+              forecast[:monthly] = {
+                current: monthly_current.round(4),
+                projected: projected_monthly.round(4),
+                limit: budget_config[:global_monthly],
+                on_track: projected_monthly <= budget_config[:global_monthly],
+                days_remaining: days_remaining,
+                rate_per_day: monthly_rate.round(4)
+              }
+            end
+
+            forecast.presence
+          end
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/agents/infrastructure/budget/spend_recorder.rb

@@ -0,0 +1,279 @@
+# frozen_string_literal: true
+
+require_relative "../cache_helper"
+
+module RubyLLM
+  module Agents
+    module Budget
+      # Records spend and token usage, and handles soft cap alerting
+      #
+      # @api private
+      module SpendRecorder
+        extend CacheHelper
+
+        class << self
+          # Records spend and checks for soft cap alerts
+          #
+          # @param agent_type [String] The agent class name
+          # @param amount [Float] The amount spent in USD
+          # @param tenant_id [String, nil] The tenant identifier
+          # @param budget_config [Hash] Budget configuration
+          # @return [void]
+          def record_spend!(agent_type, amount, tenant_id:, budget_config:)
+            return if amount.nil? || amount <= 0
+
+            # Increment all relevant counters
+            increment_spend(:global, :daily, amount, tenant_id: tenant_id)
+            increment_spend(:global, :monthly, amount, tenant_id: tenant_id)
+            increment_spend(:agent, :daily, amount, agent_type: agent_type, tenant_id: tenant_id)
+            increment_spend(:agent, :monthly, amount, agent_type: agent_type, tenant_id: tenant_id)
+
+            # Check for soft cap alerts
+            check_soft_cap_alerts(agent_type, tenant_id, budget_config) if budget_config[:enabled]
+          end
+
+          # Records token usage and checks for soft cap alerts
+          #
+          # @param agent_type [String] The agent class name
+          # @param tokens [Integer] The number of tokens used
+          # @param tenant_id [String, nil] The tenant identifier
+          # @param budget_config [Hash] Budget configuration
+          # @return [void]
+          def record_tokens!(agent_type, tokens, tenant_id:, budget_config:)
+            return if tokens.nil? || tokens <= 0
+
+            # Increment global token counters (daily and monthly)
+            # Note: We only track global token usage, not per-agent (scope is ignored in increment_tokens)
+            increment_tokens(:global, :daily, tokens, tenant_id: tenant_id)
+            increment_tokens(:global, :monthly, tokens, tenant_id: tenant_id)
+
+            # Check for soft cap alerts
+            check_soft_token_alerts(agent_type, tenant_id, budget_config) if budget_config[:enabled]
+          end
+
+          # Increments the spend counter for a scope and period
+          #
+          # @param scope [Symbol] :global or :agent
+          # @param period [Symbol] :daily or :monthly
+          # @param amount [Float] Amount to add
+          # @param agent_type [String, nil] Required when scope is :agent
+          # @param tenant_id [String, nil] The tenant identifier
+          # @return [Float] New total
+          def increment_spend(scope, period, amount, agent_type: nil, tenant_id: nil)
+            key = budget_cache_key(scope, period, agent_type: agent_type, tenant_id: tenant_id)
+            ttl = period == :daily ? 1.day : 31.days
+
+            # Read-modify-write for float values (cache increment is for integers)
+            current = (SpendRecorder.cache_read(key) || 0).to_f
+            new_total = current + amount
+            SpendRecorder.cache_write(key, new_total, expires_in: ttl)
+            new_total
+          end
+
+          # Increments the token counter for a period
+          #
+          # @param scope [Symbol] :global (only global supported for tokens)
+          # @param period [Symbol] :daily or :monthly
+          # @param tokens [Integer] Tokens to add
+          # @param agent_type [String, nil] Not used for tokens
+          # @param tenant_id [String, nil] The tenant identifier
+          # @return [Integer] New total
+          def increment_tokens(scope, period, tokens, agent_type: nil, tenant_id: nil)
+            # For now, we only track global token usage (not per-agent)
+            key = token_cache_key(period, tenant_id: tenant_id)
+            ttl = period == :daily ? 1.day : 31.days
+
+            current = (SpendRecorder.cache_read(key) || 0).to_i
+            new_total = current + tokens
+            SpendRecorder.cache_write(key, new_total, expires_in: ttl)
+            new_total
+          end
+
+          # Returns the tenant key part for cache keys
+          #
+          # @param tenant_id [String, nil] The tenant identifier
+          # @return [String] "tenant:{id}" or "global"
+          def tenant_key_part(tenant_id)
+            tenant_id.present? ? "tenant:#{tenant_id}" : "global"
+          end
+
+          # Returns the date key part for cache keys based on period
+          #
+          # @param period [Symbol] :daily or :monthly
+          # @return [String] Date string
+          def date_key_part(period)
+            period == :daily ? Date.current.to_s : Date.current.strftime("%Y-%m")
+          end
+
+          # Generates an alert cache key
+          #
+          # @param alert_type [String] Type of alert (e.g., "budget_alert", "token_alert")
+          # @param scope [Symbol] Alert scope
+          # @param tenant_id [String, nil] The tenant identifier
+          # @return [String] Cache key
+          def alert_cache_key(alert_type, scope, tenant_id)
+            SpendRecorder.cache_key(alert_type, tenant_key_part(tenant_id), scope, Date.current.to_s)
+          end
+
+          # Generates a cache key for budget tracking
+          #
+          # @param scope [Symbol] :global or :agent
+          # @param period [Symbol] :daily or :monthly
+          # @param agent_type [String, nil] Required when scope is :agent
+          # @param tenant_id [String, nil] The tenant identifier
+          # @return [String] Cache key
+          def budget_cache_key(scope, period, agent_type: nil, tenant_id: nil)
+            date_part = date_key_part(period)
+            tenant_part = tenant_key_part(tenant_id)
+
+            case scope
+            when :global
+              SpendRecorder.cache_key("budget", tenant_part, date_part)
+            when :agent
+              SpendRecorder.cache_key("budget", tenant_part, "agent", agent_type, date_part)
+            else
+              raise ArgumentError, "Unknown scope: #{scope}"
+            end
+          end
+
+          # Generates a cache key for token tracking
+          #
+          # @param period [Symbol] :daily or :monthly
+          # @param tenant_id [String, nil] The tenant identifier
+          # @return [String] Cache key
+          def token_cache_key(period, tenant_id: nil)
+            SpendRecorder.cache_key("tokens", tenant_key_part(tenant_id), date_key_part(period))
+          end
+
+          private
+
+          # Checks for soft cap alerts after recording spend
+          #
+          # @param agent_type [String] The agent class name
+          # @param tenant_id [String, nil] The tenant identifier
+          # @param budget_config [Hash] Budget configuration
+          # @return [void]
+          def check_soft_cap_alerts(agent_type, tenant_id, budget_config)
+            config = RubyLLM::Agents.configuration
+            return unless config.alerts_enabled?
+            return unless config.alert_events.include?(:budget_soft_cap) || config.alert_events.include?(:budget_hard_cap)
+
+            # Check global daily
+            check_budget_alert(:global_daily, budget_config[:global_daily],
+                              BudgetQuery.current_spend(:global, :daily, tenant_id: tenant_id),
+                              agent_type, tenant_id, budget_config)
+
+            # Check global monthly
+            check_budget_alert(:global_monthly, budget_config[:global_monthly],
+                              BudgetQuery.current_spend(:global, :monthly, tenant_id: tenant_id),
+                              agent_type, tenant_id, budget_config)
+
+            # Check per-agent daily
+            agent_daily_limit = budget_config[:per_agent_daily]&.dig(agent_type)
+            if agent_daily_limit
+              check_budget_alert(:per_agent_daily, agent_daily_limit,
+                                BudgetQuery.current_spend(:agent, :daily, agent_type: agent_type, tenant_id: tenant_id),
+                                agent_type, tenant_id, budget_config)
+            end
+
+            # Check per-agent monthly
+            agent_monthly_limit = budget_config[:per_agent_monthly]&.dig(agent_type)
+            if agent_monthly_limit
+              check_budget_alert(:per_agent_monthly, agent_monthly_limit,
+                                BudgetQuery.current_spend(:agent, :monthly, agent_type: agent_type, tenant_id: tenant_id),
+                                agent_type, tenant_id, budget_config)
+            end
+          end
+
+          # Checks if an alert should be fired for a budget
+          #
+          # @param scope [Symbol] Budget scope
+          # @param limit [Float, nil] Budget limit
+          # @param current [Float] Current spend
+          # @param agent_type [String] Agent type
+          # @param tenant_id [String, nil] The tenant identifier
+          # @param budget_config [Hash] Budget configuration
+          # @return [void]
+          def check_budget_alert(scope, limit, current, agent_type, tenant_id, budget_config)
+            return unless limit
+            return if current <= limit
+
+            event = budget_config[:enforcement] == :hard ? :budget_hard_cap : :budget_soft_cap
+            config = RubyLLM::Agents.configuration
+            return unless config.alert_events.include?(event)
+
+            # Prevent duplicate alerts by using a cache key (include tenant for isolation)
+            key = alert_cache_key("budget_alert", scope, tenant_id)
+            return if SpendRecorder.cache_exist?(key)
+
+            SpendRecorder.cache_write(key, true, expires_in: 1.hour)
+
+            AlertManager.notify(event, {
+              scope: scope,
+              limit: limit,
+              total: current.round(6),
+              agent_type: agent_type,
+              tenant_id: tenant_id,
+              timestamp: Date.current.to_s
+            })
+          end
+
+          # Checks for soft cap token alerts after recording usage
+          #
+          # @param agent_type [String] The agent class name
+          # @param tenant_id [String, nil] The tenant identifier
+          # @param budget_config [Hash] Budget configuration
+          # @return [void]
+          def check_soft_token_alerts(agent_type, tenant_id, budget_config)
+            config = RubyLLM::Agents.configuration
+            return unless config.alerts_enabled?
+            return unless config.alert_events.include?(:token_soft_cap) || config.alert_events.include?(:token_hard_cap)
+
+            # Check global daily tokens
+            check_token_alert(:global_daily_tokens, budget_config[:global_daily_tokens],
+                             BudgetQuery.current_tokens(:daily, tenant_id: tenant_id),
+                             agent_type, tenant_id, budget_config)
+
+            # Check global monthly tokens
+            check_token_alert(:global_monthly_tokens, budget_config[:global_monthly_tokens],
+                             BudgetQuery.current_tokens(:monthly, tenant_id: tenant_id),
+                             agent_type, tenant_id, budget_config)
+          end
+
+          # Checks if a token alert should be fired
+          #
+          # @param scope [Symbol] Token scope
+          # @param limit [Integer, nil] Token limit
+          # @param current [Integer] Current token usage
+          # @param agent_type [String] Agent type
+          # @param tenant_id [String, nil] The tenant identifier
+          # @param budget_config [Hash] Budget configuration
+          # @return [void]
+          def check_token_alert(scope, limit, current, agent_type, tenant_id, budget_config)
+            return unless limit
+            return if current <= limit
+
+            event = budget_config[:enforcement] == :hard ? :token_hard_cap : :token_soft_cap
+            config = RubyLLM::Agents.configuration
+            return unless config.alert_events.include?(event)
+
+            # Prevent duplicate alerts
+            key = alert_cache_key("token_alert", scope, tenant_id)
+            return if SpendRecorder.cache_exist?(key)
+
+            SpendRecorder.cache_write(key, true, expires_in: 1.hour)
+
+            AlertManager.notify(event, {
+              scope: scope,
+              limit: limit,
+              total: current,
+              agent_type: agent_type,
+              tenant_id: tenant_id,
+              timestamp: Date.current.to_s
+            })
+          end
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/agents/infrastructure/budget_tracker.rb

@@ -0,0 +1,275 @@
+# frozen_string_literal: true
+
+require_relative "cache_helper"
+require_relative "budget/config_resolver"
+require_relative "budget/spend_recorder"
+require_relative "budget/budget_query"
+require_relative "budget/forecaster"
+
+module RubyLLM
+  module Agents
+    # Cache-based budget tracking for cost governance
+    #
+    # Tracks spending against configured budget limits using cache counters.
+    # Supports daily and monthly budgets at both global and per-agent levels.
+    # In multi-tenant mode, budgets are tracked separately per tenant.
+    #
+    # Note: Uses best-effort enforcement with cache counters. In high-concurrency
+    # scenarios, slight overruns may occur due to race conditions. This is an
+    # acceptable trade-off for performance.
+    #
+    # @example Checking budget before execution
+    #   BudgetTracker.check_budget!("MyAgent")  # raises BudgetExceededError if over limit
+    #
+    # @example Recording spend after execution
+    #   BudgetTracker.record_spend!("MyAgent", 0.05)
+    #
+    # @example Multi-tenant usage
+    #   BudgetTracker.check_budget!("MyAgent", tenant_id: "acme_corp")
+    #   BudgetTracker.record_spend!("MyAgent", 0.05, tenant_id: "acme_corp")
+    #
+    # @see RubyLLM::Agents::Configuration
+    # @see RubyLLM::Agents::Reliability::BudgetExceededError
+    # @see RubyLLM::Agents::TenantBudget
+    # @api public
+    module BudgetTracker
+      extend CacheHelper
+
+      class << self
+        # Checks if the current spend exceeds budget limits
+        #
+        # @param agent_type [String] The agent class name
+        # @param tenant_id [String, nil] Optional tenant identifier (uses resolver if not provided)
+        # @param tenant_config [Hash, nil] Optional runtime tenant config (takes priority over resolver/DB)
+        # @raise [Reliability::BudgetExceededError] If hard cap is exceeded
+        # @return [void]
+        def check_budget!(agent_type, tenant_id: nil, tenant_config: nil)
+          tenant_id = Budget::ConfigResolver.resolve_tenant_id(tenant_id)
+          budget_config = Budget::ConfigResolver.resolve_budget_config(tenant_id, runtime_config: tenant_config)
+
+          return unless budget_config[:enabled]
+          return unless budget_config[:enforcement] == :hard
+
+          check_budget_limits!(agent_type, tenant_id, budget_config)
+        end
+
+        # Checks if the current token usage exceeds budget limits
+        #
+        # @param agent_type [String] The agent class name
+        # @param tenant_id [String, nil] Optional tenant identifier (uses resolver if not provided)
+        # @param tenant_config [Hash, nil] Optional runtime tenant config (takes priority over resolver/DB)
+        # @raise [Reliability::BudgetExceededError] If hard cap is exceeded
+        # @return [void]
+        def check_token_budget!(agent_type, tenant_id: nil, tenant_config: nil)
+          tenant_id = Budget::ConfigResolver.resolve_tenant_id(tenant_id)
+          budget_config = Budget::ConfigResolver.resolve_budget_config(tenant_id, runtime_config: tenant_config)
+
+          return unless budget_config[:enabled]
+          return unless budget_config[:enforcement] == :hard
+
+          check_token_limits!(agent_type, tenant_id, budget_config)
+        end
+
+        # Records spend and checks for soft cap alerts
+        #
+        # @param agent_type [String] The agent class name
+        # @param amount [Float] The amount spent in USD
+        # @param tenant_id [String, nil] Optional tenant identifier (uses resolver if not provided)
+        # @param tenant_config [Hash, nil] Optional runtime tenant config (takes priority over resolver/DB)
+        # @return [void]
+        def record_spend!(agent_type, amount, tenant_id: nil, tenant_config: nil)
+          return if amount.nil? || amount <= 0
+
+          tenant_id = Budget::ConfigResolver.resolve_tenant_id(tenant_id)
+          budget_config = Budget::ConfigResolver.resolve_budget_config(tenant_id, runtime_config: tenant_config)
+
+          Budget::SpendRecorder.record_spend!(agent_type, amount, tenant_id: tenant_id, budget_config: budget_config)
+        end
+
+        # Records token usage and checks for soft cap alerts
+        #
+        # @param agent_type [String] The agent class name
+        # @param tokens [Integer] The number of tokens used
+        # @param tenant_id [String, nil] Optional tenant identifier (uses resolver if not provided)
+        # @param tenant_config [Hash, nil] Optional runtime tenant config (takes priority over resolver/DB)
+        # @return [void]
+        def record_tokens!(agent_type, tokens, tenant_id: nil, tenant_config: nil)
+          return if tokens.nil? || tokens <= 0
+
+          tenant_id = Budget::ConfigResolver.resolve_tenant_id(tenant_id)
+          budget_config = Budget::ConfigResolver.resolve_budget_config(tenant_id, runtime_config: tenant_config)
+
+          Budget::SpendRecorder.record_tokens!(agent_type, tokens, tenant_id: tenant_id, budget_config: budget_config)
+        end
+
+        # Returns the current spend for a scope and period
+        #
+        # @param scope [Symbol] :global or :agent
+        # @param period [Symbol] :daily or :monthly
+        # @param agent_type [String, nil] Required when scope is :agent
+        # @param tenant_id [String, nil] Optional tenant identifier (uses resolver if not provided)
+        # @return [Float] Current spend in USD
+        def current_spend(scope, period, agent_type: nil, tenant_id: nil)
+          tenant_id = Budget::ConfigResolver.resolve_tenant_id(tenant_id)
+          Budget::BudgetQuery.current_spend(scope, period, agent_type: agent_type, tenant_id: tenant_id)
+        end
+
+        # Returns the current token usage for a period (global only)
+        #
+        # @param period [Symbol] :daily or :monthly
+        # @param tenant_id [String, nil] Optional tenant identifier (uses resolver if not provided)
+        # @return [Integer] Current token usage
+        def current_tokens(period, tenant_id: nil)
+          tenant_id = Budget::ConfigResolver.resolve_tenant_id(tenant_id)
+          Budget::BudgetQuery.current_tokens(period, tenant_id: tenant_id)
+        end
+
+        # Returns the remaining budget for a scope and period
+        #
+        # @param scope [Symbol] :global or :agent
+        # @param period [Symbol] :daily or :monthly
+        # @param agent_type [String, nil] Required when scope is :agent
+        # @param tenant_id [String, nil] Optional tenant identifier (uses resolver if not provided)
+        # @return [Float, nil] Remaining budget in USD, or nil if no limit configured
+        def remaining_budget(scope, period, agent_type: nil, tenant_id: nil)
+          tenant_id = Budget::ConfigResolver.resolve_tenant_id(tenant_id)
+          budget_config = Budget::ConfigResolver.resolve_budget_config(tenant_id)
+
+          Budget::BudgetQuery.remaining_budget(scope, period, agent_type: agent_type, tenant_id: tenant_id, budget_config: budget_config)
+        end
+
+        # Returns the remaining token budget for a period (global only)
+        #
+        # @param period [Symbol] :daily or :monthly
+        # @param tenant_id [String, nil] Optional tenant identifier (uses resolver if not provided)
+        # @return [Integer, nil] Remaining token budget, or nil if no limit configured
+        def remaining_token_budget(period, tenant_id: nil)
+          tenant_id = Budget::ConfigResolver.resolve_tenant_id(tenant_id)
+          budget_config = Budget::ConfigResolver.resolve_budget_config(tenant_id)
+
+          Budget::BudgetQuery.remaining_token_budget(period, tenant_id: tenant_id, budget_config: budget_config)
+        end
+
+        # Returns a summary of all budget statuses
+        #
+        # @param agent_type [String, nil] Optional agent type for per-agent budgets
+        # @param tenant_id [String, nil] Optional tenant identifier (uses resolver if not provided)
+        # @return [Hash] Budget status information
+        def status(agent_type: nil, tenant_id: nil)
+          tenant_id = Budget::ConfigResolver.resolve_tenant_id(tenant_id)
+          budget_config = Budget::ConfigResolver.resolve_budget_config(tenant_id)
+
+          Budget::BudgetQuery.status(agent_type: agent_type, tenant_id: tenant_id, budget_config: budget_config)
+        end
+
+        # Calculates budget forecasts based on current spending trends
+        #
+        # @param tenant_id [String, nil] Optional tenant identifier (uses resolver if not provided)
+        # @return [Hash, nil] Forecast information
+        def calculate_forecast(tenant_id: nil)
+          tenant_id = Budget::ConfigResolver.resolve_tenant_id(tenant_id)
+          budget_config = Budget::ConfigResolver.resolve_budget_config(tenant_id)
+
+          Budget::Forecaster.calculate_forecast(tenant_id: tenant_id, budget_config: budget_config)
+        end
+
+        # Resets all budget counters (useful for testing)
+        #
+        # @param tenant_id [String, nil] Optional tenant identifier to reset only that tenant's counters
+        # @return [void]
+        def reset!(tenant_id: nil)
+          tenant_id = Budget::ConfigResolver.resolve_tenant_id(tenant_id)
+          tenant_part = Budget::SpendRecorder.tenant_key_part(tenant_id)
+          today = Budget::SpendRecorder.date_key_part(:daily)
+          month = Budget::SpendRecorder.date_key_part(:monthly)
+
+          BudgetTracker.cache_delete(BudgetTracker.cache_key("budget", tenant_part, today))
+          BudgetTracker.cache_delete(BudgetTracker.cache_key("budget", tenant_part, month))
+
+          # Reset memoized table existence check (useful for testing)
+          Budget::ConfigResolver.reset_tenant_budget_table_check!
+        end
+
+        private
+
+        # Checks budget limits and raises error if exceeded
+        #
+        # @param agent_type [String] The agent class name
+        # @param tenant_id [String, nil] The tenant identifier
+        # @param budget_config [Hash] The budget configuration
+        # @raise [Reliability::BudgetExceededError] If limit exceeded
+        # @return [void]
+        def check_budget_limits!(agent_type, tenant_id, budget_config)
+          # Check global daily budget
+          if budget_config[:global_daily]
+            current = Budget::BudgetQuery.current_spend(:global, :daily, tenant_id: tenant_id)
+            if current >= budget_config[:global_daily]
+              raise Reliability::BudgetExceededError.new(:global_daily, budget_config[:global_daily], current, tenant_id: tenant_id)
+            end
+          end
+
+          # Check global monthly budget
+          if budget_config[:global_monthly]
+            current = Budget::BudgetQuery.current_spend(:global, :monthly, tenant_id: tenant_id)
+            if current >= budget_config[:global_monthly]
+              raise Reliability::BudgetExceededError.new(:global_monthly, budget_config[:global_monthly], current, tenant_id: tenant_id)
+            end
+          end
+
+          # Check per-agent daily budget
+          agent_daily_limit = budget_config[:per_agent_daily]&.dig(agent_type)
+          if agent_daily_limit
+            current = Budget::BudgetQuery.current_spend(:agent, :daily, agent_type: agent_type, tenant_id: tenant_id)
+            if current >= agent_daily_limit
+              raise Reliability::BudgetExceededError.new(:per_agent_daily, agent_daily_limit, current, agent_type: agent_type, tenant_id: tenant_id)
+            end
+          end
+
+          # Check per-agent monthly budget
+          agent_monthly_limit = budget_config[:per_agent_monthly]&.dig(agent_type)
+          if agent_monthly_limit
+            current = Budget::BudgetQuery.current_spend(:agent, :monthly, agent_type: agent_type, tenant_id: tenant_id)
+            if current >= agent_monthly_limit
+              raise Reliability::BudgetExceededError.new(:per_agent_monthly, agent_monthly_limit, current, agent_type: agent_type, tenant_id: tenant_id)
+            end
+          end
+        end
+
+        # Checks token limits and raises error if exceeded
+        #
+        # @param agent_type [String] The agent class name
+        # @param tenant_id [String, nil] The tenant identifier
+        # @param budget_config [Hash] The budget configuration
+        # @raise [Reliability::BudgetExceededError] If limit exceeded
+        # @return [void]
+        def check_token_limits!(agent_type, tenant_id, budget_config)
+          # Check global daily token budget
+          if budget_config[:global_daily_tokens]
+            current = Budget::BudgetQuery.current_tokens(:daily, tenant_id: tenant_id)
+            if current >= budget_config[:global_daily_tokens]
+              raise Reliability::BudgetExceededError.new(
+                :global_daily_tokens,
+                budget_config[:global_daily_tokens],
+                current,
+                tenant_id: tenant_id
+              )
+            end
+          end
+
+          # Check global monthly token budget
+          if budget_config[:global_monthly_tokens]
+            current = Budget::BudgetQuery.current_tokens(:monthly, tenant_id: tenant_id)
+            if current >= budget_config[:global_monthly_tokens]
+              raise Reliability::BudgetExceededError.new(
+                :global_monthly_tokens,
+                budget_config[:global_monthly_tokens],
+                current,
+                tenant_id: tenant_id
+              )
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/agents/{execution_logger_job.rb → infrastructure/execution_logger_job.rb}

@@ -25,7 +25,9 @@ module RubyLLM
       # @param execution_data [Hash] Execution attributes from instrumentation
       # @return [void]
       def perform(execution_data)
-        execution = Execution.create!(execution_data)
+        # Filter to only known attributes to prevent schema mismatches
+        filtered_data = filter_known_attributes(execution_data)
+        execution = Execution.create!(filtered_data)
 
         # Calculate costs if token data is available
         if execution.input_tokens && execution.output_tokens
@@ -39,6 +41,20 @@ module RubyLLM
 
       private
 
+      # Filters data to only include attributes that exist on the Execution model
+      #
+      # This provides a safety net against schema mismatches, such as when
+      # tenant_id is passed but the column doesn't exist in the database.
+      #
+      # @param data [Hash] The raw execution data
+      # @return [Hash] Filtered data with only known attributes
+      def filter_known_attributes(data)
+        return data unless defined?(Execution) && Execution.respond_to?(:column_names)
+
+        known_columns = Execution.column_names
+        data.select { |key, _| known_columns.include?(key.to_s) }
+      end
+
       # Checks if execution should be flagged as anomalous
       #
       # @param execution [Execution] The execution to check

data/lib/ruby_llm/agents/{reliability → infrastructure/reliability}/executor.rb

@@ -32,7 +32,8 @@ module RubyLLM
             backoff: retries_config[:backoff] || :exponential,
             base: retries_config[:base] || 0.4,
             max_delay: retries_config[:max_delay] || 3.0,
-            on: retries_config[:on] || []
+            on: retries_config[:on] || [],
+            patterns: config[:retryable_patterns]
           )
 
           @fallback_routing = FallbackRouting.new(