llm_cost_tracker 0.9.0 → 0.11.0

data/lib/llm_cost_tracker/capture/stream_collector.rb

@@ -69,7 +69,6 @@ module LlmCostTracker
           ensure_open!
           capture_event(data, type: type) unless data.nil?
         end
-        self
       end
 
       def usage(input_tokens:, output_tokens:, **extra)
@@ -87,7 +86,6 @@ module LlmCostTracker
             output_tokens: output_tokens
           )
         end
-        self
       end
 
       def finish!(errored: false)
@@ -104,7 +102,7 @@ module LlmCostTracker
           return nil if @finished || @recording
 
           @recording = true
-          pricing_mode = Pricing.normalize_mode(@pricing_mode)
+          pricing_mode = Pricing::Mode.normalize(@pricing_mode)
           {
             events: @events.dup,
             overflowed: @overflowed,
@@ -124,17 +122,17 @@ module LlmCostTracker
       def record_snapshot(snapshot, errored:)
         save_succeeded = false
         begin
-          capture = build_usage_capture(snapshot)
-          provider_response_id = capture.provider_response_id || snapshot[:provider_response_id]
-          capture = capture.with(provider_response_id: provider_response_id)
+          event = build_event(snapshot)
+          provider_response_id = event.provider_response_id || snapshot[:provider_response_id]
+          event = event.with(provider_response_id: provider_response_id)
 
           Tracker.record(
-            capture: capture,
+            event: event,
             latency_ms: snapshot[:latency_ms] || LlmCostTracker::Timing.elapsed_ms(@started_at),
-            pricing_mode: pricing_mode_for(capture: capture, snapshot: snapshot),
+            pricing_mode: Pricing::Mode.merge(event.pricing_mode, snapshot[:pricing_mode]),
             metadata: (errored ? { stream_errored: true } : {}).merge(snapshot[:metadata]),
             context_tags: snapshot[:context_tags]
-          ) { |stage| save_succeeded = true if stage == :after_save }
+          ) { save_succeeded = true }
         ensure
           @mutex.synchronize do
             @finished = save_succeeded
@@ -143,27 +141,8 @@ module LlmCostTracker
         end
       end
 
-      HOST_DERIVED_MODE_TOKENS = %i[data_residency].freeze
-      STANDARD_LIKE_MODE_TOKENS = %i[standard standard_only auto default].freeze
-      private_constant :HOST_DERIVED_MODE_TOKENS, :STANDARD_LIKE_MODE_TOKENS
-
-      def pricing_mode_for(capture:, snapshot:)
-        merge_pricing_modes(capture.pricing_mode, snapshot[:pricing_mode])
-      end
-
-      def merge_pricing_modes(provider_mode, request_mode)
-        return Pricing.normalize_mode(request_mode) if provider_mode.to_s.strip.empty?
-
-        provider_tokens = Pricing::Mode.tokenize(provider_mode) - STANDARD_LIKE_MODE_TOKENS
-        request_host_tokens = Pricing::Mode.tokenize(request_mode || "") & HOST_DERIVED_MODE_TOKENS
-        combined = provider_tokens | request_host_tokens
-        return nil if combined.empty?
-
-        Pricing.normalize_mode(combined.join("_"))
-      end
-
       def capture_dimensions(pricing_mode)
-        batch = @batch.nil? ? UsageCapture.batch_from_pricing_mode?(pricing_mode).presence : @batch
+        batch = @batch.nil? ? Event.batch_from_pricing_mode?(pricing_mode).presence : @batch
         {
           provider_project_id: @provider_project_id.to_s.strip.presence,
           provider_api_key_id: @provider_api_key_id.to_s.strip.presence,
@@ -178,18 +157,18 @@ module LlmCostTracker
         raise FrozenError, "can't modify finished LlmCostTracker::Capture::StreamCollector"
       end
 
-      def build_usage_capture(snapshot)
+      def build_event(snapshot)
         return build_from_explicit_usage(snapshot) if snapshot[:explicit_usage]
         return build_unknown_usage(snapshot) if snapshot[:overflowed]
 
-        capture = Parsers.find_for_provider(@provider)&.parse_stream(
+        event = Parsers.find_for_provider(@provider)&.parse_stream(
           response_status: 200,
           events: snapshot[:events],
           request_body: request_body_for(snapshot[:request])
         )
-        if capture
-          model = present_model(capture.model) || present_model(snapshot[:model]) || UsageCapture::UNKNOWN_MODEL
-          return capture.with(provider: @provider, model: model, **snapshot.fetch(:capture_dimensions))
+        if event
+          model = present_model(event.model) || present_model(snapshot[:model]) || Event::UNKNOWN_MODEL
+          return event.with(provider: @provider, model: model, **snapshot.fetch(:capture_dimensions))
         end
 
         build_unknown_usage(snapshot)
@@ -207,15 +186,15 @@ module LlmCostTracker
         return nil if value.nil?
 
         string = value.to_s.presence
-        return nil if string.nil? || string == "unknown"
+        return nil if string.nil? || string == Event::UNKNOWN_MODEL
 
         string
       end
 
       def build_from_explicit_usage(snapshot)
-        UsageCapture.build(
+        Event.build(
           provider: @provider,
-          model: snapshot[:model] || UsageCapture::UNKNOWN_MODEL,
+          model: snapshot[:model] || Event::UNKNOWN_MODEL,
           token_usage: snapshot[:explicit_usage],
           stream: true,
           usage_source: :manual,
@@ -225,9 +204,9 @@ module LlmCostTracker
       end
 
       def build_unknown_usage(snapshot)
-        UsageCapture.build(
+        Event.build(
           provider: @provider,
-          model: snapshot[:model] || UsageCapture::UNKNOWN_MODEL,
+          model: snapshot[:model] || Event::UNKNOWN_MODEL,
           token_usage: TokenUsage.build(input_tokens: 0, output_tokens: 0, total_tokens: 0),
           stream: true,
           usage_source: :unknown,
@@ -244,14 +223,14 @@ module LlmCostTracker
 
       def capture_event(data, type:)
         event = { event: type, data: strip_heavy_payload(data) }
-        size = JSON.generate(event).bytesize
+        size = approximate_bytesize(event)
         if @captured_bytes + size <= Capture::Stream::LIMIT_BYTES
-          @events << event.deep_dup
+          @events << event
           @captured_bytes += size
         else
           @overflowed = true
         end
-      rescue JSON::JSONError, TypeError, SystemStackError
+      rescue TypeError, SystemStackError
         @overflowed = true
       end
 
@@ -271,6 +250,19 @@ module LlmCostTracker
           value
         end
       end
+
+      def approximate_bytesize(value)
+        case value
+        when Hash
+          value.sum { |key, nested| approximate_bytesize(key) + approximate_bytesize(nested) + 4 }
+        when Array
+          value.sum { |nested| approximate_bytesize(nested) + 2 }
+        when Numeric, true, false, nil
+          8
+        else
+          value.to_s.bytesize + 2
+        end
+      end
     end
   end
 end

data/lib/llm_cost_tracker/capture/stream_tracker.rb

@@ -74,11 +74,7 @@ module LlmCostTracker
       end
 
       def capture(event)
-        raw_payload = event.try(:deep_to_h) || event.try(:to_h)
-        raw_payload ||= %i[type id model usage response message].each_with_object({}) do |key, attributes|
-          value = event.try(key)
-          attributes[key] = value unless value.nil?
-        end
+        raw_payload = event.try(:deep_to_h) || event.try(:to_h) || {}
         payload = normalize(raw_payload)
         type = event.try(:type) || payload["type"]
         @collector.event(payload, type: type&.to_s)
@@ -155,7 +151,7 @@ module LlmCostTracker
           next unless should_finish && active_proc.call
 
           begin
-            finish_proc.call(false)
+            Rails.application.executor.wrap { finish_proc.call(false) }
           rescue StandardError
             nil
           end

data/lib/llm_cost_tracker/configuration.rb

@@ -14,28 +14,29 @@ module LlmCostTracker
 
     BUDGET_EXCEEDED_BEHAVIORS = %i[notify raise block_requests].freeze
     UNKNOWN_PRICING_BEHAVIORS = %i[ignore warn raise].freeze
-    SHARED_SCALAR_ATTRIBUTES = %i[enabled default_tags on_budget_exceeded monthly_budget daily_budget per_call_budget
-                                  log_level prices_file max_tag_count max_tag_value_bytesize].freeze
-    SHARED_ENUM_ATTRIBUTES = {
+    INGESTION_MODES = %i[inline async].freeze
+    SCALAR_ATTRIBUTES = %i[enabled default_tags on_budget_exceeded monthly_budget daily_budget per_call_budget
+                           log_level prices_file max_tag_count max_tag_value_bytesize
+                           ingestion_pool_size auto_enable_stream_usage cache_rollups
+                           reconciliation_enabled].freeze
+    ENUM_ATTRIBUTES = {
       budget_exceeded_behavior: [BUDGET_EXCEEDED_BEHAVIORS, :notify],
-      unknown_pricing_behavior: [UNKNOWN_PRICING_BEHAVIORS, :warn]
+      unknown_pricing_behavior: [UNKNOWN_PRICING_BEHAVIORS, :warn],
+      ingestion: [INGESTION_MODES, :inline]
     }.freeze
     DEFAULT_REDACTED_TAG_KEYS = %w[api_key access_token authorization credential password refresh_token secret].freeze
 
     attr_reader(
-      *SHARED_SCALAR_ATTRIBUTES,
+      *SCALAR_ATTRIBUTES,
       :budget_exceeded_behavior,
-      :durable_ingestion,
+      :ingestion,
       :instrumented_integrations,
       :pricing_overrides,
       :report_tag_breakdowns,
       :redacted_tag_keys,
       :unknown_pricing_behavior,
       :openai_compatible_providers,
-      :reconciliation_importers,
-      :reconciliation_enabled,
-      :auto_enable_stream_usage,
-      :cache_rollups
+      :reconciliation_importers
     )
 
     def initialize
@@ -51,6 +52,7 @@ module LlmCostTracker
       @prices_file        = nil
       @max_tag_count      = 50
       @max_tag_value_bytesize = 1024
+      @ingestion_pool_size = nil
       self.pricing_overrides = {}
       @instrumented_integrations = Set.new
       @report_tag_breakdowns = []
@@ -59,33 +61,13 @@ module LlmCostTracker
       @reconciliation_importers = {}
       @reconciliation_enabled = false
       @auto_enable_stream_usage = true
-      @durable_ingestion = false
+      self.ingestion = :inline
       @cache_rollups = false
       @finalized = false
     end
 
-    def durable_ingestion=(value)
-      ensure_shared_configuration_mutable!
-      @durable_ingestion = value
-    end
-
-    def cache_rollups=(value)
-      ensure_shared_configuration_mutable!
-      @cache_rollups = value
-    end
-
-    def reconciliation_enabled=(value)
-      ensure_shared_configuration_mutable!
-      @reconciliation_enabled = value
-    end
-
-    def auto_enable_stream_usage=(value)
-      ensure_shared_configuration_mutable!
-      @auto_enable_stream_usage = value
-    end
-
     def reconciliation_importers=(importers)
-      ensure_shared_configuration_mutable!
+      ensure_mutable!
       raise Error, RECONCILIATION_DISABLED_MESSAGE unless @reconciliation_enabled
 
       @reconciliation_importers = (importers || {}).to_h do |source, importer|
@@ -96,7 +78,7 @@ module LlmCostTracker
     end
 
     def register_reconciliation_importer(source, &block)
-      ensure_shared_configuration_mutable!
+      ensure_mutable!
       raise Error, RECONCILIATION_DISABLED_MESSAGE unless @reconciliation_enabled
       raise Error, "register_reconciliation_importer requires a block" unless block
 
@@ -107,29 +89,29 @@ module LlmCostTracker
     private_constant :RECONCILIATION_DISABLED_MESSAGE
 
     def openai_compatible_providers=(providers)
-      ensure_shared_configuration_mutable!
+      ensure_mutable!
       @openai_compatible_providers = normalize_openai_compatible_providers(providers)
     end
 
     def pricing_overrides=(value)
-      ensure_shared_configuration_mutable!
+      ensure_mutable!
       @pricing_overrides = Pricing::Registry.normalize_price_table(value || {})
     rescue ArgumentError => e
       raise Error, "invalid pricing_overrides: #{e.message}"
     end
 
     def report_tag_breakdowns=(value)
-      ensure_shared_configuration_mutable!
+      ensure_mutable!
       @report_tag_breakdowns = Array(value).map { |key| Tags::Key.validate!(key, error_class: Error) }
     end
 
     def redacted_tag_keys=(value)
-      ensure_shared_configuration_mutable!
+      ensure_mutable!
       @redacted_tag_keys = Array(value).map(&:to_s)
     end
 
     def instrument(*names)
-      ensure_shared_configuration_mutable!
+      ensure_mutable!
       @instrumented_integrations.merge(normalize_instrumentation_names(names))
     end
 
@@ -137,16 +119,16 @@ module LlmCostTracker
       @instrumented_integrations.include?(name)
     end
 
-    SHARED_SCALAR_ATTRIBUTES.each do |name|
+    SCALAR_ATTRIBUTES.each do |name|
       define_method("#{name}=") do |value|
-        ensure_shared_configuration_mutable!
+        ensure_mutable!
         instance_variable_set(:"@#{name}", value)
       end
     end
 
-    SHARED_ENUM_ATTRIBUTES.each do |name, (allowed, default)|
+    ENUM_ATTRIBUTES.each do |name, (allowed, default)|
       define_method("#{name}=") do |value|
-        ensure_shared_configuration_mutable!
+        ensure_mutable!
         instance_variable_set(:"@#{name}", normalize_enum(name, value, allowed, default: default))
       end
     end
@@ -161,7 +143,11 @@ module LlmCostTracker
         normalize_openai_compatible_providers(@openai_compatible_providers)
       )
       @finalized = true
-      self
+    end
+
+    def normalized_redacted_tag_keys
+      @normalized_redacted_tag_keys ||=
+        Array(@redacted_tag_keys).map { |key| Tags::Sanitizer.normalized_key(key) }.freeze
     end
 
     def finalized?
@@ -196,7 +182,7 @@ module LlmCostTracker
       names
     end
 
-    def ensure_shared_configuration_mutable!
+    def ensure_mutable!
       return unless finalized?
 
       raise FrozenError, "can't modify frozen LlmCostTracker::Configuration"

data/lib/llm_cost_tracker/doctor/capture_verifier.rb

@@ -47,7 +47,7 @@ module LlmCostTracker
         end
 
         LlmCostTracker::Integrations.checks.map do |check|
-          Check.new(check.status, "sdk integration #{check.name}", check.message)
+          check.with(name: "sdk integration #{check.name}")
         end
       end
 

data/lib/llm_cost_tracker/doctor/ingestion_check.rb

@@ -11,14 +11,14 @@ module LlmCostTracker
 
       def call
         return unless Probe.table_exists?("llm_cost_tracker_calls")
-        return inline_check unless LlmCostTracker::Ingestion.durable?
+        return inline_check unless LlmCostTracker::Ingestion.async?
 
         missing = missing_parts
         if missing.empty?
           inbox = inbox_snapshot
           quarantined = inbox.try(:quarantined_count).to_i
           if quarantined.positive?
-            return Check.new(:warn, "durable ingestion", "#{quarantined} inbox entries quarantined after retries")
+            return Check.new(:warn, "async ingestion", "#{quarantined} inbox entries quarantined after retries")
           end
 
           pending_count = inbox.try(:pending_count).to_i
@@ -27,17 +27,17 @@ module LlmCostTracker
           if pending_count.positive? && pending_age && pending_age >= PENDING_AGE_WARNING_SECONDS
             return Check.new(
               :warn,
-              "durable ingestion",
+              "async ingestion",
               "#{pending_count} inbox entries pending; oldest pending age #{pending_age.round}s"
             )
           end
 
-          return Check.new(:ok, "durable ingestion", "inbox and ingestion lease tables available")
+          return Check.new(:ok, "async ingestion", "inbox and ingestion lease tables available")
         end
 
         Check.new(
           :error,
-          "durable ingestion",
+          "async ingestion",
           "missing #{missing.join(', ')}; see docs/upgrading.md for the recovery steps"
         )
       end
@@ -50,15 +50,15 @@ module LlmCostTracker
           return Check.new(
             :ok,
             "inline ingestion",
-            "durable_ingestion=false; events write directly to the ledger"
+            "config.ingestion = :inline; events write directly to the ledger"
           )
         end
 
         Check.new(
           :warn,
           "inline ingestion",
-          "durable_ingestion=false but found unused durable ingestion tables: #{leftovers.join(', ')}. " \
-          "Set config.durable_ingestion = true to keep the durable inbox path or drop the tables."
+          "config.ingestion = :inline but found unused async ingestion tables: #{leftovers.join(', ')}. " \
+          "Set config.ingestion = :async to keep the inbox path or drop the tables."
         )
       end
 

data/lib/llm_cost_tracker/doctor/legacy_audit_check.rb

@@ -28,8 +28,6 @@ module LlmCostTracker
         message = "#{missing}/#{total} tracked calls lack pricing_snapshot; " \
                   "stored totals remain stable but applied rates cannot be audited"
         Check.new(:warn, "pricing snapshot audit", message)
-      rescue StandardError
-        nil
       end
     end
   end

data/lib/llm_cost_tracker/doctor/legacy_billing_status_check.rb

@@ -14,8 +14,6 @@ module LlmCostTracker
         return unless LlmCostTracker::Call.where(cost_status: nil).exists?
 
         Check.new(:warn, "cost status", "legacy rows without cost_status remain; new rows will populate it")
-      rescue StandardError
-        nil
       end
     end
   end

data/lib/llm_cost_tracker/doctor.rb

@@ -14,35 +14,95 @@ require_relative "doctor/pricing_snapshot_drift_check"
 module LlmCostTracker
   class Doctor
     autoload :InvoiceReconciliationCheck, "llm_cost_tracker/doctor/invoice_reconciliation_check"
+    autoload :CaptureVerifier,            "llm_cost_tracker/doctor/capture_verifier"
+
+    STATUS_GLYPHS = { ok: "✓", warn: "!", error: "x" }.freeze
+    STATUS_COLORS = { ok: 32, warn: 33, error: 31 }.freeze
+
+    SECTIONS = ["Setup", "Schema", "Data integrity", "Operations"].freeze
+
+    SECTION_FOR_CHECK = {
+      "configuration" => "Setup",
+      "capture" => "Setup",
+      "active_record" => "Schema",
+      "llm_cost_tracker_calls" => "Schema",
+      "llm_cost_tracker_calls columns" => "Schema",
+      "call line items" => "Schema",
+      "call tags" => "Schema",
+      "provider invoices" => "Schema",
+      "provider invoice imports" => "Schema",
+      "cost drift" => "Data integrity",
+      "pricing snapshot drift" => "Data integrity",
+      "pricing snapshot audit" => "Data integrity",
+      "cost status" => "Data integrity",
+      "invoice reconciliation" => "Data integrity",
+      "call rollups" => "Operations",
+      "inline ingestion" => "Operations",
+      "async ingestion" => "Operations",
+      "prices" => "Operations",
+      "tracked calls" => "Operations"
+    }.freeze
+
+    private_constant :STATUS_GLYPHS, :STATUS_COLORS, :SECTIONS, :SECTION_FOR_CHECK
 
     class << self
       def call
         new.checks
       end
 
-      def report(checks = call)
-        (["LLM Cost Tracker doctor"] + checks.map do |check|
-          "[#{check.status}] #{check.name}: #{check.message}"
-        end).join("\n")
+      def report(checks = call, color: $stdout.tty?)
+        name_width = checks.map { |c| c.name.length }.max.to_i
+
+        lines = [bold("LLM Cost Tracker doctor", color), ""]
+        each_section(checks) do |section, members|
+          lines << bold(section, color)
+          members.each do |check|
+            status = paint_status("[#{STATUS_GLYPHS.fetch(check.status, check.status)}]", check.status, color)
+            lines << "  #{status} #{"#{check.name}:".ljust(name_width + 1)} #{check.message}"
+          end
+          lines << ""
+        end
+        lines.pop if lines.last == ""
+        lines.join("\n")
       end
 
       def healthy?(checks = call)
         checks.none? { |check| check.status == :error }
       end
+
+      private
+
+      def each_section(checks)
+        SECTIONS.each do |section|
+          members = checks.select { |c| (SECTION_FOR_CHECK[c.name] || "Setup") == section }
+          next if members.empty?
+
+          yield section, members
+        end
+      end
+
+      def paint_status(text, status, color)
+        return text unless color && STATUS_COLORS.key?(status)
+
+        "\e[#{STATUS_COLORS[status]}m#{text}\e[0m"
+      end
+
+      def bold(text, color)
+        return text unless color
+
+        "\e[1m#{text}\e[0m"
+      end
     end
 
     def checks
       [
         configuration_check,
         capture_check,
-        *integration_checks,
+        *LlmCostTracker::Integrations.checks,
         active_record_check,
         table_check,
         column_check,
-        SchemaCheck.new(name: "call line items", schema: Ledger::Schema::CallLineItems,
-                        table: "llm_cost_tracker_call_line_items").call,
-        SchemaCheck.new(name: "call tags", schema: Ledger::Schema::CallTags,
-                        table: "llm_cost_tracker_call_tags").call,
+        *dependent_core_schema_checks,
         *reconciliation_schema_checks,
         CostDriftCheck.new.call,
         PricingSnapshotDriftCheck.new.call,
@@ -58,20 +118,23 @@ module LlmCostTracker
 
     private
 
+    def dependent_core_schema_checks
+      Ledger::Schema::CORE_SCHEMAS.reject { |schema, _| schema == Ledger::Schema::Calls }.map do |schema, table|
+        SchemaCheck.new(name: table.delete_prefix("llm_cost_tracker_").tr("_", " "),
+                        schema: schema, table: table).call
+      end
+    end
+
     def reconciliation_schema_checks
       return [] unless LlmCostTracker.reconciliation_enabled?
 
-      LlmCostTracker.const_get(:Reconciliation) # autoload reconciliation + its ledger schemas
       Reconciliation::SCHEMA_TABLES.map do |schema, table|
-        SchemaCheck.new(name: humanize_table(table), schema: schema, table: table,
+        SchemaCheck.new(name: table.delete_prefix("llm_cost_tracker_").tr("_", " "),
+                        schema: schema, table: table,
                         optional: false, install_command: "llm_cost_tracker:reconciliation").call
       end.compact
     end
 
-    def humanize_table(table)
-      table.delete_prefix("llm_cost_tracker_").tr("_", " ")
-    end
-
     def reconciliation_invoice_check
       return [] unless LlmCostTracker.reconciliation_enabled?
 
@@ -104,12 +167,6 @@ module LlmCostTracker
       )
     end
 
-    def integration_checks
-      LlmCostTracker::Integrations.checks.map do |check|
-        Check.new(check.status, check.name.to_s, check.message)
-      end
-    end
-
     def active_record_check
       return Check.new(:ok, "active_record", "available") if active_record_available?
 
@@ -209,16 +266,14 @@ module LlmCostTracker
       count = snapshot.tracked_call_count.to_i
       return Check.new(:warn, "tracked calls", "none recorded yet") if count.zero?
 
-      latest_at = snapshot.latest_tracked_at
-      latest_at = latest_at.to_time if latest_at.respond_to?(:to_time)
-      latest = latest_at&.utc&.iso8601
+      latest = snapshot.latest_tracked_at.to_time.utc.iso8601
       Check.new(:ok, "tracked calls", "#{count} recorded; latest #{latest}")
     end
 
     def active_record_available?
       LlmCostTracker::Call.connection
       true
-    rescue LoadError, StandardError
+    rescue ActiveRecord::ConnectionNotEstablished, ActiveRecord::NoDatabaseError
       false
     end
 

data/lib/llm_cost_tracker/engine.rb

@@ -3,7 +3,6 @@
 require "rails"
 require_relative "../llm_cost_tracker"
 require_relative "assets"
-require_relative "dashboard_setup_state"
 require "rack/files"
 
 module LlmCostTracker
@@ -15,7 +14,7 @@ module LlmCostTracker
     end
 
     initializer "llm_cost_tracker.dashboard_setup_state" do |app|
-      app.reloader.to_prepare { LlmCostTracker::DashboardSetupState.reset! }
+      app.reloader.to_prepare { LlmCostTracker::Dashboard::SetupState.reset! }
     end
   end
 end

data/lib/llm_cost_tracker/errors.rb

@@ -6,13 +6,14 @@ module LlmCostTracker
   class InvalidFilterError < Error; end
 
   class BudgetExceededError < Error
-    attr_reader :total, :budget, :budget_type, :last_event
+    attr_reader :total, :budget, :budget_type, :last_event, :stage
 
-    def initialize(budget:, budget_type:, total:, last_event: nil)
+    def initialize(budget:, budget_type:, total:, last_event: nil, stage: :post_spend)
       @total = total
       @budget = budget
       @budget_type = budget_type
       @last_event = last_event
+      @stage = stage
 
       super(
         "LLM #{@budget_type.to_s.tr('_', '-')} budget exceeded: " \