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:
elky
2026-10-02 11:57:18 +08:00
343 changed files with 27929 additions and 2549 deletions
@@ -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));
}
}