llm_cost_tracker 0.8.0 → 0.10.0

data/lib/llm_cost_tracker/capture/stream_tracker.rb

@@ -12,8 +12,9 @@ module LlmCostTracker
         @stream = stream
         @collector = collector
         @active = active
-        @finish = finish || proc { |errored:| @collector.finish!(errored: errored) }
-        @finished = false
+        @finish = finish || proc { |errored| collector.finish!(errored: errored) }
+        @finished_ref = [false]
+        @attempted_ref = [false]
         @capture_failed = false
         @mutex = Mutex.new
       end
@@ -33,6 +34,7 @@ module LlmCostTracker
         end
         wrap_each if !iterator_wrapped && @stream.respond_to?(:each)
 
+        register_orphan_finalizer
         @stream
       rescue StandardError => e
         Logging.warn("stream integration failed to install wrapper: #{e.class}: #{e.message}")
@@ -120,14 +122,47 @@ module LlmCostTracker
 
       def finish!(errored:)
         should_finish = @mutex.synchronize do
-          next false if @finished
+          @attempted_ref[0] = true
+          next false if @finished_ref[0]
 
-          @finished = true
+          @finished_ref[0] = true
           true
         end
         return unless should_finish && @active.call
 
-        @finish.call(errored: errored)
+        begin
+          @finish.call(errored)
+        rescue StandardError
+          @mutex.synchronize { @finished_ref[0] = false }
+          raise
+        end
+      end
+
+      def register_orphan_finalizer
+        finished_ref = @finished_ref
+        attempted_ref = @attempted_ref
+        finish_proc = @finish
+        active_proc = @active
+        mutex = @mutex
+        finalizer = lambda do |_object_id|
+          should_finish = mutex.synchronize do
+            next false if finished_ref[0] || attempted_ref[0]
+
+            finished_ref[0] = true
+            attempted_ref[0] = true
+            true
+          end
+          next unless should_finish && active_proc.call
+
+          begin
+            Rails.application.executor.wrap { finish_proc.call(false) }
+          rescue StandardError
+            nil
+          end
+        end
+        ObjectSpace.define_finalizer(@stream, finalizer)
+      rescue TypeError, ArgumentError
+        nil
       end
     end
   end

data/lib/llm_cost_tracker/configuration.rb

@@ -14,23 +14,31 @@ 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].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,
+      :ingestion,
       :instrumented_integrations,
       :pricing_overrides,
       :report_tag_breakdowns,
       :redacted_tag_keys,
       :unknown_pricing_behavior,
-      :openai_compatible_providers
+      :openai_compatible_providers,
+      :reconciliation_importers,
+      :reconciliation_enabled,
+      :auto_enable_stream_usage,
+      :cache_rollups
     )
 
     def initialize
@@ -46,38 +54,81 @@ 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 = []
       @redacted_tag_keys = DEFAULT_REDACTED_TAG_KEYS.dup
       self.openai_compatible_providers = OPENAI_COMPATIBLE_PROVIDERS
+      @reconciliation_importers = {}
+      @reconciliation_enabled = false
+      @auto_enable_stream_usage = true
+      self.ingestion = :inline
+      @cache_rollups = false
       @finalized = false
     end
 
+    def cache_rollups=(value)
+      ensure_mutable!
+      @cache_rollups = value
+    end
+
+    def reconciliation_enabled=(value)
+      ensure_mutable!
+      @reconciliation_enabled = value
+    end
+
+    def auto_enable_stream_usage=(value)
+      ensure_mutable!
+      @auto_enable_stream_usage = value
+    end
+
+    def reconciliation_importers=(importers)
+      ensure_mutable!
+      raise Error, RECONCILIATION_DISABLED_MESSAGE unless @reconciliation_enabled
+
+      @reconciliation_importers = (importers || {}).to_h do |source, importer|
+        raise Error, "reconciliation_importers[#{source}] must respond to call" unless importer.respond_to?(:call)
+
+        [source.to_sym, importer]
+      end
+    end
+
+    def register_reconciliation_importer(source, &block)
+      ensure_mutable!
+      raise Error, RECONCILIATION_DISABLED_MESSAGE unless @reconciliation_enabled
+      raise Error, "register_reconciliation_importer requires a block" unless block
+
+      @reconciliation_importers[source.to_sym] = block
+    end
+
+    RECONCILIATION_DISABLED_MESSAGE = "reconciliation is disabled; set config.reconciliation_enabled = true first"
+    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
 
@@ -85,16 +136,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
@@ -109,7 +160,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?
@@ -144,7 +199,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/cost_drift_check.rb

@@ -4,6 +4,7 @@ require "bigdecimal"
 
 require_relative "check"
 require_relative "probe"
+require_relative "../ledger/rollups"
 
 module LlmCostTracker
   class Doctor
@@ -25,6 +26,7 @@ module LlmCostTracker
 
         line_item_totals = LlmCostTracker::CallLineItem
                            .where(llm_cost_tracker_call_id: sampled.map(&:first))
+                           .where(currency: Ledger::Rollups::DEFAULT_CURRENCY)
                            .group(:llm_cost_tracker_call_id)
                            .sum(:cost)
 

data/lib/llm_cost_tracker/doctor/ingestion_check.rb

@@ -11,13 +11,14 @@ module LlmCostTracker
 
       def call
         return unless Probe.table_exists?("llm_cost_tracker_calls")
+        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
@@ -26,23 +27,48 @@ 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
 
       private
 
+      def inline_check
+        leftovers = inline_leftover_tables
+        if leftovers.empty?
+          return Check.new(
+            :ok,
+            "inline ingestion",
+            "config.ingestion = :inline; events write directly to the ledger"
+          )
+        end
+
+        Check.new(
+          :warn,
+          "inline ingestion",
+          "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
+
+      def inline_leftover_tables
+        [
+          LlmCostTracker::Ingestion::InboxEntry.table_name,
+          LlmCostTracker::Ingestion::Lease.table_name
+        ].select { |table| Probe.table_exists?(table) }
+      end
+
       def missing_parts
         [
           LlmCostTracker::Ingestion::InboxEntry.table_name,

data/lib/llm_cost_tracker/doctor/invoice_reconciliation_check.rb

@@ -0,0 +1,164 @@
+# frozen_string_literal: true
+
+require "bigdecimal"
+
+require_relative "check"
+require_relative "probe"
+require_relative "../ledger/schema/adapter"
+
+module LlmCostTracker
+  class Doctor
+    class InvoiceReconciliationCheck
+      def call
+        return unless LlmCostTracker.reconciliation_enabled?
+        return unless Probe.table_exists?("llm_cost_tracker_provider_invoices")
+        return if no_imports?
+
+        scopes = imported_scopes
+        return Check.new(:ok, "invoice reconciliation", "no provider invoices imported yet") if scopes.empty?
+
+        non_canonical = non_canonical_currency_check
+        checks = scopes.map { |scope| check_scope_safely(scope) }
+        checks << non_canonical if non_canonical
+        checks
+      rescue StandardError => e
+        Check.new(:error, "invoice reconciliation", e.message)
+      end
+
+      private
+
+      def no_imports?
+        LlmCostTracker::ProviderInvoice.none?
+      end
+
+      def non_canonical_currency_check
+        legacy = LlmCostTracker::ProviderInvoice.where("currency <> UPPER(currency)").count
+        return nil if legacy.zero?
+
+        Check.new(
+          :warn,
+          "invoice reconciliation: currency canonicalisation",
+          "#{legacy} provider invoice row(s) stored with non-uppercase currency. Diff queries " \
+          "are case-sensitive — run " \
+          "`UPDATE llm_cost_tracker_provider_invoices SET currency = UPPER(currency);` to backfill."
+        )
+      end
+
+      def threshold
+        Reconciliation::DEFAULT_THRESHOLD_PERCENT
+      end
+
+      def imported_scopes
+        connection = LlmCostTracker::ProviderInvoice.connection
+        provider_expr =
+          if Ledger::Schema::Adapter.postgresql?(connection)
+            Arel.sql("metadata->>'provider'")
+          else
+            Arel.sql("JSON_UNQUOTE(JSON_EXTRACT(metadata, '$.provider'))")
+          end
+        LlmCostTracker::ProviderInvoice
+          .group(:source, provider_expr, :currency)
+          .order(:source, :currency)
+          .pluck(:source, provider_expr, :currency)
+          .map { |source, provider, currency| { source: source, provider: provider, currency: currency.upcase } }
+      end
+
+      def scope_label(scope)
+        "#{scope[:source]}/#{scope[:provider]}/#{scope[:currency]}"
+      end
+
+      def check_scope_safely(scope)
+        check_scope(scope)
+      rescue ArgumentError => e
+        Check.new(:warn, "invoice reconciliation: #{scope_label(scope)}", e.message)
+      end
+
+      def check_scope(scope)
+        window = latest_window_for(scope)
+        return stale_check(scope) if window.nil?
+
+        diff = run_diff(scope, window)
+        return ok_check(scope, window, diff) if diff.aligned?(threshold_percent: threshold)
+
+        warn_check(scope, window, diff)
+      end
+
+      def scope_relation(scope)
+        relation = LlmCostTracker::ProviderInvoice
+                   .where(source: scope[:source], currency: scope[:currency])
+        provider = scope[:provider]
+        return relation if provider.nil? || provider.to_s.empty?
+
+        connection = LlmCostTracker::ProviderInvoice.connection
+        if Ledger::Schema::Adapter.postgresql?(connection)
+          relation.where("metadata->>'provider' = ?", provider)
+        else
+          relation.where("JSON_UNQUOTE(JSON_EXTRACT(metadata, '$.provider')) = ?", provider)
+        end
+      end
+
+      def latest_window_for(scope)
+        latest = scope_relation(scope)
+                 .select(:period_start, :period_end)
+                 .order(period_end: :desc, period_start: :desc)
+                 .limit(1)
+                 .first
+        return nil unless latest
+        return nil if (Time.now.utc.to_date - latest.period_end).to_i > Reconciliation::INVOICE_FRESHNESS_DAYS
+
+        latest
+      end
+
+      def run_diff(scope, window)
+        Reconciliation.diff(
+          source: scope[:source],
+          provider: scope[:provider],
+          currency: scope[:currency],
+          period_start: window.period_start,
+          period_end: window.period_end
+        )
+      end
+
+      def stale_check(scope)
+        latest = scope_relation(scope).maximum(:period_end)
+        return scope_unreachable_check(scope) if latest.nil?
+
+        days = (Time.now.utc.to_date - latest).to_i
+        Check.new(
+          :warn,
+          "invoice reconciliation: #{scope_label(scope)}",
+          "no invoice imported in #{days} days (threshold #{Reconciliation::INVOICE_FRESHNESS_DAYS} days); " \
+          "run reconciliation import"
+        )
+      end
+
+      def scope_unreachable_check(scope)
+        Check.new(
+          :warn,
+          "invoice reconciliation: #{scope_label(scope)}",
+          "scope grouped invoices but the filter (likely currency case mismatch) matches zero rows; " \
+          "the currency-canonicalisation check below points at the backfill SQL"
+        )
+      end
+
+      def ok_check(scope, window, diff)
+        Check.new(
+          :ok,
+          "invoice reconciliation: #{scope_label(scope)}",
+          "#{window.period_start}..#{window.period_end} aligned " \
+          "(local=#{diff.local_total.to_s('F')}, provider=#{diff.provider_total.to_s('F')})"
+        )
+      end
+
+      def warn_check(scope, window, diff)
+        Check.new(
+          :warn,
+          "invoice reconciliation: #{scope_label(scope)}",
+          "#{window.period_start}..#{window.period_end} drift " \
+          "delta=#{diff.delta_amount.to_s('F')} (#{diff.delta_percent}%) " \
+          "exceeds #{threshold}% threshold"
+        )
+      end
+    end
+  end
+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/schema_check.rb

@@ -7,14 +7,17 @@ require_relative "../ledger"
 module LlmCostTracker
   class Doctor
     class SchemaCheck
-      def initialize(name:, schema:, table:)
+      def initialize(name:, schema:, table:, optional: false, install_command: "llm_cost_tracker:install")
         @name = name
         @schema = schema
         @table = table
+        @optional = optional
+        @install_command = install_command
       end
 
       def call
         return unless Probe.table_exists?("llm_cost_tracker_calls")
+        return if @optional && !Probe.table_exists?(@table)
 
         errors = @schema.current_schema_errors
         return Check.new(:ok, @name, "#{@table} exists") if errors.empty?
@@ -23,7 +26,7 @@ module LlmCostTracker
           :error,
           @name,
           "current schema required; #{errors.join('; ')}; " \
-          "run bin/rails generate llm_cost_tracker:install && bin/rails db:migrate"
+          "run bin/rails generate #{@install_command} && bin/rails db:migrate"
         )
       end
     end