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

Fix admin pool sorting and OAuth refresh

Browse Source
This commit is contained in:
fawney19
2026-05-13 09:16:52 +08:00
parent 4387a9cdd5
commit 3c2497f019
15 changed files with 498 additions and 195 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
apps/aether-gateway/src/handlers/admin/mod.rs
View File

@@ -31,7 +31,8 @@ pub(crate) use self::provider::oauth::quota::codex::refresh_codex_provider_quota
pub(crate) use self::provider::oauth::quota::kiro::refresh_kiro_provider_quota_locally;
pub(crate) use self::provider::oauth::quota::shared::provider_type_supports_quota_refresh;
pub(crate) use self::provider::oauth::runtime::{
provider_oauth_runtime_endpoint_for_provider, refresh_provider_oauth_account_state_after_update,
provider_oauth_maintenance_endpoint_for_provider, provider_oauth_runtime_endpoint_for_provider,
refresh_provider_oauth_account_state_after_update,
};
pub(crate) use self::provider::ops::providers::actions::admin_provider_ops_local_action_response;
pub(crate) use self::provider::pool::config::admin_provider_pool_config;

6
apps/aether-gateway/src/handlers/admin/provider/endpoint_keys/quota.rs
View File

@@ -21,7 +21,7 @@ use super::super::oauth::quota::shared::normalize_string_id_list;
use super::super::oauth::quota::shared::{
provider_type_supports_quota_refresh, unsupported_provider_quota_refresh_message,
};
use super::super::oauth::runtime::provider_oauth_runtime_endpoint_for_provider;
use super::super::oauth::runtime::provider_oauth_maintenance_endpoint_for_provider;
use super::super::write::provider::reconcile_admin_fixed_provider_template_endpoints;
fn unsupported_provider_quota_refresh_response(provider_type: &str) -> Response<Body> {
@@ -115,7 +115,7 @@ pub(super) async fn maybe_handle(
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await?;
let mut endpoint =
provider_oauth_runtime_endpoint_for_provider(&normalized_provider_type, &endpoints);
provider_oauth_maintenance_endpoint_for_provider(&normalized_provider_type, &endpoints);
if endpoint.is_none() && is_fixed_provider {
if !state.has_provider_catalog_data_writer() {
@@ -137,7 +137,7 @@ pub(super) async fn maybe_handle(
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await?;
endpoint =
provider_oauth_runtime_endpoint_for_provider(&normalized_provider_type, &endpoints);
provider_oauth_maintenance_endpoint_for_provider(&normalized_provider_type, &endpoints);
}
if !provider_type_supports_quota_refresh(&normalized_provider_type) {

6
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/complete/key.rs
View File

@@ -128,7 +128,6 @@ pub(super) async fn handle_admin_provider_oauth_complete_key(
};
let endpoint_resolution =
resolve_provider_oauth_runtime_endpoints(state, &provider, &provider_type).await?;
let endpoints = endpoint_resolution.endpoints;
let runtime_endpoint = endpoint_resolution.runtime_endpoint;
let request_proxy = state
.resolve_admin_provider_oauth_operation_proxy_snapshot(
@@ -223,10 +222,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_key(
let mut account_state_recheck_attempted = false;
let mut account_state_recheck_error = None::<String>;
if provider_type == "codex" {
if let Some(endpoint) = endpoints.into_iter().find(|endpoint| {
endpoint.is_active
&& crate::ai_serving::is_openai_responses_format(&endpoint.api_format)
}) {
if let Some(endpoint) = runtime_endpoint {
let refreshed_key = state
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&key_id))
.await?

178
apps/aether-gateway/src/handlers/admin/provider/oauth/runtime.rs
View File

@@ -15,98 +15,95 @@ use aether_data_contracts::repository::provider_catalog::{
pub(crate) fn provider_oauth_runtime_endpoint_for_provider(
provider_type: &str,
endpoints: &[StoredProviderCatalogEndpoint],
) -> Option<StoredProviderCatalogEndpoint> {
select_provider_oauth_runtime_endpoint(provider_type, endpoints, false)
}
pub(crate) fn provider_oauth_maintenance_endpoint_for_provider(
provider_type: &str,
endpoints: &[StoredProviderCatalogEndpoint],
) -> Option<StoredProviderCatalogEndpoint> {
select_provider_oauth_runtime_endpoint(provider_type, endpoints, true)
}
fn matching_endpoint<F>(
endpoints: &[StoredProviderCatalogEndpoint],
include_inactive: bool,
predicate: F,
) -> Option<StoredProviderCatalogEndpoint>
where
F: Fn(&StoredProviderCatalogEndpoint) -> bool,
{
endpoints
.iter()
.find(|endpoint| endpoint.is_active && predicate(endpoint))
.cloned()
.or_else(|| {
include_inactive.then(|| {
endpoints
.iter()
.find(|endpoint| !endpoint.is_active && predicate(endpoint))
.cloned()
})?
})
}
fn select_provider_oauth_runtime_endpoint(
provider_type: &str,
endpoints: &[StoredProviderCatalogEndpoint],
include_inactive: bool,
) -> Option<StoredProviderCatalogEndpoint> {
let provider_type = provider_type.trim().to_ascii_lowercase();
match provider_type.as_str() {
"codex" => endpoints
.iter()
.find(|endpoint| {
endpoint.is_active
&& crate::ai_serving::is_openai_responses_format(&endpoint.api_format)
"codex" => matching_endpoint(endpoints, include_inactive, |endpoint| {
crate::ai_serving::is_openai_responses_format(&endpoint.api_format)
}),
"chatgpt_web" => matching_endpoint(endpoints, include_inactive, |endpoint| {
endpoint
.api_format
.trim()
.eq_ignore_ascii_case("openai:image")
}),
"antigravity" => matching_endpoint(endpoints, include_inactive, |endpoint| {
endpoint
.api_format
.trim()
.eq_ignore_ascii_case("gemini:generate_content")
}),
"kiro" => matching_endpoint(endpoints, include_inactive, |endpoint| {
endpoint
.api_format
.trim()
.eq_ignore_ascii_case("claude:messages")
})
.or_else(|| matching_endpoint(endpoints, include_inactive, |_| true)),
"claude_code" => matching_endpoint(endpoints, include_inactive, |endpoint| {
endpoint
.api_format
.trim()
.eq_ignore_ascii_case("claude:messages")
}),
"gemini_cli" => matching_endpoint(endpoints, include_inactive, |endpoint| {
endpoint
.api_format
.trim()
.eq_ignore_ascii_case("gemini:generate_content")
}),
"vertex_ai" => matching_endpoint(endpoints, include_inactive, |endpoint| {
endpoint
.api_format
.trim()
.eq_ignore_ascii_case("gemini:generate_content")
})
.or_else(|| {
matching_endpoint(endpoints, include_inactive, |endpoint| {
endpoint
.api_format
.trim()
.eq_ignore_ascii_case("claude:messages")
})
.cloned(),
"chatgpt_web" => endpoints
.iter()
.find(|endpoint| {
endpoint.is_active
&& endpoint
.api_format
.trim()
.eq_ignore_ascii_case("openai:image")
})
.cloned(),
"antigravity" => endpoints
.iter()
.find(|endpoint| {
endpoint.is_active
&& endpoint
.api_format
.trim()
.eq_ignore_ascii_case("gemini:generate_content")
})
.cloned(),
"kiro" => endpoints
.iter()
.find(|endpoint| {
endpoint.is_active
&& endpoint
.api_format
.trim()
.eq_ignore_ascii_case("claude:messages")
})
.cloned()
.or_else(|| {
endpoints
.iter()
.find(|endpoint| endpoint.is_active)
.cloned()
}),
"claude_code" => endpoints
.iter()
.find(|endpoint| {
endpoint.is_active
&& endpoint
.api_format
.trim()
.eq_ignore_ascii_case("claude:messages")
})
.cloned(),
"gemini_cli" => endpoints
.iter()
.find(|endpoint| {
endpoint.is_active
&& endpoint
.api_format
.trim()
.eq_ignore_ascii_case("gemini:generate_content")
})
.cloned(),
"vertex_ai" => endpoints
.iter()
.find(|endpoint| {
endpoint.is_active
&& endpoint
.api_format
.trim()
.eq_ignore_ascii_case("gemini:generate_content")
})
.cloned()
.or_else(|| {
endpoints
.iter()
.find(|endpoint| {
endpoint.is_active
&& endpoint
.api_format
.trim()
.eq_ignore_ascii_case("claude:messages")
})
.cloned()
}),
_ => endpoints
.iter()
.find(|endpoint| endpoint.is_active)
.cloned(),
}),
_ => matching_endpoint(endpoints, include_inactive, |_| true),
}
}
@@ -125,7 +122,7 @@ pub(crate) async fn resolve_provider_oauth_runtime_endpoints(
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id))
.await?;
let mut runtime_endpoint =
provider_oauth_runtime_endpoint_for_provider(provider_type, &endpoints);
provider_oauth_maintenance_endpoint_for_provider(provider_type, &endpoints);
if runtime_endpoint.is_none()
&& state
.fixed_provider_template(&provider.provider_type)
@@ -136,7 +133,8 @@ pub(crate) async fn resolve_provider_oauth_runtime_endpoints(
endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id))
.await?;
runtime_endpoint = provider_oauth_runtime_endpoint_for_provider(provider_type, &endpoints);
runtime_endpoint =
provider_oauth_maintenance_endpoint_for_provider(provider_type, &endpoints);
}
Ok(ProviderOAuthRuntimeEndpoints {

91
apps/aether-gateway/src/handlers/admin/provider/pool_admin/read_routes/keys.rs
View File

@@ -197,6 +197,35 @@ fn admin_pool_compare_optional_unix_secs(
}
}
fn admin_pool_compare_optional_score(
left: Option<f64>,
right: Option<f64>,
direction: AdminPoolKeySortDirection,
) -> Ordering {
match (left, right) {
(Some(left), Some(right)) => {
let ordering = left.partial_cmp(&right).unwrap_or(Ordering::Equal);
match direction {
AdminPoolKeySortDirection::Asc => ordering,
AdminPoolKeySortDirection::Desc => ordering.reverse(),
}
}
(Some(_), None) => Ordering::Less,
(None, Some(_)) => Ordering::Greater,
(None, None) => Ordering::Equal,
}
}
fn admin_pool_score_for_key(
scores_by_key_id: &BTreeMap<String, StoredPoolMemberScore>,
key: &StoredProviderCatalogKey,
) -> Option<f64> {
scores_by_key_id
.get(&key.id)
.map(|score| score.score)
.filter(|score| score.is_finite())
}
fn admin_pool_sort_keys_for_request(keys: &mut [StoredProviderCatalogKey], sort: AdminPoolKeySort) {
match sort.field {
AdminPoolKeySortField::Default => pool_selection::admin_pool_sort_keys(keys),
@@ -222,9 +251,26 @@ fn admin_pool_sort_keys_for_request(keys: &mut [StoredProviderCatalogKey], sort:
.then(left.id.cmp(&right.id))
});
}
AdminPoolKeySortField::Score => {}
}
}
fn admin_pool_sort_keys_by_score(
keys: &mut [StoredProviderCatalogKey],
scores_by_key_id: &BTreeMap<String, StoredPoolMemberScore>,
direction: AdminPoolKeySortDirection,
) {
keys.sort_by(|left, right| {
admin_pool_compare_optional_score(
admin_pool_score_for_key(scores_by_key_id, left),
admin_pool_score_for_key(scores_by_key_id, right),
direction,
)
.then(left.name.cmp(&right.name))
.then(left.id.cmp(&right.id))
});
}
fn admin_pool_repository_key_order(sort: AdminPoolKeySort) -> ProviderCatalogKeyListOrder {
match (sort.field, sort.direction) {
(AdminPoolKeySortField::Default, _) => ProviderCatalogKeyListOrder::Name,
@@ -240,6 +286,7 @@ fn admin_pool_repository_key_order(sort: AdminPoolKeySort) -> ProviderCatalogKey
(AdminPoolKeySortField::LastUsedAt, AdminPoolKeySortDirection::Desc) => {
ProviderCatalogKeyListOrder::LastUsedAtDesc
}
(AdminPoolKeySortField::Score, _) => ProviderCatalogKeyListOrder::Name,
}
}
@@ -316,8 +363,9 @@ pub(super) async fn build_admin_pool_list_keys_response(
let pool_config = admin_provider_pool_config(&provider);
let page_offset = page.saturating_sub(1).saturating_mul(page_size);
let sort_by_score = matches!(sort.field, AdminPoolKeySortField::Score);
let (keys, total) = if status == "cooldown" {
let (keys, total, preloaded_pool_scores_by_key_id) = if status == "cooldown" {
let cooldown_key_ids =
read_admin_provider_pool_cooldown_key_ids(state.runtime_state(), &provider.id).await;
let mut keys = if cooldown_key_ids.is_empty() {
@@ -349,15 +397,25 @@ pub(super) async fn build_admin_pool_list_keys_response(
})
});
}
admin_pool_sort_keys_for_request(&mut keys, sort);
let preloaded_pool_scores_by_key_id = if sort_by_score {
let key_ids = keys.iter().map(|key| key.id.clone()).collect::<Vec<_>>();
let scores = read_admin_pool_scores_by_key_id(state, &provider.id, &key_ids)
.await
.unwrap_or_default();
admin_pool_sort_keys_by_score(&mut keys, &scores, sort.direction);
Some(scores)
} else {
admin_pool_sort_keys_for_request(&mut keys, sort);
None
};
let total = keys.len();
let keys = keys
.into_iter()
.skip(page_offset)
.take(page_size)
.collect::<Vec<_>>();
(keys, total)
} else if !quick_selectors.is_empty() {
(keys, total, preloaded_pool_scores_by_key_id)
} else if !quick_selectors.is_empty() || sort_by_score {
let mut keys = state
.list_provider_catalog_keys_by_provider_ids(std::slice::from_ref(&provider.id))
.await?
@@ -386,14 +444,24 @@ pub(super) async fn build_admin_pool_list_keys_response(
})
})
.collect::<Vec<_>>();
admin_pool_sort_keys_for_request(&mut keys, sort);
let preloaded_pool_scores_by_key_id = if sort_by_score {
let key_ids = keys.iter().map(|key| key.id.clone()).collect::<Vec<_>>();
let scores = read_admin_pool_scores_by_key_id(state, &provider.id, &key_ids)
.await
.unwrap_or_default();
admin_pool_sort_keys_by_score(&mut keys, &scores, sort.direction);
Some(scores)
} else {
admin_pool_sort_keys_for_request(&mut keys, sort);
None
};
let total = keys.len();
let keys = keys
.into_iter()
.skip(page_offset)
.take(page_size)
.collect::<Vec<_>>();
(keys, total)
(keys, total, preloaded_pool_scores_by_key_id)
} else {
let key_page = state
.list_provider_catalog_key_page(&ProviderCatalogKeyListQuery {
@@ -409,13 +477,16 @@ pub(super) async fn build_admin_pool_list_keys_response(
order: admin_pool_repository_key_order(sort),
})
.await?;
(key_page.items, key_page.total)
(key_page.items, key_page.total, None)
};
let key_ids = keys.iter().map(|key| key.id.clone()).collect::<Vec<_>>();
let pool_scores_by_key_id = read_admin_pool_scores_by_key_id(state, &provider.id, &key_ids)
.await
.unwrap_or_default();
let pool_scores_by_key_id = match preloaded_pool_scores_by_key_id {
Some(scores) => scores,
None => read_admin_pool_scores_by_key_id(state, &provider.id, &key_ids)
.await
.unwrap_or_default(),
};
let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id))
.await?;

6
apps/aether-gateway/src/handlers/admin/provider/pool_admin/support.rs
View File

@@ -20,6 +20,7 @@ pub(crate) enum AdminPoolKeySortField {
Default,
ImportedAt,
LastUsedAt,
Score,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
@@ -121,8 +122,11 @@ pub(crate) fn parse_admin_pool_key_sort(query: Option<&str>) -> Result<AdminPool
Some("name") => AdminPoolKeySortField::Default,
Some("imported_at") | Some("created_at") => AdminPoolKeySortField::ImportedAt,
Some("last_used_at") | Some("last_used") => AdminPoolKeySortField::LastUsedAt,
Some("score") | Some("pool_score") => AdminPoolKeySortField::Score,
Some(_) => {
return Err("sort_by must be one of: name, imported_at, last_used_at".to_string());
return Err(
"sort_by must be one of: name, imported_at, last_used_at, score".to_string(),
);
}
};
let direction = match query_param_value(query, "sort_order")