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

Fix/api key concurrency runtime miss (#309)

Browse Source 
* 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:
RWDai
2026-04-17 14:21:43 +08:00
committed by GitHub
parent e5d3722adf
commit b8702ae124
39 changed files with 1739 additions and 128 deletions
Download Patch File Download Diff File Expand all files Collapse all files

95
crates/aether-data/src/repository/auth/memory.rs
View File

@@ -417,7 +417,7 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
api_key_is_locked: false,
api_key_is_standalone: false,
api_key_rate_limit: Some(record.rate_limit),
api_key_concurrent_limit: Some(record.concurrent_limit),
api_key_concurrent_limit: record.concurrent_limit,
api_key_expires_at_unix_secs: record.expires_at_unix_secs,
api_key_allowed_providers: record.allowed_providers.clone(),
api_key_allowed_api_formats: record.allowed_api_formats.clone(),
@@ -445,7 +445,7 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
false,
false,
Some(record.rate_limit),
Some(record.concurrent_limit),
record.concurrent_limit,
record.expires_at_unix_secs.map(|value| value as i64),
record
.allowed_providers
@@ -481,7 +481,7 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
.as_ref()
.map(|value| serde_json::json!(value)),
Some(record.rate_limit),
Some(record.concurrent_limit),
record.concurrent_limit,
record.force_capabilities,
record.is_active,
record.expires_at_unix_secs.map(|value| value as i64),
@@ -537,7 +537,7 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
api_key_is_locked: false,
api_key_is_standalone: true,
api_key_rate_limit: record.rate_limit,
api_key_concurrent_limit: Some(record.concurrent_limit),
api_key_concurrent_limit: record.concurrent_limit,
api_key_expires_at_unix_secs: record.expires_at_unix_secs,
api_key_allowed_providers: record.allowed_providers.clone(),
api_key_allowed_api_formats: record.allowed_api_formats.clone(),
@@ -565,7 +565,7 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
false,
true,
record.rate_limit,
Some(record.concurrent_limit),
record.concurrent_limit,
record.expires_at_unix_secs.map(|value| value as i64),
record
.allowed_providers
@@ -601,7 +601,7 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
.as_ref()
.map(|value| serde_json::json!(value)),
record.rate_limit,
Some(record.concurrent_limit),
record.concurrent_limit,
record.force_capabilities,
record.is_active,
record.expires_at_unix_secs.map(|value| value as i64),
@@ -653,6 +653,14 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
export.rate_limit = Some(rate_limit);
}
}
if let Some(concurrent_limit) = record.concurrent_limit {
if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) {
snapshot.api_key_concurrent_limit = Some(concurrent_limit);
}
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
export.concurrent_limit = Some(concurrent_limit);
}
}
Ok(index.export_by_api_key_id.get(&record.api_key_id).cloned())
}
@@ -686,6 +694,14 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
export.rate_limit = record.rate_limit;
}
}
if record.concurrent_limit_present {
if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) {
snapshot.api_key_concurrent_limit = record.concurrent_limit;
}
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
export.concurrent_limit = record.concurrent_limit;
}
}
if let Some(allowed_providers) = record.allowed_providers {
if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) {
snapshot.api_key_allowed_providers = allowed_providers.clone();
@@ -892,6 +908,7 @@ mod tests {
use crate::repository::auth::{
AuthApiKeyLookupKey, AuthApiKeyReadRepository, AuthApiKeyWriteRepository,
StandaloneApiKeyExportListQuery, StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
UpdateStandaloneApiKeyBasicRecord, UpdateUserApiKeyBasicRecord,
};
fn sample_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
@@ -1074,4 +1091,70 @@ mod tests {
1
);
}
#[tokio::test]
async fn update_user_api_key_basic_updates_concurrent_limit() {
let repository = InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some("hash-1".to_string()),
sample_snapshot("key-1", "user-1"),
)]);
let updated = repository
.update_user_api_key_basic(UpdateUserApiKeyBasicRecord {
user_id: "user-1".to_string(),
api_key_id: "key-1".to_string(),
name: None,
rate_limit: None,
concurrent_limit: Some(11),
})
.await
.expect("update should succeed")
.expect("record should exist");
assert_eq!(updated.concurrent_limit, Some(11));
let snapshot = repository
.find_api_key_snapshot(AuthApiKeyLookupKey::ApiKeyId("key-1"))
.await
.expect("find should succeed")
.expect("snapshot should exist");
assert_eq!(snapshot.api_key_concurrent_limit, Some(11));
}
#[tokio::test]
async fn update_standalone_api_key_basic_updates_concurrent_limit_when_present() {
let mut standalone = sample_snapshot("key-standalone", "admin-1");
standalone.api_key_is_standalone = true;
let repository = InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some("hash-standalone".to_string()),
standalone,
)]);
let updated = repository
.update_standalone_api_key_basic(UpdateStandaloneApiKeyBasicRecord {
api_key_id: "key-standalone".to_string(),
name: None,
rate_limit_present: false,
rate_limit: None,
concurrent_limit_present: true,
concurrent_limit: Some(13),
allowed_providers: None,
allowed_api_formats: None,
allowed_models: None,
expires_at_present: false,
expires_at_unix_secs: None,
auto_delete_on_expiry_present: false,
auto_delete_on_expiry: false,
})
.await
.expect("update should succeed")
.expect("record should exist");
assert_eq!(updated.concurrent_limit, Some(13));
let snapshot = repository
.find_api_key_snapshot(AuthApiKeyLookupKey::ApiKeyId("key-standalone"))
.await
.expect("find should succeed")
.expect("snapshot should exist");
assert_eq!(snapshot.api_key_concurrent_limit, Some(13));
}
}

27
crates/aether-data/src/repository/auth/sql.rs
View File

@@ -464,6 +464,7 @@ UPDATE api_keys
SET
name = COALESCE($3, name),
rate_limit = COALESCE($4, rate_limit),
concurrent_limit = COALESCE($5, concurrent_limit),
updated_at = NOW()
WHERE user_id = $1
AND id = $2
@@ -493,11 +494,12 @@ UPDATE api_keys
SET
name = COALESCE($2, name),
rate_limit = CASE WHEN $3 THEN $4 ELSE rate_limit END,
allowed_providers = CASE WHEN $5 THEN $6::json ELSE allowed_providers END,
allowed_api_formats = CASE WHEN $7 THEN $8::json ELSE allowed_api_formats END,
allowed_models = CASE WHEN $9 THEN $10::json ELSE allowed_models END,
expires_at = CASE WHEN $11 THEN $12 ELSE expires_at END,
auto_delete_on_expiry = CASE WHEN $13 THEN $14 ELSE auto_delete_on_expiry END,
concurrent_limit = CASE WHEN $5 THEN $6 ELSE concurrent_limit END,
allowed_providers = CASE WHEN $7 THEN $8::json ELSE allowed_providers END,
allowed_api_formats = CASE WHEN $9 THEN $10::json ELSE allowed_api_formats END,
allowed_models = CASE WHEN $11 THEN $12::json ELSE allowed_models END,
expires_at = CASE WHEN $13 THEN $14 ELSE expires_at END,
auto_delete_on_expiry = CASE WHEN $15 THEN $16 ELSE auto_delete_on_expiry END,
updated_at = NOW()
WHERE id = $1
AND is_standalone = TRUE
@@ -1105,6 +1107,7 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.bind(record.api_key_id)
.bind(record.name)
.bind(record.rate_limit)
.bind(record.concurrent_limit)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?;
@@ -1149,6 +1152,8 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.bind(record.name)
.bind(record.rate_limit_present)
.bind(record.rate_limit)
.bind(record.concurrent_limit_present)
.bind(record.concurrent_limit)
.bind(record.allowed_providers.is_some())
.bind(allowed_providers)
.bind(record.allowed_api_formats.is_some())
@@ -1363,18 +1368,20 @@ mod tests {
#[test]
fn update_standalone_api_key_basic_sql_casts_json_case_values() {
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
.contains("allowed_providers = CASE WHEN $5 THEN $6::json ELSE allowed_providers END"));
.contains("concurrent_limit = CASE WHEN $5 THEN $6 ELSE concurrent_limit END"));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
.contains("allowed_providers = CASE WHEN $7 THEN $8::json ELSE allowed_providers END"));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL.contains(
"allowed_api_formats = CASE WHEN $7 THEN $8::json ELSE allowed_api_formats END"
"allowed_api_formats = CASE WHEN $9 THEN $10::json ELSE allowed_api_formats END"
));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
.contains("allowed_models = CASE WHEN $9 THEN $10::json ELSE allowed_models END"));
.contains("allowed_models = CASE WHEN $11 THEN $12::json ELSE allowed_models END"));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
.contains("rate_limit = CASE WHEN $3 THEN $4 ELSE rate_limit END"));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
.contains("expires_at = CASE WHEN $11 THEN $12 ELSE expires_at END"));
.contains("expires_at = CASE WHEN $13 THEN $14 ELSE expires_at END"));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL.contains(
"auto_delete_on_expiry = CASE WHEN $13 THEN $14 ELSE auto_delete_on_expiry END"
"auto_delete_on_expiry = CASE WHEN $15 THEN $16 ELSE auto_delete_on_expiry END"
));
}

7
crates/aether-data/src/repository/auth/types.rs
View File

@@ -363,7 +363,7 @@ pub struct CreateUserApiKeyRecord {
pub allowed_api_formats: Option<Vec<String>>,
pub allowed_models: Option<Vec<String>>,
pub rate_limit: i32,
pub concurrent_limit: i32,
pub concurrent_limit: Option<i32>,
pub force_capabilities: Option<serde_json::Value>,
pub is_active: bool,
pub expires_at_unix_secs: Option<u64>,
@@ -378,6 +378,7 @@ pub struct UpdateUserApiKeyBasicRecord {
pub api_key_id: String,
pub name: Option<String>,
pub rate_limit: Option<i32>,
pub concurrent_limit: Option<i32>,
}
#[derive(Debug, Clone, PartialEq)]
@@ -391,7 +392,7 @@ pub struct CreateStandaloneApiKeyRecord {
pub allowed_api_formats: Option<Vec<String>>,
pub allowed_models: Option<Vec<String>>,
pub rate_limit: Option<i32>,
pub concurrent_limit: i32,
pub concurrent_limit: Option<i32>,
pub force_capabilities: Option<serde_json::Value>,
pub is_active: bool,
pub expires_at_unix_secs: Option<u64>,
@@ -406,6 +407,8 @@ pub struct UpdateStandaloneApiKeyBasicRecord {
pub name: Option<String>,
pub rate_limit_present: bool,
pub rate_limit: Option<i32>,
pub concurrent_limit_present: bool,
pub concurrent_limit: Option<i32>,
pub allowed_providers: Option<Option<Vec<String>>>,
pub allowed_api_formats: Option<Option<Vec<String>>>,
pub allowed_models: Option<Option<Vec<String>>>,