llm_cost_tracker 0.8.0 → 0.9.0

data/lib/llm_cost_tracker/doctor.rb

@@ -13,6 +13,8 @@ require_relative "doctor/pricing_snapshot_drift_check"
 
 module LlmCostTracker
   class Doctor
+    autoload :InvoiceReconciliationCheck, "llm_cost_tracker/doctor/invoice_reconciliation_check"
+
     class << self
       def call
         new.checks
@@ -41,10 +43,10 @@ module LlmCostTracker
                         table: "llm_cost_tracker_call_line_items").call,
         SchemaCheck.new(name: "call tags", schema: Ledger::Schema::CallTags,
                         table: "llm_cost_tracker_call_tags").call,
-        SchemaCheck.new(name: "provider invoices", schema: Ledger::Schema::ProviderInvoices,
-                        table: "llm_cost_tracker_provider_invoices").call,
+        *reconciliation_schema_checks,
         CostDriftCheck.new.call,
         PricingSnapshotDriftCheck.new.call,
+        *reconciliation_invoice_check,
         LegacyBillingStatusCheck.new.call,
         LegacyAuditCheck.new.call,
         call_rollups_check,
@@ -56,6 +58,26 @@ module LlmCostTracker
 
     private
 
+    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,
+                        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?
+
+      Array(InvoiceReconciliationCheck.new.call)
+    end
+
     def configuration_check
       config = LlmCostTracker.configuration
       Check.new(:ok, "configuration", "active_record ledger enabled=#{config.enabled}")
@@ -120,9 +142,10 @@ module LlmCostTracker
 
     def call_rollups_check
       return unless llm_cost_tracker_calls_table?
+      return live_rollups_check unless LlmCostTracker.configuration.cache_rollups
 
       errors = LlmCostTracker::Ledger::Schema::CallRollups.current_schema_errors
-      return Check.new(:ok, "call rollups", "llm_cost_tracker_call_rollups exists") if errors.empty?
+      return rollups_drift_check if errors.empty?
 
       Check.new(
         :error,
@@ -131,6 +154,52 @@ module LlmCostTracker
       )
     end
 
+    ROLLUPS_DRIFT_TOLERANCE_PERCENT = 1.0
+    private_constant :ROLLUPS_DRIFT_TOLERANCE_PERCENT
+
+    def rollups_drift_check
+      drift_window = Time.now.utc.beginning_of_day
+      calls_total = LlmCostTracker::Call
+                    .where(tracked_at: drift_window..)
+                    .where.not(total_cost: nil)
+                    .sum(:total_cost)
+      rollup_total = LlmCostTracker::CallRollup
+                     .where(period: "day", period_start: drift_window.to_date)
+                     .sum(:total_cost)
+      return Check.new(:ok, "call rollups", "llm_cost_tracker_call_rollups exists") if calls_total.zero?
+
+      drift_percent = ((calls_total - rollup_total).abs * 100.0 / calls_total)
+      if drift_percent > ROLLUPS_DRIFT_TOLERANCE_PERCENT
+        return Check.new(
+          :warn, "call rollups",
+          "rollups drift detected: today's calls SUM=#{calls_total} vs rollups SUM=#{rollup_total} " \
+          "(#{drift_percent.round(2)}% > #{ROLLUPS_DRIFT_TOLERANCE_PERCENT}% threshold). " \
+          "Cached budget reads may understate spend until a rebuild."
+        )
+      end
+
+      Check.new(:ok, "call rollups", "llm_cost_tracker_call_rollups exists")
+    rescue StandardError => e
+      Check.new(:warn, "call rollups", "rollups drift check failed: #{e.class}: #{e.message}")
+    end
+
+    def live_rollups_check
+      if Probe.table_exists?("llm_cost_tracker_call_rollups")
+        Check.new(
+          :warn,
+          "call rollups",
+          "cache_rollups=false but llm_cost_tracker_call_rollups exists. " \
+          "Set config.cache_rollups = true to keep budget reads on the rollups fast path or drop the table."
+        )
+      else
+        Check.new(
+          :ok,
+          "call rollups",
+          "cache_rollups=false; budget reads aggregate from llm_cost_tracker_calls directly"
+        )
+      end
+    end
+
     def calls_check
       return unless llm_cost_tracker_calls_table?
 

data/lib/llm_cost_tracker/engine.rb

@@ -3,10 +3,19 @@
 require "rails"
 require_relative "../llm_cost_tracker"
 require_relative "assets"
+require_relative "dashboard_setup_state"
 require "rack/files"
 
 module LlmCostTracker
   class Engine < ::Rails::Engine
     isolate_namespace LlmCostTracker
+
+    initializer "llm_cost_tracker.filter_parameters" do |app|
+      app.config.filter_parameters += %i[tag tag_value]
+    end
+
+    initializer "llm_cost_tracker.dashboard_setup_state" do |app|
+      app.reloader.to_prepare { LlmCostTracker::DashboardSetupState.reset! }
+    end
   end
 end

data/lib/llm_cost_tracker/event.rb

@@ -29,7 +29,7 @@ module LlmCostTracker
     def to_h
       super.merge(
         token_usage: token_usage.to_h,
-        cost: cost&.to_h,
+        cost: cost && cost.to_h.transform_values { |v| v.is_a?(BigDecimal) ? v.to_f : v },
         tags: tags ? tags.to_h : {},
         line_items: (line_items || []).map(&:to_h)
       )

data/lib/llm_cost_tracker/generators/llm_cost_tracker/call_rollups_generator.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+require "rails/generators"
+require "rails/generators/active_record"
+
+module LlmCostTracker
+  module Generators
+    class CallRollupsGenerator < Rails::Generators::Base
+      include ActiveRecord::Generators::Migration
+
+      source_root File.expand_path("templates", __dir__)
+
+      desc "Creates the optional llm_cost_tracker_call_rollups table for fast budget reads. " \
+           "Required when config.cache_rollups = true."
+
+      def create_migration_file
+        migration_template(
+          "create_llm_cost_tracker_call_rollups.rb.erb",
+          "db/migrate/create_llm_cost_tracker_call_rollups.rb"
+        )
+      end
+
+      def warn_about_config_flag
+        say(<<~MSG, :yellow)
+          After migrating, set the following in config/initializers/llm_cost_tracker.rb:
+
+            LlmCostTracker.configure do |config|
+              config.cache_rollups = true
+            end
+
+          Without it Tracker keeps reading budget totals as live SUM aggregates over
+          llm_cost_tracker_calls. The doctor check warns about an unused rollups table.
+        MSG
+      end
+
+      private
+
+      def migration_version
+        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/durable_ingestion_generator.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+require "rails/generators"
+require "rails/generators/active_record"
+
+module LlmCostTracker
+  module Generators
+    class DurableIngestionGenerator < Rails::Generators::Base
+      include ActiveRecord::Generators::Migration
+
+      source_root File.expand_path("templates", __dir__)
+
+      desc "Creates the durable ingestion tables (llm_cost_tracker_ingestion_inbox_entries + _leases). " \
+           "Required when config.durable_ingestion = true."
+
+      def create_migration_file
+        migration_template(
+          "create_llm_cost_tracker_durable_ingestion.rb.erb",
+          "db/migrate/create_llm_cost_tracker_durable_ingestion.rb"
+        )
+      end
+
+      def warn_about_config_flag
+        say(<<~MSG, :yellow)
+          After migrating, set the following in config/initializers/llm_cost_tracker.rb:
+
+            LlmCostTracker.configure do |config|
+              config.durable_ingestion = true
+            end
+
+          Without it the durable inbox tables stay unused and Tracker keeps writing
+          inline. The doctor check warns about unused durable tables.
+        MSG
+      end
+
+      private
+
+      def migration_version
+        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/install_generator.rb

@@ -35,15 +35,25 @@ module LlmCostTracker
       def create_prices_file
         return unless options[:prices]
 
-        invoke "llm_cost_tracker:prices"
+        require_relative "prices_generator"
+        invoke LlmCostTracker::Generators::PricesGenerator
       end
 
       def mount_engine
         return unless options[:dashboard]
 
         add_engine_require
-        route %(mount LlmCostTracker::Engine => "/llm-costs")
-        say "Mount /llm-costs behind your app's admin auth before deploying.", :yellow
+        say(<<~MSG, :yellow)
+          The LLM Cost Tracker dashboard ships without authentication.
+          Mount it in config/routes.rb behind your app's admin auth, e.g.:
+
+            authenticate :admin do
+              mount LlmCostTracker::Engine => "/llm-costs"
+            end
+
+          The generator does NOT add a route automatically — leaving the dashboard
+          unauthenticated would expose spend, tags, and provider IDs to anyone.
+        MSG
       end
 
       private

data/lib/llm_cost_tracker/generators/llm_cost_tracker/reconciliation_generator.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+require "rails/generators"
+require "rails/generators/active_record"
+
+module LlmCostTracker
+  module Generators
+    class ReconciliationGenerator < Rails::Generators::Base
+      include ActiveRecord::Generators::Migration
+
+      source_root File.expand_path("templates", __dir__)
+
+      desc "Creates the optional invoice reconciliation tables. Requires provider admin/org-level API keys."
+
+      def create_migration_file
+        migration_template(
+          "create_llm_cost_tracker_reconciliation.rb.erb",
+          "db/migrate/create_llm_cost_tracker_reconciliation.rb"
+        )
+      end
+
+      def warn_about_admin_keys
+        say "Reconciliation requires admin/org-level API keys (OpenAI sk-admin-..., Anthropic admin keys, " \
+            "GCP billing.viewer service accounts). Do NOT use the runtime inference key.", :yellow
+      end
+
+      private
+
+      def migration_version
+        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/create_llm_cost_tracker_call_rollups.rb.erb

@@ -0,0 +1,15 @@
+class CreateLlmCostTrackerCallRollups < ActiveRecord::Migration<%= migration_version %>
+  def change
+    create_table :llm_cost_tracker_call_rollups do |t|
+      t.string :period, null: false
+      t.date :period_start, null: false
+      t.string :currency, null: false, default: "USD"
+      t.string :provider, null: false, default: ""
+      t.decimal :total_cost, precision: 20, scale: 8, null: false, default: 0
+
+      t.timestamps
+    end
+
+    add_index :llm_cost_tracker_call_rollups, [:period, :period_start, :currency, :provider], unique: true
+  end
+end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/create_llm_cost_tracker_calls.rb.erb

@@ -34,36 +34,6 @@ class CreateLlmCostTrackerCalls < ActiveRecord::Migration<%= migration_version %
       t.timestamps
     end
 
-    create_table :llm_cost_tracker_call_rollups do |t|
-      t.string :period, null: false
-      t.date :period_start, null: false
-      t.string :currency, null: false, default: "USD"
-      t.decimal :total_cost, precision: 20, scale: 8, null: false, default: 0
-
-      t.timestamps
-    end
-
-    create_table :llm_cost_tracker_ingestion_inbox_entries do |t|
-      t.string   :event_id, null: false
-      t.decimal  :total_cost, precision: 20, scale: 8
-      t.datetime :tracked_at, null: false
-      t.text     :payload, null: false
-      t.datetime :locked_at
-      t.string   :locked_by
-      t.integer  :attempts, null: false, default: 0
-      t.text     :last_error
-
-      t.timestamps
-    end
-
-    create_table :llm_cost_tracker_ingestion_leases do |t|
-      t.string   :name, null: false
-      t.string   :locked_by
-      t.datetime :locked_until
-
-      t.timestamps
-    end
-
     create_table :llm_cost_tracker_call_line_items do |t|
       t.references :llm_cost_tracker_call,
                    null: false,
@@ -107,25 +77,6 @@ class CreateLlmCostTrackerCalls < ActiveRecord::Migration<%= migration_version %
       t.text :value, null: false
     end
 
-    create_table :llm_cost_tracker_provider_invoices do |t|
-      t.string :source, null: false
-      t.date :period_start, null: false
-      t.date :period_end, null: false
-      t.string :external_id, null: false
-      t.decimal :billed_amount, precision: 20, scale: 8
-      t.string :currency, null: false, default: "USD"
-      if postgresql?
-        t.jsonb :metadata, null: false, default: {}
-      elsif mysql?
-        t.json :metadata, null: false
-      else
-        raise "LLM Cost Tracker supports PostgreSQL and MySQL only"
-      end
-      t.datetime :imported_at, null: false
-
-      t.timestamps
-    end
-
     add_index :llm_cost_tracker_calls, :event_id, unique: true
     add_index :llm_cost_tracker_calls, :tracked_at
     add_index :llm_cost_tracker_calls, [:provider, :tracked_at]
@@ -133,16 +84,12 @@ class CreateLlmCostTrackerCalls < ActiveRecord::Migration<%= migration_version %
     add_index :llm_cost_tracker_calls, :cost_status
     add_index :llm_cost_tracker_calls, :provider_response_id
     add_index :llm_cost_tracker_call_line_items, [:llm_cost_tracker_call_id, :position]
-    add_index :llm_cost_tracker_call_line_items, :kind
     add_index :llm_cost_tracker_call_tags, :llm_cost_tracker_call_id
-    add_index :llm_cost_tracker_call_tags, :key
-    add_index :llm_cost_tracker_call_rollups, [:period, :period_start, :currency], unique: true
-    add_index :llm_cost_tracker_ingestion_inbox_entries, :event_id, unique: true
-    add_index :llm_cost_tracker_ingestion_inbox_entries, [:tracked_at, :attempts]
-    add_index :llm_cost_tracker_ingestion_inbox_entries, [:locked_at, :id]
-    add_index :llm_cost_tracker_ingestion_leases, :name, unique: true
-    add_index :llm_cost_tracker_provider_invoices, :external_id, unique: true
-    add_index :llm_cost_tracker_provider_invoices, [:source, :period_start]
+    if postgresql?
+      add_index :llm_cost_tracker_call_tags, [:key, :value]
+    elsif mysql?
+      add_index :llm_cost_tracker_call_tags, [:key, :value], length: { value: 191 }
+    end
   end
 
   private

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/create_llm_cost_tracker_durable_ingestion.rb.erb

@@ -0,0 +1,29 @@
+class CreateLlmCostTrackerDurableIngestion < ActiveRecord::Migration<%= migration_version %>
+  def change
+    create_table :llm_cost_tracker_ingestion_inbox_entries do |t|
+      t.string   :event_id, null: false
+      t.decimal  :total_cost, precision: 20, scale: 8
+      t.datetime :tracked_at, null: false
+      t.text     :payload, null: false
+      t.datetime :locked_at
+      t.string   :locked_by
+      t.integer  :attempts, null: false, default: 0
+      t.text     :last_error
+
+      t.timestamps
+    end
+
+    create_table :llm_cost_tracker_ingestion_leases do |t|
+      t.string   :name, null: false
+      t.string   :locked_by
+      t.datetime :locked_until
+
+      t.timestamps
+    end
+
+    add_index :llm_cost_tracker_ingestion_inbox_entries, :event_id, unique: true
+    add_index :llm_cost_tracker_ingestion_inbox_entries, [:tracked_at, :attempts]
+    add_index :llm_cost_tracker_ingestion_inbox_entries, [:locked_at, :id]
+    add_index :llm_cost_tracker_ingestion_leases, :name, unique: true
+  end
+end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/create_llm_cost_tracker_reconciliation.rb.erb

@@ -0,0 +1,55 @@
+require "llm_cost_tracker/ledger/schema/adapter"
+
+class CreateLlmCostTrackerReconciliation < ActiveRecord::Migration<%= migration_version %>
+  def change
+    create_table :llm_cost_tracker_provider_invoices, if_not_exists: true do |t|
+      t.string :source, null: false
+      t.date :period_start, null: false
+      t.date :period_end, null: false
+      t.string :external_id, null: false
+      t.decimal :billed_amount, precision: 20, scale: 8
+      t.string :currency, null: false, default: "USD"
+      if postgresql?
+        t.jsonb :metadata, null: false, default: {}
+      elsif mysql?
+        t.json :metadata, null: false
+      else
+        raise "LLM Cost Tracker supports PostgreSQL and MySQL only"
+      end
+      t.datetime :imported_at, null: false
+
+      t.timestamps
+    end
+
+    create_table :llm_cost_tracker_provider_invoice_imports, if_not_exists: true do |t|
+      t.string :source, null: false
+      t.string :cursor
+      t.date :window_start
+      t.date :window_end
+      t.string :state, null: false
+      t.text :last_error
+      t.integer :rows_imported, null: false, default: 0
+      t.datetime :started_at, null: false
+      t.datetime :finished_at
+
+      t.timestamps
+    end
+
+    add_index :llm_cost_tracker_provider_invoices, :external_id, unique: true,
+              if_not_exists: true
+    add_index :llm_cost_tracker_provider_invoices, %i[source currency period_start],
+              if_not_exists: true
+    add_index :llm_cost_tracker_provider_invoice_imports, %i[source started_at],
+              if_not_exists: true
+  end
+
+  private
+
+  def postgresql?
+    LlmCostTracker::Ledger::Schema::Adapter.postgresql?(connection)
+  end
+
+  def mysql?
+    LlmCostTracker::Ledger::Schema::Adapter.mysql?(connection)
+  end
+end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/initializer.rb.erb

@@ -4,7 +4,10 @@ LlmCostTracker.configure do |config|
   # Set to false to temporarily disable tracking without removing middleware.
   config.enabled = true
 
-  # Tags are merged into every event. Use a callable for request/job-time context.
+  # LLM Cost Tracker logs warnings through Rails.logger when available.
+  config.log_level = :info
+
+  # Tags merged into every event. Use a callable for request/job-time context.
   config.default_tags = -> { { environment: Rails.env } }
 
   # Tag guardrails keep accidental high-cardinality or sensitive values out of the ledger.
@@ -18,40 +21,40 @@ LlmCostTracker.configure do |config|
   # config.instrument :anthropic
   # config.instrument :ruby_llm
 
-  # Budget behavior: :notify calls on_budget_exceeded, :raise raises after recording,
-  # :block_requests preflights monthly/daily budgets before supported requests.
-  config.budget_exceeded_behavior = :notify
-
-  # Unknown pricing records token usage with nil cost by default. Use :raise if
-  # every model must have known pricing before it can be used.
-  config.unknown_pricing_behavior = :warn
-
-  # LLM Cost Tracker logs warnings through Rails.logger when available.
-  config.log_level = :info
+  # Pricing — local file refreshed via bin/rails llm_cost_tracker:prices:refresh
+  # plus inline overrides. Prices are USD per 1M tokens.
 <% if options[:prices] -%>
-
-  # Local JSON/YAML pricing file generated by --prices. Keep it in source control
-  # and refresh it with bin/rails llm_cost_tracker:prices:refresh.
   config.prices_file = Rails.root.join("config/llm_cost_tracker_prices.yml")
+<% else -%>
+  # config.prices_file = Rails.root.join("config/llm_cost_tracker_prices.yml")
 <% end -%>
+  # config.pricing_overrides = {
+  #   "my-custom-model" => { input: 1.00, output: 2.00 }
+  # }
+  # :warn (default) records token usage with nil cost when a model has no rate.
+  # Use :raise to require known pricing for every model.
+  config.unknown_pricing_behavior = :warn
 
-  # Cumulative monthly/daily budgets and a single-call ceiling, in USD.
+  # Budget guardrails — cumulative monthly/daily and per-call ceilings in USD,
+  # plus behavior on crossing (:notify default fires on_budget_exceeded; :raise
+  # raises after recording; :block_requests preflights supported requests) and
+  # an optional callback. Cap evaluation reads from llm_cost_tracker_calls live;
+  # flip cache_rollups to true at high volume so reads hit the rollups table
+  # instead — generate the table with `bin/rails generate llm_cost_tracker:call_rollups`.
   # config.monthly_budget = 100.00
   # config.daily_budget = 10.00
   # config.per_call_budget = 1.00
-
-  # Called when :notify is selected and a monthly, daily, or per-call budget is exceeded.
+  config.budget_exceeded_behavior = :notify
   # config.on_budget_exceeded = ->(data) {
-  #   Rails.logger.warn(
-  #     "LLM #{data[:budget_type]} budget exceeded: $#{data[:total]} / $#{data[:budget]}"
-  #   )
+  #   Rails.logger.warn("LLM #{data[:budget_type]} budget exceeded: $#{data[:total]} / $#{data[:budget]}")
   # }
+  # config.cache_rollups = true
 
-  # Local pricing table and small Ruby-side overrides. Prices are USD per 1M tokens.
-  # config.prices_file = Rails.root.join("config/llm_cost_tracker_prices.yml")
-  # config.pricing_overrides = {
-  #   "my-custom-model" => { input: 1.00, output: 2.00 }
-  # }
+  # Ingestion path — false (default) writes events synchronously from the request
+  # thread. Flip to true for a write-ahead inbox + background worker that batches
+  # inserts and survives caller transaction rollbacks. Requires the optional
+  # inbox/leases tables created by `bin/rails generate llm_cost_tracker:durable_ingestion`.
+  # config.durable_ingestion = true
 
   # Register OpenAI-compatible gateway hosts and choose extra tag breakdowns
   # for bin/rails llm_cost_tracker:report.

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/upgrade_call_rollups_provider.rb.erb

@@ -0,0 +1,20 @@
+class UpgradeLlmCostTrackerCallRollupsProvider < ActiveRecord::Migration<%= migration_version %>
+  TABLE = :llm_cost_tracker_call_rollups
+  OLD_INDEX = %i[period period_start currency].freeze
+  NEW_INDEX = %i[period period_start currency provider].freeze
+
+  def up
+    unless column_exists?(TABLE, :provider)
+      execute "DELETE FROM #{TABLE}"
+      add_column TABLE, :provider, :string, null: false, default: ""
+    end
+    remove_index TABLE, column: OLD_INDEX, unique: true if index_exists?(TABLE, OLD_INDEX, unique: true)
+    add_index TABLE, NEW_INDEX, unique: true unless index_exists?(TABLE, NEW_INDEX, unique: true)
+  end
+
+  def down
+    remove_index TABLE, column: NEW_INDEX, unique: true if index_exists?(TABLE, NEW_INDEX, unique: true)
+    add_index TABLE, OLD_INDEX, unique: true unless index_exists?(TABLE, OLD_INDEX, unique: true)
+    remove_column TABLE, :provider if column_exists?(TABLE, :provider)
+  end
+end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/upgrade_call_tags_key_value_index.rb.erb

@@ -0,0 +1,32 @@
+require "llm_cost_tracker/ledger/schema/adapter"
+
+class UpgradeLlmCostTrackerCallTagsKeyValueIndex < ActiveRecord::Migration<%= migration_version %>
+  TABLE = :llm_cost_tracker_call_tags
+  INDEX_COLUMNS = %i[key value].freeze
+
+  def up
+    return if index_exists?(TABLE, INDEX_COLUMNS)
+
+    if postgresql?
+      add_index TABLE, INDEX_COLUMNS
+    elsif mysql?
+      add_index TABLE, INDEX_COLUMNS, length: { value: 191 }
+    else
+      raise "LLM Cost Tracker supports PostgreSQL and MySQL only"
+    end
+  end
+
+  def down
+    remove_index TABLE, column: INDEX_COLUMNS if index_exists?(TABLE, INDEX_COLUMNS)
+  end
+
+  private
+
+  def postgresql?
+    LlmCostTracker::Ledger::Schema::Adapter.postgresql?(connection)
+  end
+
+  def mysql?
+    LlmCostTracker::Ledger::Schema::Adapter.mysql?(connection)
+  end
+end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/upgrade_image_tokens.rb.erb

@@ -0,0 +1,18 @@
+class UpgradeLlmCostTrackerImageTokens < ActiveRecord::Migration<%= migration_version %>
+  TABLE = :llm_cost_tracker_calls
+  COLUMNS = %i[image_input_tokens image_output_tokens].freeze
+
+  def up
+    COLUMNS.each do |column|
+      next if column_exists?(TABLE, column)
+
+      add_column TABLE, column, :integer, null: false, default: 0
+    end
+  end
+
+  def down
+    COLUMNS.each do |column|
+      remove_column TABLE, column if column_exists?(TABLE, column)
+    end
+  end
+end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/upgrade_call_rollups_provider_generator.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+require "rails/generators"
+require "rails/generators/active_record"
+
+module LlmCostTracker
+  module Generators
+    class UpgradeCallRollupsProviderGenerator < Rails::Generators::Base
+      include ActiveRecord::Generators::Migration
+
+      source_root File.expand_path("templates", __dir__)
+
+      desc "Adds the v0.9 provider column and unique index to llm_cost_tracker_call_rollups."
+
+      def create_migration_file
+        migration_template(
+          "upgrade_call_rollups_provider.rb.erb",
+          "db/migrate/upgrade_llm_cost_tracker_call_rollups_provider.rb"
+        )
+      end
+
+      def warn_about_rollups_truncation
+        say(<<~MSG, :yellow)
+          The migration clears existing llm_cost_tracker_call_rollups rows before adding the
+          provider column. Budget reads fall back to live aggregation from
+          llm_cost_tracker_calls until new events repopulate the rollups under their provider
+          keys. See docs/upgrading.md for details.
+        MSG
+      end
+
+      private
+
+      def migration_version
+        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
+      end
+    end
+  end
+end