llm_cost_tracker 0.9.0 → 0.10.0

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/{create_llm_cost_tracker_durable_ingestion.rb.erb → create_llm_cost_tracker_async_ingestion.rb.erb}

@@ -1,14 +1,14 @@
-class CreateLlmCostTrackerDurableIngestion < ActiveRecord::Migration<%= migration_version %>
+class CreateLlmCostTrackerAsyncIngestion < 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.text     :payload, null: false, limit: 16.megabytes
       t.datetime :locked_at
       t.string   :locked_by
       t.integer  :attempts, null: false, default: 0
-      t.text     :last_error
+      t.text     :last_error, limit: 16.megabytes
 
       t.timestamps
     end

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

@@ -23,6 +23,7 @@ class CreateLlmCostTrackerReconciliation < ActiveRecord::Migration<%= migration_
 
     create_table :llm_cost_tracker_provider_invoice_imports, if_not_exists: true do |t|
       t.string :source, null: false
+      t.string :provider, null: false, default: ""
       t.string :cursor
       t.date :window_start
       t.date :window_end
@@ -39,7 +40,11 @@ class CreateLlmCostTrackerReconciliation < ActiveRecord::Migration<%= migration_
               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 postgresql?
+      add_index :llm_cost_tracker_provider_invoices, :metadata, using: :gin,
+                if_not_exists: true
+    end
+    add_index :llm_cost_tracker_provider_invoice_imports, %i[source provider started_at],
               if_not_exists: true
   end
 

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

@@ -22,7 +22,9 @@ LlmCostTracker.configure do |config|
   # config.instrument :ruby_llm
 
   # Pricing — local file refreshed via bin/rails llm_cost_tracker:prices:refresh
-  # plus inline overrides. Prices are USD per 1M tokens.
+  # plus inline overrides. Rates are per 1M tokens; the snapshot's currency
+  # is read from your prices_file's `metadata.currency` (USD in the bundled
+  # snapshot — set a different code per file if you maintain non-USD prices).
 <% if options[:prices] -%>
   config.prices_file = Rails.root.join("config/llm_cost_tracker_prices.yml")
 <% else -%>
@@ -37,8 +39,9 @@ LlmCostTracker.configure do |config|
 
   # 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;
+  # raises after recording; :block_requests preflights supported requests, also
+  # estimating the current call's input cost via chars/4 so it can block before
+  # send) 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
@@ -50,11 +53,15 @@ LlmCostTracker.configure do |config|
   # }
   # config.cache_rollups = true
 
-  # Ingestion path — false (default) writes events synchronously from the request
-  # thread. Flip to true for a write-ahead inbox + background worker that batches
+  # Ingestion path — :inline (default) writes events synchronously from the request
+  # thread. Set to :async 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
+  # inbox/leases tables created by `bin/rails generate llm_cost_tracker:async_ingestion`.
+  # The worker uses a dedicated ActiveRecord pool (defaults to 2 connections) so it
+  # doesn't compete with request threads — bump ingestion_pool_size if your Puma
+  # worker count outgrows that.
+  # config.ingestion          = :async
+  # config.ingestion_pool_size = 5
 
   # 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

@@ -1,20 +1,35 @@
+require "llm_cost_tracker/ledger/schema/adapter"
+
 class UpgradeLlmCostTrackerCallRollupsProvider < ActiveRecord::Migration<%= migration_version %>
+  disable_ddl_transaction!
+
   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)
+    add_column TABLE, :provider, :string, null: false, default: "" unless column_exists?(TABLE, :provider)
+    add_unique_index NEW_INDEX
+    remove_index TABLE, column: OLD_INDEX, unique: true, if_exists: 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)
+    add_unique_index OLD_INDEX
+    remove_index TABLE, column: NEW_INDEX, unique: true, if_exists: true
     remove_column TABLE, :provider if column_exists?(TABLE, :provider)
   end
+
+  private
+
+  def add_unique_index(columns)
+    if postgresql?
+      add_index TABLE, columns, unique: true, algorithm: :concurrently, if_not_exists: true
+    else
+      add_index TABLE, columns, unique: true, if_not_exists: true
+    end
+  end
+
+  def postgresql?
+    LlmCostTracker::Ledger::Schema::Adapter.postgresql?(connection)
+  end
 end

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

@@ -1,23 +1,23 @@
 require "llm_cost_tracker/ledger/schema/adapter"
 
 class UpgradeLlmCostTrackerCallTagsKeyValueIndex < ActiveRecord::Migration<%= migration_version %>
+  disable_ddl_transaction!
+
   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
+      add_index TABLE, INDEX_COLUMNS, algorithm: :concurrently, if_not_exists: true
     elsif mysql?
-      add_index TABLE, INDEX_COLUMNS, length: { value: 191 }
+      add_index TABLE, INDEX_COLUMNS, length: { value: 191 }, if_not_exists: true
     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)
+    remove_index TABLE, column: INDEX_COLUMNS, if_exists: true
   end
 
   private

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

@@ -0,0 +1,32 @@
+require "llm_cost_tracker/ledger/schema/adapter"
+
+class UpgradeLlmCostTrackerProviderInvoiceImportsProvider < ActiveRecord::Migration<%= migration_version %>
+  disable_ddl_transaction!
+
+  TABLE = :llm_cost_tracker_provider_invoice_imports
+  OLD_INDEX = %i[source started_at].freeze
+  NEW_INDEX = %i[source provider started_at].freeze
+
+  def up
+    add_column TABLE, :provider, :string, null: false, default: "" unless column_exists?(TABLE, :provider)
+    if postgresql?
+      remove_index TABLE, column: OLD_INDEX, algorithm: :concurrently, if_exists: true
+      add_index TABLE, NEW_INDEX, algorithm: :concurrently, if_not_exists: true
+    else
+      remove_index TABLE, column: OLD_INDEX, if_exists: true
+      add_index TABLE, NEW_INDEX, if_not_exists: true
+    end
+  end
+
+  def down
+    remove_index TABLE, column: NEW_INDEX, if_exists: true
+    add_index TABLE, OLD_INDEX, if_not_exists: true
+    remove_column TABLE, :provider if column_exists?(TABLE, :provider)
+  end
+
+  private
+
+  def postgresql?
+    LlmCostTracker::Ledger::Schema::Adapter.postgresql?(connection)
+  end
+end

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

@@ -0,0 +1,25 @@
+require "llm_cost_tracker/ledger/schema/adapter"
+
+class UpgradeLlmCostTrackerProviderInvoicesMetadataIndex < ActiveRecord::Migration<%= migration_version %>
+  disable_ddl_transaction!
+
+  TABLE = :llm_cost_tracker_provider_invoices
+
+  def up
+    return unless postgresql?
+
+    add_index TABLE, :metadata, using: :gin, algorithm: :concurrently, if_not_exists: true
+  end
+
+  def down
+    return unless postgresql?
+
+    remove_index TABLE, column: :metadata, if_exists: true
+  end
+
+  private
+
+  def postgresql?
+    LlmCostTracker::Ledger::Schema::Adapter.postgresql?(connection)
+  end
+end

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

@@ -19,15 +19,6 @@ module LlmCostTracker
         )
       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

data/lib/llm_cost_tracker/generators/llm_cost_tracker/upgrade_provider_invoice_imports_provider_generator.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+require "rails/generators"
+require "rails/generators/active_record"
+
+module LlmCostTracker
+  module Generators
+    class UpgradeProviderInvoiceImportsProviderGenerator < Rails::Generators::Base
+      include ActiveRecord::Generators::Migration
+
+      source_root File.expand_path("templates", __dir__)
+
+      desc "Adds a provider column to llm_cost_tracker_provider_invoice_imports and a " \
+           "(source, provider, started_at) index so resume_cursor_for and " \
+           "last_completed_window_for can isolate per-provider state on shared sources (e.g. csv)."
+
+      def create_migration_file
+        migration_template(
+          "upgrade_provider_invoice_imports_provider.rb.erb",
+          "db/migrate/upgrade_llm_cost_tracker_provider_invoice_imports_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_provider_invoices_metadata_index_generator.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+require "rails/generators"
+require "rails/generators/active_record"
+
+module LlmCostTracker
+  module Generators
+    class UpgradeProviderInvoicesMetadataIndexGenerator < Rails::Generators::Base
+      include ActiveRecord::Generators::Migration
+
+      source_root File.expand_path("templates", __dir__)
+
+      desc "Adds a GIN index on llm_cost_tracker_provider_invoices.metadata for PostgreSQL " \
+           "so Reconciliation::Diff queries that filter on metadata->>'provider' / 'row_type' / " \
+           "'match_basis' hit an index instead of a sequential scan. No-op on MySQL."
+
+      def create_migration_file
+        migration_template(
+          "upgrade_provider_invoices_metadata_index.rb.erb",
+          "db/migrate/upgrade_llm_cost_tracker_provider_invoices_metadata_index.rb"
+        )
+      end
+
+      private
+
+      def migration_version
+        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ingestion/batch.rb

@@ -40,7 +40,10 @@ module LlmCostTracker
         Ingestion::InboxEntry
           .where(id: rows.map(&:id), locked_by: identity)
           .update_all(last_error: message, locked_at: now, locked_by: nil, updated_at: now)
-      rescue StandardError
+      rescue StandardError => e
+        LlmCostTracker::Logging.warn(
+          "Inbox mark_failed failed for #{rows.size} rows: #{e.class}: #{e.message} (original error: #{error.class})"
+        )
         nil
       end
 
@@ -77,7 +80,7 @@ module LlmCostTracker
 
       def persist(rows, events)
         LlmCostTracker::Call.transaction do
-          Ledger::Store.insert_many(events)
+          Ledger::Store.insert(events)
           Ingestion::InboxEntry.where(id: rows.map(&:id), locked_by: identity).delete_all
         end
       end

data/lib/llm_cost_tracker/ingestion/inbox.rb

@@ -15,8 +15,6 @@ module LlmCostTracker
       class << self
         def save(event)
           insert_row(row_for(event))
-          Ingestion::Worker.ensure_started
-          event
         end
 
         def event_from_row(row)
@@ -54,14 +52,10 @@ module LlmCostTracker
             tracked_at: Time.iso8601(payload.fetch(:tracked_at)),
             cost_status: payload.fetch(:cost_status),
             pricing_snapshot: payload[:pricing_snapshot],
-            line_items: line_items_from(payload)
+            line_items: (payload[:line_items] || []).map { |attrs| Billing::LineItem.build(attrs) }
           }
         end
 
-        def line_items_from(payload)
-          (payload[:line_items] || []).map { |attributes| Billing::LineItem.build(attributes) }
-        end
-
         def row_for(event)
           now = Time.now.utc
           {
@@ -86,25 +80,10 @@ module LlmCostTracker
         end
 
         def insert_row(row)
-          connection = LlmCostTracker::Call.connection
-          if connection.transaction_open?
-            insert_with_separate_connection(row)
-          else
-            execute_insert(connection, row)
-          end
+          Pool.with_connection { |connection| execute_insert(connection, row) }
         rescue ActiveRecord::ConnectionTimeoutError => e
           raise LlmCostTracker::Error,
-                "ledger inbox could not checkout a separate database connection: #{e.message}"
-        end
-
-        def insert_with_separate_connection(row)
-          pool = LlmCostTracker::Call.connection_pool
-          connection = pool.checkout
-          begin
-            connection.transaction(requires_new: true) { execute_insert(connection, row) }
-          ensure
-            pool.checkin(connection)
-          end
+                "ledger inbox could not checkout a database connection: #{e.message}"
         end
 
         def execute_insert(connection, row)

data/lib/llm_cost_tracker/ingestion/pool.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Ingestion
+    module Pool
+      DEFAULT_POOL_SIZE = 2
+      MUTEX = Mutex.new
+
+      class << self
+        def with_connection(&)
+          pool.with_connection(&)
+        end
+
+        def pool
+          @pool || MUTEX.synchronize { @pool ||= connect! }
+        end
+
+        def reset!
+          MUTEX.synchronize do
+            @pool&.disconnect!
+            @pool = nil
+            @handler = nil
+          end
+        end
+
+        private
+
+        def connect!
+          @handler = ActiveRecord::ConnectionAdapters::ConnectionHandler.new
+          @handler.establish_connection(connection_config)
+        end
+
+        def connection_config
+          LlmCostTracker::Call.connection_db_config.configuration_hash.merge(pool: pool_size)
+        end
+
+        def pool_size
+          configured = LlmCostTracker.configuration.ingestion_pool_size.to_i
+          configured.positive? ? configured : DEFAULT_POOL_SIZE
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ingestion/worker.rb

@@ -16,29 +16,31 @@ module LlmCostTracker
       MAX_IDLE_INTERVAL_SECONDS = 5.0
       LEASE_SECONDS = 10
       FLUSH_TIMEOUT_SECONDS = 10
+      MUTEX = Mutex.new
+
       class << self
         def ensure_started
-          return unless Ingestion.durable?
+          return unless Ingestion.async?
 
-          thread = mutex.synchronize do
+          thread = MUTEX.synchronize do
             reset_after_fork!
-            unless @thread&.alive?
-              @stop_requested = false
-              @generation = @generation.to_i + 1
-              generation = @generation
-              @thread = Thread.new { run(generation) }
-              @thread.name = "llm_cost_tracker_ingestor"
-              @thread.report_on_exception = false
-            end
+            break @thread if @stop_requested || @thread&.alive?
+
+            @generation = @generation.to_i + 1
+            generation = @generation
+            @thread = Thread.new { run(generation) }
+            @thread.name = "llm_cost_tracker_ingestor"
+            @thread.report_on_exception = false
             @thread
           end
           wake_thread(thread)
         end
 
         def flush!(timeout: nil, require_lease: false)
-          return true unless Ingestion.durable?
+          return true unless Ingestion.async?
 
           Ingestion.ensure_current_schema!
+          MUTEX.synchronize { reset_after_fork! }
 
           deadline = Time.now.utc + flush_timeout_seconds(timeout)
           loop do
@@ -56,10 +58,10 @@ module LlmCostTracker
         end
 
         def shutdown!(timeout: nil, drain: true)
-          return true unless Ingestion.durable?
+          return true unless Ingestion.async?
 
           timeout ||= FLUSH_TIMEOUT_SECONDS
-          thread = mutex.synchronize do
+          thread = MUTEX.synchronize do
             @stop_requested = true
             @generation = @generation.to_i + 1
             @thread
@@ -71,14 +73,14 @@ module LlmCostTracker
           handle_error(e)
           false
         ensure
-          mutex.synchronize do
+          MUTEX.synchronize do
             @thread = nil if @thread.equal?(thread) && !thread&.alive?
           end
         end
 
         def reset!
-          thread = mutex.synchronize do
-            @stop_requested = true
+          thread = MUTEX.synchronize do
+            @stop_requested = false
             @generation = @generation.to_i + 1
             thread = @thread
             @thread = nil
@@ -98,6 +100,7 @@ module LlmCostTracker
 
         def ingest_once(require_lease: true)
           Ingestion.ensure_current_schema!
+          MUTEX.synchronize { reset_after_fork! }
           batch = Ingestion::Batch.new(identity: identity)
           return 0 unless batch.claimable?
           return 0 if require_lease && !Ingestion::LeaseClaim.new(identity: identity, seconds: LEASE_SECONDS).acquire
@@ -110,16 +113,12 @@ module LlmCostTracker
 
         private
 
-        def mutex
-          @mutex ||= Mutex.new
-        end
-
         def run(generation)
           idle_interval = IDLE_INTERVAL_SECONDS
           loop do
-            break if mutex.synchronize { @stop_requested || generation != @generation }
+            break if MUTEX.synchronize { @stop_requested || generation != @generation }
 
-            processed = executor_wrap { ingest_once }
+            processed = Rails.application.executor.wrap { ingest_once }
             release_connection!
             if processed.zero?
               sleep(idle_interval)
@@ -134,7 +133,7 @@ module LlmCostTracker
           end
         ensure
           release_connection!
-          mutex.synchronize { @thread = nil if @thread.equal?(Thread.current) }
+          MUTEX.synchronize { @thread = nil if @thread.equal?(Thread.current) }
         end
 
         def reset_after_fork!
@@ -151,19 +150,6 @@ module LlmCostTracker
           nil
         end
 
-        def executor_wrap(&)
-          executor = rails_executor
-          return yield unless executor
-
-          executor.wrap(&)
-        end
-
-        def rails_executor
-          Rails.application.try(:executor)
-        rescue StandardError
-          nil
-        end
-
         def identity
           @identity ||= "pid-#{Process.pid}-#{SecureRandom.hex(6)}"
         end

data/lib/llm_cost_tracker/ingestion.rb

@@ -2,11 +2,10 @@
 
 require "securerandom"
 
-require_relative "doctor/check"
 require_relative "errors"
 require_relative "ledger"
-require_relative "ingestion/inline"
 require_relative "ingestion/lease_claim"
+require_relative "ingestion/pool"
 require_relative "ingestion/inbox"
 require_relative "ingestion/batch"
 require_relative "ingestion/worker"
@@ -28,7 +27,7 @@ module LlmCostTracker
 
       ROLLUPS_SCHEMA_GUARD = ["llm_cost_tracker_call_rollups", Ledger::Schema::CallRollups].freeze
 
-      DURABLE_SCHEMA_GUARDS = [
+      ASYNC_SCHEMA_GUARDS = [
         ["llm_cost_tracker_ingestion_inbox_entries", Ledger::Schema::IngestionInboxEntries],
         ["llm_cost_tracker_ingestion_leases",        Ledger::Schema::IngestionLeases]
       ].freeze
@@ -47,8 +46,8 @@ module LlmCostTracker
         end
       end
 
-      def durable?
-        LlmCostTracker.configuration.durable_ingestion
+      def async?
+        LlmCostTracker.configuration.ingestion == :async
       end
 
       def cache_rollups?
@@ -58,7 +57,7 @@ module LlmCostTracker
       def guards_for_current_config
         guards = CORE_SCHEMA_GUARDS.dup
         guards << ROLLUPS_SCHEMA_GUARD if cache_rollups?
-        guards += DURABLE_SCHEMA_GUARDS if durable?
+        guards += ASYNC_SCHEMA_GUARDS if async?
         guards
       end
 
@@ -93,7 +92,7 @@ module LlmCostTracker
           provider_response_id: response_id,
           tags: { feature: VERIFY_TAG }
         )
-        LlmCostTracker::Ingestion::Worker.flush! if durable?
+        LlmCostTracker::Ingestion::Worker.flush! if async?
         persisted = LlmCostTracker::Call.where(provider_response_id: response_id).exists?
 
         return capture_success if persisted && notifications.any?
@@ -122,7 +121,7 @@ module LlmCostTracker
       end
 
       def capture_success
-        path = durable? ? "durable inbox" : "inline writer"
+        path = async? ? "async inbox" : "inline writer"
         LlmCostTracker::Doctor::Check.new(
           :ok,
           "active_record capture",
@@ -149,7 +148,7 @@ module LlmCostTracker
       end
 
       def cleanup_verification_inbox(event:, response_id:)
-        return unless durable? && LlmCostTracker::Ingestion::InboxEntry.table_exists?
+        return unless async? && LlmCostTracker::Ingestion::InboxEntry.table_exists?
 
         if event
           LlmCostTracker::Ingestion::InboxEntry.where(event_id: event.event_id).delete_all