clawmates/crates/cm-billing/src/lib.rs

//! Credit metering (§8.4 / §16 NFR). One credit pays for 1,000 LLM tokens
//! (rounded up, minimum one per charged run). Lots are consumed oldest
//! first and never expire; a workspace can run dry but never goes
//! negative — usage is always recorded in full either way.

use cm_domain::{AgentId, WorkspaceId};
use sqlx::PgPool;
use uuid::Uuid;

pub const TOKENS_PER_CREDIT: u64 = 1000;

#[derive(Debug, thiserror::Error)]
pub enum BillingError {
    #[error("promo code is invalid or already redeemed")]
    PromoUnavailable,
    #[error(transparent)]
    Db(#[from] sqlx::Error),
}

/// Credits owed for a token count: ceil(tokens / 1000), at least 1.
pub fn credits_for_tokens(total_tokens: u64) -> i64 {
    (total_tokens.div_ceil(TOKENS_PER_CREDIT).max(1)) as i64
}

/// Records a run's token usage and decrements credit lots oldest-first.
/// Returns the credits actually deducted (clamped at the available
/// balance; the usage event always records the full obligation).
pub async fn charge(
    pool: &PgPool,
    workspace_id: WorkspaceId,
    agent_id: AgentId,
    run_id: Uuid,
    input_tokens: u64,
    output_tokens: u64,
) -> Result<i64, BillingError> {
    let owed = credits_for_tokens(input_tokens + output_tokens);
    let mut tx = pool.begin().await?;

    sqlx::query!(
        "INSERT INTO usage_events
            (workspace_id, agent_id, run_id, kind, tokens_in, tokens_out, credits)
         VALUES ($1, $2, $3, 'llm_tokens', $4, $5, $6)",
        workspace_id.as_uuid(),
        agent_id.as_uuid(),
        run_id,
        input_tokens as i64,
        output_tokens as i64,
        sqlx::types::BigDecimal::from(owed),
    )
    .execute(&mut *tx)
    .await?;

    // Oldest lots first, locked so concurrent charges serialize.
    let lots = sqlx::query!(
        "SELECT id, remaining FROM credit_lots
         WHERE workspace_id = $1 AND remaining > 0
         ORDER BY purchased_at, id
         FOR UPDATE",
        workspace_id.as_uuid(),
    )
    .fetch_all(&mut *tx)
    .await?;

    let mut left = owed;
    for lot in lots {
        if left == 0 {
            break;
        }
        let available: i64 = lot.remaining.with_scale(0).to_string().parse().unwrap_or(0);
        let take = left.min(available);
        sqlx::query!(
            "UPDATE credit_lots SET remaining = remaining - $2 WHERE id = $1",
            lot.id,
            sqlx::types::BigDecimal::from(take),
        )
        .execute(&mut *tx)
        .await?;
        left -= take;
    }
    tx.commit().await?;
    Ok(owed - left)
}

/// Redeems a promo code exactly once (CAS) and grants its credits as a
/// new lot.
pub async fn redeem_promo(
    pool: &PgPool,
    workspace_id: WorkspaceId,
    code: &str,
) -> Result<i64, BillingError> {
    let mut tx = pool.begin().await?;
    let promo = sqlx::query!(
        "UPDATE promo_codes
         SET redeemed_by = $2, redeemed_at = now()
         WHERE code = $1 AND redeemed_by IS NULL
         RETURNING credits",
        code,
        workspace_id.as_uuid(),
    )
    .fetch_optional(&mut *tx)
    .await?
    .ok_or(BillingError::PromoUnavailable)?;
    let credits: i64 = promo.credits.with_scale(0).to_string().parse().unwrap_or(0);

    sqlx::query!(
        "INSERT INTO credit_lots (id, workspace_id, amount, remaining, source)
         VALUES ($1, $2, $3, $3, $4)",
        Uuid::now_v7(),
        workspace_id.as_uuid(),
        promo.credits,
        format!("promo:{code}"),
    )
    .execute(&mut *tx)
    .await?;
    tx.commit().await?;
    Ok(credits)
}

/// Aggregate usage over the trailing seven days (Credits page meter).
pub async fn usage_last_7_days(
    pool: &PgPool,
    workspace_id: WorkspaceId,
) -> Result<(i64, i64, i64), BillingError> {
    let row = sqlx::query!(
        r#"SELECT COALESCE(SUM(tokens_in), 0)::BIGINT AS "tokens_in!",
                  COALESCE(SUM(tokens_out), 0)::BIGINT AS "tokens_out!",
                  COALESCE(SUM(credits), 0)::BIGINT AS "credits!"
           FROM usage_events
           WHERE workspace_id = $1 AND created_at > now() - interval '7 days'"#,
        workspace_id.as_uuid(),
    )
    .fetch_one(pool)
    .await?;
    Ok((row.tokens_in, row.tokens_out, row.credits))
}