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

/// 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(
    db: &AppDatabase,
    project_uid: Uuid,
    model_id: Uuid,
    input_tokens: i64,
    output_tokens: i64,
) -> Result<BillingRecord, 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. Determine whether to bill the project or its workspace.
    let proj = project::Entity::find_by_id(project_uid)
        .one(db)
        .await?
        .ok_or_else(|| AgentError::Internal("Project not found".into()))?;

    if let Some(workspace_id) = proj.workspace_id {
        // ── Workspace-shared quota ──────────────────────────────────
        let txn = db.begin().await?;

        // SELECT FOR UPDATE to prevent race conditions
        let current = workspace_billing::Entity::find_by_id(workspace_id)
            .lock_exclusive()
            .one(&txn)
            .await?
            .ok_or_else(|| AgentError::Internal("Workspace billing account not found".into()))?;

        // Validate balance before any modifications
        if current.balance < total_cost {
            txn.rollback().await?;
            return Err(AgentError::Internal(format!(
                "Insufficient workspace billing balance. Required: {:.4} {}, Available: {:.4} {}",
                total_cost, currency, current.balance, currency
            )));
        }

        let amount_dec = -total_cost;
        let now = chrono::Utc::now();

        // Insert workspace billing history AFTER validation
        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(&txn)
        .await?;

        // Deduct from workspace balance
        let new_balance = current.balance - total_cost;
        let mut updated: workspace_billing::ActiveModel = current.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,
            workspace_id = %workspace_id.to_string(),
            "ai_usage_recorded"
        );

        Ok(BillingRecord {
            cost: cost_f64,
            currency,
            input_tokens,
            output_tokens,
        })
    } else {
        // ── Project-owned quota ─────────────────────────────────────
        let txn = db.begin().await?;

        // SELECT FOR UPDATE to prevent race conditions
        let current = project_billing::Entity::find_by_id(project_uid)
            .lock_exclusive()
            .one(&txn)
            .await?
            .ok_or_else(|| AgentError::Internal("Project billing account not found".into()))?;

        // Validate balance before any modifications
        if current.balance < total_cost {
            txn.rollback().await?;
            return Err(AgentError::Internal(format!(
                "Insufficient billing balance. Required: {:.4} {}, Available: {:.4} {}",
                total_cost, currency, current.balance, currency
            )));
        }

        let amount_dec = -total_cost;
        let now = chrono::Utc::now();

        // Insert project billing history AFTER validation
        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?;

        // Deduct from project balance
        let new_balance = current.balance - total_cost;
        let mut updated: project_billing::ActiveModel = current.into();
        updated.balance = Set(new_balance);
        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,
            "ai_usage_recorded"
        );

        Ok(BillingRecord {
            cost: cost_f64,
            currency,
            input_tokens,
            output_tokens,
        })
    }
}