use std::collections::{BTreeMap, HashSet}; use std::time::{SystemTime, UNIX_EPOCH}; use aether_ai_formats::UPSTREAM_IS_STREAM_KEY; use async_trait::async_trait; use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row}; use super::{ provider_api_key_usage_is_error, provider_api_key_usage_is_success, strip_deprecated_usage_display_fields, usage_can_recover_terminal_failure, usage_request_metadata_client_family, InMemoryUsageReadRepository, PendingUsageCleanupSummary, StoredRequestUsageAudit, StoredUsageDailySummary, StoredUsageDashboardDailyBreakdownRow, StoredUsageDashboardSummary, StoredUsageUserTotals, UpsertUsageRecord, UsageDailyHeatmapQuery, UsageDashboardDailyBreakdownQuery, UsageDashboardSummaryQuery, UsageReadRepository, UsageWriteRepository, }; use crate::driver::mysql::MysqlPool; use crate::error::SqlResultExt; use crate::DataLayerError; const USAGE_COLUMNS: &str = r#" SELECT id, request_id, user_id, api_key_id, provider_name, model, target_model, provider_id, provider_endpoint_id, provider_api_key_id, request_type, api_format, api_family, endpoint_kind, endpoint_api_format, provider_api_family, provider_endpoint_kind, has_format_conversion, is_stream, upstream_is_stream, input_tokens, output_tokens, total_tokens, cache_creation_input_tokens, cache_creation_ephemeral_5m_input_tokens, cache_creation_ephemeral_1h_input_tokens, cache_read_input_tokens, cache_creation_cost_usd, cache_read_cost_usd, output_price_per_1m, total_cost_usd, actual_total_cost_usd, status_code, error_message, error_category, response_time_ms, first_byte_time_ms, status, billing_status, request_metadata, candidate_id, candidate_index, key_name, planner_kind, route_family, route_kind, execution_path, local_execution_runtime_miss_reason, finalized_at AS finalized_at_unix_secs, created_at_unix_ms, updated_at_unix_secs FROM `usage` "#; const UPSERT_USAGE_SQL: &str = r#" INSERT INTO `usage` ( request_id, id, user_id, api_key_id, provider_name, model, target_model, provider_id, provider_endpoint_id, provider_api_key_id, request_type, api_format, api_family, endpoint_kind, endpoint_api_format, provider_api_family, provider_endpoint_kind, has_format_conversion, is_stream, upstream_is_stream, input_tokens, output_tokens, total_tokens, cache_creation_input_tokens, cache_creation_ephemeral_5m_input_tokens, cache_creation_ephemeral_1h_input_tokens, cache_read_input_tokens, cache_creation_cost_usd, cache_read_cost_usd, output_price_per_1m, total_cost_usd, actual_total_cost_usd, status_code, error_message, error_category, response_time_ms, first_byte_time_ms, status, billing_status, request_metadata, candidate_id, candidate_index, key_name, planner_kind, route_family, route_kind, execution_path, local_execution_runtime_miss_reason, finalized_at, created_at_unix_ms, updated_at_unix_secs ) VALUES ( ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ? ) ON DUPLICATE KEY UPDATE user_id = VALUES(user_id), api_key_id = VALUES(api_key_id), provider_name = VALUES(provider_name), model = VALUES(model), target_model = VALUES(target_model), provider_id = VALUES(provider_id), provider_endpoint_id = VALUES(provider_endpoint_id), provider_api_key_id = VALUES(provider_api_key_id), request_type = VALUES(request_type), api_format = VALUES(api_format), api_family = VALUES(api_family), endpoint_kind = VALUES(endpoint_kind), endpoint_api_format = VALUES(endpoint_api_format), provider_api_family = VALUES(provider_api_family), provider_endpoint_kind = VALUES(provider_endpoint_kind), has_format_conversion = VALUES(has_format_conversion), is_stream = VALUES(is_stream), upstream_is_stream = VALUES(upstream_is_stream), input_tokens = VALUES(input_tokens), output_tokens = VALUES(output_tokens), total_tokens = VALUES(total_tokens), cache_creation_input_tokens = VALUES(cache_creation_input_tokens), cache_creation_ephemeral_5m_input_tokens = VALUES(cache_creation_ephemeral_5m_input_tokens), cache_creation_ephemeral_1h_input_tokens = VALUES(cache_creation_ephemeral_1h_input_tokens), cache_read_input_tokens = VALUES(cache_read_input_tokens), cache_creation_cost_usd = VALUES(cache_creation_cost_usd), cache_read_cost_usd = VALUES(cache_read_cost_usd), output_price_per_1m = VALUES(output_price_per_1m), total_cost_usd = VALUES(total_cost_usd), actual_total_cost_usd = VALUES(actual_total_cost_usd), status_code = VALUES(status_code), error_message = VALUES(error_message), error_category = VALUES(error_category), response_time_ms = CASE WHEN VALUES(response_time_ms) IS NULL OR VALUES(response_time_ms) = 0 THEN COALESCE(response_time_ms, VALUES(response_time_ms)) ELSE VALUES(response_time_ms) END, first_byte_time_ms = CASE WHEN VALUES(first_byte_time_ms) IS NULL OR VALUES(first_byte_time_ms) = 0 THEN COALESCE(first_byte_time_ms, VALUES(first_byte_time_ms)) ELSE VALUES(first_byte_time_ms) END, status = VALUES(status), billing_status = VALUES(billing_status), request_metadata = VALUES(request_metadata), candidate_id = VALUES(candidate_id), candidate_index = VALUES(candidate_index), key_name = VALUES(key_name), planner_kind = VALUES(planner_kind), route_family = VALUES(route_family), route_kind = VALUES(route_kind), execution_path = VALUES(execution_path), local_execution_runtime_miss_reason = VALUES(local_execution_runtime_miss_reason), finalized_at = VALUES(finalized_at), updated_at_unix_secs = VALUES(updated_at_unix_secs) "#; const SELECT_STALE_PENDING_USAGE_BATCH_SQL: &str = r#" SELECT `usage`.request_id, `usage`.status, COALESCE(usage_settlement_snapshots.billing_status, `usage`.billing_status) AS billing_status FROM `usage` LEFT JOIN usage_settlement_snapshots ON usage_settlement_snapshots.request_id = `usage`.request_id WHERE `usage`.status IN ('pending', 'streaming') AND `usage`.created_at_unix_ms < ? ORDER BY `usage`.created_at_unix_ms ASC, `usage`.request_id ASC LIMIT ? "#; const SELECT_COMPLETED_REQUEST_CANDIDATES_SQL: &str = r#" SELECT status, extra_data FROM request_candidates WHERE request_id = ? AND status IN ('streaming', 'success') "#; #[derive(Debug, Clone)] pub struct MysqlUsageWriteRepository { pool: MysqlPool, } #[derive(Debug, Clone)] pub struct MysqlUsageReadRepository { pool: MysqlPool, } impl MysqlUsageReadRepository { pub fn new(pool: MysqlPool) -> Self { Self { pool } } async fn materialize_read_model(&self) -> Result { let rows = sqlx::query(&format!( "{USAGE_COLUMNS} ORDER BY created_at_unix_ms ASC, request_id ASC" )) .fetch_all(&self.pool) .await .map_sql_err()?; let items = rows .iter() .map(map_usage_row) .collect::, _>>()?; Ok(InMemoryUsageReadRepository::seed(items)) } async fn summarize_usage_daily_heatmap_raw_from_range( &self, created_from_unix_secs: u64, created_until_unix_secs: u64, user_id: Option<&str>, ) -> Result, DataLayerError> { let mut sql = String::from( r#" SELECT DATE_FORMAT(FROM_UNIXTIME(created_at_unix_ms), '%Y-%m-%d') AS date, CAST(COUNT(*) AS SIGNED) AS requests, CAST(COALESCE(SUM( GREATEST(COALESCE(input_tokens, 0), 0) + GREATEST(COALESCE(output_tokens, 0), 0) + CASE WHEN COALESCE(cache_creation_input_tokens, 0) = 0 AND (COALESCE(cache_creation_ephemeral_5m_input_tokens, 0) + COALESCE(cache_creation_ephemeral_1h_input_tokens, 0)) > 0 THEN COALESCE(cache_creation_ephemeral_5m_input_tokens, 0) + COALESCE(cache_creation_ephemeral_1h_input_tokens, 0) ELSE GREATEST(COALESCE(cache_creation_input_tokens, 0), 0) END + GREATEST(COALESCE(cache_read_input_tokens, 0), 0) ), 0) AS SIGNED) AS total_tokens, CAST(COALESCE(SUM(COALESCE(total_cost_usd, 0)), 0) AS DOUBLE) AS total_cost_usd, CAST(COALESCE(SUM(COALESCE(actual_total_cost_usd, 0)), 0) AS DOUBLE) AS actual_total_cost_usd FROM `usage` WHERE created_at_unix_ms >= ? AND created_at_unix_ms < ? AND status NOT IN ('pending', 'streaming') AND provider_name NOT IN ('unknown', 'pending') "#, ); if user_id.is_some() { sql.push_str(" AND user_id = ?\n"); } sql.push_str("GROUP BY date ORDER BY date ASC"); let mut query = sqlx::query(&sql) .bind(to_i64(created_from_unix_secs, "usage.created_at_unix_ms")?) .bind(to_i64(created_until_unix_secs, "usage.created_at_unix_ms")?); if let Some(user_id) = user_id { query = query.bind(user_id.to_string()); } let rows = query.fetch_all(&self.pool).await.map_sql_err()?; rows.iter().map(map_mysql_usage_daily_summary).collect() } async fn summarize_usage_daily_heatmap_from_daily_aggregates( &self, created_from_unix_secs: u64, created_until_unix_secs: u64, user_id: Option<&str>, ) -> Result, DataLayerError> { let rows = if let Some(user_id) = user_id { sqlx::query( r#" SELECT DATE_FORMAT(FROM_UNIXTIME(`date`), '%Y-%m-%d') AS date, total_requests AS requests, input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens AS total_tokens, total_cost AS total_cost_usd, total_cost AS actual_total_cost_usd FROM stats_user_daily WHERE user_id = ? AND `date` >= ? AND `date` < ? AND total_requests > 0 ORDER BY `date` ASC "#, ) .bind(user_id) .bind(to_i64(created_from_unix_secs, "stats_user_daily.date")?) .bind(to_i64(created_until_unix_secs, "stats_user_daily.date")?) .fetch_all(&self.pool) .await .map_sql_err()? } else { sqlx::query( r#" SELECT DATE_FORMAT(FROM_UNIXTIME(`date`), '%Y-%m-%d') AS date, total_requests AS requests, input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens AS total_tokens, total_cost AS total_cost_usd, actual_total_cost AS actual_total_cost_usd FROM stats_daily WHERE `date` >= ? AND `date` < ? AND total_requests > 0 ORDER BY `date` ASC "#, ) .bind(to_i64(created_from_unix_secs, "stats_daily.date")?) .bind(to_i64(created_until_unix_secs, "stats_daily.date")?) .fetch_all(&self.pool) .await .map_sql_err()? }; rows.iter().map(map_mysql_usage_daily_summary).collect() } async fn summarize_usage_daily_heatmap( &self, query: &UsageDailyHeatmapQuery, ) -> Result, DataLayerError> { let created_until_unix_secs = usage_current_unix_secs().saturating_add(1); let user_id = query.user_id.as_deref(); let mut summaries = BTreeMap::::new(); for item in self .summarize_usage_daily_heatmap_from_daily_aggregates( query.created_from_unix_secs, created_until_unix_secs, user_id, ) .await? { summaries.insert(item.date.clone(), item); } for item in self .summarize_usage_daily_heatmap_raw_from_range( query.created_from_unix_secs, created_until_unix_secs, user_id, ) .await? { summaries.entry(item.date.clone()).or_insert(item); } Ok(summaries.into_values().collect()) } async fn summarize_dashboard_usage_from_daily_aggregates( &self, query: &UsageDashboardSummaryQuery, ) -> Result, DataLayerError> { let row = if let Some(user_id) = query.user_id.as_deref() { sqlx::query( r#" SELECT CAST(COALESCE(SUM(total_requests), 0) AS SIGNED) AS total_requests, CAST(COALESCE(SUM(input_tokens), 0) AS SIGNED) AS input_tokens, CAST(COALESCE(SUM(input_tokens), 0) AS SIGNED) AS effective_input_tokens, CAST(COALESCE(SUM(output_tokens), 0) AS SIGNED) AS output_tokens, CAST(COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS SIGNED) AS total_tokens, CAST(COALESCE(SUM(cache_creation_tokens), 0) AS SIGNED) AS cache_creation_tokens, CAST(COALESCE(SUM(cache_read_tokens), 0) AS SIGNED) AS cache_read_tokens, CAST(COALESCE(SUM(input_tokens + cache_creation_tokens + cache_read_tokens), 0) AS SIGNED) AS total_input_context, CAST(0.0 AS DOUBLE) AS cache_creation_cost_usd, CAST(0.0 AS DOUBLE) AS cache_read_cost_usd, CAST(COALESCE(SUM(COALESCE(total_cost, 0)), 0) AS DOUBLE) AS total_cost_usd, CAST(COALESCE(SUM(COALESCE(total_cost, 0)), 0) AS DOUBLE) AS actual_total_cost_usd, CAST(COALESCE(SUM(error_requests), 0) AS SIGNED) AS error_requests, CAST(0.0 AS DOUBLE) AS response_time_sum_ms, CAST(0 AS SIGNED) AS response_time_samples FROM stats_user_daily WHERE user_id = ? AND `date` >= ? AND `date` < ? "#, ) .bind(user_id) .bind(to_i64( query.created_from_unix_secs, "stats_user_daily.date", )?) .bind(to_i64( query.created_until_unix_secs, "stats_user_daily.date", )?) .fetch_one(&self.pool) .await .map_sql_err()? } else { sqlx::query( r#" SELECT CAST(COALESCE(SUM(total_requests), 0) AS SIGNED) AS total_requests, CAST(COALESCE(SUM(input_tokens), 0) AS SIGNED) AS input_tokens, CAST(COALESCE(SUM(input_tokens), 0) AS SIGNED) AS effective_input_tokens, CAST(COALESCE(SUM(output_tokens), 0) AS SIGNED) AS output_tokens, CAST(COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS SIGNED) AS total_tokens, CAST(COALESCE(SUM(cache_creation_tokens), 0) AS SIGNED) AS cache_creation_tokens, CAST(COALESCE(SUM(cache_read_tokens), 0) AS SIGNED) AS cache_read_tokens, CAST(COALESCE(SUM(input_tokens + cache_creation_tokens + cache_read_tokens), 0) AS SIGNED) AS total_input_context, CAST(COALESCE(SUM(COALESCE(cache_creation_cost, 0)), 0) AS DOUBLE) AS cache_creation_cost_usd, CAST(COALESCE(SUM(COALESCE(cache_read_cost, 0)), 0) AS DOUBLE) AS cache_read_cost_usd, CAST(COALESCE(SUM(COALESCE(total_cost, 0)), 0) AS DOUBLE) AS total_cost_usd, CAST(COALESCE(SUM(COALESCE(actual_total_cost, 0)), 0) AS DOUBLE) AS actual_total_cost_usd, CAST(COALESCE(SUM(error_requests), 0) AS SIGNED) AS error_requests, CAST(0.0 AS DOUBLE) AS response_time_sum_ms, CAST(0 AS SIGNED) AS response_time_samples FROM stats_daily WHERE `date` >= ? AND `date` < ? "#, ) .bind(to_i64(query.created_from_unix_secs, "stats_daily.date")?) .bind(to_i64(query.created_until_unix_secs, "stats_daily.date")?) .fetch_one(&self.pool) .await .map_sql_err()? }; let total_requests = row_u64(&row, "total_requests")?; if total_requests == 0 { return Ok(None); } Ok(Some(StoredUsageDashboardSummary { total_requests, input_tokens: row_u64(&row, "input_tokens")?, effective_input_tokens: row_u64(&row, "effective_input_tokens")?, output_tokens: row_u64(&row, "output_tokens")?, total_tokens: row_u64(&row, "total_tokens")?, cache_creation_tokens: row_u64(&row, "cache_creation_tokens")?, cache_read_tokens: row_u64(&row, "cache_read_tokens")?, total_input_context: row_u64(&row, "total_input_context")?, cache_creation_cost_usd: row.try_get("cache_creation_cost_usd").map_sql_err()?, cache_read_cost_usd: row.try_get("cache_read_cost_usd").map_sql_err()?, total_cost_usd: row.try_get("total_cost_usd").map_sql_err()?, actual_total_cost_usd: row.try_get("actual_total_cost_usd").map_sql_err()?, error_requests: row_u64(&row, "error_requests")?, response_time_sum_ms: row.try_get("response_time_sum_ms").map_sql_err()?, response_time_samples: row_u64(&row, "response_time_samples")?, })) } async fn list_dashboard_daily_breakdown_from_daily_aggregates( &self, query: &UsageDashboardDailyBreakdownQuery, ) -> Result, DataLayerError> { let rows = if let Some(user_id) = query.user_id.as_deref() { sqlx::query( r#" SELECT DATE_FORMAT(FROM_UNIXTIME(`date`), '%Y-%m-%d') AS date, 'aggregate' AS model, 'aggregate' AS provider, CAST(COALESCE(SUM(total_requests), 0) AS SIGNED) AS requests, CAST(COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS SIGNED) AS total_tokens, CAST(COALESCE(SUM(COALESCE(total_cost, 0)), 0) AS DOUBLE) AS total_cost_usd, CAST(0.0 AS DOUBLE) AS response_time_sum_ms, CAST(0 AS SIGNED) AS response_time_samples FROM stats_user_daily WHERE user_id = ? AND `date` >= ? AND `date` < ? AND total_requests > 0 GROUP BY `date` ORDER BY `date` ASC "#, ) .bind(user_id) .bind(to_i64( query.created_from_unix_secs, "stats_user_daily.date", )?) .bind(to_i64( query.created_until_unix_secs, "stats_user_daily.date", )?) .fetch_all(&self.pool) .await .map_sql_err()? } else { sqlx::query( r#" SELECT DATE_FORMAT(FROM_UNIXTIME(`date`), '%Y-%m-%d') AS date, 'aggregate' AS model, 'aggregate' AS provider, CAST(COALESCE(SUM(total_requests), 0) AS SIGNED) AS requests, CAST(COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS SIGNED) AS total_tokens, CAST(COALESCE(SUM(COALESCE(total_cost, 0)), 0) AS DOUBLE) AS total_cost_usd, CAST(0.0 AS DOUBLE) AS response_time_sum_ms, CAST(0 AS SIGNED) AS response_time_samples FROM stats_daily WHERE `date` >= ? AND `date` < ? AND total_requests > 0 GROUP BY `date` ORDER BY `date` ASC "#, ) .bind(to_i64(query.created_from_unix_secs, "stats_daily.date")?) .bind(to_i64(query.created_until_unix_secs, "stats_daily.date")?) .fetch_all(&self.pool) .await .map_sql_err()? }; rows.iter() .map(|row| { Ok(StoredUsageDashboardDailyBreakdownRow { date: row.try_get("date").map_sql_err()?, model: row.try_get("model").map_sql_err()?, provider: row.try_get("provider").map_sql_err()?, requests: row_u64(row, "requests")?, total_tokens: row_u64(row, "total_tokens")?, total_cost_usd: row.try_get("total_cost_usd").map_sql_err()?, response_time_sum_ms: row.try_get("response_time_sum_ms").map_sql_err()?, response_time_samples: row_u64(row, "response_time_samples")?, }) }) .collect() } async fn summarize_dashboard_usage( &self, query: &UsageDashboardSummaryQuery, ) -> Result { if let Some(summary) = self .summarize_dashboard_usage_from_daily_aggregates(query) .await? { return Ok(summary); } let repository = self.materialize_read_model().await?; ::summarize_dashboard_usage( &repository, query, ) .await } async fn list_dashboard_daily_breakdown( &self, query: &UsageDashboardDailyBreakdownQuery, ) -> Result, DataLayerError> { let aggregate_rows = self .list_dashboard_daily_breakdown_from_daily_aggregates(query) .await?; if !aggregate_rows.is_empty() { return Ok(aggregate_rows); } let repository = self.materialize_read_model().await?; ::list_dashboard_daily_breakdown( &repository, query, ) .await } async fn summarize_usage_totals_by_user_ids( &self, user_ids: &[String], ) -> Result, DataLayerError> { if user_ids.is_empty() { return Ok(Vec::new()); } let unique_user_ids = user_ids .iter() .cloned() .collect::>(); let mut totals = BTreeMap::::new(); let mut aggregate_cutoffs = BTreeMap::::new(); let mut aggregate_builder = QueryBuilder::::new( r#" SELECT user_id, CAST(COALESCE(SUM(total_requests), 0) AS SIGNED) AS request_count, CAST(COALESCE( SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0 ) AS SIGNED) AS total_tokens, MAX(`date`) AS latest_date FROM stats_user_daily WHERE user_id IN ( "#, ); { let mut separated = aggregate_builder.separated(", "); for user_id in &unique_user_ids { separated.push_bind(user_id.clone()); } } aggregate_builder.push(") GROUP BY user_id ORDER BY user_id ASC"); let aggregate_rows = aggregate_builder .build() .fetch_all(&self.pool) .await .map_sql_err()?; for row in aggregate_rows { let user_id: String = row.try_get("user_id").map_sql_err()?; let latest_date = row.try_get::("latest_date").map_sql_err()?.max(0) as u64; aggregate_cutoffs.insert(user_id.clone(), latest_date.saturating_add(86_400)); totals.insert( user_id.clone(), StoredUsageUserTotals { user_id, request_count: row_u64(&row, "request_count")?, total_tokens: row_u64(&row, "total_tokens")?, }, ); } let mut raw_builder = QueryBuilder::::new( r#" SELECT `usage`.user_id, CAST(COUNT(*) AS SIGNED) AS request_count, CAST(COALESCE(SUM(GREATEST(COALESCE(`usage`.total_tokens, 0), 0)), 0) AS SIGNED) AS total_tokens FROM `usage` JOIN ( "#, ); for (index, user_id) in unique_user_ids.iter().enumerate() { if index > 0 { raw_builder.push(" UNION ALL "); } let cutoff = aggregate_cutoffs.get(user_id).copied().unwrap_or_default(); raw_builder .push("SELECT ") .push_bind(user_id.clone()) .push(" AS user_id, ") .push_bind(to_i64(cutoff, "usage aggregate cutoff")?) .push(" AS cutoff_unix_secs"); } raw_builder.push( r#" ) AS requested ON requested.user_id = `usage`.user_id WHERE `usage`.created_at_unix_ms >= requested.cutoff_unix_secs AND `usage`.status NOT IN ('pending', 'streaming') AND `usage`.provider_name NOT IN ('unknown', 'pending') GROUP BY `usage`.user_id ORDER BY `usage`.user_id ASC "#, ); let raw_rows = raw_builder .build() .fetch_all(&self.pool) .await .map_sql_err()?; for row in raw_rows { let user_id: String = row.try_get("user_id").map_sql_err()?; let entry = totals .entry(user_id.clone()) .or_insert_with(|| StoredUsageUserTotals { user_id, request_count: 0, total_tokens: 0, }); entry.request_count = entry .request_count .saturating_add(row_u64(&row, "request_count")?); entry.total_tokens = entry .total_tokens .saturating_add(row_u64(&row, "total_tokens")?); } Ok(totals.into_values().collect()) } } impl_materialized_usage_read_repository!(MysqlUsageReadRepository); impl MysqlUsageWriteRepository { pub fn new(pool: MysqlPool) -> Self { Self { pool } } pub async fn find_by_request_id( &self, request_id: &str, ) -> Result, DataLayerError> { let row = sqlx::query(&format!("{USAGE_COLUMNS} WHERE request_id = ? LIMIT 1")) .bind(request_id) .fetch_optional(&self.pool) .await .map_sql_err()?; row.as_ref().map(map_usage_row).transpose() } } #[async_trait] impl UsageWriteRepository for MysqlUsageWriteRepository { async fn upsert( &self, usage: UpsertUsageRecord, ) -> Result { let usage = strip_deprecated_usage_display_fields(usage); usage.validate()?; if let Some(existing) = self.find_by_request_id(&usage.request_id).await? { if (existing.billing_status == "settled" || existing.billing_status == "void") && !usage_can_recover_terminal_failure( &existing.status, &existing.billing_status, &usage.status, &usage.billing_status, ) { return Ok(existing); } } bind_upsert(sqlx::query(UPSERT_USAGE_SQL), &usage)? .execute(&self.pool) .await .map_sql_err()?; self.rebuild_api_key_usage_stats().await?; self.rebuild_provider_api_key_usage_stats().await?; self.find_by_request_id(&usage.request_id) .await? .ok_or_else(|| { DataLayerError::UnexpectedValue("usage upsert returned no row".to_string()) }) } async fn rebuild_api_key_usage_stats(&self) -> Result { sqlx::query( r#" UPDATE api_keys SET total_requests = 0, total_tokens = 0, total_cost_usd = 0, last_used_at = NULL "#, ) .execute(&self.pool) .await .map_sql_err()?; let rows = sqlx::query( r#" SELECT api_key_id, COUNT(*) AS total_requests, CAST(COALESCE(SUM(total_tokens), 0) AS SIGNED) AS total_tokens, CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE) AS total_cost_usd, MAX(updated_at_unix_secs) AS last_used_at FROM `usage` WHERE api_key_id IS NOT NULL AND api_key_id <> '' GROUP BY api_key_id "#, ) .fetch_all(&self.pool) .await .map_sql_err()?; for row in &rows { sqlx::query( r#" UPDATE api_keys SET total_requests = ?, total_tokens = ?, total_cost_usd = ?, last_used_at = ? WHERE id = ? "#, ) .bind(row.try_get::("total_requests").map_sql_err()?) .bind(row.try_get::("total_tokens").map_sql_err()?) .bind(row.try_get::("total_cost_usd").map_sql_err()?) .bind( row.try_get::, _>("last_used_at") .map_sql_err()?, ) .bind(row.try_get::("api_key_id").map_sql_err()?) .execute(&self.pool) .await .map_sql_err()?; } Ok(rows.len() as u64) } async fn rebuild_provider_api_key_usage_stats(&self) -> Result { sqlx::query( r#" UPDATE provider_api_keys SET request_count = 0, success_count = 0, error_count = 0, total_tokens = 0, total_cost_usd = 0, total_response_time_ms = 0, last_used_at = NULL "#, ) .execute(&self.pool) .await .map_sql_err()?; let rows = sqlx::query( r#" SELECT provider_api_key_id, status, status_code, error_message, total_tokens, total_cost_usd, response_time_ms, updated_at_unix_secs FROM `usage` WHERE provider_api_key_id IS NOT NULL AND provider_api_key_id <> '' "#, ) .fetch_all(&self.pool) .await .map_sql_err()?; let mut stats = BTreeMap::::new(); for row in rows { let key_id: String = row.try_get("provider_api_key_id").map_sql_err()?; let status: String = row.try_get("status").map_sql_err()?; let status_code = row.try_get::, _>("status_code").map_sql_err()?; let status_code_u16 = status_code.and_then(|value| u16::try_from(value).ok()); let error_message: Option = row.try_get("error_message").map_sql_err()?; let entry = stats.entry(key_id).or_default(); entry.request_count += 1; let is_success = provider_api_key_usage_is_success( &status, status_code_u16, error_message.as_deref(), ); let is_in_flight = matches!(status.as_str(), "pending" | "streaming"); if is_success { entry.success_count += 1; } if provider_api_key_usage_is_error(&status, status_code_u16, error_message.as_deref()) { entry.error_count += 1; } if !is_in_flight { entry.total_tokens += row.try_get::("total_tokens").map_sql_err()?; entry.total_cost_usd += row.try_get::("total_cost_usd").map_sql_err()?; } if is_success { entry.total_response_time_ms += row .try_get::, _>("response_time_ms") .map_sql_err()? .unwrap_or_default(); } entry.last_used_at = entry.last_used_at.max( row.try_get::, _>("updated_at_unix_secs") .map_sql_err()?, ); } for (key_id, stat) in &stats { sqlx::query( r#" UPDATE provider_api_keys SET request_count = ?, success_count = ?, error_count = ?, total_tokens = ?, total_cost_usd = ?, total_response_time_ms = ?, last_used_at = ? WHERE id = ? "#, ) .bind(stat.request_count) .bind(stat.success_count) .bind(stat.error_count) .bind(stat.total_tokens) .bind(stat.total_cost_usd) .bind(stat.total_response_time_ms) .bind(stat.last_used_at) .bind(key_id) .execute(&self.pool) .await .map_sql_err()?; } Ok(stats.len() as u64) } async fn cleanup_stale_pending_requests( &self, cutoff_unix_secs: u64, now_unix_secs: u64, timeout_minutes: u64, batch_size: usize, ) -> Result { if batch_size == 0 { return Ok(PendingUsageCleanupSummary::default()); } let cutoff_unix_ms = cutoff_unix_secs.saturating_mul(1000); let now_unix_ms = now_unix_secs.saturating_mul(1000); let mut summary = PendingUsageCleanupSummary::default(); let batch_size_u64 = u64::try_from(batch_size).map_err(|_| { DataLayerError::InvalidInput(format!( "invalid stale pending usage batch size: {batch_size}" )) })?; loop { let mut tx = self.pool.begin().await.map_sql_err()?; let stale_rows = sqlx::query(SELECT_STALE_PENDING_USAGE_BATCH_SQL) .bind(to_i64(cutoff_unix_ms, "stale pending usage cutoff")?) .bind(to_i64(batch_size_u64, "stale pending usage batch size")?) .fetch_all(&mut *tx) .await .map_sql_err()?; if stale_rows.is_empty() { tx.rollback().await.map_sql_err()?; break; } let stale_rows = stale_rows .iter() .map(|row| { Ok(StalePendingUsageRow { request_id: row.try_get("request_id").map_sql_err()?, status: row.try_get("status").map_sql_err()?, billing_status: row.try_get("billing_status").map_sql_err()?, }) }) .collect::, DataLayerError>>()?; let completed_request_ids = completed_request_ids_mysql(&mut tx, stale_rows.iter().map(|row| &row.request_id)) .await?; for row in stale_rows { if completed_request_ids.contains(&row.request_id) { sqlx::query( r#" UPDATE `usage` SET status = 'completed', status_code = 200, error_message = NULL WHERE request_id = ? "#, ) .bind(&row.request_id) .execute(&mut *tx) .await .map_sql_err()?; sqlx::query( r#" UPDATE request_candidates SET status = 'success', finished_at = ? WHERE request_id = ? AND status = 'streaming' "#, ) .bind(to_i64(now_unix_ms, "request candidate finished_at")?) .bind(&row.request_id) .execute(&mut *tx) .await .map_sql_err()?; summary.recovered += 1; continue; } let candidate_info = latest_failed_candidate_mysql(&mut tx, &row.request_id).await?; let (status_code, error_message) = resolve_stale_pending_failure( candidate_info.as_ref(), &row.status, timeout_minutes, ); let status_code_i64 = i64::from(status_code); if row.billing_status == "pending" { sqlx::query( r#" UPDATE `usage` SET status = 'failed', status_code = ?, error_message = ?, billing_status = 'void', finalized_at = ?, total_cost_usd = 0, actual_total_cost_usd = 0 WHERE request_id = ? "#, ) .bind(status_code_i64) .bind(&error_message) .bind(to_i64(now_unix_secs, "usage finalized_at")?) .bind(&row.request_id) .execute(&mut *tx) .await .map_sql_err()?; upsert_void_usage_settlement_snapshot_mysql( &mut tx, &row.request_id, now_unix_secs, ) .await?; } else { sqlx::query( r#" UPDATE `usage` SET status = 'failed', status_code = ?, error_message = ? WHERE request_id = ? "#, ) .bind(status_code_i64) .bind(&error_message) .bind(&row.request_id) .execute(&mut *tx) .await .map_sql_err()?; } sqlx::query( r#" UPDATE request_candidates SET status = 'failed', finished_at = ?, error_message = '请求超时（服务器可能已重启）' WHERE request_id = ? AND status IN ('pending', 'streaming') "#, ) .bind(to_i64(now_unix_ms, "request candidate finished_at")?) .bind(&row.request_id) .execute(&mut *tx) .await .map_sql_err()?; summary.failed += 1; } tx.commit().await.map_sql_err()?; } Ok(summary) } } struct StalePendingUsageRow { request_id: String, status: String, billing_status: String, } #[derive(Default)] struct ProviderKeyStats { request_count: i64, success_count: i64, error_count: i64, total_tokens: i64, total_cost_usd: f64, total_response_time_ms: i64, last_used_at: Option, } async fn completed_request_ids_mysql<'a>( tx: &mut sqlx::Transaction<'_, sqlx::MySql>, request_ids: impl Iterator, ) -> Result, DataLayerError> { let mut completed = HashSet::new(); for request_id in request_ids { let rows = sqlx::query(SELECT_COMPLETED_REQUEST_CANDIDATES_SQL) .bind(request_id) .fetch_all(&mut **tx) .await .map_sql_err()?; let mut is_completed = false; for row in &rows { if candidate_row_is_completed(row)? { is_completed = true; break; } } if is_completed { completed.insert(request_id.clone()); } } Ok(completed) } fn candidate_row_is_completed(row: &MySqlRow) -> Result { let status: String = row.try_get("status").map_sql_err()?; if status == "streaming" { return Ok(true); } if status != "success" { return Ok(false); } let Some(extra_data) = row .try_get::, _>("extra_data") .map_sql_err()? else { return Ok(false); }; let Ok(value) = serde_json::from_str::(&extra_data) else { return Ok(false); }; Ok(value .get("stream_completed") .and_then(serde_json::Value::as_bool) .unwrap_or(false)) } async fn upsert_void_usage_settlement_snapshot_mysql( tx: &mut sqlx::Transaction<'_, sqlx::MySql>, request_id: &str, now_unix_secs: u64, ) -> Result<(), DataLayerError> { let now = to_i64(now_unix_secs, "usage settlement snapshot timestamp")?; sqlx::query( r#" INSERT INTO usage_settlement_snapshots ( request_id, billing_status, finalized_at, created_at, updated_at ) VALUES (?, 'void', ?, ?, ?) ON DUPLICATE KEY UPDATE billing_status = VALUES(billing_status), finalized_at = COALESCE(usage_settlement_snapshots.finalized_at, VALUES(finalized_at)), updated_at = VALUES(updated_at) "#, ) .bind(request_id) .bind(now) .bind(now) .bind(now) .execute(&mut **tx) .await .map_sql_err()?; Ok(()) } fn stale_pending_error_message(status: &str, timeout_minutes: u64) -> String { format!("请求超时: 状态 '{status}' 超过 {timeout_minutes} 分钟未完成") } struct FailedCandidateCleanupInfo { status_code: Option, error_message: Option, } fn resolve_stale_pending_failure( candidate: Option<&FailedCandidateCleanupInfo>, status: &str, timeout_minutes: u64, ) -> (u16, String) { match candidate { Some(info) => ( info.status_code.unwrap_or(502), info.error_message .clone() .unwrap_or_else(|| stale_pending_error_message(status, timeout_minutes)), ), None => (504, stale_pending_error_message(status, timeout_minutes)), } } async fn latest_failed_candidate_mysql( tx: &mut sqlx::Transaction<'_, sqlx::MySql>, request_id: &str, ) -> Result, DataLayerError> { let row = sqlx::query( r#" SELECT status_code, error_message FROM request_candidates WHERE request_id = ? AND status IN ('failed', 'cancelled') ORDER BY COALESCE(finished_at, started_at, created_at) DESC, retry_index DESC, candidate_index DESC LIMIT 1 "#, ) .bind(request_id) .fetch_optional(&mut **tx) .await .map_sql_err()?; let Some(row) = row else { return Ok(None); }; let status_code = row .try_get::, _>("status_code") .map_sql_err()? .and_then(|value| u16::try_from(value).ok()); let error_message = row .try_get::, _>("error_message") .map_sql_err()? .map(|value| value.trim().to_string()) .filter(|value| !value.is_empty()); Ok(Some(FailedCandidateCleanupInfo { status_code, error_message, })) } fn bind_upsert<'q>( mut query: sqlx::query::Query<'q, sqlx::MySql, sqlx::mysql::MySqlArguments>, usage: &'q UpsertUsageRecord, ) -> Result, DataLayerError> { let input_tokens = usage.input_tokens.unwrap_or_default(); let output_tokens = usage.output_tokens.unwrap_or_default(); let cache_creation_tokens = usage .cache_creation_input_tokens .or_else(|| { Some( usage .cache_creation_ephemeral_5m_input_tokens .unwrap_or_default() + usage .cache_creation_ephemeral_1h_input_tokens .unwrap_or_default(), ) }) .unwrap_or_default(); let cache_read_tokens = usage.cache_read_input_tokens.unwrap_or_default(); let total_tokens = usage .total_tokens .unwrap_or(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens); let created_at = usage .created_at_unix_ms .unwrap_or(usage.updated_at_unix_secs.saturating_mul(1000)); let request_metadata = usage .request_metadata .as_ref() .map(serde_json::to_string) .transpose() .map_err(|err| DataLayerError::InvalidInput(err.to_string()))?; query = query .bind(&usage.request_id) .bind(&usage.request_id) .bind(usage.user_id.as_deref()) .bind(usage.api_key_id.as_deref()) .bind(&usage.provider_name) .bind(&usage.model) .bind(usage.target_model.as_deref()) .bind(usage.provider_id.as_deref()) .bind(usage.provider_endpoint_id.as_deref()) .bind(usage.provider_api_key_id.as_deref()) .bind(usage.request_type.as_deref()) .bind(usage.api_format.as_deref()) .bind(usage.api_family.as_deref()) .bind(usage.endpoint_kind.as_deref()) .bind(usage.endpoint_api_format.as_deref()) .bind(usage.provider_api_family.as_deref()) .bind(usage.provider_endpoint_kind.as_deref()) .bind(usage.has_format_conversion.unwrap_or(false)) .bind(usage.is_stream.unwrap_or(false)) .bind(usage_upstream_is_stream(usage)) .bind(to_i64(input_tokens, "input_tokens")?) .bind(to_i64(output_tokens, "output_tokens")?) .bind(to_i64(total_tokens, "total_tokens")?) .bind(to_i64( cache_creation_tokens, "cache_creation_input_tokens", )?) .bind(to_i64( usage .cache_creation_ephemeral_5m_input_tokens .unwrap_or_default(), "cache_creation_ephemeral_5m_input_tokens", )?) .bind(to_i64( usage .cache_creation_ephemeral_1h_input_tokens .unwrap_or_default(), "cache_creation_ephemeral_1h_input_tokens", )?) .bind(to_i64(cache_read_tokens, "cache_read_input_tokens")?) .bind(usage.cache_creation_cost_usd.unwrap_or_default()) .bind(usage.cache_read_cost_usd.unwrap_or_default()) .bind(usage.output_price_per_1m) .bind(usage.total_cost_usd.unwrap_or_default()) .bind(usage.actual_total_cost_usd.unwrap_or_default()) .bind(usage.status_code.map(i64::from)) .bind(usage.error_message.as_deref()) .bind(usage.error_category.as_deref()) .bind(usage.response_time_ms.map(|value| value as i64)) .bind(usage.first_byte_time_ms.map(|value| value as i64)) .bind(&usage.status) .bind(&usage.billing_status) .bind(request_metadata) .bind(usage.candidate_id.as_deref()) .bind(usage.candidate_index.map(|value| value as i64)) .bind(usage.key_name.as_deref()) .bind(usage.planner_kind.as_deref()) .bind(usage.route_family.as_deref()) .bind(usage.route_kind.as_deref()) .bind(usage.execution_path.as_deref()) .bind(usage.local_execution_runtime_miss_reason.as_deref()) .bind(usage.finalized_at_unix_secs.map(|value| value as i64)) .bind(to_i64(created_at, "created_at_unix_ms")?) .bind(to_i64(usage.updated_at_unix_secs, "updated_at_unix_secs")?); Ok(query) } fn map_usage_row(row: &MySqlRow) -> Result { let id = row .try_get::, _>("id") .map_sql_err()? .unwrap_or_else(|| { row.try_get::("request_id") .unwrap_or_else(|_| "unknown".to_string()) }); let mut audit = StoredRequestUsageAudit::new( id, row.try_get("request_id").map_sql_err()?, row.try_get("user_id").map_sql_err()?, row.try_get("api_key_id").map_sql_err()?, None, None, row.try_get("provider_name").map_sql_err()?, row.try_get("model").map_sql_err()?, row.try_get("target_model").map_sql_err()?, row.try_get("provider_id").map_sql_err()?, row.try_get("provider_endpoint_id").map_sql_err()?, row.try_get("provider_api_key_id").map_sql_err()?, row.try_get("request_type").map_sql_err()?, row.try_get("api_format").map_sql_err()?, row.try_get("api_family").map_sql_err()?, row.try_get("endpoint_kind").map_sql_err()?, row.try_get("endpoint_api_format").map_sql_err()?, row.try_get("provider_api_family").map_sql_err()?, row.try_get("provider_endpoint_kind").map_sql_err()?, row.try_get::("has_format_conversion") .map_sql_err()?, row.try_get::("is_stream").map_sql_err()?, row_i32(row, "input_tokens")?, row_i32(row, "output_tokens")?, row_i32(row, "total_tokens")?, row.try_get("total_cost_usd").map_sql_err()?, row.try_get("actual_total_cost_usd").map_sql_err()?, row_optional_i32(row, "status_code")?, row.try_get("error_message").map_sql_err()?, row.try_get("error_category").map_sql_err()?, row_optional_i32(row, "response_time_ms")?, row_optional_i32(row, "first_byte_time_ms")?, row.try_get("status").map_sql_err()?, row.try_get("billing_status").map_sql_err()?, row.try_get("created_at_unix_ms").map_sql_err()?, row.try_get("updated_at_unix_secs").map_sql_err()?, row.try_get("finalized_at_unix_secs").map_sql_err()?, )?; audit.cache_creation_input_tokens = row_u64(row, "cache_creation_input_tokens")?; audit.cache_creation_ephemeral_5m_input_tokens = row_u64(row, "cache_creation_ephemeral_5m_input_tokens")?; audit.cache_creation_ephemeral_1h_input_tokens = row_u64(row, "cache_creation_ephemeral_1h_input_tokens")?; audit.cache_read_input_tokens = row_u64(row, "cache_read_input_tokens")?; audit.cache_creation_cost_usd = row.try_get("cache_creation_cost_usd").map_sql_err()?; audit.cache_read_cost_usd = row.try_get("cache_read_cost_usd").map_sql_err()?; audit.output_price_per_1m = row.try_get("output_price_per_1m").map_sql_err()?; audit.request_metadata = row .try_get::, _>("request_metadata") .map_sql_err()? .map(|raw| serde_json::from_str(&raw)) .transpose() .map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?; audit.client_family = usage_request_metadata_client_family(audit.request_metadata.as_ref()) .map(ToOwned::to_owned); let upstream_is_stream = row .try_get::, _>("upstream_is_stream") .map_sql_err()?; merge_usage_stream_metadata(&mut audit.request_metadata, upstream_is_stream); audit.candidate_id = row.try_get("candidate_id").map_sql_err()?; audit.candidate_index = row .try_get::, _>("candidate_index") .map_sql_err()? .map(|value| value as u64); audit.key_name = row.try_get("key_name").map_sql_err()?; audit.planner_kind = row.try_get("planner_kind").map_sql_err()?; audit.route_family = row.try_get("route_family").map_sql_err()?; audit.route_kind = row.try_get("route_kind").map_sql_err()?; audit.execution_path = row.try_get("execution_path").map_sql_err()?; audit.local_execution_runtime_miss_reason = row .try_get("local_execution_runtime_miss_reason") .map_sql_err()?; Ok(audit) } fn to_i64(value: u64, field: &str) -> Result { i64::try_from(value).map_err(|_| DataLayerError::InvalidInput(format!("{field} overflow"))) } fn usage_upstream_is_stream(usage: &UpsertUsageRecord) -> bool { usage .request_metadata .as_ref() .and_then(serde_json::Value::as_object) .and_then(|metadata| metadata.get(UPSTREAM_IS_STREAM_KEY)) .and_then(serde_json::Value::as_bool) .unwrap_or_else(|| usage.is_stream.unwrap_or(false)) } fn merge_usage_stream_metadata(metadata: &mut Option, upstream: Option) { let Some(upstream) = upstream else { return; }; let value = metadata.get_or_insert_with(|| serde_json::json!({})); let Some(object) = value.as_object_mut() else { return; }; object .entry(UPSTREAM_IS_STREAM_KEY) .or_insert(serde_json::Value::Bool(upstream)); } fn row_i32(row: &MySqlRow, field: &str) -> Result { let value: i64 = row.try_get(field).map_sql_err()?; i32::try_from(value).map_err(|_| DataLayerError::UnexpectedValue(format!("{field} overflow"))) } fn row_optional_i32(row: &MySqlRow, field: &str) -> Result, DataLayerError> { row.try_get::, _>(field) .map_sql_err()? .map(|value| { i32::try_from(value) .map_err(|_| DataLayerError::UnexpectedValue(format!("{field} overflow"))) }) .transpose() } fn row_u64(row: &MySqlRow, field: &str) -> Result { let value: i64 = row.try_get(field).map_sql_err()?; u64::try_from(value).map_err(|_| DataLayerError::UnexpectedValue(format!("{field} negative"))) } fn map_mysql_usage_daily_summary( row: &MySqlRow, ) -> Result { Ok(StoredUsageDailySummary { date: row.try_get("date").map_sql_err()?, requests: row_u64(row, "requests")?, total_tokens: row_u64(row, "total_tokens")?, total_cost_usd: row.try_get("total_cost_usd").map_sql_err()?, actual_total_cost_usd: row.try_get("actual_total_cost_usd").map_sql_err()?, }) } fn usage_current_unix_secs() -> u64 { SystemTime::now() .duration_since(UNIX_EPOCH) .map(|duration| duration.as_secs()) .unwrap_or_default() } #[cfg(test)] mod tests { use super::{MysqlUsageReadRepository, MysqlUsageWriteRepository}; use crate::lifecycle::migrate::run_mysql_migrations; use crate::repository::usage::{ UpsertUsageRecord, UsageAuditListQuery, UsageDashboardSummaryQuery, UsageReadRepository, UsageWriteRepository, }; #[tokio::test] async fn repository_builds_from_lazy_pool() { let pool = sqlx::mysql::MySqlPoolOptions::new().connect_lazy_with( "mysql://user:pass@localhost:3306/aether" .parse() .expect("mysql options should parse"), ); let _repository = MysqlUsageWriteRepository::new(pool); } #[test] fn mysql_usage_daily_heatmap_reads_imported_daily_aggregates() { let source = include_str!("mysql.rs"); assert!(source.contains("summarize_usage_daily_heatmap_from_daily_aggregates")); assert!(source.contains("FROM stats_daily")); assert!(source.contains("FROM stats_user_daily")); assert!(source.contains("AS SIGNED) AS total_tokens")); assert!(source.contains("CAST(COUNT(*) AS SIGNED) AS requests")); assert!(source.contains("summaries.entry(item.date.clone()).or_insert(item)")); } #[test] fn mysql_usage_totals_by_user_ids_reads_imported_user_daily_aggregates() { let source = include_str!("mysql.rs"); assert!(source.contains("async fn summarize_usage_totals_by_user_ids")); assert!(source.contains("FROM stats_user_daily")); assert!(source.contains("MAX(`date`) AS latest_date")); assert!(source.contains("AS SIGNED) AS request_count")); assert!( source.contains("CAST(COALESCE(SUM(total_requests), 0) AS SIGNED) AS request_count") ); assert!(source.contains("requested.cutoff_unix_secs")); } #[test] fn mysql_dashboard_reads_imported_daily_aggregates() { let source = include_str!("mysql.rs"); assert!(source.contains("summarize_dashboard_usage_from_daily_aggregates")); assert!(source.contains("list_dashboard_daily_breakdown_from_daily_aggregates")); assert!(source.contains("FROM stats_daily")); assert!(source.contains("FROM stats_user_daily")); assert!(source.contains("'aggregate' AS model")); assert!(source.contains("AS SIGNED) AS total_requests")); assert!( source.contains("CAST(COALESCE(SUM(total_requests), 0) AS SIGNED) AS total_requests") ); assert!(source.contains("CAST(COALESCE(SUM(total_requests), 0) AS SIGNED) AS requests")); } #[tokio::test] async fn mysql_usage_write_repository_upserts_when_url_is_set() { let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL") .ok() .filter(|value| !value.trim().is_empty()) else { eprintln!( "skipping mysql usage write smoke test because AETHER_TEST_MYSQL_URL is unset" ); return; }; let pool = sqlx::mysql::MySqlPoolOptions::new() .max_connections(1) .connect(&database_url) .await .expect("mysql test pool should connect"); run_mysql_migrations(&pool) .await .expect("mysql migrations should run"); let suffix = unique_suffix(); let user_id = format!("user-{suffix}"); let api_key_id = format!("api-key-{suffix}"); let provider_id = format!("provider-{suffix}"); let provider_key_id = format!("provider-key-{suffix}"); seed_stats_targets(&pool, &user_id, &api_key_id, &provider_id, &provider_key_id).await; let repository = MysqlUsageWriteRepository::new(pool.clone()); let record = repository .upsert(sample_usage( &format!("request-{suffix}"), &user_id, &api_key_id, &provider_id, &provider_key_id, "completed", "pending", 1_000, )) .await .expect("usage should upsert"); assert_eq!(record.api_key_id.as_deref(), Some(api_key_id.as_str())); assert_eq!( record.provider_api_key_id.as_deref(), Some(provider_key_id.as_str()) ); assert_eq!(record.total_tokens, 7); assert_eq!( record.request_metadata.as_ref().unwrap()["upstream_is_stream"], true ); let upstream_is_stream: Option = sqlx::query_scalar("SELECT upstream_is_stream FROM `usage` WHERE request_id = ?") .bind(format!("request-{suffix}")) .fetch_one(&pool) .await .expect("usage stream mode should load"); assert_eq!(upstream_is_stream, Some(true)); let stats = sqlx::query_as::<_, (i64, i64, f64, Option)>( "SELECT total_requests, total_tokens, total_cost_usd, last_used_at FROM api_keys WHERE id = ?", ) .bind(&api_key_id) .fetch_one(&pool) .await .expect("api key stats should load"); assert_eq!(stats, (1, 7, 0.5, Some(1_000))); let provider_stats = sqlx::query_as::<_, (i64, i64, i64, i64, f64, i64, Option)>( "SELECT request_count, success_count, error_count, total_tokens, total_cost_usd, total_response_time_ms, last_used_at FROM provider_api_keys WHERE id = ?", ) .bind(&provider_key_id) .fetch_one(&pool) .await .expect("provider key stats should load"); assert_eq!(provider_stats, (1, 1, 0, 7, 0.5, 42, Some(1_000))); } #[tokio::test] async fn mysql_usage_read_repository_reads_usage_contract_views_when_url_is_set() { let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL") .ok() .filter(|value| !value.trim().is_empty()) else { eprintln!( "skipping mysql usage read smoke test because AETHER_TEST_MYSQL_URL is unset" ); return; }; let pool = sqlx::mysql::MySqlPoolOptions::new() .max_connections(1) .connect(&database_url) .await .expect("mysql test pool should connect"); run_mysql_migrations(&pool) .await .expect("mysql migrations should run"); let suffix = unique_suffix(); let user_id = format!("user-read-{suffix}"); let api_key_id = format!("api-key-read-{suffix}"); let provider_id = format!("provider-read-{suffix}"); let provider_key_id = format!("provider-key-read-{suffix}"); seed_stats_targets(&pool, &user_id, &api_key_id, &provider_id, &provider_key_id).await; let writer = MysqlUsageWriteRepository::new(pool.clone()); writer .upsert(sample_usage( &format!("request-read-1-{suffix}"), &user_id, &api_key_id, &provider_id, &provider_key_id, "completed", "settled", 1_000, )) .await .expect("usage should upsert"); writer .upsert(sample_usage( &format!("request-read-2-{suffix}"), &user_id, &api_key_id, &provider_id, &provider_key_id, "failed", "void", 1_010, )) .await .expect("usage should upsert"); let reader = MysqlUsageReadRepository::new(pool); let loaded = reader .find_by_request_id(&format!("request-read-1-{suffix}")) .await .expect("usage should load") .expect("usage should exist"); assert_eq!(loaded.total_tokens, 7); assert_eq!(loaded.billing_status, "settled"); let listed = reader .list_usage_audits(&UsageAuditListQuery { user_id: Some(user_id.clone()), provider_name: Some("Provider One".to_string()), newest_first: true, ..UsageAuditListQuery::default() }) .await .expect("usage list should load"); assert_eq!(listed.len(), 2); assert!(listed[0].request_id.starts_with("request-read-2-")); let summary = reader .summarize_dashboard_usage(&UsageDashboardSummaryQuery { created_from_unix_secs: 999, created_until_unix_secs: 1_020, user_id: Some(user_id), }) .await .expect("dashboard summary should load"); assert_eq!(summary.total_requests, 2); assert_eq!(summary.error_requests, 1); assert_eq!(summary.total_tokens, 10); } async fn seed_stats_targets( pool: &sqlx::MySqlPool, user_id: &str, api_key_id: &str, provider_id: &str, provider_key_id: &str, ) { sqlx::query( r#" INSERT INTO users (id, auth_source, created_at, updated_at) VALUES (?, 'local', 1, 1) "#, ) .bind(user_id) .execute(pool) .await .expect("user should seed"); sqlx::query( r#" INSERT INTO api_keys (id, user_id, key_hash, created_at, updated_at) VALUES (?, ?, ?, 1, 1) "#, ) .bind(api_key_id) .bind(user_id) .bind(format!("hash-{api_key_id}")) .execute(pool) .await .expect("api key should seed"); sqlx::query( r#" INSERT INTO providers (id, name, provider_type, created_at, updated_at) VALUES (?, ?, 'openai', 1, 1) "#, ) .bind(provider_id) .bind(format!("Provider {provider_id}")) .execute(pool) .await .expect("provider should seed"); sqlx::query( r#" INSERT INTO provider_api_keys (id, provider_id, name, created_at, updated_at) VALUES (?, ?, ?, 1, 1) "#, ) .bind(provider_key_id) .bind(provider_id) .bind(format!("Provider Key {provider_key_id}")) .execute(pool) .await .expect("provider key should seed"); } #[allow(clippy::too_many_arguments)] fn sample_usage( request_id: &str, user_id: &str, api_key_id: &str, provider_id: &str, provider_key_id: &str, status: &str, billing_status: &str, updated_at: u64, ) -> UpsertUsageRecord { UpsertUsageRecord { request_id: request_id.to_string(), user_id: Some(user_id.to_string()), api_key_id: Some(api_key_id.to_string()), username: Some("legacy-user".to_string()), api_key_name: Some("legacy-key".to_string()), provider_name: "Provider One".to_string(), model: "model-1".to_string(), target_model: Some("target-model".to_string()), provider_id: Some(provider_id.to_string()), provider_endpoint_id: Some("endpoint-1".to_string()), provider_api_key_id: Some(provider_key_id.to_string()), request_type: Some("chat".to_string()), api_format: Some("openai".to_string()), api_family: Some("chat".to_string()), endpoint_kind: Some("chat".to_string()), endpoint_api_format: Some("openai".to_string()), provider_api_family: Some("chat".to_string()), provider_endpoint_kind: Some("chat".to_string()), has_format_conversion: Some(true), is_stream: Some(false), input_tokens: Some(2), output_tokens: Some(3), total_tokens: None, cache_creation_input_tokens: None, cache_creation_ephemeral_5m_input_tokens: Some(0), cache_creation_ephemeral_1h_input_tokens: Some(0), cache_read_input_tokens: Some(2), cache_creation_cost_usd: Some(0.0), cache_read_cost_usd: Some(0.1), output_price_per_1m: Some(2.0), total_cost_usd: Some(0.5), actual_total_cost_usd: Some(0.4), status_code: Some(200), error_message: None, error_category: None, response_time_ms: Some(42), first_byte_time_ms: Some(12), status: status.to_string(), billing_status: billing_status.to_string(), request_headers: None, request_body: None, request_body_ref: None, request_body_state: None, provider_request_headers: None, provider_request_body: None, provider_request_body_ref: None, provider_request_body_state: None, response_headers: None, response_body: None, response_body_ref: None, response_body_state: None, client_response_headers: None, client_response_body: None, client_response_body_ref: None, client_response_body_state: None, candidate_id: Some("candidate-1".to_string()), candidate_index: Some(1), key_name: Some("key-one".to_string()), planner_kind: Some("default".to_string()), route_family: Some("chat".to_string()), route_kind: Some("completion".to_string()), execution_path: Some("remote".to_string()), local_execution_runtime_miss_reason: None, request_metadata: Some(serde_json::json!({ "trace_id": "trace-1", "upstream_is_stream": true, })), finalized_at_unix_secs: Some(updated_at), created_at_unix_ms: Some(updated_at), updated_at_unix_secs: updated_at, } } fn unique_suffix() -> String { let nanos = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap_or_default() .as_nanos(); format!("{}-{nanos}", std::process::id()) } }