diff --git a/apps/aether-gateway/src/handlers/admin/provider/write/keys/update.rs b/apps/aether-gateway/src/handlers/admin/provider/write/keys/update.rs index 1fa2086c3..58032d69f 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/write/keys/update.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/write/keys/update.rs @@ -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::{ @@ -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() @@ -269,11 +270,7 @@ pub(crate) fn build_admin_update_provider_key_record_with_existing_keys( .collect::>() }); 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; diff --git a/apps/aether-gateway/src/handlers/admin/provider/write/normalize.rs b/apps/aether-gateway/src/handlers/admin/provider/write/normalize.rs index 636f4d616..7439fc348 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/write/normalize.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/write/normalize.rs @@ -111,25 +111,49 @@ pub(crate) fn normalize_allow_auth_channel_mismatch_formats( field_name: &str, api_formats: &[String], ) -> Result, 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::>(); - 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>, + api_formats: &[String], +) -> Option { + let values = canonical_allow_auth_channel_mismatch_formats(values)?; + let allowed = api_formats.iter().cloned().collect::>(); + Some(json_string_array( + values + .into_iter() + .filter(|value| allowed.contains(value)) + .collect(), + )) +} + +fn canonical_allow_auth_channel_mismatch_formats( + values: Option>, +) -> Option> { + 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) -> 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 { @@ -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; @@ -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( diff --git a/apps/aether-gateway/src/tests/control/admin/pool.rs b/apps/aether-gateway/src/tests/control/admin/pool.rs index c201ce0e8..acc3b8a76 100644 --- a/apps/aether-gateway/src/tests/control/admin/pool.rs +++ b/apps/aether-gateway/src/tests/control/admin/pool.rs @@ -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], @@ -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); @@ -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); diff --git a/crates/aether-usage/runtime/src/runtime.rs b/crates/aether-usage/runtime/src/runtime.rs index 926ea1185..0c2140096 100644 --- a/crates/aether-usage/runtime/src/runtime.rs +++ b/crates/aether-usage/runtime/src/runtime.rs @@ -6368,6 +6368,7 @@ mod tests { #[derive(Clone)] struct BlockingWriteQueueConfiguredUsageStore { + records: Arc>>, queue: Arc, write_started: Arc, release_writes: Arc, @@ -7274,10 +7275,11 @@ mod tests { impl UsageRecordWriter for BlockingWriteQueueConfiguredUsageStore { async fn upsert_usage_record( &self, - _record: UpsertUsageRecord, + record: UpsertUsageRecord, ) -> Result, DataLayerError> { self.write_started.notify_one(); self.release_writes.notified().await; + self.records.lock().expect("records lock").push(record); self.writes_completed.fetch_add(1, Ordering::AcqRel); Ok(None) } @@ -10277,6 +10279,7 @@ mod tests { let release_writes = Arc::new(tokio::sync::Notify::new()); let writes_completed = Arc::new(AtomicUsize::new(0)); let store = BlockingWriteQueueConfiguredUsageStore { + records: Arc::new(Mutex::new(Vec::new())), queue: queue_runner, write_started: Arc::clone(&write_started), release_writes: Arc::clone(&release_writes), @@ -10349,6 +10352,7 @@ mod tests { let release_writes = Arc::new(tokio::sync::Notify::new()); let writes_completed = Arc::new(AtomicUsize::new(0)); let store = BlockingWriteQueueConfiguredUsageStore { + records: Arc::new(Mutex::new(Vec::new())), queue: queue_runner, write_started: Arc::clone(&write_started), release_writes: Arc::clone(&release_writes), @@ -11479,6 +11483,7 @@ mod tests { )); let queue: Arc = flaky_queue.clone(); let store = BlockingWriteQueueConfiguredUsageStore { + records: Arc::new(Mutex::new(Vec::new())), queue, write_started: Arc::new(tokio::sync::Notify::new()), release_writes: Arc::new(tokio::sync::Notify::new()), @@ -11676,9 +11681,12 @@ mod tests { }; let queue: Arc = Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default())); - let store = CloneQueueConfiguredUsageStore { + let store = BlockingWriteQueueConfiguredUsageStore { records: Arc::new(Mutex::new(Vec::new())), queue, + write_started: Arc::new(tokio::sync::Notify::new()), + release_writes: Arc::new(tokio::sync::Notify::new()), + writes_completed: Arc::new(AtomicUsize::new(0)), }; let runtime = UsageRuntime::new(config).expect("usage runtime should build"); let blocker_started = Arc::new(tokio::sync::Notify::new()); @@ -11733,7 +11741,25 @@ mod tests { .expect("terminal seed and awaited barrier should queue behind the blocker"); assert_eq!(runtime.metrics_snapshot().terminal_submission_pending, 0); - release_blocker.notify_waiters(); + let first_write_started = store.write_started.notified(); + release_blocker.notify_one(); + timeout(Duration::from_secs(1), first_write_started) + .await + .expect("terminal seed write should start"); + timeout(Duration::from_secs(1), async { + while runtime.metrics_snapshot().terminal_submission_pending != 2 { + tokio::task::yield_now().await; + } + }) + .await + .expect("awaited terminal should queue behind the terminal seed permit"); + + let second_write_started = store.write_started.notified(); + store.release_writes.notify_one(); + timeout(Duration::from_secs(1), second_write_started) + .await + .expect("awaited terminal write should start after the seed completes"); + store.release_writes.notify_one(); timeout(Duration::from_secs(2), awaited_terminal) .await .expect("awaited terminal must not deadlock behind the queued seed") @@ -11741,6 +11767,7 @@ mod tests { timeout(Duration::from_secs(2), async { while store.records.lock().expect("records lock").len() != 2 || runtime.metrics_snapshot().lifecycle_submission_pending != 0 + || store.writes_completed.load(Ordering::Acquire) != 2 { tokio::task::yield_now().await; } @@ -11771,9 +11798,12 @@ mod tests { }; let queue: Arc = Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default())); - let store = CloneQueueConfiguredUsageStore { + let store = BlockingWriteQueueConfiguredUsageStore { records: Arc::new(Mutex::new(Vec::new())), queue, + write_started: Arc::new(tokio::sync::Notify::new()), + release_writes: Arc::new(tokio::sync::Notify::new()), + writes_completed: Arc::new(AtomicUsize::new(0)), }; let runtime = UsageRuntime::new(config).expect("usage runtime should build"); let blocker_started = Arc::new(tokio::sync::Notify::new()); @@ -11815,10 +11845,30 @@ mod tests { assert_eq!(runtime.metrics_snapshot().terminal_submission_pending, 0); assert!(runtime.metrics_snapshot().lifecycle_submission_pending >= 3); - release_blocker.notify_waiters(); + + let first_write_started = store.write_started.notified(); + release_blocker.notify_one(); + timeout(Duration::from_secs(1), first_write_started) + .await + .expect("terminal seed write should start"); + timeout(Duration::from_secs(1), async { + while runtime.metrics_snapshot().terminal_submission_pending != 2 { + tokio::task::yield_now().await; + } + }) + .await + .expect("submitted terminal should queue behind the terminal seed permit"); + + let second_write_started = store.write_started.notified(); + store.release_writes.notify_one(); + timeout(Duration::from_secs(1), second_write_started) + .await + .expect("submitted terminal write should start after the seed completes"); + store.release_writes.notify_one(); timeout(Duration::from_secs(2), async { while store.records.lock().expect("records lock").len() != 2 || runtime.metrics_snapshot().lifecycle_submission_pending != 0 + || store.writes_completed.load(Ordering::Acquire) != 2 { tokio::task::yield_now().await; } @@ -12219,6 +12269,7 @@ mod tests { ..UsageRuntimeConfig::default() }; let store = BlockingWriteQueueConfiguredUsageStore { + records: Arc::new(Mutex::new(Vec::new())), queue: Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default())), write_started: Arc::new(tokio::sync::Notify::new()), release_writes: Arc::new(tokio::sync::Notify::new()),