diff --git a/crates/api/src/routes/usage.rs b/crates/api/src/routes/usage.rs index b02e8a06..8dfa4fc2 100644 --- a/crates/api/src/routes/usage.rs +++ b/crates/api/src/routes/usage.rs @@ -97,6 +97,10 @@ pub struct OrganizationBalanceResponse { pub total_tokens: i64, pub updated_at: String, pub credit_limits: Vec, + /// Spend recorded before per-type allocation was available. For multiple + /// active types, subtract this from their summed available amounts to + /// reconcile with the top-level remaining balance. + pub legacy_unattributed_amount: i64, /// Unresolved cost from completed requests that exceeded all capacity. pub unfunded_amount: i64, } @@ -176,6 +180,7 @@ pub async fn compute_organization_balance_response( total_tokens: balance.total_tokens, updated_at: balance.updated_at.to_rfc3339(), credit_limits, + legacy_unattributed_amount: balance.legacy_unattributed_amount, unfunded_amount: limit.as_ref().map_or(0, |value| value.unfunded), }) } @@ -194,6 +199,7 @@ pub async fn compute_organization_balance_response( total_tokens: 0, updated_at: Utc::now().to_rfc3339(), credit_limits, + legacy_unattributed_amount: 0, unfunded_amount: limit_info.unfunded, }) } else { diff --git a/crates/database/src/models.rs b/crates/database/src/models.rs index 26096da0..21839dde 100644 --- a/crates/database/src/models.rs +++ b/crates/database/src/models.rs @@ -739,6 +739,7 @@ pub struct OrganizationUsageLog { pub struct OrganizationBalance { pub organization_id: Uuid, pub total_spent: i64, + pub legacy_unattributed_amount: i64, pub last_usage_at: Option>, pub total_requests: i64, pub total_tokens: i64, diff --git a/crates/database/src/repositories/organization_limits_repository_impl.rs b/crates/database/src/repositories/organization_limits_repository_impl.rs index de53d9ce..1812dfd6 100644 --- a/crates/database/src/repositories/organization_limits_repository_impl.rs +++ b/crates/database/src/repositories/organization_limits_repository_impl.rs @@ -40,7 +40,21 @@ impl services::usage::ports::OrganizationLimitsRepository for OrganizationLimits &self, organization_id: Uuid, ) -> anyhow::Result> { - let (limits, _, _) = self.get_current_credit_status(organization_id).await?; + let (mut limits, _, unattributed) = self.get_current_credit_status(organization_id).await?; + + // With one active credit type, legacy spend has only one current + // capacity bucket to reduce. Fold it into that type so its breakdown + // agrees with the aggregate remaining balance. With multiple types, + // keep the legacy amount unattributed rather than inventing a split. + if let [status] = limits.as_mut_slice() { + status.consumed = status.consumed.saturating_add(unattributed); + status.available = status + .limit + .spend_limit + .saturating_sub(status.consumed) + .max(0); + } + Ok(limits .into_iter() .map(|status| OrganizationCreditLimit { diff --git a/crates/database/src/repositories/organization_usage.rs b/crates/database/src/repositories/organization_usage.rs index 3ec9cbc1..8824a71e 100644 --- a/crates/database/src/repositories/organization_usage.rs +++ b/crates/database/src/repositories/organization_usage.rs @@ -325,7 +325,7 @@ impl OrganizationUsageRepository { client .query_opt( r#" - SELECT organization_id, total_spent, last_usage_at, + SELECT organization_id, total_spent, legacy_unattributed_amount, last_usage_at, total_requests, total_tokens, updated_at FROM organization_balance WHERE organization_id = $1 @@ -678,6 +678,7 @@ impl OrganizationUsageRepository { OrganizationBalance { organization_id: row.get("organization_id"), total_spent: row.get("total_spent"), + legacy_unattributed_amount: row.get("legacy_unattributed_amount"), last_usage_at: row.get("last_usage_at"), total_requests: row.get("total_requests"), total_tokens: row.get("total_tokens"), diff --git a/crates/database/src/repositories/usage_repository_impl.rs b/crates/database/src/repositories/usage_repository_impl.rs index c7d19729..fbd47756 100644 --- a/crates/database/src/repositories/usage_repository_impl.rs +++ b/crates/database/src/repositories/usage_repository_impl.rs @@ -97,6 +97,7 @@ impl services::usage::ports::UsageRepository for OrganizationUsageRepository { Ok(balance.map(|b| OrganizationBalanceInfo { organization_id: b.organization_id, total_spent: b.total_spent, + legacy_unattributed_amount: b.legacy_unattributed_amount, last_usage_at: b.last_usage_at, total_requests: b.total_requests, total_tokens: b.total_tokens, diff --git a/crates/database/tests/credit_allocation.rs b/crates/database/tests/credit_allocation.rs index a8182a12..c80d034c 100644 --- a/crates/database/tests/credit_allocation.rs +++ b/crates/database/tests/credit_allocation.rs @@ -568,6 +568,107 @@ async fn ambiguous_legacy_spend_stays_unknown_and_cannot_restore_capacity() -> a Ok(()) } +#[tokio::test] +async fn single_credit_type_breakdown_includes_legacy_spend() -> anyhow::Result<()> { + let pool = test_pool().await?; + let limits = OrganizationLimitsRepository::new(pool.clone()); + let repository = OrganizationUsageRepository::new(pool.clone()); + let org = insert_org_fixture(&pool).await?; + let model = insert_model(&pool, "allocation-single-type-legacy-balance").await?; + set_limit(&limits, org.org_id, "payment", 100).await?; + + let client = pool.get().await?; + client + .execute( + r#"UPDATE organization_balance + SET total_spent = 90, legacy_unattributed_amount = 90, + total_requests = 1, total_tokens = 1, updated_at = NOW() + WHERE organization_id = $1"#, + &[&org.org_id], + ) + .await?; + drop(client); + + let usage = repository + .record_usage(usage(&org, &model, Uuid::new_v4(), 1)) + .await?; + assert_eq!(usage.credit_allocations.unwrap()[0].amount, 1); + + let breakdown = + services::usage::ports::OrganizationLimitsRepository::get_current_limit_breakdown( + &limits, org.org_id, + ) + .await?; + assert_eq!(breakdown.len(), 1); + assert_eq!(breakdown[0].credit_type, "payment"); + assert_eq!(breakdown[0].amount, 100); + assert_eq!(breakdown[0].consumed, 91); + assert_eq!(breakdown[0].available, 9); + + let aggregate = services::usage::ports::OrganizationLimitsRepository::get_current_limits( + &limits, org.org_id, + ) + .await? + .expect("organization has an active payment limit"); + assert_eq!(aggregate.available, 9); + + cleanup_usage_fixtures(&pool, &[org.org_id], &[model.id]).await?; + Ok(()) +} + +#[tokio::test] +async fn multi_type_breakdown_exposes_legacy_spend_without_assigning_it() -> anyhow::Result<()> { + let pool = test_pool().await?; + let limits = OrganizationLimitsRepository::new(pool.clone()); + let repository = OrganizationUsageRepository::new(pool.clone()); + let org = insert_org_fixture(&pool).await?; + set_limit(&limits, org.org_id, "grant", 100).await?; + set_limit(&limits, org.org_id, "payment", 100).await?; + + let client = pool.get().await?; + client + .execute( + r#"UPDATE organization_balance + SET total_spent = 90, legacy_unattributed_amount = 90, + total_requests = 1, total_tokens = 1, updated_at = NOW() + WHERE organization_id = $1"#, + &[&org.org_id], + ) + .await?; + drop(client); + + let breakdown = + services::usage::ports::OrganizationLimitsRepository::get_current_limit_breakdown( + &limits, org.org_id, + ) + .await?; + assert_eq!(breakdown.len(), 2); + assert!(breakdown + .iter() + .all(|credit| credit.consumed == 0 && credit.available == 100)); + + let aggregate = services::usage::ports::OrganizationLimitsRepository::get_current_limits( + &limits, org.org_id, + ) + .await? + .expect("organization has active credit limits"); + assert_eq!(aggregate.available, 110); + + let balance = repository + .get_balance(org.org_id) + .await? + .expect("organization has a balance row"); + assert_eq!(balance.legacy_unattributed_amount, 90); + assert_eq!( + breakdown.iter().map(|credit| credit.available).sum::() + - balance.legacy_unattributed_amount, + aggregate.available + ); + + cleanup_usage_fixtures(&pool, &[org.org_id], &[]).await?; + Ok(()) +} + #[tokio::test] async fn limit_replacement_preserves_consumption_and_custom_order_is_honored() -> anyhow::Result<()> { diff --git a/crates/services/src/usage/mod.rs b/crates/services/src/usage/mod.rs index 87bc89ab..71426139 100644 --- a/crates/services/src/usage/mod.rs +++ b/crates/services/src/usage/mod.rs @@ -812,6 +812,7 @@ impl UsageServiceTrait for UsageServiceImpl { Ok(balance.map(|b| OrganizationBalanceInfo { organization_id: b.organization_id, total_spent: b.total_spent, + legacy_unattributed_amount: b.legacy_unattributed_amount, last_usage_at: b.last_usage_at, total_requests: b.total_requests, total_tokens: b.total_tokens, diff --git a/crates/services/src/usage/ports.rs b/crates/services/src/usage/ports.rs index 210b8a7b..a21fa173 100644 --- a/crates/services/src/usage/ports.rs +++ b/crates/services/src/usage/ports.rs @@ -690,6 +690,7 @@ pub enum UsageCheckResult { pub struct OrganizationBalanceInfo { pub organization_id: Uuid, pub total_spent: i64, + pub legacy_unattributed_amount: i64, pub last_usage_at: Option>, pub total_requests: i64, pub total_tokens: i64,