mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
use original request sources for endpoint conditions
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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!([
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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::{
|
||||
|
||||
@@ -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)
|
||||
);
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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");
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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");
|
||||
|
||||
|
||||
Reference in New Issue
Block a user