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

@@ -119,6 +119,7 @@ pub use crate::request::standard::{
build_local_openai_responses_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_and_request_headers,
claude::{
resolve_stream_spec as resolve_claude_stream_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 ITEM_PREFIX: &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"];
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)]
enum BodyPathSegment {
Key(String),
@@ -34,6 +40,35 @@ pub fn apply_local_header_rules(
protected_keys: &[&str],
body: &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 {
let Some(rules) = rules else {
return true;
@@ -54,7 +89,8 @@ pub fn apply_local_header_rules(
if !condition_is_locally_supported(condition) {
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;
}
}
@@ -178,6 +214,29 @@ pub fn apply_local_body_rules(
body: &mut Value,
rules: 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 {
let Some(rules) = rules else {
return true;
@@ -197,7 +256,9 @@ pub fn apply_local_body_rules(
if !condition_is_locally_supported(condition) {
continue;
}
if !item_condition && !evaluate_local_condition(body, condition, original_body) {
if !item_condition
&& !evaluate_local_condition(body, condition, original_body, request_headers)
{
continue;
}
}
@@ -220,11 +281,14 @@ pub fn apply_local_body_rules(
let targets = iter_wildcard_targets(
body,
&path,
condition,
item_condition,
original_body,
false,
false,
WildcardTargetOptions {
condition,
item_condition,
original_body,
request_headers,
require_leaf: false,
reverse: false,
},
);
let value_template = rule.get("value").cloned().unwrap_or(Value::Null);
for target_path in targets {
@@ -248,11 +312,14 @@ pub fn apply_local_body_rules(
for target_path in iter_wildcard_targets(
body,
&path,
condition,
item_condition,
original_body,
true,
true,
WildcardTargetOptions {
condition,
item_condition,
original_body,
request_headers,
require_leaf: true,
reverse: true,
},
) {
let _ = delete_nested_value(body, &target_path);
}
@@ -289,11 +356,14 @@ pub fn apply_local_body_rules(
for target_path in iter_wildcard_targets(
body,
&path,
condition,
item_condition,
original_body,
true,
false,
WildcardTargetOptions {
condition,
item_condition,
original_body,
request_headers,
require_leaf: true,
reverse: false,
},
) {
if let Some(target) = get_nested_value_mut(body, &target_path) {
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(
body,
&path,
condition,
item_condition,
original_body,
true,
false,
WildcardTargetOptions {
condition,
item_condition,
original_body,
request_headers,
require_leaf: true,
reverse: false,
},
) {
if let Some(target) = get_nested_value_mut(body, &target_path) {
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")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or("current");
.unwrap_or("body");
if !CONDITION_SOURCES.contains(&source) {
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 {
return false;
};
let Some(path) = condition
.get("path")
.and_then(Value::as_str)
.map(str::trim)
.and_then(parse_body_path)
else {
let Some(path) = condition.get("path").and_then(Value::as_str) else {
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;
}
@@ -443,6 +514,7 @@ fn evaluate_local_condition(
body: &Value,
condition: &Value,
original_body: Option<&Value>,
request_headers: Option<ConditionHeaders<'_>>,
) -> bool {
let Some(condition) = condition.as_object() else {
return false;
@@ -450,41 +522,40 @@ fn evaluate_local_condition(
if let Some(children) = condition.get("all").and_then(Value::as_array) {
return !children.is_empty()
&& children
.iter()
.all(|child| evaluate_local_condition(body, child, original_body));
&& children.iter().all(|child| {
evaluate_local_condition(body, child, original_body, request_headers)
});
}
if let Some(children) = condition.get("any").and_then(Value::as_array) {
return !children.is_empty()
&& children
.iter()
.any(|child| evaluate_local_condition(body, child, original_body));
&& children.iter().any(|child| {
evaluate_local_condition(body, child, original_body, request_headers)
});
}
let source = condition
.get("source")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or("current");
let target = if source.eq_ignore_ascii_case("original") {
original_body.unwrap_or(body)
} else {
body
};
.unwrap_or("body");
let Some(op) = condition.get("op").and_then(Value::as_str).map(str::trim) else {
return false;
};
let Some(path) = condition
.get("path")
.and_then(Value::as_str)
.map(str::trim)
.and_then(parse_body_path)
else {
let Some(path) = condition.get("path").and_then(Value::as_str).map(str::trim) else {
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" {
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> {
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(
body: &Value,
path: &[BodyPathSegment],
condition: Option<&Value>,
item_condition: bool,
original_body: Option<&Value>,
require_leaf: bool,
reverse: bool,
options: WildcardTargetOptions<'_>,
) -> Vec<Vec<BodyPathSegment>> {
if !has_wildcard(path) {
return vec![path.to_vec()];
}
let mut expanded = expand_wildcard_paths(body, path, require_leaf);
if reverse {
let mut expanded = expand_wildcard_paths(body, path, options.require_leaf);
if options.reverse {
expanded.reverse();
}
if !item_condition {
if !options.item_condition {
return expanded;
}
let Some(condition) = condition else {
let Some(condition) = options.condition else {
return expanded;
};
@@ -918,7 +1012,12 @@ fn iter_wildcard_targets(
.filter(|concrete_path| {
let prefix = get_item_prefix_from_concrete(concrete_path, path);
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()
}
@@ -1194,8 +1293,10 @@ fn rename_nested_value(
#[cfg(test)]
mod tests {
use super::{
apply_local_body_rules, apply_local_header_rules, body_rules_are_locally_supported,
body_rules_handle_path, header_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,
};
#[test]
@@ -1242,7 +1343,7 @@ mod tests {
Some(&rules),
&[],
&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!(
@@ -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]
fn body_rules_support_all_runtime_actions() {
let rules = serde_json::json!([
@@ -1333,6 +1460,15 @@ mod tests {
assert!(body_rules_are_locally_supported(Some(&rules)));
let original = serde_json::json!({
"num": 10,
"text": "alpha-beta",
"tags": ["red", "green"],
"choice": "b",
"flag": true,
"maybe_null": null,
"profile": {
"name": "Ada"
},
"legacy": {
"present": 1
}
@@ -1378,6 +1514,53 @@ mod tests {
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]
fn body_rules_tolerate_invalid_regex_flags_and_negative_count() {
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,
};
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 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>,
user_api_key_id: Option<&str>,
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> {
let format_context = FormatContext::default()
.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;
}
apply_codex_openai_responses_special_body_edits(

View File

@@ -16,6 +16,7 @@ pub use family::{LocalStandardSourceFamily, LocalStandardSourceMode, LocalStanda
pub use matrix::{
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_and_request_headers,
normalize_standard_request_to_openai_chat_request,
};
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::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::url::build_gemini_files_passthrough_url;
@@ -70,6 +72,7 @@ pub fn build_gemini_files_request_body(
body_is_empty: bool,
is_upload: bool,
body_rules: Option<&Value>,
request_headers: Option<&http::HeaderMap>,
) -> Result<GeminiFilesRequestBodyParts, GeminiFilesRequestBodyError> {
let mut provider_request_body = if is_upload && !body_is_empty && body_base64.is_none() {
Some(body_json.clone())
@@ -88,7 +91,12 @@ pub fn build_gemini_files_request_body(
return Err(GeminiFilesRequestBodyError::BodyRulesUnsupportedForBinaryUpload);
}
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);
}
}
@@ -116,12 +124,13 @@ pub fn build_gemini_files_headers(
.as_ref()
.or_else(|| (!input.original_body_is_empty).then_some(input.original_request_body_json))
.unwrap_or(&null_original_request_body);
if !apply_local_header_rules(
if !apply_local_header_rules_with_request_headers(
&mut provider_request_headers,
input.header_rules,
&[input.auth_header, "content-type"],
input.provider_request_body.unwrap_or(original_request_body),
Some(original_request_body),
Some(input.headers),
) {
return None;
}
@@ -216,6 +225,7 @@ mod tests {
Some(&json!([
{"action":"set","path":"metadata.source","value":"local"}
])),
None,
)
.expect("body should build");
@@ -236,7 +246,8 @@ mod tests {
Some("YWJj"),
false,
true,
Some(&json!([{"action":"set","path":"x","value":1}]))
Some(&json!([{"action":"set","path":"x","value":1}])),
None,
),
Err(GeminiFilesRequestBodyError::BodyRulesUnsupportedForBinaryUpload)
);

View File

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

View File

@@ -75,7 +75,8 @@ pub use request_url::{
TransportRequestUrlParams,
};
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,
};
pub use same_format_provider::{
@@ -94,8 +95,9 @@ pub use snapshot::{
ProviderTransportSnapshotSource,
};
pub use standard::{
apply_standard_provider_request_body_rules, build_standard_plan_fallback_headers,
build_standard_plan_fallback_openai_chat_url,
apply_standard_provider_request_body_rules,
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,
StandardPlanFallbackAcceptPolicy, StandardPlanFallbackHeadersInput,
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::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::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("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,
input.header_rules,
&[input.auth_header, "content-type", "accept"],
input.provider_request_body,
Some(input.original_request_body),
Some(input.headers),
) {
return None;
}

View File

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

View File

@@ -6,7 +6,10 @@ use crate::auth::{
build_claude_passthrough_headers, build_complete_passthrough_headers_with_auth,
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::url::{build_openai_chat_url, build_openai_responses_url};
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)
}
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(
input: StandardProviderRequestHeadersInput<'_>,
) -> Option<StandardProviderRequestHeaders> {
@@ -193,12 +213,13 @@ pub fn build_standard_provider_request_headers(
} else {
&[input.auth_header, "content-type"][..]
};
if !apply_local_header_rules(
if !apply_local_header_rules_with_request_headers(
&mut headers,
input.header_rules,
protected_headers,
input.provider_request_body,
Some(input.original_request_body),
Some(input.headers),
) {
return None;
}

View File

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