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

fix(admin): repair system maintenance controls

Browse Source 
Implement server-side searchable user filtering for usage records, including backend search parameters and a shared frontend selector with loading, empty, and pinned-selected states.

Fix announcement deletion by cascading announcement read rows through a Postgres migration and defensive repository cleanup across supported backends.

Wire admin system purge and cleanup endpoints to real data deletion/maintenance flows, improve DataManagement messages, rebuild stats after stats purge, and add runtime OAuth token refresh maintenance when enabled.

Verified with rust-ci equivalent checks: cargo fmt, split clippy, split cargo tests, SQLite/Postgres/MySQL smoke tests, plus frontend npm ci, build, pages build, type-check, and targeted usage selector tests.
This commit is contained in:
Entropy.Xu
2026-05-07 21:26:40 +08:00
parent 2b439094c5
commit dd8c2ebec6
37 changed files with 2836 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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<()>> {