feat: allow wallet fallback after plan quota exhaustion

This commit is contained in:
elky
2026-10-09 21:21:53 +08:00
parent bad13237d2
commit 63520c8c27
39 changed files with 1474 additions and 190 deletions
+71 -7
View File
@@ -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<Option<f64>, 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(&quota_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(
+1
View File
@@ -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
@@ -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
@@ -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 {
@@ -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(
+549 -28
View File
@@ -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<str>,
policy: Arc<EffectivePlanUsagePolicy>,
admission_event_id: Arc<str>,
wallet_fallback: Arc<AtomicBool>,
}
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<Self> {
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<f64>,
) -> Result<bool, GatewayError> {
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<PlanUsageCostReservationOutcome, GatewayError> {
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<Option<AdmissionPermit>, PlanUsageAdmissionError> {
) -> Result<(Option<AdmissionPermit>, 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::<Result<Vec<_>, _>>()?;
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,
&quota_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,
&quota_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(),
@@ -642,10 +642,18 @@ impl AppState {
&self,
user_id: &str,
) -> Result<Option<UserDailyQuotaAvailabilityRecord>, 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(
@@ -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());
}
}
}
+5
View File
@@ -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<crate::data::state::StoredUserPreferenceRecord> 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<GatewayUserPreferenceView> 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,
}
}
}
@@ -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();
+1
View File
@@ -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,
@@ -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<Value>,
wallet_fallback: bool,
) -> Result<Option<Value>, 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());
}
}
@@ -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};
@@ -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)
@@ -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;
@@ -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<f64>>(
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<bool>,
now: chrono::DateTime<chrono::Utc>,
) -> Result<Vec<DailyQuotaGrant>, 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<bool> {
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<PaymentGatewayConfigRecord, DataLayerError> {
@@ -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<bool>,
now: chrono::DateTime<chrono::Utc>,
) -> Result<Vec<DailyQuotaGrant>, 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<bool> {
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<f64>>(
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::<Option<serde_json::Value>, _>("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;
@@ -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();
@@ -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()?,
})
}
@@ -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<String, Value>) -> 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 [
@@ -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);
}
}
@@ -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
);
@@ -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
);
@@ -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"
@@ -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,
]
);
}
@@ -116,6 +116,7 @@ WHERE version=20260919000000;
20261001000000,
20261004000000,
20261007000000,
20261009000000,
]
);
assert_eq!(
@@ -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;
}
@@ -92,14 +92,12 @@ fn billing_plan_from_input(
fn daily_quota_availability_from_entitlements(
entitlements: impl IntoIterator<Item = UserPlanEntitlementRecord>,
billing_plans: &BTreeMap<String, BillingPlanRecord>,
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<bool> {
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::<Vec<_>>();
let billing_plans = self
.billing_plans_by_id
.read()
.expect("billing repository lock");
Ok(Some(daily_quota_availability_from_entitlements(
entitlements,
&billing_plans,
now,
)))
}
@@ -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
@@ -103,6 +103,7 @@ pub(crate) fn retain_first_byte_request_metadata(value: Option<Value>) -> Option
| "client_requested_stream"
| "upstream_is_stream"
| "api_key_is_standalone"
| "plan_wallet_fallback"
| "plan_usage_reservation_token"
| "request_path"
| "request_query_string"
@@ -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
+6
View File
@@ -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(),
+1
View File
@@ -21,6 +21,7 @@ export interface UserPreferences {
theme?: 'light' | 'dark' | 'auto'
language?: string
notifications_enabled?: boolean
allow_wallet_overage?: boolean
[key: string]: unknown // 允许扩展其他偏好设置
}
+1
View File
@@ -42,6 +42,7 @@ export interface UserPreferences {
theme: string
language: string
timezone?: string
allow_wallet_overage?: boolean
notifications?: {
email?: boolean
usage_alerts?: boolean
+13 -2
View File
@@ -1609,11 +1609,22 @@ const mockHandlers: Record<string, (config: AxiosRequestConfig) => Promise<Axios
'GET /api/users/me/preferences': async () => {
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: '偏好设置已更新(演示模式)' })
},
@@ -746,15 +746,9 @@
每日额度套餐本身已按类型互斥;填写后还可与其他权益类型的同名组整包互斥。
</p>
</div>
<div class="flex items-center justify-between rounded-xl border border-border/60 bg-card/50 p-3">
<div>
<Label>允许超额扣钱包</Label>
<p class="mt-1 text-xs text-muted-foreground">
额度不足时继续使用钱包余额
</p>
</div>
<Switch v-model="form.allow_wallet_overage" />
</div>
<p class="rounded-xl border border-border/60 bg-card/50 p-3 text-xs leading-5 text-muted-foreground">
套餐额度不足时是否使用钱包余额,由用户在个人设置中选择,默认关闭。
</p>
<div class="flex items-center justify-between rounded-xl border border-border/60 bg-card/50 p-3 opacity-70">
<div>
<Label>额度结转</Label>
@@ -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)
+51
View File
@@ -503,6 +503,35 @@
</div>
</div>
<div class="space-y-3">
<h4 class="font-medium text-foreground">
额度使用
</h4>
<div class="flex items-center justify-between gap-4 py-2">
<div class="flex-1">
<Label
for="allow-wallet-overage"
class="text-sm font-medium cursor-pointer"
>
套餐额度不可用时使用余额
</Label>
<p
id="allow-wallet-overage-description"
class="text-xs text-muted-foreground mt-1"
>
默认关闭,套餐额度不足时不会自动扣除钱包余额。开启后优先使用套餐额度,额度耗尽、不可用或处于恢复期时使用钱包余额。
</p>
</div>
<Switch
id="allow-wallet-overage"
:model-value="allowWalletOverage"
:disabled="!billingPreferencesLoaded || savingWalletOverage"
aria-describedby="allow-wallet-overage-description"
@update:model-value="updateWalletOverage"
/>
</div>
</div>
<div class="space-y-3">
<h4 class="font-medium text-foreground">
通知设置
@@ -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<string | null>(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({
+1 -1
View File
@@ -47,7 +47,7 @@
已用 {{ formatCurrency(dailyQuotaUsed) }} / 每日 {{ formatCurrency(dailyQuotaTotal) }}
</div>
<div class="text-xs text-muted-foreground">
{{ dailyQuota?.allow_wallet_overage ? '套餐不足时继续扣钱包余额' : '套餐额度不足时会拒绝请求' }}
{{ dailyQuota?.allow_wallet_overage ? '已开启套餐不足时使用钱包余额' : '套餐额度不足时会拒绝请求,可在个人设置中开启使用余额' }}
</div>
</div>
<div
@@ -227,3 +227,72 @@ describe('Settings language preferences', () => {
expect(meApiMock.updatePreferences).not.toHaveBeenCalled()
})
})
describe('Settings wallet fallback preferences', () => {
function walletSwitch(root: HTMLElement): HTMLButtonElement {
const control = root.querySelector<HTMLButtonElement>('#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<ReturnType<typeof serverPreferences>>()
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<void>((_, 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()
})
})