test(gateway): cover chat pii redaction formats

This commit is contained in:
Kayphoon
2026-05-26 22:29:45 +08:00
parent a2f91b4108
commit 7530bec7de
18 changed files with 2049 additions and 364 deletions
@@ -17,6 +17,7 @@ mod passthrough;
mod plan_builders;
mod pool_scheduler;
pub(crate) mod pool_scores;
mod redaction;
mod report_context;
mod route;
mod runtime_miss;
@@ -70,7 +70,7 @@ pub(crate) async fn maybe_build_sync_local_same_format_provider_decision_payload
maybe_build_local_same_format_provider_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
.await
.await?
{
return Ok(Some(payload));
}
@@ -137,7 +137,7 @@ pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_paylo
maybe_build_local_same_format_provider_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
.await
.await?
{
return Ok(Some(payload));
}
@@ -22,7 +22,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 aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
@@ -40,7 +40,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
input: &LocalSameFormatProviderDecisionInput,
attempt: LocalSameFormatProviderCandidateAttempt,
spec: LocalSameFormatProviderSpec,
) -> Option<AiExecutionDecision> {
) -> Result<Option<AiExecutionDecision>, GatewayError> {
let spec_metadata = local_same_format_provider_spec_metadata(spec);
let LocalSameFormatProviderCandidateAttempt {
eligible,
@@ -51,10 +51,18 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
let candidate = &eligible.candidate;
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats(spec_metadata.api_format, spec_metadata.api_format);
let resolved = resolve_local_same_format_provider_candidate_payload_parts(
let Some(resolved) = resolve_local_same_format_provider_candidate_payload_parts(
state, parts, trace_id, body_json, input, &attempt, spec,
)
.await?;
.await?
else {
return Ok(None);
};
let original_request_body_json = if resolved.request_redacted {
Some(&resolved.provider_request_body)
} else {
Some(body_json)
};
let prompt_cache_key = resolved
.provider_request_body
@@ -116,7 +124,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
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_json,
original_request_body_base64: None,
client_session_affinity: input.client_session_affinity.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
@@ -149,9 +157,10 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
upstream_url,
provider_request_headers,
provider_request_body,
request_redacted: _,
} = resolved;
Some(build_ai_execution_decision_response(
Ok(Some(build_ai_execution_decision_response(
AiExecutionDecisionResponseParts {
decision_is_stream: spec_metadata.require_streaming,
decision_kind: spec_metadata.decision_kind.to_string(),
@@ -185,7 +194,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
report_context: Some(report_context),
auth_context: input.auth_context.clone(),
},
))
)))
}
pub(super) async fn mark_skipped_local_same_format_provider_candidate(
@@ -6,6 +6,9 @@ use serde_json::Value;
use crate::ai_serving::planner::common::{
enforce_provider_body_stream_policy, request_requires_body_stream_field,
};
use crate::ai_serving::planner::redaction::{
request_identity_response_encoding_when_redacted, resolve_provider_chat_pii_redaction,
};
use crate::ai_serving::transport::antigravity::{
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
classify_local_antigravity_request_support, AntigravityEnvelopeRequestType,
@@ -15,7 +18,7 @@ use crate::ai_serving::transport::{
build_same_format_provider_headers, SameFormatProviderHeadersInput,
};
use crate::ai_serving::{CandidateFailureDiagnostic, GatewayProviderTransportSnapshot};
use crate::AppState;
use crate::{AppState, GatewayError};
mod policy;
mod prepare;
@@ -96,6 +99,7 @@ pub(crate) struct LocalSameFormatProviderCandidatePayloadParts {
pub(super) upstream_url: String,
pub(super) provider_request_headers: BTreeMap<String, String>,
pub(super) provider_request_body: Value,
pub(super) request_redacted: bool,
}
pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
@@ -106,9 +110,9 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
input: &LocalSameFormatProviderDecisionInput,
attempt: &LocalSameFormatProviderCandidateAttempt,
spec: LocalSameFormatProviderSpec,
) -> Option<LocalSameFormatProviderCandidatePayloadParts> {
) -> Result<Option<LocalSameFormatProviderCandidatePayloadParts>, GatewayError> {
let candidate = &attempt.eligible.candidate;
let prepared = prepare_local_same_format_provider_candidate(
let Some(prepared) = prepare_local_same_format_provider_candidate(
state,
trace_id,
input,
@@ -117,7 +121,10 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
&attempt.candidate_id,
spec,
)
.await?;
.await
else {
return Ok(None);
};
let enable_model_directives =
crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model(
state,
@@ -125,6 +132,16 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
Some(&input.requested_model),
)
.await;
let redaction = resolve_provider_chat_pii_redaction(
state,
parts,
body_json,
&input.auth_context,
spec.api_format,
&attempt.candidate_id,
)
.await?;
let body_json = redaction.body_json.as_ref();
let Some(mut base_provider_request_body) =
super::super::request::build_same_format_provider_request_body(
@@ -161,7 +178,7 @@ pub(crate) async fn resolve_local_same_format_provider_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(
@@ -207,7 +224,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
"transport_unsupported",
)
.await;
return None;
return Ok(None);
}
}
} else {
@@ -239,7 +256,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
),
)
.await;
return None;
return Ok(None);
}
}
} else {
@@ -270,14 +287,14 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
),
)
.await;
return None;
return Ok(None);
};
let extra_headers = antigravity_auth
.as_ref()
.map(build_antigravity_static_identity_headers)
.unwrap_or_default();
let Some(provider_request_headers) =
let Some(mut provider_request_headers) =
build_same_format_provider_headers(SameFormatProviderHeadersInput {
headers: &parts.headers,
provider_request_body: &provider_request_body,
@@ -310,10 +327,14 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
),
)
.await;
return None;
return Ok(None);
};
request_identity_response_encoding_when_redacted(
&mut provider_request_headers,
redaction.redacted,
);
Some(LocalSameFormatProviderCandidatePayloadParts {
Ok(Some(LocalSameFormatProviderCandidatePayloadParts {
transport: prepared.transport,
is_antigravity: prepared.is_antigravity,
is_kiro: prepared.is_kiro,
@@ -326,5 +347,6 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
upstream_url,
provider_request_headers,
provider_request_body,
})
request_redacted: redaction.redacted,
}))
}
@@ -249,7 +249,7 @@ impl LocalSameFormatProviderSyncAttemptSource<'_> {
attempt,
self.spec,
)
.await
.await?
else {
return Ok(None);
};
@@ -287,7 +287,7 @@ impl LocalSameFormatProviderStreamAttemptSource<'_> {
attempt,
self.spec,
)
.await
.await?
else {
return Ok(None);
};
@@ -365,7 +365,7 @@ pub(crate) async fn build_local_sync_plan_and_reports(
let Some(payload) = maybe_build_local_same_format_provider_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
.await
.await?
else {
continue;
};
@@ -450,7 +450,7 @@ pub(crate) async fn build_local_stream_plan_and_reports(
let Some(payload) = maybe_build_local_same_format_provider_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
.await
.await?
else {
continue;
};
@@ -0,0 +1,192 @@
use std::borrow::Cow;
use std::time::{SystemTime, UNIX_EPOCH};
use serde_json::Value;
use tracing::warn;
use crate::ai_serving::ExecutionRuntimeAuthContext;
use crate::privacy::{
build_redaction_session_config, read_chat_pii_redaction_runtime_config,
try_mask_chat_pii_request_json_with_cache_options, ChatPiiRedactionRequestFormat,
MaskChatRequestOptions, RedactionMaskError, RedactionSessionSlot, RedisRedactionMappingCache,
};
use crate::{AppState, GatewayError};
pub(crate) struct ProviderRequestRedaction<'a> {
pub(crate) body_json: Cow<'a, Value>,
pub(crate) redacted: bool,
}
impl<'a> ProviderRequestRedaction<'a> {
fn disabled(body_json: &'a Value) -> Self {
Self {
body_json: Cow::Borrowed(body_json),
redacted: false,
}
}
}
#[derive(Clone, Copy, Debug, Default)]
struct ChatPiiRedactionFeatureSettings {
enabled: Option<bool>,
inject_model_instruction: Option<bool>,
}
impl ChatPiiRedactionFeatureSettings {
fn merge_from_value(&mut self, value: Option<&Value>) {
let Some(settings) = value
.and_then(Value::as_object)
.and_then(|features| features.get("chat_pii_redaction"))
.and_then(Value::as_object)
else {
return;
};
if let Some(enabled) = settings.get("enabled").and_then(Value::as_bool) {
self.enabled = Some(enabled);
}
if let Some(inject_model_instruction) = settings
.get("inject_model_instruction")
.and_then(Value::as_bool)
{
self.inject_model_instruction = Some(inject_model_instruction);
}
}
fn effective_enabled(self) -> bool {
self.enabled.unwrap_or(false)
}
fn effective_inject_model_instruction(self) -> bool {
self.inject_model_instruction.unwrap_or(true)
}
}
pub(crate) fn request_identity_response_encoding_when_redacted(
headers: &mut std::collections::BTreeMap<String, String>,
redacted: bool,
) {
if redacted {
headers.insert("accept-encoding".to_string(), "identity".to_string());
}
}
pub(crate) async fn resolve_provider_chat_pii_redaction<'a>(
state: &AppState,
parts: &http::request::Parts,
body_json: &'a Value,
auth_context: &ExecutionRuntimeAuthContext,
client_api_format: &str,
candidate_id: &str,
) -> Result<ProviderRequestRedaction<'a>, GatewayError> {
let Some(format) = ChatPiiRedactionRequestFormat::from_api_format(client_api_format) else {
return Ok(ProviderRequestRedaction::disabled(body_json));
};
let Some(slot) = parts.extensions.get::<RedactionSessionSlot>() else {
return Ok(ProviderRequestRedaction::disabled(body_json));
};
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 !runtime_config.enabled {
return Ok(ProviderRequestRedaction::disabled(body_json));
}
let feature_settings = resolve_chat_pii_redaction_feature_settings(state, auth_context).await?;
if !feature_settings.effective_enabled() {
return Ok(ProviderRequestRedaction::disabled(body_json));
}
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_pii_request_json_with_cache_options(
&body_bytes,
format,
build_redaction_session_config(hmac_key, &runtime_config, now_unix_secs),
MaskChatRequestOptions::runtime(feature_settings.effective_inject_model_instruction()),
Some(&cache),
)
.await
.map_err(redaction_mask_error_to_gateway_error)?;
if !masked.redacted {
return Ok(ProviderRequestRedaction {
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(ProviderRequestRedaction {
body_json: Cow::Owned(masked_body_json),
redacted: true,
})
}
async fn resolve_chat_pii_redaction_feature_settings(
state: &AppState,
auth_context: &ExecutionRuntimeAuthContext,
) -> Result<ChatPiiRedactionFeatureSettings, GatewayError> {
let user_settings = state
.read_user_feature_settings(&auth_context.user_id)
.await
.map_err(|err| {
warn!(
error = ?err,
"gateway failed to read user chat pii redaction feature settings"
);
GatewayError::Internal("chat pii redaction setup failed".to_string())
})?;
let key_settings = state
.read_auth_api_key_feature_settings(
&auth_context.user_id,
&auth_context.api_key_id,
auth_context.api_key_is_standalone,
)
.await
.map_err(|err| {
warn!(
error = ?err,
"gateway failed to read api key chat pii redaction feature settings"
);
GatewayError::Internal("chat pii redaction setup failed".to_string())
})?;
let mut settings = ChatPiiRedactionFeatureSettings::default();
settings.merge_from_value(user_settings.as_ref());
settings.merge_from_value(key_settings.as_ref());
Ok(settings)
}
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(),
},
}
}
@@ -233,7 +233,7 @@ impl LocalStandardSyncAttemptSource<'_> {
attempt,
self.spec,
)
.await
.await?
else {
return Ok(None);
};
@@ -270,7 +270,7 @@ impl LocalStandardStreamAttemptSource<'_> {
attempt,
self.spec,
)
.await
.await?
else {
return Ok(None);
};
@@ -331,7 +331,7 @@ pub(crate) async fn maybe_build_sync_via_standard_family_payload(
if let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
.await
.await?
{
return Ok(Some(payload));
}
@@ -380,7 +380,7 @@ pub(crate) async fn maybe_build_stream_via_standard_family_payload(
if let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
.await
.await?
{
return Ok(Some(payload));
}
@@ -438,7 +438,7 @@ pub(crate) async fn build_local_sync_plan_and_reports(
let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
.await
.await?
else {
continue;
};
@@ -512,7 +512,7 @@ pub(crate) async fn build_local_stream_plan_and_reports(
let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
.await
.await?
else {
continue;
};
@@ -24,7 +24,7 @@ use crate::ai_serving::{
};
use crate::{
append_execution_contract_fields_to_value, append_local_failover_policy_to_value,
AiExecutionDecision, AppState,
AiExecutionDecision, AppState, GatewayError,
};
use super::request::resolve_local_standard_candidate_payload_parts;
@@ -38,7 +38,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
input: &LocalStandardDecisionInput,
attempt: LocalStandardCandidateAttempt,
spec: LocalStandardSpec,
) -> Option<AiExecutionDecision> {
) -> Result<Option<AiExecutionDecision>, GatewayError> {
let spec_metadata = local_standard_spec_metadata(spec);
if api_format_alias_matches(
&attempt.eligible.provider_api_format,
@@ -70,10 +70,18 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
..
} = &attempt;
let candidate = &eligible.candidate;
let resolved = resolve_local_standard_candidate_payload_parts(
let Some(resolved) = resolve_local_standard_candidate_payload_parts(
state, parts, trace_id, body_json, input, &attempt, spec,
)
.await?;
.await?
else {
return Ok(None);
};
let original_request_body_json = if resolved.request_redacted {
Some(&resolved.provider_request_body)
} else {
Some(body_json)
};
let proxy = state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&resolved.transport)
.await;
@@ -124,7 +132,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
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_json,
original_request_body_base64: None,
client_session_affinity: input.client_session_affinity.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
@@ -157,9 +165,10 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
upstream_is_stream,
envelope_name: _,
transport,
request_redacted: _,
} = resolved;
Some(build_ai_execution_decision_response(
Ok(Some(build_ai_execution_decision_response(
AiExecutionDecisionResponseParts {
decision_is_stream: spec_metadata.require_streaming,
decision_kind: spec_metadata.decision_kind.to_string(),
@@ -193,7 +202,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
report_context: Some(report_context),
auth_context: input.auth_context.clone(),
},
))
)))
}
pub(super) async fn mark_skipped_local_standard_candidate(
@@ -512,6 +521,7 @@ mod tests {
claude_stream_spec(),
)
.await
.expect("same-format candidate should not error")
.expect("same-format candidate should build a standard-family payload");
assert_eq!(payload.endpoint_id.as_deref(), Some("endpoint-claude"));
@@ -547,6 +557,7 @@ mod tests {
claude_stream_spec(),
)
.await
.expect("cross-format candidate should not error")
.expect("cross-format candidate should still build after the same-format candidate");
assert_eq!(
@@ -11,6 +11,9 @@ use crate::ai_serving::planner::common::{
endpoint_config_forces_body_stream_field, enforce_provider_body_stream_policy,
request_requires_body_stream_field, resolve_upstream_is_stream_for_provider,
};
use crate::ai_serving::planner::redaction::{
request_identity_response_encoding_when_redacted, resolve_provider_chat_pii_redaction,
};
use crate::ai_serving::planner::spec_metadata::local_standard_spec_metadata;
use crate::ai_serving::planner::standard::{
apply_codex_openai_responses_special_headers, request_body_build_failure_extra_data,
@@ -30,7 +33,7 @@ use crate::ai_serving::{
build_openai_image_request_body_from_gemini_image_request, gemini_request_is_image_generation,
CandidateFailureDiagnostic, GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
};
use crate::AppState;
use crate::{AppState, GatewayError};
use super::payload::{
mark_skipped_local_standard_candidate, mark_skipped_local_standard_candidate_with_extra_data,
@@ -49,6 +52,7 @@ pub(crate) struct LocalStandardCandidatePayloadParts {
pub(super) upstream_is_stream: bool,
pub(super) envelope_name: Option<&'static str>,
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
pub(super) request_redacted: bool,
}
pub(crate) async fn resolve_local_standard_candidate_payload_parts(
@@ -59,7 +63,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
input: &LocalStandardDecisionInput,
attempt: &LocalStandardCandidateAttempt,
spec: LocalStandardSpec,
) -> Option<LocalStandardCandidatePayloadParts> {
) -> Result<Option<LocalStandardCandidatePayloadParts>, GatewayError> {
let spec_metadata = local_standard_spec_metadata(spec);
let planner_state = crate::ai_serving::PlannerAppState::new(state);
let candidate = &attempt.eligible.candidate;
@@ -69,16 +73,18 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
&& provider_api_format == "openai:image"
&& gemini_request_is_image_generation(body_json)
{
return resolve_local_gemini_image_to_openai_image_candidate_payload_parts(
state, parts, trace_id, body_json, input, attempt,
)
.await;
return Ok(
resolve_local_gemini_image_to_openai_image_candidate_payload_parts(
state, parts, trace_id, body_json, input, attempt,
)
.await,
);
}
let is_kiro_claude_cli = is_kiro_claude_messages_transport(transport, provider_api_format);
let Some(conversion_kind) =
crate::ai_serving::request_conversion_kind(spec_metadata.api_format, provider_api_format)
else {
return None;
return Ok(None);
};
if let Some(skip_reason) = crate::ai_serving::request_conversion_transport_unsupported_reason(
@@ -95,7 +101,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
skip_reason,
)
.await;
return None;
return Ok(None);
}
let oauth_context = OauthPreparationContext {
@@ -123,7 +129,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
"transport_auth_unavailable",
)
.await;
return None;
return Ok(None);
}
}
} else {
@@ -148,7 +154,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
skip_reason,
)
.await;
return None;
return Ok(None);
}
}
} else {
@@ -173,7 +179,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
skip_reason,
)
.await;
return None;
return Ok(None);
}
}
};
@@ -194,6 +200,16 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
Some(&input.requested_model),
)
.await;
let redaction = resolve_provider_chat_pii_redaction(
state,
parts,
body_json,
&input.auth_context,
spec_metadata.api_format,
&attempt.candidate_id,
)
.await?;
let body_json = redaction.body_json.as_ref();
let mut provider_request_body =
match crate::ai_serving::planner::standard::build_standard_request_body_with_model_directives_and_request_headers(
body_json,
@@ -229,7 +245,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
),
)
.await;
return None;
return Ok(None);
}
};
enforce_provider_body_stream_policy(
@@ -261,7 +277,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
}
if let Some(kiro_auth) = kiro_auth.as_ref() {
return build_kiro_cross_format_payload_parts(
return Ok(build_kiro_cross_format_payload_parts(
state,
parts,
trace_id,
@@ -276,8 +292,9 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
provider_request_body,
upstream_is_stream,
kiro_auth,
redaction.redacted,
)
.await;
.await);
}
let upstream_url = match crate::ai_serving::planner::standard::build_standard_upstream_url(
@@ -304,7 +321,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
),
)
.await;
return None;
return Ok(None);
}
};
let Some(resolved_headers) =
@@ -337,7 +354,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
),
)
.await;
return None;
return Ok(None);
};
let mut provider_request_headers = resolved_headers.headers;
apply_codex_openai_responses_special_headers(
@@ -349,8 +366,12 @@ pub(crate) async fn resolve_local_standard_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,
);
Some(LocalStandardCandidatePayloadParts {
Ok(Some(LocalStandardCandidatePayloadParts {
auth_header: resolved_headers.auth_header,
auth_value: resolved_headers.auth_value,
mapped_model: prepared_candidate.mapped_model,
@@ -361,7 +382,8 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
upstream_is_stream,
envelope_name: None,
transport: Arc::clone(transport),
})
request_redacted: redaction.redacted,
}))
}
async fn resolve_local_gemini_image_to_openai_image_candidate_payload_parts(
@@ -496,6 +518,7 @@ async fn resolve_local_gemini_image_to_openai_image_candidate_payload_parts(
upstream_is_stream,
envelope_name: None,
transport: Arc::clone(transport),
request_redacted: false,
})
}
@@ -515,6 +538,7 @@ async fn build_kiro_cross_format_payload_parts(
claude_request_body: Value,
upstream_is_stream: bool,
kiro_auth: &KiroRequestAuth,
request_redacted: bool,
) -> Option<LocalStandardCandidatePayloadParts> {
let candidate = &attempt.eligible.candidate;
let provider_request_body = match build_kiro_provider_request_body(
@@ -572,7 +596,7 @@ async fn build_kiro_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,
@@ -602,6 +626,10 @@ async fn build_kiro_cross_format_payload_parts(
return None;
}
};
request_identity_response_encoding_when_redacted(
&mut provider_request_headers,
request_redacted,
);
Some(LocalStandardCandidatePayloadParts {
auth_header,
@@ -614,5 +642,6 @@ async fn build_kiro_cross_format_payload_parts(
upstream_is_stream,
envelope_name: Some(KIRO_ENVELOPE_NAME),
transport: Arc::clone(transport),
request_redacted,
})
}
@@ -1,7 +1,5 @@
use std::borrow::Cow;
use std::collections::BTreeMap;
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use serde_json::Value;
@@ -14,6 +12,9 @@ use crate::ai_serving::planner::common::{
endpoint_config_forces_body_stream_field, enforce_provider_body_stream_policy,
request_requires_body_stream_field, OPENAI_CHAT_STREAM_PLAN_KIND,
};
use crate::ai_serving::planner::redaction::{
request_identity_response_encoding_when_redacted, resolve_provider_chat_pii_redaction,
};
use crate::ai_serving::planner::standard::{
apply_codex_openai_responses_special_headers, build_cross_format_openai_chat_request_body,
build_cross_format_openai_chat_upstream_url, build_local_openai_chat_request_body,
@@ -36,13 +37,7 @@ use crate::ai_serving::{
LocalResolvedOAuthRequestAuth,
};
use crate::ai_serving::{ConversionMode, ExecutionStrategy};
use crate::privacy::{
build_redaction_session_config, 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,
@@ -66,99 +61,6 @@ pub(crate) struct LocalOpenAiChatCandidatePayloadParts {
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,
}
}
}
#[derive(Clone, Copy, Debug, Default)]
struct ChatPiiRedactionFeatureSettings {
enabled: Option<bool>,
inject_model_instruction: Option<bool>,
}
impl ChatPiiRedactionFeatureSettings {
fn merge_from_value(&mut self, value: Option<&Value>) {
let Some(settings) = value
.and_then(Value::as_object)
.and_then(|features| features.get("chat_pii_redaction"))
.and_then(Value::as_object)
else {
return;
};
if let Some(enabled) = settings.get("enabled").and_then(Value::as_bool) {
self.enabled = Some(enabled);
}
if let Some(inject_model_instruction) = settings
.get("inject_model_instruction")
.and_then(Value::as_bool)
{
self.inject_model_instruction = Some(inject_model_instruction);
}
}
fn effective_enabled(self) -> bool {
self.enabled.unwrap_or(false)
}
fn effective_inject_model_instruction(self) -> bool {
self.inject_model_instruction.unwrap_or(true)
}
}
async fn resolve_chat_pii_redaction_feature_settings(
state: &AppState,
input: &LocalOpenAiChatDecisionInput,
) -> Result<ChatPiiRedactionFeatureSettings, GatewayError> {
let user_settings = state
.read_user_feature_settings(&input.auth_context.user_id)
.await
.map_err(|err| {
warn!(
error = ?err,
"gateway failed to read user chat pii redaction feature settings"
);
GatewayError::Internal("chat pii redaction setup failed".to_string())
})?;
let key_settings = state
.read_auth_api_key_feature_settings(
&input.auth_context.user_id,
&input.auth_context.api_key_id,
input.auth_context.api_key_is_standalone,
)
.await
.map_err(|err| {
warn!(
error = ?err,
"gateway failed to read api key chat pii redaction feature settings"
);
GatewayError::Internal("chat pii redaction setup failed".to_string())
})?;
let mut settings = ChatPiiRedactionFeatureSettings::default();
settings.merge_from_value(user_settings.as_ref());
settings.merge_from_value(key_settings.as_ref());
Ok(settings)
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
state: &AppState,
@@ -186,9 +88,15 @@ 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, input, candidate_id)
.await?;
let redaction = resolve_provider_chat_pii_redaction(
state,
parts,
body_json,
&input.auth_context,
"openai:chat",
candidate_id,
)
.await?;
let body_json = redaction.body_json.as_ref();
if provider_api_format == "openai:chat" {
@@ -777,87 +685,3 @@ async fn build_kiro_openai_chat_cross_format_payload_parts(
request_redacted,
})
}
async fn resolve_provider_chat_request_redaction<'a>(
state: &AppState,
parts: &http::request::Parts,
body_json: &'a Value,
input: &LocalOpenAiChatDecisionInput,
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 !runtime_config.enabled {
return Ok(ProviderChatRequestRedaction::disabled(body_json, parts));
}
let feature_settings = resolve_chat_pii_redaction_feature_settings(state, input).await?;
if !feature_settings.effective_enabled() {
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(feature_settings.effective_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(),
},
}
}
@@ -15,7 +15,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_responses_candidate_payload_parts;
@@ -30,7 +30,7 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
input: &LocalOpenAiResponsesDecisionInput,
attempt: LocalOpenAiResponsesCandidateAttempt,
spec: LocalOpenAiResponsesSpec,
) -> Option<AiExecutionDecision> {
) -> Result<Option<AiExecutionDecision>, GatewayError> {
let spec_metadata = local_openai_responses_spec_metadata(spec);
let attempt_identity = attempt.attempt_identity();
let LocalOpenAiResponsesCandidateAttempt {
@@ -39,7 +39,7 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
candidate_id,
..
} = attempt;
let resolved = resolve_local_openai_responses_candidate_payload_parts(
let Some(resolved) = resolve_local_openai_responses_candidate_payload_parts(
state,
parts,
trace_id,
@@ -50,8 +50,16 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
&candidate_id,
spec,
)
.await?;
.await?
else {
return Ok(None);
};
let candidate = &eligible.candidate;
let original_request_body_json = if resolved.request_redacted {
Some(&resolved.provider_request_body)
} else {
Some(body_json)
};
let prompt_cache_key = resolved
.provider_request_body
@@ -109,7 +117,7 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
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_json,
original_request_body_base64: None,
client_session_affinity: input.client_session_affinity.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
@@ -170,9 +178,10 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
envelope_name: _,
upstream_is_stream,
transport,
request_redacted: _,
} = resolved;
Some(build_ai_execution_decision_response(
Ok(Some(build_ai_execution_decision_response(
AiExecutionDecisionResponseParts {
decision_is_stream: spec_metadata.require_streaming,
decision_kind: spec_metadata.decision_kind.to_string(),
@@ -206,5 +215,5 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
report_context: Some(report_context),
auth_context: input.auth_context.clone(),
},
))
)))
}
@@ -13,6 +13,9 @@ use crate::ai_serving::planner::common::{
endpoint_config_forces_body_stream_field, enforce_provider_body_stream_policy,
request_requires_body_stream_field, resolve_upstream_is_stream_for_provider,
};
use crate::ai_serving::planner::redaction::{
request_identity_response_encoding_when_redacted, resolve_provider_chat_pii_redaction,
};
use crate::ai_serving::planner::spec_metadata::local_openai_responses_spec_metadata;
use crate::ai_serving::planner::standard::{
apply_codex_openai_responses_special_headers, build_cross_format_openai_responses_request_body,
@@ -44,7 +47,7 @@ use crate::ai_serving::{
LocalResolvedOAuthRequestAuth, PlannerAppState,
};
use crate::ai_serving::{ConversionMode, ExecutionStrategy};
use crate::AppState;
use crate::{AppState, GatewayError};
use super::support::{
mark_skipped_local_openai_responses_candidate,
@@ -70,6 +73,7 @@ pub(crate) struct LocalOpenAiResponsesCandidatePayloadParts {
pub(super) envelope_name: Option<&'static str>,
pub(super) upstream_is_stream: bool,
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
pub(super) request_redacted: bool,
}
#[allow(clippy::too_many_arguments)]
@@ -83,7 +87,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
candidate_index: u32,
candidate_id: &str,
spec: LocalOpenAiResponsesSpec,
) -> Option<LocalOpenAiResponsesCandidatePayloadParts> {
) -> Result<Option<LocalOpenAiResponsesCandidatePayloadParts>, GatewayError> {
let spec_metadata = local_openai_responses_spec_metadata(spec);
let client_api_format = spec_metadata.api_format.trim().to_ascii_lowercase();
let planner_state = PlannerAppState::new(state);
@@ -119,7 +123,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
skip_reason,
)
.await;
return None;
return Ok(None);
}
let oauth_context = OauthPreparationContext {
@@ -147,7 +151,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
"transport_auth_unavailable",
)
.await;
return None;
return Ok(None);
}
}
} else {
@@ -186,7 +190,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
skip_reason,
)
.await;
return None;
return Ok(None);
}
}
} else {
@@ -211,7 +215,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
skip_reason,
)
.await;
return None;
return Ok(None);
}
}
};
@@ -225,6 +229,16 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
Some(&input.requested_model),
)
.await;
let redaction = resolve_provider_chat_pii_redaction(
state,
parts,
body_json,
&input.auth_context,
spec_metadata.api_format,
candidate_id,
)
.await?;
let body_json = redaction.body_json.as_ref();
let needs_bidirectional_conversion = !same_format && conversion_kind.is_some();
let upstream_is_stream = resolve_upstream_is_stream_for_provider(
@@ -287,7 +301,7 @@ pub(crate) async fn resolve_local_openai_responses_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(
@@ -328,7 +342,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
"transport_unsupported",
)
.await;
return None;
return Ok(None);
}
}
} else {
@@ -359,7 +373,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
),
)
.await;
return None;
return Ok(None);
}
}
} else {
@@ -386,6 +400,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
upstream_is_stream,
needs_bidirectional_conversion,
kiro_auth,
redaction.redacted,
)
.await;
}
@@ -421,7 +436,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
),
)
.await;
return None;
return Ok(None);
};
let extra_headers = antigravity_auth
.as_ref()
@@ -457,7 +472,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
),
)
.await;
return None;
return Ok(None);
};
let mut provider_request_headers = resolved_headers.headers;
apply_codex_openai_responses_special_headers(
@@ -469,6 +484,10 @@ pub(crate) async fn resolve_local_openai_responses_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 (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats(spec_metadata.api_format, provider_api_format);
@@ -497,7 +516,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
"gateway resolved local openai responses upstream url"
);
Some(LocalOpenAiResponsesCandidatePayloadParts {
Ok(Some(LocalOpenAiResponsesCandidatePayloadParts {
auth_header: resolved_headers.auth_header,
auth_value: resolved_headers.auth_value,
mapped_model,
@@ -516,7 +535,8 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
},
upstream_is_stream,
transport: Arc::clone(transport),
})
request_redacted: redaction.redacted,
}))
}
fn api_format_alias_matches(left: &str, right: &str) -> bool {
@@ -543,7 +563,8 @@ async fn build_kiro_openai_responses_payload_parts(
upstream_is_stream: bool,
needs_bidirectional_conversion: bool,
kiro_auth: &KiroRequestAuth,
) -> Option<LocalOpenAiResponsesCandidatePayloadParts> {
request_redacted: bool,
) -> Result<Option<LocalOpenAiResponsesCandidatePayloadParts>, GatewayError> {
let candidate = &eligible.candidate;
let provider_request_body = match build_kiro_provider_request_body(
&claude_request_body,
@@ -569,7 +590,7 @@ async fn build_kiro_openai_responses_payload_parts(
),
)
.await;
return None;
return Ok(None);
}
};
let upstream_url = match build_kiro_cross_format_upstream_url(
@@ -597,10 +618,10 @@ async fn build_kiro_openai_responses_payload_parts(
),
)
.await;
return None;
return Ok(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,
@@ -627,7 +648,7 @@ async fn build_kiro_openai_responses_payload_parts(
),
)
.await;
return None;
return Ok(None);
}
};
let (execution_strategy, conversion_mode) =
@@ -652,7 +673,12 @@ async fn build_kiro_openai_responses_payload_parts(
"gateway resolved local openai responses kiro upstream url"
);
Some(LocalOpenAiResponsesCandidatePayloadParts {
request_identity_response_encoding_when_redacted(
&mut provider_request_headers,
request_redacted,
);
Ok(Some(LocalOpenAiResponsesCandidatePayloadParts {
auth_header,
auth_value,
mapped_model,
@@ -666,5 +692,6 @@ async fn build_kiro_openai_responses_payload_parts(
envelope_name: Some(KIRO_ENVELOPE_NAME),
upstream_is_stream,
transport: Arc::clone(transport),
})
request_redacted,
}))
}
@@ -117,7 +117,7 @@ pub(crate) async fn maybe_build_sync_local_openai_responses_decision_payload(
if let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
.await
.await?
{
return Ok(Some(payload));
}
@@ -155,7 +155,7 @@ pub(crate) async fn maybe_build_stream_local_openai_responses_decision_payload(
if let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
.await
.await?
{
return Ok(Some(payload));
}
@@ -219,7 +219,7 @@ impl LocalOpenAiResponsesSyncAttemptSource<'_> {
attempt,
self.spec,
)
.await
.await?
else {
return Ok(None);
};
@@ -257,7 +257,7 @@ impl LocalOpenAiResponsesStreamAttemptSource<'_> {
attempt,
self.spec,
)
.await
.await?
else {
return Ok(None);
};
@@ -325,7 +325,7 @@ pub(super) async fn build_local_sync_plan_and_reports(
let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
.await
.await?
else {
continue;
};
@@ -397,7 +397,7 @@ pub(super) async fn build_local_stream_plan_and_reports(
let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
.await
.await?
else {
continue;
};
File diff suppressed because it is too large Load Diff
@@ -28,3 +28,4 @@ mod claude;
mod cli;
mod gemini;
mod image;
mod pii_redaction_formats;
@@ -0,0 +1,756 @@
use super::{
any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server,
to_bytes, Arc, Body, Json, Mutex, Request, Router, StatusCode,
EXECUTION_PATH_EXECUTION_RUNTIME_SYNC, EXECUTION_PATH_HEADER, TRACE_ID_HEADER,
};
use crate::data::GatewayDataState;
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::auth::{
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeyExportRecord, 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};
const ORIGINAL_EMAIL: &str = "alice@example.com";
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum StandardFormat {
OpenAiChat,
OpenAiResponses,
ClaudeMessages,
}
#[derive(Debug, Clone)]
struct SeenExecutionRuntimeSyncRequest {
body: serde_json::Value,
headers: serde_json::Value,
url: String,
}
struct RedactionFormatCase {
test_id: &'static str,
trace_id: &'static str,
client_format: StandardFormat,
provider_format: StandardFormat,
}
impl StandardFormat {
fn api_format(self) -> &'static str {
match self {
Self::OpenAiChat => "openai:chat",
Self::OpenAiResponses => "openai:responses",
Self::ClaudeMessages => "claude:messages",
}
}
fn provider_name(self) -> &'static str {
match self {
Self::OpenAiChat | Self::OpenAiResponses => "openai",
Self::ClaudeMessages => "claude",
}
}
fn endpoint_kind(self) -> &'static str {
match self {
Self::OpenAiChat | Self::ClaudeMessages => "chat",
Self::OpenAiResponses => "cli",
}
}
fn client_path(self) -> &'static str {
match self {
Self::OpenAiChat => "/v1/chat/completions",
Self::OpenAiResponses => "/v1/responses",
Self::ClaudeMessages => "/v1/messages",
}
}
fn upstream_base_url(self) -> &'static str {
match self {
Self::OpenAiChat | Self::OpenAiResponses => "https://api.openai.example",
Self::ClaudeMessages => "https://api.anthropic.example",
}
}
fn upstream_path(self) -> &'static str {
match self {
Self::OpenAiChat => "/custom/v1/chat/completions",
Self::OpenAiResponses => "/custom/v1/responses",
Self::ClaudeMessages => "/custom/v1/messages",
}
}
fn client_model(self) -> &'static str {
match self {
Self::OpenAiChat | Self::OpenAiResponses => "gpt-5",
Self::ClaudeMessages => "claude-sonnet-4-5",
}
}
fn provider_model(self) -> &'static str {
match self {
Self::OpenAiChat | Self::OpenAiResponses => "gpt-5-upstream",
Self::ClaudeMessages => "claude-sonnet-4-5-upstream",
}
}
fn provider_auth_type(self) -> &'static str {
match self {
Self::OpenAiChat | Self::OpenAiResponses => "api_key",
Self::ClaudeMessages => "api_key",
}
}
fn client_request_body(self) -> serde_json::Value {
match self {
Self::OpenAiChat => json!({
"model": self.client_model(),
"messages": [
{"role": "system", "content": "Keep answers short."},
{"role": "user", "content": format!("Please contact {ORIGINAL_EMAIL}")}
]
}),
Self::OpenAiResponses => json!({
"model": self.client_model(),
"instructions": format!("Never expose {ORIGINAL_EMAIL}."),
"input": [{
"type": "message",
"role": "user",
"content": [{
"type": "input_text",
"text": format!("Send a status update to {ORIGINAL_EMAIL}")
}]
}],
"store": false
}),
Self::ClaudeMessages => json!({
"model": self.client_model(),
"system": format!("The private contact is {ORIGINAL_EMAIL}."),
"messages": [{
"role": "user",
"content": [{
"type": "text",
"text": format!("Draft a reply for {ORIGINAL_EMAIL}")
}]
}],
"max_tokens": 64
}),
}
}
fn execution_runtime_response_body(self, sentinel: &str) -> serde_json::Value {
let restored_text = format!("restored {sentinel}");
match self {
Self::OpenAiChat => json!({
"id": "chatcmpl-redaction-format",
"object": "chat.completion",
"model": self.provider_model(),
"choices": [{
"index": 0,
"message": {"role": "assistant", "content": restored_text},
"finish_reason": "stop"
}],
"usage": {
"prompt_tokens": 2,
"completion_tokens": 3,
"total_tokens": 5
}
}),
Self::OpenAiResponses => json!({
"id": "resp-redaction-format",
"object": "response",
"status": "completed",
"model": self.provider_model(),
"output": [{
"type": "message",
"id": "resp-redaction-format-msg",
"role": "assistant",
"status": "completed",
"content": [{
"type": "output_text",
"text": restored_text,
"annotations": []
}]
}],
"usage": {
"input_tokens": 2,
"output_tokens": 3,
"total_tokens": 5
}
}),
Self::ClaudeMessages => json!({
"id": "msg_redaction_format",
"type": "message",
"model": self.provider_model(),
"role": "assistant",
"content": [{"type": "text", "text": restored_text}],
"stop_reason": "end_turn",
"usage": {
"input_tokens": 2,
"output_tokens": 3
}
}),
}
}
}
#[tokio::test]
async fn ai_execute_openai_responses_pii_redaction_round_trip_same_format() {
let (response_json, seen) = run_redaction_format_case(RedactionFormatCase {
test_id: "openai-responses-pii-redaction-same-format",
trace_id: "trace-openai-responses-pii-redaction-same-format",
client_format: StandardFormat::OpenAiResponses,
provider_format: StandardFormat::OpenAiResponses,
})
.await;
assert_provider_request_redacted(&seen, StandardFormat::OpenAiResponses);
assert!(seen.body.get("input").is_some());
assert_restored_response(&response_json, StandardFormat::OpenAiResponses);
}
#[tokio::test]
async fn ai_execute_claude_messages_pii_redaction_round_trip_same_format() {
let (response_json, seen) = run_redaction_format_case(RedactionFormatCase {
test_id: "claude-messages-pii-redaction-same-format",
trace_id: "trace-claude-messages-pii-redaction-same-format",
client_format: StandardFormat::ClaudeMessages,
provider_format: StandardFormat::ClaudeMessages,
})
.await;
assert_provider_request_redacted(&seen, StandardFormat::ClaudeMessages);
assert!(seen.body.get("messages").is_some());
assert_restored_response(&response_json, StandardFormat::ClaudeMessages);
}
#[tokio::test]
async fn ai_execute_openai_chat_pii_redaction_before_claude_conversion() {
let (response_json, seen) = run_redaction_format_case(RedactionFormatCase {
test_id: "openai-chat-pii-redaction-before-claude-conversion",
trace_id: "trace-openai-chat-pii-redaction-before-claude-conversion",
client_format: StandardFormat::OpenAiChat,
provider_format: StandardFormat::ClaudeMessages,
})
.await;
assert_provider_request_redacted(&seen, StandardFormat::ClaudeMessages);
assert!(seen.body.get("messages").is_some());
assert_eq!(
seen.body["model"],
StandardFormat::ClaudeMessages.provider_model()
);
assert_restored_response(&response_json, StandardFormat::OpenAiChat);
}
#[tokio::test]
async fn ai_execute_openai_responses_pii_redaction_before_claude_conversion() {
let (response_json, seen) = run_redaction_format_case(RedactionFormatCase {
test_id: "openai-responses-pii-redaction-before-claude-conversion",
trace_id: "trace-openai-responses-pii-redaction-before-claude-conversion",
client_format: StandardFormat::OpenAiResponses,
provider_format: StandardFormat::ClaudeMessages,
})
.await;
assert_provider_request_redacted(&seen, StandardFormat::ClaudeMessages);
assert!(seen.body.get("messages").is_some());
assert_eq!(
seen.body["model"],
StandardFormat::ClaudeMessages.provider_model()
);
assert_restored_response(&response_json, StandardFormat::OpenAiResponses);
}
#[tokio::test]
async fn ai_execute_claude_messages_pii_redaction_before_openai_chat_conversion() {
let (response_json, seen) = run_redaction_format_case(RedactionFormatCase {
test_id: "claude-messages-pii-redaction-before-openai-chat-conversion",
trace_id: "trace-claude-messages-pii-redaction-before-openai-chat-conversion",
client_format: StandardFormat::ClaudeMessages,
provider_format: StandardFormat::OpenAiChat,
})
.await;
assert_provider_request_redacted(&seen, StandardFormat::OpenAiChat);
assert!(seen.body.get("messages").is_some());
assert_eq!(
seen.body["model"],
StandardFormat::OpenAiChat.provider_model()
);
assert_restored_response(&response_json, StandardFormat::ClaudeMessages);
}
async fn run_redaction_format_case(
case: RedactionFormatCase,
) -> (serde_json::Value, SeenExecutionRuntimeSyncRequest) {
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeSyncRequest>));
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
let provider_format = case.provider_format;
let trace_id = case.trace_id.to_string();
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |request: Request| {
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
let trace_id = trace_id.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("execution runtime payload should parse");
let provider_body = payload
.get("body")
.and_then(|value| value.get("json_body"))
.cloned()
.expect("json body should exist");
let provider_body_text =
serde_json::to_string(&provider_body).expect("json body should serialize");
let email_sentinel = collect_sentinels(&provider_body_text, "EMAIL")
.into_iter()
.next()
.expect("email sentinel should exist in provider body");
let headers = payload.get("headers").cloned().unwrap_or_else(|| json!({}));
let url = payload
.get("url")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string();
*seen_execution_runtime_inner
.lock()
.expect("mutex should lock") = Some(SeenExecutionRuntimeSyncRequest {
body: provider_body,
headers,
url,
});
Json(json!({
"request_id": trace_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": provider_format.execution_runtime_response_body(&email_sentinel)
},
"telemetry": {
"elapsed_ms": 19
}
}))
}
}),
);
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let auth_repository = auth_repository(&case);
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
candidate_row(&case),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider(&case)],
vec![endpoint(&case)],
vec![key(&case)],
));
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url.clone())
.with_data_state_for_tests(
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()),
);
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let mut request = client
.post(format!("{gateway_url}{}", case.client_format.client_path()))
.header(http::header::CONTENT_TYPE, "application/json")
.header(http::header::ACCEPT_ENCODING, "gzip")
.header(TRACE_ID_HEADER, case.trace_id)
.body(case.client_format.client_request_body().to_string());
request = match case.client_format {
StandardFormat::OpenAiChat | StandardFormat::OpenAiResponses => request.header(
http::header::AUTHORIZATION,
format!("Bearer {}", client_api_key(&case)),
),
StandardFormat::ClaudeMessages => request
.header("x-api-key", client_api_key(&case))
.header("anthropic-version", "2023-06-01"),
};
let response = request.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("response 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 mut stored_candidates = Vec::new();
for _ in 0..50 {
stored_candidates = request_candidate_repository
.list_by_request_id(case.trace_id)
.await
.expect("request candidate trace should read");
if stored_candidates.len() == 1
&& stored_candidates[0].status == RequestCandidateStatus::Success
{
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
assert_eq!(stored_candidates.len(), 1);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
let seen = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("execution runtime request should be captured");
gateway_handle.abort();
execution_runtime_handle.abort();
(response_json, seen)
}
fn auth_repository(case: &RedactionFormatCase) -> Arc<InMemoryAuthApiKeySnapshotRepository> {
let snapshot = auth_snapshot(case);
let key_hash = hash_api_key(&client_api_key(case));
Arc::new(
InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(key_hash.clone()),
snapshot.clone(),
)])
.with_export_records(vec![auth_export_record(
&snapshot,
key_hash,
Some(json!({
"chat_pii_redaction": {
"enabled": true,
"inject_model_instruction": true
}
})),
)]),
)
}
fn auth_snapshot(case: &RedactionFormatCase) -> StoredAuthApiKeySnapshot {
let allowed_providers = unique_json_array([
case.client_format.provider_name(),
case.provider_format.provider_name(),
]);
StoredAuthApiKeySnapshot::new(
format!("user-{}", case.test_id),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
Some(allowed_providers.clone()),
Some(json!([case.client_format.api_format()])),
Some(json!([case.client_format.client_model()])),
format!("api-key-{}", case.test_id),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(4_102_444_800),
Some(allowed_providers),
Some(json!([case.client_format.api_format()])),
Some(json!([case.client_format.client_model()])),
)
.expect("auth snapshot should build")
}
fn auth_export_record(
snapshot: &StoredAuthApiKeySnapshot,
key_hash: String,
feature_settings: Option<serde_json::Value>,
) -> StoredAuthApiKeyExportRecord {
StoredAuthApiKeyExportRecord::new(
snapshot.user_id.clone(),
snapshot.api_key_id.clone(),
key_hash,
None,
snapshot.api_key_name.clone(),
snapshot
.api_key_allowed_providers
.as_ref()
.map(|value| serde_json::json!(value)),
snapshot
.api_key_allowed_api_formats
.as_ref()
.map(|value| serde_json::json!(value)),
snapshot
.api_key_allowed_models
.as_ref()
.map(|value| serde_json::json!(value)),
snapshot.api_key_rate_limit,
snapshot.api_key_concurrent_limit,
None,
snapshot.api_key_is_active,
snapshot
.api_key_expires_at_unix_secs
.map(|value| value as i64),
false,
0,
0,
0.0,
snapshot.api_key_is_standalone,
)
.expect("auth api key export record should build")
.with_feature_settings(feature_settings)
}
fn candidate_row(case: &RedactionFormatCase) -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: format!("provider-{}", case.test_id),
provider_name: case.provider_format.provider_name().to_string(),
provider_type: "custom".to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: format!("endpoint-{}", case.test_id),
endpoint_api_format: case.provider_format.api_format().to_string(),
endpoint_api_family: Some(case.provider_format.provider_name().to_string()),
endpoint_kind: Some(case.provider_format.endpoint_kind().to_string()),
endpoint_is_active: true,
key_id: format!("key-{}", case.test_id),
key_name: "prod".to_string(),
key_auth_type: case.provider_format.provider_auth_type().to_string(),
key_is_active: true,
key_api_formats: Some(vec![case.provider_format.api_format().to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 5,
key_global_priority_by_format: Some(json!({case.provider_format.api_format(): 1})),
model_id: format!("model-{}", case.test_id),
global_model_id: format!("global-model-{}", case.test_id),
global_model_name: case.client_format.client_model().to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(true),
model_provider_model_name: case.provider_format.provider_model().to_string(),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: case.provider_format.provider_model().to_string(),
priority: 1,
api_formats: Some(vec![case.provider_format.api_format().to_string()]),
endpoint_ids: Some(vec![format!("endpoint-{}", case.test_id)]),
}]),
model_supports_streaming: Some(true),
model_is_active: true,
model_is_available: true,
}
}
fn provider(case: &RedactionFormatCase) -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
format!("provider-{}", case.test_id),
case.provider_format.provider_name().to_string(),
Some("https://example.com".to_string()),
"custom".to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
case.client_format != case.provider_format,
None,
Some(2),
None,
Some(20.0),
None,
None,
)
}
fn endpoint(case: &RedactionFormatCase) -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
format!("endpoint-{}", case.test_id),
format!("provider-{}", case.test_id),
case.provider_format.api_format().to_string(),
Some(case.provider_format.provider_name().to_string()),
Some(case.provider_format.endpoint_kind().to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
case.provider_format.upstream_base_url().to_string(),
None,
None,
Some(2),
Some(case.provider_format.upstream_path().to_string()),
None,
None,
None,
)
.expect("endpoint transport should build")
}
fn key(case: &RedactionFormatCase) -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
format!("key-{}", case.test_id),
format!("provider-{}", case.test_id),
"prod".to_string(),
case.provider_format.provider_auth_type().to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(json!([case.provider_format.api_format()])),
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
&format!("sk-upstream-{}", case.test_id),
)
.expect("api key should encrypt"),
None,
None,
Some(json!({case.provider_format.api_format(): 1})),
None,
None,
None,
None,
)
.expect("key transport should build")
}
fn redaction_config() -> Vec<(String, serde_json::Value)> {
vec![
("module.chat_pii_redaction.enabled".to_string(), json!(true)),
(
"module.chat_pii_redaction.rules".to_string(),
json!([{
"id": "email",
"name": "邮箱",
"pattern": r"(?i)[A-Z0-9._%+-]{1,64}@[A-Z0-9.-]{1,253}\.[A-Z]{2,63}",
"enabled": true,
"features": {"validator": "email"},
"system": true
}]),
),
(
"module.chat_pii_redaction.cache_ttl_seconds".to_string(),
json!(300),
),
]
}
fn assert_provider_request_redacted(
seen: &SeenExecutionRuntimeSyncRequest,
provider_format: StandardFormat,
) {
let provider_body_text = serde_json::to_string(&seen.body).expect("body should serialize");
assert!(
provider_body_text.contains("<AETHER:EMAIL:"),
"provider body was not redacted: {provider_body_text}"
);
assert!(
!provider_body_text.contains(ORIGINAL_EMAIL),
"provider body leaked original email: {provider_body_text}"
);
assert_eq!(seen.headers["accept-encoding"], "identity");
assert!(
seen.url.ends_with(provider_format.upstream_path()),
"unexpected provider url {}",
seen.url
);
}
fn assert_restored_response(response_json: &serde_json::Value, client_format: StandardFormat) {
match client_format {
StandardFormat::OpenAiChat => assert!(response_json["choices"].is_array()),
StandardFormat::ClaudeMessages => assert_eq!(response_json["type"], "message"),
StandardFormat::OpenAiResponses => {}
}
let mut strings = Vec::new();
collect_json_strings(response_json, &mut strings);
assert!(
strings.iter().any(|value| value.contains(ORIGINAL_EMAIL)),
"client response did not restore original email: {response_json}"
);
assert!(
strings.iter().all(|value| !value.contains("<AETHER:")),
"client response still contains redaction sentinel: {response_json}"
);
}
fn collect_json_strings<'a>(value: &'a serde_json::Value, strings: &mut Vec<&'a str>) {
match value {
serde_json::Value::String(value) => strings.push(value),
serde_json::Value::Array(items) => {
for item in items {
collect_json_strings(item, strings);
}
}
serde_json::Value::Object(map) => {
for value in map.values() {
collect_json_strings(value, strings);
}
}
_ => {}
}
}
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
}
fn hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());
format!("{:x}", hasher.finalize())
}
fn client_api_key(case: &RedactionFormatCase) -> String {
format!("sk-client-{}", case.test_id)
}
fn unique_json_array(values: [&str; 2]) -> serde_json::Value {
let mut unique = Vec::new();
for value in values {
if !unique.contains(&value) {
unique.push(value);
}
}
json!(unique)
}
@@ -42,6 +42,9 @@
<p class="max-w-3xl text-sm text-muted-foreground">
管理员只配置功能是否启用和匹配规则用户用户 Key独立余额 Key 可在各自配置中附加此功能
</p>
<p class="max-w-3xl text-xs text-muted-foreground">
当前支持 OpenAI Chat CompletionsOpenAI ResponsesClaude Messages同格式转发和已支持的跨格式转换都会在发送给供应商前替换占位符
</p>
</div>
<div class="flex items-center gap-3 rounded-xl border border-border bg-muted/40 px-4 py-3">
<div class="text-right">