llm_cost_tracker 0.14.1 → 0.15.0

data/lib/llm_cost_tracker/ingestion/inbox.rb

@@ -11,6 +11,9 @@ module LlmCostTracker
   module Ingestion
     module Inbox
       PAYLOAD_SCHEMA_VERSION = 2
+      OPTIONAL_FIELDS = %i[
+        latency_ms usage_source provider_response_id provider_project_id provider_api_key_id provider_workspace_id
+      ].freeze
 
       class << self
         def save(event)
@@ -18,7 +21,7 @@ module LlmCostTracker
         end
 
         def event_from_row(row)
-          payload = JSON.parse(row.payload, symbolize_names: true)
+          payload = JSON.parse(row.payload, symbolize_names: true, allow_duplicate_key: true)
           schema_version = payload[:schema_version]
           unless schema_version == PAYLOAD_SCHEMA_VERSION
             raise LlmCostTracker::Error, "unsupported ledger inbox payload schema version #{schema_version.inspect}"
@@ -30,9 +33,8 @@ module LlmCostTracker
         private
 
         def event_attributes_from(payload)
-          cost = payload[:cost] && Charges::Cost.from_h(payload[:cost])
-          token_usage = Usage::TokenUsage.build(**payload.fetch(:token_usage).slice(*Usage::TokenUsage.members))
-
+          cost = cost_from(payload)
+          token_usage = token_usage_from(payload)
           {
             event_id: payload.fetch(:event_id),
             provider: payload.fetch(:provider),
@@ -41,20 +43,27 @@ module LlmCostTracker
             pricing_mode: Pricing::Mode.normalize(payload[:pricing_mode]),
             cost: cost,
             tags: payload.fetch(:tags),
-            latency_ms: payload[:latency_ms],
             stream: payload.fetch(:stream),
-            usage_source: payload[:usage_source],
-            provider_response_id: payload[:provider_response_id],
-            provider_project_id: payload[:provider_project_id],
-            provider_api_key_id: payload[:provider_api_key_id],
-            provider_workspace_id: payload[:provider_workspace_id],
+            **OPTIONAL_FIELDS.to_h { |field| [field, payload[field]] },
             tracked_at: Time.iso8601(payload.fetch(:tracked_at)),
             cost_status: payload.fetch(:cost_status),
             pricing_snapshot: payload[:pricing_snapshot]&.deep_stringify_keys,
-            line_items: (payload[:line_items] || []).map { |attrs| Charges::LineItem.build(attrs) }
+            line_items: line_items_from(payload)
           }
         end
 
+        def cost_from(payload)
+          payload[:cost] && Charges::Cost.from_h(payload[:cost])
+        end
+
+        def token_usage_from(payload)
+          Usage::TokenUsage.build(**payload.fetch(:token_usage).slice(*Usage::TokenUsage.members))
+        end
+
+        def line_items_from(payload)
+          (payload[:line_items] || []).map { |attrs| Charges::LineItem.build(attrs) }
+        end
+
         def row_for(event)
           now = Time.now.utc
           {
@@ -69,17 +78,11 @@ module LlmCostTracker
         end
 
         def payload_for(event)
-          event.to_h.merge(
-            schema_version: PAYLOAD_SCHEMA_VERSION,
-            event_id: event.event_id,
-            provider: event.provider,
-            model: event.model,
-            tracked_at: event.tracked_at.iso8601(6)
-          )
+          event.to_h.merge(schema_version: PAYLOAD_SCHEMA_VERSION, tracked_at: event.tracked_at.iso8601(6))
         end
 
         def insert_row(row)
-          Pool.with_connection { |connection| execute_insert(connection, row) }
+          Ledger::Isolation.writing { Pool.with_connection { |connection| execute_insert(connection, row) } }
         rescue ActiveRecord::ConnectionTimeoutError => e
           raise LlmCostTracker::Error,
                 "ledger inbox could not checkout a database connection: #{e.message}"

data/lib/llm_cost_tracker/ingestion/worker.rb

@@ -38,39 +38,21 @@ module LlmCostTracker
         def flush!(timeout: nil, require_lease: false)
           return true unless Ingestion.async?
 
-          Ingestion.ensure_current_schema!
-          MUTEX.synchronize { reset_after_fork! }
-
-          deadline = Time.now.utc + flush_timeout_seconds(timeout)
-          loop do
-            return true unless Ingestion::Batch.new(identity: identity).pending?
-            return false if Time.now.utc >= deadline
-
-            processed = ingest_once(require_lease: require_lease)
-            next unless processed.zero?
-
-            duration = [INTERVAL_SECONDS, deadline - Time.now.utc].min
-            return false unless duration.positive?
-
-            sleep(duration)
+          prepare!
+          drained = drain_until(Time.now.utc + flush_timeout_seconds(timeout), require_lease)
+          unless drained
+            Logging.warn("Ingestion::Worker.flush! timed out with #{Ingestion::InboxEntry.pending.count} " \
+                         "inbox row(s) still pending")
           end
+          drained && Ingestion::InboxEntry.uncached { Ingestion::InboxEntry.quarantined.none? }
         end
 
         def shutdown!(timeout: nil, drain: true)
           return true unless Ingestion.async?
 
           timeout ||= FLUSH_TIMEOUT_SECONDS
-          thread = MUTEX.synchronize do
-            @stop_requested = true
-            @generation = @generation.to_i + 1
-            @thread
-          end
-          begin
-            wake_thread(thread)
-            thread&.join(timeout)
-          rescue StandardError => e
-            handle_error(e)
-          end
+          thread = request_stop
+          await_exit(thread, timeout)
           drain ? flush!(timeout: timeout, require_lease: true) : true
         rescue StandardError => e
           handle_error(e)
@@ -89,8 +71,7 @@ module LlmCostTracker
         end
 
         def ingest_once(require_lease: true)
-          Ingestion.ensure_current_schema!
-          MUTEX.synchronize { reset_after_fork! }
+          prepare!
           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
@@ -108,14 +89,7 @@ module LlmCostTracker
           loop do
             break if MUTEX.synchronize { @stop_requested || generation != @generation }
 
-            processed = Rails.application.executor.wrap { ingest_once }
-            release_connection!
-            if processed.zero?
-              sleep(idle_interval)
-              idle_interval = [idle_interval * 2, MAX_IDLE_INTERVAL_SECONDS].min
-            else
-              idle_interval = IDLE_INTERVAL_SECONDS
-            end
+            idle_interval = ingest_or_back_off(idle_interval)
           rescue StandardError => e
             handle_error(e)
             release_connection!
@@ -126,6 +100,52 @@ module LlmCostTracker
           MUTEX.synchronize { @thread = nil if @thread.equal?(Thread.current) }
         end
 
+        def ingest_or_back_off(idle_interval)
+          processed = Rails.application.executor.wrap { ingest_once }
+          release_connection!
+          return IDLE_INTERVAL_SECONDS unless processed.zero?
+
+          sleep(idle_interval)
+          [idle_interval * 2, MAX_IDLE_INTERVAL_SECONDS].min
+        end
+
+        def prepare!
+          Ingestion.ensure_current_schema!
+          MUTEX.synchronize { reset_after_fork! }
+        end
+
+        def drain_until(deadline, require_lease)
+          loop do
+            return true unless Ingestion::Batch.new(identity: identity).pending?
+            return false unless time_left(deadline).positive?
+            next unless ingest_once(require_lease: require_lease).zero?
+
+            duration = [INTERVAL_SECONDS, time_left(deadline)].min
+            return false unless duration.positive?
+
+            sleep(duration)
+          end
+        end
+
+        def time_left(deadline)
+          deadline - Time.now.utc
+        end
+
+        def request_stop
+          MUTEX.synchronize do
+            @stop_requested = true
+            @generation = @generation.to_i + 1
+            @thread
+          end
+        end
+
+        def await_exit(thread, timeout)
+          wake_thread(thread)
+          thread&.join(timeout)
+        rescue StandardError => e
+          handle_error(e)
+        end
+
         def reset_after_fork!
           return if @pid == Process.pid
 
@@ -145,7 +165,8 @@ module LlmCostTracker
         end
 
         def handle_error(error)
-          Logging.warn("ActiveRecord ingestor failed: #{error.class}: #{error.message}")
+          Logging.warn("ActiveRecord ingestor failed: #{error.class}: #{error.message}" \
+                       "#{" (#{error.backtrace.first})" if error.backtrace}")
         end
 
         def release_connection!

data/lib/llm_cost_tracker/ingestion.rb

@@ -16,6 +16,8 @@ module LlmCostTracker
     autoload :Worker, "llm_cost_tracker/ingestion/worker"
 
     VERIFY_TAG = "llm_cost_tracker_verify"
+    CAPTURE_CHECK = "active_record capture"
+    MISSING_CALLS_TABLE = "llm_cost_tracker_calls table is missing; run install generator and migrate"
 
     class << self
       def table_name_prefix
@@ -23,11 +25,11 @@ module LlmCostTracker
       end
 
       def ensure_current_schema!
-        unless LlmCostTracker::Call.table_exists?
-          raise Error, "llm_cost_tracker_calls table is missing; run install generator and migrate"
-        end
+        raise Error, MISSING_CALLS_TABLE unless LlmCostTracker::Call.table_exists?
 
         guards_for_current_config.each do |schema_module, table_name|
+          next if schema_module == Ledger::Schema::CallRollups && !Ledger::Rollups.cache_active?
+
           errors = schema_module.current_schema_errors
           next if errors.empty?
 
@@ -49,13 +51,7 @@ module LlmCostTracker
 
       def verify
         unless LlmCostTracker::Call.table_exists?
-          return [
-            LlmCostTracker::Check.new(
-              :error,
-              "active_record",
-              "llm_cost_tracker_calls table is missing; run install generator and migrate"
-            )
-          ]
+          return [LlmCostTracker::Check.new(:error, "active_record", MISSING_CALLS_TABLE)]
         end
 
         [capture_check]
@@ -70,7 +66,6 @@ module LlmCostTracker
         response_id = "lct_verify_#{SecureRandom.hex(8)}"
         notifications = []
         subscription = subscribe_to_verification(response_id, notifications)
-
         event = LlmCostTracker.track(
           provider: provider,
           model: model,
@@ -79,21 +74,9 @@ module LlmCostTracker
           tags: { feature: VERIFY_TAG }
         )
         LlmCostTracker::Ingestion::Worker.flush! if async?
-        persisted = LlmCostTracker::Call.where(provider_response_id: response_id).exists?
-
-        return capture_success if persisted && notifications.any?
-
-        LlmCostTracker::Check.new(
-          :error,
-          "active_record capture",
-          capture_failure_message(persisted, notifications)
-        )
-      rescue LlmCostTracker::BudgetExceededError => e
-        LlmCostTracker::Check.new(:error, "active_record capture", "blocked by budget guardrail: #{e.message}")
-      rescue LlmCostTracker::Error => e
-        LlmCostTracker::Check.new(:error, "active_record capture", e.message)
+        capture_result(LlmCostTracker::Call.where(provider_response_id: response_id).exists?, notifications)
       rescue StandardError => e
-        LlmCostTracker::Check.new(:error, "active_record capture", "#{e.class}: #{e.message}")
+        LlmCostTracker::Check.new(:error, CAPTURE_CHECK, capture_error_message(e))
       ensure
         cleanup_verification_call(response_id) if response_id
         cleanup_verification_inbox(event: event, response_id: response_id)
@@ -106,25 +89,29 @@ module LlmCostTracker
         end
       end
 
-      def capture_success
-        path = async? ? "async inbox" : "inline writer"
-        LlmCostTracker::Check.new(
-          :ok,
-          "active_record capture",
-          "manual event emitted and persisted through #{path}"
-        )
-      end
+      def capture_result(persisted, notifications)
+        if persisted && notifications.any?
+          path = async? ? "async inbox" : "inline writer"
+          return LlmCostTracker::Check.new(:ok, CAPTURE_CHECK, "manual event emitted and persisted through #{path}")
+        end
 
-      def capture_failure_message(persisted, notifications)
         missing = []
         missing << "notification" if notifications.empty?
         missing << "persisted row" unless persisted
-        "missing #{missing.join(' and ')} for synthetic manual event"
+        LlmCostTracker::Check.new(:error, CAPTURE_CHECK, "missing #{missing.join(' and ')} for synthetic manual event")
+      end
+
+      def capture_error_message(error)
+        case error
+        when LlmCostTracker::BudgetExceededError then "blocked by budget guardrail: #{error.message}"
+        when LlmCostTracker::Error then error.message
+        else "#{error.class}: #{error.message}"
+        end
       end
 
       def cleanup_verification_call(response_id)
         relation = LlmCostTracker::Call.where(provider_response_id: response_id)
-        records = relation.select(:id, :tracked_at, :total_cost, :pricing_snapshot, :provider).to_a
+        records = relation.select(:id, *LlmCostTracker::Ledger::Rollups::SOURCE_COLUMNS).to_a
         return if records.empty?
 
         relation.delete_all

data/lib/llm_cost_tracker/integrations/anthropic/batch_results_capture.rb

@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Integrations
+    module Anthropic
+      class BatchResultsCapture
+        include Enumerable
+
+        def initialize(raw_stream)
+          @raw_stream = raw_stream
+        end
+
+        def each(&block)
+          return enum_for(:each) unless block
+
+          deferred = nil
+          @raw_stream.each do |response|
+            begin
+              Integrations::Anthropic.record_batch_result(response)
+            rescue BudgetExceededError, UnknownPricingError => e
+              deferred ||= e
+            end
+            block.call(response)
+          end
+          raise deferred if deferred
+        end
+
+        def respond_to_missing?(name, include_private = false)
+          @raw_stream.respond_to?(name, include_private) || super
+        end
+
+        def method_missing(name, ...)
+          return super unless @raw_stream.respond_to?(name)
+
+          @raw_stream.public_send(name, ...)
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/integrations/anthropic/patches.rb

@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Integrations
+    module Anthropic
+      module MessagesPatch
+        def create(*args, **kwargs)
+          Integrations::Anthropic.wrap_blocking(args, kwargs, **Integrations::Anthropic.blocking_seam(@client)) { super }
+        end
+
+        def stream(*args, **kwargs)
+          Integrations::Anthropic.wrap_stream(args, kwargs, **Integrations::Anthropic.stream_seam(@client)) { super }
+        end
+
+        def stream_raw(*args, **kwargs)
+          Integrations::Anthropic.wrap_stream(args, kwargs, **Integrations::Anthropic.stream_seam(@client)) { super }
+        end
+      end
+
+      module FallbackMiddlewarePatch
+        def call(req, nxt)
+          return super if req.streaming? || !Integrations::Anthropic.active?
+
+          hops = []
+          response = super(req, ->(hop_req) { nxt.call(hop_req).tap { |hop| hops << [hop_req, hop] } })
+          answered = hops.select { |_hop_req, hop| hop.status < 300 }
+          answered[...-1].each { |hop| Integrations::Anthropic.record_refused_hop(*hop) }
+          response
+        end
+
+        private
+
+        def consume_hop(*args, **kwargs)
+          refused = Thread.current[:llm_cost_tracker_refused_hop]
+          Thread.current[:llm_cost_tracker_refused_hop] = nil
+          Integrations::Anthropic.record_refused_stream_hop(refused) if refused && kwargs[:splice]
+          super.tap { |hop| Thread.current[:llm_cost_tracker_refused_hop] = hop if hop[:refused] }
+        end
+      end
+
+      module BatchesPatch
+        def create(*args, **kwargs)
+          Integrations::Anthropic.enforce_budget!(request: Integrations::Anthropic.request_params(args, kwargs))
+          super
+        end
+
+        def results_streaming(*args, **kwargs)
+          raw = super
+          return raw unless Integrations::Anthropic.active?
+
+          BatchResultsCapture.new(raw)
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/integrations/anthropic.rb

@@ -1,8 +1,9 @@
 # frozen_string_literal: true
 
 require_relative "base"
-require_relative "../providers/anthropic/usage_extractor"
 require_relative "../providers/anthropic/response_parser"
+require_relative "anthropic/patches"
+require_relative "anthropic/batch_results_capture"
 
 module LlmCostTracker
   module Integrations
@@ -17,27 +18,48 @@ module LlmCostTracker
             patch_target("Anthropic::Resources::Messages", with: MessagesPatch),
             patch_target("Anthropic::Resources::Beta::Messages", with: MessagesPatch, optional: true),
             patch_target("Anthropic::Resources::Messages::Batches", with: BatchesPatch, optional: true),
-            patch_target("Anthropic::Resources::Beta::Messages::Batches", with: BatchesPatch, optional: true)
+            patch_target("Anthropic::Resources::Beta::Messages::Batches", with: BatchesPatch, optional: true),
+            patch_target("Anthropic::BetaRefusalFallbackMiddleware",
+                         with: FallbackMiddlewarePatch,
+                         optional: true,
+                         since: "1.54.0")
           ]
         end
 
-        def record_message(message, request:, latency_ms:)
+        def blocking_seam(client)
+          host = client_host(client)
+          {
+            record: lambda do |message, request, latency_ms|
+              record_message(message, request: request, latency_ms: latency_ms, host: host)
+            end
+          }
+        end
+
+        def stream_seam(client)
+          host = client_host(client)
+          { collector: ->(request) { stream_collector(request, host: host) } }
+        end
+
+        def stream_collector(request, host: nil)
+          regional = Providers::Anthropic::UsageExtractor.regional_host?(host, request[:model])
+          super(request, pricing_mode: ("data_residency" if regional))
+        end
+
+        def record_message(message, request:, latency_ms:, host: nil)
           return unless active?
 
           record_safely do
             usage = message.usage
-            next unless usage
-            next if usage.input_tokens.nil? && usage.output_tokens.nil?
-
-            usage_hash = usage.deep_to_h
+            next unless tokens_reported?(usage)
 
-            LlmCostTracker::Tracker.record(
-              event: Providers::Anthropic::ResponseParser.event_from_usage(
-                usage: usage_hash,
+            Tracker.record(
+              event: message_event(
+                message,
+                usage: usage.deep_to_h,
                 model: message.model || request[:model],
-                provider_response_id: message.id,
                 usage_source: Usage::Source::SDK_RESPONSE,
-                request: request
+                request: request,
+                host: host
               ),
               latency_ms: latency_ms
             )
@@ -48,106 +70,76 @@ module LlmCostTracker
           return unless active?
 
           record_safely do
-            next unless response.respond_to?(:result) && response.result
-
-            result = response.result
-            next unless result.respond_to?(:type) && result.type.to_s == "succeeded"
-
-            message = result.respond_to?(:message) ? result.message : nil
+            message = succeeded_message(response)
             next unless message
-            next if LlmCostTracker::Call.already_recorded?(provider: "anthropic", provider_response_id: message.id)
+            next if Call.already_recorded?(provider: "anthropic", provider_response_id: message.id)
 
             usage = message.usage
-            next unless usage
-            next if usage.input_tokens.nil? && usage.output_tokens.nil?
+            next unless tokens_reported?(usage)
 
-            usage_hash = usage.deep_to_h
-            LlmCostTracker::Tracker.record(
-              event: Providers::Anthropic::ResponseParser.event_from_usage(
-                usage: usage_hash,
+            record_once(
+              message_event(
+                message,
+                usage: usage.deep_to_h.merge(service_tier: "batch"),
                 model: message.model,
-                provider_response_id: message.id,
-                usage_source: Usage::Source::SDK_BATCH_RESULT,
-                pricing_mode: "batch"
+                usage_source: Usage::Source::SDK_BATCH_RESULT
               )
             )
           end
         end
 
-        def stream_pricing_mode(request)
-          Providers::Anthropic::UsageExtractor.pricing_mode(request: request || {}, usage: nil)
+        def record_refused_hop(request, response)
+          record_safely { record_message(response.parse, request: request.body, latency_ms: nil) }
+        rescue BudgetExceededError, UnknownPricingError
+          nil
         end
-      end
 
-      module MessagesPatch
-        def create(*args, **kwargs)
-          LlmCostTracker::Integrations::Anthropic.wrap_blocking(
-            args,
-            kwargs,
-            record: lambda do |message, request, latency_ms|
-              LlmCostTracker::Integrations::Anthropic.record_message(
-                message, request: request, latency_ms: latency_ms
-              )
-            end
-          ) { super }
-        end
-
-        def stream(*args, **kwargs)
-          LlmCostTracker::Integrations::Anthropic.wrap_stream(
-            args,
-            kwargs,
-            collector: ->(request) { LlmCostTracker::Integrations::Anthropic.stream_collector(request) }
-          ) { super }
-        end
-
-        def stream_raw(*args, **kwargs)
-          LlmCostTracker::Integrations::Anthropic.wrap_stream(
-            args,
-            kwargs,
-            collector: ->(request) { LlmCostTracker::Integrations::Anthropic.stream_collector(request) }
-          ) { super }
-        end
-      end
+        def record_refused_stream_hop(hop)
+          return unless active?
 
-      module BatchesPatch
-        def results_streaming(*args, **kwargs)
-          raw = super
-          return raw unless LlmCostTracker::Integrations::Anthropic.active?
+          usage = hop.dig(:refused, :usage).to_h.deep_symbolize_keys
+          produced_output = usage[:output_tokens].to_i.positive?
+          usage = usage.slice(:server_tool_use) if produced_output
+          return if produced_output && Providers::Anthropic::UsageExtractor.service_line_items(usage).empty?
 
-          BatchResultsCapture.new(raw)
+          record_safely do
+            Tracker.record(
+              event: Providers::Anthropic::ResponseParser.event_from_usage(
+                usage: usage,
+                model: hop[:model],
+                provider_response_id: nil,
+                usage_source: Usage::Source::STREAM_FINAL,
+                stream: true,
+                stop_reason: ("refusal" unless produced_output),
+                refusal_category: hop.dig(:refused, :stop_details, "category")
+              )
+            )
+          end
+        rescue BudgetExceededError, UnknownPricingError
+          nil
         end
-      end
-
-      class BatchResultsCapture
-        include Enumerable
 
-        def initialize(raw_stream)
-          @raw_stream = raw_stream
-        end
+        private
 
-        def each(&block)
-          return enum_for(:each) unless block
+        def succeeded_message(response)
+          result = response.result if response.respond_to?(:result)
+          return unless result.respond_to?(:type) && result.type.to_s == "succeeded"
 
-          deferred = nil
-          @raw_stream.each do |response|
-            begin
-              LlmCostTracker::Integrations::Anthropic.record_batch_result(response)
-            rescue LlmCostTracker::BudgetExceededError, LlmCostTracker::UnknownPricingError => e
-              deferred ||= e
-            end
-            block.call(response)
-          end
-          raise deferred if deferred
+          result.message if result.respond_to?(:message)
         end
 
-        def respond_to_missing?(name, include_private = false)
-          @raw_stream.respond_to?(name, include_private) || super
+        def tokens_reported?(usage)
+          usage && !(usage.input_tokens.nil? && usage.output_tokens.nil?)
         end
 
-        def method_missing(name, ...)
-          return super unless @raw_stream.respond_to?(name)
-
-          @raw_stream.public_send(name, ...)
+        def message_event(message, **attributes)
+          Providers::Anthropic::ResponseParser.event_from_usage(
+            provider_response_id: message.id,
+            stop_reason: message.stop_reason,
+            refusal_category: message.stop_details&.category,
+            content: message.deep_to_h[:content],
+            **attributes
+          )
         end
       end
     end