diff --git a/apps/aether-gateway/src/control/auth/gate.rs b/apps/aether-gateway/src/control/auth/gate.rs index 6e79bf6e6..a24f29ff6 100644 --- a/apps/aether-gateway/src/control/auth/gate.rs +++ b/apps/aether-gateway/src/control/auth/gate.rs @@ -131,7 +131,13 @@ async fn execution_plan_balance_capacity_rejection_inner( validate_execution_plan_pricing_configuration_for_plan(state, plan, report_context).await?; return Ok(None); } - let Some(available_usd) = available_balance_capacity_usd(state, auth_context).await? else { + let wallet_only = report_context + .and_then(|context| context.get("plan_wallet_fallback")) + .and_then(serde_json::Value::as_bool) + .unwrap_or(false); + let Some(available_usd) = + available_balance_capacity_usd(state, auth_context, wallet_only).await? + else { validate_execution_plan_pricing_configuration_for_plan(state, plan, report_context).await?; return Ok(None); }; @@ -177,11 +183,16 @@ async fn validate_execution_plan_pricing_configuration_for_plan( async fn available_balance_capacity_usd( state: &AppState, auth_context: &GatewayControlAuthContext, + wallet_only: bool, ) -> Result, GatewayError> { let quota_started_at = std::time::Instant::now(); - let quota_result = state - .find_user_daily_quota_availability_for_auth(&auth_context.user_id) - .await; + let quota_result = if wallet_only { + Ok(None) + } else { + state + .find_user_daily_quota_availability_for_auth(&auth_context.user_id) + .await + }; observe_gateway_stage_ms( "auth_capacity_quota", quota_started_at.elapsed().as_millis() as u64, @@ -201,6 +212,13 @@ async fn available_balance_capacity_usd( wallet_started_at.elapsed().as_millis() as u64, ); let wallet = wallet_result?; + if wallet_only + && !wallet + .as_ref() + .is_some_and(|wallet| wallet.status.eq_ignore_ascii_case("active")) + { + return Ok(Some(0.0)); + } let wallet_available_usd = wallet.as_ref().and_then(wallet_finite_available_usd); let wallet_is_unlimited = wallet .as_ref() @@ -1013,6 +1031,11 @@ mod tests { context: StoredBillingModelContext, wallet: StoredWalletSnapshot, ) -> AppState { + let mut preferences = + aether_data_contracts::repository::users::StoredUserPreferenceRecord::default_for_user( + "user-1", + ); + preferences.allow_wallet_overage = quota.allow_wallet_overage; let candidate_repository = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ sample_row(), @@ -1021,7 +1044,8 @@ mod tests { let data = GatewayDataState::with_minimal_candidate_selection_and_billing_for_tests( candidate_repository, billing_repository, - ); + ) + .with_user_preferences_for_tests([preferences]); AppState::new() .expect("state should build") .with_data_state_for_tests(data) @@ -1436,7 +1460,7 @@ mod tests { .as_ref() .expect("decision should include auth context"); - let capacity = available_balance_capacity_usd(&state, auth_context) + let capacity = available_balance_capacity_usd(&state, auth_context, false) .await .expect("capacity should resolve"); @@ -1673,10 +1697,16 @@ mod tests { Arc::clone("a_calls), Arc::clone(&model_context_calls), )); + let mut preferences = + aether_data_contracts::repository::users::StoredUserPreferenceRecord::default_for_user( + "user-1", + ); + preferences.allow_wallet_overage = true; let data = GatewayDataState::with_minimal_candidate_selection_and_billing_for_tests( candidate_repository, billing_repository, - ); + ) + .with_user_preferences_for_tests([preferences]); let state = AppState::new() .expect("state should build") .with_data_state_for_tests(data) @@ -2089,6 +2119,40 @@ mod tests { assert_eq!(rejection, None); } + #[tokio::test] + async fn plan_recovery_wallet_fallback_cannot_spend_available_daily_quota() { + let context = billing_context_with_pricing( + Some( + json!({"tiers": [{"up_to": null, "input_price_per_1m": 0.0, "output_price_per_1m": 70.0}]}), + ), + None, + None, + None, + ); + let state = state_with_quota_and_wallet(quota_availability(50.0, true), context); + let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]); + let plan = execution_plan( + json!({"model": "gpt-5", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 1_000_000}), + "openai:chat", + ); + let mut report_context = billing_report_context(); + report_context["plan_wallet_fallback"] = json!(true); + let rejection = execution_plan_balance_capacity_rejection( + &state, + &decision, + &plan, + Some(&report_context), + ) + .await + .unwrap(); + assert_eq!( + rejection, + Some(GatewayLocalAuthRejection::BalanceDenied { + remaining: Some(30.0) + }) + ); + } + #[test] fn daily_quota_estimate_falls_back_to_default_tiers_when_model_tiers_empty() { let context = billing_context_with_pricing( diff --git a/apps/aether-gateway/src/data/tests.rs b/apps/aether-gateway/src/data/tests.rs index aabe78d89..2427c2137 100644 --- a/apps/aether-gateway/src/data/tests.rs +++ b/apps/aether-gateway/src/data/tests.rs @@ -336,6 +336,7 @@ async fn data_state_checks_user_uniqueness_through_user_reader() { email_notifications: false, usage_alerts: true, announcement_notifications: false, + allow_wallet_overage: true, }; assert_eq!( state diff --git a/apps/aether-gateway/src/executor/candidate_loop.rs b/apps/aether-gateway/src/executor/candidate_loop.rs index 74afa0f41..9c0cdfb48 100644 --- a/apps/aether-gateway/src/executor/candidate_loop.rs +++ b/apps/aether-gateway/src/executor/candidate_loop.rs @@ -310,6 +310,12 @@ where attempt.report_context(), self.transfer_tracker.usage_policy_reservation_token(), ); + let report_context = crate::usage::attach_plan_wallet_fallback_context( + self.state, + report_context, + self.transfer_tracker.uses_wallet_fallback(), + ) + .await?; let balance_response = execution_plan_balance_capacity_response( self.state, self.trace_id, @@ -370,6 +376,12 @@ where { return Ok(AiAttemptExecutionOutcome::Responded(response)); } + let report_context = crate::usage::attach_plan_wallet_fallback_context( + self.state, + report_context, + self.transfer_tracker.uses_wallet_fallback(), + ) + .await?; let upstream_execution_gate_held_started_at = std::time::Instant::now(); let deferred_report_context = report_context.clone(); let execution = execute_execution_runtime_sync_with_retry_scope( @@ -754,9 +766,16 @@ impl ProviderTransferTracker { fn usage_policy_reservation_token(&self) -> Option<&str> { self.usage_policy_reservation .as_ref() + .filter(|reservation| !reservation.uses_wallet_fallback()) .map(crate::plan_usage_policy::PlanUsageReservationContext::token) } + fn uses_wallet_fallback(&self) -> bool { + self.usage_policy_reservation + .as_ref() + .is_some_and(|reservation| reservation.uses_wallet_fallback()) + } + fn record_usage_policy_reservation_plan(&self, plan: &aether_contracts::ExecutionPlan) { *self .usage_policy_reservation_plan @@ -1269,6 +1288,12 @@ where attempt.report_context(), self.transfer_tracker.usage_policy_reservation_token(), ); + let report_context = crate::usage::attach_plan_wallet_fallback_context( + self.state, + report_context, + self.transfer_tracker.uses_wallet_fallback(), + ) + .await?; let candidate_index = parse_request_candidate_report_context(report_context.as_ref()) .and_then(|context| context.candidate_index) .map(|value| value.to_string()) @@ -1321,6 +1346,24 @@ where return Ok(AiAttemptExecutionOutcome::Responded(response)); } prewarm_direct_reqwest_candidate_client(plan); + if let Some(response) = execution_plan_cost_capacity_response( + self.state, + self.trace_id, + self.decision, + plan, + report_context.as_ref(), + self.transfer_tracker, + ) + .await? + { + return Ok(AiAttemptExecutionOutcome::Responded(response)); + } + let report_context = crate::usage::attach_plan_wallet_fallback_context( + self.state, + report_context, + self.transfer_tracker.uses_wallet_fallback(), + ) + .await?; let watchdog_report_context_owned = report_context.clone(); let watchdog_report_context = watchdog_report_context_owned.as_ref(); let execution_state = self.state.clone(); @@ -1329,7 +1372,6 @@ where let execution_decision = self.decision.clone(); let execution_report_kind = attempt.report_kind(); let execution_plan = plan.clone(); - let execution_transfer_tracker = self.transfer_tracker.clone(); let stop_on_transport_errors = matches!( resolve_local_transport_failover_analysis_for_attempt( self.state, @@ -1349,18 +1391,6 @@ where watchdog_report_context, stop_on_transport_errors, move || async move { - if let Some(response) = execution_plan_cost_capacity_response( - &execution_state, - execution_trace_id.as_str(), - &execution_decision, - &execution_plan, - report_context.as_ref(), - &execution_transfer_tracker, - ) - .await? - { - return Ok(AiAttemptExecutionOutcome::Responded(response)); - } execute_execution_runtime_stream_with_retry_scope( &execution_state, execution_plan, @@ -1579,6 +1609,12 @@ async fn execution_plan_cost_capacity_response( }; let rejection = match outcome { crate::plan_usage_policy::PlanUsageCostReservationOutcome::NotRequired => return Ok(None), + crate::plan_usage_policy::PlanUsageCostReservationOutcome::WalletFallback => { + transfer_tracker + .usage_policy_cost_reserved + .store(false, Ordering::Release); + return Ok(None); + } crate::plan_usage_policy::PlanUsageCostReservationOutcome::Reserved => { transfer_tracker.record_usage_policy_reservation_plan(plan); transfer_tracker diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/turn.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/turn.rs index 55155b4fc..aeaef2010 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/turn.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/turn.rs @@ -388,7 +388,8 @@ async fn reserve_responses_plan_usage_cost_owned( ) .await?; match outcome { - crate::plan_usage_policy::PlanUsageCostReservationOutcome::NotRequired => { + crate::plan_usage_policy::PlanUsageCostReservationOutcome::NotRequired + | crate::plan_usage_policy::PlanUsageCostReservationOutcome::WalletFallback => { Ok(ResponsesPlanUsageCostReservationStart::NotRequired) } crate::plan_usage_policy::PlanUsageCostReservationOutcome::Reserved => { @@ -565,7 +566,14 @@ pub(super) async fn begin_unowned_responses_websocket_turn( )); } }; - let mut report_context = attempt.report_context; + let mut report_context = crate::usage::attach_plan_wallet_fallback_context( + state, + attempt.report_context, + plan_usage_policy_snapshot + .as_ref() + .is_some_and(PlanUsagePolicySnapshot::uses_wallet_fallback), + ) + .await?; let balance_rejection = execution_plan_balance_capacity_rejection( state, @@ -695,6 +703,14 @@ pub(super) async fn begin_unowned_responses_websocket_turn( } }; + report_context = crate::usage::attach_plan_wallet_fallback_context( + state, + report_context, + plan_usage_policy_snapshot + .as_ref() + .is_some_and(PlanUsagePolicySnapshot::uses_wallet_fallback), + ) + .await?; let lifecycle = ExecutionAttemptLifecycle::begin( state, AttemptLifecycleSeed { diff --git a/apps/aether-gateway/src/handlers/public/support/user_me_preferences.rs b/apps/aether-gateway/src/handlers/public/support/user_me_preferences.rs index bdcd90674..926d711d5 100644 --- a/apps/aether-gateway/src/handlers/public/support/user_me_preferences.rs +++ b/apps/aether-gateway/src/handlers/public/support/user_me_preferences.rs @@ -84,6 +84,7 @@ fn build_users_me_preferences_payload( "theme": preferences.theme, "language": preferences.language, "timezone": preferences.timezone, + "allow_wallet_overage": preferences.allow_wallet_overage, "notifications": { "email": preferences.email_notifications, "usage_alerts": preferences.usage_alerts, @@ -347,6 +348,17 @@ pub(super) async fn handle_users_me_preferences_put( preferences.announcement_notifications = announcement_notifications; } + let allow_wallet_overage = + match parse_users_me_optional_bool_field(payload, "allow_wallet_overage") { + Ok(value) => value, + Err(detail) => { + return build_auth_error_response(http::StatusCode::BAD_REQUEST, detail, false) + } + }; + if let Some(allow_wallet_overage) = allow_wallet_overage { + preferences.allow_wallet_overage = allow_wallet_overage; + } + match state.write_user_preferences(&preferences).await { Ok(Some(_)) => Json(json!({ "message": "偏好设置更新成功" })).into_response(), Ok(None) => build_auth_error_response( diff --git a/apps/aether-gateway/src/plan_usage_policy.rs b/apps/aether-gateway/src/plan_usage_policy.rs index 5d2ef8d16..d68405bde 100644 --- a/apps/aether-gateway/src/plan_usage_policy.rs +++ b/apps/aether-gateway/src/plan_usage_policy.rs @@ -1,4 +1,5 @@ use std::collections::BTreeMap; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::time::Duration; @@ -7,9 +8,10 @@ use aether_data_contracts::repository::billing::{ UsagePolicyWindow, UserPlanEntitlementRecord, USAGE_POLICY_COST_UNITS_PER_USD, }; use aether_data_contracts::repository::settlement::{ - ReconcileUsagePolicyCostInput, ReserveUsagePolicyCostInput, ReserveUsagePolicyCostOutcome, - ReserveUsagePolicyRequestInput, ReserveUsagePolicyRequestOutcome, - UsagePolicyCostReservationState, UsagePolicyCostWindow, UsagePolicyRequestWindow, + ReconcileUsagePolicyCostInput, ReleaseUsagePolicyRequestAdmissionInput, + ReserveUsagePolicyCostInput, ReserveUsagePolicyCostOutcome, ReserveUsagePolicyRequestInput, + ReserveUsagePolicyRequestOutcome, UsagePolicyCostReservationState, UsagePolicyCostWindow, + UsagePolicyRequestWindow, }; use aether_runtime::AdmissionPermit; use aether_runtime_state::{ @@ -33,9 +35,12 @@ pub(crate) struct PlanUsagePolicySnapshot { pub(crate) admitted_at_unix_secs: u64, subject_id: Arc, policy: Arc, + admission_event_id: Arc, + wallet_fallback: Arc, } impl PlanUsagePolicySnapshot { + #[cfg(test)] fn for_admission( subject_id: &str, policy: EffectivePlanUsagePolicy, @@ -48,9 +53,34 @@ impl PlanUsagePolicySnapshot { admitted_at_unix_secs, subject_id: subject_id.to_string().into(), policy: Arc::new(policy), + admission_event_id: "".into(), + wallet_fallback: Arc::new(AtomicBool::new(false)), }) } + fn with_admission( + subject_id: &str, + policy: EffectivePlanUsagePolicy, + admitted_at_unix_secs: u64, + event_id: &str, + wallet_fallback: bool, + ) -> Option { + if policy.cost_rules.is_empty() && !wallet_fallback { + return None; + } + Some(Self { + admitted_at_unix_secs, + subject_id: subject_id.into(), + policy: Arc::new(policy), + admission_event_id: event_id.into(), + wallet_fallback: Arc::new(AtomicBool::new(wallet_fallback)), + }) + } + + pub(crate) fn uses_wallet_fallback(&self) -> bool { + self.wallet_fallback.load(Ordering::Acquire) + } + pub(crate) fn new_reservation_context(&self) -> PlanUsageReservationContext { PlanUsageReservationContext { policy_snapshot: self.clone(), @@ -86,6 +116,8 @@ impl PlanUsageReservationContext { admitted_at_unix_secs, subject_id: subject_id.into(), policy: Arc::new(policy), + admission_event_id: "".into(), + wallet_fallback: Arc::new(AtomicBool::new(false)), }, token: token.into(), } @@ -106,6 +138,10 @@ impl PlanUsageReservationContext { pub(crate) const fn admitted_at_unix_secs(&self) -> u64 { self.policy_snapshot.admitted_at_unix_secs } + + pub(crate) fn uses_wallet_fallback(&self) -> bool { + self.policy_snapshot.uses_wallet_fallback() + } } #[derive(Debug)] @@ -182,6 +218,7 @@ pub(crate) struct PlanUsagePolicyRejection { pub(crate) enum PlanUsageCostReservationOutcome { NotRequired, Reserved, + WalletFallback, Rejected(PlanUsagePolicyRejection), } @@ -203,13 +240,12 @@ pub(crate) async fn reserve_admitted_http_plan_usage_policy_cost( "plan usage reservation subject does not match the admitted request".to_string(), )); } - reserve_plan_usage_policy_cost_with_policy( + reserve_cost_with_wallet_fallback( state, decision, plan, report_context, - reservation.policy(), - reservation.admitted_at_unix_secs(), + &reservation.policy_snapshot, reservation.token(), ) .await @@ -234,7 +270,58 @@ pub(crate) async fn reserve_admitted_plan_usage_policy_cost( "plan usage reservation subject does not match the admitted request".to_string(), )); } - reserve_plan_usage_policy_cost_with_policy( + reserve_cost_with_wallet_fallback( + state, + decision, + plan, + report_context, + snapshot, + reservation_token, + ) + .await +} + +async fn wallet_fallback_available( + state: &AppState, + auth: &crate::control::GatewayControlAuthContext, + required_cost_usd: Option, +) -> Result { + if !state + .read_user_preferences(&auth.user_id) + .await? + .is_some_and(|preferences| preferences.allow_wallet_overage) + { + return Ok(false); + } + let wallet = state + .read_wallet_snapshot_for_auth_uncached( + &auth.user_id, + &auth.api_key_id, + auth.api_key_is_standalone, + ) + .await?; + Ok(wallet.is_some_and(|wallet| { + wallet.status.eq_ignore_ascii_case("active") + && (wallet.limit_mode.eq_ignore_ascii_case("unlimited") || { + let available = wallet.balance.max(0.0) + wallet.gift_balance.max(0.0); + available > 0.000_000_01 + && required_cost_usd.is_none_or(|cost| cost <= available + 0.000_000_01) + }) + })) +} + +async fn reserve_cost_with_wallet_fallback( + state: &AppState, + decision: &GatewayControlDecision, + plan: &aether_contracts::ExecutionPlan, + report_context: Option<&serde_json::Value>, + snapshot: &PlanUsagePolicySnapshot, + reservation_token: &str, +) -> Result { + if snapshot.uses_wallet_fallback() { + return Ok(PlanUsageCostReservationOutcome::WalletFallback); + } + let outcome = reserve_plan_usage_policy_cost_with_policy( state, decision, plan, @@ -243,7 +330,52 @@ pub(crate) async fn reserve_admitted_plan_usage_policy_cost( snapshot.admitted_at_unix_secs, reservation_token, ) - .await + .await?; + if !matches!(outcome, PlanUsageCostReservationOutcome::Rejected(_)) { + return Ok(outcome); + } + let Some(auth) = plan_usage_auth_context(decision) else { + return Ok(outcome); + }; + let estimated = + crate::control::estimate_execution_plan_cost_upper_bound_usd(state, plan, report_context) + .await?; + let Some(estimated) = estimated else { + return Ok(outcome); + }; + if !wallet_fallback_available(state, auth, Some(estimated)).await? { + return Ok(outcome); + } + // A previous provider attempt may have reserved a smaller cost. A paid retry must + // release it, and must not extend the plan's long request-count recovery window. + release_plan_usage_policy_cost( + state, + decision, + plan, + reservation_token, + crate::clock::current_unix_secs(), + ) + .await?; + if !snapshot.admission_event_id.is_empty() + && snapshot + .policy + .request_rules + .iter() + .any(|rule| !request_rule_uses_runtime_state(&rule)) + { + state + .data + .release_usage_policy_request_admission(ReleaseUsagePolicyRequestAdmissionInput { + request_id: snapshot.admission_event_id.to_string(), + subject_id: auth.user_id.clone(), + event_token: snapshot.admission_event_id.to_string(), + released_at_unix_secs: crate::clock::current_unix_secs(), + }) + .await + .map_err(|error| GatewayError::Internal(error.to_string()))?; + } + snapshot.wallet_fallback.store(true, Ordering::Release); + Ok(PlanUsageCostReservationOutcome::WalletFallback) } fn plan_usage_auth_context( @@ -458,16 +590,16 @@ pub(crate) async fn check_and_acquire_plan_usage_policy_admission( 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, + let (permit, wallet_fallback) = + check_and_acquire_compiled_plan_usage_policy(state, auth, &policy, event_id, now_unix_ms) + .await?; + let policy_snapshot = PlanUsagePolicySnapshot::with_admission( &auth.user_id, - &policy, + policy, + admitted_at_unix_secs, event_id, - now_unix_ms, - ) - .await?; - let policy_snapshot = - PlanUsagePolicySnapshot::for_admission(&auth.user_id, policy, admitted_at_unix_secs); + wallet_fallback, + ); Ok(PlanUsageAdmission { permit, policy_snapshot, @@ -476,15 +608,34 @@ pub(crate) async fn check_and_acquire_plan_usage_policy_admission( async fn check_and_acquire_compiled_plan_usage_policy( state: &AppState, - subject_id: &str, + auth: &crate::control::GatewayControlAuthContext, policy: &EffectivePlanUsagePolicy, event_id: &str, now_unix_ms: u64, -) -> Result, PlanUsageAdmissionError> { +) -> Result<(Option, bool), PlanUsageAdmissionError> { + let daily_quota_exhausted = if !policy.cost_rules.is_empty() + || policy + .request_rules + .iter() + .any(|rule| !request_rule_uses_runtime_state(&rule)) + { + state + .find_user_daily_quota_availability_for_auth(&auth.user_id) + .await? + .is_some_and(|quota| { + quota.has_active_daily_quota && quota.remaining_usd <= 0.000_000_01 + }) + } else { + false + }; + let wallet_fallback = + daily_quota_exhausted && wallet_fallback_available(state, auth, None).await?; if policy.request_rules.is_empty() && policy.concurrency_limit.is_none() { - return Ok(None); + return Ok((None, wallet_fallback)); } + let subject_id = auth.user_id.as_str(); + let now_unix_secs = now_unix_ms / 1_000; let plan_permit = if let Some(limit) = policy.concurrency_limit { @@ -570,6 +721,7 @@ async fn check_and_acquire_compiled_plan_usage_policy( let durable_rules = durable_request_rules .iter() + .filter(|_| !wallet_fallback) .map(|rule| durable_request_rule(rule, now_unix_secs)) .collect::, _>>()?; if !durable_rules.is_empty() { @@ -630,6 +782,11 @@ async fn check_and_acquire_compiled_plan_usage_policy( limit_requests, .. } => { + // Only long quota windows may use paid fallback. Concurrency and + // short QPS/RPM rules above remain enforced and retain their permit. + if wallet_fallback_available(state, auth, None).await? { + return Ok((AdmissionPermit::from_parts(None, plan_permit), true)); + } release_runtime_usage_limits_best_effort( state, &runtime_inputs, @@ -678,7 +835,10 @@ async fn check_and_acquire_compiled_plan_usage_policy( } } - Ok(AdmissionPermit::from_parts(None, plan_permit)) + Ok(( + AdmissionPermit::from_parts(None, plan_permit), + wallet_fallback, + )) } pub(crate) async fn check_and_acquire_http_plan_usage_policy( @@ -705,17 +865,17 @@ pub(crate) async fn check_and_acquire_http_plan_usage_policy( 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, + let (permit, wallet_fallback) = + check_and_acquire_compiled_plan_usage_policy(state, auth, &policy, event_id, now_unix_ms) + .await?; + let reservation_context = PlanUsagePolicySnapshot::with_admission( &auth.user_id, - &policy, + policy, + admitted_at_unix_secs, event_id, - now_unix_ms, + wallet_fallback, ) - .await?; - let reservation_context = - PlanUsagePolicySnapshot::for_admission(&auth.user_id, policy, admitted_at_unix_secs) - .map(|snapshot| snapshot.new_reservation_context()); + .map(|snapshot| snapshot.new_reservation_context()); Ok(HttpPlanUsageAdmission { permit, reservation_context, @@ -1229,6 +1389,367 @@ mod tests { use super::*; use serde_json::json; + fn fallback_test_state(enabled: bool, balance: f64) -> AppState { + use aether_data::repository::billing::InMemoryBillingReadRepository; + use aether_data::repository::settlement::InMemorySettlementRepository; + use aether_data::repository::usage::InMemoryUsageReadRepository; + use aether_data::repository::wallet::{InMemoryWalletRepository, StoredWalletSnapshot}; + use aether_data_contracts::repository::billing::StoredBillingModelContext; + let wallet = StoredWalletSnapshot::new( + "wallet-user-1".into(), + Some("user-1".into()), + None, + balance, + 0.0, + "finite".into(), + "USD".into(), + "active".into(), + 0.0, + 0.0, + 0.0, + 0.0, + 100, + ) + .unwrap(); + let mut preferences = crate::GatewayUserPreferenceView::default_for_user("user-1"); + preferences.allow_wallet_overage = enabled; + let data = crate::data::GatewayDataState::with_usage_billing_and_wallet_for_tests( + Arc::new(InMemoryUsageReadRepository::default()), + Arc::new(InMemoryBillingReadRepository::seed([ + StoredBillingModelContext::new( + "provider-1".into(), + None, + Some("key-1".into()), + None, + None, + "global-model-1".into(), + "gpt-5".into(), + None, + Some(0.25), + None, + Some("model-1".into()), + Some("gpt-5".into()), + None, + None, + None, + ) + .unwrap(), + ])), + Arc::new(InMemoryWalletRepository::seed([wallet])), + ) + .with_settlement_writer_for_tests(Arc::new(InMemorySettlementRepository::default())) + .with_user_preferences_for_tests([preferences.into()]); + AppState::new() + .unwrap() + .with_data_state_for_tests(data) + .with_usage_runtime_for_tests(crate::usage::UsageRuntimeConfig { + enabled: true, + ..Default::default() + }) + } + + fn fallback_test_auth() -> crate::control::GatewayControlAuthContext { + crate::control::GatewayControlAuthContext { + user_id: "user-1".into(), + api_key_id: "api-key-1".into(), + username: None, + api_key_name: None, + balance_remaining: None, + access_allowed: true, + user_rate_limit: None, + api_key_rate_limit: None, + api_key_is_standalone: false, + admin_bypass_limits: false, + local_rejection: None, + allowed_models: None, + ip_rules: None, + verified_api_key_hash: None, + } + } + + fn fallback_test_policy(short_limit: u64) -> EffectivePlanUsagePolicy { + compile_effective_policy(&[entitlement("ent-1", json!([{ + "type":"usage_policy", "rules":[ + {"metric":"request_count","window":{"kind":"rolling","seconds":3600},"limit":1}, + {"metric":"request_count","window":{"kind":"rolling","seconds":60},"limit":short_limit}, + {"metric":"concurrency","window":{"kind":"concurrent"},"limit":1} + ] + }]))], 2_000).unwrap() + } + + #[tokio::test] + async fn wallet_fallback_requires_opt_in_and_actual_wallet_capacity() { + let auth = fallback_test_auth(); + assert!( + !wallet_fallback_available(&fallback_test_state(false, 20.0), &auth, None) + .await + .unwrap() + ); + assert!( + !wallet_fallback_available(&fallback_test_state(true, 0.0), &auth, None) + .await + .unwrap() + ); + let state = fallback_test_state(true, 20.0); + assert!(wallet_fallback_available(&state, &auth, Some(20.0)) + .await + .unwrap()); + assert!(!wallet_fallback_available(&state, &auth, Some(20.01)) + .await + .unwrap()); + } + + #[tokio::test] + async fn wallet_fallback_preserves_short_limits_and_concurrency() { + let state = fallback_test_state(true, 20.0); + let auth = fallback_test_auth(); + let policy = fallback_test_policy(2); + let mut quota_only_policy = policy.clone(); + quota_only_policy.concurrency_limit = None; + let (permit, fallback) = check_and_acquire_compiled_plan_usage_policy( + &state, + &auth, + "a_only_policy, + "first", + 2_000_000, + ) + .await + .unwrap(); + assert!(!fallback); + drop(permit); + let (permit, fallback) = check_and_acquire_compiled_plan_usage_policy( + &state, &auth, &policy, "second", 2_000_001, + ) + .await + .unwrap(); + assert!(fallback); + // The fallback turn retains the concurrency permit; this is the exact + // protection we need to assert. A third turn cannot enter concurrently. + assert!(matches!( + check_and_acquire_compiled_plan_usage_policy( + &state, + &auth, + &policy, + "concurrent", + 2_000_002 + ) + .await, + Err(PlanUsageAdmissionError::Runtime( + RuntimeSemaphoreError::Saturated { .. } + )) + )); + drop(permit); + assert!(matches!( + check_and_acquire_compiled_plan_usage_policy( + &state, + &auth, + "a_only_policy, + "third", + 2_000_003 + ) + .await, + Err(PlanUsageAdmissionError::Rejected( + PlanUsagePolicyRejection { + retry_after: 1..=60, + .. + } + )) + )); + } + + #[tokio::test] + async fn disabled_wallet_fallback_keeps_long_window_rejection() { + let state = fallback_test_state(false, 20.0); + let auth = fallback_test_auth(); + let mut policy = fallback_test_policy(10); + // This assertion exercises durable quota; concurrency is tested separately. + policy.concurrency_limit = None; + let (permit, _) = check_and_acquire_compiled_plan_usage_policy( + &state, &auth, &policy, "first", 2_000_000, + ) + .await + .unwrap(); + drop(permit); + assert!(matches!( + check_and_acquire_compiled_plan_usage_policy( + &state, &auth, &policy, "second", 2_000_001 + ) + .await, + Err(PlanUsageAdmissionError::Rejected(_)) + )); + } + + #[test] + fn wallet_fallback_snapshot_survives_http_and_websocket_retries_without_cost_rules() { + let snapshot = PlanUsagePolicySnapshot::with_admission( + "user-1", + EffectivePlanUsagePolicy::default(), + 2_000, + "event", + true, + ) + .unwrap(); + assert!(snapshot.clone().uses_wallet_fallback()); + assert!(snapshot.new_reservation_context().uses_wallet_fallback()); + } + + #[tokio::test] + async fn exhausted_daily_quota_wallet_fallback_does_not_consume_long_plan_windows() { + use aether_data_contracts::repository::billing::UserDailyQuotaAvailabilityRecord; + let state = fallback_test_state(true, 20.0); + let auth = fallback_test_auth(); + let mut policy = fallback_test_policy(10); + policy.concurrency_limit = None; + state.auth_daily_quota_availability_cache.insert( + "user-1".into(), + Some(UserDailyQuotaAvailabilityRecord { + has_active_daily_quota: true, + total_quota_usd: 10.0, + used_usd: 10.0, + remaining_usd: 0.0, + allow_wallet_overage: true, + }), + Duration::from_secs(60), + ); + let (_, fallback) = check_and_acquire_compiled_plan_usage_policy( + &state, + &auth, + &policy, + "paid-request", + 2_000_000, + ) + .await + .unwrap(); + assert!(fallback); + // After the daily allowance recovers, the paid request must not exhaust + // the independent long-window plan request allowance. + state.auth_daily_quota_availability_cache.clear(); + let (_, fallback) = check_and_acquire_compiled_plan_usage_policy( + &state, + &auth, + &policy, + "plan-request", + 2_000_001, + ) + .await + .unwrap(); + assert!(!fallback); + } + + #[tokio::test] + async fn wallet_fallback_releases_reserved_cost_and_long_request_quota() { + let state = fallback_test_state(true, 20.0); + let auth = fallback_test_auth(); + let mut decision = GatewayControlDecision::synthetic( + "/v1/chat/completions", + Some("ai_public".into()), + Some("openai".into()), + Some("chat".into()), + Some("openai:chat".into()), + ); + decision.auth_context = Some(auth.clone()); + let policy = compile_effective_policy(&[entitlement("ent-cost", json!([{ + "type":"usage_policy", "rules":[ + {"metric":"actual_cost_usd","window":{"kind":"rolling","seconds":3600},"limit":0.1}, + {"metric":"request_count","window":{"kind":"rolling","seconds":3600},"limit":1} + ] + }]))], 2_000).unwrap(); + let (_, fallback) = check_and_acquire_compiled_plan_usage_policy( + &state, &auth, &policy, "event", 2_000_000, + ) + .await + .unwrap(); + assert!(!fallback); + let snapshot = PlanUsagePolicySnapshot::with_admission( + "user-1", + policy.clone(), + 2_000, + "event", + false, + ) + .unwrap(); + let reservation = snapshot.new_reservation_context(); + let rule = runtime_cost_rule(&policy.cost_rules[0], 2_000).unwrap(); + let previous = ReserveUsagePolicyCostInput { + request_id: "request-cost".into(), + subject_id: "user-1".into(), + reservation_token: reservation.token().into(), + admitted_at_unix_secs: 2_000, + reserved_cost_units: 5_000_000, + reservation_expires_at_unix_secs: 88_400, + retain_until_unix_secs: 32 * 86_400 + 2_000, + windows: vec![rule.window], + }; + assert!(matches!( + state + .data + .reserve_usage_policy_cost(previous.clone()) + .await + .unwrap(), + Some(ReserveUsagePolicyCostOutcome::Allowed { .. }) + )); + let plan = aether_contracts::ExecutionPlan { + request_id: "request-cost".into(), + candidate_id: None, + provider_name: Some("provider".into()), + provider_id: "provider-1".into(), + endpoint_id: "endpoint-1".into(), + key_id: "key-1".into(), + method: "POST".into(), + url: "https://example.com/v1/chat/completions".into(), + headers: Default::default(), + content_type: Some("application/json".into()), + content_encoding: None, + body: aether_contracts::RequestBody::from_json( + json!({"model":"gpt-5","messages":[],"max_tokens":16}), + ), + stream: false, + client_api_format: "openai:chat".into(), + provider_api_format: "openai:chat".into(), + model_name: Some("gpt-5".into()), + proxy: None, + transport_profile: None, + timeouts: None, + }; + let context = json!({"model_id":"model-1","global_model_name":"gpt-5"}); + assert_eq!( + reserve_admitted_http_plan_usage_policy_cost( + &state, + &decision, + &plan, + Some(&context), + Some(&reservation) + ) + .await + .unwrap(), + PlanUsageCostReservationOutcome::WalletFallback + ); + assert!(snapshot.uses_wallet_fallback()); + assert!(matches!( + state + .data + .reserve_usage_policy_cost(previous) + .await + .unwrap(), + Some(ReserveUsagePolicyCostOutcome::AlreadyTerminal { + state: UsagePolicyCostReservationState::Released + }) + )); + let (_, fallback) = check_and_acquire_compiled_plan_usage_policy( + &state, + &auth, + &policy, + "next-event", + 2_000_001, + ) + .await + .unwrap(); + assert!( + !fallback, + "paid fallback must release the long plan request window" + ); + } + fn entitlement(id: &str, snapshot: serde_json::Value) -> UserPlanEntitlementRecord { UserPlanEntitlementRecord { id: id.to_string(), diff --git a/apps/aether-gateway/src/state/runtime/billing/admin.rs b/apps/aether-gateway/src/state/runtime/billing/admin.rs index caca50d0e..74259f414 100644 --- a/apps/aether-gateway/src/state/runtime/billing/admin.rs +++ b/apps/aether-gateway/src/state/runtime/billing/admin.rs @@ -642,10 +642,18 @@ impl AppState { &self, user_id: &str, ) -> Result, GatewayError> { - self.data + let mut quota = self + .data .find_user_daily_quota_availability(user_id) .await - .map_err(data_error) + .map_err(data_error)?; + if let Some(quota) = quota.as_mut() { + quota.allow_wallet_overage = self + .read_user_preferences(user_id) + .await? + .is_some_and(|preferences| preferences.allow_wallet_overage); + } + Ok(quota) } pub(crate) async fn find_user_daily_quota_availability_for_auth( diff --git a/apps/aether-gateway/src/state/runtime/user_preferences.rs b/apps/aether-gateway/src/state/runtime/user_preferences.rs index 1f8fdf615..4f1689f8d 100644 --- a/apps/aether-gateway/src/state/runtime/user_preferences.rs +++ b/apps/aether-gateway/src/state/runtime/user_preferences.rs @@ -21,10 +21,79 @@ impl AppState { { let preferences = preferences.into(); let raw_preferences: crate::data::state::StoredUserPreferenceRecord = preferences.into(); - self.data + let persisted = self + .data .write_user_preferences(&raw_preferences) .await - .map(|value| value.map(Into::into)) - .map_err(|err| GatewayError::Internal(err.to_string())) + .map_err(|err| GatewayError::Internal(err.to_string()))?; + if persisted.is_some() { + self.invalidate_auth_context_cache(); + } + Ok(persisted.map(Into::into)) + } +} + +#[cfg(test)] +mod tests { + use std::{sync::Arc, time::Duration}; + + use aether_data::repository::users::InMemoryUserReadRepository; + + use super::*; + use crate::data::GatewayDataState; + + #[tokio::test] + async fn wallet_overage_preference_changes_clear_cached_billing_access() { + let repository = Arc::new(InMemoryUserReadRepository::default()); + let state = AppState::new() + .expect("state should build") + .with_data_state_for_tests(GatewayDataState::with_user_reader_for_tests(repository)); + let user_id = "user-wallet-overage".to_string(); + let ttl = Duration::from_secs(60); + let mut preferences = GatewayUserPreferenceView::default_for_user(&user_id); + + for enabled in [true, false] { + state + .auth_daily_quota_availability_cache + .insert(user_id.clone(), None, ttl); + state + .auth_plan_usage_policy_cache + .insert(user_id.clone(), None, ttl); + state + .auth_wallet_snapshot_cache + .insert(user_id.clone(), None, ttl); + assert!(state + .auth_daily_quota_availability_cache + .get(&user_id, ttl) + .is_some()); + assert!(state + .auth_plan_usage_policy_cache + .get(&user_id, ttl) + .is_some()); + assert!(state + .auth_wallet_snapshot_cache + .get(&user_id, ttl) + .is_some()); + + preferences.allow_wallet_overage = enabled; + let persisted = state + .write_user_preferences(&preferences) + .await + .expect("preferences should persist") + .expect("preferences should exist"); + assert_eq!(persisted.allow_wallet_overage, enabled); + assert!(state + .auth_daily_quota_availability_cache + .get(&user_id, ttl) + .is_none()); + assert!(state + .auth_plan_usage_policy_cache + .get(&user_id, ttl) + .is_none()); + assert!(state + .auth_wallet_snapshot_cache + .get(&user_id, ttl) + .is_none()); + } } } diff --git a/apps/aether-gateway/src/state/types.rs b/apps/aether-gateway/src/state/types.rs index 4c1f3a318..2e1351785 100644 --- a/apps/aether-gateway/src/state/types.rs +++ b/apps/aether-gateway/src/state/types.rs @@ -302,6 +302,8 @@ pub(crate) struct GatewayUserPreferenceView { pub(crate) email_notifications: bool, pub(crate) usage_alerts: bool, pub(crate) announcement_notifications: bool, + #[serde(default)] + pub(crate) allow_wallet_overage: bool, } impl GatewayUserPreferenceView { @@ -318,6 +320,7 @@ impl GatewayUserPreferenceView { email_notifications: true, usage_alerts: true, announcement_notifications: true, + allow_wallet_overage: false, } } } @@ -336,6 +339,7 @@ impl From for GatewayUserPrefere email_notifications: value.email_notifications, usage_alerts: value.usage_alerts, announcement_notifications: value.announcement_notifications, + allow_wallet_overage: value.allow_wallet_overage, } } } @@ -360,6 +364,7 @@ impl From for crate::data::state::StoredUserPreferenc email_notifications: value.email_notifications, usage_alerts: value.usage_alerts, announcement_notifications: value.announcement_notifications, + allow_wallet_overage: value.allow_wallet_overage, } } } diff --git a/apps/aether-gateway/src/tests/frontdoor/public_support.rs b/apps/aether-gateway/src/tests/frontdoor/public_support.rs index d3dc5424b..3c2b14deb 100644 --- a/apps/aether-gateway/src/tests/frontdoor/public_support.rs +++ b/apps/aether-gateway/src/tests/frontdoor/public_support.rs @@ -6225,6 +6225,7 @@ async fn gateway_handles_users_me_preferences_locally_without_proxying_upstream( assert_eq!(get_payload["theme"], "light"); assert_eq!(get_payload["language"], "zh-CN"); assert_eq!(get_payload["timezone"], "Asia/Shanghai"); + assert_eq!(get_payload["allow_wallet_overage"], false); assert_eq!(get_payload["notifications"]["email"], true); assert_eq!(get_payload["default_provider_id"], "provider-openai"); assert!(get_payload.get("default_provider").is_none()); @@ -6242,6 +6243,7 @@ async fn gateway_handles_users_me_preferences_locally_without_proxying_upstream( "email_notifications": false, "usage_alerts": false, "announcement_notifications": true, + "allow_wallet_overage": true, })) .send() .await @@ -6267,9 +6269,46 @@ async fn gateway_handles_users_me_preferences_locally_without_proxying_upstream( assert_eq!(verify_payload["language"], "en-US"); assert_eq!(verify_payload["timezone"], "UTC"); assert_eq!(verify_payload["bio"], "hello"); + assert_eq!(verify_payload["allow_wallet_overage"], true); assert_eq!(verify_payload["notifications"]["email"], false); assert_eq!(verify_payload["notifications"]["usage_alerts"], false); assert_eq!(verify_payload["notifications"]["announcements"], true); + for (payload, expected_status, expected_overage) in [ + ( + json!({"allow_wallet_overage": "true"}), + StatusCode::BAD_REQUEST, + true, + ), + (json!({"theme": "light"}), StatusCode::OK, true), + ( + json!({"allow_wallet_overage": false}), + StatusCode::OK, + false, + ), + ] { + let response = client + .put(format!("{gateway_url}/api/users/me/preferences")) + .header("authorization", format!("Bearer {access_token}")) + .header("x-client-device-id", "device-user-pref-1") + .header("user-agent", "AetherTest/1.0") + .json(&payload) + .send() + .await + .expect("preference update should succeed"); + assert_eq!(response.status(), expected_status); + let preferences: serde_json::Value = client + .get(format!("{gateway_url}/api/users/me/preferences")) + .header("authorization", format!("Bearer {access_token}")) + .header("x-client-device-id", "device-user-pref-1") + .header("user-agent", "AetherTest/1.0") + .send() + .await + .expect("preference lookup should succeed") + .json() + .await + .expect("preferences should parse"); + assert_eq!(preferences["allow_wallet_overage"], expected_overage); + } assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0); gateway_handle.abort(); diff --git a/apps/aether-gateway/src/usage/mod.rs b/apps/aether-gateway/src/usage/mod.rs index cef83b8f5..89d922c86 100644 --- a/apps/aether-gateway/src/usage/mod.rs +++ b/apps/aether-gateway/src/usage/mod.rs @@ -11,6 +11,7 @@ pub(crate) use aether_usage_runtime::{ }; pub(crate) use aether_usage_runtime::{UsageQueueHealthSnapshot, UsageRuntimeMetricsSnapshot}; pub(crate) use reporting::{ + attach_plan_wallet_fallback_context, attach_internal_gateway_report_capability, resolve_bound_internal_gateway_report_context, spawn_sync_report, submit_stream_report, submit_sync_report, GatewayStreamReportRequest, GatewaySyncReportRequest, diff --git a/apps/aether-gateway/src/usage/reporting/context.rs b/apps/aether-gateway/src/usage/reporting/context.rs index 5e8f5b92e..263950469 100644 --- a/apps/aether-gateway/src/usage/reporting/context.rs +++ b/apps/aether-gateway/src/usage/reporting/context.rs @@ -25,6 +25,72 @@ const INTERNAL_REPORT_CAPABILITY_MINT_ATTEMPTS: usize = 4; const PLAN_USAGE_RESERVATION_TOKEN_FIELD: &str = "plan_usage_reservation_token"; const PLAN_USAGE_RESERVATION_DEFERRED_FIELD: &str = "plan_usage_reservation_deferred"; +/// The funding decision is made after planning. Bind the server's decision to the +/// already-issued capability, rather than allowing a reporting peer to select it. +pub(crate) async fn attach_plan_wallet_fallback_context( + state: &AppState, + report_context: Option, + wallet_fallback: bool, +) -> Result, crate::GatewayError> { + let Some(Value::Object(mut context)) = report_context else { + return Ok(wallet_fallback.then(|| serde_json::json!({"plan_wallet_fallback": true}))); + }; + if !wallet_fallback && !context.contains_key("plan_wallet_fallback") { + return Ok(Some(Value::Object(context))); + } + let capability = context + .get(INTERNAL_REPORT_CAPABILITY_FIELD) + .and_then(Value::as_str) + .map(str::to_string); + let previous_digest = protected_internal_report_context_sha256(&context)?; + context.remove("plan_wallet_fallback"); + if wallet_fallback { + context.insert("plan_wallet_fallback".to_string(), Value::Bool(true)); + context.remove(PLAN_USAGE_RESERVATION_TOKEN_FIELD); + context.remove(PLAN_USAGE_RESERVATION_DEFERRED_FIELD); + } + let updated_digest = protected_internal_report_context_sha256(&context)?; + if updated_digest == previous_digest { + return Ok(Some(Value::Object(context))); + } + if let Some(capability) = capability { + let key = internal_report_capability_storage_key(&capability); + let serialized = state + .runtime_state + .kv_get(&key) + .await + .map_err(|error| crate::GatewayError::Internal(error.to_string()))? + .ok_or_else(|| { + crate::GatewayError::Internal( + "report capability expired before execution".to_string(), + ) + })?; + let mut record: InternalReportCapabilityRecord = serde_json::from_str(&serialized) + .map_err(|error| crate::GatewayError::Internal(error.to_string()))?; + if record.protected_context_sha256 == updated_digest { + return Ok(Some(Value::Object(context))); + } + if record.protected_context_sha256 != previous_digest + && record.kiro_web_search_context_sha256.as_deref() != Some(previous_digest.as_str()) + { + return Err(crate::GatewayError::Internal( + "report capability context changed before funding decision".to_string(), + )); + } + record.protected_context_sha256 = updated_digest; + record.kiro_web_search_context_sha256 = + kiro_web_search_internal_report_context_sha256(&context)?; + let serialized = serde_json::to_string(&record) + .map_err(|error| crate::GatewayError::Internal(error.to_string()))?; + state + .runtime_state + .kv_set(&key, serialized, Some(INTERNAL_REPORT_CAPABILITY_TTL)) + .await + .map_err(|error| crate::GatewayError::Internal(error.to_string()))?; + } + Ok(Some(Value::Object(context))) +} + /// Fields produced while observing an upstream response. Everything else in the /// planner-issued context is immutable and covered by the capability digest. /// @@ -450,3 +516,69 @@ fn video_task_matches_requested_user( }; task.user_id.as_deref().map(str::trim) == Some(requested_user_id) } + +#[cfg(test)] +mod wallet_fallback_tests { + use super::*; + use serde_json::json; + + #[tokio::test] + async fn wallet_fallback_is_server_bound_and_cannot_be_changed_by_a_reporting_peer() { + let state = AppState::new().unwrap(); + let mut original = Some(json!({"request_id":"request-1", "user_id":"user-1"})); + attach_internal_gateway_report_capability( + &state, + "trace-1", + Some("openai_chat_sync"), + &BTreeMap::new(), + &mut original, + ) + .await + .unwrap(); + let rebound = attach_plan_wallet_fallback_context(&state, original.clone(), true) + .await + .unwrap(); + let retried = attach_plan_wallet_fallback_context(&state, original, true) + .await + .unwrap(); + assert_eq!(rebound, retried); + let mut forged = rebound.clone().unwrap(); + forged + .as_object_mut() + .unwrap() + .remove("plan_wallet_fallback"); + assert!(resolve_bound_internal_gateway_report_context( + &state, + "trace-1", + "openai_chat_sync", + Some(&forged), + ) + .await + .unwrap() + .is_none()); + let verified = resolve_bound_internal_gateway_report_context( + &state, + "trace-1", + "openai_chat_sync", + rebound.as_ref(), + ) + .await + .unwrap() + .unwrap(); + assert_eq!(verified["plan_wallet_fallback"], true); + } + + #[tokio::test] + async fn wallet_fallback_seed_is_cleared_when_server_uses_plan() { + let state = AppState::new().unwrap(); + let context = attach_plan_wallet_fallback_context( + &state, + Some(json!({"plan_wallet_fallback":true})), + false, + ) + .await + .unwrap() + .unwrap(); + assert!(context.get("plan_wallet_fallback").is_none()); + } +} diff --git a/apps/aether-gateway/src/usage/reporting/mod.rs b/apps/aether-gateway/src/usage/reporting/mod.rs index 5d23d8b8b..6d55825d9 100644 --- a/apps/aether-gateway/src/usage/reporting/mod.rs +++ b/apps/aether-gateway/src/usage/reporting/mod.rs @@ -15,6 +15,7 @@ use crate::{AppState, GatewayError}; mod context; pub(crate) mod failure; pub(crate) use context::{ + attach_plan_wallet_fallback_context, attach_internal_gateway_report_capability, resolve_bound_internal_gateway_report_context, }; use context::{report_context_is_locally_actionable, resolve_locally_actionable_report_context}; diff --git a/apps/aether-gateway/src/wallet_runtime/access.rs b/apps/aether-gateway/src/wallet_runtime/access.rs index b68824db7..f0ed8e88f 100644 --- a/apps/aether-gateway/src/wallet_runtime/access.rs +++ b/apps/aether-gateway/src/wallet_runtime/access.rs @@ -271,6 +271,69 @@ mod tests { assert_eq!(decision.remaining, Some(4.0)); } + #[tokio::test] + async fn wallet_fallback_is_opt_in_and_preference_changes_refresh_cached_admission() { + let mut wallet = empty_user_wallet(); + wallet.balance = 10.0; + // A legacy plan allowing overage cannot silently enable wallet charges. + let state = state_with_wallet_and_quota(wallet, Some(quota_availability(10.0, 0.0, true))); + let auth_snapshot = ordinary_user_api_key_snapshot(); + let denied = resolve_wallet_auth_gate(&state, &auth_snapshot) + .await + .unwrap() + .unwrap(); + assert_eq!(denied.failure, Some(WalletAccessFailure::BalanceDenied)); + + let mut preferences = crate::GatewayUserPreferenceView::default_for_user("user-1"); + preferences.allow_wallet_overage = true; + state + .write_user_preferences(&preferences) + .await + .unwrap() + .unwrap(); + let allowed = resolve_wallet_auth_gate(&state, &auth_snapshot) + .await + .unwrap() + .unwrap(); + assert!(allowed.allowed); + assert_eq!(allowed.remaining, Some(10.0)); + + preferences.allow_wallet_overage = false; + state + .write_user_preferences(&preferences) + .await + .unwrap() + .unwrap(); + let denied_again = resolve_wallet_auth_gate(&state, &auth_snapshot) + .await + .unwrap() + .unwrap(); + assert_eq!( + denied_again.failure, + Some(WalletAccessFailure::BalanceDenied) + ); + } + + #[tokio::test] + async fn wallet_fallback_requires_available_wallet_balance() { + let state = state_with_wallet_and_quota( + empty_user_wallet(), + Some(quota_availability(10.0, 0.0, false)), + ); + let mut preferences = crate::GatewayUserPreferenceView::default_for_user("user-1"); + preferences.allow_wallet_overage = true; + state + .write_user_preferences(&preferences) + .await + .unwrap() + .unwrap(); + let denied = resolve_wallet_auth_gate(&state, &ordinary_user_api_key_snapshot()) + .await + .unwrap() + .unwrap(); + assert_eq!(denied.failure, Some(WalletAccessFailure::BalanceDenied)); + } + #[tokio::test] async fn unlimited_wallet_ignores_exhausted_non_overage_quota() { let mut wallet = empty_user_wallet(); @@ -373,7 +436,8 @@ mod tests { usage_repository, billing_repository, wallet_repository, - ); + ) + .with_user_preferences_for_tests([]); AppState::new() .expect("state should build") .with_data_state_for_tests(data) diff --git a/crates/aether-data/adapters/postgres/migrations/20261009000000_user_preferences_allow_wallet_overage.sql b/crates/aether-data/adapters/postgres/migrations/20261009000000_user_preferences_allow_wallet_overage.sql new file mode 100644 index 000000000..3cc6da4b2 --- /dev/null +++ b/crates/aether-data/adapters/postgres/migrations/20261009000000_user_preferences_allow_wallet_overage.sql @@ -0,0 +1,3 @@ +-- Wallet fallback is opt-in for both existing and newly created preferences. +ALTER TABLE public.user_preferences + ADD COLUMN IF NOT EXISTS allow_wallet_overage boolean NOT NULL DEFAULT false; diff --git a/crates/aether-data/adapters/postgres/src/billing.rs b/crates/aether-data/adapters/postgres/src/billing.rs index 285ae3e55..6994fefd7 100644 --- a/crates/aether-data/adapters/postgres/src/billing.rs +++ b/crates/aether-data/adapters/postgres/src/billing.rs @@ -1193,10 +1193,8 @@ WHERE id = $1 r#" SELECT user_plan_entitlements.id, - user_plan_entitlements.entitlements_snapshot, - billing_plans.entitlements_json AS plan_entitlements_json + user_plan_entitlements.entitlements_snapshot FROM user_plan_entitlements -JOIN billing_plans ON billing_plans.id = user_plan_entitlements.plan_id WHERE user_plan_entitlements.user_id = $1 AND user_plan_entitlements.status = 'active' AND user_plan_entitlements.starts_at <= NOW() @@ -1216,12 +1214,9 @@ ORDER BY user_plan_entitlements.expires_at ASC, let entitlement_id: String = row.try_get("id").map_postgres_err()?; let entitlements: serde_json::Value = row.try_get("entitlements_snapshot").map_postgres_err()?; - let plan_entitlements: serde_json::Value = - row.try_get("plan_entitlements_json").map_postgres_err()?; grants.extend(daily_quota_grants_from_entitlement( &entitlement_id, &entitlements, - daily_quota_wallet_overage_policy(&plan_entitlements), now, )?); } @@ -1229,9 +1224,15 @@ ORDER BY user_plan_entitlements.expires_at ASC, let mut total_quota_usd = 0.0; let mut used_usd = 0.0; let mut remaining_usd = 0.0; - let mut allow_wallet_overage = true; + let allow_wallet_overage = sqlx::query_scalar::<_, bool>( + "SELECT allow_wallet_overage FROM user_preferences WHERE user_id = $1", + ) + .bind(user_id) + .fetch_optional(&self.pool) + .await + .map_postgres_err()? + .unwrap_or(false); for grant in &grants { - allow_wallet_overage &= grant.allow_wallet_overage; let used = sqlx::query_scalar::<_, Option>( r#" SELECT CAST(COALESCE(SUM(amount_usd), 0) AS DOUBLE PRECISION) @@ -1336,7 +1337,6 @@ struct DailyQuotaGrant { entitlement_id: String, daily_quota_usd: f64, usage_date: String, - allow_wallet_overage: bool, } fn daily_quota_usage_date( @@ -1355,7 +1355,6 @@ fn daily_quota_usage_date( fn daily_quota_grants_from_entitlement( entitlement_id: &str, entitlements: &serde_json::Value, - current_allow_wallet_overage: Option, now: chrono::DateTime, ) -> Result, DataLayerError> { let mut grants = Vec::new(); @@ -1381,27 +1380,11 @@ fn daily_quota_grants_from_entitlement( .and_then(serde_json::Value::as_str), now, )?, - allow_wallet_overage: current_allow_wallet_overage.unwrap_or_else(|| { - item.get("allow_wallet_overage") - .and_then(serde_json::Value::as_bool) - .unwrap_or(false) - }), }); } Ok(grants) } -fn daily_quota_wallet_overage_policy(entitlements: &serde_json::Value) -> Option { - entitlements.as_array()?.iter().find_map(|item| { - (item.get("type").and_then(serde_json::Value::as_str) == Some("daily_quota")) - .then(|| { - item.get("allow_wallet_overage") - .and_then(serde_json::Value::as_bool) - }) - .flatten() - }) -} - fn map_payment_gateway_config_row( row: &sqlx::postgres::PgRow, ) -> Result { diff --git a/crates/aether-data/adapters/postgres/src/settlement.rs b/crates/aether-data/adapters/postgres/src/settlement.rs index e31704918..d3b169719 100644 --- a/crates/aether-data/adapters/postgres/src/settlement.rs +++ b/crates/aether-data/adapters/postgres/src/settlement.rs @@ -18,6 +18,7 @@ use crate::PostgresTransactionRunner; const FIND_USAGE_FOR_SETTLEMENT_SQL: &str = r#" SELECT usage_record.request_id, + usage_record.request_metadata, COALESCE(usage_settlement_snapshots.wallet_id, usage_record.wallet_id) AS wallet_id, COALESCE(usage_settlement_snapshots.billing_status, usage_record.billing_status) AS billing_status, COALESCE( @@ -465,7 +466,6 @@ struct DailyQuotaGrant { entitlement_id: String, daily_quota_usd: f64, usage_date: String, - allow_wallet_overage: bool, } fn daily_quota_usage_date( @@ -484,7 +484,6 @@ fn daily_quota_usage_date( fn daily_quota_grants_from_entitlement( entitlement_id: &str, entitlements: &serde_json::Value, - current_allow_wallet_overage: Option, now: chrono::DateTime, ) -> Result, DataLayerError> { let mut grants = Vec::new(); @@ -511,27 +510,11 @@ fn daily_quota_grants_from_entitlement( entitlement_id: entitlement_id.to_string(), daily_quota_usd, usage_date, - allow_wallet_overage: current_allow_wallet_overage.unwrap_or_else(|| { - item.get("allow_wallet_overage") - .and_then(serde_json::Value::as_bool) - .unwrap_or(false) - }), }); } Ok(grants) } -fn daily_quota_wallet_overage_policy(entitlements: &serde_json::Value) -> Option { - entitlements.as_array()?.iter().find_map(|item| { - (item.get("type").and_then(serde_json::Value::as_str) == Some("daily_quota")) - .then(|| { - item.get("allow_wallet_overage") - .and_then(serde_json::Value::as_bool) - }) - .flatten() - }) -} - async fn consume_daily_quota_postgres( tx: &mut crate::PostgresTransaction, user_id: &str, @@ -549,16 +532,13 @@ async fn consume_daily_quota_postgres( return Ok(DailyQuotaDebitResult::default()); } let now = chrono::Utc::now(); - // Serialize each entitlement's debits. Read the shared plan's current overage policy - // from this statement's snapshot without locking every subscriber's plan row. + // Serialize each entitlement's debits without locking the shared plan row. let entitlement_rows = sqlx::query( r#" SELECT user_plan_entitlements.id, - user_plan_entitlements.entitlements_snapshot, - billing_plans.entitlements_json AS plan_entitlements_json + user_plan_entitlements.entitlements_snapshot FROM user_plan_entitlements -JOIN billing_plans ON billing_plans.id = user_plan_entitlements.plan_id WHERE user_plan_entitlements.user_id = $1 AND user_plan_entitlements.status = 'active' AND user_plan_entitlements.starts_at <= NOW() @@ -578,12 +558,9 @@ FOR UPDATE OF user_plan_entitlements let entitlement_id: String = row.try_get("id").map_postgres_err()?; let entitlements: serde_json::Value = row.try_get("entitlements_snapshot").map_postgres_err()?; - let plan_entitlements: serde_json::Value = - row.try_get("plan_entitlements_json").map_postgres_err()?; grants.extend(daily_quota_grants_from_entitlement( &entitlement_id, &entitlements, - daily_quota_wallet_overage_policy(&plan_entitlements), now, )?); } @@ -593,9 +570,16 @@ FOR UPDATE OF user_plan_entitlements let mut grants_with_remaining = Vec::new(); let mut total_remaining = 0.0; - let mut allow_wallet_overage = true; + // Wallet fallback is the user's explicit choice, never the plan's legacy default. + let allow_wallet_overage = sqlx::query_scalar::<_, bool>( + "SELECT allow_wallet_overage FROM user_preferences WHERE user_id = $1", + ) + .bind(user_id) + .fetch_optional(&mut **tx) + .await + .map_postgres_err()? + .unwrap_or(false); for grant in grants { - allow_wallet_overage &= grant.allow_wallet_overage; let used = sqlx::query_scalar::<_, Option>( r#" SELECT CAST(COALESCE(SUM(amount_usd), 0) AS DOUBLE PRECISION) @@ -1293,33 +1277,43 @@ LIMIT 1 } let billable_cost_usd = settlement_billable_cost_usd(&input); - let wallet_debit_cost_usd = if !api_key_is_standalone { - if let Some(user_id) = - input.user_id.as_deref().filter(|value| !value.is_empty()) - { - let quota = consume_daily_quota_postgres( - tx, - user_id, - &input.request_id, - billable_cost_usd, - wallet_available_usd, - wallet_can_overdraft, - ) - .await?; - if quota.insufficient { - final_billing_status = "insufficient_quota".to_string(); - settlement.billing_status = final_billing_status.clone(); - 0.0 + let plan_wallet_fallback = usage_row + .try_get::, _>("request_metadata") + .map_postgres_err()? + .and_then(|metadata| { + metadata + .get("plan_wallet_fallback") + .and_then(serde_json::Value::as_bool) + }) + .unwrap_or(false); + let wallet_debit_cost_usd = + if !api_key_is_standalone && !plan_wallet_fallback { + if let Some(user_id) = + input.user_id.as_deref().filter(|value| !value.is_empty()) + { + let quota = consume_daily_quota_postgres( + tx, + user_id, + &input.request_id, + billable_cost_usd, + wallet_available_usd, + wallet_can_overdraft, + ) + .await?; + if quota.insufficient { + final_billing_status = "insufficient_quota".to_string(); + settlement.billing_status = final_billing_status.clone(); + 0.0 + } else { + quota_covered = quota.debited_usd; + (billable_cost_usd - quota.debited_usd).max(0.0) + } } else { - quota_covered = quota.debited_usd; - (billable_cost_usd - quota.debited_usd).max(0.0) + billable_cost_usd } } else { billable_cost_usd - } - } else { - billable_cost_usd - }; + }; if final_billing_status != "settled" { sync_usage_settlement_snapshot(&mut **tx, &settlement).await?; sqlx::query(FINALIZE_USAGE_BILLING_SQL) @@ -1496,6 +1490,7 @@ mod tests { "user_plan_entitlements", "entitlement_usage_ledgers", "users", + "user_preferences", "usage_request_admissions", "usage_cost_reservations", ] { @@ -1524,17 +1519,20 @@ mod tests { for (scenario, charge, quota_covered) in [ ("wallet", 20.0, 0.0), ("quota_and_wallet", 20.0, 7.0), + ("wallet_fallback", 20.0, 0.0), ("zero_charge", 0.0, 0.0), ] { sqlx::query("INSERT INTO users (id, username, email_verified) VALUES ($1, $1, false)") .bind(scenario).execute(&pool).await.expect("user should insert"); + sqlx::query("INSERT INTO user_preferences (id, user_id, allow_wallet_overage) VALUES ($1, $1, true)") + .bind(scenario).execute(&pool).await.expect("wallet fallback preference should insert"); sqlx::query("INSERT INTO wallets (id, user_id, balance, gift_balance, total_consumed, limit_mode, created_at, updated_at) VALUES ($1, $1, 100, 0, 0, 'finite', NOW(), NOW())") .bind(scenario).execute(&pool).await.expect("wallet should insert"); // A zero-charge request must leave an active quota untouched too. if scenario != "wallet" { let grant = serde_json::json!([{ "type": "daily_quota", "daily_quota_usd": 7.0, - "reset_timezone": "UTC", "allow_wallet_overage": true, + "reset_timezone": "UTC", "allow_wallet_overage": false, }]); sqlx::query("INSERT INTO billing_plans (id, title, price_amount, duration_unit, duration_value, entitlements_json, created_at, updated_at) VALUES ($1, $1, 10, 'month', 1, $2, NOW(), NOW())") .bind(scenario).bind(&grant).execute(&pool).await.expect("plan should insert"); @@ -1542,7 +1540,7 @@ mod tests { .bind(scenario).bind(&grant).execute(&pool).await.expect("entitlement should insert"); } let multiplier = charge / 10.0; - let metadata = serde_json::json!({"billing_multiplier_snapshot": { + let metadata = serde_json::json!({"plan_wallet_fallback": scenario == "wallet_fallback", "billing_multiplier_snapshot": { "version": 1, "factors": {"routing_group": multiplier}, "multiplier": multiplier, }}); sqlx::query("INSERT INTO usage (id, request_id, user_id, provider_id, provider_name, model, status, billing_status, total_cost_usd, actual_total_cost_usd, request_metadata) VALUES ($1, $1, $1, 'provider', 'Provider', 'model', 'completed', 'pending', 10, 5, $2)") @@ -1701,6 +1699,8 @@ mod tests { .await .expect("shared plan should insert"); for user_id in ["user-a", "user-b"] { + sqlx::query("INSERT INTO users (id, username, email_verified) VALUES ($1, $1, false)") + .bind(user_id).execute(&pool).await.expect("user should insert"); sqlx::query( "INSERT INTO user_plan_entitlements (id, user_id, plan_id, payment_order_id, starts_at, expires_at, entitlements_snapshot, created_at, updated_at) VALUES ($1, $1, 'shared-plan', $1, NOW() - INTERVAL '1 hour', NOW() + INTERVAL '1 day', $2, NOW(), NOW())", ) @@ -1786,14 +1786,52 @@ mod tests { held.rollback().await.expect("held quota debit should roll back"); let mut after_edit = pool.begin().await.expect("fresh transaction should start"); + let unchanged_policy = super::consume_daily_quota_postgres( + &mut after_edit, "user-a", "request-plan-policy-after", 2.0, Some(5.0), true, + ) + .await + .expect("plan changes must not opt the user into wallet charges"); + assert!(unchanged_policy.insufficient); + after_edit.rollback().await.expect("default policy verification should roll back"); + sqlx::query("INSERT INTO user_preferences (id, user_id, allow_wallet_overage) VALUES ('user-a', 'user-a', true)") + .execute(&pool).await.expect("user should enable wallet fallback"); + let mut after_edit = pool.begin().await.expect("user preference transaction should start"); let updated_policy = super::consume_daily_quota_postgres( &mut after_edit, "user-a", "request-policy-after", 2.0, Some(5.0), true, ) .await - .expect("fresh quota read should use current plan configuration"); + .expect("fresh quota read should use the user preference"); assert!(!updated_policy.insufficient); assert_eq!(updated_policy.debited_usd, 1.0); - after_edit.rollback().await.expect("policy verification should roll back"); + after_edit.commit().await.expect("enabled wallet fallback should commit"); + let mut exhausted = pool.begin().await.expect("exhausted quota transaction should start"); + let debit = super::consume_daily_quota_postgres( + &mut exhausted, "user-a", "request-exhausted", 2.0, Some(5.0), false, + ).await.expect("exhausted quota should fall back to wallet"); + assert!(!debit.insufficient); + assert_eq!(debit.debited_usd, 0.0); + let debit = super::consume_daily_quota_postgres( + &mut exhausted, "user-a", "request-exhausted-no-wallet", 2.0, Some(0.0), false, + ).await.expect("empty wallet should reject exhausted quota fallback"); + assert!(debit.insufficient); + exhausted.rollback().await.expect("exhausted quota verification should roll back"); + + use aether_data_contracts::repository::billing::BillingReadRepository; + let billing = crate::SqlxBillingReadRepository::new(pool.clone()); + let quota = billing.find_user_daily_quota_availability("user-a").await.unwrap().unwrap(); + assert!(quota.allow_wallet_overage); + assert_eq!(quota.remaining_usd, 0.0); + sqlx::query("UPDATE user_preferences SET allow_wallet_overage = false WHERE user_id = 'user-a'") + .execute(&pool).await.expect("user should disable wallet fallback"); + let quota = billing.find_user_daily_quota_availability("user-a").await.unwrap().unwrap(); + assert!(!quota.allow_wallet_overage); + let mut disabled = pool.begin().await.expect("disabled fallback transaction should start"); + let debit = super::consume_daily_quota_postgres( + &mut disabled, "user-a", "request-disabled", 2.0, Some(5.0), true, + ).await.expect("disabled preference should prevent wallet charges"); + assert!(debit.insufficient); + assert_eq!(debit.debited_usd, 0.0); + disabled.rollback().await.expect("disabled fallback verification should roll back"); }) .catch_unwind() .await; diff --git a/crates/aether-data/adapters/postgres/src/usage/analytics_tests.rs b/crates/aether-data/adapters/postgres/src/usage/analytics_tests.rs index 9d0328016..9ec9d5565 100644 --- a/crates/aether-data/adapters/postgres/src/usage/analytics_tests.rs +++ b/crates/aether-data/adapters/postgres/src/usage/analytics_tests.rs @@ -665,6 +665,8 @@ async fn live_overview_settlement_allocations_preserve_unlimited_and_finite_wall _ => 0.0, }; if quota > 0.0 { + sqlx::query("INSERT INTO user_preferences(id,user_id,allow_wallet_overage) VALUES($1,$1,true)") + .bind(&user).execute(&pool).await.unwrap(); let grant = serde_json::json!([{"type":"daily_quota","daily_quota_usd":quota,"reset_timezone":"UTC","allow_wallet_overage":true}]); sqlx::query("INSERT INTO billing_plans(id,title,price_amount,duration_unit,duration_value,entitlements_json,created_at,updated_at) VALUES($1,'test',1,'month',1,$2,NOW(),NOW())").bind(&user).bind(&grant).execute(&pool).await.unwrap(); sqlx::query("INSERT INTO payment_orders(id,order_no,wallet_id,user_id,amount_usd,payment_method,created_at) VALUES($1,$1,$2,$1,1,'test',NOW())").bind(&user).bind(&wallet).execute(&pool).await.unwrap(); diff --git a/crates/aether-data/adapters/postgres/src/users.rs b/crates/aether-data/adapters/postgres/src/users.rs index f11afc6f3..7e65f98ca 100644 --- a/crates/aether-data/adapters/postgres/src/users.rs +++ b/crates/aether-data/adapters/postgres/src/users.rs @@ -524,7 +524,8 @@ SELECT up.timezone, up.email_notifications, up.usage_alerts, - up.announcement_notifications + up.announcement_notifications, + up.allow_wallet_overage FROM user_preferences up LEFT JOIN providers p ON p.id = up.default_provider_id @@ -546,6 +547,7 @@ WITH upserted AS ( email_notifications, usage_alerts, announcement_notifications, + allow_wallet_overage, created_at, updated_at ) VALUES ( @@ -560,6 +562,7 @@ WITH upserted AS ( $9, $10, $11, + $12, NOW(), NOW() ) @@ -573,6 +576,7 @@ WITH upserted AS ( email_notifications = EXCLUDED.email_notifications, usage_alerts = EXCLUDED.usage_alerts, announcement_notifications = EXCLUDED.announcement_notifications, + allow_wallet_overage = EXCLUDED.allow_wallet_overage, updated_at = NOW() RETURNING user_id, @@ -584,7 +588,8 @@ WITH upserted AS ( timezone, email_notifications, usage_alerts, - announcement_notifications + announcement_notifications, + allow_wallet_overage ) SELECT upserted.user_id, @@ -597,7 +602,8 @@ SELECT upserted.timezone, upserted.email_notifications, upserted.usage_alerts, - upserted.announcement_notifications + upserted.announcement_notifications, + upserted.allow_wallet_overage FROM upserted LEFT JOIN providers p ON p.id = upserted.default_provider_id @@ -2931,6 +2937,7 @@ LIMIT 1 .bind(preferences.email_notifications) .bind(preferences.usage_alerts) .bind(preferences.announcement_notifications) + .bind(preferences.allow_wallet_overage) .fetch_optional(&self.pool) .await .map_postgres_err()?; @@ -3220,6 +3227,7 @@ fn map_user_preference_row( announcement_notifications: row .try_get("announcement_notifications") .map_postgres_err()?, + allow_wallet_overage: row.try_get("allow_wallet_overage").map_postgres_err()?, }) } diff --git a/crates/aether-data/contracts/src/repository/usage/metadata_policy.rs b/crates/aether-data/contracts/src/repository/usage/metadata_policy.rs index 41f4cce75..6942e157c 100644 --- a/crates/aether-data/contracts/src/repository/usage/metadata_policy.rs +++ b/crates/aether-data/contracts/src/repository/usage/metadata_policy.rs @@ -82,13 +82,36 @@ pub fn preserve_usage_routing_group_snapshot( captured, ); } - let Some(Value::Object(snapshot)) = sanitize_usage_request_metadata_object(&snapshot) else { + let snapshot = sanitize_usage_request_metadata_object(&snapshot) + .and_then(|value| value.as_object().cloned()) + .unwrap_or_default(); + if snapshot.is_empty() + && previous + .get("plan_wallet_fallback") + .and_then(Value::as_bool) + != Some(true) + { return incoming; - }; + } let mut metadata = incoming .and_then(|value| value.as_object().cloned()) .unwrap_or_default(); metadata.extend(snapshot); + if previous + .get("plan_wallet_fallback") + .and_then(Value::as_bool) + == Some(true) + { + metadata.insert("plan_wallet_fallback".to_string(), Value::Bool(true)); + } + if metadata + .get("plan_wallet_fallback") + .and_then(Value::as_bool) + == Some(true) + { + metadata.remove(PLAN_USAGE_RESERVATION_TOKEN_KEY); + metadata.remove(PLAN_USAGE_RESERVATION_DEFERRED_METADATA_KEY); + } Some(Value::Object(metadata)) } @@ -152,6 +175,7 @@ pub fn sanitize_usage_request_metadata_object(source: &Map) -> Op "client_requested_stream", UPSTREAM_IS_STREAM_KEY, "api_key_is_standalone", + "plan_wallet_fallback", WEBSOCKET_MODE_METADATA_KEY, PLAN_USAGE_RESERVATION_DEFERRED_METADATA_KEY, "transport_error", @@ -1369,6 +1393,26 @@ mod tests { sanitize_usage_request_metadata, sanitize_usage_request_metadata_ref, }; + #[test] + fn wallet_fallback_survives_sparse_video_completion_and_removes_plan_reservation() { + let token = "550e8400-e29b-41d4-a716-446655440000"; + let previous = json!({"plan_wallet_fallback": true, "plan_usage_reservation_token": token}); + let preserved = preserve_usage_routing_group_snapshot(None, Some(&previous)).unwrap(); + assert_eq!(preserved["plan_wallet_fallback"], true); + assert!(preserved.get("plan_usage_reservation_token").is_none()); + let incoming = json!({"plan_wallet_fallback": true}); + let previous = + json!({"plan_wallet_fallback": false, "plan_usage_reservation_token": token}); + let preserved = + preserve_usage_routing_group_snapshot(Some(incoming), Some(&previous)).unwrap(); + assert_eq!(preserved["plan_wallet_fallback"], true); + assert!(preserved.get("plan_usage_reservation_token").is_none()); + assert!( + sanitize_usage_request_metadata(Some(json!({"plan_wallet_fallback": "true"}))) + .is_none() + ); + } + #[test] fn billing_multiplier_snapshot_projection_preserves_invalid_marker_and_immutable_factors() { for snapshot in [ diff --git a/crates/aether-data/contracts/src/repository/users.rs b/crates/aether-data/contracts/src/repository/users.rs index 7ecc13fee..5f8a2577b 100644 --- a/crates/aether-data/contracts/src/repository/users.rs +++ b/crates/aether-data/contracts/src/repository/users.rs @@ -776,6 +776,8 @@ pub struct StoredUserPreferenceRecord { pub email_notifications: bool, pub usage_alerts: bool, pub announcement_notifications: bool, + #[serde(default)] + pub allow_wallet_overage: bool, } #[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] @@ -925,6 +927,7 @@ impl StoredUserPreferenceRecord { email_notifications: true, usage_alerts: true, announcement_notifications: true, + allow_wallet_overage: false, } } } @@ -1967,5 +1970,14 @@ mod tests { assert!(record.email_notifications); assert!(record.usage_alerts); assert!(record.announcement_notifications); + assert!(!record.allow_wallet_overage); + let mut legacy_json = serde_json::to_value(&record).expect("preferences should serialize"); + legacy_json + .as_object_mut() + .expect("preferences object") + .remove("allow_wallet_overage"); + let restored: StoredUserPreferenceRecord = + serde_json::from_value(legacy_json).expect("legacy preferences should deserialize"); + assert!(!restored.allow_wallet_overage); } } 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 ea8c52c1e..f80838805 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 @@ -1349,6 +1349,7 @@ CREATE TABLE IF NOT EXISTS public.user_preferences ( email_notifications boolean DEFAULT true, usage_alerts boolean DEFAULT true, announcement_notifications boolean DEFAULT true, + allow_wallet_overage boolean DEFAULT false 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/generated/postgres/baseline/001_identity.sql b/crates/aether-data/runtime/schema/generated/postgres/baseline/001_identity.sql index 0f6b785c2..6ae3371b0 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 @@ -191,6 +191,7 @@ CREATE TABLE IF NOT EXISTS public.user_preferences ( email_notifications boolean DEFAULT true NOT NULL, usage_alerts boolean DEFAULT true NOT NULL, announcement_notifications boolean DEFAULT true NOT NULL, + allow_wallet_overage boolean DEFAULT false NOT NULL, created_at bigint NOT NULL, updated_at bigint NOT NULL ); diff --git a/crates/aether-data/runtime/schema/logical/001_identity.toml b/crates/aether-data/runtime/schema/logical/001_identity.toml index 15ce325e6..551b37ac1 100644 --- a/crates/aether-data/runtime/schema/logical/001_identity.toml +++ b/crates/aether-data/runtime/schema/logical/001_identity.toml @@ -788,6 +788,11 @@ name = "announcement_notifications" type = "bool" default = true +[[table.user_preferences.columns]] +name = "allow_wallet_overage" +type = "bool" +default = false + [[table.user_preferences.columns]] name = "created_at" type = "unix_seconds" diff --git a/crates/aether-data/runtime/src/lifecycle/migrate/tests.rs b/crates/aether-data/runtime/src/lifecycle/migrate/tests.rs index a59054565..1e9f1d9b1 100644 --- a/crates/aether-data/runtime/src/lifecycle/migrate/tests.rs +++ b/crates/aether-data/runtime/src/lifecycle/migrate/tests.rs @@ -37,6 +37,7 @@ mod overview_migration_safety; mod policy_nulls; mod provider_expenses; mod scoped_provider_creation; +mod user_preferences; /// A clean PostgreSQL database is bootstrapped from the schema snapshot first; /// migrations after the privacy/security frontier are intentionally left @@ -1599,6 +1600,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() { 20261001000000, 20261004000000, 20261007000000, + 20261009000000, ] ); } diff --git a/crates/aether-data/runtime/src/lifecycle/migrate/tests/legacy_overview_upgrade.rs b/crates/aether-data/runtime/src/lifecycle/migrate/tests/legacy_overview_upgrade.rs index 86da747be..49fdf1ceb 100644 --- a/crates/aether-data/runtime/src/lifecycle/migrate/tests/legacy_overview_upgrade.rs +++ b/crates/aether-data/runtime/src/lifecycle/migrate/tests/legacy_overview_upgrade.rs @@ -116,6 +116,7 @@ WHERE version=20260919000000; 20261001000000, 20261004000000, 20261007000000, + 20261009000000, ] ); assert_eq!( diff --git a/crates/aether-data/runtime/src/lifecycle/migrate/tests/user_preferences.rs b/crates/aether-data/runtime/src/lifecycle/migrate/tests/user_preferences.rs new file mode 100644 index 000000000..2e64bd7b9 --- /dev/null +++ b/crates/aether-data/runtime/src/lifecycle/migrate/tests/user_preferences.rs @@ -0,0 +1,93 @@ +use aether_data_contracts::repository::users::StoredUserPreferenceRecord; +use aether_data_postgres::SqlxUserReadRepository; +use sqlx::{query, query_scalar, PgPool}; + +use super::{ManagedPostgresServer, POSTGRES_MIGRATOR}; + +#[tokio::test] +async fn wallet_overage_preference_upgrade_defaults_off_and_round_trips() { + let Some(server) = ManagedPostgresServer::try_start() + .await + .expect("postgres should start") + else { + return; + }; + let pool = PgPool::connect(server.database_url()) + .await + .expect("database should connect"); + sqlx::raw_sql( + r#" +CREATE TABLE providers (id text PRIMARY KEY, name text NOT NULL); +CREATE TABLE user_preferences ( + id text PRIMARY KEY, user_id text UNIQUE NOT NULL, + avatar_url text, bio text, default_provider_id text, + theme text NOT NULL DEFAULT 'light', language text NOT NULL DEFAULT 'zh-CN', + timezone text NOT NULL DEFAULT 'Asia/Shanghai', + email_notifications boolean NOT NULL DEFAULT true, + usage_alerts boolean NOT NULL DEFAULT true, + announcement_notifications boolean NOT NULL DEFAULT true, + created_at timestamptz DEFAULT now(), updated_at timestamptz DEFAULT now() +); +INSERT INTO user_preferences (id, user_id, theme) VALUES ('legacy-pref', 'legacy-user', 'dark'); +"#, + ) + .execute(&pool) + .await + .expect("legacy preferences should seed"); + let migration = POSTGRES_MIGRATOR + .iter() + .find(|migration| migration.version == 20261009000000) + .expect("wallet overage migration should be embedded"); + for _ in 0..2 { + sqlx::raw_sql(&migration.sql) + .execute(&pool) + .await + .expect("migration should apply idempotently"); + } + let repository = SqlxUserReadRepository::new(pool.clone()); + let mut preferences = repository + .read_user_preferences("legacy-user") + .await + .expect("legacy preferences should read") + .expect("legacy preferences should exist"); + assert!(!preferences.allow_wallet_overage); + assert_eq!(preferences.theme, "dark"); + for enabled in [true, false] { + preferences.allow_wallet_overage = enabled; + assert_eq!( + repository + .write_user_preferences(&preferences) + .await + .expect("preferences should update"), + Some(preferences.clone()) + ); + assert_eq!( + repository + .read_user_preferences("legacy-user") + .await + .expect("preferences should read"), + Some(preferences.clone()) + ); + } + query("INSERT INTO user_preferences (id, user_id) VALUES ('new-pref', 'new-user')") + .execute(&pool) + .await + .expect("new preferences should insert"); + let default: bool = query_scalar( + "SELECT allow_wallet_overage FROM user_preferences WHERE user_id = 'new-user'", + ) + .fetch_one(&pool) + .await + .expect("new default should read"); + assert!(!default); + let mut enabled = StoredUserPreferenceRecord::default_for_user("enabled-user"); + enabled.allow_wallet_overage = true; + assert_eq!( + repository + .write_user_preferences(&enabled) + .await + .expect("new enabled preferences should insert"), + Some(enabled) + ); + pool.close().await; +} diff --git a/crates/aether-data/runtime/src/repository/billing/memory.rs b/crates/aether-data/runtime/src/repository/billing/memory.rs index c05c50246..1f57f2a36 100644 --- a/crates/aether-data/runtime/src/repository/billing/memory.rs +++ b/crates/aether-data/runtime/src/repository/billing/memory.rs @@ -92,14 +92,12 @@ fn billing_plan_from_input( fn daily_quota_availability_from_entitlements( entitlements: impl IntoIterator, - billing_plans: &BTreeMap, now: u64, ) -> UserDailyQuotaAvailabilityRecord { let mut has_active_daily_quota = false; let mut total_quota_usd = 0.0; let used_usd = 0.0; let mut remaining_usd = 0.0; - let mut allow_wallet_overage = true; for entitlement in entitlements { if entitlement.status != "active" || entitlement.starts_at_unix_secs > now @@ -110,9 +108,6 @@ fn daily_quota_availability_from_entitlements( let Some(items) = entitlement.entitlements_snapshot.as_array() else { continue; }; - let current_allow_wallet_overage = billing_plans - .get(&entitlement.plan_id) - .and_then(|plan| daily_quota_wallet_overage_policy(&plan.entitlements_json)); for item in items { if item.get("type").and_then(serde_json::Value::as_str) != Some("daily_quota") { continue; @@ -127,11 +122,6 @@ fn daily_quota_availability_from_entitlements( has_active_daily_quota = true; total_quota_usd += daily_quota_usd; remaining_usd += daily_quota_usd; - allow_wallet_overage &= current_allow_wallet_overage.unwrap_or_else(|| { - item.get("allow_wallet_overage") - .and_then(serde_json::Value::as_bool) - .unwrap_or(false) - }); } } UserDailyQuotaAvailabilityRecord { @@ -139,21 +129,11 @@ fn daily_quota_availability_from_entitlements( total_quota_usd, used_usd, remaining_usd, - allow_wallet_overage, + // The gateway combines this quota with the user's stored preference. + allow_wallet_overage: false, } } -fn daily_quota_wallet_overage_policy(entitlements: &serde_json::Value) -> Option { - entitlements.as_array()?.iter().find_map(|item| { - (item.get("type").and_then(serde_json::Value::as_str) == Some("daily_quota")) - .then(|| { - item.get("allow_wallet_overage") - .and_then(serde_json::Value::as_bool) - }) - .flatten() - }) -} - #[async_trait] impl BillingReadRepository for InMemoryBillingReadRepository { async fn find_model_context( @@ -578,13 +558,8 @@ impl BillingReadRepository for InMemoryBillingReadRepository { .filter(|item| item.user_id == user_id) .cloned() .collect::>(); - let billing_plans = self - .billing_plans_by_id - .read() - .expect("billing repository lock"); Ok(Some(daily_quota_availability_from_entitlements( entitlements, - &billing_plans, now, ))) } diff --git a/crates/aether-data/runtime/src/repository/users/memory.rs b/crates/aether-data/runtime/src/repository/users/memory.rs index 743a0d044..c4b9c5e26 100644 --- a/crates/aether-data/runtime/src/repository/users/memory.rs +++ b/crates/aether-data/runtime/src/repository/users/memory.rs @@ -3226,6 +3226,7 @@ mod tests { email_notifications: true, usage_alerts: true, announcement_notifications: true, + allow_wallet_overage: true, }; let repository = InMemoryUserReadRepository::seed_auth_users([admin.clone()]) .with_export_users([export_row]) @@ -4029,6 +4030,7 @@ mod tests { email_notifications: false, usage_alerts: true, announcement_notifications: false, + allow_wallet_overage: true, }; assert!(repository diff --git a/crates/aether-usage/runtime/src/request_metadata.rs b/crates/aether-usage/runtime/src/request_metadata.rs index 8a6ff9428..6671241a3 100644 --- a/crates/aether-usage/runtime/src/request_metadata.rs +++ b/crates/aether-usage/runtime/src/request_metadata.rs @@ -103,6 +103,7 @@ pub(crate) fn retain_first_byte_request_metadata(value: Option) -> Option | "client_requested_stream" | "upstream_is_stream" | "api_key_is_standalone" + | "plan_wallet_fallback" | "plan_usage_reservation_token" | "request_path" | "request_query_string" diff --git a/crates/aether-usage/runtime/src/settlement.rs b/crates/aether-usage/runtime/src/settlement.rs index d54d3a9bb..6e86b0a28 100644 --- a/crates/aether-usage/runtime/src/settlement.rs +++ b/crates/aether-usage/runtime/src/settlement.rs @@ -286,6 +286,15 @@ fn usage_api_key_is_standalone(usage: &StoredRequestUsageAudit) -> bool { } fn usage_policy_reservation_token(usage: &StoredRequestUsageAudit) -> Option<&str> { + if usage + .request_metadata + .as_ref() + .and_then(|metadata| metadata.get("plan_wallet_fallback")) + .and_then(serde_json::Value::as_bool) + == Some(true) + { + return None; + } usage .request_metadata .as_ref() @@ -296,6 +305,16 @@ fn usage_policy_reservation_token(usage: &StoredRequestUsageAudit) -> Option<&st } fn event_usage_policy_reservation_token(event: &UsageEvent) -> Option<&str> { + if event + .data + .request_metadata + .as_ref() + .and_then(|metadata| metadata.get("plan_wallet_fallback")) + .and_then(serde_json::Value::as_bool) + == Some(true) + { + return None; + } event .data .request_metadata diff --git a/crates/aether-usage/runtime/src/write.rs b/crates/aether-usage/runtime/src/write.rs index 3e2956352..010732fe8 100644 --- a/crates/aether-usage/runtime/src/write.rs +++ b/crates/aether-usage/runtime/src/write.rs @@ -2246,6 +2246,12 @@ fn build_runtime_request_metadata_seed_from_parts( Value::Bool(api_key_is_standalone), ); } + if let Some(wallet_fallback) = context_bool(context, "plan_wallet_fallback") { + metadata.insert( + "plan_wallet_fallback".to_string(), + Value::Bool(wallet_fallback), + ); + } if let Some(websocket_mode) = context_bool(context, WEBSOCKET_MODE_METADATA_KEY) { metadata.insert( WEBSOCKET_MODE_METADATA_KEY.to_string(), diff --git a/frontend/src/api/auth.ts b/frontend/src/api/auth.ts index 30c35610d..eb974a434 100644 --- a/frontend/src/api/auth.ts +++ b/frontend/src/api/auth.ts @@ -21,6 +21,7 @@ export interface UserPreferences { theme?: 'light' | 'dark' | 'auto' language?: string notifications_enabled?: boolean + allow_wallet_overage?: boolean [key: string]: unknown // 允许扩展其他偏好设置 } diff --git a/frontend/src/api/me.ts b/frontend/src/api/me.ts index f0f65e93e..4e26e2e8d 100644 --- a/frontend/src/api/me.ts +++ b/frontend/src/api/me.ts @@ -42,6 +42,7 @@ export interface UserPreferences { theme: string language: string timezone?: string + allow_wallet_overage?: boolean notifications?: { email?: boolean usage_alerts?: boolean diff --git a/frontend/src/mocks/handler.ts b/frontend/src/mocks/handler.ts index 586834e76..d0beff1ae 100644 --- a/frontend/src/mocks/handler.ts +++ b/frontend/src/mocks/handler.ts @@ -1609,11 +1609,22 @@ const mockHandlers: Record Promise { await delay() - return createMockResponse(getCurrentProfile().preferences || { theme: 'auto', language: 'zh-CN' }) + return createMockResponse({ + allow_wallet_overage: false, + ...(getCurrentProfile().preferences || { theme: 'auto', language: 'zh-CN' }), + }) }, - 'PUT /api/users/me/preferences': async () => { + 'PUT /api/users/me/preferences': async (config) => { await delay() + const payload = mockRequestObject(config) + if (typeof payload.allow_wallet_overage === 'boolean') { + const profile = getCurrentProfile() + profile.preferences = { + ...(profile.preferences || { theme: 'auto', language: 'zh-CN' }), + allow_wallet_overage: payload.allow_wallet_overage, + } + } return createMockResponse({ message: '偏好设置已更新（演示模式）' }) }, diff --git a/frontend/src/views/admin/BillingPlansManagement.vue b/frontend/src/views/admin/BillingPlansManagement.vue index 54524ceea..e8e1b8117 100644 --- a/frontend/src/views/admin/BillingPlansManagement.vue +++ b/frontend/src/views/admin/BillingPlansManagement.vue @@ -746,15 +746,9 @@ 每日额度套餐本身已按类型互斥；填写后还可与其他权益类型的同名组整包互斥。

-
-
- -

- 额度不足时继续使用钱包余额 -

-
- -
+

+ 套餐额度不足时是否使用钱包余额，由用户在个人设置中选择，默认关闭。 +

@@ -1272,7 +1266,6 @@ interface PlanFormState { daily_quota_usd: number reset_timezone: string carry_over: boolean - allow_wallet_overage: boolean daily_quota_replacement_group: string membership_group_enabled: boolean grant_user_groups: string[] @@ -1526,11 +1519,7 @@ const dailyQuotaSummaryText = computed(() => : '每天独立 USD 消费用量，默认不结转' ) -const dailyQuotaDetailText = computed(() => - form.allow_wallet_overage - ? '每日额度不足时会继续使用钱包余额，适合希望用户不中断请求的套餐。' - : '每日额度不足时不再继续扣钱包，适合严格封顶的月卡或体验卡。' -) +const dailyQuotaDetailText = '优先使用每日套餐额度；额度不足时是否使用钱包余额，由用户在个人设置中选择，默认关闭。' const membershipSummaryText = computed(() => planMode.value === 'mixed' @@ -1568,7 +1557,6 @@ function buildDefaultForm(): PlanFormState { daily_quota_usd: 50, reset_timezone: 'Asia/Shanghai', carry_over: false, - allow_wallet_overage: false, daily_quota_replacement_group: '', membership_group_enabled: false, grant_user_groups: [], @@ -1732,7 +1720,6 @@ function formFromPlan(plan: BillingPlan): PlanFormState { next.daily_quota_usd = Number(quota.daily_quota_usd || next.daily_quota_usd) next.reset_timezone = quota.reset_timezone || 'Asia/Shanghai' next.carry_over = Boolean(quota.carry_over) - next.allow_wallet_overage = Boolean(quota.allow_wallet_overage) next.daily_quota_replacement_group = quota.replacement_group || '' } else if (entitlement.type === 'membership_group') { const membership = entitlement as MembershipGroupEntitlement @@ -1813,7 +1800,6 @@ function buildEntitlements(): BillingEntitlement[] { daily_quota_usd: Number(form.daily_quota_usd), reset_timezone: form.reset_timezone.trim() || 'Asia/Shanghai', carry_over: false, - allow_wallet_overage: Boolean(form.allow_wallet_overage), } attachReplacementGroup(entitlement, form.daily_quota_replacement_group) entitlements.push(entitlement) diff --git a/frontend/src/views/user/Settings.vue b/frontend/src/views/user/Settings.vue index bc5c4591b..41b00e199 100644 --- a/frontend/src/views/user/Settings.vue +++ b/frontend/src/views/user/Settings.vue @@ -503,6 +503,35 @@
+
+

+ 额度使用 +

+
+
+ +

+ 默认关闭，套餐额度不足时不会自动扣除钱包余额。开启后优先使用套餐额度，额度耗尽、不可用或处于恢复期时使用钱包余额。 +

+
+ +
+
+

通知设置 @@ -744,6 +773,9 @@ const featureSettingsForm = ref({ const savingProfile = ref(false) const savingFeatureSettings = ref(false) +const allowWalletOverage = ref(false) +const billingPreferencesLoaded = ref(false) +const savingWalletOverage = ref(false) const changingPassword = ref(false) const sessionsLoading = ref(false) const sessionActionLoading = ref(null) @@ -1001,6 +1033,8 @@ async function handleUnbind(providerType: string) { async function loadPreferences() { try { const prefs = await meApi.getPreferences() + allowWalletOverage.value = prefs.allow_wallet_overage === true + billingPreferencesLoaded.value = true // 主题以本地 localStorage 为准（useDarkMode 在应用启动时已初始化） // 这样可以避免刷新页面时主题被服务端旧值覆盖 @@ -1169,6 +1203,23 @@ async function handleRevokeOtherSessions() { } } +async function updateWalletOverage(enabled: boolean) { + if (!billingPreferencesLoaded.value || savingWalletOverage.value) return + const previous = allowWalletOverage.value + allowWalletOverage.value = enabled + savingWalletOverage.value = true + try { + await meApi.updatePreferences({ allow_wallet_overage: enabled }) + success('设置已保存') + } catch (error) { + allowWalletOverage.value = previous + log.error('更新额度使用设置失败:', error) + showError('保存设置失败') + } finally { + savingWalletOverage.value = false + } +} + async function updatePreferences() { try { await meApi.updatePreferences({ diff --git a/frontend/src/views/user/WalletCenter.vue b/frontend/src/views/user/WalletCenter.vue index 4014b6e05..f7465827d 100644 --- a/frontend/src/views/user/WalletCenter.vue +++ b/frontend/src/views/user/WalletCenter.vue @@ -47,7 +47,7 @@ 已用 {{ formatCurrency(dailyQuotaUsed) }} / 每日 {{ formatCurrency(dailyQuotaTotal) }}

- {{ dailyQuota?.allow_wallet_overage ? '套餐不足时继续扣钱包余额' : '套餐额度不足时会拒绝请求' }} + {{ dailyQuota?.allow_wallet_overage ? '已开启套餐不足时使用钱包余额' : '套餐额度不足时会拒绝请求，可在个人设置中开启使用余额' }}
{ expect(meApiMock.updatePreferences).not.toHaveBeenCalled() }) }) + +describe('Settings wallet fallback preferences', () => { + function walletSwitch(root: HTMLElement): HTMLButtonElement { + const control = root.querySelector('#allow-wallet-overage') + if (!control) throw new Error('The wallet fallback switch was not rendered') + return control + } + + it('defaults to off for existing users and waits for preferences before allowing changes', async () => { + const preferences = deferred>() + meApiMock.getPreferences.mockReturnValueOnce(preferences.promise) + const root = mountSettings() + const control = walletSwitch(root) + + expect(control.getAttribute('aria-checked')).toBe('false') + expect(control.disabled).toBe(true) + control.click() + expect(meApiMock.updatePreferences).not.toHaveBeenCalled() + + preferences.resolve(serverPreferences()) + await flushPromises() + expect(control.disabled).toBe(false) + expect(control.getAttribute('aria-checked')).toBe('false') + }) + + it('loads the saved choice and persists both disabling and enabling without changing other preferences', async () => { + meApiMock.getPreferences.mockResolvedValue({ ...serverPreferences(), allow_wallet_overage: true }) + const root = mountSettings() + await flushPromises() + const control = walletSwitch(root) + + expect(control.getAttribute('aria-checked')).toBe('true') + control.click() + await flushPromises() + expect(meApiMock.updatePreferences).toHaveBeenLastCalledWith({ allow_wallet_overage: false }) + expect(control.getAttribute('aria-checked')).toBe('false') + + control.click() + await flushPromises() + expect(meApiMock.updatePreferences).toHaveBeenLastCalledWith({ allow_wallet_overage: true }) + expect(control.getAttribute('aria-checked')).toBe('true') + + chooseEnglish(root) + await flushPromises() + expect(meApiMock.updatePreferences.mock.lastCall?.[0]).not.toHaveProperty('allow_wallet_overage') + }) + + it('prevents duplicate saves and restores the saved choice when saving fails', async () => { + let reject!: (reason: Error) => void + meApiMock.updatePreferences.mockReturnValueOnce(new Promise((_, fail) => { reject = fail })) + const root = mountSettings() + await flushPromises() + const control = walletSwitch(root) + + control.click() + await nextTick() + expect(control.disabled).toBe(true) + expect(control.getAttribute('aria-checked')).toBe('true') + control.click() + expect(meApiMock.updatePreferences).toHaveBeenCalledTimes(1) + + reject(new Error('Save failed')) + await flushPromises() + expect(control.disabled).toBe(false) + expect(control.getAttribute('aria-checked')).toBe('false') + expect(toastMock.error).toHaveBeenCalledWith('保存设置失败') + expect(toastMock.success).not.toHaveBeenCalled() + }) +})