fix(provider): 模型测试按 Key 模型权限过滤

测试模型前检查 provider key 的 allowed_models:
- 空权限视为允许所有模型
- 非空权限需匹配请求模型或映射后的实际模型
- 不匹配的 key 标记为跳过,避免发起测试请求

同时补充相关单测和前端跳过原因文案。
This commit is contained in:
AAEE86
2026-05-26 16:46:21 +08:00
parent 7e76c9763d
commit 949e251b2e
4 changed files with 130 additions and 4 deletions
@@ -123,6 +123,7 @@ const ADMIN_PROVIDER_QUERY_NO_ACTIVE_TEST_CANDIDATE_DETAIL: &str =
"No active endpoint or API key found";
const ADMIN_PROVIDER_QUERY_INVALID_MAPPED_MODEL_DETAIL: &str =
"mapped_model_name is not valid for the selected model and endpoint";
const PROVIDER_QUERY_KEY_MODEL_NOT_ALLOWED_SKIP_REASON: &str = "key_model_not_allowed";
const ANTIGRAVITY_PROVIDER_CACHE_KEY_PREFIX: &str = "upstream_models_provider:";
const DEFAULT_PROVIDER_QUERY_TEST_MESSAGE: &str = "Hello! This is a test message.";
static PROVIDER_QUERY_POOL_LOAD_BALANCE_SEQUENCE: AtomicU64 = AtomicU64::new(0);
@@ -933,6 +934,35 @@ fn provider_query_selected_key_ids_all_exist(
.all(|id| keys.iter().any(|key| key.id == *id))
}
fn provider_query_model_name_matches(left: &str, right: &str) -> bool {
let left = left.trim();
let right = right.trim();
!left.is_empty() && !right.is_empty() && left.eq_ignore_ascii_case(right)
}
fn provider_query_key_allows_effective_test_model(
key: &StoredProviderCatalogKey,
requested_model: &str,
effective_model: &str,
) -> bool {
let allowed_models = json_string_list(key.allowed_models.as_ref());
if key.allowed_models.is_none() || allowed_models.is_empty() {
return true;
}
let requested_base_model = crate::ai_serving::model_directive_base_model(requested_model);
allowed_models
.iter()
.map(String::as_str)
.any(|allowed_model| {
provider_query_model_name_matches(allowed_model, requested_model)
|| provider_query_model_name_matches(allowed_model, effective_model)
|| requested_base_model.as_deref().is_some_and(|base_model| {
provider_query_model_name_matches(allowed_model, base_model)
})
})
}
fn provider_query_test_key_sort_key(
provider_type: &str,
key: &StoredProviderCatalogKey,
@@ -1382,17 +1412,43 @@ async fn provider_query_build_kiro_test_candidates(
.unwrap_or(requested_model.clone())
};
let mut keys = all_keys
let now_unix_secs = current_unix_ms() / 1000;
let mut keys = Vec::new();
let mut model_skipped_candidates = Vec::new();
for key in all_keys
.into_iter()
.filter(|key| key.is_active)
.filter(|key| provider_query_selected_key_ids_allow_key(selected_key_ids.as_ref(), &key.id))
.filter(|key| {
provider_query_key_supports_endpoint(key, &provider.provider_type, &endpoint.api_format)
})
.collect::<Vec<_>>();
let now_unix_secs = current_unix_ms() / 1000;
{
if provider_query_key_allows_effective_test_model(&key, &requested_model, &effective_model)
{
keys.push(key);
} else {
model_skipped_candidates.push(ProviderQueryTestCandidate {
endpoint: endpoint.clone(),
key,
effective_model: effective_model.clone(),
scheduler_skip_reason: Some(
PROVIDER_QUERY_KEY_MODEL_NOT_ALLOWED_SKIP_REASON.to_string(),
),
});
}
}
let candidates = if test_mode.eq_ignore_ascii_case("pool") {
model_skipped_candidates.sort_by_key(|candidate| {
provider_query_test_key_sort_key(
provider.provider_type.as_str(),
&candidate.key,
&endpoint.api_format,
now_unix_secs,
)
});
let scheduled_candidates = if test_mode.eq_ignore_ascii_case("pool") {
if let Some(pool_config) =
admin_provider_pool_config_from_config_value(provider.config.as_ref())
{
@@ -1442,6 +1498,8 @@ async fn provider_query_build_kiro_test_candidates(
})
.collect::<Vec<_>>()
};
let mut candidates = model_skipped_candidates;
candidates.extend(scheduled_candidates);
if candidates.is_empty() {
return Err(build_admin_provider_query_not_found_response(
@@ -63,6 +63,68 @@ fn sample_openai_image_transport(provider_type: &str) -> AdminGatewayProviderTra
}
}
fn sample_catalog_key_with_allowed_models(
allowed_models: Option<serde_json::Value>,
) -> aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey {
let mut key =
aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey::new(
"key-1".to_string(),
"provider-1".to_string(),
"key".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("sample provider key should build");
key.allowed_models = allowed_models;
key
}
#[test]
fn provider_query_model_test_allows_keys_without_model_restrictions() {
let unrestricted = sample_catalog_key_with_allowed_models(None);
let empty = sample_catalog_key_with_allowed_models(Some(json!([])));
assert!(provider_query_key_allows_effective_test_model(
&unrestricted,
"model-b",
"model-b-upstream",
));
assert!(provider_query_key_allows_effective_test_model(
&empty,
"model-b",
"model-b-upstream",
));
}
#[test]
fn provider_query_model_test_filters_key_disallowed_for_requested_model() {
let key = sample_catalog_key_with_allowed_models(Some(json!(["model-a"])));
assert!(!provider_query_key_allows_effective_test_model(
&key,
"model-b",
"model-b-upstream",
));
}
#[test]
fn provider_query_model_test_allows_key_for_requested_or_mapped_model() {
let requested_allowed = sample_catalog_key_with_allowed_models(Some(json!(["model-b"])));
let mapped_allowed = sample_catalog_key_with_allowed_models(Some(json!(["MODEL-B-UPSTREAM"])));
assert!(provider_query_key_allows_effective_test_model(
&requested_allowed,
"model-b",
"model-b-upstream",
));
assert!(provider_query_key_allows_effective_test_model(
&mapped_allowed,
"model-b",
"model-b-upstream",
));
}
#[test]
fn provider_query_test_request_body_preserves_custom_model() {
let payload = json!({
@@ -372,6 +372,11 @@ describe('isModelTestableEndpoint', () => {
})
describe('formatModelTestDiagnostic', () => {
it('maps model permission skips to an actionable label', () => {
expect(formatModelTestDiagnostic('key_model_not_allowed'))
.toBe('Key 未允许当前模型,已跳过')
})
it('maps pool account blocked scheduler code to an actionable label', () => {
expect(formatModelTestDiagnostic('pool_account_blocked')).toBe('账号已失效,需重新授权')
})
@@ -41,6 +41,7 @@ const MODEL_TEST_BEARER_INHERITS_PROVIDER_FORMATS = new Set([
])
const MODEL_TEST_DIAGNOSTIC_LABELS: Record<string, string> = {
key_model_not_allowed: 'Key 未允许当前模型,已跳过',
pool_account_blocked: '账号已失效,需重新授权',
}