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

17
apps/aether-gateway/src/admin_api.rs
View File

@@ -2,14 +2,15 @@ pub(crate) use crate::handlers::admin::{
admin_provider_ops_local_action_response, admin_provider_pool_config,
build_internal_control_error_response, create_provider_oauth_catalog_key,
find_duplicate_provider_oauth_key, maybe_build_local_admin_pool_response,
maybe_build_local_admin_response, provider_oauth_runtime_endpoint_for_provider,
provider_type_supports_quota_refresh, reconcile_admin_fixed_provider_template_endpoints,
refresh_antigravity_provider_quota_locally, refresh_chatgpt_web_provider_quota_locally,
refresh_codex_provider_quota_locally, refresh_kiro_provider_quota_locally,
refresh_provider_oauth_account_state_after_update, update_existing_provider_oauth_catalog_key,
AdminAppState, AdminGatewayProviderTransportSnapshot, AdminLocalOAuthRefreshError,
AdminRequestContext, AdminRouteRequest, AdminRouteResponse, AdminRouteResult,
AdminStatsTimeRange, AdminStatsUsageFilter,
maybe_build_local_admin_response, provider_oauth_maintenance_endpoint_for_provider,
provider_oauth_runtime_endpoint_for_provider, provider_type_supports_quota_refresh,
reconcile_admin_fixed_provider_template_endpoints, refresh_antigravity_provider_quota_locally,
refresh_chatgpt_web_provider_quota_locally, refresh_codex_provider_quota_locally,
refresh_kiro_provider_quota_locally, refresh_provider_oauth_account_state_after_update,
update_existing_provider_oauth_catalog_key, AdminAppState,
AdminGatewayProviderTransportSnapshot, AdminLocalOAuthRefreshError, AdminRequestContext,
AdminRouteRequest, AdminRouteResponse, AdminRouteResult, AdminStatsTimeRange,
AdminStatsUsageFilter,
};
use crate::handlers::admin::{

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")

64
apps/aether-gateway/src/maintenance/runtime/oauth_token_refresh.rs
View File

@@ -7,6 +7,7 @@ use aether_data_contracts::repository::provider_catalog::{
use serde_json::Value;
use tracing::{info, warn};
use crate::admin_api::provider_oauth_maintenance_endpoint_for_provider;
use crate::provider_key_auth::provider_key_is_oauth_managed;
use crate::{AppState, GatewayError};
@@ -75,9 +76,10 @@ pub(crate) async fn perform_oauth_token_refresh_once(
}
summary.eligible = summary.eligible.saturating_add(1);
let Some(endpoint) =
oauth_runtime_endpoint_for_provider(&provider.provider_type, provider_endpoints)
else {
let Some(endpoint) = provider_oauth_maintenance_endpoint_for_provider(
&provider.provider_type,
provider_endpoints,
) else {
summary.skipped = summary.skipped.saturating_add(1);
continue;
};
@@ -182,62 +184,6 @@ fn oauth_refresh_candidate(
&& provider_key_is_oauth_managed(key, provider.provider_type.as_str())
}
fn oauth_runtime_endpoint_for_provider(
provider_type: &str,
endpoints: &[StoredProviderCatalogEndpoint],
) -> Option<StoredProviderCatalogEndpoint> {
let provider_type = provider_type.trim().to_ascii_lowercase();
match provider_type.as_str() {
"codex" => endpoints
.iter()
.find(|endpoint| {
endpoint.is_active
&& crate::ai_serving::is_openai_responses_format(&endpoint.api_format)
})
.cloned(),
"chatgpt_web" => endpoints
.iter()
.find(|endpoint| {
endpoint.is_active
&& endpoint
.api_format
.trim()
.eq_ignore_ascii_case("openai:image")
})
.cloned(),
"antigravity" => endpoints
.iter()
.find(|endpoint| {
endpoint.is_active
&& endpoint
.api_format
.trim()
.eq_ignore_ascii_case("gemini:generate_content")
})
.cloned(),
"kiro" => endpoints
.iter()
.find(|endpoint| {
endpoint.is_active
&& endpoint
.api_format
.trim()
.eq_ignore_ascii_case("claude:messages")
})
.cloned()
.or_else(|| {
endpoints
.iter()
.find(|endpoint| endpoint.is_active)
.cloned()
}),
_ => endpoints
.iter()
.find(|endpoint| endpoint.is_active)
.cloned(),
}
}
async fn provider_key_credentials_changed(
state: &AppState,
before: &StoredProviderCatalogKey,

4
apps/aether-gateway/src/maintenance/runtime/pool_quota_probe.rs
View File

@@ -15,7 +15,7 @@ use serde_json::Value;
use tracing::{debug, info, warn};
use crate::admin_api::{
admin_provider_pool_config, provider_oauth_runtime_endpoint_for_provider,
admin_provider_pool_config, provider_oauth_maintenance_endpoint_for_provider,
provider_type_supports_quota_refresh, reconcile_admin_fixed_provider_template_endpoints,
refresh_antigravity_provider_quota_locally, refresh_chatgpt_web_provider_quota_locally,
refresh_codex_provider_quota_locally, refresh_kiro_provider_quota_locally, AdminAppState,
@@ -421,7 +421,7 @@ fn endpoint_for_probe(
provider_type: &str,
endpoints: &[StoredProviderCatalogEndpoint],
) -> Option<StoredProviderCatalogEndpoint> {
provider_oauth_runtime_endpoint_for_provider(provider_type, endpoints)
provider_oauth_maintenance_endpoint_for_provider(provider_type, endpoints)
}
async fn endpoint_for_probe_with_reconcile(

41
apps/aether-gateway/src/tests/control/admin/oauth.rs
View File

@@ -15,7 +15,7 @@ use aether_data::repository::oauth_providers::{
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::proxy_nodes::InMemoryProxyNodeRepository;
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
ProviderCatalogReadRepository, ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint,
};
use axum::body::{to_bytes, Body, Bytes};
use axum::response::{IntoResponse, Response};
@@ -5267,6 +5267,39 @@ fn gateway_manual_kiro_oauth_refresh_reconciles_missing_fixed_endpoint() {
}
async fn gateway_manual_kiro_oauth_refresh_reconciles_missing_fixed_endpoint_impl() {
run_gateway_manual_kiro_oauth_refresh_maintenance_endpoint_test(None, None, true).await;
}
#[test]
fn gateway_manual_kiro_oauth_refresh_uses_disabled_fixed_endpoint_for_maintenance() {
run_manual_kiro_oauth_refresh_test(
"gateway_manual_kiro_oauth_refresh_uses_disabled_fixed_endpoint_for_maintenance",
gateway_manual_kiro_oauth_refresh_uses_disabled_fixed_endpoint_for_maintenance_impl,
);
}
async fn gateway_manual_kiro_oauth_refresh_uses_disabled_fixed_endpoint_for_maintenance_impl() {
let mut endpoint = sample_endpoint(
"endpoint-kiro-disabled-maintenance",
"provider-kiro-oauth-refresh",
"claude:messages",
"https://q.{region}.amazonaws.com",
);
endpoint.is_active = false;
run_gateway_manual_kiro_oauth_refresh_maintenance_endpoint_test(
Some(endpoint),
Some("endpoint-kiro-disabled-maintenance"),
false,
)
.await;
}
async fn run_gateway_manual_kiro_oauth_refresh_maintenance_endpoint_test(
initial_endpoint: Option<StoredProviderCatalogEndpoint>,
expected_endpoint_id: Option<&str>,
expected_endpoint_active: bool,
) {
let refreshed_access_token = sample_kiro_device_access_token("kiro-refresh@example.com");
let expected_access_token = refreshed_access_token.clone();
let refreshed_refresh_token = "s".repeat(120);
@@ -5362,7 +5395,7 @@ async fn gateway_manual_kiro_oauth_refresh_reconciles_missing_fixed_endpoint_imp
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
initial_endpoint.into_iter().collect(),
vec![key],
));
@@ -5421,6 +5454,10 @@ async fn gateway_manual_kiro_oauth_refresh_reconciles_missing_fixed_endpoint_imp
assert_eq!(endpoints.len(), 1);
assert_eq!(endpoints[0].api_format, "claude:messages");
assert_eq!(endpoints[0].base_url, "https://q.{region}.amazonaws.com");
assert_eq!(endpoints[0].is_active, expected_endpoint_active);
if let Some(expected_endpoint_id) = expected_endpoint_id {
assert_eq!(endpoints[0].id, expected_endpoint_id);
}
assert_eq!(
*seen_endpoint_id.lock().expect("mutex should lock"),
Some(endpoints[0].id.clone())

65
apps/aether-gateway/src/tests/control/admin/pool.rs
View File

@@ -77,6 +77,36 @@ async fn local_admin_pool_response(
.expect("pool route should resolve locally")
}
fn sample_pool_member_score(provider_id: &str, key_id: &str, score: f64) -> StoredPoolMemberScore {
let score_scope = provider_key_pool_score_scope();
let score_identity = PoolMemberIdentity::provider_api_key(provider_id, key_id);
StoredPoolMemberScore {
id: provider_key_pool_score_id(&score_identity, &score_scope),
pool_kind: score_identity.pool_kind.clone(),
pool_id: score_identity.pool_id.clone(),
member_kind: score_identity.member_kind.clone(),
member_id: score_identity.member_id.clone(),
capability: score_scope.capability.clone(),
scope_kind: score_scope.scope_kind.clone(),
scope_id: score_scope.scope_id.clone(),
score,
hard_state: PoolMemberHardState::Available,
score_version: 1,
score_reason: json!({ "weights": { "manual_priority": score } }),
last_ranked_at: Some(1_700_000_000),
last_scheduled_at: None,
last_success_at: None,
last_failure_at: None,
failure_count: 0,
last_probe_attempt_at: None,
last_probe_success_at: None,
last_probe_failure_at: None,
probe_failure_count: 0,
probe_status: PoolMemberProbeStatus::Ok,
updated_at: 1_700_000_050,
}
}
fn sample_provider_key_usage_row(
id: &str,
request_id: &str,
@@ -928,11 +958,16 @@ async fn gateway_sorts_admin_pool_keys_by_imported_and_last_used_time() {
Vec::new(),
vec![old_key, fresh_key, active_key],
));
let pool_score_repository = Arc::new(InMemoryPoolMemberScoreRepository::seed(vec![
sample_pool_member_score("provider-openai", "key-openai-fresh", 0.35),
sample_pool_member_score("provider-openai", "key-openai-active", 0.92),
]));
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
));
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_reader_for_tests(provider_catalog_repository)
.with_pool_score_repository_for_tests(pool_score_repository),
);
let default_response = local_admin_pool_response(
&state,
@@ -999,6 +1034,30 @@ async fn gateway_sorts_admin_pool_keys_by_imported_and_last_used_time() {
.map(|item| item["key_name"].as_str().unwrap_or_default())
.collect::<Vec<_>>();
assert_eq!(last_used_names, vec!["active", "old", "fresh"]);
let score_response = local_admin_pool_response(
&state,
http::Method::GET,
"/api/admin/pool/provider-openai/keys?page=1&page_size=50&status=all&sort_by=score&sort_order=desc",
None,
)
.await;
assert_eq!(score_response.status(), StatusCode::OK);
let score_payload: serde_json::Value = serde_json::from_slice(
&to_bytes(score_response.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("json body should parse");
let score_names = score_payload["keys"]
.as_array()
.expect("keys should be array")
.iter()
.map(|item| item["key_name"].as_str().unwrap_or_default())
.collect::<Vec<_>>();
assert_eq!(score_names, vec!["active", "fresh", "old"]);
assert_eq!(score_payload["keys"][0]["pool_score"]["score"], json!(0.92));
assert!(score_payload["keys"][2]["pool_score"].is_null());
}
#[tokio::test]