llm_cost_tracker 0.14.2 → 0.15.0

data/lib/llm_cost_tracker/event.rb

@@ -24,33 +24,36 @@ module LlmCostTracker
     :cost_status,
     :pricing_snapshot,
     :line_items
-  ) do
-    def self.build(**attributes)
-      token_usage = attributes.fetch(:token_usage)
-      line_items = attributes[:line_items] || resolve_line_items(attributes[:service_line_items])
+  )
+
+  class Event
+    UNKNOWN_MODEL = "unknown"
+    PROVIDER_IDS = %i[provider_response_id provider_project_id provider_api_key_id provider_workspace_id].freeze
+    PASSTHROUGH_FIELDS = %i[event_id pricing_mode cost tags latency_ms tracked_at cost_status pricing_snapshot].freeze
+    private_constant :PROVIDER_IDS, :PASSTHROUGH_FIELDS
 
+    def self.build(**attributes)
       new(
-        event_id: attributes[:event_id],
+        token_usage: attributes.fetch(:token_usage),
+        line_items: attributes[:line_items] || resolve_line_items(attributes[:service_line_items]),
         provider: attributes.fetch(:provider).to_s,
-        model: attributes.fetch(:model).to_s.strip.presence || Event::UNKNOWN_MODEL,
-        token_usage: token_usage,
-        pricing_mode: attributes[:pricing_mode],
-        cost: attributes[:cost],
-        tags: attributes[:tags],
-        latency_ms: attributes[:latency_ms],
+        model: identifier(attributes.fetch(:model)) || UNKNOWN_MODEL,
         stream: attributes[:stream] || false,
         usage_source: attributes[:usage_source]&.to_s,
-        provider_response_id: attributes[:provider_response_id].to_s.strip.presence,
-        provider_project_id: attributes[:provider_project_id].to_s.strip.presence,
-        provider_api_key_id: attributes[:provider_api_key_id].to_s.strip.presence,
-        provider_workspace_id: attributes[:provider_workspace_id].to_s.strip.presence,
-        tracked_at: attributes[:tracked_at],
-        cost_status: attributes[:cost_status],
-        pricing_snapshot: attributes[:pricing_snapshot],
-        line_items: line_items
+        **PASSTHROUGH_FIELDS.to_h { |field| [field, attributes[field]] },
+        **PROVIDER_IDS.to_h { |field| [field, identifier(attributes[field])] }
       )
     end
 
+    def self.resolve_line_items(service_items)
+      Array(service_items).map do |item|
+        item.is_a?(Charges::LineItem) ? item : Charges::LineItem.build(item)
+      end
+    end
+
+    def self.identifier(value) = value.to_s.strip.presence
+    private_class_method :identifier
+
     def batch?
       Pricing::Mode.tokenize(pricing_mode.to_s).include?("batch")
     end
@@ -61,12 +64,6 @@ module LlmCostTracker
       with(event_id: Digest::UUID.uuid_v5(Digest::UUID::OID_NAMESPACE, "#{provider}/#{provider_response_id}"))
     end
 
-    def self.resolve_line_items(service_items)
-      Array(service_items).map do |item|
-        item.is_a?(Charges::LineItem) ? item : Charges::LineItem.build(item)
-      end
-    end
-
     def total_cost
       cost&.total
     end
@@ -80,8 +77,4 @@ module LlmCostTracker
       )
     end
   end
-
-  class Event
-    UNKNOWN_MODEL = "unknown"
-  end
 end

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

@@ -84,8 +84,10 @@ class CreateLlmCostTrackerCalls < ActiveRecord::Migration[<%= ActiveRecord::Migr
     add_index :llm_cost_tracker_calls, %i[model tracked_at]
     add_index :llm_cost_tracker_calls, :cost_status
     add_index :llm_cost_tracker_calls, :provider_response_id
-    add_index :llm_cost_tracker_calls, :id, name: :index_llm_cost_tracker_calls_on_unpriced,
-                                          where: "total_cost IS NULL"
+    if postgresql?
+      add_index :llm_cost_tracker_calls, :id, name: :index_llm_cost_tracker_calls_on_unpriced,
+                                            where: "total_cost IS NULL"
+    end
     add_index :llm_cost_tracker_call_line_items, [:llm_cost_tracker_call_id, :position]
     add_index :llm_cost_tracker_call_tags, :llm_cost_tracker_call_id
     if postgresql?

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

@@ -60,7 +60,8 @@ LlmCostTracker.configure do |config|
   #   user_id: { daily: 25.00, behavior: :notify }
   # }
   # config.budgets.on_exceeded = ->(data) {
-  #   Rails.logger.warn("LLM #{data[:budget_type]} budget exceeded: $#{data[:total]} / $#{data[:budget]}")
+  #   unit = data[:budget_type] == :calls ? "" : "$"
+  #   Rails.logger.warn("LLM #{data[:budget_type]} budget exceeded: #{unit}#{data[:total]} / #{unit}#{data[:budget]}")
   # }
 
   # Where budget checks read period spend from. :ledger sums llm_cost_tracker_calls on
@@ -80,9 +81,10 @@ LlmCostTracker.configure do |config|
   # config.ingestion.mode = :async
   # config.ingestion.pool_size = 5
 
-  # Register OpenAI-compatible gateway hosts beyond the built-in OpenRouter, DeepSeek and Groq;
-  # the Faraday middleware captures only listed hosts. Its calls, and openai gem calls whose
-  # base_url is a listed host, are recorded under the provider name given here.
+  # Register OpenAI-compatible gateway hosts beyond the built-in OpenRouter, DeepSeek, Groq,
+  # xAI, Mistral and Perplexity; the Faraday middleware captures only listed hosts. Its calls,
+  # and openai gem calls whose base_url is a listed host, are recorded under the provider name
+  # given here.
   # config.capture.openai_compatible_providers["llm.my-company.com"] = "internal_gateway"
 
   # Streaming calls only report token usage when the request asks for it. Leave this

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

@@ -27,6 +27,7 @@ class UpgradeLlmCostTrackerCallRollupsProvider < ActiveRecord::Migration[<%= Act
 
   def add_unique_index(columns)
     if postgresql?
+      LlmCostTracker::Ledger::Schema::Adapter.drop_invalid_index(connection, connection.index_name(TABLE, columns))
       add_index TABLE, columns, unique: true, algorithm: :concurrently, if_not_exists: true
     else
       add_index TABLE, columns, unique: true, if_not_exists: true

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

@@ -8,6 +8,7 @@ class UpgradeLlmCostTrackerCallTagsKeyValueIndex < ActiveRecord::Migration[<%= A
 
   def up
     if postgresql?
+      LlmCostTracker::Ledger::Schema::Adapter.drop_invalid_index(connection, connection.index_name(TABLE, INDEX_COLUMNS))
       add_index TABLE, INDEX_COLUMNS, algorithm: :concurrently, if_not_exists: true
     elsif mysql?
       add_index TABLE, INDEX_COLUMNS, length: { value: 191 }, if_not_exists: true

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

@@ -9,7 +9,8 @@ class UpgradeLlmCostTrackerIndexes < ActiveRecord::Migration[<%= ActiveRecord::M
   UNPRICED_INDEX = :index_llm_cost_tracker_calls_on_unpriced
 
   def up
-    add_index CALLS, :id, name: UNPRICED_INDEX, where: "total_cost IS NULL", **add_index_options
+    LlmCostTracker::Ledger::Schema::Adapter.drop_invalid_index(connection, UNPRICED_INDEX)
+    add_index CALLS, :id, name: UNPRICED_INDEX, where: "total_cost IS NULL", **add_index_options if postgresql?
     return unless table_exists?(INBOX)
 
     remove_index INBOX, column: UNUSED_INBOX_INDEX, **remove_index_options

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

@@ -10,6 +10,7 @@ class UpgradeLlmCostTrackerPerTagBudgets < ActiveRecord::Migration[<%= ActiveRec
   def up
     add_column TABLE, :total_cost, :decimal, precision: 20, scale: 8 unless column_exists?(TABLE, :total_cost)
     add_column TABLE, :tracked_at, :datetime unless column_exists?(TABLE, :tracked_at)
+    LlmCostTracker::Ledger::Schema::Adapter.drop_invalid_index(connection, connection.index_name(TABLE, INDEX))
     add_index TABLE, INDEX, **add_index_options
     remove_index TABLE, column: REPLACED_INDEX, **remove_index_options
   end

data/lib/llm_cost_tracker/ingestion/batch.rb

@@ -12,7 +12,7 @@ module LlmCostTracker
       TRANSIENT_PERSIST_ERRORS = [
         ActiveRecord::Deadlocked,
         ActiveRecord::LockWaitTimeout,
-        ActiveRecord::StatementTimeout,
+        ActiveRecord::QueryAborted,
         ActiveRecord::ConnectionNotEstablished,
         LlmCostTracker::TransactionAbortedError
       ].freeze
@@ -29,20 +29,16 @@ module LlmCostTracker
         persist_batch(valid_rows, events) if events.any?
         rows.size
       rescue StandardError => e
-        rows_to_mark = valid_rows&.any? ? valid_rows : rows
-        if rows_to_mark&.any?
-          transient = valid_rows&.any? && TRANSIENT_PERSIST_ERRORS.any? { |klass| e.is_a?(klass) }
-          mark_failed_with_message(rows_to_mark, error_message_for(e), decrement_attempts: transient)
-        end
+        release_after_failure(rows, valid_rows, e)
         raise
       end
 
       def pending?
-        Ingestion::InboxEntry.pending.exists?
+        Ingestion::InboxEntry.uncached { Ingestion::InboxEntry.pending.exists? }
       end
 
       def claimable?
-        claimable_scope(Time.now.utc - LOCK_TIMEOUT_SECONDS).exists?
+        Ingestion::InboxEntry.uncached { claimable_scope(Time.now.utc - LOCK_TIMEOUT_SECONDS).exists? }
       end
 
       def mark_failed_with_message(rows, message, decrement_attempts: false)
@@ -67,6 +63,10 @@ module LlmCostTracker
         nil
       end
 
+      private
+
+      attr_reader :identity
+
       def error_message_for(error)
         "#{error.class}: #{Redaction.text(error.message)}".byteslice(0, 1_000).scrub("")
       end
@@ -83,10 +83,6 @@ module LlmCostTracker
         )
       end
 
-      private
-
-      attr_reader :identity
-
       def id_sample(rows)
         sample = rows.first(10).map(&:id).join(", ")
         rows.size > 10 ? "#{sample}..." : sample
@@ -121,6 +117,18 @@ module LlmCostTracker
         [valid_rows, events]
       end
 
+      def release_after_failure(rows, valid_rows, error)
+        if valid_rows&.any?
+          mark_failed_with_message(valid_rows, error_message_for(error), decrement_attempts: transient?(error))
+        elsif rows&.any?
+          mark_failed_with_message(rows, error_message_for(error))
+        end
+      end
+
+      def transient?(error)
+        TRANSIENT_PERSIST_ERRORS.any? { |klass| error.is_a?(klass) }
+      end
+
       def persist_batch(rows, events)
         landed = []
         failed = Hash.new { |hash, message| hash[message] = [] }
@@ -128,14 +136,20 @@ module LlmCostTracker
       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
+        rows.zip(events) { |row, event| persist_alone(row, event, landed, failed) }
       ensure
+        settle(landed, failed)
+      end
+
+      def persist_alone(row, event, landed, failed)
+        landed.concat(persist([row], [event]))
+      rescue *TRANSIENT_PERSIST_ERRORS
+        raise
+      rescue StandardError => e
+        failed[error_message_for(e)] << row
+      end
+
+      def settle(landed, failed)
         failed.each do |message, failed_rows|
           LlmCostTracker::Logging.warn(
             "Ingestion::Batch: #{failed_rows.size} inbox row(s) could not be stored " \

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)
@@ -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
           {
@@ -73,7 +82,7 @@ module LlmCostTracker
         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,9 +25,7 @@ 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?
@@ -51,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]
@@ -72,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,
@@ -81,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)
@@ -108,20 +89,24 @@ 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)

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