Merge branch 'aether-rust-pioneer' of https://github.com/fawney19/Aether into codex/codex-image-progress-heartbeat

# Conflicts:
#	frontend/src/features/usage/components/__tests__/HorizontalRequestTimeline.spec.ts
This commit is contained in:
Entropy.Xu
2026-05-09 01:26:44 +08:00
90 changed files with 8542 additions and 1252 deletions

View File

@@ -75,6 +75,18 @@ enum LocalExecutionCandidateAttemptSourceItem<'a> {
}
impl<'a> LocalExecutionCandidateAttemptSource<'a> {
pub(crate) fn from_static_attempts_for_image_bridge(
attempts: Vec<LocalExecutionCandidateAttempt>,
) -> Self {
let mut items = VecDeque::new();
if !attempts.is_empty() {
items.push_back(LocalExecutionCandidateAttemptSourceItem::Static {
attempts: VecDeque::from(attempts),
});
}
Self { items }
}
pub(crate) async fn next_attempt(&mut self) -> Option<LocalExecutionCandidateAttempt> {
loop {
let front = self.items.front_mut()?;

View File

@@ -398,15 +398,11 @@ pub(crate) fn candidate_auth_channel_skip_reason(
) -> Option<&'static str> {
let request_auth_channel = normalize_request_auth_channel(request_auth_channel?)?;
let upstream_auth_channel = resolve_transport_request_auth_channel(transport)?;
if request_auth_channel == upstream_auth_channel
|| provider_runtime_policy(&transport.provider.provider_type)
.allow_auth_channel_mismatch_by_default
|| allow_auth_channel_mismatch_for_format(transport)
{
None
} else {
Some("auth_channel_mismatch")
if request_auth_channel == upstream_auth_channel {
return None;
}
auth_channel_mismatch_is_explicitly_disabled_for_format(transport)
.then_some("auth_channel_mismatch")
}
fn normalize_request_auth_channel(value: &str) -> Option<&'static str> {
@@ -452,19 +448,22 @@ fn resolve_transport_auth_type_for_endpoint_format(
.unwrap_or(default_auth_type)
}
fn allow_auth_channel_mismatch_for_format(transport: &GatewayProviderTransportSnapshot) -> bool {
fn auth_channel_mismatch_is_explicitly_disabled_for_format(
transport: &GatewayProviderTransportSnapshot,
) -> bool {
let api_format = crate::ai_serving::normalize_api_format_alias(&transport.endpoint.api_format);
transport
let Some(items) = transport
.key
.allow_auth_channel_mismatch_formats
.as_ref()
.and_then(serde_json::Value::as_array)
.is_some_and(|items| {
items
.iter()
.filter_map(serde_json::Value::as_str)
.any(|item| crate::ai_serving::normalize_api_format_alias(item) == api_format)
})
else {
return false;
};
!items
.iter()
.filter_map(serde_json::Value::as_str)
.any(|item| crate::ai_serving::normalize_api_format_alias(item) == api_format)
}
pub(crate) async fn read_candidate_transport_snapshot(
@@ -585,11 +584,11 @@ mod tests {
}
#[test]
fn auth_channel_gate_skips_mismatched_raw_secret_auth() {
fn auth_channel_gate_allows_mismatched_raw_secret_auth_by_default() {
let transport = sample_transport("bearer");
assert_eq!(
candidate_auth_channel_skip_reason(&transport, Some("api_key")),
Some("auth_channel_mismatch")
None
);
}
@@ -603,6 +602,16 @@ mod tests {
);
}
#[test]
fn auth_channel_gate_blocks_explicitly_disabled_mismatch_format() {
let mut transport = sample_transport("bearer");
transport.key.allow_auth_channel_mismatch_formats = Some(json!(["openai:responses"]));
assert_eq!(
candidate_auth_channel_skip_reason(&transport, Some("api_key")),
Some("auth_channel_mismatch")
);
}
#[test]
fn auth_channel_gate_treats_cli_oauth_provider_as_bearer_like() {
let mut transport = sample_transport("oauth");
@@ -613,7 +622,7 @@ mod tests {
);
assert_eq!(
candidate_auth_channel_skip_reason(&transport, Some("api_key")),
Some("auth_channel_mismatch")
None
);
}

View File

@@ -159,6 +159,40 @@ pub(crate) async fn preselect_local_execution_candidates_with_serving(
.into_iter()
.map(str::to_string)
.collect::<Vec<_>>();
preselect_local_execution_candidates_for_api_formats_with_serving(
state,
client_api_format,
requested_model,
require_streaming,
required_capabilities,
auth_snapshot,
client_session_affinity,
use_api_format_alias_match,
key_mode,
candidate_api_formats,
)
.await
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn preselect_local_execution_candidates_for_api_formats_with_serving(
state: PlannerAppState<'_>,
client_api_format: &str,
requested_model: &str,
require_streaming: bool,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: &GatewayAuthApiKeySnapshot,
client_session_affinity: Option<&ClientSessionAffinity>,
use_api_format_alias_match: bool,
key_mode: LocalCandidatePreselectionKeyMode,
candidate_api_formats: Vec<String>,
) -> Result<
AiCandidatePreselectionOutcome<
SchedulerMinimalCandidateSelectionCandidate,
SkippedLocalExecutionCandidate,
>,
GatewayError,
> {
let mut model_directive_enabled_api_formats = BTreeSet::new();
for api_format in &candidate_api_formats {
if crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model(

View File

@@ -7,6 +7,7 @@ use tracing::warn;
use crate::ai_serving::planner::candidate_materialization::LocalExecutionAttemptSource;
use crate::ai_serving::planner::plan_builders::{
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
build_passthrough_sync_plan_from_decision, build_standard_stream_plan_from_decision,
AiStreamAttempt, AiSyncAttempt,
};
@@ -314,7 +315,13 @@ impl LocalOpenAiImageSyncAttemptSource<'_> {
return Ok(None);
};
match build_passthrough_sync_plan_from_decision(self.parts, payload) {
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)
} else {
build_passthrough_sync_plan_from_decision(self.parts, payload)
};
match built {
Ok(value) => Ok(value),
Err(err) => {
warn!(
@@ -350,7 +357,13 @@ impl LocalOpenAiImageStreamAttemptSource<'_> {
return Ok(None);
};
match build_standard_stream_plan_from_decision(self.parts, self.body_json, payload, false) {
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)
} else {
build_standard_stream_plan_from_decision(self.parts, self.body_json, payload, false)
};
match built {
Ok(value) => Ok(value),
Err(err) => {
warn!(
@@ -538,7 +551,13 @@ async fn build_local_sync_plan_and_reports(
continue;
};
match build_passthrough_sync_plan_from_decision(parts, payload) {
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(parts, body_json, payload)
} else {
build_passthrough_sync_plan_from_decision(parts, payload)
};
match built {
Ok(Some(value)) => plans.push(value),
Ok(None) => {}
Err(err) => {
@@ -608,7 +627,13 @@ async fn build_local_stream_plan_and_reports(
continue;
};
match build_standard_stream_plan_from_decision(parts, body_json, payload, false) {
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(parts, body_json, payload)
} else {
build_standard_stream_plan_from_decision(parts, body_json, payload, false)
};
match built {
Ok(Some(value)) => plans.push(value),
Ok(None) => {}
Err(err) => {

View File

@@ -10,7 +10,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::{append_execution_contract_fields_to_value, AiExecutionDecision, AppState};
use super::request::resolve_local_openai_image_candidate_payload_parts;
use super::support::{LocalOpenAiImageCandidateAttempt, LocalOpenAiImageDecisionInput};
@@ -47,8 +47,12 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
} = attempt;
let candidate = eligible.candidate;
let transport = resolved.transport;
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats(spec_metadata.api_format, spec_metadata.api_format);
let provider_api_format = resolved.provider_api_format.clone();
let needs_conversion = provider_api_format != spec_metadata.api_format;
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
spec_metadata.api_format,
provider_api_format.as_str(),
);
let proxy = planner_state
.app()
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&transport)
@@ -84,41 +88,47 @@ 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 report_context = build_local_execution_report_context(LocalExecutionReportContextParts {
auth_context: &input.auth_context,
request_id: trace_id,
candidate_id: &candidate_id,
attempt_identity,
model: &resolved.requested_model,
provider_name: &transport.provider.name,
provider_id: &candidate.provider_id,
endpoint_id: &candidate.endpoint_id,
key_id: &candidate.key_id,
key_name: None,
model_id: Some(&candidate.model_id),
global_model_id: Some(&candidate.global_model_id),
global_model_name: Some(&candidate.global_model_name),
provider_api_format: spec_metadata.api_format,
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
ranking: eligible.ranking.as_ref(),
upstream_url: Some(&resolved.upstream_url),
header_rules: transport.endpoint.header_rules.as_ref(),
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,
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: body_base64,
client_session_affinity: input.client_session_affinity.as_ref(),
client_requested_stream: spec_metadata.require_streaming,
upstream_is_stream,
has_envelope: false,
needs_conversion: false,
extra_fields,
});
let report_context = append_execution_contract_fields_to_value(
build_local_execution_report_context(LocalExecutionReportContextParts {
auth_context: &input.auth_context,
request_id: trace_id,
candidate_id: &candidate_id,
attempt_identity,
model: &resolved.requested_model,
provider_name: &transport.provider.name,
provider_id: &candidate.provider_id,
endpoint_id: &candidate.endpoint_id,
key_id: &candidate.key_id,
key_name: None,
model_id: Some(&candidate.model_id),
global_model_id: Some(&candidate.global_model_id),
global_model_name: Some(&candidate.global_model_name),
provider_api_format: provider_api_format.as_str(),
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
ranking: eligible.ranking.as_ref(),
upstream_url: Some(&resolved.upstream_url),
header_rules: transport.endpoint.header_rules.as_ref(),
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,
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: body_base64,
client_session_affinity: input.client_session_affinity.as_ref(),
client_requested_stream: spec_metadata.require_streaming,
upstream_is_stream,
has_envelope: false,
needs_conversion,
extra_fields,
}),
execution_strategy,
conversion_mode,
spec_metadata.api_format,
provider_api_format.as_str(),
);
Some(build_ai_execution_decision_response(
AiExecutionDecisionResponseParts {
@@ -137,7 +147,7 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
provider_request_method: Some(parts.method.to_string()),
auth_header: Some(resolved.auth_header),
auth_value: Some(resolved.auth_value),
provider_api_format: spec_metadata.api_format.to_string(),
provider_api_format,
client_api_format: spec_metadata.api_format.to_string(),
model_name: resolved.requested_model,
mapped_model: resolved.mapped_model,

View File

@@ -9,14 +9,17 @@ use crate::ai_serving::planner::candidate_preparation::{
use crate::ai_serving::planner::spec_metadata::local_openai_image_spec_metadata;
use crate::ai_serving::transport::{
build_openai_image_headers, build_openai_image_upstream_url,
openai_image_transport_unsupported_reason, resolve_openai_image_auth,
ProviderOpenAiImageHeadersInput,
build_standard_provider_request_headers, openai_image_transport_unsupported_reason,
resolve_openai_image_auth, ProviderOpenAiImageHeadersInput,
StandardProviderRequestHeadersInput,
};
use crate::ai_serving::{
apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers,
build_chatgpt_web_image_request_body, build_openai_image_provider_request_body,
default_model_for_openai_image_operation, normalize_openai_image_request,
CandidateFailureDiagnostic, GatewayProviderTransportSnapshot, PlannerAppState,
build_chatgpt_web_image_request_body,
build_gemini_image_request_body_from_openai_image_request,
build_openai_image_provider_request_body, default_model_for_openai_image_operation,
normalize_openai_image_request, request_conversion_direct_auth, CandidateFailureDiagnostic,
GatewayProviderTransportSnapshot, PlannerAppState, RequestConversionKind,
};
use crate::AppState;
@@ -35,6 +38,7 @@ pub(super) struct LocalOpenAiImageCandidatePayloadParts {
pub(super) auth_value: String,
pub(super) requested_model: String,
pub(super) mapped_model: String,
pub(super) provider_api_format: String,
pub(super) provider_request_headers: BTreeMap<String, String>,
pub(super) provider_request_body: Value,
pub(super) upstream_url: String,
@@ -54,6 +58,21 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
let spec_metadata = local_openai_image_spec_metadata(spec);
let candidate = &attempt.eligible.candidate;
let transport = &attempt.eligible.transport;
let provider_api_format = attempt.eligible.provider_api_format.as_str();
if provider_api_format == "gemini:generate_content" {
return resolve_local_openai_image_to_gemini_candidate_payload_parts(
state,
parts,
body_json,
body_base64,
trace_id,
input,
attempt,
spec,
)
.await;
}
if let Some(skip_reason) =
openai_image_transport_unsupported_reason(transport, spec_metadata.api_format)
@@ -70,7 +89,6 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
.await;
return None;
}
let prepared_candidate = match prepare_header_authenticated_candidate(
PlannerAppState::new(state),
transport,
@@ -216,6 +234,7 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
auth_value,
requested_model,
mapped_model,
provider_api_format: spec_metadata.api_format.to_string(),
provider_request_headers,
provider_request_body,
upstream_url,
@@ -223,6 +242,190 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
})
}
async fn resolve_local_openai_image_to_gemini_candidate_payload_parts(
state: &AppState,
parts: &http::request::Parts,
body_json: &Value,
body_base64: Option<&str>,
trace_id: &str,
input: &LocalOpenAiImageDecisionInput,
attempt: &LocalOpenAiImageCandidateAttempt,
spec: LocalOpenAiImageSpec,
) -> Option<LocalOpenAiImageCandidatePayloadParts> {
let spec_metadata = local_openai_image_spec_metadata(spec);
let candidate = &attempt.eligible.candidate;
let transport = &attempt.eligible.transport;
let provider_api_format = "gemini:generate_content";
let prepared_candidate = match prepare_header_authenticated_candidate(
PlannerAppState::new(state),
transport,
candidate,
request_conversion_direct_auth(transport, RequestConversionKind::ToGeminiStandard),
OauthPreparationContext {
trace_id,
api_format: provider_api_format,
operation: "openai_image_to_gemini_candidate_request",
},
)
.await
{
Ok(prepared) => prepared,
Err(skip_reason) => {
mark_skipped_local_openai_image_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
skip_reason,
)
.await;
return None;
}
};
let Some(normalized_request) = normalize_openai_image_request(parts, body_json, body_base64)
else {
mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"provider_request_body_missing",
CandidateFailureDiagnostic::provider_request_body_missing(
spec_metadata.api_format,
provider_api_format,
"openai_image_request_normalize",
),
)
.await;
return None;
};
let Some(mut converted) = build_gemini_image_request_body_from_openai_image_request(
&normalized_request,
&prepared_candidate.mapped_model,
) else {
mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"provider_request_body_missing",
CandidateFailureDiagnostic::provider_request_body_missing(
spec_metadata.api_format,
provider_api_format,
"openai_image_to_gemini_request_body",
),
)
.await;
return None;
};
converted.body_json =
match crate::ai_serving::transport::apply_standard_provider_request_body_rules_with_request_headers(
converted.body_json,
transport.endpoint.body_rules.as_ref(),
body_json,
&parts.headers,
) {
Some(body) => body,
None => {
mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"provider_request_body_missing",
CandidateFailureDiagnostic::provider_request_body_missing(
spec_metadata.api_format,
provider_api_format,
"openai_image_to_gemini_body_rules",
),
)
.await;
return None;
}
};
let upstream_is_stream = spec_metadata.require_streaming;
let Some(upstream_url) = crate::ai_serving::planner::standard::build_standard_upstream_url(
parts,
transport,
&converted.mapped_model,
provider_api_format,
upstream_is_stream,
) else {
mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"upstream_url_missing",
CandidateFailureDiagnostic::upstream_url_missing(
spec_metadata.api_format,
provider_api_format,
"openai_image_to_gemini_url",
),
)
.await;
return None;
};
let Some(resolved_headers) =
build_standard_provider_request_headers(StandardProviderRequestHeadersInput {
transport,
provider_api_format,
same_format: false,
headers: &parts.headers,
auth_header: &prepared_candidate.auth_header,
auth_value: &prepared_candidate.auth_value,
extra_headers: &BTreeMap::new(),
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &converted.body_json,
original_request_body: body_json,
upstream_is_stream,
})
else {
mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_header_rules_apply_failed",
CandidateFailureDiagnostic::header_rules_apply_failed(
spec_metadata.api_format,
provider_api_format,
"openai_image_to_gemini_headers",
),
)
.await;
return None;
};
Some(LocalOpenAiImageCandidatePayloadParts {
transport: Arc::clone(transport),
auth_header: resolved_headers.auth_header,
auth_value: resolved_headers.auth_value,
requested_model: converted.requested_model,
mapped_model: converted.mapped_model,
provider_api_format: provider_api_format.to_string(),
provider_request_headers: resolved_headers.headers,
provider_request_body: converted.body_json,
upstream_url,
input_summary: converted.summary_json,
})
}
fn chatgpt_web_image_internal_url(base_url: &str) -> String {
let base_url = base_url.trim().trim_end_matches('/');
let base_url = if base_url.is_empty() {

View File

@@ -11,6 +11,7 @@ use crate::ai_serving::planner::candidate_metadata::{
build_local_execution_candidate_metadata_for_candidate, LocalExecutionCandidateMetadataParts,
};
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::{
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
};
@@ -19,9 +20,9 @@ use crate::ai_serving::planner::materialization_policy::{
};
use crate::ai_serving::planner::spec_metadata::local_openai_image_spec_metadata;
use crate::ai_serving::{
extract_pool_sticky_session_token, resolve_local_decision_execution_runtime_auth_context,
CandidateFailureDiagnostic, ExecutionRuntimeAuthContext, GatewayControlDecision,
PlannerAppState,
extract_pool_sticky_session_token, request_candidate_api_formats,
resolve_local_decision_execution_runtime_auth_context, CandidateFailureDiagnostic,
ExecutionRuntimeAuthContext, GatewayControlDecision, PlannerAppState,
};
use crate::client_session_affinity::client_session_affinity_from_parts;
use crate::clock::current_unix_secs;
@@ -88,52 +89,69 @@ pub(super) async fn list_local_openai_image_candidate_attempts(
api_format: &str,
decision_kind: &str,
) -> Option<Vec<LocalOpenAiImageCandidateAttempt>> {
let planner_state = PlannerAppState::new(state);
let (candidates, preselection_skipped) = match planner_state
.list_selectable_candidates_with_skip_reasons(
api_format,
&input.requested_model,
false,
input.required_capabilities.as_ref(),
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
current_unix_secs(),
)
.await
{
Ok(candidates) => candidates,
Err(err) => {
warn!(
trace_id = %trace_id,
decision_kind,
error = ?err,
"gateway local openai image decision scheduler selection failed"
);
return None;
let candidate_api_formats = image_candidate_api_formats(api_format);
let mut attempts = Vec::new();
for candidate_api_format in candidate_api_formats {
let matches_client_format = candidate_api_format == api_format;
let planner_state = PlannerAppState::new(state);
let (mut candidates, preselection_skipped) = match planner_state
.list_selectable_candidates_with_skip_reasons(
candidate_api_format,
&input.requested_model,
false,
input.required_capabilities.as_ref(),
matches_client_format.then_some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
current_unix_secs(),
)
.await
{
Ok(candidates) => candidates,
Err(err) => {
warn!(
trace_id = %trace_id,
decision_kind,
api_format = candidate_api_format,
error = ?err,
"gateway local openai image decision scheduler selection failed"
);
continue;
}
};
if !matches_client_format {
candidates.retain(|candidate| {
auth_snapshot_allows_cross_format_candidate(
&input.auth_snapshot,
&input.requested_model,
candidate,
false,
)
});
}
};
attempts.extend(
materialize_local_openai_image_candidate_attempts(
planner_state,
trace_id,
input,
body_json,
candidates,
preselection_skipped
.into_iter()
.map(|item| SkippedLocalExecutionCandidate {
candidate: item.candidate,
skip_reason: item.skip_reason,
transport: None,
ranking: None,
extra_data: None,
})
.collect(),
api_format,
)
.await,
);
}
Some(
materialize_local_openai_image_candidate_attempts(
planner_state,
trace_id,
input,
body_json,
candidates,
preselection_skipped
.into_iter()
.map(|item| SkippedLocalExecutionCandidate {
candidate: item.candidate,
skip_reason: item.skip_reason,
transport: None,
ranking: None,
extra_data: None,
})
.collect(),
api_format,
)
.await,
)
Some(attempts)
}
pub(super) async fn build_local_openai_image_candidate_attempt_source<'a>(
@@ -145,29 +163,56 @@ pub(super) async fn build_local_openai_image_candidate_attempt_source<'a>(
decision_kind: &str,
) -> Result<Option<(LocalOpenAiImageCandidateAttemptSource<'a>, usize)>, GatewayError> {
let planner_state = PlannerAppState::new(state);
let (candidates, preselection_skipped) = match planner_state
.list_selectable_candidates_with_skip_reasons(
api_format,
&input.requested_model,
false,
input.required_capabilities.as_ref(),
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
current_unix_secs(),
)
.await
{
Ok(candidates) => candidates,
Err(err) => {
warn!(
trace_id = %trace_id,
decision_kind,
error = ?err,
"gateway local openai image decision scheduler selection failed"
);
return Ok(None);
let mut candidates = Vec::new();
let mut preselection_skipped = Vec::new();
for candidate_api_format in image_candidate_api_formats(api_format) {
let matches_client_format = candidate_api_format == api_format;
match planner_state
.list_selectable_candidates_with_skip_reasons(
candidate_api_format,
&input.requested_model,
false,
input.required_capabilities.as_ref(),
matches_client_format.then_some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
current_unix_secs(),
)
.await
{
Ok((mut format_candidates, mut format_skipped)) => {
if !matches_client_format {
format_candidates.retain(|candidate| {
auth_snapshot_allows_cross_format_candidate(
&input.auth_snapshot,
&input.requested_model,
candidate,
false,
)
});
format_skipped.retain(|candidate| {
auth_snapshot_allows_cross_format_candidate(
&input.auth_snapshot,
&input.requested_model,
&candidate.candidate,
false,
)
});
}
candidates.append(&mut format_candidates);
preselection_skipped.append(&mut format_skipped);
}
Err(err) => {
warn!(
trace_id = %trace_id,
decision_kind,
api_format = candidate_api_format,
error = ?err,
"gateway local openai image decision scheduler selection failed"
);
continue;
}
}
};
}
let sticky_session_token = extract_pool_sticky_session_token(body_json);
let persistence_policy = build_local_candidate_persistence_policy(
@@ -198,23 +243,34 @@ pub(super) async fn build_local_openai_image_candidate_attempt_source<'a>(
extra_data: None,
})
.collect(),
LocalCandidateResolutionMode::Standard,
LocalCandidateResolutionMode::WithoutTransportPairGate,
|eligible| {
Some(build_local_execution_candidate_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: api_format,
provider_api_format: eligible.provider_api_format.as_str(),
client_api_format: api_format,
extra_fields: serde_json::Map::new(),
},
))
},
|mut skipped_candidate| {
let provider_api_format = skipped_candidate
.transport
.as_ref()
.map(|transport| transport.endpoint.api_format.trim().to_ascii_lowercase())
.unwrap_or_else(|| {
skipped_candidate
.candidate
.endpoint_api_format
.trim()
.to_ascii_lowercase()
});
skipped_candidate.extra_data =
Some(build_local_execution_candidate_metadata_for_candidate(
&skipped_candidate.candidate,
skipped_candidate.transport_ref(),
api_format,
provider_api_format.as_str(),
api_format,
serde_json::Map::new(),
));
@@ -254,23 +310,34 @@ async fn materialize_local_openai_image_candidate_attempts(
persistence_policy,
candidates,
preselection_skipped,
LocalCandidateResolutionMode::Standard,
LocalCandidateResolutionMode::WithoutTransportPairGate,
|eligible| {
Some(build_local_execution_candidate_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: api_format,
provider_api_format: eligible.provider_api_format.as_str(),
client_api_format: api_format,
extra_fields: serde_json::Map::new(),
},
))
},
|mut skipped_candidate| {
let provider_api_format = skipped_candidate
.transport
.as_ref()
.map(|transport| transport.endpoint.api_format.trim().to_ascii_lowercase())
.unwrap_or_else(|| {
skipped_candidate
.candidate
.endpoint_api_format
.trim()
.to_ascii_lowercase()
});
skipped_candidate.extra_data =
Some(build_local_execution_candidate_metadata_for_candidate(
&skipped_candidate.candidate,
skipped_candidate.transport_ref(),
api_format,
provider_api_format.as_str(),
api_format,
serde_json::Map::new(),
));
@@ -282,6 +349,14 @@ async fn materialize_local_openai_image_candidate_attempts(
outcome.attempts
}
fn image_candidate_api_formats(api_format: &str) -> Vec<&'static str> {
if api_format.trim().eq_ignore_ascii_case("openai:image") {
vec!["openai:image", "gemini:generate_content"]
} else {
request_candidate_api_formats(api_format, false)
}
}
pub(super) async fn mark_skipped_local_openai_image_candidate(
state: &AppState,
input: &LocalOpenAiImageDecisionInput,

View File

@@ -11,6 +11,7 @@ use crate::ai_serving::planner::candidate_metadata::{
LocalExecutionCandidateMetadataParts,
};
use crate::ai_serving::planner::candidate_source::{
preselect_local_execution_candidates_for_api_formats_with_serving,
preselect_local_execution_candidates_with_serving, LocalCandidatePreselectionKeyMode,
};
use crate::ai_serving::planner::common::extract_requested_model_from_request;
@@ -23,7 +24,8 @@ use crate::ai_serving::planner::materialization_policy::{
use crate::ai_serving::planner::spec_metadata::local_standard_spec_metadata;
use crate::ai_serving::{
ai_local_execution_contract_for_formats, extract_pool_sticky_session_token,
resolve_local_decision_execution_runtime_auth_context, GatewayControlDecision, PlannerAppState,
gemini_request_is_image_generation, resolve_local_decision_execution_runtime_auth_context,
GatewayControlDecision, PlannerAppState,
};
use crate::client_session_affinity::client_session_affinity_from_parts;
use crate::{AppState, GatewayError};
@@ -86,6 +88,8 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
spec: LocalStandardSpec,
) -> Result<(Vec<LocalStandardCandidateAttempt>, usize), GatewayError> {
let spec_metadata = local_standard_spec_metadata(spec);
let is_gemini_image_bridge = spec_metadata.api_format == "gemini:generate_content"
&& gemini_request_is_image_generation(body_json);
let planner_state = PlannerAppState::new(state);
let sticky_session_token = extract_pool_sticky_session_token(body_json);
let persistence_policy = build_local_candidate_persistence_policy(
@@ -105,6 +109,17 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
)
.await?;
let (candidates, skipped_candidates) = maybe_append_gemini_image_openai_image_preselection(
state,
planner_state,
trace_id,
input,
body_json,
spec,
preselection.candidates,
preselection.skipped_candidates,
)
.await?;
let outcome = materialize_local_execution_candidates_with_serving(
planner_state,
trace_id,
@@ -116,9 +131,13 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
sticky_session_token.as_deref(),
input.request_auth_channel.as_deref(),
persistence_policy,
preselection.candidates,
preselection.skipped_candidates,
LocalCandidateResolutionMode::Standard,
candidates,
skipped_candidates,
if is_gemini_image_bridge {
LocalCandidateResolutionMode::WithoutTransportPairGate
} else {
LocalCandidateResolutionMode::Standard
},
|eligible| {
let provider_api_format = eligible.provider_api_format.clone();
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
@@ -188,7 +207,18 @@ pub(super) async fn build_local_standard_candidate_attempt_source<'a>(
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::StandardDecision,
);
Ok(
if spec_metadata.api_format == "gemini:generate_content"
&& gemini_request_is_image_generation(body_json)
{
let (attempts, candidate_count) =
materialize_local_standard_candidate_attempts(state, trace_id, input, body_json, spec)
.await?;
let source =
LocalExecutionCandidateAttemptSource::from_static_attempts_for_image_bridge(attempts);
return Ok((source, candidate_count));
}
let (source, candidate_count) =
build_lazy_requested_model_execution_candidate_attempt_source_with_serving(
planner_state,
trace_id,
@@ -253,6 +283,51 @@ pub(super) async fn build_local_standard_candidate_attempt_source<'a>(
skipped_candidate
},
)
.await,
)
.await;
Ok((source, candidate_count))
}
#[allow(clippy::too_many_arguments)]
async fn maybe_append_gemini_image_openai_image_preselection(
state: &AppState,
planner_state: PlannerAppState<'_>,
trace_id: &str,
input: &LocalStandardDecisionInput,
body_json: &serde_json::Value,
spec: LocalStandardSpec,
mut candidates: Vec<aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate>,
mut skipped_candidates: Vec<
crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate,
>,
) -> Result<
(
Vec<aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate>,
Vec<crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate>,
),
GatewayError,
> {
let spec_metadata = local_standard_spec_metadata(spec);
if spec_metadata.api_format != "gemini:generate_content"
|| !gemini_request_is_image_generation(body_json)
{
return Ok((candidates, skipped_candidates));
}
let image_preselection = preselect_local_execution_candidates_for_api_formats_with_serving(
planner_state,
spec_metadata.api_format,
&input.requested_model,
spec_metadata.require_streaming,
input.required_capabilities.as_ref(),
&input.auth_snapshot,
input.client_session_affinity.as_ref(),
false,
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
vec!["openai:image".to_string()],
)
.await?;
candidates.extend(image_preselection.candidates);
skipped_candidates.extend(image_preselection.skipped_candidates);
let _ = (state, trace_id);
Ok((candidates, skipped_candidates))
}

View File

@@ -18,10 +18,13 @@ use crate::ai_serving::transport::kiro::{
KIRO_ENVELOPE_NAME,
};
use crate::ai_serving::transport::{
build_kiro_cross_format_upstream_url, build_standard_provider_request_headers,
StandardProviderRequestHeadersInput,
build_kiro_cross_format_upstream_url, build_openai_image_headers,
build_openai_image_upstream_url, build_standard_provider_request_headers,
openai_image_transport_unsupported_reason, resolve_openai_image_auth,
ProviderOpenAiImageHeadersInput, StandardProviderRequestHeadersInput,
};
use crate::ai_serving::{
build_openai_image_request_body_from_gemini_image_request, gemini_request_is_image_generation,
CandidateFailureDiagnostic, GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
};
use crate::AppState;
@@ -59,6 +62,15 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
let candidate = &attempt.eligible.candidate;
let transport = &attempt.eligible.transport;
let provider_api_format = attempt.eligible.provider_api_format.as_str();
if spec_metadata.api_format == "gemini:generate_content"
&& provider_api_format == "openai:image"
&& gemini_request_is_image_generation(body_json)
{
return resolve_local_gemini_image_to_openai_image_candidate_payload_parts(
state, parts, trace_id, body_json, input, attempt,
)
.await;
}
let is_kiro_claude_cli = is_kiro_claude_messages_transport(transport, provider_api_format);
let Some(conversion_kind) =
crate::ai_serving::request_conversion_kind(spec_metadata.api_format, provider_api_format)
@@ -331,6 +343,141 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
})
}
async fn resolve_local_gemini_image_to_openai_image_candidate_payload_parts(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
body_json: &serde_json::Value,
input: &LocalStandardDecisionInput,
attempt: &LocalStandardCandidateAttempt,
) -> Option<LocalStandardCandidatePayloadParts> {
let client_api_format = "gemini:generate_content";
let provider_api_format = "openai:image";
let planner_state = crate::ai_serving::PlannerAppState::new(state);
let candidate = &attempt.eligible.candidate;
let transport = &attempt.eligible.transport;
if let Some(skip_reason) =
openai_image_transport_unsupported_reason(transport, provider_api_format)
{
mark_skipped_local_standard_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
skip_reason,
)
.await;
return None;
}
let prepared_candidate = match prepare_header_authenticated_candidate(
planner_state,
transport,
candidate,
resolve_openai_image_auth(transport),
OauthPreparationContext {
trace_id,
api_format: provider_api_format,
operation: "gemini_image_to_openai_image_candidate_request",
},
)
.await
{
Ok(prepared) => prepared,
Err(skip_reason) => {
mark_skipped_local_standard_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
skip_reason,
)
.await;
return None;
}
};
let Some(converted) = build_openai_image_request_body_from_gemini_image_request(
body_json,
parts.uri.path(),
&prepared_candidate.mapped_model,
) else {
mark_skipped_local_standard_candidate_with_extra_data(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"provider_request_body_build_failed",
request_body_build_failure_extra_data(
body_json,
client_api_format,
provider_api_format,
),
)
.await;
return None;
};
let upstream_is_stream = true;
let upstream_url = build_openai_image_upstream_url(transport, None);
let Some(mut provider_request_headers) =
build_openai_image_headers(ProviderOpenAiImageHeadersInput {
headers: &parts.headers,
auth_header: &prepared_candidate.auth_header,
auth_value: &prepared_candidate.auth_value,
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &converted.body_json,
original_request_body: body_json,
})
else {
mark_skipped_local_standard_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_header_rules_apply_failed",
CandidateFailureDiagnostic::header_rules_apply_failed(
client_api_format,
provider_api_format,
"gemini_image_to_openai_image_headers",
),
)
.await;
return None;
};
apply_codex_openai_responses_special_headers(
&mut provider_request_headers,
&converted.body_json,
&parts.headers,
transport.provider.provider_type.as_str(),
provider_api_format,
Some(trace_id),
transport.key.decrypted_auth_config.as_deref(),
);
Some(LocalStandardCandidatePayloadParts {
auth_header: prepared_candidate.auth_header,
auth_value: prepared_candidate.auth_value,
mapped_model: converted.mapped_model,
provider_api_format: provider_api_format.to_string(),
provider_request_body: converted.body_json,
provider_request_headers,
upstream_url,
upstream_is_stream,
envelope_name: None,
transport: Arc::clone(transport),
})
}
#[allow(clippy::too_many_arguments)]
async fn build_kiro_cross_format_payload_parts(
state: &AppState,