mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
Fix/api key concurrency runtime miss (#309)
* test(cli): 覆盖 API key 并发等待与超时路径 * feat(scheduler): API key 并发饱和时等待可用槽位 * fix(proxy): 区分 API key 并发受限与真正的 runtime miss * fix(outcome): runtime miss 仅归因真实执行候选 * feat(api-keys): 统一 concurrent_limit 默认值与校验辅助 * feat(admin): 独立 Key 接口支持 concurrent_limit * feat(admin): 用户 API Key 路由支持 concurrent_limit * feat(public): 自助 API Key 路由支持 concurrent_limit * feat(import): 导入与存储层持久化 concurrent_limit * feat(frontend): 同步 API Key concurrent_limit 类型定义 * feat(frontend): 独立 Key 表单支持 concurrent_limit * feat(frontend): 管理员用户 API Key 表单支持 concurrent_limit * feat(frontend): 自助 API Key 页面支持 concurrent_limit * chore(fmt): 统一 runtime 归因相关 Rust 格式 * chore(fmt): 统一 admin API key 路由 Rust 格式 * chore(fmt): 统一 public 路由与相关测试 Rust 格式 * fix(test): 对齐 no-execution usage 归因断言 * test(middleware): 固定 access log tracing 用例线程模型 * fix(frontend): 提取用户 API Key payload 默认并发辅助 * fix(frontend): 保留用户 Key 的 concurrent_limit 默认值 * fix(api-keys): remove hardcoded concurrent limit default --------- Co-authored-by: fawney19 <elky0401@gmail.com>
This commit is contained in:
@@ -69,6 +69,13 @@ pub(crate) async fn list_selectable_candidates(
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) fn is_exact_all_skipped_by_auth_limit(
|
||||
selected: &[SchedulerMinimalCandidateSelectionCandidate],
|
||||
skipped: &[SchedulerSkippedCandidate],
|
||||
) -> bool {
|
||||
selection::is_exact_all_skipped_by_auth_limit(selected, skipped)
|
||||
}
|
||||
|
||||
pub(crate) async fn list_selectable_candidates_with_skip_reasons(
|
||||
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
|
||||
runtime_state: &impl SchedulerRuntimeState,
|
||||
@@ -107,9 +114,33 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
|
||||
Ok(
|
||||
list_selectable_candidates_for_required_capability_without_requested_model_with_auth_limit_signal(
|
||||
selection_row_source,
|
||||
runtime_state,
|
||||
candidate_api_format,
|
||||
required_capability,
|
||||
require_streaming,
|
||||
auth_snapshot,
|
||||
now_unix_secs,
|
||||
)
|
||||
.await?
|
||||
.0,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) async fn list_selectable_candidates_for_required_capability_without_requested_model_with_auth_limit_signal(
|
||||
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
|
||||
runtime_state: &impl SchedulerRuntimeState,
|
||||
candidate_api_format: &str,
|
||||
required_capability: &str,
|
||||
require_streaming: bool,
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<(Vec<SchedulerMinimalCandidateSelectionCandidate>, bool), GatewayError> {
|
||||
let normalized_api_format = normalize_api_format(candidate_api_format);
|
||||
if normalized_api_format.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
return Ok((Vec::new(), false));
|
||||
}
|
||||
|
||||
let capability_mode = required_capability_match_mode(required_capability);
|
||||
@@ -136,9 +167,10 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
|
||||
}
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let required_capabilities = build_required_capabilities_object(required_capability);
|
||||
let mut all_attempts_blocked_by_auth_limit = !model_names.is_empty();
|
||||
|
||||
for global_model_name in model_names {
|
||||
let mut candidates = list_selectable_candidates(
|
||||
let (mut candidates, skipped_candidates) = collect_selectable_candidates_with_skip_reasons(
|
||||
selection_row_source,
|
||||
runtime_state,
|
||||
&normalized_api_format,
|
||||
@@ -149,6 +181,8 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
|
||||
now_unix_secs,
|
||||
)
|
||||
.await?;
|
||||
all_attempts_blocked_by_auth_limit &=
|
||||
is_exact_all_skipped_by_auth_limit(&candidates, &skipped_candidates);
|
||||
match capability_mode {
|
||||
RequiredCapabilityMatchMode::Exclusive => {
|
||||
let filtered = candidates
|
||||
@@ -158,7 +192,7 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
if !filtered.is_empty() {
|
||||
return Ok(filtered);
|
||||
return Ok((filtered, false));
|
||||
}
|
||||
}
|
||||
RequiredCapabilityMatchMode::Compatible => {
|
||||
@@ -168,12 +202,12 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
|
||||
candidates.sort_by_key(|candidate| {
|
||||
!candidate_supports_required_capability(candidate, required_capability)
|
||||
});
|
||||
return Ok(candidates);
|
||||
return Ok((candidates, false));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(Vec::new())
|
||||
Ok((Vec::new(), all_attempts_blocked_by_auth_limit))
|
||||
}
|
||||
|
||||
fn required_capability_match_mode(required_capability: &str) -> RequiredCapabilityMatchMode {
|
||||
|
||||
@@ -31,6 +31,19 @@ pub(crate) struct SchedulerSkippedCandidate {
|
||||
pub(crate) skip_reason: &'static str,
|
||||
}
|
||||
|
||||
pub(super) const API_KEY_CONCURRENCY_LIMIT_SKIP_REASON: &str = "api_key_concurrency_limit_reached";
|
||||
|
||||
pub(super) fn is_exact_all_skipped_by_auth_limit(
|
||||
selected: &[SchedulerMinimalCandidateSelectionCandidate],
|
||||
skipped: &[SchedulerSkippedCandidate],
|
||||
) -> bool {
|
||||
selected.is_empty()
|
||||
&& !skipped.is_empty()
|
||||
&& skipped
|
||||
.iter()
|
||||
.all(|candidate| candidate.skip_reason == API_KEY_CONCURRENCY_LIMIT_SKIP_REASON)
|
||||
}
|
||||
|
||||
pub(super) fn reorder_candidates_by_scheduler_health(
|
||||
candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
|
||||
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
|
||||
@@ -218,7 +231,7 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
|
||||
.into_iter()
|
||||
.map(|candidate| SchedulerSkippedCandidate {
|
||||
candidate,
|
||||
skip_reason: "api_key_concurrency_limit_reached",
|
||||
skip_reason: API_KEY_CONCURRENCY_LIMIT_SKIP_REASON,
|
||||
})
|
||||
.collect(),
|
||||
));
|
||||
|
||||
@@ -1,13 +1,21 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
|
||||
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data::repository::quota::InMemoryProviderQuotaRepository;
|
||||
use aether_data_contracts::repository::candidates::{
|
||||
RequestCandidateStatus, StoredRequestCandidate,
|
||||
};
|
||||
|
||||
use crate::data::GatewayDataState;
|
||||
use crate::AppState;
|
||||
|
||||
use super::super::list_selectable_candidates_for_required_capability_without_requested_model;
|
||||
use super::support::sample_row;
|
||||
use super::super::{
|
||||
list_selectable_candidates_for_required_capability_without_requested_model,
|
||||
list_selectable_candidates_for_required_capability_without_requested_model_with_auth_limit_signal,
|
||||
};
|
||||
use super::support::{sample_auth_snapshot, sample_provider, sample_row};
|
||||
|
||||
#[tokio::test]
|
||||
async fn compatible_required_capability_prefers_matching_keys_without_hard_filtering() {
|
||||
@@ -114,3 +122,78 @@ async fn exclusive_required_capability_keeps_hard_filtering_only_matching_keys()
|
||||
assert_eq!(selection[0].provider_id, "provider-b");
|
||||
assert_eq!(selection[0].key_id, "key-b");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn required_capability_reports_auth_limit_signal_when_every_model_is_blocked_by_api_key_concurrency(
|
||||
) {
|
||||
let mut candidate = sample_row();
|
||||
candidate.key_capabilities = Some(serde_json::json!({"cache_1h": true}));
|
||||
|
||||
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
candidate,
|
||||
]));
|
||||
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-1", None)],
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
));
|
||||
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
|
||||
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
|
||||
StoredRequestCandidate::new(
|
||||
"cand-1".to_string(),
|
||||
"req-1".to_string(),
|
||||
Some("user-1".to_string()),
|
||||
Some("api-key-1".to_string()),
|
||||
None,
|
||||
None,
|
||||
0,
|
||||
0,
|
||||
Some("provider-1".to_string()),
|
||||
Some("endpoint-1".to_string()),
|
||||
Some("key-1".to_string()),
|
||||
RequestCandidateStatus::Pending,
|
||||
None,
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
95_000,
|
||||
Some(95_000),
|
||||
None,
|
||||
)
|
||||
.expect("candidate should build"),
|
||||
]));
|
||||
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 mut auth_snapshot = sample_auth_snapshot("api-key-1");
|
||||
auth_snapshot.api_key_concurrent_limit = Some(1);
|
||||
|
||||
let (selection, auth_limit_blocked) =
|
||||
list_selectable_candidates_for_required_capability_without_requested_model_with_auth_limit_signal(
|
||||
state.data.as_ref(),
|
||||
&state,
|
||||
"openai:chat",
|
||||
"cache_1h",
|
||||
false,
|
||||
Some(&auth_snapshot),
|
||||
100,
|
||||
)
|
||||
.await
|
||||
.expect("selection should succeed");
|
||||
|
||||
assert!(selection.is_empty());
|
||||
assert!(auth_limit_blocked);
|
||||
}
|
||||
|
||||
@@ -23,7 +23,7 @@ use super::super::runtime::should_skip_provider_quota;
|
||||
use super::super::selection::{
|
||||
collect_selectable_candidates as collect_selectable_candidates_impl,
|
||||
collect_selectable_candidates_with_skip_reasons as collect_selectable_candidates_with_skip_reasons_impl,
|
||||
select_minimal_candidate as select_candidate_impl,
|
||||
is_exact_all_skipped_by_auth_limit, select_minimal_candidate as select_candidate_impl,
|
||||
};
|
||||
use super::support::{sample_auth_snapshot, sample_key, sample_provider, sample_row};
|
||||
|
||||
@@ -914,6 +914,23 @@ async fn returns_none_when_auth_api_key_concurrent_limit_is_reached() {
|
||||
.expect("selection should succeed");
|
||||
|
||||
assert!(selected.is_none());
|
||||
|
||||
let (selected_candidates, skipped_candidates) =
|
||||
collect_selectable_candidates_with_skip_reasons(
|
||||
state.data.as_ref(),
|
||||
&state,
|
||||
"openai:chat",
|
||||
"gpt-4.1",
|
||||
false,
|
||||
Some(&auth_snapshot),
|
||||
100,
|
||||
)
|
||||
.await
|
||||
.expect("selection should succeed");
|
||||
assert!(is_exact_all_skipped_by_auth_limit(
|
||||
&selected_candidates,
|
||||
&skipped_candidates,
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -1138,6 +1155,7 @@ async fn exposes_runtime_skipped_candidates_with_skip_reasons() {
|
||||
assert_eq!(skipped.len(), 1);
|
||||
assert_eq!(skipped[0].candidate.provider_id, "provider-a");
|
||||
assert_eq!(skipped[0].skip_reason, "key_circuit_open");
|
||||
assert!(!is_exact_all_skipped_by_auth_limit(&selected, &skipped));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
Reference in New Issue
Block a user