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

Merge branch 'fawney19:main' into main

Browse Source
This commit is contained in:
ZheFox
2026-05-20 22:49:49 +08:00
committed by GitHub
parent 5a9f19cbf2 d49b0a8a45
commit 779877acd0
38 changed files with 1523 additions and 172 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -709,31 +709,39 @@ impl<'a> PoolKeyCursor<'a> {
}
async fn refill_queued_candidates(&mut self) -> bool {
let mut candidates = Vec::new();
let refill_target = self.window_size.max(1) as usize;
// Keep pool expansion bounded; the cursor freezes one small window at a time.
while candidates.len() < refill_target {
let Some(mut page_candidates) = self.next_page_candidates().await else {
break;
};
candidates.append(&mut page_candidates);
}
if candidates.is_empty() {
return false;
}
loop {
let mut candidates = Vec::new();
// Keep pool expansion bounded; the cursor freezes one small window at a time.
while candidates.len() < refill_target {
let Some(mut page_candidates) = self.next_page_candidates().await else {
break;
};
candidates.append(&mut page_candidates);
}
let (mut scheduled, mut skipped) = schedule_pool_page_candidates(
self.state,
candidates,
self.sticky_session_token.as_deref(),
)
.await;
scheduled.truncate(refill_target);
self.record_skipped_candidates(&skipped);
self.queued_candidates.extend(scheduled.drain(..));
self.skipped_candidates.append(&mut skipped);
!self.queued_candidates.is_empty()
if candidates.is_empty() {
return false;
}
let (mut scheduled, mut skipped) = schedule_pool_page_candidates(
self.state,
candidates,
self.sticky_session_token.as_deref(),
)
.await;
self.record_skipped_candidates(&skipped);
self.skipped_candidates.append(&mut skipped);
if scheduled.is_empty() {
continue;
}
scheduled.truncate(refill_target);
self.queued_candidates.extend(scheduled.drain(..));
return true;
}
}
async fn next_queued_candidate(&mut self) -> Option<EligibleLocalExecutionCandidate> {
@@ -2798,6 +2806,78 @@ mod tests {
assert_eq!(cursor.skip_reason_counts.get("pool_cooldown"), Some(&1));
}
#[tokio::test]
async fn pool_key_cursor_continues_after_exhausted_window() {
let provider_config = Some(json!({
"pool_advanced": {
"skip_exhausted_accounts": true
}
}));
let (provider, endpoint, mut keys, rows) = large_pool_fixture(3, provider_config.clone());
for key in keys.iter_mut().take(2) {
key.status_snapshot = Some(json!({
"quota": {
"provider_type": "openai",
"exhausted": true,
"usage_ratio": 1.0,
"windows": [
{
"code": "daily",
"used_ratio": 1.0,
"remaining_ratio": 0.0
}
]
}
}));
}
let data_state =
GatewayDataState::with_provider_catalog_and_minimal_candidate_selection_for_tests(
Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
keys,
)),
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(rows)),
)
.with_encryption_key_for_tests(aether_crypto::DEVELOPMENT_ENCRYPTION_KEY);
let app = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data_state);
let group = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"pool-group",
10,
provider_config,
);
let mut cursor = PoolKeyCursor::new(PlannerAppState::new(&app), group, None, None, None);
cursor.window_size = 2;
cursor.page_size = 2;
cursor.max_scanned_keys = 4;
let candidate = cursor
.next_key()
.await
.expect("cursor should scan past an exhausted window");
assert_eq!(candidate.candidate.key_id, "key-00002");
assert_eq!(candidate.orchestration.pool_key_index, Some(0));
assert!(candidate.orchestration.pool_key_lease.is_none());
assert_eq!(
cursor
.skip_reason_counts
.get(aether_pool_core::POOL_ACCOUNT_EXHAUSTED_SKIP_REASON),
Some(&2)
);
let skipped = cursor.take_skipped_candidates();
assert_eq!(skipped.len(), 2);
assert!(skipped.iter().all(|candidate| {
candidate.skip_reason == aether_pool_core::POOL_ACCOUNT_EXHAUSTED_SKIP_REASON
}));
}
#[tokio::test]
async fn pool_key_cursor_simulates_large_lru_pool_with_lazy_pages_and_dynamic_skips() {
const KEY_COUNT: usize = 2048;

13
apps/aether-gateway/src/handlers/admin/observability/usage/analytics_routes/aggregation.rs
View File

@@ -53,6 +53,15 @@ fn admin_usage_aggregation_by_provider_json(
json!(rows
.iter()
.map(|row| {
let identity_source = match row.secondary_name.as_deref() {
Some("legacy_name") => "legacy_name",
_ => "provider_id",
};
let provider_id = if identity_source == "provider_id" {
json!(row.group_key)
} else {
serde_json::Value::Null
};
let success_count = row.success_count.unwrap_or_default();
let error_count = row.request_count.saturating_sub(success_count);
let success_rate = if row.request_count == 0 {
@@ -67,7 +76,9 @@ fn admin_usage_aggregation_by_provider_json(
.filter(|value| !value.is_empty())
.unwrap_or(row.group_key.as_str());
json!({
"provider_id": row.group_key,
"provider_id": provider_id,
"provider_key": row.group_key,
"provider_identity_source": identity_source,
"provider": provider_name,
"request_count": row.request_count,
"total_tokens": row.total_tokens,

68
apps/aether-gateway/src/handlers/admin/provider/oauth/provisioning.rs
View File

@@ -2,7 +2,9 @@ use super::state::{
decode_jwt_claims, enrich_admin_provider_oauth_auth_config, json_non_empty_string,
json_u64_value,
};
use crate::handlers::admin::admin_provider_pool_config;
use crate::handlers::admin::request::AdminAppState;
use crate::maintenance::ensure_provider_key_pool_scores_for_keys;
use crate::provider_key_auth::provider_active_api_formats;
use crate::GatewayError;
use aether_data_contracts::repository::provider_catalog::{
@@ -168,6 +170,7 @@ pub(crate) async fn create_provider_oauth_catalog_key(
.app()
.invalidate_local_oauth_refresh_entry(&key.id)
.await;
seed_provider_oauth_pool_score(state, provider_id, key, now_unix_secs).await;
}
Ok(created)
}
@@ -221,10 +224,75 @@ pub(crate) async fn update_existing_provider_oauth_catalog_key(
.app()
.invalidate_local_oauth_refresh_entry(&key.id)
.await;
seed_provider_oauth_pool_score(state, &existing_key.provider_id, key, now_unix_secs).await;
}
Ok(persisted)
}
async fn seed_provider_oauth_pool_score(
state: &AdminAppState<'_>,
provider_id: &str,
key: &StoredProviderCatalogKey,
now_unix_secs: u64,
) {
let provider_id = provider_id.to_string();
let provider = match state
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&provider_id))
.await
{
Ok(mut providers) => providers.pop(),
Err(err) => {
tracing::debug!(
provider_id = %provider_id,
key_id = %key.id,
error = ?err,
"gateway provider oauth provisioning: failed to read provider for pool score seed"
);
return;
}
};
let Some(provider) = provider else {
return;
};
let Some(pool_config) = admin_provider_pool_config(&provider) else {
return;
};
let endpoints = match state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await
{
Ok(endpoints) => endpoints,
Err(err) => {
tracing::debug!(
provider_id = %provider_id,
key_id = %key.id,
error = ?err,
"gateway provider oauth provisioning: failed to read endpoints for pool score seed"
);
return;
}
};
let score_ensure_budget = (pool_config.score_fallback_scan_limit as usize).clamp(1, 50_000);
if let Err(err) = ensure_provider_key_pool_scores_for_keys(
state.as_ref(),
&provider,
&pool_config,
&endpoints,
std::slice::from_ref(key),
now_unix_secs,
score_ensure_budget,
)
.await
{
tracing::debug!(
provider_id = %provider_id,
key_id = %key.id,
error = ?err,
"gateway provider oauth provisioning: failed to seed pool score row"
);
}
}
fn provider_oauth_catalog_key_api_formats(
provider_type: &str,
api_formats: &[String],

24
apps/aether-gateway/src/handlers/admin/system/shared/modules.rs
View File

@@ -115,6 +115,30 @@ pub(crate) const ADMIN_MODULE_DEFINITIONS: &[AdminModuleDefinition] = &[
admin_menu_group: Some("system"),
admin_menu_order: 60,
},
AdminModuleDefinition {
name: "payment_gateways",
display_name: "支付配置",
description: "配置易支付、支付宝官方、微信支付官方和 Stripe 等支付网关",
category: "integration",
env_key: "PAYMENT_GATEWAYS_AVAILABLE",
default_available: true,
admin_route: Some("/admin/payment-gateways"),
admin_menu_icon: Some("CreditCard"),
admin_menu_group: None,
admin_menu_order: 70,
},
AdminModuleDefinition {
name: "referral",
display_name: "邀请返利",
description: "管理用户邀请关系与返利记录，支持比例返利和人头返利",
category: "integration",
env_key: "REFERRAL_AVAILABLE",
default_available: true,
admin_route: Some("/admin/referrals"),
admin_menu_icon: Some("Gift"),
admin_menu_group: Some("management"),
admin_menu_order: 75,
},
];
#[derive(Debug, Clone, serde::Deserialize)]

2
apps/aether-gateway/src/handlers/public/system_modules_helpers/keys_grouped.rs
View File

@@ -46,7 +46,7 @@ pub(crate) async fn build_admin_keys_grouped_by_format_payload(
let (endpoints_result, keys_result) = tokio::join!(
state.list_provider_catalog_endpoints_by_provider_ids(&provider_ids),
state.list_provider_catalog_key_summaries_by_provider_ids(&provider_ids),
state.list_provider_catalog_keys_by_provider_ids(&provider_ids),
);
let active_endpoints = endpoints_result

112
apps/aether-gateway/src/tests/control/admin/endpoints/keys.rs
View File

@@ -5,7 +5,12 @@ use aether_crypto::{
decrypt_python_fernet_ciphertext, encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY,
};
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::provider_catalog::ProviderCatalogReadRepository;
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, StoredProviderCatalogEndpoint,
StoredProviderCatalogKey, StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats,
StoredProviderCatalogProvider,
};
use aether_data_contracts::DataLayerError;
use axum::body::Body;
use axum::routing::any;
use axum::{extract::Request, Json, Router};
@@ -22,6 +27,103 @@ use crate::constants::{
};
use crate::data::GatewayDataState;
struct SummaryNullingProviderCatalogReadRepository {
inner: InMemoryProviderCatalogReadRepository,
}
impl SummaryNullingProviderCatalogReadRepository {
fn seed(
providers: Vec<StoredProviderCatalogProvider>,
endpoints: Vec<StoredProviderCatalogEndpoint>,
keys: Vec<StoredProviderCatalogKey>,
) -> Self {
Self {
inner: InMemoryProviderCatalogReadRepository::seed(providers, endpoints, keys),
}
}
}
#[async_trait::async_trait]
impl ProviderCatalogReadRepository for SummaryNullingProviderCatalogReadRepository {
async fn list_providers(
&self,
active_only: bool,
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
self.inner.list_providers(active_only).await
}
async fn list_providers_by_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
self.inner.list_providers_by_ids(provider_ids).await
}
async fn list_endpoints_by_ids(
&self,
endpoint_ids: &[String],
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
self.inner.list_endpoints_by_ids(endpoint_ids).await
}
async fn list_endpoints_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
self.inner
.list_endpoints_by_provider_ids(provider_ids)
.await
}
async fn list_keys_by_ids(
&self,
key_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
self.inner.list_keys_by_ids(key_ids).await
}
async fn list_keys_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
self.inner.list_keys_by_provider_ids(provider_ids).await
}
async fn list_key_summaries_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
let mut keys = self.inner.list_keys_by_provider_ids(provider_ids).await?;
for key in &mut keys {
key.internal_priority = 50;
key.global_priority_by_format = None;
key.rate_multipliers = None;
key.request_count = None;
key.success_count = None;
key.error_count = None;
key.total_response_time_ms = None;
key.circuit_breaker_by_format = None;
}
Ok(keys)
}
async fn list_keys_page(
&self,
query: &ProviderCatalogKeyListQuery,
) -> Result<StoredProviderCatalogKeyPage, DataLayerError> {
self.inner.list_keys_page(query).await
}
async fn list_key_stats_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKeyStats>, DataLayerError> {
self.inner
.list_key_stats_by_provider_ids(provider_ids)
.await
}
}
#[tokio::test]
async fn gateway_handles_admin_provider_keys_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
@@ -2195,7 +2297,7 @@ async fn gateway_handles_admin_keys_grouped_by_format_locally_with_trusted_admin
key_b.created_at_unix_ms = Some(1_711_100_000);
key_b.updated_at_unix_secs = Some(1_711_100_100);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
let provider_catalog_repository = Arc::new(SummaryNullingProviderCatalogReadRepository::seed(
vec![
sample_provider("provider-openai", "openai", 10),
sample_provider("provider-claude", "claude", 20)
@@ -2252,6 +2354,12 @@ async fn gateway_handles_admin_keys_grouped_by_format_locally_with_trusted_admin
"https://api.openai.example"
);
assert_eq!(payload["openai:chat"][0]["capabilities"], json!(["1h缓存"]));
assert_eq!(payload["openai:chat"][0]["format_priority"], 3);
assert_eq!(
payload["openai:chat"][0]["global_priority_by_format"]["openai:chat"],
3
);
assert_eq!(payload["openai:chat"][0]["internal_priority"], 10);
assert_eq!(payload["claude:messages"][0]["provider_active"], false);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);

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

@@ -12,8 +12,12 @@ use aether_data::repository::management_tokens::{
use aether_data::repository::oauth_providers::{
InMemoryOAuthProviderRepository, OAuthProviderReadRepository,
};
use aether_data::repository::pool_scores::InMemoryPoolMemberScoreRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::proxy_nodes::InMemoryProxyNodeRepository;
use aether_data_contracts::repository::pool_scores::{
GetPoolMemberScoresByIdsQuery, PoolMemberHardState, PoolMemberIdentity, PoolScoreReadRepository,
};
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogReadRepository, ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint,
};
@@ -32,6 +36,7 @@ use super::super::{
use crate::admin_api::{
maybe_build_local_admin_provider_oauth_response, AdminAppState, AdminRequestContext,
};
use crate::ai_serving::{provider_key_pool_score_id, provider_key_pool_score_scope};
use crate::audit::AdminAuditEvent;
use crate::constants::{
GATEWAY_HEADER, TRUSTED_ADMIN_MANAGEMENT_TOKEN_ID_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER,
@@ -1930,6 +1935,7 @@ async fn gateway_batch_imports_chatgpt_web_access_tokens_with_pool_hints() {
let mut provider = sample_provider("provider-chatgpt-web", "chatgpt_web", 10);
provider.provider_type = "chatgpt_web".to_string();
provider.config = Some(json!({"pool_advanced": {}}));
let endpoint = sample_endpoint(
"endpoint-chatgpt-web-image",
"provider-chatgpt-web",
@@ -1941,6 +1947,7 @@ async fn gateway_batch_imports_chatgpt_web_access_tokens_with_pool_hints() {
vec![endpoint],
vec![],
));
let pool_score_repository = Arc::new(InMemoryPoolMemberScoreRepository::default());
let (token_url, token_handle) = start_server(token_server).await;
let gateway = build_router_with_state(
@@ -1950,6 +1957,7 @@ async fn gateway_batch_imports_chatgpt_web_access_tokens_with_pool_hints() {
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_pool_score_repository_for_tests(Arc::clone(&pool_score_repository))
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_token_url_for_tests(
@@ -2030,6 +2038,20 @@ async fn gateway_batch_imports_chatgpt_web_access_tokens_with_pool_hints() {
assert_eq!(auth_config["plan_type"], "plus");
assert_eq!(auth_config["user_id"], "user-pool-image");
let score_scope = provider_key_pool_score_scope();
let score_identity =
PoolMemberIdentity::provider_api_key("provider-chatgpt-web", persisted.id.clone());
let scores = pool_score_repository
.get_pool_member_scores_by_ids(&GetPoolMemberScoresByIdsQuery {
ids: vec![provider_key_pool_score_id(&score_identity, &score_scope)],
})
.await
.expect("pool score should load");
assert_eq!(scores.len(), 1);
assert_eq!(scores[0].member_id, persisted.id);
assert_eq!(scores[0].hard_state, PoolMemberHardState::Unknown);
assert!(scores[0].score > 0.0);
gateway_handle.abort();
token_handle.abort();
}

59
apps/aether-gateway/src/tests/control/admin/usage.rs
View File

@@ -665,8 +665,14 @@ async fn gateway_handles_admin_usage_aggregation_stats_locally_with_trusted_admi
let provider_items = provider_payload.as_array().expect("array response");
assert_eq!(provider_items.len(), 2);
assert_eq!(provider_items[0]["provider"], "OpenAI");
assert_eq!(provider_items[0]["provider_id"], "provider-openai");
assert_eq!(provider_items[0]["provider_key"], "provider-openai");
assert_eq!(provider_items[0]["provider_identity_source"], "provider_id");
assert_eq!(provider_items[0]["output_tokens"], 40);
assert_eq!(provider_items[1]["provider"], "Anthropic");
assert_eq!(provider_items[1]["provider_id"], "provider-anthropic");
assert_eq!(provider_items[1]["provider_key"], "provider-anthropic");
assert_eq!(provider_items[1]["provider_identity_source"], "provider_id");
assert_eq!(provider_items[1]["output_tokens"], 20);
let api_format_response = admin_request(reqwest::Client::new().get(format!(
@@ -693,6 +699,59 @@ async fn gateway_handles_admin_usage_aggregation_stats_locally_with_trusted_admi
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_usage_aggregation_stats_for_legacy_provider_name_rows() {
let legacy_usage = {
let mut usage = sample_usage_row(
"usage-legacy",
"req-legacy",
Some("user-1"),
Some("key-1"),
Some("primary"),
"Legacy Provider",
"gpt-5",
"completed",
120,
30,
0.3,
0.36,
DAY_1_UNIX_SECS,
);
usage.provider_id = None;
usage.total_tokens = usage.input_tokens;
usage
};
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![legacy_usage]));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_usage_reader_for_tests(
usage_repository,
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = admin_request(reqwest::Client::new().get(format!(
"{gateway_url}/api/admin/usage/aggregation/stats?group_by=provider&limit=10&start_date=2024-03-21&end_date=2024-03-22&tz_offset_minutes=0"
)))
.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");
let items = payload.as_array().expect("array response");
assert_eq!(items.len(), 1);
assert_eq!(items[0]["provider"], "Legacy Provider");
assert_eq!(items[0]["provider_id"], serde_json::Value::Null);
assert_eq!(items[0]["provider_key"], "Legacy Provider");
assert_eq!(items[0]["provider_identity_source"], "legacy_name");
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_returns_service_unavailable_for_admin_usage_replay_without_provider_catalog_reader(
) {