i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-03 01:40:21 +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

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