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

feat: update gateway pool and usage flows

Browse Source
This commit is contained in:
fawney19
2026-05-12 21:05:11 +08:00
parent 38012c62ff
commit d1a47c068e
18 changed files with 582 additions and 328 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -3,6 +3,7 @@ pub(crate) use crate::handlers::admin::{
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,

339
apps/aether-gateway/src/dispatch/pool_scheduler.rs
View File

@@ -32,7 +32,8 @@ use crate::handlers::shared::provider_pool::{
admin_provider_pool_cache_affinity_enabled, admin_provider_pool_config_from_config_value,
};
use crate::handlers::shared::provider_pool::{
try_claim_admin_provider_pool_key, AdminProviderPoolConfig, AdminProviderPoolRuntimeState,
read_admin_provider_pool_key_cooldown_reason, AdminProviderPoolConfig,
AdminProviderPoolRuntimeState,
};
use crate::handlers::shared::{
parse_catalog_auth_config_json, provider_key_health_summary,
@@ -537,7 +538,7 @@ impl<'a> PoolKeyCursor<'a> {
async fn next_queued_candidate(&mut self) -> Option<EligibleLocalExecutionCandidate> {
while let Some(candidate) = self.queued_candidates.pop_front() {
let mut candidate = candidate;
if !self.attach_pool_key_lease(&mut candidate).await {
if self.skip_candidate_if_runtime_cooldown(&candidate).await {
continue;
}
candidate.orchestration.pool_key_index = Some(self.next_pool_key_index);
@@ -548,43 +549,42 @@ impl<'a> PoolKeyCursor<'a> {
None
}
async fn attach_pool_key_lease(
async fn skip_candidate_if_runtime_cooldown(
&mut self,
candidate: &mut EligibleLocalExecutionCandidate,
candidate: &EligibleLocalExecutionCandidate,
) -> bool {
if !pool_key_order_requires_exclusive_lease(&self.pool_key_order) {
return true;
}
let owner = self.state.app().tunnel.local_instance_id();
match try_claim_admin_provider_pool_key(
match read_admin_provider_pool_key_cooldown_reason(
self.state.app().runtime_state.as_ref(),
candidate.candidate.provider_id.as_str(),
candidate.candidate.key_id.as_str(),
owner,
)
.await
{
Ok(Some(lease)) => {
candidate.orchestration.pool_key_lease = Some(lease);
Ok(Some(_)) => {
self.record_skip_reason("pool_cooldown");
self.skipped_candidates
.push(SkippedLocalExecutionCandidate {
candidate: candidate.candidate.clone(),
skip_reason: "pool_cooldown",
transport: Some(candidate.transport.clone()),
ranking: candidate.ranking.clone(),
extra_data: None,
});
true
}
Ok(None) => {
self.record_skip_reason("pool_key_lease_busy");
false
}
Ok(None) => false,
Err(err) => {
warn!(
event_name = "pool_key_lease_claim_failed",
event_name = "pool_key_cooldown_check_failed",
log_type = "event",
provider_id = %candidate.candidate.provider_id,
endpoint_id = %candidate.candidate.endpoint_id,
model_id = %candidate.candidate.model_id,
key_id = %candidate.candidate.key_id,
error = ?err,
"gateway pool scheduler failed to claim pool key lease; scheduling key without lease"
"gateway pool scheduler failed to read pool key cooldown; scheduling key"
);
true
false
}
}
}
@@ -707,13 +707,6 @@ impl<'a> PoolKeyCursor<'a> {
}
}
fn pool_key_order_requires_exclusive_lease(order: &StoredPoolKeyCandidateOrder) -> bool {
!matches!(
order,
StoredPoolKeyCandidateOrder::CacheAffinity | StoredPoolKeyCandidateOrder::SingleAccount
)
}
fn pool_candidate_transport_policy_facts(
candidate: &aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate,
) -> CandidateTransportPolicyFacts<'_> {
@@ -1137,8 +1130,7 @@ fn apply_ai_pool_orchestration(
mod tests {
use super::{
apply_local_execution_pool_scheduler_with_runtime_map, build_pool_catalog_key_context,
pool_config_for_candidate, pool_key_order_requires_exclusive_lease, PoolCatalogKeyContext,
PoolKeyCursor,
pool_config_for_candidate, PoolCatalogKeyContext, PoolKeyCursor,
};
use crate::ai_serving::{
apply_local_runtime_candidate_terminal_reason, EligibleLocalExecutionCandidate,
@@ -1146,8 +1138,7 @@ mod tests {
};
use crate::data::GatewayDataState;
use crate::handlers::shared::provider_pool::{
record_admin_provider_pool_error, release_admin_provider_pool_key_lease,
try_claim_admin_provider_pool_key, AdminProviderPoolRuntimeState,
record_admin_provider_pool_error, AdminProviderPoolRuntimeState,
};
use crate::orchestration::LocalExecutionCandidateMetadata;
use crate::{AppState, LocalExecutionRuntimeMissDiagnostic};
@@ -1901,24 +1892,6 @@ mod tests {
assert_eq!(lru_cursor.pool_key_order, StoredPoolKeyCandidateOrder::Lru);
}
#[test]
fn pool_key_lease_exclusivity_follows_distribution_mode() {
assert!(!pool_key_order_requires_exclusive_lease(
&StoredPoolKeyCandidateOrder::CacheAffinity
));
assert!(!pool_key_order_requires_exclusive_lease(
&StoredPoolKeyCandidateOrder::SingleAccount
));
assert!(pool_key_order_requires_exclusive_lease(
&StoredPoolKeyCandidateOrder::Lru
));
assert!(pool_key_order_requires_exclusive_lease(
&StoredPoolKeyCandidateOrder::LoadBalance {
seed: "seed".to_string()
}
));
}
#[test]
fn pool_key_cursor_records_runtime_miss_when_exhausted_without_returning_key() {
let app = AppState::new().expect("state should build");
@@ -1941,8 +1914,8 @@ mod tests {
);
let mut cursor = PoolKeyCursor::new(PlannerAppState::new(&app), group, None, None, None)
.with_runtime_miss_diagnostic(trace_id, true);
cursor.record_skip_reason("pool_key_lease_busy");
cursor.record_skip_reason("pool_key_lease_busy");
cursor.record_skip_reason("pool_cooldown");
cursor.record_skip_reason("pool_cooldown");
cursor.record_skip_reason("transport_snapshot_missing");
cursor.log_exhausted();
@@ -1953,11 +1926,11 @@ mod tests {
.expect("runtime miss diagnostic should exist");
assert_eq!(diagnostic.reason, "all_candidates_skipped");
assert_eq!(diagnostic.skipped_candidate_count, Some(1));
assert_eq!(diagnostic.skip_reasons.get("pool_key_lease_busy"), Some(&1));
assert_eq!(diagnostic.skip_reasons.get("pool_cooldown"), Some(&1));
}
#[tokio::test]
async fn pool_key_cursor_freezes_queued_candidates_after_window_refill() {
async fn pool_key_cursor_rechecks_cooldown_for_frozen_window_candidates() {
let app = AppState::new().expect("state should build");
let provider_config = Some(json!({
"pool_advanced": {
@@ -2009,33 +1982,24 @@ mod tests {
let second = cursor
.next_key()
.await
.expect("second key should keep the frozen queue order");
assert_eq!(second.candidate.key_id, "key-b");
.expect("second key should skip the cooled-down frozen key");
assert_eq!(second.candidate.key_id, "key-c");
assert_eq!(second.orchestration.pool_key_index, Some(1));
let third = cursor
.next_key()
.await
.expect("third key should keep the frozen queue order");
assert_eq!(third.candidate.key_id, "key-c");
assert_eq!(third.orchestration.pool_key_index, Some(2));
let skipped = cursor.take_skipped_candidates();
assert!(skipped.is_empty());
assert_eq!(
skipped
.iter()
.map(|item| (item.candidate.key_id.as_str(), item.skip_reason))
.collect::<Vec<_>>(),
vec![("key-b", "pool_cooldown")]
);
}
#[tokio::test]
async fn pool_key_cursor_skips_busy_lease_and_claims_returned_key() {
async fn pool_key_cursor_allows_parallel_requests_to_use_same_healthy_key() {
let app = AppState::new().expect("state should build");
let provider_config = Some(json!({ "pool_advanced": { "lru_enabled": true } }));
let held_lease = try_claim_admin_provider_pool_key(
app.runtime_state.as_ref(),
"provider-pool",
"key-a",
"test-owner",
)
.await
.expect("lease claim should not fail")
.expect("first key should be claimed by test");
let group = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
@@ -2052,87 +2016,19 @@ mod tests {
10,
provider_config.clone(),
),
sample_eligible_candidate("provider-pool", "endpoint-1", "key-b", 10, provider_config),
]);
let candidate = cursor
.next_key()
.await
.expect("cursor should skip busy key and return next key");
assert_eq!(candidate.candidate.key_id, "key-b");
assert_eq!(candidate.orchestration.pool_key_index, Some(0));
assert_eq!(
cursor.skip_reason_counts.get("pool_key_lease_busy"),
Some(&1)
);
let returned_lease = candidate
.orchestration
.pool_key_lease
.clone()
.expect("returned key should carry its lease");
assert!(
try_claim_admin_provider_pool_key(
app.runtime_state.as_ref(),
sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-b",
"other-owner",
)
.await
.expect("second claim should not fail")
.is_none(),
"returned key should remain leased until execution releases it"
);
release_admin_provider_pool_key_lease(app.runtime_state.as_ref(), &held_lease)
.await
.expect("held lease release should not fail");
release_admin_provider_pool_key_lease(app.runtime_state.as_ref(), &returned_lease)
.await
.expect("returned lease release should not fail");
}
#[tokio::test]
async fn pool_key_cursor_allows_busy_lease_for_cache_affinity_mode() {
let app = AppState::new().expect("state should build");
let provider_config = Some(json!({
"pool_advanced": {
"scheduling_presets": [
{"preset": "cache_affinity", "enabled": true}
]
}
}));
let held_lease = try_claim_admin_provider_pool_key(
app.runtime_state.as_ref(),
"provider-pool",
"key-a",
"test-owner",
)
.await
.expect("lease claim should not fail")
.expect("first key should be claimed by test");
let group = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"pool-group",
10,
provider_config.clone(),
);
let mut cursor = PoolKeyCursor::new(PlannerAppState::new(&app), group, None, None, None);
cursor.queued_candidates = VecDeque::from([
sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-a",
10,
provider_config.clone(),
),
sample_eligible_candidate("provider-pool", "endpoint-1", "key-b", 10, provider_config),
]);
let candidate = cursor
.next_key()
.await
.expect("cache-affinity cursor should allow repeated busy key");
.expect("cursor should return the first healthy key");
assert_eq!(candidate.candidate.key_id, "key-a");
assert_eq!(candidate.orchestration.pool_key_index, Some(0));
assert!(candidate.orchestration.pool_key_lease.is_none());
@@ -2140,9 +2036,82 @@ mod tests {
.skip_reason_counts
.contains_key("pool_key_lease_busy"));
release_admin_provider_pool_key_lease(app.runtime_state.as_ref(), &held_lease)
let group = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"pool-group",
10,
provider_config,
);
let mut second_cursor =
PoolKeyCursor::new(PlannerAppState::new(&app), group, None, None, None);
second_cursor.queued_candidates = VecDeque::from([
sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-a",
10,
Some(json!({ "pool_advanced": { "lru_enabled": true } })),
),
sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-b",
10,
Some(json!({ "pool_advanced": { "lru_enabled": true } })),
),
]);
let second_candidate = second_cursor
.next_key()
.await
.expect("held lease release should not fail");
.expect("second request should also be allowed to pick the same healthy key");
assert_eq!(second_candidate.candidate.key_id, "key-a");
assert!(second_candidate.orchestration.pool_key_lease.is_none());
}
#[tokio::test]
async fn pool_key_cursor_skips_key_after_account_cooldown_is_recorded() {
let app = AppState::new().expect("state should build");
let provider_config = Some(json!({ "pool_advanced": { "lru_enabled": true } }));
let group = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"pool-group",
10,
provider_config.clone(),
);
let pool_config = pool_config_for_candidate(&group).expect("pool config should parse");
let mut cursor = PoolKeyCursor::new(PlannerAppState::new(&app), group, None, None, None);
cursor.queued_candidates = VecDeque::from([
sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-a",
10,
provider_config.clone(),
),
sample_eligible_candidate("provider-pool", "endpoint-1", "key-b", 10, provider_config),
]);
record_admin_provider_pool_error(
app.runtime_state.as_ref(),
"provider-pool",
"key-a",
&pool_config,
429,
None,
None,
)
.await;
let candidate = cursor
.next_key()
.await
.expect("cursor should skip cooled-down key and return next key");
assert_eq!(candidate.candidate.key_id, "key-b");
assert_eq!(candidate.orchestration.pool_key_index, Some(0));
assert!(candidate.orchestration.pool_key_lease.is_none());
assert_eq!(cursor.skip_reason_counts.get("pool_cooldown"), Some(&1));
}
#[tokio::test]
@@ -2177,21 +2146,7 @@ mod tests {
provider_config.clone(),
);
let pool_config = pool_config_for_candidate(&group).expect("pool config should parse");
let mut held_leases = Vec::new();
for key_id in ["key-00000", "key-00001"] {
held_leases.push(
try_claim_admin_provider_pool_key(
app.runtime_state.as_ref(),
"provider-pool",
key_id,
"large-pool-test",
)
.await
.expect("lease claim should not fail")
.expect("test key should claim"),
);
}
for key_id in ["key-00002", "key-00003"] {
record_admin_provider_pool_error(
app.runtime_state.as_ref(),
"provider-pool",
@@ -2219,22 +2174,16 @@ mod tests {
);
let mut returned_ids = Vec::new();
let mut returned_leases = Vec::new();
for _ in 0..10 {
let candidate = cursor
.next_key()
.await
.expect("large pool should return first page candidates");
returned_ids.push(candidate.candidate.key_id.clone());
returned_leases.push(
candidate
.orchestration
.pool_key_lease
.expect("lru pool candidate should carry exclusive lease"),
);
assert!(candidate.orchestration.pool_key_lease.is_none());
}
assert_eq!(returned_ids.first().map(String::as_str), Some("key-00004"));
assert_eq!(returned_ids.last().map(String::as_str), Some("key-00013"));
assert_eq!(returned_ids.first().map(String::as_str), Some("key-00002"));
assert_eq!(returned_ids.last().map(String::as_str), Some("key-00011"));
assert_eq!(cursor.scanned_keys, 64);
assert!(
cursor.queued_candidates.len() <= cursor.window_size as usize,
@@ -2251,39 +2200,18 @@ mod tests {
None,
)
.await;
held_leases.push(
try_claim_admin_provider_pool_key(
app.runtime_state.as_ref(),
"provider-pool",
"key-00015",
"large-pool-test",
)
.await
.expect("dynamic lease claim should not fail")
.expect("dynamic busy key should claim"),
);
let candidate = cursor
.next_key()
.await
.expect("cursor should keep the frozen window despite later runtime changes");
assert_eq!(candidate.candidate.key_id, "key-00014");
assert_eq!(candidate.candidate.key_id, "key-00012");
returned_ids.push(candidate.candidate.key_id.clone());
returned_leases.push(
candidate
.orchestration
.pool_key_lease
.expect("lru pool candidate should carry exclusive lease"),
);
assert!(candidate.orchestration.pool_key_lease.is_none());
while let Some(candidate) = cursor.next_key().await {
returned_ids.push(candidate.candidate.key_id.clone());
returned_leases.push(
candidate
.orchestration
.pool_key_lease
.expect("lru pool candidate should carry exclusive lease"),
);
assert!(candidate.orchestration.pool_key_lease.is_none());
}
let max_returned_windows = aether_dispatch_core::DEFAULT_POOL_MAX_SCAN
@@ -2293,34 +2221,23 @@ mod tests {
<= (max_returned_windows * aether_dispatch_core::DEFAULT_POOL_WINDOW_SIZE) as usize,
"cursor should only return bounded frozen windows per request"
);
assert_eq!(returned_ids.len(), 125);
assert_eq!(returned_ids.len(), 127);
assert_eq!(returned_ids.last().map(String::as_str), Some("key-00463"));
assert_eq!(cursor.scanned_keys, 512);
assert_eq!(
cursor.skip_reason_counts.get("pool_key_lease_busy"),
Some(&3)
);
assert_eq!(cursor.skip_reason_counts.get("pool_cooldown"), Some(&2));
for skipped in ["key-00000", "key-00001", "key-00002", "key-00003"] {
assert_eq!(cursor.skip_reason_counts.get("pool_cooldown"), Some(&3));
assert!(!cursor
.skip_reason_counts
.contains_key("pool_key_lease_busy"));
for skipped in ["key-00000", "key-00001", "key-00014"] {
assert!(
!returned_ids.iter().any(|key_id| key_id == skipped),
"{skipped} should have been skipped"
);
}
assert!(
returned_ids.iter().any(|key_id| key_id == "key-00014"),
"key-00014 was already frozen in the current window"
returned_ids.iter().any(|key_id| key_id == "key-00015"),
"key-00015 should not be blocked by request-scoped leases"
);
assert!(
!returned_ids.iter().any(|key_id| key_id == "key-00015"),
"key-00015 should still be skipped because exclusive lease is checked at dispatch time"
);
for lease in held_leases.into_iter().chain(returned_leases) {
release_admin_provider_pool_key_lease(app.runtime_state.as_ref(), &lease)
.await
.expect("lease release should not fail");
}
}
#[test]

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

@@ -29,12 +29,14 @@ pub(crate) use self::provider::oauth::quota::antigravity::refresh_antigravity_pr
pub(crate) use self::provider::oauth::quota::chatgpt_web::refresh_chatgpt_web_provider_quota_locally;
pub(crate) use self::provider::oauth::quota::codex::refresh_codex_provider_quota_locally;
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,
};
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;
pub(crate) use self::provider::pool_admin::maybe_build_local_admin_pool_response;
pub(crate) use self::provider::write::provider::reconcile_admin_fixed_provider_template_endpoints;
pub(crate) use self::provider::{
maybe_build_local_admin_provider_oauth_response, maybe_build_local_admin_providers_response,
};

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

@@ -18,6 +18,24 @@ use super::super::oauth::quota::chatgpt_web::refresh_chatgpt_web_provider_quota_
use super::super::oauth::quota::codex::refresh_codex_provider_quota_locally;
use super::super::oauth::quota::kiro::refresh_kiro_provider_quota_locally;
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::write::provider::reconcile_admin_fixed_provider_template_endpoints;
fn unsupported_provider_quota_refresh_response(provider_type: &str) -> Response<Body> {
let message = unsupported_provider_quota_refresh_message(provider_type);
Json(json!({
"success": 0,
"failed": 0,
"total": 0,
"results": [],
"message": message,
"auto_removed": 0,
}))
.into_response()
}
pub(super) async fn maybe_handle(
state: &AdminAppState<'_>,
@@ -85,41 +103,48 @@ pub(super) async fn maybe_handle(
let raw_key_ids = payload.key_ids;
let selected_key_ids = normalize_string_id_list(raw_key_ids.clone());
let explicit_key_ids_requested = raw_key_ids.is_some();
let endpoints = state
let is_fixed_provider = state
.fixed_provider_template(&provider.provider_type)
.is_some();
if !is_fixed_provider && !provider_type_supports_quota_refresh(&normalized_provider_type) {
return Ok(None);
}
let mut endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await?;
let endpoint = match normalized_provider_type.as_str() {
"codex" => endpoints.into_iter().find(|endpoint| {
endpoint.is_active
&& crate::ai_serving::is_openai_responses_format(&endpoint.api_format)
}),
"antigravity" => endpoints.into_iter().find(|endpoint| {
endpoint.is_active
&& endpoint
.api_format
.trim()
.eq_ignore_ascii_case("gemini:generate_content")
}),
"kiro" => endpoints
.iter()
.find(|endpoint| {
endpoint.is_active
&& endpoint
.api_format
.trim()
.eq_ignore_ascii_case("claude:messages")
})
.cloned()
.or_else(|| endpoints.into_iter().find(|endpoint| endpoint.is_active)),
"chatgpt_web" => endpoints.into_iter().find(|endpoint| {
endpoint.is_active
&& endpoint
.api_format
.trim()
.eq_ignore_ascii_case("openai:image")
}),
_ => return Ok(None),
};
let mut endpoint =
provider_oauth_runtime_endpoint_for_provider(&normalized_provider_type, &endpoints);
if endpoint.is_none() && is_fixed_provider {
if !state.has_provider_catalog_data_writer() {
if !provider_type_supports_quota_refresh(&normalized_provider_type) {
return Ok(Some(unsupported_provider_quota_refresh_response(
&normalized_provider_type,
)));
}
return Ok(Some(
(
http::StatusCode::BAD_REQUEST,
Json(json!({ "detail": "固定 Provider 端点缺失，且 provider catalog writer 不可用，无法自动补全端点" })),
)
.into_response(),
));
}
reconcile_admin_fixed_provider_template_endpoints(state, &provider).await?;
endpoints = state
.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);
}
if !provider_type_supports_quota_refresh(&normalized_provider_type) {
return Ok(Some(unsupported_provider_quota_refresh_response(
&normalized_provider_type,
)));
}
let Some(endpoint) = endpoint else {
let detail = match normalized_provider_type.as_str() {
@@ -127,6 +152,8 @@ pub(super) async fn maybe_handle(
"antigravity" => "找不到有效的 gemini:generate_content 端点",
"kiro" => "找不到有效的 Kiro 端点",
"chatgpt_web" => "找不到有效的 openai:image 端点",
"claude_code" => "找不到有效的 claude:messages 端点",
"gemini_cli" | "vertex_ai" => "找不到有效的 gemini:generate_content 端点",
_ => "找不到有效端点",
};
return Ok(Some(

19
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs
View File

@@ -50,6 +50,25 @@ pub(crate) fn normalize_string_id_list(values: Option<Vec<String>>) -> Option<Ve
admin_provider_quota_pure::normalize_string_id_list(values)
}
pub(crate) fn provider_type_supports_quota_refresh(provider_type: &str) -> bool {
matches!(
provider_type.trim().to_ascii_lowercase().as_str(),
"codex" | "kiro" | "antigravity" | "chatgpt_web"
)
}
pub(crate) fn unsupported_provider_quota_refresh_message(provider_type: &str) -> String {
match provider_type.trim().to_ascii_lowercase().as_str() {
"claude_code" => "Claude Code 暂不支持自动刷新额度：上游没有稳定可用的账号额度查询接口",
"gemini_cli" => {
"Gemini CLI 暂不支持自动刷新额度：当前只能通过模型同步/缓存快照展示已知配额信息"
}
"vertex_ai" => "Vertex AI 暂不支持自动刷新额度：额度属于 Google Cloud 项目/区域配额",
_ => "该 Provider 暂不支持自动刷新额度",
}
.to_string()
}
pub(super) fn coerce_json_u64(value: &serde_json::Value) -> Option<u64> {
admin_provider_quota_pure::coerce_json_u64(value)
}

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

@@ -60,6 +60,48 @@ pub(crate) fn provider_oauth_runtime_endpoint_for_provider(
.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)

4
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/keys.rs
View File

@@ -14,10 +14,6 @@ pub(super) fn pool_cooldown_key(provider_id: &str, key_id: &str) -> String {
format!("ap:{provider_id}:cooldown:{key_id}")
}
pub(super) fn pool_lease_key(provider_id: &str, key_id: &str) -> String {
format!("ap:{provider_id}:lease:{key_id}")
}
pub(super) fn pool_cooldown_index_key(provider_id: &str) -> String {
format!("ap:{provider_id}:cooldown_idx")
}

19
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/leases.rs
View File

@@ -1,23 +1,4 @@
use super::keys::pool_lease_key;
use aether_runtime_state::{DataLayerError, RuntimeLockLease, RuntimeState};
use std::time::Duration;
pub(crate) const ADMIN_PROVIDER_POOL_KEY_LEASE_TTL_MS: u64 = 15 * 60 * 1000;
pub(crate) async fn try_claim_admin_provider_pool_key(
runtime: &RuntimeState,
provider_id: &str,
key_id: &str,
owner: &str,
) -> Result<Option<RuntimeLockLease>, DataLayerError> {
runtime
.lock_try_acquire(
&pool_lease_key(provider_id, key_id),
owner,
Duration::from_millis(ADMIN_PROVIDER_POOL_KEY_LEASE_TTL_MS),
)
.await
}
pub(crate) async fn release_admin_provider_pool_key_lease(
runtime: &RuntimeState,

8
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/mod.rs
View File

@@ -5,16 +5,14 @@ mod reads;
mod status;
mod writes;
pub(crate) use self::leases::{
release_admin_provider_pool_key_lease, try_claim_admin_provider_pool_key,
ADMIN_PROVIDER_POOL_KEY_LEASE_TTL_MS,
};
pub(crate) use self::leases::release_admin_provider_pool_key_lease;
pub(crate) use self::mutations::{
clear_admin_provider_pool_cooldown, reset_admin_provider_pool_cost,
};
pub(crate) use self::reads::{
read_admin_provider_pool_cooldown_count, read_admin_provider_pool_cooldown_counts,
read_admin_provider_pool_cooldown_key_ids, read_admin_provider_pool_runtime_state,
read_admin_provider_pool_cooldown_key_ids, read_admin_provider_pool_key_cooldown_reason,
read_admin_provider_pool_runtime_state,
};
pub(crate) use self::status::build_admin_provider_pool_status_payload;
pub(crate) use self::writes::{

12
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/reads.rs
View File

@@ -7,7 +7,7 @@ use crate::handlers::admin::provider::pool::config::admin_provider_pool_cache_af
use crate::handlers::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolRuntimeState,
};
use aether_runtime_state::RuntimeState;
use aether_runtime_state::{DataLayerError, RuntimeState};
use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::warn;
@@ -202,3 +202,13 @@ pub(crate) async fn read_admin_provider_pool_cooldown_key_ids(
.await
.unwrap_or_default()
}
pub(crate) async fn read_admin_provider_pool_key_cooldown_reason(
runtime: &RuntimeState,
provider_id: &str,
key_id: &str,
) -> Result<Option<String>, DataLayerError> {
runtime
.kv_get(&pool_cooldown_key(provider_id, key_id))
.await
}

5
apps/aether-gateway/src/handlers/admin/system/management_tokens.rs
View File

@@ -328,10 +328,7 @@ pub(crate) async fn maybe_build_local_admin_management_tokens_response(
.admin_principal
.as_ref()
.and_then(|principal| principal.management_token_permissions.as_deref())
.map_or(
true,
management_token_permissions_cover_all_assignable_permissions,
);
.is_none_or(management_token_permissions_cover_all_assignable_permissions);
if is_management_token && !management_token_is_full {
return Ok(Some(
(

8
apps/aether-gateway/src/handlers/shared/provider_pool.rs
View File

@@ -2,10 +2,10 @@ pub(crate) use super::super::admin::provider::pool::config::{
admin_provider_pool_cache_affinity_enabled, admin_provider_pool_config_from_config_value,
};
pub(crate) use super::super::admin::provider::pool::runtime::{
admin_provider_pool_key_circuit_breaker_reason, read_admin_provider_pool_runtime_state,
record_admin_provider_pool_error, record_admin_provider_pool_stream_timeout,
record_admin_provider_pool_success, release_admin_provider_pool_key_lease,
try_claim_admin_provider_pool_key, ADMIN_PROVIDER_POOL_KEY_LEASE_TTL_MS,
admin_provider_pool_key_circuit_breaker_reason, read_admin_provider_pool_key_cooldown_reason,
read_admin_provider_pool_runtime_state, record_admin_provider_pool_error,
record_admin_provider_pool_stream_timeout, record_admin_provider_pool_success,
release_admin_provider_pool_key_lease,
};
pub(crate) use super::super::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolRuntimeState, AdminProviderPoolSchedulingPreset,

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

@@ -16,6 +16,7 @@ use tracing::{debug, info, warn};
use crate::admin_api::{
admin_provider_pool_config, 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, AdminAppState,
};
@@ -108,10 +109,7 @@ fn now_unix_secs() -> u64 {
}
fn provider_supports_quota_probe(provider_type: &str) -> bool {
matches!(
provider_type.trim().to_ascii_lowercase().as_str(),
"codex" | "kiro" | "antigravity" | "chatgpt_web"
)
provider_type_supports_quota_refresh(provider_type)
}
fn json_number(value: Option<&Value>) -> Option<f64> {
@@ -426,6 +424,37 @@ fn endpoint_for_probe(
provider_oauth_runtime_endpoint_for_provider(provider_type, endpoints)
}
async fn endpoint_for_probe_with_reconcile(
state: &AppState,
admin_state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
provider_type: &str,
endpoints_by_provider: &mut BTreeMap<String, Vec<StoredProviderCatalogEndpoint>>,
) -> Result<Option<StoredProviderCatalogEndpoint>, GatewayError> {
let endpoints = endpoints_by_provider
.get(&provider.id)
.map(Vec::as_slice)
.unwrap_or(&[]);
if let Some(endpoint) = endpoint_for_probe(provider_type, endpoints) {
return Ok(Some(endpoint));
}
if admin_state
.fixed_provider_template(&provider.provider_type)
.is_none()
{
return Ok(None);
}
reconcile_admin_fixed_provider_template_endpoints(admin_state, provider).await?;
let refreshed = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id))
.await?;
let endpoint = endpoint_for_probe(provider_type, &refreshed);
endpoints_by_provider.insert(provider.id.clone(), refreshed);
Ok(endpoint)
}
async fn refresh_provider_probe_keys(
admin_state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
@@ -688,10 +717,15 @@ pub(crate) async fn perform_pool_quota_probe_once_with_config(
};
for (provider, provider_type, pool_config) in providers {
let endpoints = endpoints_by_provider
.remove(&provider.id)
.unwrap_or_default();
let Some(endpoint) = endpoint_for_probe(&provider_type, &endpoints) else {
let Some(endpoint) = endpoint_for_probe_with_reconcile(
state,
&admin_state,
&provider,
&provider_type,
&mut endpoints_by_provider,
)
.await?
else {
summary.providers_skipped += 1;
debug!(
provider_id = %provider.id,
@@ -700,6 +734,9 @@ pub(crate) async fn perform_pool_quota_probe_once_with_config(
);
continue;
};
let endpoints = endpoints_by_provider
.remove(&provider.id)
.unwrap_or_else(|| vec![endpoint.clone()]);
let interval_minutes = pool_config.probing_interval_minutes;
let interval_seconds = interval_minutes.clamp(1, 1440).saturating_mul(60);

227
apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs
View File

@@ -829,6 +829,233 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_kiro_with_trusted_ad
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_refresh_kiro_quota_reconciles_missing_fixed_endpoint_before_refresh() {
let seen_endpoint_id = Arc::new(Mutex::new(None::<String>));
let seen_endpoint_id_clone = Arc::clone(&seen_endpoint_id);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |request: Request| {
let seen_endpoint_id_inner = Arc::clone(&seen_endpoint_id_clone);
async move {
let plan: aether_contracts::ExecutionPlan = serde_json::from_slice(
&to_bytes(request.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("plan should parse");
*seen_endpoint_id_inner.lock().expect("mutex should lock") =
Some(plan.endpoint_id.clone());
let result = aether_contracts::ExecutionResult {
request_id: plan.request_id,
candidate_id: None,
status_code: 200,
headers: BTreeMap::new(),
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"subscriptionInfo": {
"subscriptionTitle": "KIRO PRO"
},
"usageBreakdownList": [{
"currentUsageWithPrecision": 1.0,
"usageLimitWithPrecision": 10.0,
"nextDateReset": 1_900_000_000u64
}]
})),
body_bytes_b64: None,
}),
telemetry: None,
error: None,
};
(StatusCode::OK, Json(result))
}
}),
);
let encrypted_auth_config = encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"access_token":"kiro-access-token","api_region":"us-west-2"}"#,
)
.expect("auth config ciphertext should build");
let encrypted_api_key =
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "__placeholder__")
.expect("api key ciphertext should build");
let key = StoredProviderCatalogKey::new(
"key-kiro-reconcile".to_string(),
"provider-kiro-reconcile".to_string(),
"default".to_string(),
"bearer".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(json!(["claude:messages"])),
encrypted_api_key,
Some(encrypted_auth_config),
None,
None,
None,
None,
None,
None,
)
.expect("key transport should build");
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![StoredProviderCatalogProvider::new(
"provider-kiro-reconcile".to_string(),
"kiro".to_string(),
Some("https://example.com".to_string()),
"kiro".to_string(),
)
.expect("provider should build")],
Vec::new(),
vec![key],
));
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url.clone())
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/endpoints/providers/provider-kiro-reconcile/refresh-quota"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], 1);
assert_eq!(payload["failed"], 0);
assert_eq!(payload["results"][0]["status"], "success");
let endpoints = provider_catalog_repository
.list_endpoints_by_provider_ids(&["provider-kiro-reconcile".to_string()])
.await
.expect("endpoints should read");
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!(
*seen_endpoint_id.lock().expect("mutex should lock"),
Some(endpoints[0].id.clone())
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_refresh_quota_reconciles_unsupported_fixed_provider_endpoints_before_clear_message(
) {
let cases = [
(
"provider-claude-code-reconcile",
"claude_code",
1usize,
"claude:messages",
"https://api.anthropic.com",
"Claude Code 暂不支持自动刷新额度",
),
(
"provider-gemini-cli-reconcile",
"gemini_cli",
1usize,
"gemini:generate_content",
"https://cloudcode-pa.googleapis.com",
"Gemini CLI 暂不支持自动刷新额度",
),
(
"provider-vertex-ai-reconcile",
"vertex_ai",
2usize,
"gemini:generate_content",
"https://aiplatform.googleapis.com",
"Vertex AI 暂不支持自动刷新额度",
),
];
let providers = cases
.iter()
.map(|(provider_id, provider_type, _, _, _, _)| {
StoredProviderCatalogProvider::new(
(*provider_id).to_string(),
(*provider_type).to_string(),
Some("https://example.com".to_string()),
(*provider_type).to_string(),
)
.expect("provider should build")
})
.collect::<Vec<_>>();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
providers,
Vec::new(),
Vec::new(),
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override("http://127.0.0.1:1")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
for (provider_id, _, endpoint_count, api_format, base_url, message_prefix) in cases {
let response = client
.post(format!(
"{gateway_url}/api/admin/endpoints/providers/{provider_id}/refresh-quota"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], 0);
assert_eq!(payload["failed"], 0);
assert_eq!(payload["total"], 0);
assert!(payload["message"]
.as_str()
.expect("message should be string")
.starts_with(message_prefix));
let endpoints = provider_catalog_repository
.list_endpoints_by_provider_ids(&[provider_id.to_string()])
.await
.expect("endpoints should read");
assert_eq!(endpoints.len(), endpoint_count);
assert!(endpoints
.iter()
.any(|endpoint| endpoint.api_format == api_format && endpoint.base_url == base_url));
}
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_reports_codex_quota_runtime_failures_locally_without_falling_back_to_admin_passthrough(
) {