diff --git a/apps/aether-gateway/src/ai_serving/planner/candidate_source.rs b/apps/aether-gateway/src/ai_serving/planner/candidate_source.rs index d5ad3c92c..bdd61a702 100644 --- a/apps/aether-gateway/src/ai_serving/planner/candidate_source.rs +++ b/apps/aether-gateway/src/ai_serving/planner/candidate_source.rs @@ -1600,6 +1600,7 @@ mod tests { user_is_active: true, user_is_deleted: false, user_rate_limit: None, + user_daily_usage_limit_usd: None, user_allowed_providers: None, user_allowed_api_formats: None, user_allowed_models: None, @@ -1609,6 +1610,7 @@ mod tests { api_key_is_locked: false, api_key_is_standalone: false, api_key_rate_limit: None, + api_key_daily_usage_limit_usd: None, api_key_concurrent_limit: None, api_key_expires_at_unix_secs: None, api_key_allowed_providers: None, diff --git a/apps/aether-gateway/src/ai_serving/planner/decision_input.rs b/apps/aether-gateway/src/ai_serving/planner/decision_input.rs index 784106113..ee0bb0882 100644 --- a/apps/aether-gateway/src/ai_serving/planner/decision_input.rs +++ b/apps/aether-gateway/src/ai_serving/planner/decision_input.rs @@ -1125,6 +1125,7 @@ mod tests { user_is_active: true, user_is_deleted: false, user_rate_limit: None, + user_daily_usage_limit_usd: None, user_allowed_providers: None, user_allowed_api_formats: None, user_allowed_models: None, @@ -1134,6 +1135,7 @@ mod tests { api_key_is_locked: false, api_key_is_standalone: false, api_key_rate_limit: None, + api_key_daily_usage_limit_usd: None, api_key_concurrent_limit: None, api_key_expires_at_unix_secs: None, api_key_allowed_providers: None, diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/family/payload.rs b/apps/aether-gateway/src/ai_serving/planner/standard/family/payload.rs index cbb556095..62434a106 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/family/payload.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/family/payload.rs @@ -340,6 +340,7 @@ mod tests { user_is_active: true, user_is_deleted: false, user_rate_limit: None, + user_daily_usage_limit_usd: None, user_allowed_providers: None, user_allowed_api_formats: None, user_allowed_models: None, @@ -349,6 +350,7 @@ mod tests { api_key_is_locked: false, api_key_is_standalone: false, api_key_rate_limit: None, + api_key_daily_usage_limit_usd: None, api_key_concurrent_limit: None, api_key_expires_at_unix_secs: None, api_key_allowed_providers: None, diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs index 00c22018e..35fe649f6 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs @@ -2158,6 +2158,7 @@ mod tests { user_is_active: true, user_is_deleted: false, user_rate_limit: None, + user_daily_usage_limit_usd: None, user_allowed_providers: None, user_allowed_api_formats: None, user_allowed_models: None, @@ -2167,6 +2168,7 @@ mod tests { api_key_is_locked: false, api_key_is_standalone: false, api_key_rate_limit: None, + api_key_daily_usage_limit_usd: None, api_key_concurrent_limit: None, api_key_expires_at_unix_secs: None, api_key_allowed_providers: None, diff --git a/apps/aether-gateway/src/api/response.rs b/apps/aether-gateway/src/api/response.rs index 403037c7a..dd407e125 100644 --- a/apps/aether-gateway/src/api/response.rs +++ b/apps/aether-gateway/src/api/response.rs @@ -10,6 +10,7 @@ use crate::ai_serving::{build_core_error_body_for_client_format, LocalCoreSyncEr use crate::constants::*; use crate::control::GatewayControlDecision; use crate::control::GatewayLocalAuthRejection; +use crate::daily_usage_limit::{DailyUsageLimitedResponse, FrontdoorDailyUsageRejection}; use crate::headers::should_skip_response_header; use crate::rate_limit::FrontdoorUserRpmRejection; use crate::{insert_header_if_missing, GatewayError}; @@ -258,6 +259,74 @@ pub(crate) fn build_local_user_rpm_limited_response( ) } +pub(crate) fn build_local_daily_usage_limited_response( + trace_id: &str, + control_decision: Option<&GatewayControlDecision>, + rejection: &FrontdoorDailyUsageRejection, +) -> Result, GatewayError> { + let message = "已达到每日使用上限，请在额度重置后重试"; + let fallback_payload = json!({ + "error": { + "type": "daily_usage_limit_exceeded", + "message": message, + "details": { + "limit_usd": rejection.limit_usd, + "used_usd": rejection.used_usd, + "remaining_usd": rejection.remaining_usd, + "scope": rejection.scope, + "reset_at": rejection.reset_at_unix_secs, + "timezone": rejection.timezone, + } + } + }); + let payload = if local_error_uses_claude_format(control_decision, None) { + build_core_error_body_for_client_format( + "claude:messages", + message, + Some("daily_usage_limit_exceeded"), + LocalCoreSyncErrorKind::RateLimit, + ) + .unwrap_or(fallback_payload) + } else { + fallback_payload + }; + let body = + serde_json::to_vec(&payload).map_err(|err| GatewayError::Internal(err.to_string()))?; + let headers = BTreeMap::from([ + ("content-type".to_string(), "application/json".to_string()), + ("Retry-After".to_string(), rejection.retry_after.to_string()), + ( + "X-Daily-Usage-Limit-USD".to_string(), + format!("{:.8}", rejection.limit_usd), + ), + ( + "X-Daily-Usage-Used-USD".to_string(), + format!("{:.8}", rejection.used_usd), + ), + ( + "X-Daily-Usage-Remaining-USD".to_string(), + format!("{:.8}", rejection.remaining_usd), + ), + ( + "X-Daily-Usage-Scope".to_string(), + rejection.scope.to_string(), + ), + ( + "X-Daily-Usage-Reset".to_string(), + rejection.reset_at_unix_secs.to_string(), + ), + ]); + let mut response = build_client_response_from_parts( + StatusCode::TOO_MANY_REQUESTS.as_u16(), + &headers, + Body::from(body), + trace_id, + control_decision, + )?; + response.extensions_mut().insert(DailyUsageLimitedResponse); + Ok(response) +} + pub(crate) fn build_local_http_error_response( trace_id: &str, control_decision: Option<&GatewayControlDecision>, @@ -455,10 +524,12 @@ fn local_error_kind_for_status(status: StatusCode) -> LocalCoreSyncErrorKind { mod tests { use super::{ build_client_response_from_parts, build_local_auth_rejection_response, + build_local_daily_usage_limited_response, build_local_http_error_response_with_request_path, build_local_overloaded_response, build_local_user_rpm_limited_response, }; use crate::control::{GatewayControlDecision, GatewayLocalAuthRejection}; + use crate::daily_usage_limit::FrontdoorDailyUsageRejection; use crate::rate_limit::FrontdoorUserRpmRejection; use axum::body::{to_bytes, Body}; use std::collections::BTreeMap; @@ -547,6 +618,39 @@ mod tests { assert_eq!(overloaded["error"]["type"], "overloaded_error"); } + #[tokio::test] + async fn daily_usage_limit_response_has_dedicated_error_and_reset_headers() { + let response = build_local_daily_usage_limited_response( + "trace-daily-limit", + None, + &FrontdoorDailyUsageRejection { + scope: "user", + limit_usd: 10.0, + used_usd: 10.25, + remaining_usd: 0.0, + retry_after: 120, + reset_at_unix_secs: 1_800_000_000, + timezone: "Asia/Shanghai".to_string(), + }, + ) + .expect("daily limit response should build"); + + assert_eq!(response.status(), http::StatusCode::TOO_MANY_REQUESTS); + assert_eq!(response.headers()["retry-after"], "120"); + assert_eq!(response.headers()["x-daily-usage-limit-usd"], "10.00000000"); + assert_eq!(response.headers()["x-daily-usage-used-usd"], "10.25000000"); + assert_eq!( + response.headers()["x-daily-usage-remaining-usd"], + "0.00000000" + ); + assert_eq!(response.headers()["x-daily-usage-scope"], "user"); + assert_eq!(response.headers()["x-daily-usage-reset"], "1800000000"); + + let payload = response_json(response).await; + assert_eq!(payload["error"]["type"], "daily_usage_limit_exceeded"); + assert_eq!(payload["error"]["details"]["timezone"], "Asia/Shanghai"); + } + #[tokio::test] async fn claude_path_shapes_pre_control_http_errors_and_413() { for path in ["/v1/messages", "/v1/messages/count_tokens"] { diff --git a/apps/aether-gateway/src/app_timezone.rs b/apps/aether-gateway/src/app_timezone.rs new file mode 100644 index 000000000..09103e963 --- /dev/null +++ b/apps/aether-gateway/src/app_timezone.rs @@ -0,0 +1,104 @@ +use std::sync::LazyLock; + +use chrono::{DateTime, LocalResult, NaiveDate, TimeZone, Utc}; +use chrono_tz::Tz; +use tracing::warn; + +pub(crate) const DEFAULT_APP_TIMEZONE: &str = "Asia/Shanghai"; + +static APP_TIMEZONE: LazyLock = LazyLock::new(|| { + let configured = std::env::var("APP_TIMEZONE") + .ok() + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) + .unwrap_or_else(|| DEFAULT_APP_TIMEZONE.to_string()); + configured.parse().unwrap_or_else(|_| { + warn!( + timezone = %configured, + fallback = DEFAULT_APP_TIMEZONE, + "gateway APP_TIMEZONE invalid; falling back" + ); + DEFAULT_APP_TIMEZONE + .parse() + .expect("default application timezone should parse") + }) +}); + +pub(crate) fn app_timezone() -> Tz { + *APP_TIMEZONE +} + +pub(crate) fn local_day_window( + now_utc: DateTime, + timezone: Tz, +) -> (NaiveDate, DateTime, DateTime) { + let local_date = now_utc.with_timezone(&timezone).date_naive(); + let next_date = local_date + .succ_opt() + .expect("application local date should have a successor"); + ( + local_date, + local_midnight_utc(local_date, timezone), + local_midnight_utc(next_date, timezone), + ) +} + +fn local_midnight_utc(date: NaiveDate, timezone: Tz) -> DateTime { + let midnight = date + .and_hms_opt(0, 0, 0) + .expect("local midnight should be valid"); + match timezone.from_local_datetime(&midnight) { + LocalResult::Single(value) => value.with_timezone(&Utc), + LocalResult::Ambiguous(first, second) => first.min(second).with_timezone(&Utc), + LocalResult::None => { + for minute in 1..=180 { + let candidate = midnight + chrono::Duration::minutes(minute); + match timezone.from_local_datetime(&candidate) { + LocalResult::Single(value) => return value.with_timezone(&Utc), + LocalResult::Ambiguous(first, second) => { + return first.min(second).with_timezone(&Utc) + } + LocalResult::None => {} + } + } + panic!("local day start should resolve within three hours") + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn default_timezone_uses_shanghai_natural_day() { + let timezone: Tz = DEFAULT_APP_TIMEZONE.parse().unwrap(); + let (_, start, end) = local_day_window( + DateTime::parse_from_rfc3339("2026-08-03T12:00:00Z") + .unwrap() + .with_timezone(&Utc), + timezone, + ); + assert_eq!(start.to_rfc3339(), "2026-08-02T16:00:00+00:00"); + assert_eq!(end.to_rfc3339(), "2026-08-03T16:00:00+00:00"); + } + + #[test] + fn dst_days_use_natural_local_midnights() { + let timezone: Tz = "America/New_York".parse().unwrap(); + let spring = local_day_window( + DateTime::parse_from_rfc3339("2026-03-08T12:00:00Z") + .unwrap() + .with_timezone(&Utc), + timezone, + ); + assert_eq!((spring.2 - spring.1).num_hours(), 23); + let fall = local_day_window( + DateTime::parse_from_rfc3339("2026-11-01T12:00:00Z") + .unwrap() + .with_timezone(&Utc), + timezone, + ); + assert_eq!((fall.2 - fall.1).num_hours(), 25); + } +} diff --git a/apps/aether-gateway/src/cache/auth_context.rs b/apps/aether-gateway/src/cache/auth_context.rs index cc2c3e9dd..f9bc97319 100644 --- a/apps/aether-gateway/src/cache/auth_context.rs +++ b/apps/aether-gateway/src/cache/auth_context.rs @@ -417,8 +417,11 @@ mod tests { access_allowed: true, user_rate_limit: None, api_key_rate_limit: None, + user_daily_usage_limit_usd: None, + api_key_daily_usage_limit_usd: None, api_key_is_standalone: false, admin_bypass_limits: false, + ip_bypass_limits: false, local_rejection: None, allowed_models: None, ip_rules: None, diff --git a/apps/aether-gateway/src/cache/candidate_page.rs b/apps/aether-gateway/src/cache/candidate_page.rs index 6512fae8d..23b0e8d1a 100644 --- a/apps/aether-gateway/src/cache/candidate_page.rs +++ b/apps/aether-gateway/src/cache/candidate_page.rs @@ -542,6 +542,7 @@ mod tests { user_is_active: true, user_is_deleted: false, user_rate_limit: None, + user_daily_usage_limit_usd: None, user_allowed_providers: None, user_allowed_api_formats: None, user_allowed_models: None, @@ -551,6 +552,7 @@ mod tests { api_key_is_locked: false, api_key_is_standalone: false, api_key_rate_limit: None, + api_key_daily_usage_limit_usd: None, api_key_concurrent_limit: None, api_key_expires_at_unix_secs: None, api_key_allowed_providers: None, diff --git a/apps/aether-gateway/src/control/auth/gate.rs b/apps/aether-gateway/src/control/auth/gate.rs index 208c13cf3..fe603d815 100644 --- a/apps/aether-gateway/src/control/auth/gate.rs +++ b/apps/aether-gateway/src/control/auth/gate.rs @@ -239,6 +239,28 @@ async fn estimate_execution_plan_cost_upper_bound_usd( result } +pub(crate) async fn execution_plan_cost_is_proven_zero( + state: &AppState, + plan: &aether_contracts::ExecutionPlan, + report_context: Option<&serde_json::Value>, +) -> bool { + let model_id = report_context_string_field(report_context, "model_id"); + let global_model_name = report_context_string_field(report_context, "global_model_name"); + if model_id.is_some() || global_model_name.is_some() { + if let Ok(Some(context)) = + load_execution_plan_billing_context(state, plan, model_id, global_model_name).await + { + if aether_billing::BillingModelPricingSnapshot::from(context).is_free_tier() { + return true; + } + } + } + matches!( + estimate_execution_plan_cost_upper_bound_usd(state, plan, report_context).await, + Ok(Some(cost)) if cost <= DAILY_QUOTA_EPSILON_USD + ) +} + async fn estimate_execution_plan_cost_upper_bound_usd_inner( state: &AppState, plan: &aether_contracts::ExecutionPlan, @@ -832,9 +854,10 @@ mod tests { use serde_json::json; use super::{ - execution_plan_balance_capacity_rejection, execution_plan_cost_upper_bound_cache_key, - max_output_tokens_from_request, openai_request_input_is_self_contained, - output_choice_count_upper_bound, request_model_local_rejection, GatewayLocalAuthRejection, + execution_plan_balance_capacity_rejection, execution_plan_cost_is_proven_zero, + execution_plan_cost_upper_bound_cache_key, max_output_tokens_from_request, + openai_request_input_is_self_contained, output_choice_count_upper_bound, + request_model_local_rejection, GatewayLocalAuthRejection, }; use crate::control::{GatewayControlAuthContext, GatewayControlDecision}; use crate::data::GatewayDataState; @@ -919,8 +942,11 @@ mod tests { access_allowed: true, user_rate_limit: None, api_key_rate_limit: None, + user_daily_usage_limit_usd: None, + api_key_daily_usage_limit_usd: None, api_key_is_standalone: false, admin_bypass_limits: false, + ip_bypass_limits: false, local_rejection: None, allowed_models: Some(allowed_models), ip_rules: None, @@ -1829,6 +1855,59 @@ mod tests { assert_eq!(rejection, None); } + #[tokio::test] + async fn daily_usage_zero_cost_proof_recognizes_free_images_but_not_unknown_paid_images() { + let free_context = billing_context_with_pricing( + Some(json!({ + "tiers": [{ + "up_to": null, + "input_price_per_1m": 100.0, + "output_price_per_1m": 100.0 + }] + })), + None, + None, + Some("free_tier"), + ); + let paid_context = billing_context_with_pricing( + Some(json!({ + "tiers": [{ + "up_to": null, + "input_price_per_1m": 1.0, + "output_price_per_1m": 1.0 + }] + })), + None, + None, + None, + ); + let plan = execution_plan( + json!({ + "model": "gpt-image-1", + "prompt": "a small red circle" + }), + "openai:image", + ); + let report_context = billing_report_context(); + + assert!( + execution_plan_cost_is_proven_zero( + &state_with_quota_and_wallet(quota_availability(0.0, false), free_context,), + &plan, + Some(&report_context), + ) + .await + ); + assert!( + !execution_plan_cost_is_proven_zero( + &state_with_quota_and_wallet(quota_availability(0.0, false), paid_context,), + &plan, + Some(&report_context), + ) + .await + ); + } + #[tokio::test] async fn finalized_chat_output_fields_and_choice_count_bound_capacity() { let context = billing_context_with_pricing( diff --git a/apps/aether-gateway/src/control/auth/mod.rs b/apps/aether-gateway/src/control/auth/mod.rs index 1aa699566..ca21ed804 100644 --- a/apps/aether-gateway/src/control/auth/mod.rs +++ b/apps/aether-gateway/src/control/auth/mod.rs @@ -7,8 +7,9 @@ mod types; pub(crate) use credentials::extract_requested_model; pub(super) use credentials::resolve_gateway_credential_carrier; pub(crate) use gate::{ - execution_plan_balance_capacity_rejection, request_model_local_rejection, - should_buffer_request_for_local_auth, trusted_auth_local_rejection, GatewayLocalAuthRejection, + execution_plan_balance_capacity_rejection, execution_plan_cost_is_proven_zero, + request_model_local_rejection, should_buffer_request_for_local_auth, + trusted_auth_local_rejection, GatewayLocalAuthRejection, }; pub(crate) use resolution::{ refresh_execution_runtime_auth_context, resolve_execution_runtime_auth_context, diff --git a/apps/aether-gateway/src/control/auth/resolution.rs b/apps/aether-gateway/src/control/auth/resolution.rs index 6ebe21d03..71f8e9b19 100644 --- a/apps/aether-gateway/src/control/auth/resolution.rs +++ b/apps/aether-gateway/src/control/auth/resolution.rs @@ -84,10 +84,16 @@ pub(crate) struct GatewayControlAuthContext { #[serde(skip)] pub(crate) api_key_rate_limit: Option, #[serde(skip)] + pub(crate) user_daily_usage_limit_usd: Option, + #[serde(skip)] + pub(crate) api_key_daily_usage_limit_usd: Option, + #[serde(skip)] pub(crate) api_key_is_standalone: bool, #[serde(skip)] pub(crate) admin_bypass_limits: bool, #[serde(skip)] + pub(crate) ip_bypass_limits: bool, + #[serde(skip)] pub(crate) local_rejection: Option, #[serde(skip)] pub(crate) allowed_models: Option>, @@ -922,8 +928,11 @@ pub(super) async fn resolve_data_backed_auth_context( access_allowed: false, user_rate_limit: None, api_key_rate_limit: None, + user_daily_usage_limit_usd: None, + api_key_daily_usage_limit_usd: None, api_key_is_standalone: false, admin_bypass_limits: false, + ip_bypass_limits: false, local_rejection: Some(GatewayLocalAuthRejection::InvalidApiKey), allowed_models: None, ip_rules: None, @@ -1035,8 +1044,11 @@ async fn resolve_antigravity_bearer_bridge_auth_context( access_allowed: false, user_rate_limit: None, api_key_rate_limit: None, + user_daily_usage_limit_usd: None, + api_key_daily_usage_limit_usd: None, api_key_is_standalone: false, admin_bypass_limits: false, + ip_bypass_limits: false, local_rejection: Some(GatewayLocalAuthRejection::InvalidApiKey), allowed_models: None, ip_rules: None, @@ -1094,8 +1106,11 @@ async fn resolve_trusted_auth_context( access_allowed: false, user_rate_limit: None, api_key_rate_limit: None, + user_daily_usage_limit_usd: None, + api_key_daily_usage_limit_usd: None, api_key_is_standalone: false, admin_bypass_limits: false, + ip_bypass_limits: false, local_rejection: Some(GatewayLocalAuthRejection::InvalidApiKey), allowed_models: None, ip_rules: None, @@ -1191,9 +1206,12 @@ async fn build_data_backed_auth_context( access_allowed: key_access_allowed && local_rejection.is_none(), user_rate_limit: snapshot.user_rate_limit, api_key_rate_limit: snapshot.api_key_rate_limit, + user_daily_usage_limit_usd: snapshot.user_daily_usage_limit_usd, + api_key_daily_usage_limit_usd: snapshot.api_key_daily_usage_limit_usd, api_key_is_standalone: snapshot.api_key_is_standalone, admin_bypass_limits: snapshot.user_role.eq_ignore_ascii_case("admin") && !snapshot.api_key_is_standalone, + ip_bypass_limits: false, local_rejection, allowed_models, ip_rules: snapshot.api_key_ip_rules, diff --git a/apps/aether-gateway/src/control/mod.rs b/apps/aether-gateway/src/control/mod.rs index f62651d4a..c85023e80 100644 --- a/apps/aether-gateway/src/control/mod.rs +++ b/apps/aether-gateway/src/control/mod.rs @@ -8,8 +8,8 @@ mod public; mod route; pub(crate) use auth::{ - execution_plan_balance_capacity_rejection, extract_requested_model, - refresh_execution_runtime_auth_context, request_model_local_rejection, + execution_plan_balance_capacity_rejection, execution_plan_cost_is_proven_zero, + extract_requested_model, refresh_execution_runtime_auth_context, request_model_local_rejection, resolve_execution_runtime_auth_context, should_buffer_request_for_local_auth, trusted_auth_local_rejection, GatewayAdminPrincipalContext, GatewayControlAuthContext, GatewayCredentialCarrier, GatewayLocalAuthRejection, diff --git a/apps/aether-gateway/src/daily_usage_limit.rs b/apps/aether-gateway/src/daily_usage_limit.rs new file mode 100644 index 000000000..3abefabdc --- /dev/null +++ b/apps/aether-gateway/src/daily_usage_limit.rs @@ -0,0 +1,878 @@ +use std::collections::HashMap; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use aether_cache::ExpiringMap; +use aether_data_contracts::repository::usage::StoredRequestUsageAudit; +use aether_data_contracts::repository::usage::UsageDailyActualCostRollupQuery; +use aether_runtime_state::{ + DailyUsageLimitCountInput, DailyUsageLimitIncrementInput, DailyUsageLimitRestoreEntry, + DailyUsageLimitRestoreInput, RuntimeState, +}; +use chrono::{DateTime, SecondsFormat, Utc}; +use tracing::warn; + +use crate::app_timezone::{app_timezone, local_day_window}; +use crate::control::GatewayControlDecision; +use crate::stage_metrics::observe_gateway_stage_ms; +use crate::{AppState, GatewayError}; + +const SYSTEM_DAILY_USAGE_LIMIT_CONFIG_KEY: &str = "daily_usage_limit_usd"; +const SYSTEM_CONFIG_CACHE_TTL: Duration = Duration::from_secs(15); +const LIMIT_EPSILON_USD: f64 = 0.000_000_01; +const USD_UNITS_PER_DOLLAR: f64 = 100_000_000.0; +const COUNTER_EXPIRY_GRACE_SECONDS: u64 = 60; +const DAILY_USAGE_RUNTIME_STATE_KEY: &str = "daily_usage_limit:runtime_state"; +const DAILY_USAGE_RECOVERY_LOCK_KEY: &str = "daily_usage_limit:recovery"; +const DAILY_USAGE_RECOVERY_LOCK_OWNER: &str = "gateway-daily-usage-recovery"; +const DAILY_USAGE_RECOVERY_LOCK_TTL: Duration = Duration::from_secs(600); +const DAILY_USAGE_RECOVERY_RETRY_DELAY: Duration = Duration::from_secs(30); + +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct DailyUsageScopeStatus { + pub(crate) scope: &'static str, + pub(crate) limit_usd: f64, + pub(crate) used_usd: f64, + pub(crate) remaining_usd: f64, +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct FrontdoorDailyUsageStatus { + pub(crate) available: bool, + pub(crate) timezone: String, + pub(crate) window_start: String, + pub(crate) window_end: String, + pub(crate) reset_at_unix_secs: u64, + pub(crate) user: Option, + pub(crate) key: Option, +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct FrontdoorDailyUsageRejection { + pub(crate) scope: &'static str, + pub(crate) limit_usd: f64, + pub(crate) used_usd: f64, + pub(crate) remaining_usd: f64, + pub(crate) retry_after: u64, + pub(crate) reset_at_unix_secs: u64, + pub(crate) timezone: String, +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) enum FrontdoorDailyUsageOutcome { + NotApplicable, + Allowed, + Rejected(FrontdoorDailyUsageRejection), +} + +#[derive(Debug, Clone, Copy)] +pub(crate) struct DailyUsageLimitedResponse; + +#[derive(Debug, Clone)] +pub(crate) struct FrontdoorDailyUsageLimiter { + system_default_cache: Arc>, + recovery_inflight: Arc, + runtime_failures: Arc, + #[cfg(test)] + system_default_override: Arc>>, +} + +impl Default for FrontdoorDailyUsageLimiter { + fn default() -> Self { + Self::new() + } +} + +impl FrontdoorDailyUsageLimiter { + pub(crate) fn new() -> Self { + Self { + system_default_cache: Arc::new(ExpiringMap::default()), + recovery_inflight: Arc::new(AtomicBool::new(false)), + runtime_failures: Arc::new(AtomicU64::new(0)), + #[cfg(test)] + system_default_override: Arc::new(std::sync::Mutex::new(None)), + } + } + + pub(crate) fn clear_system_default_cache(&self) { + self.system_default_cache.clear(); + } + + pub(crate) fn runtime_failure_count(&self) -> u64 { + self.runtime_failures.load(Ordering::Relaxed) + } + + pub(crate) async fn check( + &self, + state: &AppState, + decision: &GatewayControlDecision, + ) -> FrontdoorDailyUsageOutcome { + let started_at = Instant::now(); + let status_result = self.current_status(state, decision).await; + observe_gateway_stage_ms( + "daily_usage_limit_total", + started_at.elapsed().as_millis() as u64, + ); + let status = match status_result { + Ok(Some(status)) => status, + Ok(None) => return FrontdoorDailyUsageOutcome::NotApplicable, + Err(err) => { + let failure_count = self.runtime_failures.fetch_add(1, Ordering::Relaxed) + 1; + let auth = decision.auth_context.as_ref(); + warn!( + event_name = "frontdoor_daily_usage_check_failed", + log_type = "ops", + error = ?err, + runtime_failures_total = failure_count, + user_id = auth.map(|auth| auth.user_id.as_str()).unwrap_or("-"), + api_key_id = auth.map(|auth| auth.api_key_id.as_str()).unwrap_or("-"), + "daily usage limit check failed; allowing request" + ); + return FrontdoorDailyUsageOutcome::Allowed; + } + }; + if !status.available { + return FrontdoorDailyUsageOutcome::Allowed; + } + let exceeded = status + .user + .as_ref() + .filter(|scope| scope.used_usd + LIMIT_EPSILON_USD >= scope.limit_usd) + .or_else(|| { + status + .key + .as_ref() + .filter(|scope| scope.used_usd + LIMIT_EPSILON_USD >= scope.limit_usd) + }); + let Some(exceeded) = exceeded else { + return FrontdoorDailyUsageOutcome::Allowed; + }; + let now = Utc::now().timestamp().max(0) as u64; + FrontdoorDailyUsageOutcome::Rejected(FrontdoorDailyUsageRejection { + scope: exceeded.scope, + limit_usd: exceeded.limit_usd, + used_usd: exceeded.used_usd, + remaining_usd: exceeded.remaining_usd, + retry_after: status.reset_at_unix_secs.saturating_sub(now).max(1), + reset_at_unix_secs: status.reset_at_unix_secs, + timezone: status.timezone, + }) + } + + pub(crate) async fn current_status( + &self, + state: &AppState, + decision: &GatewayControlDecision, + ) -> Result, GatewayError> { + let Some(auth) = decision.auth_context.as_ref() else { + return Ok(None); + }; + if decision.route_class.as_deref() != Some("ai_public") + || auth.local_rejection.is_some() + || auth.user_id.is_empty() + || auth.api_key_id.is_empty() + || auth.admin_bypass_limits + || auth.ip_bypass_limits + { + return Ok(None); + } + + let needs_system_default = if auth.api_key_is_standalone { + auth.api_key_daily_usage_limit_usd.is_none() + } else { + auth.user_daily_usage_limit_usd.is_none() + }; + let system_limit = if needs_system_default { + let config_started_at = Instant::now(); + let result = self.resolve_system_default_limit(state).await; + observe_gateway_stage_ms( + "daily_usage_limit_system_default", + config_started_at.elapsed().as_millis() as u64, + ); + result? + } else { + 0.0 + }; + let (user_limit, key_limit) = resolve_scope_limits( + auth.api_key_is_standalone, + auth.user_daily_usage_limit_usd, + auth.api_key_daily_usage_limit_usd, + system_limit, + ); + if user_limit.is_none() && key_limit.is_none() { + return Ok(None); + } + + let timezone = app_timezone(); + let now = Utc::now(); + let (_, start, end) = local_day_window(now, timezone); + let bucket = start.timestamp().max(0) as u64; + let user_scope_key = daily_usage_user_scope_key(&auth.user_id, bucket); + let key_scope_key = daily_usage_key_scope_key(&auth.api_key_id, bucket); + let runtime_started_at = Instant::now(); + let counts_result = state + .runtime_state + .daily_usage_limit_counts(DailyUsageLimitCountInput { + state_key: DAILY_USAGE_RUNTIME_STATE_KEY, + user_key: (!auth.api_key_is_standalone).then_some(user_scope_key.as_str()), + key_key: &key_scope_key, + bucket, + }) + .await + .map_err(|err| GatewayError::Internal(err.to_string())); + observe_gateway_stage_ms( + "daily_usage_limit_runtime_read", + runtime_started_at.elapsed().as_millis() as u64, + ); + let counts = counts_result?; + if !counts.state_ready { + self.trigger_runtime_recovery(state); + return Ok(Some(FrontdoorDailyUsageStatus { + available: false, + timezone: timezone.name().to_string(), + window_start: rfc3339(start), + window_end: rfc3339(end), + reset_at_unix_secs: end.timestamp().max(0) as u64, + user: None, + key: None, + })); + } + let user = user_limit + .map(|limit_usd| scope_status("user", limit_usd, units_to_usd(counts.user_units))); + let key = key_limit + .map(|limit_usd| scope_status("key", limit_usd, units_to_usd(counts.key_units))); + Ok(Some(FrontdoorDailyUsageStatus { + available: true, + timezone: timezone.name().to_string(), + window_start: rfc3339(start), + window_end: rfc3339(end), + reset_at_unix_secs: end.timestamp().max(0) as u64, + user, + key, + })) + } + + fn trigger_runtime_recovery(&self, state: &AppState) { + if self + .recovery_inflight + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Relaxed) + .is_err() + { + return; + } + let limiter = self.clone(); + let state = state.clone(); + tokio::spawn(async move { + let started_at = Instant::now(); + let result = recover_daily_usage_runtime(&state).await; + observe_gateway_stage_ms( + "daily_usage_limit_recovery", + started_at.elapsed().as_millis() as u64, + ); + match result { + Ok(true) => {} + Ok(false) => tokio::time::sleep(DAILY_USAGE_RECOVERY_RETRY_DELAY).await, + Err(err) => { + let failure_count = + limiter.runtime_failures.fetch_add(1, Ordering::Relaxed) + 1; + warn!( + event_name = "frontdoor_daily_usage_recovery_failed", + log_type = "ops", + error = ?err, + runtime_failures_total = failure_count, + "daily usage runtime recovery failed; limits remain fail-open" + ); + tokio::time::sleep(DAILY_USAGE_RECOVERY_RETRY_DELAY).await; + } + } + limiter.recovery_inflight.store(false, Ordering::Release); + }); + } + + async fn resolve_system_default_limit(&self, state: &AppState) -> Result { + #[cfg(test)] + if let Ok(guard) = self.system_default_override.lock() { + if let Some(limit) = *guard { + return Ok(limit); + } + } + if let Some(limit) = self + .system_default_cache + .get_fresh(SYSTEM_DAILY_USAGE_LIMIT_CONFIG_KEY, SYSTEM_CONFIG_CACHE_TTL) + { + return Ok(limit); + } + let limit = parse_system_limit( + state + .read_system_config_json_value(SYSTEM_DAILY_USAGE_LIMIT_CONFIG_KEY) + .await?, + )?; + self.system_default_cache.insert( + SYSTEM_DAILY_USAGE_LIMIT_CONFIG_KEY.to_string(), + limit, + SYSTEM_CONFIG_CACHE_TTL, + 8, + ); + Ok(limit) + } + + #[cfg(test)] + pub(crate) fn with_system_default_limit_for_tests(self, limit: f64) -> Self { + if let Ok(mut guard) = self.system_default_override.lock() { + *guard = Some(limit.max(0.0)); + } + self + } +} + +async fn recover_daily_usage_runtime(state: &AppState) -> Result { + let Some(lease) = state + .runtime_state + .lock_try_acquire( + DAILY_USAGE_RECOVERY_LOCK_KEY, + DAILY_USAGE_RECOVERY_LOCK_OWNER, + DAILY_USAGE_RECOVERY_LOCK_TTL, + ) + .await + .map_err(|err| GatewayError::Internal(err.to_string()))? + else { + return Ok(false); + }; + + let recovery_result = async { + state + .runtime_state + .kv_set(DAILY_USAGE_RUNTIME_STATE_KEY, "recovering", None) + .await + .map_err(|err| GatewayError::Internal(err.to_string()))?; + + let timezone = app_timezone(); + let now = Utc::now(); + let (_, start, end) = local_day_window(now, timezone); + let bucket = start.timestamp().max(0) as u64; + let rollups = state + .background_data + .summarize_usage_daily_actual_cost_rollups(&UsageDailyActualCostRollupQuery { + finalized_from_unix_secs: bucket, + finalized_until_unix_secs: end.timestamp().max(0) as u64, + }) + .await + .map_err(|err| GatewayError::Internal(err.to_string()))?; + + let mut user_totals = HashMap::::new(); + let mut key_totals = HashMap::, bool, f64)>::new(); + for rollup in rollups { + let Some(api_key_id) = non_empty(&rollup.api_key_id) else { + continue; + }; + let amount = rollup.actual_total_cost_usd; + if !amount.is_finite() || amount <= 0.0 { + continue; + } + let user_id = rollup + .user_id + .as_deref() + .and_then(non_empty) + .map(ToOwned::to_owned); + if !rollup.api_key_is_standalone { + if let Some(user_id) = user_id.as_ref() { + *user_totals.entry(user_id.clone()).or_default() += amount; + } + } + let key_total = key_totals.entry(api_key_id.to_string()).or_insert(( + user_id.clone(), + rollup.api_key_is_standalone, + 0.0, + )); + key_total.2 += amount; + } + + let entries = key_totals + .into_iter() + .map(|(api_key_id, (user_id, is_standalone, key_total))| { + let user_id = (!is_standalone).then_some(user_id).flatten(); + let user_units = user_id + .as_ref() + .and_then(|user_id| user_totals.get(user_id)) + .copied() + .map(usd_to_units) + .unwrap_or_default(); + DailyUsageLimitRestoreEntry { + user_key: user_id + .as_deref() + .map(|user_id| daily_usage_user_scope_key(user_id, bucket)), + key_key: daily_usage_key_scope_key(&api_key_id, bucket), + user_units, + key_units: usd_to_units(key_total), + } + }) + .collect::>(); + let ttl_seconds = (end.timestamp().max(0) as u64) + .saturating_sub(Utc::now().timestamp().max(0) as u64) + .saturating_add(COUNTER_EXPIRY_GRACE_SECONDS) + .max(1); + state + .runtime_state + .restore_daily_usage_limits(DailyUsageLimitRestoreInput { + entries: &entries, + bucket, + ttl_seconds, + }) + .await + .map_err(|err| GatewayError::Internal(err.to_string()))?; + state + .runtime_state + .kv_set(DAILY_USAGE_RUNTIME_STATE_KEY, "ready", None) + .await + .map_err(|err| GatewayError::Internal(err.to_string()))?; + Ok::<(), GatewayError>(()) + } + .await; + + if let Err(err) = state.runtime_state.lock_release(&lease).await { + warn!( + event_name = "frontdoor_daily_usage_recovery_lock_release_failed", + log_type = "ops", + error = ?err, + "daily usage runtime recovery lock release failed" + ); + } + recovery_result.map(|()| true) +} + +pub(crate) async fn record_finalized_daily_usage( + runtime_state: &RuntimeState, + usage: &StoredRequestUsageAudit, +) -> Result<(), aether_runtime_state::DataLayerError> { + if usage.status != "completed" { + return Ok(()); + } + let amount_units = usd_to_units(usage.actual_total_cost_usd); + if amount_units == 0 { + return Ok(()); + } + let Some(api_key_id) = usage.api_key_id.as_deref().and_then(non_empty) else { + return Ok(()); + }; + let finalized_at = usage + .finalized_at_unix_secs + .unwrap_or(usage.updated_at_unix_secs); + let Some(finalized_at) = DateTime::::from_timestamp(finalized_at as i64, 0) else { + return Ok(()); + }; + let timezone = app_timezone(); + let (_, start, end) = local_day_window(finalized_at, timezone); + let bucket = start.timestamp().max(0) as u64; + let key_scope_key = daily_usage_key_scope_key(api_key_id, bucket); + let is_standalone = usage + .request_metadata + .as_ref() + .and_then(|metadata| metadata.get("api_key_is_standalone")) + .and_then(serde_json::Value::as_bool) + .unwrap_or(false); + let user_scope_key = (!is_standalone) + .then(|| usage.user_id.as_deref().and_then(non_empty)) + .flatten() + .map(|user_id| daily_usage_user_scope_key(user_id, bucket)); + let now = Utc::now().timestamp().max(0) as u64; + let ttl_seconds = (end.timestamp().max(0) as u64) + .saturating_sub(now) + .saturating_add(COUNTER_EXPIRY_GRACE_SECONDS) + .max(1); + runtime_state + .increment_daily_usage_limit(DailyUsageLimitIncrementInput { + user_key: user_scope_key.as_deref(), + key_key: &key_scope_key, + bucket, + amount_units, + ttl_seconds, + }) + .await?; + Ok(()) +} + +fn scope_status(scope: &'static str, limit_usd: f64, used_usd: f64) -> DailyUsageScopeStatus { + DailyUsageScopeStatus { + scope, + limit_usd, + used_usd, + remaining_usd: (limit_usd - used_usd).max(0.0), + } +} + +fn daily_usage_user_scope_key(user_id: &str, bucket: u64) -> String { + format!("daily_usage_limit:user:{user_id}:{bucket}") +} + +fn daily_usage_key_scope_key(api_key_id: &str, bucket: u64) -> String { + format!("daily_usage_limit:key:{api_key_id}:{bucket}") +} + +fn usd_to_units(value: f64) -> u64 { + if !value.is_finite() || value <= 0.0 { + return 0; + } + (value * USD_UNITS_PER_DOLLAR) + .round() + .clamp(0.0, u64::MAX as f64) as u64 +} + +fn units_to_usd(value: u64) -> f64 { + value as f64 / USD_UNITS_PER_DOLLAR +} + +fn non_empty(value: &str) -> Option<&str> { + let value = value.trim(); + (!value.is_empty()).then_some(value) +} + +pub(crate) fn parse_system_limit(value: Option) -> Result { + let limit = match value { + None | Some(serde_json::Value::Null) => 0.0, + Some(serde_json::Value::Number(value)) => value.as_f64().ok_or_else(|| { + GatewayError::Internal("invalid system config daily_usage_limit_usd".to_string()) + })?, + Some(serde_json::Value::String(value)) => value.parse::().map_err(|_| { + GatewayError::Internal("invalid system config daily_usage_limit_usd".to_string()) + })?, + Some(_) => { + return Err(GatewayError::Internal( + "invalid system config daily_usage_limit_usd".to_string(), + )) + } + }; + if !limit.is_finite() || limit < 0.0 { + return Err(GatewayError::Internal( + "invalid system config daily_usage_limit_usd".to_string(), + )); + } + Ok(limit) +} + +fn positive_limit(value: f64) -> Option { + (value.is_finite() && value > 0.0).then_some(value) +} + +fn resolve_scope_limits( + is_standalone: bool, + user_limit: Option, + key_limit: Option, + system_limit: f64, +) -> (Option, Option) { + if is_standalone { + (None, positive_limit(key_limit.unwrap_or(system_limit))) + } else { + ( + positive_limit(user_limit.unwrap_or(system_limit)), + key_limit.and_then(positive_limit), + ) + } +} + +fn rfc3339(value: DateTime) -> String { + value.to_rfc3339_opts(SecondsFormat::Secs, true) +} + +#[cfg(test)] +mod tests { + use super::{ + daily_usage_key_scope_key, daily_usage_user_scope_key, parse_system_limit, positive_limit, + record_finalized_daily_usage, recover_daily_usage_runtime, resolve_scope_limits, + usd_to_units, FrontdoorDailyUsageLimiter, FrontdoorDailyUsageOutcome, + DAILY_USAGE_RUNTIME_STATE_KEY, + }; + use crate::control::{GatewayControlAuthContext, GatewayControlDecision}; + use crate::AppState; + use aether_data::repository::usage::InMemoryUsageReadRepository; + use aether_data_contracts::repository::usage::StoredRequestUsageAudit; + use aether_runtime_state::DailyUsageLimitIncrementInput; + use std::sync::Arc; + + fn sample_decision(user_limit: Option, key_limit: Option) -> GatewayControlDecision { + let mut decision = GatewayControlDecision::synthetic( + "/v1/chat/completions", + Some("ai_public".to_string()), + Some("openai".to_string()), + Some("chat".to_string()), + Some("openai:chat".to_string()), + ); + decision.auth_context = Some(GatewayControlAuthContext { + user_id: "user-1".to_string(), + api_key_id: "key-1".to_string(), + username: Some("alice".to_string()), + api_key_name: Some("default".to_string()), + balance_remaining: None, + access_allowed: true, + user_rate_limit: None, + api_key_rate_limit: None, + user_daily_usage_limit_usd: user_limit, + api_key_daily_usage_limit_usd: key_limit, + api_key_is_standalone: false, + admin_bypass_limits: false, + ip_bypass_limits: false, + local_rejection: None, + allowed_models: None, + ip_rules: None, + }); + decision + } + + async fn state_with_daily_usage(actual_cost_usd: f64) -> AppState { + let state = AppState::new().expect("state should build"); + let now = chrono::Utc::now(); + let (_, start, end) = + crate::app_timezone::local_day_window(now, crate::app_timezone::app_timezone()); + let bucket = start.timestamp().max(0) as u64; + let user_key = daily_usage_user_scope_key("user-1", bucket); + let key_key = daily_usage_key_scope_key("key-1", bucket); + state + .runtime_state + .increment_daily_usage_limit(DailyUsageLimitIncrementInput { + user_key: Some(&user_key), + key_key: &key_key, + bucket, + amount_units: usd_to_units(actual_cost_usd), + ttl_seconds: (end - now).num_seconds().max(1) as u64, + }) + .await + .expect("daily usage counter should update"); + state + .runtime_state + .kv_set(DAILY_USAGE_RUNTIME_STATE_KEY, "ready", None) + .await + .expect("daily usage runtime should be ready"); + state + } + + fn finalized_usage( + request_id: &str, + user_id: &str, + api_key_id: &str, + actual_cost_usd: f64, + ) -> StoredRequestUsageAudit { + let now = chrono::Utc::now(); + StoredRequestUsageAudit::new( + format!("usage-{request_id}"), + request_id.to_string(), + Some(user_id.to_string()), + Some(api_key_id.to_string()), + None, + None, + "OpenAI".to_string(), + "gpt-5".to_string(), + None, + None, + None, + None, + Some("chat".to_string()), + Some("openai:chat".to_string()), + Some("openai".to_string()), + Some("chat".to_string()), + Some("openai:chat".to_string()), + Some("openai".to_string()), + Some("chat".to_string()), + false, + false, + 10, + 10, + 20, + actual_cost_usd, + actual_cost_usd, + Some(200), + None, + None, + Some(100), + Some(20), + "completed".to_string(), + "settled".to_string(), + now.timestamp_millis(), + now.timestamp(), + Some(now.timestamp()), + ) + .expect("usage should build") + } + + #[test] + fn system_limit_defaults_to_unlimited_and_accepts_numbers_or_strings() { + assert_eq!(parse_system_limit(None).unwrap(), 0.0); + assert_eq!( + parse_system_limit(Some(serde_json::Value::Null)).unwrap(), + 0.0 + ); + assert_eq!( + parse_system_limit(Some(serde_json::json!(12.5))).unwrap(), + 12.5 + ); + assert_eq!( + parse_system_limit(Some(serde_json::json!("8.25"))).unwrap(), + 8.25 + ); + } + + #[test] + fn system_limit_rejects_invalid_or_negative_values() { + for value in [ + serde_json::json!(-1), + serde_json::json!("invalid"), + serde_json::json!({ "limit": 1 }), + ] { + assert!(parse_system_limit(Some(value)).is_err()); + } + } + + #[test] + fn zero_is_unlimited_and_positive_values_are_limits() { + assert_eq!(positive_limit(0.0), None); + assert_eq!(positive_limit(-1.0), None); + assert_eq!(positive_limit(f64::NAN), None); + assert_eq!(positive_limit(0.01), Some(0.01)); + } + + #[test] + fn normal_key_limit_only_adds_a_narrower_key_scope() { + assert_eq!( + resolve_scope_limits(false, None, None, 10.0), + (Some(10.0), None) + ); + assert_eq!( + resolve_scope_limits(false, Some(20.0), Some(5.0), 10.0), + (Some(20.0), Some(5.0)) + ); + assert_eq!( + resolve_scope_limits(false, Some(20.0), Some(0.0), 10.0), + (Some(20.0), None) + ); + } + + #[test] + fn standalone_key_inherits_or_explicitly_overrides_system_limit() { + assert_eq!( + resolve_scope_limits(true, None, None, 10.0), + (None, Some(10.0)) + ); + assert_eq!( + resolve_scope_limits(true, None, Some(0.0), 10.0), + (None, None) + ); + assert_eq!( + resolve_scope_limits(true, None, Some(3.0), 10.0), + (None, Some(3.0)) + ); + } + + #[tokio::test] + async fn daily_usage_below_limit_is_allowed_and_at_limit_is_rejected() { + let decision = sample_decision(Some(1.0), None); + + assert_eq!( + FrontdoorDailyUsageLimiter::new() + .check(&state_with_daily_usage(0.99).await, &decision) + .await, + FrontdoorDailyUsageOutcome::Allowed + ); + match FrontdoorDailyUsageLimiter::new() + .check(&state_with_daily_usage(1.0).await, &decision) + .await + { + FrontdoorDailyUsageOutcome::Rejected(rejection) => { + assert_eq!(rejection.scope, "user"); + assert_eq!(rejection.limit_usd, 1.0); + assert_eq!(rejection.used_usd, 1.0); + assert_eq!(rejection.remaining_usd, 0.0); + } + other => panic!("expected daily usage rejection, got {other:?}"), + } + } + + #[tokio::test] + async fn usage_is_accumulated_before_a_limit_is_enabled() { + let state = AppState::new().expect("state should build"); + record_finalized_daily_usage( + &state.runtime_state, + &finalized_usage("request-before-limit", "user-1", "key-1", 1.0), + ) + .await + .expect("daily usage should be recorded while unlimited"); + state + .runtime_state + .kv_set(DAILY_USAGE_RUNTIME_STATE_KEY, "ready", None) + .await + .expect("daily usage runtime should be ready"); + + match FrontdoorDailyUsageLimiter::new() + .check(&state, &sample_decision(Some(1.0), None)) + .await + { + FrontdoorDailyUsageOutcome::Rejected(rejection) => { + assert_eq!(rejection.scope, "user"); + assert_eq!(rejection.used_usd, 1.0); + } + other => panic!("expected accumulated usage rejection, got {other:?}"), + } + } + + #[tokio::test] + async fn lost_daily_usage_runtime_is_restored_with_one_grouped_recovery() { + let usage_repository = Arc::new(InMemoryUsageReadRepository::seed([finalized_usage( + "request-restore", + "user-1", + "key-1", + 1.25, + )])); + let state = AppState::new() + .expect("state should build") + .with_usage_data_reader_for_tests(usage_repository); + let decision = sample_decision(Some(1.0), None); + + assert!(recover_daily_usage_runtime(&state) + .await + .expect("daily usage runtime recovery")); + + match FrontdoorDailyUsageLimiter::new() + .check(&state, &decision) + .await + { + FrontdoorDailyUsageOutcome::Rejected(rejection) => { + assert_eq!(rejection.scope, "user"); + assert_eq!(rejection.used_usd, 1.25); + } + other => panic!("expected restored daily usage rejection, got {other:?}"), + } + } + + #[tokio::test] + async fn key_scope_can_narrow_a_normal_user_limit() { + let limiter = FrontdoorDailyUsageLimiter::new(); + let decision = sample_decision(Some(10.0), Some(0.5)); + + match limiter + .check(&state_with_daily_usage(0.5).await, &decision) + .await + { + FrontdoorDailyUsageOutcome::Rejected(rejection) => { + assert_eq!(rejection.scope, "key"); + assert_eq!(rejection.limit_usd, 0.5); + } + other => panic!("expected key daily usage rejection, got {other:?}"), + } + } + + #[tokio::test] + async fn admin_and_ip_bypass_skip_daily_usage_runtime_reads() { + for field in ["admin", "ip"] { + let limiter = FrontdoorDailyUsageLimiter::new(); + let mut decision = sample_decision(Some(1.0), None); + let auth = decision.auth_context.as_mut().unwrap(); + auth.admin_bypass_limits = field == "admin"; + auth.ip_bypass_limits = field == "ip"; + + assert_eq!( + limiter + .check(&AppState::new().expect("state should build"), &decision) + .await, + FrontdoorDailyUsageOutcome::NotApplicable + ); + } + } +} diff --git a/apps/aether-gateway/src/data/state/auth.rs b/apps/aether-gateway/src/data/state/auth.rs index be7308c0d..4b4d489db 100644 --- a/apps/aether-gateway/src/data/state/auth.rs +++ b/apps/aether-gateway/src/data/state/auth.rs @@ -1810,10 +1810,12 @@ impl GatewayDataState { allowed_models, } = resolve_group_effective_list_policies(&groups); let user_rate_limit = resolve_effective_rate_limit_policy(None, "system", &groups); + let user_daily_usage_limit_usd = resolve_group_effective_daily_usage_limit_policy(&groups); snapshot.user_allowed_providers = allowed_providers; snapshot.user_allowed_api_formats = allowed_api_formats; snapshot.user_allowed_models = allowed_models; snapshot.user_rate_limit = user_rate_limit; + snapshot.user_daily_usage_limit_usd = user_daily_usage_limit_usd; Ok(Some(GatewayAuthApiKeySnapshot::from_stored( snapshot, now_unix_secs, @@ -2016,6 +2018,23 @@ fn resolve_effective_rate_limit_policy( rate_limit_policy_value(intersect_rate_limit_policies(group_policy, user_policy)) } +pub(crate) fn resolve_group_effective_daily_usage_limit_policy( + groups: &[aether_data::repository::users::StoredUserGroup], +) -> Option { + let mut limit: Option = None; + for group in groups + .iter() + .filter(|group| group.daily_usage_limit_mode == "custom") + { + let value = group.daily_usage_limit_usd.unwrap_or(0.0).max(0.0); + if value == 0.0 { + return Some(0.0); + } + limit = Some(limit.map_or(value, |current| current.max(value))); + } + limit +} + #[derive(Clone, Copy, Debug, Eq, PartialEq)] enum RateLimitRestriction { Unlimited, @@ -2215,11 +2234,86 @@ mod tests { allowed_models_mode: allowed_models_mode.to_string(), rate_limit, rate_limit_mode: rate_limit_mode.to_string(), + daily_usage_limit_usd: None, + daily_usage_limit_mode: "inherit".to_string(), created_at: None, updated_at: None, } } + fn with_daily_usage_limit( + mut group: StoredUserGroup, + value: Option, + mode: &str, + ) -> StoredUserGroup { + group.daily_usage_limit_usd = value; + group.daily_usage_limit_mode = mode.to_string(); + group + } + + #[test] + fn daily_usage_group_policy_falls_back_without_custom_groups() { + let groups = vec![ + with_daily_usage_limit( + sample_group("system", 0, None, "unrestricted", None, "system"), + Some(10.0), + "system", + ), + with_daily_usage_limit( + sample_group("inherit", 0, None, "unrestricted", None, "system"), + Some(5.0), + "inherit", + ), + ]; + + assert_eq!( + resolve_group_effective_daily_usage_limit_policy(&groups), + None + ); + } + + #[test] + fn daily_usage_group_policy_uses_highest_custom_grant() { + let groups = vec![ + with_daily_usage_limit( + sample_group("basic", 0, None, "unrestricted", None, "system"), + Some(10.0), + "custom", + ), + with_daily_usage_limit( + sample_group("premium", 0, None, "unrestricted", None, "system"), + Some(25.0), + "custom", + ), + ]; + + assert_eq!( + resolve_group_effective_daily_usage_limit_policy(&groups), + Some(25.0) + ); + } + + #[test] + fn daily_usage_group_policy_treats_any_custom_zero_as_unlimited() { + let groups = vec![ + with_daily_usage_limit( + sample_group("limited", 0, None, "unrestricted", None, "system"), + Some(25.0), + "custom", + ), + with_daily_usage_limit( + sample_group("unlimited", 0, None, "unrestricted", None, "system"), + Some(0.0), + "custom", + ), + ]; + + assert_eq!( + resolve_group_effective_daily_usage_limit_policy(&groups), + Some(0.0) + ); + } + #[test] fn list_policy_intersects_unrestricted_group_union_with_user_restriction() { let groups = vec![ @@ -2485,6 +2579,8 @@ mod tests { allowed_models_mode: "specific".to_string(), rate_limit: Some(1), rate_limit_mode: "custom".to_string(), + daily_usage_limit_usd: None, + daily_usage_limit_mode: "inherit".to_string(), }) .await .expect("group should create") @@ -2602,6 +2698,8 @@ mod tests { allowed_models_mode: "specific".to_string(), rate_limit: Some(30), rate_limit_mode: "custom".to_string(), + daily_usage_limit_usd: None, + daily_usage_limit_mode: "inherit".to_string(), }) .await .expect("group should create") @@ -2675,6 +2773,8 @@ mod tests { allowed_models_mode: "unrestricted".to_string(), rate_limit: None, rate_limit_mode: "system".to_string(), + daily_usage_limit_usd: None, + daily_usage_limit_mode: "inherit".to_string(), }) .await .expect("group should create") diff --git a/apps/aether-gateway/src/data/state/core.rs b/apps/aether-gateway/src/data/state/core.rs index 992fc6a02..a1dc52035 100644 --- a/apps/aether-gateway/src/data/state/core.rs +++ b/apps/aether-gateway/src/data/state/core.rs @@ -2,7 +2,7 @@ use aether_data::{DataBackends, DataLayerError, DatabaseDriver}; use aether_data_contracts::repository::candidate_selection::MinimalCandidateSelectionReadRepository; use aether_data_contracts::repository::candidates::RequestCandidateReadRepository; use aether_data_contracts::repository::provider_catalog::ProviderCatalogReadRepository; -use aether_runtime_state::RuntimeQueueStore; +use aether_runtime_state::{RuntimeQueueStore, RuntimeState}; use std::sync::Arc; use std::time::Duration; @@ -433,6 +433,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, wallet_reader: None, @@ -550,6 +551,7 @@ impl GatewayDataState { user_reader, user_preferences: None, usage_worker_queue, + daily_usage_runtime_state: None, video_task_reader, video_task_writer, wallet_reader, @@ -565,11 +567,10 @@ impl GatewayDataState { self.backends.is_some() } - pub(crate) fn with_usage_worker_queue( - mut self, - queue: Option>, - ) -> Self { - self.usage_worker_queue = queue; + pub(crate) fn with_usage_runtime_state(mut self, runtime_state: Arc) -> Self { + let queue: Arc = runtime_state.clone(); + self.usage_worker_queue = Some(queue); + self.daily_usage_runtime_state = Some(runtime_state); self } diff --git a/apps/aether-gateway/src/data/state/integrations.rs b/apps/aether-gateway/src/data/state/integrations.rs index 5ad389012..91d920fed 100644 --- a/apps/aether-gateway/src/data/state/integrations.rs +++ b/apps/aether-gateway/src/data/state/integrations.rs @@ -24,6 +24,7 @@ use aether_usage_runtime::{ use aether_video_tasks_core::StoredVideoTaskReadSide; use async_trait::async_trait; use serde_json::Value; +use tracing::warn; use super::GatewayDataState; use crate::data::candidate_selection::MinimalCandidateSelectionRowSource; @@ -341,7 +342,25 @@ impl UsageRecordWriter for GatewayDataState { &self, record: UpsertUsageRecord, ) -> Result, DataLayerError> { - GatewayDataState::upsert_usage(self, record).await + let stored = GatewayDataState::upsert_usage(self, record).await?; + if let (Some(runtime_state), Some(usage)) = + (self.daily_usage_runtime_state.as_ref(), stored.as_ref()) + { + if let Err(err) = + crate::daily_usage_limit::record_finalized_daily_usage(runtime_state, usage).await + { + warn!( + event_name = "daily_usage_limit_increment_failed", + log_type = "ops", + request_id = %usage.request_id, + user_id = usage.user_id.as_deref().unwrap_or("-"), + api_key_id = usage.api_key_id.as_deref().unwrap_or("-"), + error = ?err, + "daily usage limit increment failed; usage recording continues" + ); + } + } + Ok(stored) } async fn upsert_first_byte_usage_record( diff --git a/apps/aether-gateway/src/data/state/mod.rs b/apps/aether-gateway/src/data/state/mod.rs index 685639d4f..f5a3712ed 100644 --- a/apps/aether-gateway/src/data/state/mod.rs +++ b/apps/aether-gateway/src/data/state/mod.rs @@ -140,14 +140,15 @@ use aether_data_contracts::repository::settlement::{ }; use aether_data_contracts::repository::usage::{ ApiKeyLastUsedDelta, ManagementTokenCounterDelta, PendingUsageCleanupSummary, - ProxyNodeCounterDelta, StoredProviderUsageSummary, StoredRequestUsageAudit, UpsertUsageRecord, + ProxyNodeCounterDelta, StoredProviderUsageSummary, StoredRequestUsageAudit, + StoredUsageDailyActualCostRollup, UpsertUsageRecord, UsageDailyActualCostRollupQuery, UsageReadRepository, UsageWriteRepository, }; use aether_data_contracts::repository::video_tasks::{ StoredVideoTask, UpsertVideoTask, VideoTaskLookupKey, VideoTaskModelCount, VideoTaskQueryFilter, VideoTaskReadRepository, VideoTaskStatusCount, VideoTaskWriteRepository, }; -use aether_runtime_state::RuntimeQueueStore; +use aether_runtime_state::{RuntimeQueueStore, RuntimeState}; pub(crate) use self::referrals::{ ReferralAdminStats, ReferralMutationStatus, ReferralRelationshipListQuery, @@ -194,6 +195,7 @@ pub(crate) struct GatewayDataState { user_reader: Option>, user_preferences: Option>>>, usage_worker_queue: Option>, + daily_usage_runtime_state: Option>, video_task_reader: Option>, video_task_writer: Option>, wallet_reader: Option>, @@ -376,6 +378,7 @@ impl fmt::Debug for GatewayDataState { } mod auth; +pub(crate) use auth::resolve_group_effective_daily_usage_limit_policy; mod auth_api_key_cache; mod candidate_cache; mod catalog; diff --git a/apps/aether-gateway/src/data/state/runtime.rs b/apps/aether-gateway/src/data/state/runtime.rs index 96cb79121..398a02882 100644 --- a/apps/aether-gateway/src/data/state/runtime.rs +++ b/apps/aether-gateway/src/data/state/runtime.rs @@ -37,10 +37,10 @@ use super::{ }; use aether_data_contracts::repository::usage::{ PendingUsageCleanupSummary, ProviderApiKeyWindowUsageRequest, - StoredProviderApiKeyWindowUsageSummary, StoredUsageDailySummary, UsageAuditListQuery, - UsageCleanupExecutionMode, UsageCleanupSummary, UsageCleanupTargets, UsageCleanupWindow, - UsageCounterFlushSummary, UsageCounterHealthSnapshot, UsageCounterPendingHealthSnapshot, - UsageDailyHeatmapQuery, + StoredProviderApiKeyWindowUsageSummary, StoredUsageDailyActualCostRollup, + StoredUsageDailySummary, UsageAuditListQuery, UsageCleanupExecutionMode, UsageCleanupSummary, + UsageCleanupTargets, UsageCleanupWindow, UsageCounterFlushSummary, UsageCounterHealthSnapshot, + UsageCounterPendingHealthSnapshot, UsageDailyActualCostRollupQuery, UsageDailyHeatmapQuery, }; use aether_runtime_state::RuntimeQueueStore; use aether_video_tasks_core::read_data_backed_video_task_response; @@ -1534,6 +1534,22 @@ impl GatewayDataState { } } + pub(crate) async fn summarize_usage_daily_actual_cost_rollups( + &self, + query: &UsageDailyActualCostRollupQuery, + ) -> Result, DataLayerError> { + match &self.usage_reader { + Some(repository) => { + repository + .summarize_usage_daily_actual_cost_rollups(query) + .await + } + None => Err(DataLayerError::InvalidConfiguration( + "daily usage recovery requires a usage reader".to_string(), + )), + } + } + pub(crate) async fn summarize_usage_cache_affinity_hit_summary( &self, query: &aether_data_contracts::repository::usage::UsageCacheAffinityHitSummaryQuery, diff --git a/apps/aether-gateway/src/data/state/testing/announcements.rs b/apps/aether-gateway/src/data/state/testing/announcements.rs index 1d1ac7b95..4e5e768c2 100644 --- a/apps/aether-gateway/src/data/state/testing/announcements.rs +++ b/apps/aether-gateway/src/data/state/testing/announcements.rs @@ -45,6 +45,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -103,6 +104,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, diff --git a/apps/aether-gateway/src/data/state/testing/mod.rs b/apps/aether-gateway/src/data/state/testing/mod.rs index 486d6f435..af16ecb8e 100644 --- a/apps/aether-gateway/src/data/state/testing/mod.rs +++ b/apps/aether-gateway/src/data/state/testing/mod.rs @@ -81,6 +81,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -138,6 +139,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -191,6 +193,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -344,6 +347,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -425,6 +429,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -487,6 +492,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -558,6 +564,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -611,6 +618,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -665,6 +673,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -730,6 +739,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -797,6 +807,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -848,6 +859,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -924,6 +936,7 @@ impl GatewayDataState { user_reader: Some(repository), user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -983,6 +996,7 @@ impl GatewayDataState { user_reader: Some(user_repository), user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -1043,6 +1057,7 @@ impl GatewayDataState { user_reader: Some(user_repository), user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -1107,6 +1122,7 @@ impl GatewayDataState { user_reader: Some(user_repository), user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -1172,6 +1188,7 @@ impl GatewayDataState { user_reader: Some(user_repository), user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -1236,6 +1253,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -1289,6 +1307,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -1342,6 +1361,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -1407,6 +1427,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -1477,6 +1498,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -1530,6 +1552,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -1588,6 +1611,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -1663,6 +1687,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -1733,6 +1758,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -1787,6 +1813,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -1841,6 +1868,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -1897,6 +1925,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -1951,6 +1980,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -2005,6 +2035,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -2059,6 +2090,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -2121,6 +2153,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -2184,6 +2217,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -2250,6 +2284,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -2322,6 +2357,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -2395,6 +2431,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -2472,6 +2509,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -2556,6 +2594,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -2622,6 +2661,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -2679,6 +2719,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -2732,6 +2773,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -2791,6 +2833,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -2854,6 +2897,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -2918,6 +2962,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, @@ -2975,6 +3020,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, background_task_reader: None, diff --git a/apps/aether-gateway/src/data/state/testing/video_tasks.rs b/apps/aether-gateway/src/data/state/testing/video_tasks.rs index 9db104d98..06c3c3e4a 100644 --- a/apps/aether-gateway/src/data/state/testing/video_tasks.rs +++ b/apps/aether-gateway/src/data/state/testing/video_tasks.rs @@ -49,6 +49,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: Some(repository), video_task_writer: None, background_task_reader: None, @@ -109,6 +110,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: Some(video_task_reader), video_task_writer: Some(video_task_writer), background_task_reader: None, @@ -166,6 +168,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: Some(video_task_reader), video_task_writer: Some(video_task_writer), background_task_reader: None, @@ -227,6 +230,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: Some(video_task_reader), video_task_writer: Some(video_task_writer), background_task_reader: None, @@ -292,6 +296,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: Some(video_task_reader), video_task_writer: Some(video_task_writer), background_task_reader: None, @@ -366,6 +371,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: Some(video_task_reader), video_task_writer: Some(video_task_writer), background_task_reader: None, diff --git a/apps/aether-gateway/src/data/state/testkit.rs b/apps/aether-gateway/src/data/state/testkit.rs index 0e652582d..6029b2a22 100644 --- a/apps/aether-gateway/src/data/state/testkit.rs +++ b/apps/aether-gateway/src/data/state/testkit.rs @@ -78,6 +78,7 @@ impl GatewayDataState { user_reader: None, user_preferences: None, usage_worker_queue: None, + daily_usage_runtime_state: None, video_task_reader: None, video_task_writer: None, wallet_reader: None, diff --git a/apps/aether-gateway/src/executor/candidate_loop.rs b/apps/aether-gateway/src/executor/candidate_loop.rs index 6149bc60b..3cc37965b 100644 --- a/apps/aether-gateway/src/executor/candidate_loop.rs +++ b/apps/aether-gateway/src/executor/candidate_loop.rs @@ -13,8 +13,9 @@ use async_trait::async_trait; use axum::body::Body; use axum::http::Response; use futures_util::StreamExt; +use tokio::sync::OnceCell; use tokio::time::{timeout, Duration, Instant}; -use tracing::{debug, warn, Instrument}; +use tracing::{debug, info, warn, Instrument}; use crate::ai_serving::LocalExecutionAttemptSource; use crate::clock::current_unix_ms; @@ -269,10 +270,21 @@ where ) -> Result, Self::Error> { let plan = attempt.execution_plan(); let report_context = attempt.report_context(); + let daily_usage_outcome = self + .transfer_tracker + .daily_usage_outcome + .get_or_init(|| async { + self.state + .frontdoor_daily_usage() + .check(self.state, self.decision) + .await + }) + .await; if let Some(response) = execution_plan_balance_capacity_response( self.state, self.trace_id, self.decision, + daily_usage_outcome, plan, report_context.as_ref(), ) @@ -525,6 +537,8 @@ struct ProviderTransferStateTracker { #[derive(Clone, Debug, Default)] pub(crate) struct ProviderTransferTracker { state: std::sync::Arc>, + daily_usage_outcome: + std::sync::Arc>, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -969,6 +983,16 @@ where ) -> Result, Self::Error> { let plan = attempt.execution_plan(); let report_context = attempt.report_context(); + let daily_usage_outcome = self + .transfer_tracker + .daily_usage_outcome + .get_or_init(|| async { + self.state + .frontdoor_daily_usage() + .check(self.state, self.decision) + .await + }) + .await; let candidate_index = parse_request_candidate_report_context(report_context.as_ref()) .and_then(|context| context.candidate_index) .map(|value| value.to_string()) @@ -991,6 +1015,7 @@ where self.state, self.trace_id, self.decision, + daily_usage_outcome, plan, report_context.as_ref(), ) @@ -1124,9 +1149,38 @@ async fn execution_plan_balance_capacity_response( state: &AppState, trace_id: &str, decision: &GatewayControlDecision, + daily_usage_outcome: &crate::daily_usage_limit::FrontdoorDailyUsageOutcome, plan: &aether_contracts::ExecutionPlan, report_context: Option<&serde_json::Value>, ) -> Result>, GatewayError> { + if let crate::daily_usage_limit::FrontdoorDailyUsageOutcome::Rejected(rejection) = + daily_usage_outcome + { + if !crate::control::execution_plan_cost_is_proven_zero(state, plan, report_context).await { + let auth = decision.auth_context.as_ref(); + info!( + event_name = "frontdoor_daily_usage_rejected", + log_type = "event", + trace_id, + user_id = auth.map(|auth| auth.user_id.as_str()).unwrap_or("-"), + api_key_id = auth.map(|auth| auth.api_key_id.as_str()).unwrap_or("-"), + scope = rejection.scope, + limit_usd = rejection.limit_usd, + used_usd = rejection.used_usd, + retry_after = rejection.retry_after, + "gateway rejected candidate at daily usage limit" + ); + mark_unused_local_candidate(state, plan, report_context).await; + let mut response = crate::api::response::build_local_daily_usage_limited_response( + trace_id, + Some(decision), + rejection, + )?; + attach_redaction_execution_candidate(&mut response, plan.candidate_id.as_deref()); + return Ok(Some(response)); + } + } + let rejection = match crate::control::execution_plan_balance_capacity_rejection( state, decision, @@ -2007,6 +2061,22 @@ mod tests { assert_eq!(port.unused.lock().unwrap().as_slice(), ["a-key3-retry0"]); } + #[test] + fn cloned_tracker_shares_request_daily_usage_cache() { + let tracker = ProviderTransferTracker::default(); + let cloned = tracker.clone(); + let other_request = ProviderTransferTracker::default(); + + assert!(Arc::ptr_eq( + &tracker.daily_usage_outcome, + &cloned.daily_usage_outcome + )); + assert!(!Arc::ptr_eq( + &tracker.daily_usage_outcome, + &other_request.daily_usage_outcome + )); + } + #[tokio::test] async fn cloned_tracker_preserves_transfer_budget_across_candidate_loops() { let state = AppState::new().expect("state should build"); diff --git a/apps/aether-gateway/src/handlers/admin/auth/api_keys/mutation_routes.rs b/apps/aether-gateway/src/handlers/admin/auth/api_keys/mutation_routes.rs index 23af1b9f4..ae67bc7ab 100644 --- a/apps/aether-gateway/src/handlers/admin/auth/api_keys/mutation_routes.rs +++ b/apps/aether-gateway/src/handlers/admin/auth/api_keys/mutation_routes.rs @@ -119,6 +119,14 @@ pub(super) async fn build_admin_create_api_key_response( "rate_limit 必须大于等于 0", )); } + if payload + .daily_usage_limit_usd + .is_some_and(|value| !value.is_finite() || value < 0.0) + { + return Ok(build_admin_api_keys_bad_request_response( + "daily_usage_limit_usd 必须是大于等于 0 的有限数值", + )); + } let concurrent_limit = match normalize_optional_api_key_concurrent_limit(payload.concurrent_limit) { Ok(value) => value, @@ -169,6 +177,7 @@ pub(super) async fn build_admin_create_api_key_response( allowed_models, ip_rules, rate_limit: payload.rate_limit, + daily_usage_limit_usd: payload.daily_usage_limit_usd, concurrent_limit, force_capabilities: None, is_active: true, @@ -208,6 +217,7 @@ pub(super) async fn build_admin_create_api_key_response( "is_standalone": true, "is_active": created.is_active, "rate_limit": created.rate_limit, + "daily_usage_limit_usd": created.daily_usage_limit_usd, "concurrent_limit": created.concurrent_limit, "allowed_providers": created.allowed_providers, "allowed_api_formats": created.allowed_api_formats, @@ -304,6 +314,14 @@ pub(super) async fn build_admin_update_api_key_response( "rate_limit 必须大于等于 0", )); } + if payload + .daily_usage_limit_usd + .is_some_and(|value| !value.is_finite() || value < 0.0) + { + return Ok(build_admin_api_keys_bad_request_response( + "daily_usage_limit_usd 必须是大于等于 0 的有限数值", + )); + } let concurrent_limit = match normalize_optional_api_key_concurrent_limit(payload.concurrent_limit) { Ok(value) => value, @@ -408,6 +426,8 @@ pub(super) async fn build_admin_update_api_key_response( name, rate_limit_present: field_presence.contains("rate_limit"), rate_limit: payload.rate_limit, + daily_usage_limit_present: field_presence.contains("daily_usage_limit_usd"), + daily_usage_limit_usd: payload.daily_usage_limit_usd, concurrent_limit_present: field_presence.contains("concurrent_limit"), concurrent_limit, allowed_providers, diff --git a/apps/aether-gateway/src/handlers/admin/auth/api_keys/shared.rs b/apps/aether-gateway/src/handlers/admin/auth/api_keys/shared.rs index 3adb5ce64..78cfabf65 100644 --- a/apps/aether-gateway/src/handlers/admin/auth/api_keys/shared.rs +++ b/apps/aether-gateway/src/handlers/admin/auth/api_keys/shared.rs @@ -24,6 +24,7 @@ pub(super) struct AdminStandaloneApiKeyCreateRequest { #[serde(default, alias = "allowed_ips")] pub(super) ip_rules: Option>, pub(super) rate_limit: Option, + pub(super) daily_usage_limit_usd: Option, pub(super) concurrent_limit: Option, pub(super) initial_balance_usd: Option, pub(super) unlimited_balance: Option, @@ -46,6 +47,7 @@ pub(super) struct AdminStandaloneApiKeyUpdateRequest { )] pub(super) ip_rules: Option>>, pub(super) rate_limit: Option, + pub(super) daily_usage_limit_usd: Option, pub(super) concurrent_limit: Option, pub(super) initial_balance_usd: Option, pub(super) unlimited_balance: Option, @@ -166,6 +168,7 @@ pub(super) fn build_admin_api_key_list_item_payload( "total_tokens": record.total_tokens, "total_cost_usd": record.total_cost_usd, "rate_limit": record.rate_limit, + "daily_usage_limit_usd": record.daily_usage_limit_usd, "concurrent_limit": record.concurrent_limit, "allowed_providers": record.allowed_providers, "allowed_api_formats": record.allowed_api_formats, @@ -197,6 +200,7 @@ pub(super) fn build_admin_api_key_detail_payload( "total_tokens": record.total_tokens, "total_cost_usd": record.total_cost_usd, "rate_limit": record.rate_limit, + "daily_usage_limit_usd": record.daily_usage_limit_usd, "concurrent_limit": record.concurrent_limit, "allowed_providers": record.allowed_providers, "allowed_api_formats": record.allowed_api_formats, diff --git a/apps/aether-gateway/src/handlers/admin/request/system/export.rs b/apps/aether-gateway/src/handlers/admin/request/system/export.rs index 96cc28227..56a141b91 100644 --- a/apps/aether-gateway/src/handlers/admin/request/system/export.rs +++ b/apps/aether-gateway/src/handlers/admin/request/system/export.rs @@ -244,6 +244,8 @@ impl<'a> AdminAppState<'a> { "allowed_models_mode": group.allowed_models_mode.clone(), "rate_limit": group.rate_limit, "rate_limit_mode": group.rate_limit_mode.clone(), + "daily_usage_limit_usd": group.daily_usage_limit_usd, + "daily_usage_limit_mode": group.daily_usage_limit_mode.clone(), }) }) .collect::>(); @@ -365,6 +367,10 @@ impl<'a> AdminAppState<'a> { ), ("ip_rules".to_string(), json!(key.ip_rules.clone())), ("rate_limit".to_string(), json!(key.rate_limit)), + ( + "daily_usage_limit_usd".to_string(), + json!(key.daily_usage_limit_usd), + ), ("concurrent_limit".to_string(), json!(key.concurrent_limit)), ( "force_capabilities".to_string(), diff --git a/apps/aether-gateway/src/handlers/admin/request/system/import.rs b/apps/aether-gateway/src/handlers/admin/request/system/import.rs index d02406e1e..fbd262b01 100644 --- a/apps/aether-gateway/src/handlers/admin/request/system/import.rs +++ b/apps/aether-gateway/src/handlers/admin/request/system/import.rs @@ -983,6 +983,11 @@ fn build_imported_user_group_record( let allowed_api_formats = normalize_imported_user_api_formats(group, "allowed_api_formats")?; let allowed_models = normalize_imported_user_string_list(group, "allowed_models")?; let rate_limit = imported_optional_i32(group.get("rate_limit"), "rate_limit")?; + let daily_usage_limit_usd = + imported_optional_f64(group.get("daily_usage_limit_usd"), "daily_usage_limit_usd")?; + if daily_usage_limit_usd.is_some_and(|value| value < 0.0) { + return Err(format!("{field_name}.daily_usage_limit_usd 必须大于等于 0")); + } let allowed_providers_mode = imported_optional_list_policy_mode( group.get("allowed_providers_mode"), @@ -1026,6 +1031,17 @@ fn build_imported_user_group_record( "inherit".to_string() } }); + let daily_usage_limit_mode = imported_optional_rate_limit_policy_mode( + group.get("daily_usage_limit_mode"), + "daily_usage_limit_mode", + )? + .unwrap_or_else(|| { + if daily_usage_limit_usd.is_some() { + "custom".to_string() + } else { + "inherit".to_string() + } + }); let normalized_name = name.to_ascii_lowercase(); @@ -1044,6 +1060,8 @@ fn build_imported_user_group_record( allowed_models_mode, rate_limit, rate_limit_mode, + daily_usage_limit_usd, + daily_usage_limit_mode, }, )) } @@ -2914,6 +2932,13 @@ impl<'a> AdminAppState<'a> { let rate_limit = invalid_value!(imported_optional_i32(key.get("rate_limit"), "rate_limit")) .unwrap_or(0); + let daily_usage_limit_usd = invalid_value!(imported_optional_f64( + key.get("daily_usage_limit_usd"), + "daily_usage_limit_usd" + )); + if daily_usage_limit_usd.is_some_and(|value| value < 0.0) { + return Ok(Err(invalid_request("daily_usage_limit_usd 必须大于等于 0"))); + } let concurrent_limit = invalid_value!(imported_optional_i32( key.get("concurrent_limit"), "concurrent_limit" @@ -2971,6 +2996,9 @@ impl<'a> AdminAppState<'a> { api_key_id: existing_key.api_key_id.clone(), name: name.clone(), rate_limit: Some(rate_limit), + daily_usage_limit_present: key + .contains_key("daily_usage_limit_usd"), + daily_usage_limit_usd, concurrent_limit: if key.contains_key("concurrent_limit") { concurrent_limit } else { @@ -3070,6 +3098,7 @@ impl<'a> AdminAppState<'a> { allowed_models, ip_rules, rate_limit, + daily_usage_limit_usd, concurrent_limit, force_capabilities, is_active, @@ -3170,6 +3199,13 @@ impl<'a> AdminAppState<'a> { let rate_limit = invalid_value!(imported_optional_i32(key.get("rate_limit"), "rate_limit")) .unwrap_or(0); + let daily_usage_limit_usd = invalid_value!(imported_optional_f64( + key.get("daily_usage_limit_usd"), + "daily_usage_limit_usd" + )); + if daily_usage_limit_usd.is_some_and(|value| value < 0.0) { + return Ok(Err(invalid_request("daily_usage_limit_usd 必须大于等于 0"))); + } let concurrent_limit = invalid_value!(imported_optional_i32( key.get("concurrent_limit"), "concurrent_limit" @@ -3233,6 +3269,9 @@ impl<'a> AdminAppState<'a> { name: name.clone(), rate_limit_present: true, rate_limit: Some(rate_limit), + daily_usage_limit_present: key + .contains_key("daily_usage_limit_usd"), + daily_usage_limit_usd, concurrent_limit_present: key.contains_key("concurrent_limit"), concurrent_limit, allowed_providers: Some(allowed_providers.clone()), @@ -3325,6 +3364,7 @@ impl<'a> AdminAppState<'a> { allowed_models, ip_rules, rate_limit: Some(rate_limit), + daily_usage_limit_usd, concurrent_limit, force_capabilities, is_active, diff --git a/apps/aether-gateway/src/handlers/admin/system/shared/configs.rs b/apps/aether-gateway/src/handlers/admin/system/shared/configs.rs index 8cd5ad72a..a10132011 100644 --- a/apps/aether-gateway/src/handlers/admin/system/shared/configs.rs +++ b/apps/aether-gateway/src/handlers/admin/system/shared/configs.rs @@ -122,6 +122,21 @@ pub(crate) async fn apply_admin_system_config_update( let normalized_key = update.normalized_key; let description = update.description; + if normalized_key == "daily_usage_limit_usd" { + let limit = match &value { + serde_json::Value::Number(number) => number.as_f64(), + serde_json::Value::String(raw) => raw.trim().parse::().ok(), + _ => None, + }; + let Some(limit) = limit.filter(|value| value.is_finite() && *value >= 0.0) else { + return Ok(Err(( + http::StatusCode::BAD_REQUEST, + json!({ "detail": "daily_usage_limit_usd 必须是大于等于 0 的有限数值" }), + ))); + }; + value = json!(limit); + } + if is_sensitive_admin_system_config_key(&normalized_key) && value.as_str().is_some_and(|raw| !raw.is_empty()) { diff --git a/apps/aether-gateway/src/handlers/admin/users/api_keys/helpers.rs b/apps/aether-gateway/src/handlers/admin/users/api_keys/helpers.rs index 98f9e8d3a..53c08ebf1 100644 --- a/apps/aether-gateway/src/handlers/admin/users/api_keys/helpers.rs +++ b/apps/aether-gateway/src/handlers/admin/users/api_keys/helpers.rs @@ -43,6 +43,7 @@ pub(super) fn build_admin_user_api_key_detail_payload( "total_requests": record.total_requests, "total_cost_usd": record.total_cost_usd, "rate_limit": record.rate_limit, + "daily_usage_limit_usd": record.daily_usage_limit_usd, "concurrent_limit": record.concurrent_limit, "ip_rules": record.ip_rules, "feature_settings": record.feature_settings, diff --git a/apps/aether-gateway/src/handlers/admin/users/api_keys/responses/create.rs b/apps/aether-gateway/src/handlers/admin/users/api_keys/responses/create.rs index dfe0cb37d..733c8fcc8 100644 --- a/apps/aether-gateway/src/handlers/admin/users/api_keys/responses/create.rs +++ b/apps/aether-gateway/src/handlers/admin/users/api_keys/responses/create.rs @@ -74,7 +74,7 @@ pub(crate) async fn build_admin_create_user_api_key_response( { return Ok(( http::StatusCode::BAD_REQUEST, - Json(json!({ "detail": "当前仅支持 name、rate_limit、concurrent_limit、allowed_providers、ip_rules 字段" })), + Json(json!({ "detail": "当前仅支持 name、rate_limit、daily_usage_limit_usd、concurrent_limit、allowed_providers、ip_rules 字段" })), ) .into_response()); } @@ -128,6 +128,16 @@ pub(crate) async fn build_admin_create_user_api_key_response( ) .into_response()); } + if payload + .daily_usage_limit_usd + .is_some_and(|value| !value.is_finite() || value < 0.0) + { + return Ok(( + http::StatusCode::BAD_REQUEST, + Json(json!({ "detail": "daily_usage_limit_usd 必须是大于等于 0 的有限数值" })), + ) + .into_response()); + } let concurrent_limit = match normalize_optional_api_key_concurrent_limit(payload.concurrent_limit) { Ok(value) => value, @@ -161,6 +171,7 @@ pub(crate) async fn build_admin_create_user_api_key_response( allowed_models: None, ip_rules, rate_limit, + daily_usage_limit_usd: payload.daily_usage_limit_usd, concurrent_limit, force_capabilities: None, is_active: true, @@ -213,6 +224,7 @@ pub(crate) async fn build_admin_create_user_api_key_response( "name": created.name, "key_display": masked_user_api_key_display(state, created.key_encrypted.as_deref()), "rate_limit": created.rate_limit, + "daily_usage_limit_usd": created.daily_usage_limit_usd, "concurrent_limit": created.concurrent_limit, "ip_rules": created.ip_rules, "expires_at": format_optional_unix_secs_iso8601(created.expires_at_unix_secs), diff --git a/apps/aether-gateway/src/handlers/admin/users/api_keys/responses/list.rs b/apps/aether-gateway/src/handlers/admin/users/api_keys/responses/list.rs index 50f3ee248..6eedf63b8 100644 --- a/apps/aether-gateway/src/handlers/admin/users/api_keys/responses/list.rs +++ b/apps/aether-gateway/src/handlers/admin/users/api_keys/responses/list.rs @@ -62,6 +62,7 @@ pub(crate) async fn build_admin_list_user_api_keys_response( "total_requests": record.total_requests, "total_cost_usd": record.total_cost_usd, "rate_limit": record.rate_limit, + "daily_usage_limit_usd": record.daily_usage_limit_usd, "concurrent_limit": record.concurrent_limit, "feature_settings": record.feature_settings, "expires_at": format_optional_unix_secs_iso8601(record.expires_at_unix_secs), diff --git a/apps/aether-gateway/src/handlers/admin/users/api_keys/responses/update.rs b/apps/aether-gateway/src/handlers/admin/users/api_keys/responses/update.rs index 29956628f..c7b23a0f3 100644 --- a/apps/aether-gateway/src/handlers/admin/users/api_keys/responses/update.rs +++ b/apps/aether-gateway/src/handlers/admin/users/api_keys/responses/update.rs @@ -83,6 +83,17 @@ pub(crate) async fn build_admin_update_user_api_key_response( ) .into_response()); } + if payload + .daily_usage_limit_usd + .flatten() + .is_some_and(|value| !value.is_finite() || value < 0.0) + { + return Ok(( + http::StatusCode::BAD_REQUEST, + Json(json!({ "detail": "daily_usage_limit_usd 必须是大于等于 0 的有限数值" })), + ) + .into_response()); + } let concurrent_limit = match normalize_optional_api_key_concurrent_limit(payload.concurrent_limit) { Ok(value) => value, @@ -114,6 +125,8 @@ pub(crate) async fn build_admin_update_user_api_key_response( api_key_id: api_key_id.clone(), name, rate_limit: payload.rate_limit, + daily_usage_limit_present: payload.daily_usage_limit_usd.is_some(), + daily_usage_limit_usd: payload.daily_usage_limit_usd.flatten(), concurrent_limit, ip_rules, }) diff --git a/apps/aether-gateway/src/handlers/admin/users/groups.rs b/apps/aether-gateway/src/handlers/admin/users/groups.rs index a596598b7..f1a6f81c8 100644 --- a/apps/aether-gateway/src/handlers/admin/users/groups.rs +++ b/apps/aether-gateway/src/handlers/admin/users/groups.rs @@ -37,6 +37,10 @@ struct AdminUserGroupPayload { rate_limit: Option, #[serde(default = "default_rate_limit_mode")] rate_limit_mode: String, + #[serde(default)] + daily_usage_limit_usd: Option, + #[serde(default = "default_rate_limit_mode")] + daily_usage_limit_mode: String, } #[derive(Debug, serde::Deserialize)] @@ -370,6 +374,12 @@ fn parse_group_record( if payload.rate_limit.is_some_and(|value| value < 0) { return Err("rate_limit 必须大于等于 0".to_string()); } + if payload + .daily_usage_limit_usd + .is_some_and(|value| !value.is_finite() || value < 0.0) + { + return Err("daily_usage_limit_usd 必须是大于等于 0 的有限数值".to_string()); + } let allowed_providers = normalize_admin_user_string_list(payload.allowed_providers, "allowed_providers")?; let allowed_api_formats = normalize_admin_user_api_formats(payload.allowed_api_formats)?; @@ -390,6 +400,8 @@ fn parse_group_record( allowed_models_mode: normalize_list_mode(&payload.allowed_models_mode)?, rate_limit: payload.rate_limit, rate_limit_mode: normalize_rate_mode(&payload.rate_limit_mode)?, + daily_usage_limit_usd: payload.daily_usage_limit_usd, + daily_usage_limit_mode: normalize_rate_mode(&payload.daily_usage_limit_mode)?, }) } @@ -420,6 +432,8 @@ fn user_group_payload( "allowed_models_mode": group.allowed_models_mode, "rate_limit": group.rate_limit, "rate_limit_mode": group.rate_limit_mode, + "daily_usage_limit_usd": group.daily_usage_limit_usd, + "daily_usage_limit_mode": group.daily_usage_limit_mode, "is_default": default_group_id == Some(group.id.as_str()), "created_at": format_optional_datetime_iso8601(group.created_at), "updated_at": format_optional_datetime_iso8601(group.updated_at), diff --git a/apps/aether-gateway/src/handlers/admin/users/lifecycle/create.rs b/apps/aether-gateway/src/handlers/admin/users/lifecycle/create.rs index ad5ea3f30..01d3a0090 100644 --- a/apps/aether-gateway/src/handlers/admin/users/lifecycle/create.rs +++ b/apps/aether-gateway/src/handlers/admin/users/lifecycle/create.rs @@ -4,7 +4,10 @@ use super::super::{ normalize_admin_optional_user_email, normalize_admin_user_group_ids, normalize_admin_user_role, normalize_admin_username, validate_admin_user_password, AdminCreateUserRequest, }; -use super::support::{admin_user_password_policy, build_admin_user_payload_with_groups}; +use super::support::{ + admin_system_daily_usage_limit, admin_user_password_policy, + build_admin_user_payload_with_groups, +}; use crate::handlers::admin::request::{AdminAppState, AdminRequestContext}; use crate::handlers::admin::shared::attach_admin_audit_response; use crate::GatewayError; @@ -229,8 +232,15 @@ pub(in super::super) async fn build_admin_create_user_response( None }; - let mut payload = - build_admin_user_payload_with_groups(&user, None, None, payload.unlimited, &groups); + let system_daily_usage_limit_usd = admin_system_daily_usage_limit(state).await?; + let mut payload = build_admin_user_payload_with_groups( + &user, + None, + None, + payload.unlimited, + &groups, + system_daily_usage_limit_usd, + ); payload["feature_settings"] = feature_settings.unwrap_or(Value::Null); Ok(attach_admin_audit_response( diff --git a/apps/aether-gateway/src/handlers/admin/users/lifecycle/reads.rs b/apps/aether-gateway/src/handlers/admin/users/lifecycle/reads.rs index 7e67fa755..3c0ee3a11 100644 --- a/apps/aether-gateway/src/handlers/admin/users/lifecycle/reads.rs +++ b/apps/aether-gateway/src/handlers/admin/users/lifecycle/reads.rs @@ -1,7 +1,7 @@ use super::super::{build_admin_users_bad_request_response, format_optional_datetime_iso8601}; use super::support::{ - admin_user_id_from_detail_path, build_admin_user_export_payload, - build_admin_user_payload_with_groups, find_admin_export_user, + admin_system_daily_usage_limit, admin_user_id_from_detail_path, + build_admin_user_export_payload, build_admin_user_payload_with_groups, find_admin_export_user, }; use crate::handlers::admin::request::{AdminAppState, AdminRequestContext}; use crate::handlers::admin::shared::{query_param_optional_bool, query_param_value}; @@ -70,13 +70,16 @@ pub(in super::super) async fn build_admin_list_users_response( usage_totals_result, memberships_result, groups_result, + system_daily_usage_limit_result, ) = tokio::join!( state.list_user_auth_by_ids(&user_ids), state.list_wallet_snapshots_by_user_ids(&user_ids), state.summarize_usage_totals_by_user_ids(&user_ids), state.list_user_group_memberships_by_user_ids(&user_ids), state.list_user_groups(), + admin_system_daily_usage_limit(state), ); + let system_daily_usage_limit_usd = system_daily_usage_limit_result?; let auth_by_user_id = auth_rows_result? .into_iter() .map(|user| (user.id.clone(), user)) @@ -126,6 +129,7 @@ pub(in super::super) async fn build_admin_list_users_response( .map(|item| item.total_tokens) .unwrap_or_default(), &groups, + system_daily_usage_limit_usd, )); } @@ -162,6 +166,7 @@ pub(in super::super) async fn build_admin_get_user_response( .await?; let export_row = find_admin_export_user(state, &user_id).await?; let groups = state.list_user_groups_for_user(&user_id).await?; + let system_daily_usage_limit_usd = admin_system_daily_usage_limit(state).await?; let unlimited = wallet .as_ref() .is_some_and(|wallet| wallet.limit_mode.eq_ignore_ascii_case("unlimited")); @@ -171,6 +176,7 @@ pub(in super::super) async fn build_admin_get_user_response( export_row.as_ref().map(|row| row.rate_limit_mode.as_str()), unlimited, &groups, + system_daily_usage_limit_usd, ); payload["feature_settings"] = export_row .as_ref() diff --git a/apps/aether-gateway/src/handlers/admin/users/lifecycle/support.rs b/apps/aether-gateway/src/handlers/admin/users/lifecycle/support.rs index 4b5dfbd08..8b5b42e05 100644 --- a/apps/aether-gateway/src/handlers/admin/users/lifecycle/support.rs +++ b/apps/aether-gateway/src/handlers/admin/users/lifecycle/support.rs @@ -1,4 +1,5 @@ use super::super::format_optional_datetime_iso8601; +use crate::data::state::resolve_group_effective_daily_usage_limit_policy; use crate::handlers::admin::request::AdminAppState; use crate::GatewayError; use serde_json::json; @@ -25,6 +26,16 @@ pub(super) async fn admin_user_password_policy( ) } +pub(super) async fn admin_system_daily_usage_limit( + state: &AdminAppState<'_>, +) -> Result { + crate::daily_usage_limit::parse_system_limit( + state + .read_system_config_json_value("daily_usage_limit_usd") + .await?, + ) +} + pub(super) async fn find_admin_export_user( state: &AdminAppState<'_>, user_id: &str, @@ -37,7 +48,7 @@ pub(super) fn build_admin_user_payload( rate_limit: Option, unlimited: bool, ) -> serde_json::Value { - build_admin_user_payload_with_groups(user, rate_limit, None, unlimited, &[]) + build_admin_user_payload_with_groups(user, rate_limit, None, unlimited, &[], 0.0) } pub(super) fn build_admin_user_payload_with_groups( @@ -46,6 +57,7 @@ pub(super) fn build_admin_user_payload_with_groups( rate_limit_mode: Option<&str>, unlimited: bool, groups: &[aether_data::repository::users::StoredUserGroup], + system_daily_usage_limit_usd: f64, ) -> serde_json::Value { json!({ "id": user.id, @@ -76,6 +88,7 @@ pub(super) fn build_admin_user_payload_with_groups( rate_limit, rate_limit_mode.unwrap_or("system"), groups, + system_daily_usage_limit_usd, ), }) } @@ -89,6 +102,7 @@ pub(super) fn build_admin_user_export_payload( request_count: u64, total_tokens: u64, groups: &[aether_data::repository::users::StoredUserGroup], + system_daily_usage_limit_usd: f64, ) -> serde_json::Value { json!({ "id": row.id, @@ -122,6 +136,7 @@ pub(super) fn build_admin_user_export_payload( row.rate_limit, &row.rate_limit_mode, groups, + system_daily_usage_limit_usd, ), }) } @@ -146,6 +161,7 @@ fn effective_policy_payload( rate_limit: Option, rate_limit_mode: &str, groups: &[aether_data::repository::users::StoredUserGroup], + system_daily_usage_limit_usd: f64, ) -> serde_json::Value { let mut sorted_groups = groups.to_vec(); sorted_groups.sort_by(|left, right| { @@ -173,6 +189,10 @@ fn effective_policy_payload( |group| (&group.allowed_models_mode, group.allowed_models.as_ref()), ), "rate_limit": effective_rate_limit_policy_payload(rate_limit, rate_limit_mode, &sorted_groups), + "daily_usage_limit_usd": effective_daily_usage_limit_policy_payload( + &sorted_groups, + system_daily_usage_limit_usd, + ), }) } @@ -236,6 +256,26 @@ fn effective_rate_limit_policy_payload( } } +fn effective_daily_usage_limit_policy_payload( + groups: &[aether_data::repository::users::StoredUserGroup], + system_daily_usage_limit_usd: f64, +) -> serde_json::Value { + let group_sources = groups + .iter() + .filter(|group| group.daily_usage_limit_mode == "custom") + .collect::>(); + let source = combined_policy_source(false, group_sources.len(), "system"); + match resolve_group_effective_daily_usage_limit_policy(groups) { + Some(limit) => policy_payload("custom", json!(limit), source, &group_sources), + None => policy_payload( + "system", + json!(system_daily_usage_limit_usd), + source, + &group_sources, + ), + } +} + fn policy_payload( mode: &str, value: serde_json::Value, @@ -359,3 +399,53 @@ pub(super) fn admin_user_id_from_detail_path(request_path: &str) -> Option) -> StoredUserGroup { + StoredUserGroup { + id: id.to_string(), + name: name.to_string(), + normalized_name: name.to_ascii_lowercase(), + description: None, + priority: 0, + allowed_providers: None, + allowed_providers_mode: "inherit".to_string(), + allowed_api_formats: None, + allowed_api_formats_mode: "inherit".to_string(), + allowed_models: None, + allowed_models_mode: "inherit".to_string(), + rate_limit, + rate_limit_mode: "custom".to_string(), + daily_usage_limit_usd: None, + daily_usage_limit_mode: "inherit".to_string(), + created_at: None, + updated_at: None, + } + } + + #[test] + fn effective_daily_usage_payload_reports_system_value_and_custom_sources() { + let system = effective_daily_usage_limit_policy_payload(&[], 12.5); + assert_eq!(system["mode"], "system"); + assert_eq!(system["value"], 12.5); + assert_eq!(system["source"], "system"); + + let mut basic = sample_group("group-basic", "Basic", None); + basic.daily_usage_limit_mode = "custom".to_string(); + basic.daily_usage_limit_usd = Some(10.0); + let mut pro = sample_group("group-pro", "Pro", None); + pro.daily_usage_limit_mode = "custom".to_string(); + pro.daily_usage_limit_usd = Some(25.0); + let custom = effective_daily_usage_limit_policy_payload(&[basic, pro], 12.5); + assert_eq!(custom["mode"], "custom"); + assert_eq!(custom["value"], 25.0); + assert_eq!(custom["source"], "combined"); + assert_eq!( + custom["group_ids"], + serde_json::json!(["group-basic", "group-pro"]) + ); + } +} diff --git a/apps/aether-gateway/src/handlers/admin/users/lifecycle/update.rs b/apps/aether-gateway/src/handlers/admin/users/lifecycle/update.rs index 0cb174166..42af06973 100644 --- a/apps/aether-gateway/src/handlers/admin/users/lifecycle/update.rs +++ b/apps/aether-gateway/src/handlers/admin/users/lifecycle/update.rs @@ -6,7 +6,7 @@ use super::super::{ validate_admin_user_password, AdminUpdateUserPatch, }; use super::support::{ - admin_user_id_from_detail_path, admin_user_password_policy, + admin_system_daily_usage_limit, admin_user_id_from_detail_path, admin_user_password_policy, build_admin_user_payload_with_groups, find_admin_export_user, }; use crate::handlers::admin::request::{AdminAppState, AdminRequestContext}; @@ -341,6 +341,7 @@ pub(in super::super) async fn build_admin_update_user_response( let export_row = find_admin_export_user(state, &user_id).await?; let groups = state.list_user_groups_for_user(&user_id).await?; let rate_limit = export_row.as_ref().and_then(|row| row.rate_limit); + let system_daily_usage_limit_usd = admin_system_daily_usage_limit(state).await?; let mut payload = build_admin_user_payload_with_groups( &user, @@ -348,6 +349,7 @@ pub(in super::super) async fn build_admin_update_user_response( export_row.as_ref().map(|row| row.rate_limit_mode.as_str()), unlimited, &groups, + system_daily_usage_limit_usd, ); payload["feature_settings"] = export_row .as_ref() diff --git a/apps/aether-gateway/src/handlers/admin/users/shared.rs b/apps/aether-gateway/src/handlers/admin/users/shared.rs index dbae70da8..d247a3558 100644 --- a/apps/aether-gateway/src/handlers/admin/users/shared.rs +++ b/apps/aether-gateway/src/handlers/admin/users/shared.rs @@ -1,6 +1,8 @@ use super::ADMIN_USERS_DATA_UNAVAILABLE_DETAIL; use crate::handlers::admin::shared::AdminTypedObjectPatch; -use crate::handlers::shared::{deserialize_optional_string_list_patch, normalize_ip_rules}; +use crate::handlers::shared::{ + deserialize_optional_f64_patch, deserialize_optional_string_list_patch, normalize_ip_rules, +}; use axum::{ body::Body, http, @@ -25,6 +27,8 @@ pub(super) struct AdminCreateUserApiKeyRequest { #[serde(default)] pub(super) rate_limit: Option, #[serde(default)] + pub(super) daily_usage_limit_usd: Option, + #[serde(default)] pub(super) concurrent_limit: Option, #[serde(default)] pub(super) expire_days: Option, @@ -48,6 +52,8 @@ pub(super) struct AdminUpdateUserApiKeyRequest { pub(super) name: Option, #[serde(default)] pub(super) rate_limit: Option, + #[serde(default, deserialize_with = "deserialize_optional_f64_patch")] + pub(super) daily_usage_limit_usd: Option>, #[serde(default)] pub(super) concurrent_limit: Option, #[serde(default)] diff --git a/apps/aether-gateway/src/handlers/proxy/mod.rs b/apps/aether-gateway/src/handlers/proxy/mod.rs index 8b5bd59c0..ec4f8486a 100644 --- a/apps/aether-gateway/src/handlers/proxy/mod.rs +++ b/apps/aether-gateway/src/handlers/proxy/mod.rs @@ -193,6 +193,21 @@ fn local_execution_outcome_label(outcome: &LocalExecutionRequestOutcome) -> &'st } } +fn local_execution_response_path( + response: &Response, + default_path: &'static str, +) -> &'static str { + if response + .extensions() + .get::() + .is_some() + { + EXECUTION_PATH_LOCAL_RATE_LIMITED + } else { + default_path + } +} + fn request_hits_execution_loop_guard(parts: &http::request::Parts) -> bool { request_has_execution_runtime_loop_guard(&parts.headers) && frontdoor_self_loop_public_ai_path(parts.uri.path()) @@ -1498,7 +1513,7 @@ async fn proxy_request_inner( } } - let rpm_started_at = Instant::now(); + let mut limits_control_decision = None; let ip_whitelist_applies = control_decision.and_then(|decision| decision.route_class.as_deref()) == Some("ai_public"); let ip_whitelisted = if ip_whitelist_applies { @@ -1506,14 +1521,17 @@ async fn proxy_request_inner( } else { Ok(false) }; - let rate_limit_outcome = match ip_whitelisted { - Ok(true) => FrontdoorUserRpmOutcome::NotApplicable, - Ok(false) => { - state - .frontdoor_user_rpm() - .check_and_consume(&state, control_decision) - .await? + match ip_whitelisted { + Ok(true) => { + limits_control_decision = control_decision.cloned(); + if let Some(auth) = limits_control_decision + .as_mut() + .and_then(|decision| decision.auth_context.as_mut()) + { + auth.ip_bypass_limits = true; + } } + Ok(false) => {} Err(err) => { warn!( event_name = "frontdoor_ip_whitelist_check_failed", @@ -1523,11 +1541,21 @@ async fn proxy_request_inner( error = ?err, "gateway continued with rate limiting after IP whitelist check error" ); - state - .frontdoor_user_rpm() - .check_and_consume(&state, control_decision) - .await? } + } + let control_decision = limits_control_decision.as_ref().or(control_decision); + + let rpm_started_at = Instant::now(); + let rate_limit_outcome = if control_decision + .and_then(|decision| decision.auth_context.as_ref()) + .is_some_and(|auth| auth.ip_bypass_limits) + { + FrontdoorUserRpmOutcome::NotApplicable + } else { + state + .frontdoor_user_rpm() + .check_and_consume(&state, control_decision) + .await? }; observe_gateway_stage_ms("frontdoor_rpm", rpm_started_at.elapsed().as_millis() as u64); if let FrontdoorUserRpmOutcome::Rejected(rejection) = &rate_limit_outcome { @@ -1673,13 +1701,17 @@ async fn proxy_request_inner( execution_runtime_response, &redaction_slot, )?; + let execution_path = local_execution_response_path( + &execution_runtime_response, + EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, + ); state.clear_local_execution_runtime_miss_diagnostic(&trace_id); return Ok(finalize_gateway_response_with_context( &state, execution_runtime_response, &remote_addr, &request_context, - EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, + execution_path, &started_at, request_permit.take(), )); @@ -1733,13 +1765,17 @@ async fn proxy_request_inner( &redaction_slot, ) .await?; + let execution_path = local_execution_response_path( + &execution_runtime_response, + EXECUTION_PATH_EXECUTION_RUNTIME_SYNC, + ); state.clear_local_execution_runtime_miss_diagnostic(&trace_id); return Ok(finalize_gateway_response_with_context( &state, execution_runtime_response, &remote_addr, &request_context, - EXECUTION_PATH_EXECUTION_RUNTIME_SYNC, + execution_path, &started_at, request_permit.take(), )); @@ -1793,13 +1829,17 @@ async fn proxy_request_inner( execution_runtime_response, &redaction_slot, )?; + let execution_path = local_execution_response_path( + &execution_runtime_response, + EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, + ); state.clear_local_execution_runtime_miss_diagnostic(&trace_id); return Ok(finalize_gateway_response_with_context( &state, execution_runtime_response, &remote_addr, &request_context, - EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, + execution_path, &started_at, request_permit.take(), )); diff --git a/apps/aether-gateway/src/handlers/public/support/monitoring/rate_limit_status.rs b/apps/aether-gateway/src/handlers/public/support/monitoring/rate_limit_status.rs index 5f919f440..b17880596 100644 --- a/apps/aether-gateway/src/handlers/public/support/monitoring/rate_limit_status.rs +++ b/apps/aether-gateway/src/handlers/public/support/monitoring/rate_limit_status.rs @@ -7,6 +7,8 @@ use axum::{ use chrono::Utc; use serde_json::json; +use crate::control::{GatewayControlAuthContext, GatewayControlDecision}; + use super::{ build_auth_error_response, resolve_authenticated_local_user, AppState, GatewayPublicRequestContext, @@ -19,6 +21,84 @@ fn normalize_rate_limit_value(value: Option) -> u32 { .unwrap_or(0) } +fn daily_usage_available_payload( + status: crate::daily_usage_limit::FrontdoorDailyUsageStatus, +) -> serde_json::Value { + if !status.available { + return daily_usage_empty_payload( + false, + "unavailable", + &status.timezone, + &status.window_start, + &status.window_end, + &status.window_end, + ); + } + let primary = match (status.user.as_ref(), status.key.as_ref()) { + (Some(user), Some(key)) => { + if user.remaining_usd <= key.remaining_usd { + Some(user) + } else { + Some(key) + } + } + (Some(user), None) => Some(user), + (None, Some(key)) => Some(key), + (None, None) => None, + }; + json!({ + "available": true, + "status": "available", + "limit_usd": primary.map(|scope| scope.limit_usd), + "used_usd": primary.map(|scope| scope.used_usd), + "remaining_usd": primary.map(|scope| scope.remaining_usd), + "scope": primary.map(|scope| scope.scope), + "user": status.user.map(|scope| json!({ + "limit_usd": scope.limit_usd, + "used_usd": scope.used_usd, + "remaining_usd": scope.remaining_usd, + })), + "key": status.key.map(|scope| json!({ + "limit_usd": scope.limit_usd, + "used_usd": scope.used_usd, + "remaining_usd": scope.remaining_usd, + })), + "timezone": status.timezone, + "window": "1d", + "window_start": status.window_start, + "window_end": status.window_end, + "reset_time": chrono::DateTime::::from_timestamp( + status.reset_at_unix_secs as i64, + 0, + ).map(|value| value.to_rfc3339()), + }) +} + +fn daily_usage_empty_payload( + available: bool, + status: &'static str, + timezone: &str, + window_start: &str, + window_end: &str, + reset_time: &str, +) -> serde_json::Value { + json!({ + "available": available, + "status": status, + "limit_usd": null, + "used_usd": null, + "remaining_usd": null, + "scope": null, + "user": null, + "key": null, + "timezone": timezone, + "window": "1d", + "window_start": window_start, + "window_end": window_end, + "reset_time": reset_time, + }) +} + pub(super) async fn handle_user_rate_limit_status( state: &AppState, request_context: &GatewayPublicRequestContext, @@ -46,6 +126,13 @@ pub(super) async fn handle_user_rate_limit_status( )) .to_rfc3339(); let window = format!("{}s", limiter.config().bucket_seconds()); + let daily_limiter = state.frontdoor_daily_usage(); + let daily_timezone = crate::app_timezone::app_timezone(); + let (_, daily_window_start, daily_window_end) = + crate::app_timezone::local_day_window(now, daily_timezone); + let daily_window_start = daily_window_start.to_rfc3339(); + let daily_window_end = daily_window_end.to_rfc3339(); + let daily_reset_at = daily_window_end.clone(); let export_records = match state .list_auth_api_key_export_records_by_user_ids(std::slice::from_ref(&auth.user.id)) @@ -168,6 +255,65 @@ pub(super) async fn handle_user_rate_limit_status( (None, None) => None, }; + let mut daily_decision = GatewayControlDecision::synthetic( + "/v1/monitoring/rate-limit-status", + Some("ai_public".to_string()), + Some("openai".to_string()), + Some("chat".to_string()), + Some("openai:chat".to_string()), + ); + daily_decision.auth_context = Some(GatewayControlAuthContext { + user_id: auth.user.id.clone(), + api_key_id: record.api_key_id.clone(), + username: Some(auth.user.username.clone()), + api_key_name: record.name.clone(), + balance_remaining: None, + access_allowed: true, + user_rate_limit: snapshot.as_ref().and_then(|value| value.user_rate_limit), + api_key_rate_limit: snapshot.as_ref().and_then(|value| value.api_key_rate_limit), + user_daily_usage_limit_usd: snapshot + .as_ref() + .and_then(|value| value.user_daily_usage_limit_usd), + api_key_daily_usage_limit_usd: snapshot + .as_ref() + .and_then(|value| value.api_key_daily_usage_limit_usd) + .or(record.daily_usage_limit_usd), + api_key_is_standalone: is_standalone, + admin_bypass_limits: snapshot.as_ref().is_some_and(|value| { + value.user_role.eq_ignore_ascii_case("admin") && !value.api_key_is_standalone + }), + ip_bypass_limits: false, + local_rejection: None, + allowed_models: None, + ip_rules: None, + }); + let daily_usage = match daily_limiter.current_status(state, &daily_decision).await { + Ok(Some(status)) => daily_usage_available_payload(status), + Ok(None) => daily_usage_empty_payload( + true, + "unlimited", + daily_timezone.name(), + &daily_window_start, + &daily_window_end, + &daily_reset_at, + ), + Err(err) => { + tracing::warn!( + error = ?err, + api_key_id = %record.api_key_id, + "daily usage runtime status unavailable" + ); + daily_usage_empty_payload( + false, + "unavailable", + daily_timezone.name(), + &daily_window_start, + &daily_window_end, + &daily_reset_at, + ) + } + }; + api_keys.push(json!({ "api_key_name": record .name @@ -182,6 +328,7 @@ pub(super) async fn handle_user_rate_limit_status( "user_remaining": user_remaining, "key_limit": if key_limit > 0 { Some(key_limit) } else { None:: }, "key_remaining": key_remaining, + "daily_usage": daily_usage, })); } @@ -191,3 +338,57 @@ pub(super) async fn handle_user_rate_limit_status( })) .into_response() } + +#[cfg(test)] +mod tests { + use super::daily_usage_available_payload; + use crate::daily_usage_limit::{DailyUsageScopeStatus, FrontdoorDailyUsageStatus}; + + #[test] + fn daily_usage_monitoring_payload_exposes_primary_and_nested_scopes() { + let payload = daily_usage_available_payload(FrontdoorDailyUsageStatus { + available: true, + timezone: "Asia/Shanghai".to_string(), + window_start: "2026-08-02T16:00:00Z".to_string(), + window_end: "2026-08-03T16:00:00Z".to_string(), + reset_at_unix_secs: 1_775_232_000, + user: Some(DailyUsageScopeStatus { + scope: "user", + limit_usd: 10.0, + used_usd: 4.0, + remaining_usd: 6.0, + }), + key: Some(DailyUsageScopeStatus { + scope: "key", + limit_usd: 5.0, + used_usd: 3.0, + remaining_usd: 2.0, + }), + }); + + assert_eq!(payload["status"], "available"); + assert_eq!(payload["scope"], "key"); + assert_eq!(payload["limit_usd"], 5.0); + assert_eq!(payload["user"]["used_usd"], 4.0); + assert_eq!(payload["key"]["remaining_usd"], 2.0); + assert_eq!(payload["timezone"], "Asia/Shanghai"); + } + + #[test] + fn daily_usage_monitoring_payload_marks_runtime_failures_unavailable() { + let payload = daily_usage_available_payload(FrontdoorDailyUsageStatus { + available: false, + timezone: "Asia/Shanghai".to_string(), + window_start: "start".to_string(), + window_end: "end".to_string(), + reset_at_unix_secs: 0, + user: None, + key: None, + }); + + assert_eq!(payload["available"], false); + assert_eq!(payload["status"], "unavailable"); + assert!(payload["used_usd"].is_null()); + assert!(payload["remaining_usd"].is_null()); + } +} diff --git a/apps/aether-gateway/src/handlers/public/support/user_me_api_keys.rs b/apps/aether-gateway/src/handlers/public/support/user_me_api_keys.rs index c296e5be9..d87c3d4db 100644 --- a/apps/aether-gateway/src/handlers/public/support/user_me_api_keys.rs +++ b/apps/aether-gateway/src/handlers/public/support/user_me_api_keys.rs @@ -10,7 +10,7 @@ use serde::Deserialize; use serde_json::json; use crate::handlers::shared::{ - api_key_placeholder_display, deserialize_optional_json_patch, + api_key_placeholder_display, deserialize_optional_f64_patch, deserialize_optional_json_patch, deserialize_optional_string_list_patch, generate_gateway_api_key_plaintext, masked_gateway_api_key_display, normalize_feature_settings, normalize_ip_rules, normalize_optional_api_key_concurrent_limit, @@ -32,6 +32,8 @@ struct UsersMeCreateApiKeyRequest { #[serde(default)] rate_limit: Option, #[serde(default)] + daily_usage_limit_usd: Option, + #[serde(default)] concurrent_limit: Option, #[serde(default)] feature_settings: Option, @@ -45,6 +47,8 @@ struct UsersMeUpdateApiKeyRequest { name: Option, #[serde(default)] rate_limit: Option, + #[serde(default, deserialize_with = "deserialize_optional_f64_patch")] + daily_usage_limit_usd: Option>, #[serde(default)] concurrent_limit: Option, #[serde(default, deserialize_with = "deserialize_optional_json_patch")] @@ -167,6 +171,7 @@ fn build_users_me_api_key_list_payload( "total_requests": record.total_requests, "total_cost_usd": record.total_cost_usd, "rate_limit": record.rate_limit, + "daily_usage_limit_usd": record.daily_usage_limit_usd, "concurrent_limit": record.concurrent_limit, "allowed_providers": record.allowed_providers, "ip_rules": record.ip_rules, @@ -191,6 +196,7 @@ fn build_users_me_api_key_detail_payload( "force_capabilities": record.force_capabilities, "feature_settings": record.feature_settings, "rate_limit": record.rate_limit, + "daily_usage_limit_usd": record.daily_usage_limit_usd, "concurrent_limit": record.concurrent_limit, "last_used_at": format_users_me_optional_unix_secs_iso8601(record.last_used_at_unix_secs), "expires_at": format_users_me_optional_unix_secs_iso8601(record.expires_at_unix_secs), @@ -544,6 +550,16 @@ pub(super) async fn handle_users_me_api_key_create( false, ); } + if payload + .daily_usage_limit_usd + .is_some_and(|value| !value.is_finite() || value < 0.0) + { + return build_auth_error_response( + http::StatusCode::BAD_REQUEST, + "daily_usage_limit_usd 必须是大于等于 0 的有限数值", + false, + ); + } let concurrent_limit = match normalize_optional_api_key_concurrent_limit(payload.concurrent_limit) { Ok(value) => value, @@ -583,6 +599,7 @@ pub(super) async fn handle_users_me_api_key_create( allowed_models: None, ip_rules, rate_limit, + daily_usage_limit_usd: payload.daily_usage_limit_usd, concurrent_limit, force_capabilities: None, is_active: true, @@ -635,6 +652,7 @@ pub(super) async fn handle_users_me_api_key_create( "is_active": created.is_active, "is_locked": false, "rate_limit": created.rate_limit, + "daily_usage_limit_usd": created.daily_usage_limit_usd, "concurrent_limit": created.concurrent_limit, "ip_rules": created.ip_rules, "feature_settings": created.feature_settings, @@ -702,6 +720,17 @@ pub(super) async fn handle_users_me_api_key_update( false, ); } + if payload + .daily_usage_limit_usd + .flatten() + .is_some_and(|value| !value.is_finite() || value < 0.0) + { + return build_auth_error_response( + http::StatusCode::BAD_REQUEST, + "daily_usage_limit_usd 必须是大于等于 0 的有限数值", + false, + ); + } let concurrent_limit = match normalize_optional_api_key_concurrent_limit(payload.concurrent_limit) { Ok(value) => value, @@ -734,6 +763,8 @@ pub(super) async fn handle_users_me_api_key_update( api_key_id: snapshot.api_key_id.clone(), name, rate_limit, + daily_usage_limit_present: payload.daily_usage_limit_usd.is_some(), + daily_usage_limit_usd: payload.daily_usage_limit_usd.flatten(), concurrent_limit, ip_rules, }) @@ -1144,4 +1175,25 @@ mod tests { ])), ); } + + #[test] + fn update_payload_distinguishes_missing_null_and_present_daily_limit() { + let missing = serde_json::from_value::(json!({ + "name": "unchanged-daily-limit", + })) + .expect("missing daily limit should deserialize"); + assert_eq!(missing.daily_usage_limit_usd, None); + + let cleared = serde_json::from_value::(json!({ + "daily_usage_limit_usd": null, + })) + .expect("null daily limit should deserialize"); + assert_eq!(cleared.daily_usage_limit_usd, Some(None)); + + let updated = serde_json::from_value::(json!({ + "daily_usage_limit_usd": 3.5, + })) + .expect("present daily limit should deserialize"); + assert_eq!(updated.daily_usage_limit_usd, Some(Some(3.5))); + } } diff --git a/apps/aether-gateway/src/handlers/shared/mod.rs b/apps/aether-gateway/src/handlers/shared/mod.rs index 00aabf45d..67c8d7d7f 100644 --- a/apps/aether-gateway/src/handlers/shared/mod.rs +++ b/apps/aether-gateway/src/handlers/shared/mod.rs @@ -38,10 +38,11 @@ pub(crate) use self::email_templates::{ }; pub(crate) use self::external_models::OFFICIAL_EXTERNAL_MODEL_PROVIDERS; pub(crate) use self::normalize::{ - deserialize_optional_json_patch, deserialize_optional_string_list_patch, - ip_rule_pattern_matches, ip_rules_allow, json_ip_rules_allow, normalize_feature_settings, - normalize_ip_rules, normalize_json_array, normalize_json_object, normalize_string_list, - normalize_user_self_feature_settings_update, parse_json_ip_rules, + deserialize_optional_f64_patch, deserialize_optional_json_patch, + deserialize_optional_string_list_patch, ip_rule_pattern_matches, ip_rules_allow, + json_ip_rules_allow, normalize_feature_settings, normalize_ip_rules, normalize_json_array, + normalize_json_object, normalize_string_list, normalize_user_self_feature_settings_update, + parse_json_ip_rules, }; pub(crate) use self::payloads::{ InternalGatewayAuthContextRequest, InternalGatewayExecuteRequest, diff --git a/apps/aether-gateway/src/handlers/shared/normalize.rs b/apps/aether-gateway/src/handlers/shared/normalize.rs index 597faa7ac..2d5331840 100644 --- a/apps/aether-gateway/src/handlers/shared/normalize.rs +++ b/apps/aether-gateway/src/handlers/shared/normalize.rs @@ -327,6 +327,15 @@ where > as serde::Deserialize>::deserialize(deserializer).map(Some) } +pub(crate) fn deserialize_optional_f64_patch<'de, D>( + deserializer: D, +) -> Result>, D::Error> +where + D: serde::Deserializer<'de>, +{ + as serde::Deserialize>::deserialize(deserializer).map(Some) +} + fn normalize_chat_pii_redaction_feature_settings( settings: &mut Map, ) -> Result<(), String> { diff --git a/apps/aether-gateway/src/lib.rs b/apps/aether-gateway/src/lib.rs index b59cd65ba..b1374fb62 100644 --- a/apps/aether-gateway/src/lib.rs +++ b/apps/aether-gateway/src/lib.rs @@ -29,6 +29,7 @@ mod admin_api; mod ai_serving; mod allocator_metrics; mod api; +mod app_timezone; mod async_task; mod audit; mod auth; @@ -39,6 +40,7 @@ mod client_session_affinity; mod clock; mod constants; mod control; +mod daily_usage_limit; mod data; mod dispatch; mod email_delivery; diff --git a/apps/aether-gateway/src/rate_limit.rs b/apps/aether-gateway/src/rate_limit.rs index 411513dfd..3e8684932 100644 --- a/apps/aether-gateway/src/rate_limit.rs +++ b/apps/aether-gateway/src/rate_limit.rs @@ -527,8 +527,11 @@ mod tests { access_allowed: true, user_rate_limit: Some(1), api_key_rate_limit: Some(10), + user_daily_usage_limit_usd: None, + api_key_daily_usage_limit_usd: None, api_key_is_standalone: false, admin_bypass_limits: false, + ip_bypass_limits: false, local_rejection: None, allowed_models: None, ip_rules: None, @@ -567,8 +570,11 @@ mod tests { access_allowed: true, user_rate_limit: None, api_key_rate_limit: Some(10), + user_daily_usage_limit_usd: None, + api_key_daily_usage_limit_usd: None, api_key_is_standalone: false, admin_bypass_limits: false, + ip_bypass_limits: false, local_rejection: None, allowed_models: None, ip_rules: None, @@ -607,8 +613,11 @@ mod tests { access_allowed: true, user_rate_limit: Some(1), api_key_rate_limit: Some(1), + user_daily_usage_limit_usd: None, + api_key_daily_usage_limit_usd: None, api_key_is_standalone: false, admin_bypass_limits: true, + ip_bypass_limits: false, local_rejection: None, allowed_models: None, ip_rules: None, @@ -651,8 +660,11 @@ mod tests { access_allowed: true, user_rate_limit: Some(1), api_key_rate_limit: Some(10), + user_daily_usage_limit_usd: None, + api_key_daily_usage_limit_usd: None, api_key_is_standalone: false, admin_bypass_limits: false, + ip_bypass_limits: false, local_rejection: None, allowed_models: None, ip_rules: None, diff --git a/apps/aether-gateway/src/scheduler/candidate/tests/support.rs b/apps/aether-gateway/src/scheduler/candidate/tests/support.rs index 1bf9e6f13..4647231bc 100644 --- a/apps/aether-gateway/src/scheduler/candidate/tests/support.rs +++ b/apps/aether-gateway/src/scheduler/candidate/tests/support.rs @@ -117,6 +117,7 @@ pub(super) fn sample_auth_snapshot(api_key_id: &str) -> GatewayAuthApiKeySnapsho user_is_active: true, user_is_deleted: false, user_rate_limit: None, + user_daily_usage_limit_usd: None, user_allowed_providers: None, user_allowed_api_formats: None, user_allowed_models: None, @@ -126,6 +127,7 @@ pub(super) fn sample_auth_snapshot(api_key_id: &str) -> GatewayAuthApiKeySnapsho api_key_is_locked: false, api_key_is_standalone: false, api_key_rate_limit: None, + api_key_daily_usage_limit_usd: None, api_key_concurrent_limit: None, api_key_expires_at_unix_secs: None, api_key_allowed_providers: None, diff --git a/apps/aether-gateway/src/state/app.rs b/apps/aether-gateway/src/state/app.rs index 76ad14fa6..27c7d0e16 100644 --- a/apps/aether-gateway/src/state/app.rs +++ b/apps/aether-gateway/src/state/app.rs @@ -19,6 +19,7 @@ use super::super::cache::{ AuthContextCache, AuthSnapshotCache, DashboardResponseCache, DirectPlanBypassCache, JsonValueCache, SchedulerAffinityCache, SystemConfigCache, ValueCache, }; +use super::super::daily_usage_limit::FrontdoorDailyUsageLimiter; use super::super::data::GatewayDataState; use super::super::fallback_metrics; use super::super::maintenance::UsageCounterFlushRuntimeMetrics; @@ -361,6 +362,12 @@ fn soft_fd_limit() -> Option { None } +#[derive(Debug, Clone)] +pub(crate) struct FrontdoorLimiters { + pub(crate) user_rpm: Arc, + pub(crate) daily_usage: Arc, +} + #[derive(Debug, Clone)] pub struct AppState { #[cfg(test)] @@ -428,7 +435,7 @@ pub struct AppState { pub(crate) usage_counter_exact_health_metric_refresh: Arc>, pub(crate) request_candidate_queue: Option>, pub(crate) frontdoor_cors: Option>, - pub(crate) frontdoor_user_rpm: Arc, + pub(crate) frontdoor_limiters: Arc, pub(crate) tunnel: crate::tunnel::EmbeddedTunnelState, pub(crate) provider_transport_snapshot_cache: Arc>, diff --git a/apps/aether-gateway/src/state/catalog.rs b/apps/aether-gateway/src/state/catalog.rs index a809aed38..0327f84f4 100644 --- a/apps/aether-gateway/src/state/catalog.rs +++ b/apps/aether-gateway/src/state/catalog.rs @@ -1189,6 +1189,7 @@ mod tests { user_is_active: true, user_is_deleted: false, user_rate_limit: None, + user_daily_usage_limit_usd: None, user_allowed_providers: None, user_allowed_api_formats: None, user_allowed_models: None, @@ -1198,6 +1199,7 @@ mod tests { api_key_is_locked: false, api_key_is_standalone: false, api_key_rate_limit: None, + api_key_daily_usage_limit_usd: None, api_key_concurrent_limit: None, api_key_expires_at_unix_secs: None, api_key_allowed_providers: None, diff --git a/apps/aether-gateway/src/state/core.rs b/apps/aether-gateway/src/state/core.rs index 78bd09233..72bff7638 100644 --- a/apps/aether-gateway/src/state/core.rs +++ b/apps/aether-gateway/src/state/core.rs @@ -19,15 +19,15 @@ use aether_runtime::{ MetricLabel, MetricSample, }; use aether_runtime_state::{ - MemoryRuntimeStateConfig, RedisRuntimeDiagnostics, RuntimeQueueStore, RuntimeSemaphore, - RuntimeSemaphoreError, RuntimeSemaphoreSnapshot, RuntimeState, + MemoryRuntimeStateConfig, RedisRuntimeDiagnostics, RuntimeSemaphore, RuntimeSemaphoreError, + RuntimeSemaphoreSnapshot, RuntimeState, }; use aether_scheduler_core::PROVIDER_KEY_RPM_WINDOW_SECS; use dashmap::DashMap; use tokio::sync::{Mutex as TokioMutex, RwLock as TokioRwLock}; use tracing::warn; -use super::app::METRIC_SNAPSHOT_TTL; +use super::app::{FrontdoorLimiters, METRIC_SNAPSHOT_TTL}; use super::{ AppState, FrontdoorCorsConfig, FrontdoorRuntimeGuardConfig, LocalExecutionRuntimeMissDiagnostic, }; @@ -90,6 +90,7 @@ const AUTH_AFFECTING_SYSTEM_CONFIG_KEYS: &[&str] = &[ crate::constants::ANTIGRAVITY_BEARER_BRIDGE_CONFIG_KEY, ]; const FRONTDOOR_RPM_AFFECTING_SYSTEM_CONFIG_KEYS: &[&str] = &["rate_limit_per_minute"]; +const FRONTDOOR_DAILY_USAGE_AFFECTING_SYSTEM_CONFIG_KEYS: &[&str] = &["daily_usage_limit_usd"]; const CHAT_PII_REDACTION_SYSTEM_CONFIG_PREFIX: &str = "module.chat_pii_redaction."; const METRIC_SNAPSHOT_REFRESH_TIMEOUT: Duration = Duration::from_secs(4); const METRIC_SNAPSHOT_PREWARM_TIMEOUT: Duration = Duration::from_secs(12); @@ -124,6 +125,11 @@ fn system_config_key_affects_frontdoor_rpm(key: &str) -> bool { FRONTDOOR_RPM_AFFECTING_SYSTEM_CONFIG_KEYS.contains(&key) } +fn system_config_key_affects_frontdoor_daily_usage(key: &str) -> bool { + let key = key.trim(); + FRONTDOOR_DAILY_USAGE_AFFECTING_SYSTEM_CONFIG_KEYS.contains(&key) +} + fn system_config_key_affects_chat_pii_redaction(key: &str) -> bool { key.trim() .starts_with(CHAT_PII_REDACTION_SYSTEM_CONFIG_PREFIX) @@ -141,13 +147,6 @@ impl AppState { .map_err(|err| format!("{err:?}")) } - fn usage_worker_queue_for( - runtime_state: &Arc, - ) -> Option> { - let queue: Arc = runtime_state.clone(); - Some(queue) - } - fn spawn_scheduler_affinity_runtime_write( &self, cache_key: &str, @@ -218,7 +217,7 @@ impl AppState { self.background_data = Arc::new( (*data) .clone() - .with_usage_worker_queue(Self::usage_worker_queue_for(&self.runtime_state)), + .with_usage_runtime_state(self.runtime_state.clone()), ); self.background_data_isolated = false; self.replace_foreground_data_state(data); @@ -233,7 +232,7 @@ impl AppState { self.background_data = Arc::new( (*background_data) .clone() - .with_usage_worker_queue(Self::usage_worker_queue_for(&self.runtime_state)), + .with_usage_runtime_state(self.runtime_state.clone()), ); self.background_data_isolated = background_data_isolated; self.replace_foreground_data_state(data); @@ -246,11 +245,16 @@ impl AppState { self.candidate_row_page_cache.clear(); self.candidate_resolved_page_cache.clear(); self.system_config_cache.clear(); - self.frontdoor_user_rpm.clear_system_default_cache(); + self.frontdoor_limiters + .user_rpm + .clear_system_default_cache(); + self.frontdoor_limiters + .daily_usage + .clear_system_default_cache(); let data = Arc::new( (*data) .clone() - .with_usage_worker_queue(Self::usage_worker_queue_for(&self.runtime_state)), + .with_usage_runtime_state(self.runtime_state.clone()), ); self.candidate_row_page_cache.clear(); self.candidate_page_cache.clear(); @@ -288,10 +292,8 @@ impl AppState { fn build(execution_runtime_override_base_url: Option) -> Result { let runtime_state = Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default())); - let data = Arc::new( - GatewayDataState::disabled() - .with_usage_worker_queue(Self::usage_worker_queue_for(&runtime_state)), - ); + let data = + Arc::new(GatewayDataState::disabled().with_usage_runtime_state(runtime_state.clone())); let client = build_http_client(&HttpClientConfig { connect_timeout_ms: Some(10_000), request_timeout_ms: Some(300_000), @@ -387,9 +389,12 @@ impl AppState { usage_counter_exact_health_metric_refresh: Arc::new(TokioMutex::new(())), request_candidate_queue: None, frontdoor_cors: None, - frontdoor_user_rpm: Arc::new(FrontdoorUserRpmLimiter::new( - FrontdoorUserRpmConfig::default(), - )), + frontdoor_limiters: Arc::new(FrontdoorLimiters { + user_rpm: Arc::new(FrontdoorUserRpmLimiter::new( + FrontdoorUserRpmConfig::default(), + )), + daily_usage: Arc::new(crate::daily_usage_limit::FrontdoorDailyUsageLimiter::new()), + }), tunnel: crate::tunnel::EmbeddedTunnelState::with_data_and_runtime_state( data, runtime_state.clone(), @@ -588,12 +593,12 @@ impl AppState { self.data = Arc::new( (*self.data) .clone() - .with_usage_worker_queue(Self::usage_worker_queue_for(&self.runtime_state)), + .with_usage_runtime_state(self.runtime_state.clone()), ); self.background_data = Arc::new( (*self.background_data) .clone() - .with_usage_worker_queue(Self::usage_worker_queue_for(&self.runtime_state)), + .with_usage_runtime_state(self.runtime_state.clone()), ); self.tunnel = crate::tunnel::EmbeddedTunnelState::with_data_and_runtime_state( Arc::clone(&self.data), @@ -652,7 +657,8 @@ impl AppState { } pub fn with_frontdoor_user_rpm_config(mut self, config: FrontdoorUserRpmConfig) -> Self { - self.frontdoor_user_rpm = Arc::new(FrontdoorUserRpmLimiter::new(config)); + Arc::make_mut(&mut self.frontdoor_limiters).user_rpm = + Arc::new(FrontdoorUserRpmLimiter::new(config)); self } @@ -677,7 +683,13 @@ impl AppState { } pub(crate) fn frontdoor_user_rpm(&self) -> Arc { - Arc::clone(&self.frontdoor_user_rpm) + Arc::clone(&self.frontdoor_limiters.user_rpm) + } + + pub(crate) fn frontdoor_daily_usage( + &self, + ) -> Arc { + Arc::clone(&self.frontdoor_limiters.daily_usage) } pub(crate) fn mark_provider_key_rpm_reset(&self, key_id: &str, now_unix_secs: u64) { @@ -869,7 +881,14 @@ impl AppState { self.invalidate_auth_context_cache(); } if deleted && system_config_key_affects_frontdoor_rpm(key) { - self.frontdoor_user_rpm.clear_system_default_cache(); + self.frontdoor_limiters + .user_rpm + .clear_system_default_cache(); + } + if deleted && system_config_key_affects_frontdoor_daily_usage(key) { + self.frontdoor_limiters + .daily_usage + .clear_system_default_cache(); } if deleted && system_config_key_affects_chat_pii_redaction(key) { crate::privacy::clear_chat_pii_redaction_runtime_config_cache( @@ -951,7 +970,14 @@ impl AppState { self.invalidate_auth_context_cache(); } if system_config_key_affects_frontdoor_rpm(key) { - self.frontdoor_user_rpm.clear_system_default_cache(); + self.frontdoor_limiters + .user_rpm + .clear_system_default_cache(); + } + if system_config_key_affects_frontdoor_daily_usage(key) { + self.frontdoor_limiters + .daily_usage + .clear_system_default_cache(); } if system_config_key_affects_chat_pii_redaction(key) { crate::privacy::clear_chat_pii_redaction_runtime_config_cache( @@ -1575,6 +1601,12 @@ impl AppState { MetricKind::Gauge, u64::from(self.background_data_isolated), )); + samples.push(MetricSample::new( + "frontdoor_daily_usage_runtime_failures_total", + "Total daily usage limit checks that failed open because runtime state was unavailable.", + MetricKind::Counter, + self.frontdoor_limiters.daily_usage.runtime_failure_count(), + )); if self.background_data_isolated { if let Some(summary) = self.background_data.database_pool_summary() { samples.extend(background_database_pool_metric_samples(&summary)); diff --git a/apps/aether-gateway/src/state/runtime/auth/user_lifecycle.rs b/apps/aether-gateway/src/state/runtime/auth/user_lifecycle.rs index bc5cca72d..a9b0b57c5 100644 --- a/apps/aether-gateway/src/state/runtime/auth/user_lifecycle.rs +++ b/apps/aether-gateway/src/state/runtime/auth/user_lifecycle.rs @@ -926,6 +926,8 @@ mod tests { allowed_models_mode: allowed_models_mode.to_string(), rate_limit: None, rate_limit_mode: "inherit".to_string(), + daily_usage_limit_usd: None, + daily_usage_limit_mode: "inherit".to_string(), } } @@ -939,8 +941,11 @@ mod tests { access_allowed: true, user_rate_limit: None, api_key_rate_limit: None, + user_daily_usage_limit_usd: None, + api_key_daily_usage_limit_usd: None, api_key_is_standalone: false, admin_bypass_limits: false, + ip_bypass_limits: false, local_rejection: None, allowed_models: Some(vec!["gpt-4.1".to_string()]), ip_rules: None, diff --git a/apps/aether-gateway/src/state/testing.rs b/apps/aether-gateway/src/state/testing.rs index 504ee8c6f..c19bd2fbb 100644 --- a/apps/aether-gateway/src/state/testing.rs +++ b/apps/aether-gateway/src/state/testing.rs @@ -149,8 +149,17 @@ impl AppState { } pub(crate) fn with_frontdoor_system_default_rpm_for_tests(mut self, limit: u32) -> Self { - self.frontdoor_user_rpm = Arc::new( - (*self.frontdoor_user_rpm) + Arc::make_mut(&mut self.frontdoor_limiters).user_rpm = Arc::new( + (*self.frontdoor_limiters.user_rpm) + .clone() + .with_system_default_limit_for_tests(limit), + ); + self + } + + pub(crate) fn with_frontdoor_system_daily_usage_limit_for_tests(mut self, limit: f64) -> Self { + Arc::make_mut(&mut self.frontdoor_limiters).daily_usage = Arc::new( + (*self.frontdoor_limiters.daily_usage) .clone() .with_system_default_limit_for_tests(limit), ); diff --git a/apps/aether-gateway/src/tests/control/admin/system.rs b/apps/aether-gateway/src/tests/control/admin/system.rs index d2763f71a..bfe30b556 100644 --- a/apps/aether-gateway/src/tests/control/admin/system.rs +++ b/apps/aether-gateway/src/tests/control/admin/system.rs @@ -904,6 +904,8 @@ async fn gateway_handles_admin_system_users_export_locally_with_trusted_admin_pr allowed_models_mode: "specific".to_string(), rate_limit: Some(60), rate_limit_mode: "custom".to_string(), + daily_usage_limit_usd: None, + daily_usage_limit_mode: "inherit".to_string(), }) .await .expect("user group should create") diff --git a/apps/aether-gateway/src/tests/control/admin/users.rs b/apps/aether-gateway/src/tests/control/admin/users.rs index 86558d82e..ebe08c255 100644 --- a/apps/aether-gateway/src/tests/control/admin/users.rs +++ b/apps/aether-gateway/src/tests/control/admin/users.rs @@ -611,6 +611,8 @@ async fn gateway_allows_default_user_group_access_policy_updates() { allowed_models_mode: "specific".to_string(), rate_limit: Some(60), rate_limit_mode: "custom".to_string(), + daily_usage_limit_usd: None, + daily_usage_limit_mode: "inherit".to_string(), }) .await .expect("user group should create") @@ -723,6 +725,8 @@ async fn gateway_allows_removing_default_group_members_when_other_group_remains( allowed_models_mode: "unrestricted".to_string(), rate_limit: None, rate_limit_mode: "system".to_string(), + daily_usage_limit_usd: None, + daily_usage_limit_mode: "inherit".to_string(), }) .await .expect("default group should create") @@ -740,6 +744,8 @@ async fn gateway_allows_removing_default_group_members_when_other_group_remains( allowed_models_mode: "unrestricted".to_string(), rate_limit: None, rate_limit_mode: "system".to_string(), + daily_usage_limit_usd: None, + daily_usage_limit_mode: "inherit".to_string(), }) .await .expect("team group should create") @@ -1701,6 +1707,8 @@ async fn admin_created_user_keys_inherit_owner_group_policy() { allowed_models_mode: "specific".to_string(), rate_limit: Some(100), rate_limit_mode: "custom".to_string(), + daily_usage_limit_usd: None, + daily_usage_limit_mode: "inherit".to_string(), }) .await .expect("admin group should create") @@ -1723,6 +1731,8 @@ async fn admin_created_user_keys_inherit_owner_group_policy() { allowed_models_mode: "specific".to_string(), rate_limit: Some(30), rate_limit_mode: "custom".to_string(), + daily_usage_limit_usd: None, + daily_usage_limit_mode: "inherit".to_string(), }) .await .expect("target group should create") @@ -1810,6 +1820,8 @@ async fn admin_created_user_keys_inherit_owner_group_policy() { allowed_models_mode: "specific".to_string(), rate_limit: Some(15), rate_limit_mode: "custom".to_string(), + daily_usage_limit_usd: None, + daily_usage_limit_mode: "inherit".to_string(), }, ) .await @@ -2132,6 +2144,8 @@ async fn gateway_allows_admin_update_user_to_clear_explicit_groups() { allowed_models_mode: "unrestricted".to_string(), rate_limit: None, rate_limit_mode: "system".to_string(), + daily_usage_limit_usd: None, + daily_usage_limit_mode: "inherit".to_string(), }) .await .expect("default group should create") diff --git a/apps/aether-gateway/src/tests/frontdoor/public_support.rs b/apps/aether-gateway/src/tests/frontdoor/public_support.rs index 6eb1fd7a1..bf812de0e 100644 --- a/apps/aether-gateway/src/tests/frontdoor/public_support.rs +++ b/apps/aether-gateway/src/tests/frontdoor/public_support.rs @@ -1,13 +1,13 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; use super::{ - hash_api_key, sample_endpoint, sample_key, sample_models_candidate_row, sample_provider, - sample_public_catalog_model, sample_public_global_model, - sample_public_global_model_with_capabilities, sample_request_candidate, - InMemoryAnnouncementReadRepository, InMemoryGlobalModelReadRepository, - InMemoryMinimalCandidateSelectionReadRepository, InMemoryProviderCatalogReadRepository, - InMemoryRequestCandidateRepository, RequestCandidateStatus, StoredAnnouncement, - StoredPublicGlobalModel, + hash_api_key, run_frontdoor_async_test, sample_endpoint, sample_key, + sample_models_candidate_row, sample_provider, sample_public_catalog_model, + sample_public_global_model, sample_public_global_model_with_capabilities, + sample_request_candidate, InMemoryAnnouncementReadRepository, + InMemoryGlobalModelReadRepository, InMemoryMinimalCandidateSelectionReadRepository, + InMemoryProviderCatalogReadRepository, InMemoryRequestCandidateRepository, + RequestCandidateStatus, StoredAnnouncement, StoredPublicGlobalModel, }; use crate::data::GatewayDataState; use crate::tests::{ @@ -2846,8 +2846,16 @@ async fn gateway_handles_user_monitoring_audit_logs_locally_without_proxying_ups upstream_handle.abort(); } -#[tokio::test] -async fn gateway_handles_user_monitoring_rate_limit_status_locally_without_proxying_upstream() { +#[test] +fn gateway_handles_user_monitoring_rate_limit_status_locally_without_proxying_upstream() { + run_frontdoor_async_test( + "gateway_handles_user_monitoring_rate_limit_status_locally_without_proxying_upstream", + gateway_handles_user_monitoring_rate_limit_status_locally_without_proxying_upstream_impl(), + ); +} + +async fn gateway_handles_user_monitoring_rate_limit_status_locally_without_proxying_upstream_impl() +{ let now = Utc::now(); let user = sample_auth_user(now); let access_token = build_test_auth_token( @@ -2912,6 +2920,8 @@ async fn gateway_handles_user_monitoring_rate_limit_status_locally_without_proxy allowed_models_mode: "unrestricted".to_string(), rate_limit: Some(80), rate_limit_mode: "custom".to_string(), + daily_usage_limit_usd: None, + daily_usage_limit_mode: "inherit".to_string(), }) .await .expect("group should create") @@ -8315,6 +8325,8 @@ async fn gateway_handles_users_me_providers_locally_without_proxying_upstream() allowed_models_mode: "unrestricted".to_string(), rate_limit: None, rate_limit_mode: "system".to_string(), + daily_usage_limit_usd: None, + daily_usage_limit_mode: "inherit".to_string(), }) .await .expect("group should create") @@ -9773,6 +9785,8 @@ async fn gateway_refreshes_users_me_available_models_after_group_assignment() { allowed_models_mode: "specific".to_string(), rate_limit: None, rate_limit_mode: "system".to_string(), + daily_usage_limit_usd: None, + daily_usage_limit_mode: "inherit".to_string(), }) .await .expect("group should create") @@ -9881,6 +9895,8 @@ async fn gateway_returns_no_users_me_available_models_when_group_denies_all_mode allowed_models_mode: "deny_all".to_string(), rate_limit: None, rate_limit_mode: "system".to_string(), + daily_usage_limit_usd: None, + daily_usage_limit_mode: "inherit".to_string(), }) .await .expect("group should create") @@ -9974,6 +9990,8 @@ async fn gateway_returns_service_unavailable_for_users_me_available_models_witho allowed_models_mode: "unrestricted".to_string(), rate_limit: None, rate_limit_mode: "system".to_string(), + daily_usage_limit_usd: None, + daily_usage_limit_mode: "inherit".to_string(), }) .await .expect("group should create") diff --git a/apps/aether-gateway/src/video_tasks/tests/fixtures.rs b/apps/aether-gateway/src/video_tasks/tests/fixtures.rs index 9b588b813..67f17899c 100644 --- a/apps/aether-gateway/src/video_tasks/tests/fixtures.rs +++ b/apps/aether-gateway/src/video_tasks/tests/fixtures.rs @@ -50,8 +50,11 @@ pub(super) fn sample_auth_context() -> GatewayControlAuthContext { access_allowed: true, user_rate_limit: None, api_key_rate_limit: None, + user_daily_usage_limit_usd: None, + api_key_daily_usage_limit_usd: None, api_key_is_standalone: false, admin_bypass_limits: false, + ip_bypass_limits: false, local_rejection: None, allowed_models: None, ip_rules: None, diff --git a/apps/aether-gateway/src/wallet_runtime/access.rs b/apps/aether-gateway/src/wallet_runtime/access.rs index 1d314c1d6..e7ed3841d 100644 --- a/apps/aether-gateway/src/wallet_runtime/access.rs +++ b/apps/aether-gateway/src/wallet_runtime/access.rs @@ -398,6 +398,7 @@ mod tests { user_is_active: true, user_is_deleted: false, user_rate_limit: None, + user_daily_usage_limit_usd: None, user_allowed_providers: None, user_allowed_api_formats: None, user_allowed_models: None, @@ -407,6 +408,7 @@ mod tests { api_key_is_locked: false, api_key_is_standalone: false, api_key_rate_limit: None, + api_key_daily_usage_limit_usd: None, api_key_concurrent_limit: None, api_key_expires_at_unix_secs: None, api_key_allowed_providers: None, diff --git a/crates/aether-admin/src/observability/stats.rs b/crates/aether-admin/src/observability/stats.rs index a041b6524..3aa20c1b8 100644 --- a/crates/aether-admin/src/observability/stats.rs +++ b/crates/aether-admin/src/observability/stats.rs @@ -2113,6 +2113,7 @@ mod tests { user_is_active: true, user_is_deleted: false, user_rate_limit: None, + user_daily_usage_limit_usd: None, user_allowed_providers: None, user_allowed_api_formats: None, user_allowed_models: None, @@ -2122,6 +2123,7 @@ mod tests { api_key_is_locked: false, api_key_is_standalone: false, api_key_rate_limit: None, + api_key_daily_usage_limit_usd: None, api_key_concurrent_limit: None, api_key_expires_at_unix_secs: None, api_key_allowed_providers: None, diff --git a/crates/aether-admin/src/system.rs b/crates/aether-admin/src/system.rs index b204a8d95..77eded62e 100644 --- a/crates/aether-admin/src/system.rs +++ b/crates/aether-admin/src/system.rs @@ -1689,6 +1689,7 @@ pub fn admin_system_config_default_value(key: &str) -> Option "site_name" => Some(json!("Aether")), "site_subtitle" => Some(json!("AI Gateway")), "default_user_initial_gift_usd" => Some(json!(10.0)), + "daily_usage_limit_usd" => Some(json!(0.0)), "password_policy_level" => Some(json!("weak")), REQUEST_RECORD_LEVEL_KEY => Some(json!("full")), "max_request_body_size" => Some(json!(0)), diff --git a/crates/aether-data/adapters/mysql/migrations/20260803000000_add_daily_usage_limits.sql b/crates/aether-data/adapters/mysql/migrations/20260803000000_add_daily_usage_limits.sql new file mode 100644 index 000000000..ff5091bb1 --- /dev/null +++ b/crates/aether-data/adapters/mysql/migrations/20260803000000_add_daily_usage_limits.sql @@ -0,0 +1,6 @@ +ALTER TABLE api_keys + ADD COLUMN daily_usage_limit_usd DOUBLE NULL; + +ALTER TABLE user_groups + ADD COLUMN daily_usage_limit_usd DOUBLE NULL, + ADD COLUMN daily_usage_limit_mode VARCHAR(32) NOT NULL DEFAULT 'inherit'; diff --git a/crates/aether-data/adapters/mysql/src/auth.rs b/crates/aether-data/adapters/mysql/src/auth.rs index 74527704c..fd3173aeb 100644 --- a/crates/aether-data/adapters/mysql/src/auth.rs +++ b/crates/aether-data/adapters/mysql/src/auth.rs @@ -31,6 +31,7 @@ SELECT api_keys.is_locked AS api_key_is_locked, api_keys.is_standalone AS api_key_is_standalone, api_keys.rate_limit AS api_key_rate_limit, + api_keys.daily_usage_limit_usd AS api_key_daily_usage_limit_usd, api_keys.concurrent_limit AS api_key_concurrent_limit, api_keys.expires_at AS api_key_expires_at_unix_secs, api_keys.allowed_providers AS api_key_allowed_providers, @@ -53,6 +54,7 @@ SELECT api_keys.allowed_models, api_keys.ip_rules, api_keys.rate_limit, + api_keys.daily_usage_limit_usd, api_keys.concurrent_limit, api_keys.force_capabilities, api_keys.feature_settings, @@ -115,12 +117,13 @@ impl MysqlAuthApiKeyReadRepository { r#" INSERT INTO api_keys ( id, user_id, key_hash, key_encrypted, name, allowed_providers, - allowed_api_formats, allowed_models, ip_rules, rate_limit, concurrent_limit, + allowed_api_formats, allowed_models, ip_rules, rate_limit, daily_usage_limit_usd, + concurrent_limit, force_capabilities, feature_settings, is_active, expires_at, auto_delete_on_expiry, total_requests, total_tokens, total_cost_usd, is_standalone, created_at, updated_at ) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, ?, ?, ?, ?, ?, ?, ?, ?, ?) +VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, ?, ?, ?, ?, ?, ?, ?, ?, ?) "#, ) .bind(&record.api_key_id) @@ -145,6 +148,7 @@ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, ?, ?, ?, ?, ?, ?, ?, ?, ?) "api_keys.ip_rules", )?) .bind(record.rate_limit) + .bind(record.daily_usage_limit_usd) .bind(record.concurrent_limit) .bind(optional_json_to_string( &record.force_capabilities, @@ -184,6 +188,7 @@ struct CreateApiKeyInsertRecord { allowed_models: Option>, ip_rules: Option>, rate_limit: Option, + daily_usage_limit_usd: Option, concurrent_limit: Option, force_capabilities: Option, is_active: bool, @@ -427,6 +432,7 @@ WHERE id = ? allowed_models: record.allowed_models, ip_rules: record.ip_rules, rate_limit: Some(record.rate_limit), + daily_usage_limit_usd: record.daily_usage_limit_usd, concurrent_limit: record.concurrent_limit, force_capabilities: record.force_capabilities, is_active: record.is_active, @@ -455,6 +461,7 @@ WHERE id = ? allowed_models: record.allowed_models, ip_rules: record.ip_rules, rate_limit: record.rate_limit, + daily_usage_limit_usd: record.daily_usage_limit_usd, concurrent_limit: record.concurrent_limit, force_capabilities: record.force_capabilities, is_active: record.is_active, @@ -478,6 +485,7 @@ WHERE id = ? UPDATE api_keys SET name = COALESCE(?, name), rate_limit = COALESCE(?, rate_limit), + daily_usage_limit_usd = CASE WHEN ? THEN ? ELSE daily_usage_limit_usd END, concurrent_limit = COALESCE(?, concurrent_limit), ip_rules = CASE WHEN ? THEN ? ELSE ip_rules END, updated_at = ? @@ -488,6 +496,8 @@ WHERE id = ? ) .bind(record.name.as_deref()) .bind(record.rate_limit) + .bind(record.daily_usage_limit_present) + .bind(record.daily_usage_limit_usd) .bind(record.concurrent_limit) .bind(record.ip_rules.is_some()) .bind(json_string_from_nested_string_list( @@ -513,6 +523,7 @@ WHERE id = ? UPDATE api_keys SET name = COALESCE(?, name), rate_limit = CASE WHEN ? THEN ? ELSE rate_limit END, + daily_usage_limit_usd = CASE WHEN ? THEN ? ELSE daily_usage_limit_usd END, concurrent_limit = CASE WHEN ? THEN ? ELSE concurrent_limit END, allowed_providers = CASE WHEN ? THEN ? ELSE allowed_providers END, allowed_api_formats = CASE WHEN ? THEN ? ELSE allowed_api_formats END, @@ -528,6 +539,8 @@ WHERE id = ? .bind(record.name.as_deref()) .bind(record.rate_limit_present) .bind(record.rate_limit) + .bind(record.daily_usage_limit_present) + .bind(record.daily_usage_limit_usd) .bind(record.concurrent_limit_present) .bind(record.concurrent_limit) .bind(record.allowed_providers.is_some()) @@ -977,7 +990,12 @@ fn map_auth_api_key_snapshot_row( row.try_get("api_key_ip_rules").map_sql_err()?, "api_keys.ip_rules", )?)?; - Ok(snapshot.with_user_rate_limit(row.try_get("user_rate_limit").map_sql_err()?)) + Ok(snapshot + .with_user_rate_limit(row.try_get("user_rate_limit").map_sql_err()?) + .with_daily_usage_limits( + None, + row.try_get("api_key_daily_usage_limit_usd").map_sql_err()?, + )) } fn map_auth_api_key_export_row( @@ -1025,6 +1043,9 @@ fn map_auth_api_key_export_row( "api_keys.ip_rules", )?) }) + .map(|record| { + record.with_daily_usage_limit(row.try_get("daily_usage_limit_usd").ok().flatten()) + }) .map(|record| record.with_feature_settings(feature_settings)) .and_then(|record| { record.with_activity_timestamps( diff --git a/crates/aether-data/adapters/mysql/src/usage/read.rs b/crates/aether-data/adapters/mysql/src/usage/read.rs index bca791beb..99bb4dbae 100644 --- a/crates/aether-data/adapters/mysql/src/usage/read.rs +++ b/crates/aether-data/adapters/mysql/src/usage/read.rs @@ -1,9 +1,10 @@ use aether_data_contracts::repository::usage::{ - ProviderApiKeyWindowUsageRequest, StoredRequestUsageAudit, UsageAuditKeywordSearchQuery, - UsageAuditListQuery, UsageMonitoringErrorCountQuery, UsageMonitoringErrorListQuery, + ProviderApiKeyWindowUsageRequest, StoredRequestUsageAudit, StoredUsageDailyActualCostRollup, + UsageAuditKeywordSearchQuery, UsageAuditListQuery, UsageDailyActualCostRollupQuery, + UsageMonitoringErrorCountQuery, UsageMonitoringErrorListQuery, }; use aether_data_contracts::DataLayerError; -use sqlx::{MySql, QueryBuilder}; +use sqlx::{MySql, QueryBuilder, Row}; use crate::error::SqlResultExt; @@ -135,6 +136,51 @@ impl MysqlUsageReadFilter { } impl MysqlUsageStorage { + pub async fn summarize_usage_daily_actual_cost_rollups( + &self, + query: &UsageDailyActualCostRollupQuery, + ) -> Result, DataLayerError> { + if query.finalized_from_unix_secs >= query.finalized_until_unix_secs { + return Ok(Vec::new()); + } + let rows = sqlx::query( + r#" +SELECT u.user_id, + u.api_key_id, + CASE + WHEN LOWER(COALESCE(JSON_UNQUOTE(JSON_EXTRACT( + u.request_metadata, '$.api_key_is_standalone' + )), 'false')) IN ('true', '1') + THEN 1 ELSE 0 + END AS api_key_is_standalone, + SUM(CAST(u.actual_total_cost_usd AS DOUBLE)) AS actual_total_cost_usd +FROM `usage` AS u +WHERE u.status = 'completed' + AND u.finalized_at >= ? + AND u.finalized_at < ? + AND u.api_key_id IS NOT NULL + AND TRIM(u.api_key_id) <> '' + AND CAST(u.actual_total_cost_usd AS DOUBLE) > 0 +GROUP BY u.user_id, u.api_key_id, 3 +"#, + ) + .bind(query.finalized_from_unix_secs as i64) + .bind(query.finalized_until_unix_secs as i64) + .fetch_all(&self.pool) + .await + .map_sql_err()?; + rows.into_iter() + .map(|row| { + Ok(StoredUsageDailyActualCostRollup { + user_id: row.try_get("user_id").map_sql_err()?, + api_key_id: row.try_get("api_key_id").map_sql_err()?, + api_key_is_standalone: row.try_get("api_key_is_standalone").map_sql_err()?, + actual_total_cost_usd: row.try_get("actual_total_cost_usd").map_sql_err()?, + }) + }) + .collect() + } + pub async fn find_by_id( &self, id: &str, diff --git a/crates/aether-data/adapters/mysql/src/users.rs b/crates/aether-data/adapters/mysql/src/users.rs index e36a3e8a5..cf2650c3e 100644 --- a/crates/aether-data/adapters/mysql/src/users.rs +++ b/crates/aether-data/adapters/mysql/src/users.rs @@ -159,6 +159,8 @@ SELECT allowed_models_mode, rate_limit, rate_limit_mode, + daily_usage_limit_usd, + daily_usage_limit_mode, created_at, updated_at FROM user_groups @@ -480,9 +482,10 @@ INSERT INTO user_groups ( allowed_providers, allowed_providers_mode, allowed_api_formats, allowed_api_formats_mode, allowed_models, allowed_models_mode, - rate_limit, rate_limit_mode, created_at, updated_at + rate_limit, rate_limit_mode, daily_usage_limit_usd, daily_usage_limit_mode, + created_at, updated_at ) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) "#, ) .bind(&id) @@ -502,6 +505,8 @@ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) .bind(record.allowed_models_mode) .bind(record.rate_limit) .bind(record.rate_limit_mode) + .bind(record.daily_usage_limit_usd) + .bind(record.daily_usage_limit_mode) .bind(now) .bind(now) .execute(&self.pool) @@ -538,6 +543,8 @@ SET name = ?, allowed_models_mode = ?, rate_limit = ?, rate_limit_mode = ?, + daily_usage_limit_usd = ?, + daily_usage_limit_mode = ?, updated_at = ? WHERE id = ? "#, @@ -558,6 +565,8 @@ WHERE id = ? .bind(record.allowed_models_mode) .bind(record.rate_limit) .bind(record.rate_limit_mode) + .bind(record.daily_usage_limit_usd) + .bind(record.daily_usage_limit_mode) .bind(now) .bind(group_id) .execute(&self.pool) @@ -2036,6 +2045,12 @@ fn map_user_group_row(row: &MySqlRow) -> Result optional_datetime_from_unix_secs(row.try_get("created_at").map_sql_err()?), optional_datetime_from_unix_secs(row.try_get("updated_at").map_sql_err()?), ) + .and_then(|group| { + group.with_daily_usage_limit( + row.try_get("daily_usage_limit_usd").map_sql_err()?, + row.try_get("daily_usage_limit_mode").map_sql_err()?, + ) + }) } fn map_user_group_member_row(row: &MySqlRow) -> Result { diff --git a/crates/aether-data/adapters/postgres/migrations/20260803000000_add_daily_usage_limits.sql b/crates/aether-data/adapters/postgres/migrations/20260803000000_add_daily_usage_limits.sql new file mode 100644 index 000000000..d00e3d092 --- /dev/null +++ b/crates/aether-data/adapters/postgres/migrations/20260803000000_add_daily_usage_limits.sql @@ -0,0 +1,12 @@ +ALTER TABLE api_keys + ADD COLUMN IF NOT EXISTS daily_usage_limit_usd double precision; + +ALTER TABLE user_groups + ADD COLUMN IF NOT EXISTS daily_usage_limit_usd double precision, + ADD COLUMN IF NOT EXISTS daily_usage_limit_mode text NOT NULL DEFAULT 'inherit'; + +ALTER TABLE user_groups + DROP CONSTRAINT IF EXISTS user_groups_daily_usage_limit_mode_check; +ALTER TABLE user_groups + ADD CONSTRAINT user_groups_daily_usage_limit_mode_check + CHECK (daily_usage_limit_mode IN ('inherit', 'system', 'custom')); diff --git a/crates/aether-data/adapters/postgres/src/auth.rs b/crates/aether-data/adapters/postgres/src/auth.rs index b36da3259..5142f75bb 100644 --- a/crates/aether-data/adapters/postgres/src/auth.rs +++ b/crates/aether-data/adapters/postgres/src/auth.rs @@ -31,6 +31,7 @@ SELECT api_keys.is_locked AS api_key_is_locked, api_keys.is_standalone AS api_key_is_standalone, api_keys.rate_limit AS api_key_rate_limit, + api_keys.daily_usage_limit_usd AS api_key_daily_usage_limit_usd, api_keys.concurrent_limit AS api_key_concurrent_limit, CAST(EXTRACT(EPOCH FROM api_keys.expires_at) AS BIGINT) AS api_key_expires_at_unix_secs, api_keys.allowed_providers AS api_key_allowed_providers, @@ -62,6 +63,7 @@ SELECT api_keys.is_locked AS api_key_is_locked, api_keys.is_standalone AS api_key_is_standalone, api_keys.rate_limit AS api_key_rate_limit, + api_keys.daily_usage_limit_usd AS api_key_daily_usage_limit_usd, api_keys.concurrent_limit AS api_key_concurrent_limit, CAST(EXTRACT(EPOCH FROM api_keys.expires_at) AS BIGINT) AS api_key_expires_at_unix_secs, api_keys.allowed_providers AS api_key_allowed_providers, @@ -93,6 +95,7 @@ SELECT api_keys.is_locked AS api_key_is_locked, api_keys.is_standalone AS api_key_is_standalone, api_keys.rate_limit AS api_key_rate_limit, + api_keys.daily_usage_limit_usd AS api_key_daily_usage_limit_usd, api_keys.concurrent_limit AS api_key_concurrent_limit, CAST(EXTRACT(EPOCH FROM api_keys.expires_at) AS BIGINT) AS api_key_expires_at_unix_secs, api_keys.allowed_providers AS api_key_allowed_providers, @@ -124,6 +127,7 @@ SELECT api_keys.is_locked AS api_key_is_locked, api_keys.is_standalone AS api_key_is_standalone, api_keys.rate_limit AS api_key_rate_limit, + api_keys.daily_usage_limit_usd AS api_key_daily_usage_limit_usd, api_keys.concurrent_limit AS api_key_concurrent_limit, CAST(EXTRACT(EPOCH FROM api_keys.expires_at) AS BIGINT) AS api_key_expires_at_unix_secs, api_keys.allowed_providers AS api_key_allowed_providers, @@ -148,6 +152,7 @@ SELECT api_keys.allowed_models, api_keys.ip_rules, api_keys.rate_limit, + api_keys.daily_usage_limit_usd, api_keys.concurrent_limit, api_keys.force_capabilities, api_keys.feature_settings, @@ -179,6 +184,7 @@ SELECT api_keys.allowed_models, api_keys.ip_rules, api_keys.rate_limit, + api_keys.daily_usage_limit_usd, api_keys.concurrent_limit, api_keys.force_capabilities, api_keys.feature_settings, @@ -209,6 +215,7 @@ SELECT api_keys.allowed_models, api_keys.ip_rules, api_keys.rate_limit, + api_keys.daily_usage_limit_usd, api_keys.concurrent_limit, api_keys.force_capabilities, api_keys.feature_settings, @@ -239,6 +246,7 @@ SELECT api_keys.allowed_models, api_keys.ip_rules, api_keys.rate_limit, + api_keys.daily_usage_limit_usd, api_keys.concurrent_limit, api_keys.force_capabilities, api_keys.feature_settings, @@ -269,6 +277,7 @@ SELECT api_keys.allowed_models, api_keys.ip_rules, api_keys.rate_limit, + api_keys.daily_usage_limit_usd, api_keys.concurrent_limit, api_keys.force_capabilities, api_keys.feature_settings, @@ -343,6 +352,7 @@ SELECT api_keys.allowed_models, api_keys.ip_rules, api_keys.rate_limit, + api_keys.daily_usage_limit_usd, api_keys.concurrent_limit, api_keys.force_capabilities, api_keys.feature_settings, @@ -380,6 +390,7 @@ INSERT INTO api_keys ( allowed_models, ip_rules, rate_limit, + daily_usage_limit_usd, concurrent_limit, force_capabilities, feature_settings, @@ -407,15 +418,16 @@ VALUES ( $10, $11, $12, - NULL, $13, + NULL, $14, $15, + $16, FALSE, FALSE, - $16, $17, $18, + $19, NOW(), NOW() ) @@ -430,6 +442,7 @@ RETURNING allowed_models, ip_rules, rate_limit, + daily_usage_limit_usd, concurrent_limit, force_capabilities, feature_settings, @@ -457,6 +470,7 @@ INSERT INTO api_keys ( allowed_models, ip_rules, rate_limit, + daily_usage_limit_usd, concurrent_limit, force_capabilities, feature_settings, @@ -484,15 +498,16 @@ VALUES ( $10, $11, $12, - NULL, $13, + NULL, $14, $15, + $16, FALSE, TRUE, - $16, $17, $18, + $19, NOW(), NOW() ) @@ -507,6 +522,7 @@ RETURNING allowed_models, ip_rules, rate_limit, + daily_usage_limit_usd, concurrent_limit, force_capabilities, feature_settings, @@ -527,8 +543,9 @@ UPDATE api_keys SET name = COALESCE($3, name), rate_limit = COALESCE($4, rate_limit), - concurrent_limit = COALESCE($5, concurrent_limit), - ip_rules = CASE WHEN $6 THEN $7::jsonb ELSE ip_rules END, + daily_usage_limit_usd = CASE WHEN $5 THEN $6 ELSE daily_usage_limit_usd END, + concurrent_limit = COALESCE($7, concurrent_limit), + ip_rules = CASE WHEN $8 THEN $9::jsonb ELSE ip_rules END, updated_at = NOW() WHERE user_id = $1 AND id = $2 @@ -544,6 +561,7 @@ RETURNING allowed_models, ip_rules, rate_limit, + daily_usage_limit_usd, concurrent_limit, force_capabilities, feature_settings, @@ -564,13 +582,14 @@ UPDATE api_keys SET name = COALESCE($2, name), rate_limit = CASE WHEN $3 THEN $4 ELSE rate_limit END, - concurrent_limit = CASE WHEN $5 THEN $6 ELSE concurrent_limit END, - allowed_providers = CASE WHEN $7 THEN $8::json ELSE allowed_providers END, - allowed_api_formats = CASE WHEN $9 THEN $10::json ELSE allowed_api_formats END, - allowed_models = CASE WHEN $11 THEN $12::json ELSE allowed_models END, - ip_rules = CASE WHEN $13 THEN $14::jsonb ELSE ip_rules END, - expires_at = CASE WHEN $15 THEN $16::timestamptz ELSE expires_at END, - auto_delete_on_expiry = CASE WHEN $17 THEN $18 ELSE auto_delete_on_expiry END, + daily_usage_limit_usd = CASE WHEN $5 THEN $6 ELSE daily_usage_limit_usd END, + concurrent_limit = CASE WHEN $7 THEN $8 ELSE concurrent_limit END, + allowed_providers = CASE WHEN $9 THEN $10::json ELSE allowed_providers END, + allowed_api_formats = CASE WHEN $11 THEN $12::json ELSE allowed_api_formats END, + allowed_models = CASE WHEN $13 THEN $14::json ELSE allowed_models END, + ip_rules = CASE WHEN $15 THEN $16::jsonb ELSE ip_rules END, + expires_at = CASE WHEN $17 THEN $18::timestamptz ELSE expires_at END, + auto_delete_on_expiry = CASE WHEN $19 THEN $20 ELSE auto_delete_on_expiry END, updated_at = NOW() WHERE id = $1 AND is_standalone = TRUE @@ -585,6 +604,7 @@ RETURNING allowed_models, ip_rules, rate_limit, + daily_usage_limit_usd, concurrent_limit, force_capabilities, feature_settings, @@ -619,6 +639,7 @@ RETURNING allowed_models, ip_rules, rate_limit, + daily_usage_limit_usd, concurrent_limit, force_capabilities, feature_settings, @@ -652,6 +673,7 @@ RETURNING allowed_models, ip_rules, rate_limit, + daily_usage_limit_usd, concurrent_limit, force_capabilities, feature_settings, @@ -686,6 +708,7 @@ RETURNING allowed_models, ip_rules, rate_limit, + daily_usage_limit_usd, concurrent_limit, force_capabilities, feature_settings, @@ -730,6 +753,7 @@ RETURNING allowed_models, ip_rules, rate_limit, + daily_usage_limit_usd, concurrent_limit, force_capabilities, feature_settings, @@ -763,6 +787,7 @@ RETURNING allowed_api_formats, allowed_models, rate_limit, + daily_usage_limit_usd, concurrent_limit, force_capabilities, feature_settings, @@ -1190,6 +1215,7 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository { .bind(allowed_models) .bind(ip_rules) .bind(record.rate_limit) + .bind(record.daily_usage_limit_usd) .bind(record.concurrent_limit) .bind(record.force_capabilities) .bind(record.is_active) @@ -1247,6 +1273,7 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository { .bind(allowed_models) .bind(ip_rules) .bind(record.rate_limit) + .bind(record.daily_usage_limit_usd) .bind(record.concurrent_limit) .bind(record.force_capabilities) .bind(record.is_active) @@ -1277,6 +1304,8 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository { .bind(record.api_key_id) .bind(record.name) .bind(record.rate_limit) + .bind(record.daily_usage_limit_present) + .bind(record.daily_usage_limit_usd) .bind(record.concurrent_limit) .bind(record.ip_rules.is_some()) .bind(ip_rules) @@ -1331,6 +1360,8 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository { .bind(record.name) .bind(record.rate_limit_present) .bind(record.rate_limit) + .bind(record.daily_usage_limit_present) + .bind(record.daily_usage_limit_usd) .bind(record.concurrent_limit_present) .bind(record.concurrent_limit) .bind(record.allowed_providers.is_some()) @@ -1569,7 +1600,9 @@ fn map_auth_api_key_snapshot_row( row_get(row, "api_key_allowed_models")?, )? .with_api_key_ip_rules(row_get(row, "api_key_ip_rules")?)?; - Ok(snapshot.with_user_rate_limit(row_get(row, "user_rate_limit")?)) + Ok(snapshot + .with_user_rate_limit(row_get(row, "user_rate_limit")?) + .with_daily_usage_limits(None, row_get(row, "api_key_daily_usage_limit_usd")?)) } fn map_auth_api_key_export_row( @@ -1597,6 +1630,9 @@ fn map_auth_api_key_export_row( row_get(row, "is_standalone")?, ) .and_then(|record| record.with_ip_rules(row_get(row, "ip_rules")?)) + .map(|record| { + record.with_daily_usage_limit(row_get(row, "daily_usage_limit_usd").ok().flatten()) + }) .map(|record| record.with_feature_settings(feature_settings)) .and_then(|record| { record.with_activity_timestamps( @@ -1621,40 +1657,41 @@ mod tests { assert!(CREATE_USER_API_KEY_SQL .contains("expires_at,\n auto_delete_on_expiry,\n is_locked,\n is_standalone,")); assert!( - CREATE_USER_API_KEY_SQL.contains("$13,\n $14,\n $15,\n FALSE,\n FALSE,\n $16,") + CREATE_USER_API_KEY_SQL.contains("$14,\n $15,\n $16,\n FALSE,\n FALSE,\n $17,") ); assert!(CREATE_STANDALONE_API_KEY_SQL .contains("expires_at,\n auto_delete_on_expiry,\n is_locked,\n is_standalone,")); assert!(CREATE_STANDALONE_API_KEY_SQL - .contains("$13,\n $14,\n $15,\n FALSE,\n TRUE,\n $16,")); + .contains("$14,\n $15,\n $16,\n FALSE,\n TRUE,\n $17,")); } #[test] fn update_standalone_api_key_basic_sql_casts_json_case_values() { assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL - .contains("concurrent_limit = CASE WHEN $5 THEN $6 ELSE concurrent_limit END")); - assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL - .contains("allowed_providers = CASE WHEN $7 THEN $8::json ELSE allowed_providers END")); + .contains("concurrent_limit = CASE WHEN $7 THEN $8 ELSE concurrent_limit END")); + assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL.contains( + "allowed_providers = CASE WHEN $9 THEN $10::json ELSE allowed_providers END" + )); assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL.contains( - "allowed_api_formats = CASE WHEN $9 THEN $10::json ELSE allowed_api_formats END" + "allowed_api_formats = CASE WHEN $11 THEN $12::json ELSE allowed_api_formats END" )); assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL - .contains("allowed_models = CASE WHEN $11 THEN $12::json ELSE allowed_models END")); + .contains("allowed_models = CASE WHEN $13 THEN $14::json ELSE allowed_models END")); assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL - .contains("ip_rules = CASE WHEN $13 THEN $14::jsonb ELSE ip_rules END")); + .contains("ip_rules = CASE WHEN $15 THEN $16::jsonb ELSE ip_rules END")); assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL .contains("rate_limit = CASE WHEN $3 THEN $4 ELSE rate_limit END")); assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL - .contains("expires_at = CASE WHEN $15 THEN $16::timestamptz ELSE expires_at END")); + .contains("expires_at = CASE WHEN $17 THEN $18::timestamptz ELSE expires_at END")); assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL.contains( - "auto_delete_on_expiry = CASE WHEN $17 THEN $18 ELSE auto_delete_on_expiry END" + "auto_delete_on_expiry = CASE WHEN $19 THEN $20 ELSE auto_delete_on_expiry END" )); } #[test] fn update_user_api_key_basic_sql_casts_ip_rules_as_jsonb() { assert!(UPDATE_USER_API_KEY_BASIC_SQL - .contains("ip_rules = CASE WHEN $6 THEN $7::jsonb ELSE ip_rules END")); + .contains("ip_rules = CASE WHEN $8 THEN $9::jsonb ELSE ip_rules END")); } #[tokio::test] diff --git a/crates/aether-data/adapters/postgres/src/usage/mod.rs b/crates/aether-data/adapters/postgres/src/usage/mod.rs index a95db63fe..b011c82e7 100644 --- a/crates/aether-data/adapters/postgres/src/usage/mod.rs +++ b/crates/aether-data/adapters/postgres/src/usage/mod.rs @@ -3,18 +3,19 @@ use aether_data_contracts::repository::usage::{ ProxyNodeCounterDelta, StoredUsageAuditAggregation, StoredUsageAuditSummary, StoredUsageBreakdownSummaryRow, StoredUsageCacheAffinityHitSummary, StoredUsageCacheAffinityIntervalRow, StoredUsageCacheHitSummary, StoredUsageCostSavingsSummary, - StoredUsageDashboardDailyBreakdownRow, StoredUsageDashboardProviderCount, - StoredUsageDashboardStatsSummary, StoredUsageDashboardSummary, StoredUsageErrorDistributionRow, - StoredUsageLeaderboardSummary, StoredUsagePerformancePercentilesRow, - StoredUsageProviderPerformance, StoredUsageProviderPerformanceProviderRow, - StoredUsageProviderPerformanceSummary, StoredUsageProviderPerformanceTimelineRow, - StoredUsageSettledCostSummary, StoredUsageTimeSeriesBucket, StoredUsageUserTotals, - UsageAuditAggregationGroupBy, UsageAuditAggregationQuery, UsageAuditKeywordSearchQuery, - UsageAuditSummaryQuery, UsageBodyCaptureState, UsageBodyField, UsageBreakdownGroupBy, - UsageBreakdownSummaryQuery, UsageCacheAffinityHitSummaryQuery, - UsageCacheAffinityIntervalGroupBy, UsageCacheAffinityIntervalQuery, UsageCacheHitSummaryQuery, - UsageCleanupExecutionMode, UsageCleanupSummary, UsageCleanupTargets, UsageCleanupWindow, - UsageCostSavingsSummaryQuery, UsageDashboardDailyBreakdownQuery, + StoredUsageDailyActualCostRollup, StoredUsageDashboardDailyBreakdownRow, + StoredUsageDashboardProviderCount, StoredUsageDashboardStatsSummary, + StoredUsageDashboardSummary, StoredUsageErrorDistributionRow, StoredUsageLeaderboardSummary, + StoredUsagePerformancePercentilesRow, StoredUsageProviderPerformance, + StoredUsageProviderPerformanceProviderRow, StoredUsageProviderPerformanceSummary, + StoredUsageProviderPerformanceTimelineRow, StoredUsageSettledCostSummary, + StoredUsageTimeSeriesBucket, StoredUsageUserTotals, UsageAuditAggregationGroupBy, + UsageAuditAggregationQuery, UsageAuditKeywordSearchQuery, UsageAuditSummaryQuery, + UsageBodyCaptureState, UsageBodyField, UsageBreakdownGroupBy, UsageBreakdownSummaryQuery, + UsageCacheAffinityHitSummaryQuery, UsageCacheAffinityIntervalGroupBy, + UsageCacheAffinityIntervalQuery, UsageCacheHitSummaryQuery, UsageCleanupExecutionMode, + UsageCleanupSummary, UsageCleanupTargets, UsageCleanupWindow, UsageCostSavingsSummaryQuery, + UsageDailyActualCostRollupQuery, UsageDashboardDailyBreakdownQuery, UsageDashboardProviderCountsQuery, UsageDashboardSummaryQuery, UsageErrorDistributionQuery, UsageLeaderboardGroupBy, UsageLeaderboardQuery, UsageMonitoringErrorCountQuery, UsageMonitoringErrorListQuery, UsagePerformancePercentilesQuery, UsageProviderPerformanceQuery, @@ -10251,6 +10252,54 @@ impl UsageReadRepository for SqlxUsageReadRepository { Self::summarize_usage_settled_cost(self, query).await } + async fn summarize_usage_daily_actual_cost_rollups( + &self, + query: &UsageDailyActualCostRollupQuery, + ) -> Result, DataLayerError> { + if query.finalized_from_unix_secs >= query.finalized_until_unix_secs { + return Ok(Vec::new()); + } + let rows = sqlx::query( + r#" +SELECT u.user_id, + u.api_key_id, + CASE + WHEN (u.request_metadata->>'api_key_is_standalone') IN ('true', 'false') + THEN (u.request_metadata->>'api_key_is_standalone')::boolean + ELSE FALSE + END AS api_key_is_standalone, + SUM(CAST(u.actual_total_cost_usd AS DOUBLE PRECISION)) AS actual_total_cost_usd +FROM "usage" AS u +WHERE u.status = 'completed' + AND u.finalized_at >= TO_TIMESTAMP($1::double precision) + AND u.finalized_at < TO_TIMESTAMP($2::double precision) + AND u.api_key_id IS NOT NULL + AND BTRIM(u.api_key_id) <> '' + AND CAST(u.actual_total_cost_usd AS DOUBLE PRECISION) > 0 +GROUP BY u.user_id, u.api_key_id, 3 +"#, + ) + .bind(query.finalized_from_unix_secs as f64) + .bind(query.finalized_until_unix_secs as f64) + .fetch_all(&self.pool) + .await + .map_postgres_err()?; + rows.into_iter() + .map(|row| { + Ok(StoredUsageDailyActualCostRollup { + user_id: row.try_get("user_id").map_postgres_err()?, + api_key_id: row.try_get("api_key_id").map_postgres_err()?, + api_key_is_standalone: row + .try_get("api_key_is_standalone") + .map_postgres_err()?, + actual_total_cost_usd: row + .try_get("actual_total_cost_usd") + .map_postgres_err()?, + }) + }) + .collect() + } + async fn summarize_usage_cache_affinity_hit_summary( &self, query: &UsageCacheAffinityHitSummaryQuery, diff --git a/crates/aether-data/adapters/postgres/src/users.rs b/crates/aether-data/adapters/postgres/src/users.rs index e1d5ed2fa..71bca32c4 100644 --- a/crates/aether-data/adapters/postgres/src/users.rs +++ b/crates/aether-data/adapters/postgres/src/users.rs @@ -610,6 +610,8 @@ SELECT allowed_models_mode, rate_limit, rate_limit_mode, + daily_usage_limit_usd, + daily_usage_limit_mode, created_at, updated_at FROM user_groups @@ -747,9 +749,9 @@ INSERT INTO user_groups ( allowed_providers, allowed_providers_mode, allowed_api_formats, allowed_api_formats_mode, allowed_models, allowed_models_mode, - rate_limit, rate_limit_mode + rate_limit, rate_limit_mode, daily_usage_limit_usd, daily_usage_limit_mode ) -VALUES ($1, $2, $3, $4, $5, $6::json, $7, $8::json, $9, $10::json, $11, $12, $13) +VALUES ($1, $2, $3, $4, $5, $6::json, $7, $8::json, $9, $10::json, $11, $12, $13, $14, $15) "#, ) .bind(&id) @@ -765,6 +767,8 @@ VALUES ($1, $2, $3, $4, $5, $6::json, $7, $8::json, $9, $10::json, $11, $12, $13 .bind(record.allowed_models_mode) .bind(record.rate_limit) .bind(record.rate_limit_mode) + .bind(record.daily_usage_limit_usd) + .bind(record.daily_usage_limit_mode) .execute(&self.pool) .await; match result { @@ -798,6 +802,8 @@ SET name = $2, allowed_models_mode = $11, rate_limit = $12, rate_limit_mode = $13, + daily_usage_limit_usd = $14, + daily_usage_limit_mode = $15, updated_at = now() WHERE id = $1 "#, @@ -815,6 +821,8 @@ WHERE id = $1 .bind(record.allowed_models_mode) .bind(record.rate_limit) .bind(record.rate_limit_mode) + .bind(record.daily_usage_limit_usd) + .bind(record.daily_usage_limit_mode) .execute(&self.pool) .await; match result { @@ -2285,6 +2293,12 @@ fn map_user_group_row(row: &sqlx::postgres::PgRow) -> Result>, ip_rules: Option>, rate_limit: Option, + daily_usage_limit_usd: Option, concurrent_limit: Option, force_capabilities: Option, is_active: bool, @@ -427,6 +432,7 @@ WHERE id = ? allowed_models: record.allowed_models, ip_rules: record.ip_rules, rate_limit: Some(record.rate_limit), + daily_usage_limit_usd: record.daily_usage_limit_usd, concurrent_limit: record.concurrent_limit, force_capabilities: record.force_capabilities, is_active: record.is_active, @@ -455,6 +461,7 @@ WHERE id = ? allowed_models: record.allowed_models, ip_rules: record.ip_rules, rate_limit: record.rate_limit, + daily_usage_limit_usd: record.daily_usage_limit_usd, concurrent_limit: record.concurrent_limit, force_capabilities: record.force_capabilities, is_active: record.is_active, @@ -478,6 +485,7 @@ WHERE id = ? UPDATE api_keys SET name = COALESCE(?, name), rate_limit = COALESCE(?, rate_limit), + daily_usage_limit_usd = CASE WHEN ? THEN ? ELSE daily_usage_limit_usd END, concurrent_limit = COALESCE(?, concurrent_limit), ip_rules = CASE WHEN ? THEN ? ELSE ip_rules END, updated_at = ? @@ -488,6 +496,8 @@ WHERE id = ? ) .bind(record.name.as_deref()) .bind(record.rate_limit) + .bind(record.daily_usage_limit_present) + .bind(record.daily_usage_limit_usd) .bind(record.concurrent_limit) .bind(record.ip_rules.is_some()) .bind(json_string_from_nested_string_list( @@ -513,6 +523,7 @@ WHERE id = ? UPDATE api_keys SET name = COALESCE(?, name), rate_limit = CASE WHEN ? THEN ? ELSE rate_limit END, + daily_usage_limit_usd = CASE WHEN ? THEN ? ELSE daily_usage_limit_usd END, concurrent_limit = CASE WHEN ? THEN ? ELSE concurrent_limit END, allowed_providers = CASE WHEN ? THEN ? ELSE allowed_providers END, allowed_api_formats = CASE WHEN ? THEN ? ELSE allowed_api_formats END, @@ -528,6 +539,8 @@ WHERE id = ? .bind(record.name.as_deref()) .bind(record.rate_limit_present) .bind(record.rate_limit) + .bind(record.daily_usage_limit_present) + .bind(record.daily_usage_limit_usd) .bind(record.concurrent_limit_present) .bind(record.concurrent_limit) .bind(record.allowed_providers.is_some()) @@ -977,7 +990,12 @@ fn map_auth_api_key_snapshot_row( row.try_get("api_key_ip_rules").map_sql_err()?, "api_keys.ip_rules", )?)?; - Ok(snapshot.with_user_rate_limit(row.try_get("user_rate_limit").map_sql_err()?)) + Ok(snapshot + .with_user_rate_limit(row.try_get("user_rate_limit").map_sql_err()?) + .with_daily_usage_limits( + None, + row.try_get("api_key_daily_usage_limit_usd").map_sql_err()?, + )) } fn map_auth_api_key_export_row( @@ -1025,6 +1043,9 @@ fn map_auth_api_key_export_row( "api_keys.ip_rules", )?) }) + .map(|record| { + record.with_daily_usage_limit(row.try_get("daily_usage_limit_usd").ok().flatten()) + }) .map(|record| record.with_feature_settings(feature_settings)) .and_then(|record| { record.with_activity_timestamps( @@ -1143,6 +1164,7 @@ mod tests { allowed_models: Some(vec!["gpt-4.1".to_string()]), ip_rules: Some(vec!["203.0.113.10".to_string()]), rate_limit: 100, + daily_usage_limit_usd: Some(12.5), concurrent_limit: Some(5), force_capabilities: Some(json!({"cache": true})), is_active: true, @@ -1157,6 +1179,7 @@ mod tests { .expect("user key should reload"); assert_eq!(user_key.allowed_models, Some(vec!["gpt-4.1".to_string()])); assert_eq!(user_key.total_tokens, 42); + assert_eq!(user_key.daily_usage_limit_usd, Some(12.5)); let updated_user_key = repository .update_user_api_key_basic(UpdateUserApiKeyBasicRecord { @@ -1164,6 +1187,8 @@ mod tests { api_key_id: "key-created-user".to_string(), name: Some("Updated User".to_string()), rate_limit: Some(150), + daily_usage_limit_present: true, + daily_usage_limit_usd: Some(8.0), concurrent_limit: Some(6), ip_rules: Some(Some(vec!["10.0.0.0/24".to_string()])), }) @@ -1171,8 +1196,25 @@ mod tests { .expect("user key should update") .expect("user key should reload"); assert_eq!(updated_user_key.name, Some("Updated User".to_string())); + assert_eq!(updated_user_key.daily_usage_limit_usd, Some(8.0)); assert_eq!(updated_user_key.concurrent_limit, Some(6)); + let cleared_user_key = repository + .update_user_api_key_basic(UpdateUserApiKeyBasicRecord { + user_id: "user-1".to_string(), + api_key_id: "key-created-user".to_string(), + name: None, + rate_limit: None, + daily_usage_limit_present: true, + daily_usage_limit_usd: None, + concurrent_limit: None, + ip_rules: None, + }) + .await + .expect("user key daily limit should clear") + .expect("user key should reload"); + assert_eq!(cleared_user_key.daily_usage_limit_usd, None); + assert!(repository .set_user_api_key_locked("user-1", "key-created-user", true) .await @@ -1243,6 +1285,7 @@ mod tests { allowed_models: None, ip_rules: None, rate_limit: None, + daily_usage_limit_usd: None, concurrent_limit: Some(2), force_capabilities: None, is_active: true, @@ -1263,6 +1306,8 @@ mod tests { name: Some("Updated Standalone".to_string()), rate_limit_present: true, rate_limit: Some(20), + daily_usage_limit_present: true, + daily_usage_limit_usd: Some(4.0), concurrent_limit_present: true, concurrent_limit: None, allowed_providers: Some(None), @@ -1278,6 +1323,7 @@ mod tests { .expect("standalone key should update") .expect("standalone key should reload"); assert_eq!(standalone.name, Some("Updated Standalone".to_string())); + assert_eq!(standalone.daily_usage_limit_usd, Some(4.0)); assert_eq!(standalone.allowed_providers, None); assert_eq!( standalone.allowed_api_formats, diff --git a/crates/aether-data/adapters/sqlite/src/usage.rs b/crates/aether-data/adapters/sqlite/src/usage.rs index 6030a894d..3a8540e69 100644 --- a/crates/aether-data/adapters/sqlite/src/usage.rs +++ b/crates/aether-data/adapters/sqlite/src/usage.rs @@ -17,9 +17,10 @@ use aether_data_contracts::repository::usage::{ StoredProviderApiKeyWindowUsageSummary, StoredProviderUsageSummary, StoredRequestUsageAudit, StoredUsageAuditAggregation, StoredUsageAuditSummary, StoredUsageBreakdownSummaryRow, StoredUsageCacheAffinityHitSummary, StoredUsageCacheAffinityIntervalRow, - StoredUsageCacheHitSummary, StoredUsageCostSavingsSummary, StoredUsageDailySummary, - StoredUsageDashboardDailyBreakdownRow, StoredUsageDashboardProviderCount, - StoredUsageDashboardSummary, StoredUsageErrorDistributionRow, StoredUsageLeaderboardSummary, + StoredUsageCacheHitSummary, StoredUsageCostSavingsSummary, StoredUsageDailyActualCostRollup, + StoredUsageDailySummary, StoredUsageDashboardDailyBreakdownRow, + StoredUsageDashboardProviderCount, StoredUsageDashboardSummary, + StoredUsageErrorDistributionRow, StoredUsageLeaderboardSummary, StoredUsagePerformancePercentilesRow, StoredUsageProviderPerformance, StoredUsageProviderPerformanceProviderRow, StoredUsageProviderPerformanceSummary, StoredUsageProviderPerformanceTimelineRow, StoredUsageSettledCostSummary, @@ -29,12 +30,13 @@ use aether_data_contracts::repository::usage::{ UsageCacheAffinityHitSummaryQuery, UsageCacheAffinityIntervalGroupBy, UsageCacheAffinityIntervalQuery, UsageCacheHitSummaryQuery, UsageCleanupExecutionMode, UsageCleanupPreviewCounts, UsageCleanupSummary, UsageCleanupTargets, UsageCleanupWindow, - UsageCostSavingsSummaryQuery, UsageDailyHeatmapQuery, UsageDashboardDailyBreakdownQuery, - UsageDashboardProviderCountsQuery, UsageDashboardSummaryQuery, UsageErrorDistributionQuery, - UsageLeaderboardGroupBy, UsageLeaderboardQuery, UsageMonitoringErrorCountQuery, - UsageMonitoringErrorListQuery, UsagePerformancePercentilesQuery, UsageProviderPerformanceQuery, - UsageReadRepository, UsageSettledCostSummaryQuery, UsageTimeSeriesGranularity, - UsageTimeSeriesQuery, UsageWriteRepository, + UsageCostSavingsSummaryQuery, UsageDailyActualCostRollupQuery, UsageDailyHeatmapQuery, + UsageDashboardDailyBreakdownQuery, UsageDashboardProviderCountsQuery, + UsageDashboardSummaryQuery, UsageErrorDistributionQuery, UsageLeaderboardGroupBy, + UsageLeaderboardQuery, UsageMonitoringErrorCountQuery, UsageMonitoringErrorListQuery, + UsagePerformancePercentilesQuery, UsageProviderPerformanceQuery, UsageReadRepository, + UsageSettledCostSummaryQuery, UsageTimeSeriesGranularity, UsageTimeSeriesQuery, + UsageWriteRepository, }; use aether_data_contracts::DataLayerError; @@ -2531,6 +2533,54 @@ FROM "usage" }) } + async fn summarize_usage_daily_actual_cost_rollups( + &self, + query: &UsageDailyActualCostRollupQuery, + ) -> Result, DataLayerError> { + if query.finalized_from_unix_secs >= query.finalized_until_unix_secs { + return Ok(Vec::new()); + } + let rows = sqlx::query( + r#" +SELECT u.user_id, + u.api_key_id, + CASE + WHEN LOWER(COALESCE(CAST(json_extract( + u.request_metadata, '$.api_key_is_standalone' + ) AS TEXT), 'false')) IN ('true', '1') + THEN 1 ELSE 0 + END AS api_key_is_standalone, + SUM(CAST(u.actual_total_cost_usd AS REAL)) AS actual_total_cost_usd +FROM "usage" AS u +WHERE u.status = 'completed' + AND u.finalized_at >= ? + AND u.finalized_at < ? + AND u.api_key_id IS NOT NULL + AND TRIM(u.api_key_id) <> '' + AND CAST(u.actual_total_cost_usd AS REAL) > 0 +GROUP BY u.user_id, u.api_key_id, 3 +"#, + ) + .bind(query.finalized_from_unix_secs as i64) + .bind(query.finalized_until_unix_secs as i64) + .fetch_all(&self.pool) + .await + .map_sql_err()?; + rows.into_iter() + .map(|row| { + Ok(StoredUsageDailyActualCostRollup { + user_id: row.try_get("user_id").map_sql_err()?, + api_key_id: row.try_get("api_key_id").map_sql_err()?, + api_key_is_standalone: row + .try_get::("api_key_is_standalone") + .map_sql_err()? + != 0, + actual_total_cost_usd: sqlite_real(&row, "actual_total_cost_usd")?, + }) + }) + .collect() + } + async fn summarize_usage_cache_affinity_hit_summary( &self, query: &UsageCacheAffinityHitSummaryQuery, diff --git a/crates/aether-data/adapters/sqlite/src/usage/tests.rs b/crates/aether-data/adapters/sqlite/src/usage/tests.rs index b0782b8b6..b704851bb 100644 --- a/crates/aether-data/adapters/sqlite/src/usage/tests.rs +++ b/crates/aether-data/adapters/sqlite/src/usage/tests.rs @@ -4,9 +4,10 @@ use aether_data_contracts::repository::usage::{ ProviderApiKeyWindowUsageRequest, UpsertUsageRecord, UsageAuditAggregationGroupBy, UsageAuditAggregationQuery, UsageAuditListQuery, UsageAuditSummaryQuery, UsageBodyCaptureState, UsageBreakdownGroupBy, UsageBreakdownSummaryQuery, UsageCleanupExecutionMode, - UsageCleanupTargets, UsageCleanupWindow, UsageDailyHeatmapQuery, - UsageDashboardDailyBreakdownQuery, UsageDashboardSummaryQuery, UsageProviderPerformanceQuery, - UsageReadRepository, UsageTimeSeriesGranularity, UsageWriteRepository, + UsageCleanupTargets, UsageCleanupWindow, UsageDailyActualCostRollupQuery, + UsageDailyHeatmapQuery, UsageDashboardDailyBreakdownQuery, UsageDashboardSummaryQuery, + UsageProviderPerformanceQuery, UsageReadRepository, UsageTimeSeriesGranularity, + UsageWriteRepository, }; use chrono::{DateTime, Utc}; @@ -109,6 +110,54 @@ async fn sqlite_provider_performance_can_skip_timeline() { assert!(without_timeline.timeline.is_empty()); } +#[tokio::test] +async fn sqlite_daily_actual_cost_rollups_group_normal_and_standalone_keys() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + run_migrations(&pool) + .await + .expect("sqlite migrations should run"); + seed_stats_targets(&pool).await; + + let writer = SqliteUsageWriteRepository::new(pool.clone()); + writer + .upsert(sample_usage("daily-normal", "completed", "settled", 100)) + .await + .expect("normal daily usage should seed"); + let mut standalone = sample_usage("daily-standalone", "completed", "settled", 100); + standalone.api_key_id = Some("standalone-key".to_string()); + standalone.actual_total_cost_usd = Some(0.6); + standalone.request_metadata = Some(serde_json::json!({ + "api_key_is_standalone": true, + })); + writer + .upsert(standalone) + .await + .expect("standalone daily usage should seed"); + + let rollups = SqliteUsageReadRepository::new(pool) + .summarize_usage_daily_actual_cost_rollups(&UsageDailyActualCostRollupQuery { + finalized_from_unix_secs: 50, + finalized_until_unix_secs: 150, + }) + .await + .expect("daily usage rollups should load"); + assert_eq!(rollups.len(), 2); + assert!(rollups.iter().any(|rollup| { + rollup.api_key_id == "api-key-1" + && !rollup.api_key_is_standalone + && rollup.actual_total_cost_usd == 0.4 + })); + assert!(rollups.iter().any(|rollup| { + rollup.api_key_id == "standalone-key" + && rollup.api_key_is_standalone + && rollup.actual_total_cost_usd == 0.6 + })); +} + #[tokio::test] async fn sqlite_usage_write_repository_upserts_and_flushes_counter_deltas() { let pool = sqlx::sqlite::SqlitePoolOptions::new() diff --git a/crates/aether-data/adapters/sqlite/src/users.rs b/crates/aether-data/adapters/sqlite/src/users.rs index 00fc6b42f..38fb26493 100644 --- a/crates/aether-data/adapters/sqlite/src/users.rs +++ b/crates/aether-data/adapters/sqlite/src/users.rs @@ -159,6 +159,8 @@ SELECT allowed_models_mode, rate_limit, rate_limit_mode, + daily_usage_limit_usd, + daily_usage_limit_mode, created_at, updated_at FROM user_groups @@ -480,9 +482,10 @@ INSERT INTO user_groups ( allowed_providers, allowed_providers_mode, allowed_api_formats, allowed_api_formats_mode, allowed_models, allowed_models_mode, - rate_limit, rate_limit_mode, created_at, updated_at + rate_limit, rate_limit_mode, daily_usage_limit_usd, daily_usage_limit_mode, + created_at, updated_at ) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) "#, ) .bind(&id) @@ -502,6 +505,8 @@ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) .bind(record.allowed_models_mode) .bind(record.rate_limit) .bind(record.rate_limit_mode) + .bind(record.daily_usage_limit_usd) + .bind(record.daily_usage_limit_mode) .bind(now) .bind(now) .execute(&self.pool) @@ -538,6 +543,8 @@ SET name = ?, allowed_models_mode = ?, rate_limit = ?, rate_limit_mode = ?, + daily_usage_limit_usd = ?, + daily_usage_limit_mode = ?, updated_at = ? WHERE id = ? "#, @@ -558,6 +565,8 @@ WHERE id = ? .bind(record.allowed_models_mode) .bind(record.rate_limit) .bind(record.rate_limit_mode) + .bind(record.daily_usage_limit_usd) + .bind(record.daily_usage_limit_mode) .bind(now) .bind(group_id) .execute(&self.pool) @@ -2040,6 +2049,12 @@ fn map_user_group_row(row: &SqliteRow) -> Result Result { diff --git a/crates/aether-data/contracts/src/repository/auth.rs b/crates/aether-data/contracts/src/repository/auth.rs index df5b2aff1..e1cdc317e 100644 --- a/crates/aether-data/contracts/src/repository/auth.rs +++ b/crates/aether-data/contracts/src/repository/auth.rs @@ -1,6 +1,6 @@ use async_trait::async_trait; -#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] pub struct StoredAuthApiKeySnapshot { pub user_id: String, pub username: String, @@ -10,6 +10,7 @@ pub struct StoredAuthApiKeySnapshot { pub user_is_active: bool, pub user_is_deleted: bool, pub user_rate_limit: Option, + pub user_daily_usage_limit_usd: Option, pub user_allowed_providers: Option>, pub user_allowed_api_formats: Option>, pub user_allowed_models: Option>, @@ -19,6 +20,7 @@ pub struct StoredAuthApiKeySnapshot { pub api_key_is_locked: bool, pub api_key_is_standalone: bool, pub api_key_rate_limit: Option, + pub api_key_daily_usage_limit_usd: Option, pub api_key_concurrent_limit: Option, pub api_key_expires_at_unix_secs: Option, pub api_key_allowed_providers: Option>, @@ -61,6 +63,7 @@ impl StoredAuthApiKeySnapshot { user_is_active, user_is_deleted, user_rate_limit: None, + user_daily_usage_limit_usd: None, user_allowed_providers: parse_string_list( user_allowed_providers, "users.allowed_providers", @@ -76,6 +79,7 @@ impl StoredAuthApiKeySnapshot { api_key_is_locked, api_key_is_standalone, api_key_rate_limit, + api_key_daily_usage_limit_usd: None, api_key_concurrent_limit, api_key_expires_at_unix_secs: api_key_expires_at_unix_secs .map(|value| { @@ -132,9 +136,19 @@ impl StoredAuthApiKeySnapshot { self.user_rate_limit = user_rate_limit; self } + + pub fn with_daily_usage_limits( + mut self, + user_daily_usage_limit_usd: Option, + api_key_daily_usage_limit_usd: Option, + ) -> Self { + self.user_daily_usage_limit_usd = user_daily_usage_limit_usd; + self.api_key_daily_usage_limit_usd = api_key_daily_usage_limit_usd; + self + } } -#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] pub struct ResolvedAuthApiKeySnapshot { pub user_id: String, pub username: String, @@ -144,6 +158,7 @@ pub struct ResolvedAuthApiKeySnapshot { pub user_is_active: bool, pub user_is_deleted: bool, pub user_rate_limit: Option, + pub user_daily_usage_limit_usd: Option, pub user_allowed_providers: Option>, pub user_allowed_api_formats: Option>, pub user_allowed_models: Option>, @@ -153,6 +168,7 @@ pub struct ResolvedAuthApiKeySnapshot { pub api_key_is_locked: bool, pub api_key_is_standalone: bool, pub api_key_rate_limit: Option, + pub api_key_daily_usage_limit_usd: Option, pub api_key_concurrent_limit: Option, pub api_key_expires_at_unix_secs: Option, pub api_key_allowed_providers: Option>, @@ -174,6 +190,7 @@ impl ResolvedAuthApiKeySnapshot { user_is_active: snapshot.user_is_active, user_is_deleted: snapshot.user_is_deleted, user_rate_limit: snapshot.user_rate_limit, + user_daily_usage_limit_usd: snapshot.user_daily_usage_limit_usd, user_allowed_providers: snapshot.user_allowed_providers, user_allowed_api_formats: snapshot.user_allowed_api_formats, user_allowed_models: snapshot.user_allowed_models, @@ -183,6 +200,7 @@ impl ResolvedAuthApiKeySnapshot { api_key_is_locked: snapshot.api_key_is_locked, api_key_is_standalone: snapshot.api_key_is_standalone, api_key_rate_limit: snapshot.api_key_rate_limit, + api_key_daily_usage_limit_usd: snapshot.api_key_daily_usage_limit_usd, api_key_concurrent_limit: snapshot.api_key_concurrent_limit, api_key_expires_at_unix_secs: snapshot.api_key_expires_at_unix_secs, api_key_allowed_providers: snapshot.api_key_allowed_providers, @@ -231,11 +249,13 @@ impl ResolvedAuthApiKeySnapshot { allowed_api_formats: Option>, allowed_models: Option>, rate_limit: Option, + daily_usage_limit_usd: Option, ) { self.user_allowed_providers = allowed_providers; self.user_allowed_api_formats = allowed_api_formats; self.user_allowed_models = allowed_models; self.user_rate_limit = rate_limit; + self.user_daily_usage_limit_usd = daily_usage_limit_usd; self.constrain_non_standalone_api_key_policy_to_user_policy(); } @@ -362,6 +382,7 @@ pub struct StoredAuthApiKeyExportRecord { pub allowed_models: Option>, pub ip_rules: Option>, pub rate_limit: Option, + pub daily_usage_limit_usd: Option, pub concurrent_limit: Option, pub force_capabilities: Option, pub feature_settings: Option, @@ -434,6 +455,7 @@ impl StoredAuthApiKeyExportRecord { allowed_models: parse_string_list(allowed_models, "api_keys.allowed_models")?, ip_rules: None, rate_limit, + daily_usage_limit_usd: None, concurrent_limit, force_capabilities, feature_settings: None, @@ -475,6 +497,11 @@ impl StoredAuthApiKeyExportRecord { self } + pub fn with_daily_usage_limit(mut self, daily_usage_limit_usd: Option) -> Self { + self.daily_usage_limit_usd = daily_usage_limit_usd; + self + } + pub fn with_ip_rules( mut self, ip_rules: Option, @@ -509,6 +536,7 @@ pub struct CreateUserApiKeyRecord { pub allowed_models: Option>, pub ip_rules: Option>, pub rate_limit: i32, + pub daily_usage_limit_usd: Option, pub concurrent_limit: Option, pub force_capabilities: Option, pub is_active: bool, @@ -519,12 +547,14 @@ pub struct CreateUserApiKeyRecord { pub total_cost_usd: f64, } -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq)] pub struct UpdateUserApiKeyBasicRecord { pub user_id: String, pub api_key_id: String, pub name: Option, pub rate_limit: Option, + pub daily_usage_limit_present: bool, + pub daily_usage_limit_usd: Option, pub concurrent_limit: Option, pub ip_rules: Option>>, } @@ -541,6 +571,7 @@ pub struct CreateStandaloneApiKeyRecord { pub allowed_models: Option>, pub ip_rules: Option>, pub rate_limit: Option, + pub daily_usage_limit_usd: Option, pub concurrent_limit: Option, pub force_capabilities: Option, pub is_active: bool, @@ -551,12 +582,14 @@ pub struct CreateStandaloneApiKeyRecord { pub total_cost_usd: f64, } -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq)] pub struct UpdateStandaloneApiKeyBasicRecord { pub api_key_id: String, pub name: Option, pub rate_limit_present: bool, pub rate_limit: Option, + pub daily_usage_limit_present: bool, + pub daily_usage_limit_usd: Option, pub concurrent_limit_present: bool, pub concurrent_limit: Option, pub allowed_providers: Option>>, @@ -1062,6 +1095,7 @@ mod tests { Some(vec!["openai:chat".to_string()]), Some(vec!["gpt-5".to_string()]), Some(60), + Some(12.5), ); assert_eq!( @@ -1073,6 +1107,7 @@ mod tests { Some(&["openai:chat".to_string()][..]) ); assert_eq!(resolved.effective_allowed_models(), Some(&[][..])); + assert_eq!(resolved.user_daily_usage_limit_usd, Some(12.5)); } #[test] diff --git a/crates/aether-data/contracts/src/repository/usage/mod.rs b/crates/aether-data/contracts/src/repository/usage/mod.rs index 00c784f03..3a497b676 100644 --- a/crates/aether-data/contracts/src/repository/usage/mod.rs +++ b/crates/aether-data/contracts/src/repository/usage/mod.rs @@ -14,25 +14,26 @@ pub use types::{ StoredRequestUsageAudit, StoredUsageAuditAggregation, StoredUsageAuditSummary, StoredUsageBreakdownSummaryRow, StoredUsageCacheAffinityHitSummary, StoredUsageCacheAffinityIntervalRow, StoredUsageCacheHitSummary, StoredUsageCostSavingsSummary, - StoredUsageDailySummary, StoredUsageDashboardDailyBreakdownRow, - StoredUsageDashboardProviderCount, StoredUsageDashboardStatsSummary, - StoredUsageDashboardSummary, StoredUsageErrorDistributionRow, StoredUsageLeaderboardSummary, - StoredUsagePerformancePercentilesRow, StoredUsageProviderPerformance, - StoredUsageProviderPerformanceProviderRow, StoredUsageProviderPerformanceSummary, - StoredUsageProviderPerformanceTimelineRow, StoredUsageSettledCostSummary, - StoredUsageTimeSeriesBucket, StoredUsageUserTotals, UpsertUsageRecord, - UsageAuditAggregationGroupBy, UsageAuditAggregationQuery, UsageAuditKeywordSearchQuery, - UsageAuditListQuery, UsageAuditSummaryQuery, UsageBodyCaptureResult, UsageBodyCaptureState, - UsageBodyCaptureStorage, UsageBodyField, UsageBreakdownGroupBy, UsageBreakdownSummaryQuery, - UsageCacheAffinityHitSummaryQuery, UsageCacheAffinityIntervalGroupBy, - UsageCacheAffinityIntervalQuery, UsageCacheHitSummaryQuery, UsageCleanupExecutionMode, - UsageCleanupPreviewCounts, UsageCleanupSummary, UsageCleanupTargets, UsageCleanupWindow, - UsageCostSavingsSummaryQuery, UsageCounterFlushSummary, UsageCounterHealthSnapshot, - UsageCounterPendingHealthSnapshot, UsageDailyHeatmapQuery, UsageDashboardDailyBreakdownQuery, - UsageDashboardProviderCountsQuery, UsageDashboardSummaryQuery, UsageErrorDistributionQuery, - UsageLeaderboardGroupBy, UsageLeaderboardQuery, UsageMonitoringErrorCountQuery, - UsageMonitoringErrorListQuery, UsagePerformancePercentilesQuery, UsageProviderPerformanceQuery, - UsageReadRepository, UsageRepository, UsageSettledCostSummaryQuery, UsageTimeSeriesGranularity, + StoredUsageDailyActualCostRollup, StoredUsageDailySummary, + StoredUsageDashboardDailyBreakdownRow, StoredUsageDashboardProviderCount, + StoredUsageDashboardStatsSummary, StoredUsageDashboardSummary, StoredUsageErrorDistributionRow, + StoredUsageLeaderboardSummary, StoredUsagePerformancePercentilesRow, + StoredUsageProviderPerformance, StoredUsageProviderPerformanceProviderRow, + StoredUsageProviderPerformanceSummary, StoredUsageProviderPerformanceTimelineRow, + StoredUsageSettledCostSummary, StoredUsageTimeSeriesBucket, StoredUsageUserTotals, + UpsertUsageRecord, UsageAuditAggregationGroupBy, UsageAuditAggregationQuery, + UsageAuditKeywordSearchQuery, UsageAuditListQuery, UsageAuditSummaryQuery, + UsageBodyCaptureResult, UsageBodyCaptureState, UsageBodyCaptureStorage, UsageBodyField, + UsageBreakdownGroupBy, UsageBreakdownSummaryQuery, UsageCacheAffinityHitSummaryQuery, + UsageCacheAffinityIntervalGroupBy, UsageCacheAffinityIntervalQuery, UsageCacheHitSummaryQuery, + UsageCleanupExecutionMode, UsageCleanupPreviewCounts, UsageCleanupSummary, UsageCleanupTargets, + UsageCleanupWindow, UsageCostSavingsSummaryQuery, UsageCounterFlushSummary, + UsageCounterHealthSnapshot, UsageCounterPendingHealthSnapshot, UsageDailyActualCostRollupQuery, + UsageDailyHeatmapQuery, UsageDashboardDailyBreakdownQuery, UsageDashboardProviderCountsQuery, + UsageDashboardSummaryQuery, UsageErrorDistributionQuery, UsageLeaderboardGroupBy, + UsageLeaderboardQuery, UsageMonitoringErrorCountQuery, UsageMonitoringErrorListQuery, + UsagePerformancePercentilesQuery, UsageProviderPerformanceQuery, UsageReadRepository, + UsageRepository, UsageSettledCostSummaryQuery, UsageTimeSeriesGranularity, UsageTimeSeriesQuery, UsageWriteRepository, PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY, PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY, PROVIDER_REASONING_EFFORT_METADATA_KEY, PROVIDER_SERVICE_TIER_METADATA_KEY, REQUESTED_REASONING_EFFORT_METADATA_KEY, diff --git a/crates/aether-data/contracts/src/repository/usage/types.rs b/crates/aether-data/contracts/src/repository/usage/types.rs index 8587cc5b0..4d5d53650 100644 --- a/crates/aether-data/contracts/src/repository/usage/types.rs +++ b/crates/aether-data/contracts/src/repository/usage/types.rs @@ -1083,6 +1083,20 @@ pub struct StoredUsageSettledCostSummary { pub last_finalized_at_unix_secs: Option, } +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub struct UsageDailyActualCostRollupQuery { + pub finalized_from_unix_secs: u64, + pub finalized_until_unix_secs: u64, +} + +#[derive(Debug, Clone, PartialEq, Default, serde::Serialize, serde::Deserialize)] +pub struct StoredUsageDailyActualCostRollup { + pub user_id: Option, + pub api_key_id: String, + pub api_key_is_standalone: bool, + pub actual_total_cost_usd: f64, +} + #[derive(Debug, Clone, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)] pub struct UsageCacheAffinityHitSummaryQuery { pub created_from_unix_secs: u64, @@ -1734,6 +1748,11 @@ pub trait UsageReadRepository: Send + Sync { query: &UsageSettledCostSummaryQuery, ) -> Result; + async fn summarize_usage_daily_actual_cost_rollups( + &self, + query: &UsageDailyActualCostRollupQuery, + ) -> Result, crate::DataLayerError>; + async fn summarize_usage_cache_affinity_hit_summary( &self, query: &UsageCacheAffinityHitSummaryQuery, diff --git a/crates/aether-data/contracts/src/repository/users.rs b/crates/aether-data/contracts/src/repository/users.rs index e5d8f8fac..960232c51 100644 --- a/crates/aether-data/contracts/src/repository/users.rs +++ b/crates/aether-data/contracts/src/repository/users.rs @@ -498,6 +498,8 @@ pub struct StoredUserGroup { pub allowed_models_mode: String, pub rate_limit: Option, pub rate_limit_mode: String, + pub daily_usage_limit_usd: Option, + pub daily_usage_limit_mode: String, pub created_at: Option>, pub updated_at: Option>, } @@ -568,10 +570,28 @@ impl StoredUserGroup { &rate_limit_mode, "user_groups.rate_limit_mode", )?, + daily_usage_limit_usd: None, + daily_usage_limit_mode: "inherit".to_string(), created_at, updated_at, }) } + + pub fn with_daily_usage_limit( + mut self, + daily_usage_limit_usd: Option, + daily_usage_limit_mode: String, + ) -> Result { + self.daily_usage_limit_usd = normalize_optional_non_negative_f64( + daily_usage_limit_usd, + "user_groups.daily_usage_limit_usd", + )?; + self.daily_usage_limit_mode = normalize_rate_limit_policy_mode( + &daily_usage_limit_mode, + "user_groups.daily_usage_limit_mode", + )?; + Ok(self) + } } #[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] @@ -608,6 +628,8 @@ pub struct UpsertUserGroupRecord { pub allowed_models_mode: String, pub rate_limit: Option, pub rate_limit_mode: String, + pub daily_usage_limit_usd: Option, + pub daily_usage_limit_mode: String, } impl UpsertUserGroupRecord { @@ -1052,6 +1074,18 @@ pub fn normalize_user_group_name(value: &str) -> String { value.split_whitespace().collect::>().join(" ") } +fn normalize_optional_non_negative_f64( + value: Option, + field: &str, +) -> Result, crate::DataLayerError> { + match value { + Some(value) if !value.is_finite() || value < 0.0 => Err( + crate::DataLayerError::UnexpectedValue(format!("invalid {field}: {value}")), + ), + value => Ok(value), + } +} + pub fn normalize_list_policy_mode( value: &str, field_name: &str, diff --git a/crates/aether-data/runtime/schema/bootstrap/postgres/001_types_and_tables.sql b/crates/aether-data/runtime/schema/bootstrap/postgres/001_types_and_tables.sql index 44f1df04e..bfc45e19d 100644 --- a/crates/aether-data/runtime/schema/bootstrap/postgres/001_types_and_tables.sql +++ b/crates/aether-data/runtime/schema/bootstrap/postgres/001_types_and_tables.sql @@ -178,6 +178,7 @@ CREATE TABLE IF NOT EXISTS public.api_keys ( allowed_models json, ip_rules jsonb, rate_limit integer DEFAULT 100, + daily_usage_limit_usd double precision, concurrent_limit integer, force_capabilities json, feature_settings jsonb, @@ -1499,6 +1500,8 @@ CREATE TABLE IF NOT EXISTS public.user_groups ( allowed_models_mode text DEFAULT 'inherit'::text NOT NULL, rate_limit integer, rate_limit_mode text DEFAULT 'inherit'::text NOT NULL, + daily_usage_limit_usd double precision, + daily_usage_limit_mode text DEFAULT 'inherit'::text NOT NULL, created_at timestamp with time zone DEFAULT now() NOT NULL, updated_at timestamp with time zone DEFAULT now() NOT NULL ); diff --git a/crates/aether-data/runtime/schema/bootstrap/postgres/003_constraints.sql b/crates/aether-data/runtime/schema/bootstrap/postgres/003_constraints.sql index 8975f3bbf..b0292c1a1 100644 --- a/crates/aether-data/runtime/schema/bootstrap/postgres/003_constraints.sql +++ b/crates/aether-data/runtime/schema/bootstrap/postgres/003_constraints.sql @@ -1152,6 +1152,21 @@ END $mig$; +-- +-- Name: user_groups user_groups_daily_usage_limit_mode_check; Type: CHECK CONSTRAINT; Schema: public; Owner: - +-- + +DO $mig$ BEGIN + ALTER TABLE ONLY public.user_groups + ADD CONSTRAINT user_groups_daily_usage_limit_mode_check CHECK (daily_usage_limit_mode IN ('inherit', 'system', 'custom')); +EXCEPTION + WHEN duplicate_object THEN NULL; + WHEN duplicate_table THEN NULL; + WHEN invalid_table_definition THEN NULL; +END $mig$; + + + -- -- Name: user_oauth_links user_oauth_links_pkey; Type: CONSTRAINT; Schema: public; Owner: - -- diff --git a/crates/aether-data/runtime/schema/bootstrap/postgres/006_footer.sql b/crates/aether-data/runtime/schema/bootstrap/postgres/006_footer.sql index 31636e52f..d5045baaf 100644 --- a/crates/aether-data/runtime/schema/bootstrap/postgres/006_footer.sql +++ b/crates/aether-data/runtime/schema/bootstrap/postgres/006_footer.sql @@ -9,7 +9,8 @@ INSERT INTO public.user_groups ( allowed_providers_mode, allowed_api_formats_mode, allowed_models_mode, - rate_limit_mode + rate_limit_mode, + daily_usage_limit_mode ) VALUES ( '00000000-0000-0000-0000-000000000001', @@ -20,6 +21,7 @@ VALUES ( 'unrestricted', 'unrestricted', 'unrestricted', + 'system', 'system' ) ON CONFLICT (id) DO NOTHING; diff --git a/crates/aether-data/runtime/schema/generated/mysql/baseline/001_identity.sql b/crates/aether-data/runtime/schema/generated/mysql/baseline/001_identity.sql index 5b1f0a6bc..151656329 100644 --- a/crates/aether-data/runtime/schema/generated/mysql/baseline/001_identity.sql +++ b/crates/aether-data/runtime/schema/generated/mysql/baseline/001_identity.sql @@ -49,6 +49,8 @@ CREATE TABLE IF NOT EXISTS user_groups ( `allowed_models_mode` VARCHAR(32) NOT NULL DEFAULT 'inherit', `rate_limit` INT, `rate_limit_mode` VARCHAR(32) NOT NULL DEFAULT 'inherit', + `daily_usage_limit_usd` DOUBLE, + `daily_usage_limit_mode` VARCHAR(32) NOT NULL DEFAULT 'inherit', `created_at` BIGINT NOT NULL, `updated_at` BIGINT NOT NULL, PRIMARY KEY (`id`), @@ -77,6 +79,7 @@ CREATE TABLE IF NOT EXISTS api_keys ( `allowed_api_formats` JSON, `ip_rules` JSON, `rate_limit` INT DEFAULT 100, + `daily_usage_limit_usd` DOUBLE, `concurrent_limit` INT, `force_capabilities` JSON, `feature_settings` JSON, diff --git a/crates/aether-data/runtime/schema/generated/postgres/baseline/001_identity.sql b/crates/aether-data/runtime/schema/generated/postgres/baseline/001_identity.sql index d036d8d38..992729862 100644 --- a/crates/aether-data/runtime/schema/generated/postgres/baseline/001_identity.sql +++ b/crates/aether-data/runtime/schema/generated/postgres/baseline/001_identity.sql @@ -50,6 +50,8 @@ CREATE TABLE IF NOT EXISTS public.user_groups ( allowed_models_mode character varying(32) DEFAULT 'inherit' NOT NULL, rate_limit integer, rate_limit_mode character varying(32) DEFAULT 'inherit' NOT NULL, + daily_usage_limit_usd double precision, + daily_usage_limit_mode character varying(32) DEFAULT 'inherit' NOT NULL, created_at bigint NOT NULL, updated_at bigint NOT NULL ); @@ -80,6 +82,7 @@ CREATE TABLE IF NOT EXISTS public.api_keys ( allowed_api_formats jsonb, ip_rules jsonb, rate_limit integer DEFAULT 100, + daily_usage_limit_usd double precision, concurrent_limit integer, force_capabilities jsonb, feature_settings jsonb, diff --git a/crates/aether-data/runtime/schema/generated/sqlite/baseline/001_identity.sql b/crates/aether-data/runtime/schema/generated/sqlite/baseline/001_identity.sql index 560f5c8fe..2eba29849 100644 --- a/crates/aether-data/runtime/schema/generated/sqlite/baseline/001_identity.sql +++ b/crates/aether-data/runtime/schema/generated/sqlite/baseline/001_identity.sql @@ -48,6 +48,8 @@ CREATE TABLE IF NOT EXISTS user_groups ( allowed_models_mode TEXT NOT NULL DEFAULT 'inherit', rate_limit INTEGER, rate_limit_mode TEXT NOT NULL DEFAULT 'inherit', + daily_usage_limit_usd REAL, + daily_usage_limit_mode TEXT NOT NULL DEFAULT 'inherit', created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL, UNIQUE (normalized_name) @@ -75,6 +77,7 @@ CREATE TABLE IF NOT EXISTS api_keys ( allowed_api_formats TEXT, ip_rules TEXT, rate_limit INTEGER DEFAULT 100, + daily_usage_limit_usd REAL, concurrent_limit INTEGER, force_capabilities TEXT, feature_settings TEXT, diff --git a/crates/aether-data/runtime/schema/logical/001_identity.toml b/crates/aether-data/runtime/schema/logical/001_identity.toml index e99fb74db..87362ddb6 100644 --- a/crates/aether-data/runtime/schema/logical/001_identity.toml +++ b/crates/aether-data/runtime/schema/logical/001_identity.toml @@ -236,6 +236,17 @@ type = "text" length = 32 default = "inherit" +[[table.user_groups.columns]] +name = "daily_usage_limit_usd" +type = "float64" +nullable = true + +[[table.user_groups.columns]] +name = "daily_usage_limit_mode" +type = "text" +length = 32 +default = "inherit" + [[table.user_groups.columns]] name = "created_at" type = "unix_seconds" @@ -344,6 +355,11 @@ type = "int32" nullable = true default = 100 +[[table.api_keys.columns]] +name = "daily_usage_limit_usd" +type = "float64" +nullable = true + [[table.api_keys.columns]] name = "concurrent_limit" type = "int32" diff --git a/crates/aether-data/runtime/src/lifecycle/bootstrap/postgres.rs b/crates/aether-data/runtime/src/lifecycle/bootstrap/postgres.rs index b6aa14a0a..ea869c1c2 100644 --- a/crates/aether-data/runtime/src/lifecycle/bootstrap/postgres.rs +++ b/crates/aether-data/runtime/src/lifecycle/bootstrap/postgres.rs @@ -7,7 +7,7 @@ use tracing::info; // Generated by build.rs from schema/bootstrap/postgres. pub(crate) static EMPTY_DATABASE_SNAPSHOT_SQL: &str = include_str!(concat!(env!("OUT_DIR"), "/empty_database_snapshot.sql")); -pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260731000000; +pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260803000000; const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#" SELECT COUNT(*)::BIGINT diff --git a/crates/aether-data/runtime/src/lifecycle/migrate/tests.rs b/crates/aether-data/runtime/src/lifecycle/migrate/tests.rs index 9b753bb81..40e7cca51 100644 --- a/crates/aether-data/runtime/src/lifecycle/migrate/tests.rs +++ b/crates/aether-data/runtime/src/lifecycle/migrate/tests.rs @@ -411,6 +411,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() { 20260720000000, 20260727000000, 20260731000000, + 20260803000000, ] ); } @@ -1064,6 +1065,7 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() { 20260725030000, 20260727000000, 20260731000000, + 20260803000000, ] ); assert_eq!( @@ -1098,6 +1100,7 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() { 20260725040000, 20260727000000, 20260731000000, + 20260803000000, ] ); } @@ -2204,6 +2207,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() { 20260720000000, 20260727000000, 20260731000000, + 20260803000000, ] ); } diff --git a/crates/aether-data/runtime/src/repository/auth/memory.rs b/crates/aether-data/runtime/src/repository/auth/memory.rs index c38d8ee40..7c7ec282d 100644 --- a/crates/aether-data/runtime/src/repository/auth/memory.rs +++ b/crates/aether-data/runtime/src/repository/auth/memory.rs @@ -89,6 +89,7 @@ impl InMemoryAuthApiKeySnapshotRepository { .map(|value| serde_json::json!(value)), ) }) + .map(|record| record.with_daily_usage_limit(snapshot.api_key_daily_usage_limit_usd)) .expect("derived auth api key export record should build"), ); if let Some(key_hash) = key_hash { @@ -549,6 +550,7 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository { api_key_is_locked: false, api_key_is_standalone: false, api_key_rate_limit: Some(record.rate_limit), + api_key_daily_usage_limit_usd: record.daily_usage_limit_usd, api_key_concurrent_limit: record.concurrent_limit, api_key_expires_at_unix_secs: record.expires_at_unix_secs, api_key_allowed_providers: record.allowed_providers.clone(), @@ -599,6 +601,7 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository { .as_ref() .map(|value| serde_json::json!(value)), )? + .with_daily_usage_limits(None, record.daily_usage_limit_usd) }; let now_unix_secs = current_unix_secs() as i64; @@ -637,6 +640,7 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository { .as_ref() .map(|value| serde_json::json!(value)), )? + .with_daily_usage_limit(record.daily_usage_limit_usd) .with_activity_timestamps(None, Some(now_unix_secs), Some(now_unix_secs))?; index @@ -685,6 +689,7 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository { api_key_is_locked: false, api_key_is_standalone: true, api_key_rate_limit: record.rate_limit, + api_key_daily_usage_limit_usd: record.daily_usage_limit_usd, api_key_concurrent_limit: record.concurrent_limit, api_key_expires_at_unix_secs: record.expires_at_unix_secs, api_key_allowed_providers: record.allowed_providers.clone(), @@ -735,6 +740,7 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository { .as_ref() .map(|value| serde_json::json!(value)), )? + .with_daily_usage_limits(None, record.daily_usage_limit_usd) }; let now_unix_secs = current_unix_secs() as i64; @@ -773,6 +779,7 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository { .as_ref() .map(|value| serde_json::json!(value)), )? + .with_daily_usage_limit(record.daily_usage_limit_usd) .with_activity_timestamps(None, Some(now_unix_secs), Some(now_unix_secs))?; index @@ -817,6 +824,14 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository { export.rate_limit = Some(rate_limit); } } + if record.daily_usage_limit_present { + if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) { + snapshot.api_key_daily_usage_limit_usd = record.daily_usage_limit_usd; + } + if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) { + export.daily_usage_limit_usd = record.daily_usage_limit_usd; + } + } if let Some(concurrent_limit) = record.concurrent_limit { if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) { snapshot.api_key_concurrent_limit = Some(concurrent_limit); @@ -866,6 +881,14 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository { export.rate_limit = record.rate_limit; } } + if record.daily_usage_limit_present { + if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) { + snapshot.api_key_daily_usage_limit_usd = record.daily_usage_limit_usd; + } + if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) { + export.daily_usage_limit_usd = record.daily_usage_limit_usd; + } + } if record.concurrent_limit_present { if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) { snapshot.api_key_concurrent_limit = record.concurrent_limit; @@ -1358,6 +1381,8 @@ mod tests { api_key_id: "key-1".to_string(), name: None, rate_limit: None, + daily_usage_limit_present: true, + daily_usage_limit_usd: Some(6.0), concurrent_limit: Some(11), ip_rules: None, }) @@ -1365,6 +1390,23 @@ mod tests { .expect("update should succeed") .expect("record should exist"); assert_eq!(updated.concurrent_limit, Some(11)); + assert_eq!(updated.daily_usage_limit_usd, Some(6.0)); + + let cleared = repository + .update_user_api_key_basic(UpdateUserApiKeyBasicRecord { + user_id: "user-1".to_string(), + api_key_id: "key-1".to_string(), + name: None, + rate_limit: None, + daily_usage_limit_present: true, + daily_usage_limit_usd: None, + concurrent_limit: None, + ip_rules: None, + }) + .await + .expect("daily limit clear should succeed") + .expect("record should exist"); + assert_eq!(cleared.daily_usage_limit_usd, None); let snapshot = repository .find_api_key_snapshot(AuthApiKeyLookupKey::ApiKeyId("key-1")) @@ -1389,6 +1431,8 @@ mod tests { name: None, rate_limit_present: false, rate_limit: None, + daily_usage_limit_present: true, + daily_usage_limit_usd: Some(7.0), concurrent_limit_present: true, concurrent_limit: Some(13), allowed_providers: None, @@ -1404,6 +1448,7 @@ mod tests { .expect("update should succeed") .expect("record should exist"); assert_eq!(updated.concurrent_limit, Some(13)); + assert_eq!(updated.daily_usage_limit_usd, Some(7.0)); let snapshot = repository .find_api_key_snapshot(AuthApiKeyLookupKey::ApiKeyId("key-standalone")) diff --git a/crates/aether-data/runtime/src/repository/usage/memory.rs b/crates/aether-data/runtime/src/repository/usage/memory.rs index 2db8a4133..f8c44df16 100644 --- a/crates/aether-data/runtime/src/repository/usage/memory.rs +++ b/crates/aether-data/runtime/src/repository/usage/memory.rs @@ -7,8 +7,9 @@ use aether_data_contracts::repository::usage::{ parse_usage_body_ref, usage_body_ref, StoredUsageAuditAggregation, StoredUsageAuditSummary, StoredUsageBreakdownSummaryRow, StoredUsageCacheAffinityHitSummary, StoredUsageCacheAffinityIntervalRow, StoredUsageCacheHitSummary, StoredUsageCostSavingsSummary, - StoredUsageDashboardDailyBreakdownRow, StoredUsageDashboardProviderCount, - StoredUsageDashboardSummary, StoredUsageErrorDistributionRow, StoredUsageLeaderboardSummary, + StoredUsageDailyActualCostRollup, StoredUsageDashboardDailyBreakdownRow, + StoredUsageDashboardProviderCount, StoredUsageDashboardSummary, + StoredUsageErrorDistributionRow, StoredUsageLeaderboardSummary, StoredUsagePerformancePercentilesRow, StoredUsageProviderPerformance, StoredUsageProviderPerformanceProviderRow, StoredUsageProviderPerformanceSummary, StoredUsageProviderPerformanceTimelineRow, StoredUsageSettledCostSummary, @@ -17,13 +18,13 @@ use aether_data_contracts::repository::usage::{ UsageBodyCaptureState, UsageBodyField, UsageBreakdownGroupBy, UsageBreakdownSummaryQuery, UsageCacheAffinityHitSummaryQuery, UsageCacheAffinityIntervalGroupBy, UsageCacheAffinityIntervalQuery, UsageCacheHitSummaryQuery, UsageCostSavingsSummaryQuery, - UsageDashboardDailyBreakdownQuery, UsageDashboardProviderCountsQuery, - UsageDashboardSummaryQuery, UsageErrorDistributionQuery, UsageLeaderboardGroupBy, - UsageLeaderboardQuery, UsageMonitoringErrorCountQuery, UsageMonitoringErrorListQuery, - UsagePerformancePercentilesQuery, UsageProviderPerformanceQuery, UsageSettledCostSummaryQuery, - UsageTimeSeriesGranularity, UsageTimeSeriesQuery, PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY, - PROVIDER_REASONING_EFFORT_METADATA_KEY, PROVIDER_SERVICE_TIER_METADATA_KEY, - REQUESTED_REASONING_EFFORT_METADATA_KEY, + UsageDailyActualCostRollupQuery, UsageDashboardDailyBreakdownQuery, + UsageDashboardProviderCountsQuery, UsageDashboardSummaryQuery, UsageErrorDistributionQuery, + UsageLeaderboardGroupBy, UsageLeaderboardQuery, UsageMonitoringErrorCountQuery, + UsageMonitoringErrorListQuery, UsagePerformancePercentilesQuery, UsageProviderPerformanceQuery, + UsageSettledCostSummaryQuery, UsageTimeSeriesGranularity, UsageTimeSeriesQuery, + PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY, PROVIDER_REASONING_EFFORT_METADATA_KEY, + PROVIDER_SERVICE_TIER_METADATA_KEY, REQUESTED_REASONING_EFFORT_METADATA_KEY, }; use async_trait::async_trait; use chrono::Utc; @@ -1557,6 +1558,64 @@ impl UsageReadRepository for InMemoryUsageReadRepository { Ok(summary) } + async fn summarize_usage_daily_actual_cost_rollups( + &self, + query: &UsageDailyActualCostRollupQuery, + ) -> Result, DataLayerError> { + let mut rollups = BTreeMap::<(Option, String, bool), f64>::new(); + for item in self + .by_request_id + .read() + .expect("usage repository lock") + .values() + { + let finalized_at = item + .finalized_at_unix_secs + .unwrap_or(item.updated_at_unix_secs); + if finalized_at < query.finalized_from_unix_secs + || finalized_at >= query.finalized_until_unix_secs + { + continue; + } + let Some(api_key_id) = item + .api_key_id + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + else { + continue; + }; + if item.status != "completed" + || !item.actual_total_cost_usd.is_finite() + || item.actual_total_cost_usd <= 0.0 + { + continue; + } + let is_standalone = item + .request_metadata + .as_ref() + .and_then(|metadata| metadata.get("api_key_is_standalone")) + .and_then(serde_json::Value::as_bool) + .unwrap_or(false); + *rollups + .entry((item.user_id.clone(), api_key_id.to_string(), is_standalone)) + .or_default() += item.actual_total_cost_usd; + } + Ok(rollups + .into_iter() + .map( + |((user_id, api_key_id, api_key_is_standalone), actual_total_cost_usd)| { + StoredUsageDailyActualCostRollup { + user_id, + api_key_id, + api_key_is_standalone, + actual_total_cost_usd, + } + }, + ) + .collect()) + } + async fn summarize_dashboard_usage( &self, query: &UsageDashboardSummaryQuery, diff --git a/crates/aether-data/runtime/src/repository/usage/memory/tests.rs b/crates/aether-data/runtime/src/repository/usage/memory/tests.rs index 1c9f76e08..f379ef722 100644 --- a/crates/aether-data/runtime/src/repository/usage/memory/tests.rs +++ b/crates/aether-data/runtime/src/repository/usage/memory/tests.rs @@ -15,9 +15,9 @@ use crate::repository::usage::{ }; use aether_data_contracts::repository::usage::{ usage_body_ref, ProviderApiKeyWindowUsageRequest, UsageAuditAggregationGroupBy, - UsageAuditAggregationQuery, UsageBodyCaptureState, UsageBodyField, UsageDashboardSummaryQuery, - UsageLeaderboardGroupBy, UsageLeaderboardQuery, UsageProviderPerformanceQuery, - UsageTimeSeriesGranularity, + UsageAuditAggregationQuery, UsageBodyCaptureState, UsageBodyField, + UsageDailyActualCostRollupQuery, UsageDashboardSummaryQuery, UsageLeaderboardGroupBy, + UsageLeaderboardQuery, UsageProviderPerformanceQuery, UsageTimeSeriesGranularity, }; use serde_json::json; @@ -135,6 +135,55 @@ fn sample_upsert_usage_record(request_id: &str) -> UpsertUsageRecord { } } +#[tokio::test] +async fn daily_actual_cost_rollups_filter_window_status_and_group_scopes() { + let mut normal = sample_usage("normal", 1_000); + normal.finalized_at_unix_secs = Some(100); + normal.updated_at_unix_secs = 100; + normal.actual_total_cost_usd = 0.75; + + let mut standalone = sample_usage("standalone", 1_000); + standalone.finalized_at_unix_secs = Some(100); + standalone.updated_at_unix_secs = 100; + standalone.actual_total_cost_usd = 0.5; + standalone.api_key_id = Some("standalone-key".to_string()); + standalone.request_metadata = Some(json!({ "api_key_is_standalone": true })); + + let mut pending = sample_usage("pending", 1_000); + pending.finalized_at_unix_secs = Some(100); + pending.updated_at_unix_secs = 100; + pending.actual_total_cost_usd = 9.0; + pending.status = "pending".to_string(); + + let mut outside_window = sample_usage("outside", 1_000); + outside_window.finalized_at_unix_secs = Some(200); + outside_window.updated_at_unix_secs = 200; + outside_window.actual_total_cost_usd = 7.0; + + let repository = + InMemoryUsageReadRepository::seed([normal, standalone, pending, outside_window]); + let rollups = repository + .summarize_usage_daily_actual_cost_rollups(&UsageDailyActualCostRollupQuery { + finalized_from_unix_secs: 50, + finalized_until_unix_secs: 150, + }) + .await + .expect("daily usage rollups"); + assert_eq!(rollups.len(), 2); + let normal = rollups + .iter() + .find(|rollup| rollup.api_key_id == "api-key-1") + .expect("normal key rollup"); + assert!(!normal.api_key_is_standalone); + assert_eq!(normal.actual_total_cost_usd, 0.75); + let standalone = rollups + .iter() + .find(|rollup| rollup.api_key_id == "standalone-key") + .expect("standalone key rollup"); + assert!(standalone.api_key_is_standalone); + assert_eq!(standalone.actual_total_cost_usd, 0.5); +} + #[tokio::test] async fn upsert_uses_typed_provider_capture_as_the_fast_fact_snapshot() { for (name, state, incoming_tier, expected_tier) in [ diff --git a/crates/aether-data/runtime/src/repository/usage/mod.rs b/crates/aether-data/runtime/src/repository/usage/mod.rs index 8b630b92b..d293079b1 100644 --- a/crates/aether-data/runtime/src/repository/usage/mod.rs +++ b/crates/aether-data/runtime/src/repository/usage/mod.rs @@ -15,20 +15,21 @@ pub(crate) use aether_data_contracts::repository::usage::{ StoredProviderUsageSummary, StoredProviderUsageWindow, StoredRequestUsageAudit, StoredUsageAuditAggregation, StoredUsageAuditSummary, StoredUsageBreakdownSummaryRow, StoredUsageCacheAffinityHitSummary, StoredUsageCacheAffinityIntervalRow, - StoredUsageCacheHitSummary, StoredUsageCostSavingsSummary, StoredUsageDailySummary, - StoredUsageDashboardDailyBreakdownRow, StoredUsageDashboardProviderCount, - StoredUsageDashboardStatsSummary, StoredUsageDashboardSummary, StoredUsageErrorDistributionRow, - StoredUsageLeaderboardSummary, StoredUsagePerformancePercentilesRow, - StoredUsageProviderPerformance, StoredUsageProviderPerformanceProviderRow, - StoredUsageProviderPerformanceSummary, StoredUsageProviderPerformanceTimelineRow, - StoredUsageSettledCostSummary, StoredUsageTimeSeriesBucket, StoredUsageUserTotals, - UpsertUsageRecord, UsageAuditAggregationGroupBy, UsageAuditAggregationQuery, - UsageAuditKeywordSearchQuery, UsageAuditListQuery, UsageAuditSummaryQuery, - UsageBreakdownGroupBy, UsageBreakdownSummaryQuery, UsageCacheAffinityHitSummaryQuery, - UsageCacheAffinityIntervalGroupBy, UsageCacheAffinityIntervalQuery, UsageCacheHitSummaryQuery, - UsageCleanupPreviewCounts, UsageCleanupSummary, UsageCleanupWindow, - UsageCostSavingsSummaryQuery, UsageCounterFlushSummary, UsageCounterHealthSnapshot, - UsageCounterPendingHealthSnapshot, UsageDailyHeatmapQuery, UsageDashboardDailyBreakdownQuery, + StoredUsageCacheHitSummary, StoredUsageCostSavingsSummary, StoredUsageDailyActualCostRollup, + StoredUsageDailySummary, StoredUsageDashboardDailyBreakdownRow, + StoredUsageDashboardProviderCount, StoredUsageDashboardStatsSummary, + StoredUsageDashboardSummary, StoredUsageErrorDistributionRow, StoredUsageLeaderboardSummary, + StoredUsagePerformancePercentilesRow, StoredUsageProviderPerformance, + StoredUsageProviderPerformanceProviderRow, StoredUsageProviderPerformanceSummary, + StoredUsageProviderPerformanceTimelineRow, StoredUsageSettledCostSummary, + StoredUsageTimeSeriesBucket, StoredUsageUserTotals, UpsertUsageRecord, + UsageAuditAggregationGroupBy, UsageAuditAggregationQuery, UsageAuditKeywordSearchQuery, + UsageAuditListQuery, UsageAuditSummaryQuery, UsageBreakdownGroupBy, UsageBreakdownSummaryQuery, + UsageCacheAffinityHitSummaryQuery, UsageCacheAffinityIntervalGroupBy, + UsageCacheAffinityIntervalQuery, UsageCacheHitSummaryQuery, UsageCleanupPreviewCounts, + UsageCleanupSummary, UsageCleanupWindow, UsageCostSavingsSummaryQuery, + UsageCounterFlushSummary, UsageCounterHealthSnapshot, UsageCounterPendingHealthSnapshot, + UsageDailyActualCostRollupQuery, UsageDailyHeatmapQuery, UsageDashboardDailyBreakdownQuery, UsageDashboardProviderCountsQuery, UsageDashboardSummaryQuery, UsageErrorDistributionQuery, UsageLeaderboardGroupBy, UsageLeaderboardQuery, UsageMonitoringErrorCountQuery, UsageMonitoringErrorListQuery, UsagePerformancePercentilesQuery, UsageProviderPerformanceQuery, diff --git a/crates/aether-data/runtime/src/repository/usage/mysql.rs b/crates/aether-data/runtime/src/repository/usage/mysql.rs index 74bd34c29..8fba86866 100644 --- a/crates/aether-data/runtime/src/repository/usage/mysql.rs +++ b/crates/aether-data/runtime/src/repository/usage/mysql.rs @@ -5,17 +5,18 @@ use aether_data_contracts::repository::usage::{ StoredProviderApiKeyWindowUsageSummary, StoredProviderUsageSummary, StoredRequestUsageAudit, StoredUsageAuditAggregation, StoredUsageAuditSummary, StoredUsageBreakdownSummaryRow, StoredUsageCacheAffinityHitSummary, StoredUsageCacheAffinityIntervalRow, - StoredUsageCacheHitSummary, StoredUsageCostSavingsSummary, StoredUsageDailySummary, - StoredUsageDashboardDailyBreakdownRow, StoredUsageDashboardProviderCount, - StoredUsageDashboardSummary, StoredUsageErrorDistributionRow, StoredUsageLeaderboardSummary, + StoredUsageCacheHitSummary, StoredUsageCostSavingsSummary, StoredUsageDailyActualCostRollup, + StoredUsageDailySummary, StoredUsageDashboardDailyBreakdownRow, + StoredUsageDashboardProviderCount, StoredUsageDashboardSummary, + StoredUsageErrorDistributionRow, StoredUsageLeaderboardSummary, StoredUsagePerformancePercentilesRow, StoredUsageProviderPerformance, StoredUsageSettledCostSummary, StoredUsageTimeSeriesBucket, StoredUsageUserTotals, UsageAuditAggregationQuery, UsageAuditKeywordSearchQuery, UsageAuditListQuery, UsageAuditSummaryQuery, UsageBreakdownSummaryQuery, UsageCacheAffinityHitSummaryQuery, UsageCacheAffinityIntervalQuery, UsageCacheHitSummaryQuery, UsageCostSavingsSummaryQuery, - UsageDailyHeatmapQuery, UsageDashboardDailyBreakdownQuery, UsageDashboardProviderCountsQuery, - UsageDashboardSummaryQuery, UsageErrorDistributionQuery, UsageLeaderboardQuery, - UsageMonitoringErrorCountQuery, UsageMonitoringErrorListQuery, + UsageDailyActualCostRollupQuery, UsageDailyHeatmapQuery, UsageDashboardDailyBreakdownQuery, + UsageDashboardProviderCountsQuery, UsageDashboardSummaryQuery, UsageErrorDistributionQuery, + UsageLeaderboardQuery, UsageMonitoringErrorCountQuery, UsageMonitoringErrorListQuery, UsagePerformancePercentilesQuery, UsageProviderPerformanceQuery, UsageReadRepository, UsageSettledCostSummaryQuery, UsageTimeSeriesQuery, }; @@ -167,6 +168,15 @@ impl UsageReadRepository for MysqlUsageReadRepository { repository.summarize_usage_settled_cost(query).await } + async fn summarize_usage_daily_actual_cost_rollups( + &self, + query: &UsageDailyActualCostRollupQuery, + ) -> Result, DataLayerError> { + self.storage + .summarize_usage_daily_actual_cost_rollups(query) + .await + } + async fn summarize_usage_cache_affinity_hit_summary( &self, query: &UsageCacheAffinityHitSummaryQuery, diff --git a/crates/aether-data/runtime/src/repository/users/memory.rs b/crates/aether-data/runtime/src/repository/users/memory.rs index 87929d40d..660bb56fd 100644 --- a/crates/aether-data/runtime/src/repository/users/memory.rs +++ b/crates/aether-data/runtime/src/repository/users/memory.rs @@ -276,6 +276,8 @@ fn memory_group_from_record( ) -> Result { let now = chrono::Utc::now(); let name = normalize_user_group_name(&record.name); + let daily_usage_limit_usd = record.daily_usage_limit_usd; + let daily_usage_limit_mode = record.daily_usage_limit_mode.clone(); StoredUserGroup::new( uuid::Uuid::new_v4().to_string(), name.clone(), @@ -293,6 +295,7 @@ fn memory_group_from_record( Some(now), Some(now), ) + .and_then(|group| group.with_daily_usage_limit(daily_usage_limit_usd, daily_usage_limit_mode)) } fn memory_update_group_from_record( @@ -312,7 +315,11 @@ fn memory_update_group_from_record( group.allowed_models_mode = record.allowed_models_mode; group.rate_limit = record.rate_limit; group.rate_limit_mode = record.rate_limit_mode; + group.daily_usage_limit_usd = record.daily_usage_limit_usd; + group.daily_usage_limit_mode = record.daily_usage_limit_mode; group.updated_at = Some(chrono::Utc::now()); + let daily_usage_limit_usd = group.daily_usage_limit_usd; + let daily_usage_limit_mode = group.daily_usage_limit_mode.clone(); StoredUserGroup::new( group.id, group.name, @@ -330,6 +337,7 @@ fn memory_update_group_from_record( group.created_at, group.updated_at, ) + .and_then(|group| group.with_daily_usage_limit(daily_usage_limit_usd, daily_usage_limit_mode)) } fn memory_group_members( diff --git a/crates/aether-runtime/state/src/lib.rs b/crates/aether-runtime/state/src/lib.rs index 5c5915b5a..205b7d594 100644 --- a/crates/aether-runtime/state/src/lib.rs +++ b/crates/aether-runtime/state/src/lib.rs @@ -476,6 +476,64 @@ impl RuntimeState { } } + pub async fn increment_daily_usage_limit( + &self, + input: DailyUsageLimitIncrementInput<'_>, + ) -> Result { + match self.backend.as_ref() { + RuntimeStateBackend::Memory(memory) => memory.increment_daily_usage_limit( + input.user_key, + input.key_key, + input.bucket, + input.amount_units, + Duration::from_secs(input.ttl_seconds.max(1)), + ), + RuntimeStateBackend::Redis(redis) => { + redis.runtime.increment_daily_usage_limit(input).await + } + } + } + + pub async fn restore_daily_usage_limits( + &self, + input: DailyUsageLimitRestoreInput<'_>, + ) -> Result<(), DataLayerError> { + if input.entries.is_empty() { + return Ok(()); + } + match self.backend.as_ref() { + RuntimeStateBackend::Memory(memory) => memory.restore_daily_usage_limits( + input.entries, + input.bucket, + Duration::from_secs(input.ttl_seconds.max(1)), + ), + RuntimeStateBackend::Redis(redis) => { + redis.runtime.restore_daily_usage_limits(input).await + } + } + } + + pub async fn daily_usage_limit_counts( + &self, + input: DailyUsageLimitCountInput<'_>, + ) -> Result { + match self.backend.as_ref() { + RuntimeStateBackend::Memory(memory) => { + let mut counts = + memory.daily_usage_limit_counts(input.user_key, input.key_key, input.bucket)?; + counts.state_ready = memory + .kv_get(input.state_key) + .await + .as_deref() + .is_some_and(|value| value == "ready"); + Ok(counts) + } + RuntimeStateBackend::Redis(redis) => { + redis.runtime.daily_usage_limit_counts(input).await + } + } + } + pub async fn set_add(&self, key: &str, member: &str) -> Result { match self.backend.as_ref() { RuntimeStateBackend::Memory(memory) => Ok(memory.set_add(key, member).await), @@ -761,6 +819,47 @@ pub struct RateLimitInput<'a> { pub ttl_seconds: u64, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct DailyUsageLimitCounts { + pub user_units: u64, + pub key_units: u64, + pub user_present: bool, + pub key_present: bool, + pub state_ready: bool, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DailyUsageLimitIncrementInput<'a> { + pub user_key: Option<&'a str>, + pub key_key: &'a str, + pub bucket: u64, + pub amount_units: u64, + pub ttl_seconds: u64, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DailyUsageLimitCountInput<'a> { + pub state_key: &'a str, + pub user_key: Option<&'a str>, + pub key_key: &'a str, + pub bucket: u64, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DailyUsageLimitRestoreEntry { + pub user_key: Option, + pub key_key: String, + pub user_units: u64, + pub key_units: u64, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DailyUsageLimitRestoreInput<'a> { + pub entries: &'a [DailyUsageLimitRestoreEntry], + pub bucket: u64, + pub ttl_seconds: u64, +} + #[derive(Debug, Clone, PartialEq, Eq)] pub struct RuntimeQueueEntry { pub id: String, @@ -1637,6 +1736,225 @@ mod tests { ); } + #[tokio::test] + async fn memory_daily_usage_limit_accumulates_user_and_key_scopes() { + let runtime = RuntimeState::memory(MemoryRuntimeStateConfig::default()); + let user_key = "daily_usage_limit:user:shared:1"; + let first_key = "daily_usage_limit:key:first:1"; + let second_key = "daily_usage_limit:key:second:1"; + let standalone_key = "daily_usage_limit:key:standalone:1"; + runtime + .increment_daily_usage_limit(DailyUsageLimitIncrementInput { + user_key: Some(user_key), + key_key: first_key, + bucket: 1, + amount_units: 125, + ttl_seconds: 60, + }) + .await + .expect("first daily usage increment"); + runtime + .increment_daily_usage_limit(DailyUsageLimitIncrementInput { + user_key: Some(user_key), + key_key: second_key, + bucket: 1, + amount_units: 75, + ttl_seconds: 60, + }) + .await + .expect("second daily usage increment"); + runtime + .increment_daily_usage_limit(DailyUsageLimitIncrementInput { + user_key: None, + key_key: standalone_key, + bucket: 1, + amount_units: 50, + ttl_seconds: 60, + }) + .await + .expect("standalone daily usage increment"); + + let user_counts = runtime + .daily_usage_limit_counts(DailyUsageLimitCountInput { + state_key: "daily_usage_limit:runtime_state", + user_key: Some(user_key), + key_key: first_key, + bucket: 1, + }) + .await + .expect("user daily usage count"); + assert_eq!(user_counts.user_units, 200); + assert_eq!(user_counts.key_units, 125); + + let standalone_counts = runtime + .daily_usage_limit_counts(DailyUsageLimitCountInput { + state_key: "daily_usage_limit:runtime_state", + user_key: None, + key_key: standalone_key, + bucket: 1, + }) + .await + .expect("standalone daily usage count"); + assert_eq!(standalone_counts.user_units, 0); + assert_eq!(standalone_counts.key_units, 50); + } + + #[tokio::test] + async fn memory_daily_usage_limit_batch_restore_preserves_higher_counts_and_marks_ready() { + let runtime = RuntimeState::memory(MemoryRuntimeStateConfig::default()); + let user_key = "daily_usage_limit:user:restore:1"; + let key_key = "daily_usage_limit:key:restore:1"; + + let missing = runtime + .daily_usage_limit_counts(DailyUsageLimitCountInput { + state_key: "daily_usage_limit:runtime_state", + user_key: Some(user_key), + key_key, + bucket: 1, + }) + .await + .expect("missing daily usage counts"); + assert!(!missing.user_present); + assert!(!missing.key_present); + + let restored = [DailyUsageLimitRestoreEntry { + user_key: Some(user_key.to_string()), + key_key: key_key.to_string(), + user_units: 100, + key_units: 100, + }]; + runtime + .restore_daily_usage_limits(DailyUsageLimitRestoreInput { + entries: &restored, + bucket: 1, + ttl_seconds: 60, + }) + .await + .expect("restore daily usage counts"); + + runtime + .increment_daily_usage_limit(DailyUsageLimitIncrementInput { + user_key: Some(user_key), + key_key, + bucket: 1, + amount_units: 50, + ttl_seconds: 60, + }) + .await + .expect("increment restored daily usage counts"); + let stale_restore = [DailyUsageLimitRestoreEntry { + user_key: Some(user_key.to_string()), + key_key: key_key.to_string(), + user_units: 125, + key_units: 125, + }]; + runtime + .restore_daily_usage_limits(DailyUsageLimitRestoreInput { + entries: &stale_restore, + bucket: 1, + ttl_seconds: 60, + }) + .await + .expect("preserve higher daily usage counts"); + runtime + .kv_set("daily_usage_limit:runtime_state", "ready", None) + .await + .expect("mark daily usage runtime ready"); + let preserved = runtime + .daily_usage_limit_counts(DailyUsageLimitCountInput { + state_key: "daily_usage_limit:runtime_state", + user_key: Some(user_key), + key_key, + bucket: 1, + }) + .await + .expect("read restored daily usage counts"); + assert_eq!(preserved.user_units, 150); + assert_eq!(preserved.key_units, 150); + assert!(preserved.state_ready); + } + + #[tokio::test] + async fn redis_daily_usage_limit_accumulates_user_and_key_scopes() { + let Some((_redis, runtime)) = redis_runtime_for_test("daily-usage-limit").await else { + return; + }; + let user_key = "daily_usage_limit:user:shared:1"; + let key_key = "daily_usage_limit:key:first:1"; + for amount_units in [125, 75] { + runtime + .increment_daily_usage_limit(DailyUsageLimitIncrementInput { + user_key: Some(user_key), + key_key, + bucket: 1, + amount_units, + ttl_seconds: 60, + }) + .await + .expect("redis daily usage increment"); + } + + let counts = runtime + .daily_usage_limit_counts(DailyUsageLimitCountInput { + state_key: "daily_usage_limit:runtime_state", + user_key: Some(user_key), + key_key, + bucket: 1, + }) + .await + .expect("redis daily usage counts"); + assert_eq!(counts.user_units, 200); + assert_eq!(counts.key_units, 200); + } + + #[tokio::test] + async fn redis_daily_usage_limit_batch_restore_preserves_higher_counts() { + let Some((_redis, runtime)) = redis_runtime_for_test("daily-usage-limit-restore").await + else { + return; + }; + let user_key = "daily_usage_limit:user:restore:1"; + let key_key = "daily_usage_limit:key:restore:1"; + runtime + .increment_daily_usage_limit(DailyUsageLimitIncrementInput { + user_key: Some(user_key), + key_key, + bucket: 1, + amount_units: 150, + ttl_seconds: 60, + }) + .await + .expect("redis daily usage increment"); + + let entries = [DailyUsageLimitRestoreEntry { + user_key: Some(user_key.to_string()), + key_key: key_key.to_string(), + user_units: 125, + key_units: 125, + }]; + runtime + .restore_daily_usage_limits(DailyUsageLimitRestoreInput { + entries: &entries, + bucket: 1, + ttl_seconds: 60, + }) + .await + .expect("redis daily usage restore"); + let preserved = runtime + .daily_usage_limit_counts(DailyUsageLimitCountInput { + state_key: "daily_usage_limit:runtime_state", + user_key: Some(user_key), + key_key, + bucket: 1, + }) + .await + .expect("redis restored daily usage counts"); + assert_eq!(preserved.user_units, 150); + assert_eq!(preserved.key_units, 150); + assert!(preserved.user_present); + assert!(preserved.key_present); + } + #[tokio::test] async fn memory_rate_limit_concurrent_checks_do_not_exceed_limit() { let runtime = diff --git a/crates/aether-runtime/state/src/memory.rs b/crates/aether-runtime/state/src/memory.rs index 6e0be8293..f2cca4772 100644 --- a/crates/aether-runtime/state/src/memory.rs +++ b/crates/aether-runtime/state/src/memory.rs @@ -53,7 +53,7 @@ pub(crate) struct MemoryRuntimeBackend { #[derive(Debug, Clone)] struct MemoryCounterEntry { - value: u32, + value: u64, bucket: u64, expires_at: Instant, } @@ -437,7 +437,7 @@ impl MemoryRuntimeBackend { .filter(|entry| entry.bucket == bucket) .map(|entry| entry.value) .unwrap_or_default(); - if user_count >= user_limit { + if user_count >= u64::from(user_limit) { return Ok(crate::RateLimitCheck::Rejected { scope: crate::RateLimitScope::User, limit: user_limit, @@ -452,7 +452,7 @@ impl MemoryRuntimeBackend { .filter(|entry| entry.bucket == bucket) .map(|entry| entry.value) .unwrap_or_default(); - if key_count >= key_limit { + if key_count >= u64::from(key_limit) { return Ok(crate::RateLimitCheck::Rejected { scope: crate::RateLimitScope::Key, limit: key_limit, @@ -477,7 +477,7 @@ impl MemoryRuntimeBackend { expires_at, }) .value; - remaining = Some(user_limit.saturating_sub(next)); + remaining = Some(user_limit.saturating_sub(u32::try_from(next).unwrap_or(u32::MAX))); } if key_limit > 0 { let next = shard @@ -494,7 +494,7 @@ impl MemoryRuntimeBackend { expires_at, }) .value; - let key_remaining = key_limit.saturating_sub(next); + let key_remaining = key_limit.saturating_sub(u32::try_from(next).unwrap_or(u32::MAX)); remaining = Some(remaining.map_or(key_remaining, |value| value.min(key_remaining))); } Ok(crate::RateLimitCheck::Allowed { @@ -521,13 +521,225 @@ impl MemoryRuntimeBackend { .entries .get(key) .filter(|entry| entry.bucket == bucket) - .map(|entry| entry.value) + .map(|entry| u32::try_from(entry.value).unwrap_or(u32::MAX)) .unwrap_or_default(), ); } Ok(total) } + pub(crate) fn increment_daily_usage_limit( + &self, + user_key: Option<&str>, + key_key: &str, + bucket: u64, + amount_units: u64, + ttl: Duration, + ) -> Result { + let partition_key = user_key.unwrap_or(key_key); + let shard_index = memory_rate_limit_counter_shard_index(partition_key); + let mut shard = self.counters.shards[shard_index].lock().map_err(|_| { + DataLayerError::UnexpectedValue( + "memory daily usage limit counter lock poisoned".to_string(), + ) + })?; + let now = Instant::now(); + shard.amortized_prune(now); + if let Some(user_key) = user_key { + prune_rate_limit_counter(&mut shard.entries, user_key, bucket, now); + } + prune_rate_limit_counter(&mut shard.entries, key_key, bucket, now); + + let expires_at = now + ttl; + let user_units = user_key + .map(|user_key| { + shard + .entries + .entry(user_key.to_string()) + .and_modify(|entry| { + entry.bucket = bucket; + entry.value = entry.value.saturating_add(amount_units); + entry.expires_at = expires_at; + }) + .or_insert(MemoryCounterEntry { + value: amount_units, + bucket, + expires_at, + }) + .value + }) + .unwrap_or_default(); + let key_units = shard + .entries + .entry(key_key.to_string()) + .and_modify(|entry| { + entry.bucket = bucket; + entry.value = entry.value.saturating_add(amount_units); + entry.expires_at = expires_at; + }) + .or_insert(MemoryCounterEntry { + value: amount_units, + bucket, + expires_at, + }) + .value; + Ok(crate::DailyUsageLimitCounts { + user_units, + key_units, + user_present: user_key.is_some(), + key_present: true, + state_ready: false, + }) + } + + pub(crate) fn daily_usage_limit_counts( + &self, + user_key: Option<&str>, + key_key: &str, + bucket: u64, + ) -> Result { + let partition_key = user_key.unwrap_or(key_key); + let shard_index = memory_rate_limit_counter_shard_index(partition_key); + let mut shard = self.counters.shards[shard_index].lock().map_err(|_| { + DataLayerError::UnexpectedValue( + "memory daily usage limit counter lock poisoned".to_string(), + ) + })?; + let now = Instant::now(); + shard.amortized_prune(now); + if let Some(user_key) = user_key { + prune_rate_limit_counter(&mut shard.entries, user_key, bucket, now); + } + let user_present = user_key.is_some_and(|user_key| { + shard + .entries + .get(user_key) + .is_some_and(|entry| entry.bucket == bucket) + }); + let user_units = user_key + .map(|user_key| { + shard + .entries + .get(user_key) + .filter(|entry| entry.bucket == bucket) + .map(|entry| entry.value) + .unwrap_or_default() + }) + .unwrap_or_default(); + prune_rate_limit_counter(&mut shard.entries, key_key, bucket, now); + let key_present = shard + .entries + .get(key_key) + .is_some_and(|entry| entry.bucket == bucket); + let key_units = shard + .entries + .get(key_key) + .filter(|entry| entry.bucket == bucket) + .map(|entry| entry.value) + .unwrap_or_default(); + Ok(crate::DailyUsageLimitCounts { + user_units, + key_units, + user_present, + key_present, + state_ready: false, + }) + } + + fn restore_daily_usage_limit_entry( + &self, + user_key: Option<&str>, + key_key: &str, + bucket: u64, + user_units: u64, + key_units: u64, + ttl: Duration, + ) -> Result { + let partition_key = user_key.unwrap_or(key_key); + let shard_index = memory_rate_limit_counter_shard_index(partition_key); + let mut shard = self.counters.shards[shard_index].lock().map_err(|_| { + DataLayerError::UnexpectedValue( + "memory daily usage limit counter lock poisoned".to_string(), + ) + })?; + let now = Instant::now(); + shard.amortized_prune(now); + if let Some(user_key) = user_key { + prune_rate_limit_counter(&mut shard.entries, user_key, bucket, now); + } + prune_rate_limit_counter(&mut shard.entries, key_key, bucket, now); + let expires_at = now + ttl; + let (user_value, user_present) = match user_key { + Some(user_key) => { + let entry = + shard + .entries + .entry(user_key.to_string()) + .or_insert(MemoryCounterEntry { + value: user_units, + bucket, + expires_at, + }); + if entry.bucket != bucket { + *entry = MemoryCounterEntry { + value: user_units, + bucket, + expires_at, + }; + } else { + entry.value = entry.value.max(user_units); + entry.expires_at = expires_at; + } + (entry.value, true) + } + None => (0, false), + }; + let key_entry = shard + .entries + .entry(key_key.to_string()) + .or_insert(MemoryCounterEntry { + value: key_units, + bucket, + expires_at, + }); + if key_entry.bucket != bucket { + *key_entry = MemoryCounterEntry { + value: key_units, + bucket, + expires_at, + }; + } else { + key_entry.value = key_entry.value.max(key_units); + key_entry.expires_at = expires_at; + } + Ok(crate::DailyUsageLimitCounts { + user_units: user_value, + key_units: key_entry.value, + user_present, + key_present: true, + state_ready: false, + }) + } + + pub(crate) fn restore_daily_usage_limits( + &self, + entries: &[crate::DailyUsageLimitRestoreEntry], + bucket: u64, + ttl: Duration, + ) -> Result<(), DataLayerError> { + for entry in entries { + self.restore_daily_usage_limit_entry( + entry.user_key.as_deref(), + &entry.key_key, + bucket, + entry.user_units, + entry.key_units, + ttl, + )?; + } + Ok(()) + } + pub(crate) async fn set_add(&self, key: &str, member: &str) -> bool { let mut sets = self.sets.lock().await; prune_memory_key(&mut sets, key, Instant::now()); diff --git a/crates/aether-runtime/state/src/redis/runtime.rs b/crates/aether-runtime/state/src/redis/runtime.rs index ccdfda5b1..477679671 100644 --- a/crates/aether-runtime/state/src/redis/runtime.rs +++ b/crates/aether-runtime/state/src/redis/runtime.rs @@ -6,7 +6,9 @@ use crate::redis::{ RedisKeyspace, RedisLaneDiagnostics, }; use crate::{ - DataLayerError, RateLimitCheck, RateLimitInput, RateLimitScope, RuntimeSemaphoreError, + DailyUsageLimitCountInput, DailyUsageLimitCounts, DailyUsageLimitIncrementInput, + DailyUsageLimitRestoreInput, DataLayerError, RateLimitCheck, RateLimitInput, RateLimitScope, + RuntimeSemaphoreError, }; const RATE_LIMIT_CHECK_AND_CONSUME_SCRIPT: &str = r#" @@ -51,6 +53,58 @@ end return {1, 0, 0, remaining} "#; +const DAILY_USAGE_LIMIT_INCREMENT_SCRIPT: &str = r#" +local user_key = KEYS[1] +local key_key = KEYS[2] +local include_user = tonumber(ARGV[1]) +local amount = tonumber(ARGV[2]) +local ttl = tonumber(ARGV[3]) + +local user_units = 0 +if include_user == 1 then + user_units = redis.call('INCRBY', user_key, amount) + redis.call('EXPIRE', user_key, ttl) +end + +local key_units = redis.call('INCRBY', key_key, amount) +redis.call('EXPIRE', key_key, ttl) +return {user_units, key_units} +"#; + +const DAILY_USAGE_LIMIT_RESTORE_BATCH_SCRIPT: &str = r#" +local ttl = tonumber(ARGV[1]) +local entry_count = #KEYS / 2 + +for i = 1, entry_count do + local user_key = KEYS[(i - 1) * 2 + 1] + local key_key = KEYS[(i - 1) * 2 + 2] + local arg_offset = (i - 1) * 3 + 2 + local include_user = tonumber(ARGV[arg_offset]) + local user_units = tonumber(ARGV[arg_offset + 1]) + local key_units = tonumber(ARGV[arg_offset + 2]) + + if include_user == 1 then + local current_user_units = tonumber(redis.call('GET', user_key) or '0') + if current_user_units < user_units then + redis.call('SET', user_key, user_units, 'EX', ttl) + else + redis.call('EXPIRE', user_key, ttl) + end + end + + local current_key_units = tonumber(redis.call('GET', key_key) or '0') + if current_key_units < key_units then + redis.call('SET', key_key, key_units, 'EX', ttl) + else + redis.call('EXPIRE', key_key, ttl) + end +end + +return entry_count +"#; + +const DAILY_USAGE_LIMIT_RESTORE_BATCH_SIZE: usize = 256; + #[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)] pub struct RedisRuntimeDiagnostics { pub connected_clients: Option, @@ -282,6 +336,131 @@ impl RedisRuntimeRunner { Ok(RateLimitCheck::Rejected { scope, limit }) } + pub(crate) async fn increment_daily_usage_limit( + &self, + input: DailyUsageLimitIncrementInput<'_>, + ) -> Result { + let user_key = self.keyspace.key(input.user_key.unwrap_or(input.key_key)); + let key_key = self.keyspace.key(input.key_key); + let amount = i64::try_from(input.amount_units).map_err(|_| { + DataLayerError::InvalidInput("daily usage limit increment exceeds i64".to_string()) + })?; + let raw = run_lane_with_timeout( + &self.connections, + RedisConnectionLane::Fast, + self.command_timeout_ms, + "runtime daily usage limit increment", + async { + let mut connection = self.connections.connection(RedisConnectionLane::Fast); + script(DAILY_USAGE_LIMIT_INCREMENT_SCRIPT) + .key(user_key) + .key(key_key) + .arg(i64::from(input.user_key.is_some())) + .arg(amount) + .arg(i64::try_from(input.ttl_seconds.max(1)).unwrap_or(i64::MAX)) + .invoke_async::>(&mut connection) + .await + .map_redis_err() + }, + ) + .await?; + Ok(DailyUsageLimitCounts { + user_units: raw + .first() + .copied() + .and_then(|value| u64::try_from(value).ok()) + .unwrap_or_default(), + key_units: raw + .get(1) + .copied() + .and_then(|value| u64::try_from(value).ok()) + .unwrap_or_default(), + user_present: input.user_key.is_some(), + key_present: true, + state_ready: false, + }) + } + + pub(crate) async fn restore_daily_usage_limits( + &self, + input: DailyUsageLimitRestoreInput<'_>, + ) -> Result<(), DataLayerError> { + for entries in input.entries.chunks(DAILY_USAGE_LIMIT_RESTORE_BATCH_SIZE) { + run_lane_with_timeout( + &self.connections, + RedisConnectionLane::Fast, + self.command_timeout_ms, + "runtime daily usage limit restore batch", + async { + let restore_script = script(DAILY_USAGE_LIMIT_RESTORE_BATCH_SCRIPT); + let mut invocation = restore_script.prepare_invoke(); + for entry in entries { + invocation + .key( + self.keyspace + .key(entry.user_key.as_deref().unwrap_or(&entry.key_key)), + ) + .key(self.keyspace.key(&entry.key_key)); + } + invocation.arg(i64::try_from(input.ttl_seconds.max(1)).unwrap_or(i64::MAX)); + for entry in entries { + invocation + .arg(i64::from(entry.user_key.is_some())) + .arg(i64::try_from(entry.user_units).unwrap_or(i64::MAX)) + .arg(i64::try_from(entry.key_units).unwrap_or(i64::MAX)); + } + let mut connection = self.connections.connection(RedisConnectionLane::Fast); + invocation + .invoke_async::(&mut connection) + .await + .map_redis_err() + }, + ) + .await?; + } + Ok(()) + } + + pub(crate) async fn daily_usage_limit_counts( + &self, + input: DailyUsageLimitCountInput<'_>, + ) -> Result { + let mut keys = Vec::with_capacity(3); + keys.push(input.state_key.to_string()); + if let Some(user_key) = input.user_key { + keys.push(user_key.to_string()); + } + keys.push(input.key_key.to_string()); + let values = self.kv_get_many(&keys).await?; + let parse = |value: Option<&Option>| { + value + .and_then(Option::as_deref) + .and_then(|value| value.parse::().ok()) + .unwrap_or_default() + }; + let state_ready = values + .first() + .and_then(Option::as_deref) + .is_some_and(|value| value == "ready"); + Ok(if input.user_key.is_some() { + DailyUsageLimitCounts { + user_units: parse(values.get(1)), + key_units: parse(values.get(2)), + user_present: values.get(1).is_some_and(Option::is_some), + key_present: values.get(2).is_some_and(Option::is_some), + state_ready, + } + } else { + DailyUsageLimitCounts { + user_units: 0, + key_units: parse(values.get(1)), + user_present: false, + key_present: values.get(1).is_some_and(Option::is_some), + state_ready, + } + }) + } + pub(crate) async fn set_add(&self, key: &str, member: &str) -> Result { let key = self.keyspace.key(key); let mut command = cmd("SADD"); diff --git a/crates/aether-testing/integration/src/bin/gateway_pressure_seed.rs b/crates/aether-testing/integration/src/bin/gateway_pressure_seed.rs index f6062e5f5..be47e93d4 100644 --- a/crates/aether-testing/integration/src/bin/gateway_pressure_seed.rs +++ b/crates/aether-testing/integration/src/bin/gateway_pressure_seed.rs @@ -507,6 +507,7 @@ async fn seed_api_key( allowed_models: Some(vec![config.model.clone()]), ip_rules: None, rate_limit: Some(0), + daily_usage_limit_usd: None, concurrent_limit: None, force_capabilities: None, is_active: true, @@ -525,6 +526,8 @@ async fn seed_api_key( name: Some(format!("Local pressure API key {}", key_index + 1)), rate_limit_present: true, rate_limit: Some(0), + daily_usage_limit_present: true, + daily_usage_limit_usd: None, concurrent_limit_present: true, concurrent_limit: None, allowed_providers: Some(Some(vec![config.provider_id.clone()])), diff --git a/docs/operations/redis-runtime-runbook.md b/docs/operations/redis-runtime-runbook.md index 4d3184be6..5298c2c90 100644 --- a/docs/operations/redis-runtime-runbook.md +++ b/docs/operations/redis-runtime-runbook.md @@ -31,6 +31,28 @@ crash should restore persistence in the Redis command: Expect higher tail latency when Redis persistence shares disks with Postgres or application logs. +### Daily usage limit counters + +Daily usage limits intentionally remain compatible with the default +non-persistent Redis policy. Every positive finalized request updates the user +and API key counters even when that scope is currently unlimited, so changing a +limit only changes the threshold and never starts a database query. + +The gateway reads a global daily-usage runtime-state marker together with the +applicable counters in the same Redis `MGET`. If Redis loses its runtime state, +the missing marker makes daily limits temporarily fail open and starts one +distributed background recovery. That recovery performs one grouped scan of +the current `APP_TIMEZONE` day, restores counters in Redis batches without +lowering concurrent values, and marks the runtime state ready. It does not run +a query per user or API key, does not put SQL on the request path, and does not +require usage-table indexes dedicated to this feature. + +Gateway process restarts do not trigger recovery while Redis still contains the +marker. Redis container or host restarts do trigger recovery on the first +limited request. A small amount of undercounting or over-limit traffic is +acceptable while recovery overlaps in-flight settlements; this feature is a +traffic policy, not a financial balance ledger. + ### OpenAI Responses continuation history When an OpenAI Responses request is converted to an OpenAI Chat provider, diff --git a/frontend/src/api/admin.ts b/frontend/src/api/admin.ts index 51176ae09..148a63c53 100644 --- a/frontend/src/api/admin.ts +++ b/frontend/src/api/admin.ts @@ -138,6 +138,8 @@ export interface UserGroupExport { allowed_models_mode?: 'inherit' | 'unrestricted' | 'specific' | 'deny_all' rate_limit?: number | null rate_limit_mode?: 'inherit' | 'system' | 'custom' + daily_usage_limit_usd?: number | null + daily_usage_limit_mode?: 'inherit' | 'system' | 'custom' } export interface UserExport { @@ -179,6 +181,7 @@ export interface UserApiKeyExport { allowed_models?: string[] | null ip_rules?: string[] | null rate_limit?: number | null // legacy/null 兼容；1.3+ standalone null = 跟随系统默认 + daily_usage_limit_usd?: number | null concurrent_limit?: number | null force_capabilities?: Record feature_settings?: Record | null @@ -658,6 +661,7 @@ export interface AdminApiKey { total_tokens?: number | null total_cost_usd?: number rate_limit?: number | null // null = 跟随系统默认，0 = 不限制 + daily_usage_limit_usd?: number | null // null = 跟随系统，0 = 明确不限制 concurrent_limit?: number | null // null = 跟随系统默认，0 = 不限制 allowed_providers?: string[] | null // 允许的提供商列表 allowed_api_formats?: string[] | null // 允许的 API 格式列表 @@ -679,6 +683,7 @@ export interface CreateStandaloneApiKeyRequest { allowed_models?: string[] | null ip_rules?: string[] | null rate_limit?: number | null // null = 跟随系统默认，0 = 不限制 + daily_usage_limit_usd?: number | null // null = 跟随系统，0 = 明确不限制 concurrent_limit?: number | null // null = 跟随系统默认，0 = 不限制 expires_at?: string | null // RFC3339 时间，null = 永不过期 initial_balance_usd: number | null // 初始余额，null = 无限制 diff --git a/frontend/src/api/me.ts b/frontend/src/api/me.ts index e4a312cb5..4b1bfb070 100644 --- a/frontend/src/api/me.ts +++ b/frontend/src/api/me.ts @@ -183,6 +183,7 @@ export interface ApiKey { total_requests?: number total_cost_usd?: number rate_limit?: number | null + daily_usage_limit_usd?: number | null concurrent_limit?: number | null ip_rules?: string[] | null allowed_providers?: ProviderConfig[] @@ -269,7 +270,7 @@ export const meApi = { return response.data }, - async createApiKey(data: { name: string; rate_limit?: number | null; concurrent_limit?: number | null; ip_rules?: string[] | null; feature_settings?: FeatureSettingsMap | null }): Promise { + async createApiKey(data: { name: string; rate_limit?: number | null; daily_usage_limit_usd?: number | null; concurrent_limit?: number | null; ip_rules?: string[] | null; feature_settings?: FeatureSettingsMap | null }): Promise { const response = await apiClient.post('/api/users/me/api-keys', data) return response.data }, @@ -307,7 +308,7 @@ export const meApi = { async updateApiKey( keyId: string, - data: { name?: string; rate_limit?: number | null; concurrent_limit?: number | null; ip_rules?: string[] | null; feature_settings?: FeatureSettingsMap | null | undefined } + data: { name?: string; rate_limit?: number | null; daily_usage_limit_usd?: number | null; concurrent_limit?: number | null; ip_rules?: string[] | null; feature_settings?: FeatureSettingsMap | null | undefined } ): Promise { const response = await apiClient.put( `/api/users/me/api-keys/${keyId}`, diff --git a/frontend/src/api/users.ts b/frontend/src/api/users.ts index bc4b682d0..7552010f8 100644 --- a/frontend/src/api/users.ts +++ b/frontend/src/api/users.ts @@ -30,6 +30,7 @@ export interface UserEffectivePolicy { allowed_api_formats?: EffectivePolicyField allowed_models?: EffectivePolicyField rate_limit?: EffectivePolicyField + daily_usage_limit_usd?: EffectivePolicyField } export interface User { @@ -175,6 +176,8 @@ export interface UserGroup { allowed_models_mode: ListPolicyMode rate_limit?: number | null rate_limit_mode: RateLimitPolicyMode + daily_usage_limit_usd?: number | null + daily_usage_limit_mode: RateLimitPolicyMode is_default?: boolean created_at?: string | null updated_at?: string | null @@ -191,6 +194,8 @@ export interface UpsertUserGroupRequest { allowed_models_mode?: ListPolicyMode rate_limit?: number | null rate_limit_mode?: RateLimitPolicyMode + daily_usage_limit_usd?: number | null + daily_usage_limit_mode?: RateLimitPolicyMode } export interface UserGroupMember { @@ -222,6 +227,7 @@ export interface ApiKey { is_standalone: boolean // 是否为独立余额Key feature_settings?: FeatureSettings | null rate_limit?: number | null // 普通Key: 0 = 不限制，历史 null 视为跟随系统默认 + daily_usage_limit_usd?: number | null // 普通 Key: 0/null = 不增加 Key 级限制 concurrent_limit?: number | null // 普通Key: 0 = 不限制并发，历史 null 兼容 ip_rules?: string[] | null total_requests?: number // 总请求数 @@ -231,6 +237,7 @@ export interface ApiKey { export interface UpsertUserApiKeyRequest { name?: string rate_limit?: number | null + daily_usage_limit_usd?: number | null concurrent_limit?: number | null ip_rules?: string[] | null feature_settings?: FeatureSettings | null diff --git a/frontend/src/features/api-keys/components/StandaloneKeyFormDialog.vue b/frontend/src/features/api-keys/components/StandaloneKeyFormDialog.vue index 538f91a4f..56af70144 100644 --- a/frontend/src/features/api-keys/components/StandaloneKeyFormDialog.vue +++ b/frontend/src/features/api-keys/components/StandaloneKeyFormDialog.vue @@ -273,6 +273,36 @@ +
+ +
+
+ + 跟随系统 +
+ +
+
+
+
+ + +

+ {{ legacyT('正数会进一步收窄用户额度限制；0 或留空不增加 Key 级限制。') }} +

+
+
diff --git a/frontend/src/features/users/components/UserGroupAccessControlFields.vue b/frontend/src/features/users/components/UserGroupAccessControlFields.vue index df67d0168..660a44b4e 100644 --- a/frontend/src/features/users/components/UserGroupAccessControlFields.vue +++ b/frontend/src/features/users/components/UserGroupAccessControlFields.vue @@ -128,6 +128,33 @@
+ +
+ +
+
+ + + {{ legacyT(form.daily_usage_limit_mode === 'system' ? '系统默认' : '自定义') }} + +
+
+ +
+
+
@@ -184,4 +211,12 @@ function setSystemRateLimit(value: boolean): void { function updateRateLimit(value: string | number): void { updateForm({ rate_limit: parseNumberInput(value, { min: 0, max: 10000 }) }) } + +function setSystemDailyUsageLimit(value: boolean): void { + updateForm({ daily_usage_limit_mode: value ? 'system' : 'custom' }) +} + +function updateDailyUsageLimit(value: string | number): void { + updateForm({ daily_usage_limit_usd: parseNumberInput(value, { allowFloat: true, min: 0 }) }) +} diff --git a/frontend/src/features/users/components/UserGroupsDialog.vue b/frontend/src/features/users/components/UserGroupsDialog.vue index 51ee12dfc..1ed47f522 100644 --- a/frontend/src/features/users/components/UserGroupsDialog.vue +++ b/frontend/src/features/users/components/UserGroupsDialog.vue @@ -120,9 +120,9 @@ const USER_OPTIONS_CACHE_TTL_MS = 30 * 1000 let dialogUsersLoadedAt = 0 let dialogUsersLoadedVersion = -1 -const groupPolicyHelpText = '模型、供应商和端点会在多个用户组之间叠加授权；unrestricted 仍表示不限制，deny_all 只是不授予额外权限。速率限制按付费档位取更高额度，0 表示不限速；用户/API Key 自身限制仍会收窄最终权限。' +const groupPolicyHelpText = '模型、供应商和端点会在多个用户组之间叠加授权；unrestricted 仍表示不限制，deny_all 只是不授予额外权限。速率限制和额度限制均按组取更高额度，任一自定义组为 0 表示不限；Key 自身限制仍可收窄最终权限。' const groupPolicyHelpTextLocalized = computed(() => locale.value === 'en-US' - ? 'Models, providers, and endpoints accumulate across multiple user groups. unrestricted still means no restriction, while deny_all grants no extra permission. Rate limits take the higher quota by tier, and 0 means unlimited. User/API key limits still narrow the final permissions.' + ? 'Models, providers, and endpoints accumulate across user groups. RPM and usage limits take the highest custom group value, while any custom 0 means unlimited. Key limits may narrow the result.' : groupPolicyHelpText) const form = ref(createEmptyForm()) @@ -207,6 +207,8 @@ async function selectGroup(groupId: string): Promise { allowed_models: group.allowed_models ? [...group.allowed_models] : [], rate_limit_mode: normalizeRateMode(group.rate_limit_mode), rate_limit: group.rate_limit ?? undefined, + daily_usage_limit_mode: normalizeRateMode(group.daily_usage_limit_mode), + daily_usage_limit_usd: group.daily_usage_limit_usd ?? undefined, } try { const members = await usersStore.listUserGroupMembers(group.id) @@ -236,6 +238,8 @@ function createEmptyForm(): UserGroupFormState { allowed_models: [], rate_limit_mode: 'system', rate_limit: undefined, + daily_usage_limit_mode: 'system', + daily_usage_limit_usd: undefined, } } @@ -287,6 +291,10 @@ function buildPayload(): UpsertUserGroupRequest { rate_limit: form.value.rate_limit_mode === 'custom' ? (form.value.rate_limit ?? 0) : null, + daily_usage_limit_mode: form.value.daily_usage_limit_mode, + daily_usage_limit_usd: form.value.daily_usage_limit_mode === 'custom' + ? (form.value.daily_usage_limit_usd ?? 0) + : null, } } diff --git a/frontend/src/features/users/components/UserStatusBadges.vue b/frontend/src/features/users/components/UserStatusBadges.vue index 7d7938223..37ff6851c 100644 --- a/frontend/src/features/users/components/UserStatusBadges.vue +++ b/frontend/src/features/users/components/UserStatusBadges.vue @@ -21,6 +21,14 @@ > {{ legacyT(row.rateLimitLabel) }} + + {{ legacyT('日限额') }}: {{ legacyT(row.dailyUsageLimitLabel) }} + +
+ {{ legacyT('日限额:') }} + + {{ legacyT(row.dailyUsageLimitLabel) }} + + + {{ legacyT(row.dailyUsageLimitLabel) }} + +
diff --git a/frontend/src/features/users/components/user-management-types.ts b/frontend/src/features/users/components/user-management-types.ts index 81ada6176..a8d1db429 100644 --- a/frontend/src/features/users/components/user-management-types.ts +++ b/frontend/src/features/users/components/user-management-types.ts @@ -26,6 +26,8 @@ export interface UserGroupFormState { allowed_models: string[] rate_limit_mode: RateLimitPolicyMode rate_limit: number | undefined + daily_usage_limit_mode: RateLimitPolicyMode + daily_usage_limit_usd: number | undefined } export interface UserManagementRow { @@ -46,6 +48,9 @@ export interface UserManagementRow { rateLimitLabel: string rateLimitSource: string rateLimitAsBadge: boolean + dailyUsageLimitLabel: string + dailyUsageLimitSource: string + dailyUsageLimitAsBadge: boolean createdAtLabel: string statusLabel: string statusVariant: BadgeVariant diff --git a/frontend/src/i18n/messages.ts b/frontend/src/i18n/messages.ts index b5b414f15..59b301014 100644 --- a/frontend/src/i18n/messages.ts +++ b/frontend/src/i18n/messages.ts @@ -1155,6 +1155,15 @@ const legacyExactEnglishMessages: Record = { '密钥名称': 'Key name', '例如：生产环境 Key': 'Example: Production key', '速率限制 (请求/分钟)': 'Rate limit (requests/min)', + '额度限制 (美元/日)': 'Usage limit (USD/day)', + '日限额': 'Daily limit', + '日限额:': 'Daily:', + '留空不增加 Key 级限制': 'Leave empty to add no key-level limit', + '0 = 明确不限制': '0 = explicitly unlimited', + '无日上限': 'No daily limit', + '正数会进一步收窄用户额度限制；0 或留空不增加 Key 级限制。': 'A positive value further narrows the user usage limit. Zero or empty adds no key-level limit.', + '正数会进一步收窄账户额度限制；0 或留空不增加 Key 级限制。': 'A positive value further narrows the account usage limit. Zero or empty adds no key-level limit.', + '0 表示默认不限制；未单独配置的用户和独立 Key 会跟随这里': 'Zero means unlimited by default; users and standalone keys without an individual setting inherit this value.', '留空不限': 'Leave empty for unlimited', '留空表示不限制': 'Leave empty for no limit', '留空表示保持当前值，填 0 表示不限并发': 'Leave empty to keep the current value; enter 0 for unlimited concurrency', @@ -1249,6 +1258,8 @@ const legacyExactEnglishMessages: Record = { '每日额度': 'Daily quota', '会员权益': 'Membership benefits', '继承自分组：': 'Inherited from group: ', + '组合分组：': 'Combined groups: ', + '组合分组': 'Combined groups', '用户额外限制与分组叠加：': 'User extra limit combined with groups: ', '用户额外限制与分组叠加': 'User extra limit combined with groups', '用户单独配置': 'User-specific configuration', @@ -2966,6 +2977,8 @@ const legacyFallbackTokens: Array<[string, string]> = [ ['使用提醒', 'usage reminders'], ['默认用户初始赠款(美元)', 'Default user initial grant (USD)'], ['默认速率限制 (请求/分钟)', 'Default rate limit (requests/min)'], + ['默认额度限制（美元/日）', 'Default usage limit (USD/day)'], + ['额度限制', 'Usage limit'], ['弱密码 - 至少 6 个字符', 'Weak password - at least 6 characters'], ['中等密码 - 至少 8 位，含字母和数字', 'Medium password - at least 8 characters with letters and numbers'], ['强密码 - 至少 8 位，含大小写字母、数字和特殊字符', 'Strong password - at least 8 characters with uppercase, lowercase, numbers, and symbols'], diff --git a/frontend/src/views/admin/ApiKeys.vue b/frontend/src/views/admin/ApiKeys.vue index 31718ba18..c192e995d 100644 --- a/frontend/src/views/admin/ApiKeys.vue +++ b/frontend/src/views/admin/ApiKeys.vue @@ -1202,6 +1202,7 @@ function editApiKey(apiKey: AdminApiKey) { unlimited_balance: isApiKeyUnlimited(apiKey), expires_at: expiresAt, rate_limit: apiKey.rate_limit ?? undefined, + daily_usage_limit_usd: apiKey.daily_usage_limit_usd ?? undefined, concurrent_limit: apiKey.concurrent_limit ?? undefined, auto_delete_on_expiry: apiKey.auto_delete_on_expiry || false, allowed_providers: apiKey.allowed_providers == null ? null : [...apiKey.allowed_providers], @@ -1413,6 +1414,7 @@ async function handleKeyFormSubmit(data: StandaloneKeyFormData) { name: data.name || undefined, unlimited_balance: Boolean(data.unlimited_balance), rate_limit: data.rate_limit ?? null, // undefined = 跟随系统默认，显式传 null + daily_usage_limit_usd: data.daily_usage_limit_usd ?? null, concurrent_limit: data.concurrent_limit ?? null, expires_at: serializeExpiryDate(data.expires_at), auto_delete_on_expiry: data.auto_delete_on_expiry, @@ -1445,6 +1447,7 @@ async function handleKeyFormSubmit(data: StandaloneKeyFormData) { name: data.name || undefined, initial_balance_usd: isUnlimited ? null : (data.initial_balance_usd as number), rate_limit: data.rate_limit ?? null, // undefined = 跟随系统默认，显式传 null + daily_usage_limit_usd: data.daily_usage_limit_usd ?? null, concurrent_limit: data.concurrent_limit ?? null, expires_at: serializeExpiryDate(data.expires_at), auto_delete_on_expiry: data.auto_delete_on_expiry, diff --git a/frontend/src/views/admin/SystemSettings.vue b/frontend/src/views/admin/SystemSettings.vue index c87e96aa5..40bfbb074 100644 --- a/frontend/src/views/admin/SystemSettings.vue +++ b/frontend/src/views/admin/SystemSettings.vue @@ -63,6 +63,7 @@ id="section-basic" :default-user-initial-gift-usd="systemConfig.default_user_initial_gift_usd" :rate-limit-per-minute="systemConfig.rate_limit_per_minute" + :daily-usage-limit-usd="systemConfig.daily_usage_limit_usd" :enable-registration="systemConfig.enable_registration" :password-policy-level="systemConfig.password_policy_level" :turnstile-enabled="systemConfig.turnstile_enabled" @@ -89,6 +90,7 @@ @save="saveBasicConfig" @update:default-user-initial-gift-usd="systemConfig.default_user_initial_gift_usd = $event" @update:rate-limit-per-minute="systemConfig.rate_limit_per_minute = $event" + @update:daily-usage-limit-usd="systemConfig.daily_usage_limit_usd = $event" @update:enable-registration="systemConfig.enable_registration = $event" @update:password-policy-level="systemConfig.password_policy_level = $event" @update:turnstile-enabled="systemConfig.turnstile_enabled = $event" diff --git a/frontend/src/views/admin/Users.vue b/frontend/src/views/admin/Users.vue index 2c3f68cf6..0f9671d1d 100644 --- a/frontend/src/views/admin/Users.vue +++ b/frontend/src/views/admin/Users.vue @@ -290,6 +290,7 @@ const editingUserApiKey = ref(null) const userApiKeyForm = ref({ name: '', rate_limit: undefined, + daily_usage_limit_usd: undefined, concurrent_limit: undefined, ip_rules_text: '', chat_pii_redaction_mode: 'inherit', @@ -402,6 +403,9 @@ const userRows = computed(() => rateLimitLabel: formatRateLimitInheritable(user.rate_limit), rateLimitSource: formatUserEffectiveRateLimitSource(user), rateLimitAsBadge: isRateLimitInherited(user.rate_limit) || isRateLimitUnlimited(user.rate_limit), + dailyUsageLimitLabel: formatUserEffectiveDailyUsageLimit(user), + dailyUsageLimitSource: formatUserEffectiveDailyUsageLimitSource(user), + dailyUsageLimitAsBadge: isUserDailyUsageLimitInherited(user) || isUserDailyUsageLimitUnlimited(user), createdAtLabel: formatDate(user.created_at), statusLabel: legacyT(user.is_active ? '活跃' : '禁用'), statusVariant: user.is_active ? 'success' : 'destructive', @@ -685,6 +689,38 @@ function formatUserEffectiveRateLimitSource(user: User): string { return legacyT('系统默认') } +function formatUserEffectiveDailyUsageLimit(user: User): string { + const policy = user.effective_policy?.daily_usage_limit_usd + if (!policy) return '-' + if (policy.mode === 'system') return legacyT('跟随系统') + const value = policy.value + if (typeof value !== 'number') return '-' + if (value === 0) return legacyT('无日上限') + return locale.value === 'en-US' ? `$${value.toFixed(2)}/day` : `$${value.toFixed(2)}/日` +} + +function isUserDailyUsageLimitInherited(user: User): boolean { + return user.effective_policy?.daily_usage_limit_usd?.mode === 'system' +} + +function isUserDailyUsageLimitUnlimited(user: User): boolean { + return user.effective_policy?.daily_usage_limit_usd?.mode !== 'system' + && user.effective_policy?.daily_usage_limit_usd?.value === 0 +} + +function formatUserEffectiveDailyUsageLimitSource(user: User): string { + const source = user.effective_policy?.daily_usage_limit_usd + if (!source) return '' + if (source.source === 'group' && source.group_name) { + return `${legacyT('继承自分组：')}${source.group_name}` + } + if (source.source === 'combined') { + const groupNames = Array.isArray(source.group_names) ? source.group_names.join(locale.value === 'en-US' ? ', ' : '、') : '' + return groupNames ? `${legacyT('组合分组：')}${groupNames}` : legacyT('组合分组') + } + return legacyT('系统默认') +} + function isNegativeWalletValue(value: number | null): boolean { return typeof value === 'number' && value < 0 } @@ -910,6 +946,7 @@ function openCreateUserApiKeyDialog() { userApiKeyForm.value = { name: `Key-${new Date().toISOString().split('T')[0]}`, rate_limit: undefined, + daily_usage_limit_usd: undefined, concurrent_limit: undefined, ip_rules_text: '', chat_pii_redaction_mode: redactionFeature.mode, @@ -929,6 +966,7 @@ function openEditUserApiKeyDialog(apiKey: ApiKey) { userApiKeyForm.value = { name: apiKey.name || '', rate_limit: apiKey.rate_limit ?? undefined, + daily_usage_limit_usd: apiKey.daily_usage_limit_usd ?? undefined, concurrent_limit: apiKey.concurrent_limit ?? undefined, ip_rules_text: apiKey.ip_rules?.join(', ') ?? '', chat_pii_redaction_mode: redactionFeature.mode, @@ -948,6 +986,7 @@ function closeUserApiKeyFormDialog() { userApiKeyForm.value = { name: '', rate_limit: undefined, + daily_usage_limit_usd: undefined, concurrent_limit: undefined, ip_rules_text: '', chat_pii_redaction_mode: 'inherit', @@ -989,6 +1028,7 @@ async function submitUserApiKeyForm() { await usersStore.updateApiKey(targetUserId, editingApiKey.id, { name: form.name, rate_limit: form.rate_limit ?? 0, + daily_usage_limit_usd: form.daily_usage_limit_usd ?? 0, concurrent_limit: form.concurrent_limit, ip_rules: ipRules, ...featureSettingsPatch, @@ -999,6 +1039,7 @@ async function submitUserApiKeyForm() { const response = await usersStore.createApiKey(targetUserId, { name: form.name, rate_limit: form.rate_limit ?? 0, + daily_usage_limit_usd: form.daily_usage_limit_usd ?? 0, concurrent_limit: form.concurrent_limit, ip_rules: ipRules, ...featureSettingsPatch, diff --git a/frontend/src/views/admin/system-settings/BasicConfigSection.vue b/frontend/src/views/admin/system-settings/BasicConfigSection.vue index f81177953..87531f550 100644 --- a/frontend/src/views/admin/system-settings/BasicConfigSection.vue +++ b/frontend/src/views/admin/system-settings/BasicConfigSection.vue @@ -54,6 +54,28 @@

+
+ + +

+ 0 表示默认不限制；未单独配置的用户和独立 Key 会跟随这里 +

+
+
@@ -477,6 +484,26 @@

+
+ + +

+ 正数会进一步收窄账户额度限制；0 或留空不增加 Key 级限制。 +

+
+