i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-09 12:40:20 +08:00
Code Issues Packages Projects Releases Wiki Activity

Preserve usage data in system imports

Browse Source
This commit is contained in:
fawney19
2026-05-24 21:40:48 +08:00
parent 18d9004f22
commit e2d5fc9dfb
21 changed files with 2683 additions and 309 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/data/state/auth.rs
+22
View File
@@ -1684,6 +1684,28 @@ impl GatewayDataState {
}
}
pub(crate) async fn set_api_key_usage_totals(
&self,
api_key_id: &str,
total_requests: u64,
total_tokens: u64,
total_cost_usd: f64,
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
match &self.auth_api_key_writer {
Some(repository) => {
repository
.set_api_key_usage_totals(
api_key_id,
total_requests,
total_tokens,
total_cost_usd,
)
.await
}
None => Ok(None),
}
}
pub(crate) async fn set_standalone_api_key_feature_settings(
&self,
api_key_id: &str,
apps/aether-gateway/src/data/state/core.rs
+39
View File
@@ -511,6 +511,45 @@ impl GatewayDataState {
}
}
pub(crate) async fn export_admin_system_usage_aggregates(
&self,
) -> Result<aether_data::repository::system::AdminSystemUsageAggregateSnapshot, DataLayerError>
{
match self.backends.as_ref() {
Some(backends) => backends.export_admin_system_usage_aggregates().await,
None => {
Ok(aether_data::repository::system::AdminSystemUsageAggregateSnapshot::default())
}
}
}
pub(crate) async fn import_admin_system_usage_aggregates(
&self,
snapshot: &aether_data::repository::system::AdminSystemUsageAggregateSnapshot,
user_id_map: &std::collections::BTreeMap<String, String>,
api_key_id_map: &std::collections::BTreeMap<String, String>,
mode: aether_data::repository::system::AdminSystemUsageAggregateImportMode,
) -> Result<
aether_data::repository::system::AdminSystemUsageAggregateImportSummary,
DataLayerError,
> {
match self.backends.as_ref() {
Some(backends) => {
backends
.import_admin_system_usage_aggregates(
snapshot,
user_id_map,
api_key_id_map,
mode,
)
.await
}
None => Ok(
aether_data::repository::system::AdminSystemUsageAggregateImportSummary::default(),
),
}
}
pub(crate) async fn purge_admin_request_bodies_batch(
&self,
batch_size: usize,
apps/aether-gateway/src/handlers/admin/request/system/export.rs
+4
View File
@@ -185,6 +185,7 @@ impl<'a> AdminAppState<'a> {
let standalone_wallets = self
.list_wallet_snapshots_by_api_key_ids(&standalone_api_key_ids)
.await?;
let usage_aggregates = self.export_admin_system_usage_aggregates().await?;
let wallets_by_user_id = user_wallets
.into_iter()
@@ -262,6 +263,7 @@ impl<'a> AdminAppState<'a> {
.collect::<Vec<_>>();
json!({
"id": user.id.clone(),
"email": user.email.clone(),
"email_verified": user.email_verified,
"username": user.username.clone(),
@@ -306,6 +308,7 @@ impl<'a> AdminAppState<'a> {
"user_groups": user_groups_data,
"users": users_data,
"standalone_keys": standalone_keys_data,
"usage_aggregates": usage_aggregates,
}))
}
@@ -330,6 +333,7 @@ impl<'a> AdminAppState<'a> {
include_is_standalone: bool,
) -> serde_json::Value {
let mut payload = serde_json::Map::from_iter([
("api_key_id".to_string(), json!(key.api_key_id.clone())),
("key_hash".to_string(), json!(key.key_hash.clone())),
("name".to_string(), json!(key.name.clone())),
(
apps/aether-gateway/src/handlers/admin/request/system/import.rs
+332 -263
View File
@@ -38,6 +38,10 @@ use aether_data::repository::auth_modules::StoredLdapModuleConfig;
use aether_data::repository::oauth_providers::{
EncryptedSecretUpdate, UpsertOAuthProviderConfigRecord,
};
use aether_data::repository::system::{
AdminSystemUsageAggregateImportMode, AdminSystemUsageAggregateImportSummary,
AdminSystemUsageAggregateSnapshot,
};
use aether_data::repository::wallet::WalletLookupKey;
use aether_data_contracts::repository::global_models::{
AdminGlobalModelListQuery, AdminProviderModelListQuery, CreateAdminGlobalModelRecord,
@@ -344,54 +348,6 @@ fn imported_oauth_expires_at_unix_secs(normalized_auth_config: Option<&Value>) -
None
}
fn imported_oauth_has_refresh_token(normalized_auth_config: Option<&Value>) -> bool {
normalized_auth_config
.and_then(Value::as_object)
.and_then(|object| object.get("refresh_token"))
.and_then(Value::as_str)
.map(str::trim)
.is_some_and(|value| !value.is_empty())
}
async fn refresh_imported_oauth_key_after_persist(
state: &AdminAppState<'_>,
provider: &aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider,
key_id: &str,
) -> Result<(), GatewayError> {
let Some(endpoint) =
crate::handlers::admin::provider::oauth::runtime::resolve_provider_oauth_runtime_endpoints(
state,
provider,
provider.provider_type.as_str(),
)
.await?
.runtime_endpoint
else {
return Ok(());
};
let Some(transport) = state
.read_provider_transport_snapshot(&provider.id, &endpoint.id, key_id)
.await?
else {
return Ok(());
};
if !crate::provider_transport::supports_local_oauth_request_auth_resolution(&transport) {
return Ok(());
}
if let Err(error) = state.force_local_oauth_refresh_entry(&transport).await {
tracing::warn!(
provider_id = %provider.id,
provider_type = %provider.provider_type,
key_id = %key_id,
error = ?error,
"admin system import oauth refresh after credential import failed"
);
}
Ok(())
}
fn build_import_provider_model_record(
provider_id: &str,
existing_id: Option<&str>,
@@ -435,6 +391,8 @@ struct AdminSystemUsersImportStats {
users: AdminSystemConfigImportCounter,
api_keys: AdminSystemConfigImportCounter,
standalone_keys: AdminSystemConfigImportCounter,
#[serde(skip_serializing_if = "Option::is_none")]
usage_aggregates: Option<AdminSystemUsageAggregateImportSummary>,
errors: Vec<String>,
}
@@ -488,6 +446,16 @@ fn validate_imported_system_users_export_version(version: Option<&Value>) -> Res
Ok(())
}
fn usage_aggregate_import_mode(
merge_mode: AdminImportMergeMode,
) -> AdminSystemUsageAggregateImportMode {
match merge_mode {
AdminImportMergeMode::Skip => AdminSystemUsageAggregateImportMode::Skip,
AdminImportMergeMode::Overwrite => AdminSystemUsageAggregateImportMode::Overwrite,
AdminImportMergeMode::Error => AdminSystemUsageAggregateImportMode::Error,
}
}
fn imported_object_field<'a>(
value: &'a Value,
field_name: &str,
@@ -1560,16 +1528,6 @@ impl<'a> AdminAppState<'a> {
"更新 Provider '{provider_name}' 的 Key 失败"
))));
};
if auth_type == "oauth"
&& imported_oauth_has_refresh_token(normalized_auth_config.as_ref())
{
refresh_imported_oauth_key_after_persist(
self,
&provider,
&persisted.id,
)
.await?;
}
existing_keys[existing_index] = persisted;
stats.keys.updated += 1;
}
@@ -1610,11 +1568,6 @@ impl<'a> AdminAppState<'a> {
"创建 Provider '{provider_name}' 的 Key 失败"
))));
};
if auth_type == "oauth"
&& imported_oauth_has_refresh_token(normalized_auth_config.as_ref())
{
refresh_imported_oauth_key_after_persist(self, &provider, &created.id).await?;
}
existing_keys.push(created);
stats.keys.created += 1;
}
@@ -2052,6 +2005,8 @@ impl<'a> AdminAppState<'a> {
));
let mut stats = AdminSystemUsersImportStats::default();
let mut imported_user_id_map = BTreeMap::<String, String>::new();
let mut imported_api_key_id_map = BTreeMap::<String, String>::new();
let default_group_id = self.effective_default_user_group_id().await?;
let existing_groups = self.list_user_groups().await?;
let mut groups_by_name = existing_groups
@@ -2138,6 +2093,7 @@ impl<'a> AdminAppState<'a> {
Ok(value) => value,
Err(detail) => return Ok(Err(invalid_request(detail))),
};
let source_user_id = invalid_value!(imported_optional_string(user.get("id")));
let role = invalid_value!(imported_optional_string(user.get("role")))
.unwrap_or_else(|| "user".to_string())
.to_ascii_lowercase();
@@ -2471,6 +2427,9 @@ impl<'a> AdminAppState<'a> {
stats.users.created += 1;
created.id
};
if let Some(source_user_id) = source_user_id {
imported_user_id_map.insert(source_user_id, user_id.clone());
}
let existing_api_keys = self
.list_auth_api_key_export_records_by_user_ids(std::slice::from_ref(&user_id))
@@ -2506,6 +2465,8 @@ impl<'a> AdminAppState<'a> {
));
continue;
};
let source_api_key_id =
invalid_value!(imported_optional_string(key.get("api_key_id")));
let name = invalid_value!(imported_optional_string(key.get("name")));
let allowed_providers = invalid_value!(normalize_imported_user_string_list(
key,
@@ -2538,21 +2499,21 @@ impl<'a> AdminAppState<'a> {
let auto_delete_on_expiry =
invalid_value!(imported_optional_bool(key.get("auto_delete_on_expiry")))
.unwrap_or(false);
let total_requests = invalid_value!(imported_optional_u64(
let imported_total_requests = invalid_value!(imported_optional_u64(
key.get("total_requests"),
"total_requests"
))
.unwrap_or(0);
let total_tokens = invalid_value!(imported_optional_u64(
));
let total_requests = imported_total_requests.unwrap_or(0);
let imported_total_tokens = invalid_value!(imported_optional_u64(
key.get("total_tokens"),
"total_tokens"
))
.unwrap_or(0);
let total_cost_usd = invalid_value!(imported_optional_f64(
));
let total_tokens = imported_total_tokens.unwrap_or(0);
let imported_total_cost_usd = invalid_value!(imported_optional_f64(
key.get("total_cost_usd"),
"total_cost_usd"
))
.unwrap_or(0.0);
));
let total_cost_usd = imported_total_cost_usd.unwrap_or(0.0);
let feature_settings = invalid_value!(imported_optional_json_object(
key.get("feature_settings"),
"feature_settings"
@@ -2624,13 +2585,31 @@ impl<'a> AdminAppState<'a> {
is_active,
)
.await?;
if imported_total_requests.is_some()
|| imported_total_tokens.is_some()
|| imported_total_cost_usd.is_some()
{
let updated_usage = self
.set_api_key_usage_totals(
&existing_key.api_key_id,
imported_total_requests
.unwrap_or(existing_key.total_requests),
imported_total_tokens.unwrap_or(existing_key.total_tokens),
imported_total_cost_usd
.unwrap_or(existing_key.total_cost_usd),
)
.await?;
if updated_usage.is_none() {
return Ok(Err((
http::StatusCode::SERVICE_UNAVAILABLE,
json!({ "detail": "Admin system data unavailable" }),
)));
}
}
if key.contains_key("allowed_api_formats")
|| key.contains_key("allowed_models")
|| key.contains_key("expires_at")
|| key.contains_key("auto_delete_on_expiry")
|| key.contains_key("total_requests")
|| key.contains_key("total_tokens")
|| key.contains_key("total_cost_usd")
{
stats.errors.push(format!(
"用户 '{}' 的现有 API Key 仅覆盖基础字段；高级导入字段保持原值",
@@ -2638,6 +2617,10 @@ impl<'a> AdminAppState<'a> {
));
}
stats.api_keys.updated += 1;
if let Some(source_api_key_id) = source_api_key_id.clone() {
imported_api_key_id_map
.insert(source_api_key_id, existing_key.api_key_id.clone());
}
}
}
continue;
@@ -2680,125 +2663,137 @@ impl<'a> AdminAppState<'a> {
)
.await?;
}
let created_api_key_id = created.api_key_id.clone();
existing_api_keys_by_hash.insert(key_hash, created);
if let Some(source_api_key_id) = source_api_key_id {
imported_api_key_id_map.insert(source_api_key_id, created_api_key_id);
}
stats.api_keys.created += 1;
}
}
if standalone_keys.is_empty() {
return Ok(Ok(json!({
"message": "用户数据导入成功",
"stats": stats,
})));
}
let Some(standalone_owner_id) = standalone_owner_id else {
stats.standalone_keys.skipped += standalone_keys.len() as u64;
stats
.errors
.push("无法导入独立余额 Key: 当前管理员用户记录不存在".to_string());
return Ok(Ok(json!({
"message": "用户数据导入成功",
"stats": stats,
})));
};
let existing_standalone_keys = self
.list_auth_api_key_export_standalone_records()
.await?
.into_iter()
.collect::<Vec<_>>();
let mut existing_standalone_by_hash = existing_standalone_keys
.into_iter()
.map(|record| (record.key_hash.clone(), record))
.collect::<BTreeMap<_, _>>();
for (index, raw_key) in standalone_keys.iter().enumerate() {
let key = match imported_object_field(raw_key, &format!("standalone_keys[{index}]")) {
Ok(value) => value,
Err(detail) => return Ok(Err(invalid_request(detail))),
};
let Some((key_hash, key_encrypted)) =
invalid_value!(self.resolve_imported_system_user_api_key_material(key))
else {
stats.standalone_keys.skipped += 1;
if !standalone_keys.is_empty() {
let Some(standalone_owner_id) = standalone_owner_id else {
stats.standalone_keys.skipped += standalone_keys.len() as u64;
stats
.errors
.push(format!("跳过无效独立余额 Key: standalone_keys[{index}]"));
continue;
.push("无法导入独立余额 Key: 当前管理员用户记录不存在".to_string());
if let Some(summary) = self
.import_admin_system_user_usage_aggregates(
root.get("usage_aggregates"),
&imported_user_id_map,
&imported_api_key_id_map,
merge_mode,
)
.await?
{
stats.usage_aggregates = Some(summary);
}
return Ok(Ok(json!({
"message": "用户数据导入成功",
"stats": stats,
})));
};
let name = invalid_value!(imported_optional_string(key.get("name")));
let allowed_providers = invalid_value!(normalize_imported_user_string_list(
key,
"allowed_providers"
));
let allowed_api_formats = invalid_value!(normalize_imported_user_api_formats(
key,
"allowed_api_formats"
));
let allowed_models =
invalid_value!(normalize_imported_user_string_list(key, "allowed_models"));
let ip_rules = invalid_value!(normalize_imported_user_ip_rules(key));
let rate_limit =
invalid_value!(imported_optional_i32(key.get("rate_limit"), "rate_limit"))
.unwrap_or(0);
let concurrent_limit = invalid_value!(imported_optional_i32(
key.get("concurrent_limit"),
"concurrent_limit"
));
if concurrent_limit.is_some_and(|value| value < 0) {
return Ok(Err(invalid_request("concurrent_limit 必须是非负整数")));
}
let force_capabilities = imported_optional_value(key.get("force_capabilities"));
let is_active =
invalid_value!(imported_optional_bool(key.get("is_active"))).unwrap_or(true);
let expires_at_unix_secs = invalid_value!(imported_rfc3339_to_unix_secs(
key.get("expires_at"),
"expires_at"
));
let auto_delete_on_expiry =
invalid_value!(imported_optional_bool(key.get("auto_delete_on_expiry")))
.unwrap_or(false);
let total_requests = invalid_value!(imported_optional_u64(
key.get("total_requests"),
"total_requests"
))
.unwrap_or(0);
let total_tokens = invalid_value!(imported_optional_u64(
key.get("total_tokens"),
"total_tokens"
))
.unwrap_or(0);
let total_cost_usd = invalid_value!(imported_optional_f64(
key.get("total_cost_usd"),
"total_cost_usd"
))
.unwrap_or(0.0);
let feature_settings = invalid_value!(imported_optional_json_object(
key.get("feature_settings"),
"feature_settings"
)
.and_then(normalize_admin_feature_settings));
let wallet_payload = match key.get("wallet") {
Some(Value::Object(map)) => Some(map),
Some(Value::Null) | None => None,
Some(_) => return Ok(Err(invalid_request("wallet 必须是对象"))),
};
let unlimited =
invalid_value!(imported_optional_bool(key.get("unlimited"))).unwrap_or(false);
let wallet_target =
invalid_value!(normalize_imported_wallet_target(wallet_payload, unlimited));
if let Some(existing_key) = existing_standalone_by_hash.get(&key_hash).cloned() {
match merge_mode {
AdminImportMergeMode::Skip => {
stats.standalone_keys.skipped += 1;
}
AdminImportMergeMode::Error => {
return Ok(Err(invalid_request("独立余额 Key 已存在")));
}
AdminImportMergeMode::Overwrite => {
let updated = self
let existing_standalone_keys = self
.list_auth_api_key_export_standalone_records()
.await?
.into_iter()
.collect::<Vec<_>>();
let mut existing_standalone_by_hash = existing_standalone_keys
.into_iter()
.map(|record| (record.key_hash.clone(), record))
.collect::<BTreeMap<_, _>>();
for (index, raw_key) in standalone_keys.iter().enumerate() {
let key = match imported_object_field(raw_key, &format!("standalone_keys[{index}]"))
{
Ok(value) => value,
Err(detail) => return Ok(Err(invalid_request(detail))),
};
let Some((key_hash, key_encrypted)) =
invalid_value!(self.resolve_imported_system_user_api_key_material(key))
else {
stats.standalone_keys.skipped += 1;
stats
.errors
.push(format!("跳过无效独立余额 Key: standalone_keys[{index}]"));
continue;
};
let source_api_key_id =
invalid_value!(imported_optional_string(key.get("api_key_id")));
let name = invalid_value!(imported_optional_string(key.get("name")));
let allowed_providers = invalid_value!(normalize_imported_user_string_list(
key,
"allowed_providers"
));
let allowed_api_formats = invalid_value!(normalize_imported_user_api_formats(
key,
"allowed_api_formats"
));
let allowed_models =
invalid_value!(normalize_imported_user_string_list(key, "allowed_models"));
let ip_rules = invalid_value!(normalize_imported_user_ip_rules(key));
let rate_limit =
invalid_value!(imported_optional_i32(key.get("rate_limit"), "rate_limit"))
.unwrap_or(0);
let concurrent_limit = invalid_value!(imported_optional_i32(
key.get("concurrent_limit"),
"concurrent_limit"
));
if concurrent_limit.is_some_and(|value| value < 0) {
return Ok(Err(invalid_request("concurrent_limit 必须是非负整数")));
}
let force_capabilities = imported_optional_value(key.get("force_capabilities"));
let is_active =
invalid_value!(imported_optional_bool(key.get("is_active"))).unwrap_or(true);
let expires_at_unix_secs = invalid_value!(imported_rfc3339_to_unix_secs(
key.get("expires_at"),
"expires_at"
));
let auto_delete_on_expiry =
invalid_value!(imported_optional_bool(key.get("auto_delete_on_expiry")))
.unwrap_or(false);
let imported_total_requests = invalid_value!(imported_optional_u64(
key.get("total_requests"),
"total_requests"
));
let total_requests = imported_total_requests.unwrap_or(0);
let imported_total_tokens = invalid_value!(imported_optional_u64(
key.get("total_tokens"),
"total_tokens"
));
let total_tokens = imported_total_tokens.unwrap_or(0);
let imported_total_cost_usd = invalid_value!(imported_optional_f64(
key.get("total_cost_usd"),
"total_cost_usd"
));
let total_cost_usd = imported_total_cost_usd.unwrap_or(0.0);
let feature_settings = invalid_value!(imported_optional_json_object(
key.get("feature_settings"),
"feature_settings"
)
.and_then(normalize_admin_feature_settings));
let wallet_payload = match key.get("wallet") {
Some(Value::Object(map)) => Some(map),
Some(Value::Null) | None => None,
Some(_) => return Ok(Err(invalid_request("wallet 必须是对象"))),
};
let unlimited =
invalid_value!(imported_optional_bool(key.get("unlimited"))).unwrap_or(false);
let wallet_target =
invalid_value!(normalize_imported_wallet_target(wallet_payload, unlimited));
if let Some(existing_key) = existing_standalone_by_hash.get(&key_hash).cloned() {
match merge_mode {
AdminImportMergeMode::Skip => {
stats.standalone_keys.skipped += 1;
}
AdminImportMergeMode::Error => {
return Ok(Err(invalid_request("独立余额 Key 已存在")));
}
AdminImportMergeMode::Overwrite => {
let updated = self
.update_standalone_api_key_basic(
aether_data::repository::auth::UpdateStandaloneApiKeyBasicRecord {
api_key_id: existing_key.api_key_id.clone(),
@@ -2819,94 +2814,134 @@ impl<'a> AdminAppState<'a> {
},
)
.await?;
if updated.is_none() {
return Ok(Err((
http::StatusCode::SERVICE_UNAVAILABLE,
json!({ "detail": "Admin system data unavailable" }),
)));
}
let _ = self
.set_standalone_api_key_active(&existing_key.api_key_id, is_active)
.await?;
if key.contains_key("feature_settings") {
if updated.is_none() {
return Ok(Err((
http::StatusCode::SERVICE_UNAVAILABLE,
json!({ "detail": "Admin system data unavailable" }),
)));
}
let _ = self
.set_standalone_api_key_feature_settings(
&existing_key.api_key_id,
feature_settings.clone(),
)
.set_standalone_api_key_active(&existing_key.api_key_id, is_active)
.await?;
if key.contains_key("feature_settings") {
let _ = self
.set_standalone_api_key_feature_settings(
&existing_key.api_key_id,
feature_settings.clone(),
)
.await?;
}
if imported_total_requests.is_some()
|| imported_total_tokens.is_some()
|| imported_total_cost_usd.is_some()
{
let updated_usage = self
.set_api_key_usage_totals(
&existing_key.api_key_id,
imported_total_requests
.unwrap_or(existing_key.total_requests),
imported_total_tokens.unwrap_or(existing_key.total_tokens),
imported_total_cost_usd
.unwrap_or(existing_key.total_cost_usd),
)
.await?;
if updated_usage.is_none() {
return Ok(Err((
http::StatusCode::SERVICE_UNAVAILABLE,
json!({ "detail": "Admin system data unavailable" }),
)));
}
}
if key.contains_key("expires_at")
|| key.contains_key("auto_delete_on_expiry")
|| key.contains_key("force_capabilities")
{
stats.errors.push(
"现有独立余额 Key 仅覆盖基础字段；高级导入字段保持原值"
.to_string(),
);
}
self.sync_imported_api_key_wallet(
&existing_key.api_key_id,
&wallet_target,
key.get("name")
.and_then(Value::as_str)
.unwrap_or("独立余额 Key"),
)
.await?;
stats.standalone_keys.updated += 1;
if let Some(source_api_key_id) = source_api_key_id.clone() {
imported_api_key_id_map
.insert(source_api_key_id, existing_key.api_key_id.clone());
}
}
if key.contains_key("expires_at")
|| key.contains_key("auto_delete_on_expiry")
|| key.contains_key("force_capabilities")
|| key.contains_key("total_requests")
|| key.contains_key("total_tokens")
|| key.contains_key("total_cost_usd")
{
stats.errors.push(
"现有独立余额 Key 仅覆盖基础字段；高级导入字段保持原值".to_string(),
);
}
self.sync_imported_api_key_wallet(
&existing_key.api_key_id,
&wallet_target,
key.get("name")
.and_then(Value::as_str)
.unwrap_or("独立余额 Key"),
)
.await?;
stats.standalone_keys.updated += 1;
}
continue;
}
continue;
}
let created = self
.create_standalone_api_key(
aether_data::repository::auth::CreateStandaloneApiKeyRecord {
user_id: standalone_owner_id.clone(),
api_key_id: Uuid::new_v4().to_string(),
key_hash: key_hash.clone(),
key_encrypted,
name,
allowed_providers,
allowed_api_formats,
allowed_models,
ip_rules,
rate_limit: Some(rate_limit),
concurrent_limit,
force_capabilities,
is_active,
expires_at_unix_secs,
auto_delete_on_expiry,
total_requests,
total_tokens,
total_cost_usd,
},
)
.await?;
let Some(created) = created else {
return Ok(Err((
http::StatusCode::SERVICE_UNAVAILABLE,
json!({ "detail": "Admin system data unavailable" }),
)));
};
if key.contains_key("feature_settings") {
let _ = self
.set_standalone_api_key_feature_settings(
&created.api_key_id,
feature_settings.clone(),
let created = self
.create_standalone_api_key(
aether_data::repository::auth::CreateStandaloneApiKeyRecord {
user_id: standalone_owner_id.clone(),
api_key_id: Uuid::new_v4().to_string(),
key_hash: key_hash.clone(),
key_encrypted,
name,
allowed_providers,
allowed_api_formats,
allowed_models,
ip_rules,
rate_limit: Some(rate_limit),
concurrent_limit,
force_capabilities,
is_active,
expires_at_unix_secs,
auto_delete_on_expiry,
total_requests,
total_tokens,
total_cost_usd,
},
)
.await?;
let Some(created) = created else {
return Ok(Err((
http::StatusCode::SERVICE_UNAVAILABLE,
json!({ "detail": "Admin system data unavailable" }),
)));
};
if key.contains_key("feature_settings") {
let _ = self
.set_standalone_api_key_feature_settings(
&created.api_key_id,
feature_settings.clone(),
)
.await?;
}
self.sync_imported_api_key_wallet(
&created.api_key_id,
&wallet_target,
created.name.as_deref().unwrap_or("独立余额 Key"),
)
.await?;
let created_api_key_id = created.api_key_id.clone();
existing_standalone_by_hash.insert(key_hash, created);
if let Some(source_api_key_id) = source_api_key_id {
imported_api_key_id_map.insert(source_api_key_id, created_api_key_id);
}
stats.standalone_keys.created += 1;
}
self.sync_imported_api_key_wallet(
&created.api_key_id,
&wallet_target,
created.name.as_deref().unwrap_or("独立余额 Key"),
}
if let Some(summary) = self
.import_admin_system_user_usage_aggregates(
root.get("usage_aggregates"),
&imported_user_id_map,
&imported_api_key_id_map,
merge_mode,
)
.await?;
existing_standalone_by_hash.insert(key_hash, created);
stats.standalone_keys.created += 1;
.await?
{
stats.usage_aggregates = Some(summary);
}
Ok(Ok(json!({
@@ -2915,6 +2950,40 @@ impl<'a> AdminAppState<'a> {
})))
}
async fn import_admin_system_user_usage_aggregates(
&self,
value: Option<&Value>,
user_id_map: &BTreeMap<String, String>,
api_key_id_map: &BTreeMap<String, String>,
merge_mode: AdminImportMergeMode,
) -> Result<Option<AdminSystemUsageAggregateImportSummary>, GatewayError> {
let Some(value) = value else {
return Ok(None);
};
if value.is_null() {
return Ok(None);
}
let snapshot = serde_json::from_value::<AdminSystemUsageAggregateSnapshot>(value.clone())
.map_err(|err| GatewayError::Client {
status: http::StatusCode::BAD_REQUEST,
message: format!("usage_aggregates 格式无效: {err}"),
})?;
if snapshot.stats_daily.is_empty()
&& snapshot.stats_user_daily.is_empty()
&& snapshot.stats_daily_api_key.is_empty()
{
return Ok(None);
}
self.import_admin_system_usage_aggregates(
&snapshot,
user_id_map,
api_key_id_map,
usage_aggregate_import_mode(merge_mode),
)
.await
.map(Some)
}
async fn sync_imported_user_wallet(
&self,
user_id: &str,
apps/aether-gateway/src/handlers/admin/request/system/mod.rs
+20
View File
@@ -70,6 +70,26 @@ impl<'a> AdminAppState<'a> {
self.app.purge_admin_system_data(target).await
}
pub(crate) async fn export_admin_system_usage_aggregates(
&self,
) -> Result<aether_data::repository::system::AdminSystemUsageAggregateSnapshot, GatewayError>
{
self.app.export_admin_system_usage_aggregates().await
}
pub(crate) async fn import_admin_system_usage_aggregates(
&self,
snapshot: &aether_data::repository::system::AdminSystemUsageAggregateSnapshot,
user_id_map: &std::collections::BTreeMap<String, String>,
api_key_id_map: &std::collections::BTreeMap<String, String>,
mode: aether_data::repository::system::AdminSystemUsageAggregateImportMode,
) -> Result<aether_data::repository::system::AdminSystemUsageAggregateImportSummary, GatewayError>
{
self.app
.import_admin_system_usage_aggregates(snapshot, user_id_map, api_key_id_map, mode)
.await
}
pub(crate) async fn run_admin_system_cleanup_once(
&self,
) -> Result<crate::maintenance::AdminSystemCleanupSummary, GatewayError> {
apps/aether-gateway/src/handlers/admin/request/users.rs
+13
View File
@@ -706,6 +706,19 @@ impl<'a> AdminAppState<'a> {
.await
}
pub(crate) async fn set_api_key_usage_totals(
&self,
api_key_id: &str,
total_requests: u64,
total_tokens: u64,
total_cost_usd: f64,
) -> Result<Option<aether_data::repository::auth::StoredAuthApiKeyExportRecord>, GatewayError>
{
self.app
.set_api_key_usage_totals(api_key_id, total_requests, total_tokens, total_cost_usd)
.await
}
pub(crate) async fn delete_user_api_key(
&self,
user_id: &str,
apps/aether-gateway/src/state/core.rs
+30
View File
@@ -629,6 +629,36 @@ impl AppState {
Ok(summary)
}
pub(crate) async fn export_admin_system_usage_aggregates(
&self,
) -> Result<aether_data::repository::system::AdminSystemUsageAggregateSnapshot, GatewayError>
{
self.data
.export_admin_system_usage_aggregates()
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn import_admin_system_usage_aggregates(
&self,
snapshot: &aether_data::repository::system::AdminSystemUsageAggregateSnapshot,
user_id_map: &std::collections::BTreeMap<String, String>,
api_key_id_map: &std::collections::BTreeMap<String, String>,
mode: aether_data::repository::system::AdminSystemUsageAggregateImportMode,
) -> Result<aether_data::repository::system::AdminSystemUsageAggregateImportSummary, GatewayError>
{
self.data
.import_admin_system_usage_aggregates(snapshot, user_id_map, api_key_id_map, mode)
.await
.map_err(|err| match err {
aether_data::DataLayerError::InvalidInput(detail) => GatewayError::Client {
status: http::StatusCode::BAD_REQUEST,
message: detail,
},
other => GatewayError::Internal(other.to_string()),
})
}
pub(crate) async fn run_admin_system_cleanup_once(
&self,
) -> Result<crate::maintenance::AdminSystemCleanupSummary, GatewayError> {
apps/aether-gateway/src/state/runtime/api_key_exports.rs
+19
View File
@@ -369,6 +369,25 @@ impl AppState {
Ok(api_key)
}
pub(crate) async fn set_api_key_usage_totals(
&self,
api_key_id: &str,
total_requests: u64,
total_tokens: u64,
total_cost_usd: f64,
) -> Result<Option<aether_data::repository::auth::StoredAuthApiKeyExportRecord>, GatewayError>
{
let api_key = self
.data
.set_api_key_usage_totals(api_key_id, total_requests, total_tokens, total_cost_usd)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if api_key.is_some() {
self.invalidate_auth_context_cache();
}
Ok(api_key)
}
pub(crate) async fn set_standalone_api_key_feature_settings(
&self,
api_key_id: &str,
apps/aether-gateway/src/tests/control/admin/system.rs
+15
View File
@@ -1043,6 +1043,7 @@ async fn gateway_handles_admin_system_users_export_locally_with_trusted_admin_pr
payload["users"][0]["group_names"],
json!(["Restricted GPT"])
);
assert_eq!(payload["users"][0]["id"], json!("user-1"));
assert_eq!(payload["users"][0]["wallet"]["balance"], json!(12.5));
assert_eq!(
payload["users"][0]["wallet"]["recharge_balance"],
@@ -1063,6 +1064,10 @@ async fn gateway_handles_admin_system_users_export_locally_with_trusted_admin_pr
payload["users"][0]["api_keys"][0]["is_standalone"],
json!(false)
);
assert_eq!(
payload["users"][0]["api_keys"][0]["api_key_id"],
json!("key-user-1")
);
assert_eq!(
payload["users"][0]["api_keys"][0]["total_tokens"],
json!(420)
@@ -1071,12 +1076,22 @@ async fn gateway_handles_admin_system_users_export_locally_with_trusted_admin_pr
payload["standalone_keys"][0]["key"],
json!("ak-standalone-live-1")
);
assert_eq!(
payload["standalone_keys"][0]["api_key_id"],
json!("key-standalone-1")
);
assert_eq!(payload["standalone_keys"][0]["total_tokens"], json!(84));
assert_eq!(
payload["standalone_keys"][0]["wallet"]["unlimited"],
json!(true)
);
assert_eq!(payload["standalone_keys"][0].get("is_standalone"), None,);
assert_eq!(payload["usage_aggregates"]["stats_daily"], json!([]));
assert_eq!(payload["usage_aggregates"]["stats_user_daily"], json!([]));
assert_eq!(
payload["usage_aggregates"]["stats_daily_api_key"],
json!([])
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
apps/aether-gateway/src/tests/control/admin/system_import.rs
+222 -42
View File
@@ -4,7 +4,9 @@ use aether_contracts::ExecutionPlan;
use aether_crypto::{
decrypt_python_fernet_ciphertext, encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY,
};
use aether_data::repository::auth::InMemoryAuthApiKeySnapshotRepository;
use aether_data::repository::auth::{
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
};
use aether_data::repository::auth_modules::{
AuthModuleReadRepository, InMemoryAuthModuleReadRepository, StoredOAuthProviderModuleConfig,
};
@@ -27,7 +29,7 @@ use axum::{extract::Request, Json, Router};
use http::StatusCode;
use serde_json::{json, Value};
use super::super::helpers::{sample_endpoint, sample_key, sample_provider};
use super::super::helpers::{hash_api_key, sample_endpoint, sample_key, sample_provider};
use super::super::{
build_router_with_state, build_state_with_execution_runtime_override, start_server, AppState,
};
@@ -1052,6 +1054,7 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
.expect("user api keys should load");
assert_eq!(user_api_keys.len(), 1);
assert_eq!(user_api_keys[0].name.as_deref(), Some("Alice CLI"));
assert_eq!(user_api_keys[0].total_requests, 12);
assert_eq!(user_api_keys[0].total_tokens, 3456);
assert_eq!(user_api_keys[0].total_cost_usd, 1.25);
assert_eq!(
@@ -1079,6 +1082,7 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
standalone_keys[0].name.as_deref(),
Some("Imported Standalone")
);
assert_eq!(standalone_keys[0].total_requests, 3);
assert_eq!(standalone_keys[0].total_tokens, 789);
assert_eq!(standalone_keys[0].total_cost_usd, 0.75);
assert_eq!(
@@ -1119,6 +1123,211 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
let _ = upstream_url;
}
#[tokio::test]
async fn gateway_overwrites_existing_admin_system_user_key_usage_totals() {
let user_key_hash = hash_api_key("sk-existing-user-key");
let standalone_key_hash = hash_api_key("sk-existing-standalone-key");
let existing_user = StoredUserAuthRecord::new(
"user-existing".to_string(),
Some("existing@example.com".to_string()),
true,
"existing".to_string(),
Some("existing-hash".to_string()),
"user".to_string(),
"local".to_string(),
None,
None,
None,
true,
false,
Some(chrono::Utc::now()),
Some(chrono::Utc::now()),
)
.expect("existing user should build");
let user_key_snapshot = StoredAuthApiKeySnapshot::new(
"user-existing".to_string(),
"existing".to_string(),
Some("existing@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
None,
None,
None,
"key-user-existing".to_string(),
Some("Existing User Key".to_string()),
true,
false,
false,
Some(10),
None,
None,
None,
None,
None,
)
.expect("user key snapshot should build");
let standalone_key_snapshot = StoredAuthApiKeySnapshot::new(
"admin-user-123".to_string(),
"admin".to_string(),
Some("admin@example.com".to_string()),
"admin".to_string(),
"local".to_string(),
true,
false,
None,
None,
None,
"key-standalone-existing".to_string(),
Some("Existing Standalone Key".to_string()),
true,
false,
true,
Some(20),
None,
None,
None,
None,
None,
)
.expect("standalone key snapshot should build");
let auth_repository = Arc::new(
InMemoryAuthApiKeySnapshotRepository::seed(vec![
(Some(user_key_hash.clone()), user_key_snapshot),
(Some(standalone_key_hash.clone()), standalone_key_snapshot),
])
.with_export_records(vec![
StoredAuthApiKeyExportRecord::new(
"user-existing".to_string(),
"key-user-existing".to_string(),
user_key_hash.clone(),
None,
Some("Existing User Key".to_string()),
None,
None,
None,
Some(10),
None,
None,
true,
None,
false,
1,
2,
0.03,
false,
)
.expect("existing user key export should build"),
StoredAuthApiKeyExportRecord::new(
"admin-user-123".to_string(),
"key-standalone-existing".to_string(),
standalone_key_hash.clone(),
None,
Some("Existing Standalone Key".to_string()),
None,
None,
None,
Some(20),
None,
None,
true,
None,
false,
4,
5,
0.06,
true,
)
.expect("existing standalone key export should build"),
]),
);
let user_repository =
Arc::new(aether_data::repository::users::InMemoryUserReadRepository::default());
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_auth_api_key_repository_for_tests(Arc::clone(&auth_repository))
.with_user_reader(user_repository)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_auth_users_for_tests([sample_import_admin_user("admin-user-123"), existing_user])
.with_auth_wallets_for_tests(Vec::<StoredWalletSnapshot>::new());
let gateway = build_router_with_state(state.clone());
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/system/users/import"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"version": "1.4",
"merge_mode": "overwrite",
"users": [{
"id": "source-user-existing",
"email": "existing@example.com",
"username": "existing",
"password_hash": "existing-hash",
"role": "user",
"is_active": true,
"api_keys": [{
"api_key_id": "source-user-key",
"key_hash": user_key_hash,
"name": "Imported User Key",
"is_active": true,
"total_requests": 222,
"total_tokens": 3333,
"total_cost_usd": 4.56
}]
}],
"standalone_keys": [{
"api_key_id": "source-standalone-key",
"key_hash": standalone_key_hash,
"name": "Imported Standalone Key",
"is_active": true,
"total_requests": 444,
"total_tokens": 5555,
"total_cost_usd": 6.78
}]
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
assert_eq!(payload["stats"]["users"]["updated"], json!(1));
assert_eq!(payload["stats"]["api_keys"]["updated"], json!(1));
assert_eq!(payload["stats"]["standalone_keys"]["updated"], json!(1));
let updated_records = state
.list_auth_api_key_export_records_by_ids(&[
"key-user-existing".to_string(),
"key-standalone-existing".to_string(),
])
.await
.expect("api key export records should load");
let user_key = updated_records
.iter()
.find(|record| record.api_key_id == "key-user-existing")
.expect("updated user key should exist");
assert_eq!(user_key.total_requests, 222);
assert_eq!(user_key.total_tokens, 3333);
assert_eq!(user_key.total_cost_usd, 4.56);
let standalone_key = updated_records
.iter()
.find(|record| record.api_key_id == "key-standalone-existing")
.expect("updated standalone key should exist");
assert_eq!(standalone_key.total_requests, 444);
assert_eq!(standalone_key.total_tokens, 5555);
assert_eq!(standalone_key.total_cost_usd, 6.78);
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_imports_admin_system_config_fixture_v22() {
let gateway = build_router_with_state(
@@ -1671,33 +1880,20 @@ async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_imp
}
#[tokio::test]
async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_import_and_forces_refresh(
async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_import_without_refresh(
) {
#[derive(Debug, Clone)]
struct SeenRefreshRequest {
content_type: String,
body: String,
}
let seen_refresh = Arc::new(Mutex::new(None::<SeenRefreshRequest>));
let seen_refresh = Arc::new(Mutex::new(false));
let seen_refresh_clone = Arc::clone(&seen_refresh);
let refresh_hits = Arc::new(Mutex::new(0usize));
let refresh_hits_clone = Arc::clone(&refresh_hits);
let refresh_server = Router::new().route(
"/oauth/token",
post(move |headers: HeaderMap, body: Bytes| {
post(move |_headers: HeaderMap, _body: Bytes| {
let seen_refresh_inner = Arc::clone(&seen_refresh_clone);
let refresh_hits_inner = Arc::clone(&refresh_hits_clone);
async move {
*refresh_hits_inner.lock().expect("mutex should lock") += 1;
*seen_refresh_inner.lock().expect("mutex should lock") = Some(SeenRefreshRequest {
content_type: headers
.get(http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
body: String::from_utf8(body.to_vec()).unwrap_or_default(),
});
*seen_refresh_inner.lock().expect("mutex should lock") = true;
axum::Json(json!({
"access_token": "oauth-access-token-refreshed",
"refresh_token": "oauth-refresh-token-refreshed",
@@ -1792,21 +1988,8 @@ async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_imp
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(*refresh_hits.lock().expect("mutex should lock"), 1);
let seen_refresh = seen_refresh
.lock()
.expect("mutex should lock")
.clone()
.expect("refresh request should be captured");
assert_eq!(
seen_refresh.content_type,
"application/x-www-form-urlencoded"
);
assert!(seen_refresh.body.contains("grant_type=refresh_token"));
assert!(seen_refresh
.body
.contains("refresh_token=oauth-refresh-token-new"));
assert_eq!(*refresh_hits.lock().expect("mutex should lock"), 0);
assert!(!*seen_refresh.lock().expect("mutex should lock"));
let providers = provider_catalog_repository
.list_providers(false)
@@ -1823,7 +2006,7 @@ async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_imp
assert_eq!(key.name, "oauth-primary");
assert_eq!(key.oauth_invalid_at_unix_secs, None);
assert_eq!(key.oauth_invalid_reason, None);
assert!(key.expires_at_unix_secs.is_some());
assert_eq!(key.expires_at_unix_secs, None);
assert_eq!(
decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
@@ -1832,7 +2015,7 @@ async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_imp
.expect("api key should be present"),
)
.expect("oauth access token should decrypt"),
"oauth-access-token-refreshed"
"oauth-access-token-new"
);
let auth_config = decrypt_python_fernet_ciphertext(
@@ -1845,15 +2028,12 @@ async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_imp
let auth_config: Value =
serde_json::from_str(&auth_config).expect("oauth auth config json should parse");
assert_eq!(auth_config["provider_type"], "codex");
assert_eq!(
auth_config["refresh_token"],
"oauth-refresh-token-refreshed"
);
assert_eq!(auth_config["refresh_token"], "oauth-refresh-token-new");
assert_eq!(auth_config["email"], "alice@example.com");
assert_eq!(auth_config["account_id"], "acct-codex-123");
assert_eq!(auth_config["plan_type"], "plus");
assert_eq!(auth_config["token_type"], "Bearer");
assert_eq!(auth_config["expires_at"].as_u64(), key.expires_at_unix_secs);
assert!(auth_config.get("token_type").is_none());
assert!(auth_config.get("expires_at").is_none());
gateway_handle.abort();
refresh_handle.abort();
crates/aether-data/src/backend/maintenance.rs
+85 -1
View File
@@ -8,7 +8,9 @@ use crate::maintenance::{
WalletDailyUsageAggregationInput, WalletDailyUsageAggregationResult,
};
use crate::repository::system::{
AdminSystemPurgeSummary, AdminSystemPurgeTarget, AdminSystemStats, StoredSystemConfigEntry,
AdminSystemPurgeSummary, AdminSystemPurgeTarget, AdminSystemStats,
AdminSystemUsageAggregateImportMode, AdminSystemUsageAggregateImportSummary,
AdminSystemUsageAggregateSnapshot, StoredSystemConfigEntry,
};
use crate::DataLayerError;
use sqlx::migrate::MigrateError;
@@ -195,6 +197,37 @@ impl DataBackends {
}
}
pub async fn export_admin_system_usage_aggregates(
&self,
) -> Result<AdminSystemUsageAggregateSnapshot, DataLayerError> {
match self.sql_backend() {
Some(backend) => backend.export_admin_system_usage_aggregates().await,
None => Ok(AdminSystemUsageAggregateSnapshot::default()),
}
}
pub async fn import_admin_system_usage_aggregates(
&self,
snapshot: &AdminSystemUsageAggregateSnapshot,
user_id_map: &std::collections::BTreeMap<String, String>,
api_key_id_map: &std::collections::BTreeMap<String, String>,
mode: AdminSystemUsageAggregateImportMode,
) -> Result<AdminSystemUsageAggregateImportSummary, DataLayerError> {
match self.sql_backend() {
Some(backend) => {
backend
.import_admin_system_usage_aggregates(
snapshot,
user_id_map,
api_key_id_map,
mode,
)
.await
}
None => Ok(AdminSystemUsageAggregateImportSummary::default()),
}
}
pub async fn purge_admin_request_bodies_batch(
&self,
batch_size: usize,
@@ -505,6 +538,57 @@ impl<'a> SqlBackendRef<'a> {
}
}
async fn export_admin_system_usage_aggregates(
self,
) -> Result<AdminSystemUsageAggregateSnapshot, DataLayerError> {
match self {
Self::Postgres(postgres) => postgres.export_admin_system_usage_aggregates().await,
Self::Mysql(mysql) => mysql.export_admin_system_usage_aggregates().await,
Self::Sqlite(sqlite) => sqlite.export_admin_system_usage_aggregates().await,
}
}
async fn import_admin_system_usage_aggregates(
self,
snapshot: &AdminSystemUsageAggregateSnapshot,
user_id_map: &std::collections::BTreeMap<String, String>,
api_key_id_map: &std::collections::BTreeMap<String, String>,
mode: AdminSystemUsageAggregateImportMode,
) -> Result<AdminSystemUsageAggregateImportSummary, DataLayerError> {
match self {
Self::Postgres(postgres) => {
postgres
.import_admin_system_usage_aggregates(
snapshot,
user_id_map,
api_key_id_map,
mode,
)
.await
}
Self::Mysql(mysql) => {
mysql
.import_admin_system_usage_aggregates(
snapshot,
user_id_map,
api_key_id_map,
mode,
)
.await
}
Self::Sqlite(sqlite) => {
sqlite
.import_admin_system_usage_aggregates(
snapshot,
user_id_map,
api_key_id_map,
mode,
)
.await
}
}
}
async fn purge_admin_request_bodies_batch(
self,
batch_size: usize,
crates/aether-data/src/backend/sqlite.rs
+116 -1
View File
@@ -258,9 +258,15 @@ impl SqliteBackend {
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use super::SqliteBackend;
use crate::lifecycle::migrate::run_sqlite_migrations;
use crate::repository::system::AdminSystemPurgeTarget;
use crate::repository::system::{
AdminSystemPurgeTarget, AdminSystemStatsDailyAggregate,
AdminSystemStatsDailyApiKeyAggregate, AdminSystemStatsUserDailyAggregate,
AdminSystemUsageAggregateImportMode, AdminSystemUsageAggregateSnapshot,
};
use crate::{
DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig, StatsDailyAggregationInput,
StatsHourlyAggregationInput, WalletDailyUsageAggregationInput,
@@ -456,6 +462,115 @@ VALUES
assert_eq!(admin_exists, 1);
}
#[tokio::test]
async fn admin_system_usage_aggregates_round_trip_after_sqlite_migrations() {
let config = SqlDatabaseConfig {
driver: DatabaseDriver::Sqlite,
url: "sqlite::memory:".to_string(),
pool: SqlPoolConfig {
max_connections: 1,
..SqlPoolConfig::default()
},
};
let backend = SqliteBackend::from_config(config).expect("backend should build");
run_sqlite_migrations(backend.pool())
.await
.expect("sqlite migrations should run");
sqlx::query(
r#"
INSERT INTO users (id, email, username, role, created_at, updated_at)
VALUES ('target-user-1', 'target@example.com', 'target', 'user', 1, 1)
"#,
)
.execute(backend.pool())
.await
.expect("target user should insert");
sqlx::query(
r#"
INSERT INTO api_keys (id, user_id, key_hash, name, created_at, updated_at)
VALUES ('target-key-1', 'target-user-1', 'hash-target-key', 'target key', 1, 1)
"#,
)
.execute(backend.pool())
.await
.expect("target key should insert");
let snapshot = AdminSystemUsageAggregateSnapshot {
stats_daily: vec![AdminSystemStatsDailyAggregate {
date_unix_secs: 86_400,
total_requests: 9,
success_requests: 8,
error_requests: 1,
input_tokens: 100,
output_tokens: 200,
cache_creation_tokens: 3,
cache_read_tokens: 4,
total_cost: 1.25,
actual_total_cost: 1.0,
is_complete: true,
aggregated_at_unix_secs: Some(90_000),
}],
stats_user_daily: vec![AdminSystemStatsUserDailyAggregate {
user_id: "source-user-1".to_string(),
username: Some("source".to_string()),
date_unix_secs: 86_400,
total_requests: 5,
success_requests: 5,
error_requests: 0,
input_tokens: 50,
output_tokens: 60,
cache_creation_tokens: 1,
cache_read_tokens: 2,
total_cost: 0.5,
}],
stats_daily_api_key: vec![AdminSystemStatsDailyApiKeyAggregate {
api_key_id: "source-key-1".to_string(),
api_key_name: Some("source key".to_string()),
date_unix_secs: 86_400,
total_requests: 4,
success_requests: 3,
error_requests: 1,
input_tokens: 40,
output_tokens: 30,
cache_creation_tokens: 2,
cache_read_tokens: 1,
total_cost: 0.75,
}],
};
let user_id_map =
BTreeMap::from([("source-user-1".to_string(), "target-user-1".to_string())]);
let api_key_id_map =
BTreeMap::from([("source-key-1".to_string(), "target-key-1".to_string())]);
let summary = backend
.import_admin_system_usage_aggregates(
&snapshot,
&user_id_map,
&api_key_id_map,
AdminSystemUsageAggregateImportMode::Overwrite,
)
.await
.expect("usage aggregates should import");
assert_eq!(summary.stats_daily.created, 1);
assert_eq!(summary.stats_user_daily.created, 1);
assert_eq!(summary.stats_daily_api_key.created, 1);
let exported = backend
.export_admin_system_usage_aggregates()
.await
.expect("usage aggregates should export");
assert_eq!(exported.stats_daily.len(), 1);
assert_eq!(exported.stats_daily[0].total_requests, 9);
assert_eq!(exported.stats_daily[0].actual_total_cost, 1.0);
assert_eq!(exported.stats_user_daily.len(), 1);
assert_eq!(exported.stats_user_daily[0].user_id, "target-user-1");
assert_eq!(exported.stats_user_daily[0].total_requests, 5);
assert_eq!(exported.stats_daily_api_key.len(), 1);
assert_eq!(exported.stats_daily_api_key[0].api_key_id, "target-key-1");
assert_eq!(exported.stats_daily_api_key[0].total_requests, 4);
}
#[tokio::test]
async fn admin_system_request_bodies_purge_clears_inline_usage_body_fields() {
let config = SqlDatabaseConfig {
crates/aether-data/src/backend/system.rs
+1444 -1
View File
File diff suppressed because it is too large Load Diff
crates/aether-data/src/repository/auth/memory.rs
+20
View File
@@ -1017,6 +1017,26 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
Ok(Some(export.clone()))
}
async fn set_api_key_usage_totals(
&self,
api_key_id: &str,
total_requests: u64,
total_tokens: u64,
total_cost_usd: f64,
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
let mut index = self
.index
.write()
.expect("auth api key snapshot repository lock");
let Some(export) = index.export_by_api_key_id.get_mut(api_key_id) else {
return Ok(None);
};
export.total_requests = total_requests;
export.total_tokens = total_tokens;
export.total_cost_usd = total_cost_usd;
Ok(Some(export.clone()))
}
async fn delete_user_api_key(
&self,
user_id: &str,
crates/aether-data/src/repository/auth/mysql.rs
+28
View File
@@ -692,6 +692,34 @@ WHERE id = ?
self.reload_export_by_id(api_key_id).await
}
async fn set_api_key_usage_totals(
&self,
api_key_id: &str,
total_requests: u64,
total_tokens: u64,
total_cost_usd: f64,
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
sqlx::query(
r#"
UPDATE api_keys
SET total_requests = ?,
total_tokens = ?,
total_cost_usd = ?,
updated_at = ?
WHERE id = ?
"#,
)
.bind(total_requests as i64)
.bind(total_tokens as i64)
.bind(total_cost_usd)
.bind(current_unix_secs() as i64)
.bind(api_key_id)
.execute(&self.pool)
.await
.map_sql_err()?;
self.reload_export_by_id(api_key_id).await
}
async fn delete_user_api_key(
&self,
user_id: &str,
crates/aether-data/src/repository/auth/postgres.rs
+52
View File
@@ -668,6 +668,40 @@ RETURNING
is_standalone
"#;
const SET_API_KEY_USAGE_TOTALS_SQL: &str = r#"
UPDATE api_keys
SET
total_requests = $2,
total_tokens = $3,
total_cost_usd = $4,
updated_at = NOW()
WHERE id = $1
RETURNING
user_id,
id AS api_key_id,
key_hash,
key_encrypted,
name,
allowed_providers,
allowed_api_formats,
allowed_models,
ip_rules,
rate_limit,
concurrent_limit,
force_capabilities,
feature_settings,
is_active,
CAST(EXTRACT(EPOCH FROM expires_at) AS BIGINT) AS expires_at_unix_secs,
auto_delete_on_expiry,
total_requests,
COALESCE(total_tokens, 0)::BIGINT AS total_tokens,
COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0) AS total_cost_usd,
CAST(EXTRACT(EPOCH FROM last_used_at) AS BIGINT) AS last_used_at_unix_secs,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs,
is_standalone
"#;
const SET_USER_API_KEY_LOCKED_SQL: &str = r#"
UPDATE api_keys
SET
@@ -1436,6 +1470,24 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.find(|record| record.user_id == user_id && !record.is_standalone))
}
async fn set_api_key_usage_totals(
&self,
api_key_id: &str,
total_requests: u64,
total_tokens: u64,
total_cost_usd: f64,
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
let row = sqlx::query(SET_API_KEY_USAGE_TOTALS_SQL)
.bind(api_key_id)
.bind(total_requests as i64)
.bind(total_tokens as i64)
.bind(total_cost_usd)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?;
row.as_ref().map(map_auth_api_key_export_row).transpose()
}
async fn delete_user_api_key(
&self,
user_id: &str,
crates/aether-data/src/repository/auth/sqlite.rs
+28
View File
@@ -692,6 +692,34 @@ WHERE id = ?
self.reload_export_by_id(api_key_id).await
}
async fn set_api_key_usage_totals(
&self,
api_key_id: &str,
total_requests: u64,
total_tokens: u64,
total_cost_usd: f64,
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
sqlx::query(
r#"
UPDATE api_keys
SET total_requests = ?,
total_tokens = ?,
total_cost_usd = ?,
updated_at = ?
WHERE id = ?
"#,
)
.bind(total_requests as i64)
.bind(total_tokens as i64)
.bind(total_cost_usd)
.bind(current_unix_secs() as i64)
.bind(api_key_id)
.execute(&self.pool)
.await
.map_sql_err()?;
self.reload_export_by_id(api_key_id).await
}
async fn delete_user_api_key(
&self,
user_id: &str,
crates/aether-data/src/repository/auth/types.rs
+8
View File
@@ -691,6 +691,14 @@ pub trait AuthApiKeyWriteRepository: Send + Sync {
feature_settings: Option<serde_json::Value>,
) -> Result<Option<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
async fn set_api_key_usage_totals(
&self,
api_key_id: &str,
total_requests: u64,
total_tokens: u64,
total_cost_usd: f64,
) -> Result<Option<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
async fn delete_user_api_key(
&self,
user_id: &str,
crates/aether-data/src/repository/system.rs
+79
View File
@@ -21,6 +21,85 @@ pub struct AdminSystemStats {
pub total_requests: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AdminSystemUsageAggregateImportMode {
Skip,
Overwrite,
Error,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct AdminSystemStatsDailyAggregate {
pub date_unix_secs: u64,
pub total_requests: u64,
pub success_requests: u64,
pub error_requests: u64,
pub input_tokens: u64,
pub output_tokens: u64,
pub cache_creation_tokens: u64,
pub cache_read_tokens: u64,
pub total_cost: f64,
pub actual_total_cost: f64,
pub is_complete: bool,
pub aggregated_at_unix_secs: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct AdminSystemStatsUserDailyAggregate {
pub user_id: String,
pub username: Option<String>,
pub date_unix_secs: u64,
pub total_requests: u64,
pub success_requests: u64,
pub error_requests: u64,
pub input_tokens: u64,
pub output_tokens: u64,
pub cache_creation_tokens: u64,
pub cache_read_tokens: u64,
pub total_cost: f64,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct AdminSystemStatsDailyApiKeyAggregate {
pub api_key_id: String,
pub api_key_name: Option<String>,
pub date_unix_secs: u64,
pub total_requests: u64,
pub success_requests: u64,
pub error_requests: u64,
pub input_tokens: u64,
pub output_tokens: u64,
pub cache_creation_tokens: u64,
pub cache_read_tokens: u64,
pub total_cost: f64,
}
#[derive(Debug, Clone, Default, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct AdminSystemUsageAggregateSnapshot {
#[serde(default)]
pub stats_daily: Vec<AdminSystemStatsDailyAggregate>,
#[serde(default)]
pub stats_user_daily: Vec<AdminSystemStatsUserDailyAggregate>,
#[serde(default)]
pub stats_daily_api_key: Vec<AdminSystemStatsDailyApiKeyAggregate>,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct AdminSystemUsageAggregateImportCounter {
pub created: u64,
pub updated: u64,
pub skipped: u64,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct AdminSystemUsageAggregateImportSummary {
pub stats_daily: AdminSystemUsageAggregateImportCounter,
pub stats_user_daily: AdminSystemUsageAggregateImportCounter,
pub stats_daily_api_key: AdminSystemUsageAggregateImportCounter,
pub skipped_unmapped_user_daily: u64,
pub skipped_unmapped_api_key_daily: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AdminSystemPurgeTarget {
Config,
frontend/src/api/admin.ts
+68
View File
@@ -96,6 +96,7 @@ export interface UsersExportData {
user_groups?: UserGroupExport[]
users: UserExport[]
standalone_keys?: StandaloneKeyExport[]
usage_aggregates?: UsageAggregateSnapshot
}
export interface AggregateExportData {
@@ -120,6 +121,7 @@ export interface UserGroupExport {
}
export interface UserExport {
id?: string
email: string
email_verified?: boolean
username: string
@@ -144,6 +146,7 @@ export interface UserExport {
}
export interface UserApiKeyExport {
api_key_id?: string
key?: string | null
key_hash: string
key_encrypted?: string | null
@@ -161,12 +164,76 @@ export interface UserApiKeyExport {
expires_at?: string | null
auto_delete_on_expiry?: boolean
total_requests?: number
total_tokens?: number
total_cost_usd?: number
}
// 独立余额 Key 导出结构（与 UserApiKeyExport 相同，但不包含 is_standalone）
export type StandaloneKeyExport = Omit<UserApiKeyExport, 'is_standalone'>
export interface StatsDailyAggregateExport {
date_unix_secs: number
total_requests: number
success_requests: number
error_requests: number
input_tokens: number
output_tokens: number
cache_creation_tokens: number
cache_read_tokens: number
total_cost: number
actual_total_cost: number
is_complete: boolean
aggregated_at_unix_secs?: number | null
}
export interface StatsUserDailyAggregateExport {
user_id: string
username?: string | null
date_unix_secs: number
total_requests: number
success_requests: number
error_requests: number
input_tokens: number
output_tokens: number
cache_creation_tokens: number
cache_read_tokens: number
total_cost: number
}
export interface StatsDailyApiKeyAggregateExport {
api_key_id: string
api_key_name?: string | null
date_unix_secs: number
total_requests: number
success_requests: number
error_requests: number
input_tokens: number
output_tokens: number
cache_creation_tokens: number
cache_read_tokens: number
total_cost: number
}
export interface UsageAggregateSnapshot {
stats_daily?: StatsDailyAggregateExport[]
stats_user_daily?: StatsUserDailyAggregateExport[]
stats_daily_api_key?: StatsDailyApiKeyAggregateExport[]
}
export interface UsageAggregateImportCounter {
created: number
updated: number
skipped: number
}
export interface UsageAggregateImportSummary {
stats_daily: UsageAggregateImportCounter
stats_user_daily: UsageAggregateImportCounter
stats_daily_api_key: UsageAggregateImportCounter
skipped_unmapped_user_daily: number
skipped_unmapped_api_key_daily: number
}
export interface GlobalModelExport {
name: string
display_name: string
@@ -522,6 +589,7 @@ export interface UsersImportResponse {
users: { created: number; updated: number; skipped: number }
api_keys: { created: number; updated?: number; skipped: number }
standalone_keys?: { created: number; updated?: number; skipped: number }
usage_aggregates?: UsageAggregateImportSummary
errors: string[]
}
}
frontend/src/views/admin/system-settings/AggregateImportDialog.vue
+39 -1
View File
@@ -42,6 +42,9 @@
<li v-if="aggregateImportPreview.user_data.standalone_keys?.length">
独立余额 Keys: {{ aggregateImportPreview.user_data.standalone_keys.length }} 个
</li>
<li v-if="usageAggregatePreviewCounts.total > 0">
统计聚合: {{ usageAggregatePreviewCounts.total }} 行
</li>
</ul>
</div>
</div>
@@ -84,7 +87,7 @@
</div>
<p class="text-xs text-muted-foreground">
注意：完整备份会先导入配置数据，再导入用户数据；文件包含用户、用户组、API Keys 与钱包快照，用户 API Keys 需要目标系统使用相同的 ENCRYPTION_KEY。
注意：完整备份会先导入配置数据，再导入用户数据；文件包含用户、用户组、API Keys、Key 用量、钱包快照与统计聚合。正常导出的 API Keys 会在导入时使用目标系统密钥重新加密；仅当备份中包含 key_encrypted 等未解密密文字段时，才需要目标系统使用兼容的 ENCRYPTION_KEY。
</p>
<div
@@ -149,6 +152,9 @@
用户创建 {{ aggregateImportResult.users.stats.users.created }}，
API Keys 创建 {{ aggregateImportResult.users.stats.api_keys.created }}，
跳过 {{ aggregateImportResult.users.stats.users.skipped }} 个用户
<template v-if="usageAggregateResultText">
，统计聚合 {{ usageAggregateResultText }}
</template>
</p>
</div>
</div>
@@ -217,4 +223,36 @@ const warningMessages = computed(() => {
const userErrors = props.aggregateImportResult.users.stats.errors.map((message) => `用户数据: ${message}`)
return [...configErrors, ...userErrors]
})
const usageAggregatePreviewCounts = computed(() => {
const aggregates = props.aggregateImportPreview?.user_data.usage_aggregates
const statsDaily = aggregates?.stats_daily?.length ?? 0
const statsUserDaily = aggregates?.stats_user_daily?.length ?? 0
const statsDailyApiKey = aggregates?.stats_daily_api_key?.length ?? 0
return {
statsDaily,
statsUserDaily,
statsDailyApiKey,
total: statsDaily + statsUserDaily + statsDailyApiKey,
}
})
const usageAggregateResultText = computed(() => {
const aggregates = props.aggregateImportResult?.users.stats.usage_aggregates
if (!aggregates) return ''
const counters = [
aggregates.stats_daily,
aggregates.stats_user_daily,
aggregates.stats_daily_api_key,
]
const created = counters.reduce((sum, item) => sum + item.created, 0)
const updated = counters.reduce((sum, item) => sum + item.updated, 0)
const skipped = counters.reduce((sum, item) => sum + item.skipped, 0)
const total = created + updated + skipped
if (total === 0) return ''
return `创建 ${created}，更新 ${updated}，跳过 ${skipped}`
})
</script>