use original request sources for endpoint conditions

This commit is contained in:
fawney19
2026-05-06 23:48:31 +08:00
parent e2bffbbaca
commit 4e9f063385
31 changed files with 524 additions and 135 deletions

View File

@@ -129,6 +129,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
&prepared.mapped_model, &prepared.mapped_model,
spec, spec,
prepared.transport.endpoint.body_rules.as_ref(), prepared.transport.endpoint.body_rules.as_ref(),
Some(&parts.headers),
prepared.upstream_is_stream, prepared.upstream_is_stream,
prepared.kiro_auth.as_ref(), prepared.kiro_auth.as_ref(),
prepared.is_claude_code, prepared.is_claude_code,

View File

@@ -12,6 +12,7 @@ pub(crate) fn build_same_format_provider_request_body(
mapped_model: &str, mapped_model: &str,
spec: LocalSameFormatProviderSpec, spec: LocalSameFormatProviderSpec,
body_rules: Option<&Value>, body_rules: Option<&Value>,
request_headers: Option<&http::HeaderMap>,
upstream_is_stream: bool, upstream_is_stream: bool,
kiro_auth: Option<&crate::ai_serving::transport::kiro::KiroRequestAuth>, kiro_auth: Option<&crate::ai_serving::transport::kiro::KiroRequestAuth>,
is_claude_code: bool, is_claude_code: bool,
@@ -25,6 +26,7 @@ pub(crate) fn build_same_format_provider_request_body(
source_model: body_json.get("model").and_then(Value::as_str), source_model: body_json.get("model").and_then(Value::as_str),
family: same_format_provider_family(spec.family), family: same_format_provider_family(spec.family),
body_rules, body_rules,
request_headers,
upstream_is_stream, upstream_is_stream,
kiro_auth_config: kiro_auth.map(|auth| &auth.auth_config), kiro_auth_config: kiro_auth.map(|auth| &auth.auth_config),
is_claude_code, is_claude_code,

View File

@@ -103,6 +103,7 @@ pub(super) async fn resolve_local_gemini_files_candidate_payload_parts(
body_is_empty, body_is_empty,
spec_metadata.decision_kind == GEMINI_FILES_UPLOAD_PLAN_KIND, spec_metadata.decision_kind == GEMINI_FILES_UPLOAD_PLAN_KIND,
transport.endpoint.body_rules.as_ref(), transport.endpoint.body_rules.as_ref(),
Some(&parts.headers),
) { ) {
Ok(parts) => parts, Ok(parts) => parts,
Err(GeminiFilesRequestBodyError::BodyRulesUnsupportedForBinaryUpload) => { Err(GeminiFilesRequestBodyError::BodyRulesUnsupportedForBinaryUpload) => {

View File

@@ -124,6 +124,7 @@ pub(super) async fn resolve_local_video_create_candidate_payload_parts(
provider_family, provider_family,
&mapped_model, &mapped_model,
transport.endpoint.body_rules.as_ref(), transport.endpoint.body_rules.as_ref(),
Some(&parts.headers),
) else { ) else {
mark_skipped_local_video_candidate_with_failure_diagnostic( mark_skipped_local_video_candidate_with_failure_diagnostic(
state, state,

View File

@@ -176,7 +176,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
) )
.await; .await;
let mut provider_request_body = let mut provider_request_body =
match crate::ai_serving::planner::standard::build_standard_request_body_with_model_directives( match crate::ai_serving::planner::standard::build_standard_request_body_with_model_directives_and_request_headers(
body_json, body_json,
spec_metadata.api_format, spec_metadata.api_format,
&prepared_candidate.mapped_model, &prepared_candidate.mapped_model,
@@ -190,6 +190,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
transport.endpoint.body_rules.as_ref() transport.endpoint.body_rules.as_ref()
}, },
Some(input.auth_context.api_key_id.as_str()), Some(input.auth_context.api_key_id.as_str()),
Some(&parts.headers),
enable_model_directives, enable_model_directives,
) { ) {
Some(body) => body, Some(body) => body,
@@ -353,6 +354,7 @@ async fn build_kiro_cross_format_payload_parts(
&mapped_model, &mapped_model,
&kiro_auth.auth_config, &kiro_auth.auth_config,
transport.endpoint.body_rules.as_ref(), transport.endpoint.body_rules.as_ref(),
Some(&parts.headers),
) { ) {
Some(body) => body, Some(body) => body,
None => { None => {

View File

@@ -51,6 +51,7 @@ pub(crate) use crate::ai_serving::{
}; };
pub(crate) use crate::ai_serving::{ pub(crate) use crate::ai_serving::{
build_standard_request_body, build_standard_request_body_with_model_directives, build_standard_request_body, build_standard_request_body_with_model_directives,
build_standard_request_body_with_model_directives_and_request_headers,
convert_openai_chat_request_to_claude_request, convert_openai_chat_request_to_gemini_request, convert_openai_chat_request_to_claude_request, convert_openai_chat_request_to_gemini_request,
convert_openai_chat_request_to_openai_responses_request, extract_openai_text_content, convert_openai_chat_request_to_openai_responses_request, extract_openai_text_content,
normalize_openai_responses_request_to_openai_chat_request, parse_openai_tool_result_content, normalize_openai_responses_request_to_openai_chat_request, parse_openai_tool_result_content,

View File

@@ -1,6 +1,6 @@
use serde_json::Value; use serde_json::Value;
use crate::ai_serving::transport::apply_standard_provider_request_body_rules; use crate::ai_serving::transport::apply_standard_provider_request_body_rules_with_request_headers;
use crate::ai_serving::{ use crate::ai_serving::{
apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_body_edits,
apply_openai_responses_compact_special_body_edits, apply_openai_responses_compact_special_body_edits,
@@ -14,6 +14,7 @@ pub(crate) fn build_local_openai_chat_request_body(
mapped_model: &str, mapped_model: &str,
upstream_is_stream: bool, upstream_is_stream: bool,
body_rules: Option<&Value>, body_rules: Option<&Value>,
request_headers: &http::HeaderMap,
enable_model_directives: bool, enable_model_directives: bool,
) -> Option<Value> { ) -> Option<Value> {
let provider_request_body = surface_build_local_openai_chat_request_body( let provider_request_body = surface_build_local_openai_chat_request_body(
@@ -22,7 +23,12 @@ pub(crate) fn build_local_openai_chat_request_body(
upstream_is_stream, upstream_is_stream,
enable_model_directives, enable_model_directives,
)?; )?;
apply_standard_provider_request_body_rules(provider_request_body, body_rules, body_json) apply_standard_provider_request_body_rules_with_request_headers(
provider_request_body,
body_rules,
body_json,
request_headers,
)
} }
pub(crate) fn build_local_openai_chat_upstream_url( pub(crate) fn build_local_openai_chat_upstream_url(
@@ -40,6 +46,7 @@ pub(crate) fn build_cross_format_openai_chat_request_body(
upstream_is_stream: bool, upstream_is_stream: bool,
body_rules: Option<&Value>, body_rules: Option<&Value>,
user_api_key_id: Option<&str>, user_api_key_id: Option<&str>,
request_headers: &http::HeaderMap,
enable_model_directives: bool, enable_model_directives: bool,
) -> Option<Value> { ) -> Option<Value> {
let provider_request_body = surface_build_cross_format_openai_chat_request_body( let provider_request_body = surface_build_cross_format_openai_chat_request_body(
@@ -50,7 +57,12 @@ pub(crate) fn build_cross_format_openai_chat_request_body(
enable_model_directives, enable_model_directives,
)?; )?;
let mut provider_request_body = let mut provider_request_body =
apply_standard_provider_request_body_rules(provider_request_body, body_rules, body_json)?; apply_standard_provider_request_body_rules_with_request_headers(
provider_request_body,
body_rules,
body_json,
request_headers,
)?;
apply_codex_openai_responses_special_body_edits( apply_codex_openai_responses_special_body_edits(
&mut provider_request_body, &mut provider_request_body,
provider_type, provider_type,

View File

@@ -1,6 +1,6 @@
use serde_json::Value; use serde_json::Value;
use crate::ai_serving::transport::apply_standard_provider_request_body_rules; use crate::ai_serving::transport::apply_standard_provider_request_body_rules_with_request_headers;
use crate::ai_serving::{ use crate::ai_serving::{
apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_body_edits,
apply_openai_responses_compact_special_body_edits, apply_openai_responses_compact_special_body_edits,
@@ -17,6 +17,7 @@ pub(crate) fn build_local_openai_responses_request_body(
provider_api_format: &str, provider_api_format: &str,
body_rules: Option<&Value>, body_rules: Option<&Value>,
user_api_key_id: Option<&str>, user_api_key_id: Option<&str>,
request_headers: &http::HeaderMap,
enable_model_directives: bool, enable_model_directives: bool,
) -> Option<Value> { ) -> Option<Value> {
let provider_request_body = surface_build_local_openai_responses_request_body( let provider_request_body = surface_build_local_openai_responses_request_body(
@@ -26,7 +27,12 @@ pub(crate) fn build_local_openai_responses_request_body(
enable_model_directives, enable_model_directives,
)?; )?;
let mut provider_request_body = let mut provider_request_body =
apply_standard_provider_request_body_rules(provider_request_body, body_rules, body_json)?; apply_standard_provider_request_body_rules_with_request_headers(
provider_request_body,
body_rules,
body_json,
request_headers,
)?;
apply_codex_openai_responses_special_body_edits( apply_codex_openai_responses_special_body_edits(
&mut provider_request_body, &mut provider_request_body,
provider_type, provider_type,
@@ -50,6 +56,7 @@ pub(crate) fn build_cross_format_openai_responses_request_body(
provider_type: &str, provider_type: &str,
body_rules: Option<&Value>, body_rules: Option<&Value>,
user_api_key_id: Option<&str>, user_api_key_id: Option<&str>,
request_headers: &http::HeaderMap,
enable_model_directives: bool, enable_model_directives: bool,
) -> Option<Value> { ) -> Option<Value> {
let provider_request_body = surface_build_cross_format_openai_responses_request_body( let provider_request_body = surface_build_cross_format_openai_responses_request_body(
@@ -61,7 +68,12 @@ pub(crate) fn build_cross_format_openai_responses_request_body(
enable_model_directives, enable_model_directives,
)?; )?;
let mut provider_request_body = let mut provider_request_body =
apply_standard_provider_request_body_rules(provider_request_body, body_rules, body_json)?; apply_standard_provider_request_body_rules_with_request_headers(
provider_request_body,
body_rules,
body_json,
request_headers,
)?;
apply_codex_openai_responses_special_body_edits( apply_codex_openai_responses_special_body_edits(
&mut provider_request_body, &mut provider_request_body,
provider_type, provider_type,

View File

@@ -90,6 +90,7 @@ fn builds_openai_chat_cross_format_request_body_from_openai_responses_source() {
"openai", "openai",
None, None,
None, None,
&http::HeaderMap::new(),
false, false,
) )
.expect("openai responses to openai chat body should build"); .expect("openai responses to openai chat body should build");
@@ -124,6 +125,7 @@ fn local_openai_responses_wrapper_preserves_body_order_after_edits() {
"openai:responses", "openai:responses",
None, None,
Some("key-123"), Some("key-123"),
&http::HeaderMap::new(),
false, false,
) )
.expect("local openai responses body should build"); .expect("local openai responses body should build");
@@ -162,6 +164,7 @@ fn local_openai_responses_compact_wrapper_strips_store_for_same_format_requests(
"openai:responses:compact", "openai:responses:compact",
None, None,
None, None,
&http::HeaderMap::new(),
false, false,
) )
.expect("local openai compact body should build"); .expect("local openai compact body should build");
@@ -188,6 +191,7 @@ fn local_openai_responses_wrapper_applies_model_directive_before_body_rules() {
"openai:responses", "openai:responses",
Some(&body_rules), Some(&body_rules),
None, None,
&http::HeaderMap::new(),
true, true,
) )
.expect("local openai responses body should build"); .expect("local openai responses body should build");
@@ -236,6 +240,7 @@ fn strips_metadata_for_codex_openai_responses_requests() {
"codex", "codex",
None, None,
None, None,
&http::HeaderMap::new(),
false, false,
) )
.expect("claude cli to codex request should build"); .expect("claude cli to codex request should build");
@@ -269,6 +274,7 @@ fn applies_codex_defaults_unless_body_rules_handle_the_field() {
"codex", "codex",
Some(&body_rules), Some(&body_rules),
None, None,
&http::HeaderMap::new(),
false, false,
) )
.expect("claude cli to codex request should build"); .expect("claude cli to codex request should build");
@@ -297,6 +303,7 @@ fn injects_codex_prompt_cache_key_for_openai_responses_cross_format_requests() {
"codex", "codex",
None, None,
Some("key-123"), Some("key-123"),
&http::HeaderMap::new(),
false, false,
) )
.expect("claude cli to codex request should build"); .expect("claude cli to codex request should build");
@@ -325,6 +332,7 @@ fn injects_codex_prompt_cache_key_for_openai_chat_cross_format_requests() {
false, false,
None, None,
Some("key-123"), Some("key-123"),
&http::HeaderMap::new(),
false, false,
) )
.expect("openai chat to codex request should build"); .expect("openai chat to codex request should build");

View File

@@ -129,6 +129,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
&prepared_candidate.mapped_model, &prepared_candidate.mapped_model,
upstream_is_stream, upstream_is_stream,
transport.endpoint.body_rules.as_ref(), transport.endpoint.body_rules.as_ref(),
&parts.headers,
enable_model_directives, enable_model_directives,
) else { ) else {
mark_skipped_local_openai_chat_candidate_with_extra_data( mark_skipped_local_openai_chat_candidate_with_extra_data(
@@ -354,6 +355,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
transport.endpoint.body_rules.as_ref() transport.endpoint.body_rules.as_ref()
}, },
Some(input.auth_context.api_key_id.as_str()), Some(input.auth_context.api_key_id.as_str()),
&parts.headers,
enable_model_directives, enable_model_directives,
) else { ) else {
mark_skipped_local_openai_chat_candidate_with_extra_data( mark_skipped_local_openai_chat_candidate_with_extra_data(
@@ -527,6 +529,7 @@ async fn build_kiro_openai_chat_cross_format_payload_parts(
&mapped_model, &mapped_model,
&kiro_auth.auth_config, &kiro_auth.auth_config,
transport.endpoint.body_rules.as_ref(), transport.endpoint.body_rules.as_ref(),
Some(&parts.headers),
) { ) {
Some(body) => body, Some(body) => body,
None => { None => {

View File

@@ -244,6 +244,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
transport.endpoint.body_rules.as_ref() transport.endpoint.body_rules.as_ref()
}, },
Some(input.auth_context.api_key_id.as_str()), Some(input.auth_context.api_key_id.as_str()),
&parts.headers,
enable_model_directives, enable_model_directives,
) )
} else { } else {
@@ -259,6 +260,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
transport.endpoint.body_rules.as_ref() transport.endpoint.body_rules.as_ref()
}, },
Some(input.auth_context.api_key_id.as_str()), Some(input.auth_context.api_key_id.as_str()),
&parts.headers,
enable_model_directives, enable_model_directives,
) )
}) else { }) else {
@@ -532,6 +534,7 @@ async fn build_kiro_openai_responses_payload_parts(
&mapped_model, &mapped_model,
&kiro_auth.auth_config, &kiro_auth.auth_config,
transport.endpoint.body_rules.as_ref(), transport.endpoint.body_rules.as_ref(),
Some(&parts.headers),
) { ) {
Some(body) => body, Some(body) => body,
None => { None => {

View File

@@ -20,11 +20,12 @@ pub(crate) use aether_ai_formats::api::{
build_openai_image_provider_request_body, build_openai_responses_response, build_openai_image_provider_request_body, build_openai_responses_response,
build_standard_request_body, build_standard_request_body_from_canonical, build_standard_request_body, build_standard_request_body_from_canonical,
build_standard_request_body_from_canonical_with_model_directives, build_standard_request_body_from_canonical_with_model_directives,
build_standard_request_body_with_model_directives, calculate_kiro_context_input_tokens, build_standard_request_body_with_model_directives,
canonicalize_tool_arguments, convert_claude_chat_response_to_openai_chat, build_standard_request_body_with_model_directives_and_request_headers,
convert_claude_response_to_openai_responses, convert_gemini_chat_response_to_openai_chat, calculate_kiro_context_input_tokens, canonicalize_tool_arguments,
convert_gemini_response_to_openai_responses, convert_openai_chat_request_to_claude_request, convert_claude_chat_response_to_openai_chat, convert_claude_response_to_openai_responses,
convert_openai_chat_request_to_gemini_request, convert_gemini_chat_response_to_openai_chat, convert_gemini_response_to_openai_responses,
convert_openai_chat_request_to_claude_request, convert_openai_chat_request_to_gemini_request,
convert_openai_chat_request_to_openai_responses_request, convert_openai_chat_request_to_openai_responses_request,
convert_openai_chat_response_to_claude_chat, convert_openai_chat_response_to_gemini_chat, convert_openai_chat_response_to_claude_chat, convert_openai_chat_response_to_gemini_chat,
convert_openai_chat_response_to_openai_responses, convert_openai_chat_response_to_openai_responses,

View File

@@ -47,11 +47,13 @@ pub(crate) mod vertex {
} }
pub(crate) use aether_provider_transport::{ pub(crate) use aether_provider_transport::{
append_transport_diagnostics_to_value, apply_local_body_rules, apply_local_header_rules, append_transport_diagnostics_to_value, apply_local_body_rules,
apply_standard_provider_request_body_rules, body_rules_are_locally_supported, apply_local_body_rules_with_request_headers, apply_local_header_rules,
body_rules_handle_path, build_cross_format_openai_chat_upstream_url, apply_local_header_rules_with_request_headers, apply_standard_provider_request_body_rules,
build_cross_format_openai_responses_upstream_url, build_gemini_files_headers, apply_standard_provider_request_body_rules_with_request_headers,
build_gemini_files_request_body, build_gemini_files_upstream_url, body_rules_are_locally_supported, body_rules_handle_path,
build_cross_format_openai_chat_upstream_url, build_cross_format_openai_responses_upstream_url,
build_gemini_files_headers, build_gemini_files_request_body, build_gemini_files_upstream_url,
build_kiro_cross_format_upstream_url, build_local_openai_chat_upstream_url, build_kiro_cross_format_upstream_url, build_local_openai_chat_upstream_url,
build_local_openai_responses_upstream_url, build_openai_image_headers, build_local_openai_responses_upstream_url, build_openai_image_headers,
build_openai_image_upstream_url, build_passthrough_headers, build_request_trace_proxy_value, build_openai_image_upstream_url, build_passthrough_headers, build_request_trace_proxy_value,

View File

@@ -245,11 +245,35 @@ fn provider_query_extract_request_headers(payload: &Value) -> HeaderMap {
} }
fn provider_query_build_test_request_body(payload: &Value, model: &str) -> Value { fn provider_query_build_test_request_body(payload: &Value, model: &str) -> Value {
provider_query_build_test_request_body_with_model_policy(payload, model, false)
}
fn provider_query_build_test_request_body_for_route(
payload: &Value,
model: &str,
route_path: &str,
) -> Value {
provider_query_build_test_request_body_with_model_policy(
payload,
model,
route_path.ends_with("/test-model-failover"),
)
}
fn provider_query_build_test_request_body_with_model_policy(
payload: &Value,
model: &str,
override_custom_model: bool,
) -> Value {
if let Some(mut body) = provider_query_extract_request_body(payload) { if let Some(mut body) = provider_query_extract_request_body(payload) {
if let Some(object) = body.as_object_mut() { if let Some(object) = body.as_object_mut() {
object if override_custom_model {
.entry("model".to_string()) object.insert("model".to_string(), Value::String(model.to_string()));
.or_insert_with(|| Value::String(model.to_string())); } else {
object
.entry("model".to_string())
.or_insert_with(|| Value::String(model.to_string()));
}
} }
return body; return body;
} }
@@ -791,7 +815,12 @@ async fn provider_query_execute_kiro_test_candidate(
}); });
}; };
let request_body = provider_query_build_test_request_body(payload, &candidate.effective_model); let request_body = provider_query_build_test_request_body_for_route(
payload,
&candidate.effective_model,
route_path,
);
let incoming_request_headers = provider_query_extract_request_headers(payload);
let request_model = let request_model =
provider_query_request_body_model(&request_body, &candidate.effective_model); provider_query_request_body_model(&request_body, &candidate.effective_model);
let provider_request_body = match build_kiro_provider_request_body( let provider_request_body = match build_kiro_provider_request_body(
@@ -799,6 +828,7 @@ async fn provider_query_execute_kiro_test_candidate(
request_model, request_model,
&kiro_auth.auth_config, &kiro_auth.auth_config,
transport.endpoint.body_rules.as_ref(), transport.endpoint.body_rules.as_ref(),
Some(&incoming_request_headers),
) { ) {
Some(body) => body, Some(body) => body,
None => { None => {
@@ -821,7 +851,7 @@ async fn provider_query_execute_kiro_test_candidate(
.uri(route_path) .uri(route_path)
.body(()) .body(())
.map_err(|err| GatewayError::Internal(err.to_string()))?; .map_err(|err| GatewayError::Internal(err.to_string()))?;
*synthetic_request.headers_mut() = provider_query_extract_request_headers(payload); *synthetic_request.headers_mut() = incoming_request_headers;
let (parts, _) = synthetic_request.into_parts(); let (parts, _) = synthetic_request.into_parts();
let request_url = build_kiro_generate_assistant_response_url( let request_url = build_kiro_generate_assistant_response_url(
@@ -949,8 +979,12 @@ async fn provider_query_execute_standard_test_candidate(
}); });
} }
let original_request_body = let original_request_body = provider_query_build_test_request_body_for_route(
provider_query_build_test_request_body(payload, &candidate.effective_model); payload,
&candidate.effective_model,
route_path,
);
let incoming_request_headers = provider_query_extract_request_headers(payload);
let mut request_body = original_request_body.clone(); let mut request_body = original_request_body.clone();
if let Some(object) = request_body.as_object_mut() { if let Some(object) = request_body.as_object_mut() {
object.insert("stream".to_string(), Value::Bool(false)); object.insert("stream".to_string(), Value::Bool(false));
@@ -975,10 +1009,11 @@ async fn provider_query_execute_standard_test_candidate(
format!("Provider request body could not be built for {provider_api_format}"), format!("Provider request body could not be built for {provider_api_format}"),
)); ));
}; };
if !crate::provider_transport::apply_local_body_rules( if !crate::provider_transport::apply_local_body_rules_with_request_headers(
&mut provider_request_body, &mut provider_request_body,
transport.endpoint.body_rules.as_ref(), transport.endpoint.body_rules.as_ref(),
Some(&request_body), Some(&request_body),
Some(&incoming_request_headers),
) { ) {
return Ok(provider_query_skipped_execution_outcome( return Ok(provider_query_skipped_execution_outcome(
request_body.clone(), request_body.clone(),
@@ -1001,10 +1036,11 @@ async fn provider_query_execute_standard_test_candidate(
format!("Provider request body could not be built for {provider_api_format}"), format!("Provider request body could not be built for {provider_api_format}"),
)); ));
}; };
if !crate::provider_transport::apply_local_body_rules( if !crate::provider_transport::apply_local_body_rules_with_request_headers(
&mut provider_request_body, &mut provider_request_body,
transport.endpoint.body_rules.as_ref(), transport.endpoint.body_rules.as_ref(),
Some(&request_body), Some(&request_body),
Some(&incoming_request_headers),
) { ) {
return Ok(provider_query_skipped_execution_outcome( return Ok(provider_query_skipped_execution_outcome(
request_body.clone(), request_body.clone(),
@@ -1027,10 +1063,11 @@ async fn provider_query_execute_standard_test_candidate(
format!("Provider request body could not be built for {provider_api_format}"), format!("Provider request body could not be built for {provider_api_format}"),
)); ));
}; };
if !crate::provider_transport::apply_local_body_rules( if !crate::provider_transport::apply_local_body_rules_with_request_headers(
&mut provider_request_body, &mut provider_request_body,
transport.endpoint.body_rules.as_ref(), transport.endpoint.body_rules.as_ref(),
Some(&request_body), Some(&request_body),
Some(&incoming_request_headers),
) { ) {
return Ok(provider_query_skipped_execution_outcome( return Ok(provider_query_skipped_execution_outcome(
request_body.clone(), request_body.clone(),
@@ -1104,7 +1141,7 @@ async fn provider_query_execute_standard_test_candidate(
.uri(route_path) .uri(route_path)
.body(()) .body(())
.map_err(|err| GatewayError::Internal(err.to_string()))?; .map_err(|err| GatewayError::Internal(err.to_string()))?;
*synthetic_request.headers_mut() = provider_query_extract_request_headers(payload); *synthetic_request.headers_mut() = incoming_request_headers;
let (parts, _) = synthetic_request.into_parts(); let (parts, _) = synthetic_request.into_parts();
let request_url = crate::provider_transport::build_transport_request_url( let request_url = crate::provider_transport::build_transport_request_url(
@@ -1166,12 +1203,13 @@ async fn provider_query_execute_standard_test_candidate(
} else { } else {
vec![auth_header.as_deref().unwrap_or_default(), "content-type"] vec![auth_header.as_deref().unwrap_or_default(), "content-type"]
}; };
if !state.apply_local_header_rules( if !crate::provider_transport::apply_local_header_rules_with_request_headers(
&mut request_headers, &mut request_headers,
transport.endpoint.header_rules.as_ref(), transport.endpoint.header_rules.as_ref(),
&protected_headers, &protected_headers,
&provider_request_body, &provider_request_body,
Some(&request_body), Some(&request_body),
Some(&parts.headers),
) { ) {
return Ok(ProviderQueryExecutionOutcome { return Ok(ProviderQueryExecutionOutcome {
status: "failed", status: "failed",
@@ -1983,6 +2021,24 @@ mod tests {
assert_eq!(body["model"], json!("fallback-model")); assert_eq!(body["model"], json!("fallback-model"));
} }
#[test]
fn provider_query_failover_request_body_overrides_custom_model() {
let payload = json!({
"request_body": {
"model": "custom-upstream-model",
"messages": []
}
});
let body = provider_query_build_test_request_body_for_route(
&payload,
"failover-model",
"/api/admin/provider-query/test-model-failover",
);
assert_eq!(body["model"], json!("failover-model"));
}
#[test] #[test]
fn provider_query_request_body_model_uses_non_empty_string_only() { fn provider_query_request_body_model_uses_non_empty_string_only() {
let custom = json!({ "model": " custom-model " }); let custom = json!({ "model": " custom-model " });

View File

@@ -2632,7 +2632,7 @@ fn ai_serving_standard_attempts_consume_eligible_local_candidates_without_transp
"build_complete_passthrough_headers_with_auth(", "build_complete_passthrough_headers_with_auth(",
"build_claude_passthrough_headers(", "build_claude_passthrough_headers(",
"build_openai_passthrough_headers(", "build_openai_passthrough_headers(",
"apply_local_header_rules(", "apply_local_header_rules_with_request_headers(",
"uses_vertex_api_key_query_auth(", "uses_vertex_api_key_query_auth(",
] { ] {
assert!( assert!(

View File

@@ -119,6 +119,7 @@ pub use crate::request::standard::{
build_local_openai_responses_request_body, build_local_openai_responses_request_body,
build_local_openai_responses_request_body_with_model_directives, build_standard_request_body, build_local_openai_responses_request_body_with_model_directives, build_standard_request_body,
build_standard_request_body_with_model_directives, build_standard_request_body_with_model_directives,
build_standard_request_body_with_model_directives_and_request_headers,
claude::{ claude::{
resolve_stream_spec as resolve_claude_stream_spec, resolve_stream_spec as resolve_claude_stream_spec,
resolve_sync_spec as resolve_claude_sync_spec, resolve_sync_spec as resolve_claude_sync_spec,

View File

@@ -9,11 +9,17 @@ use serde_json::{Map, Value};
const ORIGINAL_PLACEHOLDER: &str = "{{$original}}"; const ORIGINAL_PLACEHOLDER: &str = "{{$original}}";
const ITEM_PREFIX: &str = "$item."; const ITEM_PREFIX: &str = "$item.";
const ITEM_EXACT: &str = "$item"; const ITEM_EXACT: &str = "$item";
const CONDITION_SOURCES: &[&str] = &["current", "original"]; const CONDITION_SOURCES: &[&str] = &["body", "request_headers", "headers", "original", "current"];
const CONDITION_TYPE_VALUES: &[&str] = &["string", "number", "boolean", "array", "object", "null"]; const CONDITION_TYPE_VALUES: &[&str] = &["string", "number", "boolean", "array", "object", "null"];
static RANGE_RE: OnceLock<Regex> = OnceLock::new(); static RANGE_RE: OnceLock<Regex> = OnceLock::new();
#[derive(Clone, Copy)]
enum ConditionHeaders<'a> {
Request(&'a http::HeaderMap),
Map(&'a BTreeMap<String, String>),
}
#[derive(Debug, Clone, PartialEq, Eq)] #[derive(Debug, Clone, PartialEq, Eq)]
enum BodyPathSegment { enum BodyPathSegment {
Key(String), Key(String),
@@ -34,6 +40,35 @@ pub fn apply_local_header_rules(
protected_keys: &[&str], protected_keys: &[&str],
body: &Value, body: &Value,
original_body: Option<&Value>, original_body: Option<&Value>,
) -> bool {
apply_local_header_rules_inner(headers, rules, protected_keys, body, original_body, None)
}
pub fn apply_local_header_rules_with_request_headers(
headers: &mut BTreeMap<String, String>,
rules: Option<&Value>,
protected_keys: &[&str],
body: &Value,
original_body: Option<&Value>,
request_headers: Option<&http::HeaderMap>,
) -> bool {
apply_local_header_rules_inner(
headers,
rules,
protected_keys,
body,
original_body,
request_headers.map(ConditionHeaders::Request),
)
}
fn apply_local_header_rules_inner(
headers: &mut BTreeMap<String, String>,
rules: Option<&Value>,
protected_keys: &[&str],
body: &Value,
original_body: Option<&Value>,
request_headers: Option<ConditionHeaders<'_>>,
) -> bool { ) -> bool {
let Some(rules) = rules else { let Some(rules) = rules else {
return true; return true;
@@ -54,7 +89,8 @@ pub fn apply_local_header_rules(
if !condition_is_locally_supported(condition) { if !condition_is_locally_supported(condition) {
continue; continue;
} }
if !evaluate_local_condition(body, condition, original_body) { let condition_headers = request_headers.or(Some(ConditionHeaders::Map(&*headers)));
if !evaluate_local_condition(body, condition, original_body, condition_headers) {
continue; continue;
} }
} }
@@ -178,6 +214,29 @@ pub fn apply_local_body_rules(
body: &mut Value, body: &mut Value,
rules: Option<&Value>, rules: Option<&Value>,
original_body: Option<&Value>, original_body: Option<&Value>,
) -> bool {
apply_local_body_rules_inner(body, rules, original_body, None)
}
pub fn apply_local_body_rules_with_request_headers(
body: &mut Value,
rules: Option<&Value>,
original_body: Option<&Value>,
request_headers: Option<&http::HeaderMap>,
) -> bool {
apply_local_body_rules_inner(
body,
rules,
original_body,
request_headers.map(ConditionHeaders::Request),
)
}
fn apply_local_body_rules_inner(
body: &mut Value,
rules: Option<&Value>,
original_body: Option<&Value>,
request_headers: Option<ConditionHeaders<'_>>,
) -> bool { ) -> bool {
let Some(rules) = rules else { let Some(rules) = rules else {
return true; return true;
@@ -197,7 +256,9 @@ pub fn apply_local_body_rules(
if !condition_is_locally_supported(condition) { if !condition_is_locally_supported(condition) {
continue; continue;
} }
if !item_condition && !evaluate_local_condition(body, condition, original_body) { if !item_condition
&& !evaluate_local_condition(body, condition, original_body, request_headers)
{
continue; continue;
} }
} }
@@ -220,11 +281,14 @@ pub fn apply_local_body_rules(
let targets = iter_wildcard_targets( let targets = iter_wildcard_targets(
body, body,
&path, &path,
condition, WildcardTargetOptions {
item_condition, condition,
original_body, item_condition,
false, original_body,
false, request_headers,
require_leaf: false,
reverse: false,
},
); );
let value_template = rule.get("value").cloned().unwrap_or(Value::Null); let value_template = rule.get("value").cloned().unwrap_or(Value::Null);
for target_path in targets { for target_path in targets {
@@ -248,11 +312,14 @@ pub fn apply_local_body_rules(
for target_path in iter_wildcard_targets( for target_path in iter_wildcard_targets(
body, body,
&path, &path,
condition, WildcardTargetOptions {
item_condition, condition,
original_body, item_condition,
true, original_body,
true, request_headers,
require_leaf: true,
reverse: true,
},
) { ) {
let _ = delete_nested_value(body, &target_path); let _ = delete_nested_value(body, &target_path);
} }
@@ -289,11 +356,14 @@ pub fn apply_local_body_rules(
for target_path in iter_wildcard_targets( for target_path in iter_wildcard_targets(
body, body,
&path, &path,
condition, WildcardTargetOptions {
item_condition, condition,
original_body, item_condition,
true, original_body,
false, request_headers,
require_leaf: true,
reverse: false,
},
) { ) {
if let Some(target) = get_nested_value_mut(body, &target_path) { if let Some(target) = get_nested_value_mut(body, &target_path) {
if let Some(values) = target.as_array_mut() { if let Some(values) = target.as_array_mut() {
@@ -352,11 +422,14 @@ pub fn apply_local_body_rules(
for target_path in iter_wildcard_targets( for target_path in iter_wildcard_targets(
body, body,
&path, &path,
condition, WildcardTargetOptions {
item_condition, condition,
original_body, item_condition,
true, original_body,
false, request_headers,
require_leaf: true,
reverse: false,
},
) { ) {
if let Some(target) = get_nested_value_mut(body, &target_path) { if let Some(target) = get_nested_value_mut(body, &target_path) {
let Some(current) = target.as_str().map(str::to_owned) else { let Some(current) = target.as_str().map(str::to_owned) else {
@@ -394,7 +467,7 @@ fn condition_is_locally_supported(condition: &Value) -> bool {
.get("source") .get("source")
.and_then(Value::as_str) .and_then(Value::as_str)
.map(str::trim) .map(str::trim)
.unwrap_or("current"); .unwrap_or("body");
if !CONDITION_SOURCES.contains(&source) { if !CONDITION_SOURCES.contains(&source) {
return false; return false;
} }
@@ -402,15 +475,13 @@ fn condition_is_locally_supported(condition: &Value) -> bool {
let Some(op) = condition.get("op").and_then(Value::as_str).map(str::trim) else { let Some(op) = condition.get("op").and_then(Value::as_str).map(str::trim) else {
return false; return false;
}; };
let Some(path) = condition let Some(path) = condition.get("path").and_then(Value::as_str) else {
.get("path")
.and_then(Value::as_str)
.map(str::trim)
.and_then(parse_body_path)
else {
return false; return false;
}; };
if path.is_empty() { if path.trim().is_empty() {
return false;
}
if !condition_source_is_headers(source) && parse_body_path(path).is_none() {
return false; return false;
} }
@@ -443,6 +514,7 @@ fn evaluate_local_condition(
body: &Value, body: &Value,
condition: &Value, condition: &Value,
original_body: Option<&Value>, original_body: Option<&Value>,
request_headers: Option<ConditionHeaders<'_>>,
) -> bool { ) -> bool {
let Some(condition) = condition.as_object() else { let Some(condition) = condition.as_object() else {
return false; return false;
@@ -450,41 +522,40 @@ fn evaluate_local_condition(
if let Some(children) = condition.get("all").and_then(Value::as_array) { if let Some(children) = condition.get("all").and_then(Value::as_array) {
return !children.is_empty() return !children.is_empty()
&& children && children.iter().all(|child| {
.iter() evaluate_local_condition(body, child, original_body, request_headers)
.all(|child| evaluate_local_condition(body, child, original_body)); });
} }
if let Some(children) = condition.get("any").and_then(Value::as_array) { if let Some(children) = condition.get("any").and_then(Value::as_array) {
return !children.is_empty() return !children.is_empty()
&& children && children.iter().any(|child| {
.iter() evaluate_local_condition(body, child, original_body, request_headers)
.any(|child| evaluate_local_condition(body, child, original_body)); });
} }
let source = condition let source = condition
.get("source") .get("source")
.and_then(Value::as_str) .and_then(Value::as_str)
.map(str::trim) .map(str::trim)
.unwrap_or("current"); .unwrap_or("body");
let target = if source.eq_ignore_ascii_case("original") {
original_body.unwrap_or(body)
} else {
body
};
let Some(op) = condition.get("op").and_then(Value::as_str).map(str::trim) else { let Some(op) = condition.get("op").and_then(Value::as_str).map(str::trim) else {
return false; return false;
}; };
let Some(path) = condition let Some(path) = condition.get("path").and_then(Value::as_str).map(str::trim) else {
.get("path")
.and_then(Value::as_str)
.map(str::trim)
.and_then(parse_body_path)
else {
return false; return false;
}; };
let current_value = get_nested_value(target, &path); let current_value = if condition_source_is_headers(source) {
request_headers.and_then(|headers| get_header_condition_value(headers, path))
} else {
let target = if source.eq_ignore_ascii_case("current") {
body
} else {
original_body.unwrap_or(body)
};
parse_body_path(path).and_then(|path| get_nested_value(target, &path))
};
if op == "exists" { if op == "exists" {
return current_value.is_some(); return current_value.is_some();
} }
@@ -556,6 +627,24 @@ fn evaluate_local_condition(
} }
} }
fn condition_source_is_headers(source: &str) -> bool {
source.eq_ignore_ascii_case("request_headers") || source.eq_ignore_ascii_case("headers")
}
fn get_header_condition_value(headers: ConditionHeaders<'_>, path: &str) -> Option<Value> {
let key = path.trim().to_ascii_lowercase();
if key.is_empty() {
return None;
}
match headers {
ConditionHeaders::Request(headers) => headers
.get(key.as_str())
.and_then(|value| value.to_str().ok())
.map(|value| Value::String(value.trim().to_string())),
ConditionHeaders::Map(headers) => headers.get(&key).cloned().map(Value::String),
}
}
fn json_number(value: &Value) -> Option<f64> { fn json_number(value: &Value) -> Option<f64> {
value.as_f64().filter(|_| !value.is_boolean()) value.as_f64().filter(|_| !value.is_boolean())
} }
@@ -887,29 +976,34 @@ fn expand_wildcard_paths_recursive(
} }
} }
struct WildcardTargetOptions<'a> {
condition: Option<&'a Value>,
item_condition: bool,
original_body: Option<&'a Value>,
request_headers: Option<ConditionHeaders<'a>>,
require_leaf: bool,
reverse: bool,
}
fn iter_wildcard_targets( fn iter_wildcard_targets(
body: &Value, body: &Value,
path: &[BodyPathSegment], path: &[BodyPathSegment],
condition: Option<&Value>, options: WildcardTargetOptions<'_>,
item_condition: bool,
original_body: Option<&Value>,
require_leaf: bool,
reverse: bool,
) -> Vec<Vec<BodyPathSegment>> { ) -> Vec<Vec<BodyPathSegment>> {
if !has_wildcard(path) { if !has_wildcard(path) {
return vec![path.to_vec()]; return vec![path.to_vec()];
} }
let mut expanded = expand_wildcard_paths(body, path, require_leaf); let mut expanded = expand_wildcard_paths(body, path, options.require_leaf);
if reverse { if options.reverse {
expanded.reverse(); expanded.reverse();
} }
if !item_condition { if !options.item_condition {
return expanded; return expanded;
} }
let Some(condition) = condition else { let Some(condition) = options.condition else {
return expanded; return expanded;
}; };
@@ -918,7 +1012,12 @@ fn iter_wildcard_targets(
.filter(|concrete_path| { .filter(|concrete_path| {
let prefix = get_item_prefix_from_concrete(concrete_path, path); let prefix = get_item_prefix_from_concrete(concrete_path, path);
let resolved = resolve_item_condition(condition, &prefix); let resolved = resolve_item_condition(condition, &prefix);
evaluate_local_condition(body, &resolved, original_body) evaluate_local_condition(
body,
&resolved,
options.original_body,
options.request_headers,
)
}) })
.collect() .collect()
} }
@@ -1194,8 +1293,10 @@ fn rename_nested_value(
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::{ use super::{
apply_local_body_rules, apply_local_header_rules, body_rules_are_locally_supported, apply_local_body_rules, apply_local_body_rules_with_request_headers,
body_rules_handle_path, header_rules_are_locally_supported, apply_local_header_rules, apply_local_header_rules_with_request_headers,
body_rules_are_locally_supported, body_rules_handle_path,
header_rules_are_locally_supported,
}; };
#[test] #[test]
@@ -1242,7 +1343,7 @@ mod tests {
Some(&rules), Some(&rules),
&[], &[],
&serde_json::json!({"metadata":{"mode":"safe"}}), &serde_json::json!({"metadata":{"mode":"safe"}}),
Some(&serde_json::json!({"metadata":{"client":"desktop"}})), Some(&serde_json::json!({"metadata":{"mode":"safe","client":"desktop"}})),
)); ));
assert_eq!(headers.get("x-added").map(String::as_str), Some("1")); assert_eq!(headers.get("x-added").map(String::as_str), Some("1"));
assert_eq!( assert_eq!(
@@ -1251,6 +1352,32 @@ mod tests {
); );
} }
#[test]
fn header_rules_can_read_original_request_header_conditions() {
let rules = serde_json::json!([
{"action":"set","key":"x-applied","value":"yes","condition":{"source":"request_headers","path":"X-Mode","op":"eq","value":"debug"}},
{"action":"set","key":"x-skipped","value":"yes","condition":{"source":"request_headers","path":"x-missing","op":"exists"}}
]);
let mut request_headers = http::HeaderMap::new();
request_headers.insert("x-mode", "debug".parse().unwrap());
let mut headers = std::collections::BTreeMap::from([(
"x-mode".to_string(),
"provider-value".to_string(),
)]);
assert!(apply_local_header_rules_with_request_headers(
&mut headers,
Some(&rules),
&[],
&serde_json::json!({}),
None,
Some(&request_headers),
));
assert_eq!(headers.get("x-applied").map(String::as_str), Some("yes"));
assert!(!headers.contains_key("x-skipped"));
}
#[test] #[test]
fn body_rules_support_all_runtime_actions() { fn body_rules_support_all_runtime_actions() {
let rules = serde_json::json!([ let rules = serde_json::json!([
@@ -1333,6 +1460,15 @@ mod tests {
assert!(body_rules_are_locally_supported(Some(&rules))); assert!(body_rules_are_locally_supported(Some(&rules)));
let original = serde_json::json!({ let original = serde_json::json!({
"num": 10,
"text": "alpha-beta",
"tags": ["red", "green"],
"choice": "b",
"flag": true,
"maybe_null": null,
"profile": {
"name": "Ada"
},
"legacy": { "legacy": {
"present": 1 "present": 1
} }
@@ -1378,6 +1514,53 @@ mod tests {
assert_eq!(body["results"], expected); assert_eq!(body["results"], expected);
} }
#[test]
fn body_conditions_default_to_original_request_body() {
let rules = serde_json::json!([
{"action":"set","path":"model","value":"provider-model"},
{"action":"set","path":"metadata.original_hit","value":true,"condition":{"path":"model","op":"eq","value":"client-model"}},
{"action":"set","path":"metadata.current_hit","value":true,"condition":{"path":"model","op":"eq","value":"provider-model"}}
]);
let original = serde_json::json!({
"model": "client-model"
});
let mut body = original.clone();
assert!(apply_local_body_rules(
&mut body,
Some(&rules),
Some(&original)
));
assert_eq!(body["model"], "provider-model");
assert_eq!(body["metadata"]["original_hit"], true);
assert!(body["metadata"].get("current_hit").is_none());
}
#[test]
fn body_rules_can_read_original_request_header_conditions() {
let rules = serde_json::json!([
{"action":"set","path":"metadata.from_header","value":true,"condition":{"source":"request_headers","path":"X-Mode","op":"eq","value":"debug"}},
{"action":"set","path":"metadata.contains","value":true,"condition":{"source":"headers","path":"x-feature","op":"contains","value":"beta"}},
{"action":"set","path":"metadata.skipped","value":true,"condition":{"source":"request_headers","path":"x-mode","op":"eq","value":"prod"}}
]);
let mut request_headers = http::HeaderMap::new();
request_headers.insert("x-mode", "debug".parse().unwrap());
request_headers.insert("x-feature", "alpha,beta".parse().unwrap());
let mut body = serde_json::json!({});
assert!(apply_local_body_rules_with_request_headers(
&mut body,
Some(&rules),
None,
Some(&request_headers),
));
assert_eq!(body["metadata"]["from_header"], true);
assert_eq!(body["metadata"]["contains"], true);
assert!(body["metadata"].get("skipped").is_none());
}
#[test] #[test]
fn body_rules_tolerate_invalid_regex_flags_and_negative_count() { fn body_rules_tolerate_invalid_regex_flags_and_negative_count() {
let invalid_flags = serde_json::json!([ let invalid_flags = serde_json::json!([

View File

@@ -8,7 +8,7 @@ use aether_ai_formats::protocol::conversion::request::{
normalize_openai_responses_request_to_openai_chat_request, normalize_openai_responses_request_to_openai_chat_request,
}; };
use aether_ai_formats::protocol::registry::{convert_request, FormatContext}; use aether_ai_formats::protocol::registry::{convert_request, FormatContext};
use aether_ai_formats::provider_compat::proxy::rules::apply_local_body_rules; use aether_ai_formats::provider_compat::proxy::rules::apply_local_body_rules_with_request_headers;
use serde_json::Value; use serde_json::Value;
use crate::request::model_directives::apply_model_directive_overrides_from_request; use crate::request::model_directives::apply_model_directive_overrides_from_request;
@@ -57,6 +57,35 @@ pub fn build_standard_request_body_with_model_directives(
body_rules: Option<&Value>, body_rules: Option<&Value>,
user_api_key_id: Option<&str>, user_api_key_id: Option<&str>,
enable_model_directives: bool, enable_model_directives: bool,
) -> Option<Value> {
build_standard_request_body_with_model_directives_and_request_headers(
body_json,
client_api_format,
mapped_model,
provider_type,
provider_api_format,
request_path,
upstream_is_stream,
body_rules,
user_api_key_id,
None,
enable_model_directives,
)
}
#[allow(clippy::too_many_arguments)]
pub fn build_standard_request_body_with_model_directives_and_request_headers(
body_json: &Value,
client_api_format: &str,
mapped_model: &str,
provider_type: &str,
provider_api_format: &str,
request_path: &str,
upstream_is_stream: bool,
body_rules: Option<&Value>,
user_api_key_id: Option<&str>,
request_headers: Option<&http::HeaderMap>,
enable_model_directives: bool,
) -> Option<Value> { ) -> Option<Value> {
let format_context = FormatContext::default() let format_context = FormatContext::default()
.with_mapped_model(mapped_model) .with_mapped_model(mapped_model)
@@ -80,7 +109,12 @@ pub fn build_standard_request_body_with_model_directives(
); );
} }
if !apply_local_body_rules(&mut provider_request_body, body_rules, Some(body_json)) { if !apply_local_body_rules_with_request_headers(
&mut provider_request_body,
body_rules,
Some(body_json),
request_headers,
) {
return None; return None;
} }
apply_codex_openai_responses_special_body_edits( apply_codex_openai_responses_special_body_edits(

View File

@@ -16,6 +16,7 @@ pub use family::{LocalStandardSourceFamily, LocalStandardSourceMode, LocalStanda
pub use matrix::{ pub use matrix::{
build_standard_request_body, build_standard_request_body_from_canonical_with_model_directives, build_standard_request_body, build_standard_request_body_from_canonical_with_model_directives,
build_standard_request_body_with_model_directives, build_standard_request_body_with_model_directives,
build_standard_request_body_with_model_directives_and_request_headers,
normalize_standard_request_to_openai_chat_request, normalize_standard_request_to_openai_chat_request,
}; };
pub use normalize::{ pub use normalize::{

View File

@@ -4,7 +4,9 @@ use serde_json::{json, Value};
use crate::auth::{build_passthrough_headers_with_auth, resolve_local_gemini_auth}; use crate::auth::{build_passthrough_headers_with_auth, resolve_local_gemini_auth};
use crate::policy::local_gemini_transport_unsupported_reason_with_network; use crate::policy::local_gemini_transport_unsupported_reason_with_network;
use crate::rules::{apply_local_body_rules, apply_local_header_rules}; use crate::rules::{
apply_local_body_rules_with_request_headers, apply_local_header_rules_with_request_headers,
};
use crate::snapshot::GatewayProviderTransportSnapshot; use crate::snapshot::GatewayProviderTransportSnapshot;
use crate::url::build_gemini_files_passthrough_url; use crate::url::build_gemini_files_passthrough_url;
@@ -70,6 +72,7 @@ pub fn build_gemini_files_request_body(
body_is_empty: bool, body_is_empty: bool,
is_upload: bool, is_upload: bool,
body_rules: Option<&Value>, body_rules: Option<&Value>,
request_headers: Option<&http::HeaderMap>,
) -> Result<GeminiFilesRequestBodyParts, GeminiFilesRequestBodyError> { ) -> Result<GeminiFilesRequestBodyParts, GeminiFilesRequestBodyError> {
let mut provider_request_body = if is_upload && !body_is_empty && body_base64.is_none() { let mut provider_request_body = if is_upload && !body_is_empty && body_base64.is_none() {
Some(body_json.clone()) Some(body_json.clone())
@@ -88,7 +91,12 @@ pub fn build_gemini_files_request_body(
return Err(GeminiFilesRequestBodyError::BodyRulesUnsupportedForBinaryUpload); return Err(GeminiFilesRequestBodyError::BodyRulesUnsupportedForBinaryUpload);
} }
if let Some(body) = provider_request_body.as_mut() { if let Some(body) = provider_request_body.as_mut() {
if !apply_local_body_rules(body, body_rules, Some(body_json)) { if !apply_local_body_rules_with_request_headers(
body,
body_rules,
Some(body_json),
request_headers,
) {
return Err(GeminiFilesRequestBodyError::BodyRulesApplyFailed); return Err(GeminiFilesRequestBodyError::BodyRulesApplyFailed);
} }
} }
@@ -116,12 +124,13 @@ pub fn build_gemini_files_headers(
.as_ref() .as_ref()
.or_else(|| (!input.original_body_is_empty).then_some(input.original_request_body_json)) .or_else(|| (!input.original_body_is_empty).then_some(input.original_request_body_json))
.unwrap_or(&null_original_request_body); .unwrap_or(&null_original_request_body);
if !apply_local_header_rules( if !apply_local_header_rules_with_request_headers(
&mut provider_request_headers, &mut provider_request_headers,
input.header_rules, input.header_rules,
&[input.auth_header, "content-type"], &[input.auth_header, "content-type"],
input.provider_request_body.unwrap_or(original_request_body), input.provider_request_body.unwrap_or(original_request_body),
Some(original_request_body), Some(original_request_body),
Some(input.headers),
) { ) {
return None; return None;
} }
@@ -216,6 +225,7 @@ mod tests {
Some(&json!([ Some(&json!([
{"action":"set","path":"metadata.source","value":"local"} {"action":"set","path":"metadata.source","value":"local"}
])), ])),
None,
) )
.expect("body should build"); .expect("body should build");
@@ -236,7 +246,8 @@ mod tests {
Some("YWJj"), Some("YWJj"),
false, false,
true, true,
Some(&json!([{"action":"set","path":"x","value":1}])) Some(&json!([{"action":"set","path":"x","value":1}])),
None,
), ),
Err(GeminiFilesRequestBodyError::BodyRulesUnsupportedForBinaryUpload) Err(GeminiFilesRequestBodyError::BodyRulesUnsupportedForBinaryUpload)
); );

View File

@@ -25,10 +25,10 @@ pub use policy::{
}; };
pub use refresh::KiroOAuthRefreshAdapter; pub use refresh::KiroOAuthRefreshAdapter;
pub use request::{ pub use request::{
apply_local_body_rules, apply_local_header_rules, body_rules_are_locally_supported, apply_local_body_rules_with_request_headers, apply_local_header_rules_with_request_headers,
build_kiro_provider_headers, build_kiro_provider_request_body, body_rules_are_locally_supported, build_kiro_provider_headers,
header_rules_are_locally_supported, supports_local_kiro_request_shape, build_kiro_provider_request_body, header_rules_are_locally_supported,
KiroProviderHeadersInput, supports_local_kiro_request_shape, KiroProviderHeadersInput,
}; };
pub use url::{ pub use url::{
build_kiro_generate_assistant_response_url, build_kiro_mcp_url, build_kiro_generate_assistant_response_url, build_kiro_mcp_url,

View File

@@ -3,8 +3,8 @@ use std::collections::BTreeMap;
use serde_json::{json, Value}; use serde_json::{json, Value};
pub use super::super::rules::{ pub use super::super::rules::{
apply_local_body_rules, apply_local_header_rules, body_rules_are_locally_supported, apply_local_body_rules_with_request_headers, apply_local_header_rules_with_request_headers,
header_rules_are_locally_supported, body_rules_are_locally_supported, header_rules_are_locally_supported,
}; };
use super::super::should_skip_upstream_passthrough_header; use super::super::should_skip_upstream_passthrough_header;
use super::converter::convert_claude_messages_to_conversation_state; use super::converter::convert_claude_messages_to_conversation_state;
@@ -23,6 +23,7 @@ pub fn build_kiro_provider_request_body(
mapped_model: &str, mapped_model: &str,
auth_config: &KiroAuthConfig, auth_config: &KiroAuthConfig,
body_rules: Option<&Value>, body_rules: Option<&Value>,
request_headers: Option<&http::HeaderMap>,
) -> Option<Value> { ) -> Option<Value> {
let conversation_state = let conversation_state =
convert_claude_messages_to_conversation_state(body_json, mapped_model)?; convert_claude_messages_to_conversation_state(body_json, mapped_model)?;
@@ -70,7 +71,12 @@ pub fn build_kiro_provider_request_body(
); );
} }
if !apply_local_body_rules(&mut provider_request_body, body_rules, Some(body_json)) { if !apply_local_body_rules_with_request_headers(
&mut provider_request_body,
body_rules,
Some(body_json),
request_headers,
) {
return None; return None;
} }
@@ -119,12 +125,13 @@ pub fn build_kiro_provider_headers(
out.insert(key, value.to_string()); out.insert(key, value.to_string());
} }
if !apply_local_header_rules( if !apply_local_header_rules_with_request_headers(
&mut out, &mut out,
header_rules, header_rules,
&[auth_header, "content-type"], &[auth_header, "content-type"],
provider_request_body, provider_request_body,
Some(original_request_body), Some(original_request_body),
Some(headers),
) { ) {
return None; return None;
} }
@@ -202,6 +209,7 @@ mod tests {
Some(&json!([ Some(&json!([
{"action":"set","path":"debugTag","value":"kiro-local"} {"action":"set","path":"debugTag","value":"kiro-local"}
])), ])),
None,
) )
.expect("payload should build"); .expect("payload should build");
@@ -292,6 +300,7 @@ mod tests {
"claude-sonnet-4-upstream", "claude-sonnet-4-upstream",
&auth_config, &auth_config,
None, None,
None,
) )
.expect("payload should build"); .expect("payload should build");

View File

@@ -75,7 +75,8 @@ pub use request_url::{
TransportRequestUrlParams, TransportRequestUrlParams,
}; };
pub use rules::{ pub use rules::{
apply_local_body_rules, apply_local_header_rules, body_rules_are_locally_supported, apply_local_body_rules, apply_local_body_rules_with_request_headers, apply_local_header_rules,
apply_local_header_rules_with_request_headers, body_rules_are_locally_supported,
body_rules_handle_path, header_rules_are_locally_supported, body_rules_handle_path, header_rules_are_locally_supported,
}; };
pub use same_format_provider::{ pub use same_format_provider::{
@@ -94,8 +95,9 @@ pub use snapshot::{
ProviderTransportSnapshotSource, ProviderTransportSnapshotSource,
}; };
pub use standard::{ pub use standard::{
apply_standard_provider_request_body_rules, build_standard_plan_fallback_headers, apply_standard_provider_request_body_rules,
build_standard_plan_fallback_openai_chat_url, apply_standard_provider_request_body_rules_with_request_headers,
build_standard_plan_fallback_headers, build_standard_plan_fallback_openai_chat_url,
build_standard_plan_fallback_openai_responses_url, build_standard_provider_request_headers, build_standard_plan_fallback_openai_responses_url, build_standard_provider_request_headers,
StandardPlanFallbackAcceptPolicy, StandardPlanFallbackHeadersInput, StandardPlanFallbackAcceptPolicy, StandardPlanFallbackHeadersInput,
StandardProviderRequestHeaders, StandardProviderRequestHeadersInput, StandardProviderRequestHeaders, StandardProviderRequestHeadersInput,

View File

@@ -4,7 +4,7 @@ use serde_json::Value;
use crate::auth::{build_passthrough_headers_with_auth, resolve_local_openai_bearer_auth}; use crate::auth::{build_passthrough_headers_with_auth, resolve_local_openai_bearer_auth};
use crate::policy::local_standard_transport_unsupported_reason_with_network; use crate::policy::local_standard_transport_unsupported_reason_with_network;
use crate::rules::apply_local_header_rules; use crate::rules::apply_local_header_rules_with_request_headers;
use crate::snapshot::GatewayProviderTransportSnapshot; use crate::snapshot::GatewayProviderTransportSnapshot;
use crate::url::build_openai_responses_url; use crate::url::build_openai_responses_url;
@@ -59,12 +59,13 @@ pub fn build_openai_image_headers(
); );
provider_request_headers.insert("content-type".to_string(), "application/json".to_string()); provider_request_headers.insert("content-type".to_string(), "application/json".to_string());
provider_request_headers.insert("accept".to_string(), "text/event-stream".to_string()); provider_request_headers.insert("accept".to_string(), "text/event-stream".to_string());
if !apply_local_header_rules( if !apply_local_header_rules_with_request_headers(
&mut provider_request_headers, &mut provider_request_headers,
input.header_rules, input.header_rules,
&[input.auth_header, "content-type", "accept"], &[input.auth_header, "content-type", "accept"],
input.provider_request_body, input.provider_request_body,
Some(input.original_request_body), Some(input.original_request_body),
Some(input.headers),
) { ) {
return None; return None;
} }

View File

@@ -18,7 +18,9 @@ use crate::policy::{
local_gemini_transport_unsupported_reason_with_network, local_gemini_transport_unsupported_reason_with_network,
local_standard_transport_unsupported_reason_with_network, local_standard_transport_unsupported_reason_with_network,
}; };
use crate::rules::{apply_local_body_rules, apply_local_header_rules}; use crate::rules::{
apply_local_body_rules_with_request_headers, apply_local_header_rules_with_request_headers,
};
use crate::snapshot::GatewayProviderTransportSnapshot; use crate::snapshot::GatewayProviderTransportSnapshot;
use crate::vertex::{ use crate::vertex::{
is_vertex_api_key_transport_context, is_vertex_api_key_transport_context,
@@ -57,6 +59,7 @@ pub struct SameFormatProviderRequestBodyInput<'a> {
pub source_model: Option<&'a str>, pub source_model: Option<&'a str>,
pub family: SameFormatProviderFamily, pub family: SameFormatProviderFamily,
pub body_rules: Option<&'a Value>, pub body_rules: Option<&'a Value>,
pub request_headers: Option<&'a http::HeaderMap>,
pub upstream_is_stream: bool, pub upstream_is_stream: bool,
pub kiro_auth_config: Option<&'a KiroAuthConfig>, pub kiro_auth_config: Option<&'a KiroAuthConfig>,
pub is_claude_code: bool, pub is_claude_code: bool,
@@ -131,6 +134,7 @@ pub fn build_same_format_provider_request_body(
input.mapped_model, input.mapped_model,
kiro_auth_config, kiro_auth_config,
input.body_rules, input.body_rules,
input.request_headers,
); );
} }
@@ -183,10 +187,11 @@ pub fn build_same_format_provider_request_body(
); );
} }
} }
if !apply_local_body_rules( if !apply_local_body_rules_with_request_headers(
&mut provider_request_body, &mut provider_request_body,
input.body_rules, input.body_rules,
Some(input.body_json), Some(input.body_json),
input.request_headers,
) { ) {
return None; return None;
} }
@@ -257,12 +262,13 @@ pub fn build_same_format_provider_headers(
.filter(|value| !value.trim().is_empty()) .filter(|value| !value.trim().is_empty())
.map(|value| vec![value, "content-type"]) .map(|value| vec![value, "content-type"])
.unwrap_or_else(|| vec!["content-type"]); .unwrap_or_else(|| vec!["content-type"]);
if !apply_local_header_rules( if !apply_local_header_rules_with_request_headers(
&mut provider_request_headers, &mut provider_request_headers,
input.header_rules, input.header_rules,
&protected_headers, &protected_headers,
input.provider_request_body, input.provider_request_body,
Some(input.original_request_body), Some(input.original_request_body),
Some(input.headers),
) { ) {
return None; return None;
} }
@@ -509,6 +515,7 @@ mod tests {
source_model: Some("client-model"), source_model: Some("client-model"),
family: SameFormatProviderFamily::Standard, family: SameFormatProviderFamily::Standard,
body_rules: None, body_rules: None,
request_headers: None,
upstream_is_stream: true, upstream_is_stream: true,
kiro_auth_config: None, kiro_auth_config: None,
is_claude_code: false, is_claude_code: false,
@@ -536,6 +543,7 @@ mod tests {
body_rules: Some(&json!([ body_rules: Some(&json!([
{"action":"set","path":"metadata.body_rule_seen","value":true} {"action":"set","path":"metadata.body_rule_seen","value":true}
])), ])),
request_headers: None,
upstream_is_stream: false, upstream_is_stream: false,
kiro_auth_config: None, kiro_auth_config: None,
is_claude_code: false, is_claude_code: false,

View File

@@ -6,7 +6,10 @@ use crate::auth::{
build_claude_passthrough_headers, build_complete_passthrough_headers_with_auth, build_claude_passthrough_headers, build_complete_passthrough_headers_with_auth,
build_openai_passthrough_headers, build_passthrough_headers, ensure_upstream_auth_header, build_openai_passthrough_headers, build_passthrough_headers, ensure_upstream_auth_header,
}; };
use crate::rules::{apply_local_body_rules, apply_local_header_rules}; use crate::rules::{
apply_local_body_rules, apply_local_body_rules_with_request_headers,
apply_local_header_rules_with_request_headers,
};
use crate::snapshot::GatewayProviderTransportSnapshot; use crate::snapshot::GatewayProviderTransportSnapshot;
use crate::url::{build_openai_chat_url, build_openai_responses_url}; use crate::url::{build_openai_chat_url, build_openai_responses_url};
use crate::vertex::uses_vertex_api_key_query_auth; use crate::vertex::uses_vertex_api_key_query_auth;
@@ -157,6 +160,23 @@ pub fn apply_standard_provider_request_body_rules(
Some(provider_request_body) Some(provider_request_body)
} }
pub fn apply_standard_provider_request_body_rules_with_request_headers(
mut provider_request_body: Value,
body_rules: Option<&Value>,
original_request_body: &Value,
request_headers: &http::HeaderMap,
) -> Option<Value> {
if !apply_local_body_rules_with_request_headers(
&mut provider_request_body,
body_rules,
Some(original_request_body),
Some(request_headers),
) {
return None;
}
Some(provider_request_body)
}
pub fn build_standard_provider_request_headers( pub fn build_standard_provider_request_headers(
input: StandardProviderRequestHeadersInput<'_>, input: StandardProviderRequestHeadersInput<'_>,
) -> Option<StandardProviderRequestHeaders> { ) -> Option<StandardProviderRequestHeaders> {
@@ -193,12 +213,13 @@ pub fn build_standard_provider_request_headers(
} else { } else {
&[input.auth_header, "content-type"][..] &[input.auth_header, "content-type"][..]
}; };
if !apply_local_header_rules( if !apply_local_header_rules_with_request_headers(
&mut headers, &mut headers,
input.header_rules, input.header_rules,
protected_headers, protected_headers,
input.provider_request_body, input.provider_request_body,
Some(input.original_request_body), Some(input.original_request_body),
Some(input.headers),
) { ) {
return None; return None;
} }

View File

@@ -17,7 +17,9 @@ use super::policy::{
local_standard_transport_unsupported_reason_with_network, supports_local_gemini_transport, local_standard_transport_unsupported_reason_with_network, supports_local_gemini_transport,
supports_local_standard_transport, supports_local_standard_transport,
}; };
use super::rules::{apply_local_body_rules, apply_local_header_rules}; use super::rules::{
apply_local_body_rules_with_request_headers, apply_local_header_rules_with_request_headers,
};
use super::snapshot::GatewayProviderTransportSnapshot; use super::snapshot::GatewayProviderTransportSnapshot;
use super::url::{build_gemini_video_predict_long_running_url, build_passthrough_path_url}; use super::url::{build_gemini_video_predict_long_running_url, build_passthrough_path_url};
@@ -117,6 +119,7 @@ pub fn build_video_create_request_body(
family: ProviderVideoCreateFamily, family: ProviderVideoCreateFamily,
mapped_model: &str, mapped_model: &str,
body_rules: Option<&Value>, body_rules: Option<&Value>,
request_headers: Option<&http::HeaderMap>,
) -> Option<Value> { ) -> Option<Value> {
let mut provider_request_body = match family { let mut provider_request_body = match family {
ProviderVideoCreateFamily::OpenAi => { ProviderVideoCreateFamily::OpenAi => {
@@ -127,7 +130,12 @@ pub fn build_video_create_request_body(
} }
ProviderVideoCreateFamily::Gemini => body_json.clone(), ProviderVideoCreateFamily::Gemini => body_json.clone(),
}; };
if !apply_local_body_rules(&mut provider_request_body, body_rules, Some(body_json)) { if !apply_local_body_rules_with_request_headers(
&mut provider_request_body,
body_rules,
Some(body_json),
request_headers,
) {
return None; return None;
} }
Some(provider_request_body) Some(provider_request_body)
@@ -184,12 +192,13 @@ pub fn build_video_create_headers(
input.auth_value, input.auth_value,
&BTreeMap::new(), &BTreeMap::new(),
); );
if !apply_local_header_rules( if !apply_local_header_rules_with_request_headers(
&mut provider_request_headers, &mut provider_request_headers,
input.header_rules, input.header_rules,
&[input.auth_header, "content-type"], &[input.auth_header, "content-type"],
input.provider_request_body, input.provider_request_body,
Some(input.original_request_body), Some(input.original_request_body),
Some(input.headers),
) { ) {
return None; return None;
} }
@@ -409,6 +418,7 @@ mod tests {
ProviderVideoCreateFamily::OpenAi, ProviderVideoCreateFamily::OpenAi,
"upstream-video-model", "upstream-video-model",
None, None,
None,
) )
.expect("body should build"); .expect("body should build");

View File

@@ -140,7 +140,7 @@ export interface BodyRuleConditionLeaf {
path: string path: string
op: BodyRuleConditionOp op: BodyRuleConditionOp
value?: unknown // exists / not_exists 不需要 value value?: unknown // exists / not_exists 不需要 value
source?: 'original' | 'current' source?: 'request_headers' // 不填表示请求体;填 request_headers 表示请求头
} }
export interface BodyRuleConditionAll { export interface BodyRuleConditionAll {

View File

@@ -127,17 +127,17 @@
<SelectValue /> <SelectValue />
</SelectTrigger> </SelectTrigger>
<SelectContent> <SelectContent>
<SelectItem value="current"> <SelectItem value="body">
Current 请求体
</SelectItem> </SelectItem>
<SelectItem value="original"> <SelectItem value="request_headers">
Original 请求头
</SelectItem> </SelectItem>
</SelectContent> </SelectContent>
</Select> </Select>
<Input <Input
:model-value="modelValue.path" :model-value="modelValue.path"
:placeholder="pathHint || '字段路径'" :placeholder="modelValue.source === 'request_headers' ? 'Header 名称' : (pathHint || '字段路径')"
size="sm" size="sm"
class="flex-1 min-w-[120px] h-7 text-xs" class="flex-1 min-w-[120px] h-7 text-xs"
@update:model-value="(value) => updateLeafField('path', value)" @update:model-value="(value) => updateLeafField('path', value)"

View File

@@ -635,14 +635,14 @@
<code>in</code> 在列表中(值填 <code>["a","b"]</code><br> <code>in</code> 在列表中(值填 <code>["a","b"]</code><br>
<code>type_is</code> 类型判断string/number/boolean/array/object/null<br> <code>type_is</code> 类型判断string/number/boolean/array/object/null<br>
条件路径支持 <code>$item.xxx</code> 引用通配符当前元素<br> 条件路径支持 <code>$item.xxx</code> 引用通配符当前元素<br>
可切换 <code>Current</code>/<code>Original</code> 数据源,并支持 <code>ALL</code>/<code>ANY</code> 组合条件 可切换 <code>请求体</code>/<code>请求头</code> 数据源,并支持 <code>ALL</code>/<code>ANY</code> 组合条件
</div> </div>
</div> </div>
<div class="text-muted-foreground"> <div class="text-muted-foreground">
规则按顺序执行,前面的修改对后续规则可见 规则按顺序执行;条件判断使用客户端原始请求体或原始请求头,不受前面规则修改影响
</div> </div>
<div class="text-muted-foreground"> <div class="text-muted-foreground">
规则默认在格式转换后按目标提供商结构匹配;条件切到 <code>Original</code> 时按客户端原始请求匹配。 条件默认按客户端原始请求体匹配;切到 <code>请求头</code> 时按客户端原始请求匹配。
</div> </div>
</div> </div>
</PopoverContent> </PopoverContent>
@@ -1484,8 +1484,8 @@ function validateJsonConditionShape(condition: Record<string, unknown>, label: s
if (typeof condition.op !== 'string' || !CONDITION_JSON_OPS.has(condition.op)) { if (typeof condition.op !== 'string' || !CONDITION_JSON_OPS.has(condition.op)) {
return `${label}.op 无效` return `${label}.op 无效`
} }
if (condition.source !== undefined && condition.source !== 'original' && condition.source !== 'current') { if (condition.source !== undefined && condition.source !== 'request_headers') {
return `${label}.source 只能是 original/current` return `${label}.source 只能是 request_headers请求体条件不要填写 source`
} }
return null return null
} }

View File

@@ -1,6 +1,6 @@
import type { BodyRuleCondition, BodyRuleConditionOp } from '@/api/endpoints' import type { BodyRuleCondition, BodyRuleConditionOp } from '@/api/endpoints'
export type ConditionSource = 'current' | 'original' export type ConditionSource = 'body' | 'request_headers'
export type ConditionGroupMode = 'all' | 'any' export type ConditionGroupMode = 'all' | 'any'
export interface EditableConditionLeaf { export interface EditableConditionLeaf {
@@ -46,7 +46,7 @@ export function createEmptyConditionLeaf(): EditableConditionLeaf {
path: '', path: '',
op: 'eq', op: 'eq',
value: '', value: '',
source: 'current', source: 'body',
} }
} }
@@ -86,6 +86,7 @@ export function conditionToEditable(condition?: BodyRuleCondition | null): Edita
condition.any.map(child => conditionToEditable(child) || createEmptyConditionLeaf()), condition.any.map(child => conditionToEditable(child) || createEmptyConditionLeaf()),
) )
} }
const source = (condition as { source?: unknown }).source
return { return {
kind: 'leaf', kind: 'leaf',
path: condition.path || '', path: condition.path || '',
@@ -93,7 +94,9 @@ export function conditionToEditable(condition?: BodyRuleCondition | null): Edita
value: condition.value !== undefined value: condition.value !== undefined
? (typeof condition.value === 'string' ? condition.value : JSON.stringify(condition.value)) ? (typeof condition.value === 'string' ? condition.value : JSON.stringify(condition.value))
: '', : '',
source: condition.source === 'original' ? 'original' : 'current', source: source === 'request_headers' || source === 'headers'
? 'request_headers'
: 'body',
} }
} }
@@ -114,7 +117,7 @@ export function editableConditionToApi(node: EditableConditionNode | null): Body
const base = { const base = {
path, path,
op: node.op, op: node.op,
...(node.source === 'original' ? { source: 'original' as const } : {}), ...(node.source === 'request_headers' ? { source: 'request_headers' as const } : {}),
} }
if (node.op === 'exists' || node.op === 'not_exists') { if (node.op === 'exists' || node.op === 'not_exists') {