//! AI usage billing — records token costs against a project or workspace balance. //! //! Called by service-layer code after each successful AI call. If the project //! belongs to a workspace, the cost is deducted from the workspace's shared quota //! (workspace_billing). Otherwise it is deducted from the project's own quota. //! //! 1. Queries the most recent active price for `model_id`. //! 2. Computes `cost = (input/1000)*input_price + (output/1000)*output_price`. //! 3. Determines whether to bill the project or its workspace. //! 4. Writes a billing_history entry and decrements the appropriate balance. use crate::AppService; use crate::error::AppError; use models::agents::model_pricing; use models::projects::project; use models::projects::project_billing; use models::projects::project_billing_history; use models::workspaces::workspace_billing; use models::workspaces::workspace_billing_history; use rust_decimal::Decimal; use sea_orm::*; use serde::{Deserialize, Serialize}; use utoipa::ToSchema; use uuid::Uuid; /// Breakdown of a billing record. #[derive(Debug, Clone, Serialize, Deserialize, ToSchema)] pub struct BillingRecord { /// Total cost in the billing currency. pub cost: f64, pub currency: String, pub input_tokens: i64, pub output_tokens: i64, } impl AppService { /// Record AI usage for a project. /// /// If the project belongs to a workspace, the cost is deducted from the /// workspace's shared quota. Otherwise it is deducted from the project's own /// billing balance. /// /// Returns an error if there is insufficient balance. pub async fn record_ai_usage( &self, project_uid: Uuid, model_id: Uuid, input_tokens: i64, output_tokens: i64, ) -> Result { // 1. Look up the active price for this model. let pricing = model_pricing::Entity::find() .filter(model_pricing::Column::ModelVersionId.eq(model_id)) .order_by_desc(model_pricing::Column::EffectiveFrom) .one(&self.db) .await? .ok_or_else(|| { AppError::InternalServerError( "No pricing record found for this model. Please configure AI model pricing first." .into(), ) })?; // 2. Compute cost using Decimal arithmetic. let input_price: Decimal = pricing .input_price_per_1k_tokens .parse() .unwrap_or(Decimal::ZERO); let output_price: Decimal = pricing .output_price_per_1k_tokens .parse() .unwrap_or(Decimal::ZERO); let tokens_i = Decimal::from(input_tokens); let tokens_o = Decimal::from(output_tokens); let thousand = Decimal::from(1000); let total_cost: f64 = ((tokens_i / thousand) * input_price + (tokens_o / thousand) * output_price) .to_string() .parse() .unwrap_or(0.0); let currency = pricing.currency.clone(); // 3. Determine whether to bill the project or its workspace. let proj = project::Entity::find_by_id(project_uid) .one(&self.db) .await? .ok_or_else(|| AppError::NotFound("Project not found".to_string()))?; if let Some(workspace_id) = proj.workspace_id { // ── Workspace-shared quota ────────────────────────────────── let current = workspace_billing::Entity::find_by_id(workspace_id) .one(&self.db) .await? .ok_or_else(|| { AppError::NotFound("Workspace billing account not found".to_string()) })?; let current_balance: f64 = current.balance.to_string().parse().unwrap_or(0.0); if current_balance < total_cost { return Err(AppError::BadRequest(format!( "Insufficient workspace billing balance. Required: {:.4} {}, Available: {:.4} {}", total_cost, currency, current_balance, currency ))); } let amount_dec = Decimal::from_f64_retain(-total_cost).unwrap_or(Decimal::ZERO); let now = chrono::Utc::now(); // Insert workspace billing history. let _ = workspace_billing_history::ActiveModel { uid: Set(Uuid::new_v4()), workspace_id: Set(workspace_id), user_id: Set(Some(proj.created_by)), amount: Set(amount_dec), currency: Set(currency.clone()), reason: Set(format!("ai_usage:{}", project_uid)), extra: Set(Some(serde_json::json!({ "project_id": project_uid.to_string(), "model_id": model_id.to_string(), "input_tokens": input_tokens, "output_tokens": output_tokens, }))), created_at: Set(now), } .insert(&self.db) .await; // Deduct from workspace balance. let new_balance = Decimal::from_f64_retain(current_balance - total_cost).unwrap_or(Decimal::ZERO); let mut updated: workspace_billing::ActiveModel = current.into(); updated.balance = Set(new_balance); updated.updated_at = Set(now); updated.update(&self.db).await?; Ok(BillingRecord { cost: total_cost, currency, input_tokens, output_tokens, }) } else { // ── Project-owned quota ───────────────────────────────────── let amount_dec = Decimal::from_f64_retain(-total_cost).unwrap_or(Decimal::ZERO); let _ = project_billing_history::ActiveModel { uid: Set(Uuid::new_v4()), project: Set(project_uid), user: Set(None), amount: Set(amount_dec), currency: Set(currency.clone()), 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, }))), created_at: Set(chrono::Utc::now()), ..Default::default() } .insert(&self.db) .await; let current = project_billing::Entity::find_by_id(project_uid) .one(&self.db) .await? .ok_or_else(|| { AppError::NotFound("Project billing account not found".to_string()) })?; let current_balance: f64 = current.balance.to_string().parse().unwrap_or(0.0); if current_balance < total_cost { return Err(AppError::BadRequest(format!( "Insufficient billing balance. Required: {:.4} {}, Available: {:.4} {}", total_cost, currency, current_balance, currency ))); } let new_balance = Decimal::from_f64_retain(current_balance - total_cost).unwrap_or(Decimal::ZERO); let mut updated: project_billing::ActiveModel = current.into(); updated.balance = Set(new_balance); updated.update(&self.db).await?; Ok(BillingRecord { cost: total_cost, currency, input_tokens, output_tokens, }) } } }