Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,8 @@ use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyUpdatePa
use crate::handlers::admin::provider::write::normalize::{
normalize_allow_auth_channel_mismatch_formats, normalize_api_format_json_object_keys,
normalize_api_format_list, normalize_auth_type, normalize_auth_type_by_format,
normalize_max_probe_interval_minutes, normalize_rate_multipliers, validate_vertex_api_formats,
normalize_max_probe_interval_minutes, normalize_rate_multipliers,
reconcile_allow_auth_channel_mismatch_formats, validate_vertex_api_formats,
};
use crate::handlers::admin::request::AdminAppState;
use crate::handlers::admin::shared::{
Expand Down Expand Up @@ -256,7 +257,7 @@ pub(crate) fn build_admin_update_provider_key_record_with_existing_keys(
"allow_auth_channel_mismatch_formats",
&effective_api_formats,
)?;
} else if fields.contains("api_formats") {
} else if fields.contains("api_formats") && !managed_fixed_oauth_key {
let existing = updated
.allow_auth_channel_mismatch_formats
.as_ref()
Expand All @@ -269,11 +270,7 @@ pub(crate) fn build_admin_update_provider_key_record_with_existing_keys(
.collect::<Vec<_>>()
});
updated.allow_auth_channel_mismatch_formats =
normalize_allow_auth_channel_mismatch_formats(
existing,
"allow_auth_channel_mismatch_formats",
&effective_api_formats,
)?;
reconcile_allow_auth_channel_mismatch_formats(existing, &effective_api_formats);
}

updated.auth_type = target_auth_type;
Expand Down
77 changes: 62 additions & 15 deletions apps/aether-gateway/src/handlers/admin/provider/write/normalize.rs
Original file line number Diff line number Diff line change
Expand Up @@ -111,25 +111,49 @@ pub(crate) fn normalize_allow_auth_channel_mismatch_formats(
field_name: &str,
api_formats: &[String],
) -> Result<Option<serde_json::Value>, String> {
let Some(values) = values else {
let Some(values) = canonical_allow_auth_channel_mismatch_formats(values) else {
return Ok(None);
};
let allowed = api_formats.iter().cloned().collect::<BTreeSet<_>>();
let mut seen = BTreeSet::new();
let mut normalized = Vec::new();
for value in values {
let canonical = crate::ai_serving::normalize_api_format_alias(&value);
if canonical.is_empty() {
continue;
}
if !allowed.is_empty() && !allowed.contains(&canonical) {
return Err(format!("{field_name} 包含未选择的 API 格式: {canonical}"));
}
if seen.insert(canonical.clone()) {
normalized.push(serde_json::Value::String(canonical));
for value in &values {
if !allowed.is_empty() && !allowed.contains(value) {
return Err(format!("{field_name} 包含未选择的 API 格式: {value}"));
}
}
Ok(Some(serde_json::Value::Array(normalized)))
Ok(Some(json_string_array(values)))
}

pub(crate) fn reconcile_allow_auth_channel_mismatch_formats(
values: Option<Vec<String>>,
api_formats: &[String],
) -> Option<serde_json::Value> {
let values = canonical_allow_auth_channel_mismatch_formats(values)?;
let allowed = api_formats.iter().cloned().collect::<BTreeSet<_>>();
Some(json_string_array(
values
.into_iter()
.filter(|value| allowed.contains(value))
.collect(),
))
}

fn canonical_allow_auth_channel_mismatch_formats(
values: Option<Vec<String>>,
) -> Option<Vec<String>> {
let values = values?;
let mut seen = BTreeSet::new();
Some(
values
.into_iter()
.map(|value| crate::ai_serving::normalize_api_format_alias(&value))
.filter(|value| !value.is_empty())
.filter(|value| seen.insert(value.clone()))
.collect(),
)
}

fn json_string_array(values: Vec<String>) -> serde_json::Value {
serde_json::Value::Array(values.into_iter().map(serde_json::Value::String).collect())
}

pub(crate) fn normalize_auth_type(value: Option<&str>) -> Result<String, String> {
Expand Down Expand Up @@ -233,7 +257,8 @@ mod tests {
normalize_allow_auth_channel_mismatch_formats, normalize_api_format_json_object_keys,
normalize_api_format_list, normalize_auth_type, normalize_auth_type_by_format,
normalize_chat_pii_redaction_config, normalize_pool_advanced_config,
normalize_provider_type_input, normalize_rate_multipliers, validate_vertex_api_formats,
normalize_provider_type_input, normalize_rate_multipliers,
reconcile_allow_auth_channel_mismatch_formats, validate_vertex_api_formats,
};
use serde_json::json;

Expand Down Expand Up @@ -405,6 +430,28 @@ mod tests {
);
}

#[test]
fn reconcile_allow_auth_channel_mismatch_formats_keeps_only_selected_formats() {
assert_eq!(
reconcile_allow_auth_channel_mismatch_formats(
Some(vec![
"OPENAI:EMBEDDING".to_string(),
"gemini:generate_content".to_string(),
" GEMINI:GENERATE_CONTENT ".to_string(),
]),
&["gemini:generate_content".to_string()],
),
Some(json!(["gemini:generate_content"]))
);
assert_eq!(
reconcile_allow_auth_channel_mismatch_formats(
Some(vec!["openai:embedding".to_string()]),
&["gemini:generate_content".to_string()],
),
Some(json!([]))
);
}

#[test]
fn validate_vertex_api_formats_rejects_unimplemented_anthropic_transport() {
assert!(validate_vertex_api_formats(
Expand Down
115 changes: 115 additions & 0 deletions apps/aether-gateway/src/tests/control/admin/pool.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3606,9 +3606,12 @@ async fn gateway_batch_updates_shared_pool_key_configuration() {
first_key.name = "alpha".to_string();
first_key.auto_fetch_models = true;
first_key.allowed_models = Some(json!(["legacy-model"]));
first_key.allow_auth_channel_mismatch_formats = Some(json!(["openai:embedding"]));
first_key.learned_rpm_limit = Some(18);
let mut second_key = sample_key("key-openai-b", "provider-openai", "openai:chat", "sk-b");
second_key.name = "beta".to_string();
second_key.allow_auth_channel_mismatch_formats =
Some(json!(["openai:chat", "openai:embedding"]));
second_key.learned_rpm_limit = Some(24);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Expand Down Expand Up @@ -3662,6 +3665,7 @@ async fn gateway_batch_updates_shared_pool_key_configuration() {
assert_eq!(stored.len(), 2);
for key in stored {
assert_eq!(key.api_formats, Some(json!(["openai:responses"])));
assert_eq!(key.allow_auth_channel_mismatch_formats, Some(json!([])));
assert_eq!(key.internal_priority, 7);
assert_eq!(key.rpm_limit, None);
assert_eq!(key.learned_rpm_limit, None);
Expand All @@ -3677,6 +3681,117 @@ async fn gateway_batch_updates_shared_pool_key_configuration() {
gateway_handle.abort();
}

#[tokio::test]
async fn gateway_preserves_inherited_fixed_oauth_mismatch_formats_on_batch_update() {
let mut provider = sample_provider("provider-gemini-cli", "gemini_cli", 10);
provider.provider_type = "gemini_cli".to_string();
let mut key = sample_key(
"key-gemini-cli-a",
"provider-gemini-cli",
"gemini:generate_content",
"oauth-placeholder",
);
key.auth_type = "oauth".to_string();
key.internal_priority = 3;
key.allow_auth_channel_mismatch_formats = Some(json!(["gemini:generate_content"]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![key],
));
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone(
&provider_catalog_repository,
)),
);

let response = local_admin_pool_response(
&state,
http::Method::PATCH,
"/api/admin/pool/provider-gemini-cli/keys/batch-update",
Some(json!({
"key_ids": ["key-gemini-cli-a"],
"patch": {
"api_formats": ["openai:responses"],
"internal_priority": 9
}
})),
)
.await;
assert_eq!(response.status(), StatusCode::OK);

let stored = provider_catalog_repository
.list_keys_by_ids(&["key-gemini-cli-a".to_string()])
.await
.expect("key should load");
assert_eq!(stored[0].api_formats, None);
assert_eq!(
stored[0].allow_auth_channel_mismatch_formats,
Some(json!(["gemini:generate_content"]))
);
assert_eq!(stored[0].internal_priority, 9);
}

#[tokio::test]
async fn gateway_rejects_explicit_invalid_mismatch_format_without_writing_any_key() {
let provider = sample_provider("provider-openai", "openai", 10);
let mut first_key = sample_key("key-openai-a", "provider-openai", "openai:chat", "sk-a");
first_key.name = "alpha".to_string();
first_key.internal_priority = 3;
first_key.allow_auth_channel_mismatch_formats = Some(json!(["openai:chat"]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![first_key],
));
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone(
&provider_catalog_repository,
)),
);

let response = local_admin_pool_response(
&state,
http::Method::PATCH,
"/api/admin/pool/provider-openai/keys/batch-update",
Some(json!({
"key_ids": ["key-openai-a"],
"patch": {
"api_formats": ["openai:responses"],
"allow_auth_channel_mismatch_formats": ["openai:embedding"],
"internal_priority": 9
}
})),
)
.await;
assert_eq!(response.status(), StatusCode::BAD_REQUEST);

let body = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.expect("response body should load");
let payload: serde_json::Value =
serde_json::from_slice(&body).expect("response body should be json");
assert_eq!(
payload["detail"],
json!("密钥 alpha 配置无效: allow_auth_channel_mismatch_formats 包含未选择的 API 格式: openai:embedding")
);

let stored = provider_catalog_repository
.list_keys_by_ids(&["key-openai-a".to_string()])
.await
.expect("key should load");
assert_eq!(stored[0].api_formats, Some(json!(["openai:chat"])));
assert_eq!(
stored[0].allow_auth_channel_mismatch_formats,
Some(json!(["openai:chat"]))
);
assert_eq!(stored[0].internal_priority, 3);
}

#[tokio::test]
async fn gateway_rejects_pool_batch_update_before_writing_any_key() {
let provider = sample_provider("provider-openai", "openai", 10);
Expand Down
Loading
Loading