mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
Merge remote-tracking branch 'origin/pr-484'
This commit is contained in:
@@ -157,7 +157,9 @@ pub(crate) fn resolve_local_decision_execution_runtime_auth_context(
|
|||||||
decision: &GatewayControlDecision,
|
decision: &GatewayControlDecision,
|
||||||
) -> Option<ExecutionRuntimeAuthContext> {
|
) -> Option<ExecutionRuntimeAuthContext> {
|
||||||
resolve_decision_execution_runtime_auth_context(decision).filter(|auth_context| {
|
resolve_decision_execution_runtime_auth_context(decision).filter(|auth_context| {
|
||||||
!auth_context.user_id.trim().is_empty() && !auth_context.api_key_id.trim().is_empty()
|
auth_context.access_allowed
|
||||||
|
&& !auth_context.user_id.trim().is_empty()
|
||||||
|
&& !auth_context.api_key_id.trim().is_empty()
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -106,7 +106,7 @@ async fn balance_capacity_rejection(
|
|||||||
requested_model: Option<&str>,
|
requested_model: Option<&str>,
|
||||||
body: &Bytes,
|
body: &Bytes,
|
||||||
) -> Result<Option<GatewayLocalAuthRejection>, GatewayError> {
|
) -> Result<Option<GatewayLocalAuthRejection>, GatewayError> {
|
||||||
if auth_context.api_key_is_standalone || auth_context.admin_bypass_limits {
|
if auth_context.api_key_is_standalone {
|
||||||
return Ok(None);
|
return Ok(None);
|
||||||
}
|
}
|
||||||
if auth_context.local_rejection.is_some() {
|
if auth_context.local_rejection.is_some() {
|
||||||
@@ -816,6 +816,43 @@ mod tests {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn admin_bypass_limits_does_not_skip_exhausted_daily_quota_capacity() {
|
||||||
|
let context = billing_context_with_pricing(
|
||||||
|
Some(json!({
|
||||||
|
"tiers": [{
|
||||||
|
"up_to": null,
|
||||||
|
"input_price_per_1m": 1.0,
|
||||||
|
"output_price_per_1m": 2.0
|
||||||
|
}]
|
||||||
|
})),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
);
|
||||||
|
let state = state_with_quota_and_wallet(quota_availability(0.0, false), context);
|
||||||
|
let mut decision = decision_with_allowed_models(vec!["gpt-5".to_string()]);
|
||||||
|
if let Some(auth_context) = decision.auth_context.as_mut() {
|
||||||
|
auth_context.admin_bypass_limits = true;
|
||||||
|
}
|
||||||
|
let uri: Uri = "/v1/chat/completions".parse().expect("uri should parse");
|
||||||
|
let body = Bytes::from_static(
|
||||||
|
br#"{"model":"gpt-5","messages":[{"role":"user","content":"hi"}],"stream":true}"#,
|
||||||
|
);
|
||||||
|
|
||||||
|
let rejection =
|
||||||
|
request_model_local_rejection(&state, Some(&decision), &uri, &json_headers(), &body)
|
||||||
|
.await
|
||||||
|
.expect("quota rejection should resolve");
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
rejection,
|
||||||
|
Some(GatewayLocalAuthRejection::BalanceDenied {
|
||||||
|
remaining: Some(0.0),
|
||||||
|
})
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn positive_balance_still_denies_known_cost_above_available_capacity() {
|
async fn positive_balance_still_denies_known_cost_above_available_capacity() {
|
||||||
let context = billing_context_with_pricing(
|
let context = billing_context_with_pricing(
|
||||||
|
|||||||
@@ -9,7 +9,8 @@ pub(crate) use gate::{
|
|||||||
request_model_local_rejection, should_buffer_request_for_local_auth,
|
request_model_local_rejection, should_buffer_request_for_local_auth,
|
||||||
trusted_auth_local_rejection, GatewayLocalAuthRejection,
|
trusted_auth_local_rejection, GatewayLocalAuthRejection,
|
||||||
};
|
};
|
||||||
pub(super) use resolution::{resolve_control_decision_auth, ControlDecisionAuthResolution};
|
|
||||||
pub(crate) use resolution::{
|
pub(crate) use resolution::{
|
||||||
resolve_execution_runtime_auth_context, GatewayAdminPrincipalContext, GatewayControlAuthContext,
|
refresh_execution_runtime_auth_context, resolve_execution_runtime_auth_context,
|
||||||
|
GatewayAdminPrincipalContext, GatewayControlAuthContext,
|
||||||
};
|
};
|
||||||
|
pub(super) use resolution::{resolve_control_decision_auth, ControlDecisionAuthResolution};
|
||||||
|
|||||||
@@ -433,7 +433,14 @@ pub(crate) async fn resolve_execution_runtime_auth_context(
|
|||||||
let _ = trace_id;
|
let _ = trace_id;
|
||||||
|
|
||||||
if let Some(auth_context) = decision.auth_context.clone() {
|
if let Some(auth_context) = decision.auth_context.clone() {
|
||||||
return Ok(Some(auth_context));
|
return Ok(Some(
|
||||||
|
refresh_execution_runtime_auth_context(
|
||||||
|
state,
|
||||||
|
auth_context,
|
||||||
|
decision.auth_endpoint_signature.as_deref(),
|
||||||
|
)
|
||||||
|
.await?,
|
||||||
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
let Some(auth_endpoint_signature) = decision.auth_endpoint_signature.as_deref() else {
|
let Some(auth_endpoint_signature) = decision.auth_endpoint_signature.as_deref() else {
|
||||||
@@ -445,7 +452,14 @@ pub(crate) async fn resolve_execution_runtime_auth_context(
|
|||||||
};
|
};
|
||||||
|
|
||||||
if let Some(auth_context) = get_cached_auth_context(state, &cache_key) {
|
if let Some(auth_context) = get_cached_auth_context(state, &cache_key) {
|
||||||
return Ok(Some(auth_context));
|
let refreshed = refresh_execution_runtime_auth_context(
|
||||||
|
state,
|
||||||
|
auth_context,
|
||||||
|
Some(auth_endpoint_signature),
|
||||||
|
)
|
||||||
|
.await?;
|
||||||
|
put_cached_auth_context(state, cache_key, refreshed.clone());
|
||||||
|
return Ok(Some(refreshed));
|
||||||
}
|
}
|
||||||
|
|
||||||
if let Some(auth_context) =
|
if let Some(auth_context) =
|
||||||
@@ -461,6 +475,56 @@ pub(crate) async fn resolve_execution_runtime_auth_context(
|
|||||||
Ok(None)
|
Ok(None)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn refresh_execution_runtime_auth_context(
|
||||||
|
state: &AppState,
|
||||||
|
auth_context: GatewayControlAuthContext,
|
||||||
|
auth_endpoint_signature: Option<&str>,
|
||||||
|
) -> Result<GatewayControlAuthContext, GatewayError> {
|
||||||
|
if auth_context.local_rejection.is_some() || !auth_context.access_allowed {
|
||||||
|
return Ok(auth_context);
|
||||||
|
}
|
||||||
|
let Some(auth_endpoint_signature) = auth_endpoint_signature
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
else {
|
||||||
|
return Ok(auth_context);
|
||||||
|
};
|
||||||
|
if !state.has_auth_api_key_reader()
|
||||||
|
|| auth_context.user_id.trim().is_empty()
|
||||||
|
|| auth_context.api_key_id.trim().is_empty()
|
||||||
|
{
|
||||||
|
return Ok(auth_context);
|
||||||
|
}
|
||||||
|
|
||||||
|
let snapshot = state
|
||||||
|
.data
|
||||||
|
.read_auth_api_key_snapshot(
|
||||||
|
&auth_context.user_id,
|
||||||
|
&auth_context.api_key_id,
|
||||||
|
current_unix_secs(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||||
|
let Some(snapshot) = snapshot else {
|
||||||
|
let mut denied = auth_context;
|
||||||
|
denied.access_allowed = false;
|
||||||
|
denied.local_rejection = Some(GatewayLocalAuthRejection::InvalidApiKey);
|
||||||
|
denied.balance_remaining = None;
|
||||||
|
return Ok(denied);
|
||||||
|
};
|
||||||
|
|
||||||
|
let wallet_access = resolve_wallet_auth_gate(state, &snapshot).await?;
|
||||||
|
Ok(build_data_backed_auth_context(
|
||||||
|
state,
|
||||||
|
snapshot,
|
||||||
|
auth_endpoint_signature,
|
||||||
|
Some(true),
|
||||||
|
auth_context.balance_remaining,
|
||||||
|
wallet_access,
|
||||||
|
)
|
||||||
|
.await)
|
||||||
|
}
|
||||||
|
|
||||||
fn put_cached_auth_context(
|
fn put_cached_auth_context(
|
||||||
state: &AppState,
|
state: &AppState,
|
||||||
cache_key: String,
|
cache_key: String,
|
||||||
@@ -609,7 +673,7 @@ async fn build_data_backed_auth_context(
|
|||||||
.api_key_expires_at_unix_secs
|
.api_key_expires_at_unix_secs
|
||||||
.is_some_and(|expires_at| expires_at < current_unix_secs());
|
.is_some_and(|expires_at| expires_at < current_unix_secs());
|
||||||
let locked_api_key = snapshot.api_key_is_locked && !snapshot.api_key_is_standalone;
|
let locked_api_key = snapshot.api_key_is_locked && !snapshot.api_key_is_standalone;
|
||||||
let access_allowed = header_access_allowed
|
let key_access_allowed = header_access_allowed
|
||||||
.map(|value| value && snapshot.currently_usable)
|
.map(|value| value && snapshot.currently_usable)
|
||||||
.unwrap_or(snapshot.currently_usable);
|
.unwrap_or(snapshot.currently_usable);
|
||||||
let wallet_remaining = wallet_access
|
let wallet_remaining = wallet_access
|
||||||
@@ -656,7 +720,7 @@ async fn build_data_backed_auth_context(
|
|||||||
user_id: snapshot.user_id,
|
user_id: snapshot.user_id,
|
||||||
api_key_id: snapshot.api_key_id,
|
api_key_id: snapshot.api_key_id,
|
||||||
balance_remaining: wallet_remaining.or(balance_remaining),
|
balance_remaining: wallet_remaining.or(balance_remaining),
|
||||||
access_allowed,
|
access_allowed: key_access_allowed && local_rejection.is_none(),
|
||||||
user_rate_limit: snapshot.user_rate_limit,
|
user_rate_limit: snapshot.user_rate_limit,
|
||||||
api_key_rate_limit: snapshot.api_key_rate_limit,
|
api_key_rate_limit: snapshot.api_key_rate_limit,
|
||||||
api_key_is_standalone: snapshot.api_key_is_standalone,
|
api_key_is_standalone: snapshot.api_key_is_standalone,
|
||||||
@@ -835,13 +899,20 @@ mod tests {
|
|||||||
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
|
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
|
||||||
};
|
};
|
||||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||||
|
use aether_data::repository::wallet::{
|
||||||
|
InMemoryWalletRepository, StoredWalletSnapshot, WalletReadRepository,
|
||||||
|
};
|
||||||
use aether_data_contracts::repository::provider_catalog::{
|
use aether_data_contracts::repository::provider_catalog::{
|
||||||
StoredProviderCatalogEndpoint, StoredProviderCatalogProvider,
|
StoredProviderCatalogEndpoint, StoredProviderCatalogProvider,
|
||||||
};
|
};
|
||||||
use axum::http::{HeaderMap, Uri};
|
use axum::http::{HeaderMap, Uri};
|
||||||
|
|
||||||
use super::{resolve_data_backed_auth_context, GatewayLocalAuthRejection};
|
use super::{
|
||||||
|
resolve_data_backed_auth_context, resolve_execution_runtime_auth_context,
|
||||||
|
GatewayLocalAuthRejection,
|
||||||
|
};
|
||||||
use crate::control::auth::credentials::hash_api_key;
|
use crate::control::auth::credentials::hash_api_key;
|
||||||
|
use crate::control::GatewayControlDecision;
|
||||||
use crate::data::GatewayDataState;
|
use crate::data::GatewayDataState;
|
||||||
use crate::AppState;
|
use crate::AppState;
|
||||||
|
|
||||||
@@ -946,6 +1017,154 @@ mod tests {
|
|||||||
assert_eq!(repository.touch_count("key-1"), 1);
|
assert_eq!(repository.touch_count("key-1"), 1);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn data_backed_auth_context_marks_wallet_denial_as_not_allowed() {
|
||||||
|
let api_key = "sk-test-empty-wallet";
|
||||||
|
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||||
|
Some(hash_api_key(api_key)),
|
||||||
|
sample_snapshot("key-empty-wallet", "user-empty-wallet"),
|
||||||
|
)]));
|
||||||
|
let wallet_repository = Arc::new(InMemoryWalletRepository::seed(vec![
|
||||||
|
StoredWalletSnapshot::new(
|
||||||
|
"wallet-empty".to_string(),
|
||||||
|
Some("user-empty-wallet".to_string()),
|
||||||
|
None,
|
||||||
|
0.0,
|
||||||
|
0.0,
|
||||||
|
"finite".to_string(),
|
||||||
|
"USD".to_string(),
|
||||||
|
"active".to_string(),
|
||||||
|
0.0,
|
||||||
|
0.0,
|
||||||
|
0.0,
|
||||||
|
0.0,
|
||||||
|
100,
|
||||||
|
)
|
||||||
|
.expect("wallet should build"),
|
||||||
|
]));
|
||||||
|
let data =
|
||||||
|
GatewayDataState::with_auth_and_wallet_for_tests(auth_repository, wallet_repository);
|
||||||
|
let state = AppState::new()
|
||||||
|
.expect("state should build")
|
||||||
|
.with_data_state_for_tests(data);
|
||||||
|
|
||||||
|
let mut headers = HeaderMap::new();
|
||||||
|
headers.insert(
|
||||||
|
http::header::AUTHORIZATION,
|
||||||
|
format!("Bearer {api_key}").parse().unwrap(),
|
||||||
|
);
|
||||||
|
|
||||||
|
let auth_context = resolve_data_backed_auth_context(
|
||||||
|
&state,
|
||||||
|
&headers,
|
||||||
|
&uri("/v1/chat/completions"),
|
||||||
|
Some("openai:chat"),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.expect("resolution should succeed")
|
||||||
|
.expect("auth context should exist");
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
auth_context.local_rejection,
|
||||||
|
Some(GatewayLocalAuthRejection::BalanceDenied {
|
||||||
|
remaining: Some(0.0),
|
||||||
|
})
|
||||||
|
);
|
||||||
|
assert!(!auth_context.access_allowed);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn execution_runtime_auth_context_revalidates_cached_wallet_state() {
|
||||||
|
let api_key = "sk-test-runtime-wallet-cache";
|
||||||
|
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||||
|
Some(hash_api_key(api_key)),
|
||||||
|
sample_snapshot("key-runtime-wallet-cache", "user-runtime-wallet-cache"),
|
||||||
|
)]));
|
||||||
|
let wallet_repository = Arc::new(InMemoryWalletRepository::seed(vec![
|
||||||
|
StoredWalletSnapshot::new(
|
||||||
|
"wallet-runtime-cache".to_string(),
|
||||||
|
Some("user-runtime-wallet-cache".to_string()),
|
||||||
|
None,
|
||||||
|
10.0,
|
||||||
|
0.0,
|
||||||
|
"finite".to_string(),
|
||||||
|
"USD".to_string(),
|
||||||
|
"active".to_string(),
|
||||||
|
10.0,
|
||||||
|
0.0,
|
||||||
|
0.0,
|
||||||
|
0.0,
|
||||||
|
100,
|
||||||
|
)
|
||||||
|
.expect("wallet should build"),
|
||||||
|
]));
|
||||||
|
let data = GatewayDataState::with_auth_and_wallet_for_tests(
|
||||||
|
auth_repository,
|
||||||
|
Arc::clone(&wallet_repository),
|
||||||
|
);
|
||||||
|
let state = AppState::new()
|
||||||
|
.expect("state should build")
|
||||||
|
.with_data_state_for_tests(data);
|
||||||
|
let decision = GatewayControlDecision::synthetic(
|
||||||
|
"/v1/chat/completions",
|
||||||
|
Some("ai_public".to_string()),
|
||||||
|
Some("openai".to_string()),
|
||||||
|
Some("chat".to_string()),
|
||||||
|
Some("openai:chat".to_string()),
|
||||||
|
);
|
||||||
|
let mut headers = HeaderMap::new();
|
||||||
|
headers.insert("x-api-key", api_key.parse().unwrap());
|
||||||
|
|
||||||
|
let first = resolve_execution_runtime_auth_context(
|
||||||
|
&state,
|
||||||
|
&decision,
|
||||||
|
&headers,
|
||||||
|
&uri("/v1/chat/completions"),
|
||||||
|
"trace-runtime-wallet-cache",
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.expect("resolution should succeed")
|
||||||
|
.expect("auth context should exist");
|
||||||
|
assert!(first.access_allowed);
|
||||||
|
|
||||||
|
wallet_repository
|
||||||
|
.update_auth_user_wallet_snapshot(
|
||||||
|
"user-runtime-wallet-cache",
|
||||||
|
0.0,
|
||||||
|
0.0,
|
||||||
|
"finite",
|
||||||
|
"USD",
|
||||||
|
"active",
|
||||||
|
10.0,
|
||||||
|
10.0,
|
||||||
|
0.0,
|
||||||
|
0.0,
|
||||||
|
Some(101),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.expect("wallet update should succeed")
|
||||||
|
.expect("wallet should exist");
|
||||||
|
|
||||||
|
let second = resolve_execution_runtime_auth_context(
|
||||||
|
&state,
|
||||||
|
&decision,
|
||||||
|
&headers,
|
||||||
|
&uri("/v1/chat/completions"),
|
||||||
|
"trace-runtime-wallet-cache",
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.expect("resolution should succeed")
|
||||||
|
.expect("auth context should exist");
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
second.local_rejection,
|
||||||
|
Some(GatewayLocalAuthRejection::BalanceDenied {
|
||||||
|
remaining: Some(0.0),
|
||||||
|
})
|
||||||
|
);
|
||||||
|
assert!(!second.access_allowed);
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn data_backed_auth_context_allows_provider_id_for_matching_provider_type() {
|
async fn data_backed_auth_context_allows_provider_id_for_matching_provider_type() {
|
||||||
let api_key = "sk-test-provider-id";
|
let api_key = "sk-test-provider-id";
|
||||||
|
|||||||
@@ -8,9 +8,10 @@ mod public;
|
|||||||
mod route;
|
mod route;
|
||||||
|
|
||||||
pub(crate) use auth::{
|
pub(crate) use auth::{
|
||||||
extract_requested_model, request_model_local_rejection, resolve_execution_runtime_auth_context,
|
extract_requested_model, refresh_execution_runtime_auth_context, request_model_local_rejection,
|
||||||
should_buffer_request_for_local_auth, trusted_auth_local_rejection,
|
resolve_execution_runtime_auth_context, should_buffer_request_for_local_auth,
|
||||||
GatewayAdminPrincipalContext, GatewayControlAuthContext, GatewayLocalAuthRejection,
|
trusted_auth_local_rejection, GatewayAdminPrincipalContext, GatewayControlAuthContext,
|
||||||
|
GatewayLocalAuthRejection,
|
||||||
};
|
};
|
||||||
pub(crate) use execute::{allows_control_execute_emergency, maybe_execute_via_control};
|
pub(crate) use execute::{allows_control_execute_emergency, maybe_execute_via_control};
|
||||||
pub(crate) use management_token_permissions::{
|
pub(crate) use management_token_permissions::{
|
||||||
|
|||||||
@@ -30,6 +30,25 @@ use axum::response::IntoResponse;
|
|||||||
use axum::Json;
|
use axum::Json;
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
|
|
||||||
|
async fn apply_supplied_auth_context(
|
||||||
|
state: &AppState,
|
||||||
|
decision: &mut GatewayControlDecision,
|
||||||
|
auth_context: Option<crate::control::GatewayControlAuthContext>,
|
||||||
|
) -> Result<bool, GatewayError> {
|
||||||
|
let Some(auth_context) = auth_context else {
|
||||||
|
return Ok(false);
|
||||||
|
};
|
||||||
|
let refreshed = crate::control::refresh_execution_runtime_auth_context(
|
||||||
|
state,
|
||||||
|
auth_context,
|
||||||
|
decision.auth_endpoint_signature.as_deref(),
|
||||||
|
)
|
||||||
|
.await?;
|
||||||
|
decision.local_auth_rejection = refreshed.local_rejection.clone();
|
||||||
|
decision.auth_context = Some(refreshed);
|
||||||
|
Ok(true)
|
||||||
|
}
|
||||||
|
|
||||||
pub(crate) async fn maybe_build_local_internal_proxy_response_impl(
|
pub(crate) async fn maybe_build_local_internal_proxy_response_impl(
|
||||||
state: &AppState,
|
state: &AppState,
|
||||||
request_context: &GatewayPublicRequestContext,
|
request_context: &GatewayPublicRequestContext,
|
||||||
@@ -206,11 +225,8 @@ pub(crate) async fn maybe_build_local_internal_proxy_response_impl(
|
|||||||
Json(build_internal_gateway_fallback_plan_payload(None)).into_response(),
|
Json(build_internal_gateway_fallback_plan_payload(None)).into_response(),
|
||||||
));
|
));
|
||||||
};
|
};
|
||||||
let provided_auth_context = payload.auth_context.is_some();
|
let provided_auth_context =
|
||||||
if let Some(auth_context) = payload.auth_context {
|
apply_supplied_auth_context(state, &mut resolved, payload.auth_context).await?;
|
||||||
resolved.auth_context = Some(auth_context);
|
|
||||||
resolved.local_auth_rejection = None;
|
|
||||||
}
|
|
||||||
let auth_context = resolved.auth_context.as_ref();
|
let auth_context = resolved.auth_context.as_ref();
|
||||||
if auth_context
|
if auth_context
|
||||||
.map(|value| !value.access_allowed)
|
.map(|value| !value.access_allowed)
|
||||||
@@ -308,11 +324,8 @@ pub(crate) async fn maybe_build_local_internal_proxy_response_impl(
|
|||||||
Json(build_internal_gateway_fallback_plan_payload(None)).into_response(),
|
Json(build_internal_gateway_fallback_plan_payload(None)).into_response(),
|
||||||
));
|
));
|
||||||
};
|
};
|
||||||
let provided_auth_context = payload.auth_context.is_some();
|
let provided_auth_context =
|
||||||
if let Some(auth_context) = payload.auth_context {
|
apply_supplied_auth_context(state, &mut resolved, payload.auth_context).await?;
|
||||||
resolved.auth_context = Some(auth_context);
|
|
||||||
resolved.local_auth_rejection = None;
|
|
||||||
}
|
|
||||||
let auth_context = resolved.auth_context.as_ref();
|
let auth_context = resolved.auth_context.as_ref();
|
||||||
if auth_context
|
if auth_context
|
||||||
.map(|value| !value.access_allowed)
|
.map(|value| !value.access_allowed)
|
||||||
@@ -407,11 +420,8 @@ pub(crate) async fn maybe_build_local_internal_proxy_response_impl(
|
|||||||
else {
|
else {
|
||||||
return Ok(Some(build_internal_gateway_proxy_public_response()));
|
return Ok(Some(build_internal_gateway_proxy_public_response()));
|
||||||
};
|
};
|
||||||
let provided_auth_context = payload.auth_context.is_some();
|
let provided_auth_context =
|
||||||
if let Some(auth_context) = payload.auth_context {
|
apply_supplied_auth_context(state, &mut resolved, payload.auth_context).await?;
|
||||||
resolved.auth_context = Some(auth_context);
|
|
||||||
resolved.local_auth_rejection = None;
|
|
||||||
}
|
|
||||||
if let Some(mut planned) = api::maybe_build_sync_plan_payload(
|
if let Some(mut planned) = api::maybe_build_sync_plan_payload(
|
||||||
state,
|
state,
|
||||||
&parts,
|
&parts,
|
||||||
@@ -474,11 +484,8 @@ pub(crate) async fn maybe_build_local_internal_proxy_response_impl(
|
|||||||
else {
|
else {
|
||||||
return Ok(Some(build_internal_gateway_proxy_public_response()));
|
return Ok(Some(build_internal_gateway_proxy_public_response()));
|
||||||
};
|
};
|
||||||
let provided_auth_context = payload.auth_context.is_some();
|
let provided_auth_context =
|
||||||
if let Some(auth_context) = payload.auth_context {
|
apply_supplied_auth_context(state, &mut resolved, payload.auth_context).await?;
|
||||||
resolved.auth_context = Some(auth_context);
|
|
||||||
resolved.local_auth_rejection = None;
|
|
||||||
}
|
|
||||||
if let Some(mut planned) = api::maybe_build_stream_plan_payload(
|
if let Some(mut planned) = api::maybe_build_stream_plan_payload(
|
||||||
state,
|
state,
|
||||||
&parts,
|
&parts,
|
||||||
@@ -546,10 +553,7 @@ pub(crate) async fn maybe_build_local_internal_proxy_response_impl(
|
|||||||
else {
|
else {
|
||||||
return Ok(None);
|
return Ok(None);
|
||||||
};
|
};
|
||||||
if let Some(auth_context) = payload.auth_context {
|
apply_supplied_auth_context(state, &mut resolved, payload.auth_context).await?;
|
||||||
resolved.auth_context = Some(auth_context);
|
|
||||||
resolved.local_auth_rejection = None;
|
|
||||||
}
|
|
||||||
if let Some(plan_payload) = api::maybe_build_sync_plan_payload(
|
if let Some(plan_payload) = api::maybe_build_sync_plan_payload(
|
||||||
state,
|
state,
|
||||||
&parts,
|
&parts,
|
||||||
@@ -630,10 +634,7 @@ pub(crate) async fn maybe_build_local_internal_proxy_response_impl(
|
|||||||
else {
|
else {
|
||||||
return Ok(None);
|
return Ok(None);
|
||||||
};
|
};
|
||||||
if let Some(auth_context) = payload.auth_context {
|
apply_supplied_auth_context(state, &mut resolved, payload.auth_context).await?;
|
||||||
resolved.auth_context = Some(auth_context);
|
|
||||||
resolved.local_auth_rejection = None;
|
|
||||||
}
|
|
||||||
if let Some(plan_payload) = api::maybe_build_stream_plan_payload(
|
if let Some(plan_payload) = api::maybe_build_stream_plan_payload(
|
||||||
state,
|
state,
|
||||||
&parts,
|
&parts,
|
||||||
|
|||||||
@@ -30,10 +30,10 @@ mod refunds;
|
|||||||
use self::flow::handle_wallet_flow;
|
use self::flow::handle_wallet_flow;
|
||||||
pub(in crate::handlers::public::support) use self::reads::build_wallet_balance_payload_for_user;
|
pub(in crate::handlers::public::support) use self::reads::build_wallet_balance_payload_for_user;
|
||||||
use self::reads::{
|
use self::reads::{
|
||||||
build_wallet_daily_usage_payload, build_wallet_payload, build_wallet_zero_today_entry,
|
build_wallet_daily_usage_payload, build_wallet_live_today_usage_payload_for_user,
|
||||||
handle_wallet_balance, handle_wallet_today_cost, handle_wallet_transactions,
|
build_wallet_payload, build_wallet_zero_today_entry, handle_wallet_balance,
|
||||||
parse_wallet_limit, parse_wallet_offset, wallet_fixed_offset, wallet_today_billing_date_string,
|
handle_wallet_today_cost, handle_wallet_transactions, parse_wallet_limit, parse_wallet_offset,
|
||||||
wallet_transaction_payload_from_record,
|
wallet_fixed_offset, wallet_transaction_payload_from_record,
|
||||||
};
|
};
|
||||||
pub(crate) use self::recharge::sanitize_wallet_gateway_response;
|
pub(crate) use self::recharge::sanitize_wallet_gateway_response;
|
||||||
use self::recharge::{
|
use self::recharge::{
|
||||||
|
|||||||
@@ -1,9 +1,10 @@
|
|||||||
use super::{
|
use super::{
|
||||||
build_auth_error_response, build_auth_json_response, build_wallet_daily_usage_payload,
|
build_auth_error_response, build_auth_json_response, build_wallet_daily_usage_payload,
|
||||||
build_wallet_payload, build_wallet_zero_today_entry, http, parse_wallet_limit,
|
build_wallet_live_today_usage_payload_for_user, build_wallet_payload,
|
||||||
parse_wallet_offset, resolve_authenticated_local_user, unix_secs_to_rfc3339,
|
build_wallet_zero_today_entry, http, parse_wallet_limit, parse_wallet_offset,
|
||||||
wallet_fixed_offset, wallet_today_billing_date_string, wallet_transaction_payload_from_record,
|
resolve_authenticated_local_user, unix_secs_to_rfc3339, wallet_fixed_offset,
|
||||||
AppState, Body, GatewayPublicRequestContext, Response, WALLET_LEGACY_TIMEZONE,
|
wallet_transaction_payload_from_record, AppState, Body, GatewayPublicRequestContext, Response,
|
||||||
|
WALLET_LEGACY_TIMEZONE,
|
||||||
};
|
};
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
|
|
||||||
@@ -98,7 +99,17 @@ pub(super) async fn handle_wallet_flow(
|
|||||||
return build_auth_json_response(http::StatusCode::OK, payload, None);
|
return build_auth_json_response(http::StatusCode::OK, payload, None);
|
||||||
};
|
};
|
||||||
|
|
||||||
let mut today_entry = build_wallet_zero_today_entry();
|
let mut today_entry =
|
||||||
|
match build_wallet_live_today_usage_payload_for_user(state, &auth.user.id).await {
|
||||||
|
Ok(Some(today_usage)) => today_usage,
|
||||||
|
_ => build_wallet_zero_today_entry(),
|
||||||
|
};
|
||||||
|
if today_entry
|
||||||
|
.get("total_requests")
|
||||||
|
.and_then(serde_json::Value::as_u64)
|
||||||
|
.unwrap_or_default()
|
||||||
|
== 0
|
||||||
|
{
|
||||||
if let Ok(Some(today_usage)) = state
|
if let Ok(Some(today_usage)) = state
|
||||||
.find_wallet_today_usage(&wallet.id, WALLET_LEGACY_TIMEZONE)
|
.find_wallet_today_usage(&wallet.id, WALLET_LEGACY_TIMEZONE)
|
||||||
.await
|
.await
|
||||||
@@ -125,6 +136,7 @@ pub(super) async fn handle_wallet_flow(
|
|||||||
true,
|
true,
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
|
||||||
let fetch_size = offset.saturating_add(limit).min(5200);
|
let fetch_size = offset.saturating_add(limit).min(5200);
|
||||||
let (transactions, tx_total) = state
|
let (transactions, tx_total) = state
|
||||||
|
|||||||
@@ -1,11 +1,11 @@
|
|||||||
use super::{
|
use super::{
|
||||||
build_auth_error_response, build_auth_json_response, build_auth_wallet_summary_payload, http,
|
build_auth_error_response, build_auth_json_response, build_auth_wallet_summary_payload, http,
|
||||||
query_param_value, resolve_authenticated_local_user, unix_secs_to_rfc3339, AppState, Body,
|
query_param_value, resolve_authenticated_local_user, unix_secs_to_rfc3339, AppState, Body,
|
||||||
GatewayError, GatewayPublicRequestContext, Response, WALLET_LEGACY_TIMEZONE,
|
GatewayPublicRequestContext, Response, WALLET_LEGACY_TIMEZONE,
|
||||||
};
|
};
|
||||||
use crate::handlers::shared::round_to;
|
use crate::handlers::shared::round_to;
|
||||||
use aether_data_contracts::repository::usage::UsageSettledCostSummaryQuery;
|
use aether_data_contracts::repository::usage::UsageSettledCostSummaryQuery;
|
||||||
use chrono::Utc;
|
use chrono::{TimeZone, Utc};
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
|
|
||||||
const WALLET_TODAY_COST_UNAVAILABLE_DETAIL: &str = "钱包今日费用数据暂不可用";
|
const WALLET_TODAY_COST_UNAVAILABLE_DETAIL: &str = "钱包今日费用数据暂不可用";
|
||||||
@@ -143,6 +143,26 @@ pub(super) fn wallet_today_billing_date_string() -> String {
|
|||||||
.to_string()
|
.to_string()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn wallet_today_usage_window() -> Result<(String, String, u64, u64), String> {
|
||||||
|
let offset = wallet_fixed_offset();
|
||||||
|
let today = Utc::now().with_timezone(&offset).date_naive();
|
||||||
|
let Some(local_start_naive) = today.and_hms_opt(0, 0, 0) else {
|
||||||
|
return Err("wallet today start is invalid".to_string());
|
||||||
|
};
|
||||||
|
let Some(local_start) = offset.from_local_datetime(&local_start_naive).single() else {
|
||||||
|
return Err("wallet today local start is ambiguous".to_string());
|
||||||
|
};
|
||||||
|
let local_end = local_start + chrono::Duration::days(1);
|
||||||
|
let start_unix_secs = local_start.timestamp().max(0) as u64;
|
||||||
|
let end_unix_secs = local_end.timestamp().max(0) as u64;
|
||||||
|
Ok((
|
||||||
|
today.to_string(),
|
||||||
|
WALLET_LEGACY_TIMEZONE.to_string(),
|
||||||
|
start_unix_secs,
|
||||||
|
end_unix_secs,
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
pub(super) fn build_wallet_daily_usage_payload(
|
pub(super) fn build_wallet_daily_usage_payload(
|
||||||
id: Option<String>,
|
id: Option<String>,
|
||||||
date: String,
|
date: String,
|
||||||
@@ -193,6 +213,43 @@ pub(super) fn build_wallet_zero_today_entry() -> serde_json::Value {
|
|||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(super) async fn build_wallet_live_today_usage_payload_for_user(
|
||||||
|
state: &AppState,
|
||||||
|
user_id: &str,
|
||||||
|
) -> Result<Option<serde_json::Value>, String> {
|
||||||
|
if !state.has_usage_data_reader() {
|
||||||
|
return Ok(None);
|
||||||
|
}
|
||||||
|
let (date, timezone, start_unix_secs, end_unix_secs) = wallet_today_usage_window()?;
|
||||||
|
let summary = state
|
||||||
|
.summarize_usage_settled_cost(&UsageSettledCostSummaryQuery {
|
||||||
|
created_from_unix_secs: start_unix_secs,
|
||||||
|
created_until_unix_secs: end_unix_secs,
|
||||||
|
user_id: Some(user_id.to_string()),
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.map_err(|err| format!("wallet today cost lookup failed: {err:?}"))?;
|
||||||
|
Ok(Some(build_wallet_daily_usage_payload(
|
||||||
|
None,
|
||||||
|
date,
|
||||||
|
timezone,
|
||||||
|
summary.total_cost_usd,
|
||||||
|
summary.total_requests,
|
||||||
|
summary.input_tokens,
|
||||||
|
summary.output_tokens,
|
||||||
|
summary.cache_creation_tokens,
|
||||||
|
summary.cache_read_tokens,
|
||||||
|
summary
|
||||||
|
.first_finalized_at_unix_secs
|
||||||
|
.and_then(unix_secs_to_rfc3339),
|
||||||
|
summary
|
||||||
|
.last_finalized_at_unix_secs
|
||||||
|
.and_then(unix_secs_to_rfc3339),
|
||||||
|
Some(Utc::now().to_rfc3339()),
|
||||||
|
true,
|
||||||
|
)))
|
||||||
|
}
|
||||||
|
|
||||||
pub(super) fn wallet_transaction_payload_from_record(
|
pub(super) fn wallet_transaction_payload_from_record(
|
||||||
record: &aether_data::repository::wallet::StoredAdminWalletTransaction,
|
record: &aether_data::repository::wallet::StoredAdminWalletTransaction,
|
||||||
) -> serde_json::Value {
|
) -> serde_json::Value {
|
||||||
@@ -253,65 +310,17 @@ pub(super) async fn handle_wallet_today_cost(
|
|||||||
Ok(value) => value,
|
Ok(value) => value,
|
||||||
Err(response) => return response,
|
Err(response) => return response,
|
||||||
};
|
};
|
||||||
let today = Utc::now().date_naive();
|
match build_wallet_live_today_usage_payload_for_user(state, &auth.user.id).await {
|
||||||
let Some(start_of_day) = today.and_hms_opt(0, 0, 0) else {
|
Ok(Some(payload)) => build_auth_json_response(http::StatusCode::OK, payload, None),
|
||||||
return build_auth_error_response(
|
Ok(None) => build_auth_error_response(
|
||||||
http::StatusCode::INTERNAL_SERVER_ERROR,
|
http::StatusCode::SERVICE_UNAVAILABLE,
|
||||||
"wallet today start is invalid",
|
WALLET_TODAY_COST_UNAVAILABLE_DETAIL,
|
||||||
false,
|
false,
|
||||||
);
|
),
|
||||||
};
|
Err(detail) => {
|
||||||
let start_unix_secs = u64::try_from(
|
build_auth_error_response(http::StatusCode::INTERNAL_SERVER_ERROR, detail, false)
|
||||||
chrono::DateTime::<chrono::Utc>::from_naive_utc_and_offset(start_of_day, chrono::Utc)
|
}
|
||||||
.timestamp(),
|
|
||||||
)
|
|
||||||
.unwrap_or_default();
|
|
||||||
let end_unix_secs = start_unix_secs.saturating_add(24 * 3600);
|
|
||||||
|
|
||||||
let summary = match state
|
|
||||||
.summarize_usage_settled_cost(&UsageSettledCostSummaryQuery {
|
|
||||||
created_from_unix_secs: start_unix_secs,
|
|
||||||
created_until_unix_secs: end_unix_secs,
|
|
||||||
user_id: Some(auth.user.id.clone()),
|
|
||||||
})
|
|
||||||
.await
|
|
||||||
{
|
|
||||||
Ok(value) => value,
|
|
||||||
Err(err) => {
|
|
||||||
return build_auth_error_response(
|
|
||||||
http::StatusCode::INTERNAL_SERVER_ERROR,
|
|
||||||
format!("wallet today cost lookup failed: {err:?}"),
|
|
||||||
false,
|
|
||||||
)
|
|
||||||
}
|
}
|
||||||
};
|
|
||||||
|
|
||||||
let first_finalized_at = summary
|
|
||||||
.first_finalized_at_unix_secs
|
|
||||||
.and_then(unix_secs_to_rfc3339);
|
|
||||||
let last_finalized_at = summary
|
|
||||||
.last_finalized_at_unix_secs
|
|
||||||
.and_then(unix_secs_to_rfc3339);
|
|
||||||
|
|
||||||
build_auth_json_response(
|
|
||||||
http::StatusCode::OK,
|
|
||||||
json!({
|
|
||||||
"id": serde_json::Value::Null,
|
|
||||||
"date": today.to_string(),
|
|
||||||
"timezone": "UTC",
|
|
||||||
"total_cost": round_to(summary.total_cost_usd, 6),
|
|
||||||
"total_requests": summary.total_requests,
|
|
||||||
"input_tokens": summary.input_tokens,
|
|
||||||
"output_tokens": summary.output_tokens,
|
|
||||||
"cache_creation_tokens": summary.cache_creation_tokens,
|
|
||||||
"cache_read_tokens": summary.cache_read_tokens,
|
|
||||||
"first_finalized_at": first_finalized_at,
|
|
||||||
"last_finalized_at": last_finalized_at,
|
|
||||||
"aggregated_at": Utc::now().to_rfc3339(),
|
|
||||||
"is_today": true,
|
|
||||||
}),
|
|
||||||
None,
|
|
||||||
)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(super) async fn handle_wallet_transactions(
|
pub(super) async fn handle_wallet_transactions(
|
||||||
|
|||||||
@@ -11,6 +11,9 @@ use crate::tests::{
|
|||||||
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
|
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
|
||||||
EXECUTION_PATH_HEADER,
|
EXECUTION_PATH_HEADER,
|
||||||
};
|
};
|
||||||
|
use aether_data::repository::billing::InMemoryBillingReadRepository;
|
||||||
|
use aether_data::repository::usage::InMemoryUsageReadRepository;
|
||||||
|
use aether_data::repository::wallet::{InMemoryWalletRepository, StoredWalletSnapshot};
|
||||||
use base64::Engine as _;
|
use base64::Engine as _;
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
@@ -1154,6 +1157,118 @@ async fn gateway_handles_internal_gateway_decision_sync_locally_with_supplied_au
|
|||||||
upstream_handle.abort();
|
upstream_handle.abort();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn gateway_internal_decision_sync_revalidates_supplied_auth_context_wallet() {
|
||||||
|
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||||
|
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||||
|
let upstream = Router::new().route(
|
||||||
|
"/{*path}",
|
||||||
|
any(move |_request: Request| {
|
||||||
|
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||||
|
async move {
|
||||||
|
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||||
|
(StatusCode::OK, Body::from("proxied"))
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
|
||||||
|
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||||
|
None,
|
||||||
|
unrestricted_models_snapshot("api-key-empty-wallet", "user-empty-wallet"),
|
||||||
|
)]));
|
||||||
|
let candidate_repository =
|
||||||
|
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||||
|
sample_models_candidate_row("provider-1", "openai", "openai:chat", "gpt-5", 10),
|
||||||
|
]));
|
||||||
|
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||||
|
vec![sample_provider("provider-1", "openai", 10)],
|
||||||
|
vec![sample_endpoint(
|
||||||
|
"endpoint-provider-1",
|
||||||
|
"provider-1",
|
||||||
|
"openai:chat",
|
||||||
|
"https://api.openai.example",
|
||||||
|
)],
|
||||||
|
vec![sample_key(
|
||||||
|
"key-provider-1",
|
||||||
|
"provider-1",
|
||||||
|
"openai:chat",
|
||||||
|
"sk-upstream-openai",
|
||||||
|
)],
|
||||||
|
));
|
||||||
|
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||||
|
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
||||||
|
let billing_repository = Arc::new(InMemoryBillingReadRepository::seed(Vec::new()));
|
||||||
|
let wallet_repository = Arc::new(InMemoryWalletRepository::seed(vec![
|
||||||
|
StoredWalletSnapshot::new(
|
||||||
|
"wallet-empty".to_string(),
|
||||||
|
Some("user-empty-wallet".to_string()),
|
||||||
|
None,
|
||||||
|
0.0,
|
||||||
|
0.0,
|
||||||
|
"finite".to_string(),
|
||||||
|
"USD".to_string(),
|
||||||
|
"active".to_string(),
|
||||||
|
0.0,
|
||||||
|
0.0,
|
||||||
|
0.0,
|
||||||
|
0.0,
|
||||||
|
100,
|
||||||
|
)
|
||||||
|
.expect("wallet should build"),
|
||||||
|
]));
|
||||||
|
|
||||||
|
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||||
|
let gateway = build_router_with_state(
|
||||||
|
AppState::new()
|
||||||
|
.expect("gateway should build")
|
||||||
|
.with_data_state_for_tests(
|
||||||
|
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_usage_billing_and_wallet_for_tests(
|
||||||
|
auth_repository,
|
||||||
|
candidate_repository,
|
||||||
|
provider_catalog_repository,
|
||||||
|
request_candidate_repository,
|
||||||
|
usage_repository,
|
||||||
|
billing_repository,
|
||||||
|
wallet_repository,
|
||||||
|
DEVELOPMENT_ENCRYPTION_KEY,
|
||||||
|
),
|
||||||
|
),
|
||||||
|
);
|
||||||
|
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||||
|
|
||||||
|
let response = reqwest::Client::new()
|
||||||
|
.post(format!("{gateway_url}/api/internal/gateway/decision-sync"))
|
||||||
|
.json(&json!({
|
||||||
|
"trace_id": "trace-internal-decision-sync-empty-wallet",
|
||||||
|
"method": "POST",
|
||||||
|
"path": "/v1/chat/completions",
|
||||||
|
"headers": {
|
||||||
|
"content-type": "application/json",
|
||||||
|
},
|
||||||
|
"body_json": {
|
||||||
|
"model": "gpt-5",
|
||||||
|
"messages": [],
|
||||||
|
},
|
||||||
|
"auth_context": {
|
||||||
|
"user_id": "user-empty-wallet",
|
||||||
|
"api_key_id": "api-key-empty-wallet",
|
||||||
|
"access_allowed": true,
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.expect("request should succeed");
|
||||||
|
|
||||||
|
assert_eq!(response.status(), StatusCode::OK);
|
||||||
|
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||||
|
assert_eq!(payload["action"], "fallback_plan");
|
||||||
|
assert_eq!(payload["auth_context"], serde_json::Value::Null);
|
||||||
|
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||||
|
|
||||||
|
gateway_handle.abort();
|
||||||
|
upstream_handle.abort();
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn gateway_returns_internal_gateway_decision_sync_fallback_with_resolved_auth_context() {
|
async fn gateway_returns_internal_gateway_decision_sync_fallback_with_resolved_auth_context() {
|
||||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||||
|
|||||||
@@ -4133,13 +4133,7 @@ async fn gateway_handles_wallet_balance_locally_without_proxying_upstream() {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn gateway_handles_wallet_today_cost_locally_without_proxying_upstream() {
|
async fn gateway_handles_wallet_today_cost_locally_without_proxying_upstream() {
|
||||||
let auth_now = Utc::now();
|
let auth_now = Utc::now();
|
||||||
let usage_now = chrono::DateTime::<chrono::Utc>::from_naive_utc_and_offset(
|
let usage_now = auth_now - chrono::Duration::minutes(30);
|
||||||
auth_now
|
|
||||||
.date_naive()
|
|
||||||
.and_hms_opt(12, 0, 0)
|
|
||||||
.expect("midday should be valid"),
|
|
||||||
chrono::Utc,
|
|
||||||
);
|
|
||||||
let user = sample_auth_user(auth_now);
|
let user = sample_auth_user(auth_now);
|
||||||
let access_token = build_test_auth_token(
|
let access_token = build_test_auth_token(
|
||||||
"access",
|
"access",
|
||||||
@@ -4211,6 +4205,72 @@ async fn gateway_handles_wallet_today_cost_locally_without_proxying_upstream() {
|
|||||||
upstream_handle.abort();
|
upstream_handle.abort();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn gateway_wallet_flow_today_entry_uses_live_settled_usage() {
|
||||||
|
let auth_now = Utc::now();
|
||||||
|
let user = sample_auth_user(auth_now);
|
||||||
|
let access_token = build_test_auth_token(
|
||||||
|
"access",
|
||||||
|
serde_json::Map::from_iter([
|
||||||
|
("user_id".to_string(), json!(user.id)),
|
||||||
|
("role".to_string(), json!(user.role)),
|
||||||
|
(
|
||||||
|
"created_at".to_string(),
|
||||||
|
json!(user.created_at.map(|value| value.to_rfc3339())),
|
||||||
|
),
|
||||||
|
("session_id".to_string(), json!("session-wallet-flow-today")),
|
||||||
|
]),
|
||||||
|
auth_now + chrono::Duration::hours(1),
|
||||||
|
);
|
||||||
|
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
|
||||||
|
sample_user_usage_audit(
|
||||||
|
"usage-wallet-flow-today",
|
||||||
|
"req-wallet-flow-today",
|
||||||
|
"user-auth-1",
|
||||||
|
"gpt-4.1",
|
||||||
|
"OpenAI",
|
||||||
|
"completed",
|
||||||
|
auth_now - chrono::Duration::minutes(5),
|
||||||
|
),
|
||||||
|
]));
|
||||||
|
let (gateway_url, upstream_hits, gateway_handle, upstream_handle) =
|
||||||
|
start_auth_gateway_with_usage_state(
|
||||||
|
user,
|
||||||
|
sample_auth_wallet("user-auth-1", auth_now),
|
||||||
|
[sample_auth_session(
|
||||||
|
"user-auth-1",
|
||||||
|
"session-wallet-flow-today",
|
||||||
|
"device-wallet-flow-today",
|
||||||
|
"refresh-token-placeholder",
|
||||||
|
auth_now,
|
||||||
|
)],
|
||||||
|
usage_repository,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
let response = reqwest::Client::new()
|
||||||
|
.get(format!("{gateway_url}/api/wallet/flow?limit=20&offset=0"))
|
||||||
|
.header("authorization", format!("Bearer {access_token}"))
|
||||||
|
.header("x-client-device-id", "device-wallet-flow-today")
|
||||||
|
.header("user-agent", "AetherTest/1.0")
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.expect("request should succeed");
|
||||||
|
|
||||||
|
assert_eq!(response.status(), StatusCode::OK);
|
||||||
|
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||||
|
assert_eq!(payload["today_entry"]["is_today"], true);
|
||||||
|
assert_eq!(payload["today_entry"]["timezone"], "Asia/Shanghai");
|
||||||
|
assert_eq!(payload["today_entry"]["total_requests"], 1);
|
||||||
|
assert_eq!(payload["today_entry"]["input_tokens"], 120);
|
||||||
|
assert_eq!(payload["today_entry"]["cache_read_tokens"], 15);
|
||||||
|
assert_eq!(payload["today_entry"]["total_cost"], 1.25);
|
||||||
|
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||||
|
|
||||||
|
gateway_handle.abort();
|
||||||
|
upstream_handle.abort();
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn gateway_returns_service_unavailable_for_wallet_today_cost_without_usage_reader() {
|
async fn gateway_returns_service_unavailable_for_wallet_today_cost_without_usage_reader() {
|
||||||
let now = Utc::now();
|
let now = Utc::now();
|
||||||
|
|||||||
@@ -24,14 +24,9 @@ pub(crate) async fn resolve_wallet_auth_gate(
|
|||||||
auth_snapshot.api_key_is_standalone,
|
auth_snapshot.api_key_is_standalone,
|
||||||
)
|
)
|
||||||
.await?;
|
.await?;
|
||||||
let is_admin = wallet_auth_allows_admin_bypass(
|
|
||||||
&auth_snapshot.user_role,
|
|
||||||
auth_snapshot.api_key_is_standalone,
|
|
||||||
);
|
|
||||||
|
|
||||||
let decision = match wallet.as_ref() {
|
let decision = match wallet.as_ref() {
|
||||||
Some(wallet) => map_wallet_snapshot(wallet).access_decision(is_admin),
|
Some(wallet) => map_wallet_snapshot(wallet).access_decision(false),
|
||||||
None if is_admin => WalletAccessDecision::allowed(None),
|
|
||||||
None => WalletAccessDecision::wallet_unavailable(None),
|
None => WalletAccessDecision::wallet_unavailable(None),
|
||||||
};
|
};
|
||||||
if !auth_snapshot.api_key_is_standalone {
|
if !auth_snapshot.api_key_is_standalone {
|
||||||
@@ -83,19 +78,50 @@ fn map_wallet_snapshot(snapshot: &StoredWalletSnapshot) -> WalletSnapshot {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn wallet_auth_allows_admin_bypass(user_role: &str, api_key_is_standalone: bool) -> bool {
|
|
||||||
user_role.eq_ignore_ascii_case("admin") && !api_key_is_standalone
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use aether_data::repository::wallet::StoredWalletSnapshot;
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use aether_data::repository::usage::InMemoryUsageReadRepository;
|
||||||
|
use aether_data::repository::wallet::{InMemoryWalletRepository, StoredWalletSnapshot};
|
||||||
|
use aether_data_contracts::repository::billing::{
|
||||||
|
BillingReadRepository, StoredBillingModelContext, UserDailyQuotaAvailabilityRecord,
|
||||||
|
};
|
||||||
|
use aether_data_contracts::DataLayerError;
|
||||||
use aether_wallet::{WalletAccessFailure, WalletLimitMode, WalletSnapshot, WalletStatus};
|
use aether_wallet::{WalletAccessFailure, WalletLimitMode, WalletSnapshot, WalletStatus};
|
||||||
|
use async_trait::async_trait;
|
||||||
|
|
||||||
use super::{
|
use super::{
|
||||||
local_rejection_from_wallet_access, map_wallet_snapshot, wallet_auth_allows_admin_bypass,
|
local_rejection_from_wallet_access, map_wallet_snapshot, resolve_wallet_auth_gate,
|
||||||
};
|
};
|
||||||
use crate::control::GatewayLocalAuthRejection;
|
use crate::control::GatewayLocalAuthRejection;
|
||||||
|
use crate::data::auth::GatewayAuthApiKeySnapshot;
|
||||||
|
use crate::data::GatewayDataState;
|
||||||
|
use crate::AppState;
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
struct FixedQuotaBillingReadRepository {
|
||||||
|
quota: Option<UserDailyQuotaAvailabilityRecord>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl BillingReadRepository for FixedQuotaBillingReadRepository {
|
||||||
|
async fn find_model_context(
|
||||||
|
&self,
|
||||||
|
_provider_id: &str,
|
||||||
|
_provider_api_key_id: Option<&str>,
|
||||||
|
_global_model_name: &str,
|
||||||
|
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
|
||||||
|
Ok(None)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn find_user_daily_quota_availability(
|
||||||
|
&self,
|
||||||
|
_user_id: &str,
|
||||||
|
) -> Result<Option<UserDailyQuotaAvailabilityRecord>, DataLayerError> {
|
||||||
|
Ok(self.quota.clone())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn maps_wallet_snapshot_and_derives_balance_denied() {
|
fn maps_wallet_snapshot_and_derives_balance_denied() {
|
||||||
@@ -145,8 +171,143 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn standalone_key_never_uses_admin_wallet_bypass() {
|
fn admin_user_with_empty_finite_wallet_is_balance_denied() {
|
||||||
assert!(wallet_auth_allows_admin_bypass("admin", false));
|
let stored = StoredWalletSnapshot::new(
|
||||||
assert!(!wallet_auth_allows_admin_bypass("admin", true));
|
"wallet-1".to_string(),
|
||||||
|
Some("admin-1".to_string()),
|
||||||
|
None,
|
||||||
|
0.0,
|
||||||
|
0.0,
|
||||||
|
"finite".to_string(),
|
||||||
|
"USD".to_string(),
|
||||||
|
"active".to_string(),
|
||||||
|
0.0,
|
||||||
|
0.0,
|
||||||
|
0.0,
|
||||||
|
0.0,
|
||||||
|
100,
|
||||||
|
)
|
||||||
|
.expect("wallet should build");
|
||||||
|
|
||||||
|
let decision = map_wallet_snapshot(&stored).access_decision(true);
|
||||||
|
|
||||||
|
assert_eq!(decision.failure, Some(WalletAccessFailure::BalanceDenied));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn ordinary_user_key_without_quota_denies_empty_wallet() {
|
||||||
|
let state = state_with_wallet_and_quota(empty_user_wallet(), None);
|
||||||
|
let auth_snapshot = ordinary_user_api_key_snapshot();
|
||||||
|
|
||||||
|
let decision = resolve_wallet_auth_gate(&state, &auth_snapshot)
|
||||||
|
.await
|
||||||
|
.expect("wallet gate should resolve")
|
||||||
|
.expect("wallet gate should return a decision");
|
||||||
|
|
||||||
|
assert!(!decision.allowed);
|
||||||
|
assert_eq!(decision.failure, Some(WalletAccessFailure::BalanceDenied));
|
||||||
|
assert_eq!(
|
||||||
|
local_rejection_from_wallet_access(&decision),
|
||||||
|
Some(GatewayLocalAuthRejection::BalanceDenied {
|
||||||
|
remaining: Some(0.0),
|
||||||
|
})
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn ordinary_user_key_with_remaining_quota_allows_empty_wallet() {
|
||||||
|
let state = state_with_wallet_and_quota(
|
||||||
|
empty_user_wallet(),
|
||||||
|
Some(quota_availability(10.0, 4.0, false)),
|
||||||
|
);
|
||||||
|
let auth_snapshot = ordinary_user_api_key_snapshot();
|
||||||
|
|
||||||
|
let decision = resolve_wallet_auth_gate(&state, &auth_snapshot)
|
||||||
|
.await
|
||||||
|
.expect("wallet gate should resolve")
|
||||||
|
.expect("wallet gate should return a decision");
|
||||||
|
|
||||||
|
assert!(decision.allowed);
|
||||||
|
assert_eq!(decision.failure, None);
|
||||||
|
assert_eq!(decision.remaining, Some(4.0));
|
||||||
|
}
|
||||||
|
|
||||||
|
fn state_with_wallet_and_quota(
|
||||||
|
wallet: StoredWalletSnapshot,
|
||||||
|
quota: Option<UserDailyQuotaAvailabilityRecord>,
|
||||||
|
) -> AppState {
|
||||||
|
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
||||||
|
let billing_repository: Arc<dyn BillingReadRepository> =
|
||||||
|
Arc::new(FixedQuotaBillingReadRepository { quota });
|
||||||
|
let wallet_repository = Arc::new(InMemoryWalletRepository::seed(vec![wallet]));
|
||||||
|
let data = GatewayDataState::with_usage_billing_and_wallet_for_tests(
|
||||||
|
usage_repository,
|
||||||
|
billing_repository,
|
||||||
|
wallet_repository,
|
||||||
|
);
|
||||||
|
AppState::new()
|
||||||
|
.expect("state should build")
|
||||||
|
.with_data_state_for_tests(data)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn empty_user_wallet() -> StoredWalletSnapshot {
|
||||||
|
StoredWalletSnapshot::new(
|
||||||
|
"wallet-user-1".to_string(),
|
||||||
|
Some("user-1".to_string()),
|
||||||
|
None,
|
||||||
|
0.0,
|
||||||
|
0.0,
|
||||||
|
"finite".to_string(),
|
||||||
|
"USD".to_string(),
|
||||||
|
"active".to_string(),
|
||||||
|
0.0,
|
||||||
|
0.0,
|
||||||
|
0.0,
|
||||||
|
0.0,
|
||||||
|
100,
|
||||||
|
)
|
||||||
|
.expect("wallet should build")
|
||||||
|
}
|
||||||
|
|
||||||
|
fn quota_availability(
|
||||||
|
total_quota_usd: f64,
|
||||||
|
remaining_usd: f64,
|
||||||
|
allow_wallet_overage: bool,
|
||||||
|
) -> UserDailyQuotaAvailabilityRecord {
|
||||||
|
UserDailyQuotaAvailabilityRecord {
|
||||||
|
has_active_daily_quota: true,
|
||||||
|
total_quota_usd,
|
||||||
|
used_usd: total_quota_usd - remaining_usd,
|
||||||
|
remaining_usd,
|
||||||
|
allow_wallet_overage,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn ordinary_user_api_key_snapshot() -> GatewayAuthApiKeySnapshot {
|
||||||
|
GatewayAuthApiKeySnapshot {
|
||||||
|
user_id: "user-1".to_string(),
|
||||||
|
username: "ordinary-user".to_string(),
|
||||||
|
email: Some("ordinary@example.com".to_string()),
|
||||||
|
user_role: "user".to_string(),
|
||||||
|
user_auth_source: "local".to_string(),
|
||||||
|
user_is_active: true,
|
||||||
|
user_is_deleted: false,
|
||||||
|
user_rate_limit: None,
|
||||||
|
user_allowed_providers: None,
|
||||||
|
user_allowed_api_formats: None,
|
||||||
|
user_allowed_models: None,
|
||||||
|
api_key_id: "api-key-1".to_string(),
|
||||||
|
api_key_name: Some("admin-created-key".to_string()),
|
||||||
|
api_key_is_active: true,
|
||||||
|
api_key_is_locked: false,
|
||||||
|
api_key_is_standalone: false,
|
||||||
|
api_key_rate_limit: None,
|
||||||
|
api_key_concurrent_limit: None,
|
||||||
|
api_key_expires_at_unix_secs: None,
|
||||||
|
api_key_allowed_providers: None,
|
||||||
|
api_key_allowed_api_formats: None,
|
||||||
|
api_key_allowed_models: None,
|
||||||
|
currently_usable: true,
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -961,7 +961,7 @@ ORDER BY expires_at ASC, created_at ASC, id ASC
|
|||||||
allow_wallet_overage &= grant.allow_wallet_overage;
|
allow_wallet_overage &= grant.allow_wallet_overage;
|
||||||
let used = sqlx::query_scalar::<_, f64>(
|
let used = sqlx::query_scalar::<_, f64>(
|
||||||
r#"
|
r#"
|
||||||
SELECT COALESCE(SUM(amount_usd), 0)
|
SELECT CAST(COALESCE(SUM(amount_usd), 0) AS REAL)
|
||||||
FROM entitlement_usage_ledgers
|
FROM entitlement_usage_ledgers
|
||||||
WHERE user_entitlement_id = ?
|
WHERE user_entitlement_id = ?
|
||||||
AND usage_date = ?
|
AND usage_date = ?
|
||||||
|
|||||||
@@ -464,6 +464,19 @@ FOR UPDATE
|
|||||||
}
|
}
|
||||||
None => Some(0.0),
|
None => Some(0.0),
|
||||||
};
|
};
|
||||||
|
if let Some(row) = wallet_row.as_ref() {
|
||||||
|
let wallet_id: String = row.try_get("id").map_sql_err()?;
|
||||||
|
let before_recharge: f64 = row.try_get("balance").map_sql_err()?;
|
||||||
|
let before_gift: f64 = row.try_get("gift_balance").map_sql_err()?;
|
||||||
|
let before_total = before_recharge + before_gift;
|
||||||
|
settlement.wallet_id = Some(wallet_id);
|
||||||
|
settlement.wallet_balance_before = Some(before_total);
|
||||||
|
settlement.wallet_balance_after = Some(before_total);
|
||||||
|
settlement.wallet_recharge_balance_before = Some(before_recharge);
|
||||||
|
settlement.wallet_recharge_balance_after = Some(before_recharge);
|
||||||
|
settlement.wallet_gift_balance_before = Some(before_gift);
|
||||||
|
settlement.wallet_gift_balance_after = Some(before_gift);
|
||||||
|
}
|
||||||
|
|
||||||
let wallet_debit_cost_usd = if !api_key_is_standalone {
|
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()) {
|
if let Some(user_id) = input.user_id.as_deref().filter(|value| !value.is_empty()) {
|
||||||
|
|||||||
@@ -530,6 +530,20 @@ LIMIT 1
|
|||||||
}
|
}
|
||||||
None => Some(0.0),
|
None => Some(0.0),
|
||||||
};
|
};
|
||||||
|
if let Some(row) = wallet_row.as_ref() {
|
||||||
|
let wallet_id: String = row.try_get("id").map_postgres_err()?;
|
||||||
|
let before_recharge: f64 = row.try_get("balance").map_postgres_err()?;
|
||||||
|
let before_gift: f64 =
|
||||||
|
row.try_get("gift_balance").map_postgres_err()?;
|
||||||
|
let before_total = before_recharge + before_gift;
|
||||||
|
settlement.wallet_id = Some(wallet_id);
|
||||||
|
settlement.wallet_balance_before = Some(before_total);
|
||||||
|
settlement.wallet_balance_after = Some(before_total);
|
||||||
|
settlement.wallet_recharge_balance_before = Some(before_recharge);
|
||||||
|
settlement.wallet_recharge_balance_after = Some(before_recharge);
|
||||||
|
settlement.wallet_gift_balance_before = Some(before_gift);
|
||||||
|
settlement.wallet_gift_balance_after = Some(before_gift);
|
||||||
|
}
|
||||||
|
|
||||||
let wallet_debit_cost_usd = if !api_key_is_standalone {
|
let wallet_debit_cost_usd = if !api_key_is_standalone {
|
||||||
if let Some(user_id) =
|
if let Some(user_id) =
|
||||||
|
|||||||
@@ -270,7 +270,7 @@ ORDER BY expires_at ASC, created_at ASC, id ASC
|
|||||||
allow_wallet_overage &= grant.allow_wallet_overage;
|
allow_wallet_overage &= grant.allow_wallet_overage;
|
||||||
let used = sqlx::query_scalar::<_, f64>(
|
let used = sqlx::query_scalar::<_, f64>(
|
||||||
r#"
|
r#"
|
||||||
SELECT COALESCE(SUM(amount_usd), 0)
|
SELECT CAST(COALESCE(SUM(amount_usd), 0) AS REAL)
|
||||||
FROM entitlement_usage_ledgers
|
FROM entitlement_usage_ledgers
|
||||||
WHERE user_entitlement_id = ?
|
WHERE user_entitlement_id = ?
|
||||||
AND usage_date = ?
|
AND usage_date = ?
|
||||||
@@ -475,6 +475,19 @@ LIMIT 1
|
|||||||
}
|
}
|
||||||
None => Some(0.0),
|
None => Some(0.0),
|
||||||
};
|
};
|
||||||
|
if let Some(row) = wallet_row.as_ref() {
|
||||||
|
let wallet_id: String = row.try_get("id").map_sql_err()?;
|
||||||
|
let before_recharge = sqlite_real(row, "balance")?;
|
||||||
|
let before_gift = sqlite_real(row, "gift_balance")?;
|
||||||
|
let before_total = before_recharge + before_gift;
|
||||||
|
settlement.wallet_id = Some(wallet_id);
|
||||||
|
settlement.wallet_balance_before = Some(before_total);
|
||||||
|
settlement.wallet_balance_after = Some(before_total);
|
||||||
|
settlement.wallet_recharge_balance_before = Some(before_recharge);
|
||||||
|
settlement.wallet_recharge_balance_after = Some(before_recharge);
|
||||||
|
settlement.wallet_gift_balance_before = Some(before_gift);
|
||||||
|
settlement.wallet_gift_balance_after = Some(before_gift);
|
||||||
|
}
|
||||||
|
|
||||||
let wallet_debit_cost_usd = if !api_key_is_standalone {
|
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()) {
|
if let Some(user_id) = input.user_id.as_deref().filter(|value| !value.is_empty()) {
|
||||||
@@ -802,6 +815,58 @@ mod tests {
|
|||||||
assert_eq!(wallet_total, 12.0);
|
assert_eq!(wallet_total, 12.0);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn sqlite_repository_records_wallet_for_quota_covered_user_usage() {
|
||||||
|
let pool = sqlx::sqlite::SqlitePoolOptions::new()
|
||||||
|
.max_connections(1)
|
||||||
|
.connect("sqlite::memory:")
|
||||||
|
.await
|
||||||
|
.expect("sqlite pool should connect");
|
||||||
|
run_sqlite_migrations(&pool)
|
||||||
|
.await
|
||||||
|
.expect("sqlite migrations should run");
|
||||||
|
seed_quota_covered_settlement_rows(&pool).await;
|
||||||
|
|
||||||
|
let repository = SqliteSettlementRepository::new(pool.clone());
|
||||||
|
let settlement = repository
|
||||||
|
.settle_usage(UsageSettlementInput {
|
||||||
|
request_id: "request-quota-covered".to_string(),
|
||||||
|
user_id: Some("user-quota".to_string()),
|
||||||
|
api_key_id: Some("key-quota".to_string()),
|
||||||
|
api_key_is_standalone: false,
|
||||||
|
provider_id: None,
|
||||||
|
status: "completed".to_string(),
|
||||||
|
billing_status: "pending".to_string(),
|
||||||
|
total_cost_usd: 3.0,
|
||||||
|
actual_total_cost_usd: 2.0,
|
||||||
|
finalized_at_unix_secs: Some(1_260),
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.expect("settlement should run")
|
||||||
|
.expect("usage should exist");
|
||||||
|
|
||||||
|
assert_eq!(settlement.billing_status, "settled");
|
||||||
|
assert_eq!(settlement.wallet_id.as_deref(), Some("wallet-quota"));
|
||||||
|
assert_eq!(settlement.wallet_balance_before, Some(0.0));
|
||||||
|
assert_eq!(settlement.wallet_balance_after, Some(0.0));
|
||||||
|
|
||||||
|
let wallet_total: f64 = sqlx::query_scalar(
|
||||||
|
"SELECT balance + gift_balance FROM wallets WHERE id = 'wallet-quota'",
|
||||||
|
)
|
||||||
|
.fetch_one(&pool)
|
||||||
|
.await
|
||||||
|
.expect("wallet should load");
|
||||||
|
assert_eq!(wallet_total, 0.0);
|
||||||
|
|
||||||
|
let quota_used: f64 = sqlx::query_scalar(
|
||||||
|
"SELECT CAST(COALESCE(SUM(amount_usd), 0) AS REAL) FROM entitlement_usage_ledgers WHERE request_id = 'request-quota-covered'",
|
||||||
|
)
|
||||||
|
.fetch_one(&pool)
|
||||||
|
.await
|
||||||
|
.expect("quota ledger should load");
|
||||||
|
assert_eq!(quota_used, 3.0);
|
||||||
|
}
|
||||||
|
|
||||||
async fn seed_settlement_rows(pool: &sqlx::SqlitePool) {
|
async fn seed_settlement_rows(pool: &sqlx::SqlitePool) {
|
||||||
sqlx::query(
|
sqlx::query(
|
||||||
r#"
|
r#"
|
||||||
@@ -827,4 +892,62 @@ VALUES
|
|||||||
.await
|
.await
|
||||||
.expect("settlement rows should seed");
|
.expect("settlement rows should seed");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn seed_quota_covered_settlement_rows(pool: &sqlx::SqlitePool) {
|
||||||
|
sqlx::query(
|
||||||
|
r#"
|
||||||
|
INSERT INTO users (
|
||||||
|
id, username, email, role, auth_source, password_hash, is_active,
|
||||||
|
is_deleted, created_at, updated_at
|
||||||
|
) VALUES (
|
||||||
|
'user-quota', 'quota-user', 'quota@example.com', 'user', 'local',
|
||||||
|
'hash', 1, 0, 1, 1
|
||||||
|
);
|
||||||
|
|
||||||
|
INSERT INTO wallets (
|
||||||
|
id, user_id, balance, gift_balance, limit_mode, created_at, updated_at
|
||||||
|
) VALUES (
|
||||||
|
'wallet-quota', 'user-quota', 0.0, 0.0, 'finite', 1, 1
|
||||||
|
);
|
||||||
|
|
||||||
|
INSERT INTO "usage" (
|
||||||
|
request_id, user_id, api_key_id, status, billing_status,
|
||||||
|
total_cost_usd, actual_total_cost_usd
|
||||||
|
) VALUES (
|
||||||
|
'request-quota-covered', 'user-quota', 'key-quota', 'completed',
|
||||||
|
'pending', 3.0, 2.0
|
||||||
|
);
|
||||||
|
|
||||||
|
INSERT INTO billing_plans (
|
||||||
|
id, title, price_amount, price_currency, duration_unit,
|
||||||
|
duration_value, entitlements_json, created_at, updated_at
|
||||||
|
) VALUES (
|
||||||
|
'plan-quota', 'Quota Plan', 0.0, 'USD', 'month', 1,
|
||||||
|
'[{"type":"daily_quota","daily_quota_usd":10.0,"reset_timezone":"Asia/Shanghai","allow_wallet_overage":false}]',
|
||||||
|
1, 1
|
||||||
|
);
|
||||||
|
|
||||||
|
INSERT INTO payment_orders (
|
||||||
|
id, order_no, wallet_id, user_id, amount_usd, refunded_amount_usd,
|
||||||
|
refundable_amount_usd, payment_method, gateway_response, status, created_at
|
||||||
|
) VALUES (
|
||||||
|
'order-quota', 'order-quota', 'wallet-quota', 'user-quota', 0.0, 0.0,
|
||||||
|
0.0, 'admin_manual', '{}', 'credited', 1
|
||||||
|
);
|
||||||
|
|
||||||
|
INSERT INTO user_plan_entitlements (
|
||||||
|
id, user_id, plan_id, payment_order_id, status, starts_at, expires_at,
|
||||||
|
entitlements_snapshot, created_at, updated_at
|
||||||
|
) VALUES (
|
||||||
|
'entitlement-quota', 'user-quota', 'plan-quota', 'order-quota',
|
||||||
|
'active', 1, 9999999999,
|
||||||
|
'[{"type":"daily_quota","daily_quota_usd":10.0,"reset_timezone":"Asia/Shanghai","allow_wallet_overage":false}]',
|
||||||
|
1, 1
|
||||||
|
);
|
||||||
|
"#,
|
||||||
|
)
|
||||||
|
.execute(pool)
|
||||||
|
.await
|
||||||
|
.expect("quota settlement rows should seed");
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -63,10 +63,7 @@ impl WalletSnapshot {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn access_decision(&self, is_admin: bool) -> WalletAccessDecision {
|
pub fn access_decision(&self, _is_admin: bool) -> WalletAccessDecision {
|
||||||
if is_admin {
|
|
||||||
return WalletAccessDecision::allowed(None);
|
|
||||||
}
|
|
||||||
if self.status != WalletStatus::Active {
|
if self.status != WalletStatus::Active {
|
||||||
return WalletAccessDecision::wallet_unavailable(self.balance_snapshot());
|
return WalletAccessDecision::wallet_unavailable(self.balance_snapshot());
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user