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

@@ -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");