llm_cost_tracker 0.14.0 → 0.14.1

data/lib/llm_cost_tracker/capture/sse.rb

@@ -11,52 +11,83 @@ module LlmCostTracker
 
       class << self
         def parse(body)
-          return [] if body.blank?
+          events = []
+          reader = Reader.new { |event| events << event }
+          reader << body
+          reader.finish
+          events
+        end
+      end
 
-          return parse_json_array(body) if body.match?(/\A\s*\[/)
+      class Reader
+        def initialize(&on_event)
+          @on_event = on_event
+          @pending = String.new(encoding: Encoding::BINARY)
+          @mode = nil
+          @event_name = nil
+          @data_lines = []
+          @scan_pos = 0
+          @depth = 0
+          @in_string = false
+          @object_start = nil
+        end
 
-          parse_event_stream(body)
+        def <<(chunk)
+          @pending << chunk.to_s.b
+          @mode ||= detect_mode
+          case @mode
+          when :sse then consume_lines
+          when :array then consume_array
+          end
         end
 
-        private
+        def pending_bytesize
+          @pending.bytesize
+        end
 
-        def parse_event_stream(body)
-          events = []
-          current_event = nil
-          data_lines = []
+        def finish
+          return unless @mode == :sse
 
-          body.each_line do |raw|
-            line = raw.chomp
+          consume_line(@pending) unless @pending.empty?
+          @pending = String.new(encoding: Encoding::BINARY)
+          dispatch
+        end
 
-            if line.empty?
-              events << finalize_event(current_event, data_lines) if data_lines.any?
-              current_event = nil
-              data_lines = []
-              next
-            end
+        private
 
-            next if line.start_with?(":")
+        def detect_mode
+          first = @pending[/\S/]
+          return nil unless first
 
-            field, _, value = line.partition(":")
-            value = value[1..] if value.start_with?(" ")
+          first == "[" ? :array : :sse
+        end
 
-            case field
-            when "event" then current_event = value
-            when "data"  then data_lines << value
-            end
+        def consume_lines
+          start = 0
+          while (newline = @pending.index("\n", start))
+            consume_line(@pending.byteslice(start, newline - start))
+            start = newline + 1
           end
-
-          events << finalize_event(current_event, data_lines) if data_lines.any?
-          events.compact
+          @pending = @pending.byteslice(start..)
         end
 
-        def parse_json_array(body)
-          parsed = JSON.parse(body)
-          return [] unless parsed.is_a?(Array)
+        def consume_line(raw)
+          line = raw.chomp("\r").force_encoding(Encoding::UTF_8)
+          return dispatch if line.empty?
+          return if line.start_with?(":")
 
-          parsed.map { |entry| { event: nil, data: entry } }
-        rescue JSON::ParserError
-          []
+          field, _, value = line.partition(":")
+          value = value[1..] if value.start_with?(" ")
+          case field
+          when "event" then @event_name = value
+          when "data" then @data_lines << value
+          end
+        end
+
+        def dispatch
+          emit(finalize_event(@event_name, @data_lines)) if @data_lines.any?
+          @event_name = nil
+          @data_lines = []
         end
 
         def finalize_event(event_name, data_lines)
@@ -73,6 +104,51 @@ module LlmCostTracker
         rescue JSON::ParserError
           payload
         end
+
+        def consume_array
+          pos = @scan_pos
+          while (index = @pending.index(@in_string ? /["\\]/ : /[{}"]/, pos))
+            pos = index + 1
+            pos = advance_array(@pending.getbyte(index).chr, index, pos)
+          end
+          trim_array_buffer(pos)
+        end
+
+        def advance_array(char, index, pos)
+          if @in_string
+            return pos + 1 if char == "\\"
+
+            @in_string = false
+          elsif char == '"'
+            @in_string = true
+          elsif char == "{"
+            @object_start = index if @depth.zero?
+            @depth += 1
+          elsif @depth.positive?
+            @depth -= 1
+            emit_object(index) if @depth.zero?
+          end
+          pos
+        end
+
+        def emit_object(end_index)
+          text = @pending.byteslice(@object_start, end_index - @object_start + 1)
+          @object_start = nil
+          emit({ event: nil, data: JSON.parse(text.force_encoding(Encoding::UTF_8)) })
+        rescue JSON::ParserError
+          nil
+        end
+
+        def trim_array_buffer(pos)
+          keep_from = @object_start || [pos, @pending.bytesize].min
+          @pending = @pending.byteslice(keep_from..)
+          @scan_pos = pos - keep_from
+          @object_start &&= 0
+        end
+
+        def emit(event)
+          @on_event.call(event) if event
+        end
       end
     end
   end

data/lib/llm_cost_tracker/capture/stream_collector.rb

@@ -4,7 +4,7 @@ require "active_support/core_ext/object/blank"
 require "active_support/core_ext/object/deep_dup"
 require "json"
 
-require_relative "sse"
+require_relative "event_window"
 require_relative "../timing"
 
 module LlmCostTracker
@@ -34,9 +34,7 @@ module LlmCostTracker
         @metadata = (metadata || {}).deep_dup
         @context_tags = (context_tags || LlmCostTracker::Tags::Context.tags).deep_dup
         @request = request
-        @events = []
-        @captured_bytes = 0
-        @overflowed = false
+        @window = EventWindow.new(notable: Parsers.find_for_provider(@provider)&.method(:retain_stream_event?))
         @explicit_usage = nil
         @started_at = LlmCostTracker::Timing.now_monotonic
         @finished = false
@@ -61,7 +59,7 @@ module LlmCostTracker
       def event(data, type: nil)
         @mutex.synchronize do
           ensure_open!
-          capture_event(data, type: type) unless data.nil?
+          @window.push(data, type: type) unless data.nil?
         end
       end
 
@@ -91,6 +89,12 @@ module LlmCostTracker
         return if snapshot.nil?
 
         record_snapshot(snapshot, errored: errored)
+      rescue TransactionAbortedError
+        raise
+      rescue StandardError => e
+        raise unless errored
+
+        Logging.warn("Recording an errored stream raised #{e.class}: #{e.message}; kept the stream's own exception")
       end
 
       private
@@ -102,8 +106,8 @@ module LlmCostTracker
           @recording = true
           pricing_mode = Pricing::Mode.normalize(@pricing_mode)
           {
-            events: @events.dup,
-            overflowed: @overflowed,
+            events: @window.events,
+            overflowed: @window.overflowed?,
             explicit_usage: @explicit_usage,
             model: @model,
             latency_ms: @latency_ms,
@@ -136,10 +140,16 @@ module LlmCostTracker
           @mutex.synchronize do
             @finished = save_succeeded
             @recording = false
+            release_buffers if save_succeeded
           end
         end
       end
 
+      def release_buffers
+        @window = EventWindow.new
+        @request = nil
+      end
+
       def capture_dimensions
         {
           provider_project_id: @provider_project_id.to_s.strip.presence,
@@ -207,55 +217,6 @@ module LlmCostTracker
           **snapshot.fetch(:capture_dimensions)
         )
       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: strip_heavy_payload(data) }
-        size = approximate_bytesize(event)
-        if @captured_bytes + size <= Capture::SSE::LIMIT_BYTES
-          @events << event
-          @captured_bytes += size
-        else
-          @overflowed = true
-        end
-      rescue 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
-
-      def approximate_bytesize(value)
-        case value
-        when Hash
-          value.sum { |key, nested| approximate_bytesize(key) + approximate_bytesize(nested) + 4 }
-        when Array
-          value.sum { |nested| approximate_bytesize(nested) + 2 }
-        when Numeric, true, false, nil
-          8
-        else
-          value.to_s.bytesize + 2
-        end
-      end
     end
   end
 end

data/lib/llm_cost_tracker/capture/stream_tap.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+require_relative "event_window"
+require_relative "sse"
+
+module LlmCostTracker
+  module Capture
+    class StreamTap
+      MAX_PENDING_BYTES = 16 * 1_048_576
+
+      def initialize(notable: nil)
+        @window = EventWindow.new(notable: notable)
+        @reader = SSE::Reader.new { |event| @window.push(event[:data], type: event[:event]) }
+        @received = false
+        @failed = false
+      end
+
+      def <<(chunk)
+        return if @failed
+
+        @received = true
+        @reader << chunk
+        fail! if @reader.pending_bytesize > MAX_PENDING_BYTES
+      rescue StandardError
+        fail!
+      end
+
+      def received?
+        @received
+      end
+
+      def failed?
+        @failed
+      end
+
+      def overflowed?
+        @window.overflowed?
+      end
+
+      def events
+        @reader&.finish
+        @failed ? [] : @window.events
+      rescue StandardError
+        fail!
+        []
+      end
+
+      private
+
+      def fail!
+        @failed = true
+        @reader = nil
+        @window = EventWindow.new
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/capture/stream_tracker.rb

@@ -13,8 +13,7 @@ module LlmCostTracker
         @collector = collector
         @active = active
         @finish = finish || proc { |errored| collector.finish!(errored: errored) }
-        @finished_ref = [false]
-        @attempted_ref = [false]
+        @finished = false
         @capture_failed = false
         @mutex = Mutex.new
       end
@@ -45,7 +44,6 @@ module LlmCostTracker
           )
         end
 
-        register_orphan_finalizer
         @stream
       rescue StandardError => e
         Logging.warn("stream integration failed to install wrapper: #{e.class}: #{e.message}")
@@ -77,7 +75,7 @@ module LlmCostTracker
             yield event
           end
         end
-      rescue StandardError
+      rescue Exception # rubocop:disable Lint/RescueException -- mark Interrupt/shutdown too, so it is not masked
         errored = true
         raise
       ensure
@@ -107,10 +105,9 @@ module LlmCostTracker
 
       def finish!(errored:)
         should_finish = @mutex.synchronize do
-          @attempted_ref[0] = true
-          next false if @finished_ref[0]
+          next false if @finished
 
-          @finished_ref[0] = true
+          @finished = true
           true
         end
         return unless should_finish && @active.call
@@ -118,37 +115,10 @@ module LlmCostTracker
         begin
           @finish.call(errored)
         rescue StandardError
-          @mutex.synchronize { @finished_ref[0] = false }
+          @mutex.synchronize { @finished = 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
 end

data/lib/llm_cost_tracker/doctor/price_check.rb

@@ -13,6 +13,7 @@ module LlmCostTracker
       def call
         path = LlmCostTracker.configuration.pricing.file
         return bundled_check unless path
+        return Check.new(:error, "prices", "#{path} does not exist; #{REFRESH_COMMAND}") unless File.exist?(path)
 
         count = LlmCostTracker::Pricing::Registry.file_prices(path).size
         metadata = LlmCostTracker::Pricing::Registry.file_metadata(path)

data/lib/llm_cost_tracker/errors.rb

@@ -1,10 +1,21 @@
 # frozen_string_literal: true
 
+require_relative "redaction"
+
 module LlmCostTracker
   class Error < StandardError; end
 
   class InvalidFilterError < Error; end
 
+  class TransactionAbortedError < Error
+    def initialize(error)
+      super(
+        "The database rolled back the whole surrounding transaction while recording LLM usage " \
+        "(#{error.class}: #{Redaction.text(error.message)}); the caller's transaction no longer exists"
+      )
+    end
+  end
+
   class BudgetExceededError < Error
     attr_reader :total, :budget, :budget_type, :last_event, :stage, :scope
 
@@ -40,4 +51,6 @@ module LlmCostTracker
       super("No pricing configured for LLM model: #{model.inspect}")
     end
   end
+
+  CALLER_ERRORS = [BudgetExceededError, UnknownPricingError, TransactionAbortedError].freeze
 end

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

@@ -34,8 +34,8 @@ LlmCostTracker.configure do |config|
   #   "my-custom-model" => { input: 1.00, output: 2.00 }
   # }
   # :warn (default) records token usage with nil cost when a model has no rate.
-  # Use :raise to reject models that have no rate at all; a model priced with one
-  # component rate missing lands as partial and does not trigger it.
+  # Use :raise to record and then raise for models that have no rate at all; a model
+  # priced with one component rate missing lands as partial and does not trigger it.
   config.pricing.unknown_model_behavior = :warn
 
   # Budget guardrails — cumulative monthly/daily and per-call ceilings in USD,
@@ -63,9 +63,9 @@ LlmCostTracker.configure do |config|
   # }
 
   # Where budget checks read period spend from. :ledger sums llm_cost_tracker_calls on
-  # every check; :cache keeps running totals in llm_cost_tracker_call_rollups so checks
-  # read one row instead — switch once the live SUM gets slow, and generate the table
-  # first with `bin/rails generate llm_cost_tracker:call_rollups`.
+  # every check; :cache also keeps running totals in llm_cost_tracker_call_rollups and
+  # takes the greater of the two, so it guards against drift but does not make checks
+  # faster. Generate the table first with `bin/rails generate llm_cost_tracker:call_rollups`.
   # config.budgets.totals_source = :cache
 
   # Ingestion path — :inline (default) writes events synchronously from the request
@@ -85,7 +85,7 @@ LlmCostTracker.configure do |config|
 
   # Streaming calls only report token usage when the request asks for it. Leave this
   # on and the Faraday middleware adds stream_options: { include_usage: true } to
-  # OpenAI-compatible streaming request bodies that don't already set it.
+  # streaming chat-completions requests on hosts known to accept it.
   # config.capture.request_stream_usage = false
 
 end

data/lib/llm_cost_tracker/ingestion/batch.rb

@@ -13,7 +13,8 @@ module LlmCostTracker
         ActiveRecord::Deadlocked,
         ActiveRecord::LockWaitTimeout,
         ActiveRecord::StatementTimeout,
-        ActiveRecord::ConnectionNotEstablished
+        ActiveRecord::ConnectionNotEstablished,
+        LlmCostTracker::TransactionAbortedError
       ].freeze
 
       def initialize(identity:)
@@ -25,7 +26,7 @@ module LlmCostTracker
         return 0 if rows.empty?
 
         valid_rows, events = decode(rows)
-        persist(valid_rows, events) if events.any?
+        persist_batch(valid_rows, events) if events.any?
         rows.size
       rescue StandardError => e
         rows_to_mark = valid_rows&.any? ? valid_rows : rows
@@ -67,7 +68,7 @@ module LlmCostTracker
       end
 
       def error_message_for(error)
-        "#{error.class}: #{error.message}".byteslice(0, 1_000)
+        "#{error.class}: #{Redaction.text(error.message)}".byteslice(0, 1_000).scrub("")
       end
 
       def warn_on_quarantine(rows)
@@ -75,12 +76,10 @@ module LlmCostTracker
         quarantined = rows.select { |row| row.attempts.to_i + 1 >= threshold }
         return if quarantined.empty?
 
-        sample = quarantined.first(10).map(&:id).join(", ")
-        sample += "..." if quarantined.size > 10
         LlmCostTracker::Logging.warn(
           "Ingestion::Batch: #{quarantined.size} inbox row(s) reached " \
           "MAX_ATTEMPTS_BEFORE_QUARANTINE=#{threshold} and will be skipped " \
-          "on the next claim cycle (ids: #{sample})"
+          "on the next claim cycle (ids: #{id_sample(quarantined)})"
         )
       end
 
@@ -88,6 +87,11 @@ module LlmCostTracker
 
       attr_reader :identity
 
+      def id_sample(rows)
+        sample = rows.first(10).map(&:id).join(", ")
+        rows.size > 10 ? "#{sample}..." : sample
+      end
+
       def claim
         now = Time.now.utc
         cutoff = now - LOCK_TIMEOUT_SECONDS
@@ -117,13 +121,38 @@ module LlmCostTracker
         [valid_rows, events]
       end
 
+      def persist_batch(rows, events)
+        landed = []
+        failed = Hash.new { |hash, message| hash[message] = [] }
+        landed.concat(persist(rows, events))
+      rescue *TRANSIENT_PERSIST_ERRORS
+        raise
+      rescue StandardError
+        rows.zip(events) do |row, event|
+          landed.concat(persist([row], [event]))
+        rescue *TRANSIENT_PERSIST_ERRORS
+          raise
+        rescue StandardError => e
+          failed[error_message_for(e)] << row
+        end
+      ensure
+        failed.each do |message, failed_rows|
+          LlmCostTracker::Logging.warn(
+            "Ingestion::Batch: #{failed_rows.size} inbox row(s) could not be stored " \
+            "(ids: #{id_sample(failed_rows)}): #{message}"
+          )
+          mark_failed_with_message(failed_rows, message)
+        end
+        Ledger::Rollups.increment_safely!(landed)
+        Budget.notify_persisted_safely!(landed)
+      end
+
       def persist(rows, events, retry_on_conflict: true)
         LlmCostTracker::Call.transaction do
           Ledger::Store.persist_records(events)
           Ingestion::InboxEntry.where(id: rows.map(&:id), locked_by: identity).delete_all
         end
-        Ledger::Rollups.increment_safely!(events)
-        Budget.notify_persisted_safely!(events)
+        events
       rescue ActiveRecord::RecordNotUnique
         raise unless retry_on_conflict
 

data/lib/llm_cost_tracker/integrations/anthropic.rb

@@ -128,10 +128,16 @@ module LlmCostTracker
         def each(&block)
           return enum_for(:each) unless block
 
+          deferred = nil
           @raw_stream.each do |response|
-            LlmCostTracker::Integrations::Anthropic.record_batch_result(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
         end
 
         def respond_to_missing?(name, include_private = false)

data/lib/llm_cost_tracker/integrations/base.rb

@@ -27,6 +27,7 @@ module LlmCostTracker
 
       def install
         validate_contract!
+        Logging.warn(untested_version_message) if untested_version?
         patch_targets.each do |target|
           target_class = target.fetch(:constant_name).to_s.safe_constantize
           install_patch(target_class, target.fetch(:patch)) if target_class
@@ -39,6 +40,7 @@ module LlmCostTracker
         if problems.any?
           return Check.new(:warn, name, "#{name} integration cannot be installed: #{problems.join('; ')}")
         end
+        return Check.new(:warn, name, untested_version_message) if untested_version?
 
         installed = patch_targets.reject { |target| target.fetch(:optional) }.all? do |target|
           target.fetch(:constant_name).to_s.safe_constantize&.ancestors&.include?(target.fetch(:patch))
@@ -60,7 +62,7 @@ module LlmCostTracker
 
       def record_safely
         yield
-      rescue LlmCostTracker::Error
+      rescue *LlmCostTracker::CALLER_ERRORS
         raise
       rescue StandardError => e
         Logging.warn("#{integration_name} integration failed to record usage: #{e.class}: #{e.message}")
@@ -124,6 +126,11 @@ module LlmCostTracker
         @minimum_version
       end
 
+      def maximum_version(value = nil)
+        @maximum_version = value if value
+        @maximum_version
+      end
+
       def gem_version
         Gem.loaded_specs[integration_name.to_s]&.version
       end
@@ -160,6 +167,15 @@ module LlmCostTracker
         ["#{name} >= #{minimum_version} is required, detected #{version}"]
       end
 
+      def untested_version?
+        maximum_version && gem_version && gem_version >= Gem::Version.new(maximum_version)
+      end
+
+      def untested_version_message
+        "#{integration_name} #{gem_version} is newer than the tested range (< #{maximum_version}); " \
+          "its calls may not be recorded"
+      end
+
       def target_problems
         patch_targets.flat_map do |target|
           constant_name = target.fetch(:constant_name)

data/lib/llm_cost_tracker/integrations/openai/batch_capture.rb

@@ -21,8 +21,9 @@ module LlmCostTracker
             host = Openai.client_host_for(resource)
             Openai.record_safely do
               io = client.files.content(batch.output_file_id)
-              capture_jsonl(io.respond_to?(:read) ? io.read : io.to_s, host: host)
+              deferred = capture_jsonl(io.respond_to?(:read) ? io.read : io.to_s, host: host)
               mark_captured(batch.id)
+              raise deferred if deferred
             end
           end
 
@@ -41,6 +42,7 @@ module LlmCostTracker
           end
 
           def capture_jsonl(jsonl, host:)
+            deferred = nil
             jsonl.each_line do |line|
               line = line.strip
               next if line.empty?
@@ -52,7 +54,10 @@ module LlmCostTracker
               next unless response.is_a?(Hash) && response["usage"]
 
               record_result(response, host: host)
+            rescue LlmCostTracker::BudgetExceededError, LlmCostTracker::UnknownPricingError => e
+              deferred ||= e
             end
+            deferred
           end
 
           def parse_line(line)