mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-08 02:17:46 +08:00
fix(gateway): add provider key concurrency and cache affinity modes
This commit is contained in:
@@ -205,8 +205,8 @@ async fn available_balance_capacity_usd(
|
||||
.as_ref()
|
||||
.is_some_and(|wallet| wallet.limit_mode.eq_ignore_ascii_case("unlimited"));
|
||||
Ok(match quota.as_ref() {
|
||||
Some(quota) if !quota.allow_wallet_overage => Some(quota.remaining_usd.max(0.0)),
|
||||
Some(_) if wallet_is_unlimited => None,
|
||||
Some(quota) if !quota.allow_wallet_overage => Some(quota.remaining_usd.max(0.0)),
|
||||
Some(quota) => Some(quota.remaining_usd.max(0.0) + wallet_available_usd.unwrap_or(0.0)),
|
||||
None if wallet_is_unlimited => None,
|
||||
None => wallet_available_usd,
|
||||
@@ -832,9 +832,10 @@ mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::{
|
||||
execution_plan_balance_capacity_rejection, execution_plan_cost_upper_bound_cache_key,
|
||||
max_output_tokens_from_request, openai_request_input_is_self_contained,
|
||||
output_choice_count_upper_bound, request_model_local_rejection, GatewayLocalAuthRejection,
|
||||
available_balance_capacity_usd, execution_plan_balance_capacity_rejection,
|
||||
execution_plan_cost_upper_bound_cache_key, max_output_tokens_from_request,
|
||||
openai_request_input_is_self_contained, output_choice_count_upper_bound,
|
||||
request_model_local_rejection, GatewayLocalAuthRejection,
|
||||
};
|
||||
use crate::control::{GatewayControlAuthContext, GatewayControlDecision};
|
||||
use crate::data::GatewayDataState;
|
||||
@@ -939,6 +940,14 @@ mod tests {
|
||||
fn state_with_quota_and_wallet(
|
||||
quota: UserDailyQuotaAvailabilityRecord,
|
||||
context: StoredBillingModelContext,
|
||||
) -> AppState {
|
||||
state_with_quota_context_and_wallet(quota, context, sample_wallet("user-1", 30.0))
|
||||
}
|
||||
|
||||
fn state_with_quota_context_and_wallet(
|
||||
quota: UserDailyQuotaAvailabilityRecord,
|
||||
context: StoredBillingModelContext,
|
||||
wallet: StoredWalletSnapshot,
|
||||
) -> AppState {
|
||||
let candidate_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
@@ -952,7 +961,7 @@ mod tests {
|
||||
AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(data)
|
||||
.with_auth_wallets_for_tests(vec![sample_wallet("user-1", 30.0)])
|
||||
.with_auth_wallets_for_tests(vec![wallet])
|
||||
}
|
||||
|
||||
fn state_with_model_mapping() -> AppState {
|
||||
@@ -1350,6 +1359,26 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unlimited_wallet_capacity_ignores_exhausted_non_overage_quota() {
|
||||
let context = billing_context_with_pricing(None, None, None, None);
|
||||
let mut wallet = sample_wallet("user-1", 0.0);
|
||||
wallet.limit_mode = "unlimited".to_string();
|
||||
let state =
|
||||
state_with_quota_context_and_wallet(quota_availability(0.0, false), context, wallet);
|
||||
let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]);
|
||||
let auth_context = decision
|
||||
.auth_context
|
||||
.as_ref()
|
||||
.expect("decision should include auth context");
|
||||
|
||||
let capacity = available_balance_capacity_usd(&state, auth_context)
|
||||
.await
|
||||
.expect("capacity should resolve");
|
||||
|
||||
assert_eq!(capacity, None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn positive_balance_does_not_allow_historical_invalid_processing_pricing() {
|
||||
let context = billing_context_with_pricing(
|
||||
|
||||
Reference in New Issue
Block a user