feat: add routing profile scheduling policies

This commit is contained in:
fawney19
2026-05-18 11:03:49 +08:00
parent a2f91b4108
commit 92813e6122
124 changed files with 11681 additions and 578 deletions
@@ -29,7 +29,7 @@ use self::support::{
pub(crate) struct LocalGeminiFilesSyncAttemptSource<'a> {
state: &'a AppState,
parts: &'a http::request::Parts,
body_json: &'a serde_json::Value,
body_json: serde_json::Value,
body_base64: Option<&'a str>,
body_is_empty: bool,
trace_id: &'a str,
@@ -110,10 +110,11 @@ pub(crate) async fn build_local_gemini_files_sync_attempt_source_for_kind<'a>(
trace_id,
decision,
)
.await
.await?
else {
return Ok(None);
};
let effective_body_json = input.effective_body_json(body_json).clone();
let (candidates, candidate_count) =
build_local_gemini_files_candidate_attempt_source(state, trace_id, &input).await?;
if candidate_count == 0 {
@@ -124,7 +125,7 @@ pub(crate) async fn build_local_gemini_files_sync_attempt_source_for_kind<'a>(
LocalGeminiFilesSyncAttemptSource {
state,
parts,
body_json,
body_json: effective_body_json,
body_base64,
body_is_empty,
trace_id,
@@ -148,7 +149,7 @@ pub(crate) async fn build_local_gemini_files_stream_attempt_source_for_kind<'a>(
};
let Some(input) =
resolve_local_gemini_files_decision_input(state, parts, None, trace_id, decision).await
resolve_local_gemini_files_decision_input(state, parts, None, trace_id, decision).await?
else {
return Ok(None);
};
@@ -226,7 +227,7 @@ impl LocalGeminiFilesSyncAttemptSource<'_> {
let Some(payload) = maybe_build_local_gemini_files_decision_payload_for_candidate(
self.state,
self.parts,
self.body_json,
&self.body_json,
self.body_base64,
self.body_is_empty,
self.trace_id,
@@ -234,7 +235,7 @@ impl LocalGeminiFilesSyncAttemptSource<'_> {
attempt,
self.spec,
)
.await
.await?
else {
return Ok(None);
};
@@ -272,7 +273,7 @@ impl LocalGeminiFilesStreamAttemptSource<'_> {
attempt,
self.spec,
)
.await
.await?
else {
return Ok(None);
};
@@ -313,10 +314,11 @@ pub(crate) async fn maybe_build_sync_local_gemini_files_decision_payload(
trace_id,
decision,
)
.await
.await?
else {
return Ok(None);
};
let body_json = input.effective_body_json(body_json);
let (mut source, _) =
build_local_gemini_files_candidate_attempt_source(state, trace_id, &input).await?;
@@ -333,7 +335,7 @@ pub(crate) async fn maybe_build_sync_local_gemini_files_decision_payload(
attempt,
spec,
)
.await
.await?
{
return Ok(Some(payload));
}
@@ -354,7 +356,7 @@ pub(crate) async fn maybe_build_stream_local_gemini_files_decision_payload(
};
let Some(input) =
resolve_local_gemini_files_decision_input(state, parts, None, trace_id, decision).await
resolve_local_gemini_files_decision_input(state, parts, None, trace_id, decision).await?
else {
return Ok(None);
};
@@ -375,7 +377,7 @@ pub(crate) async fn maybe_build_stream_local_gemini_files_decision_payload(
attempt,
spec,
)
.await
.await?
{
return Ok(Some(payload));
}
@@ -402,10 +404,11 @@ async fn build_local_sync_plan_and_reports(
trace_id,
decision,
)
.await
.await?
else {
return Ok(Vec::new());
};
let body_json = input.effective_body_json(body_json);
let (mut source, _) =
build_local_gemini_files_candidate_attempt_source(state, trace_id, &input).await?;
@@ -423,7 +426,7 @@ async fn build_local_sync_plan_and_reports(
attempt,
spec,
)
.await
.await?
else {
continue;
};
@@ -454,7 +457,7 @@ async fn build_local_stream_plan_and_reports(
) -> Result<Vec<AiStreamAttempt>, GatewayError> {
let spec_metadata = local_gemini_files_spec_metadata(spec);
let Some(input) =
resolve_local_gemini_files_decision_input(state, parts, None, trace_id, decision).await
resolve_local_gemini_files_decision_input(state, parts, None, trace_id, decision).await?
else {
return Ok(Vec::new());
};
@@ -476,7 +479,7 @@ async fn build_local_stream_plan_and_reports(
attempt,
spec,
)
.await
.await?
else {
continue;
};
@@ -1,6 +1,7 @@
use serde_json::json;
use crate::ai_serving::build_request_trace_proxy_value;
use crate::ai_serving::planner::decision_input::apply_provider_request_routing_policy_to_decision;
use crate::ai_serving::planner::report_context::{
build_local_execution_report_context, LocalExecutionReportContextParts,
};
@@ -12,7 +13,7 @@ use crate::ai_serving::transport::{
resolve_transport_execution_timeouts, resolve_transport_profile,
};
use crate::ai_serving::{ai_local_execution_contract_for_formats, PlannerAppState};
use crate::{AiExecutionDecision, AppState};
use crate::{AiExecutionDecision, AppState, GatewayError};
use super::request::resolve_local_gemini_files_candidate_payload_parts;
use super::support::{
@@ -31,7 +32,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
input: &LocalGeminiFilesDecisionInput,
attempt: LocalGeminiFilesCandidateAttempt,
spec: LocalGeminiFilesSpec,
) -> Option<AiExecutionDecision> {
) -> Result<Option<AiExecutionDecision>, GatewayError> {
let spec_metadata = local_gemini_files_spec_metadata(spec);
let planner_state = PlannerAppState::new(state);
let attempt_identity = attempt.attempt_identity();
@@ -46,7 +47,10 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
&attempt,
spec,
)
.await?;
.await;
let Some(resolved) = resolved else {
return Ok(None);
};
let LocalGeminiFilesCandidateAttempt {
eligible,
candidate_id,
@@ -69,6 +73,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
}
extra_fields.insert("file_key_id".to_string(), json!(candidate.key_id));
extra_fields.insert("file_name".to_string(), json!(resolved.file_name));
let effective_headers = input.effective_headers(&parts.headers);
let report_context = build_local_execution_report_context(LocalExecutionReportContextParts {
auth_context: &input.auth_context,
request_id: trace_id,
@@ -94,7 +99,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
body_rules: transport.endpoint.body_rules.as_ref(),
provider_request_method: None,
provider_request_headers: None,
original_headers: &parts.headers,
original_headers: effective_headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
@@ -119,45 +124,44 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
file_name: _,
} = resolved;
Some(build_ai_execution_decision_response(
AiExecutionDecisionResponseParts {
decision_is_stream: spec_metadata.require_streaming,
decision_kind: spec_metadata.decision_kind.to_string(),
execution_strategy,
conversion_mode,
request_id: trace_id.to_string(),
candidate_id: candidate_id.clone(),
provider_name: transport.provider.name.clone(),
provider_id: candidate.provider_id.clone(),
endpoint_id: candidate.endpoint_id.clone(),
key_id: candidate.key_id.clone(),
upstream_base_url: transport.endpoint.base_url.clone(),
upstream_url,
provider_request_method: Some(parts.method.to_string()),
auth_header: Some(auth_header),
auth_value: Some(auth_value),
provider_api_format: GEMINI_FILES_CLIENT_API_FORMAT.to_string(),
client_api_format: GEMINI_FILES_CLIENT_API_FORMAT.to_string(),
model_name: "gemini-files".to_string(),
mapped_model: candidate.selected_provider_model_name.clone(),
prompt_cache_key: None,
provider_request_headers,
provider_request_body,
provider_request_body_base64,
content_type: parts
.headers
.get(http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
proxy,
transport_profile,
timeouts: resolve_transport_execution_timeouts(&transport),
upstream_is_stream: spec_metadata.require_streaming,
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
report_context: Some(report_context),
auth_context: input.auth_context.clone(),
},
))
let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
decision_is_stream: spec_metadata.require_streaming,
decision_kind: spec_metadata.decision_kind.to_string(),
execution_strategy,
conversion_mode,
request_id: trace_id.to_string(),
candidate_id: candidate_id.clone(),
provider_name: transport.provider.name.clone(),
provider_id: candidate.provider_id.clone(),
endpoint_id: candidate.endpoint_id.clone(),
key_id: candidate.key_id.clone(),
upstream_base_url: transport.endpoint.base_url.clone(),
upstream_url,
provider_request_method: Some(parts.method.to_string()),
auth_header: Some(auth_header),
auth_value: Some(auth_value),
provider_api_format: GEMINI_FILES_CLIENT_API_FORMAT.to_string(),
client_api_format: GEMINI_FILES_CLIENT_API_FORMAT.to_string(),
model_name: "gemini-files".to_string(),
mapped_model: candidate.selected_provider_model_name.clone(),
prompt_cache_key: None,
provider_request_headers,
provider_request_body,
provider_request_body_base64,
content_type: effective_headers
.get(http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
proxy,
transport_profile,
timeouts: resolve_transport_execution_timeouts(&transport),
upstream_is_stream: spec_metadata.require_streaming,
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
report_context: Some(report_context),
auth_context: input.auth_context.clone(),
});
apply_provider_request_routing_policy_to_decision(input, &mut decision)?;
Ok(Some(decision))
}
@@ -45,6 +45,7 @@ pub(super) async fn resolve_local_gemini_files_candidate_payload_parts(
let spec_metadata = local_gemini_files_spec_metadata(spec);
let candidate = &attempt.eligible.candidate;
let transport = &attempt.eligible.transport;
let effective_headers = input.effective_headers(&parts.headers);
if let Some(skip_reason) =
gemini_files_transport_unsupported_reason(transport, GEMINI_FILES_CANDIDATE_API_FORMAT)
@@ -103,7 +104,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),
Some(effective_headers),
) {
Ok(parts) => parts,
Err(GeminiFilesRequestBodyError::BodyRulesUnsupportedForBinaryUpload) => {
@@ -145,7 +146,7 @@ pub(super) async fn resolve_local_gemini_files_candidate_payload_parts(
};
let Some(provider_request_headers) = build_gemini_files_headers(GeminiFilesHeadersInput {
headers: &parts.headers,
headers: effective_headers,
auth_header: &auth_header,
auth_value: &auth_value,
header_rules: transport.endpoint.header_rules.as_ref(),
@@ -14,7 +14,8 @@ use crate::ai_serving::planner::candidate_metadata::{
build_local_execution_candidate_metadata_for_candidate, LocalExecutionCandidateMetadataParts,
};
use crate::ai_serving::planner::decision_input::{
build_local_authenticated_decision_input, resolve_local_authenticated_decision_input,
attach_routing_policy_to_local_requested_model_input,
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
};
use crate::ai_serving::planner::materialization_policy::{
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
@@ -29,11 +30,12 @@ use crate::{AppState, GatewayError};
pub(super) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalGeminiFilesCandidateAttempt;
pub(super) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttemptSource as LocalGeminiFilesCandidateAttemptSource;
pub(super) use crate::ai_serving::planner::decision_input::LocalAuthenticatedDecisionInput as LocalGeminiFilesDecisionInput;
pub(super) use crate::ai_serving::planner::decision_input::LocalRequestedModelDecisionInput as LocalGeminiFilesDecisionInput;
pub(super) const GEMINI_FILES_CANDIDATE_API_FORMAT: &str = "gemini:files";
pub(super) const GEMINI_FILES_CLIENT_API_FORMAT: &str = "gemini:files";
pub(super) const GEMINI_FILES_REQUIRED_CAPABILITY: &str = "gemini_files";
pub(super) const GEMINI_FILES_ROUTING_MODEL: &str = "gemini-files";
pub(super) async fn resolve_local_gemini_files_decision_input(
state: &AppState,
@@ -41,9 +43,9 @@ pub(super) async fn resolve_local_gemini_files_decision_input(
body_json: Option<&serde_json::Value>,
trace_id: &str,
decision: &GatewayControlDecision,
) -> Option<LocalGeminiFilesDecisionInput> {
) -> Result<Option<LocalGeminiFilesDecisionInput>, GatewayError> {
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
return None;
return Ok(None);
};
let explicit_required_capabilities = json!({ "gemini_files": true });
@@ -56,20 +58,33 @@ pub(super) async fn resolve_local_gemini_files_decision_input(
.await
{
Ok(Some(resolved_input)) => resolved_input,
Ok(None) => return None,
Ok(None) => return Ok(None),
Err(err) => {
warn!(
trace_id = %trace_id,
error = ?err,
"gateway local gemini files decision auth snapshot read failed"
);
return None;
return Err(err);
}
};
let mut input = build_local_authenticated_decision_input(resolved_input);
let routing_body_json = body_json.cloned().unwrap_or(serde_json::Value::Null);
let mut input = build_local_requested_model_decision_input(
resolved_input,
GEMINI_FILES_ROUTING_MODEL.to_string(),
);
input.request_auth_channel = decision.request_auth_channel.clone();
input.client_session_affinity = client_session_affinity_from_parts(parts, body_json);
Some(input)
attach_routing_policy_to_local_requested_model_input(
state,
parts,
&mut input,
&routing_body_json,
GEMINI_FILES_CLIENT_API_FORMAT,
)
.await?;
Ok(Some(input))
}
pub(super) async fn materialize_local_gemini_files_candidate_attempts(
@@ -101,8 +116,9 @@ pub(super) async fn materialize_local_gemini_files_candidate_attempts(
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
input.required_capabilities.as_ref(),
input.routing_policy.as_ref(),
None,
None,
input.request_auth_channel.as_deref(),
persistence_policy,
candidates,
Vec::new(),
@@ -173,8 +189,9 @@ pub(super) async fn build_local_gemini_files_candidate_attempt_source<'a>(
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
input.required_capabilities.as_ref(),
input.routing_policy.as_ref(),
None,
None,
input.request_auth_channel.as_deref(),
persistence_policy,
candidates,
Vec::new(),
@@ -32,7 +32,7 @@ pub(super) use crate::ai_serving::LocalOpenAiImageSpec;
pub(crate) struct LocalOpenAiImageSyncAttemptSource<'a> {
state: &'a AppState,
parts: &'a http::request::Parts,
body_json: &'a serde_json::Value,
body_json: serde_json::Value,
body_base64: Option<&'a str>,
trace_id: &'a str,
input: LocalOpenAiImageDecisionInput,
@@ -43,7 +43,7 @@ pub(crate) struct LocalOpenAiImageSyncAttemptSource<'a> {
pub(crate) struct LocalOpenAiImageStreamAttemptSource<'a> {
state: &'a AppState,
parts: &'a http::request::Parts,
body_json: &'a serde_json::Value,
body_json: serde_json::Value,
body_base64: Option<&'a str>,
trace_id: &'a str,
input: LocalOpenAiImageDecisionInput,
@@ -152,16 +152,17 @@ pub(crate) async fn build_local_image_sync_attempt_source_for_kind<'a>(
trace_id,
decision,
)
.await
.await?
else {
return Ok(None);
};
let effective_body_json = input.effective_body_json(body_json).clone();
let Some((candidates, candidate_count)) = build_local_openai_image_candidate_attempt_source(
state,
trace_id,
&input,
body_json,
&effective_body_json,
spec_metadata.api_format,
spec_metadata.decision_kind,
)
@@ -178,7 +179,7 @@ pub(crate) async fn build_local_image_sync_attempt_source_for_kind<'a>(
LocalOpenAiImageSyncAttemptSource {
state,
parts,
body_json,
body_json: effective_body_json,
body_base64,
trace_id,
input,
@@ -211,16 +212,17 @@ pub(crate) async fn build_local_image_stream_attempt_source_for_kind<'a>(
trace_id,
decision,
)
.await
.await?
else {
return Ok(None);
};
let effective_body_json = input.effective_body_json(body_json).clone();
let Some((candidates, candidate_count)) = build_local_openai_image_candidate_attempt_source(
state,
trace_id,
&input,
body_json,
&effective_body_json,
spec_metadata.api_format,
spec_metadata.decision_kind,
)
@@ -237,7 +239,7 @@ pub(crate) async fn build_local_image_stream_attempt_source_for_kind<'a>(
LocalOpenAiImageStreamAttemptSource {
state,
parts,
body_json,
body_json: effective_body_json,
body_base64,
trace_id,
input,
@@ -303,21 +305,21 @@ impl LocalOpenAiImageSyncAttemptSource<'_> {
let Some(payload) = maybe_build_local_openai_image_decision_payload_for_candidate(
self.state,
self.parts,
self.body_json,
&self.body_json,
self.body_base64,
self.trace_id,
&self.input,
attempt,
self.spec,
)
.await
.await?
else {
return Ok(None);
};
let provider_api_format = payload.provider_api_format.as_deref().unwrap_or_default();
let built = if provider_api_format == "gemini:generate_content" {
build_gemini_sync_plan_from_decision(self.parts, self.body_json, payload)
build_gemini_sync_plan_from_decision(self.parts, &self.body_json, payload)
} else {
build_passthrough_sync_plan_from_decision(self.parts, payload)
};
@@ -345,23 +347,23 @@ impl LocalOpenAiImageStreamAttemptSource<'_> {
let Some(payload) = maybe_build_local_openai_image_decision_payload_for_candidate(
self.state,
self.parts,
self.body_json,
&self.body_json,
self.body_base64,
self.trace_id,
&self.input,
attempt,
self.spec,
)
.await
.await?
else {
return Ok(None);
};
let provider_api_format = payload.provider_api_format.as_deref().unwrap_or_default();
let built = if provider_api_format == "gemini:generate_content" {
build_gemini_stream_plan_from_decision(self.parts, self.body_json, payload)
build_gemini_stream_plan_from_decision(self.parts, &self.body_json, payload)
} else {
build_standard_stream_plan_from_decision(self.parts, self.body_json, payload, false)
build_standard_stream_plan_from_decision(self.parts, &self.body_json, payload, false)
};
match built {
Ok(value) => Ok(value),
@@ -400,10 +402,11 @@ pub(crate) async fn maybe_build_sync_local_image_decision_payload(
trace_id,
decision,
)
.await
.await?
else {
return Ok(None);
};
let body_json = input.effective_body_json(body_json);
let Some((mut source, _)) = build_local_openai_image_candidate_attempt_source(
state,
@@ -429,7 +432,7 @@ pub(crate) async fn maybe_build_sync_local_image_decision_payload(
attempt,
spec,
)
.await
.await?
{
return Ok(Some(payload));
}
@@ -460,10 +463,11 @@ pub(crate) async fn maybe_build_stream_local_image_decision_payload(
trace_id,
decision,
)
.await
.await?
else {
return Ok(None);
};
let body_json = input.effective_body_json(body_json);
let Some((mut source, _)) = build_local_openai_image_candidate_attempt_source(
state,
@@ -489,7 +493,7 @@ pub(crate) async fn maybe_build_stream_local_image_decision_payload(
attempt,
spec,
)
.await
.await?
{
return Ok(Some(payload));
}
@@ -516,10 +520,11 @@ async fn build_local_sync_plan_and_reports(
trace_id,
decision,
)
.await
.await?
else {
return Ok(Vec::new());
};
let body_json = input.effective_body_json(body_json);
let Some((mut source, _)) = build_local_openai_image_candidate_attempt_source(
state,
@@ -546,7 +551,7 @@ async fn build_local_sync_plan_and_reports(
attempt,
spec,
)
.await
.await?
else {
continue;
};
@@ -592,10 +597,11 @@ async fn build_local_stream_plan_and_reports(
trace_id,
decision,
)
.await
.await?
else {
return Ok(Vec::new());
};
let body_json = input.effective_body_json(body_json);
let Some((mut source, _)) = build_local_openai_image_candidate_attempt_source(
state,
@@ -622,7 +628,7 @@ async fn build_local_stream_plan_and_reports(
attempt,
spec,
)
.await
.await?
else {
continue;
};
@@ -1,4 +1,5 @@
use crate::ai_serving::build_request_trace_proxy_value;
use crate::ai_serving::planner::decision_input::apply_provider_request_routing_policy_to_decision;
use crate::ai_serving::planner::report_context::{
build_local_execution_report_context, LocalExecutionReportContextParts,
};
@@ -10,7 +11,9 @@ use crate::ai_serving::transport::{
resolve_transport_execution_timeouts, resolve_transport_profile,
};
use crate::ai_serving::{ai_local_execution_contract_for_formats, PlannerAppState};
use crate::{append_execution_contract_fields_to_value, AiExecutionDecision, AppState};
use crate::{
append_execution_contract_fields_to_value, AiExecutionDecision, AppState, GatewayError,
};
use super::request::resolve_local_openai_image_candidate_payload_parts;
use super::support::{LocalOpenAiImageCandidateAttempt, LocalOpenAiImageDecisionInput};
@@ -25,11 +28,11 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
input: &LocalOpenAiImageDecisionInput,
attempt: LocalOpenAiImageCandidateAttempt,
spec: LocalOpenAiImageSpec,
) -> Option<AiExecutionDecision> {
) -> Result<Option<AiExecutionDecision>, GatewayError> {
let spec_metadata = local_openai_image_spec_metadata(spec);
let planner_state = PlannerAppState::new(state);
let attempt_identity = attempt.attempt_identity();
let resolved = resolve_local_openai_image_candidate_payload_parts(
let Some(resolved) = resolve_local_openai_image_candidate_payload_parts(
state,
parts,
body_json,
@@ -39,7 +42,10 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
&attempt,
spec,
)
.await?;
.await
else {
return Ok(None);
};
let LocalOpenAiImageCandidateAttempt {
eligible,
candidate_id,
@@ -88,6 +94,7 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
.get("stream")
.and_then(serde_json::Value::as_bool)
.unwrap_or(spec_metadata.require_streaming);
let effective_headers = input.effective_headers(&parts.headers);
let report_context = append_execution_contract_fields_to_value(
build_local_execution_report_context(LocalExecutionReportContextParts {
auth_context: &input.auth_context,
@@ -114,7 +121,7 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
body_rules: transport.endpoint.body_rules.as_ref(),
provider_request_method: Some(serde_json::Value::String(parts.method.to_string())),
provider_request_headers: Some(&resolved.provider_request_headers),
original_headers: &parts.headers,
original_headers: effective_headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
@@ -134,39 +141,39 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
provider_api_format.as_str(),
);
Some(build_ai_execution_decision_response(
AiExecutionDecisionResponseParts {
decision_is_stream: spec_metadata.require_streaming,
decision_kind: spec_metadata.decision_kind.to_string(),
execution_strategy,
conversion_mode,
request_id: trace_id.to_string(),
candidate_id: candidate_id.clone(),
provider_name: transport.provider.name.clone(),
provider_id: candidate.provider_id.clone(),
endpoint_id: candidate.endpoint_id.clone(),
key_id: candidate.key_id.clone(),
upstream_base_url: transport.endpoint.base_url.clone(),
upstream_url: resolved.upstream_url,
provider_request_method: Some(parts.method.to_string()),
auth_header: Some(resolved.auth_header),
auth_value: Some(resolved.auth_value),
provider_api_format,
client_api_format: spec_metadata.api_format.to_string(),
model_name: resolved.requested_model,
mapped_model: resolved.mapped_model,
prompt_cache_key: None,
provider_request_headers: resolved.provider_request_headers,
provider_request_body: Some(resolved.provider_request_body),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
proxy,
transport_profile,
timeouts: resolve_transport_execution_timeouts(&transport),
upstream_is_stream,
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
report_context: Some(report_context),
auth_context: input.auth_context.clone(),
},
))
let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
decision_is_stream: spec_metadata.require_streaming,
decision_kind: spec_metadata.decision_kind.to_string(),
execution_strategy,
conversion_mode,
request_id: trace_id.to_string(),
candidate_id: candidate_id.clone(),
provider_name: transport.provider.name.clone(),
provider_id: candidate.provider_id.clone(),
endpoint_id: candidate.endpoint_id.clone(),
key_id: candidate.key_id.clone(),
upstream_base_url: transport.endpoint.base_url.clone(),
upstream_url: resolved.upstream_url,
provider_request_method: Some(parts.method.to_string()),
auth_header: Some(resolved.auth_header),
auth_value: Some(resolved.auth_value),
provider_api_format,
client_api_format: spec_metadata.api_format.to_string(),
model_name: resolved.requested_model,
mapped_model: resolved.mapped_model,
prompt_cache_key: None,
provider_request_headers: resolved.provider_request_headers,
provider_request_body: Some(resolved.provider_request_body),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
proxy,
transport_profile,
timeouts: resolve_transport_execution_timeouts(&transport),
upstream_is_stream,
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
report_context: Some(report_context),
auth_context: input.auth_context.clone(),
});
apply_provider_request_routing_policy_to_decision(input, &mut decision)?;
Ok(Some(decision))
}
@@ -59,6 +59,7 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
let candidate = &attempt.eligible.candidate;
let transport = &attempt.eligible.transport;
let provider_api_format = attempt.eligible.provider_api_format.as_str();
let effective_headers = input.effective_headers(&parts.headers);
if provider_api_format == "gemini:generate_content" {
return resolve_local_openai_image_to_gemini_candidate_payload_parts(
@@ -170,7 +171,7 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
let Some(mut provider_request_headers) =
build_openai_image_headers(ProviderOpenAiImageHeadersInput {
headers: &parts.headers,
headers: effective_headers,
auth_header: &auth_header,
auth_value: &auth_value,
header_rules: transport.endpoint.header_rules.as_ref(),
@@ -201,7 +202,7 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
apply_codex_openai_responses_special_headers(
&mut provider_request_headers,
&provider_request_body,
&parts.headers,
effective_headers,
transport.provider.provider_type.as_str(),
spec_metadata.api_format,
Some(trace_id),
@@ -256,6 +257,7 @@ async fn resolve_local_openai_image_to_gemini_candidate_payload_parts(
let candidate = &attempt.eligible.candidate;
let transport = &attempt.eligible.transport;
let provider_api_format = "gemini:generate_content";
let effective_headers = input.effective_headers(&parts.headers);
let prepared_candidate = match prepare_header_authenticated_candidate(
PlannerAppState::new(state),
@@ -332,7 +334,7 @@ async fn resolve_local_openai_image_to_gemini_candidate_payload_parts(
converted.body_json,
transport.endpoint.body_rules.as_ref(),
body_json,
&parts.headers,
effective_headers,
) {
Some(body) => body,
None => {
@@ -384,7 +386,7 @@ async fn resolve_local_openai_image_to_gemini_candidate_payload_parts(
transport,
provider_api_format,
same_format: false,
headers: &parts.headers,
headers: effective_headers,
auth_header: &prepared_candidate.auth_header,
auth_value: &prepared_candidate.auth_value,
extra_headers: &BTreeMap::new(),
@@ -13,6 +13,7 @@ use crate::ai_serving::planner::candidate_metadata::{
use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate;
use crate::ai_serving::planner::candidate_source::auth_snapshot_allows_cross_format_candidate;
use crate::ai_serving::planner::decision_input::{
attach_routing_policy_to_local_requested_model_input,
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
};
use crate::ai_serving::planner::materialization_policy::{
@@ -42,12 +43,16 @@ pub(super) async fn resolve_local_openai_image_decision_input(
body_base64: Option<&str>,
trace_id: &str,
decision: &GatewayControlDecision,
) -> Option<LocalOpenAiImageDecisionInput> {
) -> Result<Option<LocalOpenAiImageDecisionInput>, GatewayError> {
let Some(auth_context) = resolve_local_openai_image_auth_context(decision) else {
return None;
return Ok(None);
};
let requested_model = resolve_requested_image_model_for_request(parts, body_json, body_base64)?;
let Some(requested_model) =
resolve_requested_image_model_for_request(parts, body_json, body_base64)
else {
return Ok(None);
};
let resolved_input = match resolve_local_authenticated_decision_input(
state,
@@ -58,21 +63,37 @@ pub(super) async fn resolve_local_openai_image_decision_input(
.await
{
Ok(Some(resolved_input)) => resolved_input,
Ok(None) => return None,
Ok(None) => return Ok(None),
Err(err) => {
warn!(
trace_id = %trace_id,
error = ?err,
"gateway local openai image decision auth snapshot read failed"
);
return None;
return Err(err);
}
};
let mut input = build_local_requested_model_decision_input(resolved_input, requested_model);
input.request_auth_channel = decision.request_auth_channel.clone();
input.client_session_affinity = client_session_affinity_from_parts(parts, Some(body_json));
Some(input)
if let Err(err) = attach_routing_policy_to_local_requested_model_input(
state,
parts,
&mut input,
body_json,
"openai:image",
)
.await
{
warn!(
trace_id = %trace_id,
error = ?err,
"gateway local openai image decision routing profile resolution failed"
);
return Err(err);
}
Ok(Some(input))
}
fn resolve_local_openai_image_auth_context(
@@ -229,6 +250,7 @@ pub(super) async fn build_local_openai_image_candidate_attempt_source<'a>(
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
input.required_capabilities.as_ref(),
input.routing_policy.as_ref(),
sticky_session_token.as_deref(),
input.request_auth_channel.as_deref(),
persistence_policy,
@@ -305,6 +327,7 @@ async fn materialize_local_openai_image_candidate_attempts(
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
input.required_capabilities.as_ref(),
input.routing_policy.as_ref(),
sticky_session_token.as_deref(),
input.request_auth_channel.as_deref(),
persistence_policy,
@@ -27,7 +27,7 @@ use self::support::{
pub(crate) struct LocalVideoCreateSyncAttemptSource<'a> {
state: &'a AppState,
parts: &'a http::request::Parts,
body_json: &'a serde_json::Value,
body_json: serde_json::Value,
trace_id: &'a str,
input: LocalVideoCreateDecisionInput,
spec: LocalVideoCreateSpec,
@@ -65,16 +65,17 @@ pub(crate) async fn build_local_video_sync_attempt_source_for_kind<'a>(
let Some(input) = resolve_local_video_create_decision_input(
state, parts, trace_id, decision, body_json, spec,
)
.await
.await?
else {
return Ok(None);
};
let effective_body_json = input.effective_body_json(body_json).clone();
let Some((candidates, candidate_count)) = build_local_video_create_candidate_attempt_source(
state,
trace_id,
&input,
body_json,
&effective_body_json,
spec_metadata.api_format,
spec_metadata.decision_kind,
)
@@ -91,7 +92,7 @@ pub(crate) async fn build_local_video_sync_attempt_source_for_kind<'a>(
LocalVideoCreateSyncAttemptSource {
state,
parts,
body_json,
body_json: effective_body_json,
trace_id,
input,
spec,
@@ -133,13 +134,13 @@ impl LocalVideoCreateSyncAttemptSource<'_> {
let Some(payload) = maybe_build_local_video_create_decision_payload_for_candidate(
self.state,
self.parts,
self.body_json,
&self.body_json,
self.trace_id,
&self.input,
attempt,
self.spec,
)
.await
.await?
else {
return Ok(None);
};
@@ -175,10 +176,11 @@ pub(crate) async fn maybe_build_sync_local_video_decision_payload(
let Some(input) = resolve_local_video_create_decision_input(
state, parts, trace_id, decision, body_json, spec,
)
.await
.await?
else {
return Ok(None);
};
let body_json = input.effective_body_json(body_json);
let Some((mut source, _)) = build_local_video_create_candidate_attempt_source(
state,
@@ -197,7 +199,7 @@ pub(crate) async fn maybe_build_sync_local_video_decision_payload(
if let Some(payload) = maybe_build_local_video_create_decision_payload_for_candidate(
state, parts, body_json, trace_id, &input, attempt, spec,
)
.await
.await?
{
return Ok(Some(payload));
}
@@ -218,10 +220,11 @@ async fn build_local_sync_plan_and_reports(
let Some(input) = resolve_local_video_create_decision_input(
state, parts, trace_id, decision, body_json, spec,
)
.await
.await?
else {
return Ok(Vec::new());
};
let body_json = input.effective_body_json(body_json);
let Some((mut source, _)) = build_local_video_create_candidate_attempt_source(
state,
@@ -241,7 +244,7 @@ async fn build_local_sync_plan_and_reports(
let Some(payload) = maybe_build_local_video_create_decision_payload_for_candidate(
state, parts, body_json, trace_id, &input, attempt, spec,
)
.await
.await?
else {
continue;
};
@@ -1,4 +1,5 @@
use crate::ai_serving::build_request_trace_proxy_value;
use crate::ai_serving::planner::decision_input::apply_provider_request_routing_policy_to_decision;
use crate::ai_serving::planner::report_context::{
build_local_execution_report_context, LocalExecutionReportContextParts,
};
@@ -10,7 +11,7 @@ use crate::ai_serving::transport::{
resolve_transport_execution_timeouts, resolve_transport_profile,
};
use crate::ai_serving::{ai_local_execution_contract_for_formats, PlannerAppState};
use crate::{AiExecutionDecision, AppState};
use crate::{AiExecutionDecision, AppState, GatewayError};
use super::request::resolve_local_video_create_candidate_payload_parts;
use super::support::{LocalVideoCreateCandidateAttempt, LocalVideoCreateDecisionInput};
@@ -24,14 +25,17 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
input: &LocalVideoCreateDecisionInput,
attempt: LocalVideoCreateCandidateAttempt,
spec: LocalVideoCreateSpec,
) -> Option<AiExecutionDecision> {
) -> Result<Option<AiExecutionDecision>, GatewayError> {
let spec_metadata = local_video_create_spec_metadata(spec);
let planner_state = PlannerAppState::new(state);
let attempt_identity = attempt.attempt_identity();
let resolved = resolve_local_video_create_candidate_payload_parts(
let Some(resolved) = resolve_local_video_create_candidate_payload_parts(
state, parts, body_json, trace_id, input, &attempt, spec,
)
.await?;
.await
else {
return Ok(None);
};
let LocalVideoCreateCandidateAttempt {
eligible,
candidate_id,
@@ -50,6 +54,7 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
if let Some(proxy_value) = build_request_trace_proxy_value(Some(&transport), proxy.as_ref()) {
extra_fields.insert("proxy".to_string(), proxy_value);
}
let effective_headers = input.effective_headers(&parts.headers);
let report_context = build_local_execution_report_context(LocalExecutionReportContextParts {
auth_context: &input.auth_context,
request_id: trace_id,
@@ -75,7 +80,7 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
body_rules: transport.endpoint.body_rules.as_ref(),
provider_request_method: None,
provider_request_headers: None,
original_headers: &parts.headers,
original_headers: effective_headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
@@ -99,45 +104,45 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
upstream_url,
} = resolved;
Some(build_ai_execution_decision_response(
AiExecutionDecisionResponseParts {
decision_is_stream: false,
decision_kind: spec_metadata.decision_kind.to_string(),
execution_strategy,
conversion_mode,
request_id: trace_id.to_string(),
candidate_id: candidate_id.clone(),
provider_name: transport.provider.name.clone(),
provider_id: candidate.provider_id.clone(),
endpoint_id: candidate.endpoint_id.clone(),
key_id: candidate.key_id.clone(),
upstream_base_url: transport.endpoint.base_url.clone(),
upstream_url,
provider_request_method: Some(parts.method.to_string()),
auth_header: Some(auth_header),
auth_value: Some(auth_value),
provider_api_format: spec_metadata.api_format.to_string(),
client_api_format: spec_metadata.api_format.to_string(),
model_name: input.requested_model.clone(),
mapped_model,
prompt_cache_key: None,
provider_request_headers,
provider_request_body: Some(provider_request_body),
provider_request_body_base64: None,
content_type: parts
.headers
.get(http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
proxy,
transport_profile,
timeouts: resolve_transport_execution_timeouts(&transport),
upstream_is_stream: false,
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
report_context: Some(report_context),
auth_context: input.auth_context.clone(),
},
))
let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
decision_is_stream: false,
decision_kind: spec_metadata.decision_kind.to_string(),
execution_strategy,
conversion_mode,
request_id: trace_id.to_string(),
candidate_id: candidate_id.clone(),
provider_name: transport.provider.name.clone(),
provider_id: candidate.provider_id.clone(),
endpoint_id: candidate.endpoint_id.clone(),
key_id: candidate.key_id.clone(),
upstream_base_url: transport.endpoint.base_url.clone(),
upstream_url,
provider_request_method: Some(parts.method.to_string()),
auth_header: Some(auth_header),
auth_value: Some(auth_value),
provider_api_format: spec_metadata.api_format.to_string(),
client_api_format: spec_metadata.api_format.to_string(),
model_name: input.requested_model.clone(),
mapped_model,
prompt_cache_key: None,
provider_request_headers,
provider_request_body: Some(provider_request_body),
provider_request_body_base64: None,
content_type: parts
.headers
.get(http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
proxy,
transport_profile,
timeouts: resolve_transport_execution_timeouts(&transport),
upstream_is_stream: false,
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
report_context: Some(report_context),
auth_context: input.auth_context.clone(),
});
apply_provider_request_routing_policy_to_decision(input, &mut decision)?;
Ok(Some(decision))
}
@@ -41,6 +41,7 @@ pub(super) async fn resolve_local_video_create_candidate_payload_parts(
let spec_metadata = local_video_create_spec_metadata(spec);
let candidate = &attempt.eligible.candidate;
let transport = &attempt.eligible.transport;
let effective_headers = input.effective_headers(&parts.headers);
let provider_family = provider_video_create_family(spec.family);
let transport_unsupported_reason = video_create_transport_unsupported_reason(
@@ -124,7 +125,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),
Some(effective_headers),
) else {
mark_skipped_local_video_candidate_with_failure_diagnostic(
state,
@@ -146,7 +147,7 @@ pub(super) async fn resolve_local_video_create_candidate_payload_parts(
let Some(provider_request_headers) =
build_video_create_headers(ProviderVideoCreateHeadersInput {
headers: &parts.headers,
headers: effective_headers,
auth_header: &auth_header,
auth_value: &auth_value,
header_rules: transport.endpoint.header_rules.as_ref(),
@@ -15,6 +15,7 @@ use crate::ai_serving::planner::candidate_metadata::{
use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate;
use crate::ai_serving::planner::common::extract_requested_model_from_request;
use crate::ai_serving::planner::decision_input::{
attach_routing_policy_to_local_requested_model_input,
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
};
use crate::ai_serving::planner::materialization_policy::{
@@ -41,19 +42,21 @@ pub(super) async fn resolve_local_video_create_decision_input(
decision: &GatewayControlDecision,
body_json: &serde_json::Value,
spec: LocalVideoCreateSpec,
) -> Option<LocalVideoCreateDecisionInput> {
) -> Result<Option<LocalVideoCreateDecisionInput>, GatewayError> {
let spec_metadata = local_video_create_spec_metadata(spec);
let Some(auth_context) = resolve_local_video_create_auth_context(decision, spec.family) else {
return None;
return Ok(None);
};
let requested_model = extract_requested_model_from_request(
let Some(requested_model) = extract_requested_model_from_request(
parts,
body_json,
spec_metadata
.requested_model_family
.expect("video specs should declare requested-model family"),
)?;
) else {
return Ok(None);
};
let resolved_input = match resolve_local_authenticated_decision_input(
state,
@@ -64,7 +67,7 @@ pub(super) async fn resolve_local_video_create_decision_input(
.await
{
Ok(Some(resolved_input)) => resolved_input,
Ok(None) => return None,
Ok(None) => return Ok(None),
Err(err) => {
warn!(
trace_id = %trace_id,
@@ -72,14 +75,31 @@ pub(super) async fn resolve_local_video_create_decision_input(
error = ?err,
"gateway local video decision auth snapshot read failed"
);
return None;
return Err(err);
}
};
let mut input = build_local_requested_model_decision_input(resolved_input, requested_model);
input.request_auth_channel = decision.request_auth_channel.clone();
input.client_session_affinity = client_session_affinity_from_parts(parts, Some(body_json));
Some(input)
if let Err(err) = attach_routing_policy_to_local_requested_model_input(
state,
parts,
&mut input,
body_json,
spec_metadata.api_format,
)
.await
{
warn!(
trace_id = %trace_id,
decision_kind = spec_metadata.decision_kind,
error = ?err,
"gateway local video decision routing profile resolution failed"
);
return Err(err);
}
Ok(Some(input))
}
fn resolve_local_video_create_auth_context(
@@ -196,6 +216,7 @@ pub(super) async fn build_local_video_create_candidate_attempt_source<'a>(
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
input.required_capabilities.as_ref(),
input.routing_policy.as_ref(),
sticky_session_token.as_deref(),
input.request_auth_channel.as_deref(),
persistence_policy,
@@ -261,6 +282,7 @@ async fn materialize_local_video_create_candidate_attempts(
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
input.required_capabilities.as_ref(),
input.routing_policy.as_ref(),
sticky_session_token.as_deref(),
input.request_auth_channel.as_deref(),
persistence_policy,