669 lines
21 KiB
Rust
669 lines
21 KiB
Rust
//! 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<BillingResult, AgentError> {
|
|
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<BillingResult, AgentError> {
|
|
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<bool, AgentError> {
|
|
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<bool, AgentError> {
|
|
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<Decimal, AgentError> {
|
|
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<String, AgentError> {
|
|
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<serde_json::Value>,
|
|
) -> 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)
|
|
}
|