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::(&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") } } }