mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-04 16:37:46 +08:00
test(gateway): cover chat pii redaction formats
This commit is contained in:
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
+12
-188
@@ -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(),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
+16
-7
@@ -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(),
|
||||
},
|
||||
))
|
||||
)))
|
||||
}
|
||||
|
||||
+47
-20
@@ -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;
|
||||
};
|
||||
|
||||
Reference in New Issue
Block a user