llm_cost_tracker 0.8.0 → 0.9.0

data/lib/llm_cost_tracker/capture/stream_collector.rb

@@ -5,6 +5,7 @@ require "active_support/core_ext/object/deep_dup"
 require "json"
 
 require_relative "stream"
+require_relative "../pricing/mode"
 require_relative "../timing"
 
 module LlmCostTracker
@@ -14,7 +15,7 @@ module LlmCostTracker
 
       def initialize(provider:, model:, latency_ms: nil, provider_response_id: nil, provider_project_id: nil,
                      provider_api_key_id: nil, provider_workspace_id: nil, batch: nil, pricing_mode: nil,
-                     metadata: {}, context_tags: nil)
+                     metadata: {}, context_tags: nil, request: nil)
         @provider = provider.to_s
         @model = model
         @latency_ms = latency_ms
@@ -26,12 +27,14 @@ module LlmCostTracker
         @pricing_mode = pricing_mode
         @metadata = (metadata || {}).deep_dup
         @context_tags = (context_tags || LlmCostTracker::Tags::Context.tags).deep_dup
+        @request = request
         @events = []
         @captured_bytes = 0
         @overflowed = false
         @explicit_usage = nil
         @started_at = LlmCostTracker::Timing.now_monotonic
         @finished = false
+        @recording = false
         @mutex = Mutex.new
       end
 
@@ -88,10 +91,19 @@ module LlmCostTracker
       end
 
       def finish!(errored: false)
-        snapshot = @mutex.synchronize do
-          return if @finished
+        snapshot = claim_recording_slot
+        return if snapshot.nil?
 
-          @finished = true
+        record_snapshot(snapshot, errored: errored)
+      end
+
+      private
+
+      def claim_recording_slot
+        @mutex.synchronize do
+          return nil if @finished || @recording
+
+          @recording = true
           pricing_mode = Pricing.normalize_mode(@pricing_mode)
           {
             events: @events.dup,
@@ -103,24 +115,52 @@ module LlmCostTracker
             capture_dimensions: capture_dimensions(pricing_mode),
             pricing_mode: pricing_mode,
             metadata: @metadata.deep_dup,
-            context_tags: @context_tags.deep_dup
+            context_tags: @context_tags.deep_dup,
+            request: @request
           }
         end
+      end
+
+      def record_snapshot(snapshot, errored:)
+        save_succeeded = false
+        begin
+          capture = build_usage_capture(snapshot)
+          provider_response_id = capture.provider_response_id || snapshot[:provider_response_id]
+          capture = capture.with(provider_response_id: provider_response_id)
+
+          Tracker.record(
+            capture: capture,
+            latency_ms: snapshot[:latency_ms] || LlmCostTracker::Timing.elapsed_ms(@started_at),
+            pricing_mode: pricing_mode_for(capture: capture, snapshot: snapshot),
+            metadata: (errored ? { stream_errored: true } : {}).merge(snapshot[:metadata]),
+            context_tags: snapshot[:context_tags]
+          ) { |stage| save_succeeded = true if stage == :after_save }
+        ensure
+          @mutex.synchronize do
+            @finished = save_succeeded
+            @recording = false
+          end
+        end
+      end
 
-        capture = build_usage_capture(snapshot)
-        provider_response_id = capture.provider_response_id || snapshot[:provider_response_id]
-        capture = capture.with(provider_response_id: provider_response_id)
+      HOST_DERIVED_MODE_TOKENS = %i[data_residency].freeze
+      STANDARD_LIKE_MODE_TOKENS = %i[standard standard_only auto default].freeze
+      private_constant :HOST_DERIVED_MODE_TOKENS, :STANDARD_LIKE_MODE_TOKENS
 
-        Tracker.record(
-          capture: capture,
-          latency_ms: snapshot[:latency_ms] || LlmCostTracker::Timing.elapsed_ms(@started_at),
-          pricing_mode: snapshot[:pricing_mode],
-          metadata: (errored ? { stream_errored: true } : {}).merge(snapshot[:metadata]),
-          context_tags: snapshot[:context_tags]
-        )
+      def pricing_mode_for(capture:, snapshot:)
+        merge_pricing_modes(capture.pricing_mode, snapshot[:pricing_mode])
       end
 
-      private
+      def merge_pricing_modes(provider_mode, request_mode)
+        return Pricing.normalize_mode(request_mode) if provider_mode.to_s.strip.empty?
+
+        provider_tokens = Pricing::Mode.tokenize(provider_mode) - STANDARD_LIKE_MODE_TOKENS
+        request_host_tokens = Pricing::Mode.tokenize(request_mode || "") & HOST_DERIVED_MODE_TOKENS
+        combined = provider_tokens | request_host_tokens
+        return nil if combined.empty?
+
+        Pricing.normalize_mode(combined.join("_"))
+      end
 
       def capture_dimensions(pricing_mode)
         batch = @batch.nil? ? UsageCapture.batch_from_pricing_mode?(pricing_mode).presence : @batch
@@ -140,12 +180,14 @@ module LlmCostTracker
 
       def build_usage_capture(snapshot)
         return build_from_explicit_usage(snapshot) if snapshot[:explicit_usage]
+        return build_unknown_usage(snapshot) if snapshot[:overflowed]
 
         capture = Parsers.find_for_provider(@provider)&.parse_stream(
           response_status: 200,
-          events: snapshot[:events]
+          events: snapshot[:events],
+          request_body: request_body_for(snapshot[:request])
         )
-        if capture && (capture.usage_source != :unknown || !snapshot[:overflowed])
+        if capture
           model = present_model(capture.model) || present_model(snapshot[:model]) || UsageCapture::UNKNOWN_MODEL
           return capture.with(provider: @provider, model: model, **snapshot.fetch(:capture_dimensions))
         end
@@ -153,6 +195,14 @@ module LlmCostTracker
         build_unknown_usage(snapshot)
       end
 
+      def request_body_for(request)
+        return nil unless request
+
+        JSON.generate(request)
+      rescue StandardError
+        nil
+      end
+
       def present_model(value)
         return nil if value.nil?
 
@@ -186,8 +236,14 @@ module LlmCostTracker
         )
       end
 
+      IGNORED_PAYLOAD_KEYS = %w[b64_json partial_image_b64].freeze
+      private_constant :IGNORED_PAYLOAD_KEYS
+
+      HEAVY_STRING_BYTES = 8 * 1024
+      private_constant :HEAVY_STRING_BYTES
+
       def capture_event(data, type:)
-        event = { event: type, data: data }
+        event = { event: type, data: strip_heavy_payload(data) }
         size = JSON.generate(event).bytesize
         if @captured_bytes + size <= Capture::Stream::LIMIT_BYTES
           @events << event.deep_dup
@@ -195,9 +251,26 @@ module LlmCostTracker
         else
           @overflowed = true
         end
-      rescue JSON::JSONError, TypeError
+      rescue JSON::JSONError, TypeError, SystemStackError
         @overflowed = true
       end
+
+      def strip_heavy_payload(value)
+        case value
+        when Hash
+          value.each_with_object({}) do |(key, nested), out|
+            next if IGNORED_PAYLOAD_KEYS.include?(key.to_s)
+
+            out[key] = strip_heavy_payload(nested)
+          end
+        when Array
+          value.map { |nested| strip_heavy_payload(nested) }
+        when String
+          value.bytesize > HEAVY_STRING_BYTES ? "" : value
+        else
+          value
+        end
+      end
     end
   end
 end

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
+            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

@@ -25,12 +25,17 @@ module LlmCostTracker
     attr_reader(
       *SHARED_SCALAR_ATTRIBUTES,
       :budget_exceeded_behavior,
+      :durable_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
@@ -51,9 +56,56 @@ module LlmCostTracker
       @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
+      @durable_ingestion = false
+      @cache_rollups = false
       @finalized = false
     end
 
+    def durable_ingestion=(value)
+      ensure_shared_configuration_mutable!
+      @durable_ingestion = value
+    end
+
+    def cache_rollups=(value)
+      ensure_shared_configuration_mutable!
+      @cache_rollups = value
+    end
+
+    def reconciliation_enabled=(value)
+      ensure_shared_configuration_mutable!
+      @reconciliation_enabled = value
+    end
+
+    def auto_enable_stream_usage=(value)
+      ensure_shared_configuration_mutable!
+      @auto_enable_stream_usage = value
+    end
+
+    def reconciliation_importers=(importers)
+      ensure_shared_configuration_mutable!
+      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_shared_configuration_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!
       @openai_compatible_providers = normalize_openai_compatible_providers(providers)

data/lib/llm_cost_tracker/dashboard_setup_state.rb

@@ -0,0 +1,109 @@
+# frozen_string_literal: true
+
+require_relative "ledger/schema/calls"
+require_relative "ledger/schema/call_line_items"
+require_relative "ledger/schema/call_tags"
+require_relative "ledger/schema/call_rollups"
+
+module LlmCostTracker
+  module DashboardSetupState
+    SetupRequired = Data.define(:message, :details)
+    DOCS_HINT = "See docs/upgrading.md for the migration path."
+    MUTEX = Mutex.new
+
+    CORE_SCHEMA_CHECKS = [
+      [
+        LlmCostTracker::Ledger::Schema::Calls,
+        "The llm_cost_tracker_calls table does not match the current LLM Cost Tracker schema."
+      ],
+      [
+        LlmCostTracker::Ledger::Schema::CallLineItems,
+        "The llm_cost_tracker_call_line_items table does not match the current LLM Cost Tracker schema."
+      ],
+      [
+        LlmCostTracker::Ledger::Schema::CallTags,
+        "The llm_cost_tracker_call_tags table does not match the current LLM Cost Tracker schema."
+      ]
+    ].freeze
+
+    OPTIONAL_CALL_ROLLUPS_CHECK = [
+      LlmCostTracker::Ledger::Schema::CallRollups,
+      "The llm_cost_tracker_call_rollups table does not match the current LLM Cost Tracker schema."
+    ].freeze
+
+    private_constant :MUTEX, :CORE_SCHEMA_CHECKS, :OPTIONAL_CALL_ROLLUPS_CHECK, :DOCS_HINT
+
+    class << self
+      def current
+        return @cached if defined?(@cached)
+
+        MUTEX.synchronize do
+          @cached = compute unless defined?(@cached)
+        end
+        @cached
+      end
+
+      def reset!
+        MUTEX.synchronize do
+          remove_instance_variable(:@cached) if defined?(@cached)
+        end
+      end
+
+      private
+
+      def compute
+        LlmCostTracker::Logging.debug("DashboardSetupState recomputing")
+        return calls_table_missing unless LlmCostTracker::Call.table_exists?
+
+        core_drift = drift_in(schema_checks_for_current_config)
+        return core_drift if core_drift
+        return nil unless LlmCostTracker.reconciliation_enabled?
+
+        reconciliation_drift
+      end
+
+      def schema_checks_for_current_config
+        return CORE_SCHEMA_CHECKS unless LlmCostTracker.configuration.cache_rollups
+
+        CORE_SCHEMA_CHECKS + [OPTIONAL_CALL_ROLLUPS_CHECK]
+      end
+
+      def drift_in(checks)
+        checks.each do |schema, message|
+          errors = schema.current_schema_errors
+          next if errors.empty?
+
+          return SetupRequired.new(message: message, details: errors + [DOCS_HINT])
+        end
+        nil
+      end
+
+      def reconciliation_drift
+        LlmCostTracker.const_get(:Reconciliation) # autoload reconciliation + its ledger schemas
+        connection = ActiveRecord::Base.connection
+        LlmCostTracker::Reconciliation::SCHEMA_TABLES.each do |schema, table|
+          unless connection.data_source_exists?(table)
+            return SetupRequired.new(
+              message: "The #{table} table is required when reconciliation is enabled.",
+              details: ["run bin/rails generate llm_cost_tracker:reconciliation && bin/rails db:migrate", DOCS_HINT]
+            )
+          end
+
+          errors = schema.current_schema_errors
+          next if errors.empty?
+
+          message = "The #{table} table does not match the current LLM Cost Tracker schema."
+          return SetupRequired.new(message: message, details: errors + [DOCS_HINT])
+        end
+        nil
+      end
+
+      def calls_table_missing
+        SetupRequired.new(
+          message: "The llm_cost_tracker_calls table is not available yet.",
+          details: nil
+        )
+      end
+    end
+  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,6 +11,7 @@ module LlmCostTracker
 
       def call
         return unless Probe.table_exists?("llm_cost_tracker_calls")
+        return inline_check unless LlmCostTracker::Ingestion.durable?
 
         missing = missing_parts
         if missing.empty?
@@ -43,6 +44,31 @@ module LlmCostTracker
 
       private
 
+      def inline_check
+        leftovers = inline_leftover_tables
+        if leftovers.empty?
+          return Check.new(
+            :ok,
+            "inline ingestion",
+            "durable_ingestion=false; events write directly to the ledger"
+          )
+        end
+
+        Check.new(
+          :warn,
+          "inline ingestion",
+          "durable_ingestion=false but found unused durable ingestion tables: #{leftovers.join(', ')}. " \
+          "Set config.durable_ingestion = true to keep the durable inbox path or drop the tables."
+        )
+      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/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