gitdataai/libs/agent/billing.rs

//! AI usage billing — records token costs against a project or workspace balance.
//!
//! All functions take `&DatabaseConnection` instead of `&AppService`.

use db::database::AppDatabase;
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 uuid::Uuid;

use crate::error::AgentError;

#[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,
}

/// Extended result that includes insufficient balance flag for system message creation.
#[derive(Debug)]
pub enum BillingResult {
    Success(BillingRecord),
    InsufficientBalance { message: String },
}

/// Record AI usage for a project with cascading billing.
///
/// Billing strategy:
/// 1. Try to deduct from project balance first
/// 2. If insufficient, fallback to workspace balance (if project belongs to workspace)
/// 3. If both insufficient or no workspace, return InsufficientBalance error with room_id
///
/// Returns BillingError::InsufficientBalance with room_id for system message creation.
pub async fn record_ai_usage(
    db: &AppDatabase,
    project_uid: Uuid,
    model_id: Uuid,
    input_tokens: i64,
    output_tokens: i64,
) -> Result<BillingResult, AgentError> {
    // 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(db)
        .await?
        .ok_or_else(|| {
            AgentError::Internal(
                "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()
        .map_err(|e| AgentError::Internal(format!("Invalid input price format: {}", e)))?;
    let output_price: Decimal = pricing
        .output_price_per_1k_tokens
        .parse()
        .map_err(|e| AgentError::Internal(format!("Invalid output price format: {}", e)))?;
    let tokens_i = Decimal::from(input_tokens);
    let tokens_o = Decimal::from(output_tokens);
    let thousand = Decimal::from(1000);

    let total_cost = (tokens_i / thousand) * input_price
        + (tokens_o / thousand) * output_price;

    let currency = pricing.currency.clone();

    // 3. Cascading billing: project balance first, then workspace if insufficient.
    let proj = project::Entity::find_by_id(project_uid)
        .one(db)
        .await?
        .ok_or_else(|| AgentError::Internal("Project not found".into()))?;

    let txn = db.begin().await?;

    // Always check project balance first
    let project_billing = project_billing::Entity::find_by_id(project_uid)
        .lock_exclusive()
        .one(&txn)
        .await?
        .ok_or_else(|| AgentError::Internal("Project billing account not found".into()))?;

    let now = chrono::Utc::now();

    if project_billing.balance >= total_cost {
        // ── Project has sufficient balance ──────────────────────────
        let amount_dec = -total_cost;

        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(now),
            ..Default::default()
        }
        .insert(&txn)
        .await?;

        let new_balance = project_billing.balance - total_cost;
        let mut updated: project_billing::ActiveModel = project_billing.into();
        updated.balance = Set(new_balance);
        updated.updated_at = Set(now);
        updated.update(&txn).await?;

        txn.commit().await?;

        let cost_f64 = total_cost.to_string().parse().unwrap_or(0.0);

        tracing::info!(
            project_id = %project_uid,
            model_id = %model_id,
            input_tokens = input_tokens,
            output_tokens = output_tokens,
            cost = %cost_f64,
            currency = %currency,
            source = "project",
            "ai_usage_recorded"
        );

        Ok(BillingResult::Success(BillingRecord {
            cost: cost_f64,
            currency,
            input_tokens,
            output_tokens,
        }))
    } else if let Some(workspace_id) = proj.workspace_id {
        // ── Project insufficient, fallback to workspace ─────────────
        let workspace_billing = workspace_billing::Entity::find_by_id(workspace_id)
            .lock_exclusive()
            .one(&txn)
            .await?
            .ok_or_else(|| AgentError::Internal("Workspace billing account not found".into()))?;

        if workspace_billing.balance < total_cost {
            txn.rollback().await?;
            return Ok(BillingResult::InsufficientBalance {
                message: format!(
                    "Insufficient balance. Project: {:.4} {}, Workspace: {:.4} {}, Required: {:.4} {}",
                    project_billing.balance, currency,
                    workspace_billing.balance, currency,
                    total_cost, currency
                ),
            });
        }

        let amount_dec = -total_cost;

        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,
                "fallback_reason": "project_balance_insufficient"
            }))),
            created_at: Set(now),
        }
        .insert(&txn)
        .await?;

        let new_balance = workspace_billing.balance - total_cost;
        let new_total_spent = workspace_billing.total_spent + total_cost;
        let mut updated: workspace_billing::ActiveModel = workspace_billing.into();
        updated.balance = Set(new_balance);
        updated.total_spent = Set(new_total_spent);
        updated.updated_at = Set(now);
        updated.update(&txn).await?;

        txn.commit().await?;

        let cost_f64 = total_cost.to_string().parse().unwrap_or(0.0);

        tracing::info!(
            project_id = %project_uid,
            model_id = %model_id,
            input_tokens = input_tokens,
            output_tokens = output_tokens,
            cost = %cost_f64,
            currency = %currency,
            workspace_id = %workspace_id.to_string(),
            source = "workspace_fallback",
            "ai_usage_recorded"
        );

        Ok(BillingResult::Success(BillingRecord {
            cost: cost_f64,
            currency,
            input_tokens,
            output_tokens,
        }))
    } else {
        // ── Project insufficient and no workspace ───────────────────
        txn.rollback().await?;
        Ok(BillingResult::InsufficientBalance {
            message: format!(
                "Insufficient balance. Required: {:.4} {}, Available: {:.4} {}",
                total_cost, currency, project_billing.balance, currency
            ),
        })
    }
}