use std::collections::BTreeMap; use std::sync::Arc; use std::time::Duration; use aether_data_contracts::repository::billing::{ nonnegative_usd_to_usage_policy_cost_units, parse_usage_policy_entitlements, UsagePolicyMetric, UsagePolicyWindow, UserPlanEntitlementRecord, USAGE_POLICY_COST_UNITS_PER_USD, }; use aether_data_contracts::repository::settlement::{ ReconcileUsagePolicyCostInput, ReserveUsagePolicyCostInput, ReserveUsagePolicyCostOutcome, ReserveUsagePolicyRequestInput, ReserveUsagePolicyRequestOutcome, UsagePolicyCostReservationState, UsagePolicyCostWindow, UsagePolicyRequestWindow, }; use aether_runtime::AdmissionPermit; use aether_runtime_state::{ RuntimeSemaphoreConfig, RuntimeSemaphoreError, UsageLimitCheck, UsageLimitInput, UsageLimitReleaseInput, UsageLimitRule, }; use chrono::{Datelike, TimeZone, Utc, Weekday}; use tracing::warn; use crate::control::GatewayControlDecision; use crate::{AppState, GatewayError}; const POLICY_CACHE_TTL: Duration = Duration::from_secs(5); const CONCURRENCY_GATE: &str = "plan_usage_concurrency"; const DEFAULT_CALENDAR_TIMEZONE: &str = "Asia/Shanghai"; const COST_RESERVATION_TTL_SECS: u64 = 24 * 60 * 60; const COST_RESERVATION_SAFE_HISTORY_SECS: u64 = 32 * 24 * 60 * 60; #[derive(Debug, Clone)] pub(crate) struct PlanUsagePolicySnapshot { pub(crate) admitted_at_unix_secs: u64, subject_id: Arc, policy: Arc, } impl PlanUsagePolicySnapshot { fn for_admission( subject_id: &str, policy: EffectivePlanUsagePolicy, admitted_at_unix_secs: u64, ) -> Option { if policy.cost_rules.is_empty() { return None; } Some(Self { admitted_at_unix_secs, subject_id: subject_id.to_string().into(), policy: Arc::new(policy), }) } pub(crate) fn new_reservation_context(&self) -> PlanUsageReservationContext { PlanUsageReservationContext { policy_snapshot: self.clone(), token: uuid::Uuid::new_v4().to_string().into(), } } pub(crate) fn subject_id(&self) -> &str { self.subject_id.as_ref() } pub(crate) fn policy(&self) -> &EffectivePlanUsagePolicy { self.policy.as_ref() } } #[derive(Debug, Clone)] pub(crate) struct PlanUsageReservationContext { policy_snapshot: PlanUsagePolicySnapshot, token: Arc, } impl PlanUsageReservationContext { #[cfg(test)] pub(crate) fn for_test( subject_id: impl Into>, token: impl Into>, admitted_at_unix_secs: u64, policy: EffectivePlanUsagePolicy, ) -> Self { Self { policy_snapshot: PlanUsagePolicySnapshot { admitted_at_unix_secs, subject_id: subject_id.into(), policy: Arc::new(policy), }, token: token.into(), } } pub(crate) fn subject_id(&self) -> &str { self.policy_snapshot.subject_id() } pub(crate) fn token(&self) -> &str { self.token.as_ref() } pub(crate) fn policy(&self) -> &EffectivePlanUsagePolicy { self.policy_snapshot.policy() } pub(crate) const fn admitted_at_unix_secs(&self) -> u64 { self.policy_snapshot.admitted_at_unix_secs } } #[derive(Debug)] pub(crate) struct HttpPlanUsageAdmission { pub(crate) permit: Option, pub(crate) reservation_context: Option, } #[derive(Debug)] pub(crate) struct PlanUsageAdmission { pub(crate) permit: Option, pub(crate) policy_snapshot: Option, } #[derive(Debug, Clone, PartialEq, Default)] pub(crate) struct EffectivePlanUsagePolicy { request_rules: Vec, cost_rules: Vec, concurrency_limit: Option, valid_until_unix_secs: Option, } #[derive(Debug, Clone, PartialEq)] struct EffectiveRequestRule { identity: String, window: UsagePolicyWindow, limit: u64, entitlement_period: Option<(u64, u64)>, } #[derive(Debug, Clone, PartialEq)] struct EffectiveCostRule { identity: String, window: UsagePolicyWindow, limit_cost_units: u64, entitlement_period: Option<(u64, u64)>, } #[derive(Debug, Clone, PartialEq, Eq)] struct RuntimeRequestRule { key: String, limit: u64, window_seconds: u64, retention_seconds: u64, retry_after: u64, window: &'static str, } #[derive(Debug, Clone, PartialEq, Eq)] struct DurableRequestRule { window: UsagePolicyRequestWindow, retry_after: u64, label: &'static str, influence_ends_at_unix_secs: u64, } #[derive(Debug, Clone, PartialEq, Eq)] struct RuntimeCostRule { window: UsagePolicyCostWindow, retry_after: u64, label: &'static str, influence_ends_at_unix_secs: u64, } #[derive(Debug, Clone, PartialEq)] pub(crate) struct PlanUsagePolicyRejection { pub(crate) metric: &'static str, pub(crate) limit: f64, pub(crate) retry_after: u64, pub(crate) window: &'static str, } #[derive(Debug, Clone, PartialEq)] pub(crate) enum PlanUsageCostReservationOutcome { NotRequired, Reserved, Rejected(PlanUsagePolicyRejection), } pub(crate) async fn reserve_admitted_http_plan_usage_policy_cost( state: &AppState, decision: &GatewayControlDecision, plan: &aether_contracts::ExecutionPlan, report_context: Option<&serde_json::Value>, reservation: Option<&PlanUsageReservationContext>, ) -> Result { let Some(auth) = plan_usage_auth_context(decision) else { return Ok(PlanUsageCostReservationOutcome::NotRequired); }; let Some(reservation) = reservation else { return Ok(PlanUsageCostReservationOutcome::NotRequired); }; if reservation.subject_id() != auth.user_id { return Err(GatewayError::Internal( "plan usage reservation subject does not match the admitted request".to_string(), )); } reserve_plan_usage_policy_cost_with_policy( state, decision, plan, report_context, reservation.policy(), reservation.admitted_at_unix_secs(), reservation.token(), ) .await } pub(crate) async fn reserve_admitted_plan_usage_policy_cost( state: &AppState, decision: &GatewayControlDecision, plan: &aether_contracts::ExecutionPlan, report_context: Option<&serde_json::Value>, snapshot: Option<&PlanUsagePolicySnapshot>, reservation_token: &str, ) -> Result { let Some(auth) = plan_usage_auth_context(decision) else { return Ok(PlanUsageCostReservationOutcome::NotRequired); }; let Some(snapshot) = snapshot else { return Ok(PlanUsageCostReservationOutcome::NotRequired); }; if snapshot.subject_id() != auth.user_id { return Err(GatewayError::Internal( "plan usage reservation subject does not match the admitted request".to_string(), )); } reserve_plan_usage_policy_cost_with_policy( state, decision, plan, report_context, snapshot.policy(), snapshot.admitted_at_unix_secs, reservation_token, ) .await } fn plan_usage_auth_context( decision: &GatewayControlDecision, ) -> Option<&crate::control::GatewayControlAuthContext> { decision.auth_context.as_ref().filter(|auth| { decision.route_class.as_deref() == Some("ai_public") && !auth.admin_bypass_limits && !auth.api_key_is_standalone }) } #[allow(clippy::too_many_arguments)] async fn reserve_plan_usage_policy_cost_with_policy( state: &AppState, decision: &GatewayControlDecision, plan: &aether_contracts::ExecutionPlan, report_context: Option<&serde_json::Value>, policy: &EffectivePlanUsagePolicy, admitted_at_unix_secs: u64, reservation_token: &str, ) -> Result { let Some(auth) = plan_usage_auth_context(decision) else { return Ok(PlanUsageCostReservationOutcome::NotRequired); }; if policy.cost_rules.is_empty() { return Ok(PlanUsageCostReservationOutcome::NotRequired); } ensure_cost_policy_usage_runtime_enabled(state.usage_runtime.is_enabled())?; let estimated_cost_usd = crate::control::estimate_execution_plan_cost_upper_bound_usd(state, plan, report_context) .await? .ok_or_else(|| { GatewayError::Internal( "a hard plan cost limit is active, but this request cost cannot be estimated safely" .to_string(), ) })?; let reserved_cost_units = nonnegative_usd_to_usage_policy_cost_units(estimated_cost_usd) .ok_or_else(|| { GatewayError::Internal( "estimated request cost is outside the supported plan limit range".to_string(), ) })?; let runtime_rules = policy .cost_rules .iter() .map(|rule| runtime_cost_rule(rule, admitted_at_unix_secs)) .collect::, _>>()?; let reservation_expires_at_unix_secs = admitted_at_unix_secs.saturating_add(COST_RESERVATION_TTL_SECS); let retain_until_unix_secs = runtime_rules .iter() .map(|rule| rule.influence_ends_at_unix_secs) .chain(std::iter::once(reservation_expires_at_unix_secs)) .chain(std::iter::once( admitted_at_unix_secs.saturating_add(COST_RESERVATION_SAFE_HISTORY_SECS), )) .max() .unwrap_or(reservation_expires_at_unix_secs); let outcome = state .data .reserve_usage_policy_cost(ReserveUsagePolicyCostInput { request_id: plan.request_id.clone(), reservation_token: reservation_token.to_string(), subject_id: auth.user_id.clone(), admitted_at_unix_secs, reserved_cost_units, reservation_expires_at_unix_secs, retain_until_unix_secs, windows: runtime_rules .iter() .map(|rule| rule.window.clone()) .collect(), }) .await .map_err(|error| GatewayError::Internal(error.to_string()))? .ok_or_else(|| { GatewayError::Internal( "plan cost limits require a settlement write repository".to_string(), ) })?; match outcome { ReserveUsagePolicyCostOutcome::Allowed { .. } => { Ok(PlanUsageCostReservationOutcome::Reserved) } ReserveUsagePolicyCostOutcome::Rejected { window_index, limit_cost_units, .. } => { let rule = runtime_rules.get(window_index).ok_or_else(|| { GatewayError::Internal( "usage policy cost reservation returned an invalid window index".to_string(), ) })?; Ok(PlanUsageCostReservationOutcome::Rejected( PlanUsagePolicyRejection { metric: "actual_cost_usd", limit: limit_cost_units as f64 / USAGE_POLICY_COST_UNITS_PER_USD as f64, retry_after: rule.retry_after, window: rule.label, }, )) } ReserveUsagePolicyCostOutcome::AlreadyTerminal { state } => { Err(GatewayError::Internal(format!( "plan cost reservation for request {} is already {}", plan.request_id, state.as_str() ))) } ReserveUsagePolicyCostOutcome::Conflict => Err(GatewayError::Internal( "plan cost reservation request identity conflict".to_string(), )), } } pub(crate) async fn release_plan_usage_policy_cost( state: &AppState, decision: &GatewayControlDecision, plan: &aether_contracts::ExecutionPlan, reservation_token: &str, finalized_at_unix_secs: u64, ) -> Result<(), GatewayError> { let Some(auth) = decision .auth_context .as_ref() .filter(|_| decision.route_class.as_deref() == Some("ai_public")) else { return Ok(()); }; if auth.admin_bypass_limits || auth.api_key_is_standalone { return Ok(()); } state .data .reconcile_usage_policy_cost(build_usage_policy_cost_release_input( plan.request_id.as_str(), auth.user_id.as_str(), reservation_token, finalized_at_unix_secs, )) .await .map_err(|error| GatewayError::Internal(error.to_string()))?; Ok(()) } fn build_usage_policy_cost_release_input( request_id: &str, subject_id: &str, reservation_token: &str, finalized_at_unix_secs: u64, ) -> ReconcileUsagePolicyCostInput { ReconcileUsagePolicyCostInput { request_id: request_id.to_string(), subject_id: subject_id.to_string(), reservation_token: reservation_token.to_string(), actual_cost_units: 0, terminal_state: UsagePolicyCostReservationState::Released, finalized_at_unix_secs, } } fn ensure_cost_policy_usage_runtime_enabled(enabled: bool) -> Result<(), GatewayError> { if enabled { Ok(()) } else { Err(GatewayError::Internal( "plan cost limits require the usage runtime to be enabled for terminal reconciliation" .to_string(), )) } } #[derive(Debug)] pub(crate) enum PlanUsageAdmissionError { Rejected(PlanUsagePolicyRejection), Runtime(RuntimeSemaphoreError), Gateway(GatewayError), } impl From for PlanUsageAdmissionError { fn from(value: GatewayError) -> Self { Self::Gateway(value) } } pub(crate) async fn check_and_acquire_plan_usage_policy_admission( state: &AppState, decision: Option<&GatewayControlDecision>, event_id: &str, now_unix_ms: u64, ) -> Result { let Some(auth) = decision .filter(|decision| decision.route_class.as_deref() == Some("ai_public")) .and_then(|decision| decision.auth_context.as_ref()) else { return Ok(PlanUsageAdmission { permit: None, policy_snapshot: None, }); }; if auth.admin_bypass_limits || auth.api_key_is_standalone { return Ok(PlanUsageAdmission { permit: None, policy_snapshot: None, }); } let admitted_at_unix_secs = now_unix_ms / 1_000; let policy = load_effective_policy(state, &auth.user_id, admitted_at_unix_secs).await?; let permit = check_and_acquire_compiled_plan_usage_policy( state, &auth.user_id, &policy, event_id, now_unix_ms, ) .await?; let policy_snapshot = PlanUsagePolicySnapshot::for_admission(&auth.user_id, policy, admitted_at_unix_secs); Ok(PlanUsageAdmission { permit, policy_snapshot, }) } async fn check_and_acquire_compiled_plan_usage_policy( state: &AppState, subject_id: &str, policy: &EffectivePlanUsagePolicy, event_id: &str, now_unix_ms: u64, ) -> Result, PlanUsageAdmissionError> { if policy.request_rules.is_empty() && policy.concurrency_limit.is_none() { return Ok(None); } let now_unix_secs = now_unix_ms / 1_000; let plan_permit = if let Some(limit) = policy.concurrency_limit { let limit = usize::try_from(limit).map_err(|_| { PlanUsageAdmissionError::Gateway(GatewayError::Internal( "plan usage concurrency limit exceeds platform capacity".to_string(), )) })?; let gate = state .runtime_state .keyed_semaphore( CONCURRENCY_GATE, format!("admission:{CONCURRENCY_GATE}:user:{{{subject_id}}}"), limit, RuntimeSemaphoreConfig::default(), ) .map_err(PlanUsageAdmissionError::Runtime)?; Some( gate.try_acquire() .await .map_err(PlanUsageAdmissionError::Runtime)?, ) } else { None }; let (runtime_request_rules, durable_request_rules): (Vec<_>, Vec<_>) = policy .request_rules .iter() .partition(|rule| request_rule_uses_runtime_state(rule)); let runtime_rules = runtime_request_rules .iter() .map(|rule| runtime_request_rule(subject_id, rule, now_unix_secs)) .collect::, _>>()?; let runtime_inputs = runtime_rules .iter() .map(|rule| UsageLimitRule { key: rule.key.as_str(), limit: rule.limit, window_seconds: rule.window_seconds, retention_seconds: rule.retention_seconds, }) .collect::>(); // Consume the short-lived runtime windows first. If the durable admission below rejects or // errors, remove this event from every short window. This ordering bounds a process-crash // compensation gap to the QPS/RPM retention (at most 60 seconds) instead of leaking a // calendar/subscription-period admission for weeks or months. if !runtime_inputs.is_empty() { match state .runtime_state .check_and_consume_usage_limits(UsageLimitInput { rules: &runtime_inputs, event_id, now_unix_ms, }) .await .map_err(|error| { PlanUsageAdmissionError::Gateway(GatewayError::Internal(error.to_string())) })? { UsageLimitCheck::Allowed => {} UsageLimitCheck::Rejected { rule_index, limit, retry_after, } => { let rule = runtime_rules.get(rule_index).ok_or_else(|| { PlanUsageAdmissionError::Gateway(GatewayError::Internal( "usage policy runtime returned an invalid rule index".to_string(), )) })?; return Err(PlanUsageAdmissionError::Rejected( PlanUsagePolicyRejection { metric: "request_count", limit: limit as f64, retry_after: retry_after.min(rule.retry_after).max(1), window: rule.window, }, )); } } } let durable_rules = durable_request_rules .iter() .map(|rule| durable_request_rule(rule, now_unix_secs)) .collect::, _>>()?; if !durable_rules.is_empty() { let retain_until_unix_secs = durable_rules .iter() .map(|rule| rule.influence_ends_at_unix_secs) .chain(std::iter::once( now_unix_secs.saturating_add(COST_RESERVATION_SAFE_HISTORY_SECS), )) .max() .unwrap_or(now_unix_secs.saturating_add(COST_RESERVATION_SAFE_HISTORY_SECS)); let outcome = state .data .reserve_usage_policy_request(ReserveUsagePolicyRequestInput { request_id: event_id.to_string(), subject_id: subject_id.to_string(), event_token: event_id.to_string(), admitted_at_unix_secs: now_unix_secs, retain_until_unix_secs, windows: durable_rules .iter() .map(|rule| rule.window.clone()) .collect(), }) .await; let outcome = match outcome { Ok(Some(outcome)) => outcome, Ok(None) => { release_runtime_usage_limits_best_effort( state, &runtime_inputs, event_id, "settlement_writer_unavailable", ) .await; return Err(PlanUsageAdmissionError::Gateway(GatewayError::Internal( "long-window plan request limits require a settlement write repository" .to_string(), ))); } Err(error) => { release_runtime_usage_limits_best_effort( state, &runtime_inputs, event_id, "durable_admission_error", ) .await; return Err(PlanUsageAdmissionError::Gateway(GatewayError::Internal( error.to_string(), ))); } }; match outcome { ReserveUsagePolicyRequestOutcome::Allowed => {} ReserveUsagePolicyRequestOutcome::Rejected { window_index, limit_requests, .. } => { release_runtime_usage_limits_best_effort( state, &runtime_inputs, event_id, "durable_admission_rejected", ) .await; let rule = durable_rules.get(window_index).ok_or_else(|| { PlanUsageAdmissionError::Gateway(GatewayError::Internal( "usage policy request ledger returned an invalid rule index".to_string(), )) })?; return Err(PlanUsageAdmissionError::Rejected( PlanUsagePolicyRejection { metric: "request_count", limit: limit_requests as f64, retry_after: rule.retry_after, window: rule.label, }, )); } ReserveUsagePolicyRequestOutcome::AlreadyReleased => { release_runtime_usage_limits_best_effort( state, &runtime_inputs, event_id, "durable_admission_already_released", ) .await; return Err(PlanUsageAdmissionError::Gateway(GatewayError::Internal( "plan request admission token is already released".to_string(), ))); } ReserveUsagePolicyRequestOutcome::Conflict => { release_runtime_usage_limits_best_effort( state, &runtime_inputs, event_id, "durable_admission_conflict", ) .await; return Err(PlanUsageAdmissionError::Gateway(GatewayError::Internal( "plan request admission identity conflict".to_string(), ))); } } } Ok(AdmissionPermit::from_parts(None, plan_permit)) } pub(crate) async fn check_and_acquire_http_plan_usage_policy( state: &AppState, decision: Option<&GatewayControlDecision>, event_id: &str, now_unix_ms: u64, ) -> Result { let Some(auth) = decision .filter(|decision| decision.route_class.as_deref() == Some("ai_public")) .and_then(|decision| decision.auth_context.as_ref()) else { return Ok(HttpPlanUsageAdmission { permit: None, reservation_context: None, }); }; if auth.admin_bypass_limits || auth.api_key_is_standalone { return Ok(HttpPlanUsageAdmission { permit: None, reservation_context: None, }); } let admitted_at_unix_secs = now_unix_ms / 1_000; let policy = load_effective_policy(state, &auth.user_id, admitted_at_unix_secs).await?; let permit = check_and_acquire_compiled_plan_usage_policy( state, &auth.user_id, &policy, event_id, now_unix_ms, ) .await?; let reservation_context = PlanUsagePolicySnapshot::for_admission(&auth.user_id, policy, admitted_at_unix_secs) .map(|snapshot| snapshot.new_reservation_context()); Ok(HttpPlanUsageAdmission { permit, reservation_context, }) } fn request_rule_uses_runtime_state(rule: &&EffectiveRequestRule) -> bool { matches!(rule.window, UsagePolicyWindow::Rolling { seconds } if seconds <= 60) } async fn release_runtime_usage_limits_best_effort( state: &AppState, rules: &[UsageLimitRule<'_>], event_id: &str, reason: &'static str, ) { if rules.is_empty() { return; } if let Err(error) = state .runtime_state .release_usage_limits(UsageLimitReleaseInput { rules, event_id }) .await { warn!( event_name = "plan_usage_runtime_compensation_failed", log_type = "ops", event_id, reason, error = %error, "gateway failed to compensate short-window plan usage after durable admission failed" ); } } async fn load_effective_policy( state: &AppState, user_id: &str, now_unix_secs: u64, ) -> Result { let cache_key = user_id.to_string(); let policy = state .auth_plan_usage_policy_cache .get_or_load(cache_key.clone(), POLICY_CACHE_TTL, || async move { let _permit = state.acquire_auth_snapshot_load_gate().await?; let entitlements = state .list_user_plan_entitlements(user_id) .await? .unwrap_or_default(); compile_effective_policy(&entitlements, now_unix_secs).map(Some) }) .await .map(|policy| policy.unwrap_or_default())?; if policy .valid_until_unix_secs .is_some_and(|valid_until| valid_until <= now_unix_secs) { state.auth_plan_usage_policy_cache.clear(); let _permit = state.acquire_auth_snapshot_load_gate().await?; let entitlements = state .list_user_plan_entitlements(user_id) .await? .unwrap_or_default(); let policy = compile_effective_policy(&entitlements, now_unix_secs)?; state.auth_plan_usage_policy_cache.insert( cache_key, Some(policy.clone()), POLICY_CACHE_TTL, ); return Ok(policy); } Ok(policy) } fn compile_effective_policy( entitlements: &[UserPlanEntitlementRecord], now_unix_secs: u64, ) -> Result { let mut request_rules = BTreeMap::::new(); let mut cost_rules = BTreeMap::::new(); let mut concurrency_limit = None::; let mut valid_until_unix_secs = None::; for entitlement in entitlements.iter().filter(|entitlement| { entitlement.status == "active" && entitlement.starts_at_unix_secs <= now_unix_secs && entitlement.expires_at_unix_secs > now_unix_secs }) { valid_until_unix_secs = Some( valid_until_unix_secs.map_or(entitlement.expires_at_unix_secs, |current| { current.min(entitlement.expires_at_unix_secs) }), ); let policies = parse_usage_policy_entitlements(&entitlement.entitlements_snapshot) .map_err(|error| GatewayError::Internal(error.to_string()))?; for policy in policies { for rule in policy.rules { match rule.metric { UsagePolicyMetric::Concurrency => { let limit = rule.request_limit().ok_or_else(|| { GatewayError::Internal( "validated concurrency usage policy lost its integer limit" .to_string(), ) })?; concurrency_limit = Some(concurrency_limit.map_or(limit, |current| current.min(limit))); } UsagePolicyMetric::RequestCount => { let limit = rule.request_limit().ok_or_else(|| { GatewayError::Internal( "validated request-count usage policy lost its integer limit" .to_string(), ) })?; let (identity, entitlement_period) = match &rule.window { UsagePolicyWindow::SubscriptionPeriod => ( format!("subscription:{}", entitlement.id), Some(( entitlement.starts_at_unix_secs, entitlement.expires_at_unix_secs, )), ), window => (window_identity(window), None), }; let candidate = EffectiveRequestRule { identity: identity.clone(), window: rule.window, limit, entitlement_period, }; request_rules .entry(identity) .and_modify(|current| { if candidate.limit < current.limit { *current = candidate.clone(); } }) .or_insert(candidate); } UsagePolicyMetric::ActualCostUsd => { let limit_cost_units = rule.cost_limit_units().ok_or_else(|| { GatewayError::Internal( "validated cost usage policy lost its fixed-point limit" .to_string(), ) })?; let (identity, entitlement_period) = match &rule.window { UsagePolicyWindow::SubscriptionPeriod => ( format!("subscription:{}", entitlement.id), Some(( entitlement.starts_at_unix_secs, entitlement.expires_at_unix_secs, )), ), window => (window_identity(window), None), }; let candidate = EffectiveCostRule { identity: identity.clone(), window: rule.window, limit_cost_units, entitlement_period, }; cost_rules .entry(identity) .and_modify(|current| { if candidate.limit_cost_units < current.limit_cost_units { *current = candidate.clone(); } }) .or_insert(candidate); } } } } } Ok(EffectivePlanUsagePolicy { request_rules: request_rules.into_values().collect(), cost_rules: cost_rules.into_values().collect(), concurrency_limit, valid_until_unix_secs, }) } fn window_identity(window: &UsagePolicyWindow) -> String { match window { UsagePolicyWindow::Rolling { seconds } => format!("rolling:{seconds}"), UsagePolicyWindow::CalendarDay { timezone } => { format!( "calendar_day:{}", effective_timezone_name(timezone.as_deref()) ) } UsagePolicyWindow::CalendarWeek { timezone, week_start, } => format!( "calendar_week:{}:{week_start}", effective_timezone_name(timezone.as_deref()) ), UsagePolicyWindow::CalendarMonth { timezone } => { format!( "calendar_month:{}", effective_timezone_name(timezone.as_deref()) ) } UsagePolicyWindow::SubscriptionPeriod => "subscription_period".to_string(), UsagePolicyWindow::Concurrent => "concurrent".to_string(), } } fn runtime_request_rule( user_id: &str, rule: &EffectiveRequestRule, now_unix_secs: u64, ) -> Result { let (bucket, window_seconds, retry_after, window) = match &rule.window { UsagePolicyWindow::Rolling { seconds } => ( "sliding".to_string(), *seconds, *seconds, rolling_label(*seconds), ), UsagePolicyWindow::CalendarDay { timezone } => { calendar_bucket(now_unix_secs, timezone.as_deref(), CalendarWindow::Day, 1)? } UsagePolicyWindow::CalendarWeek { timezone, week_start, } => calendar_bucket( now_unix_secs, timezone.as_deref(), CalendarWindow::Week, *week_start, )?, UsagePolicyWindow::CalendarMonth { timezone } => { calendar_bucket(now_unix_secs, timezone.as_deref(), CalendarWindow::Month, 1)? } UsagePolicyWindow::SubscriptionPeriod => { let (starts_at, expires_at) = rule.entitlement_period.ok_or_else(|| { PlanUsageAdmissionError::Gateway(GatewayError::Internal( "subscription period usage policy is missing its entitlement period" .to_string(), )) })?; let window_seconds = expires_at.saturating_sub(starts_at).max(1); let remaining = expires_at.saturating_sub(now_unix_secs).max(1); ( starts_at.to_string(), window_seconds, remaining, "subscription_period", ) } UsagePolicyWindow::Concurrent => { return Err(PlanUsageAdmissionError::Gateway(GatewayError::Internal( "concurrent window reached request counter compiler".to_string(), ))); } }; Ok(RuntimeRequestRule { key: format!("plan-usage:user:{{{user_id}}}:{}:{bucket}", rule.identity), limit: rule.limit, window_seconds: window_seconds.max(1), retention_seconds: retry_after.max(1), retry_after: retry_after.max(1), window, }) } fn durable_request_rule( rule: &EffectiveRequestRule, admitted_at_unix_secs: u64, ) -> Result { let (starts_at, ends_at, retry_after, label) = usage_window_bounds(&rule.window, rule.entitlement_period, admitted_at_unix_secs) .map_err(PlanUsageAdmissionError::Gateway)?; Ok(DurableRequestRule { window: UsagePolicyRequestWindow { starts_at_unix_secs: starts_at, ends_at_unix_secs: ends_at, limit_requests: rule.limit, }, retry_after, label, influence_ends_at_unix_secs: match &rule.window { UsagePolicyWindow::Rolling { seconds } => admitted_at_unix_secs .saturating_add(*seconds) .saturating_add(1), _ => ends_at, }, }) } fn runtime_cost_rule( rule: &EffectiveCostRule, admitted_at_unix_secs: u64, ) -> Result { let (starts_at, ends_at, retry_after, label) = usage_window_bounds(&rule.window, rule.entitlement_period, admitted_at_unix_secs)?; Ok(RuntimeCostRule { window: UsagePolicyCostWindow { window_id: format!("{}:{starts_at}", rule.identity), starts_at_unix_secs: starts_at, ends_at_unix_secs: ends_at, limit_cost_units: rule.limit_cost_units, }, retry_after, label, influence_ends_at_unix_secs: match &rule.window { UsagePolicyWindow::Rolling { seconds } => admitted_at_unix_secs .saturating_add(*seconds) .saturating_add(1), _ => ends_at, }, }) } fn usage_window_bounds( window: &UsagePolicyWindow, entitlement_period: Option<(u64, u64)>, now_unix_secs: u64, ) -> Result<(u64, u64, u64, &'static str), GatewayError> { match window { UsagePolicyWindow::Rolling { seconds } => Ok(( now_unix_secs.saturating_sub(*seconds), now_unix_secs.saturating_add(1), *seconds, rolling_label(*seconds), )), UsagePolicyWindow::CalendarDay { timezone } => { calendar_window_bounds(now_unix_secs, timezone.as_deref(), CalendarWindow::Day, 1) } UsagePolicyWindow::CalendarWeek { timezone, week_start, } => calendar_window_bounds( now_unix_secs, timezone.as_deref(), CalendarWindow::Week, *week_start, ), UsagePolicyWindow::CalendarMonth { timezone } => { calendar_window_bounds(now_unix_secs, timezone.as_deref(), CalendarWindow::Month, 1) } UsagePolicyWindow::SubscriptionPeriod => { let (starts_at, ends_at) = entitlement_period.ok_or_else(|| { GatewayError::Internal( "subscription period cost policy is missing its entitlement period".to_string(), ) })?; Ok(( starts_at, ends_at, ends_at.saturating_sub(now_unix_secs).max(1), "subscription_period", )) } UsagePolicyWindow::Concurrent => Err(GatewayError::Internal( "concurrent window reached cost policy compiler".to_string(), )), } } fn rolling_label(seconds: u64) -> &'static str { match seconds { 1 => "qps", 60 => "rpm", _ => "rolling", } } #[derive(Clone, Copy)] enum CalendarWindow { Day, Week, Month, } fn calendar_bucket( now_unix_secs: u64, timezone: Option<&str>, kind: CalendarWindow, week_start: u8, ) -> Result<(String, u64, u64, &'static str), PlanUsageAdmissionError> { let timezone = effective_timezone_name(timezone) .parse::() .map_err(|_| { PlanUsageAdmissionError::Gateway(GatewayError::Internal( "usage policy contains an invalid timezone".to_string(), )) })?; let now = Utc .timestamp_opt(i64::try_from(now_unix_secs).unwrap_or(i64::MAX), 0) .single() .ok_or_else(|| { PlanUsageAdmissionError::Gateway(GatewayError::Internal( "usage policy timestamp is out of range".to_string(), )) })? .with_timezone(&timezone); let current_date = now.date_naive(); let (start_date, end_date, label) = match kind { CalendarWindow::Day => ( current_date, current_date.succ_opt().ok_or_else(date_overflow)?, "calendar_day", ), CalendarWindow::Week => { let current = weekday_number(current_date.weekday()); let elapsed = (current + 7 - u32::from(week_start)) % 7; let start = current_date - chrono::Duration::days(i64::from(elapsed)); ( start, start .checked_add_days(chrono::Days::new(7)) .ok_or_else(date_overflow)?, "calendar_week", ) } CalendarWindow::Month => { let start = current_date.with_day(1).ok_or_else(date_overflow)?; let (year, month) = if start.month() == 12 { (start.year() + 1, 1) } else { (start.year(), start.month() + 1) }; ( start, chrono::NaiveDate::from_ymd_opt(year, month, 1).ok_or_else(date_overflow)?, "calendar_month", ) } }; let start = timezone .from_local_datetime(&start_date.and_hms_opt(0, 0, 0).ok_or_else(date_overflow)?) .earliest() .ok_or_else(date_overflow)?; let end = timezone .from_local_datetime(&end_date.and_hms_opt(0, 0, 0).ok_or_else(date_overflow)?) .latest() .ok_or_else(date_overflow)?; let end_unix = u64::try_from(end.timestamp()).map_err(|_| date_overflow())?; let start_unix = u64::try_from(start.timestamp()).map_err(|_| date_overflow())?; let window_seconds = end_unix.saturating_sub(start_unix).max(1); let remaining = end_unix.saturating_sub(now_unix_secs).max(1); Ok(( start.timestamp().to_string(), window_seconds, remaining, label, )) } fn calendar_window_bounds( now_unix_secs: u64, timezone: Option<&str>, kind: CalendarWindow, week_start: u8, ) -> Result<(u64, u64, u64, &'static str), GatewayError> { let (start, window_seconds, retry_after, label) = calendar_bucket(now_unix_secs, timezone, kind, week_start).map_err( |error| match error { PlanUsageAdmissionError::Gateway(error) => error, PlanUsageAdmissionError::Runtime(error) => { GatewayError::Internal(error.to_string()) } PlanUsageAdmissionError::Rejected(_) => GatewayError::Internal( "calendar usage policy compilation unexpectedly rejected a request".to_string(), ), }, )?; let starts_at_unix_secs = start.parse::().map_err(|_| { GatewayError::Internal("calendar usage policy produced an invalid epoch".to_string()) })?; Ok(( starts_at_unix_secs, starts_at_unix_secs.saturating_add(window_seconds), retry_after, label, )) } fn effective_timezone_name(explicit: Option<&str>) -> String { explicit .map(str::trim) .filter(|value| !value.is_empty()) .map(ToOwned::to_owned) .or_else(|| { std::env::var("APP_TIMEZONE") .ok() .map(|value| value.trim().to_string()) .filter(|value| value.parse::().is_ok()) }) .unwrap_or_else(|| DEFAULT_CALENDAR_TIMEZONE.to_string()) } fn weekday_number(weekday: Weekday) -> u32 { weekday.num_days_from_monday() + 1 } fn date_overflow() -> PlanUsageAdmissionError { PlanUsageAdmissionError::Gateway(GatewayError::Internal( "usage policy calendar window is out of range".to_string(), )) } #[cfg(test)] mod tests { use super::*; use serde_json::json; fn entitlement(id: &str, snapshot: serde_json::Value) -> UserPlanEntitlementRecord { UserPlanEntitlementRecord { id: id.to_string(), user_id: "user-1".to_string(), plan_id: "plan-1".to_string(), payment_order_id: "order-1".to_string(), status: "active".to_string(), starts_at_unix_secs: 1_000, expires_at_unix_secs: 10_000, entitlements_snapshot: snapshot, created_at_unix_secs: 1_000, updated_at_unix_secs: 1_000, } } #[test] fn compiles_week_only_and_combined_rules() { let policy = compile_effective_policy( &[entitlement( "ent-1", json!([{ "type": "usage_policy", "rules": [ {"metric":"request_count","window":{"kind":"rolling","seconds":18000},"limit":500}, {"metric":"request_count","window":{"kind":"calendar_week","timezone":"Asia/Shanghai"},"limit":10000}, {"metric":"concurrency","window":{"kind":"concurrent"},"limit":4} ] }]), )], 2_000, ) .expect("policy"); assert_eq!(policy.request_rules.len(), 2); assert_eq!(policy.concurrency_limit, Some(4)); } #[test] fn policy_snapshot_exists_only_for_cost_rules_and_derives_unique_tokens() { assert!(PlanUsagePolicySnapshot::for_admission( "user-1", EffectivePlanUsagePolicy::default(), 2_000, ) .is_none()); let policy = compile_effective_policy( &[entitlement( "ent-cost", json!([{"type":"usage_policy","rules":[ {"metric":"actual_cost_usd","window":{"kind":"subscription_period"},"limit":10.0} ]}]), )], 2_000, ) .expect("cost policy"); let snapshot = PlanUsagePolicySnapshot::for_admission("user-1", policy.clone(), 2_000) .expect("cost rules require a policy snapshot"); let reservation = snapshot.new_reservation_context(); let retry_reservation = snapshot.new_reservation_context(); assert_eq!(reservation.subject_id(), "user-1"); assert_eq!(reservation.admitted_at_unix_secs(), 2_000); assert_eq!(reservation.policy(), &policy); assert!(!reservation.token().trim().is_empty()); assert_ne!(reservation.token(), retry_reservation.token()); assert_eq!(retry_reservation.admitted_at_unix_secs(), 2_000); assert!(std::ptr::eq( reservation.policy(), retry_reservation.policy() )); } #[test] fn admitted_cost_snapshot_keeps_original_window_after_entitlement_expiry() { let policy = compile_effective_policy( &[entitlement( "ent-cost", json!([{"type":"usage_policy","rules":[ {"metric":"actual_cost_usd","window":{"kind":"subscription_period"},"limit":10.0} ]}]), )], 9_999, ) .expect("policy before expiry"); let snapshot = PlanUsagePolicySnapshot::for_admission("user-1", policy, 9_999) .expect("cost policy snapshot"); let rule = snapshot .policy() .cost_rules .first() .expect("snapshotted cost rule"); let runtime = runtime_cost_rule(rule, snapshot.admitted_at_unix_secs) .expect("snapshot remains compilable at its admitted timestamp"); assert_eq!(runtime.window.starts_at_unix_secs, 1_000); assert_eq!(runtime.window.ends_at_unix_secs, 10_000); assert_eq!(runtime.retry_after, 1); assert!(compile_effective_policy( &[entitlement( "ent-cost", json!([{"type":"usage_policy","rules":[ {"metric":"actual_cost_usd","window":{"kind":"subscription_period"},"limit":10.0} ]}]), )], 10_000, ) .expect("policy at expiry") .cost_rules .is_empty()); } #[test] fn same_window_uses_strictest_limit_but_subscription_periods_stay_independent() { let first = entitlement( "ent-a", json!([{"type":"usage_policy","rules":[ {"metric":"request_count","window":{"kind":"rolling","seconds":60},"limit":100}, {"metric":"request_count","window":{"kind":"subscription_period"},"limit":1000} ]}]), ); let second = entitlement( "ent-b", json!([{"type":"usage_policy","rules":[ {"metric":"request_count","window":{"kind":"rolling","seconds":60},"limit":40}, {"metric":"request_count","window":{"kind":"subscription_period"},"limit":2000} ]}]), ); let policy = compile_effective_policy(&[first, second], 2_000).expect("policy"); assert_eq!(policy.request_rules.len(), 3); assert_eq!( policy .request_rules .iter() .find(|rule| rule.identity == "rolling:60") .map(|rule| rule.limit), Some(40) ); } #[test] fn rolling_windows_use_stable_sliding_keys_and_calendar_windows_use_epoch_keys() { let rolling = EffectiveRequestRule { identity: "rolling:18000".to_string(), window: UsagePolicyWindow::Rolling { seconds: 18_000 }, limit: 10, entitlement_period: None, }; let first = runtime_request_rule("user-1", &rolling, 20_000).expect("first rolling"); let second = runtime_request_rule("user-1", &rolling, 40_000).expect("second rolling"); assert_eq!(first.key, second.key); assert_eq!(first.window_seconds, 18_000); assert_eq!(first.retention_seconds, 18_000); let weekly = EffectiveRequestRule { identity: "calendar_week:Asia/Shanghai:1".to_string(), window: UsagePolicyWindow::CalendarWeek { timezone: Some("Asia/Shanghai".to_string()), week_start: 1, }, limit: 20, entitlement_period: None, }; let sunday = chrono::DateTime::parse_from_rfc3339("2026-08-16T12:00:00+08:00") .unwrap() .timestamp() as u64; let monday = chrono::DateTime::parse_from_rfc3339("2026-08-17T12:00:00+08:00") .unwrap() .timestamp() as u64; let first = runtime_request_rule("user-1", &weekly, sunday).expect("first week"); let second = runtime_request_rule("user-1", &weekly, monday).expect("second week"); assert_ne!(first.key, second.key); assert_eq!(first.window, "calendar_week"); assert_eq!(first.retention_seconds, first.retry_after); assert!(first.retention_seconds < first.window_seconds); } #[test] fn subscription_period_counter_keeps_the_entitlement_epoch_isolated() { let rule = EffectiveRequestRule { identity: "subscription:ent-1".to_string(), window: UsagePolicyWindow::SubscriptionPeriod, limit: 5, entitlement_period: Some((1_000, 10_000)), }; let runtime = runtime_request_rule("user-1", &rule, 9_500).expect("subscription rule"); assert!(runtime.key.ends_with(":subscription:ent-1:1000")); assert_eq!(runtime.window_seconds, 9_000); assert_eq!(runtime.retention_seconds, 500); assert_eq!(runtime.retry_after, 500); } #[test] fn calendar_bucket_retention_uses_only_the_short_remaining_period() { let daily = EffectiveRequestRule { identity: "calendar_day:Asia/Shanghai".to_string(), window: UsagePolicyWindow::CalendarDay { timezone: Some("Asia/Shanghai".to_string()), }, limit: 5, entitlement_period: None, }; let near_midnight = chrono::DateTime::parse_from_rfc3339("2026-08-18T23:59:59+08:00") .unwrap() .timestamp() as u64; let runtime = runtime_request_rule("user-1", &daily, near_midnight).expect("daily rule"); assert_eq!(runtime.window_seconds, 24 * 60 * 60); assert_eq!(runtime.retention_seconds, 1); assert_eq!(runtime.retry_after, 1); } #[test] fn cost_policy_fails_closed_when_usage_runtime_is_disabled() { assert!(ensure_cost_policy_usage_runtime_enabled(true).is_ok()); assert!(matches!( ensure_cost_policy_usage_runtime_enabled(false), Err(GatewayError::Internal(message)) if message.contains("usage runtime") && message.contains("terminal reconciliation") )); } #[test] fn cost_retention_uses_exclusive_rolling_influence_end() { let rule = EffectiveCostRule { identity: "rolling:18000".to_string(), window: UsagePolicyWindow::Rolling { seconds: 18_000 }, limit_cost_units: 100, entitlement_period: None, }; let runtime = runtime_cost_rule(&rule, 20_000).expect("runtime cost rule"); assert_eq!(runtime.influence_ends_at_unix_secs, 38_001); } #[test] fn explicit_cost_release_uses_server_token_and_zero_actual_cost() { let input = build_usage_policy_cost_release_input( "request-1", "user-1", "server-reservation-token", 12_345, ); assert_eq!(input.request_id, "request-1"); assert_eq!(input.subject_id, "user-1"); assert_eq!(input.reservation_token, "server-reservation-token"); assert_eq!(input.actual_cost_units, 0); assert_eq!( input.terminal_state, UsagePolicyCostReservationState::Released ); assert_eq!(input.finalized_at_unix_secs, 12_345); input.validate().expect("release input should validate"); } }