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: add provider-key concurrent limit (#352)

Browse Source 
* feat: add provider-key concurrent limit

* Fix provider key concurrent limit checks

---------

Co-authored-by: fawney19 <elky0401@gmail.com>
This commit is contained in:
Kayphoon
2026-05-03 01:55:23 +08:00
committed by GitHub
parent 11c5884d4f
commit fe27fb17fb
23 changed files with 1827 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
apps/aether-gateway/src/handlers/admin/provider/shared/payloads.rs
View File

@@ -25,6 +25,8 @@ pub(crate) struct AdminProviderKeyCreateRequest {
#[serde(default)]
pub(crate) rpm_limit: Option<u32>,
#[serde(default)]
pub(crate) concurrent_limit: Option<i32>,
#[serde(default)]
pub(crate) allowed_models: Option<Vec<String>>,
#[serde(default)]
pub(crate) capabilities: Option<serde_json::Value>,
@@ -69,6 +71,8 @@ pub(crate) struct AdminProviderKeyUpdateRequest {
#[serde(default)]
pub(crate) rpm_limit: Option<u32>,
#[serde(default)]
pub(crate) concurrent_limit: Option<i32>,
#[serde(default)]
pub(crate) allowed_models: Option<Vec<String>>,
#[serde(default)]
pub(crate) capabilities: Option<serde_json::Value>,

2
apps/aether-gateway/src/handlers/admin/provider/write/keys/create.rs
View File

@@ -9,6 +9,7 @@ use crate::handlers::admin::shared::{
decrypt_catalog_secret_with_fallbacks, encrypt_catalog_secret_with_fallbacks,
normalize_json_object, normalize_string_list, parse_catalog_auth_config_json,
};
use crate::handlers::shared::normalize_optional_api_key_concurrent_limit;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
@@ -178,6 +179,7 @@ pub(crate) async fn build_admin_create_provider_key_record(
.filter(|value| !value.is_empty());
key.internal_priority = payload.internal_priority.unwrap_or(50);
key.rpm_limit = payload.rpm_limit;
key.concurrent_limit = normalize_optional_api_key_concurrent_limit(payload.concurrent_limit)?;
key.cache_ttl_minutes = payload.cache_ttl_minutes.unwrap_or(5);
key.max_probe_interval_minutes = payload.max_probe_interval_minutes.unwrap_or(32);
key.request_count = Some(0);

5
apps/aether-gateway/src/handlers/admin/provider/write/keys/update.rs
View File

@@ -9,6 +9,7 @@ use crate::handlers::admin::shared::{
decrypt_catalog_secret_with_fallbacks, encrypt_catalog_secret_with_fallbacks, json_string_list,
normalize_json_object, normalize_string_list, parse_catalog_auth_config_json,
};
use crate::handlers::shared::normalize_optional_api_key_concurrent_limit;
use crate::provider_key_auth::provider_key_is_oauth_managed;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogKey, StoredProviderCatalogProvider,
@@ -277,6 +278,10 @@ pub(crate) async fn build_admin_update_provider_key_record(
updated.learned_rpm_limit = None;
}
}
if fields.contains("concurrent_limit") {
updated.concurrent_limit =
normalize_optional_api_key_concurrent_limit(payload.concurrent_limit)?;
}
if fields.contains("allowed_models") {
updated.allowed_models =
normalize_string_list(payload.allowed_models).map(|value| json!(value));

1
apps/aether-gateway/src/handlers/shared/catalog.rs
View File

@@ -1351,6 +1351,7 @@ pub(crate) fn build_admin_provider_key_response(
json!(key.global_priority_by_format),
);
payload.insert("rpm_limit".to_string(), json!(key.rpm_limit));
payload.insert("concurrent_limit".to_string(), json!(key.concurrent_limit));
payload.insert(
"allowed_models".to_string(),
serde_json::Value::Array(

24
apps/aether-gateway/src/scheduler/candidate/affinity.rs
View File

@@ -1,5 +1,6 @@
use aether_scheduler_core::{
build_scheduler_affinity_cache_key_for_api_key_id, SchedulerAffinityTarget,
build_scheduler_affinity_cache_key_for_api_key_id, candidate_affinity_hash, candidate_key,
matches_affinity_target, SchedulerAffinityTarget,
};
use crate::data::auth::GatewayAuthApiKeySnapshot;
@@ -21,6 +22,20 @@ pub(super) fn build_scheduler_affinity_cache_key(
build_scheduler_affinity_cache_key_for_api_key_id(api_key_id, api_format, global_model_name)
}
pub(super) fn scheduler_candidate_affinity_hash(
affinity_key: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> u64 {
candidate_affinity_hash(affinity_key, candidate)
}
pub(super) fn scheduler_candidate_matches_affinity_target(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
target: &SchedulerAffinityTarget,
) -> bool {
matches_affinity_target(candidate, target)
}
#[cfg_attr(not(test), allow(dead_code))]
pub(super) fn remember_scheduler_affinity(
affinity_cache_key: Option<&str>,
@@ -30,13 +45,14 @@ pub(super) fn remember_scheduler_affinity(
let Some(cache_key) = affinity_cache_key else {
return;
};
let (provider_id, endpoint_id, key_id) = candidate_key(candidate);
state.remember_scheduler_affinity_target(
cache_key,
SchedulerAffinityTarget {
provider_id: candidate.provider_id.clone(),
endpoint_id: candidate.endpoint_id.clone(),
key_id: candidate.key_id.clone(),
provider_id,
endpoint_id,
key_id,
},
SCHEDULER_AFFINITY_TTL,
SCHEDULER_AFFINITY_MAX_ENTRIES,

21
apps/aether-gateway/src/scheduler/candidate/ranking.rs
View File

@@ -1,12 +1,15 @@
use aether_scheduler_core::{
apply_scheduler_candidate_ranking, candidate_affinity_hash,
effective_provider_key_health_score, matches_affinity_target, provider_key_health_bucket,
requested_capability_priority_for_candidate, SchedulerAffinityTarget,
SchedulerRankableCandidate, SchedulerRankingContext, SchedulerRankingMode,
apply_scheduler_candidate_ranking, effective_provider_key_health_score,
provider_key_health_bucket, requested_capability_priority_for_candidate,
SchedulerAffinityTarget, SchedulerRankableCandidate, SchedulerRankingContext,
SchedulerRankingMode,
};
use crate::scheduler::config::{SchedulerOrderingConfig, SchedulerSchedulingMode};
use super::affinity::{
scheduler_candidate_affinity_hash, scheduler_candidate_matches_affinity_target,
};
use super::runtime::CandidateRuntimeSelectionSnapshot;
use super::SchedulerMinimalCandidateSelectionCandidate;
@@ -31,12 +34,12 @@ pub(super) fn rank_scheduler_candidates(
required_capabilities,
candidate,
))
.with_cached_affinity_match(
cached_affinity_target
.is_some_and(|target| matches_affinity_target(candidate, target)),
)
.with_cached_affinity_match(cached_affinity_target.is_some_and(|target| {
scheduler_candidate_matches_affinity_target(candidate, target)
}))
.with_affinity_hash(
priority_affinity_key.map(|key| candidate_affinity_hash(key, candidate)),
priority_affinity_key
.map(|key| scheduler_candidate_affinity_hash(key, candidate)),
)
.with_health(
provider_key.and_then(|key| {

389
apps/aether-gateway/src/scheduler/candidate/tests/selection.rs
View File

@@ -5,10 +5,13 @@ use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelect
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::quota::InMemoryProviderQuotaRepository;
use aether_data_contracts::repository::candidate_selection::StoredProviderModelMapping;
use aether_data_contracts::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};
use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, StoredRequestCandidate,
};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use aether_data_contracts::repository::quota::StoredProviderQuotaSnapshot;
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use serde_json::json;
@@ -99,6 +102,102 @@ async fn collect_selectable_candidates_with_skip_reasons(
.await
}
fn provider_key_concurrency_row(
provider_id: &str,
endpoint_id: &str,
key_id: &str,
key_name: &str,
provider_priority: i32,
key_priority: i32,
) -> StoredMinimalCandidateSelectionRow {
let mut row = sample_row();
row.provider_id = provider_id.to_string();
row.provider_name = provider_id.to_string();
row.endpoint_id = endpoint_id.to_string();
row.key_id = key_id.to_string();
row.key_name = key_name.to_string();
row.provider_priority = provider_priority;
row.key_internal_priority = key_priority;
row.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": key_priority}));
row
}
fn provider_key_with_concurrent_limit(
key_id: &str,
provider_id: &str,
concurrent_limit: Option<i32>,
) -> StoredProviderCatalogKey {
let mut key = sample_key(key_id, provider_id, Some(10));
key.concurrent_limit = concurrent_limit;
key
}
fn active_provider_key_candidate(
candidate_id: &str,
request_id: &str,
provider_id: &str,
endpoint_id: &str,
key_id: &str,
status: RequestCandidateStatus,
) -> StoredRequestCandidate {
StoredRequestCandidate::new(
candidate_id.to_string(),
request_id.to_string(),
None,
None,
None,
None,
0,
0,
Some(provider_id.to_string()),
Some(endpoint_id.to_string()),
Some(key_id.to_string()),
status,
None,
false,
None,
None,
None,
None,
None,
None,
None,
95_000,
Some(95_000),
None,
)
.expect("candidate should build")
}
fn provider_key_concurrency_state(
rows: Vec<StoredMinimalCandidateSelectionRow>,
keys: Vec<StoredProviderCatalogKey>,
request_candidates: Vec<StoredRequestCandidate>,
) -> AppState {
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(rows));
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider("test-provider-a", None),
sample_provider("test-provider-b", None),
],
Vec::new(),
keys,
));
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(request_candidates));
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,
),
)
}
#[test]
fn skips_only_exhausted_monthly_quota_provider() {
let inactive = StoredProviderQuotaSnapshot::new(
@@ -860,6 +959,294 @@ async fn selects_next_candidate_when_first_provider_concurrent_limit_is_reached(
assert_eq!(selected.key_id, "key-b");
}
#[tokio::test]
async fn provider_key_concurrency_selects_next_key_when_first_provider_key_concurrent_limit_is_reached(
) {
let state = provider_key_concurrency_state(
vec![
provider_key_concurrency_row(
"test-provider-a",
"endpoint-a",
"provider-key-a",
"alpha",
0,
0,
),
provider_key_concurrency_row(
"test-provider-a",
"endpoint-a",
"provider-key-b",
"beta",
0,
1,
),
],
vec![
provider_key_with_concurrent_limit("provider-key-a", "test-provider-a", Some(1)),
provider_key_with_concurrent_limit("provider-key-b", "test-provider-a", Some(1)),
],
vec![active_provider_key_candidate(
"cand-provider-key-a",
"req-provider-key-a",
"test-provider-a",
"endpoint-a",
"provider-key-a",
RequestCandidateStatus::Streaming,
)],
);
let selected = select_candidate(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
None,
100,
)
.await
.expect("selection should succeed")
.expect("candidate should exist");
assert_eq!(selected.provider_id, "test-provider-a");
assert_eq!(selected.key_id, "provider-key-b");
}
#[tokio::test]
async fn provider_key_concurrency_selects_next_provider_when_all_provider_keys_concurrent_limit_reached(
) {
let state = provider_key_concurrency_state(
vec![
provider_key_concurrency_row(
"test-provider-a",
"endpoint-a",
"provider-key-a",
"alpha",
0,
0,
),
provider_key_concurrency_row(
"test-provider-a",
"endpoint-a",
"provider-key-b",
"beta",
0,
1,
),
provider_key_concurrency_row(
"test-provider-b",
"endpoint-b",
"provider-key-c",
"gamma",
1,
0,
),
],
vec![
provider_key_with_concurrent_limit("provider-key-a", "test-provider-a", Some(1)),
provider_key_with_concurrent_limit("provider-key-b", "test-provider-a", Some(1)),
provider_key_with_concurrent_limit("provider-key-c", "test-provider-b", Some(1)),
],
vec![
active_provider_key_candidate(
"cand-provider-key-a",
"req-provider-key-a",
"test-provider-a",
"endpoint-a",
"provider-key-a",
RequestCandidateStatus::Pending,
),
active_provider_key_candidate(
"cand-provider-key-b",
"req-provider-key-b",
"test-provider-a",
"endpoint-a",
"provider-key-b",
RequestCandidateStatus::Streaming,
),
],
);
let selected = select_candidate(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
None,
100,
)
.await
.expect("selection should succeed")
.expect("candidate should exist");
assert_eq!(selected.provider_id, "test-provider-b");
assert_eq!(selected.key_id, "provider-key-c");
}
#[tokio::test]
async fn provider_key_concurrency_returns_none_when_all_provider_keys_concurrent_limit_reached() {
let state = provider_key_concurrency_state(
vec![
provider_key_concurrency_row(
"test-provider-a",
"endpoint-a",
"provider-key-a",
"alpha",
0,
0,
),
provider_key_concurrency_row(
"test-provider-a",
"endpoint-a",
"provider-key-b",
"beta",
0,
1,
),
provider_key_concurrency_row(
"test-provider-b",
"endpoint-b",
"provider-key-c",
"gamma",
1,
0,
),
],
vec![
provider_key_with_concurrent_limit("provider-key-a", "test-provider-a", Some(1)),
provider_key_with_concurrent_limit("provider-key-b", "test-provider-a", Some(1)),
provider_key_with_concurrent_limit("provider-key-c", "test-provider-b", Some(1)),
],
vec![
active_provider_key_candidate(
"cand-provider-key-a",
"req-provider-key-a",
"test-provider-a",
"endpoint-a",
"provider-key-a",
RequestCandidateStatus::Pending,
),
active_provider_key_candidate(
"cand-provider-key-b",
"req-provider-key-b",
"test-provider-a",
"endpoint-a",
"provider-key-b",
RequestCandidateStatus::Streaming,
),
active_provider_key_candidate(
"cand-provider-key-c",
"req-provider-key-c",
"test-provider-b",
"endpoint-b",
"provider-key-c",
RequestCandidateStatus::Streaming,
),
],
);
let selected = select_candidate(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
None,
100,
)
.await
.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,
None,
100,
)
.await
.expect("selection should succeed");
assert!(selected_candidates.is_empty());
assert_eq!(skipped_candidates.len(), 3);
assert!(skipped_candidates
.iter()
.all(|skipped| { skipped.skip_reason == "provider_key_concurrency_limit_reached" }));
assert!(!is_exact_all_skipped_by_auth_limit(
&selected_candidates,
&skipped_candidates,
));
}
#[tokio::test]
async fn provider_key_concurrency_collects_exact_skip_reason_for_saturated_provider_keys() {
let state = provider_key_concurrency_state(
vec![
provider_key_concurrency_row(
"test-provider-a",
"endpoint-a",
"provider-key-a",
"alpha",
0,
0,
),
provider_key_concurrency_row(
"test-provider-a",
"endpoint-a",
"provider-key-b",
"beta",
0,
1,
),
],
vec![
provider_key_with_concurrent_limit("provider-key-a", "test-provider-a", Some(1)),
provider_key_with_concurrent_limit("provider-key-b", "test-provider-a", Some(1)),
],
vec![active_provider_key_candidate(
"cand-provider-key-a",
"req-provider-key-a",
"test-provider-a",
"endpoint-a",
"provider-key-a",
RequestCandidateStatus::Pending,
)],
);
let (selected_candidates, skipped_candidates) =
collect_selectable_candidates_with_skip_reasons(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
None,
100,
)
.await
.expect("selection should succeed");
assert_eq!(selected_candidates.len(), 1);
assert_eq!(selected_candidates[0].provider_id, "test-provider-a");
assert_eq!(selected_candidates[0].key_id, "provider-key-b");
assert_eq!(skipped_candidates.len(), 1);
assert_eq!(
skipped_candidates[0].candidate.provider_id,
"test-provider-a"
);
assert_eq!(skipped_candidates[0].candidate.key_id, "provider-key-a");
assert_eq!(
skipped_candidates[0].skip_reason,
"provider_key_concurrency_limit_reached",
);
}
#[tokio::test]
async fn returns_none_when_auth_api_key_concurrent_limit_is_reached() {
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![

12
apps/aether-gateway/src/scheduler/candidate/tests/support.rs
View File

@@ -90,7 +90,17 @@ pub(super) fn sample_key(
true,
)
.expect("key should build")
.with_rate_limit_fields(rpm_limit, None, None, None, None, None, Some(20), Some(20))
.with_rate_limit_fields(
rpm_limit,
None,
None,
None,
None,
None,
None,
Some(20),
Some(20),
)
}
pub(super) fn sample_auth_snapshot(api_key_id: &str) -> GatewayAuthApiKeySnapshot {

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

@@ -282,6 +282,232 @@ async fn gateway_creates_admin_provider_key_locally_with_trusted_admin_principal
upstream_handle.abort();
}
#[tokio::test]
async fn provider_key_concurrent_limit_create_and_list_responses() {
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
vec![],
vec![],
));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.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 create_with_limit_response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/endpoints/providers/provider-openai/keys"
))
.header(crate::constants::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")
.json(&json!({
"api_formats": ["openai:chat"],
"api_key": "sk-created-openai-concurrent",
"name": "created key with concurrency",
"rpm_limit": 60,
"concurrent_limit": 3
}))
.send()
.await
.expect("request should succeed");
assert_eq!(create_with_limit_response.status(), StatusCode::OK);
let create_payload: serde_json::Value = create_with_limit_response
.json()
.await
.expect("json body should parse");
let create_with_limit_id = create_payload["id"]
.as_str()
.expect("created key id should be returned")
.to_string();
assert_eq!(create_payload["rpm_limit"], 60);
assert_eq!(create_payload["concurrent_limit"], 3);
let create_null_response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/endpoints/providers/provider-openai/keys"
))
.header(crate::constants::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")
.json(&json!({
"api_formats": ["openai:chat"],
"api_key": "sk-created-openai-concurrent-null",
"name": "created key with null concurrency",
"concurrent_limit": null
}))
.send()
.await
.expect("request should succeed");
assert_eq!(create_null_response.status(), StatusCode::OK);
let null_payload: serde_json::Value = create_null_response
.json()
.await
.expect("json body should parse");
let create_null_id = null_payload["id"]
.as_str()
.expect("created null-limit key id should be returned")
.to_string();
assert_eq!(null_payload["concurrent_limit"], serde_json::Value::Null);
let create_negative_response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/endpoints/providers/provider-openai/keys"
))
.header(crate::constants::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")
.json(&json!({
"api_formats": ["openai:chat"],
"api_key": "sk-created-openai-concurrent-negative",
"name": "created key with negative concurrency",
"concurrent_limit": -1
}))
.send()
.await
.expect("request should succeed");
assert_eq!(create_negative_response.status(), StatusCode::BAD_REQUEST);
let negative_payload: serde_json::Value = create_negative_response
.json()
.await
.expect("json body should parse");
assert!(negative_payload["detail"]
.as_str()
.expect("detail should be string")
.contains("concurrent_limit"));
let list_response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/endpoints/providers/provider-openai/keys?skip=0&limit=50"
))
.header(crate::constants::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!(list_response.status(), StatusCode::OK);
let list_payload: serde_json::Value =
list_response.json().await.expect("json body should parse");
let items = list_payload.as_array().expect("payload should be an array");
assert_eq!(items.len(), 2);
let with_limit = items
.iter()
.find(|item| item["name"].as_str() == Some("created key with concurrency"))
.expect("created key with concurrency should be listed");
let with_null = items
.iter()
.find(|item| item["name"].as_str() == Some("created key with null concurrency"))
.expect("created key with null concurrency should be listed");
assert_eq!(with_limit["concurrent_limit"], 3);
assert_eq!(with_null["concurrent_limit"], serde_json::Value::Null);
let read_back = provider_catalog_repository
.list_keys_by_ids(&[create_with_limit_id, create_null_id])
.await
.expect("created keys should read by id");
assert_eq!(read_back.len(), 2);
assert!(read_back.iter().any(|key| {
key.name == "created key with concurrency" && key.concurrent_limit == Some(3)
}));
assert!(read_back.iter().any(|key| {
key.name == "created key with null concurrency" && key.concurrent_limit.is_none()
}));
let keys = provider_catalog_repository
.list_keys_by_provider_ids(&["provider-openai".to_string()])
.await
.expect("keys should read");
assert_eq!(keys.len(), 2);
assert!(keys
.iter()
.any(|key| key.name == "created key with concurrency" && key.concurrent_limit == Some(3)));
assert!(keys.iter().any(
|key| key.name == "created key with null concurrency" && key.concurrent_limit.is_none()
));
gateway_handle.abort();
}
#[tokio::test]
async fn provider_key_concurrent_limit_reads_existing_list_response() {
let mut key_a = sample_key(
"provider-key-a",
"test-provider-a",
"openai:chat",
"sk-provider-key-a",
);
key_a.concurrent_limit = Some(1);
let mut key_b = sample_key(
"provider-key-b",
"test-provider-a",
"openai:chat",
"sk-provider-key-b",
);
key_b.concurrent_limit = None;
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("test-provider-a", "openai", 10)],
vec![],
vec![key_a, key_b],
));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/endpoints/providers/test-provider-a/keys?skip=0&limit=50"
))
.header(crate::constants::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");
let items = payload.as_array().expect("payload should be an array");
assert_eq!(items.len(), 2);
let provider_key_a = items
.iter()
.find(|item| item["id"].as_str() == Some("provider-key-a"))
.expect("provider-key-a should be listed");
let provider_key_b = items
.iter()
.find(|item| item["id"].as_str() == Some("provider-key-b"))
.expect("provider-key-b should be listed");
assert_eq!(provider_key_a["concurrent_limit"], 1);
assert_eq!(provider_key_b["concurrent_limit"], serde_json::Value::Null);
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_fetches_allowed_models_immediately_when_creating_key_with_auto_fetch() {
let execution_runtime_hits = Arc::new(Mutex::new(0usize));
@@ -885,6 +1111,132 @@ async fn gateway_updates_admin_provider_key_locally_with_trusted_admin_principal
upstream_handle.abort();
}
#[tokio::test]
async fn provider_key_concurrent_limit_update_presence_semantics() {
let mut key = sample_key(
"key-openai-a",
"provider-openai",
"openai:chat",
"sk-test-a",
);
key.concurrent_limit = Some(4);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
vec![],
vec![key],
));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.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 omitted_response = reqwest::Client::new()
.put(format!(
"{gateway_url}/api/admin/endpoints/keys/key-openai-a"
))
.header(crate::constants::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")
.json(&json!({
"name": "renamed without concurrency"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(omitted_response.status(), StatusCode::OK);
let omitted_payload: serde_json::Value = omitted_response
.json()
.await
.expect("json body should parse");
assert_eq!(omitted_payload["name"], "renamed without concurrency");
assert_eq!(omitted_payload["concurrent_limit"], 4);
let set_response = reqwest::Client::new()
.put(format!(
"{gateway_url}/api/admin/endpoints/keys/key-openai-a"
))
.header(crate::constants::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")
.json(&json!({
"concurrent_limit": 7
}))
.send()
.await
.expect("request should succeed");
assert_eq!(set_response.status(), StatusCode::OK);
let set_payload: serde_json::Value = set_response.json().await.expect("json body should parse");
assert_eq!(set_payload["concurrent_limit"], 7);
let clear_response = reqwest::Client::new()
.put(format!(
"{gateway_url}/api/admin/endpoints/keys/key-openai-a"
))
.header(crate::constants::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")
.json(&json!({
"concurrent_limit": null
}))
.send()
.await
.expect("request should succeed");
assert_eq!(clear_response.status(), StatusCode::OK);
let clear_payload: serde_json::Value =
clear_response.json().await.expect("json body should parse");
assert_eq!(clear_payload["concurrent_limit"], serde_json::Value::Null);
let negative_response = reqwest::Client::new()
.put(format!(
"{gateway_url}/api/admin/endpoints/keys/key-openai-a"
))
.header(crate::constants::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")
.json(&json!({
"concurrent_limit": -1
}))
.send()
.await
.expect("request should succeed");
assert_eq!(negative_response.status(), StatusCode::BAD_REQUEST);
let negative_payload: serde_json::Value = negative_response
.json()
.await
.expect("json body should parse");
assert!(negative_payload["detail"]
.as_str()
.expect("detail should be string")
.contains("concurrent_limit"));
let reloaded = provider_catalog_repository
.list_keys_by_ids(&["key-openai-a".to_string()])
.await
.expect("keys should read");
assert_eq!(reloaded.len(), 1);
assert_eq!(reloaded[0].name, "renamed without concurrency");
assert_eq!(reloaded[0].concurrent_limit, None);
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_clears_allowed_models_when_disabling_auto_fetch_on_provider_key_update() {
let upstream_hits = Arc::new(Mutex::new(0usize));

2
apps/aether-gateway/src/tests/control/admin/health_access.rs
View File

@@ -334,7 +334,7 @@ async fn gateway_handles_admin_key_health_locally_with_trusted_admin_principal()
)],
vec![
sample_key("key-openai", "provider-openai", "openai:chat", "sk-test")
.with_rate_limit_fields(None, None, None, None, None, None, Some(10), Some(7))
.with_rate_limit_fields(None, None, None, None, None, None, None, Some(10), Some(7))
.with_usage_fields(Some(3), Some(2100))
.with_health_fields(
Some(json!({"openai:chat": {

4
apps/aether-gateway/src/tests/control/admin/system.rs
View File

@@ -1424,7 +1424,7 @@ async fn gateway_handles_admin_key_rpm_locally_with_trusted_admin_principal() {
)],
vec![
sample_key("key-openai", "provider-openai", "openai:chat", "sk-test")
.with_rate_limit_fields(Some(60), None, None, None, None, None, None, None),
.with_rate_limit_fields(Some(60), None, None, None, None, None, None, None, None),
],
));
let now_unix_secs = SystemTime::now()
@@ -1511,7 +1511,7 @@ async fn gateway_resets_admin_key_rpm_locally_with_trusted_admin_principal() {
)],
vec![
sample_key("key-openai", "provider-openai", "openai:chat", "sk-test")
.with_rate_limit_fields(Some(60), None, None, None, None, None, None, None),
.with_rate_limit_fields(Some(60), None, None, None, None, None, None, None, None),
],
));
let now_unix_secs = SystemTime::now()