llm_cost_tracker 0.14.1 → 0.15.0

data/lib/llm_cost_tracker/doctor.rb

@@ -37,19 +37,12 @@ module LlmCostTracker
       end
 
       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 << ""
+        name_width = checks.map { |check| check.name.length }.max.to_i
+        sections = checks.group_by { |check| SECTION_FOR_CHECK.fetch(check.name, "Setup") }.slice(*SECTIONS)
+        blocks = sections.map do |section, members|
+          [bold(section, color), *members.map { |check| check_line(check, name_width, color) }].join("\n")
         end
-        lines.pop if lines.last == ""
-        lines.join("\n")
+        [bold("LLM Cost Tracker doctor", color), *blocks].join("\n\n")
       end
 
       def healthy?(checks = call)
@@ -58,13 +51,9 @@ module LlmCostTracker
 
       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
+      def check_line(check, name_width, color)
+        status = paint_status("[#{STATUS_GLYPHS.fetch(check.status, check.status)}]", check.status, color)
+        "  #{status} #{"#{check.name}:".ljust(name_width + 1)} #{check.message}"
       end
 
       def paint_status(text, status, color)
@@ -89,6 +78,7 @@ module LlmCostTracker
         table_check,
         column_check,
         *dependent_core_schema_checks,
+        per_tag_budget_check,
         call_rollups_check,
         IngestionCheck.new.call,
         PriceCheck.new.call,
@@ -106,6 +96,19 @@ module LlmCostTracker
       end
     end
 
+    def per_tag_budget_check
+      return if LlmCostTracker.configuration.budgets.per_tag.empty?
+      return unless Probe.table_exists?("llm_cost_tracker_call_tags")
+      return if Budget::PerTag.columns?
+
+      Check.new(
+        :error,
+        "per-tag budgets",
+        "budgets.per_tag is set but llm_cost_tracker_call_tags has no cost columns, so no per-tag budget is " \
+        "enforced; run bin/rails generate llm_cost_tracker:upgrade_per_tag_budgets and migrate"
+      )
+    end
+
     def configuration_check
       config = LlmCostTracker.configuration
       Check.new(:ok, "configuration", "active_record ledger enabled=#{config.enabled}")
@@ -167,7 +170,7 @@ module LlmCostTracker
       return live_rollups_check unless LlmCostTracker.configuration.budgets.totals_source == :cache
 
       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,
@@ -176,13 +179,30 @@ module LlmCostTracker
       )
     end
 
+    def rollups_drift_check
+      month_start = LlmCostTracker::Ledger::Period.range_start(:month, Time.now)
+      today = LlmCostTracker::Ledger::Period.range_start(:day, Time.now)
+      rolled = LlmCostTracker::CallRollup.where(period: "day", period_start: month_start.to_date...today.to_date)
+                                         .sum(:total_cost)
+      recorded = LlmCostTracker::Call.where(tracked_at: month_start...today).sum(:total_cost)
+      return Check.new(:ok, "call rollups", "llm_cost_tracker_call_rollups exists") if rolled == recorded
+
+      Check.new(
+        :warn,
+        "call rollups",
+        "this month's day rollups before today total #{rolled.to_d.to_s('F')} but the calls total " \
+        "#{recorded.to_d.to_s('F')}, and monthly budgets read the rollups; " \
+        "run bin/rails llm_cost_tracker:rebuild_rollups"
+      )
+    end
+
     def live_rollups_check
       if Probe.table_exists?("llm_cost_tracker_call_rollups")
         Check.new(
           :warn,
           "call rollups",
           "budgets.totals_source=:ledger but llm_cost_tracker_call_rollups exists. " \
-          "Set config.budgets.totals_source = :cache to keep budget reads on the rollups fast path or drop the table."
+          "Set config.budgets.totals_source = :cache to use it or drop the table."
         )
       else
         Check.new(
@@ -207,7 +227,7 @@ module LlmCostTracker
     end
 
     def active_record_available?
-      LlmCostTracker::Call.connection
+      LlmCostTracker::Call.lease_connection
       true
     rescue ActiveRecord::ConnectionNotEstablished, ActiveRecord::NoDatabaseError
       false

data/lib/llm_cost_tracker/engine.rb

@@ -3,12 +3,19 @@
 require "rails"
 require_relative "../llm_cost_tracker"
 require_relative "assets"
-require "rack/files"
 
 module LlmCostTracker
   class Engine < ::Rails::Engine
     isolate_namespace LlmCostTracker
 
+    initializer "llm_cost_tracker.host_inflections" do |app|
+      Rails.autoloaders.each { |autoloader| autoloader.inflector.inflect("llm_cost_tracker" => "LlmCostTracker") }
+      app.reloader.to_prepare do
+        host_name = "llm_cost_tracker".camelize
+        Object.const_set(host_name, LlmCostTracker) unless Object.const_defined?(host_name)
+      end
+    end
+
     initializer "llm_cost_tracker.deprecator" do |app|
       app.deprecators[:llm_cost_tracker] = LlmCostTracker.deprecator
     end

data/lib/llm_cost_tracker/errors.rb

@@ -27,10 +27,8 @@ module LlmCostTracker
       @stage = stage
       @scope = scope
 
-      super(
-        "LLM #{@budget_type.to_s.tr('_', '-')} budget exceeded#{scope_suffix}: " \
-        "$#{format('%.6f', @total)} / $#{format('%.6f', budget)}"
-      )
+      amounts = budget_type == :calls ? "#{total} / #{budget}" : format("$%<total>.6f / $%<budget>.6f", total:, budget:)
+      super("LLM #{@budget_type.to_s.tr('_', '-')} budget exceeded#{scope_suffix}: #{amounts}")
     end
 
     private

data/lib/llm_cost_tracker/event.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require "active_support/core_ext/digest/uuid"
+
 require_relative "pricing/mode"
 
 module LlmCostTracker
@@ -22,43 +24,46 @@ module LlmCostTracker
     :cost_status,
     :pricing_snapshot,
     :line_items
-  ) do
-    def self.build(**attributes)
-      token_usage = attributes.fetch(:token_usage)
-      line_items = attributes[:line_items] || resolve_line_items(attributes[:service_line_items])
+  )
+
+  class Event
+    UNKNOWN_MODEL = "unknown"
+    PROVIDER_IDS = %i[provider_response_id provider_project_id provider_api_key_id provider_workspace_id].freeze
+    PASSTHROUGH_FIELDS = %i[event_id pricing_mode cost tags latency_ms tracked_at cost_status pricing_snapshot].freeze
+    private_constant :PROVIDER_IDS, :PASSTHROUGH_FIELDS
 
+    def self.build(**attributes)
       new(
-        event_id: attributes[:event_id],
+        token_usage: attributes.fetch(:token_usage),
+        line_items: attributes[:line_items] || resolve_line_items(attributes[:service_line_items]),
         provider: attributes.fetch(:provider).to_s,
-        model: attributes.fetch(:model).to_s.strip.presence || Event::UNKNOWN_MODEL,
-        token_usage: token_usage,
-        pricing_mode: attributes[:pricing_mode],
-        cost: attributes[:cost],
-        tags: attributes[:tags],
-        latency_ms: attributes[:latency_ms],
+        model: identifier(attributes.fetch(:model)) || UNKNOWN_MODEL,
         stream: attributes[:stream] || false,
         usage_source: attributes[:usage_source]&.to_s,
-        provider_response_id: attributes[:provider_response_id].to_s.strip.presence,
-        provider_project_id: attributes[:provider_project_id].to_s.strip.presence,
-        provider_api_key_id: attributes[:provider_api_key_id].to_s.strip.presence,
-        provider_workspace_id: attributes[:provider_workspace_id].to_s.strip.presence,
-        tracked_at: attributes[:tracked_at],
-        cost_status: attributes[:cost_status],
-        pricing_snapshot: attributes[:pricing_snapshot],
-        line_items: line_items
+        **PASSTHROUGH_FIELDS.to_h { |field| [field, attributes[field]] },
+        **PROVIDER_IDS.to_h { |field| [field, identifier(attributes[field])] }
       )
     end
 
-    def batch?
-      Pricing::Mode.tokenize(pricing_mode.to_s).include?("batch")
-    end
-
     def self.resolve_line_items(service_items)
       Array(service_items).map do |item|
         item.is_a?(Charges::LineItem) ? item : Charges::LineItem.build(item)
       end
     end
 
+    def self.identifier(value) = value.to_s.strip.presence
+    private_class_method :identifier
+
+    def batch?
+      Pricing::Mode.tokenize(pricing_mode.to_s).include?("batch")
+    end
+
+    def keyed_by_response_id
+      return self unless provider_response_id
+
+      with(event_id: Digest::UUID.uuid_v5(Digest::UUID::OID_NAMESPACE, "#{provider}/#{provider_response_id}"))
+    end
+
     def total_cost
       cost&.total
     end
@@ -72,8 +77,4 @@ module LlmCostTracker
       )
     end
   end
-
-  class Event
-    UNKNOWN_MODEL = "unknown"
-  end
 end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/async_ingestion_generator.rb

@@ -32,12 +32,6 @@ module LlmCostTracker
           inline. The doctor check warns about unused async ingestion 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/call_rollups_generator.rb

@@ -10,7 +10,7 @@ module LlmCostTracker
 
       source_root File.expand_path("templates", __dir__)
 
-      desc "Creates the optional llm_cost_tracker_call_rollups table for fast budget reads. " \
+      desc "Creates the optional llm_cost_tracker_call_rollups table of running budget totals. " \
            "Required when config.budgets.totals_source = :cache."
 
       def create_migration_file
@@ -28,16 +28,13 @@ module LlmCostTracker
               config.budgets.totals_source = :cache
             end
 
-          Without it Tracker keeps reading budget totals as live SUM aggregates over
+          Once that is deployed, run bin/rails llm_cost_tracker:rebuild_rollups. Until then
+          calls recorded before the switch are missing from monthly budget totals.
+
+          Without the setting 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/install_generator.rb

@@ -53,12 +53,6 @@ module LlmCostTracker
           unauthenticated would expose spend, tags, and provider IDs to anyone.
         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/templates/create_llm_cost_tracker_async_ingestion.rb.erb

@@ -1,4 +1,4 @@
-class CreateLlmCostTrackerAsyncIngestion < ActiveRecord::Migration<%= migration_version %>
+class CreateLlmCostTrackerAsyncIngestion < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
   def change
     create_table :llm_cost_tracker_ingestion_inbox_entries do |t|
       t.string   :event_id, null: false

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

@@ -1,4 +1,4 @@
-class CreateLlmCostTrackerCallRollups < ActiveRecord::Migration<%= migration_version %>
+class CreateLlmCostTrackerCallRollups < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
   def change
     create_table :llm_cost_tracker_call_rollups do |t|
       t.string :period, null: false

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

@@ -1,7 +1,7 @@
 require "llm_cost_tracker/charges/cost_status"
 require "llm_cost_tracker/ledger/schema/adapter"
 
-class CreateLlmCostTrackerCalls < ActiveRecord::Migration<%= migration_version %>
+class CreateLlmCostTrackerCalls < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
   def change
     create_table :llm_cost_tracker_calls do |t|
       t.string  :event_id,      null: false
@@ -84,8 +84,10 @@ class CreateLlmCostTrackerCalls < ActiveRecord::Migration<%= migration_version %
     add_index :llm_cost_tracker_calls, %i[model tracked_at]
     add_index :llm_cost_tracker_calls, :cost_status
     add_index :llm_cost_tracker_calls, :provider_response_id
-    add_index :llm_cost_tracker_calls, :id, name: :index_llm_cost_tracker_calls_on_unpriced,
-                                          where: "total_cost IS NULL"
+    if postgresql?
+      add_index :llm_cost_tracker_calls, :id, name: :index_llm_cost_tracker_calls_on_unpriced,
+                                            where: "total_cost IS NULL"
+    end
     add_index :llm_cost_tracker_call_line_items, [:llm_cost_tracker_call_id, :position]
     add_index :llm_cost_tracker_call_tags, :llm_cost_tracker_call_id
     if postgresql?

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

@@ -50,8 +50,9 @@ LlmCostTracker.configure do |config|
 
   # One budget applied to every distinct value of a tag: each tenant gets its own
   # 1000 a month, not 1000 shared between them. Declare as many tags as you need; a rule
-  # can set its own behavior and callback instead of following the global ones. Needs the
-  # cost columns from `bin/rails generate llm_cost_tracker:upgrade_per_tag_budgets`.
+  # can set its own behavior and callback instead of following the global ones. Installs created
+  # before v0.14 add the cost columns with `bin/rails generate llm_cost_tracker:upgrade_per_tag_budgets`
+  # and `bin/rails db:migrate`, then run `bin/rails llm_cost_tracker:backfill_tag_costs` so earlier calls count.
   # Budget high-cardinality tags only — a tag with a handful of values covers most of the
   # ledger, so its check cannot use an index and every call pays for a near-full scan.
   # config.budgets.per_tag = {
@@ -59,13 +60,14 @@ LlmCostTracker.configure do |config|
   #   user_id: { daily: 25.00, behavior: :notify }
   # }
   # config.budgets.on_exceeded = ->(data) {
-  #   Rails.logger.warn("LLM #{data[:budget_type]} budget exceeded: $#{data[:total]} / $#{data[:budget]}")
+  #   unit = data[:budget_type] == :calls ? "" : "$"
+  #   Rails.logger.warn("LLM #{data[:budget_type]} budget exceeded: #{unit}#{data[:total]} / #{unit}#{data[:budget]}")
   # }
 
   # Where budget checks read period spend from. :ledger sums llm_cost_tracker_calls on
-  # every check; :cache also keeps running totals in llm_cost_tracker_call_rollups and
-  # takes the greater of the two, so it guards against drift but does not make checks
-  # faster. Generate the table first with `bin/rails generate llm_cost_tracker:call_rollups`.
+  # every check; :cache reads earlier days from llm_cost_tracker_call_rollups and sums
+  # only today's calls. Run `bin/rails generate llm_cost_tracker:call_rollups`, deploy
+  # this setting, then run `bin/rails llm_cost_tracker:rebuild_rollups`.
   # config.budgets.totals_source = :cache
 
   # Ingestion path — :inline (default) writes events synchronously from the request
@@ -74,13 +76,15 @@ LlmCostTracker.configure do |config|
   # inbox/leases tables created by `bin/rails generate llm_cost_tracker:async_ingestion`.
   # Synchronous inbox writes use a dedicated ActiveRecord pool (defaults to 2 connections)
   # so they don't compete with request threads for the default pool when a tracked call
-  # happens inside an open caller transaction. Bump ingestion.pool_size if your Puma
-  # worker count outgrows that.
+  # happens inside an open caller transaction. The pool is per process: bump ingestion.pool_size
+  # if concurrent tracked calls in one process (e.g. Puma threads) outgrow that.
   # config.ingestion.mode = :async
   # config.ingestion.pool_size = 5
 
-  # Register OpenAI-compatible gateway hosts so their calls are attributed to a
-  # provider name instead of the generic openai_compatible bucket.
+  # Register OpenAI-compatible gateway hosts beyond the built-in OpenRouter, DeepSeek, Groq,
+  # xAI, Mistral and Perplexity; the Faraday middleware captures only listed hosts. Its calls,
+  # and openai gem calls whose base_url is a listed host, are recorded under the provider name
+  # given here.
   # config.capture.openai_compatible_providers["llm.my-company.com"] = "internal_gateway"
 
   # Streaming calls only report token usage when the request asks for it. Leave this

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

@@ -1,6 +1,6 @@
 require "llm_cost_tracker/ledger/schema/adapter"
 
-class UpgradeLlmCostTrackerCallRollupsProvider < ActiveRecord::Migration<%= migration_version %>
+class UpgradeLlmCostTrackerCallRollupsProvider < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
   disable_ddl_transaction!
 
   TABLE = :llm_cost_tracker_call_rollups
@@ -27,6 +27,7 @@ class UpgradeLlmCostTrackerCallRollupsProvider < ActiveRecord::Migration<%= migr
 
   def add_unique_index(columns)
     if postgresql?
+      LlmCostTracker::Ledger::Schema::Adapter.drop_invalid_index(connection, connection.index_name(TABLE, columns))
       add_index TABLE, columns, unique: true, algorithm: :concurrently, if_not_exists: true
     else
       add_index TABLE, columns, unique: true, if_not_exists: true

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

@@ -1,6 +1,6 @@
 require "llm_cost_tracker/ledger/schema/adapter"
 
-class UpgradeLlmCostTrackerCallTagsKeyValueIndex < ActiveRecord::Migration<%= migration_version %>
+class UpgradeLlmCostTrackerCallTagsKeyValueIndex < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
   disable_ddl_transaction!
 
   TABLE = :llm_cost_tracker_call_tags
@@ -8,6 +8,7 @@ class UpgradeLlmCostTrackerCallTagsKeyValueIndex < ActiveRecord::Migration<%= mi
 
   def up
     if postgresql?
+      LlmCostTracker::Ledger::Schema::Adapter.drop_invalid_index(connection, connection.index_name(TABLE, INDEX_COLUMNS))
       add_index TABLE, INDEX_COLUMNS, algorithm: :concurrently, if_not_exists: true
     elsif mysql?
       add_index TABLE, INDEX_COLUMNS, length: { value: 191 }, if_not_exists: true

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

@@ -1,4 +1,4 @@
-class UpgradeLlmCostTrackerImageTokens < ActiveRecord::Migration<%= migration_version %>
+class UpgradeLlmCostTrackerImageTokens < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
   TABLE = :llm_cost_tracker_calls
   COLUMNS = %i[image_input_tokens image_output_tokens].freeze
 

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

@@ -1,6 +1,6 @@
 require "llm_cost_tracker/ledger/schema/adapter"
 
-class UpgradeLlmCostTrackerIndexes < ActiveRecord::Migration<%= migration_version %>
+class UpgradeLlmCostTrackerIndexes < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
   disable_ddl_transaction!
 
   CALLS = :llm_cost_tracker_calls
@@ -9,7 +9,8 @@ class UpgradeLlmCostTrackerIndexes < ActiveRecord::Migration<%= migration_versio
   UNPRICED_INDEX = :index_llm_cost_tracker_calls_on_unpriced
 
   def up
-    add_index CALLS, :id, name: UNPRICED_INDEX, where: "total_cost IS NULL", **add_index_options
+    LlmCostTracker::Ledger::Schema::Adapter.drop_invalid_index(connection, UNPRICED_INDEX)
+    add_index CALLS, :id, name: UNPRICED_INDEX, where: "total_cost IS NULL", **add_index_options if postgresql?
     return unless table_exists?(INBOX)
 
     remove_index INBOX, column: UNUSED_INBOX_INDEX, **remove_index_options

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

@@ -1,6 +1,6 @@
 require "llm_cost_tracker/ledger/schema/adapter"
 
-class UpgradeLlmCostTrackerPerTagBudgets < ActiveRecord::Migration<%= migration_version %>
+class UpgradeLlmCostTrackerPerTagBudgets < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
   disable_ddl_transaction!
 
   TABLE = :llm_cost_tracker_call_tags
@@ -10,6 +10,7 @@ class UpgradeLlmCostTrackerPerTagBudgets < ActiveRecord::Migration<%= migration_
   def up
     add_column TABLE, :total_cost, :decimal, precision: 20, scale: 8 unless column_exists?(TABLE, :total_cost)
     add_column TABLE, :tracked_at, :datetime unless column_exists?(TABLE, :tracked_at)
+    LlmCostTracker::Ledger::Schema::Adapter.drop_invalid_index(connection, connection.index_name(TABLE, INDEX))
     add_index TABLE, INDEX, **add_index_options
     remove_index TABLE, column: REPLACED_INDEX, **remove_index_options
   end

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

@@ -18,12 +18,6 @@ module LlmCostTracker
           "db/migrate/upgrade_llm_cost_tracker_call_rollups_provider.rb"
         )
       end
-
-      private
-
-      def migration_version
-        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
-      end
     end
   end
 end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/upgrade_call_tags_key_value_index_generator.rb

@@ -19,12 +19,6 @@ module LlmCostTracker
           "db/migrate/upgrade_llm_cost_tracker_call_tags_key_value_index.rb"
         )
       end
-
-      private
-
-      def migration_version
-        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
-      end
     end
   end
 end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/upgrade_image_tokens_generator.rb

@@ -18,12 +18,6 @@ module LlmCostTracker
           "db/migrate/upgrade_llm_cost_tracker_image_tokens.rb"
         )
       end
-
-      private
-
-      def migration_version
-        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
-      end
     end
   end
 end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/upgrade_indexes_generator.rb

@@ -10,8 +10,9 @@ module LlmCostTracker
 
       source_root File.expand_path("templates", __dir__)
 
-      desc "Adds a partial index over unpriced calls and drops the unused ingestion inbox lock index" \
-           "scanning the whole ledger."
+      desc "Adds the unpriced-calls index (partial on PostgreSQL, so backfill_unknown_pricing stops scanning " \
+           "the whole ledger; a plain id index on MySQL, which has no partial indexes) and drops the unused " \
+           "ingestion inbox lock index."
 
       def create_migration_file
         migration_template(
@@ -19,12 +20,6 @@ module LlmCostTracker
           "db/migrate/upgrade_llm_cost_tracker_indexes.rb"
         )
       end
-
-      private
-
-      def migration_version
-        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
-      end
     end
   end
 end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/upgrade_per_tag_budgets_generator.rb

@@ -10,8 +10,9 @@ module LlmCostTracker
 
       source_root File.expand_path("templates", __dir__)
 
-      desc "Carries each call's cost and time onto its tag rows so per-tag budgets read " \
-           "llm_cost_tracker_call_tags without a join. Required when config.budgets.per_tag is set."
+      desc "Adds cost and time columns to llm_cost_tracker_call_tags so per-tag budgets read it " \
+           "without a join; run llm_cost_tracker:backfill_tag_costs afterwards so earlier calls count. " \
+           "Required for config.budgets.per_tag on installs created before v0.14."
 
       def create_migration_file
         migration_template(
@@ -19,12 +20,6 @@ module LlmCostTracker
           "db/migrate/upgrade_llm_cost_tracker_per_tag_budgets.rb"
         )
       end
-
-      private
-
-      def migration_version
-        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
-      end
     end
   end
 end

data/lib/llm_cost_tracker/ingestion/batch.rb

@@ -12,7 +12,7 @@ module LlmCostTracker
       TRANSIENT_PERSIST_ERRORS = [
         ActiveRecord::Deadlocked,
         ActiveRecord::LockWaitTimeout,
-        ActiveRecord::StatementTimeout,
+        ActiveRecord::QueryAborted,
         ActiveRecord::ConnectionNotEstablished,
         LlmCostTracker::TransactionAbortedError
       ].freeze
@@ -29,20 +29,16 @@ module LlmCostTracker
         persist_batch(valid_rows, events) if events.any?
         rows.size
       rescue StandardError => e
-        rows_to_mark = valid_rows&.any? ? valid_rows : rows
-        if rows_to_mark&.any?
-          transient = valid_rows&.any? && TRANSIENT_PERSIST_ERRORS.any? { |klass| e.is_a?(klass) }
-          mark_failed_with_message(rows_to_mark, error_message_for(e), decrement_attempts: transient)
-        end
+        release_after_failure(rows, valid_rows, e)
         raise
       end
 
       def pending?
-        Ingestion::InboxEntry.pending.exists?
+        Ingestion::InboxEntry.uncached { Ingestion::InboxEntry.pending.exists? }
       end
 
       def claimable?
-        claimable_scope(Time.now.utc - LOCK_TIMEOUT_SECONDS).exists?
+        Ingestion::InboxEntry.uncached { claimable_scope(Time.now.utc - LOCK_TIMEOUT_SECONDS).exists? }
       end
 
       def mark_failed_with_message(rows, message, decrement_attempts: false)
@@ -67,6 +63,10 @@ module LlmCostTracker
         nil
       end
 
+      private
+
+      attr_reader :identity
+
       def error_message_for(error)
         "#{error.class}: #{Redaction.text(error.message)}".byteslice(0, 1_000).scrub("")
       end
@@ -83,10 +83,6 @@ module LlmCostTracker
         )
       end
 
-      private
-
-      attr_reader :identity
-
       def id_sample(rows)
         sample = rows.first(10).map(&:id).join(", ")
         rows.size > 10 ? "#{sample}..." : sample
@@ -121,6 +117,18 @@ module LlmCostTracker
         [valid_rows, events]
       end
 
+      def release_after_failure(rows, valid_rows, error)
+        if valid_rows&.any?
+          mark_failed_with_message(valid_rows, error_message_for(error), decrement_attempts: transient?(error))
+        elsif rows&.any?
+          mark_failed_with_message(rows, error_message_for(error))
+        end
+      end
+
+      def transient?(error)
+        TRANSIENT_PERSIST_ERRORS.any? { |klass| error.is_a?(klass) }
+      end
+
       def persist_batch(rows, events)
         landed = []
         failed = Hash.new { |hash, message| hash[message] = [] }
@@ -128,14 +136,20 @@ module LlmCostTracker
       rescue *TRANSIENT_PERSIST_ERRORS
         raise
       rescue StandardError
-        rows.zip(events) do |row, event|
-          landed.concat(persist([row], [event]))
-        rescue *TRANSIENT_PERSIST_ERRORS
-          raise
-        rescue StandardError => e
-          failed[error_message_for(e)] << row
-        end
+        rows.zip(events) { |row, event| persist_alone(row, event, landed, failed) }
       ensure
+        settle(landed, failed)
+      end
+
+      def persist_alone(row, event, landed, failed)
+        landed.concat(persist([row], [event]))
+      rescue *TRANSIENT_PERSIST_ERRORS
+        raise
+      rescue StandardError => e
+        failed[error_message_for(e)] << row
+      end
+
+      def settle(landed, failed)
         failed.each do |message, failed_rows|
           LlmCostTracker::Logging.warn(
             "Ingestion::Batch: #{failed_rows.size} inbox row(s) could not be stored " \