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,
spec,
prepared.transport.endpoint.body_rules.as_ref(),
Some(&parts.headers),
prepared.upstream_is_stream,
prepared.kiro_auth.as_ref(),
prepared.is_claude_code,

View File

@@ -12,6 +12,7 @@ pub(crate) fn build_same_format_provider_request_body(
mapped_model: &str,
spec: LocalSameFormatProviderSpec,
body_rules: Option<&Value>,
request_headers: Option<&http::HeaderMap>,
upstream_is_stream: bool,
kiro_auth: Option<&crate::ai_serving::transport::kiro::KiroRequestAuth>,
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),
family: same_format_provider_family(spec.family),
body_rules,
request_headers,
upstream_is_stream,
kiro_auth_config: kiro_auth.map(|auth| &auth.auth_config),
is_claude_code,

View File

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

View File

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

View File

@@ -176,7 +176,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
)
.await;
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,
spec_metadata.api_format,
&prepared_candidate.mapped_model,
@@ -190,6 +190,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
transport.endpoint.body_rules.as_ref()
},
Some(input.auth_context.api_key_id.as_str()),
Some(&parts.headers),
enable_model_directives,
) {
Some(body) => body,
@@ -353,6 +354,7 @@ async fn build_kiro_cross_format_payload_parts(
&mapped_model,
&kiro_auth.auth_config,
transport.endpoint.body_rules.as_ref(),
Some(&parts.headers),
) {
Some(body) => body,
None => {

View File

@@ -51,6 +51,7 @@ 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_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_openai_responses_request, extract_openai_text_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 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::{
apply_codex_openai_responses_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,
upstream_is_stream: bool,
body_rules: Option<&Value>,
request_headers: &http::HeaderMap,
enable_model_directives: bool,
) -> Option<Value> {
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,
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(
@@ -40,6 +46,7 @@ pub(crate) fn build_cross_format_openai_chat_request_body(
upstream_is_stream: bool,
body_rules: Option<&Value>,
user_api_key_id: Option<&str>,
request_headers: &http::HeaderMap,
enable_model_directives: bool,
) -> Option<Value> {
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,
)?;
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(
&mut provider_request_body,
provider_type,

View File

@@ -1,6 +1,6 @@
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::{
apply_codex_openai_responses_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,
body_rules: Option<&Value>,
user_api_key_id: Option<&str>,
request_headers: &http::HeaderMap,
enable_model_directives: bool,
) -> Option<Value> {
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,
)?;
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(
&mut provider_request_body,
provider_type,
@@ -50,6 +56,7 @@ pub(crate) fn build_cross_format_openai_responses_request_body(
provider_type: &str,
body_rules: Option<&Value>,
user_api_key_id: Option<&str>,
request_headers: &http::HeaderMap,
enable_model_directives: bool,
) -> Option<Value> {
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,
)?;
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(
&mut provider_request_body,
provider_type,

View File

@@ -90,6 +90,7 @@ fn builds_openai_chat_cross_format_request_body_from_openai_responses_source() {
"openai",
None,
None,
&http::HeaderMap::new(),
false,
)
.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",
None,
Some("key-123"),
&http::HeaderMap::new(),
false,
)
.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",
None,
None,
&http::HeaderMap::new(),
false,
)
.expect("local openai compact body should build");
@@ -188,6 +191,7 @@ fn local_openai_responses_wrapper_applies_model_directive_before_body_rules() {
"openai:responses",
Some(&body_rules),
None,
&http::HeaderMap::new(),
true,
)
.expect("local openai responses body should build");
@@ -236,6 +240,7 @@ fn strips_metadata_for_codex_openai_responses_requests() {
"codex",
None,
None,
&http::HeaderMap::new(),
false,
)
.expect("claude cli to codex request should build");
@@ -269,6 +274,7 @@ fn applies_codex_defaults_unless_body_rules_handle_the_field() {
"codex",
Some(&body_rules),
None,
&http::HeaderMap::new(),
false,
)
.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",
None,
Some("key-123"),
&http::HeaderMap::new(),
false,
)
.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,
None,
Some("key-123"),
&http::HeaderMap::new(),
false,
)
.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,
upstream_is_stream,
transport.endpoint.body_rules.as_ref(),
&parts.headers,
enable_model_directives,
) else {
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()
},
Some(input.auth_context.api_key_id.as_str()),
&parts.headers,
enable_model_directives,
) else {
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,
&kiro_auth.auth_config,
transport.endpoint.body_rules.as_ref(),
Some(&parts.headers),
) {
Some(body) => body,
None => {

View File

@@ -244,6 +244,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
transport.endpoint.body_rules.as_ref()
},
Some(input.auth_context.api_key_id.as_str()),
&parts.headers,
enable_model_directives,
)
} else {
@@ -259,6 +260,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
transport.endpoint.body_rules.as_ref()
},
Some(input.auth_context.api_key_id.as_str()),
&parts.headers,
enable_model_directives,
)
}) else {
@@ -532,6 +534,7 @@ async fn build_kiro_openai_responses_payload_parts(
&mapped_model,
&kiro_auth.auth_config,
transport.endpoint.body_rules.as_ref(),
Some(&parts.headers),
) {
Some(body) => body,
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_standard_request_body, build_standard_request_body_from_canonical,
build_standard_request_body_from_canonical_with_model_directives,
build_standard_request_body_with_model_directives, calculate_kiro_context_input_tokens,
canonicalize_tool_arguments, convert_claude_chat_response_to_openai_chat,
convert_claude_response_to_openai_responses, 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,
build_standard_request_body_with_model_directives,
build_standard_request_body_with_model_directives_and_request_headers,
calculate_kiro_context_input_tokens, canonicalize_tool_arguments,
convert_claude_chat_response_to_openai_chat, convert_claude_response_to_openai_responses,
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_response_to_claude_chat, convert_openai_chat_response_to_gemini_chat,
convert_openai_chat_response_to_openai_responses,

View File

@@ -47,11 +47,13 @@ pub(crate) mod vertex {
}
pub(crate) use aether_provider_transport::{
append_transport_diagnostics_to_value, apply_local_body_rules, apply_local_header_rules,
apply_standard_provider_request_body_rules, 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,
append_transport_diagnostics_to_value, apply_local_body_rules,
apply_local_body_rules_with_request_headers, apply_local_header_rules,
apply_local_header_rules_with_request_headers, apply_standard_provider_request_body_rules,
apply_standard_provider_request_body_rules_with_request_headers,
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_local_openai_responses_upstream_url, build_openai_image_headers,
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 {
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(object) = body.as_object_mut() {
object
.entry("model".to_string())
.or_insert_with(|| Value::String(model.to_string()));
if override_custom_model {
object.insert("model".to_string(), Value::String(model.to_string()));
} else {
object
.entry("model".to_string())
.or_insert_with(|| Value::String(model.to_string()));
}
}
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 =
provider_query_request_body_model(&request_body, &candidate.effective_model);
let provider_request_body = match build_kiro_provider_request_body(
@@ -799,6 +828,7 @@ async fn provider_query_execute_kiro_test_candidate(
request_model,
&kiro_auth.auth_config,
transport.endpoint.body_rules.as_ref(),
Some(&incoming_request_headers),
) {
Some(body) => body,
None => {
@@ -821,7 +851,7 @@ async fn provider_query_execute_kiro_test_candidate(
.uri(route_path)
.body(())
.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 request_url = build_kiro_generate_assistant_response_url(
@@ -949,8 +979,12 @@ async fn provider_query_execute_standard_test_candidate(
});
}
let original_request_body =
provider_query_build_test_request_body(payload, &candidate.effective_model);
let original_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 mut request_body = original_request_body.clone();
if let Some(object) = request_body.as_object_mut() {
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}"),
));
};
if !crate::provider_transport::apply_local_body_rules(
if !crate::provider_transport::apply_local_body_rules_with_request_headers(
&mut provider_request_body,
transport.endpoint.body_rules.as_ref(),
Some(&request_body),
Some(&incoming_request_headers),
) {
return Ok(provider_query_skipped_execution_outcome(
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}"),
));
};
if !crate::provider_transport::apply_local_body_rules(
if !crate::provider_transport::apply_local_body_rules_with_request_headers(
&mut provider_request_body,
transport.endpoint.body_rules.as_ref(),
Some(&request_body),
Some(&incoming_request_headers),
) {
return Ok(provider_query_skipped_execution_outcome(
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}"),
));
};
if !crate::provider_transport::apply_local_body_rules(
if !crate::provider_transport::apply_local_body_rules_with_request_headers(
&mut provider_request_body,
transport.endpoint.body_rules.as_ref(),
Some(&request_body),
Some(&incoming_request_headers),
) {
return Ok(provider_query_skipped_execution_outcome(
request_body.clone(),
@@ -1104,7 +1141,7 @@ async fn provider_query_execute_standard_test_candidate(
.uri(route_path)
.body(())
.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 request_url = crate::provider_transport::build_transport_request_url(
@@ -1166,12 +1203,13 @@ async fn provider_query_execute_standard_test_candidate(
} else {
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,
transport.endpoint.header_rules.as_ref(),
&protected_headers,
&provider_request_body,
Some(&request_body),
Some(&parts.headers),
) {
return Ok(ProviderQueryExecutionOutcome {
status: "failed",
@@ -1983,6 +2021,24 @@ mod tests {
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]
fn provider_query_request_body_model_uses_non_empty_string_only() {
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_claude_passthrough_headers(",
"build_openai_passthrough_headers(",
"apply_local_header_rules(",
"apply_local_header_rules_with_request_headers(",
"uses_vertex_api_key_query_auth(",
] {
assert!(