//! Billing service — handles user-level and project-level billing, deduction, //! credit initialization, and error persistence. //! //! Architecture: //! - Each user gets $10 personal balance on signup. //! - Each project gets $20 balance only if it's the creator's first project, //! $0 otherwise. //! - AI usage is deducted from the project balance first; if insufficient, //! falls through to the user's personal balance. //! - Monthly quota only applies to pro users (is_pro = true). //! - If both project and user balance are insufficient, a billing_error //! record is persisted and an error is returned to the caller. use db::database::AppDatabase; use models::agents::model_pricing; use models::ai::billing_error; use models::projects::{project, project_billing, project_billing_history}; use models::users::{user_billing, user_billing_history}; use rust_decimal::Decimal; use sea_orm::*; use uuid::Uuid; use crate::error::AgentError; fn default_user_balance() -> Decimal { Decimal::new(100_000, 4) } // $10.0000 fn first_project_credit() -> Decimal { Decimal::new(200_000, 4) } // $20.0000 const SUBSEQUENT_PROJECT_BALANCE: Decimal = Decimal::ZERO; #[derive(Debug, Clone, serde::Serialize, serde::Deserialize, utoipa::ToSchema)] pub struct BillingRecord { pub cost: f64, pub currency: String, pub input_tokens: i64, pub output_tokens: i64, pub deducted_from: String, // "project" or "user" } #[derive(Debug)] pub enum BillingResult { Success(BillingRecord), InsufficientBalance { message: String }, } /// Record AI usage: deduct from project balance first, fall through to user balance. /// /// Returns `InsufficientBalance` if neither account can cover the cost. /// On insufficient balance, a `billing_error` record is persisted for frontend display. pub async fn record_ai_usage( db: &AppDatabase, project_uid: Uuid, user_uid: Uuid, model_id: Uuid, input_tokens: i64, output_tokens: i64, ) -> Result { let total_cost = compute_cost(db, model_id, input_tokens, output_tokens).await?; let currency = get_currency(db, model_id).await?; // Verify project exists let _ = project::Entity::find_by_id(project_uid) .one(db) .await? .ok_or_else(|| AgentError::Internal("Project not found".into()))?; // Attempt project-level deduction first let project_result = deduct_from_project( db, project_uid, total_cost, ¤cy, model_id, input_tokens, output_tokens, ) .await; match project_result { Ok(()) => { let cost_f64 = decimal_to_f64(total_cost); tracing::info!( project_id = %project_uid, model_id = %model_id, input_tokens, output_tokens, cost = %cost_f64, currency = %currency, deducted_from = "project", "ai_usage_recorded" ); Ok(BillingResult::Success(BillingRecord { cost: cost_f64, currency, input_tokens, output_tokens, deducted_from: "project".to_string(), })) } Err(_) => { // Project balance insufficient — try user personal balance let user_result = deduct_from_user( db, user_uid, total_cost, ¤cy, project_uid, model_id, input_tokens, output_tokens, ) .await; match user_result { Ok(()) => { let cost_f64 = decimal_to_f64(total_cost); tracing::info!( user_id = %user_uid, project_id = %project_uid, model_id = %model_id, input_tokens, output_tokens, cost = %cost_f64, currency = %currency, deducted_from = "user", "ai_usage_recorded" ); Ok(BillingResult::Success(BillingRecord { cost: cost_f64, currency, input_tokens, output_tokens, deducted_from: "user".to_string(), })) } Err(insufficient_msg) => { // Both project and user balance insufficient — persist error persist_billing_error( db, "project", project_uid, "insufficient_balance", &insufficient_msg, Some(serde_json::json!({ "user_id": user_uid.to_string(), "model_id": model_id.to_string(), "input_tokens": input_tokens, "output_tokens": output_tokens, "cost": decimal_to_f64(total_cost), "currency": currency, })), ) .await?; Ok(BillingResult::InsufficientBalance { message: insufficient_msg, }) } } } } } /// Record personal AI usage against the user's own billing balance. pub async fn record_user_ai_usage( db: &AppDatabase, user_uid: Uuid, model_id: Uuid, input_tokens: i64, output_tokens: i64, ) -> Result { let total_cost = compute_cost(db, model_id, input_tokens, output_tokens).await?; let currency = get_currency(db, model_id).await?; match deduct_from_user_personal( db, user_uid, total_cost, ¤cy, model_id, input_tokens, output_tokens, ) .await { Ok(()) => { let cost_f64 = decimal_to_f64(total_cost); tracing::info!( user_id = %user_uid, model_id = %model_id, input_tokens, output_tokens, cost = %cost_f64, currency = %currency, deducted_from = "user", scope = "personal", "ai_usage_recorded" ); Ok(BillingResult::Success(BillingRecord { cost: cost_f64, currency, input_tokens, output_tokens, deducted_from: "user".to_string(), })) } Err(insufficient_msg) => { persist_billing_error( db, "user", user_uid, "insufficient_balance", &insufficient_msg, Some(serde_json::json!({ "user_id": user_uid.to_string(), "model_id": model_id.to_string(), "input_tokens": input_tokens, "output_tokens": output_tokens, "cost": decimal_to_f64(total_cost), "currency": currency, "scope": "personal", })), ) .await?; Ok(BillingResult::InsufficientBalance { message: insufficient_msg, }) } } } /// Check whether a project + user has sufficient combined balance for a potential AI call. /// Called before starting AI processing to avoid wasted compute. pub async fn check_balance( db: &AppDatabase, project_uid: Uuid, user_uid: Uuid, model_id: Uuid, estimated_input_tokens: i64, estimated_output_tokens: i64, ) -> Result { let estimated_cost = compute_cost( db, model_id, estimated_input_tokens, estimated_output_tokens, ) .await?; let project_balance = get_project_balance(db, project_uid).await; let user_balance = get_user_balance(db, user_uid).await; Ok(project_balance + user_balance >= estimated_cost) } /// Check whether a user's personal balance can cover a potential AI call. pub async fn check_user_balance( db: &AppDatabase, user_uid: Uuid, model_id: Uuid, estimated_input_tokens: i64, estimated_output_tokens: i64, ) -> Result { let estimated_cost = compute_cost( db, model_id, estimated_input_tokens, estimated_output_tokens, ) .await?; let user_balance = get_user_balance(db, user_uid).await; Ok(user_balance >= estimated_cost) } // ── Initialization ── /// Initialize a user billing account with the default $10 balance. /// Called on user signup / first login. pub async fn initialize_user_billing(db: &AppDatabase, user_uid: Uuid) -> Result<(), AgentError> { let now = chrono::Utc::now(); user_billing::ActiveModel { user: Set(user_uid), balance: Set(default_user_balance()), currency: Set("USD".to_string()), is_pro: Set(false), monthly_quota: Set(Decimal::ZERO), month_used: Set(Decimal::ZERO), cycle_start: Set(None), cycle_end: Set(None), updated_at: Set(now), created_at: Set(now), } .insert(db) .await .map_err(|e| AgentError::Internal(format!("failed to create user billing: {}", e)))?; tracing::info!(user_id = %user_uid, balance = "$10", "user_billing_initialized"); Ok(()) } /// Initialize a project billing account. /// Grants $20 only if this is the creator's first project; $0 otherwise. pub async fn initialize_project_billing( db: &AppDatabase, project_uid: Uuid, creator_uid: Uuid, ) -> Result<(), AgentError> { // Check how many projects this user has already created let existing_count = project::Entity::find() .filter(project::Column::CreatedBy.eq(creator_uid)) .filter(project::Column::Id.ne(project_uid)) .count(db) .await .map_err(|e| AgentError::Internal(format!("failed to count user projects: {}", e)))?; let is_first = existing_count == 0; let initial_balance = if is_first { first_project_credit() } else { SUBSEQUENT_PROJECT_BALANCE }; let now = chrono::Utc::now(); project_billing::ActiveModel { project: Set(project_uid), balance: Set(initial_balance), currency: Set("USD".to_string()), user: Set(Some(creator_uid)), initial_credit_granted: Set(is_first), is_pro: Set(false), monthly_quota: Set(Decimal::ZERO), month_used: Set(Decimal::ZERO), cycle_start: Set(None), cycle_end: Set(None), updated_at: Set(now), created_at: Set(now), } .insert(db) .await .map_err(|e| AgentError::Internal(format!("failed to create project billing: {}", e)))?; if is_first { // Record the credit in billing history project_billing_history::ActiveModel { uid: Set(Uuid::new_v4()), project: Set(project_uid), user: Set(Some(creator_uid)), amount: Set(first_project_credit()), currency: Set("USD".to_string()), reason: Set("first_project_credit".to_string()), extra: Set(Some(serde_json::json!({ "is_first_project": true, }))), created_at: Set(now), ..Default::default() } .insert(db) .await .map_err(|e| AgentError::Internal(format!("failed to record credit history: {}", e)))?; } tracing::info!( project_id = %project_uid, creator_id = %creator_uid, is_first_project = is_first, balance = if is_first { "$20" } else { "$0" }, "project_billing_initialized" ); Ok(()) } // ── Internal helpers ── async fn compute_cost( db: &AppDatabase, model_id: Uuid, input_tokens: i64, output_tokens: i64, ) -> Result { let pricing = model_pricing::Entity::find() .filter(model_pricing::Column::ModelVersionId.eq(model_id)) .order_by_desc(model_pricing::Column::EffectiveFrom) .one(db) .await? .ok_or_else(|| { AgentError::Internal( "No pricing record found for this model. Please configure AI model pricing first." .into(), ) })?; let input_price: Decimal = pricing .input_price_per_1k_tokens .parse() .map_err(|e| AgentError::Internal(format!("Invalid input price: {}", e)))?; let output_price: Decimal = pricing .output_price_per_1k_tokens .parse() .map_err(|e| AgentError::Internal(format!("Invalid output price: {}", e)))?; if input_price <= Decimal::ZERO && output_price <= Decimal::ZERO { return Err(AgentError::Internal( "Model pricing is not configured or is zero. Please configure non-zero AI model pricing first." .into(), )); } // DB stores per-1M-token prices; divide tokens by 1M to compute cost. let million = Decimal::from(1_000_000); Ok((Decimal::from(input_tokens) / million) * input_price + (Decimal::from(output_tokens) / million) * output_price) } async fn get_currency(db: &AppDatabase, model_id: Uuid) -> Result { let pricing = model_pricing::Entity::find() .filter(model_pricing::Column::ModelVersionId.eq(model_id)) .one(db) .await? .ok_or_else(|| AgentError::Internal("No pricing found".into()))?; Ok(pricing.currency.clone()) } async fn get_project_balance(db: &AppDatabase, project_uid: Uuid) -> Decimal { project_billing::Entity::find_by_id(project_uid) .one(db) .await .ok() .flatten() .map(|b| b.balance) .unwrap_or(Decimal::ZERO) } async fn get_user_balance(db: &AppDatabase, user_uid: Uuid) -> Decimal { user_billing::Entity::find_by_id(user_uid) .one(db) .await .ok() .flatten() .map(|b| b.balance) .unwrap_or(Decimal::ZERO) } async fn deduct_from_project( db: &AppDatabase, project_uid: Uuid, cost: Decimal, currency: &str, model_id: Uuid, input_tokens: i64, output_tokens: i64, ) -> Result<(), String> { let txn = db .begin() .await .map_err(|e| format!("db txn error: {}", e))?; let billing = project_billing::Entity::find_by_id(project_uid) .lock_exclusive() .one(&txn) .await .map_err(|e| format!("db error: {}", e))? .ok_or_else(|| "Project billing account not found".to_string())?; if billing.balance < cost { txn.rollback().await.ok(); return Err(format!( "Project balance insufficient. Required: {:.4} {}, Available: {:.4} {}", cost, currency, billing.balance, currency )); } let now = chrono::Utc::now(); project_billing_history::ActiveModel { uid: Set(Uuid::new_v4()), project: Set(project_uid), user: Set(None), amount: Set(-cost), currency: Set(currency.to_string()), reason: Set("ai_usage".to_string()), extra: Set(Some(serde_json::json!({ "model_id": model_id.to_string(), "input_tokens": input_tokens, "output_tokens": output_tokens, "deducted_from": "project", }))), created_at: Set(now), ..Default::default() } .insert(&txn) .await .map_err(|e| format!("failed to insert history: {}", e))?; let mut updated: project_billing::ActiveModel = billing.into(); updated.balance = Set(updated.balance.unwrap() - cost); updated.updated_at = Set(now); updated .update(&txn) .await .map_err(|e| format!("failed to update balance: {}", e))?; txn.commit() .await .map_err(|e| format!("commit error: {}", e))?; Ok(()) } async fn deduct_from_user( db: &AppDatabase, user_uid: Uuid, cost: Decimal, currency: &str, project_uid: Uuid, model_id: Uuid, input_tokens: i64, output_tokens: i64, ) -> Result<(), String> { let txn = db .begin() .await .map_err(|e| format!("db txn error: {}", e))?; let billing = user_billing::Entity::find_by_id(user_uid) .lock_exclusive() .one(&txn) .await .map_err(|e| format!("db error: {}", e))? .ok_or_else(|| "User billing account not found".to_string())?; if billing.balance < cost { txn.rollback().await.ok(); return Err(format!( "Insufficient balance (project + user). Project: unavailable, User: {:.4} {}. Required: {:.4} {}", billing.balance, currency, cost, currency )); } let now = chrono::Utc::now(); // Record in project billing history (but deducted from user) project_billing_history::ActiveModel { uid: Set(Uuid::new_v4()), project: Set(project_uid), user: Set(Some(user_uid)), amount: Set(-cost), currency: Set(currency.to_string()), reason: Set("ai_usage_user_fallback".to_string()), extra: Set(Some(serde_json::json!({ "model_id": model_id.to_string(), "input_tokens": input_tokens, "output_tokens": output_tokens, "deducted_from": "user", }))), created_at: Set(now), ..Default::default() } .insert(&txn) .await .map_err(|e| format!("failed to insert history: {}", e))?; let mut updated: user_billing::ActiveModel = billing.into(); updated.balance = Set(updated.balance.unwrap() - cost); updated.updated_at = Set(now); updated .update(&txn) .await .map_err(|e| format!("failed to update user balance: {}", e))?; txn.commit() .await .map_err(|e| format!("commit error: {}", e))?; Ok(()) } async fn deduct_from_user_personal( db: &AppDatabase, user_uid: Uuid, cost: Decimal, currency: &str, model_id: Uuid, input_tokens: i64, output_tokens: i64, ) -> Result<(), String> { let txn = db .begin() .await .map_err(|e| format!("db txn error: {}", e))?; let billing = user_billing::Entity::find_by_id(user_uid) .lock_exclusive() .one(&txn) .await .map_err(|e| format!("db error: {}", e))? .ok_or_else(|| "User billing account not found".to_string())?; if billing.balance < cost { txn.rollback().await.ok(); return Err(format!( "Insufficient balance. User: {:.4} {}. Required: {:.4} {}", billing.balance, billing.currency, cost, billing.currency )); } let now = chrono::Utc::now(); user_billing_history::ActiveModel { uid: Set(Uuid::new_v4()), user: Set(user_uid), amount: Set(-cost), currency: Set(currency.to_string()), reason: Set("ai_usage_personal".to_string()), extra: Set(Some(serde_json::json!({ "model_id": model_id.to_string(), "input_tokens": input_tokens, "output_tokens": output_tokens, "deducted_from": "user", "scope": "personal", }))), created_at: Set(now), ..Default::default() } .insert(&txn) .await .map_err(|e| format!("failed to insert user history: {}", e))?; let mut updated: user_billing::ActiveModel = billing.into(); updated.balance = Set(updated.balance.unwrap() - cost); updated.updated_at = Set(now); updated .update(&txn) .await .map_err(|e| format!("failed to update user balance: {}", e))?; txn.commit() .await .map_err(|e| format!("commit error: {}", e))?; Ok(()) } pub async fn persist_billing_error( db: &AppDatabase, scope: &str, scope_id: Uuid, error_type: &str, message: &str, details: Option, ) -> Result<(), AgentError> { billing_error::ActiveModel { id: Set(Uuid::new_v4()), scope: Set(scope.to_string()), scope_id: Set(scope_id), error_type: Set(error_type.to_string()), message: Set(message.to_string()), details: Set(details), resolved: Set(false), created_at: Set(chrono::Utc::now()), } .insert(db) .await .map_err(|e| AgentError::Internal(format!("failed to persist billing error: {}", e)))?; tracing::warn!(scope, %scope_id, error_type, "billing_error_persisted"); Ok(()) } fn decimal_to_f64(d: Decimal) -> f64 { d.round_dp(10).to_string().parse().unwrap_or(0.0) }