mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
feat(gateway): add reversible chat pii redaction
This commit is contained in:
@@ -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(),
|
||||
},
|
||||
))
|
||||
)))
|
||||
}
|
||||
|
||||
@@ -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(),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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));
|
||||
}
|
||||
|
||||
@@ -134,7 +134,7 @@ impl LocalOpenAiChatStreamAttemptSource<'_> {
|
||||
"openai_chat_stream_success",
|
||||
upstream_is_stream,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
@@ -134,7 +134,7 @@ impl LocalOpenAiChatSyncAttemptSource<'_> {
|
||||
"openai_chat_sync_success",
|
||||
upstream_is_stream,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
@@ -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!({
|
||||
|
||||
@@ -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> {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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!(
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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: "异常通知",
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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() {
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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),
|
||||
}
|
||||
}
|
||||
|
||||
4319
apps/aether-gateway/src/privacy/mod.rs
Normal file
4319
apps/aether-gateway/src/privacy/mod.rs
Normal file
File diff suppressed because it is too large
Load Diff
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1414,6 +1414,7 @@ impl AppState {
|
||||
message: match err {
|
||||
GatewayError::UpstreamUnavailable { message, .. }
|
||||
| GatewayError::ControlUnavailable { message, .. }
|
||||
| GatewayError::Client { message, .. }
|
||||
| GatewayError::Internal(message) => message,
|
||||
},
|
||||
},
|
||||
|
||||
@@ -26,3 +26,4 @@ use super::{
|
||||
|
||||
mod decision;
|
||||
mod image;
|
||||
mod pii_redaction;
|
||||
|
||||
340
apps/aether-gateway/src/tests/ai_execute/stream/pii_redaction.rs
Normal file
340
apps/aether-gateway/src/tests/ai_execute/stream/pii_redaction.rs
Normal 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();
|
||||
}
|
||||
@@ -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(
|
||||
) {
|
||||
|
||||
@@ -42,3 +42,4 @@ use sha2::{Digest, Sha256};
|
||||
|
||||
mod failover;
|
||||
mod local_decision;
|
||||
mod pii_redaction;
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
@@ -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));
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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));
|
||||
|
||||
Reference in New Issue
Block a user