feat(accounting): add executable monetary settlement API

This commit is contained in:
Yujong Lee 2026-09-28 15:09:55 -07:00
parent 7e383c9f6a
commit 0923e2947a
10 changed files with 1131 additions and 0 deletions

View file

@ -97,6 +97,12 @@ version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "03918c3dbd7701a85c6b9887732e2921175f26c350b4563841d0958c21d57e6d"
[[package]]
name = "arrayvec"
version = "0.7.8"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56"
[[package]]
name = "asn1-rs"
version = "0.7.2"
@ -3345,6 +3351,17 @@ checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53"
name = "litellm"
version = "0.0.1"
[[package]]
name = "litellm-accounting"
version = "0.1.0"
dependencies = [
"rstest",
"rust_decimal",
"rusty-money",
"thiserror 2.0.19",
"tokio",
]
[[package]]
name = "litellm-auth"
version = "0.1.0"
@ -5869,6 +5886,16 @@ dependencies = [
"sqlite-wasm-rs",
]
[[package]]
name = "rust_decimal"
version = "1.43.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7653272e75dcac41dc199fbea6f5797633994fafd339943c06c9af16bf29cd3a"
dependencies = [
"arrayvec",
"num-traits",
]
[[package]]
name = "rustc-hash"
version = "2.1.3"
@ -6057,6 +6084,16 @@ dependencies = [
"wait-timeout",
]
[[package]]
name = "rusty-money"
version = "0.5.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e9826983e8a9a2f301a3810bb98618948ae7e7f2ba073eaaa16ab19104b33f2c"
dependencies = [
"arrayvec",
"rust_decimal",
]
[[package]]
name = "ryu"
version = "1.0.23"

View file

@ -9,6 +9,7 @@ license = "MIT"
repository = "https://github.com/BerriAI/litellm"
[workspace.dependencies]
litellm-accounting = { path = "crates/accounting" }
litellm-config = { path = "crates/config" }
litellm-router = { path = "crates/router" }
litellm-tracing = { path = "crates/tracing" }
@ -79,6 +80,7 @@ reqwest = { version = "0.12", default-features = false, features = ["json", "mul
qdrant-client = { version = "1.19.0", default-features = false }
uuid = { version = "1", features = ["v4"] }
rstest = "0.26.1"
rusty-money = "0.5.1"
rstest_reuse = "0.7.0"
rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }
rustify = "=0.7.0"

View file

@ -0,0 +1,15 @@
[package]
name = "litellm-accounting"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
rusty-money.workspace = true
thiserror.workspace = true
[dev-dependencies]
rstest.workspace = true
rust_decimal = { version = "1.43.0", default-features = false }
tokio.workspace = true

View file

@ -0,0 +1,76 @@
# Accounting API
This first version implements a per-call settlement state machine with injected storage operations. It does not implement prices, admission policy, Redis operations, database writes, or Python callback dispatch. Production calls are not connected to it yet
## Ownership
Core supplies execution facts and reported usage. Accounting preserves those inputs and tracks settlement progress. The host drives asynchronous operations and owns cancellation cleanup. An accounting backend implements monetary budget reconciliation, spend recording, and release of remaining monetary reservations
The Python integration belongs at the `python-bridge` boundary. `host-python` provides runtime mechanics, and `callbacks-legacy-python` preserves callback delivery. No crate is renamed. Existing Python accounting remains active until a real adapter replaces its orchestration. That adapter must settle accounting before terminal callback fan-out, supply callbacks their accounting payloads, and preserve public entrypoints and ordering without giving callbacks ownership of settlement
## Rate limiting boundary
Rate limiting owns request and token windows, deployment quotas, and concurrency permits. It consumes execution facts and usage directly, independently of pricing or monetary settlement. It does not use accounting's backend, receipts, effects, or settlement status
Python currently has deployment rate checks in the router and key/user/concurrency limits in proxy hooks. `crates/router` is the Rust counterpart of `litellm.Router`, with deployment quota checks belonging near selection and retries. Axum's HTTP router mounts gateway routes and can inject request-level admission through middleware. SDK calls need an injected limiter without HTTP middleware. These implementations remain separate follow-up work
Prefer a library for the limiter algorithm. [governor](https://docs.rs/governor/latest/governor/) provides a transport-independent GCRA engine, already used for gateway UI login throttling. [axum-governor](https://docs.rs/axum-governor/latest/axum_governor/) adds HTTP middleware; [axum-limit](https://docs.rs/axum-limit/latest/axum_limit/) offers extractor integration and a Redis backend. Selection requires tests against LiteLLM's window, token reservation, usage reconciliation, and multi-node policies. None of these choices requires an accounting dependency
The host composes accounting and limiter cleanup around the same call lifecycle. It must attempt both even when either fails. Streaming reconciliation and permit release follow body completion, failure, or cancellation rather than response header delivery. The eventual lifecycle composition must preserve Python ordering while providing both services the shared facts they need
Coordination stores for budgets and rate limits must be injected independently of response-cache storage. Changing or disabling the response-cache backend must not replace either service's coordination dependency
## Inputs
`BudgetAdmission<R>` contains an optional, already-acquired monetary budget receipt. Receipt type `R` belongs to the backend and can identify a monetary reservation or retain its budget coordination context. Constructing a session does not enforce a budget or acquire a reservation. Rejected admission never creates an admitted session
`ReportedUsage<U>` distinguishes unknown usage from a reported value, including a reported zero. The consumer supplies its typed usage representation. Known cost does not imply known usage, and reported usage does not imply known pricing
`Usd` uses decimal-backed [rusty-money](https://docs.rs/rusty-money/latest/rusty_money/) amounts. `Charges::new` validates nonnegative USD components, and totals use the library’s checked arithmetic. Fractional cents remain intact without currency-minor-unit rounding. Callers supply decimal amounts rather than floating-point values. Signed reservation adjustments belong to the backend, not this charge type
`Charges` preserves provider and independently incurred service costs. Each component is known or unknown. An unknown component leaves the total unknown without erasing the known component
`Terminal<U>` records success, failure, or cancellation, reported usage, chargeable provider work, and charges. `ProviderWork::NotStarted` requires an established fact that this call incurred no provider generation work. It makes the provider charge zero while retaining usage and service charges. A response-cache hit can use this classification even when the cached response reports usage. Provider prompt-cache usage still represents started provider work
The inspected main branch has no shared provider/model/source/cache-hit execution-facts contract yet. Production wiring requires a core fact establishing whether this call started provider generation, retained reported usage on cache hits, and terminal facts for partial failure and cancellation. Reuse the cache work’s shared contract when it lands, without importing that branch or duplicating its provider/model/source/cache-key types here. Missing facts must not be classified as avoided provider work
Failure and cancellation never automatically zero charges. Pricing and admission policy must supply incurred costs, including partial-stream costs or an input-cost estimate where the established policy requires one. This API preserves those decisions rather than guessing charges from a terminal outcome
## Settlement
`finish` selects one immutable terminal value. A second delivery returns `AlreadyTerminal` and cannot replace the first value. `settle` requires a selected terminal and attempts pending effects in this order: monetary budget reconciliation, spend recording, remaining budget reservation release
Budget reconciliation and release run only when a monetary budget receipt exists. Spend recording runs even without a monetary reservation. This API has no response-cache dependency or callback registry
`Backend::apply` receives the terminal, admission receipts, and read-only progress of all effects. Its future is awaited inline and need not be `Send`, allowing a Python adapter to preserve the caller's task and context
| Backend result | Meaning | Session behavior |
| --- | --- | --- |
| `Accepted` | Responsibility transferred to a queue or another owner, not committed | Retain the acknowledgement and do not resubmit |
| `Committed` | The effect completed according to the backend's persistence contract | Retain completion and do not repeat it |
| `NotApplied(error)` | The effect made no changes and transferred no responsibility | Preserve the error; allow an explicit retry |
| `Indeterminate(error)` | The effect may have applied or transferred responsibility | Preserve uncertainty and forbid automatic retry |
Every effect records its result before the next effect starts. A failed effect does not skip later effects, including release. A backend must report partial application or a lost acknowledgement as indeterminate, never not-applied. Returning from an existing Python writer is insufficient evidence of commitment
Release must use the supplied progress to release remaining resources without undoing reconciled charges or independently refunding an uncertain adjustment. Release covers only monetary budget reservations. The backend must define how failed or uncertain monetary reservations are repaired; this crate does not invent that storage policy
`retry_not_applied` makes only the selected failed effect pending again. A backend must implement effects against stable admission identities and account for previously completed cleanup when explicitly retried. Accepted, committed, pending, interrupted, and indeterminate effects cannot be reset by this method
`SettlementStatus::Accepted` distinguishes queue acceptance from `Committed`. `Unpriced` means effects have been acknowledged but at least one cost remains unknown. `NeedsAttention` means an effect failed or its completion is uncertain. Per-effect progress remains available in every case
## Cancellation and recovery
The state becomes `InFlight` before invoking a backend. Dropping the settlement future leaves that effect in flight because its external result is unknown. A later `settle` attempts only the remaining pending effects, including release, and does not replay the interrupted operation
The host must retain the session and resume remaining cleanup through its cancellation-safe finalization path. Dropping the entire session performs no asynchronous cleanup. Cancellation during release itself leaves release uncertain and requires backend-specific recovery
This is local protection within one retained session. There is no persistent settlement key, remote deduplication, crash recovery, queue acknowledgement reconciliation, or automatic resolution of unknown prices. A new session can repeat the same remote operations. Persistent idempotency and recovery require a later backend implementation
## Validation
Public API tests inject backend outcomes at every effect, drop futures at every suspension boundary, and verify completed work is not replayed. They cover optional monetary admission receipts, unknown and zero values, partial usage on failure and cancellation, avoided provider work with service charges, overflow, queue acceptance, and release progress
The next Python adapter must consume shared execution facts and separate existing built-in settlement from CustomLogger delivery. Native gateway accounting later supplies pricing, monetary budget policy, storage operations, and persistent idempotency behind the same contract. Both paths must avoid duplicate cost calculation and spend updates
Python callback compatibility, real pricing, coordination-store behavior, and native gateway accounting require production adapters and their own integration tests. The API tests do not claim those paths are migrated or validated

View file

@ -0,0 +1,81 @@
use rusty_money::{Money, iso};
use crate::Error;
pub type Usd = Money<'static, iso::Currency>;
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum Cost {
Unknown,
Known(Usd),
}
impl Cost {
fn validate(self) -> Result<(), Error> {
let Self::Known(money) = self else {
return Ok(());
};
if money.currency() != iso::USD {
return Err(Error::InvalidCurrency {
currency: money.currency().iso_alpha_code,
});
}
if money.is_negative() {
return Err(Error::NegativeCost);
}
Ok(())
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ProviderWork {
NotStarted,
Started,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum ReportedUsage<U> {
Unknown,
Known(U),
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Charges {
provider: Cost,
services: Cost,
}
impl Charges {
pub fn new(provider: Cost, services: Cost) -> Result<Self, Error> {
provider.validate()?;
services.validate()?;
Ok(Self { provider, services })
}
pub fn provider(self) -> Cost {
self.provider
}
pub fn services(self) -> Cost {
self.services
}
pub fn total(self) -> Result<Cost, Error> {
match (self.provider, self.services) {
(Cost::Known(provider), Cost::Known(services)) => {
provider.add(services).map(Cost::Known).map_err(Error::from)
}
_ => Ok(Cost::Unknown),
}
}
pub(crate) fn for_work(self, work: ProviderWork) -> Self {
match work {
ProviderWork::Started => self,
ProviderWork::NotStarted => Self {
provider: Cost::Known(Money::from_major(0, iso::USD)),
services: self.services,
},
}
}
}

View file

@ -0,0 +1,17 @@
use crate::Effect;
#[derive(Debug, PartialEq, thiserror::Error)]
pub enum Error {
#[error("accounting requires USD charges, received {currency}")]
InvalidCurrency { currency: &'static str },
#[error("USD charges must be nonnegative")]
NegativeCost,
#[error(transparent)]
Money(#[from] rusty_money::MoneyError),
#[error("the call already has a terminal outcome")]
AlreadyTerminal,
#[error("the call has no terminal outcome")]
NotTerminal,
#[error("{effect:?} can only be retried after a confirmed not-applied result")]
UnsafeRetry { effect: Effect },
}

View file

@ -0,0 +1,10 @@
mod charges;
mod error;
mod session;
pub use charges::{Charges, Cost, ProviderWork, ReportedUsage, Usd};
pub use error::Error;
pub use session::{
ApplyResult, Backend, BudgetAdmission, Effect, EffectState, Outcome, Progress, Session,
Settlement, SettlementStatus, Terminal,
};

View file

@ -0,0 +1,267 @@
use std::future::Future;
use crate::{Charges, Cost, Error, ProviderWork, ReportedUsage};
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct BudgetAdmission<R> {
budget: Option<R>,
}
impl<R> BudgetAdmission<R> {
pub fn new(budget: Option<R>) -> Self {
Self { budget }
}
pub fn budget(&self) -> Option<&R> {
self.budget.as_ref()
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Outcome {
Succeeded,
Failed,
Cancelled,
}
#[derive(Clone, Debug, PartialEq)]
pub struct Terminal<U> {
outcome: Outcome,
work: ProviderWork,
usage: ReportedUsage<U>,
charges: Charges,
}
impl<U> Terminal<U> {
pub fn new(
outcome: Outcome,
work: ProviderWork,
usage: ReportedUsage<U>,
charges: Charges,
) -> Result<Self, Error> {
let charges = charges.for_work(work);
charges.total()?;
Ok(Self {
outcome,
work,
usage,
charges,
})
}
pub fn outcome(&self) -> Outcome {
self.outcome
}
pub fn work(&self) -> ProviderWork {
self.work
}
pub fn usage(&self) -> &ReportedUsage<U> {
&self.usage
}
pub fn charges(&self) -> Charges {
self.charges
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Effect {
ReconcileBudget,
RecordSpend,
ReleaseBudgetReservation,
}
impl Effect {
const ORDER: [Self; 3] = [
Self::ReconcileBudget,
Self::RecordSpend,
Self::ReleaseBudgetReservation,
];
fn index(self) -> usize {
match self {
Self::ReconcileBudget => 0,
Self::RecordSpend => 1,
Self::ReleaseBudgetReservation => 2,
}
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum ApplyResult<E> {
Accepted,
Committed,
NotApplied(E),
Indeterminate(E),
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum EffectState<E> {
NotRequired,
Pending,
InFlight,
Accepted,
Committed,
NotApplied(E),
Indeterminate(E),
}
impl<E> From<ApplyResult<E>> for EffectState<E> {
fn from(result: ApplyResult<E>) -> Self {
match result {
ApplyResult::Accepted => Self::Accepted,
ApplyResult::Committed => Self::Committed,
ApplyResult::NotApplied(error) => Self::NotApplied(error),
ApplyResult::Indeterminate(error) => Self::Indeterminate(error),
}
}
}
pub struct Progress<E> {
effects: [EffectState<E>; 3],
}
impl<E> Progress<E> {
pub fn effect(&self, effect: Effect) -> &EffectState<E> {
&self.effects[effect.index()]
}
}
pub struct Settlement<'a, R, U, E> {
pub admission: &'a BudgetAdmission<R>,
pub terminal: &'a Terminal<U>,
pub progress: &'a Progress<E>,
}
pub trait Backend<R, U> {
type Error;
fn apply(
&mut self,
effect: Effect,
settlement: Settlement<'_, R, U, Self::Error>,
) -> impl Future<Output = ApplyResult<Self::Error>>;
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum SettlementStatus {
Open,
Pending,
NeedsAttention,
Unpriced,
Accepted,
Committed,
}
pub struct Session<R, U, E> {
admission: BudgetAdmission<R>,
terminal: Option<Terminal<U>>,
progress: Progress<E>,
}
impl<R, U, E> Session<R, U, E> {
pub fn new(admission: BudgetAdmission<R>) -> Self {
let effects = Effect::ORDER.map(|effect| {
let required = match effect {
Effect::RecordSpend => true,
Effect::ReconcileBudget => admission.budget.is_some(),
Effect::ReleaseBudgetReservation => admission.budget.is_some(),
};
if required {
EffectState::Pending
} else {
EffectState::NotRequired
}
});
Self {
admission,
terminal: None,
progress: Progress { effects },
}
}
pub fn finish(&mut self, terminal: Terminal<U>) -> Result<(), Error> {
if self.terminal.is_some() {
return Err(Error::AlreadyTerminal);
}
self.terminal = Some(terminal);
Ok(())
}
pub fn terminal(&self) -> Option<&Terminal<U>> {
self.terminal.as_ref()
}
pub fn progress(&self) -> &Progress<E> {
&self.progress
}
pub fn status(&self) -> SettlementStatus {
let Some(terminal) = &self.terminal else {
return SettlementStatus::Open;
};
if self.progress.effects.iter().any(|state| {
matches!(
state,
EffectState::InFlight | EffectState::NotApplied(_) | EffectState::Indeterminate(_)
)
}) {
return SettlementStatus::NeedsAttention;
}
if self
.progress
.effects
.iter()
.any(|state| matches!(state, EffectState::Pending))
{
return SettlementStatus::Pending;
}
if terminal.charges.total() == Ok(Cost::Unknown) {
return SettlementStatus::Unpriced;
}
if self
.progress
.effects
.iter()
.any(|state| matches!(state, EffectState::Accepted))
{
return SettlementStatus::Accepted;
}
SettlementStatus::Committed
}
pub fn retry_not_applied(&mut self, effect: Effect) -> Result<(), Error> {
if !matches!(self.progress.effect(effect), EffectState::NotApplied(_)) {
return Err(Error::UnsafeRetry { effect });
}
self.progress.effects[effect.index()] = EffectState::Pending;
Ok(())
}
pub async fn settle<B: Backend<R, U, Error = E>>(
&mut self,
backend: &mut B,
) -> Result<SettlementStatus, Error> {
let terminal = self.terminal.as_ref().ok_or(Error::NotTerminal)?;
for effect in Effect::ORDER {
if !matches!(self.progress.effects[effect.index()], EffectState::Pending) {
continue;
}
self.progress.effects[effect.index()] = EffectState::InFlight;
let result = backend
.apply(
effect,
Settlement {
admission: &self.admission,
terminal,
progress: &self.progress,
},
)
.await;
self.progress.effects[effect.index()] = result.into();
}
Ok(self.status())
}
}

View file

@ -0,0 +1,135 @@
use litellm_accounting::{
Charges, Cost, Error, Outcome, ProviderWork, ReportedUsage, Terminal, Usd,
};
use rstest::{fixture, rstest};
use rust_decimal::Decimal;
use rusty_money::{Money, MoneyError, iso};
fn usd(value: &str) -> Usd {
Money::from_str(value, iso::USD).unwrap()
}
#[fixture]
fn services() -> Cost {
Cost::Known(usd("2"))
}
#[rstest]
#[case::provider(Cost::Known(usd("-1")), Cost::Known(usd("0")))]
#[case::services(Cost::Known(usd("0")), Cost::Known(usd("-1")))]
fn negative_charges_cannot_enter_accounting(#[case] provider: Cost, #[case] services: Cost) {
assert_eq!(Charges::new(provider, services), Err(Error::NegativeCost));
}
#[rstest]
#[case::provider(Cost::Known(Money::from_major(1, iso::EUR)), Cost::Known(usd("0")))]
#[case::services(Cost::Known(usd("0")), Cost::Known(Money::from_major(1, iso::EUR)))]
fn non_usd_charges_cannot_enter_accounting(#[case] provider: Cost, #[case] services: Cost) {
assert_eq!(
Charges::new(provider, services),
Err(Error::InvalidCurrency {
currency: iso::EUR.iso_alpha_code,
})
);
}
#[rstest]
fn overflowing_totals_cannot_be_selected_for_settlement() {
let maximum = Cost::Known(Money::from_decimal(Decimal::MAX, iso::USD));
let charges = Charges::new(maximum, maximum).unwrap();
assert_eq!(charges.total(), Err(Error::Money(MoneyError::Overflow)));
assert_eq!(
Terminal::<()>::new(
Outcome::Succeeded,
ProviderWork::Started,
ReportedUsage::Unknown,
charges,
),
Err(Error::Money(MoneyError::Overflow))
);
}
#[rstest]
#[case::unknown_provider(Cost::Unknown, Cost::Known(usd("0")))]
#[case::unknown_services(Cost::Known(usd("0")), Cost::Unknown)]
#[case::both_unknown(Cost::Unknown, Cost::Unknown)]
fn unknown_cost_is_never_a_known_zero(#[case] provider: Cost, #[case] services: Cost) {
let charges = Charges::new(provider, services).unwrap();
assert_eq!(charges.total(), Ok(Cost::Unknown));
assert_eq!(charges.provider(), provider);
assert_eq!(charges.services(), services);
}
#[rstest]
#[case::decimal("0.1", "0.2", "0.3")]
#[case::fractional_cent("0.000000001", "0.000000002", "0.000000003")]
#[case::zero("0", "0", "0")]
fn decimal_totals_preserve_amounts_and_breakdown(
#[case] provider_amount: &str,
#[case] service_amount: &str,
#[case] expected_amount: &str,
) {
let provider = Cost::Known(usd(provider_amount));
let services = Cost::Known(usd(service_amount));
let charges = Charges::new(provider, services).unwrap();
assert_eq!(charges.total(), Ok(Cost::Known(usd(expected_amount))));
assert_eq!(charges.provider(), provider);
assert_eq!(charges.services(), services);
}
#[rstest]
#[case::success(Outcome::Succeeded)]
#[case::failure(Outcome::Failed)]
#[case::cancellation(Outcome::Cancelled)]
fn avoided_provider_work_preserves_reported_usage_and_service_charges(
services: Cost,
#[case] outcome: Outcome,
) {
let usage = ReportedUsage::Known((100_u64, 20_u64));
let terminal = Terminal::new(
outcome,
ProviderWork::NotStarted,
usage.clone(),
Charges::new(Cost::Known(usd("3")), services).unwrap(),
)
.unwrap();
assert_eq!(terminal.usage(), &usage);
assert_eq!(terminal.work(), ProviderWork::NotStarted);
assert_eq!(terminal.charges().provider(), Cost::Known(usd("0")));
assert_eq!(terminal.charges().services(), services);
assert_eq!(terminal.charges().total(), Ok(services));
assert_eq!(terminal.outcome(), outcome);
}
#[rstest]
fn avoided_provider_work_does_not_make_unknown_service_cost_free() {
let terminal = Terminal::<()>::new(
Outcome::Succeeded,
ProviderWork::NotStarted,
ReportedUsage::Unknown,
Charges::new(Cost::Unknown, Cost::Unknown).unwrap(),
)
.unwrap();
assert_eq!(terminal.charges().provider(), Cost::Known(usd("0")));
assert_eq!(terminal.charges().services(), Cost::Unknown);
assert_eq!(terminal.charges().total(), Ok(Cost::Unknown));
assert_eq!(terminal.usage(), &ReportedUsage::Unknown);
}
#[rstest]
#[case::success(Outcome::Succeeded)]
#[case::failure(Outcome::Failed)]
#[case::cancellation(Outcome::Cancelled)]
fn terminal_outcomes_do_not_erase_incurred_charges(services: Cost, #[case] outcome: Outcome) {
let charges = Charges::new(Cost::Known(usd("3")), services).unwrap();
let terminal = Terminal::new(
outcome,
ProviderWork::Started,
ReportedUsage::Known((10_u64, 1_u64)),
charges,
)
.unwrap();
assert_eq!(terminal.outcome(), outcome);
assert_eq!(terminal.charges(), charges);
assert_eq!(terminal.usage(), &ReportedUsage::Known((10, 1)));
}

View file

@ -0,0 +1,491 @@
use std::{future::Future, task::Context};
use litellm_accounting::{
ApplyResult, Backend, BudgetAdmission, Charges, Cost, Effect, EffectState, Error, Outcome,
ProviderWork, ReportedUsage, Session, Settlement, SettlementStatus, Terminal, Usd,
};
use rstest::{fixture, rstest};
use rusty_money::{Money, iso};
fn usd(value: &str) -> Usd {
Money::from_str(value, iso::USD).unwrap()
}
const EFFECTS: [Effect; 3] = [
Effect::ReconcileBudget,
Effect::RecordSpend,
Effect::ReleaseBudgetReservation,
];
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum Failure {
StoreUnavailable,
LostAcknowledgement,
}
#[derive(Clone, Debug, Eq, PartialEq)]
struct Usage {
input: u64,
output: u64,
}
struct Store {
calls: Vec<Effect>,
result: Option<(Effect, ApplyResult<Failure>)>,
blocked: Option<Effect>,
recorded: Vec<Charges>,
observations: Vec<(Outcome, ReportedUsage<Usage>)>,
admissions: Vec<Option<String>>,
release_progress: Vec<[EffectState<Failure>; 2]>,
}
impl Backend<String, Usage> for Store {
type Error = Failure;
async fn apply(
&mut self,
effect: Effect,
settlement: Settlement<'_, String, Usage, Failure>,
) -> ApplyResult<Failure> {
self.calls.push(effect);
self.admissions.push(settlement.admission.budget().cloned());
self.observations.push((
settlement.terminal.outcome(),
settlement.terminal.usage().clone(),
));
if effect == Effect::ReleaseBudgetReservation {
self.release_progress.push([
settlement.progress.effect(Effect::ReconcileBudget).clone(),
settlement.progress.effect(Effect::RecordSpend).clone(),
]);
}
if self.blocked == Some(effect) {
if effect == Effect::RecordSpend {
self.recorded.push(settlement.terminal.charges());
}
return std::future::pending().await;
}
if let Some((target, result)) = &self.result
&& *target == effect
{
return result.clone();
}
if effect == Effect::RecordSpend {
self.recorded.push(settlement.terminal.charges());
}
ApplyResult::Committed
}
}
#[fixture]
fn store() -> Store {
Store {
calls: Vec::new(),
result: None,
blocked: None,
recorded: Vec::new(),
observations: Vec::new(),
admissions: Vec::new(),
release_progress: Vec::new(),
}
}
#[fixture]
fn terminal() -> Terminal<Usage> {
Terminal::new(
Outcome::Succeeded,
ProviderWork::Started,
ReportedUsage::Known(Usage {
input: 10,
output: 2,
}),
Charges::new(Cost::Known(usd("4")), Cost::Known(usd("1"))).unwrap(),
)
.unwrap()
}
#[fixture]
fn session() -> Session<String, Usage, Failure> {
Session::new(BudgetAdmission::new(Some("budget-receipt".to_string())))
}
#[rstest]
#[tokio::test]
async fn admission_cannot_settle_before_a_terminal_outcome(
mut session: Session<String, Usage, Failure>,
mut store: Store,
terminal: Terminal<Usage>,
) {
assert_eq!(session.status(), SettlementStatus::Open);
assert_eq!(session.settle(&mut store).await, Err(Error::NotTerminal));
assert!(store.calls.is_empty());
session.finish(terminal).unwrap();
assert_eq!(session.status(), SettlementStatus::Pending);
assert!(store.calls.is_empty());
assert_eq!(
session.settle(&mut store).await,
Ok(SettlementStatus::Committed)
);
assert_eq!(store.calls, EFFECTS);
assert_eq!(store.recorded.len(), 1);
assert_eq!(
store.release_progress,
[[const { EffectState::Committed }; 2]]
);
}
#[rstest]
#[tokio::test]
async fn unreserved_calls_still_record_spend(mut store: Store, terminal: Terminal<Usage>) {
let mut session = Session::new(BudgetAdmission::new(None));
session.finish(terminal.clone()).unwrap();
assert_eq!(
session.settle(&mut store).await,
Ok(SettlementStatus::Committed)
);
assert_eq!(store.calls, [Effect::RecordSpend]);
assert_eq!(store.recorded, [terminal.charges()]);
assert_eq!(
session.progress().effect(Effect::ReleaseBudgetReservation),
&EffectState::NotRequired
);
}
#[rstest]
#[tokio::test]
async fn duplicate_terminal_delivery_cannot_replace_or_repeat_settlement(
mut session: Session<String, Usage, Failure>,
mut store: Store,
terminal: Terminal<Usage>,
) {
session.finish(terminal.clone()).unwrap();
session.settle(&mut store).await.unwrap();
assert_eq!(
session.finish(terminal.clone()),
Err(Error::AlreadyTerminal)
);
let conflicting = Terminal::new(
Outcome::Failed,
ProviderWork::NotStarted,
ReportedUsage::Unknown,
Charges::new(Cost::Known(usd("0")), Cost::Known(usd("0"))).unwrap(),
)
.unwrap();
assert_eq!(session.finish(conflicting), Err(Error::AlreadyTerminal));
assert_eq!(session.terminal(), Some(&terminal));
assert_eq!(
session.settle(&mut store).await,
Ok(SettlementStatus::Committed)
);
assert_eq!(store.calls, EFFECTS);
assert_eq!(store.recorded, [terminal.charges()]);
}
#[rstest]
#[case::budget(Effect::ReconcileBudget)]
#[case::spend(Effect::RecordSpend)]
#[case::release(Effect::ReleaseBudgetReservation)]
#[tokio::test]
async fn a_failed_effect_does_not_skip_cleanup_or_repeat_completed_work(
mut session: Session<String, Usage, Failure>,
mut store: Store,
terminal: Terminal<Usage>,
#[case] failed: Effect,
) {
store.result = Some((failed, ApplyResult::NotApplied(Failure::StoreUnavailable)));
session.finish(terminal).unwrap();
assert_eq!(
session.settle(&mut store).await,
Ok(SettlementStatus::NeedsAttention)
);
assert_eq!(store.calls, EFFECTS);
assert_eq!(
session.progress().effect(failed),
&EffectState::NotApplied(Failure::StoreUnavailable)
);
session.settle(&mut store).await.unwrap();
assert_eq!(store.calls, EFFECTS);
session.retry_not_applied(failed).unwrap();
store.result = None;
assert_eq!(
session.settle(&mut store).await,
Ok(SettlementStatus::Committed)
);
assert_eq!(store.calls, [EFFECTS.as_slice(), &[failed]].concat());
assert_eq!(store.recorded.len(), 1);
assert_eq!(
session.retry_not_applied(failed),
Err(Error::UnsafeRetry { effect: failed })
);
}
#[rstest]
#[case::budget(Effect::ReconcileBudget)]
#[case::spend(Effect::RecordSpend)]
#[case::release(Effect::ReleaseBudgetReservation)]
#[tokio::test]
async fn uncertain_remote_results_are_not_retried(
mut session: Session<String, Usage, Failure>,
mut store: Store,
terminal: Terminal<Usage>,
#[case] uncertain: Effect,
) {
store.result = Some((
uncertain,
ApplyResult::Indeterminate(Failure::LostAcknowledgement),
));
session.finish(terminal).unwrap();
session.settle(&mut store).await.unwrap();
assert_eq!(
session.progress().effect(uncertain),
&EffectState::Indeterminate(Failure::LostAcknowledgement)
);
assert_eq!(
session.retry_not_applied(uncertain),
Err(Error::UnsafeRetry { effect: uncertain })
);
assert_eq!(
session.settle(&mut store).await,
Ok(SettlementStatus::NeedsAttention)
);
assert_eq!(store.calls, EFFECTS);
}
#[rstest]
#[case::budget(Effect::ReconcileBudget)]
#[case::spend(Effect::RecordSpend)]
#[case::release(Effect::ReleaseBudgetReservation)]
#[tokio::test]
async fn queue_acceptance_is_distinct_from_committed_settlement(
mut session: Session<String, Usage, Failure>,
mut store: Store,
terminal: Terminal<Usage>,
#[case] queued: Effect,
) {
store.result = Some((queued, ApplyResult::Accepted));
session.finish(terminal).unwrap();
assert_eq!(
session.settle(&mut store).await,
Ok(SettlementStatus::Accepted)
);
assert_eq!(session.progress().effect(queued), &EffectState::Accepted);
assert_eq!(
session.retry_not_applied(queued),
Err(Error::UnsafeRetry { effect: queued })
);
session.settle(&mut store).await.unwrap();
assert_eq!(store.calls, EFFECTS);
}
#[rstest]
#[case::budget(Effect::ReconcileBudget)]
#[case::spend(Effect::RecordSpend)]
#[case::release(Effect::ReleaseBudgetReservation)]
#[tokio::test]
async fn dropping_settlement_retains_uncertainty_and_remaining_cleanup(
mut session: Session<String, Usage, Failure>,
mut store: Store,
terminal: Terminal<Usage>,
#[case] interrupted: Effect,
) {
store.blocked = Some(interrupted);
session.finish(terminal).unwrap();
{
let mut future = Box::pin(session.settle(&mut store));
let mut context = Context::from_waker(std::task::Waker::noop());
assert!(future.as_mut().poll(&mut context).is_pending());
}
assert_eq!(
session.progress().effect(interrupted),
&EffectState::InFlight
);
assert_eq!(session.status(), SettlementStatus::NeedsAttention);
assert_eq!(
session.retry_not_applied(interrupted),
Err(Error::UnsafeRetry {
effect: interrupted
})
);
assert_eq!(
session.settle(&mut store).await,
Ok(SettlementStatus::NeedsAttention)
);
assert_eq!(store.calls, EFFECTS);
if interrupted != Effect::ReleaseBudgetReservation {
assert_eq!(
session.progress().effect(Effect::ReleaseBudgetReservation),
&EffectState::Committed
);
}
assert_eq!(store.recorded.len(), 1);
}
#[rstest]
#[case::failure(Outcome::Failed)]
#[case::cancellation(Outcome::Cancelled)]
#[tokio::test]
async fn interrupted_execution_preserves_partial_usage_and_charges(
mut session: Session<String, Usage, Failure>,
mut store: Store,
#[case] outcome: Outcome,
) {
let usage = ReportedUsage::Known(Usage {
input: 10,
output: 1,
});
let charges = Charges::new(Cost::Known(usd("1")), Cost::Known(usd("0"))).unwrap();
session
.finish(Terminal::new(outcome, ProviderWork::Started, usage.clone(), charges).unwrap())
.unwrap();
assert_eq!(
session.settle(&mut store).await,
Ok(SettlementStatus::Committed)
);
assert_eq!(store.recorded, [charges]);
assert_eq!(store.observations, vec![(outcome, usage); EFFECTS.len()]);
assert_eq!(store.calls.last(), Some(&Effect::ReleaseBudgetReservation));
}
#[rstest]
#[case::unknown(ReportedUsage::Unknown, Cost::Unknown, SettlementStatus::Unpriced)]
#[case::known_zero(
ReportedUsage::Known(Usage { input: 0, output: 0 }),
Cost::Known(usd("0")),
SettlementStatus::Committed
)]
#[case::usage_unknown_but_cost_known(
ReportedUsage::Unknown,
Cost::Known(usd("2")),
SettlementStatus::Committed
)]
#[case::usage_reported_but_price_unknown(
ReportedUsage::Known(Usage { input: 10, output: 2 }),
Cost::Unknown,
SettlementStatus::Unpriced
)]
#[tokio::test]
async fn unknown_usage_and_cost_remain_distinct_from_reported_zero(
mut session: Session<String, Usage, Failure>,
mut store: Store,
#[case] usage: ReportedUsage<Usage>,
#[case] cost: Cost,
#[case] expected: SettlementStatus,
) {
let terminal = Terminal::new(
Outcome::Failed,
ProviderWork::Started,
usage.clone(),
Charges::new(cost, Cost::Known(usd("0"))).unwrap(),
)
.unwrap();
session.finish(terminal).unwrap();
assert_eq!(session.settle(&mut store).await, Ok(expected));
assert_eq!(store.recorded[0].provider(), cost);
assert_eq!(store.recorded[0].total(), Ok(cost));
assert_eq!(session.terminal().unwrap().usage(), &usage);
assert_eq!(
session.progress().effect(Effect::ReleaseBudgetReservation),
&EffectState::Committed
);
}
#[rstest]
#[case::reserved(true)]
#[case::unreserved(false)]
#[tokio::test]
async fn budget_receipt_is_preserved_for_every_effect(
mut store: Store,
terminal: Terminal<Usage>,
#[case] reserved: bool,
) {
let receipt = reserved.then(|| "budget-receipt".to_string());
let admission = BudgetAdmission::new(receipt.clone());
assert_eq!(admission.budget(), receipt.as_ref());
let mut session = Session::new(admission);
session.finish(terminal).unwrap();
assert_eq!(
session.settle(&mut store).await,
Ok(SettlementStatus::Committed)
);
let expected = if reserved {
EFFECTS.as_slice()
} else {
&[Effect::RecordSpend]
};
assert_eq!(store.calls, expected);
assert_eq!(store.admissions, vec![receipt; expected.len()]);
assert_eq!(store.recorded.len(), 1);
}
#[rstest]
#[case::not_applied(ApplyResult::NotApplied(Failure::StoreUnavailable))]
#[case::uncertain(ApplyResult::Indeterminate(Failure::LostAcknowledgement))]
#[case::queued(ApplyResult::Accepted)]
#[tokio::test]
async fn cleanup_receives_the_actual_settlement_progress(
mut session: Session<String, Usage, Failure>,
mut store: Store,
terminal: Terminal<Usage>,
#[case] result: ApplyResult<Failure>,
) {
store.result = Some((Effect::ReconcileBudget, result.clone()));
session.finish(terminal).unwrap();
session.settle(&mut store).await.unwrap();
assert_eq!(
store.release_progress,
[[result.into(), EffectState::Committed]]
);
}
#[rstest]
#[tokio::test]
async fn retries_require_a_confirmed_not_applied_result(
mut session: Session<String, Usage, Failure>,
terminal: Terminal<Usage>,
) {
assert_eq!(
session.retry_not_applied(Effect::RecordSpend),
Err(Error::UnsafeRetry {
effect: Effect::RecordSpend
})
);
session.finish(terminal).unwrap();
assert_eq!(
session.retry_not_applied(Effect::RecordSpend),
Err(Error::UnsafeRetry {
effect: Effect::RecordSpend
})
);
}
#[rstest]
#[tokio::test]
async fn unknown_charges_cannot_become_successful_settlement_through_queue_acceptance(
mut session: Session<String, Usage, Failure>,
mut store: Store,
) {
store.result = Some((Effect::RecordSpend, ApplyResult::Accepted));
session
.finish(
Terminal::new(
Outcome::Cancelled,
ProviderWork::Started,
ReportedUsage::Unknown,
Charges::new(Cost::Unknown, Cost::Known(usd("0"))).unwrap(),
)
.unwrap(),
)
.unwrap();
assert_eq!(
session.settle(&mut store).await,
Ok(SettlementStatus::Unpriced)
);
assert_eq!(
session.progress().effect(Effect::RecordSpend),
&EffectState::Accepted
);
assert_eq!(
session.progress().effect(Effect::ReleaseBudgetReservation),
&EffectState::Committed
);
}