gitdataai/libs/service/agent/billing.rs

//! 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<BillingRecord, AppError> {
        // 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,
            })
        }
    }
}