Module budget 

Cost-bounded coordination primitives.

The kernel often needs to gate dispatch on a finite budget — rows parsed, LLM tokens spent, dollars burned, etc. This module exposes two domain-neutral traits and lock-free reference implementations:

• BudgetGuard — generic unit-cost reservations (rows, requests, queue slots).
• TokenBudget — LLM-token accounting with optimistic reservation and after-the-fact reconciliation against provider-reported usage.

Both traits are intentionally narrow so a single coordinator can compose multiple budget policies (e.g. a row guard and a token guard) without coupling to any particular backend.

Why two traits?

BudgetGuard is a symmetric reserve/release pair: callers know the cost up front and either commit or roll back. TokenBudget is optimistic: callers reserve an estimate and receive a TokenReservation, send the prompt, then pass that reservation to TokenBudget::record_usage with the provider’s reported totals to reconcile the over- or under-estimate for that specific call.

Reference implementations

AtomicBudget and AtomicTokenBudget are lock-free token-bucket counters built on AtomicU64::compare_exchange_weak. They are safe for high-contention dispatch loops and refill on demand via AtomicBudget::refill.

use rig_compose::budget::{AtomicBudget, BudgetGuard};

let budget = AtomicBudget::new(1_000);
if budget.try_reserve(250).await? {
    // dispatch happens here ...
    budget.release(250).await;
}

Structs

AtomicBudget

Atomic, lock-free token-bucket budget.

AtomicTokenBudget

Atomic counterpart of AtomicBudget that tracks LLM token spend.

TokenReservation

Reservation handle returned by TokenBudget::try_reserve_tokens.

Enums

BudgetError

Errors a budget implementation may surface.

Traits

BudgetGuard

Symmetric reserve/release budget guard.

TokenBudget

Soft-cap on cumulative LLM token spend.

Type Aliases

TokenRefund

Refund channel invoked when a TokenReservation is dropped without being passed to TokenBudget::record_usage.