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

Merge remote-tracking branch 'origin/pr/395' into aether-rust-pioneer

Browse Source
This commit is contained in:
fawney19
2026-05-08 02:47:07 +08:00
parent 66d8e90647 dd8c2ebec6
commit 31d142effc
37 changed files with 2836 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
apps/aether-gateway/src/data/state/core.rs
View File

@@ -423,4 +423,28 @@ impl GatewayDataState {
None => Ok(super::AdminSystemStats::default()),
}
}
pub(crate) async fn purge_admin_system_data(
&self,
target: aether_data::repository::system::AdminSystemPurgeTarget,
) -> Result<aether_data::repository::system::AdminSystemPurgeSummary, DataLayerError> {
if matches!(
target,
aether_data::repository::system::AdminSystemPurgeTarget::Config
) {
if let Some(values) = &self.system_config_values {
let mut values = values.write().expect("system config values lock");
let deleted = values.len() as u64;
values.clear();
let mut summary =
aether_data::repository::system::AdminSystemPurgeSummary::default();
summary.add("system_configs", deleted);
return Ok(summary);
}
}
match self.backends.as_ref() {
Some(backends) => backends.purge_admin_system_data(target).await,
None => Ok(aether_data::repository::system::AdminSystemPurgeSummary::default()),
}
}
}

19
apps/aether-gateway/src/handlers/admin/request/system/mod.rs
View File

@@ -60,6 +60,25 @@ impl<'a> AdminAppState<'a> {
self.app.read_admin_system_stats().await
}
pub(crate) async fn purge_admin_system_data(
&self,
target: aether_data::repository::system::AdminSystemPurgeTarget,
) -> Result<aether_data::repository::system::AdminSystemPurgeSummary, GatewayError> {
self.app.purge_admin_system_data(target).await
}
pub(crate) async fn run_admin_system_cleanup_once(
&self,
) -> Result<crate::maintenance::AdminSystemCleanupSummary, GatewayError> {
self.app.run_admin_system_cleanup_once().await
}
pub(crate) async fn rebuild_admin_stats_once(
&self,
) -> Result<crate::maintenance::AdminStatsRebuildSummary, GatewayError> {
self.app.rebuild_admin_stats_once().await
}
pub(crate) async fn find_proxy_node(
&self,
node_id: &str,

167
apps/aether-gateway/src/handlers/admin/system/core/system_routes.rs
View File

@@ -18,6 +18,7 @@ use crate::handlers::admin::system::shared::settings::{
};
use crate::handlers::admin::system::shared::smtp::build_admin_smtp_test_payload;
use crate::GatewayError;
use aether_data::repository::system::AdminSystemPurgeTarget;
use axum::{
body::{Body, Bytes},
http,
@@ -183,26 +184,29 @@ pub(super) async fn maybe_build_local_admin_core_system_response(
));
}
if matches!(
decision.route_kind.as_deref(),
Some(
"cleanup"
| "purge_config"
| "purge_users"
| "purge_usage"
| "purge_audit_logs"
| "purge_request_bodies"
| "purge_stats"
)
) && request_method == http::Method::POST
if decision.route_kind.as_deref() == Some("cleanup") && request_method == http::Method::POST {
return Ok(Some(attach_admin_audit_response(
Json(build_admin_system_cleanup_payload(state).await?).into_response(),
"admin_system_cleanup_completed",
"cleanup_system_data",
"system_cleanup",
"global",
)));
}
if let Some((target, action, object_type, object_id)) =
admin_system_purge_target_for_route_kind(decision.route_kind.as_deref())
{
return Ok(Some(
(
http::StatusCode::SERVICE_UNAVAILABLE,
Json(json!({ "detail": "Admin system data unavailable" })),
)
.into_response(),
));
if request_method != http::Method::POST {
return Ok(None);
}
return Ok(Some(attach_admin_audit_response(
Json(build_admin_system_purge_payload(state, target).await?).into_response(),
"admin_system_data_purged",
action,
object_type,
object_id,
)));
}
if decision.route_kind.as_deref() == Some("settings_set")
@@ -430,3 +434,128 @@ pub(super) async fn maybe_build_local_admin_core_system_response(
Ok(None)
}
fn admin_system_purge_target_for_route_kind(
route_kind: Option<&str>,
) -> Option<(
AdminSystemPurgeTarget,
&'static str,
&'static str,
&'static str,
)> {
match route_kind {
Some("purge_config") => Some((
AdminSystemPurgeTarget::Config,
"purge_system_config",
"system_config",
"global",
)),
Some("purge_users") => Some((
AdminSystemPurgeTarget::Users,
"purge_non_admin_users",
"users",
"non_admin",
)),
Some("purge_usage") => Some((
AdminSystemPurgeTarget::Usage,
"purge_usage_records",
"usage",
"all",
)),
Some("purge_audit_logs") => Some((
AdminSystemPurgeTarget::AuditLogs,
"purge_audit_logs",
"audit_logs",
"all",
)),
Some("purge_request_bodies") => Some((
AdminSystemPurgeTarget::RequestBodies,
"purge_request_bodies",
"request_bodies",
"all",
)),
Some("purge_stats") => Some((AdminSystemPurgeTarget::Stats, "purge_stats", "stats", "all")),
_ => None,
}
}
async fn build_admin_system_purge_payload(
state: &AdminAppState<'_>,
target: AdminSystemPurgeTarget,
) -> Result<serde_json::Value, GatewayError> {
let summary = state.purge_admin_system_data(target).await?;
let total = summary.total();
let affected = summary.affected.clone();
if target == AdminSystemPurgeTarget::Stats {
let rebuild = state.rebuild_admin_stats_once().await?;
let message = if rebuild.capped {
format!(
"统计聚合已清空，已重建 {} 个小时桶和 {} 个日桶，仍有历史统计待后台任务继续重建",
rebuild.hourly_buckets, rebuild.daily_buckets
)
} else {
format!(
"统计聚合已清空并重建，删除 {} 行，重建 {} 个小时桶和 {} 个日桶",
total, rebuild.hourly_buckets, rebuild.daily_buckets
)
};
return Ok(json!({
"message": message,
"deleted": affected,
"rebuilt": {
"hourly_buckets": rebuild.hourly_buckets,
"daily_buckets": rebuild.daily_buckets,
"capped": rebuild.capped,
},
}));
}
let (message, count_key) = match target {
AdminSystemPurgeTarget::Config => ("系统配置已清空", "deleted"),
AdminSystemPurgeTarget::Users => ("非管理员用户已清空", "deleted"),
AdminSystemPurgeTarget::Usage => ("使用记录已清空", "deleted"),
AdminSystemPurgeTarget::AuditLogs => ("审计日志已清空", "deleted"),
AdminSystemPurgeTarget::RequestBodies => ("请求/响应体已清空", "cleaned"),
AdminSystemPurgeTarget::Stats => unreachable!("stats handled above"),
};
Ok(json!({
"message": format!("{message}，影响 {} 行", total),
count_key: affected,
}))
}
async fn build_admin_system_cleanup_payload(
state: &AdminAppState<'_>,
) -> Result<serde_json::Value, GatewayError> {
let summary = state.run_admin_system_cleanup_once().await?;
let cleaned = json!({
"audit_logs": summary.audit_logs_deleted,
"request_candidates": summary.request_candidates_deleted,
"pending_failed": summary.pending_failed,
"pending_recovered": summary.pending_recovered,
"usage_body_externalized": summary.usage.body_externalized,
"usage_legacy_body_refs_migrated": summary.usage.legacy_body_refs_migrated,
"usage_body_cleaned": summary.usage.body_cleaned,
"usage_header_cleaned": summary.usage.header_cleaned,
"usage_keys_cleaned": summary.usage.keys_cleaned,
"usage_records_deleted": summary.usage.records_deleted,
});
let total = summary
.audit_logs_deleted
.saturating_add(summary.request_candidates_deleted)
.saturating_add(summary.pending_failed)
.saturating_add(summary.pending_recovered)
.saturating_add(summary.usage.body_externalized)
.saturating_add(summary.usage.legacy_body_refs_migrated)
.saturating_add(summary.usage.body_cleaned)
.saturating_add(summary.usage.header_cleaned)
.saturating_add(summary.usage.keys_cleaned)
.saturating_add(summary.usage.records_deleted);
Ok(json!({
"message": format!("系统清理已执行，影响 {} 项", total),
"cleaned": cleaned,
}))
}

4
apps/aether-gateway/src/handlers/admin/users/lifecycle/reads.rs
View File

@@ -29,6 +29,9 @@ pub(in super::super) async fn build_admin_list_users_response(
.map(|value| value.trim().to_ascii_lowercase())
.filter(|value| !value.is_empty());
let is_active = query_param_optional_bool(request_context.query_string(), "is_active");
let search = query_param_value(request_context.query_string(), "search")
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty());
let paged_rows = state
.list_export_users_page(&aether_data::repository::users::UserExportListQuery {
@@ -36,6 +39,7 @@ pub(in super::super) async fn build_admin_list_users_response(
limit,
role: role.clone(),
is_active,
search,
})
.await?;
let user_ids = paged_rows

28
apps/aether-gateway/src/maintenance/mod.rs
View File

@@ -4,18 +4,20 @@ mod tests;
pub(crate) use runtime::{
cancel_proxy_upgrade_rollout, clear_proxy_upgrade_rollout_conflicts,
inspect_proxy_upgrade_rollout, perform_pool_quota_probe_once, perform_provider_checkin_once,
record_proxy_upgrade_traffic_success, restore_proxy_upgrade_rollout_skipped_nodes,
retry_proxy_upgrade_rollout_node, skip_proxy_upgrade_rollout_node, spawn_audit_cleanup_worker,
inspect_proxy_upgrade_rollout, perform_oauth_token_refresh_once, perform_pool_quota_probe_once,
perform_provider_checkin_once, rebuild_admin_stats_once, record_proxy_upgrade_traffic_success,
restore_proxy_upgrade_rollout_skipped_nodes, retry_proxy_upgrade_rollout_node,
run_admin_system_cleanup_once, skip_proxy_upgrade_rollout_node, spawn_audit_cleanup_worker,
spawn_db_maintenance_worker, spawn_gemini_file_mapping_cleanup_worker,
spawn_pending_cleanup_worker, spawn_pool_monitor_worker, spawn_pool_quota_probe_worker,
spawn_provider_checkin_worker, spawn_proxy_node_stale_cleanup_worker,
spawn_proxy_upgrade_rollout_worker, spawn_request_candidate_cleanup_worker,
spawn_stats_aggregation_worker, spawn_stats_hourly_aggregation_worker,
spawn_usage_cleanup_worker, spawn_wallet_daily_usage_aggregation_worker,
start_proxy_upgrade_rollout, PoolQuotaProbeRunSummary, ProviderCheckinRunSummary,
ProxyUpgradeRolloutCancelSummary, ProxyUpgradeRolloutConflictClearSummary,
ProxyUpgradeRolloutNodeActionSummary, ProxyUpgradeRolloutProbeConfig,
ProxyUpgradeRolloutSkippedRestoreSummary, ProxyUpgradeRolloutStatus,
ProxyUpgradeRolloutTrackedNodeState,
spawn_oauth_token_refresh_worker, spawn_pending_cleanup_worker, spawn_pool_monitor_worker,
spawn_pool_quota_probe_worker, spawn_provider_checkin_worker,
spawn_proxy_node_stale_cleanup_worker, spawn_proxy_upgrade_rollout_worker,
spawn_request_candidate_cleanup_worker, spawn_stats_aggregation_worker,
spawn_stats_hourly_aggregation_worker, spawn_usage_cleanup_worker,
spawn_wallet_daily_usage_aggregation_worker, start_proxy_upgrade_rollout,
AdminStatsRebuildSummary, AdminSystemCleanupSummary, OAuthTokenRefreshRunSummary,
PoolQuotaProbeRunSummary, ProviderCheckinRunSummary, ProxyUpgradeRolloutCancelSummary,
ProxyUpgradeRolloutConflictClearSummary, ProxyUpgradeRolloutNodeActionSummary,
ProxyUpgradeRolloutProbeConfig, ProxyUpgradeRolloutSkippedRestoreSummary,
ProxyUpgradeRolloutStatus, ProxyUpgradeRolloutTrackedNodeState,
};

77
apps/aether-gateway/src/maintenance/runtime.rs
View File

@@ -12,6 +12,8 @@ mod audit_cleanup;
mod config;
#[path = "runtime/db_maintenance.rs"]
mod db_maintenance;
#[path = "runtime/oauth_token_refresh.rs"]
mod oauth_token_refresh;
#[path = "runtime/pending_cleanup.rs"]
mod pending_cleanup;
#[path = "runtime/pool_quota_probe.rs"]
@@ -47,6 +49,9 @@ pub(crate) use aether_data_contracts::repository::usage::{
use audit_cleanup::*;
use config::*;
use db_maintenance::*;
pub(crate) use oauth_token_refresh::{
perform_oauth_token_refresh_once, OAuthTokenRefreshRunSummary,
};
use pending_cleanup::*;
pub(crate) use pool_quota_probe::{
perform_pool_quota_probe_once, perform_pool_quota_probe_once_with_config,
@@ -89,6 +94,7 @@ const PROXY_UPGRADE_ROLLOUT_INTERVAL: Duration = Duration::from_secs(15);
const PROXY_NODE_STALE_MIN_GRACE_SECS: u64 = 15;
const PROXY_NODE_STALE_MISSED_HEARTBEATS: u64 = 3;
const POOL_MONITOR_INTERVAL: Duration = Duration::from_secs(5 * 60);
const OAUTH_TOKEN_REFRESH_INTERVAL: Duration = Duration::from_secs(60);
const PROVIDER_CHECKIN_CONCURRENCY: usize = 3;
const PROVIDER_CHECKIN_DEFAULT_TIME: &str = "01:05";
const REQUEST_CANDIDATE_CLEANUP_INTERVAL: Duration = Duration::from_secs(24 * 60 * 60);
@@ -105,6 +111,7 @@ const DB_MAINTENANCE_HOUR: u32 = 5;
const DB_MAINTENANCE_MINUTE: u32 = 0;
const MAINTENANCE_DEFAULT_TIMEZONE: &str = "Asia/Shanghai";
const DB_MAINTENANCE_TABLES: &[&str] = &["usage", "request_candidates", "audit_logs"];
const MAX_ADMIN_STATS_REBUILD_BUCKETS: usize = 100_000;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct UsageCleanupSettings {
@@ -116,6 +123,76 @@ struct UsageCleanupSettings {
auto_delete_expired_keys: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Default, serde::Serialize)]
pub(crate) struct AdminSystemCleanupSummary {
pub(crate) audit_logs_deleted: usize,
pub(crate) request_candidates_deleted: usize,
pub(crate) pending_failed: usize,
pub(crate) pending_recovered: usize,
pub(crate) usage: UsageCleanupSummary,
}
#[derive(Debug, Clone, PartialEq, Eq, Default, serde::Serialize)]
pub(crate) struct AdminStatsRebuildSummary {
pub(crate) hourly_buckets: usize,
pub(crate) daily_buckets: usize,
pub(crate) capped: bool,
}
pub(crate) async fn run_admin_system_cleanup_once(
data: &GatewayDataState,
) -> Result<AdminSystemCleanupSummary, aether_data::DataLayerError> {
let audit_logs_deleted = cleanup_audit_logs_once(data).await?;
let request_candidates_deleted = cleanup_request_candidates_once(data).await?;
let pending = cleanup_stale_pending_requests_once(data).await?;
let usage = perform_usage_cleanup_once(data).await?;
Ok(AdminSystemCleanupSummary {
audit_logs_deleted,
request_candidates_deleted,
pending_failed: pending.failed,
pending_recovered: pending.recovered,
usage,
})
}
pub(crate) async fn rebuild_admin_stats_once(
data: &GatewayDataState,
) -> Result<AdminStatsRebuildSummary, aether_data::DataLayerError> {
let now_utc = chrono::Utc::now();
let mut summary = AdminStatsRebuildSummary::default();
if data.has_stats_hourly_aggregation_backend() {
let input = aether_data::StatsHourlyAggregationInput {
target_hour_utc: stats_hourly_aggregation_target_hour(now_utc),
aggregated_at: now_utc,
};
while data.aggregate_stats_hourly(&input).await?.is_some() {
summary.hourly_buckets = summary.hourly_buckets.saturating_add(1);
if summary.hourly_buckets >= MAX_ADMIN_STATS_REBUILD_BUCKETS {
summary.capped = true;
break;
}
}
}
if data.has_stats_daily_aggregation_backend() {
let input = aether_data::StatsDailyAggregationInput {
target_day_utc: stats_aggregation_target_day(now_utc),
aggregated_at: now_utc,
};
while data.aggregate_stats_daily(&input).await?.is_some() {
summary.daily_buckets = summary.daily_buckets.saturating_add(1);
if summary.daily_buckets >= MAX_ADMIN_STATS_REBUILD_BUCKETS {
summary.capped = true;
break;
}
}
}
Ok(summary)
}
pub(crate) async fn cleanup_expired_gemini_file_mappings_once(
data: &GatewayDataState,
) -> Result<usize, aether_data::DataLayerError> {

278
apps/aether-gateway/src/maintenance/runtime/oauth_token_refresh.rs Normal file
View File

@@ -0,0 +1,278 @@
use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use serde_json::Value;
use tracing::{info, warn};
use crate::provider_key_auth::provider_key_is_oauth_managed;
use crate::{AppState, GatewayError};
use super::system_config_bool;
const OAUTH_TOKEN_REFRESH_LOOKAHEAD_SECS: u64 = 120;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize)]
pub(crate) struct OAuthTokenRefreshRunSummary {
pub(crate) scanned: usize,
pub(crate) eligible: usize,
pub(crate) refreshed: usize,
pub(crate) resolved: usize,
pub(crate) skipped: usize,
pub(crate) failed: usize,
}
pub(crate) async fn perform_oauth_token_refresh_once(
state: &AppState,
) -> Result<OAuthTokenRefreshRunSummary, GatewayError> {
if !state.has_provider_catalog_data_reader() || !state.has_provider_catalog_data_writer() {
return Ok(OAuthTokenRefreshRunSummary::default());
}
if !system_config_bool(&state.data, "enable_oauth_token_refresh", true)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?
{
return Ok(OAuthTokenRefreshRunSummary::default());
}
let providers = state.list_provider_catalog_providers(true).await?;
let provider_ids = providers
.iter()
.map(|provider| provider.id.clone())
.collect::<Vec<_>>();
if provider_ids.is_empty() {
return Ok(OAuthTokenRefreshRunSummary::default());
}
let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(&provider_ids)
.await?;
let keys = state
.list_provider_catalog_keys_by_provider_ids(&provider_ids)
.await?;
let endpoints_by_provider = group_endpoints_by_provider(endpoints);
let keys_by_provider = group_keys_by_provider(keys);
let mut summary = OAuthTokenRefreshRunSummary::default();
let refresh_cutoff_unix_secs =
now_unix_secs().saturating_add(OAUTH_TOKEN_REFRESH_LOOKAHEAD_SECS);
for provider in providers {
let provider_keys = keys_by_provider
.get(provider.id.as_str())
.map(Vec::as_slice)
.unwrap_or(&[]);
let provider_endpoints = endpoints_by_provider
.get(provider.id.as_str())
.map(Vec::as_slice)
.unwrap_or(&[]);
for key in provider_keys {
summary.scanned = summary.scanned.saturating_add(1);
if !oauth_refresh_candidate(&provider, key, refresh_cutoff_unix_secs) {
summary.skipped = summary.skipped.saturating_add(1);
continue;
}
summary.eligible = summary.eligible.saturating_add(1);
let Some(endpoint) =
oauth_runtime_endpoint_for_provider(&provider.provider_type, provider_endpoints)
else {
summary.skipped = summary.skipped.saturating_add(1);
continue;
};
let Some(transport) = state
.read_provider_transport_snapshot(&provider.id, &endpoint.id, &key.id)
.await?
else {
summary.skipped = summary.skipped.saturating_add(1);
continue;
};
if !auth_config_has_refresh_token(transport.key.decrypted_auth_config.as_deref()) {
summary.skipped = summary.skipped.saturating_add(1);
continue;
}
match state.resolve_local_oauth_request_auth(&transport).await {
Ok(Some(_auth)) => {
summary.resolved = summary.resolved.saturating_add(1);
if provider_key_credentials_changed(state, key).await? {
summary.refreshed = summary.refreshed.saturating_add(1);
}
}
Ok(None) => {
summary.skipped = summary.skipped.saturating_add(1);
}
Err(err) => {
summary.failed = summary.failed.saturating_add(1);
warn!(
event_name = "oauth_token_refresh_failed",
log_type = "ops",
worker = "oauth_token_refresh",
provider_id = %provider.id,
key_id = %key.id,
error = ?err,
"gateway oauth token auto refresh failed"
);
}
}
}
}
if summary.eligible > 0 || summary.refreshed > 0 || summary.failed > 0 {
info!(
event_name = "oauth_token_refresh_completed",
log_type = "ops",
worker = "oauth_token_refresh",
scanned = summary.scanned,
eligible = summary.eligible,
refreshed = summary.refreshed,
resolved = summary.resolved,
skipped = summary.skipped,
failed = summary.failed,
"gateway completed oauth token auto refresh scan"
);
}
Ok(summary)
}
fn group_endpoints_by_provider(
endpoints: Vec<StoredProviderCatalogEndpoint>,
) -> BTreeMap<String, Vec<StoredProviderCatalogEndpoint>> {
let mut grouped = BTreeMap::new();
for endpoint in endpoints {
grouped
.entry(endpoint.provider_id.clone())
.or_insert_with(Vec::new)
.push(endpoint);
}
grouped
}
fn group_keys_by_provider(
keys: Vec<StoredProviderCatalogKey>,
) -> BTreeMap<String, Vec<StoredProviderCatalogKey>> {
let mut grouped = BTreeMap::new();
for key in keys {
grouped
.entry(key.provider_id.clone())
.or_insert_with(Vec::new)
.push(key);
}
grouped
}
fn oauth_refresh_candidate(
provider: &StoredProviderCatalogProvider,
key: &StoredProviderCatalogKey,
refresh_cutoff_unix_secs: u64,
) -> bool {
key.is_active
&& key.oauth_invalid_at_unix_secs.is_none()
&& key
.encrypted_auth_config
.as_deref()
.map(str::trim)
.is_some_and(|value| !value.is_empty())
&& key
.expires_at_unix_secs
.is_some_and(|expires_at| expires_at <= refresh_cutoff_unix_secs)
&& provider_key_is_oauth_managed(key, provider.provider_type.as_str())
}
fn oauth_runtime_endpoint_for_provider(
provider_type: &str,
endpoints: &[StoredProviderCatalogEndpoint],
) -> Option<StoredProviderCatalogEndpoint> {
let provider_type = provider_type.trim().to_ascii_lowercase();
match provider_type.as_str() {
"codex" => endpoints
.iter()
.find(|endpoint| {
endpoint.is_active
&& crate::ai_serving::is_openai_responses_format(&endpoint.api_format)
})
.cloned(),
"chatgpt_web" => endpoints
.iter()
.find(|endpoint| {
endpoint.is_active
&& endpoint
.api_format
.trim()
.eq_ignore_ascii_case("openai:image")
})
.cloned(),
"antigravity" => endpoints
.iter()
.find(|endpoint| {
endpoint.is_active
&& endpoint
.api_format
.trim()
.eq_ignore_ascii_case("gemini:generate_content")
})
.cloned(),
"kiro" => endpoints
.iter()
.find(|endpoint| {
endpoint.is_active
&& endpoint
.api_format
.trim()
.eq_ignore_ascii_case("claude:messages")
})
.cloned()
.or_else(|| {
endpoints
.iter()
.find(|endpoint| endpoint.is_active)
.cloned()
}),
_ => endpoints
.iter()
.find(|endpoint| endpoint.is_active)
.cloned(),
}
}
async fn provider_key_credentials_changed(
state: &AppState,
before: &StoredProviderCatalogKey,
) -> Result<bool, GatewayError> {
let Some(after) = state
.list_provider_catalog_keys_by_ids(std::slice::from_ref(&before.id))
.await?
.into_iter()
.next()
else {
return Ok(false);
};
Ok(after.encrypted_api_key != before.encrypted_api_key
|| after.encrypted_auth_config != before.encrypted_auth_config
|| after.expires_at_unix_secs != before.expires_at_unix_secs)
}
fn auth_config_has_refresh_token(auth_config: Option<&str>) -> bool {
let Some(auth_config) = auth_config.map(str::trim).filter(|value| !value.is_empty()) else {
return false;
};
let Ok(value) = serde_json::from_str::<Value>(auth_config) else {
return false;
};
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())
}
fn now_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs())
.unwrap_or_default()
}

30
apps/aether-gateway/src/maintenance/runtime/tests.rs
View File

@@ -20,17 +20,17 @@ use super::{
pending_cleanup_batch_size, pending_cleanup_timeout_minutes, plan_pending_cleanup_batch,
provider_checkin_schedule, record_proxy_upgrade_traffic_success, run_db_maintenance_with,
run_proxy_upgrade_rollout_once, spawn_audit_cleanup_worker, spawn_db_maintenance_worker,
spawn_pending_cleanup_worker, spawn_pool_monitor_worker, spawn_pool_quota_probe_worker,
spawn_provider_checkin_worker, spawn_proxy_node_stale_cleanup_worker,
spawn_proxy_upgrade_rollout_worker, spawn_stats_aggregation_worker,
spawn_stats_hourly_aggregation_worker, spawn_usage_cleanup_worker,
spawn_wallet_daily_usage_aggregation_worker, start_proxy_upgrade_rollout,
stats_aggregation_target_day, stats_hourly_aggregation_target_hour, summarize_database_pool,
usage_cleanup_settings, usage_cleanup_window, wallet_daily_usage_aggregation_target, AppState,
DbMaintenanceRunSummary, FailedPendingUsageRow, GatewayDataState,
ProxyUpgradeRolloutProbeConfig, StalePendingUsageRow, UsageCleanupSettings, USAGE_CLEANUP_HOUR,
USAGE_CLEANUP_MINUTE, WALLET_DAILY_USAGE_AGGREGATION_HOUR,
WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
spawn_oauth_token_refresh_worker, spawn_pending_cleanup_worker, spawn_pool_monitor_worker,
spawn_pool_quota_probe_worker, spawn_provider_checkin_worker,
spawn_proxy_node_stale_cleanup_worker, spawn_proxy_upgrade_rollout_worker,
spawn_stats_aggregation_worker, spawn_stats_hourly_aggregation_worker,
spawn_usage_cleanup_worker, spawn_wallet_daily_usage_aggregation_worker,
start_proxy_upgrade_rollout, stats_aggregation_target_day,
stats_hourly_aggregation_target_hour, summarize_database_pool, usage_cleanup_settings,
usage_cleanup_window, wallet_daily_usage_aggregation_target, AppState, DbMaintenanceRunSummary,
FailedPendingUsageRow, GatewayDataState, ProxyUpgradeRolloutProbeConfig, StalePendingUsageRow,
UsageCleanupSettings, USAGE_CLEANUP_HOUR, USAGE_CLEANUP_MINUTE,
WALLET_DAILY_USAGE_AGGREGATION_HOUR, WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
};
#[tokio::test]
@@ -63,6 +63,14 @@ async fn spawn_proxy_upgrade_rollout_worker_skips_when_proxy_nodes_unavailable()
assert!(spawn_proxy_upgrade_rollout_worker(state).is_none());
}
#[tokio::test]
async fn spawn_oauth_token_refresh_worker_skips_when_provider_catalog_unavailable() {
let state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(GatewayDataState::disabled());
assert!(spawn_oauth_token_refresh_worker(state).is_none());
}
#[tokio::test]
async fn spawn_proxy_upgrade_rollout_worker_skips_when_system_config_unavailable() {
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![]));

41
apps/aether-gateway/src/maintenance/runtime/workers.rs
View File

@@ -9,16 +9,18 @@ use crate::AppState;
use super::{
duration_until_next_daily_run, duration_until_next_db_maintenance_run,
duration_until_next_stats_aggregation_run, duration_until_next_stats_hourly_aggregation_run,
maintenance_timezone, parse_hhmm_time, provider_checkin_schedule, run_audit_cleanup_once,
run_db_maintenance_once, run_gemini_file_mapping_cleanup_once, run_pending_cleanup_once,
run_pool_monitor_once, run_provider_checkin_once, run_proxy_node_stale_cleanup_once,
run_proxy_upgrade_rollout_once, run_request_candidate_cleanup_once, run_stats_aggregation_once,
maintenance_timezone, parse_hhmm_time, perform_oauth_token_refresh_once,
provider_checkin_schedule, run_audit_cleanup_once, run_db_maintenance_once,
run_gemini_file_mapping_cleanup_once, run_pending_cleanup_once, run_pool_monitor_once,
run_provider_checkin_once, run_proxy_node_stale_cleanup_once, run_proxy_upgrade_rollout_once,
run_request_candidate_cleanup_once, run_stats_aggregation_once,
run_stats_hourly_aggregation_once, run_usage_cleanup_once,
run_wallet_daily_usage_aggregation_once, AUDIT_LOG_CLEANUP_INTERVAL,
GEMINI_FILE_MAPPING_CLEANUP_INTERVAL, PENDING_CLEANUP_INTERVAL, POOL_MONITOR_INTERVAL,
PROVIDER_CHECKIN_DEFAULT_TIME, PROXY_NODE_STALE_SWEEP_INTERVAL, PROXY_UPGRADE_ROLLOUT_INTERVAL,
REQUEST_CANDIDATE_CLEANUP_INTERVAL, USAGE_CLEANUP_HOUR, USAGE_CLEANUP_MINUTE,
WALLET_DAILY_USAGE_AGGREGATION_HOUR, WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
GEMINI_FILE_MAPPING_CLEANUP_INTERVAL, OAUTH_TOKEN_REFRESH_INTERVAL, PENDING_CLEANUP_INTERVAL,
POOL_MONITOR_INTERVAL, PROVIDER_CHECKIN_DEFAULT_TIME, PROXY_NODE_STALE_SWEEP_INTERVAL,
PROXY_UPGRADE_ROLLOUT_INTERVAL, REQUEST_CANDIDATE_CLEANUP_INTERVAL, USAGE_CLEANUP_HOUR,
USAGE_CLEANUP_MINUTE, WALLET_DAILY_USAGE_AGGREGATION_HOUR,
WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
};
const STATS_DAILY_CATCH_UP_BURST_LIMIT: usize = 14;
@@ -198,6 +200,29 @@ pub(crate) fn spawn_provider_checkin_worker(
}))
}
pub(crate) fn spawn_oauth_token_refresh_worker(
state: AppState,
) -> Option<tokio::task::JoinHandle<()>> {
if !state.has_provider_catalog_data_reader() || !state.has_provider_catalog_data_writer() {
return None;
}
Some(tokio::spawn(async move {
if let Err(err) = perform_oauth_token_refresh_once(&state).await {
log_maintenance_worker_failure("oauth_token_refresh", "startup", &err);
}
let mut interval = tokio::time::interval(OAUTH_TOKEN_REFRESH_INTERVAL);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
loop {
interval.tick().await;
if let Err(err) = perform_oauth_token_refresh_once(&state).await {
log_maintenance_worker_failure("oauth_token_refresh", "tick", &err);
}
}
}))
}
pub(crate) fn spawn_gemini_file_mapping_cleanup_worker(
data: Arc<GatewayDataState>,
) -> Option<tokio::task::JoinHandle<()>> {

42
apps/aether-gateway/src/state/core.rs
View File

@@ -41,6 +41,7 @@ use super::super::{provider_transport, usage};
use crate::maintenance::spawn_audit_cleanup_worker;
use crate::maintenance::spawn_db_maintenance_worker;
use crate::maintenance::spawn_gemini_file_mapping_cleanup_worker;
use crate::maintenance::spawn_oauth_token_refresh_worker;
use crate::maintenance::spawn_pending_cleanup_worker;
use crate::maintenance::spawn_pool_monitor_worker;
use crate::maintenance::spawn_pool_quota_probe_worker;
@@ -526,6 +527,44 @@ impl AppState {
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn purge_admin_system_data(
&self,
target: aether_data::repository::system::AdminSystemPurgeTarget,
) -> Result<aether_data::repository::system::AdminSystemPurgeSummary, GatewayError> {
let summary = self
.data
.purge_admin_system_data(target)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if matches!(
target,
aether_data::repository::system::AdminSystemPurgeTarget::Config
| aether_data::repository::system::AdminSystemPurgeTarget::Users
| aether_data::repository::system::AdminSystemPurgeTarget::Usage
| aether_data::repository::system::AdminSystemPurgeTarget::Stats
) {
self.system_config_cache.clear();
self.clear_provider_transport_snapshot_cache();
}
Ok(summary)
}
pub(crate) async fn run_admin_system_cleanup_once(
&self,
) -> Result<crate::maintenance::AdminSystemCleanupSummary, GatewayError> {
crate::maintenance::run_admin_system_cleanup_once(&self.data)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn rebuild_admin_stats_once(
&self,
) -> Result<crate::maintenance::AdminStatsRebuildSummary, GatewayError> {
crate::maintenance::rebuild_admin_stats_once(&self.data)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn find_proxy_node(
&self,
node_id: &str,
@@ -946,6 +985,9 @@ impl AppState {
if let Some(handle) = spawn_provider_checkin_worker(self.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_oauth_token_refresh_worker(self.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_request_candidate_cleanup_worker(self.data.clone()) {
tasks.push(handle);
}

28
apps/aether-gateway/src/tests/control/admin/system.rs
View File

@@ -791,9 +791,11 @@ async fn gateway_handles_admin_system_unavailable_write_routes_locally_with_trus
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let paths = [
let unavailable_paths = [
"/api/admin/system/config/import",
"/api/admin/system/users/import",
];
let local_paths = [
"/api/admin/system/cleanup",
"/api/admin/system/purge/config",
"/api/admin/system/purge/users",
@@ -803,7 +805,7 @@ async fn gateway_handles_admin_system_unavailable_write_routes_locally_with_trus
"/api/admin/system/purge/stats",
];
for path in paths {
for path in unavailable_paths {
let response = client
.post(format!("{gateway_url}{path}"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
@@ -820,6 +822,28 @@ async fn gateway_handles_admin_system_unavailable_write_routes_locally_with_trus
assert_eq!(payload["detail"], DETAIL, "path={path}");
}
for path in local_paths {
let response = client
.post(format!("{gateway_url}{path}"))
.header(crate::constants::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!({}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK, "path={path}");
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert!(
payload["message"]
.as_str()
.is_some_and(|value| !value.is_empty()),
"path={path}"
);
}
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();

68
apps/aether-gateway/src/tests/control/admin/users.rs
View File

@@ -313,6 +313,74 @@ async fn gateway_handles_admin_users_root_locally_with_trusted_admin_principal()
assert_eq!(items[0]["is_active"], true);
assert_eq!(items[0]["request_count"], 2);
assert_eq!(items[0]["total_tokens"], 100);
let search_response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/users?skip=0&limit=20&search=carol"
))
.header(crate::constants::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")
.send()
.await
.expect("search request should succeed");
assert_eq!(search_response.status(), StatusCode::OK);
let search_payload: serde_json::Value = search_response
.json()
.await
.expect("search json body should parse");
let search_items = search_payload
.as_array()
.expect("search list payload should be array");
assert_eq!(search_items.len(), 1);
assert_eq!(search_items[0]["id"], "user-3");
assert_eq!(search_items[0]["email"], "carol@example.com");
let id_search_response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/users?skip=0&limit=20&search=user-3"
))
.header(crate::constants::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")
.send()
.await
.expect("id search request should succeed");
assert_eq!(id_search_response.status(), StatusCode::OK);
let id_search_payload: serde_json::Value = id_search_response
.json()
.await
.expect("id search json body should parse");
let id_search_items = id_search_payload
.as_array()
.expect("id search list payload should be array");
assert_eq!(id_search_items.len(), 1);
assert_eq!(id_search_items[0]["id"], "user-3");
let limited_search_response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/users?skip=0&limit=2&search=example.com"
))
.header(crate::constants::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")
.send()
.await
.expect("limited search request should succeed");
assert_eq!(limited_search_response.status(), StatusCode::OK);
let limited_search_payload: serde_json::Value = limited_search_response
.json()
.await
.expect("limited search json body should parse");
let limited_search_items = limited_search_payload
.as_array()
.expect("limited search list payload should be array");
assert_eq!(limited_search_items.len(), 2);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();