llm_cost_tracker 0.14.2 → 0.15.0

data/app/services/llm_cost_tracker/dashboard/spend_chart.rb

@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Dashboard
+    class SpendChart
+      WIDTH = 1180
+      PADDING = { top: 16, right: 16, bottom: 28, left: 56 }.freeze
+      MIN_MAX_COST = 0.0001
+
+      attr_reader :points, :height, :y_ticks, :max_cost, :coords, :comparison_coords
+
+      def initialize(points, comparison_points: nil, height: 180, y_ticks: 3)
+        @points = points
+        @height = height
+        @y_ticks = y_ticks
+        costs = points.map { |point| point[:cost].to_f } + Array(comparison_points).map { |point| point[:cost].to_f }
+        @max_cost = [costs.max.to_f, MIN_MAX_COST].max
+        @coords = coords_for(points)
+        @comparison_coords = coords_for(comparison_points) if comparison_points.present?
+      end
+
+      def peak_index
+        @peak_index ||= points.each_with_index.max_by { |point, _| point[:cost].to_f }.last
+      end
+
+      def left = PADDING[:left]
+      def right = left + plot_width
+      def baseline = PADDING[:top] + plot_height
+      def tick_y(index) = PADDING[:top] + (plot_height * index.to_f / y_ticks)
+      def tick_value(index) = max_cost * (y_ticks - index).to_f / y_ticks
+
+      def label_indexes
+        count = points.size
+        count <= 2 ? (0...count).to_a : [0, count / 2, count - 1].uniq
+      end
+
+      private
+
+      def plot_width = WIDTH - PADDING[:left] - PADDING[:right]
+      def plot_height = height - PADDING[:top] - PADDING[:bottom]
+
+      def coords_for(series)
+        step = series.size > 1 ? plot_width.to_f / (series.size - 1) : 0.0
+        series.each_with_index.map do |point, index|
+          [left + (index * step), baseline - ((point[:cost].to_f / max_cost) * plot_height)]
+        end
+      end
+    end
+  end
+end

data/app/services/llm_cost_tracker/dashboard/tag_breakdown.rb

@@ -20,8 +20,7 @@ module LlmCostTracker
         @key = LlmCostTracker::Tags::Key.validate!(key, error_class: LlmCostTracker::InvalidFilterError)
         limit = limit.to_i
         @limit = limit.positive? ? [limit, DEFAULT_LIMIT].min : DEFAULT_LIMIT
-        @sort = DEFAULT_DIRECTIONS.key?(sort.to_s) ? sort.to_s : "cost"
-        @direction = Sort::DIRECTIONS.include?(direction.to_s) ? direction.to_s : DEFAULT_DIRECTIONS[@sort]
+        @order = Sort.resolve(sort, direction, natural_directions: DEFAULT_DIRECTIONS, fallback: "cost")
       end
 
       def rows
@@ -34,7 +33,7 @@ module LlmCostTracker
               calls: calls,
               total_cost: row.total_cost,
               average_cost_per_call: row.average_cost_per_call,
-              share_percent: total.positive? ? (calls.to_f / total) * 100.0 : 0.0
+              share_percent: Percent.of(calls, total)
             )
           end
         end
@@ -54,7 +53,7 @@ module LlmCostTracker
 
       private
 
-      attr_reader :scope, :key, :sort, :direction
+      attr_reader :scope, :key, :order
 
       def summary_counts
         @summary_counts ||= scope.klass.find_by_sql(summary_sql).first
@@ -76,8 +75,8 @@ module LlmCostTracker
       end
 
       def order_clause
-        dir = direction.upcase
-        case sort
+        dir = order.direction.upcase
+        case order.column
         when "value"    then "#{tag_value_column} #{dir}"
         when "calls"    then "COUNT(*) #{dir}, total_cost DESC"
         when "avg_cost" then "average_cost_per_call #{dir}, total_cost DESC"
@@ -108,11 +107,11 @@ module LlmCostTracker
       end
 
       def quote_column(name)
-        scope.connection.quote_column_name(name)
+        scope.lease_connection.quote_column_name(name)
       end
 
       def quoted_key
-        scope.connection.quote(key)
+        scope.lease_connection.quote(key)
       end
     end
   end

data/app/services/llm_cost_tracker/dashboard/tag_key_explorer.rb

@@ -13,16 +13,13 @@ module LlmCostTracker
 
       def initialize(scope:, limit:)
         @scope = scope
-        @connection = LlmCostTracker::Call.connection
+        @connection = LlmCostTracker::Call.lease_connection
         limit = limit.to_i
         @limit = limit.positive? ? [limit, DEFAULT_LIMIT].min : DEFAULT_LIMIT
       end
 
       def rows
         scope.klass.find_by_sql(build_sql)
-      rescue ActiveRecord::StatementInvalid => e
-        LlmCostTracker::Logging.warn("Tag key discovery failed (#{connection.adapter_name}): #{e.class}: #{e.message}")
-        []
       end
 
       private

data/app/services/llm_cost_tracker/dashboard/top_models.rb

@@ -42,21 +42,20 @@ module LlmCostTracker
       def initialize(scope:, limit:, sort: DEFAULT_SORT, direction: nil)
         @scope = scope
         @limit = limit
-        @sort = DEFAULT_DIRECTIONS.key?(sort.to_s) ? sort.to_s : DEFAULT_SORT
-        @direction = Sort::DIRECTIONS.include?(direction.to_s) ? direction.to_s : DEFAULT_DIRECTIONS[@sort]
+        @order = Sort.resolve(sort, direction, natural_directions: DEFAULT_DIRECTIONS, fallback: DEFAULT_SORT)
       end
 
       def rows
         scope
           .group(:provider, :model)
           .select(selects)
-          .order(*ORDER_NODES.fetch([sort, direction]))
+          .order(*ORDER_NODES.fetch([order.column, order.direction]))
           .then { |r| limit ? r.limit(limit) : r }
       end
 
       private
 
-      attr_reader :scope, :limit, :sort, :direction
+      attr_reader :scope, :limit, :order
 
       def selects
         columns = [

data/app/views/layouts/llm_cost_tracker/application.html.erb

@@ -54,28 +54,28 @@
         <aside class="lct-sidebar" aria-label="Dashboard">
           <div class="lct-sidebar-inner">
           <div class="lct-sidebar-section">Insights</div>
-          <%= link_to root_path, class: ("lct-sidebar-link lct-active" if dashboard_section == :overview) || "lct-sidebar-link", aria: (dashboard_section == :overview ? { current: "page" } : {}) do %>
+          <%= link_to root_path, sidebar_link_options(:overview) do %>
             <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><rect x="3" y="3" width="7" height="9"/><rect x="14" y="3" width="7" height="5"/><rect x="14" y="12" width="7" height="9"/><rect x="3" y="16" width="7" height="5"/></svg>
             Overview
           <% end %>
-          <%= link_to models_path, class: ("lct-sidebar-link lct-active" if dashboard_section == :models) || "lct-sidebar-link", aria: (dashboard_section == :models ? { current: "page" } : {}) do %>
+          <%= link_to models_path, sidebar_link_options(:models) do %>
             <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="M21 8a2 2 0 0 0-1-1.73l-7-4a2 2 0 0 0-2 0l-7 4A2 2 0 0 0 3 8v8a2 2 0 0 0 1 1.73l7 4a2 2 0 0 0 2 0l7-4A2 2 0 0 0 21 16Z"/><path d="m3.3 7 8.7 5 8.7-5"/><path d="M12 22V12"/></svg>
             Models
           <% end %>
-          <%= link_to calls_path, class: ("lct-sidebar-link lct-active" if dashboard_section == :calls) || "lct-sidebar-link", aria: (dashboard_section == :calls ? { current: "page" } : {}) do %>
+          <%= link_to calls_path, sidebar_link_options(:calls) do %>
             <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><path d="M21 15a2 2 0 0 1-2 2H7l-4 4V5a2 2 0 0 1 2-2h14a2 2 0 0 1 2 2z"/></svg>
             Calls
           <% end %>
-          <%= link_to tags_path, class: ("lct-sidebar-link lct-active" if dashboard_section == :tags) || "lct-sidebar-link", aria: (dashboard_section == :tags ? { current: "page" } : {}) do %>
+          <%= link_to tags_path, sidebar_link_options(:tags) do %>
             <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><path d="M20.59 13.41l-7.17 7.17a2 2 0 0 1-2.83 0L2 12V2h10l8.59 8.59a2 2 0 0 1 0 2.82z"/><line x1="7" y1="7" x2="7.01" y2="7"/></svg>
             Tags
           <% end %>
-          <%= link_to data_quality_path, class: ("lct-sidebar-link lct-active" if dashboard_section == :data_quality) || "lct-sidebar-link", aria: (dashboard_section == :data_quality ? { current: "page" } : {}) do %>
+          <%= link_to data_quality_path, sidebar_link_options(:data_quality) do %>
             <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><path d="M22 11.08V12a10 10 0 1 1-5.93-9.14"/><polyline points="22 4 12 14.01 9 11.01"/></svg>
             Data Quality
           <% end %>
           <div class="lct-sidebar-section">Reference</div>
-          <%= link_to pricing_path, class: ("lct-sidebar-link lct-active" if dashboard_section == :pricing) || "lct-sidebar-link", aria: (dashboard_section == :pricing ? { current: "page" } : {}) do %>
+          <%= link_to pricing_path, sidebar_link_options(:pricing) do %>
             <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><line x1="12" y1="1" x2="12" y2="23"/><path d="M17 5H9.5a3.5 3.5 0 0 0 0 7h5a3.5 3.5 0 0 1 0 7H6"/></svg>
             Pricing
           <% end %>

data/app/views/llm_cost_tracker/calls/index.html.erb

@@ -19,7 +19,8 @@
   <% end %>
 
   <span class="lct-filter-row-meta"><%= number_with_delimiter(@calls_count) %> call<%= "s" unless @calls_count == 1 %> · <%= money(@calls_total_cost) %></span>
-  <%= link_to "Export CSV", calls_path(current_query(format: :csv)), class: "lct-page-link", title: "Capped at #{number_with_delimiter(LlmCostTracker::CallsController::CSV_EXPORT_LIMIT)} rows per request" %>
+  <% export_limit = LlmCostTracker::CallsController::CSV_EXPORT_LIMIT %>
+  <%= link_to(@calls_count > export_limit ? "Export CSV (first #{number_with_delimiter(export_limit)})" : "Export CSV", calls_path(current_query(format: :csv)), class: "lct-page-link", title: "Capped at #{number_with_delimiter(export_limit)} rows per request") %>
 </div>
 
 <% if @calls_count.zero? %>

data/app/views/llm_cost_tracker/calls/show.html.erb

@@ -1,15 +1,4 @@
-<% priced_components = token_usage_stack_components %>
-<% line_item_costs_by_component = call_line_item_costs_by_component(@call) %>
-<% token_segments = priced_components.map do |component|
-  token_key = component.fetch(:token_key)
-  value = @call[token_key]
-  { label: component.fetch(:label), value: value, formatted_value: number_with_delimiter(value), css_class: component.fetch(:css_class) }
-end %>
-<% cost_segments = priced_components.map do |component|
-  value = line_item_costs_by_component[component.fetch(:price_key)]
-  { label: component.fetch(:label), value: value, formatted_value: optional_money(value), css_class: component.fetch(:css_class) }
-end %>
-<% service_line_items = @call.line_items.reject { |li| li.unit == "token" }.sort_by(&:position) %>
+<% service_line_items = call_service_line_items(@call) %>
 
 <p class="lct-breadcrumb-back"><%= link_to "← All calls", calls_path %></p>
 <h2 class="lct-page-title"><code class="lct-code-id"><%= @call.model %></code> <span class="lct-page-title-meta">#<%= @call.id %></span></h2>
@@ -42,7 +31,7 @@ end %>
   <section class="lct-panel">
     <div class="lct-panel-head"><h2 class="lct-panel-title">Token mix</h2></div>
     <div class="lct-panel-body">
-      <%= render "llm_cost_tracker/shared/metric_stack", segments: token_segments, empty_message: "No token data for this call." %>
+      <%= render "llm_cost_tracker/shared/metric_stack", segments: call_token_segments(@call), empty_message: "No token data for this call." %>
     </div>
   </section>
 
@@ -52,7 +41,7 @@ end %>
       <% if @call.total_cost.nil? %>
         <p class="lct-stack-empty">Pricing not available for this call.</p>
       <% else %>
-        <%= render "llm_cost_tracker/shared/metric_stack", segments: cost_segments, empty_message: "No cost breakdown for this call." %>
+        <%= render "llm_cost_tracker/shared/metric_stack", segments: call_cost_segments(@call), empty_message: "No cost breakdown for this call." %>
       <% end %>
     </div>
   </section>

data/app/views/llm_cost_tracker/errors/query.html.erb

@@ -0,0 +1,6 @@
+<section class="lct-panel lct-empty">
+  <h2 class="lct-state-title">Query failed</h2>
+  <p class="lct-state-copy">
+    A dashboard query failed, for example on a statement timeout or a lock. The error is in your application log.
+  </p>
+</section>

data/app/views/llm_cost_tracker/shared/_filter_pill_stream.html.erb

@@ -1,5 +1,5 @@
 <% active = params[:stream].presence %>
-<% display_value = case active when "yes" then "Streaming" when "no" then "Non-streaming" else "All" end %>
+<% display_value = LlmCostTracker::Dashboard::Filter::STREAM_FILTER_LABELS.fetch(active, "All") %>
 <details class="lct-filter-pop" name="lct-filter">
   <summary class="lct-filter-pill <%= 'lct-active' if active %>">
     <svg width="12" height="12" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><polyline points="22 12 18 12 15 21 9 3 6 12 2 12"/></svg>

data/app/views/llm_cost_tracker/tags/show.html.erb

@@ -1,6 +1,6 @@
 <% form_url = tag_path(params[:key]) %>
 
-<% if @value.present? %>
+<% if @breakdown.nil? %>
   <p class="lct-breadcrumb-back"><%= link_to "← All values for #{params[:key]}", tag_path(params[:key], current_query.except(:tag_value)) %></p>
   <h2 class="lct-page-title">Tag: <code class="lct-code-id"><%= params[:key] %></code> = <code class="lct-code-id"><%= @value %></code></h2>
 

data/lib/llm_cost_tracker/budget/estimate.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+require "bigdecimal"
+
+require_relative "../pricing/estimator"
+
+module LlmCostTracker
+  module Budget
+    Estimate = Data.define(:total, :largest) do
+      def self.for(provider:, model:, request:)
+        costs = request_costs(provider, model, request)
+        total = costs.sum(BigDecimal("0"))
+        new(total: total, largest: costs.max || total)
+      end
+
+      def self.request_costs(provider, model, request)
+        batch = request["requests"] if request && !model
+        return [cost(provider, model, request)] unless batch.is_a?(Array)
+
+        batch.map do |entry|
+          params = entry["params"] if entry.is_a?(Hash)
+          params.is_a?(Hash) ? cost(provider, params["model"], params) : 0
+        end
+      end
+
+      def self.cost(provider, model, request)
+        return BigDecimal("0") unless provider && model && request
+
+        Pricing::Estimator.call(provider: provider, model: model, request: request) || BigDecimal("0")
+      end
+
+      private_class_method :request_costs, :cost
+    end
+  end
+end

data/lib/llm_cost_tracker/budget/global.rb

@@ -0,0 +1,69 @@
+# frozen_string_literal: true
+
+require "bigdecimal"
+
+require_relative "limit"
+require_relative "../ledger"
+
+module LlmCostTracker
+  module Budget
+    module Global
+      PERIODS = { monthly: :month, daily: :day }.freeze
+
+      class << self
+        def enforce!(estimate, time:)
+          budgets = LlmCostTracker.configuration.budgets
+          per_call = budgets.per_call
+          largest = estimate.largest
+          Limit.global(:per_call, per_call).block!(largest) if per_call && largest.positive? && largest >= per_call
+
+          windows = { monthly: budgets.monthly, daily: budgets.daily }
+          each_exceeded(windows, time: time, estimate: estimate.total) { |limit, total| limit.block!(total) }
+        end
+
+        def post_spend_errors(event, behavior_override)
+          budgets = LlmCostTracker.configuration.budgets
+          errors = [per_call_error(event, budgets.per_call, behavior_override)]
+          each_exceeded({ daily: budgets.daily, monthly: budgets.monthly }, time: event.tracked_at) do |limit, total|
+            errors << limit.handle_exceeded(total: total,
+                                            previous_total: total - event.total_cost,
+                                            last_event: event,
+                                            behavior_override: behavior_override)
+          end
+          errors
+        end
+
+        def notify_repriced(changes, time:)
+          budgets = LlmCostTracker.configuration.budgets
+          each_exceeded({ daily: budgets.daily, monthly: budgets.monthly }, time: time) do |limit, total|
+            limit.handle_exceeded(total: total,
+                                  previous_total: total - limit.amount_in_window(changes, time),
+                                  behavior_override: :notify)
+          end
+        end
+
+        private
+
+        def per_call_error(event, per_call, behavior_override)
+          total = event.total_cost
+          return unless per_call && total >= per_call
+
+          Limit.global(:per_call, per_call).handle_exceeded(
+            total: total, previous_total: nil, last_event: event, behavior_override: behavior_override
+          )
+        end
+
+        def each_exceeded(windows, time:, estimate: BigDecimal("0"))
+          windows = windows.compact
+          return if windows.empty?
+
+          totals = Ledger::Period::Totals.call(windows.keys.map { |type| PERIODS.fetch(type) }, time: time)
+          windows.each do |type, budget|
+            total = totals.fetch(PERIODS.fetch(type)) + estimate
+            yield Limit.global(type, budget), total if total >= budget
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/budget/limit.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+require "bigdecimal"
+
+module LlmCostTracker
+  module Budget
+    Limit = Data.define(:budget_type, :budget, :scope, :behavior, :on_exceeded) do
+      def self.global(budget_type, budget)
+        new(budget_type: budget_type, budget: budget, scope: nil, behavior: nil, on_exceeded: nil)
+      end
+
+      def over?(total)
+        budget_type == :calls ? total > budget : total >= budget
+      end
+
+      def charge(cost)
+        budget_type == :calls ? 1 : cost
+      end
+
+      def counts?(event)
+        budget_type == :calls || !event.total_cost.nil?
+      end
+
+      def amount_in_window(changes, time)
+        start = PerTag.window_start(budget_type, time)
+        changes.sum(BigDecimal("0")) { |change| start.nil? || change.tracked_at >= start ? change.total_cost : 0 }
+      end
+
+      def block!(total)
+        raise BudgetExceededError.new(
+          budget_type: budget_type, total: total, budget: budget, stage: :pre_send, scope: scope
+        )
+      end
+
+      def handle_exceeded(total:, previous_total:, last_event: nil, behavior_override: nil)
+        budgets = LlmCostTracker.configuration.budgets
+        raising = %i[raise block_requests].include?(behavior_override || behavior || budgets.exceeded_behavior)
+        callback = on_exceeded || budgets.on_exceeded
+        payload = {
+          budget_type: budget_type,
+          total: total,
+          budget: budget,
+          last_event: last_event,
+          stage: :post_spend,
+          scope: scope
+        }
+        callback.call(payload) if callback && (previous_total.nil? || !over?(previous_total))
+        BudgetExceededError.new(**payload) if raising
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/budget/per_tag.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require_relative "limit"
+require_relative "../timing"
 require_relative "../ledger/isolation"
 require_relative "../ledger/schema/adapter"
 require_relative "../ledger/tags/encoding"
@@ -14,7 +16,15 @@ module LlmCostTracker
       SLOW_READ_SECONDS = 0.1
       DEFAULT_BACKFILL_BATCH = 5_000
 
-      Rule = Data.define(:key, :value, :windows, :behavior, :on_exceeded)
+      Rule = Data.define(:key, :value, :windows, :behavior, :on_exceeded) do
+        def limit(window)
+          Limit.new(budget_type: window,
+                    budget: windows.fetch(window),
+                    scope: { key: key, value: value },
+                    behavior: behavior,
+                    on_exceeded: on_exceeded)
+        end
+      end
 
       class << self
         def configured
@@ -53,45 +63,42 @@ module LlmCostTracker
         end
 
         def rules_for_events(events)
-          events.select(&:total_cost).each_with_object({}) do |event, grouped|
+          events.each_with_object({}) do |event, grouped|
             rules_for(event.tags).each { |rule| (grouped[rule] ||= []) << event }
           end
         end
 
         def spend(key, value, window, time:)
-          spend_by_value(key, [value], window, window_start(window, time)).fetch(value, [0]).first.to_d
+          spend_by_value(key, [value], window, window_start(window, time)).fetch(value, [0]).first
         end
 
         def spend_by_value(key, values, window, bucket, upto = nil)
-          started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
-          range = window_range(window, bucket)
-          upto_sum = LlmCostTracker::CallTag.sanitize_sql_array(
-            ["SUM(CASE WHEN #{TIME_COLUMN} <= ? THEN #{COST_COLUMN} ELSE 0 END)", upto || range.end]
-          )
+          started_at = Timing.now_monotonic
+          aggregates = sum_columns(window, upto)
           totals = Ledger::Isolation.guard(LlmCostTracker::CallTag) do
-            LlmCostTracker::CallTag
-              .where(key: key, value: values, TIME_COLUMN => range)
-              .group(:value)
-              .pluck(:value, Arel.sql("SUM(#{COST_COLUMN})"), Arel.sql(upto_sum))
-              .to_h { |value, *sums| [value, sums.map(&:to_d)] }
+            pluck_by_exact_value(spend_rows(key, values, window, bucket), values, aggregates).to_h do |value, *sums|
+              [value.dup.force_encoding(Encoding::UTF_8), sums.map { |sum| window == :calls ? sum.to_i : sum.to_d }]
+            end
           end
-          warn_slow_read(key, window, Process.clock_gettime(Process::CLOCK_MONOTONIC) - started_at)
+          warn_slow_read(key, window, Timing.now_monotonic - started_at)
           totals
         end
 
         def window_start(window, time)
-          time.to_time.utc.public_send(WINDOW_STARTS.fetch(window))
+          WINDOW_STARTS[window]&.then { |start| time.to_time.utc.public_send(start) }
         end
 
         def backfill(batch_size: DEFAULT_BACKFILL_BATCH)
           return 0 unless columns?
 
           filled = 0
+          after = 0
           loop do
-            copied = copy_next_batch(batch_size)
-            break if copied.zero?
+            ids = unfilled_ids(after, batch_size)
+            break if ids.empty?
 
-            filled += copied
+            filled += LlmCostTracker::CallTag.lease_connection.update(copy_sql(ids))
+            after = ids.last
           end
           filled
         end
@@ -103,22 +110,42 @@ module LlmCostTracker
 
         private
 
+        def sum_columns(window, upto)
+          measure = window == :calls ? "1" : COST_COLUMN
+          total = "SUM(#{measure})"
+          upto_sum = "SUM(CASE WHEN #{TIME_COLUMN} <= ? THEN #{measure} ELSE 0 END)"
+          upto_total = upto ? LlmCostTracker::CallTag.sanitize_sql_array([upto_sum, upto]) : total
+          [Arel.sql(total), Arel.sql(upto_total)]
+        end
+
+        def spend_rows(key, values, window, bucket)
+          rows = LlmCostTracker::CallTag.where(key: key, value: values)
+          bucket ? rows.where(TIME_COLUMN => window_range(window, bucket)) : rows.where.not(TIME_COLUMN => nil)
+        end
+
+        def pluck_by_exact_value(rows, values, aggregates)
+          column = :value
+          if Ledger::Schema::Adapter.mysql?(LlmCostTracker::CallTag.lease_connection)
+            column = Arel.sql("CAST(value AS BINARY)")
+            rows = rows.where(column.in(values))
+          end
+          rows.group(column).pluck(column, *aggregates)
+        end
+
         def window_range(window, bucket)
           bucket...bucket.public_send(WINDOW_NEXTS.fetch(window))
         end
 
-        def copy_next_batch(batch_size)
-          ids = LlmCostTracker::CallTag.where(TIME_COLUMN => nil).limit(batch_size).pluck(:id)
-          return 0 if ids.empty?
-
-          LlmCostTracker::CallTag.connection.update(copy_sql(ids))
+        def unfilled_ids(after, batch_size)
+          LlmCostTracker::CallTag.where(TIME_COLUMN => nil).where(id: (after + 1)..)
+                                 .order(:id).limit(batch_size).pluck(:id)
         end
 
         def copy_sql(ids)
           tags = LlmCostTracker::CallTag.quoted_table_name
           calls = LlmCostTracker::Call.quoted_table_name
           list = ids.join(",")
-          return <<~SQL.squish unless Ledger::Schema::Adapter.postgresql?(LlmCostTracker::CallTag.connection)
+          return <<~SQL.squish unless Ledger::Schema::Adapter.postgresql?(LlmCostTracker::CallTag.lease_connection)
             UPDATE #{tags} t JOIN #{calls} c ON c.id = t.llm_cost_tracker_call_id
                SET t.#{COST_COLUMN} = c.#{COST_COLUMN}, t.#{TIME_COLUMN} = c.#{TIME_COLUMN}
              WHERE t.id IN (#{list})
@@ -146,10 +173,15 @@ module LlmCostTracker
           @slow_keys ||= Set.new
           return unless @slow_keys.add?(key)
 
+          reason = if WINDOW_STARTS.key?(window)
+                     "A tag with few distinct values covers most of the ledger, so its budget check cannot use " \
+                       "an index effectively. Budget high-cardinality tags such as a tenant or user id."
+                   else
+                     "A #{window} limit sums every call the value has recorded, so it suits short-lived values " \
+                       "such as a run."
+                   end
           Logging.warn(
-            "config.budgets.per_tag[#{key.inspect}] #{window} read took #{(seconds * 1000).round} ms. " \
-            "A tag with few distinct values covers most of the ledger, so its budget check cannot use " \
-            "an index effectively. Budget high-cardinality tags such as a tenant or user id."
+            "config.budgets.per_tag[#{key.inspect}] #{window} read took #{(seconds * 1000).round} ms. #{reason}"
           )
         end
 

data/lib/llm_cost_tracker/budget/tagged.rb

@@ -0,0 +1,74 @@
+# frozen_string_literal: true
+
+require_relative "per_tag"
+require_relative "../tags/context"
+
+module LlmCostTracker
+  module Budget
+    module Tagged
+      class << self
+        def enforce!(tags, estimated_cost, time:, blocking_only:)
+          PerTag.rules_for(tags || Tags::Context.tags, blocking_only: blocking_only).each do |rule|
+            rule.windows.each_key do |window|
+              limit = rule.limit(window)
+              total = PerTag.spend(rule.key, rule.value, window, time: time) + limit.charge(estimated_cost)
+              limit.block!(total) if limit.over?(total)
+            end
+          end
+        end
+
+        def post_spend_errors(events, behavior_override)
+          window_buckets(scored_rules(events, behavior_override)).flat_map do |(key, window, bucket), recorded_by_rule|
+            upto = recorded_by_rule.values.flatten.map(&:tracked_at).max unless Ingestion.async?
+            totals = PerTag.spend_by_value(key, recorded_by_rule.keys.map(&:value), window, bucket, upto)
+            recorded_by_rule.filter_map do |rule, recorded|
+              total, total_through_batch = totals.fetch(rule.value, [0, 0])
+              bucket_error(rule.limit(window), recorded, total, total_through_batch, behavior_override)
+            end
+          end
+        end
+
+        def notify_repriced(changes, time:)
+          PerTag.rules_for_events(changes).each do |rule, repriced|
+            rule.windows.except(:calls).each_key do |window|
+              limit = rule.limit(window)
+              amount = limit.amount_in_window(repriced, time)
+              total = PerTag.spend(rule.key, rule.value, window, time: time) if rule.on_exceeded && amount.positive?
+              next unless total && limit.over?(total)
+
+              limit.handle_exceeded(total: total, previous_total: total - amount, behavior_override: :notify)
+            end
+          end
+        end
+
+        private
+
+        def scored_rules(events, behavior_override)
+          by_rule = PerTag.rules_for_events(events)
+          behavior_override == :notify ? by_rule.reject { |rule, _| rule.on_exceeded.nil? } : by_rule
+        end
+
+        def bucket_error(limit, recorded, total, total_through_batch, behavior_override)
+          return unless limit.over?(total)
+
+          reference = limit.over?(total_through_batch) ? total_through_batch : total
+          limit.handle_exceeded(total: total,
+                                previous_total: reference - recorded.sum { |event| limit.charge(event.total_cost) },
+                                last_event: recorded.last,
+                                behavior_override: behavior_override)
+        end
+
+        def window_buckets(by_rule)
+          by_rule.each_with_object({}) do |(rule, events), grouped|
+            rule.windows.each_key do |window|
+              limit = rule.limit(window)
+              events.select { |event| limit.counts?(event) }
+                    .group_by { |event| PerTag.window_start(window, event.tracked_at) }
+                    .each { |bucket, bucket_events| (grouped[[rule.key, window, bucket]] ||= {})[rule] = bucket_events }
+            end
+          end
+        end
+      end
+    end
+  end
+end