mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-09 18:59:50 +08:00
Merge origin/main into main
Integrate upstream updates while preserving the local analytics dashboards and schema-only migration changes. Combine user account analysis with upstream user/group usage statistics in separate tabs, retain all migration versions, and keep the deleted audit document removed. Validation: gateway all-target cargo check, frontend type check and 57 focused tests, 48 migration tests, schema composition checks, and diff whitespace checks.
This commit is contained in:
@@ -63,13 +63,14 @@ pub(in super::super) async fn build_admin_wallet_adjust_response(
|
||||
}
|
||||
let operator_id = admin_wallet_operator_id(request_context);
|
||||
let has_wallet_writer = state.has_wallet_data_writer();
|
||||
let Some((wallet, transaction)) = state
|
||||
let Some((wallet, Some(transaction))) = state
|
||||
.admin_adjust_wallet_balance(
|
||||
&wallet_id,
|
||||
amount_usd,
|
||||
&balance_type,
|
||||
operator_id.as_deref(),
|
||||
description.as_deref(),
|
||||
false,
|
||||
)
|
||||
.await?
|
||||
else {
|
||||
|
||||
@@ -12,3 +12,138 @@ pub(crate) use self::stats::{
|
||||
};
|
||||
pub(crate) use self::stats::{AdminStatsTimeRange, AdminStatsUsageFilter};
|
||||
pub(crate) use self::usage::maybe_build_local_admin_usage_response;
|
||||
|
||||
pub(crate) async fn resolve_usage_user_group_scope(
|
||||
state: &crate::handlers::admin::request::AdminAppState<'_>,
|
||||
query: Option<&str>,
|
||||
include_inactive: bool,
|
||||
exclude_admin: bool,
|
||||
) -> Result<Result<Option<Vec<String>>, String>, crate::GatewayError> {
|
||||
let group_id = crate::handlers::admin::shared::query_param_value(query, "user_group_id");
|
||||
let Some(group_id) = group_id else {
|
||||
return Ok(Ok(None));
|
||||
};
|
||||
if crate::handlers::admin::shared::query_param_value(query, "user_id").is_some() {
|
||||
return Ok(Err(
|
||||
"user_id and user_group_id cannot be used together".to_string()
|
||||
));
|
||||
}
|
||||
if !state.has_user_data_reader() {
|
||||
return Ok(Err("user group data is unavailable".to_string()));
|
||||
}
|
||||
|
||||
if group_id == UNGROUPED_USAGE_ID {
|
||||
let ids = ungrouped_usage_users(state)
|
||||
.await?
|
||||
.into_iter()
|
||||
.filter(|user| include_inactive || user.is_active)
|
||||
.filter(|user| !exclude_admin || !user.role.eq_ignore_ascii_case("admin"))
|
||||
.map(|user| user.id)
|
||||
.collect();
|
||||
return Ok(Ok(Some(ids)));
|
||||
}
|
||||
match state
|
||||
.resolve_usage_user_group_member_ids(&group_id, include_inactive, exclude_admin)
|
||||
.await?
|
||||
{
|
||||
Some(user_ids) => Ok(Ok(Some(user_ids))),
|
||||
None => Ok(Err("user_group_id does not exist".to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
/// Reserved statistics-only scope; never a permission group.
|
||||
pub(crate) const UNGROUPED_USAGE_ID: &str = "__ungrouped__";
|
||||
|
||||
pub(crate) async fn ungrouped_usage_users(
|
||||
state: &crate::handlers::admin::request::AdminAppState<'_>,
|
||||
) -> Result<Vec<aether_data::repository::users::StoredUserSummary>, crate::GatewayError> {
|
||||
use aether_data::repository::users::UserExportListQuery;
|
||||
let mut users = Vec::new();
|
||||
let mut skip = 0;
|
||||
loop {
|
||||
let page = state
|
||||
.list_export_users_page(&UserExportListQuery {
|
||||
skip,
|
||||
limit: 500,
|
||||
..Default::default()
|
||||
})
|
||||
.await?;
|
||||
let count = page.len();
|
||||
if count == 0 {
|
||||
break;
|
||||
}
|
||||
let ids = page.into_iter().map(|user| user.id).collect::<Vec<_>>();
|
||||
let grouped = state
|
||||
.list_user_group_memberships_by_user_ids(&ids)
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(|membership| membership.user_id)
|
||||
.collect::<std::collections::BTreeSet<_>>();
|
||||
let ids = ids
|
||||
.into_iter()
|
||||
.filter(|id| !grouped.contains(id))
|
||||
.collect::<Vec<_>>();
|
||||
users.extend(
|
||||
state
|
||||
.list_users_by_ids(&ids)
|
||||
.await?
|
||||
.into_iter()
|
||||
.filter(|user| !user.is_deleted),
|
||||
);
|
||||
skip += count;
|
||||
if count < 500 {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(users)
|
||||
}
|
||||
|
||||
/// Current group provider policy, resolved to the provider-name dimension used by usage rollups.
|
||||
/// None is unrestricted; Some(empty) deliberately matches no usage.
|
||||
pub(crate) async fn usage_group_provider_names(
|
||||
state: &crate::handlers::admin::request::AdminAppState<'_>,
|
||||
group: &aether_data::repository::users::StoredUserGroup,
|
||||
) -> Result<Option<Vec<String>>, crate::GatewayError> {
|
||||
if matches!(
|
||||
group.allowed_providers_mode.as_str(),
|
||||
"unrestricted" | "inherit"
|
||||
) {
|
||||
return Ok(None);
|
||||
}
|
||||
if group.allowed_providers_mode != "specific" {
|
||||
return Ok(Some(Vec::new()));
|
||||
}
|
||||
let allowed = group.allowed_providers.as_deref().unwrap_or_default();
|
||||
let providers = state.list_provider_catalog_providers(false).await?;
|
||||
let mut names = providers
|
||||
.into_iter()
|
||||
.filter(|provider| {
|
||||
allowed.iter().any(|value| {
|
||||
let value = value.trim();
|
||||
value.eq_ignore_ascii_case(&provider.id)
|
||||
|| value.eq_ignore_ascii_case(&provider.name)
|
||||
|| value.eq_ignore_ascii_case(&provider.provider_type)
|
||||
})
|
||||
})
|
||||
.map(|provider| provider.name)
|
||||
.collect::<Vec<_>>();
|
||||
names.sort();
|
||||
names.dedup();
|
||||
Ok(Some(names))
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_usage_group_provider_names(
|
||||
state: &crate::handlers::admin::request::AdminAppState<'_>,
|
||||
query: Option<&str>,
|
||||
) -> Result<Option<Vec<String>>, crate::GatewayError> {
|
||||
let Some(id) = crate::handlers::admin::shared::query_param_value(query, "user_group_id") else {
|
||||
return Ok(None);
|
||||
};
|
||||
if id == UNGROUPED_USAGE_ID {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(group) = state.find_user_group_by_id(&id).await? else {
|
||||
return Ok(Some(Vec::new()));
|
||||
};
|
||||
usage_group_provider_names(state, &group).await
|
||||
}
|
||||
|
||||
@@ -169,9 +169,11 @@ pub(super) async fn build_admin_monitoring_system_status_response(
|
||||
|
||||
let today_usage = state
|
||||
.summarize_usage_audits(&UsageAuditSummaryQuery {
|
||||
provider_names: None,
|
||||
created_from_unix_secs: today_start.timestamp().max(0) as u64,
|
||||
created_until_unix_secs: now_unix_secs.saturating_add(1),
|
||||
user_id: None,
|
||||
user_ids: None,
|
||||
provider_name: None,
|
||||
model: None,
|
||||
})
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
use super::super::resolve_usage_user_group_scope;
|
||||
use super::range::{build_comparison_range, parse_bounded_u32};
|
||||
use super::resolve_admin_usage_time_range;
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||
@@ -109,6 +110,7 @@ pub(super) async fn maybe_build_local_admin_stats_analytics_response(
|
||||
};
|
||||
let current_summary = state
|
||||
.summarize_usage_audits(&UsageAuditSummaryQuery {
|
||||
provider_names: None,
|
||||
created_from_unix_secs: current_from_unix_secs,
|
||||
created_until_unix_secs: current_until_unix_secs,
|
||||
..Default::default()
|
||||
@@ -116,6 +118,7 @@ pub(super) async fn maybe_build_local_admin_stats_analytics_response(
|
||||
.await?;
|
||||
let comparison_summary = state
|
||||
.summarize_usage_audits(&UsageAuditSummaryQuery {
|
||||
provider_names: None,
|
||||
created_from_unix_secs: comparison_from_unix_secs,
|
||||
created_until_unix_secs: comparison_until_unix_secs,
|
||||
..Default::default()
|
||||
@@ -294,6 +297,17 @@ pub(super) async fn maybe_build_local_admin_stats_analytics_response(
|
||||
}
|
||||
|
||||
let filters = AdminStatsUsageFilter::from_query(request_context.query_string());
|
||||
let user_ids = match resolve_usage_user_group_scope(
|
||||
state,
|
||||
request_context.query_string(),
|
||||
false,
|
||||
false,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
Ok(value) => value,
|
||||
Err(detail) => return Ok(Some(admin_stats_bad_request_response(detail))),
|
||||
};
|
||||
let query_granularity = match granularity {
|
||||
AdminStatsGranularity::Hour => UsageTimeSeriesGranularity::Hour,
|
||||
AdminStatsGranularity::Day
|
||||
@@ -306,11 +320,17 @@ pub(super) async fn maybe_build_local_admin_stats_analytics_response(
|
||||
};
|
||||
let buckets = state
|
||||
.summarize_usage_time_series(&UsageTimeSeriesQuery {
|
||||
provider_names: super::super::resolve_usage_group_provider_names(
|
||||
state,
|
||||
request_context.query_string(),
|
||||
)
|
||||
.await?,
|
||||
created_from_unix_secs,
|
||||
created_until_unix_secs,
|
||||
granularity: query_granularity,
|
||||
tz_offset_minutes: time_range.tz_offset_minutes,
|
||||
user_id: filters.user_id,
|
||||
user_ids,
|
||||
provider_name: filters.provider_name,
|
||||
model: filters.model,
|
||||
})
|
||||
|
||||
@@ -72,11 +72,13 @@ pub(super) async fn maybe_build_local_admin_stats_cost_response(
|
||||
};
|
||||
let buckets = state
|
||||
.summarize_usage_time_series(&UsageTimeSeriesQuery {
|
||||
provider_names: None,
|
||||
created_from_unix_secs,
|
||||
created_until_unix_secs,
|
||||
granularity: UsageTimeSeriesGranularity::Day,
|
||||
tz_offset_minutes: time_range.tz_offset_minutes,
|
||||
user_id: None,
|
||||
user_ids: None,
|
||||
provider_name: None,
|
||||
model: None,
|
||||
})
|
||||
|
||||
@@ -3,12 +3,13 @@ use crate::GatewayError;
|
||||
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
|
||||
|
||||
pub(super) use aether_admin::observability::stats::{
|
||||
build_admin_stats_leaderboard_response, build_api_key_leaderboard_items,
|
||||
build_api_key_leaderboard_items_from_summaries, build_model_leaderboard_items,
|
||||
build_model_leaderboard_items_from_summaries, build_user_leaderboard_items,
|
||||
build_user_leaderboard_items_from_summaries, compare_leaderboard_items, compute_dense_rank,
|
||||
AdminStatsLeaderboardItem, AdminStatsLeaderboardMetric, AdminStatsLeaderboardNameMode,
|
||||
AdminStatsSortOrder, AdminStatsUserMetadata,
|
||||
build_admin_stats_leaderboard_response, build_admin_stats_user_group_leaderboard_response,
|
||||
build_api_key_leaderboard_items, build_api_key_leaderboard_items_from_summaries,
|
||||
build_model_leaderboard_items, build_model_leaderboard_items_from_summaries,
|
||||
build_user_leaderboard_items, build_user_leaderboard_items_from_summaries,
|
||||
compare_leaderboard_items, compute_dense_rank, AdminStatsLeaderboardItem,
|
||||
AdminStatsLeaderboardMetric, AdminStatsLeaderboardNameMode, AdminStatsSortOrder,
|
||||
AdminStatsUserMetadata,
|
||||
};
|
||||
|
||||
pub(super) async fn load_user_leaderboard_metadata(
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
use super::super::resolve_usage_user_group_scope;
|
||||
use super::leaderboard::{
|
||||
build_admin_stats_leaderboard_response, build_api_key_leaderboard_items_from_summaries,
|
||||
build_model_leaderboard_items_from_summaries, build_user_leaderboard_items_from_summaries,
|
||||
compare_leaderboard_items, load_user_leaderboard_metadata, AdminStatsLeaderboardNameMode,
|
||||
build_admin_stats_leaderboard_response, build_admin_stats_user_group_leaderboard_response,
|
||||
build_api_key_leaderboard_items_from_summaries, build_model_leaderboard_items_from_summaries,
|
||||
build_user_leaderboard_items_from_summaries, compare_leaderboard_items,
|
||||
load_user_leaderboard_metadata, AdminStatsLeaderboardItem, AdminStatsLeaderboardNameMode,
|
||||
};
|
||||
use super::range::{parse_bounded_u32, parse_nonnegative_usize};
|
||||
use super::resolve_admin_usage_time_range;
|
||||
@@ -14,6 +16,7 @@ use aether_admin::observability::stats::{
|
||||
};
|
||||
use aether_data_contracts::repository::usage::{UsageLeaderboardGroupBy, UsageLeaderboardQuery};
|
||||
use axum::{body::Body, http, response::Response};
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
pub(super) async fn maybe_build_local_admin_stats_leaderboard_response(
|
||||
state: &AdminAppState<'_>,
|
||||
@@ -75,10 +78,12 @@ pub(super) async fn maybe_build_local_admin_stats_leaderboard_response(
|
||||
};
|
||||
let summaries = state
|
||||
.summarize_usage_leaderboard(&UsageLeaderboardQuery {
|
||||
provider_names: None,
|
||||
created_from_unix_secs,
|
||||
created_until_unix_secs,
|
||||
group_by: UsageLeaderboardGroupBy::Model,
|
||||
user_id: filters.user_id,
|
||||
user_ids: None,
|
||||
provider_name: filters.provider_name,
|
||||
model: filters.model,
|
||||
})
|
||||
@@ -152,10 +157,12 @@ pub(super) async fn maybe_build_local_admin_stats_leaderboard_response(
|
||||
};
|
||||
let summaries = state
|
||||
.summarize_usage_leaderboard(&UsageLeaderboardQuery {
|
||||
provider_names: None,
|
||||
created_from_unix_secs,
|
||||
created_until_unix_secs,
|
||||
group_by: UsageLeaderboardGroupBy::ApiKey,
|
||||
user_id: filters.user_id,
|
||||
user_ids: None,
|
||||
provider_name: filters.provider_name,
|
||||
model: filters.model,
|
||||
})
|
||||
@@ -206,6 +213,181 @@ pub(super) async fn maybe_build_local_admin_stats_leaderboard_response(
|
||||
)));
|
||||
}
|
||||
|
||||
if request_context
|
||||
.decision()
|
||||
.and_then(|decision| decision.route_kind.as_deref())
|
||||
== Some("leaderboard_user_groups")
|
||||
&& request_context.method() == http::Method::GET
|
||||
&& matches!(
|
||||
request_context.path(),
|
||||
"/api/admin/stats/leaderboard/user-groups"
|
||||
| "/api/admin/stats/leaderboard/user-groups/"
|
||||
)
|
||||
{
|
||||
let time_range = match resolve_admin_usage_time_range(query) {
|
||||
Ok(value) => value,
|
||||
Err(detail) => return Ok(Some(admin_stats_bad_request_response(detail))),
|
||||
};
|
||||
let metric = match AdminStatsLeaderboardMetric::parse(query) {
|
||||
Ok(value) => value,
|
||||
Err(detail) => return Ok(Some(admin_stats_bad_request_response(detail))),
|
||||
};
|
||||
let order = match AdminStatsSortOrder::parse(query) {
|
||||
Ok(value) => value,
|
||||
Err(detail) => return Ok(Some(admin_stats_bad_request_response(detail))),
|
||||
};
|
||||
let limit = match query_param_value(query, "limit")
|
||||
.map(|value| parse_bounded_u32("limit", &value, 1, 100))
|
||||
.transpose()
|
||||
{
|
||||
Ok(Some(value)) => value as usize,
|
||||
Ok(None) => 10,
|
||||
Err(detail) => return Ok(Some(admin_stats_bad_request_response(detail))),
|
||||
};
|
||||
let offset = match query_param_value(query, "offset")
|
||||
.map(|value| parse_nonnegative_usize("offset", &value))
|
||||
.transpose()
|
||||
{
|
||||
Ok(Some(value)) => value,
|
||||
Ok(None) => 0,
|
||||
Err(detail) => return Ok(Some(admin_stats_bad_request_response(detail))),
|
||||
};
|
||||
let empty_counts = BTreeMap::new();
|
||||
if !state.has_usage_data_reader() || !state.has_user_data_reader() {
|
||||
return Ok(Some(build_admin_stats_user_group_leaderboard_response(
|
||||
metric,
|
||||
Some(&time_range),
|
||||
&[],
|
||||
&empty_counts,
|
||||
&empty_counts,
|
||||
offset,
|
||||
limit,
|
||||
)));
|
||||
}
|
||||
let include_inactive = query_param_bool(query, "include_inactive", false);
|
||||
let exclude_admin = query_param_bool(query, "exclude_admin", false);
|
||||
let filters = AdminStatsUsageFilter::from_query(query);
|
||||
if filters.user_id.is_some() {
|
||||
return Ok(Some(admin_stats_bad_request_response(
|
||||
"user_id is not supported for the user group leaderboard".to_string(),
|
||||
)));
|
||||
}
|
||||
let Some((created_from_unix_secs, created_until_unix_secs)) = time_range.to_unix_bounds()
|
||||
else {
|
||||
return Ok(Some(build_admin_stats_user_group_leaderboard_response(
|
||||
metric,
|
||||
Some(&time_range),
|
||||
&[],
|
||||
&empty_counts,
|
||||
&empty_counts,
|
||||
offset,
|
||||
limit,
|
||||
)));
|
||||
};
|
||||
|
||||
let mut leaderboard = Vec::new();
|
||||
let mut member_counts = BTreeMap::new();
|
||||
let mut active_member_counts = BTreeMap::new();
|
||||
for group in state.list_user_groups().await? {
|
||||
let members = state.list_user_group_members(&group.id).await?;
|
||||
let member_count = members.iter().filter(|member| !member.is_deleted).count();
|
||||
let active_member_count = members
|
||||
.iter()
|
||||
.filter(|member| !member.is_deleted && member.is_active)
|
||||
.count();
|
||||
let user_ids = members
|
||||
.iter()
|
||||
.filter(|member| !member.is_deleted)
|
||||
.filter(|member| include_inactive || member.is_active)
|
||||
.filter(|member| !exclude_admin || !member.role.eq_ignore_ascii_case("admin"))
|
||||
.map(|member| member.user_id.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let summaries = state
|
||||
.summarize_usage_leaderboard(&UsageLeaderboardQuery {
|
||||
created_from_unix_secs,
|
||||
created_until_unix_secs,
|
||||
group_by: UsageLeaderboardGroupBy::User,
|
||||
user_id: None,
|
||||
user_ids: Some(user_ids),
|
||||
provider_names: super::super::usage_group_provider_names(state, &group).await?,
|
||||
provider_name: filters.provider_name.clone(),
|
||||
model: filters.model.clone(),
|
||||
})
|
||||
.await?;
|
||||
let user_ids = summaries
|
||||
.iter()
|
||||
.map(|row| row.group_key.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let metadata = load_user_leaderboard_metadata(state, &user_ids).await?;
|
||||
let users = build_user_leaderboard_items_from_summaries(
|
||||
&summaries,
|
||||
&metadata,
|
||||
state.has_auth_user_data_reader(),
|
||||
state.has_user_data_reader(),
|
||||
include_inactive,
|
||||
exclude_admin,
|
||||
);
|
||||
let mut item = AdminStatsLeaderboardItem {
|
||||
id: group.id.clone(),
|
||||
name: group.name,
|
||||
requests: 0,
|
||||
tokens: 0,
|
||||
cost: 0.0,
|
||||
};
|
||||
for user in users {
|
||||
item.requests = item.requests.saturating_add(user.requests);
|
||||
item.tokens = item.tokens.saturating_add(user.tokens);
|
||||
item.cost += user.cost;
|
||||
}
|
||||
member_counts.insert(group.id.clone(), member_count);
|
||||
active_member_counts.insert(group.id, active_member_count);
|
||||
leaderboard.push(item);
|
||||
}
|
||||
let ungrouped = super::super::ungrouped_usage_users(state).await?;
|
||||
let id = super::super::UNGROUPED_USAGE_ID.to_string();
|
||||
member_counts.insert(id.clone(), ungrouped.len());
|
||||
active_member_counts.insert(
|
||||
id.clone(),
|
||||
ungrouped.iter().filter(|user| user.is_active).count(),
|
||||
);
|
||||
let user_ids = ungrouped
|
||||
.into_iter()
|
||||
.filter(|user| include_inactive || user.is_active)
|
||||
.filter(|user| !exclude_admin || !user.role.eq_ignore_ascii_case("admin"))
|
||||
.map(|user| user.id)
|
||||
.collect();
|
||||
let rows = state
|
||||
.summarize_usage_leaderboard(&UsageLeaderboardQuery {
|
||||
created_from_unix_secs,
|
||||
created_until_unix_secs,
|
||||
group_by: UsageLeaderboardGroupBy::User,
|
||||
user_id: None,
|
||||
user_ids: Some(user_ids),
|
||||
provider_names: None,
|
||||
provider_name: filters.provider_name.clone(),
|
||||
model: filters.model.clone(),
|
||||
})
|
||||
.await?;
|
||||
leaderboard.push(AdminStatsLeaderboardItem {
|
||||
id,
|
||||
name: "Ungrouped".to_string(),
|
||||
requests: rows.iter().map(|row| row.request_count).sum(),
|
||||
tokens: rows.iter().map(|row| row.total_tokens).sum(),
|
||||
cost: rows.iter().map(|row| row.total_cost_usd).sum(),
|
||||
});
|
||||
leaderboard.sort_by(|left, right| compare_leaderboard_items(metric, order, left, right));
|
||||
|
||||
return Ok(Some(build_admin_stats_user_group_leaderboard_response(
|
||||
metric,
|
||||
Some(&time_range),
|
||||
&leaderboard,
|
||||
&member_counts,
|
||||
&active_member_counts,
|
||||
offset,
|
||||
limit,
|
||||
)));
|
||||
}
|
||||
|
||||
if request_context
|
||||
.decision()
|
||||
.and_then(|decision| decision.route_kind.as_deref())
|
||||
@@ -253,6 +435,13 @@ pub(super) async fn maybe_build_local_admin_stats_leaderboard_response(
|
||||
let include_inactive = query_param_bool(query, "include_inactive", false);
|
||||
let exclude_admin = query_param_bool(query, "exclude_admin", false);
|
||||
let filters = AdminStatsUsageFilter::from_query(query);
|
||||
let scoped_user_ids =
|
||||
match resolve_usage_user_group_scope(state, query, include_inactive, exclude_admin)
|
||||
.await?
|
||||
{
|
||||
Ok(value) => value,
|
||||
Err(detail) => return Ok(Some(admin_stats_bad_request_response(detail))),
|
||||
};
|
||||
let Some((created_from_unix_secs, created_until_unix_secs)) = time_range.to_unix_bounds()
|
||||
else {
|
||||
return Ok(Some(admin_stats_leaderboard_empty_response(
|
||||
@@ -262,10 +451,13 @@ pub(super) async fn maybe_build_local_admin_stats_leaderboard_response(
|
||||
};
|
||||
let summaries = state
|
||||
.summarize_usage_leaderboard(&UsageLeaderboardQuery {
|
||||
provider_names: super::super::resolve_usage_group_provider_names(state, query)
|
||||
.await?,
|
||||
created_from_unix_secs,
|
||||
created_until_unix_secs,
|
||||
group_by: UsageLeaderboardGroupBy::User,
|
||||
user_id: filters.user_id,
|
||||
user_ids: scoped_user_ids,
|
||||
provider_name: filters.provider_name,
|
||||
model: filters.model,
|
||||
})
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
use super::super::resolve_usage_user_group_scope;
|
||||
use super::super::stats::resolve_admin_usage_time_range;
|
||||
use super::analytics::admin_usage_api_key_names;
|
||||
use super::analytics::admin_usage_provider_key_names;
|
||||
@@ -121,21 +122,30 @@ fn apply_admin_usage_status_filter(query: &mut UsageAuditListQuery, status: Opti
|
||||
"pending" | "streaming" | "completed" | "cancelled" => {
|
||||
query.statuses = Some(vec![status]);
|
||||
}
|
||||
"has_fallback" | "has_retry" => {}
|
||||
"has_fallback" | "has_retry" | "has_skipped_candidate" => {}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default)]
|
||||
#[derive(Clone, Debug, Default)]
|
||||
struct AdminUsageAttemptFlags {
|
||||
has_fallback: bool,
|
||||
has_retry: bool,
|
||||
/// 是否存在"被调度跳过"的候选(调度阶段判定本次不可用,从未向上游发起请求)。
|
||||
///
|
||||
/// 这是与 has_fallback 正交的信号:has_fallback 表示"更靠前的候选真的失败并被换掉",
|
||||
/// 而本字段表示"更靠前的候选压根没被发出去"。两者在日志列表里观感都是"换了提供商",
|
||||
/// 但用户拿不到 has_fallback 小图标时容易误判为调度错误,故单独暴露。
|
||||
has_skipped_candidate: bool,
|
||||
/// 跳过原因(去重、保持出现顺序),用于前端 tooltip 直接说明"为什么没用它"。
|
||||
skipped_candidate_reasons: Vec<String>,
|
||||
}
|
||||
|
||||
fn admin_usage_attempt_status_filter(status: Option<&str>) -> Option<&'static str> {
|
||||
match status?.trim().to_ascii_lowercase().as_str() {
|
||||
"has_fallback" => Some("has_fallback"),
|
||||
"has_retry" => Some("has_retry"),
|
||||
"has_skipped_candidate" => Some("has_skipped_candidate"),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
@@ -190,25 +200,52 @@ fn admin_usage_attempt_flags_from_candidates(
|
||||
})
|
||||
});
|
||||
let has_retry = candidates.iter().any(admin_usage_candidate_was_retried);
|
||||
let skipped_candidate_reasons = admin_usage_skipped_candidate_reasons(candidates);
|
||||
|
||||
AdminUsageAttemptFlags {
|
||||
has_fallback,
|
||||
has_retry,
|
||||
has_skipped_candidate: !skipped_candidate_reasons.is_empty(),
|
||||
skipped_candidate_reasons,
|
||||
}
|
||||
}
|
||||
|
||||
/// 收集被跳过候选的原因,去重并保持候选顺序(决定性的在前,便于阅读)。
|
||||
fn admin_usage_skipped_candidate_reasons(candidates: &[StoredRequestCandidate]) -> Vec<String> {
|
||||
let mut reasons = Vec::new();
|
||||
for candidate in candidates
|
||||
.iter()
|
||||
.filter(|candidate| candidate.status == RequestCandidateStatus::Skipped)
|
||||
{
|
||||
let Some(reason) = candidate
|
||||
.skip_reason
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|reason| !reason.is_empty())
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
if !reasons.iter().any(|existing| existing == reason) {
|
||||
reasons.push(reason.to_string());
|
||||
}
|
||||
}
|
||||
reasons
|
||||
}
|
||||
|
||||
fn admin_usage_attempt_flags_for_item(
|
||||
item: &StoredRequestUsageAudit,
|
||||
flags_by_usage_id: &BTreeMap<String, AdminUsageAttemptFlags>,
|
||||
request_candidate_reader_available: bool,
|
||||
) -> AdminUsageAttemptFlags {
|
||||
flags_by_usage_id.get(&item.id).copied().unwrap_or_else(|| {
|
||||
flags_by_usage_id.get(&item.id).cloned().unwrap_or_else(|| {
|
||||
if request_candidate_reader_available {
|
||||
AdminUsageAttemptFlags::default()
|
||||
} else {
|
||||
AdminUsageAttemptFlags {
|
||||
has_fallback: admin_usage_has_fallback(item),
|
||||
has_retry: false,
|
||||
has_skipped_candidate: false,
|
||||
skipped_candidate_reasons: Vec::new(),
|
||||
}
|
||||
}
|
||||
})
|
||||
@@ -477,6 +514,8 @@ fn admin_usage_matches_attempt_status(
|
||||
match status {
|
||||
"has_fallback" => flags.has_fallback,
|
||||
"has_retry" => flags.has_retry,
|
||||
// 与 has_fallback 区分:这里是"更靠前的候选被调度跳过、根本没发出去"
|
||||
"has_skipped_candidate" => flags.has_skipped_candidate,
|
||||
_ => true,
|
||||
}
|
||||
}
|
||||
@@ -548,6 +587,9 @@ fn build_admin_usage_records_response_with_attempt_flags(
|
||||
);
|
||||
record["has_fallback"] = json!(flags.has_fallback);
|
||||
record["has_retry"] = json!(flags.has_retry);
|
||||
// 被跳过的候选:前端据此提示"这次没用某个提供商,是因为它在调度阶段就被排除了"。
|
||||
record["has_skipped_candidate"] = json!(flags.has_skipped_candidate);
|
||||
record["skipped_candidate_reasons"] = json!(flags.skipped_candidate_reasons);
|
||||
record
|
||||
})
|
||||
.collect();
|
||||
@@ -799,11 +841,18 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
|
||||
&Default::default(),
|
||||
)));
|
||||
};
|
||||
let user_ids = match resolve_usage_user_group_scope(state, query, false, false).await? {
|
||||
Ok(value) => value,
|
||||
Err(detail) => return Ok(Some(admin_usage_bad_request_response(detail))),
|
||||
};
|
||||
let summary = state
|
||||
.summarize_usage_audits(&UsageAuditSummaryQuery {
|
||||
provider_names: super::super::resolve_usage_group_provider_names(state, query)
|
||||
.await?,
|
||||
created_from_unix_secs,
|
||||
created_until_unix_secs,
|
||||
user_id: query_param_value(query, "user_id"),
|
||||
user_ids,
|
||||
provider_name: query_param_value(query, "provider"),
|
||||
model: query_param_value(query, "model"),
|
||||
})
|
||||
@@ -1202,9 +1251,11 @@ mod tests {
|
||||
use aether_data_contracts::repository::candidates::{
|
||||
RequestCandidateStatus, StoredRequestCandidate,
|
||||
};
|
||||
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
|
||||
use serde_json::json;
|
||||
|
||||
use super::{
|
||||
admin_usage_attempt_flags_from_candidates, admin_usage_skipped_candidate_reasons,
|
||||
admin_usage_terminal_candidate_state_override, build_admin_usage_keyword_search_query,
|
||||
build_admin_usage_records_query, latest_admin_usage_image_progress,
|
||||
AdminUsageSearchContext,
|
||||
@@ -1246,6 +1297,144 @@ mod tests {
|
||||
.expect("candidate should build")
|
||||
}
|
||||
|
||||
/// 构造一条"被调度跳过"的候选(从未向上游发起请求)。
|
||||
fn skipped_candidate(candidate_index: i32, reason: &str) -> StoredRequestCandidate {
|
||||
let mut candidate = sample_candidate(
|
||||
candidate_index,
|
||||
RequestCandidateStatus::Skipped,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
);
|
||||
candidate.skip_reason = Some(reason.to_string());
|
||||
// 跳过候选没有开始时间,is_attempted 因此为 false
|
||||
candidate.started_at_unix_ms = None;
|
||||
candidate
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn skipped_candidate_reasons_are_deduplicated_in_candidate_order() {
|
||||
let reasons = admin_usage_skipped_candidate_reasons(&[
|
||||
skipped_candidate(0, "key_rpm_exhausted"),
|
||||
skipped_candidate(1, "provider_inactive"),
|
||||
skipped_candidate(2, "key_rpm_exhausted"),
|
||||
]);
|
||||
|
||||
assert_eq!(
|
||||
reasons,
|
||||
vec![
|
||||
"key_rpm_exhausted".to_string(),
|
||||
"provider_inactive".to_string()
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn skipped_candidate_reasons_ignore_attempted_candidates() {
|
||||
// 真正发起过请求的失败候选不属于"被跳过",避免与 has_fallback 语义混淆
|
||||
let failed = sample_candidate(
|
||||
0,
|
||||
RequestCandidateStatus::Failed,
|
||||
Some(503),
|
||||
Some(1_000),
|
||||
Some("upstream exploded"),
|
||||
);
|
||||
assert!(admin_usage_skipped_candidate_reasons(&[failed]).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn attempt_flags_report_skipped_candidates_without_fallback() {
|
||||
let candidates = vec![
|
||||
skipped_candidate(0, "key_rpm_exhausted"),
|
||||
sample_candidate(
|
||||
1,
|
||||
RequestCandidateStatus::Success,
|
||||
Some(200),
|
||||
Some(900),
|
||||
None,
|
||||
),
|
||||
];
|
||||
|
||||
let flags = admin_usage_attempt_flags_from_candidates(&sample_usage_audit(), &candidates);
|
||||
|
||||
// 这正是用户遇到的场景:换了提供商,但没有任何候选失败过
|
||||
assert!(flags.has_skipped_candidate);
|
||||
assert!(!flags.has_fallback);
|
||||
assert_eq!(
|
||||
flags.skipped_candidate_reasons,
|
||||
vec!["key_rpm_exhausted".to_string()]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn attempt_flags_keep_fallback_and_skipped_candidate_independent() {
|
||||
let candidates = vec![
|
||||
skipped_candidate(0, "provider_inactive"),
|
||||
sample_candidate(
|
||||
1,
|
||||
RequestCandidateStatus::Failed,
|
||||
Some(503),
|
||||
Some(500),
|
||||
None,
|
||||
),
|
||||
sample_candidate(
|
||||
2,
|
||||
RequestCandidateStatus::Success,
|
||||
Some(200),
|
||||
Some(700),
|
||||
None,
|
||||
),
|
||||
];
|
||||
|
||||
let flags = admin_usage_attempt_flags_from_candidates(&sample_usage_audit(), &candidates);
|
||||
|
||||
assert!(flags.has_skipped_candidate);
|
||||
assert!(flags.has_fallback);
|
||||
}
|
||||
|
||||
/// 最小可用的用量审计行,仅用于驱动 flags 计算(其中候选 id 为空即可)。
|
||||
fn sample_usage_audit() -> StoredRequestUsageAudit {
|
||||
StoredRequestUsageAudit::new(
|
||||
"usage-1".to_string(),
|
||||
"req-1".to_string(),
|
||||
Some("user-1".to_string()),
|
||||
Some("api-key-1".to_string()),
|
||||
Some("alice".to_string()),
|
||||
Some("default".to_string()),
|
||||
"OpenAI".to_string(),
|
||||
"gpt-4.1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some("openai:chat".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
Some("openai:chat".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
false,
|
||||
false,
|
||||
10,
|
||||
20,
|
||||
30,
|
||||
0.0,
|
||||
0.0,
|
||||
Some(200),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
"completed".to_string(),
|
||||
"settled".to_string(),
|
||||
1_000,
|
||||
1_001,
|
||||
None,
|
||||
)
|
||||
.expect("usage should build")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn admin_usage_active_override_uses_current_terminal_candidate_latency() {
|
||||
let candidate = sample_candidate(
|
||||
|
||||
@@ -74,7 +74,8 @@ fn validate_batch_access_token_import(
|
||||
) -> Result<(), String> {
|
||||
if !provider_type_supports_access_token_import(provider_type) {
|
||||
return Err(
|
||||
"Access Token 导入仅支持 Claude Code / Codex / ChatGPT Web / Grok Provider".to_string(),
|
||||
"Access Token 导入仅支持 Claude Code / Codex / ChatGPT Web / Grok / xAI Provider"
|
||||
.to_string(),
|
||||
);
|
||||
}
|
||||
if provider_type.eq_ignore_ascii_case("claude_code") {
|
||||
|
||||
@@ -214,7 +214,12 @@ fn extract_admin_provider_oauth_batch_import_entry(
|
||||
} else {
|
||||
let sso_from_cookie = grok_cookie_session_token(provider_type, raw_token);
|
||||
let token_input = sso_from_cookie.as_deref().unwrap_or(raw_token);
|
||||
let (refresh_token, access_token) = import_tokens_from_raw_token(token_input);
|
||||
let (refresh_token, access_token) =
|
||||
if provider_type.trim().eq_ignore_ascii_case("xai") {
|
||||
(None, Some(token_input.to_string()))
|
||||
} else {
|
||||
import_tokens_from_raw_token(token_input)
|
||||
};
|
||||
let (refresh_token, access_token) = normalize_provider_import_tokens(
|
||||
provider_type,
|
||||
refresh_token.as_deref(),
|
||||
@@ -262,6 +267,7 @@ fn extract_admin_provider_oauth_batch_import_entry(
|
||||
let object = normalized_claude_object.as_ref().unwrap_or(object);
|
||||
let is_grok = provider_type.trim().eq_ignore_ascii_case("grok");
|
||||
let is_windsurf = provider_type.trim().eq_ignore_ascii_case("windsurf");
|
||||
let is_xai = provider_type.trim().eq_ignore_ascii_case("xai");
|
||||
let is_codex_agent_identity = provider_type.trim().eq_ignore_ascii_case("codex")
|
||||
&& aether_provider_transport::is_codex_agent_identity_auth_config_value(item);
|
||||
if is_codex_agent_identity {
|
||||
@@ -336,14 +342,6 @@ fn extract_admin_provider_oauth_batch_import_entry(
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let (refresh_token, access_token) = normalize_provider_import_tokens(
|
||||
provider_type,
|
||||
refresh_token.as_deref(),
|
||||
access_token
|
||||
.as_deref()
|
||||
.or(session_token.as_deref())
|
||||
.or(header_bearer_token.as_deref()),
|
||||
);
|
||||
let windsurf_api_key = is_windsurf
|
||||
.then(|| {
|
||||
coerce_admin_provider_oauth_import_str(
|
||||
@@ -351,6 +349,22 @@ fn extract_admin_provider_oauth_batch_import_entry(
|
||||
)
|
||||
})
|
||||
.flatten();
|
||||
let xai_api_key = is_xai
|
||||
.then(|| {
|
||||
coerce_admin_provider_oauth_import_str(
|
||||
object.get("api_key").or_else(|| object.get("apiKey")),
|
||||
)
|
||||
})
|
||||
.flatten();
|
||||
let (refresh_token, access_token) = normalize_provider_import_tokens(
|
||||
provider_type,
|
||||
refresh_token.as_deref(),
|
||||
access_token
|
||||
.as_deref()
|
||||
.or(session_token.as_deref())
|
||||
.or(header_bearer_token.as_deref())
|
||||
.or(xai_api_key.as_deref()),
|
||||
);
|
||||
let windsurf_token = is_windsurf
|
||||
.then(|| {
|
||||
coerce_admin_provider_oauth_import_str(
|
||||
@@ -1577,4 +1591,23 @@ mod tests {
|
||||
assert!(entries[1].access_token.is_none());
|
||||
assert!(entries[1].raw_credentials.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_xai_api_key_json_and_raw_lines_as_access_token() {
|
||||
let entries = parse_admin_provider_oauth_batch_import_entries(
|
||||
"xai",
|
||||
r#"{"api_key":"xai-api-key","email":"[email protected]"}
|
||||
{"refresh_token":"xai-refresh"}
|
||||
xai-raw-api-key"#,
|
||||
);
|
||||
|
||||
assert_eq!(entries.len(), 3);
|
||||
assert!(entries[0].refresh_token.is_none());
|
||||
assert_eq!(entries[0].access_token.as_deref(), Some("xai-api-key"));
|
||||
assert_eq!(entries[0].email.as_deref(), Some("[email protected]"));
|
||||
assert_eq!(entries[1].refresh_token.as_deref(), Some("xai-refresh"));
|
||||
assert!(entries[1].access_token.is_none());
|
||||
assert!(entries[2].refresh_token.is_none());
|
||||
assert_eq!(entries[2].access_token.as_deref(), Some("xai-raw-api-key"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -186,10 +186,10 @@ pub(super) async fn handle_admin_provider_oauth_device_authorize(
|
||||
));
|
||||
};
|
||||
let provider_type = provider.provider_type.trim().to_ascii_lowercase();
|
||||
if provider_type != "kiro" && provider_type != "windsurf" {
|
||||
if provider_type != "kiro" && provider_type != "windsurf" && provider_type != "xai" {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
"设备授权仅支持 Kiro / Windsurf provider",
|
||||
"设备授权仅支持 Kiro / Windsurf / xAI provider",
|
||||
));
|
||||
}
|
||||
let Some(principal) = request_context
|
||||
@@ -219,6 +219,19 @@ pub(super) async fn handle_admin_provider_oauth_device_authorize(
|
||||
)
|
||||
.await;
|
||||
|
||||
if provider_type == "xai" {
|
||||
return super::xai::handle_admin_provider_oauth_xai_device_authorize(
|
||||
state,
|
||||
&provider_id,
|
||||
&provider,
|
||||
principal,
|
||||
runtime_endpoint.as_ref(),
|
||||
request_proxy,
|
||||
payload.proxy_node_id.as_deref(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
if provider_type == "windsurf" {
|
||||
let session_id = generate_provider_oauth_nonce();
|
||||
let login_option = payload
|
||||
|
||||
@@ -2,6 +2,7 @@ mod authorize;
|
||||
mod lease;
|
||||
mod poll;
|
||||
mod session;
|
||||
mod xai;
|
||||
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||
use crate::GatewayError;
|
||||
|
||||
@@ -479,6 +479,18 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
|
||||
)
|
||||
.await;
|
||||
|
||||
if provider_type == "xai" {
|
||||
return super::xai::handle_admin_provider_oauth_xai_device_poll(
|
||||
state,
|
||||
&provider,
|
||||
&endpoints,
|
||||
request_proxy,
|
||||
session_id,
|
||||
session,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
if provider_type == "windsurf" {
|
||||
return handle_admin_provider_oauth_windsurf_browser_device_poll(
|
||||
state,
|
||||
|
||||
@@ -0,0 +1,368 @@
|
||||
use super::session::attach_admin_provider_oauth_device_poll_terminal_response;
|
||||
use crate::control::GatewayAdminPrincipalContext;
|
||||
use crate::handlers::admin::provider::oauth::dispatch::helpers::admin_provider_oauth_key_name_from_auth_config;
|
||||
use crate::handlers::admin::provider::oauth::errors::build_internal_control_error_response;
|
||||
use crate::handlers::admin::provider::oauth::provisioning::{
|
||||
provider_oauth_active_api_formats, provider_oauth_key_proxy_value,
|
||||
};
|
||||
use crate::handlers::admin::provider::oauth::runtime::spawn_provider_oauth_account_state_refresh_after_update;
|
||||
use crate::handlers::admin::provider::oauth::state::{
|
||||
current_unix_secs, generate_provider_oauth_nonce,
|
||||
};
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::GatewayError;
|
||||
use aether_contracts::ProxySnapshot;
|
||||
use aether_data::repository::provider_oauth::{
|
||||
StoredAdminProviderOAuthDeviceSession, KIRO_DEVICE_AUTH_SESSION_TTL_BUFFER_SECS,
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_oauth::core::OAuthError;
|
||||
use aether_oauth::provider::providers::{
|
||||
XaiDevicePollOutcome, XaiProviderOAuthAdapter, XAI_CLIENT_ID, XAI_DEVICE_CODE_URL,
|
||||
XAI_TOKEN_URL,
|
||||
};
|
||||
use aether_oauth::provider::ProviderOAuthTransportContext;
|
||||
use axum::{
|
||||
body::Body,
|
||||
http,
|
||||
response::{IntoResponse, Response},
|
||||
Json,
|
||||
};
|
||||
use serde_json::{json, Value};
|
||||
|
||||
pub(super) async fn handle_admin_provider_oauth_xai_device_authorize(
|
||||
state: &AdminAppState<'_>,
|
||||
provider_id: &str,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
principal: &GatewayAdminPrincipalContext,
|
||||
runtime_endpoint: Option<&StoredProviderCatalogEndpoint>,
|
||||
request_proxy: Option<ProxySnapshot>,
|
||||
proxy_node_id: Option<&str>,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
let device_url = state.provider_oauth_token_url("xai_device", XAI_DEVICE_CODE_URL);
|
||||
let token_url = state.provider_oauth_token_url("xai", XAI_TOKEN_URL);
|
||||
let adapter =
|
||||
XaiProviderOAuthAdapter::default().with_endpoint_overrides(&device_url, &token_url);
|
||||
let ctx = ProviderOAuthTransportContext {
|
||||
provider_id: provider_id.to_string(),
|
||||
provider_type: provider.provider_type.clone(),
|
||||
endpoint_id: runtime_endpoint.map(|endpoint| endpoint.id.clone()),
|
||||
key_id: None,
|
||||
auth_type: Some("oauth".to_string()),
|
||||
decrypted_api_key: None,
|
||||
decrypted_auth_config: None,
|
||||
provider_config: provider.config.clone(),
|
||||
endpoint_config: runtime_endpoint.and_then(|endpoint| endpoint.config.clone()),
|
||||
key_config: None,
|
||||
network: aether_oauth::network::OAuthNetworkContext::provider_operation(
|
||||
request_proxy.clone(),
|
||||
),
|
||||
};
|
||||
let executor = crate::oauth::GatewayOAuthHttpExecutor::new(*state);
|
||||
let authorization = match adapter.start_device_flow(&executor, &ctx).await {
|
||||
Ok(authorization) => authorization,
|
||||
Err(error) => {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
sanitize_xai_oauth_error(&error),
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
let now_unix_secs = current_unix_secs();
|
||||
let session_id = generate_provider_oauth_nonce();
|
||||
let session = StoredAdminProviderOAuthDeviceSession {
|
||||
session_id: session_id.clone(),
|
||||
provider_id: provider_id.to_string(),
|
||||
initiated_by_user_id: principal.user_id.clone(),
|
||||
initiated_by_session_id: principal.session_id.clone(),
|
||||
initiated_by_management_token_id: principal.management_token_id.clone(),
|
||||
region: String::new(),
|
||||
client_id: XAI_CLIENT_ID.to_string(),
|
||||
client_secret: String::new(),
|
||||
device_code: authorization.device_code.clone(),
|
||||
auth_type: Some("device".to_string()),
|
||||
social_provider: None,
|
||||
code_verifier: None,
|
||||
redirect_uri: Some(token_url),
|
||||
machine_id: None,
|
||||
interval: authorization.interval,
|
||||
expires_at_unix_secs: now_unix_secs.saturating_add(authorization.expires_in),
|
||||
status: "pending".to_string(),
|
||||
proxy_node_id: proxy_node_id
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
created_at_unix_ms: now_unix_secs,
|
||||
key_id: None,
|
||||
email: None,
|
||||
replaced: false,
|
||||
error_msg: None,
|
||||
};
|
||||
if let Err(response) = state
|
||||
.save_provider_oauth_device_session(
|
||||
&session_id,
|
||||
&session,
|
||||
authorization
|
||||
.expires_in
|
||||
.saturating_add(KIRO_DEVICE_AUTH_SESSION_TTL_BUFFER_SECS),
|
||||
)
|
||||
.await
|
||||
{
|
||||
return Ok(response);
|
||||
}
|
||||
|
||||
Ok(Json(json!({
|
||||
"session_id": session_id,
|
||||
"user_code": authorization.user_code,
|
||||
"verification_uri": authorization.verification_uri,
|
||||
"verification_uri_complete": authorization.verification_uri_complete,
|
||||
"expires_in": authorization.expires_in,
|
||||
"interval": authorization.interval,
|
||||
"auth_type": "device",
|
||||
}))
|
||||
.into_response())
|
||||
}
|
||||
|
||||
pub(super) async fn handle_admin_provider_oauth_xai_device_poll(
|
||||
state: &AdminAppState<'_>,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
endpoints: &[StoredProviderCatalogEndpoint],
|
||||
request_proxy: Option<ProxySnapshot>,
|
||||
session_id: &str,
|
||||
mut session: StoredAdminProviderOAuthDeviceSession,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
let token_url = session
|
||||
.redirect_uri
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
.unwrap_or_else(|| state.provider_oauth_token_url("xai", XAI_TOKEN_URL));
|
||||
let adapter =
|
||||
XaiProviderOAuthAdapter::default().with_endpoint_overrides(XAI_DEVICE_CODE_URL, token_url);
|
||||
let ctx = ProviderOAuthTransportContext {
|
||||
provider_id: provider.id.clone(),
|
||||
provider_type: provider.provider_type.clone(),
|
||||
endpoint_id: None,
|
||||
key_id: None,
|
||||
auth_type: Some("oauth".to_string()),
|
||||
decrypted_api_key: None,
|
||||
decrypted_auth_config: None,
|
||||
provider_config: provider.config.clone(),
|
||||
endpoint_config: None,
|
||||
key_config: None,
|
||||
network: aether_oauth::network::OAuthNetworkContext::provider_operation(
|
||||
request_proxy.clone(),
|
||||
),
|
||||
};
|
||||
let executor = crate::oauth::GatewayOAuthHttpExecutor::new(*state);
|
||||
let outcome = match adapter
|
||||
.poll_device_token(&executor, &ctx, &session.device_code)
|
||||
.await
|
||||
{
|
||||
Ok(outcome) => outcome,
|
||||
Err(error) => {
|
||||
return Ok(xai_device_poll_terminal_from_error(
|
||||
state,
|
||||
session_id,
|
||||
&mut session,
|
||||
&error,
|
||||
)
|
||||
.await);
|
||||
}
|
||||
};
|
||||
|
||||
match outcome {
|
||||
XaiDevicePollOutcome::Pending => {
|
||||
Ok(Json(json!({"status": "pending", "replaced": false})).into_response())
|
||||
}
|
||||
XaiDevicePollOutcome::SlowDown => {
|
||||
Ok(Json(json!({"status": "slow_down", "replaced": false})).into_response())
|
||||
}
|
||||
XaiDevicePollOutcome::Authorized(result) => {
|
||||
persist_xai_device_authorization(
|
||||
state,
|
||||
provider,
|
||||
endpoints,
|
||||
request_proxy,
|
||||
session_id,
|
||||
session,
|
||||
*result,
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn persist_xai_device_authorization(
|
||||
state: &AdminAppState<'_>,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
endpoints: &[StoredProviderCatalogEndpoint],
|
||||
request_proxy: Option<ProxySnapshot>,
|
||||
session_id: &str,
|
||||
mut session: StoredAdminProviderOAuthDeviceSession,
|
||||
result: aether_oauth::provider::ProviderOAuthTokenSet,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
let access_token = result.token_set.access_token.trim().to_string();
|
||||
if access_token.is_empty() {
|
||||
return Ok(Json(json!({
|
||||
"status": "error",
|
||||
"error": "xAI token 响应缺少 access_token",
|
||||
"replaced": false,
|
||||
}))
|
||||
.into_response());
|
||||
}
|
||||
let mut auth_config = result.auth_config.as_object().cloned().unwrap_or_default();
|
||||
auth_config.insert("provider_type".to_string(), json!("xai"));
|
||||
auth_config.insert("auth_method".to_string(), json!("oauth"));
|
||||
auth_config.insert("using_api".to_string(), json!(false));
|
||||
|
||||
let duplicate = match state
|
||||
.find_duplicate_provider_oauth_key(&provider.id, &auth_config, None)
|
||||
.await
|
||||
{
|
||||
Ok(duplicate) => duplicate,
|
||||
Err(detail) => {
|
||||
return Ok(Json(json!({
|
||||
"status": "error",
|
||||
"error": detail,
|
||||
"replaced": false,
|
||||
}))
|
||||
.into_response());
|
||||
}
|
||||
};
|
||||
|
||||
let api_formats = provider_oauth_active_api_formats(endpoints);
|
||||
let key_proxy = provider_oauth_key_proxy_value(session.proxy_node_id.as_deref());
|
||||
let expires_at = result.token_set.expires_at_unix_secs;
|
||||
let email = auth_config
|
||||
.get("email")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
let mut replaced = false;
|
||||
let persisted_key = if let Some(existing_key) = duplicate {
|
||||
replaced = true;
|
||||
match state
|
||||
.update_existing_provider_oauth_catalog_key(
|
||||
&existing_key,
|
||||
&provider.provider_type,
|
||||
&access_token,
|
||||
&auth_config,
|
||||
&api_formats,
|
||||
key_proxy.clone(),
|
||||
expires_at,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
Some(key) => key,
|
||||
None => {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::SERVICE_UNAVAILABLE,
|
||||
"provider oauth write unavailable",
|
||||
));
|
||||
}
|
||||
}
|
||||
} else {
|
||||
let key_name = admin_provider_oauth_key_name_from_auth_config(
|
||||
&provider.provider_type,
|
||||
&auth_config,
|
||||
None,
|
||||
);
|
||||
match state
|
||||
.create_provider_oauth_catalog_key(
|
||||
&provider.id,
|
||||
&provider.provider_type,
|
||||
&key_name,
|
||||
&access_token,
|
||||
&auth_config,
|
||||
&api_formats,
|
||||
key_proxy,
|
||||
expires_at,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
Some(key) => key,
|
||||
None => {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::SERVICE_UNAVAILABLE,
|
||||
"provider oauth write unavailable",
|
||||
));
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
spawn_provider_oauth_account_state_refresh_after_update(
|
||||
state.cloned_app(),
|
||||
provider.clone(),
|
||||
persisted_key.id.clone(),
|
||||
request_proxy.clone(),
|
||||
);
|
||||
|
||||
session.status = "authorized".to_string();
|
||||
session.key_id = Some(persisted_key.id.clone());
|
||||
session.email = email.clone();
|
||||
session.replaced = replaced;
|
||||
session.error_msg = None;
|
||||
let _ = state
|
||||
.save_provider_oauth_device_session(session_id, &session, 60)
|
||||
.await;
|
||||
|
||||
Ok(attach_admin_provider_oauth_device_poll_terminal_response(
|
||||
session_id,
|
||||
"authorized",
|
||||
Json(json!({
|
||||
"status": "authorized",
|
||||
"key_id": persisted_key.id,
|
||||
"email": email,
|
||||
"replaced": replaced,
|
||||
}))
|
||||
.into_response(),
|
||||
))
|
||||
}
|
||||
|
||||
async fn xai_device_poll_terminal_from_error(
|
||||
state: &AdminAppState<'_>,
|
||||
session_id: &str,
|
||||
session: &mut StoredAdminProviderOAuthDeviceSession,
|
||||
error: &OAuthError,
|
||||
) -> Response<Body> {
|
||||
let (status, message) = match error {
|
||||
OAuthError::InvalidRequest(detail) if detail.to_ascii_lowercase().contains("expired") => {
|
||||
("expired", "设备码已过期".to_string())
|
||||
}
|
||||
OAuthError::InvalidRequest(detail) if detail.to_ascii_lowercase().contains("denied") => {
|
||||
("error", "用户拒绝授权".to_string())
|
||||
}
|
||||
_ => ("error", sanitize_xai_oauth_error(error)),
|
||||
};
|
||||
session.status = status.to_string();
|
||||
session.error_msg = Some(message.clone());
|
||||
let _ = state
|
||||
.save_provider_oauth_device_session(session_id, session, 30)
|
||||
.await;
|
||||
attach_admin_provider_oauth_device_poll_terminal_response(
|
||||
session_id,
|
||||
status,
|
||||
Json(json!({
|
||||
"status": status,
|
||||
"error": message,
|
||||
"replaced": false,
|
||||
}))
|
||||
.into_response(),
|
||||
)
|
||||
}
|
||||
|
||||
fn sanitize_xai_oauth_error(error: &OAuthError) -> String {
|
||||
match error {
|
||||
OAuthError::InvalidRequest(_) => "xAI 设备授权失败: 请求参数无效".to_string(),
|
||||
OAuthError::HttpStatus { status_code, .. } => {
|
||||
format!("xAI 设备授权失败: HTTP {status_code}")
|
||||
}
|
||||
_ => "xAI 设备授权失败".to_string(),
|
||||
}
|
||||
}
|
||||
@@ -715,7 +715,7 @@ async fn resolve_admin_provider_oauth_single_import_tokens(
|
||||
if !provider_type_supports_access_token_import(provider_type) {
|
||||
return Err(build_internal_control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
"Access Token 导入仅支持 Claude Code / Codex / ChatGPT Web / Grok Provider",
|
||||
"Access Token 导入仅支持 Claude Code / Codex / ChatGPT Web / Grok / xAI Provider",
|
||||
));
|
||||
}
|
||||
|
||||
@@ -867,7 +867,7 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
|
||||
flatten_claude_code_credentials_payload(&mut raw_payload);
|
||||
}
|
||||
let refresh_token_input = import_payload_string(&raw_payload, "refresh_token", "refreshToken");
|
||||
let access_token_input = import_payload_string_any(
|
||||
let mut access_token_input = import_payload_string_any(
|
||||
&raw_payload,
|
||||
&[
|
||||
"access_token",
|
||||
@@ -879,6 +879,9 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
|
||||
],
|
||||
)
|
||||
.or_else(|| provider_oauth_import_authorization_bearer_token_from_object(&raw_payload));
|
||||
if provider_type == "xai" && access_token_input.is_none() {
|
||||
access_token_input = import_payload_string(&raw_payload, "api_key", "apiKey");
|
||||
}
|
||||
let imported_expires_at =
|
||||
import_payload_u64_any(&raw_payload, &["expires_at", "expiresAt", "expired"]);
|
||||
let (refresh_token_input, access_token_input) = normalize_provider_import_tokens(
|
||||
@@ -901,7 +904,11 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
|
||||
if !create_agent_identity && refresh_token_input.is_none() && access_token_input.is_none() {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
"Refresh Token、Access Token 或 sso_token 不能为空",
|
||||
if provider_type == "xai" {
|
||||
"Refresh Token、Access Token 或 api_key 不能为空"
|
||||
} else {
|
||||
"Refresh Token、Access Token 或 sso_token 不能为空"
|
||||
},
|
||||
));
|
||||
}
|
||||
if !is_fixed_provider_type_for_provider_oauth(&provider_type) {
|
||||
|
||||
@@ -70,6 +70,12 @@ pub(super) async fn handle_admin_provider_oauth_start_key(
|
||||
"Windsurf 请使用浏览器登录或导入凭据。",
|
||||
));
|
||||
}
|
||||
if provider_type == "xai" {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
"xAI 请使用设备授权或导入凭据。",
|
||||
));
|
||||
}
|
||||
let Some(template) = admin_provider_oauth_template(&provider_type) else {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
@@ -167,6 +173,12 @@ pub(super) async fn handle_admin_provider_oauth_start_provider(
|
||||
"Windsurf 请使用浏览器登录或导入凭据。",
|
||||
));
|
||||
}
|
||||
if provider_type == "xai" {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
"xAI 请使用设备授权或导入凭据。",
|
||||
));
|
||||
}
|
||||
let Some(template) = admin_provider_oauth_template(&provider_type) else {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
|
||||
@@ -121,6 +121,9 @@ pub(super) fn normalize_provider_import_tokens(
|
||||
if provider_type == "grok" {
|
||||
return (None, access_token.or(refresh_token));
|
||||
}
|
||||
if provider_type == "xai" {
|
||||
return (refresh_token, access_token);
|
||||
}
|
||||
if provider_type == "claude_code" {
|
||||
if access_token.is_none() && refresh_token.as_deref().is_some_and(is_claude_access_token) {
|
||||
return (None, refresh_token);
|
||||
@@ -237,7 +240,7 @@ pub(super) fn provider_oauth_import_authorization_bearer_token_from_object(
|
||||
pub(super) fn provider_type_supports_access_token_import(provider_type: &str) -> bool {
|
||||
matches!(
|
||||
provider_type.trim().to_ascii_lowercase().as_str(),
|
||||
"claude_code" | "codex" | "chatgpt_web" | "grok"
|
||||
"claude_code" | "codex" | "chatgpt_web" | "grok" | "xai"
|
||||
)
|
||||
}
|
||||
|
||||
@@ -331,6 +334,15 @@ pub(super) fn build_provider_access_token_import_auth_config(
|
||||
auth_config.insert("sso_token".to_string(), json!(access_token));
|
||||
auth_config.insert("auth_method".to_string(), json!("sso_token"));
|
||||
}
|
||||
if provider_type.trim().eq_ignore_ascii_case("xai") {
|
||||
if refresh_token.is_some() {
|
||||
auth_config.insert("auth_method".to_string(), json!("oauth"));
|
||||
auth_config.insert("using_api".to_string(), json!(false));
|
||||
} else {
|
||||
auth_config.insert("auth_method".to_string(), json!("api_key"));
|
||||
auth_config.insert("using_api".to_string(), json!(true));
|
||||
}
|
||||
}
|
||||
|
||||
auth_config.insert(
|
||||
"access_token_import_temporary".to_string(),
|
||||
@@ -532,6 +544,41 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_xai_import_keeps_refresh_token_separate_from_api_key() {
|
||||
let (refresh_token, access_token) =
|
||||
normalize_provider_import_tokens("xai", Some("xai-refresh-token"), None);
|
||||
assert_eq!(refresh_token.as_deref(), Some("xai-refresh-token"));
|
||||
assert!(access_token.is_none());
|
||||
|
||||
let (refresh_token, access_token) =
|
||||
normalize_provider_import_tokens("xai", None, Some("xai-api-key"));
|
||||
assert!(refresh_token.is_none());
|
||||
assert_eq!(access_token.as_deref(), Some("xai-api-key"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_xai_auth_config_from_api_key_and_oauth_tokens() {
|
||||
let (api_key_config, _) =
|
||||
build_provider_access_token_import_auth_config("xai", "xai-api-key", None, None, None);
|
||||
assert_eq!(api_key_config.get("auth_method"), Some(&json!("api_key")));
|
||||
assert_eq!(api_key_config.get("using_api"), Some(&json!(true)));
|
||||
|
||||
let (oauth_config, _) = build_provider_access_token_import_auth_config(
|
||||
"xai",
|
||||
"xai-access-token",
|
||||
Some("xai-refresh-token"),
|
||||
None,
|
||||
None,
|
||||
);
|
||||
assert_eq!(oauth_config.get("auth_method"), Some(&json!("oauth")));
|
||||
assert_eq!(oauth_config.get("using_api"), Some(&json!(false)));
|
||||
assert_eq!(
|
||||
oauth_config.get("refresh_token"),
|
||||
Some(&json!("xai-refresh-token"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn flattens_only_claude_ai_oauth_credentials_and_converts_expiry_ms() {
|
||||
let mut payload = json!({
|
||||
|
||||
@@ -0,0 +1,221 @@
|
||||
use super::shared::{
|
||||
build_provider_quota_execution_plan, build_quota_snapshot_payload,
|
||||
default_provider_quota_execution_timeouts, execute_provider_quota_plan,
|
||||
extract_execution_error_message, oauth_refresh_auto_removed_result,
|
||||
persist_provider_quota_refresh_state, quota_key_auto_removed,
|
||||
quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
|
||||
};
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
|
||||
use crate::GatewayError;
|
||||
use aether_admin::provider::quota::parse_claude_code_oauth_usage_response;
|
||||
use aether_admin::provider::redaction::admin_provider_metadata_bucket_safe_json;
|
||||
use aether_contracts::ProxySnapshot;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_provider_pool::build_claude_code_pool_quota_request;
|
||||
use serde_json::json;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
async fn execute_claude_code_quota_plan(
|
||||
state: &AdminAppState<'_>,
|
||||
transport: &AdminGatewayProviderTransportSnapshot,
|
||||
authorization: (String, String),
|
||||
proxy_override: Option<&ProxySnapshot>,
|
||||
) -> Result<ProviderQuotaExecutionOutcome, GatewayError> {
|
||||
let proxy = match proxy_override {
|
||||
Some(proxy) => Some(proxy.clone()),
|
||||
None => {
|
||||
state
|
||||
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
|
||||
.await
|
||||
}
|
||||
};
|
||||
let timeouts = state
|
||||
.resolve_transport_execution_timeouts(transport)
|
||||
.or(Some(default_provider_quota_execution_timeouts(
|
||||
proxy.as_ref(),
|
||||
)));
|
||||
let spec = build_claude_code_pool_quota_request(&transport.key.id, authorization);
|
||||
let plan = build_provider_quota_execution_plan(
|
||||
transport,
|
||||
spec,
|
||||
proxy,
|
||||
state.resolve_transport_profile(transport),
|
||||
timeouts,
|
||||
);
|
||||
|
||||
execute_provider_quota_plan(state, transport, plan, "claude_code").await
|
||||
}
|
||||
|
||||
pub(crate) async fn refresh_claude_code_provider_quota_locally(
|
||||
state: &AdminAppState<'_>,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
endpoint: &StoredProviderCatalogEndpoint,
|
||||
keys: Vec<StoredProviderCatalogKey>,
|
||||
proxy_override: Option<ProxySnapshot>,
|
||||
) -> Result<Option<serde_json::Value>, GatewayError> {
|
||||
let mut results = Vec::new();
|
||||
let mut success_count = 0usize;
|
||||
let mut failed_count = 0usize;
|
||||
let mut auto_removed_count = 0usize;
|
||||
|
||||
for key in keys {
|
||||
let transport = match state
|
||||
.read_provider_transport_snapshot(&provider.id, &endpoint.id, &key.id)
|
||||
.await?
|
||||
{
|
||||
Some(transport) => transport,
|
||||
None => {
|
||||
failed_count += 1;
|
||||
results.push(json!({
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": "Provider transport snapshot unavailable",
|
||||
}));
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let authorization = match state.resolve_local_oauth_header_auth(&transport).await? {
|
||||
Some(auth) => auth,
|
||||
_ => {
|
||||
if quota_key_auto_removed(state, &key.id).await? {
|
||||
auto_removed_count += 1;
|
||||
results.push(oauth_refresh_auto_removed_result(&key));
|
||||
continue;
|
||||
}
|
||||
failed_count += 1;
|
||||
results.push(json!({
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": "缺少 OAuth 认证信息,请先授权/刷新 Token",
|
||||
}));
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let result = match execute_claude_code_quota_plan(
|
||||
state,
|
||||
&transport,
|
||||
authorization,
|
||||
proxy_override.as_ref(),
|
||||
)
|
||||
.await?
|
||||
{
|
||||
ProviderQuotaExecutionOutcome::Response(result) => result,
|
||||
ProviderQuotaExecutionOutcome::Failure(_) => {
|
||||
failed_count += 1;
|
||||
results.push(json!({
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": "oauth/usage 请求执行失败",
|
||||
"status_code": 502,
|
||||
}));
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let now_unix_secs = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
.map(|duration| duration.as_secs())
|
||||
.unwrap_or(0);
|
||||
let mut metadata_update = None::<serde_json::Value>;
|
||||
let (oauth_invalid_at_unix_secs, oauth_invalid_reason) =
|
||||
quota_refresh_success_invalid_state(&key);
|
||||
let mut status = "error".to_string();
|
||||
let mut message = None::<String>;
|
||||
|
||||
if result.status_code == 200 {
|
||||
if let Some(body_json) = result
|
||||
.body
|
||||
.as_ref()
|
||||
.and_then(|body| body.json_body.as_ref())
|
||||
{
|
||||
metadata_update = parse_claude_code_oauth_usage_response(body_json, now_unix_secs)
|
||||
.map(|metadata| json!({ "claude_code": metadata }));
|
||||
if metadata_update.is_some() {
|
||||
status = "success".to_string();
|
||||
} else {
|
||||
status = "no_metadata".to_string();
|
||||
message = Some("响应中未包含额度窗口".to_string());
|
||||
}
|
||||
} else {
|
||||
status = "no_metadata".to_string();
|
||||
message = Some("响应中未包含配额信息".to_string());
|
||||
}
|
||||
} else {
|
||||
message = Some(match result.status_code {
|
||||
401 => "oauth/usage 返回 401,Token 可能已失效,请刷新 Token".to_string(),
|
||||
403 => "oauth/usage 返回 403,该账号缺少 user:profile 权限(如 Setup Token),无法查询额度"
|
||||
.to_string(),
|
||||
429 => "oauth/usage 被限流,请稍后重试".to_string(),
|
||||
code => format!("oauth/usage 返回状态码 {code}"),
|
||||
});
|
||||
}
|
||||
|
||||
if !persist_provider_quota_refresh_state(
|
||||
state,
|
||||
&key.id,
|
||||
metadata_update.as_ref(),
|
||||
oauth_invalid_at_unix_secs,
|
||||
oauth_invalid_reason,
|
||||
None,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
failed_count += 1;
|
||||
results.push(json!({
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": "Key 状态写入失败",
|
||||
}));
|
||||
continue;
|
||||
}
|
||||
|
||||
if status == "success" {
|
||||
success_count += 1;
|
||||
} else {
|
||||
failed_count += 1;
|
||||
}
|
||||
|
||||
let mut payload = serde_json::Map::new();
|
||||
payload.insert("key_id".to_string(), json!(key.id));
|
||||
payload.insert("key_name".to_string(), json!(key.name));
|
||||
payload.insert("status".to_string(), json!(status));
|
||||
if let Some(message) = message {
|
||||
payload.insert("message".to_string(), json!(message));
|
||||
}
|
||||
if let Some(metadata) = metadata_update
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("claude_code"))
|
||||
{
|
||||
payload.insert(
|
||||
"metadata".to_string(),
|
||||
admin_provider_metadata_bucket_safe_json("claude_code", Some(metadata)),
|
||||
);
|
||||
}
|
||||
if let Some(quota_snapshot) = build_quota_snapshot_payload(
|
||||
"claude_code",
|
||||
key.status_snapshot.as_ref(),
|
||||
metadata_update.as_ref(),
|
||||
) {
|
||||
payload.insert("quota_snapshot".to_string(), quota_snapshot);
|
||||
}
|
||||
results.push(serde_json::Value::Object(payload));
|
||||
}
|
||||
|
||||
Ok(Some(json!({
|
||||
"success": success_count,
|
||||
"failed": failed_count,
|
||||
"total": results.len(),
|
||||
"results": results,
|
||||
"message": format!("已处理 {} 个 Key", results.len()),
|
||||
"auto_removed": auto_removed_count,
|
||||
})))
|
||||
}
|
||||
@@ -3,11 +3,13 @@ use std::pin::Pin;
|
||||
|
||||
use super::antigravity::refresh_antigravity_provider_quota_locally;
|
||||
use super::chatgpt_web::refresh_chatgpt_web_provider_quota_locally;
|
||||
use super::claude_code::refresh_claude_code_provider_quota_locally;
|
||||
use super::codex::refresh_codex_provider_quota_locally;
|
||||
use super::gemini_cli::refresh_gemini_cli_provider_quota_locally;
|
||||
use super::grok::refresh_grok_provider_quota_locally;
|
||||
use super::kiro::refresh_kiro_provider_quota_locally;
|
||||
use super::windsurf::refresh_windsurf_provider_quota_locally;
|
||||
use super::xai::refresh_xai_provider_quota_locally;
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::GatewayError;
|
||||
use aether_contracts::ProxySnapshot;
|
||||
@@ -35,6 +37,10 @@ const PROVIDER_QUOTA_REFRESH_HANDLERS: &[(&str, ProviderQuotaRefreshHandler)] =
|
||||
"chatgpt_web",
|
||||
refresh_chatgpt_web_provider_quota_locally_boxed,
|
||||
),
|
||||
(
|
||||
"claude_code",
|
||||
refresh_claude_code_provider_quota_locally_boxed,
|
||||
),
|
||||
("codex", refresh_codex_provider_quota_locally_boxed),
|
||||
(
|
||||
"gemini_cli",
|
||||
@@ -43,6 +49,7 @@ const PROVIDER_QUOTA_REFRESH_HANDLERS: &[(&str, ProviderQuotaRefreshHandler)] =
|
||||
("grok", refresh_grok_provider_quota_locally_boxed),
|
||||
("kiro", refresh_kiro_provider_quota_locally_boxed),
|
||||
("windsurf", refresh_windsurf_provider_quota_locally_boxed),
|
||||
("xai", refresh_xai_provider_quota_locally_boxed),
|
||||
];
|
||||
|
||||
pub(crate) async fn refresh_provider_pool_quota_locally(
|
||||
@@ -111,6 +118,22 @@ fn refresh_codex_provider_quota_locally_boxed<'a>(
|
||||
))
|
||||
}
|
||||
|
||||
fn refresh_claude_code_provider_quota_locally_boxed<'a>(
|
||||
state: &'a AdminAppState<'a>,
|
||||
provider: &'a StoredProviderCatalogProvider,
|
||||
endpoint: &'a StoredProviderCatalogEndpoint,
|
||||
keys: Vec<StoredProviderCatalogKey>,
|
||||
proxy_override: Option<ProxySnapshot>,
|
||||
) -> ProviderQuotaRefreshFuture<'a> {
|
||||
Box::pin(refresh_claude_code_provider_quota_locally(
|
||||
state,
|
||||
provider,
|
||||
endpoint,
|
||||
keys,
|
||||
proxy_override,
|
||||
))
|
||||
}
|
||||
|
||||
fn refresh_gemini_cli_provider_quota_locally_boxed<'a>(
|
||||
state: &'a AdminAppState<'a>,
|
||||
provider: &'a StoredProviderCatalogProvider,
|
||||
@@ -174,3 +197,19 @@ fn refresh_windsurf_provider_quota_locally_boxed<'a>(
|
||||
proxy_override,
|
||||
))
|
||||
}
|
||||
|
||||
fn refresh_xai_provider_quota_locally_boxed<'a>(
|
||||
state: &'a AdminAppState<'a>,
|
||||
provider: &'a StoredProviderCatalogProvider,
|
||||
endpoint: &'a StoredProviderCatalogEndpoint,
|
||||
keys: Vec<StoredProviderCatalogKey>,
|
||||
proxy_override: Option<ProxySnapshot>,
|
||||
) -> ProviderQuotaRefreshFuture<'a> {
|
||||
Box::pin(refresh_xai_provider_quota_locally(
|
||||
state,
|
||||
provider,
|
||||
endpoint,
|
||||
keys,
|
||||
proxy_override,
|
||||
))
|
||||
}
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
pub(crate) mod antigravity;
|
||||
pub(crate) mod chatgpt_web;
|
||||
pub(crate) mod claude_code;
|
||||
pub(crate) mod codex;
|
||||
pub(crate) mod dispatch;
|
||||
pub(crate) mod gemini_cli;
|
||||
@@ -7,3 +8,4 @@ pub(crate) mod grok;
|
||||
pub(crate) mod kiro;
|
||||
pub(crate) mod shared;
|
||||
pub(crate) mod windsurf;
|
||||
pub(crate) mod xai;
|
||||
|
||||
@@ -1713,8 +1713,10 @@ fn provider_quota_url_has_allowed_origin(provider_name: &str, value: &str) -> bo
|
||||
| "daily-cloudcode-pa.sandbox.googleapis.com"
|
||||
),
|
||||
"gemini_cli" => host == "cloudcode-pa.googleapis.com",
|
||||
"claude_code" => host == "api.anthropic.com",
|
||||
"chatgpt_web" | "codex" => host == "chatgpt.com",
|
||||
"grok" => host == "grok.com",
|
||||
"xai" => host == "cli-chat-proxy.grok.com",
|
||||
"windsurf" => host == "server.codeium.com",
|
||||
"kiro" => kiro_quota_host_is_allowed(host),
|
||||
_ => false,
|
||||
@@ -1814,6 +1816,14 @@ mod tests {
|
||||
),
|
||||
("codex", "https://chatgpt.com/backend-api/wham/usage"),
|
||||
("grok", "https://grok.com/rest/rate-limits"),
|
||||
(
|
||||
"xai",
|
||||
"https://cli-chat-proxy.grok.com/v1/billing?format=credits",
|
||||
),
|
||||
(
|
||||
"xai",
|
||||
"https://cli-chat-proxy.grok.com/v1/user",
|
||||
),
|
||||
(
|
||||
"windsurf",
|
||||
"https://server.codeium.com/exa.seat_management_pb.SeatManagementService/GetUserStatus",
|
||||
@@ -1847,6 +1857,11 @@ mod tests {
|
||||
"https://chatgpt.com.attacker.test/backend-api/wham/usage",
|
||||
),
|
||||
("grok", "https://grok.com.attacker.test/rest/rate-limits"),
|
||||
(
|
||||
"xai",
|
||||
"https://cli-chat-proxy.grok.com.attacker.test/v1/billing",
|
||||
),
|
||||
("xai", "https://api.x.ai/v1/billing?format=credits"),
|
||||
("windsurf", "https://server.codeium.com.attacker.test/quota"),
|
||||
(
|
||||
"gemini_cli",
|
||||
|
||||
@@ -0,0 +1,313 @@
|
||||
use super::shared::{
|
||||
build_provider_quota_execution_plan, build_quota_snapshot_payload,
|
||||
default_provider_quota_execution_timeouts, execute_provider_quota_plan,
|
||||
extract_execution_error_message, oauth_refresh_auto_removed_result,
|
||||
persist_provider_quota_refresh_state, quota_key_auto_removed,
|
||||
quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
|
||||
};
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
|
||||
use crate::GatewayError;
|
||||
use aether_admin::provider::quota::parse_xai_billing_response;
|
||||
use aether_admin::provider::redaction::admin_provider_metadata_bucket_safe_json;
|
||||
use aether_contracts::ProxySnapshot;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_provider_pool::{build_xai_pool_billing_request, build_xai_pool_user_request};
|
||||
use aether_provider_transport::xai::{
|
||||
extract_xai_user_id_from_auth_config, extract_xai_user_id_from_value, xai_auth_uses_api,
|
||||
};
|
||||
use serde_json::{json, Value};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
async fn execute_xai_quota_plan(
|
||||
state: &AdminAppState<'_>,
|
||||
transport: &AdminGatewayProviderTransportSnapshot,
|
||||
spec: aether_provider_pool::ProviderPoolQuotaRequestSpec,
|
||||
proxy_override: Option<&ProxySnapshot>,
|
||||
) -> Result<ProviderQuotaExecutionOutcome, GatewayError> {
|
||||
let proxy = match proxy_override {
|
||||
Some(proxy) => Some(proxy.clone()),
|
||||
None => {
|
||||
state
|
||||
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
|
||||
.await
|
||||
}
|
||||
};
|
||||
let timeouts = state
|
||||
.resolve_transport_execution_timeouts(transport)
|
||||
.or(Some(default_provider_quota_execution_timeouts(
|
||||
proxy.as_ref(),
|
||||
)));
|
||||
let plan = build_provider_quota_execution_plan(
|
||||
transport,
|
||||
spec,
|
||||
proxy,
|
||||
state.resolve_transport_profile(transport),
|
||||
timeouts,
|
||||
);
|
||||
|
||||
execute_provider_quota_plan(state, transport, plan, "xai").await
|
||||
}
|
||||
|
||||
fn xai_authorization_from_header(authorization: &(String, String)) -> (String, String) {
|
||||
authorization.clone()
|
||||
}
|
||||
|
||||
fn enrich_xai_subscription_title(mut metadata: Value, auth_config: Option<&str>) -> Value {
|
||||
if metadata
|
||||
.get("subscription_title")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty())
|
||||
{
|
||||
return metadata;
|
||||
}
|
||||
let Some(config) = auth_config
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.and_then(|value| serde_json::from_str::<Value>(value).ok())
|
||||
else {
|
||||
return metadata;
|
||||
};
|
||||
let title = ["subscription_tier", "subscriptionTier", "tier", "plan"]
|
||||
.iter()
|
||||
.find_map(|field| {
|
||||
config
|
||||
.get(*field)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
});
|
||||
if let Some(title) = title {
|
||||
if let Some(object) = metadata.as_object_mut() {
|
||||
object.insert("subscription_title".to_string(), json!(title));
|
||||
}
|
||||
}
|
||||
metadata
|
||||
}
|
||||
|
||||
pub(crate) async fn refresh_xai_provider_quota_locally(
|
||||
state: &AdminAppState<'_>,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
endpoint: &StoredProviderCatalogEndpoint,
|
||||
keys: Vec<StoredProviderCatalogKey>,
|
||||
proxy_override: Option<ProxySnapshot>,
|
||||
) -> Result<Option<serde_json::Value>, GatewayError> {
|
||||
let mut results = Vec::new();
|
||||
let mut success_count = 0usize;
|
||||
let mut failed_count = 0usize;
|
||||
let mut auto_removed_count = 0usize;
|
||||
|
||||
for key in keys {
|
||||
let transport = match state
|
||||
.read_provider_transport_snapshot(&provider.id, &endpoint.id, &key.id)
|
||||
.await?
|
||||
{
|
||||
Some(transport) => transport,
|
||||
None => {
|
||||
failed_count += 1;
|
||||
results.push(json!({
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": "Provider transport snapshot unavailable",
|
||||
}));
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
if xai_auth_uses_api(
|
||||
transport.key.auth_type.as_str(),
|
||||
transport.key.decrypted_auth_config.as_deref(),
|
||||
) {
|
||||
results.push(json!({
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "skipped",
|
||||
"message": "xAI API Key 账号没有 Grok Build 订阅额度接口,请使用设备授权账号查询额度。",
|
||||
}));
|
||||
continue;
|
||||
}
|
||||
|
||||
let authorization = match state.resolve_local_oauth_header_auth(&transport).await? {
|
||||
Some(auth) => auth,
|
||||
_ => {
|
||||
if quota_key_auto_removed(state, &key.id).await? {
|
||||
auto_removed_count += 1;
|
||||
results.push(oauth_refresh_auto_removed_result(&key));
|
||||
continue;
|
||||
}
|
||||
failed_count += 1;
|
||||
results.push(json!({
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": "缺少 OAuth 认证信息,请先授权/刷新 Token",
|
||||
}));
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let fallback_user_id =
|
||||
extract_xai_user_id_from_auth_config(transport.key.decrypted_auth_config.as_deref());
|
||||
let user_id = match execute_xai_quota_plan(
|
||||
state,
|
||||
&transport,
|
||||
build_xai_pool_user_request(
|
||||
&transport.key.id,
|
||||
xai_authorization_from_header(&authorization),
|
||||
),
|
||||
proxy_override.as_ref(),
|
||||
)
|
||||
.await?
|
||||
{
|
||||
ProviderQuotaExecutionOutcome::Response(result) if result.status_code == 200 => result
|
||||
.body
|
||||
.as_ref()
|
||||
.and_then(|body| body.json_body.as_ref())
|
||||
.and_then(extract_xai_user_id_from_value)
|
||||
.or(fallback_user_id),
|
||||
_ => fallback_user_id,
|
||||
};
|
||||
|
||||
let result = match execute_xai_quota_plan(
|
||||
state,
|
||||
&transport,
|
||||
build_xai_pool_billing_request(
|
||||
&transport.key.id,
|
||||
xai_authorization_from_header(&authorization),
|
||||
user_id.as_deref(),
|
||||
),
|
||||
proxy_override.as_ref(),
|
||||
)
|
||||
.await?
|
||||
{
|
||||
ProviderQuotaExecutionOutcome::Response(result) => result,
|
||||
ProviderQuotaExecutionOutcome::Failure(_) => {
|
||||
failed_count += 1;
|
||||
results.push(json!({
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": "xAI billing 请求执行失败",
|
||||
"status_code": 502,
|
||||
}));
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let now_unix_secs = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
.map(|duration| duration.as_secs())
|
||||
.unwrap_or(0);
|
||||
let mut metadata_update = None::<serde_json::Value>;
|
||||
let (mut oauth_invalid_at_unix_secs, mut oauth_invalid_reason) =
|
||||
quota_refresh_success_invalid_state(&key);
|
||||
let mut status = "error".to_string();
|
||||
let mut message = None::<String>;
|
||||
|
||||
if result.status_code == 200 {
|
||||
if let Some(body_json) = result
|
||||
.body
|
||||
.as_ref()
|
||||
.and_then(|body| body.json_body.as_ref())
|
||||
{
|
||||
metadata_update =
|
||||
parse_xai_billing_response(body_json, now_unix_secs).map(|metadata| {
|
||||
json!({
|
||||
"xai": enrich_xai_subscription_title(
|
||||
metadata,
|
||||
transport.key.decrypted_auth_config.as_deref(),
|
||||
)
|
||||
})
|
||||
});
|
||||
if metadata_update.is_some() {
|
||||
status = "success".to_string();
|
||||
} else {
|
||||
status = "no_metadata".to_string();
|
||||
message = Some("响应中未包含可用的 Grok Build 额度信息".to_string());
|
||||
}
|
||||
} else {
|
||||
status = "no_metadata".to_string();
|
||||
message = Some("响应中未包含配额信息".to_string());
|
||||
}
|
||||
} else {
|
||||
message = Some(
|
||||
extract_execution_error_message(&result)
|
||||
.unwrap_or_else(|| format!("xAI billing 返回状态码 {}", result.status_code)),
|
||||
);
|
||||
if result.status_code == 401 || result.status_code == 403 {
|
||||
let reason = message
|
||||
.clone()
|
||||
.unwrap_or_else(|| "账户访问被禁止".to_string());
|
||||
oauth_invalid_at_unix_secs = Some(now_unix_secs);
|
||||
oauth_invalid_reason = Some(format!("账户访问被禁止: {reason}"));
|
||||
status = if result.status_code == 401 {
|
||||
"unauthorized".to_string()
|
||||
} else {
|
||||
"forbidden".to_string()
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
if !persist_provider_quota_refresh_state(
|
||||
state,
|
||||
&key.id,
|
||||
metadata_update.as_ref(),
|
||||
oauth_invalid_at_unix_secs,
|
||||
oauth_invalid_reason,
|
||||
None,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
failed_count += 1;
|
||||
results.push(json!({
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": "Key 状态写入失败",
|
||||
}));
|
||||
continue;
|
||||
}
|
||||
|
||||
if status == "success" {
|
||||
success_count += 1;
|
||||
} else {
|
||||
failed_count += 1;
|
||||
}
|
||||
|
||||
let mut payload = serde_json::Map::new();
|
||||
payload.insert("key_id".to_string(), json!(key.id));
|
||||
payload.insert("key_name".to_string(), json!(key.name));
|
||||
payload.insert("status".to_string(), json!(status));
|
||||
if let Some(message) = message {
|
||||
payload.insert("message".to_string(), json!(message));
|
||||
}
|
||||
if let Some(metadata) = metadata_update.as_ref().and_then(|value| value.get("xai")) {
|
||||
payload.insert(
|
||||
"metadata".to_string(),
|
||||
admin_provider_metadata_bucket_safe_json("xai", Some(metadata)),
|
||||
);
|
||||
}
|
||||
if let Some(quota_snapshot) = build_quota_snapshot_payload(
|
||||
"xai",
|
||||
key.status_snapshot.as_ref(),
|
||||
metadata_update.as_ref(),
|
||||
) {
|
||||
payload.insert("quota_snapshot".to_string(), quota_snapshot);
|
||||
}
|
||||
results.push(serde_json::Value::Object(payload));
|
||||
}
|
||||
|
||||
Ok(Some(json!({
|
||||
"success": success_count,
|
||||
"failed": failed_count,
|
||||
"total": results.len(),
|
||||
"results": results,
|
||||
"message": format!("已处理 {} 个 Key", results.len()),
|
||||
"auto_removed": auto_removed_count,
|
||||
})))
|
||||
}
|
||||
@@ -411,6 +411,7 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
|
||||
unschedulable_rules: Vec::new(),
|
||||
lru_enabled: false,
|
||||
skip_exhausted_accounts: false,
|
||||
reserve_minimum_quota: false,
|
||||
sticky_session_ttl_seconds: 3600,
|
||||
latency_window_seconds: 3600,
|
||||
latency_sample_limit: 50,
|
||||
@@ -446,6 +447,10 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
|
||||
.get("skip_exhausted_accounts")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false),
|
||||
reserve_minimum_quota: pool_advanced
|
||||
.get("reserve_minimum_quota")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false),
|
||||
sticky_session_ttl_seconds: pool_advanced
|
||||
.get("sticky_session_ttl_seconds")
|
||||
.and_then(json_u64)
|
||||
@@ -574,6 +579,22 @@ mod tests {
|
||||
let config = admin_provider_pool_config(&provider).expect("pool config should exist");
|
||||
|
||||
assert!(!config.skip_exhausted_accounts);
|
||||
assert!(!config.reserve_minimum_quota);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_reserve_minimum_quota_independently_of_skip_exhausted_accounts() {
|
||||
for enabled in [false, true] {
|
||||
let provider = sample_provider(json!({
|
||||
"pool_advanced": {
|
||||
"reserve_minimum_quota": enabled,
|
||||
"skip_exhausted_accounts": false
|
||||
}
|
||||
}));
|
||||
let config = admin_provider_pool_config(&provider).expect("pool config should exist");
|
||||
assert_eq!(config.reserve_minimum_quota, enabled);
|
||||
assert!(!config.skip_exhausted_accounts);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -651,6 +651,7 @@ mod tests {
|
||||
unschedulable_rules: Vec::new(),
|
||||
lru_enabled: true,
|
||||
skip_exhausted_accounts: false,
|
||||
reserve_minimum_quota: false,
|
||||
sticky_session_ttl_seconds: 120,
|
||||
latency_window_seconds: 600,
|
||||
latency_sample_limit: 10,
|
||||
|
||||
@@ -932,6 +932,13 @@ fn admin_pool_build_account_quota(
|
||||
return Some(account_quota);
|
||||
}
|
||||
}
|
||||
"xai" => {
|
||||
if let Some(account_quota) =
|
||||
admin_pool_build_kiro_account_quota_from_snapshot(quota_snapshot)
|
||||
{
|
||||
return Some(account_quota);
|
||||
}
|
||||
}
|
||||
"chatgpt_web" => {
|
||||
if let Some(account_quota) =
|
||||
admin_pool_build_chatgpt_web_account_quota_from_snapshot(quota_snapshot)
|
||||
@@ -1110,10 +1117,16 @@ pub(super) fn build_admin_pool_key_payload(
|
||||
let health_score = admin_pool_health_score(key);
|
||||
let circuit_breaker_open = false;
|
||||
let auth_semantics = provider_key_auth_semantics(key, provider_type);
|
||||
let account_quota_exhausted = pool_config
|
||||
.as_ref()
|
||||
.is_some_and(|config| config.skip_exhausted_accounts)
|
||||
&& admin_provider_pool_pure::admin_pool_key_account_quota_exhausted(key, provider_type);
|
||||
let account_quota_exhausted = pool_config.as_ref().is_some_and(|config| {
|
||||
(config.skip_exhausted_accounts
|
||||
&& admin_provider_pool_pure::admin_pool_key_account_quota_exhausted(key, provider_type))
|
||||
|| (config.reserve_minimum_quota
|
||||
&& admin_provider_pool_pure::admin_pool_key_minimum_quota_reached(
|
||||
key,
|
||||
provider_type,
|
||||
None,
|
||||
))
|
||||
});
|
||||
let auth_config = state.parse_catalog_auth_config_json(key);
|
||||
let oauth_expires_at =
|
||||
admin_pool_derive_oauth_expires_at(provider_type, key, auth_config.as_ref());
|
||||
@@ -1591,4 +1604,29 @@ mod tests {
|
||||
Some("Auto剩余 40.0% (60/150) | Heavy剩余 0.0% (0/20)".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn xai_account_quota_is_rendered_as_remaining_percent() {
|
||||
let quota_snapshot = json!({
|
||||
"provider_type": "xai",
|
||||
"code": "ok",
|
||||
"exhausted": false,
|
||||
"plan_type": "SuperGrok",
|
||||
"windows": [
|
||||
{
|
||||
"code": "usage",
|
||||
"label": "周额度",
|
||||
"scope": "account",
|
||||
"used_ratio": 0.46,
|
||||
"remaining_ratio": 0.54
|
||||
}
|
||||
]
|
||||
});
|
||||
let quota_snapshot = quota_snapshot.as_object().unwrap();
|
||||
|
||||
assert_eq!(
|
||||
admin_pool_build_account_quota("xai", Some(quota_snapshot)),
|
||||
Some("剩余 54.0%".to_string())
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -488,9 +488,16 @@ pub(super) fn admin_pool_key_visible_status_filter(
|
||||
) {
|
||||
return status;
|
||||
}
|
||||
if pool_config.is_some_and(|config| config.skip_exhausted_accounts)
|
||||
&& admin_provider_pool_pure::admin_pool_key_account_quota_exhausted(key, provider_type)
|
||||
{
|
||||
if pool_config.is_some_and(|config| {
|
||||
(config.skip_exhausted_accounts
|
||||
&& admin_provider_pool_pure::admin_pool_key_account_quota_exhausted(key, provider_type))
|
||||
|| (config.reserve_minimum_quota
|
||||
&& admin_provider_pool_pure::admin_pool_key_minimum_quota_reached(
|
||||
key,
|
||||
provider_type,
|
||||
None,
|
||||
))
|
||||
}) {
|
||||
return "quota_exhausted";
|
||||
}
|
||||
if !key.is_active {
|
||||
|
||||
@@ -594,11 +594,12 @@ async fn provider_query_fetch_models_for_key(
|
||||
});
|
||||
}
|
||||
|
||||
let dynamic_client_version = crate::ai_serving::api::codex_client_version();
|
||||
let client_version = is_codex.then(|| {
|
||||
codex_catalog
|
||||
.as_ref()
|
||||
.map(|catalog| catalog.client_version.as_str())
|
||||
.unwrap_or(crate::ai_serving::CODEX_CLIENT_VERSION)
|
||||
.unwrap_or(dynamic_client_version.as_str())
|
||||
});
|
||||
let outcome =
|
||||
match fetch_models_from_transports_for_management(state.app(), &transports, client_version)
|
||||
|
||||
@@ -1340,6 +1340,7 @@ async fn provider_query_apply_pool_scheduler_to_test_candidates(
|
||||
provider.id.clone(),
|
||||
provider_query_ai_pool_runtime_state(&runtime),
|
||||
);
|
||||
let reserve_minimum_quota = pool_config.reserve_minimum_quota;
|
||||
let pool_config =
|
||||
provider_query_ai_pool_scheduling_config(pool_config, provider.provider_type.as_str());
|
||||
let inputs = keys
|
||||
@@ -1351,6 +1352,14 @@ async fn provider_query_apply_pool_scheduler_to_test_candidates(
|
||||
effective_model: effective_model.to_string(),
|
||||
scheduler_skip_reason: None,
|
||||
};
|
||||
let mut key_context =
|
||||
provider_query_pool_catalog_key_context(state, &key, &provider.provider_type);
|
||||
key_context.quota_exhausted |= reserve_minimum_quota
|
||||
&& admin_provider_pool_pure::admin_pool_key_minimum_quota_reached(
|
||||
&key,
|
||||
&provider.provider_type,
|
||||
Some(effective_model),
|
||||
);
|
||||
AiPoolCandidateInput {
|
||||
facts: AiPoolCandidateFacts {
|
||||
provider_id: provider.id.clone(),
|
||||
@@ -1362,11 +1371,7 @@ async fn provider_query_apply_pool_scheduler_to_test_candidates(
|
||||
key_internal_priority: key.internal_priority,
|
||||
},
|
||||
pool_config: Some(pool_config.clone()),
|
||||
key_context: provider_query_pool_catalog_key_context(
|
||||
state,
|
||||
&key,
|
||||
&provider.provider_type,
|
||||
),
|
||||
key_context,
|
||||
candidate,
|
||||
}
|
||||
})
|
||||
@@ -3511,6 +3516,11 @@ async fn provider_query_execute_standard_test_candidate(
|
||||
codex_model_capabilities.as_ref(),
|
||||
);
|
||||
}
|
||||
crate::provider_transport::insert_cli_identity_headers_if_needed(
|
||||
&transport,
|
||||
provider_api_format,
|
||||
&mut request_headers,
|
||||
);
|
||||
if !uses_vertex_query_auth {
|
||||
if let (Some(auth_header), Some(auth_value)) =
|
||||
(auth_header.as_deref(), auth_value.as_deref())
|
||||
|
||||
@@ -65,6 +65,7 @@ pub(crate) struct AdminProviderPoolConfig {
|
||||
pub(crate) unschedulable_rules: Vec<AdminProviderPoolUnschedulableRule>,
|
||||
pub(crate) lru_enabled: bool,
|
||||
pub(crate) skip_exhausted_accounts: bool,
|
||||
pub(crate) reserve_minimum_quota: bool,
|
||||
pub(crate) sticky_session_ttl_seconds: u64,
|
||||
pub(crate) latency_window_seconds: u64,
|
||||
pub(crate) latency_sample_limit: u64,
|
||||
|
||||
@@ -4,9 +4,9 @@ pub(crate) fn normalize_provider_type_input(value: &str) -> Result<String, Strin
|
||||
let normalized = value.trim().to_ascii_lowercase();
|
||||
match normalized.as_str() {
|
||||
"custom" | "claude_code" | "kiro" | "codex" | "chatgpt_web" | "gemini_cli"
|
||||
| "antigravity" | "vertex_ai" | "grok" | "windsurf" => Ok(normalized),
|
||||
| "antigravity" | "vertex_ai" | "grok" | "windsurf" | "xai" => Ok(normalized),
|
||||
_ => Err(
|
||||
"provider_type 仅支持 custom / claude_code / kiro / codex / chatgpt_web / gemini_cli / antigravity / vertex_ai / grok / windsurf"
|
||||
"provider_type 仅支持 custom / claude_code / kiro / codex / chatgpt_web / gemini_cli / antigravity / vertex_ai / grok / windsurf / xai"
|
||||
.to_string(),
|
||||
),
|
||||
}
|
||||
@@ -405,6 +405,14 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_provider_type_supports_xai() {
|
||||
assert_eq!(
|
||||
normalize_provider_type_input(" xAI ").expect("type should normalize"),
|
||||
"xai"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_api_format_list_dedupes_canonical_formats() {
|
||||
assert_eq!(
|
||||
|
||||
@@ -388,10 +388,11 @@ impl<'a> AdminAppState<'a> {
|
||||
balance_type: &str,
|
||||
operator_id: Option<&str>,
|
||||
description: Option<&str>,
|
||||
clamp_deduction_to_available_balance: bool,
|
||||
) -> Result<
|
||||
Option<(
|
||||
aether_data::repository::wallet::StoredWalletSnapshot,
|
||||
crate::AdminWalletTransactionRecord,
|
||||
Option<crate::AdminWalletTransactionRecord>,
|
||||
)>,
|
||||
GatewayError,
|
||||
> {
|
||||
@@ -402,6 +403,7 @@ impl<'a> AdminAppState<'a> {
|
||||
balance_type,
|
||||
operator_id,
|
||||
description,
|
||||
clamp_deduction_to_available_balance,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
@@ -17,6 +17,64 @@ impl<'a> AdminAppState<'a> {
|
||||
pub(crate) fn cloned_app(&self) -> AppState {
|
||||
self.app.clone()
|
||||
}
|
||||
|
||||
pub(crate) async fn get_admin_user_wallet_balance_batch(
|
||||
&self,
|
||||
admin_user_id: &str,
|
||||
idempotency_key: &str,
|
||||
request_fingerprint: &str,
|
||||
) -> Result<
|
||||
Option<aether_data::repository::wallet::PrepareAdminUserWalletBalanceBatchOutcome>,
|
||||
GatewayError,
|
||||
> {
|
||||
self.app
|
||||
.get_admin_user_wallet_balance_batch(
|
||||
admin_user_id,
|
||||
idempotency_key,
|
||||
request_fingerprint,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn prepare_admin_user_wallet_balance_batch(
|
||||
&self,
|
||||
input: aether_data::repository::wallet::PrepareAdminUserWalletBalanceBatchInput,
|
||||
) -> Result<
|
||||
aether_data::repository::wallet::PrepareAdminUserWalletBalanceBatchOutcome,
|
||||
GatewayError,
|
||||
> {
|
||||
self.app
|
||||
.prepare_admin_user_wallet_balance_batch(input)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn record_admin_user_wallet_balance_batch_failure(
|
||||
&self,
|
||||
admin_user_id: &str,
|
||||
idempotency_key: &str,
|
||||
user_id: &str,
|
||||
reason: &str,
|
||||
) -> Result<aether_data::repository::wallet::AdminUserWalletBalanceBatchUserOutcome, GatewayError>
|
||||
{
|
||||
self.app
|
||||
.record_admin_user_wallet_balance_batch_failure(
|
||||
admin_user_id,
|
||||
idempotency_key,
|
||||
user_id,
|
||||
reason,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn adjust_admin_user_wallet_balance_batch_user(
|
||||
&self,
|
||||
input: aether_data::repository::wallet::AdjustWalletBalanceInBatchInput,
|
||||
) -> Result<aether_data::repository::wallet::AdminUserWalletBalanceBatchUserOutcome, GatewayError>
|
||||
{
|
||||
self.app
|
||||
.adjust_admin_user_wallet_balance_batch_user(input)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> AsRef<AppState> for AdminAppState<'a> {
|
||||
|
||||
@@ -126,6 +126,30 @@ impl<'a> AdminAppState<'a> {
|
||||
self.app.list_user_group_members(group_id).await
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_usage_user_group_member_ids(
|
||||
&self,
|
||||
group_id: &str,
|
||||
include_inactive: bool,
|
||||
exclude_admin: bool,
|
||||
) -> Result<Option<Vec<String>>, GatewayError> {
|
||||
if self.find_user_group_by_id(group_id).await?.is_none() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let mut user_ids = self
|
||||
.list_user_group_members(group_id)
|
||||
.await?
|
||||
.into_iter()
|
||||
.filter(|member| !member.is_deleted)
|
||||
.filter(|member| include_inactive || member.is_active)
|
||||
.filter(|member| !exclude_admin || !member.role.eq_ignore_ascii_case("admin"))
|
||||
.map(|member| member.user_id)
|
||||
.collect::<Vec<_>>();
|
||||
user_ids.sort();
|
||||
user_ids.dedup();
|
||||
Ok(Some(user_ids))
|
||||
}
|
||||
|
||||
pub(crate) async fn replace_user_group_members(
|
||||
&self,
|
||||
group_id: &str,
|
||||
|
||||
@@ -64,6 +64,7 @@ pub(crate) async fn build_admin_list_user_api_keys_response(
|
||||
"rate_limit": record.rate_limit,
|
||||
"concurrent_limit": record.concurrent_limit,
|
||||
"feature_settings": record.feature_settings,
|
||||
"ip_rules": record.ip_rules,
|
||||
"expires_at": format_optional_unix_secs_iso8601(record.expires_at_unix_secs),
|
||||
"last_used_at": format_optional_unix_secs_iso8601(record.last_used_at_unix_secs),
|
||||
"created_at": format_optional_unix_secs_iso8601(record.created_at_unix_secs),
|
||||
|
||||
@@ -1,11 +1,16 @@
|
||||
use super::{
|
||||
build_admin_users_bad_request_response, build_admin_users_permission_denied_response,
|
||||
build_admin_users_read_only_response, disabled_user_policy_detail, disabled_user_policy_field,
|
||||
build_admin_users_read_only_response, build_admin_users_wallet_permission_denied_response,
|
||||
disabled_user_policy_detail, disabled_user_policy_field,
|
||||
management_token_may_adjust_admin_wallet_balance,
|
||||
management_token_may_administer_user_accounts, normalize_admin_user_role,
|
||||
};
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||
use crate::handlers::admin::shared::attach_admin_audit_response;
|
||||
use crate::GatewayError;
|
||||
use aether_data::repository::wallet::{
|
||||
AdminUserWalletBalanceBatchUserOutcome, PrepareAdminUserWalletBalanceBatchOutcome,
|
||||
};
|
||||
use axum::{
|
||||
body::{Body, Bytes},
|
||||
http,
|
||||
@@ -13,9 +18,10 @@ use axum::{
|
||||
Json,
|
||||
};
|
||||
use serde_json::{json, Value};
|
||||
use sha2::{Digest as _, Sha256};
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
#[derive(Debug, Clone, Default, serde::Deserialize)]
|
||||
#[derive(Debug, Clone, Default, serde::Deserialize, serde::Serialize)]
|
||||
struct AdminUserSelectionFilters {
|
||||
#[serde(default)]
|
||||
search: Option<String>,
|
||||
@@ -27,7 +33,7 @@ struct AdminUserSelectionFilters {
|
||||
group_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
#[derive(Debug, Clone, Default, serde::Serialize)]
|
||||
struct AdminUserSelectionRequest {
|
||||
user_ids: Vec<String>,
|
||||
group_ids: Vec<String>,
|
||||
@@ -40,6 +46,7 @@ struct AdminUserBatchActionRequest {
|
||||
selection: AdminUserSelectionRequest,
|
||||
action: String,
|
||||
payload: Option<Value>,
|
||||
idempotency_key: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
@@ -48,6 +55,8 @@ struct RawAdminUserBatchActionRequest {
|
||||
action: String,
|
||||
#[serde(default)]
|
||||
payload: Option<Value>,
|
||||
#[serde(default)]
|
||||
idempotency_key: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
@@ -68,7 +77,7 @@ struct AdminUserSelectionItem {
|
||||
matched_by: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
#[derive(Debug, Clone, serde::Deserialize, serde::Serialize)]
|
||||
struct AdminUserSelectionWarning {
|
||||
#[serde(rename = "type")]
|
||||
warning_type: String,
|
||||
@@ -88,6 +97,7 @@ struct AdminUserBatchMutation {
|
||||
role: Option<String>,
|
||||
is_active: Option<bool>,
|
||||
unlimited: Option<bool>,
|
||||
wallet_balance_adjustment: Option<AdminUserWalletBalanceAdjustment>,
|
||||
modified_fields: Vec<&'static str>,
|
||||
}
|
||||
|
||||
@@ -97,6 +107,28 @@ impl AdminUserBatchMutation {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
struct AdminUserWalletBalanceAdjustment {
|
||||
operation: AdminUserWalletBalanceOperation,
|
||||
amount: f64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
enum AdminUserWalletBalanceOperation {
|
||||
Add,
|
||||
Deduct,
|
||||
}
|
||||
|
||||
enum AdminBatchWalletBalanceAdjustmentError {
|
||||
WalletLookup,
|
||||
BalanceAdjustment,
|
||||
}
|
||||
|
||||
enum AdminBatchWalletLimitModeError {
|
||||
WalletLookup,
|
||||
Mutation,
|
||||
}
|
||||
|
||||
pub(in super::super) async fn build_admin_resolve_user_selection_response(
|
||||
state: &AdminAppState<'_>,
|
||||
_request_context: &AdminRequestContext<'_>,
|
||||
@@ -128,10 +160,29 @@ pub(in super::super) async fn build_admin_user_batch_action_response(
|
||||
Ok(value) => value,
|
||||
Err(detail) => return Ok(build_admin_user_batch_bad_request_response(detail)),
|
||||
};
|
||||
let mutation = match parse_batch_mutation(&request.action, request.payload) {
|
||||
let mutation = match parse_batch_mutation(&request.action, request.payload.clone()) {
|
||||
Ok(value) => value,
|
||||
Err(detail) => return Ok(build_admin_user_batch_bad_request_response(detail)),
|
||||
};
|
||||
if mutation.wallet_balance_adjustment.is_some() {
|
||||
if !management_token_may_adjust_admin_wallet_balance(request_context) {
|
||||
return Ok(build_admin_users_wallet_permission_denied_response(
|
||||
request_context,
|
||||
));
|
||||
}
|
||||
if !state.has_auth_wallet_write_capability() {
|
||||
return Ok(build_admin_users_read_only_response(
|
||||
"当前为只读模式,无法批量调整用户钱包余额",
|
||||
));
|
||||
}
|
||||
return build_admin_user_wallet_balance_batch_response(
|
||||
state,
|
||||
request_context,
|
||||
request,
|
||||
mutation,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
let resolved = match resolve_admin_user_selection(state, request.selection).await {
|
||||
Ok(value) => value,
|
||||
Err(detail) => return Ok(build_admin_user_batch_bad_request_response(detail)),
|
||||
@@ -177,9 +228,29 @@ pub(in super::super) async fn build_admin_user_batch_action_response(
|
||||
.iter()
|
||||
.map(|user_id| json!({ "user_id": user_id, "reason": "用户不存在或已删除" }))
|
||||
.collect::<Vec<_>>();
|
||||
let mut completed_user_ids = Vec::new();
|
||||
let mut uncertain_user_ids = Vec::new();
|
||||
let mut unprocessed_user_ids = Vec::new();
|
||||
let mut interrupted = false;
|
||||
|
||||
for item in &resolved.items {
|
||||
if state.find_user_auth_by_id(&item.user_id).await?.is_none() {
|
||||
for (item_index, item) in resolved.items.iter().enumerate() {
|
||||
let user = match state.find_user_auth_by_id(&item.user_id).await {
|
||||
Ok(user) => user,
|
||||
Err(_) => {
|
||||
record_batch_action_interruption(
|
||||
&resolved.items,
|
||||
item_index,
|
||||
false,
|
||||
"读取用户状态失败,批次已中止,该用户未执行",
|
||||
&mut failures,
|
||||
&mut uncertain_user_ids,
|
||||
&mut unprocessed_user_ids,
|
||||
);
|
||||
interrupted = true;
|
||||
break;
|
||||
}
|
||||
};
|
||||
if user.is_none() {
|
||||
failures.push(json!({
|
||||
"user_id": item.user_id,
|
||||
"reason": "用户不存在或已删除",
|
||||
@@ -202,17 +273,92 @@ pub(in super::super) async fn build_admin_user_batch_action_response(
|
||||
}
|
||||
|
||||
if let Some(unlimited) = mutation.unlimited {
|
||||
if !apply_batch_user_wallet_limit_mode(state, &item.user_id, unlimited).await? {
|
||||
failures.push(json!({
|
||||
"user_id": item.user_id,
|
||||
"reason": "用户钱包不可用",
|
||||
}));
|
||||
continue;
|
||||
match apply_batch_user_wallet_limit_mode(state, &item.user_id, unlimited).await {
|
||||
Ok(true) => {}
|
||||
Ok(false) => {
|
||||
failures.push(json!({
|
||||
"user_id": item.user_id,
|
||||
"reason": "用户钱包不可用",
|
||||
}));
|
||||
continue;
|
||||
}
|
||||
Err(AdminBatchWalletLimitModeError::WalletLookup) => {
|
||||
record_batch_action_interruption(
|
||||
&resolved.items,
|
||||
item_index,
|
||||
false,
|
||||
"读取用户钱包失败,批次已中止,该用户未执行",
|
||||
&mut failures,
|
||||
&mut uncertain_user_ids,
|
||||
&mut unprocessed_user_ids,
|
||||
);
|
||||
interrupted = true;
|
||||
break;
|
||||
}
|
||||
Err(AdminBatchWalletLimitModeError::Mutation) => {
|
||||
record_batch_action_interruption(
|
||||
&resolved.items,
|
||||
item_index,
|
||||
true,
|
||||
"用户钱包更新结果未确认,批次已中止,请核对钱包后再重试",
|
||||
&mut failures,
|
||||
&mut uncertain_user_ids,
|
||||
&mut unprocessed_user_ids,
|
||||
);
|
||||
interrupted = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if mutation.has_auth_user_fields()
|
||||
&& state
|
||||
if let Some(adjustment) = mutation.wallet_balance_adjustment {
|
||||
match apply_batch_user_wallet_balance_adjustment(
|
||||
state,
|
||||
&item.user_id,
|
||||
adjustment,
|
||||
current_admin_user_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(true) => {}
|
||||
Ok(false) => {
|
||||
failures.push(json!({
|
||||
"user_id": item.user_id,
|
||||
"reason": "用户钱包不可用",
|
||||
}));
|
||||
continue;
|
||||
}
|
||||
Err(AdminBatchWalletBalanceAdjustmentError::WalletLookup) => {
|
||||
record_batch_action_interruption(
|
||||
&resolved.items,
|
||||
item_index,
|
||||
false,
|
||||
"读取用户钱包失败,批次已中止,该用户未执行",
|
||||
&mut failures,
|
||||
&mut uncertain_user_ids,
|
||||
&mut unprocessed_user_ids,
|
||||
);
|
||||
interrupted = true;
|
||||
break;
|
||||
}
|
||||
Err(AdminBatchWalletBalanceAdjustmentError::BalanceAdjustment) => {
|
||||
record_batch_action_interruption(
|
||||
&resolved.items,
|
||||
item_index,
|
||||
true,
|
||||
"余额调整结果未确认,批次已中止,请核对钱包后再重试",
|
||||
&mut failures,
|
||||
&mut uncertain_user_ids,
|
||||
&mut unprocessed_user_ids,
|
||||
);
|
||||
interrupted = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if mutation.has_auth_user_fields() {
|
||||
let updated_user = match state
|
||||
.update_local_auth_user_admin_fields(
|
||||
&item.user_id,
|
||||
mutation.role.clone(),
|
||||
@@ -226,22 +372,39 @@ pub(in super::super) async fn build_admin_user_batch_action_response(
|
||||
None,
|
||||
mutation.is_active,
|
||||
)
|
||||
.await?
|
||||
.is_none()
|
||||
{
|
||||
failures.push(json!({
|
||||
"user_id": item.user_id,
|
||||
"reason": "用户不存在或已删除",
|
||||
}));
|
||||
continue;
|
||||
.await
|
||||
{
|
||||
Ok(user) => user,
|
||||
Err(_) => {
|
||||
record_batch_action_interruption(
|
||||
&resolved.items,
|
||||
item_index,
|
||||
true,
|
||||
"用户更新结果未确认,批次已中止,请核对后再重试",
|
||||
&mut failures,
|
||||
&mut uncertain_user_ids,
|
||||
&mut unprocessed_user_ids,
|
||||
);
|
||||
interrupted = true;
|
||||
break;
|
||||
}
|
||||
};
|
||||
if updated_user.is_none() {
|
||||
failures.push(json!({
|
||||
"user_id": item.user_id,
|
||||
"reason": "用户不存在或已删除",
|
||||
}));
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
success += 1;
|
||||
completed_user_ids.push(item.user_id.clone());
|
||||
}
|
||||
|
||||
let failed = failures.len();
|
||||
let total = success + failed;
|
||||
let response = Json(json!({
|
||||
let mut response_payload = json!({
|
||||
"total": total,
|
||||
"success": success,
|
||||
"failed": failed,
|
||||
@@ -249,8 +412,14 @@ pub(in super::super) async fn build_admin_user_batch_action_response(
|
||||
"warnings": resolved.warnings,
|
||||
"action": request.action.trim().to_ascii_lowercase(),
|
||||
"modified_fields": mutation.modified_fields,
|
||||
}))
|
||||
.into_response();
|
||||
"interrupted": interrupted,
|
||||
});
|
||||
if interrupted {
|
||||
response_payload["completed_user_ids"] = json!(completed_user_ids);
|
||||
response_payload["uncertain_user_ids"] = json!(uncertain_user_ids);
|
||||
response_payload["unprocessed_user_ids"] = json!(unprocessed_user_ids);
|
||||
}
|
||||
let response = Json(response_payload).into_response();
|
||||
|
||||
Ok(attach_admin_audit_response(
|
||||
response,
|
||||
@@ -261,6 +430,408 @@ pub(in super::super) async fn build_admin_user_batch_action_response(
|
||||
))
|
||||
}
|
||||
|
||||
fn record_batch_action_interruption(
|
||||
items: &[AdminUserSelectionItem],
|
||||
item_index: usize,
|
||||
current_result_uncertain: bool,
|
||||
reason: &str,
|
||||
failures: &mut Vec<Value>,
|
||||
uncertain_user_ids: &mut Vec<String>,
|
||||
unprocessed_user_ids: &mut Vec<String>,
|
||||
) {
|
||||
let current_item = &items[item_index];
|
||||
failures.push(json!({
|
||||
"user_id": current_item.user_id,
|
||||
"reason": reason,
|
||||
}));
|
||||
if current_result_uncertain {
|
||||
uncertain_user_ids.push(current_item.user_id.clone());
|
||||
} else {
|
||||
unprocessed_user_ids.push(current_item.user_id.clone());
|
||||
}
|
||||
|
||||
for item in items.iter().skip(item_index + 1) {
|
||||
failures.push(json!({
|
||||
"user_id": item.user_id,
|
||||
"reason": "因前序错误未执行",
|
||||
}));
|
||||
unprocessed_user_ids.push(item.user_id.clone());
|
||||
}
|
||||
}
|
||||
|
||||
async fn build_admin_user_wallet_balance_batch_response(
|
||||
state: &AdminAppState<'_>,
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
request: AdminUserBatchActionRequest,
|
||||
mutation: AdminUserBatchMutation,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
let Some(idempotency_key) = request
|
||||
.idempotency_key
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| {
|
||||
!value.is_empty()
|
||||
&& value.len() <= 128
|
||||
&& value.bytes().all(|byte| (0x21..=0x7e).contains(&byte))
|
||||
})
|
||||
else {
|
||||
return Ok(build_admin_user_batch_bad_request_response(
|
||||
"余额批量操作必须提供有效的 idempotency_key".to_string(),
|
||||
));
|
||||
};
|
||||
let Some(admin_user_id) = request_context
|
||||
.decision()
|
||||
.and_then(|decision| decision.admin_principal.as_ref())
|
||||
.map(|principal| principal.user_id.clone())
|
||||
else {
|
||||
return Ok(build_admin_users_permission_denied_response(
|
||||
request_context,
|
||||
));
|
||||
};
|
||||
let action = request.action.trim().to_ascii_lowercase();
|
||||
let fingerprint_payload = json!({
|
||||
"selection": &request.selection,
|
||||
"action": &action,
|
||||
"payload": &request.payload,
|
||||
});
|
||||
let encoded = serde_json::to_vec(&fingerprint_payload)
|
||||
.map_err(|error| GatewayError::Internal(error.to_string()))?;
|
||||
let request_fingerprint = format!("{:x}", Sha256::digest(encoded));
|
||||
|
||||
let existing = state
|
||||
.get_admin_user_wallet_balance_batch(&admin_user_id, idempotency_key, &request_fingerprint)
|
||||
.await?;
|
||||
let batch = match existing {
|
||||
Some(PrepareAdminUserWalletBalanceBatchOutcome::Conflict) => {
|
||||
return Ok(build_admin_user_batch_idempotency_conflict_response());
|
||||
}
|
||||
Some(PrepareAdminUserWalletBalanceBatchOutcome::Ready(batch)) => batch,
|
||||
None => {
|
||||
let resolved =
|
||||
match resolve_admin_user_selection(state, request.selection.clone()).await {
|
||||
Ok(value) => value,
|
||||
Err(detail) => return Ok(build_admin_user_batch_bad_request_response(detail)),
|
||||
};
|
||||
let warnings = serde_json::to_value(&resolved.warnings)
|
||||
.ok()
|
||||
.and_then(|value| value.as_array().cloned())
|
||||
.unwrap_or_default();
|
||||
let prepared = state
|
||||
.prepare_admin_user_wallet_balance_batch(
|
||||
aether_data::repository::wallet::PrepareAdminUserWalletBalanceBatchInput {
|
||||
admin_user_id: admin_user_id.clone(),
|
||||
idempotency_key: idempotency_key.to_string(),
|
||||
request_fingerprint: request_fingerprint.clone(),
|
||||
target_user_ids: resolved
|
||||
.items
|
||||
.iter()
|
||||
.map(|item| item.user_id.clone())
|
||||
.collect(),
|
||||
missing_user_ids: resolved.missing_user_ids,
|
||||
warnings,
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
match prepared {
|
||||
PrepareAdminUserWalletBalanceBatchOutcome::Conflict => {
|
||||
return Ok(build_admin_user_batch_idempotency_conflict_response());
|
||||
}
|
||||
PrepareAdminUserWalletBalanceBatchOutcome::Ready(batch) => batch,
|
||||
}
|
||||
}
|
||||
};
|
||||
let warnings: Vec<AdminUserSelectionWarning> =
|
||||
serde_json::from_value(Value::Array(batch.warnings.clone()))
|
||||
.map_err(|error| GatewayError::Internal(error.to_string()))?;
|
||||
let resolved = ResolvedAdminUserSelection {
|
||||
items: batch
|
||||
.target_user_ids
|
||||
.iter()
|
||||
.map(|user_id| AdminUserSelectionItem {
|
||||
user_id: user_id.clone(),
|
||||
username: String::new(),
|
||||
email: None,
|
||||
role: "user".to_string(),
|
||||
is_active: true,
|
||||
matched_by: Vec::new(),
|
||||
})
|
||||
.collect(),
|
||||
missing_user_ids: batch.missing_user_ids.clone(),
|
||||
warnings,
|
||||
};
|
||||
let adjustment = mutation
|
||||
.wallet_balance_adjustment
|
||||
.expect("wallet balance action should have an adjustment");
|
||||
let signed_amount = match adjustment.operation {
|
||||
AdminUserWalletBalanceOperation::Add => adjustment.amount,
|
||||
AdminUserWalletBalanceOperation::Deduct => -adjustment.amount,
|
||||
};
|
||||
|
||||
let mut outcomes = batch.user_outcomes;
|
||||
let mut completed_user_ids = outcomes
|
||||
.iter()
|
||||
.filter_map(|(user_id, outcome)| {
|
||||
matches!(outcome, AdminUserWalletBalanceBatchUserOutcome::Succeeded)
|
||||
.then_some(user_id.clone())
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let mut success = completed_user_ids.len();
|
||||
let mut failures = resolved
|
||||
.missing_user_ids
|
||||
.iter()
|
||||
.map(|user_id| json!({ "user_id": user_id, "reason": "用户不存在或已删除" }))
|
||||
.collect::<Vec<_>>();
|
||||
for (user_id, outcome) in &outcomes {
|
||||
if let AdminUserWalletBalanceBatchUserOutcome::Failed(reason) = outcome {
|
||||
failures.push(json!({ "user_id": user_id, "reason": reason }));
|
||||
}
|
||||
}
|
||||
let mut uncertain_user_ids = Vec::new();
|
||||
let mut unprocessed_user_ids = Vec::new();
|
||||
let mut interrupted = false;
|
||||
|
||||
for (item_index, item) in resolved.items.iter().enumerate() {
|
||||
if outcomes.contains_key(&item.user_id) {
|
||||
continue;
|
||||
}
|
||||
match state.find_user_auth_by_id(&item.user_id).await {
|
||||
Err(_) => {
|
||||
failures.push(json!({
|
||||
"user_id": item.user_id,
|
||||
"reason": "读取用户状态失败,批次已中止,该用户未执行",
|
||||
}));
|
||||
unprocessed_user_ids.push(item.user_id.clone());
|
||||
interrupted = true;
|
||||
}
|
||||
Ok(None) => {
|
||||
let outcome = match state
|
||||
.record_admin_user_wallet_balance_batch_failure(
|
||||
&admin_user_id,
|
||||
idempotency_key,
|
||||
&item.user_id,
|
||||
"用户不存在或已删除",
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(outcome) => outcome,
|
||||
Err(_) => {
|
||||
failures.push(json!({
|
||||
"user_id": item.user_id,
|
||||
"reason": "记录用户状态失败,批次已中止,该用户未执行",
|
||||
}));
|
||||
unprocessed_user_ids.push(item.user_id.clone());
|
||||
interrupted = true;
|
||||
append_wallet_batch_unprocessed_suffix(
|
||||
&resolved.items,
|
||||
item_index + 1,
|
||||
&outcomes,
|
||||
&mut failures,
|
||||
&mut unprocessed_user_ids,
|
||||
);
|
||||
break;
|
||||
}
|
||||
};
|
||||
outcomes.insert(item.user_id.clone(), outcome.clone());
|
||||
match outcome {
|
||||
AdminUserWalletBalanceBatchUserOutcome::Succeeded => {
|
||||
completed_user_ids.push(item.user_id.clone());
|
||||
success += 1;
|
||||
}
|
||||
AdminUserWalletBalanceBatchUserOutcome::Failed(reason) => {
|
||||
failures.push(json!({ "user_id": item.user_id, "reason": reason }));
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(Some(_)) => {
|
||||
let wallet = match state
|
||||
.find_wallet(aether_data::repository::wallet::WalletLookupKey::UserId(
|
||||
&item.user_id,
|
||||
))
|
||||
.await
|
||||
{
|
||||
Ok(Some(wallet)) => wallet,
|
||||
Ok(None) => {
|
||||
let outcome = match state
|
||||
.record_admin_user_wallet_balance_batch_failure(
|
||||
&admin_user_id,
|
||||
idempotency_key,
|
||||
&item.user_id,
|
||||
"用户钱包不可用",
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(outcome) => outcome,
|
||||
Err(_) => {
|
||||
failures.push(json!({
|
||||
"user_id": item.user_id,
|
||||
"reason": "记录用户钱包状态失败,批次已中止,该用户未执行",
|
||||
}));
|
||||
unprocessed_user_ids.push(item.user_id.clone());
|
||||
interrupted = true;
|
||||
append_wallet_batch_unprocessed_suffix(
|
||||
&resolved.items,
|
||||
item_index + 1,
|
||||
&outcomes,
|
||||
&mut failures,
|
||||
&mut unprocessed_user_ids,
|
||||
);
|
||||
break;
|
||||
}
|
||||
};
|
||||
outcomes.insert(item.user_id.clone(), outcome.clone());
|
||||
match outcome {
|
||||
AdminUserWalletBalanceBatchUserOutcome::Succeeded => {
|
||||
completed_user_ids.push(item.user_id.clone());
|
||||
success += 1;
|
||||
}
|
||||
AdminUserWalletBalanceBatchUserOutcome::Failed(reason) => {
|
||||
failures.push(json!({ "user_id": item.user_id, "reason": reason }));
|
||||
}
|
||||
}
|
||||
continue;
|
||||
}
|
||||
Err(_) => {
|
||||
failures.push(json!({
|
||||
"user_id": item.user_id,
|
||||
"reason": "读取用户钱包失败,批次已中止,该用户未执行",
|
||||
}));
|
||||
unprocessed_user_ids.push(item.user_id.clone());
|
||||
interrupted = true;
|
||||
append_wallet_batch_unprocessed_suffix(
|
||||
&resolved.items,
|
||||
item_index + 1,
|
||||
&outcomes,
|
||||
&mut failures,
|
||||
&mut unprocessed_user_ids,
|
||||
);
|
||||
break;
|
||||
}
|
||||
};
|
||||
let result = state
|
||||
.adjust_admin_user_wallet_balance_batch_user(
|
||||
aether_data::repository::wallet::AdjustWalletBalanceInBatchInput {
|
||||
admin_user_id: admin_user_id.clone(),
|
||||
idempotency_key: idempotency_key.to_string(),
|
||||
user_id: item.user_id.clone(),
|
||||
adjustment: aether_data::repository::wallet::AdjustWalletBalanceInput {
|
||||
wallet_id: wallet.id,
|
||||
amount_usd: signed_amount,
|
||||
balance_type: "recharge".to_string(),
|
||||
operator_id: Some(admin_user_id.clone()),
|
||||
description: Some("管理员批量调整用户余额".to_string()),
|
||||
clamp_deduction_to_available_balance: true,
|
||||
batch_context: None,
|
||||
},
|
||||
},
|
||||
)
|
||||
.await;
|
||||
match result {
|
||||
Ok(AdminUserWalletBalanceBatchUserOutcome::Succeeded) => {
|
||||
outcomes.insert(
|
||||
item.user_id.clone(),
|
||||
AdminUserWalletBalanceBatchUserOutcome::Succeeded,
|
||||
);
|
||||
completed_user_ids.push(item.user_id.clone());
|
||||
success += 1;
|
||||
}
|
||||
Ok(AdminUserWalletBalanceBatchUserOutcome::Failed(reason)) => {
|
||||
outcomes.insert(
|
||||
item.user_id.clone(),
|
||||
AdminUserWalletBalanceBatchUserOutcome::Failed(reason.clone()),
|
||||
);
|
||||
failures.push(json!({ "user_id": item.user_id, "reason": reason }));
|
||||
}
|
||||
Err(_) => {
|
||||
failures.push(json!({
|
||||
"user_id": item.user_id,
|
||||
"reason": "余额调整结果未确认,批次已中止,请使用同一批次重试以核对结果",
|
||||
}));
|
||||
uncertain_user_ids.push(item.user_id.clone());
|
||||
interrupted = true;
|
||||
append_wallet_batch_unprocessed_suffix(
|
||||
&resolved.items,
|
||||
item_index + 1,
|
||||
&outcomes,
|
||||
&mut failures,
|
||||
&mut unprocessed_user_ids,
|
||||
);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if interrupted {
|
||||
for pending in resolved.items.iter().skip(item_index + 1) {
|
||||
if outcomes.contains_key(&pending.user_id) {
|
||||
continue;
|
||||
}
|
||||
failures.push(json!({
|
||||
"user_id": pending.user_id,
|
||||
"reason": "因前序错误未执行",
|
||||
}));
|
||||
unprocessed_user_ids.push(pending.user_id.clone());
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
let failed = failures.len();
|
||||
let total = success + failed;
|
||||
let mut response_payload = json!({
|
||||
"total": total,
|
||||
"success": success,
|
||||
"failed": failed,
|
||||
"failures": failures,
|
||||
"warnings": resolved.warnings,
|
||||
"action": action,
|
||||
"modified_fields": mutation.modified_fields,
|
||||
"interrupted": interrupted,
|
||||
});
|
||||
if interrupted {
|
||||
response_payload["completed_user_ids"] = json!(completed_user_ids);
|
||||
response_payload["uncertain_user_ids"] = json!(uncertain_user_ids);
|
||||
response_payload["unprocessed_user_ids"] = json!(unprocessed_user_ids);
|
||||
}
|
||||
let response = Json(response_payload).into_response();
|
||||
Ok(attach_admin_audit_response(
|
||||
response,
|
||||
"admin_users_batch_action_executed",
|
||||
"batch_update_users",
|
||||
"user_batch",
|
||||
"users",
|
||||
))
|
||||
}
|
||||
|
||||
fn append_wallet_batch_unprocessed_suffix(
|
||||
items: &[AdminUserSelectionItem],
|
||||
start_index: usize,
|
||||
outcomes: &BTreeMap<String, AdminUserWalletBalanceBatchUserOutcome>,
|
||||
failures: &mut Vec<Value>,
|
||||
unprocessed_user_ids: &mut Vec<String>,
|
||||
) {
|
||||
for pending in items.iter().skip(start_index) {
|
||||
if outcomes.contains_key(&pending.user_id) {
|
||||
continue;
|
||||
}
|
||||
failures.push(json!({
|
||||
"user_id": pending.user_id,
|
||||
"reason": "因前序错误未执行",
|
||||
}));
|
||||
unprocessed_user_ids.push(pending.user_id.clone());
|
||||
}
|
||||
}
|
||||
|
||||
fn build_admin_user_batch_idempotency_conflict_response() -> Response<Body> {
|
||||
(
|
||||
http::StatusCode::CONFLICT,
|
||||
Json(json!({
|
||||
"detail": "idempotency_key was already used with a different request",
|
||||
"error_code": "idempotency_key_conflict",
|
||||
})),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
|
||||
fn parse_resolve_selection_request(
|
||||
request_body: Option<&Bytes>,
|
||||
) -> Result<AdminUserSelectionRequest, String> {
|
||||
@@ -286,6 +857,7 @@ fn parse_batch_action_request(
|
||||
selection: parse_selection_request_value(raw.selection)?,
|
||||
action: raw.action,
|
||||
payload: raw.payload,
|
||||
idempotency_key: raw.idempotency_key,
|
||||
})
|
||||
}
|
||||
_ => Err("Invalid JSON request body".to_string()),
|
||||
@@ -604,10 +1176,37 @@ fn parse_batch_mutation(
|
||||
}),
|
||||
"update_access_control" => parse_access_control_mutation(payload),
|
||||
"update_role" => parse_role_mutation(payload),
|
||||
"adjust_wallet_balance" => parse_wallet_balance_adjustment_mutation(payload),
|
||||
_ => Err("不支持的批量操作".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_wallet_balance_adjustment_mutation(
|
||||
payload: Option<Value>,
|
||||
) -> Result<AdminUserBatchMutation, String> {
|
||||
let Some(Value::Object(payload)) = payload else {
|
||||
return Err("payload 必须是对象".to_string());
|
||||
};
|
||||
let operation = match payload.get("operation").and_then(Value::as_str) {
|
||||
Some("add") => AdminUserWalletBalanceOperation::Add,
|
||||
Some("deduct") => AdminUserWalletBalanceOperation::Deduct,
|
||||
_ => return Err("operation 必须为 add 或 deduct".to_string()),
|
||||
};
|
||||
let amount = payload
|
||||
.get("amount")
|
||||
.and_then(Value::as_f64)
|
||||
.ok_or_else(|| "amount 必须为大于 0 的有限数字".to_string())?;
|
||||
if !amount.is_finite() || amount <= 0.0 {
|
||||
return Err("amount 必须为大于 0 的有限数字".to_string());
|
||||
}
|
||||
|
||||
Ok(AdminUserBatchMutation {
|
||||
wallet_balance_adjustment: Some(AdminUserWalletBalanceAdjustment { operation, amount }),
|
||||
modified_fields: vec!["wallet_balance"],
|
||||
..AdminUserBatchMutation::default()
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_role_mutation(payload: Option<Value>) -> Result<AdminUserBatchMutation, String> {
|
||||
let Some(Value::Object(payload)) = payload else {
|
||||
return Err("payload 必须是对象".to_string());
|
||||
@@ -700,30 +1299,68 @@ async fn apply_batch_user_wallet_limit_mode(
|
||||
state: &AdminAppState<'_>,
|
||||
user_id: &str,
|
||||
unlimited: bool,
|
||||
) -> Result<bool, GatewayError> {
|
||||
) -> Result<bool, AdminBatchWalletLimitModeError> {
|
||||
let desired_limit_mode = if unlimited { "unlimited" } else { "finite" };
|
||||
match state
|
||||
let wallet = state
|
||||
.find_wallet(aether_data::repository::wallet::WalletLookupKey::UserId(
|
||||
user_id,
|
||||
))
|
||||
.await?
|
||||
{
|
||||
.await
|
||||
.map_err(|_| AdminBatchWalletLimitModeError::WalletLookup)?;
|
||||
match wallet {
|
||||
Some(wallet) => {
|
||||
if wallet.limit_mode.eq_ignore_ascii_case(desired_limit_mode) {
|
||||
return Ok(true);
|
||||
}
|
||||
Ok(state
|
||||
.update_auth_user_wallet_limit_mode(user_id, desired_limit_mode)
|
||||
.await?
|
||||
.await
|
||||
.map_err(|_| AdminBatchWalletLimitModeError::Mutation)?
|
||||
.is_some())
|
||||
}
|
||||
None => Ok(state
|
||||
.initialize_auth_user_wallet(user_id, 0.0, unlimited)
|
||||
.await?
|
||||
.await
|
||||
.map_err(|_| AdminBatchWalletLimitModeError::Mutation)?
|
||||
.is_some()),
|
||||
}
|
||||
}
|
||||
|
||||
async fn apply_batch_user_wallet_balance_adjustment(
|
||||
state: &AdminAppState<'_>,
|
||||
user_id: &str,
|
||||
adjustment: AdminUserWalletBalanceAdjustment,
|
||||
operator_id: Option<&str>,
|
||||
) -> Result<bool, AdminBatchWalletBalanceAdjustmentError> {
|
||||
// Resolve only the wallet ID; the repository clamps the deduction under its row lock.
|
||||
let Some(wallet) = state
|
||||
.find_wallet(aether_data::repository::wallet::WalletLookupKey::UserId(
|
||||
user_id,
|
||||
))
|
||||
.await
|
||||
.map_err(|_| AdminBatchWalletBalanceAdjustmentError::WalletLookup)?
|
||||
else {
|
||||
return Ok(false);
|
||||
};
|
||||
let amount = match adjustment.operation {
|
||||
AdminUserWalletBalanceOperation::Add => adjustment.amount,
|
||||
AdminUserWalletBalanceOperation::Deduct => -adjustment.amount,
|
||||
};
|
||||
|
||||
state
|
||||
.admin_adjust_wallet_balance(
|
||||
&wallet.id,
|
||||
amount,
|
||||
"recharge",
|
||||
operator_id,
|
||||
Some("管理员批量调整用户余额"),
|
||||
true,
|
||||
)
|
||||
.await
|
||||
.map_err(|_| AdminBatchWalletBalanceAdjustmentError::BalanceAdjustment)
|
||||
.map(|result| result.is_some())
|
||||
}
|
||||
|
||||
fn build_admin_user_batch_bad_request_response(detail: String) -> Response<Body> {
|
||||
if detail.as_str() == "缺少 user_id" {
|
||||
return build_admin_users_bad_request_response("缺少 user_id");
|
||||
|
||||
@@ -49,12 +49,14 @@ use self::shared::AdminUpdateUserPatch;
|
||||
use self::shared::{
|
||||
admin_default_user_initial_gift, build_admin_users_bad_request_response,
|
||||
build_admin_users_data_unavailable_response, build_admin_users_permission_denied_response,
|
||||
build_admin_users_read_only_response, disabled_user_policy_detail, disabled_user_policy_field,
|
||||
format_optional_datetime_iso8601, legacy_admin_list_policy_mode,
|
||||
legacy_admin_rate_limit_policy_mode, management_token_may_administer_user_accounts,
|
||||
normalize_admin_optional_user_email, normalize_admin_user_group_ids, normalize_admin_user_role,
|
||||
normalize_admin_username, validate_admin_user_password, AdminCreateUserApiKeyRequest,
|
||||
AdminCreateUserRequest, AdminToggleUserApiKeyLockRequest, AdminUpdateUserApiKeyRequest,
|
||||
build_admin_users_read_only_response, build_admin_users_wallet_permission_denied_response,
|
||||
disabled_user_policy_detail, disabled_user_policy_field, format_optional_datetime_iso8601,
|
||||
legacy_admin_list_policy_mode, legacy_admin_rate_limit_policy_mode,
|
||||
management_token_may_adjust_admin_wallet_balance,
|
||||
management_token_may_administer_user_accounts, normalize_admin_optional_user_email,
|
||||
normalize_admin_user_group_ids, normalize_admin_user_role, normalize_admin_username,
|
||||
validate_admin_user_password, AdminCreateUserApiKeyRequest, AdminCreateUserRequest,
|
||||
AdminToggleUserApiKeyLockRequest, AdminUpdateUserApiKeyRequest,
|
||||
};
|
||||
pub(crate) use self::shared::{
|
||||
normalize_admin_list_policy_mode, normalize_admin_rate_limit_policy_mode,
|
||||
|
||||
@@ -165,6 +165,18 @@ pub(super) fn management_token_may_administer_user_accounts(
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn management_token_may_adjust_admin_wallet_balance(
|
||||
request_context: &crate::handlers::admin::request::AdminRequestContext<'_>,
|
||||
) -> bool {
|
||||
request_context.decision().is_some_and(|decision| {
|
||||
crate::control::management_token_principal_has_permission(decision, "admin:wallets:write")
|
||||
|| crate::control::management_token_principal_has_permission(
|
||||
decision,
|
||||
"admin:wallets:admin",
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn build_admin_users_permission_denied_response(
|
||||
request_context: &crate::handlers::admin::request::AdminRequestContext<'_>,
|
||||
) -> Response<Body> {
|
||||
@@ -192,6 +204,34 @@ pub(super) fn build_admin_users_permission_denied_response(
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn build_admin_users_wallet_permission_denied_response(
|
||||
request_context: &crate::handlers::admin::request::AdminRequestContext<'_>,
|
||||
) -> Response<Body> {
|
||||
let actor_id = request_context
|
||||
.decision()
|
||||
.and_then(|decision| decision.admin_principal.as_ref())
|
||||
.and_then(|principal| principal.management_token_id.as_deref())
|
||||
.unwrap_or("unknown");
|
||||
crate::handlers::admin::shared::attach_admin_audit_response(
|
||||
(
|
||||
http::StatusCode::FORBIDDEN,
|
||||
Json(json!({
|
||||
"detail": "management token permission denied",
|
||||
"required_permissions": ["admin:wallets:write", "admin:wallets:admin"],
|
||||
"permission_mode": "any_of",
|
||||
"route_family": request_context.route_family(),
|
||||
"route_kind": request_context.route_kind(),
|
||||
"request_path": request_context.path(),
|
||||
})),
|
||||
)
|
||||
.into_response(),
|
||||
"admin_user_wallet_balance_permission_denied",
|
||||
"permission_denied",
|
||||
"admin_user_wallet_balance",
|
||||
actor_id,
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn normalize_admin_optional_user_email(
|
||||
value: Option<&str>,
|
||||
) -> Result<Option<String>, String> {
|
||||
@@ -397,9 +437,52 @@ pub(super) fn format_optional_datetime_iso8601(
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{normalize_admin_user_api_formats, AdminUpdateUserApiKeyRequest};
|
||||
use super::{
|
||||
build_admin_users_wallet_permission_denied_response, normalize_admin_user_api_formats,
|
||||
AdminUpdateUserApiKeyRequest,
|
||||
};
|
||||
use crate::control::{GatewayControlDecision, GatewayPublicRequestContext};
|
||||
use crate::handlers::admin::request::AdminRequestContext;
|
||||
use axum::http::{HeaderMap, Method, Uri};
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn wallet_permission_denial_uses_wallet_audit_category() {
|
||||
let uri: Uri = "/api/admin/users/batch-action"
|
||||
.parse()
|
||||
.expect("uri should parse");
|
||||
let method = Method::POST;
|
||||
let headers = HeaderMap::new();
|
||||
let decision = GatewayControlDecision::synthetic(
|
||||
uri.path(),
|
||||
Some("admin_proxy".to_string()),
|
||||
Some("users_manage".to_string()),
|
||||
Some("batch_user_action".to_string()),
|
||||
Some("admin:users".to_string()),
|
||||
);
|
||||
let context = GatewayPublicRequestContext::from_request_parts(
|
||||
"trace-wallet-permission-denied",
|
||||
&method,
|
||||
&uri,
|
||||
&headers,
|
||||
Some(decision),
|
||||
);
|
||||
let request_context = AdminRequestContext::new(&context);
|
||||
|
||||
let response = build_admin_users_wallet_permission_denied_response(&request_context);
|
||||
let event = response
|
||||
.extensions()
|
||||
.get::<crate::audit::AdminAuditEvent>()
|
||||
.expect("wallet denial should attach an audit event");
|
||||
|
||||
assert_eq!(
|
||||
event.event_name,
|
||||
"admin_user_wallet_balance_permission_denied"
|
||||
);
|
||||
assert_eq!(event.action, "permission_denied");
|
||||
assert_eq!(event.target_type, "admin_user_wallet_balance");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn admin_user_api_formats_accept_current_canonical_signatures() {
|
||||
assert_eq!(
|
||||
|
||||
@@ -90,6 +90,8 @@ pub(super) struct ResponsesWebSocketContinuationRecord {
|
||||
/// request JSON can never set it.
|
||||
#[serde(default)]
|
||||
deepseek_opaque_reasoning_replay: bool,
|
||||
#[serde(default)]
|
||||
xai_encrypted_reasoning_replay: bool,
|
||||
/// A prior turn stored PII sentinels whose restore mapping exists only on
|
||||
/// the original downstream socket. Such a chain cannot safely resume on a
|
||||
/// new socket without leaking sentinels, so lookup succeeds but bootstrap
|
||||
@@ -122,6 +124,10 @@ impl ResponsesWebSocketContinuationRecord {
|
||||
normalization.reasoning_replay_policy(),
|
||||
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque
|
||||
),
|
||||
xai_encrypted_reasoning_replay: matches!(
|
||||
normalization.reasoning_replay_policy(),
|
||||
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::XaiEncrypted
|
||||
),
|
||||
has_connection_local_redaction,
|
||||
responses_lite_static_config,
|
||||
};
|
||||
@@ -156,7 +162,9 @@ impl ResponsesWebSocketContinuationRecord {
|
||||
pub(super) fn reasoning_replay_policy(
|
||||
&self,
|
||||
) -> crate::ai_serving::OpenAiResponsesReasoningReplayPolicy {
|
||||
if self.deepseek_opaque_reasoning_replay {
|
||||
if self.xai_encrypted_reasoning_replay {
|
||||
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::XaiEncrypted
|
||||
} else if self.deepseek_opaque_reasoning_replay {
|
||||
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque
|
||||
} else {
|
||||
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::OpenAiItemIds
|
||||
@@ -476,6 +484,7 @@ mod tests {
|
||||
binding_fingerprint: [7; 32],
|
||||
normalization_fingerprint: [9; 32],
|
||||
deepseek_opaque_reasoning_replay: false,
|
||||
xai_encrypted_reasoning_replay: false,
|
||||
has_connection_local_redaction: false,
|
||||
responses_lite_static_config: Some(ResponsesLiteStaticConfig::from_response_create(
|
||||
&json!({
|
||||
@@ -714,6 +723,29 @@ mod tests {
|
||||
assert_eq!(decoded, record());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serialized_record_preserves_xai_replay_policy_and_reads_legacy_records() {
|
||||
let mut expected = record();
|
||||
expected.xai_encrypted_reasoning_replay = true;
|
||||
let mut serialized = serde_json::to_value(&expected).unwrap();
|
||||
let decoded: ResponsesWebSocketContinuationRecord =
|
||||
serde_json::from_value(serialized.clone()).unwrap();
|
||||
assert_eq!(
|
||||
decoded.reasoning_replay_policy(),
|
||||
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::XaiEncrypted
|
||||
);
|
||||
serialized
|
||||
.as_object_mut()
|
||||
.unwrap()
|
||||
.remove("xai_encrypted_reasoning_replay");
|
||||
let legacy: ResponsesWebSocketContinuationRecord =
|
||||
serde_json::from_value(serialized).unwrap();
|
||||
assert_eq!(
|
||||
legacy.reasoning_replay_policy(),
|
||||
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::OpenAiItemIds
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serialized_record_preserves_only_the_server_derived_reasoning_replay_policy_bit() {
|
||||
let mut expected = record();
|
||||
|
||||
@@ -567,6 +567,7 @@ fn build_users_me_usage_record_payload(
|
||||
"id": item.id,
|
||||
"model": item.model,
|
||||
"target_model": serde_json::Value::Null,
|
||||
"response_model": item.provider_response_model(),
|
||||
"api_format": item.api_format,
|
||||
"endpoint_api_format": item.endpoint_api_format,
|
||||
"has_format_conversion": item.has_format_conversion,
|
||||
@@ -681,6 +682,7 @@ fn build_users_me_usage_active_payload(item: &StoredRequestUsageAudit) -> serde_
|
||||
"client_ip": users_me_usage_metadata_string(item, "client_ip"),
|
||||
"user_agent": users_me_usage_metadata_string(item, "user_agent"),
|
||||
"target_model": item.target_model,
|
||||
"response_model": item.provider_response_model(),
|
||||
"has_fallback": item.has_fallback(),
|
||||
});
|
||||
payload["end_to_end_time_ms"] = json!(users_me_usage_metadata_u64(item, "end_to_end_time_ms"));
|
||||
@@ -1865,6 +1867,25 @@ mod tests {
|
||||
assert_eq!(payload["cache_creation_ephemeral_1h_input_tokens"], 6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_usage_payloads_expose_response_model_separately_from_mapping() {
|
||||
let item = StoredRequestUsageAudit {
|
||||
target_model: Some("provider-mapped-model".to_string()),
|
||||
request_metadata: Some(json!({
|
||||
"provider_response_model": "gpt-5.1"
|
||||
})),
|
||||
..sample_usage("completed")
|
||||
};
|
||||
|
||||
let record = build_users_me_usage_record_payload(&item, false, &BTreeMap::new(), false);
|
||||
let active = build_users_me_usage_active_payload(&item);
|
||||
|
||||
for payload in [&record, &active] {
|
||||
assert_eq!(payload["target_model"], "provider-mapped-model");
|
||||
assert_eq!(payload["response_model"], "gpt-5.1");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_usage_payloads_project_end_to_end_timings_from_metadata() {
|
||||
let item = StoredRequestUsageAudit {
|
||||
|
||||
@@ -21,6 +21,7 @@ use aether_crypto::{
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||
use aether_provider_pool::{
|
||||
grok_pool_tier_from_quota_bucket, grok_supported_quota_windows_for_tier,
|
||||
provider_pool_codex_metadata_has_account_quota,
|
||||
};
|
||||
use aether_scheduler_core::provider_key_circuit_payload_is_active_open_at;
|
||||
use serde_json::{json, Map, Value};
|
||||
@@ -1111,9 +1112,8 @@ fn build_codex_quota_status_snapshot(
|
||||
source: &str,
|
||||
) -> Option<Value> {
|
||||
let metadata = provider_quota_metadata_bucket(upstream_metadata, "codex")?;
|
||||
let observed_at_unix_secs = metadata
|
||||
.get("updated_at")
|
||||
.and_then(admin_provider_quota_pure::coerce_json_u64);
|
||||
let observed_at_unix_secs = provider_quota_timestamp_unix_secs(metadata.get("observed_at"))
|
||||
.or_else(|| provider_quota_timestamp_unix_secs(metadata.get("updated_at")));
|
||||
let plan_type = metadata
|
||||
.get("plan_type")
|
||||
.and_then(Value::as_str)
|
||||
@@ -1388,6 +1388,168 @@ fn build_kiro_quota_status_snapshot(
|
||||
}))
|
||||
}
|
||||
|
||||
fn build_xai_quota_status_snapshot(
|
||||
upstream_metadata: Option<&Value>,
|
||||
source: &str,
|
||||
) -> Option<Value> {
|
||||
let metadata = provider_quota_metadata_bucket(upstream_metadata, "xai")?;
|
||||
let observed_at_unix_secs = provider_quota_timestamp_unix_secs(metadata.get("updated_at"));
|
||||
let usage_limit = metadata
|
||||
.get("usage_limit")
|
||||
.and_then(admin_provider_quota_pure::coerce_json_f64);
|
||||
let current_usage = metadata
|
||||
.get("current_usage")
|
||||
.and_then(admin_provider_quota_pure::coerce_json_f64);
|
||||
let remaining = metadata
|
||||
.get("remaining")
|
||||
.and_then(admin_provider_quota_pure::coerce_json_f64);
|
||||
let usage_ratio = metadata
|
||||
.get("usage_percentage")
|
||||
.and_then(admin_provider_quota_pure::coerce_json_f64)
|
||||
.map(|value| (value / 100.0).clamp(0.0, 1.0))
|
||||
.or_else(|| {
|
||||
current_usage
|
||||
.zip(usage_limit)
|
||||
.and_then(|(current_usage, usage_limit)| {
|
||||
(usage_limit > 0.0).then_some((current_usage / usage_limit).clamp(0.0, 1.0))
|
||||
})
|
||||
});
|
||||
let remaining_ratio = usage_ratio.map(|value| (1.0 - value).max(0.0));
|
||||
let next_reset_at = provider_quota_timestamp_unix_secs(metadata.get("next_reset_at"));
|
||||
let reset_seconds = quota_window_reset_seconds(observed_at_unix_secs, next_reset_at);
|
||||
let plan_type = metadata
|
||||
.get("subscription_title")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
let period_type = metadata
|
||||
.get("period_type")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
let usage_label = match period_type.as_deref() {
|
||||
Some("monthly") => "月额度",
|
||||
Some("weekly") => "周额度",
|
||||
_ => "额度",
|
||||
};
|
||||
|
||||
let mut windows = Vec::new();
|
||||
if usage_ratio.is_some()
|
||||
|| remaining.is_some()
|
||||
|| usage_limit.is_some()
|
||||
|| current_usage.is_some()
|
||||
|| next_reset_at.is_some()
|
||||
{
|
||||
windows.push(json!({
|
||||
"code": "usage",
|
||||
"label": usage_label,
|
||||
"scope": "account",
|
||||
"unit": if usage_limit.is_some() { "usd" } else { "percent" },
|
||||
"used_ratio": usage_ratio,
|
||||
"remaining_ratio": remaining_ratio,
|
||||
"used_value": current_usage,
|
||||
"remaining_value": remaining,
|
||||
"limit_value": usage_limit,
|
||||
"reset_at": next_reset_at,
|
||||
"reset_seconds": reset_seconds,
|
||||
}));
|
||||
}
|
||||
|
||||
let prepaid_balance = metadata
|
||||
.get("prepaid_balance")
|
||||
.and_then(admin_provider_quota_pure::coerce_json_f64);
|
||||
if prepaid_balance.is_some_and(|value| value > 0.0) {
|
||||
windows.push(json!({
|
||||
"code": "prepaid",
|
||||
"label": "预付额度",
|
||||
"scope": "account",
|
||||
"unit": "usd",
|
||||
"used_ratio": serde_json::Value::Null,
|
||||
"remaining_ratio": serde_json::Value::Null,
|
||||
"remaining_value": prepaid_balance,
|
||||
"reset_at": serde_json::Value::Null,
|
||||
"reset_seconds": serde_json::Value::Null,
|
||||
}));
|
||||
}
|
||||
|
||||
let on_demand_cap = metadata
|
||||
.get("on_demand_cap")
|
||||
.and_then(admin_provider_quota_pure::coerce_json_f64);
|
||||
let on_demand_used = metadata
|
||||
.get("on_demand_used")
|
||||
.and_then(admin_provider_quota_pure::coerce_json_f64);
|
||||
let on_demand_enabled = metadata
|
||||
.get("on_demand_enabled")
|
||||
.and_then(admin_provider_quota_pure::coerce_json_bool)
|
||||
!= Some(false);
|
||||
if on_demand_enabled && on_demand_cap.is_some_and(|value| value > 0.0) {
|
||||
let on_demand_remaining = on_demand_cap
|
||||
.zip(on_demand_used)
|
||||
.map(|(cap, used)| (cap - used).max(0.0));
|
||||
let on_demand_ratio = on_demand_cap
|
||||
.zip(on_demand_used)
|
||||
.and_then(|(cap, used)| (cap > 0.0).then_some((used / cap).clamp(0.0, 1.0)));
|
||||
windows.push(json!({
|
||||
"code": "on_demand",
|
||||
"label": "按需额度",
|
||||
"scope": "account",
|
||||
"unit": "usd",
|
||||
"used_ratio": on_demand_ratio,
|
||||
"remaining_ratio": on_demand_ratio.map(|value| (1.0 - value).max(0.0)),
|
||||
"used_value": on_demand_used,
|
||||
"remaining_value": on_demand_remaining,
|
||||
"limit_value": on_demand_cap,
|
||||
"reset_at": serde_json::Value::Null,
|
||||
"reset_seconds": serde_json::Value::Null,
|
||||
}));
|
||||
}
|
||||
|
||||
if windows.is_empty() && plan_type.is_none() && observed_at_unix_secs.is_none() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let prepaid_available = prepaid_balance.is_some_and(|value| value > 0.0);
|
||||
let on_demand_available = on_demand_enabled
|
||||
&& on_demand_cap.is_some_and(|value| value > 0.0)
|
||||
&& on_demand_used
|
||||
.zip(on_demand_cap)
|
||||
.is_some_and(|(used, cap)| used < cap);
|
||||
let usage_exhausted = remaining.is_some_and(|value| value <= 0.0)
|
||||
|| usage_ratio.is_some_and(|value| value >= 1.0 - 1e-6);
|
||||
let exhausted = usage_exhausted && !prepaid_available && !on_demand_available;
|
||||
let reason = if exhausted {
|
||||
Some("额度已耗尽".to_string())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let label = if exhausted {
|
||||
Some("额度耗尽")
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let code = if exhausted { "exhausted" } else { "ok" };
|
||||
|
||||
Some(json!({
|
||||
"version": 2,
|
||||
"provider_type": "xai",
|
||||
"code": code,
|
||||
"label": label,
|
||||
"reason": reason,
|
||||
"freshness": "fresh",
|
||||
"source": source,
|
||||
"observed_at": observed_at_unix_secs,
|
||||
"exhausted": exhausted,
|
||||
"usage_ratio": usage_ratio,
|
||||
"updated_at": observed_at_unix_secs,
|
||||
"reset_at": next_reset_at,
|
||||
"reset_seconds": reset_seconds,
|
||||
"plan_type": plan_type,
|
||||
"windows": windows,
|
||||
}))
|
||||
}
|
||||
|
||||
fn build_chatgpt_web_quota_status_snapshot(
|
||||
upstream_metadata: Option<&Value>,
|
||||
source: &str,
|
||||
@@ -2132,6 +2294,111 @@ fn build_gemini_cli_quota_status_snapshot(
|
||||
}))
|
||||
}
|
||||
|
||||
fn build_claude_code_quota_status_snapshot(
|
||||
upstream_metadata: Option<&Value>,
|
||||
source: &str,
|
||||
) -> Option<Value> {
|
||||
let metadata = provider_quota_metadata_bucket(upstream_metadata, "claude_code")?;
|
||||
let observed_at_unix_secs = provider_quota_timestamp_unix_secs(metadata.get("updated_at"));
|
||||
// (metadata prefix, window code, window minutes, account-wide?). Display labels are
|
||||
// resolved by the frontend from `code` so they follow the UI locale.
|
||||
let definitions: [(&str, &str, u64, bool); 4] = [
|
||||
("five_hour", "5h", 300, true),
|
||||
("seven_day", "weekly", 10_080, true),
|
||||
("seven_day_sonnet", "weekly_sonnet", 10_080, false),
|
||||
("seven_day_fable", "weekly_fable", 10_080, false),
|
||||
];
|
||||
let mut windows = Vec::new();
|
||||
for (prefix, code, window_minutes, account_wide) in definitions {
|
||||
let used_percent = metadata
|
||||
.get(&format!("{prefix}_used_percent"))
|
||||
.and_then(Value::as_f64);
|
||||
let reset_at =
|
||||
provider_quota_timestamp_unix_secs(metadata.get(&format!("{prefix}_reset_at")));
|
||||
// A window whose reset time already passed no longer describes current usage.
|
||||
let expired = reset_at
|
||||
.zip(observed_at_unix_secs)
|
||||
.is_some_and(|(reset_at, observed_at)| reset_at <= observed_at);
|
||||
let Some(used_percent) = used_percent else {
|
||||
continue;
|
||||
};
|
||||
let used_ratio = if expired {
|
||||
0.0
|
||||
} else {
|
||||
(used_percent / 100.0).clamp(0.0, 1.0)
|
||||
};
|
||||
let reset_seconds = reset_at
|
||||
.zip(observed_at_unix_secs)
|
||||
.map(|(reset_at, observed_at)| reset_at.saturating_sub(observed_at));
|
||||
let mut window = json!({
|
||||
"code": code,
|
||||
"scope": if account_wide { "account" } else { "model" },
|
||||
"unit": "percent",
|
||||
"used_ratio": used_ratio,
|
||||
"remaining_ratio": 1.0 - used_ratio,
|
||||
"reset_at": reset_at,
|
||||
"reset_seconds": reset_seconds,
|
||||
"window_minutes": window_minutes,
|
||||
"is_exhausted": used_ratio >= 1.0 - 1e-6,
|
||||
});
|
||||
if !account_wide {
|
||||
window["quota_group"] = json!(code);
|
||||
}
|
||||
windows.push(window);
|
||||
}
|
||||
if windows.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let account_windows = windows
|
||||
.iter()
|
||||
.filter(|window| window.get("scope").and_then(Value::as_str) == Some("account"))
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
let blocking_windows = account_windows
|
||||
.iter()
|
||||
.filter(|window| window.get("is_exhausted").and_then(Value::as_bool) == Some(true))
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
let exhausted = !blocking_windows.is_empty();
|
||||
// The account is usable again only once every exhausted window resets.
|
||||
let reset_at = if exhausted {
|
||||
blocking_windows
|
||||
.iter()
|
||||
.filter_map(|window| provider_quota_timestamp_unix_secs(window.get("reset_at")))
|
||||
.max()
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let reset_seconds = if exhausted {
|
||||
blocking_windows
|
||||
.iter()
|
||||
.filter_map(|window| window.get("reset_seconds").and_then(Value::as_u64))
|
||||
.max()
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
Some(json!({
|
||||
"version": 2,
|
||||
"provider_type": "claude_code",
|
||||
"code": if exhausted { "exhausted" } else { "ok" },
|
||||
"freshness": "fresh",
|
||||
"source": source,
|
||||
"observed_at": observed_at_unix_secs,
|
||||
"exhausted": exhausted,
|
||||
"usage_ratio": quota_windows_usage_ratio(&account_windows),
|
||||
"updated_at": observed_at_unix_secs,
|
||||
"reset_at": reset_at,
|
||||
"reset_seconds": reset_seconds,
|
||||
"reset_credits": build_codex_reset_credits_status_snapshot(
|
||||
metadata,
|
||||
observed_at_unix_secs,
|
||||
),
|
||||
"windows": windows,
|
||||
}))
|
||||
}
|
||||
|
||||
fn build_codex_reset_credits_status_snapshot(
|
||||
metadata: &Map<String, Value>,
|
||||
observed_at_unix_secs: Option<u64>,
|
||||
@@ -2255,11 +2522,13 @@ pub(crate) fn sync_provider_key_quota_status_snapshot(
|
||||
let mut quota = match normalized_provider_type.as_str() {
|
||||
"codex" => build_codex_quota_status_snapshot(upstream_metadata, source),
|
||||
"kiro" => build_kiro_quota_status_snapshot(upstream_metadata, source),
|
||||
"xai" => build_xai_quota_status_snapshot(upstream_metadata, source),
|
||||
"chatgpt_web" => build_chatgpt_web_quota_status_snapshot(upstream_metadata, source),
|
||||
"windsurf" => build_windsurf_quota_status_snapshot(upstream_metadata, source),
|
||||
"antigravity" => build_antigravity_quota_status_snapshot(upstream_metadata, source),
|
||||
"grok" => build_grok_quota_status_snapshot(upstream_metadata, source),
|
||||
"gemini_cli" => build_gemini_cli_quota_status_snapshot(upstream_metadata, source),
|
||||
"claude_code" => build_claude_code_quota_status_snapshot(upstream_metadata, source),
|
||||
_ => None,
|
||||
}?;
|
||||
if normalized_provider_type == "codex" {
|
||||
@@ -2339,17 +2608,19 @@ fn codex_upstream_metadata_is_at_least_as_fresh(
|
||||
let Some(metadata) = provider_quota_metadata_bucket(upstream_metadata, "codex") else {
|
||||
return false;
|
||||
};
|
||||
let Some(metadata_updated_at) = metadata
|
||||
.get("updated_at")
|
||||
.and_then(admin_provider_quota_pure::coerce_json_u64)
|
||||
// Identity, reset-credit, and model-only updates do not replace the
|
||||
// account's quota observation, even when their timestamp is newer.
|
||||
if !provider_pool_codex_metadata_has_account_quota(metadata) {
|
||||
return false;
|
||||
}
|
||||
let Some(metadata_updated_at) = provider_quota_timestamp_unix_secs(metadata.get("observed_at"))
|
||||
.or_else(|| provider_quota_timestamp_unix_secs(metadata.get("updated_at")))
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
let snapshot_updated_at = quota_snapshot.and_then(|quota| {
|
||||
quota
|
||||
.get("updated_at")
|
||||
.or_else(|| quota.get("observed_at"))
|
||||
.and_then(admin_provider_quota_pure::coerce_json_u64)
|
||||
provider_quota_timestamp_unix_secs(quota.get("observed_at"))
|
||||
.or_else(|| provider_quota_timestamp_unix_secs(quota.get("updated_at")))
|
||||
});
|
||||
|
||||
snapshot_updated_at.is_none_or(|updated_at| metadata_updated_at >= updated_at)
|
||||
@@ -2448,6 +2719,19 @@ pub(crate) fn provider_key_status_snapshot_payload(
|
||||
let mut snapshot = provider_key_status_snapshot_object(Some(&payload))
|
||||
.or_else(|| default_provider_key_status_snapshot().as_object().cloned())
|
||||
.unwrap_or_default();
|
||||
// Legacy snapshots can retain an exhausted summary after a window reset or
|
||||
// newer quota observation. Use the same decision as scheduling so the
|
||||
// account list and its status filter do not keep displaying that stale block.
|
||||
if provider_type.trim().eq_ignore_ascii_case("codex")
|
||||
&& !aether_provider_pool::provider_pool_key_account_quota_exhausted(key, provider_type)
|
||||
{
|
||||
if let Some(quota) = snapshot.get_mut("quota").and_then(Value::as_object_mut) {
|
||||
quota.insert("exhausted".to_string(), json!(false));
|
||||
if quota.get("code").and_then(Value::as_str) == Some("exhausted") {
|
||||
quota.insert("code".to_string(), json!("ok"));
|
||||
}
|
||||
}
|
||||
}
|
||||
snapshot.insert(
|
||||
"oauth".to_string(),
|
||||
build_provider_key_oauth_status_snapshot(key),
|
||||
@@ -3559,6 +3843,59 @@ mod tests {
|
||||
assert_eq!(window.get("used_value"), Some(&json!(0.0)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_key_status_snapshot_payload_backfills_claude_code_usage_windows() {
|
||||
let mut key = sample_catalog_key();
|
||||
key.upstream_metadata = Some(json!({
|
||||
"claude_code": {
|
||||
"updated_at": 1_800_000_000u64,
|
||||
"five_hour_used_percent": 100.0,
|
||||
"five_hour_reset_at": 1_800_003_600u64,
|
||||
"seven_day_used_percent": 40.0,
|
||||
"seven_day_reset_at": 1_800_400_000u64,
|
||||
"seven_day_sonnet_used_percent": 10.0,
|
||||
"seven_day_sonnet_reset_at": 1_800_400_000u64,
|
||||
"reset_credits": {
|
||||
"available_count": 2,
|
||||
"updated_at": 1_800_000_000u64,
|
||||
"detail_source": "claude_oauth_usage",
|
||||
"credits": [{
|
||||
"display_key": "Key-1",
|
||||
"status": "available",
|
||||
"expires_at": 1_800_144_000u64
|
||||
}]
|
||||
}
|
||||
}
|
||||
}));
|
||||
|
||||
let payload = provider_key_status_snapshot_payload(&key, "claude_code");
|
||||
let quota = payload
|
||||
.get("quota")
|
||||
.and_then(Value::as_object)
|
||||
.expect("quota snapshot should be object");
|
||||
assert_eq!(quota.get("provider_type"), Some(&json!("claude_code")));
|
||||
// An exhausted 5h window blocks the whole account until it resets.
|
||||
assert_eq!(quota.get("exhausted"), Some(&json!(true)));
|
||||
assert_eq!(quota.get("reset_at"), Some(&json!(1_800_003_600u64)));
|
||||
let windows = quota
|
||||
.get("windows")
|
||||
.and_then(Value::as_array)
|
||||
.expect("windows should exist");
|
||||
assert_eq!(windows.len(), 3);
|
||||
assert_eq!(windows[0]["code"], json!("5h"));
|
||||
assert_eq!(windows[0]["scope"], json!("account"));
|
||||
assert_eq!(windows[0]["window_minutes"], json!(300));
|
||||
assert_eq!(windows[1]["code"], json!("weekly"));
|
||||
assert_eq!(windows[1]["used_ratio"], json!(0.4));
|
||||
assert_eq!(windows[2]["code"], json!("weekly_sonnet"));
|
||||
assert_eq!(windows[2]["scope"], json!("model"));
|
||||
assert_eq!(quota["reset_credits"]["available_count"], json!(2));
|
||||
assert_eq!(
|
||||
quota["reset_credits"]["credits"][0]["remaining_seconds"],
|
||||
json!(144_000u64)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_key_status_snapshot_payload_backfills_grok_model_quota() {
|
||||
let mut key = sample_catalog_key();
|
||||
@@ -3622,6 +3959,43 @@ mod tests {
|
||||
assert_eq!(auto.get("used_value"), Some(&json!(90.0)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_key_status_snapshot_payload_backfills_xai_weekly_credits() {
|
||||
let mut key = sample_catalog_key();
|
||||
key.upstream_metadata = Some(json!({
|
||||
"xai": {
|
||||
"updated_at": 1_778_067_246u64,
|
||||
"usage_percentage": 46.0,
|
||||
"period_type": "weekly",
|
||||
"next_reset_at": 1_778_157_172u64,
|
||||
"subscription_title": "SuperGrok",
|
||||
"prepaid_balance": 0.0,
|
||||
"on_demand_cap": 0.0,
|
||||
"on_demand_used": 0.0
|
||||
}
|
||||
}));
|
||||
|
||||
let payload = provider_key_status_snapshot_payload(&key, "xai");
|
||||
let quota = payload
|
||||
.get("quota")
|
||||
.and_then(Value::as_object)
|
||||
.expect("quota snapshot should be object");
|
||||
let windows = quota
|
||||
.get("windows")
|
||||
.and_then(Value::as_array)
|
||||
.expect("xai quota windows should exist");
|
||||
|
||||
assert_eq!(quota.get("provider_type"), Some(&json!("xai")));
|
||||
assert_eq!(quota.get("code"), Some(&json!("ok")));
|
||||
assert_eq!(quota.get("exhausted"), Some(&json!(false)));
|
||||
assert_eq!(quota.get("plan_type"), Some(&json!("SuperGrok")));
|
||||
assert_eq!(quota.get("usage_ratio"), Some(&json!(0.46)));
|
||||
assert_eq!(quota.get("reset_at"), Some(&json!(1_778_157_172u64)));
|
||||
assert_eq!(windows.len(), 1);
|
||||
assert_eq!(windows[0].get("code"), Some(&json!("usage")));
|
||||
assert_eq!(windows[0].get("label"), Some(&json!("周额度")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_key_status_snapshot_payload_backfills_gemini_cli_account_credits() {
|
||||
let mut key = sample_catalog_key();
|
||||
@@ -4013,6 +4387,116 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_key_status_snapshot_payload_clears_stale_codex_exhaustion_summary() {
|
||||
let mut key = sample_catalog_key();
|
||||
key.status_snapshot = Some(json!({
|
||||
"quota": {
|
||||
"provider_type": "codex",
|
||||
"code": "exhausted",
|
||||
"exhausted": true,
|
||||
"windows": [{
|
||||
"code": "weekly",
|
||||
"scope": "account",
|
||||
"used_ratio": 0.83,
|
||||
"remaining_ratio": 0.17,
|
||||
"reset_at": 4_102_444_800u64
|
||||
}]
|
||||
}
|
||||
}));
|
||||
let payload = provider_key_status_snapshot_payload(&key, "codex");
|
||||
assert_eq!(payload["quota"]["code"], "ok");
|
||||
assert_eq!(payload["quota"]["exhausted"], false);
|
||||
assert!(payload["quota"]["label"].is_null());
|
||||
|
||||
// An explicit current upstream refusal is not a stale percentage summary.
|
||||
key.status_snapshot.as_mut().unwrap()["quota"]["allowed"] = json!(false);
|
||||
let payload = provider_key_status_snapshot_payload(&key, "codex");
|
||||
assert_eq!(payload["quota"]["code"], "exhausted");
|
||||
assert_eq!(payload["quota"]["exhausted"], true);
|
||||
|
||||
// Missing capacity evidence must not clear an exhausted summary either.
|
||||
key.status_snapshot = Some(json!({"quota": {
|
||||
"provider_type": "codex", "code": "exhausted", "exhausted": true
|
||||
}}));
|
||||
let payload = provider_key_status_snapshot_payload(&key, "codex");
|
||||
assert_eq!(payload["quota"]["code"], "exhausted");
|
||||
assert_eq!(payload["quota"]["exhausted"], true);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_key_status_snapshot_payload_refreshes_codex_timestamp_formats() {
|
||||
for updated_at in [
|
||||
json!(1_900_000_000u64),
|
||||
json!(1_900_000_000_000u64),
|
||||
json!("2030-03-17T17:46:40Z"),
|
||||
] {
|
||||
let mut key = sample_catalog_key();
|
||||
key.upstream_metadata = Some(json!({
|
||||
"codex": {
|
||||
"updated_at": updated_at,
|
||||
"primary_used_percent": 83.0,
|
||||
"primary_reset_at": 4_102_444_800u64
|
||||
}
|
||||
}));
|
||||
key.status_snapshot = Some(json!({
|
||||
"quota": {
|
||||
"provider_type": "codex",
|
||||
"updated_at": 1_899_999_000u64,
|
||||
"code": "exhausted",
|
||||
"exhausted": true,
|
||||
"allowed": false,
|
||||
"windows": [{
|
||||
"code": "weekly",
|
||||
"scope": "account",
|
||||
"used_ratio": 1.0,
|
||||
"remaining_ratio": 0.0,
|
||||
"reset_at": 4_102_444_800u64
|
||||
}]
|
||||
}
|
||||
}));
|
||||
let payload = provider_key_status_snapshot_payload(&key, "codex");
|
||||
assert_eq!(payload["quota"]["code"], "ok");
|
||||
assert_eq!(payload["quota"]["updated_at"], 1_900_000_000u64);
|
||||
assert_eq!(payload["quota"]["windows"][0]["used_ratio"], 0.83);
|
||||
assert!(payload["quota"]["allowed"].is_null());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_key_status_snapshot_payload_preserves_codex_account_quota_on_unrelated_updates() {
|
||||
for patch in [
|
||||
json!({"plan_type": "pro"}),
|
||||
json!({"spark_primary_used_percent": 83.0}),
|
||||
json!({"credits_unlimited": false}),
|
||||
json!({"windows": [{"code": "weekly", "reset_at": 4_102_444_800u64}]}),
|
||||
] {
|
||||
let mut metadata = patch;
|
||||
metadata["updated_at"] = json!("2030-03-17T17:46:40Z");
|
||||
let mut key = sample_catalog_key();
|
||||
key.upstream_metadata = Some(json!({"codex": metadata}));
|
||||
key.status_snapshot = Some(json!({"quota": {
|
||||
"provider_type": "codex", "updated_at": 200,
|
||||
"code": "exhausted", "exhausted": true, "allowed": false,
|
||||
"windows": [{"code": "weekly", "scope": "account", "used_ratio": 1.0,
|
||||
"remaining_ratio": 0.0, "reset_at": 4_102_444_800u64}]
|
||||
}}));
|
||||
let payload = provider_key_status_snapshot_payload(&key, "codex");
|
||||
assert_eq!(payload["quota"]["code"], "exhausted", "{metadata}");
|
||||
assert_eq!(payload["quota"]["exhausted"], true, "{metadata}");
|
||||
assert_eq!(payload["quota"]["allowed"], false, "{metadata}");
|
||||
assert_eq!(payload["quota"]["updated_at"], 200, "{metadata}");
|
||||
assert_eq!(
|
||||
payload["quota"]["windows"][0]["code"], "weekly",
|
||||
"{metadata}"
|
||||
);
|
||||
assert_eq!(
|
||||
payload["quota"]["windows"][0]["used_ratio"], 1.0,
|
||||
"{metadata}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_key_status_snapshot_payload_restores_complete_codex_cache() {
|
||||
let mut key = sample_catalog_key();
|
||||
|
||||
@@ -276,6 +276,10 @@ pub(crate) fn admin_proxy_local_requires_buffered_body(
|
||||
| (Some("system_manage"), http::Method::POST, Some("config_import"))
|
||||
| (Some("system_manage"), http::Method::POST, Some("users_import"))
|
||||
| (Some("system_manage"), http::Method::POST, Some("data_import"))
|
||||
| (Some("system_manage"), http::Method::POST, Some("cleanup_usage_manual"))
|
||||
| (Some("system_manage"), http::Method::POST, Some("smtp_test"))
|
||||
| (Some("system_manage"), http::Method::POST, Some("prepare_update"))
|
||||
| (Some("system_manage"), http::Method::POST, Some("apply_update"))
|
||||
| (Some("system_manage"), http::Method::PUT, Some("settings_set"))
|
||||
| (Some("system_manage"), http::Method::PUT, Some("config_set"))
|
||||
| (Some("system_manage"), http::Method::PUT, Some("email_template_set"))
|
||||
@@ -611,4 +615,27 @@ mod tests {
|
||||
"/v1/chat/completions?key=passthrough"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn manual_cleanup_route_requires_buffered_body() {
|
||||
use crate::control::GatewayPublicRequestContext;
|
||||
let mut decision = GatewayControlDecision::synthetic(
|
||||
"/api/admin/system/cleanup/usage/manual",
|
||||
Some("admin_proxy".to_string()),
|
||||
Some("system_manage".to_string()),
|
||||
Some("cleanup_usage_manual".to_string()),
|
||||
Some("system_manage:cleanup_usage_manual".to_string()),
|
||||
);
|
||||
decision.route_class = Some("admin_proxy".to_string());
|
||||
let uri: http::Uri = "/api/admin/system/cleanup/usage/manual".parse().unwrap();
|
||||
let headers = http::HeaderMap::new();
|
||||
let context = GatewayPublicRequestContext::from_request_parts(
|
||||
"trace-manual-cleanup",
|
||||
&http::Method::POST,
|
||||
&uri,
|
||||
&headers,
|
||||
Some(decision),
|
||||
);
|
||||
assert!(super::admin_proxy_local_requires_buffered_body(&context));
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user