i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-10 11:19:50 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: allow wallet fallback after plan quota exhaustion

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/control/auth/gate.rs
+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(
apps/aether-gateway/src/data/tests.rs
+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
apps/aether-gateway/src/executor/candidate_loop.rs
+49 -13
View File
@@ -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
apps/aether-gateway/src/handlers/proxy/websocket/responses/turn.rs
+18 -2
View File
@@ -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 {
apps/aether-gateway/src/handlers/public/support/user_me_preferences.rs
+12
View File
@@ -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(
apps/aether-gateway/src/plan_usage_policy.rs
+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(),
apps/aether-gateway/src/state/runtime/billing/admin.rs
+10 -2
View File
@@ -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(
apps/aether-gateway/src/state/runtime/user_preferences.rs
+72 -3
View File
@@ -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());
}
}
}
apps/aether-gateway/src/state/types.rs
+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,
}
}
}
apps/aether-gateway/src/tests/frontdoor/public_support.rs
+39
View File
@@ -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();
apps/aether-gateway/src/usage/mod.rs
+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,
apps/aether-gateway/src/usage/reporting/context.rs
+132
View File
@@ -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());
}
}
apps/aether-gateway/src/usage/reporting/mod.rs
+1
View File
@@ -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};
apps/aether-gateway/src/wallet_runtime/access.rs
+65 -1
View File
@@ -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)