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

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,