Class: AgentHarness::TokenUsageTracker

Inherits:
Object
  • Object
show all
Defined in:
lib/agent_harness/token_usage_tracker.rb

Overview

Fallback quota tracker for providers that do not expose a public quota API.

When Providers::Base#check_quota returns a QuotaStatus with available: false (the default for providers like Aider or Kilocode that route through opaque upstreams), the orchestration layer falls back to TokenUsageTracker. The tracker records the same per-request token usage surfaced by Providers::Base#track_tokens, then derives an estimated QuotaStatus by subtracting usage from a caller-configured billing period limit.

The tracker holds no connection to the provider and never makes an HTTP request; it is purely an in-memory aggregate that mirrors what TokenTracker records but with the additional concept of a billing-period limit per provider.

Examples:

Configure a daily token limit and record usage

tracker = AgentHarness::TokenUsageTracker.new
tracker.set_quota(provider: :aider, limit: 2_000_000, reset_at: Time.utc(2026, 8, 1), unit: :tokens)
tracker.record(provider: :aider, model: "gpt-4o", input_tokens: 500, output_tokens: 100)
tracker.estimated_usage(provider: :aider)
# => #<AgentHarness::QuotaStatus available=true remaining=1999400 limit=2000000 unit=:tokens>

Defined Under Namespace

Classes: QuotaConfig, UsageEvent

Instance Method Summary collapse

Constructor Details

#initialize ⇒ TokenUsageTracker

Returns a new instance of TokenUsageTracker.



36
37
38
39
40
41
 
# File 'lib/agent_harness/token_usage_tracker.rb', line 36

def initialize
  @events = []
  @quota_configs = Hash.new { |hash, key| hash[key] = QuotaConfig.new }
  @callbacks = []
  @mutex = Mutex.new
end

Instance Method Details

#clear!(provider: nil) ⇒ void

This method returns an undefined value.

Drop all recorded events for a provider. Useful when a billing period rolls over or in tests.

Parameters:

  • provider (Symbol, String, nil) (defaults to: nil) —

    provider name; clears everything when nil



131
132
133
134
135
136
137
138
139
140
 
# File 'lib/agent_harness/token_usage_tracker.rb', line 131

def clear!(provider: nil)
  @mutex.synchronize do
    if provider.nil?
      @events.clear
    else
      provider_key = normalize_provider(provider)
      @events.reject! { |event| event.provider == provider_key }
    end
  end
end

#estimated_usage(provider:, since: nil) ⇒ AgentHarness::QuotaStatus

Return an estimated QuotaStatus for the given provider.

Returns a QuotaStatus with available: true as long as the caller has configured a quota for the provider (via #set_quota). The remaining value is the configured limit minus recorded usage since since: (or every recorded event when since: is omitted — callers are responsible for clearing events with #clear! when a billing period rolls over).

Parameters:

  • provider (Symbol, String) —

    provider name

  • since (Time, nil) (defaults to: nil) —

    cutoff for usage events

Returns:



101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
 
# File 'lib/agent_harness/token_usage_tracker.rb', line 101

def estimated_usage(provider:, since: nil)
  provider_key = normalize_provider(provider)
  config = @mutex.synchronize { @quota_configs[provider_key] }
  return QuotaStatus.unavailable unless config&.configured?

  used = usage_total(provider_key, since:)
  remaining = config.limit.nil? ? nil : (config.limit - used)

  QuotaStatus.new(
    available: true,
    remaining: remaining,
    limit: config.limit,
    reset_at: config.reset_at,
    unit: config.unit
  )
end

#event_count(provider = nil, since: nil) ⇒ Integer

Number of usage events recorded for a provider since an optional cutoff.

Parameters:

  • provider (Symbol, String) (defaults to: nil) —

    provider name

  • since (Time, nil) (defaults to: nil) —

    cutoff

Returns:

  • (Integer) 


158
159
160
161
162
163
164
165
 
# File 'lib/agent_harness/token_usage_tracker.rb', line 158

def event_count(provider = nil, since: nil)
  @mutex.synchronize do
    events = @events.dup
    events = events.select { |event| event.provider == normalize_provider(provider) } if provider
    events = events.select { |event| since.nil? || event.timestamp >= since }
    events.size
  end
end

#on_usage_recorded {|UsageEvent| ... } ⇒ void

This method returns an undefined value.

Register a callback invoked whenever usage is recorded.

Yields:



122
123
124
 
# File 'lib/agent_harness/token_usage_tracker.rb', line 122

def on_usage_recorded(&block)
  @callbacks << block
end

#record(provider:, model: nil, input_tokens: 0, output_tokens: 0, total_tokens: nil, request_id: nil) ⇒ UsageEvent

Record token usage for a provider.

Mirrors the AgentHarness::TokenTracker#record signature so the tracker can be wired into the same track_tokens hook on Providers::Base.

Parameters:

  • provider (Symbol, String) —

    provider name

  • model (String, nil) (defaults to: nil) —

    model identifier

  • input_tokens (Integer) (defaults to: 0) —

    input tokens used

  • output_tokens (Integer) (defaults to: 0) —

    output tokens used

  • total_tokens (Integer, nil) (defaults to: nil) —

    total tokens (calculated if nil)

  • request_id (String, nil) (defaults to: nil) —

    unique request ID (generated if nil)

Returns:



73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
 
# File 'lib/agent_harness/token_usage_tracker.rb', line 73

def record(provider:, model: nil, input_tokens: 0, output_tokens: 0, total_tokens: nil, request_id: nil)
  total = total_tokens || (input_tokens + output_tokens)
  event = UsageEvent.new(
    provider: normalize_provider(provider),
    model: model,
    input_tokens: input_tokens.to_i,
    output_tokens: output_tokens.to_i,
    total_tokens: total.to_i,
    timestamp: Time.now.utc,
    request_id: request_id || SecureRandom.uuid
  )

  @mutex.synchronize { @events << event }
  notify_callbacks(event)
  event
end

#set_quota(provider:, limit:, reset_at:, unit:) ⇒ void

This method returns an undefined value.

Configure the billing-period quota for a provider.

The orchestration layer uses this to seed tracker state from Paid's provider config before any usage is recorded.

Parameters:

  • provider (Symbol, String) —

    provider name

  • limit (Integer, Float, nil) —

    total quota for the period; pass nil to mark the provider as untracked

  • reset_at (Time, nil) —

    when the period resets

  • unit (Symbol) —

    quota unit (:tokens, :requests, :credits, :cost_cents)



54
55
56
57
58
59
 
# File 'lib/agent_harness/token_usage_tracker.rb', line 54

def set_quota(provider:, limit:, reset_at:, unit:)
  provider_key = normalize_provider(provider)
  @mutex.synchronize do
    @quota_configs[provider_key] = QuotaConfig.new(limit: limit, reset_at: reset_at, unit: normalize_unit(unit))
  end
end

#usage_total(provider, since: nil) ⇒ Integer

Total recorded token usage for a provider since an optional cutoff.

Parameters:

  • provider (Symbol, String) —

    provider name

  • since (Time, nil) (defaults to: nil) —

    cutoff

Returns:

  • (Integer) 


147
148
149
150
151
 
# File 'lib/agent_harness/token_usage_tracker.rb', line 147

def usage_total(provider, since: nil)
  provider_key = normalize_provider(provider)
  events = filtered_events(provider_key, since:)
  events.sum(&:total_tokens)
end