gitdataai/libs/agent/chat/session_recording.rs

use db::cache::AppCache;
use db::database::AppDatabase;
use sea_orm::{ActiveModelTrait, ColumnTrait, EntityTrait, QueryFilter, QueryOrder};
use uuid::Uuid;

/// Record an AI session with cost calculation.
pub async fn record_ai_session(
    cache: &AppCache,
    db: &AppDatabase,
    project_id: Uuid,
    user_id: Uuid,
    session_id: Uuid,
    room_id: Uuid,
    model_id: Uuid,
    version_id: Uuid,
    input_tokens: i64,
    output_tokens: i64,
    latency_ms: i64,
) {
    metrics::histogram!("ai_call_latency_ms", "model" => model_id.to_string())
        .record(latency_ms as f64);

    let session = models::ai::ai_session::ActiveModel {
        id: sea_orm::Set(session_id),
        room: sea_orm::Set(room_id),
        model: sea_orm::Set(model_id),
        version: sea_orm::Set(version_id),
        token_input: sea_orm::Set(input_tokens),
        token_output: sea_orm::Set(output_tokens),
        latency_ms: sea_orm::Set(Some(latency_ms)),
        cost: sea_orm::Set(None),
        currency: sea_orm::Set(None),
        error_message: sea_orm::Set(None),
        error_code: sea_orm::Set(None),
        created_at: sea_orm::Set(chrono::Utc::now()),
    };

    if let Err(e) = session.insert(db).await {
        tracing::error!(error = %e, session_id = %session_id, "failed to insert ai session record");
        return;
    }

    let (cost, currency, error_msg) = match crate::billing::record_ai_usage(
        db,
        project_id,
        user_id,
        version_id,
        input_tokens,
        output_tokens,
    )
    .await
    {
        Ok(crate::billing::BillingResult::Success(record)) => {
            (Some(record.cost), Some(record.currency), None)
        }
        Ok(crate::billing::BillingResult::InsufficientBalance { message }) => {
            create_billing_error_system_message(cache, db, room_id, &message).await;
            (None, None, Some(message))
        }
        Err(e) => (None, None, Some(e.to_string())),
    };

    use sea_orm::sea_query::Expr;
    let _ = models::ai::ai_session::Entity::update_many()
        .col_expr(models::ai::ai_session::Column::Cost, Expr::value(cost))
        .col_expr(
            models::ai::ai_session::Column::Currency,
            Expr::value(currency),
        )
        .col_expr(
            models::ai::ai_session::Column::ErrorMessage,
            Expr::value(error_msg),
        )
        .filter(models::ai::ai_session::Column::Id.eq(session_id))
        .exec(db)
        .await;
}

/// Create a system message in the room for billing errors.
async fn create_billing_error_system_message(
    cache: &AppCache,
    db: &AppDatabase,
    room_id: Uuid,
    message: &str,
) {
    use models::rooms::{MessageContentType, MessageSenderType, room_message};
    use sea_orm::Set;

    let seq_key = format!("seq:room:{}", room_id);
    let seq = match cache.conn().await {
        Ok(mut conn) => {
            match redis::cmd("INCR")
                .arg(&seq_key)
                .query_async::<i64>(&mut conn)
                .await
            {
                Ok(s) => s,
                Err(e) => {
                    tracing::warn!(error = %e, "cache INCR failed for system message seq, falling back to DB");
                    let last_seq = match room_message::Entity::find()
                        .filter(room_message::Column::Room.eq(room_id))
                        .order_by_desc(room_message::Column::Seq)
                        .one(db)
                        .await
                    {
                        Ok(Some(m)) => m.seq,
                        Ok(None) => 0,
                        Err(e) => {
                            tracing::warn!(error = %e, "Failed to get last seq for system message");
                            return;
                        }
                    };
                    last_seq + 1
                }
            }
        }
        Err(e) => {
            tracing::warn!(error = %e, "Failed to get Redis connection for system message seq");
            return;
        }
    };

    let now = chrono::Utc::now();
    let result = room_message::ActiveModel {
        id: Set(Uuid::now_v7()),
        seq: Set(seq),
        room: Set(room_id),
        sender_type: Set(MessageSenderType::System),
        sender_id: Set(None),
        model_id: Set(None),
        thread: Set(None),
        in_reply_to: Set(None),
        content: Set(message.to_string()),
        content_type: Set(MessageContentType::Text),
        thinking_content: Set(None),
        edited_at: Set(None),
        send_at: Set(now),
        revoked: Set(None),
        revoked_by: Set(None),
    }
    .insert(db)
    .await;

    match result {
        Ok(_) => {
            tracing::info!(room_id = %room_id, message = %message, "system_message_created_for_billing_error")
        }
        Err(e) => {
            tracing::warn!(error = %e, room_id = %room_id, "Failed to create system message for billing error")
        }
    }
}