Merge remote-tracking branch 'origin/pr/451' into aether-rust-pioneer

This commit is contained in:
fawney19
2026-05-14 02:10:37 +08:00
43 changed files with 7770 additions and 126 deletions

View File

@@ -257,6 +257,7 @@ pub(super) async fn execute_provider_quota_plan(
let error = match err {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
};
let proxy_node_id = plan

View File

@@ -304,6 +304,7 @@ fn admin_provider_ops_gateway_error_message(error: GatewayError) -> String {
match error {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
}
}

View File

@@ -176,6 +176,7 @@ pub(crate) fn build_admin_provider_summary_value(
"claude_code_advanced": config.and_then(|cfg| cfg.get("claude_code_advanced")).cloned(),
"pool_advanced": config.and_then(|cfg| cfg.get("pool_advanced")).cloned(),
"failover_rules": config.and_then(|cfg| cfg.get("failover_rules")).cloned(),
"chat_pii_redaction": config.and_then(|cfg| cfg.get("chat_pii_redaction")).cloned(),
"total_endpoints": total_endpoints,
"active_endpoints": active_endpoints,
"total_keys": total_keys,

View File

@@ -129,6 +129,28 @@ pub(crate) fn normalize_pool_advanced_config(
}
}
pub(crate) fn normalize_chat_pii_redaction_config(
value: Option<serde_json::Value>,
) -> Result<Option<serde_json::Value>, String> {
let Some(value) = value else {
return Ok(None);
};
match value {
serde_json::Value::Null => Ok(None),
serde_json::Value::Object(mut map) => {
if map.len() != 1 || !map.contains_key("enabled") {
return Err("chat_pii_redaction 仅支持 enabled 布尔配置".to_string());
}
let enabled = map
.remove("enabled")
.and_then(|value| value.as_bool())
.ok_or_else(|| "chat_pii_redaction.enabled 必须是布尔值".to_string())?;
Ok(Some(serde_json::json!({ "enabled": enabled })))
}
_ => Err("chat_pii_redaction 必须是 JSON 对象".to_string()),
}
}
pub(crate) fn validate_vertex_api_formats(
provider_type: &str,
auth_type: &str,
@@ -177,7 +199,8 @@ mod tests {
use super::{
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_pool_advanced_config, normalize_provider_type_input, validate_vertex_api_formats,
normalize_chat_pii_redaction_config, normalize_pool_advanced_config,
normalize_provider_type_input, validate_vertex_api_formats,
};
use serde_json::json;
@@ -201,6 +224,26 @@ mod tests {
);
}
#[test]
fn normalize_chat_pii_redaction_requires_enabled_boolean_only() {
assert_eq!(
normalize_chat_pii_redaction_config(Some(json!({ "enabled": true })))
.expect("chat pii redaction should normalize"),
Some(json!({ "enabled": true }))
);
assert_eq!(
normalize_chat_pii_redaction_config(Some(
json!({ "enabled": true, "entities": ["email"] })
))
.unwrap_err(),
"chat_pii_redaction 仅支持 enabled 布尔配置"
);
assert_eq!(
normalize_chat_pii_redaction_config(Some(json!({ "enabled": "yes" }))).unwrap_err(),
"chat_pii_redaction.enabled 必须是布尔值"
);
}
#[test]
fn normalize_auth_type_supports_bearer() {
assert_eq!(

View File

@@ -2,6 +2,7 @@ use crate::handlers::admin::provider::shared::payloads::AdminProviderCreateReque
use crate::handlers::admin::provider::shared::support::{
normalize_provider_billing_type, parse_optional_rfc3339_unix_secs,
};
use crate::handlers::admin::provider::write::normalize::normalize_chat_pii_redaction_config;
use crate::handlers::admin::provider::write::normalize::normalize_pool_advanced_config;
use crate::handlers::admin::provider::write::normalize::normalize_provider_type_input;
use crate::handlers::admin::request::AdminAppState;
@@ -134,6 +135,12 @@ pub(crate) async fn build_admin_create_provider_record(
}
config_map.insert("claude_code_advanced".to_string(), value);
}
if config_map.contains_key("chat_pii_redaction") {
let value = normalize_chat_pii_redaction_config(config_map.remove("chat_pii_redaction"))?;
if let Some(value) = value {
config_map.insert("chat_pii_redaction".to_string(), value);
}
}
let config = (!config_map.is_empty()).then_some(serde_json::Value::Object(config_map));
let now_unix_secs = SystemTime::now()

View File

@@ -2,6 +2,7 @@ use crate::handlers::admin::provider::shared::payloads::AdminProviderUpdatePatch
use crate::handlers::admin::provider::shared::support::{
normalize_provider_billing_type, parse_optional_rfc3339_unix_secs,
};
use crate::handlers::admin::provider::write::normalize::normalize_chat_pii_redaction_config;
use crate::handlers::admin::provider::write::normalize::normalize_pool_advanced_config;
use crate::handlers::admin::provider::write::normalize::normalize_provider_type_input;
use crate::handlers::admin::request::AdminAppState;
@@ -225,14 +226,29 @@ pub(crate) async fn build_admin_update_provider_record(
updated.enable_format_conversion = enable_format_conversion;
}
let config_seed = if fields.contains("config") {
normalize_json_object(payload.config, "config")?
} else {
updated.config.clone()
};
let mut config_map = config_seed
let mut config_map = updated
.config
.clone()
.and_then(|value| value.as_object().cloned())
.unwrap_or_default();
if fields.contains("config") {
if fields.is_null("config") {
config_map.clear();
} else {
let value = normalize_json_object(payload.config, "config")?
.ok_or_else(|| "config 必须是 JSON 对象".to_string())?;
let serde_json::Value::Object(patch_map) = value else {
return Err("config 必须是 JSON 对象".to_string());
};
for (key, value) in patch_map {
if value.is_null() {
config_map.remove(&key);
} else {
config_map.insert(key, value);
}
}
}
}
if fields.contains("claude_code_advanced") {
if fields.is_null("claude_code_advanced") {
@@ -270,6 +286,13 @@ pub(crate) async fn build_admin_update_provider_record(
}
}
if config_map.contains_key("chat_pii_redaction") {
let value = normalize_chat_pii_redaction_config(config_map.remove("chat_pii_redaction"))?;
if let Some(value) = value {
config_map.insert("chat_pii_redaction".to_string(), value);
}
}
updated.config = (!config_map.is_empty()).then_some(serde_json::Value::Object(config_map));
updated.updated_at_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)

View File

@@ -666,6 +666,7 @@ fn admin_provider_oauth_gateway_error_message(error: GatewayError) -> String {
match error {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
}
}

View File

@@ -55,6 +55,18 @@ pub(crate) const ADMIN_MODULE_DEFINITIONS: &[AdminModuleDefinition] = &[
admin_menu_group: None,
admin_menu_order: 0,
},
AdminModuleDefinition {
name: "chat_pii_redaction",
display_name: "敏感信息替换保护",
description: "发送给供应商前将聊天消息中的敏感信息替换为占位符,返回客户端前自动还原。",
category: "security",
env_key: "CHAT_PII_REDACTION_AVAILABLE",
default_available: true,
admin_route: Some("/admin/modules/chat-pii-redaction"),
admin_menu_icon: Some("ShieldCheck"),
admin_menu_group: Some("system"),
admin_menu_order: 59,
},
AdminModuleDefinition {
name: "notification_email",
display_name: "异常通知",

View File

@@ -357,6 +357,7 @@ pub(crate) fn gateway_error_message(error: GatewayError) -> String {
match error {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
}
}

View File

@@ -16,7 +16,7 @@ use crate::api::response::{
build_local_user_rpm_limited_response,
};
use crate::constants::{
DEPENDENCY_REASON_HEADER, EXECUTION_PATH_CONTROL_EXECUTE_STREAM,
CONTROL_CANDIDATE_ID_HEADER, DEPENDENCY_REASON_HEADER, EXECUTION_PATH_CONTROL_EXECUTE_STREAM,
EXECUTION_PATH_CONTROL_EXECUTE_SYNC, EXECUTION_PATH_DISTRIBUTED_OVERLOADED,
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
EXECUTION_PATH_LOCAL_AI_PUBLIC, EXECUTION_PATH_LOCAL_API_KEY_CONCURRENCY_LIMITED,
@@ -62,6 +62,7 @@ use crate::{
use axum::body::{to_bytes, Body, Bytes};
use axum::extract::{ConnectInfo, Request, State};
use axum::http::{self, header::HeaderName, header::HeaderValue, Response};
use futures_util::StreamExt;
use sha2::{Digest, Sha256};
use std::{collections::BTreeMap, time::Instant};
use tracing::{debug, info, warn};
@@ -494,6 +495,109 @@ fn collect_upstream_response_headers(
.collect()
}
fn collect_response_headers(headers: &http::HeaderMap) -> BTreeMap<String, String> {
headers
.iter()
.map(|(name, value)| {
(
name.as_str().to_string(),
value.to_str().unwrap_or_default().to_string(),
)
})
.collect()
}
fn replace_response_headers(
headers: &mut http::HeaderMap,
values: &BTreeMap<String, String>,
) -> Result<(), GatewayError> {
headers.clear();
for (name, value) in values {
headers.insert(
HeaderName::from_bytes(name.as_bytes())
.map_err(|err| GatewayError::Internal(err.to_string()))?,
HeaderValue::from_str(value).map_err(|err| GatewayError::Internal(err.to_string()))?,
);
}
Ok(())
}
fn take_redaction_session_for_response(
headers: &http::HeaderMap,
redaction_slot: &crate::privacy::RedactionSessionSlot,
) -> Option<crate::privacy::RedactionSession> {
let candidate_id = headers
.get(CONTROL_CANDIDATE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty());
redaction_slot.take_for_candidate(candidate_id)
}
async fn restore_redacted_sync_execution_response(
response: Response<Body>,
redaction_slot: &crate::privacy::RedactionSessionSlot,
) -> Result<Response<Body>, GatewayError> {
let (mut parts, body) = response.into_parts();
let Some(session) = take_redaction_session_for_response(&parts.headers, redaction_slot) else {
return Ok(Response::from_parts(parts, body));
};
let mut headers = collect_response_headers(&parts.headers);
let body_bytes = to_bytes(body, usize::MAX)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let restored =
crate::privacy::restore_sync_response_body(&mut headers, body_bytes.as_ref(), &session)?;
replace_response_headers(&mut parts.headers, &headers)?;
Ok(Response::from_parts(parts, Body::from(restored.body)))
}
fn restore_redacted_stream_execution_response(
response: Response<Body>,
redaction_slot: &crate::privacy::RedactionSessionSlot,
) -> Result<Response<Body>, GatewayError> {
let (mut parts, body) = response.into_parts();
let Some(session) = take_redaction_session_for_response(&parts.headers, redaction_slot) else {
return Ok(Response::from_parts(parts, body));
};
let headers = collect_response_headers(&parts.headers);
let _ = crate::privacy::StreamingResponseRestorer::new(&headers, &session)?;
parts.headers.remove(http::header::CONTENT_LENGTH);
let stream_headers = headers;
let stream = async_stream::stream! {
let mut restorer = match crate::privacy::StreamingResponseRestorer::new(&stream_headers, &session) {
Ok(restorer) => restorer,
Err(err) => {
yield Err(std::io::Error::other(format!("{err:?}")));
return;
}
};
let mut body_stream = body.into_data_stream();
while let Some(chunk) = body_stream.next().await {
match chunk {
Ok(chunk) => match restorer.push_chunk(chunk.as_ref()) {
Ok(restored) if restored.is_empty() => {}
Ok(restored) => yield Ok(Bytes::from(restored)),
Err(err) => {
yield Err(std::io::Error::other(format!("{err:?}")));
return;
}
},
Err(err) => {
yield Err(std::io::Error::other(err.to_string()));
return;
}
}
}
match restorer.finish() {
Ok(restored) if restored.is_empty() => {}
Ok(restored) => yield Ok(Bytes::from(restored)),
Err(err) => yield Err(std::io::Error::other(format!("{err:?}"))),
}
};
Ok(Response::from_parts(parts, Body::from_stream(stream)))
}
fn aggregate_sync_sse_response_for_client(
decision: &GatewayControlDecision,
public_path: &str,
@@ -749,6 +853,8 @@ pub(crate) async fn proxy_request(
};
let request_admission_ms = started_at.elapsed().as_millis() as u64;
let (mut parts, body) = request.into_parts();
let redaction_slot = crate::privacy::RedactionSessionSlot::default();
parts.extensions.insert(redaction_slot.clone());
parts
.extensions
.insert(request_origin_from_headers_and_remote_addr(
@@ -1181,6 +1287,10 @@ pub(crate) async fn proxy_request(
);
match stream_outcome {
LocalExecutionRequestOutcome::Responded(execution_runtime_response) => {
let execution_runtime_response = restore_redacted_stream_execution_response(
execution_runtime_response,
&redaction_slot,
)?;
state.clear_local_execution_runtime_miss_diagnostic(&trace_id);
return Ok(finalize_gateway_response_with_context(
&state,
@@ -1202,6 +1312,11 @@ pub(crate) async fn proxy_request(
.await?
{
LocalExecutionRequestOutcome::Responded(execution_runtime_response) => {
let execution_runtime_response = restore_redacted_sync_execution_response(
execution_runtime_response,
&redaction_slot,
)
.await?;
state.clear_local_execution_runtime_miss_diagnostic(&trace_id);
return Ok(finalize_gateway_response_with_context(
&state,
@@ -1229,6 +1344,10 @@ pub(crate) async fn proxy_request(
.await?
{
LocalExecutionRequestOutcome::Responded(execution_runtime_response) => {
let execution_runtime_response = restore_redacted_stream_execution_response(
execution_runtime_response,
&redaction_slot,
)?;
state.clear_local_execution_runtime_miss_diagnostic(&trace_id);
return Ok(finalize_gateway_response_with_context(
&state,
@@ -1278,6 +1397,15 @@ pub(crate) async fn proxy_request(
Some(control_execution_path),
reason,
);
let control_response = if stream_request {
restore_redacted_stream_execution_response(
control_response,
&redaction_slot,
)?
} else {
restore_redacted_sync_execution_response(control_response, &redaction_slot)
.await?
};
let mut control_response = control_response;
state.clear_local_execution_runtime_miss_diagnostic(&trace_id);
control_response.headers_mut().insert(
@@ -1778,8 +1906,109 @@ fn local_execution_runtime_miss_route_detail(
mod tests {
use super::{
diagnostic_is_auth_api_key_concurrency_limited, local_execution_runtime_miss_detail,
restore_redacted_stream_execution_response, restore_redacted_sync_execution_response,
GatewayControlDecision, LocalExecutionRuntimeMissDiagnostic,
};
use axum::body::{to_bytes, Body};
use axum::http::{header, Response};
use serde_json::json;
fn redaction_slot_for_email() -> (crate::privacy::RedactionSessionSlot, String) {
let masked = crate::privacy::mask_chat_request_json(
br#"{"messages":[{"role":"user","content":"Email alice@example.com"}]}"#,
crate::privacy::RedactionSessionConfig::new(
b"proxy-wrapper-test-key".to_vec(),
300,
600,
),
);
let sentinel = masked
.session
.sentinel_for_original("alice@example.com")
.expect("email sentinel should exist")
.to_string();
let slot = crate::privacy::RedactionSessionSlot::default();
slot.put(masked.session);
(slot, sentinel)
}
#[tokio::test]
async fn proxy_pii_redaction_sync_response_wrapper_restores_current_request_sentinel() {
let (slot, sentinel) = redaction_slot_for_email();
let body = serde_json::to_vec(&json!({
"choices": [{"message": {"role": "assistant", "content": format!("hello {sentinel}")}}]
}))
.expect("response should serialize");
let response = Response::builder()
.header(header::CONTENT_TYPE, "application/json")
.header(header::CONTENT_LENGTH, body.len().to_string())
.body(Body::from(body))
.expect("response should build");
let restored = restore_redacted_sync_execution_response(response, &slot)
.await
.expect("sync wrapper should restore");
let restored_content_length = restored
.headers()
.get(header::CONTENT_LENGTH)
.and_then(|value| value.to_str().ok())
.map(ToOwned::to_owned);
let body = to_bytes(restored.into_body(), usize::MAX)
.await
.expect("body should read");
let value: serde_json::Value = serde_json::from_slice(&body).expect("body should parse");
assert_eq!(restored_content_length, Some(body.len().to_string()));
assert_eq!(
value["choices"][0]["message"]["content"],
"hello alice@example.com"
);
}
#[tokio::test]
async fn proxy_pii_redaction_stream_response_wrapper_restores_current_request_sentinel() {
let (slot, sentinel) = redaction_slot_for_email();
let response = Response::builder()
.header(header::CONTENT_TYPE, "text/event-stream")
.header(header::CONTENT_LENGTH, "999")
.body(Body::from(format!(
"data: {{\"choices\":[{{\"delta\":{{\"content\":\"hello {sentinel}\"}}}}]}}\n\n"
)))
.expect("response should build");
let restored = restore_redacted_stream_execution_response(response, &slot)
.expect("stream wrapper should restore");
assert!(restored.headers().get(header::CONTENT_LENGTH).is_none());
let body = to_bytes(restored.into_body(), usize::MAX)
.await
.expect("body should read");
let text = String::from_utf8(body.to_vec()).expect("body should be utf8");
assert!(text.contains("hello alice@example.com"));
assert!(!text.contains(&sentinel));
}
#[tokio::test]
async fn proxy_pii_redaction_compressed_response_safe_error() {
let (slot, sentinel) = redaction_slot_for_email();
let response = Response::builder()
.header(header::CONTENT_TYPE, "application/json")
.header(header::CONTENT_ENCODING, "gzip")
.header(header::CONTENT_LENGTH, "999")
.body(Body::from(format!(
"{{\"choices\":[{{\"message\":{{\"content\":\"hello {sentinel}\"}}}}]}}"
)))
.expect("response should build");
let err = restore_redacted_sync_execution_response(response, &slot)
.await
.expect_err("compressed active redaction should fail safely");
let message = format!("{err:?}");
assert!(message.contains("encoded response bodies"));
assert!(!message.contains("alice@example.com"));
assert!(!message.contains(&sentinel));
}
#[test]
fn runtime_miss_detail_returns_model_specific_stream_message_when_candidates_are_unavailable() {

View File

@@ -138,6 +138,7 @@ pub(super) fn announcements_internal_detail(err: GatewayError) -> String {
match err {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
}
}