use crate::AppService; use crate::error::AppError; use chrono::{DateTime, Datelike, NaiveDate, Utc}; use models::Decimal; use models::workspaces::{workspace_billing, workspace_billing_history, workspace_membership}; use sea_orm::sea_query::prelude::rust_decimal::prelude::ToPrimitive; use sea_orm::*; use serde::{Deserialize, Serialize}; use session::Session; use utoipa::{IntoParams, ToSchema}; use uuid::Uuid; /// Default monthly AI quota for workspace (shared across all its projects). const DEFAULT_MONTHLY_QUOTA: f64 = 100.0; #[derive(Debug, Clone, Serialize, Deserialize, ToSchema)] pub struct WorkspaceBillingCurrentResponse { pub workspace_id: Uuid, pub currency: String, pub monthly_quota: f64, pub balance: f64, pub total_spent: f64, pub month_used: f64, pub cycle_start_utc: DateTime, pub cycle_end_utc: DateTime, pub updated_at: DateTime, pub created_at: DateTime, } #[derive(Debug, Clone, Serialize, Deserialize, ToSchema, IntoParams)] pub struct WorkspaceBillingHistoryQuery { pub page: Option, pub per_page: Option, } #[derive(Debug, Clone, Serialize, Deserialize, ToSchema)] pub struct WorkspaceBillingHistoryItem { pub uid: Uuid, pub workspace_id: Uuid, pub user_id: Option, pub amount: f64, pub currency: String, pub reason: String, pub extra: Option, pub created_at: DateTime, } #[derive(Debug, Clone, Serialize, Deserialize, ToSchema)] pub struct WorkspaceBillingHistoryResponse { pub page: u64, pub per_page: u64, pub total: u64, pub list: Vec, } #[derive(Debug, Clone, Serialize, Deserialize, ToSchema)] pub struct WorkspaceBillingAddCreditParams { pub amount: f64, pub reason: Option, } impl AppService { /// Get current workspace billing info. pub async fn workspace_billing_current( &self, ctx: &Session, workspace_slug: String, ) -> Result { let user_uid = ctx.user().ok_or(AppError::Unauthorized)?; let ws = self .utils_find_workspace_by_slug(workspace_slug.clone()) .await?; let _ = workspace_membership::Entity::find() .filter(workspace_membership::Column::WorkspaceId.eq(ws.id)) .filter(workspace_membership::Column::UserId.eq(user_uid)) .filter(workspace_membership::Column::Status.eq("active")) .one(&self.db) .await? .ok_or(AppError::NotWorkspaceMember)?; let billing = self.ensure_workspace_billing(ws.id).await?; let now_utc = Utc::now(); let (month_start, next_month_start) = utc_month_bounds(now_utc)?; let month_used = workspace_billing_history::Entity::find() .filter(workspace_billing_history::Column::WorkspaceId.eq(ws.id)) .filter(workspace_billing_history::Column::Reason.like("ai_usage%")) .filter(workspace_billing_history::Column::CreatedAt.gte(month_start)) .filter(workspace_billing_history::Column::CreatedAt.lt(next_month_start)) .all(&self.db) .await? .into_iter() .map(|m| m.amount.to_f64().unwrap_or_default()) .sum::(); Ok(WorkspaceBillingCurrentResponse { workspace_id: ws.id, currency: billing.currency.clone(), monthly_quota: billing .monthly_quota .to_f64() .unwrap_or(DEFAULT_MONTHLY_QUOTA), balance: billing.balance.to_f64().unwrap_or_default(), total_spent: billing.total_spent.to_f64().unwrap_or_default(), month_used, cycle_start_utc: month_start, cycle_end_utc: next_month_start, updated_at: billing.updated_at, created_at: billing.created_at, }) } /// Get workspace billing history. pub async fn workspace_billing_history( &self, ctx: &Session, workspace_slug: String, query: WorkspaceBillingHistoryQuery, ) -> Result { let user_uid = ctx.user().ok_or(AppError::Unauthorized)?; let ws = self .utils_find_workspace_by_slug(workspace_slug.clone()) .await?; let _ = workspace_membership::Entity::find() .filter(workspace_membership::Column::WorkspaceId.eq(ws.id)) .filter(workspace_membership::Column::UserId.eq(user_uid)) .filter(workspace_membership::Column::Status.eq("active")) .one(&self.db) .await? .ok_or(AppError::NotWorkspaceMember)?; let page = std::cmp::max(query.page.unwrap_or(1), 1); let per_page = query.per_page.unwrap_or(20).clamp(1, 200); self.ensure_workspace_billing(ws.id).await?; let paginator = workspace_billing_history::Entity::find() .filter(workspace_billing_history::Column::WorkspaceId.eq(ws.id)) .order_by_desc(workspace_billing_history::Column::CreatedAt) .paginate(&self.db, per_page); let total = paginator.num_items().await?; let rows = paginator.fetch_page(page - 1).await?; let list = rows .into_iter() .map(|x| WorkspaceBillingHistoryItem { uid: x.uid, workspace_id: x.workspace_id, user_id: x.user_id, amount: x.amount.to_f64().unwrap_or_default(), currency: x.currency, reason: x.reason, extra: x.extra.map(|v| v.into()), created_at: x.created_at, }) .collect(); Ok(WorkspaceBillingHistoryResponse { page, per_page, total, list, }) } /// Add credit to workspace billing (admin action). pub async fn workspace_billing_add_credit( &self, ctx: &Session, workspace_slug: String, params: WorkspaceBillingAddCreditParams, ) -> Result { let user_uid = ctx.user().ok_or(AppError::Unauthorized)?; let ws = self .utils_find_workspace_by_slug(workspace_slug.clone()) .await?; let _ = workspace_membership::Entity::find() .filter(workspace_membership::Column::WorkspaceId.eq(ws.id)) .filter(workspace_membership::Column::UserId.eq(user_uid)) .filter(workspace_membership::Column::Status.eq("active")) .one(&self.db) .await? .ok_or(AppError::NotWorkspaceMember)?; if params.amount <= 0.0 { return Err(AppError::BadRequest("Amount must be positive".to_string())); } let billing = self.ensure_workspace_billing(ws.id).await?; let now_utc = Utc::now(); let new_balance = Decimal::from_f64_retain(billing.balance.to_f64().unwrap_or_default() + params.amount) .unwrap_or(Decimal::ZERO); let _ = workspace_billing::ActiveModel { workspace_id: Unchanged(ws.id), balance: Set(new_balance), updated_at: Set(now_utc), ..Default::default() } .update(&self.db) .await; let _ = workspace_billing_history::ActiveModel { uid: Set(Uuid::now_v7()), workspace_id: Set(ws.id), user_id: Set(Some(user_uid)), amount: Set(Decimal::from_f64_retain(params.amount).unwrap_or(Decimal::ZERO)), currency: Set(billing.currency.clone()), reason: Set(params.reason.unwrap_or_else(|| "credit_added".to_string())), extra: Set(None), created_at: Set(now_utc), } .insert(&self.db) .await; self.workspace_billing_current(ctx, workspace_slug).await } /// Ensure workspace billing record exists (create with defaults if not). pub async fn ensure_workspace_billing( &self, workspace_id: Uuid, ) -> Result { if let Some(billing) = workspace_billing::Entity::find_by_id(workspace_id) .one(&self.db) .await? { return Ok(billing); } let now_utc = Utc::now(); let created = workspace_billing::ActiveModel { workspace_id: Set(workspace_id), balance: Set(Decimal::ZERO), currency: Set("USD".to_string()), monthly_quota: Set( Decimal::from_f64_retain(DEFAULT_MONTHLY_QUOTA).unwrap_or(Decimal::ZERO) ), total_spent: Set(Decimal::ZERO), updated_at: Set(now_utc), created_at: Set(now_utc), }; Ok(created.insert(&self.db).await?) } } fn utc_month_bounds(now_utc: DateTime) -> Result<(DateTime, DateTime), AppError> { let year = now_utc.year(); let month = now_utc.month(); let month_start = NaiveDate::from_ymd_opt(year, month, 1) .and_then(|d| d.and_hms_opt(0, 0, 0)) .map(|d| chrono::TimeZone::from_utc_datetime(&Utc, &d)) .ok_or_else(|| AppError::InternalServerError("Invalid UTC month start".to_string()))?; let (next_year, next_month) = if month == 12 { (year + 1, 1) } else { (year, month + 1) }; let next_month_start = NaiveDate::from_ymd_opt(next_year, next_month, 1) .and_then(|d| d.and_hms_opt(0, 0, 0)) .map(|d| chrono::TimeZone::from_utc_datetime(&Utc, &d)) .ok_or_else(|| AppError::InternalServerError("Invalid UTC next month start".to_string()))?; Ok((month_start, next_month_start)) }