mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
fix(kiro,pool,model): 对齐 Kiro 管理链路并修复全局模型删除行为 (#305)
* feat(pool): 号池支持跳过额度耗尽账号 - 新增 pool_advanced.skip_exhausted_accounts 开关及高级设置 UI, 默认关闭并兼容旧配置 - 为 Codex/Kiro 增加额度耗尽判定, 接入请求侧候选跳过并新增 account_quota_exhausted skip reason - 号池列表将额度耗尽账号标记为 blocked/额度耗尽, 并补充前后端相关测试 * fix(kiro): 对齐账号管理与 provider-query 的 Rust 行为 - 修复 Kiro 单条导入误走 import-refresh-token 的前端分流, 并为误用路径返回明确错误提示 - 为 Kiro 导入与本地请求链补齐 bearer 兼容, 同步放开账号启停等 Key 更新操作的 auth_type 校验 - 实现 Kiro provider-query 本地模型测试与 failover 执行链, 并修复结果弹窗在无 trace 时无法展示 attempts/响应体的问题 * fix(model): 删除全局模型时级联清理关联提供商模型 - 对齐 Python 版本删除逻辑, GlobalModel 删除前先在事务内清理关联的 Provider Model 记录 - 修复已绑定 Provider 的模型在 Rust SQL 仓库下会被外键约束拦住、无法正常删除的问题 - 增加管理端回归测试, 覆盖绑定 Provider Model 的 GlobalModel 删除场景 * fix(kiro,ci): 恢复 Kiro OAuth 持久化并修复 Rust CI * Fix oauth-managed provider key semantics --------- Co-authored-by: fawney19 <elky0401@gmail.com>
This commit is contained in:
Entropy.Xu
@@ -1,5 +1,6 @@
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
use aether_admin::provider::pool as admin_provider_pool_pure;
|
||||
use aether_data_contracts::repository::candidates::StoredRequestCandidate;
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||
use aether_scheduler_core::{
|
||||
@@ -20,6 +21,7 @@ pub(super) struct CandidateRuntimeSelectionSnapshot {
|
||||
pub(super) provider_concurrent_limits: BTreeMap<String, usize>,
|
||||
pub(super) provider_key_rpm_states: BTreeMap<String, StoredProviderCatalogKey>,
|
||||
provider_quota_blocks_requests: BTreeMap<String, bool>,
|
||||
key_account_quota_exhausted: BTreeMap<String, bool>,
|
||||
provider_key_rpm_reset_ats: BTreeMap<String, Option<u64>>,
|
||||
}
|
||||
|
||||
@@ -30,7 +32,14 @@ pub(super) async fn read_candidate_runtime_selection_snapshot(
|
||||
) -> Result<CandidateRuntimeSelectionSnapshot, GatewayError> {
|
||||
let recent_candidates = state.read_recent_request_candidates(128).await?;
|
||||
let provider_concurrent_limits = read_provider_concurrent_limits(state, candidates).await?;
|
||||
let provider_skip_exhausted_accounts =
|
||||
read_provider_skip_exhausted_account_map(state, candidates).await?;
|
||||
let provider_key_rpm_states = read_provider_key_rpm_states(state, candidates).await?;
|
||||
let key_account_quota_exhausted = read_key_account_quota_exhaustion_map(
|
||||
candidates,
|
||||
&provider_key_rpm_states,
|
||||
&provider_skip_exhausted_accounts,
|
||||
);
|
||||
let provider_quota_blocks_requests =
|
||||
read_provider_quota_block_map(state, candidates, now_unix_secs).await?;
|
||||
let provider_key_rpm_reset_ats =
|
||||
@@ -41,6 +50,7 @@ pub(super) async fn read_candidate_runtime_selection_snapshot(
|
||||
provider_concurrent_limits,
|
||||
provider_key_rpm_states,
|
||||
provider_quota_blocks_requests,
|
||||
key_account_quota_exhausted,
|
||||
provider_key_rpm_reset_ats,
|
||||
})
|
||||
}
|
||||
@@ -89,6 +99,11 @@ pub(super) fn is_candidate_selectable(
|
||||
.get(candidate.provider_id.as_str())
|
||||
.copied()
|
||||
.unwrap_or(false),
|
||||
account_quota_exhausted: snapshot
|
||||
.key_account_quota_exhausted
|
||||
.get(candidate.key_id.as_str())
|
||||
.copied()
|
||||
.unwrap_or(false),
|
||||
rpm_reset_at: snapshot
|
||||
.provider_key_rpm_reset_ats
|
||||
.get(candidate.key_id.as_str())
|
||||
@@ -122,6 +137,11 @@ pub(super) fn current_candidate_runtime_skip_reason(
|
||||
now_unix_secs,
|
||||
cached_affinity_target,
|
||||
provider_quota_blocks_requests,
|
||||
account_quota_exhausted: snapshot
|
||||
.key_account_quota_exhausted
|
||||
.get(candidate.key_id.as_str())
|
||||
.copied()
|
||||
.unwrap_or(false),
|
||||
rpm_reset_at,
|
||||
})
|
||||
}
|
||||
@@ -192,6 +212,64 @@ async fn read_provider_quota_block_map(
|
||||
Ok(quota_blocks)
|
||||
}
|
||||
|
||||
async fn read_provider_skip_exhausted_account_map(
|
||||
state: &(impl SchedulerRuntimeState + ?Sized),
|
||||
candidates: &[SchedulerMinimalCandidateSelectionCandidate],
|
||||
) -> Result<BTreeMap<String, bool>, GatewayError> {
|
||||
let provider_ids = candidates
|
||||
.iter()
|
||||
.map(|candidate| candidate.provider_id.clone())
|
||||
.collect::<BTreeSet<_>>()
|
||||
.into_iter()
|
||||
.collect::<Vec<_>>();
|
||||
if provider_ids.is_empty() {
|
||||
return Ok(BTreeMap::new());
|
||||
}
|
||||
|
||||
let providers = state
|
||||
.read_provider_catalog_providers_by_ids(&provider_ids)
|
||||
.await?;
|
||||
Ok(providers
|
||||
.into_iter()
|
||||
.map(|provider| {
|
||||
let skip_exhausted_accounts = provider
|
||||
.config
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("pool_advanced"))
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.and_then(|value| value.get("skip_exhausted_accounts"))
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
.unwrap_or(false);
|
||||
(provider.id, skip_exhausted_accounts)
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn read_key_account_quota_exhaustion_map(
|
||||
candidates: &[SchedulerMinimalCandidateSelectionCandidate],
|
||||
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
|
||||
provider_skip_exhausted_accounts: &BTreeMap<String, bool>,
|
||||
) -> BTreeMap<String, bool> {
|
||||
candidates
|
||||
.iter()
|
||||
.map(|candidate| {
|
||||
let exhausted = provider_skip_exhausted_accounts
|
||||
.get(candidate.provider_id.as_str())
|
||||
.copied()
|
||||
.unwrap_or(false)
|
||||
&& provider_key_rpm_states
|
||||
.get(candidate.key_id.as_str())
|
||||
.is_some_and(|key| {
|
||||
admin_provider_pool_pure::admin_pool_key_account_quota_exhausted(
|
||||
key,
|
||||
candidate.provider_type.as_str(),
|
||||
)
|
||||
});
|
||||
(candidate.key_id.clone(), exhausted)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn read_provider_key_rpm_reset_at_map(
|
||||
state: &(impl SchedulerRuntimeState + ?Sized),
|
||||
candidates: &[SchedulerMinimalCandidateSelectionCandidate],
|
||||
|
||||
@@ -1140,6 +1140,253 @@ async fn exposes_runtime_skipped_candidates_with_skip_reasons() {
|
||||
assert_eq!(skipped[0].skip_reason, "key_circuit_open");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn skips_codex_candidate_when_account_quota_is_exhausted_and_pool_flag_enabled() {
|
||||
let mut first = sample_row();
|
||||
first.provider_id = "provider-codex".to_string();
|
||||
first.provider_name = "codex".to_string();
|
||||
first.provider_type = "codex".to_string();
|
||||
first.endpoint_id = "endpoint-codex".to_string();
|
||||
first.endpoint_api_format = "openai:cli".to_string();
|
||||
first.key_id = "key-codex".to_string();
|
||||
first.key_name = "codex-exhausted".to_string();
|
||||
first.key_auth_type = "oauth".to_string();
|
||||
first.key_api_formats = Some(vec!["openai:cli".to_string()]);
|
||||
first.key_global_priority_by_format = Some(serde_json::json!({"openai:cli": 1}));
|
||||
|
||||
let mut second = sample_row();
|
||||
second.provider_id = "provider-openai".to_string();
|
||||
second.provider_name = "openai".to_string();
|
||||
second.endpoint_id = "endpoint-openai".to_string();
|
||||
second.endpoint_api_format = "openai:cli".to_string();
|
||||
second.key_id = "key-openai".to_string();
|
||||
second.key_name = "fallback".to_string();
|
||||
second.key_api_formats = Some(vec!["openai:cli".to_string()]);
|
||||
second.key_global_priority_by_format = Some(serde_json::json!({"openai:cli": 2}));
|
||||
|
||||
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
first, second,
|
||||
]));
|
||||
let mut codex_provider = sample_provider("provider-codex", None);
|
||||
codex_provider.provider_type = "codex".to_string();
|
||||
codex_provider.config = Some(serde_json::json!({
|
||||
"pool_advanced": {
|
||||
"skip_exhausted_accounts": true
|
||||
}
|
||||
}));
|
||||
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![codex_provider, sample_provider("provider-openai", None)],
|
||||
Vec::new(),
|
||||
vec![
|
||||
{
|
||||
let mut key = sample_key("key-codex", "provider-codex", Some(10));
|
||||
key.auth_type = "oauth".to_string();
|
||||
key.upstream_metadata = Some(serde_json::json!({
|
||||
"codex": {
|
||||
"secondary_used_percent": 100.0
|
||||
}
|
||||
}));
|
||||
key
|
||||
},
|
||||
sample_key("key-openai", "provider-openai", Some(10)),
|
||||
],
|
||||
));
|
||||
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
|
||||
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![]));
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_candidate_selection_provider_catalog_quota_and_request_candidates_for_tests(
|
||||
candidates,
|
||||
provider_catalog,
|
||||
quotas,
|
||||
request_candidates,
|
||||
),
|
||||
);
|
||||
|
||||
let (selected, skipped) = collect_selectable_candidates_with_skip_reasons(
|
||||
state.data.as_ref(),
|
||||
&state,
|
||||
"openai:cli",
|
||||
"gpt-4.1",
|
||||
false,
|
||||
None,
|
||||
100,
|
||||
)
|
||||
.await
|
||||
.expect("selection should succeed");
|
||||
|
||||
assert_eq!(selected.len(), 1);
|
||||
assert_eq!(selected[0].provider_id, "provider-openai");
|
||||
assert_eq!(skipped.len(), 1);
|
||||
assert_eq!(skipped[0].candidate.provider_id, "provider-codex");
|
||||
assert_eq!(skipped[0].skip_reason, "account_quota_exhausted");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn keeps_codex_candidate_selectable_when_exhausted_account_flag_is_disabled() {
|
||||
let mut first = sample_row();
|
||||
first.provider_id = "provider-codex".to_string();
|
||||
first.provider_name = "codex".to_string();
|
||||
first.provider_type = "codex".to_string();
|
||||
first.endpoint_id = "endpoint-codex".to_string();
|
||||
first.endpoint_api_format = "openai:cli".to_string();
|
||||
first.key_id = "key-codex".to_string();
|
||||
first.key_name = "codex-exhausted".to_string();
|
||||
first.key_auth_type = "oauth".to_string();
|
||||
first.key_api_formats = Some(vec!["openai:cli".to_string()]);
|
||||
first.key_global_priority_by_format = Some(serde_json::json!({"openai:cli": 1}));
|
||||
|
||||
let mut second = sample_row();
|
||||
second.provider_id = "provider-openai".to_string();
|
||||
second.provider_name = "openai".to_string();
|
||||
second.endpoint_id = "endpoint-openai".to_string();
|
||||
second.endpoint_api_format = "openai:cli".to_string();
|
||||
second.key_id = "key-openai".to_string();
|
||||
second.key_name = "fallback".to_string();
|
||||
second.key_api_formats = Some(vec!["openai:cli".to_string()]);
|
||||
second.key_global_priority_by_format = Some(serde_json::json!({"openai:cli": 2}));
|
||||
|
||||
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
first, second,
|
||||
]));
|
||||
let mut codex_provider = sample_provider("provider-codex", None);
|
||||
codex_provider.provider_type = "codex".to_string();
|
||||
codex_provider.config = Some(serde_json::json!({
|
||||
"pool_advanced": {}
|
||||
}));
|
||||
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![codex_provider, sample_provider("provider-openai", None)],
|
||||
Vec::new(),
|
||||
vec![
|
||||
{
|
||||
let mut key = sample_key("key-codex", "provider-codex", Some(10));
|
||||
key.auth_type = "oauth".to_string();
|
||||
key.upstream_metadata = Some(serde_json::json!({
|
||||
"codex": {
|
||||
"secondary_used_percent": 100.0
|
||||
}
|
||||
}));
|
||||
key
|
||||
},
|
||||
sample_key("key-openai", "provider-openai", Some(10)),
|
||||
],
|
||||
));
|
||||
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
|
||||
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![]));
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_candidate_selection_provider_catalog_quota_and_request_candidates_for_tests(
|
||||
candidates,
|
||||
provider_catalog,
|
||||
quotas,
|
||||
request_candidates,
|
||||
),
|
||||
);
|
||||
|
||||
let (selected, skipped) = collect_selectable_candidates_with_skip_reasons(
|
||||
state.data.as_ref(),
|
||||
&state,
|
||||
"openai:cli",
|
||||
"gpt-4.1",
|
||||
false,
|
||||
None,
|
||||
100,
|
||||
)
|
||||
.await
|
||||
.expect("selection should succeed");
|
||||
|
||||
assert_eq!(selected.len(), 2);
|
||||
assert!(selected
|
||||
.iter()
|
||||
.any(|candidate| candidate.provider_id == "provider-codex"));
|
||||
assert!(skipped.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn skips_kiro_candidate_when_account_quota_is_exhausted_and_pool_flag_enabled() {
|
||||
let mut first = sample_row();
|
||||
first.provider_id = "provider-kiro".to_string();
|
||||
first.provider_name = "kiro".to_string();
|
||||
first.provider_type = "kiro".to_string();
|
||||
first.endpoint_id = "endpoint-kiro".to_string();
|
||||
first.endpoint_api_format = "claude:cli".to_string();
|
||||
first.key_id = "key-kiro".to_string();
|
||||
first.key_name = "kiro-exhausted".to_string();
|
||||
first.key_auth_type = "oauth".to_string();
|
||||
first.key_api_formats = Some(vec!["claude:cli".to_string()]);
|
||||
first.key_global_priority_by_format = Some(serde_json::json!({"claude:cli": 1}));
|
||||
|
||||
let mut second = sample_row();
|
||||
second.provider_id = "provider-openai".to_string();
|
||||
second.provider_name = "openai".to_string();
|
||||
second.endpoint_id = "endpoint-openai".to_string();
|
||||
second.endpoint_api_format = "claude:cli".to_string();
|
||||
second.key_id = "key-openai".to_string();
|
||||
second.key_name = "fallback".to_string();
|
||||
second.key_api_formats = Some(vec!["claude:cli".to_string()]);
|
||||
second.key_global_priority_by_format = Some(serde_json::json!({"claude:cli": 2}));
|
||||
|
||||
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
first, second,
|
||||
]));
|
||||
let mut kiro_provider = sample_provider("provider-kiro", None);
|
||||
kiro_provider.provider_type = "kiro".to_string();
|
||||
kiro_provider.config = Some(serde_json::json!({
|
||||
"pool_advanced": {
|
||||
"skip_exhausted_accounts": true
|
||||
}
|
||||
}));
|
||||
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![kiro_provider, sample_provider("provider-openai", None)],
|
||||
Vec::new(),
|
||||
vec![
|
||||
{
|
||||
let mut key = sample_key("key-kiro", "provider-kiro", Some(10));
|
||||
key.auth_type = "oauth".to_string();
|
||||
key.upstream_metadata = Some(serde_json::json!({
|
||||
"kiro": {
|
||||
"remaining": 0
|
||||
}
|
||||
}));
|
||||
key
|
||||
},
|
||||
sample_key("key-openai", "provider-openai", Some(10)),
|
||||
],
|
||||
));
|
||||
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
|
||||
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![]));
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_candidate_selection_provider_catalog_quota_and_request_candidates_for_tests(
|
||||
candidates,
|
||||
provider_catalog,
|
||||
quotas,
|
||||
request_candidates,
|
||||
),
|
||||
);
|
||||
|
||||
let (selected, skipped) = collect_selectable_candidates_with_skip_reasons(
|
||||
state.data.as_ref(),
|
||||
&state,
|
||||
"claude:cli",
|
||||
"gpt-4.1",
|
||||
false,
|
||||
None,
|
||||
100,
|
||||
)
|
||||
.await
|
||||
.expect("selection should succeed");
|
||||
|
||||
assert_eq!(selected.len(), 1);
|
||||
assert_eq!(selected[0].provider_id, "provider-openai");
|
||||
assert_eq!(skipped.len(), 1);
|
||||
assert_eq!(skipped[0].candidate.provider_id, "provider-kiro");
|
||||
assert_eq!(skipped[0].skip_reason, "account_quota_exhausted");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn same_priority_candidates_prefer_healthier_provider_key_before_id_order() {
|
||||
let mut first = sample_row();
|
||||
|
||||
Reference in New Issue
Block a user