feat(gateway): add reversible chat pii redaction

This commit is contained in:
Kayphoon
2026-05-13 18:25:13 +08:00
parent 3c2497f019
commit 2958041dc7
43 changed files with 7770 additions and 126 deletions

View File

@@ -12,7 +12,7 @@ use crate::ai_serving::transport::{
};
use crate::{
append_execution_contract_fields_to_value, append_local_failover_policy_to_value,
AiExecutionDecision, AppState,
AiExecutionDecision, AppState, GatewayError,
};
use super::request::resolve_local_openai_chat_candidate_payload_parts;
@@ -29,7 +29,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
decision_kind: &str,
report_kind: &str,
upstream_is_stream: bool,
) -> Option<AiExecutionDecision> {
) -> Result<Option<AiExecutionDecision>, GatewayError> {
let decision_is_stream = decision_kind == OPENAI_CHAT_STREAM_PLAN_KIND;
let attempt_identity = attempt.attempt_identity();
let LocalOpenAiChatCandidateAttempt {
@@ -38,7 +38,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
candidate_id,
..
} = attempt;
let resolved = resolve_local_openai_chat_candidate_payload_parts(
let Some(resolved) = resolve_local_openai_chat_candidate_payload_parts(
state,
parts,
trace_id,
@@ -51,7 +51,10 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
report_kind,
upstream_is_stream,
)
.await?;
.await?
else {
return Ok(None);
};
let candidate = &eligible.candidate;
let prompt_cache_key = resolved
@@ -82,60 +85,6 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
&mut extra_fields,
resolved.transport.provider.provider_type.as_str(),
);
let report_context = append_local_failover_policy_to_value(
append_execution_contract_fields_to_value(
build_local_execution_report_context(LocalExecutionReportContextParts {
auth_context: &input.auth_context,
request_id: trace_id,
candidate_id: &candidate_id,
attempt_identity,
model: &input.requested_model,
provider_name: &resolved.transport.provider.name,
provider_id: &candidate.provider_id,
endpoint_id: &candidate.endpoint_id,
key_id: &candidate.key_id,
key_name: Some(&candidate.key_name),
model_id: Some(&candidate.model_id),
global_model_id: Some(&candidate.global_model_id),
global_model_name: Some(&candidate.global_model_name),
provider_api_format: &resolved.provider_api_format,
client_api_format: "openai:chat",
mapped_model: Some(&resolved.mapped_model),
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
pool_key_lease: eligible.orchestration.pool_key_lease.as_ref(),
ranking: eligible.ranking.as_ref(),
upstream_url: Some(&resolved.upstream_url),
header_rules: resolved.transport.endpoint.header_rules.as_ref(),
body_rules: resolved.transport.endpoint.body_rules.as_ref(),
provider_request_method: Some(serde_json::Value::Null),
provider_request_headers: Some(&resolved.provider_request_headers),
original_headers: &parts.headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: None,
client_session_affinity: input.client_session_affinity.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
client_requested_stream: body_json
.get("stream")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false),
upstream_is_stream,
has_envelope: resolved.envelope_name.is_some(),
needs_conversion: matches!(
resolved.conversion_mode,
crate::ai_serving::ConversionMode::Bidirectional
),
extra_fields,
}),
resolved.execution_strategy,
resolved.conversion_mode,
"openai:chat",
candidate.endpoint_api_format.as_str(),
),
&resolved.transport,
);
let super::request::LocalOpenAiChatCandidatePayloadParts {
auth_header,
auth_value,
@@ -147,11 +96,71 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
execution_strategy,
conversion_mode,
report_kind,
envelope_name: _,
envelope_name,
transport,
request_redacted,
} = resolved;
let original_request_body_json = if request_redacted {
Some(&provider_request_body)
} else {
Some(body_json)
};
let report_context = append_local_failover_policy_to_value(
append_execution_contract_fields_to_value(
build_local_execution_report_context(LocalExecutionReportContextParts {
auth_context: &input.auth_context,
request_id: trace_id,
candidate_id: &candidate_id,
attempt_identity,
model: &input.requested_model,
provider_name: &transport.provider.name,
provider_id: &candidate.provider_id,
endpoint_id: &candidate.endpoint_id,
key_id: &candidate.key_id,
key_name: Some(&candidate.key_name),
model_id: Some(&candidate.model_id),
global_model_id: Some(&candidate.global_model_id),
global_model_name: Some(&candidate.global_model_name),
provider_api_format: &provider_api_format,
client_api_format: "openai:chat",
mapped_model: Some(&mapped_model),
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
pool_key_lease: eligible.orchestration.pool_key_lease.as_ref(),
ranking: eligible.ranking.as_ref(),
upstream_url: Some(&upstream_url),
header_rules: transport.endpoint.header_rules.as_ref(),
body_rules: transport.endpoint.body_rules.as_ref(),
provider_request_method: Some(serde_json::Value::Null),
provider_request_headers: Some(&provider_request_headers),
original_headers: &parts.headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json,
original_request_body_base64: None,
client_session_affinity: input.client_session_affinity.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
client_requested_stream: body_json
.get("stream")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false),
upstream_is_stream,
has_envelope: envelope_name.is_some(),
needs_conversion: matches!(
conversion_mode,
crate::ai_serving::ConversionMode::Bidirectional
),
extra_fields,
}),
execution_strategy,
conversion_mode,
"openai:chat",
candidate.endpoint_api_format.as_str(),
),
&transport,
);
Some(build_ai_execution_decision_response(
Ok(Some(build_ai_execution_decision_response(
AiExecutionDecisionResponseParts {
decision_is_stream,
decision_kind: decision_kind.to_string(),
@@ -185,5 +194,5 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
report_context: Some(report_context),
auth_context: input.auth_context.clone(),
},
))
)))
}

View File

@@ -1,5 +1,7 @@
use std::borrow::Cow;
use std::collections::BTreeMap;
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use serde_json::Value;
@@ -34,7 +36,13 @@ use crate::ai_serving::{
LocalResolvedOAuthRequestAuth,
};
use crate::ai_serving::{ConversionMode, ExecutionStrategy};
use crate::AppState;
use crate::privacy::{
build_redaction_session_config, provider_chat_pii_redaction_enabled,
read_chat_pii_redaction_runtime_config, try_mask_chat_request_json_with_cache_options,
MaskChatRequestOptions, RedactionMaskError, RedactionSessionSlot, RedisRedactionMappingCache,
};
use crate::{AppState, GatewayError};
use tracing::warn;
use super::support::{
mark_skipped_local_openai_chat_candidate,
@@ -55,6 +63,30 @@ pub(crate) struct LocalOpenAiChatCandidatePayloadParts {
pub(super) report_kind: String,
pub(super) envelope_name: Option<&'static str>,
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
pub(super) request_redacted: bool,
}
fn request_identity_response_encoding_when_redacted(
headers: &mut BTreeMap<String, String>,
redacted: bool,
) {
if redacted {
headers.insert("accept-encoding".to_string(), "identity".to_string());
}
}
struct ProviderChatRequestRedaction<'a> {
body_json: Cow<'a, Value>,
redacted: bool,
}
impl<'a> ProviderChatRequestRedaction<'a> {
fn disabled(body_json: &'a Value, _parts: &http::request::Parts) -> Self {
Self {
body_json: Cow::Borrowed(body_json),
redacted: false,
}
}
}
#[allow(clippy::too_many_arguments)]
@@ -70,7 +102,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
decision_kind: &str,
report_kind: &str,
upstream_is_stream: bool,
) -> Option<LocalOpenAiChatCandidatePayloadParts> {
) -> Result<Option<LocalOpenAiChatCandidatePayloadParts>, GatewayError> {
let planner_state = crate::ai_serving::PlannerAppState::new(state);
let candidate = &eligible.candidate;
let provider_api_format = eligible.provider_api_format.as_str();
@@ -84,6 +116,10 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
Some(&input.requested_model),
)
.await;
let redaction =
resolve_provider_chat_request_redaction(state, parts, body_json, transport, candidate_id)
.await?;
let body_json = redaction.body_json.as_ref();
if provider_api_format == "openai:chat" {
if let Some(skip_reason) = local_openai_chat_transport_unsupported_reason(transport) {
@@ -97,8 +133,8 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
skip_reason,
)
.await;
return None;
}
return Ok(None);
};
let prepared_candidate = match prepare_header_authenticated_candidate(
planner_state,
@@ -125,7 +161,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
skip_reason,
)
.await;
return None;
return Ok(None);
}
};
@@ -153,7 +189,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
),
)
.await;
return None;
return Ok(None);
};
let Some(upstream_url) = build_local_openai_chat_upstream_url(parts, transport) else {
@@ -172,7 +208,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
),
)
.await;
return None;
return Ok(None);
};
let Some(resolved_headers) =
@@ -205,7 +241,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
),
)
.await;
return None;
return Ok(None);
};
let mut provider_request_headers = resolved_headers.headers;
apply_codex_openai_responses_special_headers(
@@ -217,7 +253,6 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
Some(trace_id),
transport.key.decrypted_auth_config.as_deref(),
);
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats("openai:chat", "openai:chat");
let resolved_report_kind =
@@ -227,7 +262,12 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
"openai_chat_sync_finalize".to_string()
};
return Some(LocalOpenAiChatCandidatePayloadParts {
request_identity_response_encoding_when_redacted(
&mut provider_request_headers,
redaction.redacted,
);
return Ok(Some(LocalOpenAiChatCandidatePayloadParts {
auth_header: resolved_headers.auth_header,
auth_value: resolved_headers.auth_value,
mapped_model: prepared_candidate.mapped_model,
@@ -240,8 +280,9 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
report_kind: resolved_report_kind,
envelope_name: None,
transport: Arc::clone(transport),
});
}
request_redacted: redaction.redacted,
}));
};
let provider_api_format = provider_api_format.trim().to_ascii_lowercase();
let Some(conversion_kind) =
@@ -257,7 +298,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
"transport_api_format_unsupported",
)
.await;
return None;
return Ok(None);
};
if let Some(skip_reason) = crate::ai_serving::request_conversion_transport_unsupported_reason(
transport,
@@ -273,7 +314,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
skip_reason,
)
.await;
return None;
return Ok(None);
}
let is_kiro_claude_cli =
is_kiro_claude_messages_transport(transport, provider_api_format.as_str());
@@ -302,7 +343,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
"transport_auth_unavailable",
)
.await;
return None;
return Ok(None);
}
}
} else {
@@ -326,7 +367,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
skip_reason,
)
.await;
return None;
return Ok(None);
}
}
} else {
@@ -351,7 +392,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
skip_reason,
)
.await;
return None;
return Ok(None);
}
}
};
@@ -387,7 +428,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
),
)
.await;
return None;
return Ok(None);
};
if let Some(mapping) =
crate::system_features::reasoning_model_directive_mapping_for_api_format_and_model(
@@ -412,7 +453,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
}
if let Some(kiro_auth) = kiro_auth.as_ref() {
return build_kiro_openai_chat_cross_format_payload_parts(
return Ok(build_kiro_openai_chat_cross_format_payload_parts(
state,
parts,
trace_id,
@@ -430,8 +471,9 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
provider_request_body,
upstream_is_stream,
kiro_auth,
redaction.redacted,
)
.await;
.await);
}
let Some(upstream_url) = build_cross_format_openai_chat_upstream_url(
@@ -456,7 +498,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
),
)
.await;
return None;
return Ok(None);
};
let Some(resolved_headers) =
build_standard_provider_request_headers(StandardProviderRequestHeadersInput {
@@ -488,7 +530,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
),
)
.await;
return None;
return Ok(None);
};
let mut provider_request_headers = resolved_headers.headers;
apply_codex_openai_responses_special_headers(
@@ -500,6 +542,10 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
Some(trace_id),
transport.key.decrypted_auth_config.as_deref(),
);
request_identity_response_encoding_when_redacted(
&mut provider_request_headers,
redaction.redacted,
);
let resolved_report_kind = if decision_kind == OPENAI_CHAT_STREAM_PLAN_KIND {
"openai_chat_stream_success".to_string()
@@ -509,7 +555,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats("openai:chat", provider_api_format.as_str());
Some(LocalOpenAiChatCandidatePayloadParts {
Ok(Some(LocalOpenAiChatCandidatePayloadParts {
auth_header: resolved_headers.auth_header,
auth_value: resolved_headers.auth_value,
mapped_model: prepared_candidate.mapped_model,
@@ -522,7 +568,8 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
report_kind: resolved_report_kind,
envelope_name: None,
transport: Arc::clone(transport),
})
request_redacted: redaction.redacted,
}))
}
#[allow(clippy::too_many_arguments)]
@@ -544,6 +591,7 @@ async fn build_kiro_openai_chat_cross_format_payload_parts(
claude_request_body: Value,
upstream_is_stream: bool,
kiro_auth: &KiroRequestAuth,
request_redacted: bool,
) -> Option<LocalOpenAiChatCandidatePayloadParts> {
let candidate = &eligible.candidate;
let provider_request_body = match build_kiro_provider_request_body(
@@ -601,7 +649,7 @@ async fn build_kiro_openai_chat_cross_format_payload_parts(
return None;
}
};
let provider_request_headers = match build_kiro_provider_headers(KiroProviderHeadersInput {
let mut provider_request_headers = match build_kiro_provider_headers(KiroProviderHeadersInput {
headers: &parts.headers,
provider_request_body: &provider_request_body,
original_request_body: original_body_json,
@@ -631,6 +679,10 @@ async fn build_kiro_openai_chat_cross_format_payload_parts(
return None;
}
};
request_identity_response_encoding_when_redacted(
&mut provider_request_headers,
request_redacted,
);
let resolved_report_kind = if decision_kind == OPENAI_CHAT_STREAM_PLAN_KIND {
"openai_chat_stream_success".to_string()
} else {
@@ -652,5 +704,86 @@ async fn build_kiro_openai_chat_cross_format_payload_parts(
report_kind: resolved_report_kind,
envelope_name: Some(KIRO_ENVELOPE_NAME),
transport: Arc::clone(transport),
request_redacted,
})
}
async fn resolve_provider_chat_request_redaction<'a>(
state: &AppState,
parts: &http::request::Parts,
body_json: &'a Value,
transport: &GatewayProviderTransportSnapshot,
candidate_id: &str,
) -> Result<ProviderChatRequestRedaction<'a>, GatewayError> {
if parts.uri.path() != "/v1/chat/completions" {
return Ok(ProviderChatRequestRedaction::disabled(body_json, parts));
}
let Some(slot) = parts.extensions.get::<RedactionSessionSlot>() else {
return Ok(ProviderChatRequestRedaction::disabled(body_json, parts));
};
let runtime_config = read_chat_pii_redaction_runtime_config(state)
.await
.map_err(|err| {
warn!(
error = ?err,
"gateway failed to read chat pii redaction runtime config"
);
GatewayError::Internal("chat pii redaction setup failed".to_string())
})?;
if !provider_chat_pii_redaction_enabled(transport.provider.config.as_ref(), &runtime_config) {
return Ok(ProviderChatRequestRedaction::disabled(body_json, parts));
}
let Some(hmac_key) = state.encryption_key().map(str::as_bytes).map(Vec::from) else {
warn!("gateway chat pii redaction is enabled but encryption key is unavailable");
return Err(GatewayError::Internal(
"chat pii redaction setup failed".to_string(),
));
};
let body_bytes = serde_json::to_vec(body_json).map_err(|err| {
warn!(
error = ?err,
"gateway failed to serialize provider chat pii redaction body"
);
GatewayError::Internal("chat pii redaction setup failed".to_string())
})?;
let now_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
let cache = RedisRedactionMappingCache::new(state.runtime_state.as_ref());
let masked = try_mask_chat_request_json_with_cache_options(
&body_bytes,
build_redaction_session_config(hmac_key, &runtime_config, now_unix_secs),
MaskChatRequestOptions::runtime(runtime_config.inject_model_instruction),
Some(&cache),
)
.await
.map_err(redaction_mask_error_to_gateway_error)?;
if !masked.redacted {
return Ok(ProviderChatRequestRedaction {
body_json: Cow::Borrowed(body_json),
redacted: false,
});
}
let masked_body_json = serde_json::from_slice::<Value>(&masked.body).map_err(|err| {
warn!(
error = ?err,
"gateway failed to decode redacted provider chat pii body"
);
GatewayError::Internal("chat pii redaction setup failed".to_string())
})?;
slot.put_for_candidate(candidate_id, masked.session);
Ok(ProviderChatRequestRedaction {
body_json: Cow::Owned(masked_body_json),
redacted: true,
})
}
fn redaction_mask_error_to_gateway_error(error: RedactionMaskError) -> GatewayError {
match error {
RedactionMaskError::Limit(limit) => GatewayError::Client {
status: limit.client_status(),
message: limit.safe_message().to_string(),
},
}
}

View File

@@ -163,7 +163,7 @@ pub(crate) async fn maybe_build_sync_local_decision_payload(
"openai_chat_sync_success",
upstream_is_stream,
)
.await
.await?
{
return Ok(Some(payload));
}
@@ -214,7 +214,7 @@ pub(crate) async fn maybe_build_stream_local_decision_payload(
"openai_chat_stream_success",
upstream_is_stream,
)
.await
.await?
{
return Ok(Some(payload));
}

View File

@@ -134,7 +134,7 @@ impl LocalOpenAiChatStreamAttemptSource<'_> {
"openai_chat_stream_success",
upstream_is_stream,
)
.await
.await?
else {
return Ok(None);
};

View File

@@ -134,7 +134,7 @@ impl LocalOpenAiChatSyncAttemptSource<'_> {
"openai_chat_sync_success",
upstream_is_stream,
)
.await
.await?
else {
return Ok(None);
};

View File

@@ -13,6 +13,7 @@ use crate::insert_header_if_missing;
pub(crate) enum GatewayError {
UpstreamUnavailable { trace_id: String, message: String },
ControlUnavailable { trace_id: String, message: String },
Client { status: StatusCode, message: String },
Internal(String),
}
@@ -55,6 +56,15 @@ impl IntoResponse for GatewayError {
);
response
}
Self::Client { status, message } => (
status,
Json(json!({
"error": {
"message": message,
}
})),
)
.into_response(),
Self::Internal(message) => (
StatusCode::INTERNAL_SERVER_ERROR,
Json(json!({

View File

@@ -19,6 +19,7 @@ use crate::executor::{build_local_execution_exhaustion, LocalExecutionRequestOut
use crate::handlers::shared::provider_pool::release_admin_provider_pool_key_lease;
use crate::log_ids::short_request_id;
use crate::orchestration::local_execution_candidate_metadata_from_report_context;
use crate::privacy::RedactionExecutionCandidateId;
use crate::request_candidate_runtime::{
record_local_request_candidate_status, RequestCandidateRuntimeWriter,
};
@@ -26,6 +27,17 @@ use crate::{AppState, GatewayError};
const DEFAULT_STREAM_CANDIDATE_WATCHDOG_TIMEOUT_MS: u64 = 300_000;
fn attach_redaction_execution_candidate(response: &mut Response<Body>, candidate_id: Option<&str>) {
if let Some(candidate_id) = candidate_id
.map(str::trim)
.filter(|value| !value.is_empty())
{
response
.extensions_mut()
.insert(RedactionExecutionCandidateId::new(candidate_id));
}
}
pub(crate) async fn execute_sync_plan_and_reports<T>(
state: &AppState,
parts: &http::request::Parts,
@@ -136,7 +148,7 @@ where
type Error = GatewayError;
async fn execute_attempt(&self, attempt: &T) -> Result<Option<Self::Response>, Self::Error> {
execute_execution_runtime_sync(
let mut response = execute_execution_runtime_sync(
self.state,
self.parts.uri.path(),
attempt.execution_plan().clone(),
@@ -146,7 +158,14 @@ where
attempt.report_kind(),
attempt.report_context(),
)
.await
.await?;
if let Some(response) = response.as_mut() {
attach_redaction_execution_candidate(
response,
attempt.execution_plan().candidate_id.as_deref(),
);
}
Ok(response)
}
async fn mark_unused_attempts(&self, attempts: Vec<T>) -> Result<(), Self::Error> {
@@ -346,7 +365,7 @@ where
let execution_plan_kind = self.plan_kind.to_string();
let execution_decision = self.decision.clone();
let execution_report_kind = attempt.report_kind();
execute_stream_candidate_with_watchdog(
let mut response = execute_stream_candidate_with_watchdog(
self.state,
self.trace_id,
self.plan_kind,
@@ -365,7 +384,11 @@ where
.await
},
)
.await
.await?;
if let Some(response) = response.as_mut() {
attach_redaction_execution_candidate(response, watchdog_plan.candidate_id.as_deref());
}
Ok(response)
}
async fn mark_unused_attempts(&self, attempts: Vec<T>) -> Result<(), Self::Error> {

View File

@@ -191,6 +191,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,
}
}

View File

@@ -51,6 +51,7 @@ pub(crate) mod middleware;
mod model_fetch;
mod oauth;
mod orchestration;
mod privacy;
mod provider_key_auth;
pub(crate) use aether_provider_transport as provider_transport;
mod rate_limit;

View File

@@ -160,6 +160,7 @@ fn gateway_error_to_oauth_error(error: GatewayError) -> OAuthError {
match error {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => OAuthError::Transport(message),
}
}

File diff suppressed because it is too large Load Diff

View File

@@ -63,6 +63,7 @@ impl provider_transport::VideoTaskTransportSnapshotLookup for AppState {
.map_err(|err| match err {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
})
}
@@ -79,6 +80,7 @@ impl ModelFetchTransportRuntime for AppState {
.map_err(|err| match err {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
})
}
@@ -100,6 +102,7 @@ impl ModelFetchTransportRuntime for AppState {
.map_err(|err| match err {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
})
}

View File

@@ -1414,6 +1414,7 @@ impl AppState {
message: match err {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
},
},

View File

@@ -26,3 +26,4 @@ use super::{
mod decision;
mod image;
mod pii_redaction;

View File

@@ -0,0 +1,340 @@
use super::*;
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::auth::{
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
};
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};
use aether_data_contracts::repository::candidates::{
RequestCandidateReadRepository, RequestCandidateStatus,
};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use sha2::{Digest, Sha256};
#[derive(Debug, Clone)]
struct SeenProviderStreamRequest {
body: serde_json::Value,
authorization: String,
accept_encoding: String,
accept: String,
}
fn hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());
format!("{:x}", hasher.finalize())
}
fn auth_snapshot() -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
"user-ai-execute-stream-pii-redaction".to_string(),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-5"])),
"api-key-ai-execute-stream-pii-redaction".to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(4_102_444_800),
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-5"])),
)
.expect("auth snapshot should build")
}
fn candidate_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-ai-execute-stream-pii-redaction".to_string(),
provider_name: "openai".to_string(),
provider_type: "custom".to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: "endpoint-ai-execute-stream-pii-redaction".to_string(),
endpoint_api_format: "openai:chat".to_string(),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("chat".to_string()),
endpoint_is_active: true,
key_id: "key-ai-execute-stream-pii-redaction".to_string(),
key_name: "prod".to_string(),
key_auth_type: "api_key".to_string(),
key_is_active: true,
key_api_formats: Some(vec!["openai:chat".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 5,
key_global_priority_by_format: Some(serde_json::json!({"openai:chat": 1})),
model_id: "model-ai-execute-stream-pii-redaction".to_string(),
global_model_id: "global-model-ai-execute-stream-pii-redaction".to_string(),
global_model_name: "gpt-5".to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(true),
model_provider_model_name: "gpt-5-upstream".to_string(),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: "gpt-5-upstream".to_string(),
priority: 1,
api_formats: Some(vec!["openai:chat".to_string()]),
endpoint_ids: Some(vec!["endpoint-ai-execute-stream-pii-redaction".to_string()]),
}]),
model_supports_streaming: Some(true),
model_is_active: true,
model_is_available: true,
}
}
fn provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-ai-execute-stream-pii-redaction".to_string(),
"openai".to_string(),
Some("https://example.com".to_string()),
"custom".to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
true,
None,
Some(2),
None,
Some(20.0),
None,
Some(serde_json::json!({"chat_pii_redaction": {"enabled": true}})),
)
}
fn endpoint(base_url: String) -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-ai-execute-stream-pii-redaction".to_string(),
"provider-ai-execute-stream-pii-redaction".to_string(),
"openai:chat".to_string(),
Some("openai".to_string()),
Some("chat".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(base_url, None, None, Some(2), None, None, None, None)
.expect("endpoint transport should build")
}
fn key() -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
"key-ai-execute-stream-pii-redaction".to_string(),
"provider-ai-execute-stream-pii-redaction".to_string(),
"prod".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["openai:chat"])),
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-upstream-stream-pii")
.expect("api key should encrypt"),
None,
None,
Some(serde_json::json!({"openai:chat": 1})),
None,
None,
None,
None,
)
.expect("key transport should build")
}
fn collect_email_sentinel(text: &str) -> String {
let start = text
.find("<AETHER:EMAIL:")
.expect("email sentinel should exist");
let end = text[start..]
.find('>')
.map(|index| start + index + 1)
.expect("sentinel should close");
text[start..end].to_string()
}
#[tokio::test]
async fn ai_execute_stream_pii_redaction_round_trip() {
let seen_provider_request = Arc::new(Mutex::new(None::<SeenProviderStreamRequest>));
let seen_provider_request_clone = Arc::clone(&seen_provider_request);
let provider_app = Router::new().route(
"/v1/chat/completions",
any(move |request: Request| {
let seen_provider_request_inner = Arc::clone(&seen_provider_request_clone);
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).expect("provider payload should parse");
let payload_text = serde_json::to_string(&payload).expect("payload should serialize");
let sentinel = collect_email_sentinel(&payload_text);
*seen_provider_request_inner
.lock()
.expect("mutex should lock") = Some(SeenProviderStreamRequest {
body: payload,
authorization: parts
.headers
.get(http::header::AUTHORIZATION)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
accept_encoding: parts
.headers
.get(http::header::ACCEPT_ENCODING)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
accept: parts
.headers
.get(http::header::ACCEPT)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
});
let split_at = sentinel.len() / 2;
let (first_half, second_half) = sentinel.split_at(split_at);
let first_chunk = format!(
"data: {{\"id\":\"chatcmpl-stream-pii\",\"object\":\"chat.completion.chunk\",\"model\":\"gpt-5-upstream\",\"choices\":[{{\"index\":0,\"delta\":{{\"role\":\"assistant\",\"content\":\"stream {first_half}"
);
let second_chunk = format!(
"{second_half} restored\"}},\"finish_reason\":null}}]}}\n\n"
);
let stream = futures_util::stream::iter([
Ok::<_, Infallible>(Bytes::from(first_chunk)),
Ok::<_, Infallible>(Bytes::from(second_chunk)),
Ok::<_, Infallible>(Bytes::from_static(b"data: [DONE]\n\n")),
]);
let mut response = Response::builder()
.status(StatusCode::OK)
.body(Body::from_stream(stream))
.expect("response should build");
response.headers_mut().insert(
http::header::CONTENT_TYPE,
HeaderValue::from_static("text/event-stream"),
);
response
}
}),
);
let (provider_url, provider_handle) = start_server(provider_app).await;
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-ai-execute-stream-pii-redaction")),
auth_snapshot(),
)]));
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
candidate_row(),
]));
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider()],
vec![endpoint(provider_url)],
vec![key()],
));
let data_state = crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
DEVELOPMENT_ENCRYPTION_KEY,
)
.with_system_config_values_for_tests(vec![
("module.chat_pii_redaction.enabled".to_string(), json!(true)),
(
"module.chat_pii_redaction.provider_scope".to_string(),
json!("selected_providers"),
),
(
"module.chat_pii_redaction.entities".to_string(),
json!(["email"]),
),
(
"module.chat_pii_redaction.cache_ttl_seconds".to_string(),
json!(300),
),
(
"module.chat_pii_redaction.inject_model_instruction".to_string(),
json!(true),
),
]);
let gateway_state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(data_state);
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
"Bearer sk-client-ai-execute-stream-pii-redaction",
)
.header(http::header::ACCEPT_ENCODING, "gzip")
.header(TRACE_ID_HEADER, "trace-ai-execute-stream-pii-redaction")
.body(
json!({
"model": "gpt-5",
"messages": [{"role": "user", "content": "Email stream.user@example.com"}],
"stream": true
})
.to_string(),
)
.send()
.await
.expect("request should succeed");
let status = response.status();
let execution_path = response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok())
.map(ToOwned::to_owned);
let response_text = response.text().await.expect("body should read");
assert_eq!(status, StatusCode::OK, "{response_text}");
assert_eq!(
execution_path.as_deref(),
Some(EXECUTION_PATH_EXECUTION_RUNTIME_STREAM)
);
assert!(response_text.contains("stream stream.user@example.com restored"));
assert!(response_text.contains("data: [DONE]"));
assert!(!response_text.contains("<AETHER:EMAIL:"));
let seen = seen_provider_request
.lock()
.expect("mutex should lock")
.clone()
.expect("provider stream request should be captured");
assert_eq!(seen.authorization, "Bearer sk-upstream-stream-pii");
assert_eq!(seen.accept, "text/event-stream");
assert_eq!(seen.accept_encoding, "identity");
let provider_body_text = serde_json::to_string(&seen.body).expect("body should serialize");
assert!(!provider_body_text.contains("stream.user@example.com"));
assert!(provider_body_text.contains("<AETHER:EMAIL:"));
let stored_candidates = request_candidate_repository
.list_by_request_id("trace-ai-execute-stream-pii-redaction")
.await
.expect("request candidate trace should read");
assert_eq!(stored_candidates.len(), 1);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
gateway_handle.abort();
provider_handle.abort();
}

View File

@@ -10,6 +10,315 @@ use super::{
TRACE_ID_HEADER,
};
#[tokio::test]
async fn proxy_pii_redaction_local_openai_chat_runtime_masks_headers_and_restores_sync_response() {
#[derive(Debug, Clone)]
struct SeenProviderRequest {
body: serde_json::Value,
authorization: String,
accept_encoding: String,
}
fn hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());
format!("{:x}", hasher.finalize())
}
fn auth_snapshot() -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
"user-redaction-1".to_string(),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-5"])),
"api-key-redaction-1".to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(4_102_444_800),
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-5"])),
)
.expect("auth snapshot should build")
}
fn candidate_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-redaction-1".to_string(),
provider_name: "openai".to_string(),
provider_type: "custom".to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: "endpoint-redaction-1".to_string(),
endpoint_api_format: "openai:chat".to_string(),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("chat".to_string()),
endpoint_is_active: true,
key_id: "key-redaction-1".to_string(),
key_name: "prod".to_string(),
key_auth_type: "api_key".to_string(),
key_is_active: true,
key_api_formats: Some(vec!["openai:chat".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 5,
key_global_priority_by_format: Some(serde_json::json!({"openai:chat": 1})),
model_id: "model-redaction-1".to_string(),
global_model_id: "global-model-redaction-1".to_string(),
global_model_name: "gpt-5".to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(true),
model_provider_model_name: "gpt-5-upstream".to_string(),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: "gpt-5-upstream".to_string(),
priority: 1,
api_formats: Some(vec!["openai:chat".to_string()]),
endpoint_ids: Some(vec!["endpoint-redaction-1".to_string()]),
}]),
model_supports_streaming: Some(true),
model_is_active: true,
model_is_available: true,
}
}
fn provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-redaction-1".to_string(),
"openai".to_string(),
Some("https://example.com".to_string()),
"custom".to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
true,
None,
Some(2),
None,
Some(20.0),
None,
Some(serde_json::json!({"chat_pii_redaction": {"enabled": true}})),
)
}
fn endpoint(base_url: String) -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-redaction-1".to_string(),
"provider-redaction-1".to_string(),
"openai:chat".to_string(),
Some("openai".to_string()),
Some("chat".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(base_url, None, None, Some(2), None, None, None, None)
.expect("endpoint transport should build")
}
fn key() -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
"key-redaction-1".to_string(),
"provider-redaction-1".to_string(),
"prod".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["openai:chat"])),
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-upstream-redaction")
.expect("api key should encrypt"),
None,
None,
Some(serde_json::json!({"openai:chat": 1})),
None,
None,
None,
None,
)
.expect("key transport should build")
}
fn sentinel_from_body(body: &serde_json::Value) -> String {
let messages = body["messages"].as_array().expect("messages should exist");
let user_content = messages
.iter()
.find(|message| message["role"] == "user")
.and_then(|message| message["content"].as_str())
.expect("user content should be text");
let start = user_content
.find("<AETHER:EMAIL:")
.expect("redacted content should include email sentinel");
let end = user_content[start..]
.find('>')
.map(|index| start + index + 1)
.expect("sentinel should close");
user_content[start..end].to_string()
}
let seen_provider_request = Arc::new(Mutex::new(None::<SeenProviderRequest>));
let seen_provider_request_clone = Arc::clone(&seen_provider_request);
let provider_app = Router::new().route(
"/v1/chat/completions",
any(move |request: Request| {
let seen_provider_request_inner = Arc::clone(&seen_provider_request_clone);
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).expect("provider payload should parse");
let sentinel = sentinel_from_body(&payload);
*seen_provider_request_inner
.lock()
.expect("mutex should lock") = Some(SeenProviderRequest {
body: payload,
authorization: parts
.headers
.get(http::header::AUTHORIZATION)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
accept_encoding: parts
.headers
.get(http::header::ACCEPT_ENCODING)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
});
Json(json!({
"id": "chatcmpl-redaction-1",
"object": "chat.completion",
"model": "gpt-5-upstream",
"choices": [{
"index": 0,
"message": {"role": "assistant", "content": format!("restored {sentinel}")},
"finish_reason": "stop"
}],
"usage": {"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5}
}))
}
}),
);
let (provider_url, provider_handle) = start_server(provider_app).await;
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-redaction")),
auth_snapshot(),
)]));
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
candidate_row(),
]));
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider()],
vec![endpoint(provider_url)],
vec![key()],
));
let data_state = crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
DEVELOPMENT_ENCRYPTION_KEY,
)
.with_system_config_values_for_tests(vec![
("module.chat_pii_redaction.enabled".to_string(), json!(true)),
(
"module.chat_pii_redaction.provider_scope".to_string(),
json!("selected_providers"),
),
(
"module.chat_pii_redaction.entities".to_string(),
json!(["email"]),
),
(
"module.chat_pii_redaction.cache_ttl_seconds".to_string(),
json!(300),
),
(
"module.chat_pii_redaction.inject_model_instruction".to_string(),
json!(true),
),
]);
let gateway_state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(data_state);
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(http::header::AUTHORIZATION, "Bearer sk-client-redaction")
.header(http::header::ACCEPT_ENCODING, "gzip")
.header(TRACE_ID_HEADER, "trace-proxy-pii-redaction-sync")
.body(
r#"{"model":"gpt-5","messages":[{"role":"user","content":"Email alice@example.com"}]}"#,
)
.send()
.await
.expect("request should succeed");
let status = response.status();
let execution_path = response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok())
.map(ToOwned::to_owned);
let response_text = response.text().await.expect("body should read");
assert_eq!(status, StatusCode::OK, "{response_text}");
assert_eq!(
execution_path.as_deref(),
Some(EXECUTION_PATH_EXECUTION_RUNTIME_SYNC)
);
let response_json: serde_json::Value =
serde_json::from_str(&response_text).expect("body should parse");
assert_eq!(
response_json["choices"][0]["message"]["content"],
"restored alice@example.com"
);
let seen = seen_provider_request
.lock()
.expect("mutex should lock")
.clone()
.expect("provider request should be captured");
assert_eq!(seen.authorization, "Bearer sk-upstream-redaction");
assert_eq!(seen.accept_encoding, "identity");
let provider_body_text = serde_json::to_string(&seen.body).expect("body should serialize");
assert!(!provider_body_text.contains("alice@example.com"));
assert!(provider_body_text.contains("<AETHER:EMAIL:"));
assert_eq!(seen.body["messages"][0]["role"], "assistant");
let notice = seen.body["messages"][0]["content"]
.as_str()
.expect("notice should be text");
assert!(notice.contains("not a user request"));
assert!(notice.contains("do not answer"));
assert_eq!(seen.body["messages"][1]["role"], "user");
let stored_candidates = request_candidate_repository
.list_by_request_id("trace-proxy-pii-redaction-sync")
.await
.expect("request candidate trace should read");
assert_eq!(stored_candidates.len(), 1);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
gateway_handle.abort();
provider_handle.abort();
}
#[tokio::test]
async fn gateway_executes_openai_chat_sync_via_local_decision_gate_without_execution_runtime_override(
) {

View File

@@ -42,3 +42,4 @@ use sha2::{Digest, Sha256};
mod failover;
mod local_decision;
mod pii_redaction;

View File

@@ -0,0 +1,836 @@
use super::*;
use std::sync::atomic::{AtomicUsize, Ordering};
#[derive(Debug, Clone)]
struct SeenProviderRequest {
body: serde_json::Value,
authorization: String,
accept_encoding: String,
}
fn hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());
format!("{:x}", hasher.finalize())
}
fn auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
user_id.to_string(),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-5"])),
api_key_id.to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(4_102_444_800),
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-5"])),
)
.expect("auth snapshot should build")
}
fn candidate_row(test_id: &str) -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: format!("provider-{test_id}"),
provider_name: "openai".to_string(),
provider_type: "custom".to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: format!("endpoint-{test_id}"),
endpoint_api_format: "openai:chat".to_string(),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("chat".to_string()),
endpoint_is_active: true,
key_id: format!("key-{test_id}"),
key_name: "prod".to_string(),
key_auth_type: "api_key".to_string(),
key_is_active: true,
key_api_formats: Some(vec!["openai:chat".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 5,
key_global_priority_by_format: Some(serde_json::json!({"openai:chat": 1})),
model_id: format!("model-{test_id}"),
global_model_id: format!("global-model-{test_id}"),
global_model_name: "gpt-5".to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(true),
model_provider_model_name: "gpt-5-upstream".to_string(),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: "gpt-5-upstream".to_string(),
priority: 1,
api_formats: Some(vec!["openai:chat".to_string()]),
endpoint_ids: Some(vec![format!("endpoint-{test_id}")]),
}]),
model_supports_streaming: Some(true),
model_is_active: true,
model_is_available: true,
}
}
fn provider(test_id: &str, redaction_enabled: bool) -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
format!("provider-{test_id}"),
"openai".to_string(),
Some("https://example.com".to_string()),
"custom".to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
true,
None,
Some(2),
None,
Some(20.0),
None,
Some(serde_json::json!({"chat_pii_redaction": {"enabled": redaction_enabled}})),
)
}
fn endpoint(test_id: &str, base_url: String) -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
format!("endpoint-{test_id}"),
format!("provider-{test_id}"),
"openai:chat".to_string(),
Some("openai".to_string()),
Some("chat".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(base_url, None, None, Some(2), None, None, None, None)
.expect("endpoint transport should build")
}
fn key(test_id: &str) -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
format!("key-{test_id}"),
format!("provider-{test_id}"),
"prod".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["openai:chat"])),
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-upstream-pii-redaction")
.expect("api key should encrypt"),
None,
None,
Some(serde_json::json!({"openai:chat": 1})),
None,
None,
None,
None,
)
.expect("key transport should build")
}
fn redaction_config(module_enabled: bool) -> Vec<(String, serde_json::Value)> {
redaction_config_with_entities(
module_enabled,
json!(["email", "access_token", "secret_key"]),
)
}
fn redaction_config_with_entities(
module_enabled: bool,
entities: serde_json::Value,
) -> Vec<(String, serde_json::Value)> {
vec![
(
"module.chat_pii_redaction.enabled".to_string(),
json!(module_enabled),
),
(
"module.chat_pii_redaction.provider_scope".to_string(),
json!("selected_providers"),
),
("module.chat_pii_redaction.entities".to_string(), entities),
(
"module.chat_pii_redaction.cache_ttl_seconds".to_string(),
json!(300),
),
(
"module.chat_pii_redaction.inject_model_instruction".to_string(),
json!(true),
),
]
}
fn collect_sentinels(text: &str, kind: &str) -> Vec<String> {
let prefix = format!("<AETHER:{kind}:");
let mut sentinels = Vec::new();
let mut offset = 0;
while let Some(relative_start) = text[offset..].find(&prefix) {
let start = offset + relative_start;
let Some(relative_end) = text[start..].find('>') else {
break;
};
let end = start + relative_end + 1;
sentinels.push(text[start..end].to_string());
offset = end;
}
sentinels
}
async fn run_sync_redaction_case(
test_id: &str,
module_enabled: bool,
provider_enabled: bool,
provider_response: &'static str,
request_body: serde_json::Value,
) -> (serde_json::Value, SeenProviderRequest) {
run_sync_redaction_case_with_system_config(
test_id,
provider_enabled,
provider_response,
request_body,
redaction_config(module_enabled),
)
.await
}
async fn run_sync_redaction_case_with_system_config(
test_id: &str,
provider_enabled: bool,
provider_response: &'static str,
request_body: serde_json::Value,
system_config: Vec<(String, serde_json::Value)>,
) -> (serde_json::Value, SeenProviderRequest) {
let seen_provider_request = Arc::new(Mutex::new(None::<SeenProviderRequest>));
let seen_provider_request_clone = Arc::clone(&seen_provider_request);
let provider_app = Router::new().route(
"/v1/chat/completions",
any(move |request: Request| {
let seen_provider_request_inner = Arc::clone(&seen_provider_request_clone);
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).expect("provider payload should parse");
let payload_text =
serde_json::to_string(&payload).expect("payload should serialize");
let email_sentinels = collect_sentinels(&payload_text, "EMAIL");
let access_token_sentinels = collect_sentinels(&payload_text, "ACCESS_TOKEN");
let secret_key_sentinels = collect_sentinels(&payload_text, "SECRET_KEY");
*seen_provider_request_inner
.lock()
.expect("mutex should lock") = Some(SeenProviderRequest {
body: payload,
authorization: parts
.headers
.get(http::header::AUTHORIZATION)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
accept_encoding: parts
.headers
.get(http::header::ACCEPT_ENCODING)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
});
let content = match provider_response {
"known" => format!(
"restored {} {} {}",
email_sentinels
.first()
.expect("user email sentinel should exist"),
access_token_sentinels
.first()
.expect("access token sentinel should exist"),
secret_key_sentinels
.first()
.expect("secret key sentinel should exist")
),
"unknown" => "unknown <AETHER:EMAIL:TSRQPONMLKJIHGFEDCBA>".to_string(),
"pass_through" => "pass alice@example.com".to_string(),
_ => unreachable!("provider response mode should be known"),
};
Json(json!({
"id": format!("chatcmpl-{provider_response}"),
"object": "chat.completion",
"model": "gpt-5-upstream",
"choices": [{
"index": 0,
"message": {"role": "assistant", "content": content},
"finish_reason": "stop"
}],
"usage": {"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5}
}))
}
}),
);
let (provider_url, provider_handle) = start_server(provider_app).await;
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key(&format!("sk-client-{test_id}"))),
auth_snapshot(&format!("api-key-{test_id}"), &format!("user-{test_id}")),
)]));
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
candidate_row(test_id),
]));
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider(test_id, provider_enabled)],
vec![endpoint(test_id, provider_url)],
vec![key(test_id)],
));
let data_state = crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
DEVELOPMENT_ENCRYPTION_KEY,
)
.with_system_config_values_for_tests(system_config);
let gateway_state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(data_state);
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
format!("Bearer sk-client-{test_id}"),
)
.header(http::header::ACCEPT_ENCODING, "gzip")
.header(TRACE_ID_HEADER, format!("trace-{test_id}"))
.body(request_body.to_string())
.send()
.await
.expect("request should succeed");
let status = response.status();
let execution_path = response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok())
.map(ToOwned::to_owned);
let response_text = response.text().await.expect("body should read");
assert_eq!(status, StatusCode::OK, "{response_text}");
assert_eq!(
execution_path.as_deref(),
Some(EXECUTION_PATH_EXECUTION_RUNTIME_SYNC)
);
let response_json: serde_json::Value =
serde_json::from_str(&response_text).expect("response body should parse");
let stored_candidates = request_candidate_repository
.list_by_request_id(&format!("trace-{test_id}"))
.await
.expect("request candidate trace should read");
assert_eq!(stored_candidates.len(), 1);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
let seen = seen_provider_request
.lock()
.expect("mutex should lock")
.clone()
.expect("provider request should be captured");
gateway_handle.abort();
provider_handle.abort();
(response_json, seen)
}
fn rich_pii_request() -> serde_json::Value {
json!({
"model": "gpt-5",
"messages": [
{"role": "system", "content": "Be concise."},
{
"role": "user",
"content": [
{"type": "text", "text": "Contact alice@example.com now"},
{"type": "input_audio", "input_audio": {"data": "AAAA", "format": "wav"}}
]
},
{
"role": "assistant",
"content": "Preparing lookup.",
"tool_calls": [{
"id": "call_1",
"type": "function",
"function": {
"name": "lookup_contact",
"arguments": "{\"email\":\"bob@example.net\"}"
}
}]
},
{
"role": "tool",
"tool_call_id": "call_1",
"content": "Tool returned access_token=accessValueABCDEF1234567890abcdef secret_key=secretValueABCDEF1234567890abcdef"
}
],
"tools": [{
"type": "function",
"function": {
"name": "lookup_contact",
"parameters": {
"type": "object",
"properties": {"email": {"type": "string"}}
}
}
}]
})
}
#[tokio::test]
async fn ai_execute_sync_pii_redaction_round_trip() {
let (response_json, seen) = run_sync_redaction_case(
"ai-execute-sync-pii-redaction-round-trip",
true,
true,
"known",
rich_pii_request(),
)
.await;
assert_eq!(seen.authorization, "Bearer sk-upstream-pii-redaction");
assert_eq!(seen.accept_encoding, "identity");
let provider_body_text = serde_json::to_string(&seen.body).expect("body should serialize");
for original in [
"alice@example.com",
"bob@example.net",
"access_token=accessValueABCDEF1234567890abcdef",
"secret_key\\\":\\\"secretValueABCDEF1234567890abcdef",
] {
assert!(
!provider_body_text.contains(original),
"leaked {original} in {provider_body_text}"
);
}
assert!(provider_body_text.contains("<AETHER:EMAIL:"));
assert!(provider_body_text.contains("<AETHER:ACCESS_TOKEN:"));
assert!(provider_body_text.contains("<AETHER:SECRET_KEY:"));
assert_eq!(seen.body["messages"][0]["role"], "system");
assert_eq!(seen.body["messages"][1]["role"], "assistant");
let notice = seen.body["messages"][1]["content"]
.as_str()
.expect("notice should be text");
assert!(notice.contains("not a user request"));
assert_eq!(seen.body["messages"][2]["role"], "user");
assert_eq!(seen.body["messages"][3]["role"], "assistant");
assert_eq!(seen.body["messages"][4]["role"], "tool");
let response_content = response_json["choices"][0]["message"]["content"]
.as_str()
.expect("assistant content should be text");
assert!(response_content.contains("alice@example.com"));
assert!(response_content.contains("access_token=accessValueABCDEF1234567890abcdef"));
assert!(response_content.contains("secretValueABCDEF1234567890abcdef"));
assert!(!response_content.contains("<AETHER:"));
}
#[tokio::test]
async fn ai_execute_pii_redaction_disabled_module_passes_original_chat_through() {
let (response_json, seen) = run_sync_redaction_case(
"ai-execute-pii-redaction-disabled-module",
false,
true,
"pass_through",
rich_pii_request(),
)
.await;
let provider_body_text = serde_json::to_string(&seen.body).expect("body should serialize");
assert!(provider_body_text.contains("alice@example.com"));
assert!(provider_body_text.contains("bob@example.net"));
assert!(provider_body_text.contains("access_token=accessValueABCDEF1234567890abcdef"));
assert!(provider_body_text.contains("secretValueABCDEF1234567890abcdef"));
assert!(!provider_body_text.contains("<AETHER:"));
assert_eq!(
response_json["choices"][0]["message"]["content"],
"pass alice@example.com"
);
}
#[tokio::test]
async fn ai_execute_pii_redaction_disabled_provider_passes_original_chat_through() {
let (response_json, seen) = run_sync_redaction_case(
"ai-execute-pii-redaction-disabled-provider",
true,
false,
"pass_through",
rich_pii_request(),
)
.await;
let provider_body_text = serde_json::to_string(&seen.body).expect("body should serialize");
assert!(provider_body_text.contains("alice@example.com"));
assert!(provider_body_text.contains("bob@example.net"));
assert!(provider_body_text.contains("access_token=accessValueABCDEF1234567890abcdef"));
assert!(provider_body_text.contains("secretValueABCDEF1234567890abcdef"));
assert!(!provider_body_text.contains("<AETHER:"));
assert_eq!(
response_json["choices"][0]["message"]["content"],
"pass alice@example.com"
);
}
#[tokio::test]
async fn ai_execute_pii_redaction_empty_entities_passes_original_chat_through() {
let (response_json, seen) = run_sync_redaction_case_with_system_config(
"ai-execute-pii-redaction-empty-entities",
true,
"pass_through",
rich_pii_request(),
redaction_config_with_entities(true, json!([])),
)
.await;
let provider_body_text = serde_json::to_string(&seen.body).expect("body should serialize");
assert!(provider_body_text.contains("alice@example.com"));
assert!(provider_body_text.contains("bob@example.net"));
assert!(provider_body_text.contains("access_token=accessValueABCDEF1234567890abcdef"));
assert!(provider_body_text.contains("secretValueABCDEF1234567890abcdef"));
assert!(!provider_body_text.contains("<AETHER:"));
assert_eq!(
response_json["choices"][0]["message"]["content"],
"pass alice@example.com"
);
}
#[tokio::test]
async fn ai_execute_pii_redaction_unknown_sentinel_like_output_is_not_restored() {
let (response_json, seen) = run_sync_redaction_case(
"ai-execute-pii-redaction-unknown-sentinel",
true,
true,
"unknown",
rich_pii_request(),
)
.await;
let provider_body_text = serde_json::to_string(&seen.body).expect("body should serialize");
assert!(!provider_body_text.contains("alice@example.com"));
assert!(provider_body_text.contains("<AETHER:EMAIL:"));
assert_eq!(
response_json["choices"][0]["message"]["content"],
"unknown <AETHER:EMAIL:TSRQPONMLKJIHGFEDCBA>"
);
}
#[tokio::test]
async fn ai_execute_pii_redaction_restores_executed_candidate_session_after_later_candidate_planning(
) {
let seen_provider_request = Arc::new(Mutex::new(None::<SeenProviderRequest>));
let seen_provider_request_clone = Arc::clone(&seen_provider_request);
let provider_app = Router::new().route(
"/v1/chat/completions",
any(move |request: Request| {
let seen_provider_request_inner = Arc::clone(&seen_provider_request_clone);
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).expect("provider payload should parse");
let payload_text = serde_json::to_string(&payload).expect("payload should serialize");
let email_sentinel = collect_sentinels(&payload_text, "EMAIL")
.into_iter()
.next()
.expect("email sentinel should exist");
*seen_provider_request_inner.lock().expect("mutex should lock") = Some(
SeenProviderRequest {
body: payload,
authorization: parts
.headers
.get(http::header::AUTHORIZATION)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
accept_encoding: parts
.headers
.get(http::header::ACCEPT_ENCODING)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
},
);
Json(json!({
"id": "chatcmpl-redaction-candidate-session",
"object": "chat.completion",
"model": "gpt-5-upstream",
"choices": [{
"index": 0,
"message": {"role": "assistant", "content": format!("restored {email_sentinel}")},
"finish_reason": "stop"
}],
"usage": {"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5}
}))
}
}),
);
let (provider_url, provider_handle) = start_server(provider_app).await;
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-redaction-candidate-session")),
auth_snapshot(
"api-key-redaction-candidate-session",
"user-redaction-candidate-session",
),
)]));
let mut later_candidate = candidate_row("redaction-candidate-session");
later_candidate.provider_id = "provider-redaction-candidate-session-later".to_string();
later_candidate.endpoint_id = "endpoint-redaction-candidate-session-later".to_string();
later_candidate.key_id = "key-redaction-candidate-session-later".to_string();
later_candidate.provider_priority = 20;
later_candidate.key_internal_priority = 6;
later_candidate.model_id = "model-redaction-candidate-session-later".to_string();
later_candidate.global_model_id = "global-model-redaction-candidate-session-later".to_string();
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
candidate_row("redaction-candidate-session"),
later_candidate,
]));
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let mut later_provider = provider("redaction-candidate-session", false);
later_provider.id = "provider-redaction-candidate-session-later".to_string();
let mut later_endpoint = endpoint(
"redaction-candidate-session",
"http://127.0.0.1:9".to_string(),
);
later_endpoint.id = "endpoint-redaction-candidate-session-later".to_string();
later_endpoint.provider_id = "provider-redaction-candidate-session-later".to_string();
let mut later_key = key("redaction-candidate-session");
later_key.id = "key-redaction-candidate-session-later".to_string();
later_key.provider_id = "provider-redaction-candidate-session-later".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![
provider("redaction-candidate-session", true),
later_provider,
],
vec![
endpoint("redaction-candidate-session", provider_url),
later_endpoint,
],
vec![key("redaction-candidate-session"), later_key],
));
let data_state = crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
DEVELOPMENT_ENCRYPTION_KEY,
)
.with_system_config_values_for_tests(redaction_config(true));
let gateway_state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(data_state);
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
"Bearer sk-client-redaction-candidate-session",
)
.header(TRACE_ID_HEADER, "trace-redaction-candidate-session")
.body(
json!({
"model": "gpt-5",
"messages": [{"role": "user", "content": "Contact alice@example.com"}]
})
.to_string(),
)
.send()
.await
.expect("request should succeed");
let status = response.status();
let response_text = response.text().await.expect("body should read");
assert_eq!(status, StatusCode::OK, "{response_text}");
let response_json: serde_json::Value =
serde_json::from_str(&response_text).expect("response body should parse");
assert_eq!(
response_json["choices"][0]["message"]["content"],
"restored alice@example.com"
);
let seen = seen_provider_request
.lock()
.expect("mutex should lock")
.clone()
.expect("provider request should be captured");
let provider_body_text = serde_json::to_string(&seen.body).expect("body should serialize");
assert!(!provider_body_text.contains("alice@example.com"));
assert!(provider_body_text.contains("<AETHER:EMAIL:"));
gateway_handle.abort();
provider_handle.abort();
}
#[tokio::test]
async fn pii_redaction_performance_limits_do_not_forward_unredacted_body_upstream() {
let provider_hits = Arc::new(AtomicUsize::new(0));
let provider_hits_clone = Arc::clone(&provider_hits);
let provider_app = Router::new().route(
"/v1/chat/completions",
any(move |_request: Request| {
let provider_hits_inner = Arc::clone(&provider_hits_clone);
async move {
provider_hits_inner.fetch_add(1, Ordering::SeqCst);
Json(json!({
"id": "unexpected",
"choices": [{"message": {"role": "assistant", "content": "unexpected"}}]
}))
}
}),
);
let (provider_url, provider_handle) = start_server(provider_app).await;
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-pii-redaction-limit")),
auth_snapshot("api-key-pii-redaction-limit", "user-pii-redaction-limit"),
)]));
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
candidate_row("pii-redaction-limit"),
]));
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider("pii-redaction-limit", true)],
vec![endpoint("pii-redaction-limit", provider_url)],
vec![key("pii-redaction-limit")],
));
let data_state = crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
DEVELOPMENT_ENCRYPTION_KEY,
)
.with_system_config_values_for_tests(redaction_config(true));
let gateway_state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(data_state);
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let original = format!("alice@example.com {}", "x".repeat(2 * 1024 * 1024));
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
"Bearer sk-client-pii-redaction-limit",
)
.header(TRACE_ID_HEADER, "trace-pii-redaction-limit")
.body(
json!({
"model": "gpt-5",
"messages": [{"role": "user", "content": original}]
})
.to_string(),
)
.send()
.await
.expect("request should complete");
let status = response.status();
let response_text = response.text().await.expect("body should read");
assert_eq!(status, StatusCode::PAYLOAD_TOO_LARGE, "{response_text}");
assert!(response_text.contains("scanned text limit exceeded"));
assert!(!response_text.contains("alice@example.com"));
assert_eq!(provider_hits.load(Ordering::SeqCst), 0);
gateway_handle.abort();
provider_handle.abort();
}
#[tokio::test]
async fn ai_execute_pii_redaction_missing_encryption_key_fails_closed_before_provider() {
let execution_runtime_hits = Arc::new(AtomicUsize::new(0));
let execution_runtime_hits_clone = Arc::clone(&execution_runtime_hits);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |_request: Request| {
let execution_runtime_hits_inner = Arc::clone(&execution_runtime_hits_clone);
async move {
execution_runtime_hits_inner.fetch_add(1, Ordering::SeqCst);
Json(json!({"ok": true}))
}
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let test_id = "ai-execute-pii-redaction-missing-encryption-key";
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key(&format!("sk-client-{test_id}"))),
auth_snapshot(&format!("api-key-{test_id}"), &format!("user-{test_id}")),
)]));
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
candidate_row(test_id),
]));
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider(test_id, true)],
vec![endpoint(test_id, "https://example.com".to_string())],
vec![key(test_id)],
));
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url.clone())
.with_data_state_for_tests(
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
"",
)
.with_system_config_values_for_tests(redaction_config(true)),
);
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
format!("Bearer sk-client-{test_id}"),
)
.header(TRACE_ID_HEADER, format!("trace-{test_id}"))
.body(rich_pii_request().to_string())
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
let response_text = response.text().await.expect("body should read");
for original in [
"alice@example.com",
"bob@example.net",
"accessValueABCDEF1234567890abcdef",
"secretValueABCDEF1234567890abcdef",
] {
assert!(!response_text.contains(original));
}
assert!(!response_text.contains("<AETHER:"));
assert_eq!(execution_runtime_hits.load(Ordering::SeqCst), 0);
gateway_handle.abort();
execution_runtime_handle.abort();
}

View File

@@ -782,6 +782,19 @@ async fn gateway_handles_admin_modules_status_locally_with_trusted_admin_princip
assert_eq!(payload["oauth"]["active"], json!(true));
assert_eq!(payload["oauth"]["config_validated"], json!(true));
assert_eq!(payload["management_tokens"]["active"], json!(true));
assert_eq!(payload["chat_pii_redaction"]["enabled"], json!(false));
assert_eq!(
payload["chat_pii_redaction"]["display_name"],
"敏感信息替换保护"
);
assert_eq!(
payload["chat_pii_redaction"]["config_validated"],
json!(true)
);
assert_eq!(
payload["chat_pii_redaction"]["admin_route"],
"/admin/modules/chat-pii-redaction"
);
assert_eq!(
payload["notification_email"]["config_validated"],
json!(true)
@@ -908,6 +921,63 @@ async fn gateway_handles_admin_module_status_detail_locally_with_trusted_admin_p
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_chat_pii_redaction_module_status_detail_locally_with_trusted_admin_principal(
) {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/modules/status/chat_pii_redaction",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let auth_module_repository = Arc::new(InMemoryAuthModuleReadRepository::default());
let data_state = GatewayDataState::with_auth_module_reader_for_tests(auth_module_repository)
.with_system_config_values_for_tests(vec![(
"module.chat_pii_redaction.enabled".to_string(),
json!(true),
)]);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data_state),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/modules/status/chat_pii_redaction"
))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["name"], "chat_pii_redaction");
assert_eq!(payload["display_name"], "敏感信息替换保护");
assert_eq!(payload["enabled"], json!(true));
assert_eq!(payload["active"], json!(true));
assert_eq!(payload["config_validated"], json!(true));
assert_eq!(payload["admin_route"], "/admin/modules/chat-pii-redaction");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_sets_admin_module_enabled_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));

View File

@@ -131,7 +131,7 @@ async fn gateway_handles_admin_providers_locally_with_trusted_admin_principal()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository,
provider_catalog_repository.clone(),
),
),
);
@@ -240,6 +240,7 @@ async fn gateway_handles_admin_provider_summary_locally_with_trusted_admin_princ
"claude_code_advanced": {"pool_size": 3},
"pool_advanced": {"enabled": true},
"failover_rules": {"strategy": "ordered"},
"chat_pii_redaction": {"enabled": true},
"provider_ops": {"architecture_id": "anyrouter"}
})),
);
@@ -351,6 +352,7 @@ async fn gateway_handles_admin_provider_summary_locally_with_trusted_admin_princ
);
assert_eq!(payload["ops_configured"], true);
assert_eq!(payload["ops_architecture_id"], "anyrouter");
assert_eq!(payload["chat_pii_redaction"], json!({"enabled": true}));
assert_eq!(payload["created_at"], "2024-03-21T05:46:40Z");
assert_eq!(payload["updated_at"], "2024-03-21T05:48:20Z");
assert_eq!(
@@ -774,6 +776,20 @@ async fn gateway_updates_admin_provider_locally_with_trusted_admin_principal() {
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider("provider-openai", "openai", 10)
.with_transport_fields(
true,
false,
false,
None,
None,
None,
None,
None,
Some(json!({
"pool_advanced": {},
"failover_rules": {"strategy": "ordered"}
})),
)
.with_timestamps(Some(1_711_000_000), Some(1_711_000_100)),
sample_provider("provider-other", "other", 20),
],
@@ -792,7 +808,7 @@ async fn gateway_updates_admin_provider_locally_with_trusted_admin_principal() {
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository,
provider_catalog_repository.clone(),
),
),
);
@@ -817,10 +833,11 @@ async fn gateway_updates_admin_provider_locally_with_trusted_admin_principal() {
"request_timeout": 55.0,
"stream_first_byte_timeout": 11.0,
"enable_format_conversion": false,
"config": {"provider_ops": {"architecture_id": "cubence"}},
"config": {
"provider_ops": {"architecture_id": "cubence"},
"chat_pii_redaction": {"enabled": true}
},
"claude_code_advanced": {"pool_size": 2},
"pool_advanced": {},
"failover_rules": {"strategy": "ordered"},
"proxy": {"url": "https://proxy.example"}
}))
.send()
@@ -847,8 +864,88 @@ async fn gateway_updates_admin_provider_locally_with_trusted_admin_principal() {
assert_eq!(payload["claude_code_advanced"], json!({"pool_size": 2}));
assert_eq!(payload["pool_advanced"], json!({}));
assert_eq!(payload["failover_rules"], json!({"strategy": "ordered"}));
assert_eq!(payload["chat_pii_redaction"], json!({"enabled": true}));
assert_eq!(payload["ops_configured"], true);
assert_eq!(payload["ops_architecture_id"], "cubence");
let disable_response = reqwest::Client::new()
.patch(format!("{gateway_url}/api/admin/providers/provider-openai"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"config": {
"chat_pii_redaction": {"enabled": false}
}
}))
.send()
.await
.expect("request should succeed");
let disable_status = disable_response.status();
let disable_body = disable_response.text().await.expect("body should read");
assert_eq!(disable_status, StatusCode::OK, "body={disable_body}");
let disable_payload: serde_json::Value =
serde_json::from_str(&disable_body).expect("json body should parse");
assert_eq!(
disable_payload["chat_pii_redaction"],
json!({"enabled": false})
);
assert_eq!(disable_payload["pool_advanced"], json!({}));
assert_eq!(
disable_payload["failover_rules"],
json!({"strategy": "ordered"})
);
assert_eq!(disable_payload["ops_architecture_id"], "cubence");
let invalid_response = reqwest::Client::new()
.patch(format!("{gateway_url}/api/admin/providers/provider-openai"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"config": {
"chat_pii_redaction": {"enabled": true, "entities": ["email"]}
}
}))
.send()
.await
.expect("request should succeed");
assert_eq!(invalid_response.status(), StatusCode::BAD_REQUEST);
let providers = provider_catalog_repository
.list_providers(false)
.await
.expect("providers should list");
let updated_provider = providers
.iter()
.find(|provider| provider.id == "provider-openai")
.expect("provider should exist");
assert_eq!(
updated_provider
.config
.as_ref()
.and_then(|value| value.get("chat_pii_redaction"))
.cloned(),
Some(json!({"enabled": false}))
);
assert_eq!(
updated_provider
.config
.as_ref()
.and_then(|value| value.get("pool_advanced"))
.cloned(),
Some(json!({}))
);
assert_eq!(
updated_provider
.config
.as_ref()
.and_then(|value| value.get("failover_rules"))
.cloned(),
Some(json!({"strategy": "ordered"}))
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
@@ -901,6 +998,7 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() {
"website": "codex.example",
"keep_priority_on_conversion": true,
"max_retries": 7,
"config": {"chat_pii_redaction": {"enabled": true}},
"pool_advanced": {},
"failover_rules": {"strategy": "ordered"},
"proxy": {"url": "https://proxy.example"}
@@ -943,6 +1041,38 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() {
.cloned(),
Some(json!({}))
);
assert_eq!(
created
.config
.as_ref()
.and_then(|value| value.get("chat_pii_redaction"))
.cloned(),
Some(json!({"enabled": true}))
);
assert_eq!(
created
.config
.as_ref()
.and_then(|value| value.get("failover_rules"))
.cloned(),
Some(json!({"strategy": "ordered"}))
);
let invalid_response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/providers/"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"name": "invalid-redaction-provider",
"provider_type": "custom",
"config": {"chat_pii_redaction": {"enabled": true, "entities": ["email"]}}
}))
.send()
.await
.expect("request should succeed");
assert_eq!(invalid_response.status(), StatusCode::BAD_REQUEST);
let endpoints = provider_catalog_repository
.list_endpoints_by_provider_ids(std::slice::from_ref(&created.id))

View File

@@ -1324,6 +1324,260 @@ async fn gateway_handles_admin_system_model_directives_default_as_disabled() {
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_validates_chat_pii_redaction_system_config_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/system/configs/module.chat_pii_redaction.cache_ttl_seconds",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let data_state =
GatewayDataState::disabled()
.with_system_config_values_for_tests(Vec::<(String, serde_json::Value)>::new());
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data_state),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let get_config = |key: &'static str| {
let client = client.clone();
let gateway_url = gateway_url.clone();
async move {
let response = client
.get(format!("{gateway_url}/api/admin/system/configs/{key}"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK, "key={key}");
response
.json::<serde_json::Value>()
.await
.expect("json body should parse")
}
};
let put_config = |key: &'static str, value: serde_json::Value| {
let client = client.clone();
let gateway_url = gateway_url.clone();
async move {
client
.put(format!("{gateway_url}/api/admin/system/configs/{key}"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({ "value": value }))
.send()
.await
.expect("request should succeed")
}
};
assert_eq!(
get_config("module.chat_pii_redaction.enabled").await["value"],
json!(false)
);
assert_eq!(
get_config("module.chat_pii_redaction.provider_scope").await["value"],
json!("selected_providers")
);
assert_eq!(
get_config("module.chat_pii_redaction.inject_model_instruction").await["value"],
json!(true)
);
assert_eq!(
get_config("module.chat_pii_redaction.cache_ttl_seconds").await["value"],
json!(300)
);
assert_eq!(
get_config("module.chat_pii_redaction.entities").await["value"],
json!([
"email",
"cn_phone",
"global_phone",
"cn_id",
"payment_card",
"ipv4",
"ipv6",
"api_key",
"access_token",
"secret_key",
"bearer_token",
"jwt"
])
);
let enabled_response = put_config("module.chat_pii_redaction.enabled", json!(true)).await;
assert_eq!(enabled_response.status(), StatusCode::OK);
let enabled_payload: serde_json::Value = enabled_response
.json()
.await
.expect("json body should parse");
assert_eq!(enabled_payload["value"], json!(true));
let scope_response = put_config(
"module.chat_pii_redaction.provider_scope",
json!("all_providers"),
)
.await;
assert_eq!(scope_response.status(), StatusCode::OK);
let scope_payload: serde_json::Value =
scope_response.json().await.expect("json body should parse");
assert_eq!(scope_payload["value"], json!("all_providers"));
let selected_entities_response = put_config(
"module.chat_pii_redaction.entities",
json!(["email", "jwt", "cn_phone"]),
)
.await;
assert_eq!(selected_entities_response.status(), StatusCode::OK);
let selected_entities_payload: serde_json::Value = selected_entities_response
.json()
.await
.expect("json body should parse");
assert_eq!(
selected_entities_payload["value"],
json!(["email", "cn_phone", "jwt"])
);
let ttl_response = put_config("module.chat_pii_redaction.cache_ttl_seconds", json!(3600)).await;
assert_eq!(ttl_response.status(), StatusCode::OK);
let ttl_payload: serde_json::Value = ttl_response.json().await.expect("json body should parse");
assert_eq!(ttl_payload["value"], json!(3600));
let instruction_response = put_config(
"module.chat_pii_redaction.inject_model_instruction",
json!(false),
)
.await;
assert_eq!(instruction_response.status(), StatusCode::OK);
let instruction_payload: serde_json::Value = instruction_response
.json()
.await
.expect("json body should parse");
assert_eq!(instruction_payload["value"], json!(false));
let invalid_scope_response = put_config(
"module.chat_pii_redaction.provider_scope",
json!("enabled_providers"),
)
.await;
assert_eq!(invalid_scope_response.status(), StatusCode::BAD_REQUEST);
let invalid_entities_response = put_config(
"module.chat_pii_redaction.entities",
json!(["email", "name"]),
)
.await;
assert_eq!(invalid_entities_response.status(), StatusCode::BAD_REQUEST);
let invalid_ttl_response =
put_config("module.chat_pii_redaction.cache_ttl_seconds", json!(600)).await;
assert_eq!(invalid_ttl_response.status(), StatusCode::BAD_REQUEST);
let invalid_instruction_response = put_config(
"module.chat_pii_redaction.inject_model_instruction",
json!("yes"),
)
.await;
assert_eq!(
invalid_instruction_response.status(),
StatusCode::BAD_REQUEST
);
let enabled_default_response =
put_config("module.chat_pii_redaction.enabled", serde_json::Value::Null).await;
assert_eq!(enabled_default_response.status(), StatusCode::OK);
let enabled_default_payload: serde_json::Value = enabled_default_response
.json()
.await
.expect("json body should parse");
assert_eq!(enabled_default_payload["value"], json!(false));
let scope_default_response = put_config(
"module.chat_pii_redaction.provider_scope",
serde_json::Value::Null,
)
.await;
assert_eq!(scope_default_response.status(), StatusCode::OK);
let scope_default_payload: serde_json::Value = scope_default_response
.json()
.await
.expect("json body should parse");
assert_eq!(scope_default_payload["value"], json!("selected_providers"));
let entities_default_response = put_config(
"module.chat_pii_redaction.entities",
serde_json::Value::Null,
)
.await;
assert_eq!(entities_default_response.status(), StatusCode::OK);
let entities_default_payload: serde_json::Value = entities_default_response
.json()
.await
.expect("json body should parse");
assert_eq!(
entities_default_payload["value"],
json!([
"email",
"cn_phone",
"global_phone",
"cn_id",
"payment_card",
"ipv4",
"ipv6",
"api_key",
"access_token",
"secret_key",
"bearer_token",
"jwt"
])
);
let ttl_default_response = put_config(
"module.chat_pii_redaction.cache_ttl_seconds",
serde_json::Value::Null,
)
.await;
assert_eq!(ttl_default_response.status(), StatusCode::OK);
let ttl_default_payload: serde_json::Value = ttl_default_response
.json()
.await
.expect("json body should parse");
assert_eq!(ttl_default_payload["value"], json!(300));
let instruction_default_response = put_config(
"module.chat_pii_redaction.inject_model_instruction",
serde_json::Value::Null,
)
.await;
assert_eq!(instruction_default_response.status(), StatusCode::OK);
let instruction_default_payload: serde_json::Value = instruction_default_response
.json()
.await
.expect("json body should parse");
assert_eq!(instruction_default_payload["value"], json!(true));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_system_provider_priority_mode_locally_with_bearer_admin_session() {
let upstream_hits = Arc::new(Mutex::new(0usize));