mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-09 10:57:03 +08:00
feat(gateway): harden provider request execution
Preserve exact request payloads and model client surface and API operation explicitly. Add Anthropic compatibility profiles, bounded stream commitment, and scoped OAuth retry behavior across provider transports.
This commit is contained in:
@@ -287,6 +287,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_key(
|
||||
&ProviderCatalogKeyOAuthRuntimeStateCasUpdate {
|
||||
key_id: key_id.clone(),
|
||||
expected_encrypted_auth_config: state_data.expected_encrypted_auth_config,
|
||||
expected_credential: None,
|
||||
encrypted_auth_config: persisted_encrypted_auth_config.clone(),
|
||||
encrypted_api_key_update: Some(encrypted_api_key),
|
||||
expires_at_unix_secs_update: Some(expires_at),
|
||||
|
||||
@@ -370,6 +370,7 @@ pub(crate) async fn persist_fenced_provider_quota_refresh_state(
|
||||
&ProviderCatalogKeyOAuthRuntimeStateCasUpdate {
|
||||
key_id: key_id.to_string(),
|
||||
expected_encrypted_auth_config: Some(expected_encrypted_auth_config.to_string()),
|
||||
expected_credential: None,
|
||||
encrypted_auth_config: expected_encrypted_auth_config.to_string(),
|
||||
encrypted_api_key_update: None,
|
||||
expires_at_unix_secs_update: None,
|
||||
|
||||
@@ -100,14 +100,6 @@ fn select_provider_oauth_runtime_endpoint(
|
||||
.api_format
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("gemini:generate_content")
|
||||
})
|
||||
.or_else(|| {
|
||||
matching_endpoint(endpoints, include_inactive, |endpoint| {
|
||||
endpoint
|
||||
.api_format
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("claude:messages")
|
||||
})
|
||||
}),
|
||||
_ => matching_endpoint(endpoints, include_inactive, |_| true),
|
||||
}
|
||||
@@ -255,3 +247,44 @@ pub(crate) fn spawn_provider_oauth_account_state_refresh_after_update(
|
||||
.await;
|
||||
});
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
provider_oauth_maintenance_endpoint_for_provider,
|
||||
provider_oauth_runtime_endpoint_for_provider,
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogEndpoint;
|
||||
|
||||
fn endpoint(id: &str, api_format: &str, is_active: bool) -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
id.to_string(),
|
||||
"provider-1".to_string(),
|
||||
api_format.to_string(),
|
||||
None,
|
||||
None,
|
||||
is_active,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn vertex_oauth_runtime_never_falls_back_to_retired_claude_endpoint() {
|
||||
let endpoints = vec![endpoint("claude", "claude:messages", true)];
|
||||
|
||||
assert!(provider_oauth_runtime_endpoint_for_provider("vertex_ai", &endpoints).is_none());
|
||||
assert!(
|
||||
provider_oauth_maintenance_endpoint_for_provider("vertex_ai", &endpoints).is_none()
|
||||
);
|
||||
|
||||
let endpoints = vec![
|
||||
endpoint("claude", "claude:messages", true),
|
||||
endpoint("gemini", "gemini:generate_content", true),
|
||||
];
|
||||
assert_eq!(
|
||||
provider_oauth_runtime_endpoint_for_provider("vertex_ai", &endpoints)
|
||||
.map(|endpoint| endpoint.id),
|
||||
Some("gemini".to_string())
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2668,6 +2668,7 @@ async fn provider_query_execute_antigravity_test_candidate(
|
||||
upstream_is_stream: false,
|
||||
request_query: parts.uri.query(),
|
||||
kiro_api_region: None,
|
||||
api_operation: None,
|
||||
},
|
||||
);
|
||||
let Some(request_url) = request_url else {
|
||||
@@ -3364,6 +3365,7 @@ async fn provider_query_execute_standard_test_candidate(
|
||||
upstream_is_stream,
|
||||
request_query: parts.uri.query(),
|
||||
kiro_api_region: None,
|
||||
api_operation: None,
|
||||
},
|
||||
Some(&provider_request_body),
|
||||
);
|
||||
|
||||
@@ -195,11 +195,7 @@ pub(crate) fn validate_vertex_api_formats(
|
||||
|
||||
let allowed = match auth_type {
|
||||
"api_key" => &["gemini:generate_content", "gemini:embedding"][..],
|
||||
"service_account" | "vertex_ai" => &[
|
||||
"claude:messages",
|
||||
"gemini:generate_content",
|
||||
"gemini:embedding",
|
||||
][..],
|
||||
"service_account" | "vertex_ai" => &["gemini:generate_content", "gemini:embedding"][..],
|
||||
_ => return Ok(()),
|
||||
};
|
||||
let invalid = api_formats
|
||||
@@ -410,20 +406,11 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validate_vertex_api_formats_uses_canonical_message_formats() {
|
||||
fn validate_vertex_api_formats_rejects_unimplemented_anthropic_transport() {
|
||||
assert!(validate_vertex_api_formats(
|
||||
"vertex_ai",
|
||||
"service_account",
|
||||
&[
|
||||
"claude:messages".to_string(),
|
||||
"gemini:generate_content".to_string()
|
||||
],
|
||||
)
|
||||
.is_ok());
|
||||
assert!(validate_vertex_api_formats(
|
||||
"vertex_ai",
|
||||
"service_account",
|
||||
&["claude:chat".to_string()],
|
||||
&["claude:messages".to_string()],
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
@@ -443,7 +430,6 @@ mod tests {
|
||||
"vertex_ai",
|
||||
"service_account",
|
||||
&[
|
||||
"claude:messages".to_string(),
|
||||
"gemini:generate_content".to_string(),
|
||||
"gemini:embedding".to_string()
|
||||
],
|
||||
|
||||
@@ -158,6 +158,8 @@ pub(crate) async fn build_admin_create_provider_record(
|
||||
}
|
||||
}
|
||||
let config = (!config_map.is_empty()).then_some(serde_json::Value::Object(config_map));
|
||||
crate::provider_transport::validate_anthropic_compatibility_profile_config(config.as_ref())
|
||||
.map_err(|_| "无效的 Anthropic compatibility profile".to_string())?;
|
||||
|
||||
let now_unix_secs = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
|
||||
@@ -312,6 +312,10 @@ pub(crate) async fn build_admin_update_provider_record(
|
||||
}
|
||||
|
||||
updated.config = (!config_map.is_empty()).then_some(serde_json::Value::Object(config_map));
|
||||
crate::provider_transport::validate_anthropic_compatibility_profile_config(
|
||||
updated.config.as_ref(),
|
||||
)
|
||||
.map_err(|_| "无效的 Anthropic compatibility profile".to_string())?;
|
||||
updated.updated_at_unix_secs = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
|
||||
@@ -259,6 +259,10 @@ impl<'a> AdminAppState<'a> {
|
||||
admin_endpoint_signature_parts(&payload.api_format)
|
||||
.ok_or_else(|| format!("无效的 api_format: {}", payload.api_format))?;
|
||||
validate_admin_endpoint_stream_policy(normalized_api_format, payload.config.as_ref())?;
|
||||
crate::provider_transport::validate_anthropic_compatibility_profile_config(
|
||||
payload.config.as_ref(),
|
||||
)
|
||||
.map_err(|_| "无效的 Anthropic compatibility profile".to_string())?;
|
||||
let base_url = normalize_admin_base_url(&payload.base_url)?;
|
||||
|
||||
let existing_endpoints = self
|
||||
@@ -369,6 +373,10 @@ impl<'a> AdminAppState<'a> {
|
||||
existing_endpoint.api_format.as_str(),
|
||||
updated.config.as_ref(),
|
||||
)?;
|
||||
crate::provider_transport::validate_anthropic_compatibility_profile_config(
|
||||
updated.config.as_ref(),
|
||||
)
|
||||
.map_err(|_| "无效的 Anthropic compatibility profile".to_string())?;
|
||||
}
|
||||
|
||||
if provider_type == "codex"
|
||||
|
||||
@@ -192,6 +192,15 @@ fn normalize_import_endpoint_format(value: &str) -> Result<String, String> {
|
||||
.ok_or_else(|| format!("无效的 api_format: {value}"))
|
||||
}
|
||||
|
||||
fn fixed_provider_import_endpoint_supported(provider_type: &str, api_format: &str) -> bool {
|
||||
crate::provider_transport::provider_types::fixed_provider_template(provider_type).is_none()
|
||||
|| crate::provider_transport::provider_types::fixed_provider_endpoint_template_by_api_format(
|
||||
provider_type,
|
||||
api_format,
|
||||
)
|
||||
.is_some()
|
||||
}
|
||||
|
||||
fn normalize_import_key_formats(
|
||||
item: &ImportedProviderKey,
|
||||
provider_endpoint_formats: &BTreeSet<String>,
|
||||
@@ -1419,6 +1428,12 @@ impl<'a> AdminAppState<'a> {
|
||||
for imported_provider_item in imported_providers {
|
||||
let (raw_provider, imported_provider) = imported_provider_item.into_parts();
|
||||
let provider_name = invalid!(trim_required(&imported_provider.name, "name"));
|
||||
invalid!(
|
||||
crate::provider_transport::validate_anthropic_compatibility_profile_config(
|
||||
imported_provider.config.as_ref(),
|
||||
)
|
||||
.map_err(|_| "无效的 Anthropic compatibility profile".to_string())
|
||||
);
|
||||
let existing_provider = providers_by_name.get(&provider_name).cloned();
|
||||
|
||||
let provider = if let Some(existing) = existing_provider {
|
||||
@@ -1513,6 +1528,42 @@ impl<'a> AdminAppState<'a> {
|
||||
let normalized_api_format = invalid!(normalize_import_endpoint_format(
|
||||
&imported_endpoint.api_format
|
||||
));
|
||||
invalid!(
|
||||
crate::provider_transport::validate_anthropic_compatibility_profile_config(
|
||||
imported_endpoint.config.as_ref(),
|
||||
)
|
||||
.map_err(|_| "无效的 Anthropic compatibility profile".to_string())
|
||||
);
|
||||
if !fixed_provider_import_endpoint_supported(
|
||||
&provider.provider_type,
|
||||
&normalized_api_format,
|
||||
) {
|
||||
let retired = existing_endpoints_by_format.remove(&normalized_api_format);
|
||||
if let Some(mut retired) = retired {
|
||||
if retired.is_active {
|
||||
retired.is_active = false;
|
||||
retired.updated_at_unix_secs = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
.map(|duration| duration.as_secs());
|
||||
let Some(_) = self.update_provider_catalog_endpoint(&retired).await?
|
||||
else {
|
||||
return Ok(Err(invalid_request(format!(
|
||||
"停用 Provider '{provider_name}' 的已移除 Endpoint '{normalized_api_format}' 失败"
|
||||
))));
|
||||
};
|
||||
stats.endpoints.updated += 1;
|
||||
} else {
|
||||
stats.endpoints.skipped += 1;
|
||||
}
|
||||
} else {
|
||||
stats.endpoints.skipped += 1;
|
||||
}
|
||||
stats.errors.push(format!(
|
||||
"固定 Provider '{provider_name}' 不再支持 Endpoint '{normalized_api_format}',已跳过或停用"
|
||||
));
|
||||
continue;
|
||||
}
|
||||
let existing_endpoint = existing_endpoints_by_format
|
||||
.get(&normalized_api_format)
|
||||
.cloned();
|
||||
|
||||
Reference in New Issue
Block a user