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,
    // Optional: `usage_events.run_id` references `agent_runs`, and a topology
    // turn has no row there — its id lives in `topology_runs`. Passing that id
    // was a foreign-key violation, so mission usage went unrecorded. NULL is
    // the honest value for a charge that is not an agent_run.
    run_id: Option<Uuid>,
    input_tokens: u64,
    output_tokens: u64,
    // Who was paid, and for which mission. `None` where the executor did not
    // say. Until 2026-09-14 no agent-side row carried these, so the larger
    // half of the spend could not be asked per provider — the judge's half
    // could, and that is how a plan emptying twice went unexplained.
    provider: Option<&str>,
    model: Option<&str>,
    mission_id: Option<Uuid>,
) -> Result<i64, BillingError> {
    let owed = credits_for_tokens(input_tokens + output_tokens);
    let mut tx = pool.begin().await?;

    // `sqlx::query`, not `query!`: the macro pins this statement to offline
    // metadata that a schema change then has to regenerate against a live
    // database, and the columns added by migration 0085 are nullable text
    // and uuid — nothing here that a compile-time check would catch.
    sqlx::query(
        "INSERT INTO usage_events
            (workspace_id, agent_id, run_id, kind, tokens_in, tokens_out, credits,
             provider, model, mission_id)
         VALUES ($1, $2, $3, 'llm_tokens', $4, $5, $6, $7, $8, $9)",
    )
    .bind(workspace_id.as_uuid())
    .bind(agent_id.as_uuid())
    .bind(run_id)
    .bind(input_tokens as i64)
    .bind(output_tokens as i64)
    .bind(sqlx::types::BigDecimal::from(owed))
    .bind(provider)
    .bind(model)
    .bind(mission_id)
    .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))
}

#[cfg(test)]
mod pricing_tests {
    use super::*;

    /// Pricing is the one place a rounding slip bills a real person.
    ///
    /// Pure, cheap to test, and previously untested — the crate's four tests
    /// all exercise the database path, so the arithmetic underneath them was
    /// never checked directly.
    #[test]
    fn credits_round_up_and_never_charge_zero() {
        // A run that used tokens always costs at least one credit; charging
        // zero for real work is how usage silently stops being metered.
        assert_eq!(credits_for_tokens(1), 1);
        assert_eq!(credits_for_tokens(TOKENS_PER_CREDIT - 1), 1);
        assert_eq!(credits_for_tokens(TOKENS_PER_CREDIT), 1);
        // Round UP, not to nearest: one token into the next bracket is a
        // whole credit, which is the documented contract.
        assert_eq!(credits_for_tokens(TOKENS_PER_CREDIT + 1), 2);
        assert_eq!(credits_for_tokens(TOKENS_PER_CREDIT * 3), 3);
        assert_eq!(credits_for_tokens(TOKENS_PER_CREDIT * 3 + 1), 4);
    }

    /// Zero tokens is the odd case: the `.max(1)` floor means it still costs a
    /// credit. That is deliberate, and worth pinning so a future "fix" to it
    /// is a decision rather than an accident.
    #[test]
    fn a_zero_token_run_still_costs_one_credit() {
        assert_eq!(credits_for_tokens(0), 1);
    }

    /// The cast to i64 must not wrap into a negative charge — a negative
    /// credit is a refund, and a refund granted by an overflow is the worst
    /// shape this bug could take.
    #[test]
    fn an_absurd_token_count_does_not_wrap_negative() {
        assert!(credits_for_tokens(u64::MAX / 2) > 0);
        assert!(credits_for_tokens(u64::MAX) > 0);
    }
}