synops/maskinrommet/src/ai_budget.rs

// AI-budsjettsjekk — sjekker ai_budget-metadata mot ai_usage_log aggregat.
//
// Samlings- og brukernoder kan ha budsjettgrense i metadata:
//   { "ai_budget": { "monthly_limit_usd": 50.0 } }
//
// Sjekken aggregerer denne månedens forbruk fra ai_usage_log og
// sammenligner med grensen. Ved overskridelse returneres feil
// og en work_item-node opprettes.
//
// Ref: docs/infra/ai_gateway.md § 6.3, oppgave 28.3

use sqlx::PgPool;
use uuid::Uuid;

/// Estimerer kostnad fra tokens (samme formel som metrics.rs).
/// Input: $3/MTok, Output: $15/MTok (konservativt gjennomsnitt).
pub fn estimate_cost_usd(prompt_tokens: i64, completion_tokens: i64) -> f64 {
    let input_cost = prompt_tokens as f64 * 3.0 / 1_000_000.0;
    let output_cost = completion_tokens as f64 * 15.0 / 1_000_000.0;
    ((input_cost + output_cost) * 100.0).round() / 100.0
}

/// Resultat av en budsjettsjekk.
pub struct BudgetStatus {
    pub current_cost: f64,
    pub limit: f64,
    pub node_title: String,
}

/// Sjekker ai_budget for en samlings-node.
///
/// Returnerer Ok(()) hvis budsjettet ikke er overskredet eller ikke er satt.
/// Returnerer Err med BudgetStatus hvis overskredet.
pub async fn check_collection_budget(
    db: &PgPool,
    collection_id: Uuid,
) -> Result<(), BudgetStatus> {
    check_budget_for_node(db, collection_id, "collection_node_id").await
}

/// Sjekker ai_budget for en bruker-node.
pub async fn check_user_budget(
    db: &PgPool,
    user_id: Uuid,
) -> Result<(), BudgetStatus> {
    check_budget_for_node(db, user_id, "requested_by").await
}

async fn check_budget_for_node(
    db: &PgPool,
    node_id: Uuid,
    column: &str,
) -> Result<(), BudgetStatus> {
    // Hent metadata fra noden
    let metadata: Option<serde_json::Value> =
        match sqlx::query_scalar("SELECT metadata FROM nodes WHERE id = $1")
            .bind(node_id)
            .fetch_optional(db)
            .await
        {
            Ok(m) => m,
            Err(e) => {
                tracing::warn!(error = %e, "Kunne ikke hente metadata for budsjettsjekk");
                return Ok(()); // Fail open — ikke blokker AI-kall ved DB-feil
            }
        };

    let monthly_limit = metadata
        .as_ref()
        .and_then(|m| m.get("ai_budget"))
        .and_then(|b| b.get("monthly_limit_usd"))
        .and_then(|v| v.as_f64());

    let limit = match monthly_limit {
        Some(l) => l,
        None => return Ok(()), // Ingen budsjettgrense satt
    };

    // Aggreger denne månedens forbruk
    // Bruker dynamisk kolonne via to separate queries for sikkerhet (unngår SQL injection)
    let (total_prompt, total_completion): (i64, i64) = match column {
        "collection_node_id" => {
            match sqlx::query_as::<_, (i64, i64)>(
                r#"SELECT
                    COALESCE(SUM(prompt_tokens)::BIGINT, 0),
                    COALESCE(SUM(completion_tokens)::BIGINT, 0)
                FROM ai_usage_log
                WHERE collection_node_id = $1
                  AND created_at >= date_trunc('month', now())"#,
            )
            .bind(node_id)
            .fetch_one(db)
            .await
            {
                Ok(r) => r,
                Err(e) => {
                    tracing::warn!(error = %e, "Kunne ikke aggregere AI-forbruk");
                    return Ok(()); // Fail open
                }
            }
        }
        "requested_by" => {
            match sqlx::query_as::<_, (i64, i64)>(
                r#"SELECT
                    COALESCE(SUM(prompt_tokens)::BIGINT, 0),
                    COALESCE(SUM(completion_tokens)::BIGINT, 0)
                FROM ai_usage_log
                WHERE requested_by = $1
                  AND created_at >= date_trunc('month', now())"#,
            )
            .bind(node_id)
            .fetch_one(db)
            .await
            {
                Ok(r) => r,
                Err(e) => {
                    tracing::warn!(error = %e, "Kunne ikke aggregere bruker-AI-forbruk");
                    return Ok(()); // Fail open
                }
            }
        }
        _ => return Ok(()),
    };

    let current_cost = estimate_cost_usd(total_prompt, total_completion);

    if current_cost >= limit {
        let node_title: String = sqlx::query_scalar("SELECT title FROM nodes WHERE id = $1")
            .bind(node_id)
            .fetch_optional(db)
            .await
            .ok()
            .flatten()
            .unwrap_or_else(|| node_id.to_string());

        return Err(BudgetStatus {
            current_cost,
            limit,
            node_title,
        });
    }

    tracing::debug!(
        node_id = %node_id,
        current_cost,
        limit,
        "Budsjettsjekk OK"
    );

    Ok(())
}

/// Oppretter work_item-node når budsjett er overskredet.
pub async fn create_budget_work_item(
    db: &PgPool,
    collection_id: Uuid,
    created_by: Option<Uuid>,
    status: &BudgetStatus,
) {
    let work_item_id = Uuid::now_v7();
    let title = format!("AI-budsjett overskredet: {}", status.node_title);
    let content = format!(
        "AI-budsjettet for \"{}\" er overskredet.\n\
         Forbruk denne måneden: ${:.2}\n\
         Budsjettgrense: ${:.2}\n\n\
         AI-kall er blokkert inntil budsjettet økes eller ny måned starter.",
        status.node_title, status.current_cost, status.limit
    );
    let creator = created_by.unwrap_or(collection_id);

    let metadata = serde_json::json!({
        "work_item_type": "budget_exceeded",
        "collection_id": collection_id.to_string(),
        "current_cost_usd": status.current_cost,
        "monthly_limit_usd": status.limit,
    });

    if let Err(e) = sqlx::query(
        r#"INSERT INTO nodes (id, node_kind, title, content, visibility, metadata, created_by)
        VALUES ($1, 'content', $2, $3, 'hidden'::visibility, $4, $5)"#,
    )
    .bind(work_item_id)
    .bind(&title)
    .bind(&content)
    .bind(&metadata)
    .bind(creator)
    .execute(db)
    .await
    {
        tracing::warn!(error = %e, "Kunne ikke opprette work_item for budsjettoverskridelse");
        return;
    }

    // tagged-edge: budget_exceeded
    let tag_edge_id = Uuid::now_v7();
    if let Err(e) = sqlx::query(
        r#"INSERT INTO edges (id, source_id, target_id, edge_type, metadata, system, created_by)
        VALUES ($1, $2, $2, 'tagged', '{"tag": "budget_exceeded"}'::jsonb, true, $3)"#,
    )
    .bind(tag_edge_id)
    .bind(work_item_id)
    .bind(creator)
    .execute(db)
    .await
    {
        tracing::warn!(error = %e, "Kunne ikke opprette tagged-edge for work_item");
    }

    // belongs_to-edge til samlingen
    let belongs_edge_id = Uuid::now_v7();
    if let Err(e) = sqlx::query(
        r#"INSERT INTO edges (id, source_id, target_id, edge_type, metadata, system, created_by)
        VALUES ($1, $2, $3, 'belongs_to', '{}'::jsonb, true, $4)"#,
    )
    .bind(belongs_edge_id)
    .bind(work_item_id)
    .bind(collection_id)
    .bind(creator)
    .execute(db)
    .await
    {
        tracing::warn!(error = %e, "Kunne ikke opprette belongs_to-edge for work_item");
    }

    tracing::info!(
        work_item_id = %work_item_id,
        collection_id = %collection_id,
        current_cost = status.current_cost,
        limit = status.limit,
        "Work item opprettet for budsjettoverskridelse"
    );
}