fix(wallet): 修复额度耗尽后仍可消费

This commit is contained in:
Entropy.Xu
2026-05-17 11:49:08 +08:00
parent a2f91b4108
commit 1b24e1c22a
17 changed files with 925 additions and 160 deletions

View File

@@ -30,6 +30,25 @@ use axum::response::IntoResponse;
use axum::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(
state: &AppState,
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(),
));
};
let provided_auth_context = payload.auth_context.is_some();
if let Some(auth_context) = payload.auth_context {
resolved.auth_context = Some(auth_context);
resolved.local_auth_rejection = None;
}
let provided_auth_context =
apply_supplied_auth_context(state, &mut resolved, payload.auth_context).await?;
let auth_context = resolved.auth_context.as_ref();
if auth_context
.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(),
));
};
let provided_auth_context = payload.auth_context.is_some();
if let Some(auth_context) = payload.auth_context {
resolved.auth_context = Some(auth_context);
resolved.local_auth_rejection = None;
}
let provided_auth_context =
apply_supplied_auth_context(state, &mut resolved, payload.auth_context).await?;
let auth_context = resolved.auth_context.as_ref();
if auth_context
.map(|value| !value.access_allowed)
@@ -407,11 +420,8 @@ pub(crate) async fn maybe_build_local_internal_proxy_response_impl(
else {
return Ok(Some(build_internal_gateway_proxy_public_response()));
};
let provided_auth_context = payload.auth_context.is_some();
if let Some(auth_context) = payload.auth_context {
resolved.auth_context = Some(auth_context);
resolved.local_auth_rejection = None;
}
let provided_auth_context =
apply_supplied_auth_context(state, &mut resolved, payload.auth_context).await?;
if let Some(mut planned) = api::maybe_build_sync_plan_payload(
state,
&parts,
@@ -474,11 +484,8 @@ pub(crate) async fn maybe_build_local_internal_proxy_response_impl(
else {
return Ok(Some(build_internal_gateway_proxy_public_response()));
};
let provided_auth_context = payload.auth_context.is_some();
if let Some(auth_context) = payload.auth_context {
resolved.auth_context = Some(auth_context);
resolved.local_auth_rejection = None;
}
let provided_auth_context =
apply_supplied_auth_context(state, &mut resolved, payload.auth_context).await?;
if let Some(mut planned) = api::maybe_build_stream_plan_payload(
state,
&parts,
@@ -546,10 +553,7 @@ pub(crate) async fn maybe_build_local_internal_proxy_response_impl(
else {
return Ok(None);
};
if let Some(auth_context) = payload.auth_context {
resolved.auth_context = Some(auth_context);
resolved.local_auth_rejection = None;
}
apply_supplied_auth_context(state, &mut resolved, payload.auth_context).await?;
if let Some(plan_payload) = api::maybe_build_sync_plan_payload(
state,
&parts,
@@ -630,10 +634,7 @@ pub(crate) async fn maybe_build_local_internal_proxy_response_impl(
else {
return Ok(None);
};
if let Some(auth_context) = payload.auth_context {
resolved.auth_context = Some(auth_context);
resolved.local_auth_rejection = None;
}
apply_supplied_auth_context(state, &mut resolved, payload.auth_context).await?;
if let Some(plan_payload) = api::maybe_build_stream_plan_payload(
state,
&parts,

View File

@@ -30,10 +30,10 @@ mod refunds;
use self::flow::handle_wallet_flow;
pub(in crate::handlers::public::support) use self::reads::build_wallet_balance_payload_for_user;
use self::reads::{
build_wallet_daily_usage_payload, build_wallet_payload, build_wallet_zero_today_entry,
handle_wallet_balance, handle_wallet_today_cost, handle_wallet_transactions,
parse_wallet_limit, parse_wallet_offset, wallet_fixed_offset, wallet_today_billing_date_string,
wallet_transaction_payload_from_record,
build_wallet_daily_usage_payload, build_wallet_live_today_usage_payload_for_user,
build_wallet_payload, build_wallet_zero_today_entry, handle_wallet_balance,
handle_wallet_today_cost, handle_wallet_transactions, parse_wallet_limit, parse_wallet_offset,
wallet_fixed_offset, wallet_transaction_payload_from_record,
};
pub(crate) use self::recharge::sanitize_wallet_gateway_response;
use self::recharge::{

View File

@@ -1,9 +1,10 @@
use super::{
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,
parse_wallet_offset, resolve_authenticated_local_user, unix_secs_to_rfc3339,
wallet_fixed_offset, wallet_today_billing_date_string, wallet_transaction_payload_from_record,
AppState, Body, GatewayPublicRequestContext, Response, WALLET_LEGACY_TIMEZONE,
build_wallet_live_today_usage_payload_for_user, build_wallet_payload,
build_wallet_zero_today_entry, http, parse_wallet_limit, parse_wallet_offset,
resolve_authenticated_local_user, unix_secs_to_rfc3339, wallet_fixed_offset,
wallet_transaction_payload_from_record, AppState, Body, GatewayPublicRequestContext, Response,
WALLET_LEGACY_TIMEZONE,
};
use serde_json::json;
@@ -98,32 +99,43 @@ pub(super) async fn handle_wallet_flow(
return build_auth_json_response(http::StatusCode::OK, payload, None);
};
let mut today_entry = build_wallet_zero_today_entry();
if let Ok(Some(today_usage)) = state
.find_wallet_today_usage(&wallet.id, WALLET_LEGACY_TIMEZONE)
.await
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
{
today_entry = build_wallet_daily_usage_payload(
today_usage.id,
today_usage.billing_date,
today_usage.billing_timezone,
today_usage.total_cost_usd,
today_usage.total_requests,
today_usage.input_tokens,
today_usage.output_tokens,
today_usage.cache_creation_tokens,
today_usage.cache_read_tokens,
today_usage
.first_finalized_at_unix_secs
.and_then(unix_secs_to_rfc3339),
today_usage
.last_finalized_at_unix_secs
.and_then(unix_secs_to_rfc3339),
today_usage
.aggregated_at_unix_secs
.and_then(unix_secs_to_rfc3339),
true,
);
if let Ok(Some(today_usage)) = state
.find_wallet_today_usage(&wallet.id, WALLET_LEGACY_TIMEZONE)
.await
{
today_entry = build_wallet_daily_usage_payload(
today_usage.id,
today_usage.billing_date,
today_usage.billing_timezone,
today_usage.total_cost_usd,
today_usage.total_requests,
today_usage.input_tokens,
today_usage.output_tokens,
today_usage.cache_creation_tokens,
today_usage.cache_read_tokens,
today_usage
.first_finalized_at_unix_secs
.and_then(unix_secs_to_rfc3339),
today_usage
.last_finalized_at_unix_secs
.and_then(unix_secs_to_rfc3339),
today_usage
.aggregated_at_unix_secs
.and_then(unix_secs_to_rfc3339),
true,
);
}
}
let fetch_size = offset.saturating_add(limit).min(5200);

View File

@@ -1,11 +1,11 @@
use super::{
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,
GatewayError, GatewayPublicRequestContext, Response, WALLET_LEGACY_TIMEZONE,
GatewayPublicRequestContext, Response, WALLET_LEGACY_TIMEZONE,
};
use crate::handlers::shared::round_to;
use aether_data_contracts::repository::usage::UsageSettledCostSummaryQuery;
use chrono::Utc;
use chrono::{TimeZone, Utc};
use serde_json::json;
const WALLET_TODAY_COST_UNAVAILABLE_DETAIL: &str = "钱包今日费用数据暂不可用";
@@ -143,6 +143,26 @@ pub(super) fn wallet_today_billing_date_string() -> 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(
id: Option<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(
record: &aether_data::repository::wallet::StoredAdminWalletTransaction,
) -> serde_json::Value {
@@ -253,65 +310,17 @@ pub(super) async fn handle_wallet_today_cost(
Ok(value) => value,
Err(response) => return response,
};
let today = Utc::now().date_naive();
let Some(start_of_day) = today.and_hms_opt(0, 0, 0) else {
return build_auth_error_response(
http::StatusCode::INTERNAL_SERVER_ERROR,
"wallet today start is invalid",
match build_wallet_live_today_usage_payload_for_user(state, &auth.user.id).await {
Ok(Some(payload)) => build_auth_json_response(http::StatusCode::OK, payload, None),
Ok(None) => build_auth_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
WALLET_TODAY_COST_UNAVAILABLE_DETAIL,
false,
);
};
let start_unix_secs = u64::try_from(
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,
)
),
Err(detail) => {
build_auth_error_response(http::StatusCode::INTERNAL_SERVER_ERROR, detail, 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(