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

Add usage queue worker autoscaling

Browse Source
This commit is contained in:
elky
2026-06-26 14:02:57 +08:00
parent 6c5e70ccb1
commit 7e9424008f
21 changed files with 1801 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/mod.rs
+29 -3
View File
@@ -15,7 +15,8 @@ use self::plan::{build_codex_quota_request_spec, execute_codex_quota_plan};
use super::shared::{
build_quota_snapshot_payload, extract_execution_error_message,
oauth_refresh_auto_removed_result, persist_provider_quota_refresh_state,
provider_auto_remove_banned_keys, quota_key_auto_removed, quota_refresh_success_invalid_state,
provider_auto_remove_banned_keys, provider_auto_remove_quota_exhausted_keys,
quota_key_auto_removed, quota_refresh_success_invalid_state,
should_auto_remove_structured_reason, ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::request::AdminAppState;
@@ -61,6 +62,8 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
proxy_override: Option<ProxySnapshot>,
) -> Result<Option<serde_json::Value>, GatewayError> {
let auto_remove_abnormal_keys = provider_auto_remove_banned_keys(provider.config.as_ref());
let auto_remove_quota_exhausted_keys =
provider_auto_remove_quota_exhausted_keys(provider.config.as_ref());
let mut results = Vec::new();
let mut success_count = 0usize;
let mut failed_count = 0usize;
@@ -312,7 +315,7 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}));
continue;
}
let auto_removed = if auto_remove_candidate {
let auto_removed_hard_banned = if auto_remove_candidate {
state
.cleanup_provider_catalog_key_if_current(provider, &key.id, |latest_key| {
should_auto_remove_structured_reason(latest_key.oauth_invalid_reason.as_deref())
@@ -321,10 +324,28 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
} else {
false
};
if auto_removed {
if auto_removed_hard_banned {
auto_removed_count += 1;
auto_removed_hard_banned_count += 1;
}
let auto_removed_quota_exhausted =
if !auto_removed_hard_banned && auto_remove_quota_exhausted_keys {
state
.cleanup_provider_catalog_key_if_current(provider, &key.id, |latest_key| {
aether_admin::provider::pool::admin_pool_key_account_quota_exhausted(
latest_key,
provider.provider_type.as_str(),
)
})
.await?
} else {
false
};
if auto_removed_quota_exhausted {
auto_removed_count += 1;
status = "quota_exhausted".to_string();
}
let auto_removed = auto_removed_hard_banned || auto_removed_quota_exhausted;
let refresh_fixed =
status == "success" && had_oauth_refresh_issue && oauth_invalid_reason.is_none();
if refresh_fixed {
@@ -370,8 +391,13 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}
if auto_removed {
payload.insert("auto_removed".to_string(), json!(true));
}
if auto_removed_hard_banned {
payload.insert("auto_removed_hard_banned".to_string(), json!(true));
}
if auto_removed_quota_exhausted {
payload.insert("auto_removed_quota_exhausted".to_string(), json!(true));
}
if refresh_fixed {
payload.insert("refresh_fixed".to_string(), json!(true));
}
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/kiro/mod.rs
+25 -2
View File
@@ -6,8 +6,8 @@ use self::plan::execute_kiro_quota_plan;
use super::shared::{
build_quota_snapshot_payload, extract_execution_error_message,
oauth_refresh_auto_removed_result, persist_provider_quota_refresh_state,
persist_quota_oauth_refresh_failure_state, quota_refresh_success_invalid_state,
ProviderQuotaExecutionOutcome,
persist_quota_oauth_refresh_failure_state, provider_auto_remove_quota_exhausted_keys,
quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::request::{AdminAppState, AdminLocalOAuthRefreshError};
use crate::GatewayError;
@@ -97,6 +97,8 @@ pub(crate) async fn refresh_kiro_provider_quota_locally(
let mut success_count = 0usize;
let mut failed_count = 0usize;
let mut auto_removed_count = 0usize;
let auto_remove_quota_exhausted_keys =
provider_auto_remove_quota_exhausted_keys(provider.config.as_ref());
for key in keys {
let transport = match state
@@ -304,6 +306,23 @@ pub(crate) async fn refresh_kiro_provider_quota_locally(
continue;
}
let auto_removed_quota_exhausted = if auto_remove_quota_exhausted_keys {
state
.cleanup_provider_catalog_key_if_current(provider, &key.id, |latest_key| {
aether_admin::provider::pool::admin_pool_key_account_quota_exhausted(
latest_key,
provider.provider_type.as_str(),
)
})
.await?
} else {
false
};
if auto_removed_quota_exhausted {
auto_removed_count += 1;
status = "quota_exhausted".to_string();
}
if status == "success" {
success_count += 1;
} else {
@@ -331,6 +350,10 @@ pub(crate) async fn refresh_kiro_provider_quota_locally(
) {
payload.insert("quota_snapshot".to_string(), quota_snapshot);
}
if auto_removed_quota_exhausted {
payload.insert("auto_removed".to_string(), json!(true));
payload.insert("auto_removed_quota_exhausted".to_string(), json!(true));
}
results.push(serde_json::Value::Object(payload));
}
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs
+6
View File
@@ -67,6 +67,12 @@ pub(crate) fn provider_auto_remove_banned_keys(config: Option<&serde_json::Value
admin_provider_quota_pure::provider_auto_remove_banned_keys(config)
}
pub(crate) fn provider_auto_remove_quota_exhausted_keys(
config: Option<&serde_json::Value>,
) -> bool {
admin_provider_quota_pure::provider_auto_remove_quota_exhausted_keys(config)
}
pub(super) fn should_auto_remove_structured_reason(reason: Option<&str>) -> bool {
admin_provider_quota_pure::should_auto_remove_structured_reason(reason)
}
apps/aether-gateway/src/handlers/admin/provider/write/provider/template.rs
+93 -2
View File
@@ -16,6 +16,8 @@ use std::collections::{BTreeMap, BTreeSet};
const FIXED_PROVIDER_TEMPLATE_METADATA_KEY: &str = "_aether_fixed_provider_template";
const OVERRIDE_BODY_RULES: &str = "body_rules";
const OVERRIDE_FORMAT_ACCEPTANCE_CONFIG: &str = "format_acceptance_config";
const OVERRIDE_BASE_URL: &str = "base_url";
const OVERRIDE_CUSTOM_PATH: &str = "custom_path";
const OVERRIDE_HEADER_RULES: &str = "header_rules";
const OVERRIDE_IS_ACTIVE: &str = "is_active";
const OVERRIDE_MAX_RETRIES: &str = "max_retries";
@@ -158,6 +160,20 @@ pub(crate) fn apply_admin_fixed_provider_endpoint_template_overrides(
.unwrap_or_else(|| managed_fixed_provider_endpoint_metadata(template, endpoint_template));
let mut overrides = metadata.overrides.clone();
sync_override_if_changed(
&mut overrides,
OVERRIDE_BASE_URL,
&existing_endpoint.base_url,
&updated_endpoint.base_url,
&defaults.base_url,
);
sync_override_if_changed(
&mut overrides,
OVERRIDE_CUSTOM_PATH,
&existing_endpoint.custom_path,
&updated_endpoint.custom_path,
&defaults.custom_path,
);
sync_override_if_changed(
&mut overrides,
OVERRIDE_HEADER_RULES,
@@ -249,9 +265,13 @@ fn reconcile_fixed_provider_endpoint(
updated.api_format = defaults.api_format.clone();
updated.api_family = Some(defaults.api_family.clone());
updated.endpoint_kind = Some(defaults.endpoint_kind.clone());
updated.base_url = defaults.base_url;
updated.custom_path = defaults.custom_path;
if !metadata.overrides.contains(OVERRIDE_BASE_URL) {
updated.base_url = defaults.base_url;
}
if !metadata.overrides.contains(OVERRIDE_CUSTOM_PATH) {
updated.custom_path = defaults.custom_path;
}
if !metadata.overrides.contains(OVERRIDE_HEADER_RULES) {
updated.header_rules = defaults.header_rules;
}
@@ -534,3 +554,74 @@ fn sync_override_if_changed<T>(
}
sync_override(overrides, key, actual, desired);
}
#[cfg(test)]
mod tests {
use super::{
apply_admin_fixed_provider_endpoint_template_overrides, fixed_provider_endpoint_metadata,
reconcile_fixed_provider_endpoint,
};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogProvider,
};
use aether_provider_transport::provider_types::fixed_provider_template;
fn sample_codex_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-codex".to_string(),
"Codex".to_string(),
None,
"codex".to_string(),
)
.expect("provider should build")
}
fn sample_codex_endpoint(base_url: &str) -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-codex-responses".to_string(),
"provider-codex".to_string(),
"openai:responses".to_string(),
Some("openai".to_string()),
Some("responses".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
base_url.to_string(),
None,
None,
Some(2),
None,
None,
None,
None,
)
.expect("endpoint transport should build")
}
#[test]
fn fixed_provider_endpoint_reconcile_preserves_base_url_override() {
let provider = sample_codex_provider();
let template = fixed_provider_template("codex").expect("codex template should exist");
let endpoint_template = template
.endpoints
.iter()
.find(|endpoint| endpoint.api_format == "openai:responses")
.expect("responses endpoint template should exist");
let existing = sample_codex_endpoint("https://chatgpt.com/backend-api/codex");
let mut updated = existing.clone();
updated.base_url = "http://127.0.0.1:18181/v1".to_string();
apply_admin_fixed_provider_endpoint_template_overrides(&provider, &existing, &mut updated)
.expect("override metadata should apply");
let metadata = fixed_provider_endpoint_metadata(&updated)
.expect("fixed provider metadata should exist");
assert!(metadata.overrides.contains("base_url"));
let reconciled =
reconcile_fixed_provider_endpoint(&provider, &updated, template, endpoint_template)
.expect("endpoint should reconcile");
assert_eq!(reconciled.base_url, "http://127.0.0.1:18181/v1");
}
}