agent-orchestrator/agor-pro/src/budget.rs

// SPDX-License-Identifier: LicenseRef-Commercial
// Budget Governor — per-project monthly token budgets with soft/hard limits.

use rusqlite::params;
use serde::Serialize;

#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct BudgetStatus {
    pub project_id: String,
    pub limit: i64,
    pub used: i64,
    pub remaining: i64,
    pub percent: f64,
    pub reset_date: i64,
}

#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct BudgetDecision {
    pub allowed: bool,
    pub reason: String,
    pub remaining: i64,
}

#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct BudgetEntry {
    pub project_id: String,
    pub monthly_limit_tokens: i64,
    pub used_tokens: i64,
    pub reset_date: i64,
}

fn ensure_tables(conn: &rusqlite::Connection) -> Result<(), String> {
    conn.execute_batch(
        "CREATE TABLE IF NOT EXISTS pro_budgets (
            project_id TEXT PRIMARY KEY,
            monthly_limit_tokens INTEGER NOT NULL,
            used_tokens INTEGER NOT NULL DEFAULT 0,
            reset_date INTEGER NOT NULL
        );
        CREATE TABLE IF NOT EXISTS pro_budget_log (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            project_id TEXT NOT NULL,
            session_id TEXT NOT NULL,
            tokens_used INTEGER NOT NULL,
            timestamp INTEGER NOT NULL
        );
        CREATE INDEX IF NOT EXISTS idx_budget_log_project ON pro_budget_log(project_id, timestamp);"
    ).map_err(|e| format!("Failed to create budget tables: {e}"))
}

fn now_epoch() -> i64 {
    super::analytics::now_epoch()
}

/// Calculate reset date: first day of next calendar month as epoch.
fn next_month_epoch() -> i64 {
    let now = now_epoch();
    let days_since_epoch = now / 86400;
    let mut year = 1970i64;
    let mut remaining_days = days_since_epoch;
    loop {
        let days_in_year = if is_leap_year(year) { 366 } else { 365 };
        if remaining_days < days_in_year { break; }
        remaining_days -= days_in_year;
        year += 1;
    }
    let month_days: [i64; 12] = if is_leap_year(year) {
        [31,29,31,30,31,30,31,31,30,31,30,31]
    } else {
        [31,28,31,30,31,30,31,31,30,31,30,31]
    };
    let mut month = 0usize;
    for (i, &d) in month_days.iter().enumerate() {
        if remaining_days < d { month = i; break; }
        remaining_days -= d;
    }
    // Advance to first of next month
    let (next_year, next_month) = if month >= 11 {
        (year + 1, 0usize)
    } else {
        (year, month + 1)
    };
    // Calculate epoch for first of next_month in next_year
    let mut epoch: i64 = 0;
    for y in 1970..next_year {
        epoch += if is_leap_year(y) { 366 } else { 365 };
    }
    let nm_days: [i64; 12] = if is_leap_year(next_year) {
        [31,29,31,30,31,30,31,31,30,31,30,31]
    } else {
        [31,28,31,30,31,30,31,31,30,31,30,31]
    };
    for i in 0..next_month {
        epoch += nm_days[i];
    }
    epoch * 86400
}

fn is_leap_year(y: i64) -> bool {
    (y % 4 == 0 && y % 100 != 0) || y % 400 == 0
}

#[tauri::command]
pub fn pro_budget_set(project_id: String, monthly_limit_tokens: i64) -> Result<(), String> {
    if monthly_limit_tokens <= 0 {
        return Err("monthly_limit_tokens must be positive".into());
    }
    let conn = super::open_sessions_db()?;
    ensure_tables(&conn)?;
    let reset = next_month_epoch();
    conn.execute(
        "INSERT INTO pro_budgets (project_id, monthly_limit_tokens, used_tokens, reset_date)
         VALUES (?1, ?2, 0, ?3)
         ON CONFLICT(project_id) DO UPDATE SET monthly_limit_tokens = ?2",
        params![project_id, monthly_limit_tokens, reset],
    ).map_err(|e| format!("Failed to set budget: {e}"))?;
    Ok(())
}

#[tauri::command]
pub fn pro_budget_get(project_id: String) -> Result<BudgetStatus, String> {
    let conn = super::open_sessions_db()?;
    ensure_tables(&conn)?;
    auto_reset_if_expired(&conn, &project_id)?;

    let mut stmt = conn.prepare(
        "SELECT monthly_limit_tokens, used_tokens, reset_date FROM pro_budgets WHERE project_id = ?1"
    ).map_err(|e| format!("Query failed: {e}"))?;

    stmt.query_row(params![project_id], |row| {
        let limit: i64 = row.get("monthly_limit_tokens")?;
        let used: i64 = row.get("used_tokens")?;
        let reset_date: i64 = row.get("reset_date")?;
        let remaining = (limit - used).max(0);
        let percent = if limit > 0 { (used as f64 / limit as f64) * 100.0 } else { 0.0 };
        Ok(BudgetStatus { project_id: project_id.clone(), limit, used, remaining, percent, reset_date })
    }).map_err(|e| format!("Budget not found for project '{}': {e}", project_id))
}

#[tauri::command]
pub fn pro_budget_check(project_id: String, estimated_tokens: i64) -> Result<BudgetDecision, String> {
    if estimated_tokens < 0 {
        return Err("estimated_tokens must be non-negative".into());
    }
    let conn = super::open_sessions_db()?;
    ensure_tables(&conn)?;
    auto_reset_if_expired(&conn, &project_id)?;

    let result = conn.prepare(
        "SELECT monthly_limit_tokens, used_tokens FROM pro_budgets WHERE project_id = ?1"
    ).map_err(|e| format!("Query failed: {e}"))?
    .query_row(params![project_id], |row| {
        Ok((row.get::<_, i64>("monthly_limit_tokens")?, row.get::<_, i64>("used_tokens")?))
    });

    match result {
        Ok((limit, used)) => {
            let remaining = (limit - used).max(0);
            if used + estimated_tokens > limit {
                Ok(BudgetDecision {
                    allowed: false,
                    reason: format!("Would exceed budget: {} remaining, {} requested", remaining, estimated_tokens),
                    remaining,
                })
            } else {
                Ok(BudgetDecision { allowed: true, reason: "Within budget".into(), remaining })
            }
        }
        Err(_) => {
            // No budget set — allow by default
            Ok(BudgetDecision { allowed: true, reason: "No budget configured".into(), remaining: i64::MAX })
        }
    }
}

#[tauri::command]
pub fn pro_budget_log_usage(project_id: String, session_id: String, tokens_used: i64) -> Result<(), String> {
    let conn = super::open_sessions_db()?;
    ensure_tables(&conn)?;
    let ts = now_epoch();

    let tx = conn.unchecked_transaction()
        .map_err(|e| format!("Transaction failed: {e}"))?;
    tx.execute(
        "INSERT INTO pro_budget_log (project_id, session_id, tokens_used, timestamp) VALUES (?1, ?2, ?3, ?4)",
        params![project_id, session_id, tokens_used, ts],
    ).map_err(|e| format!("Failed to log usage: {e}"))?;
    tx.execute(
        "UPDATE pro_budgets SET used_tokens = used_tokens + ?2 WHERE project_id = ?1",
        params![project_id, tokens_used],
    ).map_err(|e| format!("Failed to update used tokens: {e}"))?;
    tx.commit().map_err(|e| format!("Commit failed: {e}"))?;
    Ok(())
}

#[tauri::command]
pub fn pro_budget_reset(project_id: String) -> Result<(), String> {
    let conn = super::open_sessions_db()?;
    ensure_tables(&conn)?;
    let reset = next_month_epoch();
    conn.execute(
        "UPDATE pro_budgets SET used_tokens = 0, reset_date = ?2 WHERE project_id = ?1",
        params![project_id, reset],
    ).map_err(|e| format!("Failed to reset budget: {e}"))?;
    Ok(())
}

#[tauri::command]
pub fn pro_budget_list() -> Result<Vec<BudgetEntry>, String> {
    let conn = super::open_sessions_db()?;
    ensure_tables(&conn)?;
    let mut stmt = conn.prepare(
        "SELECT project_id, monthly_limit_tokens, used_tokens, reset_date FROM pro_budgets ORDER BY project_id"
    ).map_err(|e| format!("Query failed: {e}"))?;

    let rows = stmt.query_map([], |row| {
        Ok(BudgetEntry {
            project_id: row.get("project_id")?,
            monthly_limit_tokens: row.get("monthly_limit_tokens")?,
            used_tokens: row.get("used_tokens")?,
            reset_date: row.get("reset_date")?,
        })
    }).map_err(|e| format!("Query failed: {e}"))?
    .collect::<Result<Vec<_>, _>>()
    .map_err(|e| format!("Row read failed: {e}"))?;

    Ok(rows)
}

fn auto_reset_if_expired(conn: &rusqlite::Connection, project_id: &str) -> Result<(), String> {
    let now = now_epoch();
    conn.execute(
        "UPDATE pro_budgets SET used_tokens = 0, reset_date = ?3
         WHERE project_id = ?1 AND reset_date < ?2",
        params![project_id, now, now + 30 * 86400],
    ).map_err(|e| format!("Auto-reset failed: {e}"))?;
    Ok(())
}

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

    #[test]
    fn test_budget_status_serializes_camel_case() {
        let s = BudgetStatus {
            project_id: "proj1".into(),
            limit: 100_000,
            used: 25_000,
            remaining: 75_000,
            percent: 25.0,
            reset_date: 1710000000,
        };
        let json = serde_json::to_string(&s).unwrap();
        assert!(json.contains("projectId"));
        assert!(json.contains("resetDate"));
        assert!(json.contains("\"remaining\":75000"));
    }

    #[test]
    fn test_budget_decision_serializes_camel_case() {
        let d = BudgetDecision {
            allowed: true,
            reason: "Within budget".into(),
            remaining: 50_000,
        };
        let json = serde_json::to_string(&d).unwrap();
        assert!(json.contains("\"allowed\":true"));
        assert!(json.contains("\"remaining\":50000"));
    }

    #[test]
    fn test_budget_entry_serializes_camel_case() {
        let e = BudgetEntry {
            project_id: "p".into(),
            monthly_limit_tokens: 200_000,
            used_tokens: 10_000,
            reset_date: 1710000000,
        };
        let json = serde_json::to_string(&e).unwrap();
        assert!(json.contains("monthlyLimitTokens"));
        assert!(json.contains("usedTokens"));
    }
}