Merge remote-tracking branch 'upstream/aether-rust-pioneer' into fix-management-token-oauth-jsonb

# Conflicts:
#	crates/aether-data/src/lifecycle/bootstrap/postgres.rs
#	crates/aether-data/src/lifecycle/migrate/tests.rs
This commit is contained in:
Entropy.Xu
2026-05-10 16:21:09 +08:00
93 changed files with 7604 additions and 443 deletions

View File

@@ -23,6 +23,7 @@ aether-provider-transport.workspace = true
aether-scheduler-core.workspace = true
aether-runtime.workspace = true
aether-runtime-state.workspace = true
aether-task-runtime.workspace = true
aether-usage-runtime.workspace = true
aether-video-tasks-core.workspace = true
aether-wallet.workspace = true

View File

@@ -152,6 +152,11 @@ const PERMISSION_GROUPS: &[PermissionGroup] = &[
label: "系统",
assignable: true,
},
PermissionGroup {
scope: "tasks",
label: "后台任务",
assignable: true,
},
PermissionGroup {
scope: "usage",
label: "用量",
@@ -434,6 +439,9 @@ fn permission_key(scope: &str, access: &str) -> &'static str {
("system", "read") => "admin:system:read",
("system", "write") => "admin:system:write",
("system", "admin") => "admin:system:admin",
("tasks", "read") => "admin:tasks:read",
("tasks", "write") => "admin:tasks:write",
("tasks", "admin") => "admin:tasks:admin",
("usage", "read") => "admin:usage:read",
("usage", "write") => "admin:usage:write",
("usage", "admin") => "admin:usage:admin",
@@ -492,6 +500,7 @@ mod tests {
"security",
"stats",
"system",
"tasks",
"usage",
"users",
"video_tasks",

View File

@@ -46,6 +46,83 @@ pub(super) fn classify_admin_operations_family_route(
"admin:video_tasks",
false,
))
} else if method == http::Method::GET
&& matches!(normalized_path, "/api/admin/tasks" | "/api/admin/tasks/")
{
Some(classified(
"admin_proxy",
"tasks_manage",
"list_tasks",
"admin:tasks",
false,
))
} else if method == http::Method::GET
&& matches!(
normalized_path,
"/api/admin/tasks/stats" | "/api/admin/tasks/stats/"
)
{
Some(classified(
"admin_proxy",
"tasks_manage",
"stats",
"admin:tasks",
false,
))
} else if method == http::Method::GET
&& normalized_path.starts_with("/api/admin/tasks/")
&& normalized_path.ends_with("/events")
&& normalized_path.matches('/').count() == 5
{
Some(classified(
"admin_proxy",
"tasks_manage",
"events",
"admin:tasks",
false,
))
} else if method == http::Method::POST
&& normalized_path.starts_with("/api/admin/tasks/")
&& normalized_path.ends_with("/cancel")
&& normalized_path.matches('/').count() == 5
{
Some(classified(
"admin_proxy",
"tasks_manage",
"cancel",
"admin:tasks",
false,
))
} else if method == http::Method::POST
&& normalized_path.starts_with("/api/admin/tasks/")
&& normalized_path.ends_with("/trigger")
&& normalized_path.matches('/').count() == 5
{
Some(classified(
"admin_proxy",
"tasks_manage",
"trigger",
"admin:tasks",
false,
))
} else if method == http::Method::GET
&& normalized_path.starts_with("/api/admin/tasks/")
&& normalized_path["/api/admin/tasks/".len()..]
.split('/')
.count()
== 1
&& !matches!(
normalized_path,
"/api/admin/tasks/stats" | "/api/admin/tasks/stats/"
)
{
Some(classified(
"admin_proxy",
"tasks_manage",
"detail",
"admin:tasks",
false,
))
} else if method == http::Method::GET
&& normalized_path.starts_with("/api/admin/video-tasks/")
&& normalized_path.ends_with("/video")

View File

@@ -95,6 +95,14 @@ pub(super) fn classify_admin_system_family_route(
"admin:system",
false,
))
} else if method == http::Method::GET && normalized_path == "/api/admin/system/cleanup/runs" {
Some(classified(
"admin_proxy",
"system_manage",
"cleanup_runs",
"admin:system",
false,
))
} else if method == http::Method::POST && normalized_path == "/api/admin/system/purge/config" {
Some(classified(
"admin_proxy",
@@ -139,6 +147,16 @@ pub(super) fn classify_admin_system_family_route(
"admin:system",
false,
))
} else if method == http::Method::POST
&& normalized_path == "/api/admin/system/purge/request-bodies/task"
{
Some(classified(
"admin_proxy",
"system_manage",
"purge_request_bodies_task",
"admin:system",
false,
))
} else if method == http::Method::POST && normalized_path == "/api/admin/system/purge/stats" {
Some(classified(
"admin_proxy",

View File

@@ -172,6 +172,10 @@ fn classifies_admin_system_maintenance_write_routes_as_admin_proxy_route() {
"/api/admin/system/purge/request-bodies",
"purge_request_bodies",
),
(
"/api/admin/system/purge/request-bodies/task",
"purge_request_bodies_task",
),
("/api/admin/system/purge/stats", "purge_stats"),
];
@@ -246,6 +250,25 @@ fn classifies_admin_system_stats_as_admin_proxy_route() {
assert!(!decision.is_execution_runtime_candidate());
}
#[test]
fn classifies_admin_system_cleanup_runs_as_admin_proxy_route() {
let headers = headers(&[]);
let uri: Uri = "/api/admin/system/cleanup/runs"
.parse()
.expect("uri should parse");
let decision =
classify_control_route(&http::Method::GET, &uri, &headers).expect("route should classify");
assert_eq!(decision.route_class.as_deref(), Some("admin_proxy"));
assert_eq!(decision.route_family.as_deref(), Some("system_manage"));
assert_eq!(decision.route_kind.as_deref(), Some("cleanup_runs"));
assert_eq!(
decision.auth_endpoint_signature.as_deref(),
Some("admin:system")
);
assert!(!decision.is_execution_runtime_candidate());
}
#[test]
fn classifies_admin_system_settings_set_as_admin_proxy_route() {
let headers = headers(&[]);

View File

@@ -34,6 +34,8 @@ impl GatewayDataState {
proxy_node_reader: None,
proxy_node_writer: None,
billing_reader: None,
background_task_reader: None,
background_task_writer: None,
gemini_file_mapping_reader: None,
gemini_file_mapping_writer: None,
global_model_reader: None,
@@ -73,6 +75,8 @@ impl GatewayDataState {
let proxy_node_reader = backends.read().proxy_nodes();
let proxy_node_writer = backends.write().proxy_nodes();
let billing_reader = backends.read().billing();
let background_task_reader = backends.read().background_tasks();
let background_task_writer = backends.write().background_tasks();
let gemini_file_mapping_reader = backends.read().gemini_file_mappings();
let global_model_reader = backends.read().global_models();
let global_model_writer = backends.write().global_models();
@@ -110,6 +114,8 @@ impl GatewayDataState {
proxy_node_reader,
proxy_node_writer,
billing_reader,
background_task_reader,
background_task_writer,
gemini_file_mapping_reader,
gemini_file_mapping_writer,
global_model_reader,
@@ -197,6 +203,14 @@ impl GatewayDataState {
self.announcement_writer.is_some()
}
pub(crate) fn has_background_task_reader(&self) -> bool {
self.background_task_reader.is_some()
}
pub(crate) fn has_background_task_writer(&self) -> bool {
self.background_task_writer.is_some()
}
pub(crate) fn has_audit_log_reader(&self) -> bool {
self.backends
.as_ref()
@@ -447,4 +461,14 @@ impl GatewayDataState {
None => Ok(aether_data::repository::system::AdminSystemPurgeSummary::default()),
}
}
pub(crate) async fn purge_admin_request_bodies_batch(
&self,
batch_size: usize,
) -> Result<aether_data::repository::system::AdminSystemPurgeSummary, DataLayerError> {
match self.backends.as_ref() {
Some(backends) => backends.purge_admin_request_bodies_batch(batch_size).await,
None => Ok(aether_data::repository::system::AdminSystemPurgeSummary::default()),
}
}
}

View File

@@ -81,6 +81,11 @@ use aether_data::{
DataBackends, DataLayerError, DatabaseMaintenanceSummary, WalletDailyUsageAggregationInput,
WalletDailyUsageAggregationResult,
};
use aether_data_contracts::repository::background_tasks::{
BackgroundTaskListQuery, BackgroundTaskReadRepository, BackgroundTaskSummary,
BackgroundTaskWriteRepository, StoredBackgroundTaskEvent, StoredBackgroundTaskRun,
StoredBackgroundTaskRunPage, UpsertBackgroundTaskEvent, UpsertBackgroundTaskRun,
};
use aether_data_contracts::repository::billing::{
AdminBillingCollectorRecord, AdminBillingCollectorWriteInput, AdminBillingMutationOutcome,
AdminBillingPresetApplyResult, AdminBillingRuleRecord, AdminBillingRuleWriteInput,
@@ -140,6 +145,8 @@ pub(crate) struct GatewayDataState {
proxy_node_reader: Option<Arc<dyn ProxyNodeReadRepository>>,
proxy_node_writer: Option<Arc<dyn ProxyNodeWriteRepository>>,
billing_reader: Option<Arc<dyn BillingReadRepository>>,
background_task_reader: Option<Arc<dyn BackgroundTaskReadRepository>>,
background_task_writer: Option<Arc<dyn BackgroundTaskWriteRepository>>,
gemini_file_mapping_reader: Option<Arc<dyn GeminiFileMappingReadRepository>>,
gemini_file_mapping_writer: Option<Arc<dyn GeminiFileMappingWriteRepository>>,
global_model_reader: Option<Arc<dyn GlobalModelReadRepository>>,
@@ -206,6 +213,14 @@ impl fmt::Debug for GatewayDataState {
.field("has_proxy_node_reader", &self.proxy_node_reader.is_some())
.field("has_proxy_node_writer", &self.proxy_node_writer.is_some())
.field("has_billing_reader", &self.billing_reader.is_some())
.field(
"has_background_task_reader",
&self.background_task_reader.is_some(),
)
.field(
"has_background_task_writer",
&self.background_task_writer.is_some(),
)
.field(
"has_gemini_file_mapping_reader",
&self.gemini_file_mapping_reader.is_some(),

View File

@@ -5,11 +5,12 @@ use super::{
AdminBillingRuleWriteInput, AdminPaymentOrderListQuery, AdminRedeemCodeBatchListQuery,
AdminRedeemCodeListQuery, AdminWalletLedgerQuery, AdminWalletListQuery,
AdminWalletRefundRequestListQuery, AnnouncementListQuery, AuditLogListQuery,
CompleteAdminWalletRefundInput, CreateAdminRedeemCodeBatchInput,
CreateAdminRedeemCodeBatchResult, CreateAnnouncementRecord, CreateManualWalletRechargeInput,
CreateWalletRechargeOrderInput, CreateWalletRechargeOrderOutcome,
CreateWalletRefundRequestInput, CreateWalletRefundRequestOutcome, CreditAdminPaymentOrderInput,
DataLayerError, DatabaseMaintenanceSummary, DecisionTrace, DeleteAdminRedeemCodeBatchInput,
BackgroundTaskListQuery, BackgroundTaskSummary, CompleteAdminWalletRefundInput,
CreateAdminRedeemCodeBatchInput, CreateAdminRedeemCodeBatchResult, CreateAnnouncementRecord,
CreateManualWalletRechargeInput, CreateWalletRechargeOrderInput,
CreateWalletRechargeOrderOutcome, CreateWalletRefundRequestInput,
CreateWalletRefundRequestOutcome, CreditAdminPaymentOrderInput, DataLayerError,
DatabaseMaintenanceSummary, DecisionTrace, DeleteAdminRedeemCodeBatchInput,
DisableAdminRedeemCodeBatchInput, DisableAdminRedeemCodeInput, FailAdminWalletRefundInput,
GatewayDataState, GatewayProviderTransportSnapshot, LocalVideoTaskReadResponse,
ProcessAdminWalletRefundInput, ProcessPaymentCallbackInput, ProcessPaymentCallbackOutcome,
@@ -19,15 +20,16 @@ use super::{
StoredAdminRedeemCodePage, StoredAdminWalletLedgerPage, StoredAdminWalletListPage,
StoredAdminWalletRefund, StoredAdminWalletRefundPage, StoredAdminWalletRefundRequestPage,
StoredAdminWalletTransaction, StoredAdminWalletTransactionPage, StoredAnnouncement,
StoredAnnouncementPage, StoredBillingModelContext, StoredProviderQuotaSnapshot,
StoredAnnouncementPage, StoredBackgroundTaskEvent, StoredBackgroundTaskRun,
StoredBackgroundTaskRunPage, StoredBillingModelContext, StoredProviderQuotaSnapshot,
StoredProviderUsageSummary, StoredRequestUsageAudit, StoredSuspiciousActivity,
StoredUsageSettlement, StoredUserAuditLogPage, StoredUserAuthRecord, StoredUserExportRow,
StoredUserSummary, StoredVideoTask, StoredWalletDailyUsageLedger,
StoredWalletDailyUsageLedgerPage, StoredWalletSnapshot, UpdateAnnouncementRecord,
UpsertUsageRecord, UpsertVideoTask, UsageSettlementInput, VideoTaskLookupKey,
VideoTaskModelCount, VideoTaskQueryFilter, VideoTaskStatusCount,
WalletDailyUsageAggregationInput, WalletDailyUsageAggregationResult, WalletLookupKey,
WalletMutationOutcome,
UpsertBackgroundTaskEvent, UpsertBackgroundTaskRun, UpsertUsageRecord, UpsertVideoTask,
UsageSettlementInput, VideoTaskLookupKey, VideoTaskModelCount, VideoTaskQueryFilter,
VideoTaskStatusCount, WalletDailyUsageAggregationInput, WalletDailyUsageAggregationResult,
WalletLookupKey, WalletMutationOutcome,
};
use aether_data_contracts::repository::usage::{
PendingUsageCleanupSummary, ProviderApiKeyWindowUsageRequest,
@@ -1716,6 +1718,82 @@ impl GatewayDataState {
) -> Result<Option<LocalVideoTaskReadResponse>, DataLayerError> {
read_data_backed_video_task_response(self, route_family, request_path).await
}
pub(crate) async fn find_background_task_run(
&self,
run_id: &str,
) -> Result<Option<StoredBackgroundTaskRun>, DataLayerError> {
match &self.background_task_reader {
Some(repository) => repository.find_run(run_id).await,
None => Ok(None),
}
}
pub(crate) async fn list_background_task_runs(
&self,
query: &BackgroundTaskListQuery,
) -> Result<StoredBackgroundTaskRunPage, DataLayerError> {
match &self.background_task_reader {
Some(repository) => repository.list_runs(query).await,
None => Ok(StoredBackgroundTaskRunPage::default()),
}
}
pub(crate) async fn list_background_task_events(
&self,
run_id: &str,
offset: usize,
limit: usize,
) -> Result<Vec<StoredBackgroundTaskEvent>, DataLayerError> {
match &self.background_task_reader {
Some(repository) => repository.list_events(run_id, offset, limit).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn summarize_background_task_runs(
&self,
) -> Result<BackgroundTaskSummary, DataLayerError> {
match &self.background_task_reader {
Some(repository) => repository.summarize_runs().await,
None => Ok(BackgroundTaskSummary::default()),
}
}
pub(crate) async fn upsert_background_task_run(
&self,
run: UpsertBackgroundTaskRun,
) -> Result<Option<StoredBackgroundTaskRun>, DataLayerError> {
match &self.background_task_writer {
Some(repository) => repository.upsert_run(run).await.map(Some),
None => Ok(None),
}
}
pub(crate) async fn request_cancel_background_task_run(
&self,
run_id: &str,
updated_at_unix_secs: u64,
) -> Result<bool, DataLayerError> {
match &self.background_task_writer {
Some(repository) => {
repository
.request_cancel(run_id, updated_at_unix_secs)
.await
}
None => Ok(false),
}
}
pub(crate) async fn upsert_background_task_event(
&self,
event: UpsertBackgroundTaskEvent,
) -> Result<Option<StoredBackgroundTaskEvent>, DataLayerError> {
match &self.background_task_writer {
Some(repository) => repository.upsert_event(event).await.map(Some),
None => Ok(None),
}
}
}
#[cfg(test)]

View File

@@ -43,6 +43,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -93,6 +95,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,

View File

@@ -76,6 +76,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -125,6 +127,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -170,6 +174,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -303,6 +309,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -364,6 +372,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -418,6 +428,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -481,6 +493,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -526,6 +540,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -572,6 +588,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -629,6 +647,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -688,6 +708,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -731,6 +753,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -789,6 +813,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -840,6 +866,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: Some(wallet_reader),
wallet_writer: Some(wallet_writer),
settlement_writer: None,
@@ -896,6 +924,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: Some(wallet_reader),
wallet_writer: Some(wallet_writer),
settlement_writer: None,
@@ -953,6 +983,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: Some(wallet_reader),
wallet_writer: Some(wallet_writer),
settlement_writer: None,
@@ -1009,6 +1041,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1054,6 +1088,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1099,6 +1135,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1156,6 +1194,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1218,6 +1258,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1263,6 +1305,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1313,6 +1357,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1380,6 +1426,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1442,6 +1490,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1488,6 +1538,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1534,6 +1586,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1582,6 +1636,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1628,6 +1684,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1674,6 +1732,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1728,6 +1788,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1783,6 +1845,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1841,6 +1905,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1905,6 +1971,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1970,6 +2038,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -2039,6 +2109,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -2115,6 +2187,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: Some(wallet_reader),
wallet_writer: Some(wallet_writer),
settlement_writer: None,
@@ -2173,6 +2247,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -2222,6 +2298,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -2267,6 +2345,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -2318,6 +2398,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: Some(wallet_reader),
wallet_writer: Some(wallet_writer),
settlement_writer: None,
@@ -2373,6 +2455,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: Some(wallet_reader),
wallet_writer: Some(wallet_writer),
settlement_writer: None,
@@ -2429,6 +2513,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: Some(wallet_reader),
wallet_writer: Some(wallet_writer),
settlement_writer: None,
@@ -2478,6 +2564,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,

View File

@@ -47,6 +47,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: Some(repository),
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -99,6 +101,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: Some(video_task_reader),
video_task_writer: Some(video_task_writer),
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -148,6 +152,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: Some(video_task_reader),
video_task_writer: Some(video_task_writer),
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -201,6 +207,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: Some(video_task_reader),
video_task_writer: Some(video_task_writer),
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -258,6 +266,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: Some(video_task_reader),
video_task_writer: Some(video_task_writer),
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -324,6 +334,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: Some(video_task_reader),
video_task_writer: Some(video_task_writer),
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,

View File

@@ -0,0 +1,15 @@
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::GatewayError;
use axum::body::{Body, Bytes};
use axum::response::Response;
mod routes;
pub(crate) async fn maybe_build_local_admin_background_tasks_response(
state: &AdminAppState<'_>,
request_context: &AdminRequestContext<'_>,
request_body: Option<&Bytes>,
) -> Result<Option<Response<Body>>, GatewayError> {
routes::maybe_build_local_admin_background_tasks_response(state, request_context, request_body)
.await
}

View File

@@ -0,0 +1,373 @@
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::handlers::admin::shared::{
attach_admin_audit_response, query_param_value, unix_secs_to_rfc3339,
};
use crate::task_runtime::{
self, set_cancel_signal, TASK_KEY_PROVIDER_DELETE, TASK_KEY_PROVIDER_OAUTH_BATCH_IMPORT,
};
use crate::GatewayError;
use aether_data_contracts::repository::background_tasks::{
BackgroundTaskKind, BackgroundTaskListQuery, BackgroundTaskStatus,
};
use axum::{
body::{Body, Bytes},
http,
response::{IntoResponse, Response},
Json,
};
use serde_json::json;
const DEFAULT_PAGE_SIZE: usize = 20;
const MAX_PAGE_SIZE: usize = 100;
const DEFAULT_EVENTS_PAGE_SIZE: usize = 50;
pub(super) async fn maybe_build_local_admin_background_tasks_response(
state: &AdminAppState<'_>,
request_context: &AdminRequestContext<'_>,
request_body: Option<&Bytes>,
) -> Result<Option<Response<Body>>, GatewayError> {
if request_context.route_family() != Some("tasks_manage") {
return Ok(None);
}
match request_context.route_kind() {
Some("list_tasks") if request_context.method() == http::Method::GET => {
let query = request_context.query_string();
let page = query_param_value(query, "page")
.and_then(|value| value.parse::<usize>().ok())
.unwrap_or(1)
.max(1);
let page_size = query_param_value(query, "page_size")
.and_then(|value| value.parse::<usize>().ok())
.unwrap_or(DEFAULT_PAGE_SIZE)
.clamp(1, MAX_PAGE_SIZE);
let kind = query_param_value(query, "kind")
.map(|value| BackgroundTaskKind::from_database(value.as_str()))
.transpose()
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let status = query_param_value(query, "status")
.map(|value| BackgroundTaskStatus::from_database(value.as_str()))
.transpose()
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let trigger = query_param_value(query, "trigger");
let task_key_substring = query_param_value(query, "task_key");
let offset = (page - 1).saturating_mul(page_size);
let response = state
.list_background_task_runs(&BackgroundTaskListQuery {
task_key_substring,
kind,
status,
trigger,
offset,
limit: page_size,
})
.await?;
let pages = if response.total == 0 {
0
} else {
(response.total + page_size - 1) / page_size
};
let items = response
.items
.iter()
.map(|run| {
json!({
"id": run.id,
"task_key": run.task_key,
"kind": run.kind.as_database(),
"trigger": run.trigger,
"status": run.status.as_database(),
"attempt": run.attempt,
"max_attempts": run.max_attempts,
"owner_instance": run.owner_instance,
"progress_percent": run.progress_percent,
"progress_message": run.progress_message,
"payload": run.payload_json,
"result": run.result_json,
"error_message": run.error_message,
"cancel_requested": run.cancel_requested,
"created_by": run.created_by,
"created_at": unix_secs_to_rfc3339(run.created_at_unix_secs),
"started_at": run.started_at_unix_secs.and_then(unix_secs_to_rfc3339),
"finished_at": run.finished_at_unix_secs.and_then(unix_secs_to_rfc3339),
"updated_at": unix_secs_to_rfc3339(run.updated_at_unix_secs),
})
})
.collect::<Vec<_>>();
let definitions = task_runtime::task_definitions()
.iter()
.map(|definition| {
json!({
"task_key": definition.key,
"kind": definition.kind.as_str(),
"trigger": definition.trigger,
"max_attempts": definition.retry_policy.max_attempts,
"singleton": definition.singleton,
"persist_history": definition.persist_history,
})
})
.collect::<Vec<_>>();
return Ok(Some(
Json(json!({
"items": items,
"total": response.total,
"page": page,
"page_size": page_size,
"pages": pages,
"definitions": definitions,
}))
.into_response(),
));
}
Some("stats") if request_context.method() == http::Method::GET => {
let stats = state.summarize_background_task_runs().await?;
return Ok(Some(
Json(json!({
"total": stats.total,
"running_count": stats.running_count,
"by_status": stats.by_status,
"by_kind": stats.by_kind,
"registered_tasks": task_runtime::task_definitions().len(),
}))
.into_response(),
));
}
Some("detail") if request_context.method() == http::Method::GET => {
let Some(run_id) = task_id_from_path(request_context.path()) else {
return Ok(Some(
(
http::StatusCode::NOT_FOUND,
Json(json!({"detail":"Task not found"})),
)
.into_response(),
));
};
let Some(run) = state.find_background_task_run(run_id).await? else {
return Ok(Some(
(
http::StatusCode::NOT_FOUND,
Json(json!({"detail":"Task not found"})),
)
.into_response(),
));
};
return Ok(Some(attach_admin_audit_response(
Json(json!({
"id": run.id,
"task_key": run.task_key,
"kind": run.kind.as_database(),
"trigger": run.trigger,
"status": run.status.as_database(),
"attempt": run.attempt,
"max_attempts": run.max_attempts,
"owner_instance": run.owner_instance,
"progress_percent": run.progress_percent,
"progress_message": run.progress_message,
"payload": run.payload_json,
"result": run.result_json,
"error_message": run.error_message,
"cancel_requested": run.cancel_requested,
"created_by": run.created_by,
"created_at": unix_secs_to_rfc3339(run.created_at_unix_secs),
"started_at": run.started_at_unix_secs.and_then(unix_secs_to_rfc3339),
"finished_at": run.finished_at_unix_secs.and_then(unix_secs_to_rfc3339),
"updated_at": unix_secs_to_rfc3339(run.updated_at_unix_secs),
}))
.into_response(),
"admin_task_detail_viewed",
"view_task_detail",
"background_task",
run_id,
)));
}
Some("events") if request_context.method() == http::Method::GET => {
let Some(run_id) = nested_task_id_from_path(request_context.path(), "/events") else {
return Ok(Some(
(
http::StatusCode::NOT_FOUND,
Json(json!({"detail":"Task not found"})),
)
.into_response(),
));
};
let query = request_context.query_string();
let page = query_param_value(query, "page")
.and_then(|value| value.parse::<usize>().ok())
.unwrap_or(1)
.max(1);
let page_size = query_param_value(query, "page_size")
.and_then(|value| value.parse::<usize>().ok())
.unwrap_or(DEFAULT_EVENTS_PAGE_SIZE)
.clamp(1, MAX_PAGE_SIZE);
let offset = (page - 1).saturating_mul(page_size);
let events = state
.list_background_task_events(run_id, offset, page_size)
.await?;
return Ok(Some(
Json(json!({
"items": events.into_iter().map(|event| {
json!({
"id": event.id,
"run_id": event.run_id,
"event_type": event.event_type,
"message": event.message,
"payload": event.payload_json,
"created_at": unix_secs_to_rfc3339(event.created_at_unix_secs),
})
}).collect::<Vec<_>>(),
"page": page,
"page_size": page_size,
}))
.into_response(),
));
}
Some("cancel") if request_context.method() == http::Method::POST => {
let Some(run_id) = nested_task_id_from_path(request_context.path(), "/cancel") else {
return Ok(Some(
(
http::StatusCode::NOT_FOUND,
Json(json!({"detail":"Task not found"})),
)
.into_response(),
));
};
let now = task_runtime::now_unix_secs();
let cancelled = state
.request_cancel_background_task_run(run_id, now)
.await?;
if !cancelled {
return Ok(Some(
(
http::StatusCode::NOT_FOUND,
Json(json!({ "detail": "Task not found" })),
)
.into_response(),
));
}
let _ = set_cancel_signal(state.app(), run_id).await;
task_runtime::append_event_with_logging(
state.app(),
run_id,
"cancel_requested",
"cancel requested by admin",
None,
)
.await;
return Ok(Some(attach_admin_audit_response(
Json(json!({
"id": run_id,
"status": "cancel_requested",
"message": "Task cancellation requested",
}))
.into_response(),
"admin_task_cancel_requested",
"cancel_task",
"background_task",
run_id,
)));
}
Some("trigger") if request_context.method() == http::Method::POST => {
let Some(task_key) = nested_task_id_from_path(request_context.path(), "/trigger")
else {
return Ok(Some(
(
http::StatusCode::NOT_FOUND,
Json(json!({"detail":"Task not found"})),
)
.into_response(),
));
};
let payload = parse_json_payload(request_body)?;
if task_key == TASK_KEY_PROVIDER_DELETE {
let provider_id = payload
.get("provider_id")
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or_else(|| {
GatewayError::Internal(
"admin task trigger provider delete requires provider_id".to_string(),
)
})?;
let Some(run_id) =
task_runtime::submit_provider_delete_task(state, provider_id, Some("admin"))
.await?
else {
return Ok(Some(
(
http::StatusCode::NOT_FOUND,
Json(json!({"detail":"Provider 不存在"})),
)
.into_response(),
));
};
return Ok(Some(attach_admin_audit_response(
Json(json!({
"task_key": task_key,
"run_id": run_id,
"status": "queued",
}))
.into_response(),
"admin_task_triggered",
"trigger_task",
"background_task",
task_key,
)));
}
if task_key == TASK_KEY_PROVIDER_OAUTH_BATCH_IMPORT {
return Ok(Some(
(
http::StatusCode::BAD_REQUEST,
Json(json!({
"detail": "请使用 provider oauth batch import 专用接口触发该任务",
})),
)
.into_response(),
));
}
return Ok(Some(
(
http::StatusCode::BAD_REQUEST,
Json(json!({
"detail": format!("Unsupported task_key: {task_key}"),
})),
)
.into_response(),
));
}
_ => {}
}
Ok(None)
}
fn task_id_from_path(request_path: &str) -> Option<&str> {
let task_id = request_path.strip_prefix("/api/admin/tasks/")?;
if task_id.is_empty() || task_id.contains('/') || task_id == "stats" {
return None;
}
Some(task_id)
}
fn nested_task_id_from_path<'a>(request_path: &'a str, suffix: &str) -> Option<&'a str> {
let task_id = request_path
.strip_prefix("/api/admin/tasks/")?
.strip_suffix(suffix)?;
if task_id.is_empty() || task_id.contains('/') {
return None;
}
Some(task_id)
}
fn parse_json_payload(request_body: Option<&Bytes>) -> Result<serde_json::Value, GatewayError> {
let Some(body) = request_body else {
return Ok(json!({}));
};
if body.is_empty() {
return Ok(json!({}));
}
serde_json::from_slice::<serde_json::Value>(body)
.map_err(|err| GatewayError::Internal(format!("invalid json body: {err}")))
}

View File

@@ -1,7 +1,9 @@
mod background_tasks;
mod gemini_files;
mod routes;
mod video_tasks;
pub(super) use self::background_tasks::maybe_build_local_admin_background_tasks_response;
pub(super) use self::gemini_files::maybe_build_local_admin_gemini_files_response;
pub(super) use self::routes::maybe_build_local_admin_features_response;
pub(crate) use self::video_tasks::maybe_build_local_admin_video_tasks_response;

View File

@@ -1,9 +1,19 @@
use super::{gemini_files, video_tasks};
use super::{background_tasks, gemini_files, video_tasks};
use crate::handlers::admin::request::{AdminRouteRequest, AdminRouteResult};
pub(crate) async fn maybe_build_local_admin_features_response(
request: AdminRouteRequest<'_>,
) -> AdminRouteResult {
if let Some(response) = background_tasks::maybe_build_local_admin_background_tasks_response(
&request.state(),
&request.request_context(),
request.request_body(),
)
.await?
{
return Ok(Some(response));
}
if let Some(response) = video_tasks::maybe_build_local_admin_video_tasks_response(
&request.state(),
&request.request_context(),

View File

@@ -1,11 +1,10 @@
use crate::handlers::admin::provider::delete_task::run_admin_provider_delete_task;
use crate::handlers::admin::provider::shared::paths::{
admin_provider_delete_task_parts, admin_provider_id_for_manage_path,
};
use crate::handlers::admin::provider::shared::support::build_admin_provider_delete_task_payload;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::handlers::admin::shared::attach_admin_audit_response;
use crate::{GatewayError, LocalProviderDeleteTaskState};
use crate::GatewayError;
use axum::{
body::Body,
http,
@@ -13,8 +12,6 @@ use axum::{
Json,
};
use serde_json::json;
use tracing::warn;
use uuid::Uuid;
fn build_admin_provider_not_found_response(detail: impl Into<String>) -> Response<Body> {
(
@@ -60,52 +57,18 @@ pub(crate) async fn maybe_build_local_admin_provider_delete_task_response(
"Provider 不存在",
)));
};
let Some(provider) = state
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&provider_id))
.await?
.into_iter()
.next()
let Some(task_id) =
crate::task_runtime::submit_provider_delete_task(state, &provider_id, Some("admin"))
.await?
else {
return Ok(Some(build_admin_provider_not_found_response(
"提供商不存在",
)));
};
let task_id = Uuid::new_v4().simple().to_string()[..16].to_string();
let pending_task = LocalProviderDeleteTaskState {
task_id: task_id.clone(),
provider_id: provider.id.clone(),
status: "pending".to_string(),
stage: "queued".to_string(),
total_keys: 0,
deleted_keys: 0,
total_endpoints: 0,
deleted_endpoints: 0,
message: "delete task submitted".to_string(),
};
state.put_provider_delete_task(pending_task.clone());
if let Err(err) = state
.run_admin_provider_delete_task(&provider.id, &task_id)
.await
{
warn!(
"gateway admin provider delete task failed for provider {}: {:?}",
provider.id, err
);
state.put_provider_delete_task(LocalProviderDeleteTaskState {
task_id: task_id.clone(),
provider_id: provider.id.clone(),
status: "failed".to_string(),
stage: "failed".to_string(),
total_keys: 0,
deleted_keys: 0,
total_endpoints: 0,
deleted_endpoints: 0,
message: format!("provider delete failed: {err:?}"),
});
}
return Ok(Some(attach_admin_audit_response(
Json(json!({
"task_id": task_id,
"run_id": task_id,
"status": "pending",
"message": "删除任务已提交,提供商已进入后台删除队列",
}))
@@ -113,7 +76,7 @@ pub(crate) async fn maybe_build_local_admin_provider_delete_task_response(
"admin_provider_delete_queued",
"delete_provider",
"provider",
&provider.id,
&provider_id,
)));
}

View File

@@ -15,7 +15,14 @@ use crate::handlers::admin::provider::oauth::state::{
};
use crate::handlers::admin::provider::shared::paths::admin_provider_oauth_batch_import_task_provider_id;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::task_runtime::{
append_event_with_logging, now_unix_secs, task_definition, update_run_status,
upsert_run_with_logging, TASK_KEY_PROVIDER_OAUTH_BATCH_IMPORT,
};
use crate::GatewayError;
use aether_data_contracts::repository::background_tasks::{
BackgroundTaskKind, BackgroundTaskStatus, UpsertBackgroundTaskRun,
};
use axum::{
body::Bytes,
http,
@@ -133,11 +140,7 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
}
let task_id = Uuid::new_v4().to_string();
let created_at = SystemTime::now()
.duration_since(UNIX_EPOCH)
.ok()
.map(|duration| duration.as_secs())
.unwrap_or(0);
let created_at = now_unix_secs();
let submitted_state = build_admin_provider_oauth_batch_task_state(
&task_id,
&provider_id,
@@ -167,6 +170,50 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
));
}
if state.has_background_task_data_writer() {
let max_attempts = task_definition(TASK_KEY_PROVIDER_OAUTH_BATCH_IMPORT)
.map(|item| item.retry_policy.max_attempts)
.unwrap_or(1);
let run = UpsertBackgroundTaskRun {
id: task_id.clone(),
task_key: TASK_KEY_PROVIDER_OAUTH_BATCH_IMPORT.to_string(),
kind: BackgroundTaskKind::OnDemand,
trigger: "manual".to_string(),
status: BackgroundTaskStatus::Queued,
attempt: 1,
max_attempts,
owner_instance: Some(state.app().tunnel.local_instance_id().to_string()),
progress_percent: 0,
progress_message: Some("provider oauth batch import queued".to_string()),
payload_json: Some(json!({
"provider_id": provider_id.clone(),
"provider_type": provider_type.clone(),
"total": total,
})),
result_json: None,
error_message: None,
cancel_requested: false,
created_by: Some("admin".to_string()),
created_at_unix_secs: created_at,
started_at_unix_secs: None,
finished_at_unix_secs: None,
updated_at_unix_secs: created_at,
};
let _ = upsert_run_with_logging(state.app(), run).await;
append_event_with_logging(
state.app(),
&task_id,
"queued",
"provider oauth batch import queued",
Some(json!({
"provider_id": provider_id.clone(),
"provider_type": provider_type.clone(),
"total": total,
})),
)
.await;
}
let task_state = state.cloned_app();
let task_id_for_worker = task_id.clone();
let provider_id_for_worker = provider_id.clone();
@@ -201,6 +248,27 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
.save_provider_oauth_batch_task_payload(&task_id_for_worker, &processing_state)
.await;
let _ = update_run_status(
&task_state,
&task_id_for_worker,
BackgroundTaskStatus::Running,
Some(1),
Some("provider oauth batch import started".to_string()),
None,
None,
Some(started_at),
None,
)
.await;
append_event_with_logging(
&task_state,
&task_id_for_worker,
"running",
"provider oauth batch import started",
None,
)
.await;
let mut progress_reporter = BatchTaskProgressReporter {
app: task_state.clone(),
task_id: task_id_for_worker.clone(),
@@ -270,6 +338,34 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
let _ = AdminAppState::new(&task_state)
.save_provider_oauth_batch_task_payload(&task_id_for_worker, &completed_state)
.await;
let _ = update_run_status(
&task_state,
&task_id_for_worker,
BackgroundTaskStatus::Succeeded,
Some(100),
Some(message),
Some(json!({
"provider_id": provider_id_for_worker,
"provider_type": provider_type_for_worker,
"total": outcome.total,
"success": outcome.success,
"failed": outcome.failed,
"created_count": created_count,
"replaced_count": replaced_count,
})),
None,
None,
Some(finished_at),
)
.await;
append_event_with_logging(
&task_state,
&task_id_for_worker,
"succeeded",
"provider oauth batch import completed",
None,
)
.await;
}
Err(err) => {
let finished_at = SystemTime::now()
@@ -299,6 +395,26 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
let _ = AdminAppState::new(&task_state)
.save_provider_oauth_batch_task_payload(&task_id_for_worker, &failed_state)
.await;
let _ = update_run_status(
&task_state,
&task_id_for_worker,
BackgroundTaskStatus::Failed,
Some(100),
Some("provider oauth batch import failed".to_string()),
None,
Some(error_message.clone()),
None,
Some(finished_at),
)
.await;
append_event_with_logging(
&task_state,
&task_id_for_worker,
"failed",
"provider oauth batch import failed",
Some(json!({ "error": error_message.clone() })),
)
.await;
tracing::warn!(
task_id = %task_id_for_worker,
provider_id = %provider_id_for_worker,

View File

@@ -4,6 +4,7 @@ use super::quota::codex::refresh_codex_provider_quota_locally;
use super::quota::kiro::refresh_kiro_provider_quota_locally;
use crate::handlers::admin::request::AdminAppState;
use crate::provider_key_auth::provider_key_is_oauth_managed;
use crate::task_runtime::{spawn_fire_and_forget, TASK_KEY_PROVIDER_OAUTH_ACCOUNT_REFRESH};
use crate::{AppState, GatewayError};
use aether_contracts::ProxySnapshot;
use aether_data_contracts::repository::provider_catalog::{
@@ -170,7 +171,7 @@ pub(crate) fn spawn_provider_oauth_account_state_refresh_after_update(
key_id: String,
proxy_override: Option<ProxySnapshot>,
) {
tokio::spawn(async move {
spawn_fire_and_forget(TASK_KEY_PROVIDER_OAUTH_ACCOUNT_REFRESH, async move {
let _ = refresh_provider_oauth_account_state_after_update(
&AdminAppState::new(&app),
&provider,

View File

@@ -1,5 +1,6 @@
use super::actions::admin_provider_ops_local_action_response;
use crate::handlers::admin::request::AdminAppState;
use crate::task_runtime::{spawn_fire_and_forget, TASK_KEY_PROVIDER_BALANCE_REFRESH};
use serde_json::{json, Value};
use std::collections::HashSet;
use std::time::Duration;
@@ -130,7 +131,7 @@ pub(super) async fn spawn_admin_provider_ops_balance_refresh(
let app = state.cloned_app();
let provider_id = provider_id.to_string();
tokio::spawn(async move {
spawn_fire_and_forget(TASK_KEY_PROVIDER_BALANCE_REFRESH, async move {
let permit = match tokio::time::timeout(
Duration::from_secs(5),
ADMIN_PROVIDER_OPS_BALANCE_REFRESH_SEMAPHORE.acquire(),

View File

@@ -76,6 +76,14 @@ impl<'a> AdminAppState<'a> {
self.app.has_gemini_file_mapping_data_writer()
}
pub(crate) fn has_background_task_data_reader(&self) -> bool {
self.app.has_background_task_data_reader()
}
pub(crate) fn has_background_task_data_writer(&self) -> bool {
self.app.has_background_task_data_writer()
}
pub(crate) fn has_auth_api_key_data_reader(&self) -> bool {
self.app.has_auth_api_key_data_reader()
}

View File

@@ -2,6 +2,59 @@ use super::{AdminAppState, AdminCancelVideoTaskError};
use crate::GatewayError;
impl<'a> AdminAppState<'a> {
pub(crate) async fn list_background_task_runs(
&self,
query: &aether_data_contracts::repository::background_tasks::BackgroundTaskListQuery,
) -> Result<
aether_data_contracts::repository::background_tasks::StoredBackgroundTaskRunPage,
GatewayError,
> {
self.app.list_background_task_runs(query).await
}
pub(crate) async fn find_background_task_run(
&self,
run_id: &str,
) -> Result<
Option<aether_data_contracts::repository::background_tasks::StoredBackgroundTaskRun>,
GatewayError,
> {
self.app.find_background_task_run(run_id).await
}
pub(crate) async fn list_background_task_events(
&self,
run_id: &str,
offset: usize,
limit: usize,
) -> Result<
Vec<aether_data_contracts::repository::background_tasks::StoredBackgroundTaskEvent>,
GatewayError,
> {
self.app
.list_background_task_events(run_id, offset, limit)
.await
}
pub(crate) async fn summarize_background_task_runs(
&self,
) -> Result<
aether_data_contracts::repository::background_tasks::BackgroundTaskSummary,
GatewayError,
> {
self.app.summarize_background_task_runs().await
}
pub(crate) async fn request_cancel_background_task_run(
&self,
run_id: &str,
updated_at_unix_secs: u64,
) -> Result<bool, GatewayError> {
self.app
.request_cancel_background_task_run(run_id, updated_at_unix_secs)
.await
}
pub(crate) async fn list_gemini_file_mappings(
&self,
query: &aether_data::repository::gemini_file_mappings::GeminiFileMappingListQuery,

View File

@@ -18,7 +18,6 @@ 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,
@@ -26,6 +25,7 @@ use axum::{
Json,
};
use serde_json::json;
use std::time::Instant;
pub(super) async fn maybe_build_local_admin_core_system_response(
state: &AdminAppState<'_>,
@@ -194,15 +194,31 @@ pub(super) async fn maybe_build_local_admin_core_system_response(
)));
}
if let Some((target, action, object_type, object_id)) =
admin_system_purge_target_for_route_kind(decision.route_kind.as_deref())
if decision.route_kind.as_deref() == Some("cleanup_runs")
&& request_method == http::Method::GET
&& request_path == "/api/admin/system/cleanup/runs"
{
let records = crate::maintenance::list_admin_cleanup_run_records(&state.app().data)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
return Ok(Some(Json(json!({ "items": records })).into_response()));
}
if let Some((task_kind, action, object_type, object_id)) =
admin_system_purge_task_for_route_kind(decision.route_kind.as_deref())
{
if request_method != http::Method::POST {
return Ok(None);
}
let task = crate::maintenance::start_admin_system_purge_task(state.cloned_app(), task_kind)
.await?;
return Ok(Some(attach_admin_audit_response(
Json(build_admin_system_purge_payload(state, target).await?).into_response(),
"admin_system_data_purged",
Json(json!({
"message": task.message.clone(),
"task": task,
}))
.into_response(),
"admin_system_purge_task_started",
action,
object_type,
object_id,
@@ -435,100 +451,60 @@ pub(super) async fn maybe_build_local_admin_core_system_response(
Ok(None)
}
fn admin_system_purge_target_for_route_kind(
fn admin_system_purge_task_for_route_kind(
route_kind: Option<&str>,
) -> Option<(
AdminSystemPurgeTarget,
crate::maintenance::AdminCleanupTaskKind,
&'static str,
&'static str,
&'static str,
)> {
match route_kind {
Some("purge_config") => Some((
AdminSystemPurgeTarget::Config,
"purge_system_config",
crate::maintenance::AdminCleanupTaskKind::Config,
"purge_system_config_async",
"system_config",
"global",
)),
Some("purge_users") => Some((
AdminSystemPurgeTarget::Users,
"purge_non_admin_users",
crate::maintenance::AdminCleanupTaskKind::Users,
"purge_non_admin_users_async",
"users",
"non_admin",
)),
Some("purge_usage") => Some((
AdminSystemPurgeTarget::Usage,
"purge_usage_records",
crate::maintenance::AdminCleanupTaskKind::Usage,
"purge_usage_records_async",
"usage",
"all",
)),
Some("purge_audit_logs") => Some((
AdminSystemPurgeTarget::AuditLogs,
"purge_audit_logs",
crate::maintenance::AdminCleanupTaskKind::AuditLogs,
"purge_audit_logs_async",
"audit_logs",
"all",
)),
Some("purge_request_bodies") => Some((
AdminSystemPurgeTarget::RequestBodies,
"purge_request_bodies",
Some("purge_request_bodies") | Some("purge_request_bodies_task") => Some((
crate::maintenance::AdminCleanupTaskKind::RequestBodies,
"purge_request_bodies_async",
"request_bodies",
"all",
)),
Some("purge_stats") => Some((AdminSystemPurgeTarget::Stats, "purge_stats", "stats", "all")),
Some("purge_stats") => Some((
crate::maintenance::AdminCleanupTaskKind::Stats,
"purge_stats_async",
"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 started_at_unix_secs = chrono::Utc::now().timestamp().max(0) as u64;
let started_at = Instant::now();
let summary = state.run_admin_system_cleanup_once().await?;
let cleaned = json!({
"audit_logs": summary.audit_logs_deleted,
@@ -558,6 +534,17 @@ async fn build_admin_system_cleanup_payload(
.saturating_add(summary.usage.keys_cleaned)
.saturating_add(summary.usage.records_deleted);
crate::maintenance::record_completed_cleanup_run(
&state.app().data,
"system_cleanup",
"manual",
started_at_unix_secs,
started_at,
cleaned.clone(),
format!("系统清理已执行,影响 {total}"),
)
.await;
Ok(json!({
"message": format!("系统清理已执行,影响 {} 项", total),
"cleaned": cleaned,

View File

@@ -635,9 +635,9 @@ fn preserve_quota_window_usage_state(current_status_snapshot: Option<&Value>, qu
}
fn codex_default_window_minutes(code: &str) -> Option<u64> {
if code.eq_ignore_ascii_case("5h") {
if code.eq_ignore_ascii_case("5h") || code.eq_ignore_ascii_case("spark_5h") {
Some(300)
} else if code.eq_ignore_ascii_case("weekly") {
} else if code.eq_ignore_ascii_case("weekly") || code.eq_ignore_ascii_case("spark_weekly") {
Some(10_080)
} else {
None
@@ -735,6 +735,20 @@ fn build_codex_quota_status_snapshot(
let windows = [
codex_quota_window_snapshot(metadata, "primary", "weekly", "", observed_at_unix_secs),
codex_quota_window_snapshot(metadata, "secondary", "5h", "5H", observed_at_unix_secs),
codex_quota_window_snapshot(
metadata,
"spark_primary",
"spark_5h",
"Spark 5H",
observed_at_unix_secs,
),
codex_quota_window_snapshot(
metadata,
"spark_secondary",
"spark_weekly",
"Spark 周",
observed_at_unix_secs,
),
]
.into_iter()
.flatten()
@@ -750,21 +764,34 @@ fn build_codex_quota_status_snapshot(
return None;
}
let usage_ratio = windows
let primary_windows = windows
.iter()
.filter(|window| {
window
.get("code")
.and_then(Value::as_str)
.is_some_and(|code| {
code.eq_ignore_ascii_case("weekly") || code.eq_ignore_ascii_case("5h")
})
})
.cloned()
.collect::<Vec<_>>();
let usage_ratio = primary_windows
.iter()
.filter_map(Value::as_object)
.filter_map(|window| window.get("used_ratio"))
.filter_map(Value::as_f64)
.max_by(f64::total_cmp);
let reset_seconds = windows
let reset_seconds = primary_windows
.iter()
.filter_map(Value::as_object)
.filter_map(|window| window.get("reset_seconds"))
.filter_map(admin_provider_quota_pure::coerce_json_u64)
.min();
let reset_at = quota_windows_min_reset_at(&windows);
let exhausted_by_credits =
windows.is_empty() && credits_unlimited != Some(true) && credits_has_credits == Some(false);
let reset_at = quota_windows_min_reset_at(&primary_windows);
let exhausted_by_credits = primary_windows.is_empty()
&& credits_unlimited != Some(true)
&& credits_has_credits == Some(false);
let exhausted_by_window = usage_ratio.is_some_and(|value| value >= 1.0 - 1e-6);
let exhausted = exhausted_by_credits || exhausted_by_window;
@@ -1926,6 +1953,48 @@ mod tests {
);
}
#[test]
fn provider_key_status_snapshot_payload_backfills_codex_spark_windows() {
let mut key = sample_catalog_key();
key.upstream_metadata = Some(json!({
"codex": {
"updated_at": 1_775_553_285u64,
"plan_type": "plus",
"primary_used_percent": 55.0,
"primary_reset_at": 1_900_000_000u64,
"secondary_used_percent": 12.5,
"secondary_reset_at": 1_900_500_000u64,
"spark_primary_used_percent": 40.0,
"spark_primary_reset_at": 1_900_100_000u64,
"spark_primary_window_minutes": 300u64,
"spark_secondary_used_percent": 5.0,
"spark_secondary_reset_at": 1_900_600_000u64,
"spark_secondary_window_minutes": 10_080u64
}
}));
let payload = provider_key_status_snapshot_payload(&key, "codex");
let windows = payload["quota"]["windows"]
.as_array()
.expect("quota windows should exist");
let spark_5h = windows
.iter()
.filter_map(Value::as_object)
.find(|window| window.get("code") == Some(&json!("spark_5h")))
.expect("Spark 5H window should exist");
let spark_weekly = windows
.iter()
.filter_map(Value::as_object)
.find(|window| window.get("code") == Some(&json!("spark_weekly")))
.expect("Spark weekly window should exist");
assert_eq!(windows.len(), 4);
assert_eq!(spark_5h.get("label"), Some(&json!("Spark 5H")));
assert_eq!(spark_5h.get("remaining_ratio"), Some(&json!(0.6)));
assert_eq!(spark_weekly.get("label"), Some(&json!("Spark 周")));
assert_eq!(spark_weekly.get("remaining_ratio"), Some(&json!(0.95)));
}
#[test]
fn provider_key_status_snapshot_payload_keeps_codex_free_window_quota_available() {
let mut key = sample_catalog_key();

View File

@@ -58,6 +58,7 @@ mod router;
mod scheduler;
mod state;
mod system_features;
mod task_runtime;
mod tunnel;
mod usage;
mod video_tasks;

View File

@@ -1202,13 +1202,13 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
state.bootstrap_admin_from_env().await?;
let background_tasks = if args.node_role.spawns_background_tasks() {
state.spawn_background_tasks()
Some(state.spawn_background_tasks())
} else {
info!(
node_role = args.node_role.as_str(),
"background workers disabled for this node role"
);
Vec::new()
None
};
let listener = tokio::net::TcpListener::bind(bind_addr).await?;
let public_base_url = resolve_local_http_base_url(app_port)?;
@@ -1241,8 +1241,8 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
router.into_make_service_with_connect_info::<std::net::SocketAddr>(),
)
.await?;
for handle in background_tasks {
handle.abort();
if let Some(background_tasks) = background_tasks {
background_tasks.shutdown().await;
}
Ok(())
}

View File

@@ -4,8 +4,9 @@ mod tests;
pub(crate) use runtime::{
cancel_proxy_upgrade_rollout, clear_proxy_upgrade_rollout_conflicts,
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,
inspect_proxy_upgrade_rollout, list_admin_cleanup_run_records,
perform_oauth_token_refresh_once, perform_pool_quota_probe_once, perform_provider_checkin_once,
rebuild_admin_stats_once, record_completed_cleanup_run, 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,
@@ -15,10 +16,11 @@ pub(crate) use runtime::{
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,
start_admin_request_body_cleanup_task, start_admin_system_purge_task,
start_proxy_upgrade_rollout, AdminCleanupRunRecord, AdminCleanupTaskKind,
AdminStatsRebuildSummary, AdminSystemCleanupSummary, OAuthTokenRefreshRunSummary,
PoolQuotaProbeRunSummary, ProviderCheckinRunSummary, ProxyUpgradeRolloutCancelSummary,
ProxyUpgradeRolloutConflictClearSummary, ProxyUpgradeRolloutNodeActionSummary,
ProxyUpgradeRolloutProbeConfig, ProxyUpgradeRolloutSkippedRestoreSummary,
ProxyUpgradeRolloutStatus, ProxyUpgradeRolloutTrackedNodeState,
};

View File

@@ -8,6 +8,8 @@ use crate::{AppState, GatewayError};
#[path = "runtime/audit_cleanup.rs"]
mod audit_cleanup;
#[path = "runtime/cleanup_runs.rs"]
mod cleanup_runs;
#[path = "runtime/config.rs"]
mod config;
#[path = "runtime/db_maintenance.rs"]
@@ -49,6 +51,11 @@ pub(crate) use aether_data_contracts::repository::usage::{
UsageCleanupSummary, UsageCleanupWindow,
};
use audit_cleanup::*;
pub(crate) use cleanup_runs::{
list_admin_cleanup_run_records, record_completed_cleanup_run, record_failed_cleanup_run,
start_admin_request_body_cleanup_task, start_admin_system_purge_task, AdminCleanupRunRecord,
AdminCleanupTaskKind,
};
use config::*;
use db_maintenance::*;
pub(crate) use oauth_token_refresh::{

View File

@@ -0,0 +1,518 @@
use std::time::Instant;
use aether_data::repository::system::{AdminSystemPurgeSummary, AdminSystemPurgeTarget};
use aether_data_contracts::DataLayerError;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use tracing::{info, warn};
use crate::data::GatewayDataState;
use crate::{AppState, GatewayError};
use super::{now_unix_secs, system_config_usize};
const CLEANUP_RUN_HISTORY_KEY: &str = "admin_cleanup_run_history";
const CLEANUP_RUN_HISTORY_LIMIT: usize = 50;
const REQUEST_BODY_PROGRESS_UPDATE_BATCHES: usize = 10;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub(crate) struct AdminCleanupRunRecord {
pub(crate) id: String,
pub(crate) kind: String,
pub(crate) trigger: String,
pub(crate) status: String,
pub(crate) message: String,
pub(crate) started_at_unix_secs: u64,
pub(crate) completed_at_unix_secs: Option<u64>,
pub(crate) duration_ms: Option<u64>,
pub(crate) summary: Value,
pub(crate) error: Option<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum AdminCleanupTaskKind {
Config,
Users,
RequestBodies,
Usage,
AuditLogs,
Stats,
}
impl AdminCleanupTaskKind {
fn record_kind(self) -> &'static str {
match self {
Self::Config => "config_purge",
Self::Users => "users_purge",
Self::RequestBodies => "request_bodies",
Self::Usage => "usage_purge",
Self::AuditLogs => "audit_logs_purge",
Self::Stats => "stats_purge",
}
}
fn target(self) -> Option<AdminSystemPurgeTarget> {
match self {
Self::RequestBodies => None,
Self::Config => Some(AdminSystemPurgeTarget::Config),
Self::Users => Some(AdminSystemPurgeTarget::Users),
Self::Usage => Some(AdminSystemPurgeTarget::Usage),
Self::AuditLogs => Some(AdminSystemPurgeTarget::AuditLogs),
Self::Stats => Some(AdminSystemPurgeTarget::Stats),
}
}
fn start_message(self, batch_size: Option<usize>) -> String {
match self {
Self::Config => "系统配置后台清空已开始".to_string(),
Self::Users => "非管理员用户后台清空已开始".to_string(),
Self::RequestBodies => format!(
"请求/响应体后台清理已开始,每批 {}",
batch_size.unwrap_or(1)
),
Self::Usage => "使用记录后台清空已开始".to_string(),
Self::AuditLogs => "审计日志后台清空已开始".to_string(),
Self::Stats => "统计聚合后台清空和重建已开始".to_string(),
}
}
fn failure_message(self) -> &'static str {
match self {
Self::Config => "系统配置后台清空失败",
Self::Users => "非管理员用户后台清空失败",
Self::RequestBodies => "请求/响应体后台清理失败",
Self::Usage => "使用记录后台清空失败",
Self::AuditLogs => "审计日志后台清空失败",
Self::Stats => "统计聚合后台清空或重建失败",
}
}
}
pub(crate) async fn list_admin_cleanup_run_records(
data: &GatewayDataState,
) -> Result<Vec<AdminCleanupRunRecord>, DataLayerError> {
let Some(value) = data
.find_system_config_value(CLEANUP_RUN_HISTORY_KEY)
.await?
else {
return Ok(Vec::new());
};
Ok(parse_cleanup_run_records(value))
}
pub(crate) async fn start_admin_request_body_cleanup_task(
app: AppState,
) -> Result<AdminCleanupRunRecord, GatewayError> {
start_admin_system_purge_task(app, AdminCleanupTaskKind::RequestBodies).await
}
pub(crate) async fn start_admin_system_purge_task(
app: AppState,
kind: AdminCleanupTaskKind,
) -> Result<AdminCleanupRunRecord, GatewayError> {
let data = app.data.clone();
let records = list_admin_cleanup_run_records(&data)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if let Some(existing) = records
.iter()
.find(|record| record.kind == kind.record_kind() && record.status == "processing")
.cloned()
{
return Ok(existing);
}
let batch_size = if kind == AdminCleanupTaskKind::RequestBodies {
Some(
system_config_usize(&data, "cleanup_batch_size", 1_000)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?
.max(1),
)
} else {
None
};
let started_at = now_unix_secs();
let record = AdminCleanupRunRecord {
id: uuid::Uuid::new_v4().to_string(),
kind: kind.record_kind().to_string(),
trigger: "manual".to_string(),
status: "processing".to_string(),
message: kind.start_message(batch_size),
started_at_unix_secs: started_at,
completed_at_unix_secs: None,
duration_ms: None,
summary: initial_task_summary(kind, batch_size),
error: None,
};
record_cleanup_run(&data, record.clone())
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
match kind {
AdminCleanupTaskKind::RequestBodies => {
tokio::spawn(run_request_body_cleanup_task(
data,
record.clone(),
batch_size.unwrap_or(1),
));
}
AdminCleanupTaskKind::Usage
| AdminCleanupTaskKind::AuditLogs
| AdminCleanupTaskKind::Config
| AdminCleanupTaskKind::Users
| AdminCleanupTaskKind::Stats => {
tokio::spawn(run_admin_system_purge_task(app, record.clone(), kind));
}
}
Ok(record)
}
pub(crate) async fn record_completed_cleanup_run(
data: &GatewayDataState,
kind: &str,
trigger: &str,
started_at_unix_secs: u64,
started_at: Instant,
summary: Value,
message: impl Into<String>,
) {
let record = AdminCleanupRunRecord {
id: uuid::Uuid::new_v4().to_string(),
kind: kind.to_string(),
trigger: trigger.to_string(),
status: "completed".to_string(),
message: message.into(),
started_at_unix_secs,
completed_at_unix_secs: Some(now_unix_secs()),
duration_ms: Some(
started_at
.elapsed()
.as_millis()
.try_into()
.unwrap_or(u64::MAX),
),
summary,
error: None,
};
if let Err(err) = record_cleanup_run(data, record).await {
warn!(error = %err, kind, "failed to record cleanup run");
}
}
pub(crate) async fn record_failed_cleanup_run(
data: &GatewayDataState,
kind: &str,
trigger: &str,
started_at_unix_secs: u64,
started_at: Instant,
error: &DataLayerError,
) {
let record = AdminCleanupRunRecord {
id: uuid::Uuid::new_v4().to_string(),
kind: kind.to_string(),
trigger: trigger.to_string(),
status: "failed".to_string(),
message: "清理执行失败".to_string(),
started_at_unix_secs,
completed_at_unix_secs: Some(now_unix_secs()),
duration_ms: Some(
started_at
.elapsed()
.as_millis()
.try_into()
.unwrap_or(u64::MAX),
),
summary: json!({}),
error: Some(error.to_string()),
};
if let Err(err) = record_cleanup_run(data, record).await {
warn!(error = %err, kind, "failed to record failed cleanup run");
}
}
async fn run_request_body_cleanup_task(
data: std::sync::Arc<GatewayDataState>,
initial_record: AdminCleanupRunRecord,
batch_size: usize,
) {
let started_at = Instant::now();
let mut total = AdminSystemPurgeSummary::default();
let mut batches = 0usize;
loop {
match data.purge_admin_request_bodies_batch(batch_size).await {
Ok(batch) if batch.total() == 0 => break,
Ok(batch) => {
batches = batches.saturating_add(1);
total.merge(&batch);
if batches.is_multiple_of(REQUEST_BODY_PROGRESS_UPDATE_BATCHES) {
let progress = request_body_cleanup_record(
&initial_record,
"processing",
format!(
"请求/响应体后台清理中,已处理 {} 批,影响 {}",
batches,
total.total()
),
batches,
batch_size,
&total,
Some(started_at),
None,
);
if let Err(err) = record_cleanup_run(&data, progress).await {
warn!(error = %err, "failed to update request body cleanup progress");
}
}
tokio::task::yield_now().await;
}
Err(err) => {
let failed = request_body_cleanup_record(
&initial_record,
"failed",
"请求/响应体后台清理失败".to_string(),
batches,
batch_size,
&total,
Some(started_at),
Some(err.to_string()),
);
if let Err(record_err) = record_cleanup_run(&data, failed).await {
warn!(error = %record_err, "failed to record request body cleanup failure");
}
warn!(error = %err, "request body cleanup task failed");
return;
}
}
}
let completed = request_body_cleanup_record(
&initial_record,
"completed",
format!(
"请求/响应体后台清理完成,影响 {} 行,共 {}",
total.total(),
batches
),
batches,
batch_size,
&total,
Some(started_at),
None,
);
if let Err(err) = record_cleanup_run(&data, completed).await {
warn!(error = %err, "failed to record request body cleanup completion");
}
info!(
event_name = "request_body_cleanup_task_completed",
log_type = "ops",
worker = "request_body_cleanup_task",
batches,
affected = total.total(),
"gateway finished request body cleanup task"
);
}
async fn run_admin_system_purge_task(
app: AppState,
initial_record: AdminCleanupRunRecord,
kind: AdminCleanupTaskKind,
) {
let started_at = Instant::now();
let data = app.data.clone();
match run_admin_system_purge_task_once(&app, kind).await {
Ok((summary, message)) => {
let completed = cleanup_task_record(
&initial_record,
"completed",
message,
summary,
Some(started_at),
None,
);
if let Err(err) = record_cleanup_run(&data, completed).await {
warn!(error = %err, "failed to record admin system purge task completion");
}
info!(
event_name = "admin_system_purge_task_completed",
log_type = "ops",
worker = "admin_system_purge_task",
kind = ?kind,
"gateway finished admin system purge task"
);
}
Err(err) => {
let failed = cleanup_task_record(
&initial_record,
"failed",
kind.failure_message().to_string(),
json!({}),
Some(started_at),
Some(format!("{err:?}")),
);
if let Err(record_err) = record_cleanup_run(&data, failed).await {
warn!(error = %record_err, "failed to record admin system purge task failure");
}
warn!(error = ?err, kind = ?kind, "admin system purge task failed");
}
}
}
async fn run_admin_system_purge_task_once(
app: &AppState,
kind: AdminCleanupTaskKind,
) -> Result<(Value, String), GatewayError> {
let target = kind
.target()
.expect("non-request-body cleanup task should map to purge target");
let purge = app.purge_admin_system_data(target).await?;
let deleted_total = purge.total();
let affected = purge.affected.clone();
if kind == AdminCleanupTaskKind::Stats {
let rebuild = app.rebuild_admin_stats_once().await?;
let message = if rebuild.capped {
format!(
"统计聚合后台清空完成,已重建 {} 个小时桶和 {} 个日桶,仍有历史统计待后台任务继续重建",
rebuild.hourly_buckets, rebuild.daily_buckets
)
} else {
format!(
"统计聚合后台清空并重建完成,删除 {} 行,重建 {} 个小时桶和 {} 个日桶",
deleted_total, rebuild.hourly_buckets, rebuild.daily_buckets
)
};
return Ok((
json!({
"deleted": affected,
"rebuilt": {
"hourly_buckets": rebuild.hourly_buckets,
"daily_buckets": rebuild.daily_buckets,
"capped": rebuild.capped,
},
"total": deleted_total,
}),
message,
));
}
let message = match kind {
AdminCleanupTaskKind::Config => format!("系统配置后台清空完成,影响 {deleted_total}"),
AdminCleanupTaskKind::Users => format!("非管理员用户后台清空完成,影响 {deleted_total}"),
AdminCleanupTaskKind::Usage => format!("使用记录后台清空完成,影响 {deleted_total}"),
AdminCleanupTaskKind::AuditLogs => format!("审计日志后台清空完成,影响 {deleted_total}"),
AdminCleanupTaskKind::RequestBodies | AdminCleanupTaskKind::Stats => unreachable!(),
};
Ok((
json!({
"deleted": affected,
"total": deleted_total,
}),
message,
))
}
fn initial_task_summary(kind: AdminCleanupTaskKind, batch_size: Option<usize>) -> Value {
match kind {
AdminCleanupTaskKind::RequestBodies => json!({
"batch_size": batch_size.unwrap_or(1),
"batches": 0,
"cleaned": {},
}),
AdminCleanupTaskKind::Stats => json!({
"deleted": {},
"rebuilt": {},
}),
AdminCleanupTaskKind::Config
| AdminCleanupTaskKind::Users
| AdminCleanupTaskKind::Usage
| AdminCleanupTaskKind::AuditLogs => json!({
"deleted": {},
}),
}
}
fn cleanup_task_record(
initial: &AdminCleanupRunRecord,
status: &str,
message: String,
summary: Value,
started_at: Option<Instant>,
error: Option<String>,
) -> AdminCleanupRunRecord {
let completed = matches!(status, "completed" | "failed");
AdminCleanupRunRecord {
id: initial.id.clone(),
kind: initial.kind.clone(),
trigger: initial.trigger.clone(),
status: status.to_string(),
message,
started_at_unix_secs: initial.started_at_unix_secs,
completed_at_unix_secs: completed.then(now_unix_secs),
duration_ms: started_at
.map(|value| value.elapsed().as_millis().try_into().unwrap_or(u64::MAX)),
summary,
error,
}
}
fn request_body_cleanup_record(
initial: &AdminCleanupRunRecord,
status: &str,
message: String,
batches: usize,
batch_size: usize,
total: &AdminSystemPurgeSummary,
started_at: Option<Instant>,
error: Option<String>,
) -> AdminCleanupRunRecord {
let completed = matches!(status, "completed" | "failed");
AdminCleanupRunRecord {
id: initial.id.clone(),
kind: initial.kind.clone(),
trigger: initial.trigger.clone(),
status: status.to_string(),
message,
started_at_unix_secs: initial.started_at_unix_secs,
completed_at_unix_secs: completed.then(now_unix_secs),
duration_ms: started_at
.map(|value| value.elapsed().as_millis().try_into().unwrap_or(u64::MAX)),
summary: json!({
"batch_size": batch_size,
"batches": batches,
"cleaned": total.affected,
"total": total.total(),
}),
error,
}
}
async fn record_cleanup_run(
data: &GatewayDataState,
record: AdminCleanupRunRecord,
) -> Result<(), DataLayerError> {
let mut records = list_admin_cleanup_run_records(data).await?;
records.retain(|existing| existing.id != record.id);
records.insert(0, record);
records.truncate(CLEANUP_RUN_HISTORY_LIMIT);
let value = serde_json::to_value(records).map_err(|err| {
DataLayerError::UnexpectedValue(format!("invalid cleanup run history: {err}"))
})?;
data.upsert_system_config_entry(
CLEANUP_RUN_HISTORY_KEY,
&value,
Some("最近的系统清理执行记录"),
)
.await?;
Ok(())
}
fn parse_cleanup_run_records(value: Value) -> Vec<AdminCleanupRunRecord> {
value
.as_array()
.into_iter()
.flat_map(|items| items.iter())
.filter_map(|item| serde_json::from_value::<AdminCleanupRunRecord>(item.clone()).ok())
.collect()
}

View File

@@ -1,9 +1,14 @@
use aether_data_contracts::DataLayerError;
use serde_json::json;
use std::time::Instant;
use tracing::info;
use crate::data::GatewayDataState;
use super::{now_unix_secs, system_config_bool, system_config_u64, system_config_usize};
use super::{
now_unix_secs, record_completed_cleanup_run, record_failed_cleanup_run, system_config_bool,
system_config_u64, system_config_usize,
};
pub(crate) async fn cleanup_request_candidates_once(
data: &GatewayDataState,
@@ -47,7 +52,33 @@ pub(crate) async fn cleanup_request_candidates_once(
pub(super) async fn run_request_candidate_cleanup_once(
data: &GatewayDataState,
) -> Result<(), DataLayerError> {
let deleted = cleanup_request_candidates_once(data).await?;
let started_at_unix_secs = now_unix_secs();
let started_at = Instant::now();
let deleted = match cleanup_request_candidates_once(data).await {
Ok(deleted) => deleted,
Err(err) => {
record_failed_cleanup_run(
data,
"request_candidate_cleanup",
"auto",
started_at_unix_secs,
started_at,
&err,
)
.await;
return Err(err);
}
};
record_completed_cleanup_run(
data,
"request_candidate_cleanup",
"auto",
started_at_unix_secs,
started_at,
json!({ "request_candidates_deleted": deleted }),
format!("候选记录自动清理完成,删除 {deleted}"),
)
.await;
if deleted > 0 {
info!(deleted, "gateway deleted expired request candidates");
}

View File

@@ -1,4 +1,6 @@
use aether_data_contracts::DataLayerError;
use serde_json::json;
use std::time::Instant;
use tracing::{info, warn};
use crate::data::GatewayDataState;
@@ -8,15 +10,41 @@ use super::{
advance_proxy_upgrade_rollout_once, cleanup_audit_logs_once,
cleanup_expired_gemini_file_mappings_once, cleanup_proxy_node_metrics_once,
cleanup_request_candidates_once, cleanup_stale_pending_requests_once,
cleanup_stale_proxy_nodes_once, collect_proxy_upgrade_rollout_probes,
cleanup_stale_proxy_nodes_once, collect_proxy_upgrade_rollout_probes, now_unix_secs,
perform_db_maintenance_once, perform_provider_checkin_once, perform_stats_aggregation_once,
perform_stats_hourly_aggregation_once, perform_usage_cleanup_once,
perform_wallet_daily_usage_aggregation_once, record_proxy_upgrade_traffic_success,
summarize_database_pool,
perform_wallet_daily_usage_aggregation_once, record_completed_cleanup_run,
record_failed_cleanup_run, record_proxy_upgrade_traffic_success, summarize_database_pool,
};
pub(super) async fn run_audit_cleanup_once(data: &GatewayDataState) -> Result<(), DataLayerError> {
let deleted = cleanup_audit_logs_once(data).await?;
let started_at_unix_secs = now_unix_secs();
let started_at = Instant::now();
let deleted = match cleanup_audit_logs_once(data).await {
Ok(deleted) => deleted,
Err(err) => {
record_failed_cleanup_run(
data,
"audit_cleanup",
"auto",
started_at_unix_secs,
started_at,
&err,
)
.await;
return Err(err);
}
};
record_completed_cleanup_run(
data,
"audit_cleanup",
"auto",
started_at_unix_secs,
started_at,
json!({ "audit_logs_deleted": deleted }),
format!("审计日志自动清理完成,删除 {deleted}"),
)
.await;
if deleted > 0 {
info!(
event_name = "audit_cleanup_completed",
@@ -204,7 +232,49 @@ pub(super) async fn run_stats_aggregation_once(
}
pub(super) async fn run_usage_cleanup_once(data: &GatewayDataState) -> Result<(), DataLayerError> {
let summary = perform_usage_cleanup_once(data).await?;
let started_at_unix_secs = now_unix_secs();
let started_at = Instant::now();
let summary = match perform_usage_cleanup_once(data).await {
Ok(summary) => summary,
Err(err) => {
record_failed_cleanup_run(
data,
"usage_cleanup",
"auto",
started_at_unix_secs,
started_at,
&err,
)
.await;
return Err(err);
}
};
record_completed_cleanup_run(
data,
"usage_cleanup",
"auto",
started_at_unix_secs,
started_at,
json!({
"body_externalized": summary.body_externalized,
"legacy_body_refs_migrated": summary.legacy_body_refs_migrated,
"body_cleaned": summary.body_cleaned,
"header_cleaned": summary.header_cleaned,
"keys_cleaned": summary.keys_cleaned,
"records_deleted": summary.records_deleted,
}),
format!(
"请求记录自动清理完成,影响 {}",
summary
.body_externalized
.saturating_add(summary.legacy_body_refs_migrated)
.saturating_add(summary.body_cleaned)
.saturating_add(summary.header_cleaned)
.saturating_add(summary.keys_cleaned)
.saturating_add(summary.records_deleted)
),
)
.await;
if summary.body_externalized > 0
|| summary.legacy_body_refs_migrated > 0
|| summary.body_cleaned > 0

View File

@@ -119,9 +119,7 @@ async fn gateway_background_request_candidate_cleanup_deletes_expired_entries_in
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
for handle in background_tasks {
handle.abort();
}
background_tasks.shutdown().await;
}
#[tokio::test]

View File

@@ -539,8 +539,6 @@ async fn gateway_background_model_fetch_updates_key_and_syncs_provider_model_whi
.expect("execution runtime plan should be captured");
assert_eq!(seen_plan.url, "https://api.openai.example/v1/models");
for handle in background_tasks {
handle.abort();
}
background_tasks.shutdown().await;
execution_runtime_handle.abort();
}

View File

@@ -19,7 +19,6 @@ use aether_runtime_state::{
RuntimeSemaphoreSnapshot, RuntimeState,
};
use aether_scheduler_core::PROVIDER_KEY_RPM_WINDOW_SECS;
use tokio::task::JoinHandle;
use super::{AppState, FrontdoorCorsConfig, LocalExecutionRuntimeMissDiagnostic};
@@ -997,71 +996,112 @@ impl AppState {
Ok(self)
}
pub fn spawn_background_tasks(&self) -> Vec<JoinHandle<()>> {
let mut tasks = Vec::new();
if let Some(handle) = self.usage_runtime.spawn_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) =
crate::wallet_runtime::spawn_provider_quota_reset_worker(self.data.clone())
{
tasks.push(handle);
}
if let Some(handle) = spawn_audit_cleanup_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_db_maintenance_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_wallet_daily_usage_aggregation_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_stats_aggregation_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_usage_cleanup_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_pool_monitor_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_pool_quota_probe_worker(self.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_stats_hourly_aggregation_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_pending_cleanup_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_proxy_node_stale_cleanup_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_proxy_node_metrics_cleanup_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_proxy_upgrade_rollout_worker(self.clone()) {
tasks.push(handle);
}
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);
}
if let Some(handle) = spawn_gemini_file_mapping_cleanup_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_model_fetch_worker(self.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_video_task_poller(self.clone()) {
tasks.push(handle);
}
tasks
pub fn spawn_background_tasks(&self) -> crate::task_runtime::TaskSupervisor {
let mut supervisor = crate::task_runtime::TaskSupervisor::new();
let record_boot = |task_key: &'static str| {
if !self.has_background_task_data_writer() {
return;
}
let Some(definition) = crate::task_runtime::task_definition(task_key) else {
return;
};
std::mem::drop(crate::task_runtime::spawn_record_worker_boot(
self.clone(),
task_key,
crate::task_runtime::background_task_kind(definition.kind),
definition.trigger,
));
};
let mut supervise_worker =
|task_key: &'static str, handle: Option<tokio::task::JoinHandle<()>>| {
if let Some(handle) = handle {
supervisor.supervise_handle(task_key, handle);
record_boot(task_key);
}
};
supervise_worker(
crate::task_runtime::TASK_KEY_USAGE_QUEUE_WORKER,
self.usage_runtime.spawn_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_PROVIDER_QUOTA_RESET,
crate::wallet_runtime::spawn_provider_quota_reset_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_AUDIT_CLEANUP,
spawn_audit_cleanup_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_DB_MAINTENANCE,
spawn_db_maintenance_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_WALLET_DAILY_USAGE_AGG,
spawn_wallet_daily_usage_aggregation_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_STATS_DAILY_AGG,
spawn_stats_aggregation_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_USAGE_CLEANUP,
spawn_usage_cleanup_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_POOL_MONITOR,
spawn_pool_monitor_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_POOL_QUOTA_PROBE,
spawn_pool_quota_probe_worker(self.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_STATS_HOURLY_AGG,
spawn_stats_hourly_aggregation_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_PENDING_CLEANUP,
spawn_pending_cleanup_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_PROXY_NODE_STALE_CLEANUP,
spawn_proxy_node_stale_cleanup_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_PROXY_NODE_METRICS_CLEANUP,
spawn_proxy_node_metrics_cleanup_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_PROXY_UPGRADE_ROLLOUT,
spawn_proxy_upgrade_rollout_worker(self.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_PROVIDER_CHECKIN,
spawn_provider_checkin_worker(self.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_OAUTH_TOKEN_REFRESH,
spawn_oauth_token_refresh_worker(self.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_REQUEST_CANDIDATE_CLEANUP,
spawn_request_candidate_cleanup_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_GEMINI_FILES_CLEANUP,
spawn_gemini_file_mapping_cleanup_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_MODEL_FETCH_WORKER,
spawn_model_fetch_worker(self.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_VIDEO_TASK_POLLER,
spawn_video_task_poller(self.clone()),
);
supervisor
}
}

View File

@@ -0,0 +1,80 @@
use crate::{AppState, GatewayError};
use aether_data_contracts::repository::background_tasks::{
BackgroundTaskListQuery, BackgroundTaskSummary, StoredBackgroundTaskEvent,
StoredBackgroundTaskRun, StoredBackgroundTaskRunPage, UpsertBackgroundTaskEvent,
UpsertBackgroundTaskRun,
};
impl AppState {
pub(crate) async fn find_background_task_run(
&self,
run_id: &str,
) -> Result<Option<StoredBackgroundTaskRun>, GatewayError> {
self.data
.find_background_task_run(run_id)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn list_background_task_runs(
&self,
query: &BackgroundTaskListQuery,
) -> Result<StoredBackgroundTaskRunPage, GatewayError> {
self.data
.list_background_task_runs(query)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn list_background_task_events(
&self,
run_id: &str,
offset: usize,
limit: usize,
) -> Result<Vec<StoredBackgroundTaskEvent>, GatewayError> {
self.data
.list_background_task_events(run_id, offset, limit)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn summarize_background_task_runs(
&self,
) -> Result<BackgroundTaskSummary, GatewayError> {
self.data
.summarize_background_task_runs()
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn upsert_background_task_run(
&self,
run: UpsertBackgroundTaskRun,
) -> Result<Option<StoredBackgroundTaskRun>, GatewayError> {
self.data
.upsert_background_task_run(run)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn request_cancel_background_task_run(
&self,
run_id: &str,
updated_at_unix_secs: u64,
) -> Result<bool, GatewayError> {
self.data
.request_cancel_background_task_run(run_id, updated_at_unix_secs)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn upsert_background_task_event(
&self,
event: UpsertBackgroundTaskEvent,
) -> Result<Option<StoredBackgroundTaskEvent>, GatewayError> {
self.data
.upsert_background_task_event(event)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
}

View File

@@ -10,6 +10,7 @@ use super::{
mod announcements;
mod api_key_exports;
mod auth;
mod background_tasks;
mod billing;
mod candidate_queries;
mod gemini_files;
@@ -29,6 +30,14 @@ impl AppState {
self.data.has_announcement_writer()
}
pub fn has_background_task_data_reader(&self) -> bool {
self.data.has_background_task_reader()
}
pub fn has_background_task_data_writer(&self) -> bool {
self.data.has_background_task_writer()
}
pub fn has_video_task_data_reader(&self) -> bool {
self.data.has_video_task_reader()
}

View File

@@ -0,0 +1,647 @@
use std::future::Future;
use std::time::{SystemTime, UNIX_EPOCH};
use aether_data_contracts::repository::background_tasks::{
BackgroundTaskKind, BackgroundTaskStatus, StoredBackgroundTaskRun, UpsertBackgroundTaskEvent,
UpsertBackgroundTaskRun,
};
use aether_runtime::task::spawn_named;
pub(crate) use aether_task_runtime::TaskSupervisor;
use aether_task_runtime::{RetryPolicy, TaskDefinition, TaskKind};
use serde_json::Value;
use tokio::task::JoinHandle;
use tracing::warn;
use uuid::Uuid;
use crate::{AppState, GatewayError};
pub(crate) const TASK_KEY_PROVIDER_DELETE: &str = "admin.provider.delete";
pub(crate) const TASK_KEY_PROVIDER_OAUTH_BATCH_IMPORT: &str = "admin.provider.oauth.batch_import";
pub(crate) const TASK_KEY_USAGE_QUEUE_WORKER: &str = "usage.queue.worker";
pub(crate) const TASK_KEY_VIDEO_TASK_POLLER: &str = "video.task.poller";
pub(crate) const TASK_KEY_MODEL_FETCH_WORKER: &str = "model.fetch.worker";
pub(crate) const TASK_KEY_PROVIDER_QUOTA_RESET: &str = "provider.quota.reset.worker";
pub(crate) const TASK_KEY_POOL_QUOTA_PROBE: &str = "pool.quota.probe.worker";
pub(crate) const TASK_KEY_POOL_MONITOR: &str = "pool.monitor.worker";
pub(crate) const TASK_KEY_AUDIT_CLEANUP: &str = "maintenance.audit.cleanup";
pub(crate) const TASK_KEY_DB_MAINTENANCE: &str = "maintenance.database";
pub(crate) const TASK_KEY_PENDING_CLEANUP: &str = "maintenance.pending.cleanup";
pub(crate) const TASK_KEY_REQUEST_CANDIDATE_CLEANUP: &str = "maintenance.request.candidate.cleanup";
pub(crate) const TASK_KEY_GEMINI_FILES_CLEANUP: &str = "maintenance.gemini.files.cleanup";
pub(crate) const TASK_KEY_OAUTH_TOKEN_REFRESH: &str = "maintenance.oauth.token.refresh";
pub(crate) const TASK_KEY_PROXY_NODE_STALE_CLEANUP: &str = "maintenance.proxy.node.stale.cleanup";
pub(crate) const TASK_KEY_PROXY_NODE_METRICS_CLEANUP: &str =
"maintenance.proxy.node.metrics.cleanup";
pub(crate) const TASK_KEY_PROXY_UPGRADE_ROLLOUT: &str = "maintenance.proxy.upgrade.rollout";
pub(crate) const TASK_KEY_PROVIDER_CHECKIN: &str = "maintenance.provider.checkin";
pub(crate) const TASK_KEY_USAGE_CLEANUP: &str = "maintenance.usage.cleanup";
pub(crate) const TASK_KEY_WALLET_DAILY_USAGE_AGG: &str = "maintenance.wallet.daily.usage.agg";
pub(crate) const TASK_KEY_STATS_DAILY_AGG: &str = "maintenance.stats.daily.agg";
pub(crate) const TASK_KEY_STATS_HOURLY_AGG: &str = "maintenance.stats.hourly.agg";
pub(crate) const TASK_KEY_USAGE_SYNC_REPORT: &str = "usage.sync.report";
pub(crate) const TASK_KEY_PROVIDER_OAUTH_ACCOUNT_REFRESH: &str = "provider.oauth.account.refresh";
pub(crate) const TASK_KEY_PROVIDER_BALANCE_REFRESH: &str = "provider.ops.balance.refresh";
const RETRY_ONCE: RetryPolicy = RetryPolicy { max_attempts: 1 };
const TASK_DEFINITIONS: &[TaskDefinition] = &[
TaskDefinition::new(
TASK_KEY_PROVIDER_DELETE,
TaskKind::OnDemand,
"manual",
false,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_PROVIDER_OAUTH_BATCH_IMPORT,
TaskKind::OnDemand,
"manual",
false,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_USAGE_QUEUE_WORKER,
TaskKind::Daemon,
"daemon",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_VIDEO_TASK_POLLER,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_MODEL_FETCH_WORKER,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_PROVIDER_QUOTA_RESET,
TaskKind::Scheduled,
"daily",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_POOL_QUOTA_PROBE,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_POOL_MONITOR,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_AUDIT_CLEANUP,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_DB_MAINTENANCE,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_PENDING_CLEANUP,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_REQUEST_CANDIDATE_CLEANUP,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_GEMINI_FILES_CLEANUP,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_OAUTH_TOKEN_REFRESH,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_PROXY_NODE_STALE_CLEANUP,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_PROXY_NODE_METRICS_CLEANUP,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_PROXY_UPGRADE_ROLLOUT,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_PROVIDER_CHECKIN,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_USAGE_CLEANUP,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_WALLET_DAILY_USAGE_AGG,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_STATS_DAILY_AGG,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_STATS_HOURLY_AGG,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_USAGE_SYNC_REPORT,
TaskKind::FireAndForget,
"internal",
false,
false,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_PROVIDER_OAUTH_ACCOUNT_REFRESH,
TaskKind::FireAndForget,
"internal",
false,
false,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_PROVIDER_BALANCE_REFRESH,
TaskKind::FireAndForget,
"internal",
false,
false,
RETRY_ONCE,
),
];
pub(crate) fn task_definitions() -> &'static [TaskDefinition] {
TASK_DEFINITIONS
}
pub(crate) fn task_definition(task_key: &str) -> Option<TaskDefinition> {
task_definitions()
.iter()
.copied()
.find(|definition| definition.key == task_key)
}
pub(crate) const fn background_task_kind(kind: TaskKind) -> BackgroundTaskKind {
match kind {
TaskKind::Scheduled => BackgroundTaskKind::Scheduled,
TaskKind::Daemon => BackgroundTaskKind::Daemon,
TaskKind::OnDemand => BackgroundTaskKind::OnDemand,
TaskKind::FireAndForget => BackgroundTaskKind::FireAndForget,
}
}
pub(crate) fn task_cancel_kv_key(run_id: &str) -> String {
format!("task_runtime:run:{run_id}:cancel")
}
pub(crate) fn now_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs())
.unwrap_or_default()
}
pub(crate) fn build_task_run_id() -> String {
Uuid::new_v4().to_string()
}
pub(crate) fn spawn_fire_and_forget<F>(task_name: &'static str, future: F) -> JoinHandle<()>
where
F: Future<Output = ()> + Send + 'static,
{
spawn_named(task_name, future)
}
fn stored_run_to_upsert(run: StoredBackgroundTaskRun) -> UpsertBackgroundTaskRun {
UpsertBackgroundTaskRun {
id: run.id,
task_key: run.task_key,
kind: run.kind,
trigger: run.trigger,
status: run.status,
attempt: run.attempt,
max_attempts: run.max_attempts,
owner_instance: run.owner_instance,
progress_percent: run.progress_percent,
progress_message: run.progress_message,
payload_json: run.payload_json,
result_json: run.result_json,
error_message: run.error_message,
cancel_requested: run.cancel_requested,
created_by: run.created_by,
created_at_unix_secs: run.created_at_unix_secs,
started_at_unix_secs: run.started_at_unix_secs,
finished_at_unix_secs: run.finished_at_unix_secs,
updated_at_unix_secs: run.updated_at_unix_secs,
}
}
pub(crate) async fn upsert_run_with_logging(
app: &AppState,
run: UpsertBackgroundTaskRun,
) -> Option<StoredBackgroundTaskRun> {
match app.upsert_background_task_run(run).await {
Ok(result) => result,
Err(error) => {
warn!(error = ?error, "failed to upsert background task run");
None
}
}
}
pub(crate) async fn update_run_status(
app: &AppState,
run_id: &str,
status: BackgroundTaskStatus,
progress_percent: Option<u16>,
progress_message: Option<String>,
result_json: Option<Value>,
error_message: Option<String>,
started_at_unix_secs: Option<u64>,
finished_at_unix_secs: Option<u64>,
) -> Option<StoredBackgroundTaskRun> {
let Some(mut existing) = app.find_background_task_run(run_id).await.ok().flatten() else {
return None;
};
existing.status = status;
if let Some(progress_percent) = progress_percent {
existing.progress_percent = progress_percent.min(100);
}
if let Some(progress_message) = progress_message {
existing.progress_message = Some(progress_message);
}
if result_json.is_some() {
existing.result_json = result_json;
}
if error_message.is_some() {
existing.error_message = error_message;
}
if started_at_unix_secs.is_some() {
existing.started_at_unix_secs = started_at_unix_secs;
}
if finished_at_unix_secs.is_some() {
existing.finished_at_unix_secs = finished_at_unix_secs;
}
existing.updated_at_unix_secs = now_unix_secs();
upsert_run_with_logging(app, stored_run_to_upsert(existing)).await
}
pub(crate) async fn append_event_with_logging(
app: &AppState,
run_id: &str,
event_type: &str,
message: &str,
payload_json: Option<Value>,
) {
let event = UpsertBackgroundTaskEvent {
id: Uuid::new_v4().to_string(),
run_id: run_id.to_string(),
event_type: event_type.to_string(),
message: message.to_string(),
payload_json,
created_at_unix_secs: now_unix_secs(),
};
if let Err(error) = app.upsert_background_task_event(event).await {
warn!(error = ?error, run_id = %run_id, "failed to upsert background task event");
}
}
pub(crate) fn spawn_record_worker_boot(
app: AppState,
task_key: &'static str,
kind: BackgroundTaskKind,
trigger: &'static str,
) -> JoinHandle<()> {
spawn_named("task-runtime-record-worker-boot", async move {
let now = now_unix_secs();
let run_id = format!("boot:{}:{}", task_key, app.tunnel.local_instance_id());
let run = UpsertBackgroundTaskRun {
id: run_id.clone(),
task_key: task_key.to_string(),
kind,
trigger: trigger.to_string(),
status: BackgroundTaskStatus::Running,
attempt: 1,
max_attempts: 1,
owner_instance: Some(app.tunnel.local_instance_id().to_string()),
progress_percent: 0,
progress_message: Some("worker booted".to_string()),
payload_json: None,
result_json: None,
error_message: None,
cancel_requested: false,
created_by: Some("system".to_string()),
created_at_unix_secs: now,
started_at_unix_secs: Some(now),
finished_at_unix_secs: None,
updated_at_unix_secs: now,
};
let _ = upsert_run_with_logging(&app, run).await;
append_event_with_logging(
&app,
&run_id,
"worker_boot",
"background worker started",
None,
)
.await;
})
}
pub(crate) async fn set_cancel_signal(app: &AppState, run_id: &str) -> Result<(), GatewayError> {
app.runtime_kv_setex(&task_cancel_kv_key(run_id), "1", 60 * 60)
.await
}
pub(crate) async fn is_cancel_requested(app: &AppState, run_id: &str) -> bool {
if let Ok(Some(run)) = app.find_background_task_run(run_id).await {
if run.cancel_requested {
return true;
}
}
app.runtime_kv_exists(&task_cancel_kv_key(run_id))
.await
.unwrap_or(false)
}
pub(crate) async fn submit_provider_delete_task(
state: &crate::admin_api::AdminAppState<'_>,
provider_id: &str,
created_by: Option<&str>,
) -> Result<Option<String>, GatewayError> {
let Some(provider) = state
.read_provider_catalog_providers_by_ids(&[provider_id.to_string()])
.await?
.into_iter()
.next()
else {
return Ok(None);
};
let task_id = Uuid::new_v4().simple().to_string()[..16].to_string();
state.put_provider_delete_task(crate::LocalProviderDeleteTaskState {
task_id: task_id.clone(),
provider_id: provider.id.clone(),
status: "pending".to_string(),
stage: "queued".to_string(),
total_keys: 0,
deleted_keys: 0,
total_endpoints: 0,
deleted_endpoints: 0,
message: "delete task submitted".to_string(),
});
let app = state.cloned_app();
let provider_id = provider.id.clone();
let run_id = task_id.clone();
let created_at = now_unix_secs();
let max_attempts = task_definition(TASK_KEY_PROVIDER_DELETE)
.map(|item| item.retry_policy.max_attempts)
.unwrap_or(1);
if state.has_background_task_data_writer() {
let run = UpsertBackgroundTaskRun {
id: run_id.clone(),
task_key: TASK_KEY_PROVIDER_DELETE.to_string(),
kind: BackgroundTaskKind::OnDemand,
trigger: "manual".to_string(),
status: BackgroundTaskStatus::Queued,
attempt: 1,
max_attempts,
owner_instance: Some(app.tunnel.local_instance_id().to_string()),
progress_percent: 0,
progress_message: Some("delete task queued".to_string()),
payload_json: Some(serde_json::json!({ "provider_id": provider_id.clone() })),
result_json: None,
error_message: None,
cancel_requested: false,
created_by: Some(created_by.unwrap_or("admin").to_string()),
created_at_unix_secs: created_at,
started_at_unix_secs: None,
finished_at_unix_secs: None,
updated_at_unix_secs: created_at,
};
let _ = upsert_run_with_logging(&app, run).await;
append_event_with_logging(
&app,
&run_id,
"queued",
"provider delete task queued",
Some(serde_json::json!({ "provider_id": provider_id.clone() })),
)
.await;
}
spawn_named("task-runtime-provider-delete", async move {
let lock_key = format!("task_runtime:lock:{TASK_KEY_PROVIDER_DELETE}:{provider_id}");
let lock_ttl = std::time::Duration::from_secs(60 * 15);
let lock = app
.runtime_state
.lock_try_acquire(&lock_key, app.tunnel.local_instance_id(), lock_ttl)
.await
.ok()
.flatten();
if lock.is_none() {
let _ = update_run_status(
&app,
&run_id,
BackgroundTaskStatus::Skipped,
Some(0),
Some("provider delete skipped: another node is running this task".to_string()),
None,
None,
None,
Some(now_unix_secs()),
)
.await;
append_event_with_logging(
&app,
&run_id,
"skipped",
"provider delete skipped by singleton lock",
None,
)
.await;
return;
}
let started_at = now_unix_secs();
let _ = update_run_status(
&app,
&run_id,
BackgroundTaskStatus::Running,
Some(5),
Some("provider delete task started".to_string()),
None,
None,
Some(started_at),
None,
)
.await;
append_event_with_logging(
&app,
&run_id,
"running",
"provider delete task started",
None,
)
.await;
let admin_state = crate::admin_api::AdminAppState::new(&app);
let result = admin_state
.run_admin_provider_delete_task(&provider_id, &run_id)
.await;
match result {
Ok(task_state) => {
let _ = update_run_status(
&app,
&run_id,
BackgroundTaskStatus::Succeeded,
Some(100),
Some(task_state.message.clone()),
Some(serde_json::json!({
"provider_id": task_state.provider_id,
"status": task_state.status,
"stage": task_state.stage,
"deleted_keys": task_state.deleted_keys,
"total_keys": task_state.total_keys,
"deleted_endpoints": task_state.deleted_endpoints,
"total_endpoints": task_state.total_endpoints,
"message": task_state.message,
})),
None,
None,
Some(now_unix_secs()),
)
.await;
append_event_with_logging(
&app,
&run_id,
"succeeded",
"provider delete task completed",
None,
)
.await;
}
Err(err) => {
warn!(
"gateway admin provider delete task failed for provider {}: {:?}",
provider_id, err
);
app.put_provider_delete_task(crate::LocalProviderDeleteTaskState {
task_id: run_id.clone(),
provider_id: provider_id.clone(),
status: "failed".to_string(),
stage: "failed".to_string(),
total_keys: 0,
deleted_keys: 0,
total_endpoints: 0,
deleted_endpoints: 0,
message: format!("provider delete failed: {err:?}"),
});
let _ = update_run_status(
&app,
&run_id,
BackgroundTaskStatus::Failed,
Some(100),
Some("provider delete task failed".to_string()),
None,
Some(format!("{err:?}")),
None,
Some(now_unix_secs()),
)
.await;
append_event_with_logging(
&app,
&run_id,
"failed",
"provider delete task failed",
Some(serde_json::json!({ "error": format!("{err:?}") })),
)
.await;
}
}
if let Some(lock) = lock {
let _ = app.runtime_state.lock_release(&lock).await;
}
});
Ok(Some(task_id))
}

View File

@@ -72,7 +72,5 @@ async fn gateway_background_gemini_file_mapping_cleanup_deletes_expired_entries(
"active mapping should remain"
);
for handle in background_tasks {
handle.abort();
}
background_tasks.shutdown().await;
}

View File

@@ -231,9 +231,7 @@ async fn gateway_background_video_task_poller_refreshes_due_openai_task_from_rep
}]
);
for handle in background_tasks {
handle.abort();
}
background_tasks.shutdown().await;
execution_runtime_handle.abort();
}
@@ -321,8 +319,6 @@ async fn gateway_background_video_task_poller_refreshes_due_openai_task_from_rep
}]
);
for handle in background_tasks {
handle.abort();
}
background_tasks.shutdown().await;
upstream_handle.abort();
}

View File

@@ -9,6 +9,7 @@ use crate::clock::current_unix_ms;
use crate::log_ids::short_request_id;
use crate::orchestration::{apply_local_report_effect, LocalReportEffect};
use crate::request_candidate_runtime::record_report_request_candidate_status;
use crate::task_runtime::{spawn_fire_and_forget, TASK_KEY_USAGE_SYNC_REPORT};
use crate::{AppState, GatewayError};
mod context;
@@ -142,7 +143,7 @@ pub(crate) async fn submit_sync_report(
pub(crate) fn spawn_sync_report(state: AppState, payload: GatewaySyncReportRequest) {
let report_request_id_for_log =
short_request_id(report_request_id(payload.report_context.as_ref()));
tokio::spawn(async move {
spawn_fire_and_forget(TASK_KEY_USAGE_SYNC_REPORT, async move {
let trace_id = payload.trace_id.clone();
if let Err(err) = submit_sync_report(&state, payload).await {
warn!(