refactor(gateway): 重构 ai pipeline 规划链路

This commit is contained in:
fawney19
2026-04-14 01:27:04 +08:00
parent 37bb120d20
commit 53acfbabf6
74 changed files with 7357 additions and 4380 deletions

View File

@@ -67,6 +67,20 @@ pub(crate) fn is_json_request(headers: &http::HeaderMap) -> bool {
crate::headers::is_json_request(headers)
}
pub(crate) fn extract_gemini_model_from_path(path: &str) -> Option<String> {
let (_, suffix) = path.split_once("/models/")?;
let model = suffix
.split_once(':')
.map(|(value, _)| value)
.unwrap_or(suffix);
let model = model.trim();
if model.is_empty() {
None
} else {
Some(model.to_string())
}
}
pub(crate) fn build_execution_runtime_auth_context(
auth_context: &crate::control::GatewayControlAuthContext,
) -> ExecutionRuntimeAuthContext {
@@ -107,7 +121,7 @@ pub(crate) fn maybe_build_local_sync_finalize_response(
#[cfg(test)]
mod tests {
use super::build_report_context_original_request_echo;
use super::{build_report_context_original_request_echo, extract_gemini_model_from_path};
use serde_json::json;
#[test]
@@ -126,4 +140,12 @@ mod tests {
assert_eq!(echo, body);
}
#[test]
fn extract_gemini_model_from_path_trims_method_suffix() {
let model =
extract_gemini_model_from_path("/v1beta/models/gemini-2.5-pro:streamGenerateContent");
assert_eq!(model.as_deref(), Some("gemini-2.5-pro"));
}
}

View File

@@ -12,6 +12,10 @@ use aether_scheduler_core::{
requested_capability_priority_for_candidate, SchedulerAffinityTarget,
};
use super::candidate_eligibility::{
read_candidate_transport_snapshot, EligibleLocalExecutionCandidate,
};
const PLANNER_SCHEDULER_AFFINITY_MAX_ENTRIES: usize = 10_000;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
@@ -43,7 +47,8 @@ pub(crate) async fn prefer_local_tunnel_owner_candidates(
.collect()
}
pub(crate) async fn rank_local_execution_candidates(
#[cfg(test)]
async fn rank_local_execution_candidates(
state: PlannerAppState<'_>,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
client_api_format: &str,
@@ -96,6 +101,62 @@ pub(crate) async fn rank_local_execution_candidates(
.collect()
}
pub(crate) async fn rank_eligible_local_execution_candidates(
state: PlannerAppState<'_>,
candidates: Vec<EligibleLocalExecutionCandidate>,
client_api_format: &str,
required_capabilities: Option<&serde_json::Value>,
) -> Vec<EligibleLocalExecutionCandidate> {
let normalized_client_api_format = client_api_format.trim().to_ascii_lowercase();
let ordering_config = read_scheduler_ordering_config_or_default(state).await;
let mut ranked = Vec::with_capacity(candidates.len());
for (original_index, eligible) in candidates.into_iter().enumerate() {
let ordering = resolve_candidate_execution_ordering_from_transport(
state,
&eligible.transport,
ordering_config,
)
.await;
let is_same_format = eligible
.provider_api_format
.eq_ignore_ascii_case(normalized_client_api_format.as_str());
let demote_cross_format = !is_same_format && !ordering.keep_priority_on_conversion;
let capability_priority =
requested_capability_priority_for_candidate(required_capabilities, &eligible.candidate);
ranked.push((
capability_priority.0,
capability_priority.1,
ordering.tunnel_bucket,
demote_cross_format,
original_index,
eligible,
));
}
ranked.sort_by(|left, right| {
left.0
.cmp(&right.0)
.then(left.1.cmp(&right.1))
.then(left.2.cmp(&right.2))
.then(left.3.cmp(&right.3))
.then_with(|| {
compare_candidates_by_priority_mode(
&left.5.candidate,
&right.5.candidate,
ordering_config.priority_mode,
None,
)
})
.then(left.4.cmp(&right.4))
});
ranked
.into_iter()
.map(|(_, _, _, _, _, eligible)| eligible)
.collect()
}
pub(crate) fn remember_scheduler_affinity_for_candidate(
state: PlannerAppState<'_>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
@@ -152,39 +213,18 @@ async fn resolve_candidate_execution_ordering(
};
};
CandidateExecutionOrdering {
tunnel_bucket: resolve_tunnel_owner_affinity_from_transport(state, &transport).await,
keep_priority_on_conversion: ordering_config.keep_priority_on_conversion
|| transport.provider.keep_priority_on_conversion,
}
resolve_candidate_execution_ordering_from_transport(state, &transport, ordering_config).await
}
async fn read_candidate_transport_snapshot(
async fn resolve_candidate_execution_ordering_from_transport(
state: PlannerAppState<'_>,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> Option<GatewayProviderTransportSnapshot> {
match state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
{
Ok(Some(transport)) => Some(transport),
Ok(None) => None,
Err(error) => {
warn!(
event_name = "candidate_affinity_transport_load_failed",
log_type = "event",
provider_id = %candidate.provider_id,
endpoint_id = %candidate.endpoint_id,
key_id = %candidate.key_id,
error = ?error,
"failed to load provider transport while evaluating execution ordering"
);
None
}
transport: &GatewayProviderTransportSnapshot,
ordering_config: SchedulerOrderingConfig,
) -> CandidateExecutionOrdering {
CandidateExecutionOrdering {
tunnel_bucket: resolve_tunnel_owner_affinity_from_transport(state, transport).await,
keep_priority_on_conversion: ordering_config.keep_priority_on_conversion
|| transport.provider.keep_priority_on_conversion,
}
}
@@ -267,6 +307,7 @@ mod tests {
remember_scheduler_affinity_for_candidate, PlannerAppState,
SchedulerMinimalCandidateSelectionCandidate,
};
use crate::ai_pipeline::planner::candidate_eligibility::filter_and_rank_local_execution_candidates;
use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::data::GatewayDataState;
use crate::tunnel::TunnelAttachmentRecord;
@@ -373,6 +414,24 @@ mod tests {
provider_id: &str,
id: &str,
node_id: &str,
) -> StoredProviderCatalogKey {
sample_key_for_provider_with_options(
provider_id,
id,
node_id,
true,
Some(json!(["openai:chat"])),
None,
)
}
fn sample_key_for_provider_with_options(
provider_id: &str,
id: &str,
node_id: &str,
is_active: bool,
api_formats: Option<serde_json::Value>,
allowed_models: Option<serde_json::Value>,
) -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
id.to_string(),
@@ -380,16 +439,16 @@ mod tests {
id.to_string(),
"api_key".to_string(),
None,
true,
is_active,
)
.expect("key should build")
.with_transport_fields(
Some(json!(["openai:chat"])),
api_formats,
"plain-upstream-key".to_string(),
None,
None,
Some(json!({"openai:chat": 1})),
None,
allowed_models,
None,
Some(json!({
"enabled": true,
@@ -812,6 +871,227 @@ mod tests {
assert_eq!(ranked[1].endpoint_id, "endpoint-miss");
}
#[tokio::test]
async fn realtime_gate_skips_inactive_candidates_before_ranking() {
let provider_catalog = InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider_with_options("provider-disabled", false, 0),
sample_provider_with_options("provider-active", false, 10),
],
vec![
sample_endpoint_for_provider(
"provider-disabled",
"endpoint-disabled",
"openai:chat",
),
sample_endpoint_for_provider("provider-active", "endpoint-active", "openai:chat"),
],
vec![
sample_key_for_provider_with_options(
"provider-disabled",
"key-disabled",
"",
false,
Some(json!(["openai:chat"])),
None,
),
sample_key_for_provider_with_options(
"provider-active",
"key-active",
"",
true,
Some(json!(["openai:chat"])),
None,
),
],
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
);
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data_state);
let (ranked, skipped) = filter_and_rank_local_execution_candidates(
PlannerAppState::new(&state),
vec![
sample_priority_candidate(
"provider-disabled",
"endpoint-disabled",
"key-disabled",
"openai:chat",
Some(0),
0,
),
sample_priority_candidate(
"provider-active",
"endpoint-active",
"key-active",
"openai:chat",
Some(10),
10,
),
],
"openai:chat",
"gpt-4.1",
None,
)
.await;
assert_eq!(ranked.len(), 1);
assert_eq!(ranked[0].candidate.endpoint_id, "endpoint-active");
assert_eq!(skipped.len(), 1);
assert_eq!(skipped[0].candidate.endpoint_id, "endpoint-disabled");
assert_eq!(skipped[0].skip_reason, "key_inactive");
}
#[tokio::test]
async fn realtime_gate_skips_candidates_when_key_model_binding_is_disabled() {
let provider_catalog = InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider_with_options("provider-restricted", false, 0),
sample_provider_with_options("provider-open", false, 10),
],
vec![
sample_endpoint_for_provider(
"provider-restricted",
"endpoint-restricted",
"openai:chat",
),
sample_endpoint_for_provider("provider-open", "endpoint-open", "openai:chat"),
],
vec![
sample_key_for_provider_with_options(
"provider-restricted",
"key-restricted",
"",
true,
Some(json!(["openai:chat"])),
Some(json!(["gpt-4o"])),
),
sample_key_for_provider_with_options(
"provider-open",
"key-open",
"",
true,
Some(json!(["openai:chat"])),
None,
),
],
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
);
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data_state);
let (ranked, skipped) = filter_and_rank_local_execution_candidates(
PlannerAppState::new(&state),
vec![
sample_priority_candidate(
"provider-restricted",
"endpoint-restricted",
"key-restricted",
"openai:chat",
Some(0),
0,
),
sample_priority_candidate(
"provider-open",
"endpoint-open",
"key-open",
"openai:chat",
Some(10),
10,
),
],
"openai:chat",
"gpt-4.1",
None,
)
.await;
assert_eq!(ranked.len(), 1);
assert_eq!(ranked[0].candidate.endpoint_id, "endpoint-open");
assert_eq!(skipped.len(), 1);
assert_eq!(skipped[0].candidate.endpoint_id, "endpoint-restricted");
assert_eq!(skipped[0].skip_reason, "key_model_disabled");
}
#[tokio::test]
async fn realtime_gate_skips_cross_format_candidates_when_conversion_is_disabled() {
let provider_catalog = InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider_with_options("provider-cross", true, 0),
sample_provider_with_options("provider-same", false, 10),
],
vec![
sample_endpoint_for_provider("provider-cross", "endpoint-cross", "claude:chat"),
sample_endpoint_for_provider("provider-same", "endpoint-same", "openai:chat"),
],
vec![
sample_key_for_provider_with_options(
"provider-cross",
"key-cross",
"",
true,
Some(json!(["claude:chat"])),
None,
),
sample_key_for_provider_with_options(
"provider-same",
"key-same",
"",
true,
Some(json!(["openai:chat"])),
None,
),
],
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
);
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data_state);
let (ranked, skipped) = filter_and_rank_local_execution_candidates(
PlannerAppState::new(&state),
vec![
sample_priority_candidate(
"provider-cross",
"endpoint-cross",
"key-cross",
"claude:chat",
Some(0),
0,
),
sample_priority_candidate(
"provider-same",
"endpoint-same",
"key-same",
"openai:chat",
Some(10),
10,
),
],
"openai:chat",
"gpt-4.1",
None,
)
.await;
assert_eq!(ranked.len(), 1);
assert_eq!(ranked[0].candidate.endpoint_id, "endpoint-same");
assert_eq!(skipped.len(), 1);
assert_eq!(skipped[0].candidate.endpoint_id, "endpoint-cross");
assert_eq!(skipped[0].skip_reason, "format_conversion_disabled");
}
#[tokio::test]
async fn remembers_scheduler_affinity_for_candidate_using_requested_model_key() {
let state = AppState::new().expect("state should build");

View File

@@ -0,0 +1,279 @@
use tracing::warn;
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use crate::ai_pipeline::{GatewayProviderTransportSnapshot, PlannerAppState};
use super::candidate_affinity::rank_eligible_local_execution_candidates;
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct EligibleLocalExecutionCandidate {
pub(crate) candidate: SchedulerMinimalCandidateSelectionCandidate,
pub(crate) transport: GatewayProviderTransportSnapshot,
pub(crate) provider_api_format: String,
}
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct SkippedLocalExecutionCandidate {
pub(crate) candidate: SchedulerMinimalCandidateSelectionCandidate,
pub(crate) skip_reason: &'static str,
}
pub(crate) async fn filter_and_rank_local_execution_candidates(
state: PlannerAppState<'_>,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
client_api_format: &str,
requested_model: &str,
required_capabilities: Option<&serde_json::Value>,
) -> (
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
) {
filter_and_rank_local_execution_candidates_with_gate(
state,
candidates,
client_api_format,
required_capabilities,
|candidate, transport| {
current_local_execution_candidate_skip_reason_with_transport(
candidate,
transport,
client_api_format,
requested_model,
)
},
)
.await
}
pub(crate) async fn filter_and_rank_local_execution_candidates_without_transport_pair_gate(
state: PlannerAppState<'_>,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
client_api_format: &str,
requested_model: Option<&str>,
required_capabilities: Option<&serde_json::Value>,
) -> (
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
) {
filter_and_rank_local_execution_candidates_with_gate(
state,
candidates,
client_api_format,
required_capabilities,
|candidate, transport| {
current_local_execution_candidate_common_skip_reason_with_transport(
candidate,
transport,
requested_model,
)
},
)
.await
}
async fn filter_and_rank_local_execution_candidates_with_gate<F>(
state: PlannerAppState<'_>,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
client_api_format: &str,
required_capabilities: Option<&serde_json::Value>,
runtime_skip_reason: F,
) -> (
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
)
where
F: Fn(
&SchedulerMinimalCandidateSelectionCandidate,
&GatewayProviderTransportSnapshot,
) -> Option<&'static str>,
{
let mut selectable = Vec::with_capacity(candidates.len());
let mut skipped = Vec::new();
for candidate in candidates {
let Some(transport) = read_candidate_transport_snapshot(state, &candidate).await else {
skipped.push(SkippedLocalExecutionCandidate {
candidate,
skip_reason: "transport_snapshot_missing",
});
continue;
};
match runtime_skip_reason(&candidate, &transport) {
Some(skip_reason) => skipped.push(SkippedLocalExecutionCandidate {
candidate,
skip_reason,
}),
None => selectable.push(EligibleLocalExecutionCandidate {
provider_api_format: transport.endpoint.api_format.trim().to_ascii_lowercase(),
candidate,
transport,
}),
}
}
let ranked = rank_eligible_local_execution_candidates(
state,
selectable,
client_api_format,
required_capabilities,
)
.await;
(ranked, skipped)
}
fn current_local_execution_candidate_common_skip_reason_with_transport(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
transport: &GatewayProviderTransportSnapshot,
requested_model: Option<&str>,
) -> Option<&'static str> {
let requested_model = requested_model.unwrap_or_default();
if !transport.provider.is_active {
return Some("provider_inactive");
}
if !transport.endpoint.is_active {
return Some("endpoint_inactive");
}
if !transport.key.is_active {
return Some("key_inactive");
}
let candidate_api_format = candidate.endpoint_api_format.trim().to_ascii_lowercase();
let endpoint_api_format = transport.endpoint.api_format.trim().to_ascii_lowercase();
if endpoint_api_format != candidate_api_format {
return Some("endpoint_api_format_changed");
}
if !transport_key_supports_api_format(transport, endpoint_api_format.as_str()) {
return Some("key_api_format_disabled");
}
if !transport_key_allows_candidate_model(transport, requested_model, candidate) {
return Some("key_model_disabled");
}
None
}
fn current_local_execution_candidate_skip_reason_with_transport(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
transport: &GatewayProviderTransportSnapshot,
client_api_format: &str,
requested_model: &str,
) -> Option<&'static str> {
if let Some(skip_reason) = current_local_execution_candidate_common_skip_reason_with_transport(
candidate,
transport,
Some(requested_model),
) {
return Some(skip_reason);
}
let endpoint_api_format = transport.endpoint.api_format.trim().to_ascii_lowercase();
let client_api_format = client_api_format.trim().to_ascii_lowercase();
if client_api_format == endpoint_api_format {
return None;
}
if !crate::ai_pipeline::conversion::request_pair_allowed_for_transport(
transport,
client_api_format.as_str(),
endpoint_api_format.as_str(),
) {
let skip_reason = if crate::ai_pipeline::conversion::request_conversion_kind(
client_api_format.as_str(),
endpoint_api_format.as_str(),
)
.is_some()
&& crate::ai_pipeline::conversion::request_conversion_requires_enable_flag(
client_api_format.as_str(),
endpoint_api_format.as_str(),
)
&& !transport.provider.enable_format_conversion
{
"format_conversion_disabled"
} else {
"transport_unsupported"
};
return Some(skip_reason);
}
None
}
fn transport_key_supports_api_format(
transport: &GatewayProviderTransportSnapshot,
endpoint_api_format: &str,
) -> bool {
match transport.key.api_formats.as_deref() {
None => true,
Some(formats) => formats
.iter()
.any(|value| value.trim().eq_ignore_ascii_case(endpoint_api_format)),
}
}
fn transport_key_allows_candidate_model(
transport: &GatewayProviderTransportSnapshot,
requested_model: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> bool {
let Some(allowed_models) = transport.key.allowed_models.as_deref() else {
return true;
};
let allowed_models = allowed_models
.iter()
.map(|value| value.trim())
.filter(|value| !value.is_empty())
.collect::<Vec<_>>();
if allowed_models.is_empty() {
return false;
}
let requested_model = requested_model.trim();
let global_model_name = candidate.global_model_name.trim();
let selected_provider_model_name = candidate.selected_provider_model_name.trim();
let mapping_matched_model = candidate
.mapping_matched_model
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
allowed_models.iter().any(|allowed_model| {
*allowed_model == requested_model
|| *allowed_model == global_model_name
|| *allowed_model == selected_provider_model_name
|| mapping_matched_model.is_some_and(|value| value == *allowed_model)
})
}
pub(crate) async fn read_candidate_transport_snapshot(
state: PlannerAppState<'_>,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> Option<GatewayProviderTransportSnapshot> {
match state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
{
Ok(Some(transport)) => Some(transport),
Ok(None) => None,
Err(error) => {
warn!(
event_name = "candidate_eligibility_transport_load_failed",
log_type = "event",
provider_id = %candidate.provider_id,
endpoint_id = %candidate.endpoint_id,
key_id = %candidate.key_id,
error = ?error,
"failed to load provider transport while evaluating local candidate eligibility"
);
None
}
}
}

View File

@@ -0,0 +1,237 @@
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use serde_json::Value;
use uuid::Uuid;
use crate::ai_pipeline::planner::candidate_affinity::remember_scheduler_affinity_for_candidate;
use crate::ai_pipeline::planner::candidate_eligibility::{
EligibleLocalExecutionCandidate, SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::runtime_miss::record_local_runtime_candidate_skip_reason;
use crate::ai_pipeline::{GatewayAuthApiKeySnapshot, PlannerAppState};
use crate::clock::current_unix_ms;
use crate::AppState;
#[derive(Debug, Clone)]
pub(crate) struct LocalExecutionCandidateAttempt {
pub(crate) eligible: EligibleLocalExecutionCandidate,
pub(crate) candidate_index: u32,
pub(crate) candidate_id: String,
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct LocalAvailableCandidatePersistenceContext<'a> {
pub(crate) user_id: &'a str,
pub(crate) api_key_id: &'a str,
pub(crate) required_capabilities: Option<&'a Value>,
pub(crate) error_context: &'static str,
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct LocalSkippedCandidatePersistenceContext<'a> {
pub(crate) user_id: &'a str,
pub(crate) api_key_id: &'a str,
pub(crate) required_capabilities: Option<&'a Value>,
pub(crate) error_context: &'static str,
pub(crate) record_runtime_miss_diagnostic: bool,
}
pub(crate) fn remember_first_local_candidate_affinity(
state: PlannerAppState<'_>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_api_format: &str,
requested_model: Option<&str>,
candidates: &[EligibleLocalExecutionCandidate],
) {
let Some(first_candidate) = candidates.first() else {
return;
};
let affinity_requested_model = requested_model
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(first_candidate.candidate.global_model_name.as_str());
remember_scheduler_affinity_for_candidate(
state,
auth_snapshot,
client_api_format,
affinity_requested_model,
&first_candidate.candidate,
);
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn persist_available_local_execution_candidates<F>(
state: PlannerAppState<'_>,
trace_id: &str,
user_id: &str,
api_key_id: &str,
required_capabilities: Option<&Value>,
candidates: Vec<EligibleLocalExecutionCandidate>,
error_context: &'static str,
build_extra_data: F,
) -> Vec<LocalExecutionCandidateAttempt>
where
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value>,
{
let created_at_unix_ms = current_unix_ms();
let mut materialized = Vec::with_capacity(candidates.len());
for (candidate_index, eligible) in candidates.into_iter().enumerate() {
let generated_candidate_id = Uuid::new_v4().to_string();
let candidate_id = state
.persist_available_local_candidate(
trace_id,
user_id,
api_key_id,
&eligible.candidate,
candidate_index as u32,
&generated_candidate_id,
required_capabilities,
build_extra_data(&eligible),
created_at_unix_ms,
error_context,
)
.await;
materialized.push(LocalExecutionCandidateAttempt {
eligible,
candidate_index: candidate_index as u32,
candidate_id,
});
}
materialized
}
pub(crate) async fn persist_available_local_execution_candidates_with_context<F>(
state: PlannerAppState<'_>,
trace_id: &str,
context: LocalAvailableCandidatePersistenceContext<'_>,
candidates: Vec<EligibleLocalExecutionCandidate>,
build_extra_data: F,
) -> Vec<LocalExecutionCandidateAttempt>
where
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value>,
{
persist_available_local_execution_candidates(
state,
trace_id,
context.user_id,
context.api_key_id,
context.required_capabilities,
candidates,
context.error_context,
build_extra_data,
)
.await
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn persist_skipped_local_execution_candidate(
state: &AppState,
trace_id: &str,
user_id: &str,
api_key_id: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
candidate_index: u32,
candidate_id: &str,
required_capabilities: Option<&Value>,
skip_reason: &'static str,
error_context: &'static str,
record_runtime_miss_diagnostic: bool,
) {
if record_runtime_miss_diagnostic {
record_local_runtime_candidate_skip_reason(state, trace_id, skip_reason);
}
PlannerAppState::new(state)
.persist_skipped_local_candidate(
trace_id,
user_id,
api_key_id,
candidate,
candidate_index,
candidate_id,
required_capabilities,
skip_reason,
current_unix_ms(),
error_context,
)
.await;
}
pub(crate) async fn mark_skipped_local_execution_candidate(
state: &AppState,
trace_id: &str,
context: LocalSkippedCandidatePersistenceContext<'_>,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
candidate_index: u32,
candidate_id: &str,
skip_reason: &'static str,
) {
persist_skipped_local_execution_candidate(
state,
trace_id,
context.user_id,
context.api_key_id,
candidate,
candidate_index,
candidate_id,
context.required_capabilities,
skip_reason,
context.error_context,
context.record_runtime_miss_diagnostic,
)
.await;
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn persist_skipped_local_execution_candidates(
state: &AppState,
trace_id: &str,
user_id: &str,
api_key_id: &str,
required_capabilities: Option<&Value>,
starting_candidate_index: u32,
skipped_candidates: Vec<SkippedLocalExecutionCandidate>,
error_context: &'static str,
record_runtime_miss_diagnostic: bool,
) {
for (skipped_offset, skipped_candidate) in skipped_candidates.into_iter().enumerate() {
let generated_candidate_id = Uuid::new_v4().to_string();
persist_skipped_local_execution_candidate(
state,
trace_id,
user_id,
api_key_id,
&skipped_candidate.candidate,
starting_candidate_index + skipped_offset as u32,
&generated_candidate_id,
required_capabilities,
skipped_candidate.skip_reason,
error_context,
record_runtime_miss_diagnostic,
)
.await;
}
}
pub(crate) async fn persist_skipped_local_execution_candidates_with_context(
state: &AppState,
trace_id: &str,
context: LocalSkippedCandidatePersistenceContext<'_>,
starting_candidate_index: u32,
skipped_candidates: Vec<SkippedLocalExecutionCandidate>,
) {
persist_skipped_local_execution_candidates(
state,
trace_id,
context.user_id,
context.api_key_id,
context.required_capabilities,
starting_candidate_index,
skipped_candidates,
context.error_context,
context.record_runtime_miss_diagnostic,
)
.await;
}

View File

@@ -0,0 +1,77 @@
use serde_json::{Map, Value};
use crate::ai_pipeline::planner::candidate_eligibility::EligibleLocalExecutionCandidate;
use crate::ai_pipeline::{ConversionMode, ExecutionStrategy};
use crate::append_execution_contract_fields_to_value;
pub(crate) struct LocalExecutionCandidateMetadataParts<'a> {
pub(crate) eligible: &'a EligibleLocalExecutionCandidate,
pub(crate) provider_api_format: &'a str,
pub(crate) client_api_format: &'a str,
pub(crate) extra_fields: Map<String, Value>,
}
pub(crate) fn build_local_execution_candidate_metadata(
parts: LocalExecutionCandidateMetadataParts<'_>,
) -> Value {
let candidate = &parts.eligible.candidate;
let mut object = Map::new();
object.insert(
"provider_api_format".to_string(),
Value::String(parts.provider_api_format.to_string()),
);
object.insert(
"client_api_format".to_string(),
Value::String(parts.client_api_format.to_string()),
);
object.insert(
"global_model_id".to_string(),
Value::String(candidate.global_model_id.clone()),
);
object.insert(
"global_model_name".to_string(),
Value::String(candidate.global_model_name.clone()),
);
object.insert(
"model_id".to_string(),
Value::String(candidate.model_id.clone()),
);
object.insert(
"selected_provider_model_name".to_string(),
Value::String(candidate.selected_provider_model_name.clone()),
);
object.insert(
"mapping_matched_model".to_string(),
candidate
.mapping_matched_model
.clone()
.map(Value::String)
.unwrap_or(Value::Null),
);
object.insert(
"provider_name".to_string(),
Value::String(candidate.provider_name.clone()),
);
object.insert(
"key_name".to_string(),
Value::String(candidate.key_name.clone()),
);
object.extend(parts.extra_fields);
Value::Object(object)
}
pub(crate) fn build_local_execution_candidate_contract_metadata(
parts: LocalExecutionCandidateMetadataParts<'_>,
execution_strategy: ExecutionStrategy,
conversion_mode: ConversionMode,
provider_contract: &str,
) -> Value {
let client_api_format = parts.client_api_format;
append_execution_contract_fields_to_value(
build_local_execution_candidate_metadata(parts),
execution_strategy,
conversion_mode,
client_api_format,
provider_contract,
)
}

View File

@@ -0,0 +1,84 @@
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use tracing::warn;
use crate::ai_pipeline::{
GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth, PlannerAppState,
};
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct PreparedHeaderAuthenticatedCandidate {
pub(crate) auth_header: String,
pub(crate) auth_value: String,
pub(crate) mapped_model: String,
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct OauthPreparationContext<'a> {
pub(crate) trace_id: &'a str,
pub(crate) api_format: &'a str,
pub(crate) operation: &'a str,
}
pub(crate) async fn prepare_header_authenticated_candidate(
state: PlannerAppState<'_>,
transport: &GatewayProviderTransportSnapshot,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
direct_auth: Option<(String, String)>,
context: OauthPreparationContext<'_>,
) -> Result<PreparedHeaderAuthenticatedCandidate, &'static str> {
let oauth_auth = if direct_auth.is_none() {
match resolve_candidate_oauth_auth(state, transport, context).await {
Some(LocalResolvedOAuthRequestAuth::Header { name, value }) => Some((name, value)),
Some(LocalResolvedOAuthRequestAuth::Kiro(_)) => None,
None => None,
}
} else {
None
};
let Some((auth_header, auth_value)) = direct_auth.or(oauth_auth) else {
return Err("transport_auth_unavailable");
};
let mapped_model = resolve_candidate_mapped_model(candidate)?;
Ok(PreparedHeaderAuthenticatedCandidate {
auth_header,
auth_value,
mapped_model,
})
}
pub(crate) fn resolve_candidate_mapped_model(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> Result<String, &'static str> {
let mapped_model = candidate.selected_provider_model_name.trim().to_string();
if mapped_model.is_empty() {
return Err("mapped_model_missing");
}
Ok(mapped_model)
}
pub(crate) async fn resolve_candidate_oauth_auth(
state: PlannerAppState<'_>,
transport: &GatewayProviderTransportSnapshot,
context: OauthPreparationContext<'_>,
) -> Option<LocalResolvedOAuthRequestAuth> {
match state.resolve_local_oauth_request_auth(transport).await {
Ok(Some(auth)) => Some(auth),
Ok(None) => None,
Err(err) => {
warn!(
event_name = "candidate_preparation_oauth_auth_resolution_failed",
log_type = "event",
trace_id = %context.trace_id,
api_format = %context.api_format,
operation = %context.operation,
provider_type = %transport.provider.provider_type,
error = ?err,
"failed to resolve oauth auth while preparing local candidate"
);
None
}
}
}

View File

@@ -0,0 +1,34 @@
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use crate::ai_pipeline::GatewayAuthApiKeySnapshot;
pub(crate) fn auth_snapshot_allows_cross_format_candidate(
auth_snapshot: &GatewayAuthApiKeySnapshot,
requested_model: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> bool {
if let Some(allowed_providers) = auth_snapshot.effective_allowed_providers() {
let provider_allowed = allowed_providers.iter().any(|value| {
value
.trim()
.eq_ignore_ascii_case(candidate.provider_id.trim())
|| value
.trim()
.eq_ignore_ascii_case(candidate.provider_name.trim())
});
if !provider_allowed {
return false;
}
}
if let Some(allowed_models) = auth_snapshot.effective_allowed_models() {
let model_allowed = allowed_models
.iter()
.any(|value| value == requested_model || value == &candidate.global_model_name);
if !model_allowed {
return false;
}
}
true
}

View File

@@ -14,10 +14,19 @@ pub(crate) use crate::ai_pipeline::contracts::{
OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND,
OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
};
use crate::ai_pipeline::GatewayControlDecision;
use crate::ai_pipeline::{
extract_gemini_model_from_path as extract_gemini_model_from_path_impl,
force_upstream_streaming_for_provider as force_upstream_streaming_for_provider_impl,
is_json_request, parse_direct_request_body as parse_direct_request_body_impl,
};
use crate::LocalExecutionRuntimeMissDiagnostic;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum RequestedModelFamily {
Standard,
Gemini,
}
pub(crate) fn parse_direct_request_body(
parts: &http::request::Parts,
@@ -33,9 +42,82 @@ pub(crate) fn force_upstream_streaming_for_provider(
force_upstream_streaming_for_provider_impl(provider_type, provider_api_format)
}
pub(crate) fn extract_standard_requested_model(body_json: &serde_json::Value) -> Option<String> {
body_json
.get("model")
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
pub(crate) fn extract_requested_model_from_request(
parts: &http::request::Parts,
body_json: &serde_json::Value,
family: RequestedModelFamily,
) -> Option<String> {
match family {
RequestedModelFamily::Standard => extract_standard_requested_model(body_json),
RequestedModelFamily::Gemini => extract_gemini_model_from_path_impl(parts.uri.path()),
}
}
pub(crate) fn build_local_runtime_miss_diagnostic(
decision: &GatewayControlDecision,
plan_kind: &str,
requested_model: Option<&str>,
reason: &str,
) -> LocalExecutionRuntimeMissDiagnostic {
LocalExecutionRuntimeMissDiagnostic {
reason: reason.to_string(),
route_family: decision.route_family.clone(),
route_kind: decision.route_kind.clone(),
public_path: Some(decision.public_path.clone()),
plan_kind: Some(plan_kind.to_string()),
requested_model: requested_model.map(ToOwned::to_owned),
candidate_count: None,
skipped_candidate_count: None,
skip_reasons: std::collections::BTreeMap::new(),
}
}
pub(crate) fn apply_local_candidate_evaluation_progress(
diagnostic: &mut LocalExecutionRuntimeMissDiagnostic,
candidate_count: usize,
) {
diagnostic.candidate_count = Some(candidate_count);
diagnostic.reason = if candidate_count == 0 {
"candidate_list_empty".to_string()
} else {
"candidate_evaluation_incomplete".to_string()
};
}
pub(crate) fn apply_local_candidate_terminal_plan_reason(
diagnostic: &mut LocalExecutionRuntimeMissDiagnostic,
no_plan_reason: &'static str,
) {
let candidate_count = diagnostic.candidate_count.unwrap_or(0);
let skipped_candidate_count = diagnostic.skipped_candidate_count.unwrap_or(0);
diagnostic.reason = if candidate_count == 0 {
"candidate_list_empty".to_string()
} else if skipped_candidate_count >= candidate_count {
"all_candidates_skipped".to_string()
} else {
no_plan_reason.to_string()
};
}
#[cfg(test)]
mod tests {
use super::force_upstream_streaming_for_provider;
use super::{
apply_local_candidate_evaluation_progress, apply_local_candidate_terminal_plan_reason,
build_local_runtime_miss_diagnostic, extract_requested_model_from_request,
extract_standard_requested_model, force_upstream_streaming_for_provider,
RequestedModelFamily,
};
use axum::http::Request;
use serde_json::json;
#[test]
fn forces_streaming_for_codex_openai_cli() {
@@ -53,4 +135,95 @@ mod tests {
"openai:cli"
));
}
#[test]
fn extracts_standard_requested_model_from_request_body() {
let requested_model =
extract_standard_requested_model(&json!({ "model": " claude-sonnet-4 " }));
assert_eq!(requested_model.as_deref(), Some("claude-sonnet-4"));
}
#[test]
fn request_family_helper_delegates_standard_model_extraction() {
let request = Request::builder()
.uri("https://example.test/v1/chat/completions")
.body(())
.expect("request should build");
let (parts, _) = request.into_parts();
let requested_model = extract_requested_model_from_request(
&parts,
&json!({ "model": " claude-sonnet-4 " }),
RequestedModelFamily::Standard,
);
assert_eq!(requested_model.as_deref(), Some("claude-sonnet-4"));
}
#[test]
fn extracts_gemini_requested_model_from_request_path() {
let request = Request::builder()
.uri("https://example.test/v1beta/models/gemini-2.5-pro:streamGenerateContent?alt=sse")
.body(())
.expect("request should build");
let (parts, _) = request.into_parts();
let requested_model =
extract_requested_model_from_request(&parts, &json!({}), RequestedModelFamily::Gemini);
assert_eq!(requested_model.as_deref(), Some("gemini-2.5-pro"));
}
#[test]
fn candidate_evaluation_progress_sets_candidate_count_and_reason() {
let mut diagnostic = build_local_runtime_miss_diagnostic(
&crate::ai_pipeline::GatewayControlDecision::synthetic(
"/v1/test",
Some("passthrough".to_string()),
Some("ai".to_string()),
Some("chat".to_string()),
Some("test#sync".to_string()),
),
"test_plan",
Some("test-model"),
"seed",
);
apply_local_candidate_evaluation_progress(&mut diagnostic, 0);
assert_eq!(diagnostic.candidate_count, Some(0));
assert_eq!(diagnostic.reason, "candidate_list_empty");
apply_local_candidate_evaluation_progress(&mut diagnostic, 3);
assert_eq!(diagnostic.candidate_count, Some(3));
assert_eq!(diagnostic.reason, "candidate_evaluation_incomplete");
}
#[test]
fn candidate_terminal_reason_prefers_empty_then_skipped_then_fallback() {
let mut diagnostic = build_local_runtime_miss_diagnostic(
&crate::ai_pipeline::GatewayControlDecision::synthetic(
"/v1/test",
Some("passthrough".to_string()),
Some("ai".to_string()),
Some("chat".to_string()),
Some("test#sync".to_string()),
),
"test_plan",
Some("test-model"),
"seed",
);
apply_local_candidate_terminal_plan_reason(&mut diagnostic, "no_local_sync_plans");
assert_eq!(diagnostic.reason, "candidate_list_empty");
diagnostic.candidate_count = Some(2);
diagnostic.skipped_candidate_count = Some(2);
apply_local_candidate_terminal_plan_reason(&mut diagnostic, "no_local_sync_plans");
assert_eq!(diagnostic.reason, "all_candidates_skipped");
diagnostic.skipped_candidate_count = Some(1);
apply_local_candidate_terminal_plan_reason(&mut diagnostic, "no_local_sync_plans");
assert_eq!(diagnostic.reason, "no_local_sync_plans");
}
}

View File

@@ -0,0 +1,83 @@
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::{GatewayAuthApiKeySnapshot, PlannerAppState};
use crate::clock::current_unix_secs;
use crate::{AppState, GatewayError};
#[derive(Debug, Clone)]
pub(crate) struct ResolvedLocalDecisionAuthInput {
pub(crate) auth_context: ExecutionRuntimeAuthContext,
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
pub(crate) required_capabilities: Option<serde_json::Value>,
}
#[derive(Debug, Clone)]
pub(crate) struct LocalRequestedModelDecisionInput {
pub(crate) auth_context: ExecutionRuntimeAuthContext,
pub(crate) requested_model: String,
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
pub(crate) required_capabilities: Option<serde_json::Value>,
}
#[derive(Debug, Clone)]
pub(crate) struct LocalAuthenticatedDecisionInput {
pub(crate) auth_context: ExecutionRuntimeAuthContext,
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
pub(crate) required_capabilities: Option<serde_json::Value>,
}
pub(crate) fn build_local_requested_model_decision_input(
resolved_input: ResolvedLocalDecisionAuthInput,
requested_model: String,
) -> LocalRequestedModelDecisionInput {
LocalRequestedModelDecisionInput {
auth_context: resolved_input.auth_context,
requested_model,
auth_snapshot: resolved_input.auth_snapshot,
required_capabilities: resolved_input.required_capabilities,
}
}
pub(crate) fn build_local_authenticated_decision_input(
resolved_input: ResolvedLocalDecisionAuthInput,
) -> LocalAuthenticatedDecisionInput {
LocalAuthenticatedDecisionInput {
auth_context: resolved_input.auth_context,
auth_snapshot: resolved_input.auth_snapshot,
required_capabilities: resolved_input.required_capabilities,
}
}
pub(crate) async fn resolve_local_authenticated_decision_input(
state: &AppState,
auth_context: ExecutionRuntimeAuthContext,
requested_model: Option<&str>,
explicit_required_capabilities: Option<&serde_json::Value>,
) -> Result<Option<ResolvedLocalDecisionAuthInput>, GatewayError> {
let planner_state = PlannerAppState::new(state);
let auth_snapshot = match planner_state
.read_auth_api_key_snapshot(
&auth_context.user_id,
&auth_context.api_key_id,
current_unix_secs(),
)
.await?
{
Some(snapshot) => snapshot,
None => return Ok(None),
};
let required_capabilities = planner_state
.resolve_request_candidate_required_capabilities(
&auth_context.user_id,
&auth_context.api_key_id,
requested_model,
explicit_required_capabilities,
)
.await;
Ok(Some(ResolvedLocalDecisionAuthInput {
auth_context,
auth_snapshot,
required_capabilities,
}))
}

View File

@@ -0,0 +1,79 @@
use serde_json::Value;
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::planner::candidate_materialization::{
LocalAvailableCandidatePersistenceContext, LocalSkippedCandidatePersistenceContext,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[allow(clippy::enum_variant_names)]
pub(crate) enum LocalCandidatePersistencePolicyKind {
StandardDecision,
SameFormatProviderDecision,
OpenAiChatDecision,
OpenAiCliDecision,
GeminiFilesDecision,
VideoDecision,
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct LocalCandidatePersistencePolicy<'a> {
pub(crate) available: LocalAvailableCandidatePersistenceContext<'a>,
pub(crate) skipped: LocalSkippedCandidatePersistenceContext<'a>,
}
pub(crate) fn build_local_candidate_persistence_policy<'a>(
auth_context: &'a ExecutionRuntimeAuthContext,
required_capabilities: Option<&'a Value>,
kind: LocalCandidatePersistencePolicyKind,
) -> LocalCandidatePersistencePolicy<'a> {
let (available_error_context, skipped_error_context, record_runtime_miss_diagnostic) =
match kind {
LocalCandidatePersistencePolicyKind::StandardDecision => (
"gateway local standard decision request candidate upsert failed",
"gateway local standard decision failed to persist skipped candidate",
true,
),
LocalCandidatePersistencePolicyKind::SameFormatProviderDecision => (
"gateway local same-format decision request candidate upsert failed",
"gateway local same-format decision failed to persist skipped candidate",
true,
),
LocalCandidatePersistencePolicyKind::OpenAiChatDecision => (
"gateway local openai chat decision request candidate upsert failed",
"gateway local openai chat decision failed to persist skipped candidate",
true,
),
LocalCandidatePersistencePolicyKind::OpenAiCliDecision => (
"gateway local openai cli decision request candidate upsert failed",
"gateway local openai cli decision failed to persist skipped candidate",
false,
),
LocalCandidatePersistencePolicyKind::GeminiFilesDecision => (
"gateway local gemini files request candidate upsert failed",
"gateway local gemini files failed to persist skipped candidate",
false,
),
LocalCandidatePersistencePolicyKind::VideoDecision => (
"gateway local video decision request candidate upsert failed",
"gateway local video decision failed to persist skipped candidate",
false,
),
};
LocalCandidatePersistencePolicy {
available: LocalAvailableCandidatePersistenceContext {
user_id: &auth_context.user_id,
api_key_id: &auth_context.api_key_id,
required_capabilities,
error_context: available_error_context,
},
skipped: LocalSkippedCandidatePersistenceContext {
user_id: &auth_context.user_id,
api_key_id: &auth_context.api_key_id,
required_capabilities,
error_context: skipped_error_context,
record_runtime_miss_diagnostic,
},
}
}

View File

@@ -5,11 +5,22 @@ use crate::ai_pipeline::GatewayControlDecision;
use crate::{AppState, GatewayError};
mod candidate_affinity;
mod candidate_eligibility;
mod candidate_materialization;
mod candidate_metadata;
mod candidate_preparation;
mod candidate_source;
mod common;
mod decision;
mod decision_input;
mod materialization_policy;
mod passthrough;
mod payload_metadata;
mod plan_builders;
mod report_context;
mod route;
mod runtime_miss;
mod spec_metadata;
mod specialized;
mod standard;
mod state;

View File

@@ -78,9 +78,5 @@ pub(crate) use self::family::{
pub(crate) use self::plans::{
build_local_stream_plan_and_reports, build_local_sync_plan_and_reports,
};
use self::request::{
build_same_format_provider_request_body, build_same_format_upstream_url,
extract_gemini_model_from_path,
};
const ANTIGRAVITY_ENVELOPE_NAME: &str = "antigravity:v1internal";

View File

@@ -1,7 +1,11 @@
use crate::ai_pipeline::GatewayControlDecision;
use crate::{
AppState, GatewayControlSyncDecisionResponse, GatewayError, LocalExecutionRuntimeMissDiagnostic,
use crate::ai_pipeline::planner::common::extract_requested_model_from_request;
use crate::ai_pipeline::planner::runtime_miss::{
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal,
apply_local_runtime_candidate_terminal_reason, set_local_runtime_miss_diagnostic_reason,
};
use crate::ai_pipeline::planner::spec_metadata::local_same_format_provider_spec_metadata;
use crate::ai_pipeline::GatewayControlDecision;
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
use super::super::plans::{resolve_stream_spec, resolve_sync_spec};
use super::candidates::{
@@ -10,52 +14,6 @@ use super::candidates::{
};
use super::payload::maybe_build_local_same_format_provider_decision_payload_for_candidate;
fn extract_requested_model(
parts: &http::request::Parts,
body_json: &serde_json::Value,
spec: crate::ai_pipeline::LocalSameFormatProviderSpec,
) -> Option<String> {
match spec.family {
crate::ai_pipeline::LocalSameFormatProviderFamily::Standard => body_json
.get("model")
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
crate::ai_pipeline::LocalSameFormatProviderFamily::Gemini => {
let marker = "/models/";
let start = parts.uri.path().find(marker)? + marker.len();
let tail = &parts.uri.path()[start..];
let end = tail.find(':').unwrap_or(tail.len());
let model = tail[..end].trim();
if model.is_empty() {
None
} else {
Some(model.to_string())
}
}
}
}
fn build_local_same_format_miss_diagnostic(
decision: &GatewayControlDecision,
spec: crate::ai_pipeline::LocalSameFormatProviderSpec,
requested_model: Option<&str>,
reason: &str,
) -> LocalExecutionRuntimeMissDiagnostic {
LocalExecutionRuntimeMissDiagnostic {
reason: reason.to_string(),
route_family: decision.route_family.clone(),
route_kind: decision.route_kind.clone(),
public_path: Some(decision.public_path.clone()),
plan_kind: Some(spec.decision_kind.to_string()),
requested_model: requested_model.map(ToOwned::to_owned),
candidate_count: None,
skipped_candidate_count: None,
skip_reasons: std::collections::BTreeMap::new(),
}
}
pub(crate) async fn maybe_build_sync_local_same_format_provider_decision_payload(
state: &AppState,
parts: &http::request::Parts,
@@ -67,48 +25,44 @@ pub(crate) async fn maybe_build_sync_local_same_format_provider_decision_payload
let Some(spec) = resolve_sync_spec(plan_kind) else {
return Ok(None);
};
let spec_metadata = local_same_format_provider_spec_metadata(spec);
let requested_model_family = spec_metadata
.requested_model_family
.expect("same-format provider spec metadata should include requested-model family");
let Some(input) = resolve_local_same_format_provider_decision_input(
state, parts, trace_id, decision, body_json, spec,
)
.await
else {
state.set_local_execution_runtime_miss_diagnostic(
set_local_runtime_miss_diagnostic_reason(
state,
trace_id,
build_local_same_format_miss_diagnostic(
decision,
spec,
extract_requested_model(parts, body_json, spec).as_deref(),
"decision_input_unavailable",
),
decision,
spec_metadata.decision_kind,
extract_requested_model_from_request(parts, body_json, requested_model_family)
.as_deref(),
"decision_input_unavailable",
);
return Ok(None);
};
state.set_local_execution_runtime_miss_diagnostic(
set_local_runtime_miss_diagnostic_reason(
state,
trace_id,
build_local_same_format_miss_diagnostic(
decision,
spec,
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
),
decision,
spec_metadata.decision_kind,
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) =
materialize_local_same_format_provider_candidate_attempts(state, trace_id, &input, spec)
.await?;
let preserve_existing_candidate_signal = candidate_count == 0
&& state.local_execution_runtime_miss_diagnostic_has_candidate_signal(trace_id);
if !preserve_existing_candidate_signal {
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
diagnostic.candidate_count = Some(candidate_count);
diagnostic.reason = if candidate_count == 0 {
"candidate_list_empty".to_string()
} else {
"candidate_evaluation_incomplete".to_string()
};
});
}
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
state,
trace_id,
candidate_count,
);
for attempt in attempts {
if let Some(payload) =
@@ -121,17 +75,7 @@ pub(crate) async fn maybe_build_sync_local_same_format_provider_decision_payload
}
}
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
let candidate_count = diagnostic.candidate_count.unwrap_or(0);
let skipped_candidate_count = diagnostic.skipped_candidate_count.unwrap_or(0);
diagnostic.reason = if candidate_count == 0 {
"candidate_list_empty".to_string()
} else if skipped_candidate_count >= candidate_count {
"all_candidates_skipped".to_string()
} else {
"no_local_sync_plans".to_string()
};
});
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_sync_plans");
Ok(None)
}
@@ -147,48 +91,44 @@ pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_paylo
let Some(spec) = resolve_stream_spec(plan_kind) else {
return Ok(None);
};
let spec_metadata = local_same_format_provider_spec_metadata(spec);
let requested_model_family = spec_metadata
.requested_model_family
.expect("same-format provider spec metadata should include requested-model family");
let Some(input) = resolve_local_same_format_provider_decision_input(
state, parts, trace_id, decision, body_json, spec,
)
.await
else {
state.set_local_execution_runtime_miss_diagnostic(
set_local_runtime_miss_diagnostic_reason(
state,
trace_id,
build_local_same_format_miss_diagnostic(
decision,
spec,
extract_requested_model(parts, body_json, spec).as_deref(),
"decision_input_unavailable",
),
decision,
spec_metadata.decision_kind,
extract_requested_model_from_request(parts, body_json, requested_model_family)
.as_deref(),
"decision_input_unavailable",
);
return Ok(None);
};
state.set_local_execution_runtime_miss_diagnostic(
set_local_runtime_miss_diagnostic_reason(
state,
trace_id,
build_local_same_format_miss_diagnostic(
decision,
spec,
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
),
decision,
spec_metadata.decision_kind,
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) =
materialize_local_same_format_provider_candidate_attempts(state, trace_id, &input, spec)
.await?;
let preserve_existing_candidate_signal = candidate_count == 0
&& state.local_execution_runtime_miss_diagnostic_has_candidate_signal(trace_id);
if !preserve_existing_candidate_signal {
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
diagnostic.candidate_count = Some(candidate_count);
diagnostic.reason = if candidate_count == 0 {
"candidate_list_empty".to_string()
} else {
"candidate_evaluation_incomplete".to_string()
};
});
}
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
state,
trace_id,
candidate_count,
);
for attempt in attempts {
if let Some(payload) =
@@ -201,17 +141,7 @@ pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_paylo
}
}
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
let candidate_count = diagnostic.candidate_count.unwrap_or(0);
let skipped_candidate_count = diagnostic.skipped_candidate_count.unwrap_or(0);
diagnostic.reason = if candidate_count == 0 {
"candidate_list_empty".to_string()
} else if skipped_candidate_count >= candidate_count {
"all_candidates_skipped".to_string()
} else {
"no_local_stream_plans".to_string()
};
});
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_stream_plans");
Ok(None)
}

View File

@@ -1,20 +1,32 @@
use serde_json::json;
use tracing::warn;
use uuid::Uuid;
use crate::ai_pipeline::planner::candidate_affinity::{
rank_local_execution_candidates, remember_scheduler_affinity_for_candidate,
use crate::ai_pipeline::planner::candidate_eligibility::filter_and_rank_local_execution_candidates;
use crate::ai_pipeline::planner::candidate_materialization::{
persist_available_local_execution_candidates_with_context,
persist_skipped_local_execution_candidates_with_context,
remember_first_local_candidate_affinity,
};
use crate::ai_pipeline::planner::candidate_metadata::{
build_local_execution_candidate_contract_metadata, LocalExecutionCandidateMetadataParts,
};
use crate::ai_pipeline::planner::common::extract_requested_model_from_request;
use crate::ai_pipeline::planner::decision_input::{
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
};
use crate::ai_pipeline::planner::materialization_policy::{
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
};
use crate::ai_pipeline::planner::spec_metadata::local_same_format_provider_spec_metadata;
use crate::ai_pipeline::{
resolve_local_decision_execution_runtime_auth_context, ConversionMode, ExecutionStrategy,
GatewayControlDecision, PlannerAppState,
};
use crate::clock::{current_unix_ms, current_unix_secs};
use crate::{append_execution_contract_fields_to_value, AppState, GatewayError};
use crate::clock::current_unix_secs;
use crate::{AppState, GatewayError};
use super::{
LocalSameFormatProviderCandidateAttempt, LocalSameFormatProviderDecisionInput,
LocalSameFormatProviderFamily, LocalSameFormatProviderSpec,
LocalSameFormatProviderSpec,
};
pub(crate) async fn resolve_local_same_format_provider_decision_input(
@@ -25,37 +37,33 @@ pub(crate) async fn resolve_local_same_format_provider_decision_input(
body_json: &serde_json::Value,
spec: LocalSameFormatProviderSpec,
) -> Option<LocalSameFormatProviderDecisionInput> {
let planner_state = PlannerAppState::new(state);
let spec_metadata = local_same_format_provider_spec_metadata(spec);
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
return None;
};
let requested_model = match spec.family {
LocalSameFormatProviderFamily::Standard => body_json
.get("model")
.and_then(|value| value.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)?,
LocalSameFormatProviderFamily::Gemini => {
super::super::request::extract_gemini_model_from_path(parts.uri.path())?
}
};
let requested_model = extract_requested_model_from_request(
parts,
body_json,
spec_metadata
.requested_model_family
.expect("same-format provider specs should declare requested-model family"),
)?;
let auth_snapshot = match planner_state
.read_auth_api_key_snapshot(
&auth_context.user_id,
&auth_context.api_key_id,
current_unix_secs(),
)
.await
let resolved_input = match resolve_local_authenticated_decision_input(
state,
auth_context,
Some(requested_model.as_str()),
None,
)
.await
{
Ok(Some(snapshot)) => snapshot,
Ok(Some(resolved_input)) => resolved_input,
Ok(None) => return None,
Err(err) => {
warn!(
trace_id = %trace_id,
api_format = spec.api_format,
api_format = spec_metadata.api_format,
error = ?err,
"gateway local same-format decision auth snapshot read failed"
);
@@ -63,21 +71,10 @@ pub(crate) async fn resolve_local_same_format_provider_decision_input(
}
};
let required_capabilities = planner_state
.resolve_request_candidate_required_capabilities(
&auth_context.user_id,
&auth_context.api_key_id,
Some(requested_model.as_str()),
None,
)
.await;
Some(LocalSameFormatProviderDecisionInput {
auth_context,
Some(build_local_requested_model_decision_input(
resolved_input,
requested_model,
auth_snapshot,
required_capabilities,
})
))
}
pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
@@ -86,80 +83,69 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
input: &LocalSameFormatProviderDecisionInput,
spec: LocalSameFormatProviderSpec,
) -> Result<(Vec<LocalSameFormatProviderCandidateAttempt>, usize), GatewayError> {
let spec_metadata = local_same_format_provider_spec_metadata(spec);
let planner_state = PlannerAppState::new(state);
let persistence_policy = build_local_candidate_persistence_policy(
&input.auth_context,
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::SameFormatProviderDecision,
);
let candidates = planner_state
.list_selectable_candidates(
spec.api_format,
spec_metadata.api_format,
&input.requested_model,
spec.require_streaming,
spec_metadata.require_streaming,
input.required_capabilities.as_ref(),
Some(&input.auth_snapshot),
current_unix_secs(),
)
.await?;
let candidates = rank_local_execution_candidates(
let (candidates, skipped_candidates) = filter_and_rank_local_execution_candidates(
planner_state,
candidates,
spec.api_format,
spec_metadata.api_format,
&input.requested_model,
input.required_capabilities.as_ref(),
)
.await;
let candidate_count = candidates.len();
let candidate_count = candidates.len() + skipped_candidates.len();
let created_at_unix_ms = current_unix_ms();
let mut attempts = Vec::with_capacity(candidates.len());
let mut affinity_remembered = false;
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
let generated_candidate_id = Uuid::new_v4().to_string();
if !affinity_remembered {
remember_scheduler_affinity_for_candidate(
planner_state,
Some(&input.auth_snapshot),
spec.api_format,
&input.requested_model,
&candidate,
);
affinity_remembered = true;
}
let extra_data = append_execution_contract_fields_to_value(
json!({
"provider_api_format": spec.api_format,
"client_api_format": spec.api_format,
"global_model_id": candidate.global_model_id.clone(),
"global_model_name": candidate.global_model_name.clone(),
"model_id": candidate.model_id.clone(),
"selected_provider_model_name": candidate.selected_provider_model_name.clone(),
"mapping_matched_model": candidate.mapping_matched_model.clone(),
"provider_name": candidate.provider_name.clone(),
"key_name": candidate.key_name.clone(),
}),
ExecutionStrategy::LocalSameFormat,
ConversionMode::None,
spec.api_format,
spec.api_format,
);
remember_first_local_candidate_affinity(
planner_state,
Some(&input.auth_snapshot),
spec_metadata.api_format,
Some(&input.requested_model),
&candidates,
);
let attempts = persist_available_local_execution_candidates_with_context(
planner_state,
trace_id,
persistence_policy.available,
candidates,
|eligible| {
Some(build_local_execution_candidate_contract_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: spec_metadata.api_format,
client_api_format: spec_metadata.api_format,
extra_fields: serde_json::Map::new(),
},
ExecutionStrategy::LocalSameFormat,
ConversionMode::None,
spec_metadata.api_format,
))
},
)
.await;
let candidate_id = planner_state
.persist_available_local_candidate(
trace_id,
&input.auth_context.user_id,
&input.auth_context.api_key_id,
&candidate,
candidate_index as u32,
&generated_candidate_id,
input.required_capabilities.as_ref(),
Some(extra_data),
created_at_unix_ms,
"gateway local same-format decision request candidate upsert failed",
)
.await;
attempts.push(LocalSameFormatProviderCandidateAttempt {
candidate,
candidate_index: candidate_index as u32,
candidate_id,
});
}
persist_skipped_local_execution_candidates_with_context(
state,
trace_id,
persistence_policy.skipped,
attempts.len() as u32,
skipped_candidates,
)
.await;
Ok((attempts, candidate_count))
}

View File

@@ -1,9 +1,7 @@
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::GatewayAuthApiKeySnapshot;
mod build;
mod candidates;
mod payload;
mod request;
pub(crate) use self::build::{
maybe_build_stream_local_same_format_provider_decision_payload,
@@ -14,19 +12,6 @@ pub(crate) use self::candidates::{
resolve_local_same_format_provider_decision_input,
};
pub(crate) use self::payload::maybe_build_local_same_format_provider_decision_payload_for_candidate;
pub(crate) use crate::ai_pipeline::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalSameFormatProviderCandidateAttempt;
pub(crate) use crate::ai_pipeline::planner::decision_input::LocalRequestedModelDecisionInput as LocalSameFormatProviderDecisionInput;
pub(crate) use crate::ai_pipeline::{LocalSameFormatProviderFamily, LocalSameFormatProviderSpec};
#[derive(Debug, Clone)]
pub(crate) struct LocalSameFormatProviderDecisionInput {
pub(crate) auth_context: ExecutionRuntimeAuthContext,
pub(crate) requested_model: String,
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
pub(crate) required_capabilities: Option<serde_json::Value>,
}
#[derive(Debug, Clone)]
pub(crate) struct LocalSameFormatProviderCandidateAttempt {
pub(crate) candidate: aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate,
pub(crate) candidate_index: u32,
pub(crate) candidate_id: String,
}

View File

@@ -1,40 +1,28 @@
use std::collections::BTreeMap;
use serde_json::json;
use crate::ai_pipeline::transport::antigravity::{
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
classify_local_antigravity_request_support, AntigravityEnvelopeRequestType,
AntigravityRequestEnvelopeSupport, AntigravityRequestSideSupport,
use crate::ai_pipeline::planner::candidate_materialization::mark_skipped_local_execution_candidate;
use crate::ai_pipeline::planner::materialization_policy::{
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
};
use crate::ai_pipeline::transport::auth::{
build_complete_passthrough_headers, build_complete_passthrough_headers_with_auth,
use crate::ai_pipeline::planner::payload_metadata::{
build_local_execution_decision_response, LocalExecutionDecisionResponseParts,
};
use crate::ai_pipeline::transport::claude_code::build_claude_code_passthrough_headers;
use crate::ai_pipeline::transport::kiro::{
build_kiro_provider_headers, KiroProviderHeadersInput, KIRO_ENVELOPE_NAME,
use crate::ai_pipeline::planner::report_context::{
build_local_execution_report_context, LocalExecutionReportContextParts,
};
use crate::ai_pipeline::planner::spec_metadata::local_same_format_provider_spec_metadata;
use crate::ai_pipeline::transport::{
apply_local_header_rules, ensure_upstream_auth_header, resolve_transport_execution_timeouts,
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_transport_tls_profile,
resolve_transport_execution_timeouts, resolve_transport_proxy_snapshot_with_tunnel_affinity,
resolve_transport_tls_profile,
};
use crate::ai_pipeline::{
collect_control_headers, ConversionMode, ExecutionStrategy, PlannerAppState,
};
use crate::clock::current_unix_ms;
use crate::ai_pipeline::{ConversionMode, ExecutionStrategy};
use crate::{
append_execution_contract_fields_to_value, append_local_failover_policy_to_value, AppState,
GatewayControlSyncDecisionResponse, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
GatewayControlSyncDecisionResponse,
};
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
#[path = "payload/prepare.rs"]
mod prepare;
use self::prepare::{
prepare_local_same_format_provider_candidate, PreparedSameFormatProviderCandidate,
};
use super::request::resolve_local_same_format_provider_candidate_payload_parts;
use super::{
LocalSameFormatProviderCandidateAttempt, LocalSameFormatProviderDecisionInput,
LocalSameFormatProviderSpec,
@@ -49,304 +37,109 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
attempt: LocalSameFormatProviderCandidateAttempt,
spec: LocalSameFormatProviderSpec,
) -> Option<GatewayControlSyncDecisionResponse> {
let planner_state = PlannerAppState::new(state);
let spec_metadata = local_same_format_provider_spec_metadata(spec);
let LocalSameFormatProviderCandidateAttempt {
candidate,
eligible,
candidate_index,
candidate_id,
} = attempt;
let PreparedSameFormatProviderCandidate {
transport,
is_antigravity,
is_claude_code,
is_vertex,
is_kiro,
kiro_auth,
auth_header,
auth_value,
mapped_model,
report_kind,
upstream_is_stream,
} = prepare_local_same_format_provider_candidate(
planner_state.app(),
trace_id,
input,
&candidate,
candidate_index,
&candidate_id,
spec,
} = &attempt;
let candidate = &eligible.candidate;
let resolved = resolve_local_same_format_provider_candidate_payload_parts(
state, parts, trace_id, body_json, input, &attempt, spec,
)
.await?;
let Some(base_provider_request_body) =
super::super::request::build_same_format_provider_request_body(
body_json,
&mapped_model,
spec,
transport.endpoint.body_rules.as_ref(),
upstream_is_stream,
kiro_auth.as_ref(),
is_claude_code,
)
else {
mark_skipped_local_same_format_provider_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"provider_request_body_missing",
)
.await;
return None;
};
let antigravity_auth = if is_antigravity {
match classify_local_antigravity_request_support(
&transport,
&base_provider_request_body,
AntigravityEnvelopeRequestType::Agent,
) {
AntigravityRequestSideSupport::Supported(spec) => Some(spec.auth),
AntigravityRequestSideSupport::Unsupported(_) => {
mark_skipped_local_same_format_provider_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"transport_unsupported",
)
.await;
return None;
}
}
} else {
None
};
let provider_request_body = if let Some(antigravity_auth) = antigravity_auth.as_ref() {
match build_antigravity_safe_v1internal_request(
antigravity_auth,
trace_id,
&mapped_model,
&base_provider_request_body,
AntigravityEnvelopeRequestType::Agent,
) {
AntigravityRequestEnvelopeSupport::Supported(envelope) => envelope,
AntigravityRequestEnvelopeSupport::Unsupported(_) => {
mark_skipped_local_same_format_provider_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"provider_request_body_missing",
)
.await;
return None;
}
}
} else {
base_provider_request_body
};
let Some(upstream_url) = super::super::request::build_same_format_upstream_url(
parts,
&transport,
&mapped_model,
spec,
upstream_is_stream,
kiro_auth.as_ref(),
) else {
mark_skipped_local_same_format_provider_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"upstream_url_missing",
)
.await;
return None;
};
let Some(provider_request_headers) = (if let Some(kiro_auth) = kiro_auth.as_ref() {
build_kiro_provider_headers(KiroProviderHeadersInput {
headers: &parts.headers,
provider_request_body: &provider_request_body,
original_request_body: body_json,
header_rules: transport.endpoint.header_rules.as_ref(),
auth_header: auth_header.as_deref().unwrap_or_default(),
auth_value: auth_value.as_deref().unwrap_or_default(),
auth_config: &kiro_auth.auth_config,
machine_id: kiro_auth.machine_id.as_str(),
})
} else {
let extra_headers = antigravity_auth
.as_ref()
.map(build_antigravity_static_identity_headers)
.unwrap_or_default();
let mut provider_request_headers = if is_claude_code {
build_claude_code_passthrough_headers(
&parts.headers,
auth_header.as_deref().unwrap_or_default(),
auth_value.as_deref().unwrap_or_default(),
&extra_headers,
upstream_is_stream,
transport.key.fingerprint.as_ref(),
)
} else if is_vertex {
build_complete_passthrough_headers(
&parts.headers,
&extra_headers,
Some("application/json"),
)
} else {
build_complete_passthrough_headers_with_auth(
&parts.headers,
auth_header.as_deref().unwrap_or_default(),
auth_value.as_deref().unwrap_or_default(),
&extra_headers,
Some("application/json"),
)
};
let protected_headers = auth_header
.as_deref()
.filter(|value| !value.trim().is_empty())
.map(|value| vec![value, "content-type"])
.unwrap_or_else(|| vec!["content-type"]);
if !apply_local_header_rules(
&mut provider_request_headers,
transport.endpoint.header_rules.as_ref(),
&protected_headers,
&provider_request_body,
Some(body_json),
) {
None
} else {
if let (Some(auth_header), Some(auth_value)) =
(auth_header.as_deref(), auth_value.as_deref())
{
ensure_upstream_auth_header(&mut provider_request_headers, auth_header, auth_value);
}
if upstream_is_stream {
provider_request_headers
.insert("accept".to_string(), "text/event-stream".to_string());
}
Some(provider_request_headers)
}
}) else {
mark_skipped_local_same_format_provider_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"transport_header_rules_apply_failed",
)
.await;
return None;
};
let prompt_cache_key = provider_request_body
let prompt_cache_key = resolved
.provider_request_body
.get("prompt_cache_key")
.and_then(|value| value.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let proxy =
resolve_transport_proxy_snapshot_with_tunnel_affinity(planner_state.app(), &transport)
.await;
let tls_profile = resolve_transport_tls_profile(&transport);
resolve_transport_proxy_snapshot_with_tunnel_affinity(state, &resolved.transport).await;
let tls_profile = resolve_transport_tls_profile(&resolved.transport);
let mut extra_fields = serde_json::Map::new();
if resolved.is_kiro {
extra_fields.insert(
"envelope_name".to_string(),
json!(crate::ai_pipeline::transport::kiro::KIRO_ENVELOPE_NAME),
);
} else if resolved.is_antigravity {
extra_fields.insert(
"envelope_name".to_string(),
json!(super::super::ANTIGRAVITY_ENVELOPE_NAME),
);
}
let report_context = append_local_failover_policy_to_value(
append_execution_contract_fields_to_value(
json!({
"user_id": input.auth_context.user_id,
"api_key_id": input.auth_context.api_key_id,
"username": input.auth_context.username,
"api_key_name": input.auth_context.api_key_name,
"request_id": trace_id,
"candidate_id": candidate_id,
"candidate_index": candidate_index,
"retry_index": 0,
"model": input.requested_model,
"provider_name": transport.provider.name,
"provider_id": candidate.provider_id,
"endpoint_id": candidate.endpoint_id,
"key_id": candidate.key_id,
"key_name": candidate.key_name,
"provider_api_format": spec.api_format,
"client_api_format": spec.api_format,
"mapped_model": mapped_model,
"upstream_url": upstream_url,
"provider_request_method": serde_json::Value::Null,
"provider_request_headers": provider_request_headers,
"original_headers": collect_control_headers(&parts.headers),
"original_request_body": crate::ai_pipeline::build_report_context_original_request_echo(body_json),
"has_envelope": is_kiro || is_antigravity,
"envelope_name": if is_kiro {
Some(KIRO_ENVELOPE_NAME)
} else if is_antigravity {
Some(super::super::ANTIGRAVITY_ENVELOPE_NAME)
} else {
None
},
"needs_conversion": false,
build_local_execution_report_context(LocalExecutionReportContextParts {
auth_context: &input.auth_context,
request_id: trace_id,
candidate_id,
candidate_index: *candidate_index,
retry_index: 0,
model: &input.requested_model,
provider_name: &resolved.transport.provider.name,
provider_id: &candidate.provider_id,
endpoint_id: &candidate.endpoint_id,
key_id: &candidate.key_id,
key_name: Some(&candidate.key_name),
provider_api_format: spec_metadata.api_format,
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),
upstream_url: Some(&resolved.upstream_url),
provider_request_method: Some(serde_json::Value::Null),
provider_request_headers: Some(&resolved.provider_request_headers),
original_headers: &parts.headers,
original_request_body: body_json,
has_envelope: resolved.is_kiro || resolved.is_antigravity,
needs_conversion: false,
extra_fields,
}),
ExecutionStrategy::LocalSameFormat,
ConversionMode::None,
spec.api_format,
spec.api_format,
spec_metadata.api_format,
spec_metadata.api_format,
),
&transport,
&resolved.transport,
);
Some(GatewayControlSyncDecisionResponse {
action: if spec.require_streaming {
EXECUTION_RUNTIME_STREAM_DECISION_ACTION.to_string()
} else {
EXECUTION_RUNTIME_SYNC_DECISION_ACTION.to_string()
Some(build_local_execution_decision_response(
LocalExecutionDecisionResponseParts {
decision_is_stream: spec_metadata.require_streaming,
decision_kind: spec_metadata.decision_kind.to_string(),
execution_strategy: ExecutionStrategy::LocalSameFormat,
conversion_mode: ConversionMode::None,
request_id: trace_id.to_string(),
candidate_id: candidate_id.to_string(),
provider_name: resolved.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: resolved.transport.endpoint.base_url.clone(),
upstream_url: resolved.upstream_url.clone(),
provider_request_method: None,
auth_header: resolved.auth_header.clone(),
auth_value: resolved.auth_value.clone(),
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: resolved.mapped_model.clone(),
prompt_cache_key,
provider_request_headers: resolved.provider_request_headers.clone(),
provider_request_body: Some(resolved.provider_request_body.clone()),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
proxy,
tls_profile,
timeouts: resolve_transport_execution_timeouts(&resolved.transport),
upstream_is_stream: resolved.upstream_is_stream,
report_kind: Some(resolved.report_kind.to_string()),
report_context: Some(report_context),
auth_context: input.auth_context.clone(),
},
decision_kind: Some(spec.decision_kind.to_string()),
execution_strategy: Some(ExecutionStrategy::LocalSameFormat.as_str().to_string()),
conversion_mode: Some(ConversionMode::None.as_str().to_string()),
request_id: Some(trace_id.to_string()),
candidate_id: Some(candidate_id.clone()),
provider_name: Some(transport.provider.name.clone()),
provider_id: Some(candidate.provider_id.clone()),
endpoint_id: Some(candidate.endpoint_id.clone()),
key_id: Some(candidate.key_id.clone()),
upstream_base_url: Some(transport.endpoint.base_url.clone()),
upstream_url: Some(upstream_url.clone()),
provider_request_method: None,
auth_header,
auth_value,
provider_api_format: Some(spec.api_format.to_string()),
client_api_format: Some(spec.api_format.to_string()),
provider_contract: Some(spec.api_format.to_string()),
client_contract: Some(spec.api_format.to_string()),
model_name: Some(input.requested_model.clone()),
mapped_model: Some(mapped_model.clone()),
prompt_cache_key,
extra_headers: BTreeMap::new(),
provider_request_headers: provider_request_headers.clone(),
provider_request_body: Some(provider_request_body.clone()),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
proxy,
tls_profile,
timeouts: resolve_transport_execution_timeouts(&transport),
upstream_is_stream,
report_kind: Some(report_kind.to_string()),
report_context: Some(report_context),
auth_context: Some(input.auth_context.clone()),
})
))
}
pub(super) async fn mark_skipped_local_same_format_provider_candidate(
@@ -358,25 +151,19 @@ pub(super) async fn mark_skipped_local_same_format_provider_candidate(
candidate_id: &str,
skip_reason: &'static str,
) {
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
*diagnostic
.skip_reasons
.entry(skip_reason.to_string())
.or_insert(0) += 1;
*diagnostic.skipped_candidate_count.get_or_insert(0) += 1;
});
PlannerAppState::new(state)
.persist_skipped_local_candidate(
trace_id,
&input.auth_context.user_id,
&input.auth_context.api_key_id,
candidate,
candidate_index,
candidate_id,
input.required_capabilities.as_ref(),
skip_reason,
current_unix_ms(),
"gateway local same-format decision failed to persist skipped candidate",
)
.await;
let persistence_policy = build_local_candidate_persistence_policy(
&input.auth_context,
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::SameFormatProviderDecision,
);
mark_skipped_local_execution_candidate(
state,
trace_id,
persistence_policy.skipped,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
}

View File

@@ -1,270 +0,0 @@
use tracing::warn;
use crate::ai_pipeline::transport::auth::{resolve_local_gemini_auth, resolve_local_standard_auth};
use crate::ai_pipeline::transport::claude_code::supports_local_claude_code_transport_with_network;
use crate::ai_pipeline::transport::kiro::{
supports_local_kiro_request_transport_with_network, KiroRequestAuth,
};
use crate::ai_pipeline::transport::policy::{
supports_local_gemini_transport_with_network, supports_local_standard_transport_with_network,
};
use crate::ai_pipeline::transport::vertex::{
resolve_local_vertex_api_key_query_auth,
supports_local_vertex_api_key_gemini_transport_with_network,
};
use crate::ai_pipeline::{
GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth, PlannerAppState,
};
use crate::AppState;
use super::super::{
LocalSameFormatProviderDecisionInput, LocalSameFormatProviderFamily,
LocalSameFormatProviderSpec,
};
pub(super) struct PreparedSameFormatProviderCandidate {
pub(super) transport: GatewayProviderTransportSnapshot,
pub(super) is_antigravity: bool,
pub(super) is_claude_code: bool,
pub(super) is_vertex: bool,
pub(super) is_kiro: bool,
pub(super) kiro_auth: Option<KiroRequestAuth>,
pub(super) auth_header: Option<String>,
pub(super) auth_value: Option<String>,
pub(super) mapped_model: String,
pub(super) report_kind: &'static str,
pub(super) upstream_is_stream: bool,
}
pub(super) async fn prepare_local_same_format_provider_candidate(
state: &AppState,
trace_id: &str,
input: &LocalSameFormatProviderDecisionInput,
candidate: &aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate,
candidate_index: u32,
candidate_id: &str,
spec: LocalSameFormatProviderSpec,
) -> Option<PreparedSameFormatProviderCandidate> {
let planner_state = PlannerAppState::new(state);
let transport = match planner_state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
{
Ok(Some(snapshot)) => snapshot,
Ok(None) => {
super::mark_skipped_local_same_format_provider_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_snapshot_missing",
)
.await;
return None;
}
Err(err) => {
warn!(
trace_id = %trace_id,
api_format = spec.api_format,
error = ?err,
"gateway local same-format decision provider transport read failed"
);
super::mark_skipped_local_same_format_provider_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_snapshot_read_failed",
)
.await;
return None;
}
};
let is_antigravity = transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("antigravity");
let is_claude_code = transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("claude_code");
let is_vertex = transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("vertex_ai");
let is_kiro = transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("kiro");
let transport_supported = if is_kiro {
supports_local_kiro_request_transport_with_network(&transport)
} else if is_antigravity {
true
} else if is_claude_code {
supports_local_claude_code_transport_with_network(&transport, spec.api_format)
} else if is_vertex {
supports_local_vertex_api_key_gemini_transport_with_network(&transport)
} else {
match spec.family {
LocalSameFormatProviderFamily::Standard => {
supports_local_standard_transport_with_network(&transport, spec.api_format)
}
LocalSameFormatProviderFamily::Gemini => {
supports_local_gemini_transport_with_network(&transport, spec.api_format)
}
}
};
if !transport_supported {
super::mark_skipped_local_same_format_provider_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_unsupported",
)
.await;
return None;
}
let vertex_query_auth = if is_vertex {
resolve_local_vertex_api_key_query_auth(&transport)
} else {
None
};
let should_try_oauth_auth = is_kiro
|| matches!(spec.family, LocalSameFormatProviderFamily::Standard)
&& resolve_local_standard_auth(&transport).is_none()
|| matches!(spec.family, LocalSameFormatProviderFamily::Gemini)
&& !is_vertex
&& resolve_local_gemini_auth(&transport).is_none();
let oauth_auth = if should_try_oauth_auth {
match planner_state
.resolve_local_oauth_request_auth(&transport)
.await
{
Ok(Some(LocalResolvedOAuthRequestAuth::Kiro(auth))) => {
Some(LocalResolvedOAuthRequestAuth::Kiro(auth))
}
Ok(Some(LocalResolvedOAuthRequestAuth::Header { name, value })) => {
Some(LocalResolvedOAuthRequestAuth::Header { name, value })
}
Ok(None) => None,
Err(err) => {
warn!(
trace_id = %trace_id,
api_format = spec.api_format,
provider_type = %transport.provider.provider_type,
error = ?err,
"gateway local same-format oauth auth resolution failed"
);
None
}
}
} else {
None
};
let kiro_auth = match oauth_auth.as_ref() {
Some(LocalResolvedOAuthRequestAuth::Kiro(auth)) => Some(auth.clone()),
_ => None,
};
let auth = if let Some(kiro_auth) = kiro_auth.as_ref() {
Some((kiro_auth.name.to_string(), kiro_auth.value.clone()))
} else if let Some(LocalResolvedOAuthRequestAuth::Header { name, value }) = oauth_auth.as_ref()
{
Some((name.clone(), value.clone()))
} else if is_vertex {
None
} else {
match spec.family {
LocalSameFormatProviderFamily::Standard => resolve_local_standard_auth(&transport),
LocalSameFormatProviderFamily::Gemini => resolve_local_gemini_auth(&transport),
}
};
let (auth_header, auth_value) = match auth {
Some((name, value)) => (Some(name), Some(value)),
None if is_vertex && vertex_query_auth.is_some() => (None, None),
None => {
super::mark_skipped_local_same_format_provider_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_auth_unavailable",
)
.await;
return None;
}
};
if is_vertex && vertex_query_auth.is_none() {
super::mark_skipped_local_same_format_provider_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_auth_unavailable",
)
.await;
return None;
}
let mapped_model = candidate.selected_provider_model_name.trim().to_string();
if mapped_model.is_empty() {
super::mark_skipped_local_same_format_provider_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"mapped_model_missing",
)
.await;
return None;
}
let upstream_is_stream = is_kiro || is_antigravity || spec.require_streaming;
let report_kind = if is_kiro && !spec.require_streaming {
"claude_cli_sync_finalize"
} else if is_antigravity && !spec.require_streaming {
match spec.api_format {
"gemini:chat" => "gemini_chat_sync_finalize",
"gemini:cli" => "gemini_cli_sync_finalize",
_ => spec.report_kind,
}
} else {
spec.report_kind
};
Some(PreparedSameFormatProviderCandidate {
transport,
is_antigravity,
is_claude_code,
is_vertex,
is_kiro,
kiro_auth,
auth_header,
auth_value,
mapped_model,
report_kind,
upstream_is_stream,
})
}

View File

@@ -0,0 +1,254 @@
use std::collections::BTreeMap;
use serde_json::Value;
use crate::ai_pipeline::transport::antigravity::{
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
classify_local_antigravity_request_support, AntigravityEnvelopeRequestType,
AntigravityRequestEnvelopeSupport, AntigravityRequestSideSupport,
};
use crate::ai_pipeline::transport::auth::{
build_complete_passthrough_headers, build_complete_passthrough_headers_with_auth,
};
use crate::ai_pipeline::transport::claude_code::build_claude_code_passthrough_headers;
use crate::ai_pipeline::transport::kiro::{build_kiro_provider_headers, KiroProviderHeadersInput};
use crate::ai_pipeline::transport::{apply_local_header_rules, ensure_upstream_auth_header};
use crate::ai_pipeline::GatewayProviderTransportSnapshot;
use crate::AppState;
mod policy;
mod prepare;
use self::prepare::prepare_local_same_format_provider_candidate;
use super::payload::mark_skipped_local_same_format_provider_candidate;
use super::{
LocalSameFormatProviderCandidateAttempt, LocalSameFormatProviderDecisionInput,
LocalSameFormatProviderSpec,
};
pub(crate) struct LocalSameFormatProviderCandidatePayloadParts {
pub(super) transport: GatewayProviderTransportSnapshot,
pub(super) is_antigravity: bool,
pub(super) is_kiro: bool,
pub(super) auth_header: Option<String>,
pub(super) auth_value: Option<String>,
pub(super) mapped_model: String,
pub(super) report_kind: &'static str,
pub(super) upstream_is_stream: bool,
pub(super) upstream_url: String,
pub(super) provider_request_headers: BTreeMap<String, String>,
pub(super) provider_request_body: Value,
}
pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
body_json: &serde_json::Value,
input: &LocalSameFormatProviderDecisionInput,
attempt: &LocalSameFormatProviderCandidateAttempt,
spec: LocalSameFormatProviderSpec,
) -> Option<LocalSameFormatProviderCandidatePayloadParts> {
let candidate = &attempt.eligible.candidate;
let prepared = prepare_local_same_format_provider_candidate(
state,
trace_id,
input,
&attempt.eligible,
attempt.candidate_index,
&attempt.candidate_id,
spec,
)
.await?;
let Some(base_provider_request_body) =
super::super::request::build_same_format_provider_request_body(
body_json,
&prepared.mapped_model,
spec,
prepared.transport.endpoint.body_rules.as_ref(),
prepared.upstream_is_stream,
prepared.kiro_auth.as_ref(),
prepared.is_claude_code,
)
else {
mark_skipped_local_same_format_provider_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"provider_request_body_missing",
)
.await;
return None;
};
let antigravity_auth = if prepared.is_antigravity {
match classify_local_antigravity_request_support(
&prepared.transport,
&base_provider_request_body,
AntigravityEnvelopeRequestType::Agent,
) {
AntigravityRequestSideSupport::Supported(spec) => Some(spec.auth),
AntigravityRequestSideSupport::Unsupported(_) => {
mark_skipped_local_same_format_provider_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_unsupported",
)
.await;
return None;
}
}
} else {
None
};
let provider_request_body = if let Some(antigravity_auth) = antigravity_auth.as_ref() {
match build_antigravity_safe_v1internal_request(
antigravity_auth,
trace_id,
&prepared.mapped_model,
&base_provider_request_body,
AntigravityEnvelopeRequestType::Agent,
) {
AntigravityRequestEnvelopeSupport::Supported(envelope) => envelope,
AntigravityRequestEnvelopeSupport::Unsupported(_) => {
mark_skipped_local_same_format_provider_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"provider_request_body_missing",
)
.await;
return None;
}
}
} else {
base_provider_request_body
};
let Some(upstream_url) = super::super::request::build_same_format_upstream_url(
parts,
&prepared.transport,
&prepared.mapped_model,
spec,
prepared.upstream_is_stream,
prepared.kiro_auth.as_ref(),
) else {
mark_skipped_local_same_format_provider_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"upstream_url_missing",
)
.await;
return None;
};
let Some(provider_request_headers) = (if let Some(kiro_auth) = prepared.kiro_auth.as_ref() {
build_kiro_provider_headers(KiroProviderHeadersInput {
headers: &parts.headers,
provider_request_body: &provider_request_body,
original_request_body: body_json,
header_rules: prepared.transport.endpoint.header_rules.as_ref(),
auth_header: prepared.auth_header.as_deref().unwrap_or_default(),
auth_value: prepared.auth_value.as_deref().unwrap_or_default(),
auth_config: &kiro_auth.auth_config,
machine_id: kiro_auth.machine_id.as_str(),
})
} else {
let extra_headers = antigravity_auth
.as_ref()
.map(build_antigravity_static_identity_headers)
.unwrap_or_default();
let mut provider_request_headers = if prepared.is_claude_code {
build_claude_code_passthrough_headers(
&parts.headers,
prepared.auth_header.as_deref().unwrap_or_default(),
prepared.auth_value.as_deref().unwrap_or_default(),
&extra_headers,
prepared.upstream_is_stream,
prepared.transport.key.fingerprint.as_ref(),
)
} else if prepared.is_vertex {
build_complete_passthrough_headers(
&parts.headers,
&extra_headers,
Some("application/json"),
)
} else {
build_complete_passthrough_headers_with_auth(
&parts.headers,
prepared.auth_header.as_deref().unwrap_or_default(),
prepared.auth_value.as_deref().unwrap_or_default(),
&extra_headers,
Some("application/json"),
)
};
let protected_headers = prepared
.auth_header
.as_deref()
.filter(|value| !value.trim().is_empty())
.map(|value| vec![value, "content-type"])
.unwrap_or_else(|| vec!["content-type"]);
if !apply_local_header_rules(
&mut provider_request_headers,
prepared.transport.endpoint.header_rules.as_ref(),
&protected_headers,
&provider_request_body,
Some(body_json),
) {
None
} else {
if let (Some(auth_header), Some(auth_value)) = (
prepared.auth_header.as_deref(),
prepared.auth_value.as_deref(),
) {
ensure_upstream_auth_header(&mut provider_request_headers, auth_header, auth_value);
}
if prepared.upstream_is_stream {
provider_request_headers
.insert("accept".to_string(), "text/event-stream".to_string());
}
Some(provider_request_headers)
}
}) else {
mark_skipped_local_same_format_provider_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_header_rules_apply_failed",
)
.await;
return None;
};
Some(LocalSameFormatProviderCandidatePayloadParts {
transport: prepared.transport,
is_antigravity: prepared.is_antigravity,
is_kiro: prepared.is_kiro,
auth_header: prepared.auth_header,
auth_value: prepared.auth_value,
mapped_model: prepared.mapped_model,
report_kind: prepared.report_kind,
upstream_is_stream: prepared.upstream_is_stream,
upstream_url,
provider_request_headers,
provider_request_body,
})
}

View File

@@ -0,0 +1,124 @@
use crate::ai_pipeline::planner::spec_metadata::LocalExecutionSurfaceSpecMetadata;
use crate::ai_pipeline::transport::auth::{resolve_local_gemini_auth, resolve_local_standard_auth};
use crate::ai_pipeline::transport::claude_code::supports_local_claude_code_transport_with_network;
use crate::ai_pipeline::transport::kiro::supports_local_kiro_request_transport_with_network;
use crate::ai_pipeline::transport::policy::{
supports_local_gemini_transport_with_network, supports_local_standard_transport_with_network,
};
use crate::ai_pipeline::transport::vertex::supports_local_vertex_api_key_gemini_transport_with_network;
use crate::ai_pipeline::GatewayProviderTransportSnapshot;
use super::super::LocalSameFormatProviderFamily;
pub(super) struct SameFormatProviderRequestBehavior {
pub(super) is_antigravity: bool,
pub(super) is_claude_code: bool,
pub(super) is_vertex: bool,
pub(super) is_kiro: bool,
pub(super) upstream_is_stream: bool,
pub(super) report_kind: &'static str,
}
pub(super) fn classify_same_format_provider_request_behavior(
transport: &GatewayProviderTransportSnapshot,
spec_metadata: LocalExecutionSurfaceSpecMetadata,
) -> SameFormatProviderRequestBehavior {
let is_antigravity = transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("antigravity");
let is_claude_code = transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("claude_code");
let is_vertex = transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("vertex_ai");
let is_kiro = transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("kiro");
let default_report_kind = spec_metadata
.report_kind
.expect("same-format provider specs should declare report kind");
let upstream_is_stream = is_kiro || is_antigravity || spec_metadata.require_streaming;
let report_kind = if is_kiro && !spec_metadata.require_streaming {
"claude_cli_sync_finalize"
} else if is_antigravity && !spec_metadata.require_streaming {
match spec_metadata.api_format {
"gemini:chat" => "gemini_chat_sync_finalize",
"gemini:cli" => "gemini_cli_sync_finalize",
_ => default_report_kind,
}
} else {
default_report_kind
};
SameFormatProviderRequestBehavior {
is_antigravity,
is_claude_code,
is_vertex,
is_kiro,
upstream_is_stream,
report_kind,
}
}
pub(super) fn same_format_provider_transport_supported(
behavior: &SameFormatProviderRequestBehavior,
transport: &GatewayProviderTransportSnapshot,
family: LocalSameFormatProviderFamily,
api_format: &str,
) -> bool {
if behavior.is_kiro {
supports_local_kiro_request_transport_with_network(transport)
} else if behavior.is_antigravity {
true
} else if behavior.is_claude_code {
supports_local_claude_code_transport_with_network(transport, api_format)
} else if behavior.is_vertex {
supports_local_vertex_api_key_gemini_transport_with_network(transport)
} else {
match family {
LocalSameFormatProviderFamily::Standard => {
supports_local_standard_transport_with_network(transport, api_format)
}
LocalSameFormatProviderFamily::Gemini => {
supports_local_gemini_transport_with_network(transport, api_format)
}
}
}
}
pub(super) fn should_try_same_format_provider_oauth_auth(
behavior: &SameFormatProviderRequestBehavior,
transport: &GatewayProviderTransportSnapshot,
family: LocalSameFormatProviderFamily,
) -> bool {
behavior.is_kiro
|| matches!(family, LocalSameFormatProviderFamily::Standard)
&& resolve_local_standard_auth(transport).is_none()
|| matches!(family, LocalSameFormatProviderFamily::Gemini)
&& !behavior.is_vertex
&& resolve_local_gemini_auth(transport).is_none()
}
pub(super) fn resolve_same_format_provider_direct_auth(
behavior: &SameFormatProviderRequestBehavior,
transport: &GatewayProviderTransportSnapshot,
family: LocalSameFormatProviderFamily,
) -> Option<(String, String)> {
if behavior.is_vertex {
None
} else {
match family {
LocalSameFormatProviderFamily::Standard => resolve_local_standard_auth(transport),
LocalSameFormatProviderFamily::Gemini => resolve_local_gemini_auth(transport),
}
}
}

View File

@@ -0,0 +1,162 @@
use crate::ai_pipeline::planner::candidate_eligibility::EligibleLocalExecutionCandidate;
use crate::ai_pipeline::planner::candidate_preparation::{
resolve_candidate_mapped_model, resolve_candidate_oauth_auth, OauthPreparationContext,
};
use crate::ai_pipeline::planner::spec_metadata::local_same_format_provider_spec_metadata;
use crate::ai_pipeline::transport::kiro::KiroRequestAuth;
use crate::ai_pipeline::transport::vertex::resolve_local_vertex_api_key_query_auth;
use crate::ai_pipeline::{
GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth, PlannerAppState,
};
use crate::AppState;
use super::super::LocalSameFormatProviderDecisionInput;
use super::super::LocalSameFormatProviderSpec;
use super::policy::{
classify_same_format_provider_request_behavior, resolve_same_format_provider_direct_auth,
same_format_provider_transport_supported, should_try_same_format_provider_oauth_auth,
};
pub(super) struct PreparedSameFormatProviderCandidate {
pub(super) transport: GatewayProviderTransportSnapshot,
pub(super) is_antigravity: bool,
pub(super) is_claude_code: bool,
pub(super) is_vertex: bool,
pub(super) is_kiro: bool,
pub(super) kiro_auth: Option<KiroRequestAuth>,
pub(super) auth_header: Option<String>,
pub(super) auth_value: Option<String>,
pub(super) mapped_model: String,
pub(super) report_kind: &'static str,
pub(super) upstream_is_stream: bool,
}
pub(super) async fn prepare_local_same_format_provider_candidate(
state: &AppState,
trace_id: &str,
input: &LocalSameFormatProviderDecisionInput,
eligible: &EligibleLocalExecutionCandidate,
candidate_index: u32,
candidate_id: &str,
spec: LocalSameFormatProviderSpec,
) -> Option<PreparedSameFormatProviderCandidate> {
let spec_metadata = local_same_format_provider_spec_metadata(spec);
let planner_state = PlannerAppState::new(state);
let candidate = &eligible.candidate;
let transport = eligible.transport.clone();
let behavior = classify_same_format_provider_request_behavior(&transport, spec_metadata);
if !same_format_provider_transport_supported(
&behavior,
&transport,
spec.family,
spec_metadata.api_format,
) {
super::super::payload::mark_skipped_local_same_format_provider_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_unsupported",
)
.await;
return None;
}
let vertex_query_auth = if behavior.is_vertex {
resolve_local_vertex_api_key_query_auth(&transport)
} else {
None
};
let should_try_oauth_auth =
should_try_same_format_provider_oauth_auth(&behavior, &transport, spec.family);
let oauth_auth = if should_try_oauth_auth {
resolve_candidate_oauth_auth(
planner_state,
&transport,
OauthPreparationContext {
trace_id,
api_format: spec_metadata.api_format,
operation: "same_format_provider_prepare",
},
)
.await
} else {
None
};
let kiro_auth = match oauth_auth.as_ref() {
Some(LocalResolvedOAuthRequestAuth::Kiro(auth)) => Some(auth.clone()),
_ => None,
};
let auth = if let Some(kiro_auth) = kiro_auth.as_ref() {
Some((kiro_auth.name.to_string(), kiro_auth.value.clone()))
} else if let Some(LocalResolvedOAuthRequestAuth::Header { name, value }) = oauth_auth.as_ref()
{
Some((name.clone(), value.clone()))
} else {
resolve_same_format_provider_direct_auth(&behavior, &transport, spec.family)
};
let (auth_header, auth_value) = match auth {
Some((name, value)) => (Some(name), Some(value)),
None if behavior.is_vertex && vertex_query_auth.is_some() => (None, None),
None => {
super::super::payload::mark_skipped_local_same_format_provider_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_auth_unavailable",
)
.await;
return None;
}
};
if behavior.is_vertex && vertex_query_auth.is_none() {
super::super::payload::mark_skipped_local_same_format_provider_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_auth_unavailable",
)
.await;
return None;
}
let mapped_model = match resolve_candidate_mapped_model(candidate) {
Ok(mapped_model) => mapped_model,
Err(skip_reason) => {
super::super::payload::mark_skipped_local_same_format_provider_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
return None;
}
};
Some(PreparedSameFormatProviderCandidate {
transport,
is_antigravity: behavior.is_antigravity,
is_claude_code: behavior.is_claude_code,
is_vertex: behavior.is_vertex,
is_kiro: behavior.is_kiro,
kiro_auth,
auth_header,
auth_value,
mapped_model,
report_kind: behavior.report_kind,
upstream_is_stream: behavior.upstream_is_stream,
})
}

View File

@@ -1,5 +1,14 @@
use tracing::warn;
use crate::ai_pipeline::planner::common::extract_requested_model_from_request;
use crate::ai_pipeline::planner::runtime_miss::{
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal,
apply_local_runtime_candidate_terminal_reason, set_local_runtime_miss_diagnostic_reason,
};
use crate::ai_pipeline::planner::spec_metadata::{
build_stream_plan_from_requested_model_family, build_sync_plan_from_requested_model_family,
local_same_format_provider_spec_metadata,
};
pub(crate) use crate::ai_pipeline::{
resolve_local_same_format_stream_spec as resolve_stream_spec,
resolve_local_same_format_sync_spec as resolve_sync_spec,
@@ -9,60 +18,8 @@ use super::{
materialize_local_same_format_provider_candidate_attempts,
maybe_build_local_same_format_provider_decision_payload_for_candidate,
resolve_local_same_format_provider_decision_input, AppState, GatewayControlDecision,
GatewayError, LocalSameFormatProviderFamily, LocalSameFormatProviderSpec,
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
GatewayError, LocalSameFormatProviderSpec, LocalStreamPlanAndReport, LocalSyncPlanAndReport,
};
use crate::ai_pipeline::planner::plan_builders::{
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
};
use crate::LocalExecutionRuntimeMissDiagnostic;
fn extract_requested_model(
parts: &http::request::Parts,
body_json: &serde_json::Value,
spec: LocalSameFormatProviderSpec,
) -> Option<String> {
match spec.family {
LocalSameFormatProviderFamily::Standard => body_json
.get("model")
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
LocalSameFormatProviderFamily::Gemini => {
let marker = "/models/";
let start = parts.uri.path().find(marker)? + marker.len();
let tail = &parts.uri.path()[start..];
let end = tail.find(':').unwrap_or(tail.len());
let model = tail[..end].trim();
if model.is_empty() {
None
} else {
Some(model.to_string())
}
}
}
}
fn build_local_same_format_miss_diagnostic(
decision: &GatewayControlDecision,
spec: LocalSameFormatProviderSpec,
requested_model: Option<&str>,
reason: &str,
) -> LocalExecutionRuntimeMissDiagnostic {
LocalExecutionRuntimeMissDiagnostic {
reason: reason.to_string(),
route_family: decision.route_family.clone(),
route_kind: decision.route_kind.clone(),
public_path: Some(decision.public_path.clone()),
plan_kind: Some(spec.decision_kind.to_string()),
requested_model: requested_model.map(ToOwned::to_owned),
candidate_count: None,
skipped_candidate_count: None,
skip_reasons: std::collections::BTreeMap::new(),
}
}
pub(crate) async fn build_local_sync_plan_and_reports(
state: &AppState,
@@ -72,46 +29,42 @@ pub(crate) async fn build_local_sync_plan_and_reports(
body_json: &serde_json::Value,
spec: LocalSameFormatProviderSpec,
) -> Result<Vec<LocalSyncPlanAndReport>, GatewayError> {
let spec_metadata = local_same_format_provider_spec_metadata(spec);
let requested_model_family = spec_metadata
.requested_model_family
.expect("same-format provider spec metadata should include requested-model family");
let Some(input) = resolve_local_same_format_provider_decision_input(
state, parts, trace_id, decision, body_json, spec,
)
.await
else {
state.set_local_execution_runtime_miss_diagnostic(
set_local_runtime_miss_diagnostic_reason(
state,
trace_id,
build_local_same_format_miss_diagnostic(
decision,
spec,
extract_requested_model(parts, body_json, spec).as_deref(),
"decision_input_unavailable",
),
decision,
spec_metadata.decision_kind,
extract_requested_model_from_request(parts, body_json, requested_model_family)
.as_deref(),
"decision_input_unavailable",
);
return Ok(Vec::new());
};
state.set_local_execution_runtime_miss_diagnostic(
set_local_runtime_miss_diagnostic_reason(
state,
trace_id,
build_local_same_format_miss_diagnostic(
decision,
spec,
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
),
decision,
spec_metadata.decision_kind,
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) =
materialize_local_same_format_provider_candidate_attempts(state, trace_id, &input, spec)
.await?;
let preserve_existing_candidate_signal = candidate_count == 0
&& state.local_execution_runtime_miss_diagnostic_has_candidate_signal(trace_id);
if !preserve_existing_candidate_signal {
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
diagnostic.candidate_count = Some(candidate_count);
diagnostic.reason = if candidate_count == 0 {
"candidate_list_empty".to_string()
} else {
"candidate_evaluation_incomplete".to_string()
};
});
}
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
state,
trace_id,
candidate_count,
);
if candidate_count == 0 {
return Ok(Vec::new());
}
@@ -126,14 +79,12 @@ pub(crate) async fn build_local_sync_plan_and_reports(
continue;
};
let built = match spec.family {
LocalSameFormatProviderFamily::Standard => {
build_standard_sync_plan_from_decision(parts, body_json, payload)
}
LocalSameFormatProviderFamily::Gemini => {
build_gemini_sync_plan_from_decision(parts, body_json, payload)
}
};
let built = build_sync_plan_from_requested_model_family(
requested_model_family,
parts,
body_json,
payload,
);
match built {
Ok(Some(value)) => plans.push(value),
@@ -141,7 +92,7 @@ pub(crate) async fn build_local_sync_plan_and_reports(
Err(err) => {
warn!(
trace_id = %trace_id,
api_format = spec.api_format,
api_format = spec_metadata.api_format,
error = ?err,
"gateway local same-format sync decision plan build failed"
);
@@ -149,15 +100,7 @@ pub(crate) async fn build_local_sync_plan_and_reports(
}
}
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
let candidate_count = diagnostic.candidate_count.unwrap_or(0);
let skipped_candidate_count = diagnostic.skipped_candidate_count.unwrap_or(0);
diagnostic.reason = if candidate_count > 0 && skipped_candidate_count >= candidate_count {
"all_candidates_skipped".to_string()
} else {
"no_local_sync_plans".to_string()
};
});
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_sync_plans");
Ok(plans)
}
@@ -170,46 +113,42 @@ pub(crate) async fn build_local_stream_plan_and_reports(
body_json: &serde_json::Value,
spec: LocalSameFormatProviderSpec,
) -> Result<Vec<LocalStreamPlanAndReport>, GatewayError> {
let spec_metadata = local_same_format_provider_spec_metadata(spec);
let requested_model_family = spec_metadata
.requested_model_family
.expect("same-format provider spec metadata should include requested-model family");
let Some(input) = resolve_local_same_format_provider_decision_input(
state, parts, trace_id, decision, body_json, spec,
)
.await
else {
state.set_local_execution_runtime_miss_diagnostic(
set_local_runtime_miss_diagnostic_reason(
state,
trace_id,
build_local_same_format_miss_diagnostic(
decision,
spec,
extract_requested_model(parts, body_json, spec).as_deref(),
"decision_input_unavailable",
),
decision,
spec_metadata.decision_kind,
extract_requested_model_from_request(parts, body_json, requested_model_family)
.as_deref(),
"decision_input_unavailable",
);
return Ok(Vec::new());
};
state.set_local_execution_runtime_miss_diagnostic(
set_local_runtime_miss_diagnostic_reason(
state,
trace_id,
build_local_same_format_miss_diagnostic(
decision,
spec,
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
),
decision,
spec_metadata.decision_kind,
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) =
materialize_local_same_format_provider_candidate_attempts(state, trace_id, &input, spec)
.await?;
let preserve_existing_candidate_signal = candidate_count == 0
&& state.local_execution_runtime_miss_diagnostic_has_candidate_signal(trace_id);
if !preserve_existing_candidate_signal {
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
diagnostic.candidate_count = Some(candidate_count);
diagnostic.reason = if candidate_count == 0 {
"candidate_list_empty".to_string()
} else {
"candidate_evaluation_incomplete".to_string()
};
});
}
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
state,
trace_id,
candidate_count,
);
if candidate_count == 0 {
return Ok(Vec::new());
}
@@ -224,14 +163,12 @@ pub(crate) async fn build_local_stream_plan_and_reports(
continue;
};
let built = match spec.family {
LocalSameFormatProviderFamily::Standard => {
build_standard_stream_plan_from_decision(parts, body_json, payload, false)
}
LocalSameFormatProviderFamily::Gemini => {
build_gemini_stream_plan_from_decision(parts, body_json, payload)
}
};
let built = build_stream_plan_from_requested_model_family(
requested_model_family,
parts,
body_json,
payload,
);
match built {
Ok(Some(value)) => plans.push(value),
@@ -239,7 +176,7 @@ pub(crate) async fn build_local_stream_plan_and_reports(
Err(err) => {
warn!(
trace_id = %trace_id,
api_format = spec.api_format,
api_format = spec_metadata.api_format,
error = ?err,
"gateway local same-format stream decision plan build failed"
);
@@ -247,15 +184,7 @@ pub(crate) async fn build_local_stream_plan_and_reports(
}
}
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
let candidate_count = diagnostic.candidate_count.unwrap_or(0);
let skipped_candidate_count = diagnostic.skipped_candidate_count.unwrap_or(0);
diagnostic.reason = if candidate_count > 0 && skipped_candidate_count >= candidate_count {
"all_candidates_skipped".to_string()
} else {
"no_local_stream_plans".to_string()
};
});
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_stream_plans");
Ok(plans)
}

View File

@@ -1,200 +1,5 @@
use std::collections::BTreeMap;
mod body;
mod url;
use serde_json::Value;
use url::form_urlencoded;
use crate::ai_pipeline::GatewayProviderTransportSnapshot;
use super::{
apply_local_body_rules, build_antigravity_v1internal_url, build_claude_code_messages_url,
build_claude_messages_url, build_gemini_content_url,
build_kiro_generate_assistant_response_url, build_kiro_provider_request_body,
build_passthrough_path_url, build_vertex_api_key_gemini_content_url,
resolve_local_vertex_api_key_query_auth, sanitize_claude_code_request_body,
AntigravityRequestUrlAction, LocalSameFormatProviderFamily, LocalSameFormatProviderSpec,
};
pub(super) fn build_same_format_provider_request_body(
body_json: &Value,
mapped_model: &str,
spec: LocalSameFormatProviderSpec,
body_rules: Option<&Value>,
upstream_is_stream: bool,
kiro_auth: Option<&crate::ai_pipeline::transport::kiro::KiroRequestAuth>,
is_claude_code: bool,
) -> Option<Value> {
if let Some(kiro_auth) = kiro_auth {
return build_kiro_provider_request_body(
body_json,
mapped_model,
&kiro_auth.auth_config,
body_rules,
);
}
let request_body_object = body_json.as_object()?;
let mut provider_request_body = serde_json::Map::from_iter(
request_body_object
.iter()
.map(|(key, value)| (key.clone(), value.clone())),
);
match spec.family {
LocalSameFormatProviderFamily::Standard => {
provider_request_body
.insert("model".to_string(), Value::String(mapped_model.to_string()));
if upstream_is_stream {
provider_request_body.insert("stream".to_string(), Value::Bool(true));
}
}
LocalSameFormatProviderFamily::Gemini => {
provider_request_body.remove("model");
}
}
let mut provider_request_body = Value::Object(provider_request_body);
if is_claude_code {
sanitize_claude_code_request_body(&mut provider_request_body);
}
if !apply_local_body_rules(&mut provider_request_body, body_rules, Some(body_json)) {
return None;
}
Some(provider_request_body)
}
pub(super) fn build_same_format_upstream_url(
parts: &http::request::Parts,
transport: &GatewayProviderTransportSnapshot,
mapped_model: &str,
spec: LocalSameFormatProviderSpec,
upstream_is_stream: bool,
kiro_auth: Option<&crate::ai_pipeline::transport::kiro::KiroRequestAuth>,
) -> Option<String> {
if let Some(kiro_auth) = kiro_auth {
return build_kiro_generate_assistant_response_url(
&transport.endpoint.base_url,
parts.uri.query(),
Some(kiro_auth.auth_config.effective_api_region()),
);
}
if transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("claude_code")
{
return Some(build_claude_code_messages_url(
&transport.endpoint.base_url,
parts.uri.query(),
));
}
if transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("vertex_ai")
{
let auth = resolve_local_vertex_api_key_query_auth(transport)?;
return build_vertex_api_key_gemini_content_url(
mapped_model,
upstream_is_stream,
&auth.value,
parts.uri.query(),
);
}
if transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("antigravity")
{
let query = parts.uri.query().map(|query| {
form_urlencoded::parse(query.as_bytes())
.into_owned()
.collect::<BTreeMap<String, String>>()
});
return build_antigravity_v1internal_url(
&transport.endpoint.base_url,
if upstream_is_stream {
AntigravityRequestUrlAction::StreamGenerateContent
} else {
AntigravityRequestUrlAction::GenerateContent
},
query.as_ref(),
);
}
let custom_path = transport
.endpoint
.custom_path
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
if let Some(path) = custom_path {
let blocked_keys = match spec.family {
LocalSameFormatProviderFamily::Standard => &[][..],
LocalSameFormatProviderFamily::Gemini => &["key"][..],
};
let url = build_passthrough_path_url(
&transport.endpoint.base_url,
path,
parts.uri.query(),
blocked_keys,
)?;
return Some(maybe_add_gemini_stream_alt_sse(url, spec));
}
let url = match spec.family {
LocalSameFormatProviderFamily::Standard => Some(build_claude_messages_url(
&transport.endpoint.base_url,
parts.uri.query(),
)),
LocalSameFormatProviderFamily::Gemini => build_gemini_content_url(
&transport.endpoint.base_url,
mapped_model,
spec.require_streaming,
parts.uri.query(),
),
}?;
Some(maybe_add_gemini_stream_alt_sse(url, spec))
}
pub(super) fn extract_gemini_model_from_path(path: &str) -> Option<String> {
let (_, suffix) = path.split_once("/models/")?;
let model = suffix
.split_once(':')
.map(|(value, _)| value)
.unwrap_or(suffix);
let model = model.trim();
if model.is_empty() {
None
} else {
Some(model.to_string())
}
}
fn maybe_add_gemini_stream_alt_sse(
upstream_url: String,
spec: LocalSameFormatProviderSpec,
) -> String {
if spec.family != LocalSameFormatProviderFamily::Gemini || !spec.require_streaming {
return upstream_url;
}
let has_alt = upstream_url
.split_once('?')
.map(|(_, query)| {
form_urlencoded::parse(query.as_bytes())
.any(|(key, _)| key.as_ref().eq_ignore_ascii_case("alt"))
})
.unwrap_or(false);
if has_alt {
return upstream_url;
}
if upstream_url.contains('?') {
format!("{upstream_url}&alt=sse")
} else {
format!("{upstream_url}?alt=sse")
}
}
pub(super) use self::body::build_same_format_provider_request_body;
pub(super) use self::url::build_same_format_upstream_url;

View File

@@ -0,0 +1,52 @@
use serde_json::Value;
use super::super::{
apply_local_body_rules, build_kiro_provider_request_body, sanitize_claude_code_request_body,
LocalSameFormatProviderFamily, LocalSameFormatProviderSpec,
};
pub(crate) fn build_same_format_provider_request_body(
body_json: &Value,
mapped_model: &str,
spec: LocalSameFormatProviderSpec,
body_rules: Option<&Value>,
upstream_is_stream: bool,
kiro_auth: Option<&crate::ai_pipeline::transport::kiro::KiroRequestAuth>,
is_claude_code: bool,
) -> Option<Value> {
if let Some(kiro_auth) = kiro_auth {
return build_kiro_provider_request_body(
body_json,
mapped_model,
&kiro_auth.auth_config,
body_rules,
);
}
let request_body_object = body_json.as_object()?;
let mut provider_request_body = serde_json::Map::from_iter(
request_body_object
.iter()
.map(|(key, value)| (key.clone(), value.clone())),
);
match spec.family {
LocalSameFormatProviderFamily::Standard => {
provider_request_body
.insert("model".to_string(), Value::String(mapped_model.to_string()));
if upstream_is_stream {
provider_request_body.insert("stream".to_string(), Value::Bool(true));
}
}
LocalSameFormatProviderFamily::Gemini => {
provider_request_body.remove("model");
}
}
let mut provider_request_body = Value::Object(provider_request_body);
if is_claude_code {
sanitize_claude_code_request_body(&mut provider_request_body);
}
if !apply_local_body_rules(&mut provider_request_body, body_rules, Some(body_json)) {
return None;
}
Some(provider_request_body)
}

View File

@@ -0,0 +1,138 @@
use std::collections::BTreeMap;
use url::form_urlencoded;
use crate::ai_pipeline::GatewayProviderTransportSnapshot;
use super::super::{
build_antigravity_v1internal_url, build_claude_code_messages_url, build_claude_messages_url,
build_gemini_content_url, build_kiro_generate_assistant_response_url,
build_passthrough_path_url, build_vertex_api_key_gemini_content_url,
resolve_local_vertex_api_key_query_auth, AntigravityRequestUrlAction,
LocalSameFormatProviderFamily, LocalSameFormatProviderSpec,
};
pub(crate) fn build_same_format_upstream_url(
parts: &http::request::Parts,
transport: &GatewayProviderTransportSnapshot,
mapped_model: &str,
spec: LocalSameFormatProviderSpec,
upstream_is_stream: bool,
kiro_auth: Option<&crate::ai_pipeline::transport::kiro::KiroRequestAuth>,
) -> Option<String> {
if let Some(kiro_auth) = kiro_auth {
return build_kiro_generate_assistant_response_url(
&transport.endpoint.base_url,
parts.uri.query(),
Some(kiro_auth.auth_config.effective_api_region()),
);
}
if transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("claude_code")
{
return Some(build_claude_code_messages_url(
&transport.endpoint.base_url,
parts.uri.query(),
));
}
if transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("vertex_ai")
{
let auth = resolve_local_vertex_api_key_query_auth(transport)?;
return build_vertex_api_key_gemini_content_url(
mapped_model,
upstream_is_stream,
&auth.value,
parts.uri.query(),
);
}
if transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("antigravity")
{
let query = parts.uri.query().map(|query| {
form_urlencoded::parse(query.as_bytes())
.into_owned()
.collect::<BTreeMap<String, String>>()
});
return build_antigravity_v1internal_url(
&transport.endpoint.base_url,
if upstream_is_stream {
AntigravityRequestUrlAction::StreamGenerateContent
} else {
AntigravityRequestUrlAction::GenerateContent
},
query.as_ref(),
);
}
let custom_path = transport
.endpoint
.custom_path
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
if let Some(path) = custom_path {
let blocked_keys = match spec.family {
LocalSameFormatProviderFamily::Standard => &[][..],
LocalSameFormatProviderFamily::Gemini => &["key"][..],
};
let url = build_passthrough_path_url(
&transport.endpoint.base_url,
path,
parts.uri.query(),
blocked_keys,
)?;
return Some(maybe_add_gemini_stream_alt_sse(url, spec));
}
let url = match spec.family {
LocalSameFormatProviderFamily::Standard => Some(build_claude_messages_url(
&transport.endpoint.base_url,
parts.uri.query(),
)),
LocalSameFormatProviderFamily::Gemini => build_gemini_content_url(
&transport.endpoint.base_url,
mapped_model,
spec.require_streaming,
parts.uri.query(),
),
}?;
Some(maybe_add_gemini_stream_alt_sse(url, spec))
}
fn maybe_add_gemini_stream_alt_sse(
upstream_url: String,
spec: LocalSameFormatProviderSpec,
) -> String {
if spec.family != LocalSameFormatProviderFamily::Gemini || !spec.require_streaming {
return upstream_url;
}
let has_alt = upstream_url
.split_once('?')
.map(|(_, query)| {
form_urlencoded::parse(query.as_bytes())
.any(|(key, _)| key.as_ref().eq_ignore_ascii_case("alt"))
})
.unwrap_or(false);
if has_alt {
return upstream_url;
}
if upstream_url.contains('?') {
format!("{upstream_url}&alt=sse")
} else {
format!("{upstream_url}?alt=sse")
}
}

View File

@@ -0,0 +1,90 @@
use std::collections::BTreeMap;
use aether_contracts::{ExecutionTimeouts, ProxySnapshot};
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::{ConversionMode, ExecutionStrategy, GatewayControlSyncDecisionResponse};
use crate::{EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION};
pub(crate) struct LocalExecutionDecisionResponseParts {
pub(crate) decision_is_stream: bool,
pub(crate) decision_kind: String,
pub(crate) execution_strategy: ExecutionStrategy,
pub(crate) conversion_mode: ConversionMode,
pub(crate) request_id: String,
pub(crate) candidate_id: String,
pub(crate) provider_name: String,
pub(crate) provider_id: String,
pub(crate) endpoint_id: String,
pub(crate) key_id: String,
pub(crate) upstream_base_url: String,
pub(crate) upstream_url: String,
pub(crate) provider_request_method: Option<String>,
pub(crate) auth_header: Option<String>,
pub(crate) auth_value: Option<String>,
pub(crate) provider_api_format: String,
pub(crate) client_api_format: String,
pub(crate) model_name: String,
pub(crate) mapped_model: String,
pub(crate) prompt_cache_key: Option<String>,
pub(crate) provider_request_headers: BTreeMap<String, String>,
pub(crate) provider_request_body: Option<serde_json::Value>,
pub(crate) provider_request_body_base64: Option<String>,
pub(crate) content_type: Option<String>,
pub(crate) proxy: Option<ProxySnapshot>,
pub(crate) tls_profile: Option<String>,
pub(crate) timeouts: Option<ExecutionTimeouts>,
pub(crate) upstream_is_stream: bool,
pub(crate) report_kind: Option<String>,
pub(crate) report_context: Option<serde_json::Value>,
pub(crate) auth_context: ExecutionRuntimeAuthContext,
}
pub(crate) fn build_local_execution_decision_response(
parts: LocalExecutionDecisionResponseParts,
) -> GatewayControlSyncDecisionResponse {
GatewayControlSyncDecisionResponse {
action: local_execution_decision_action(parts.decision_is_stream).to_string(),
decision_kind: Some(parts.decision_kind),
execution_strategy: Some(parts.execution_strategy.as_str().to_string()),
conversion_mode: Some(parts.conversion_mode.as_str().to_string()),
request_id: Some(parts.request_id),
candidate_id: Some(parts.candidate_id),
provider_name: Some(parts.provider_name),
provider_id: Some(parts.provider_id),
endpoint_id: Some(parts.endpoint_id),
key_id: Some(parts.key_id),
upstream_base_url: Some(parts.upstream_base_url),
upstream_url: Some(parts.upstream_url),
provider_request_method: parts.provider_request_method,
auth_header: parts.auth_header,
auth_value: parts.auth_value,
provider_api_format: Some(parts.provider_api_format.clone()),
client_api_format: Some(parts.client_api_format.clone()),
provider_contract: Some(parts.provider_api_format),
client_contract: Some(parts.client_api_format),
model_name: Some(parts.model_name),
mapped_model: Some(parts.mapped_model),
prompt_cache_key: parts.prompt_cache_key,
extra_headers: BTreeMap::new(),
provider_request_headers: parts.provider_request_headers,
provider_request_body: parts.provider_request_body,
provider_request_body_base64: parts.provider_request_body_base64,
content_type: parts.content_type,
proxy: parts.proxy,
tls_profile: parts.tls_profile,
timeouts: parts.timeouts,
upstream_is_stream: parts.upstream_is_stream,
report_kind: parts.report_kind,
report_context: parts.report_context,
auth_context: Some(parts.auth_context),
}
}
pub(crate) fn local_execution_decision_action(decision_is_stream: bool) -> &'static str {
if decision_is_stream {
EXECUTION_RUNTIME_STREAM_DECISION_ACTION
} else {
EXECUTION_RUNTIME_SYNC_DECISION_ACTION
}
}

View File

@@ -0,0 +1,152 @@
use std::collections::BTreeMap;
use serde_json::{Map, Value};
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
pub(crate) struct LocalExecutionReportContextParts<'a> {
pub(crate) auth_context: &'a ExecutionRuntimeAuthContext,
pub(crate) request_id: &'a str,
pub(crate) candidate_id: &'a str,
pub(crate) candidate_index: u32,
pub(crate) retry_index: u32,
pub(crate) model: &'a str,
pub(crate) provider_name: &'a str,
pub(crate) provider_id: &'a str,
pub(crate) endpoint_id: &'a str,
pub(crate) key_id: &'a str,
pub(crate) key_name: Option<&'a str>,
pub(crate) provider_api_format: &'a str,
pub(crate) client_api_format: &'a str,
pub(crate) mapped_model: Option<&'a str>,
pub(crate) upstream_url: Option<&'a str>,
pub(crate) provider_request_method: Option<Value>,
pub(crate) provider_request_headers: Option<&'a BTreeMap<String, String>>,
pub(crate) original_headers: &'a http::HeaderMap,
pub(crate) original_request_body: &'a Value,
pub(crate) has_envelope: bool,
pub(crate) needs_conversion: bool,
pub(crate) extra_fields: Map<String, Value>,
}
pub(crate) fn build_local_execution_report_context(
parts: LocalExecutionReportContextParts<'_>,
) -> Value {
let mut object = Map::new();
object.insert(
"user_id".to_string(),
Value::String(parts.auth_context.user_id.clone()),
);
object.insert(
"api_key_id".to_string(),
Value::String(parts.auth_context.api_key_id.clone()),
);
object.insert(
"username".to_string(),
parts
.auth_context
.username
.clone()
.map(Value::String)
.unwrap_or(Value::Null),
);
object.insert(
"api_key_name".to_string(),
parts
.auth_context
.api_key_name
.clone()
.map(Value::String)
.unwrap_or(Value::Null),
);
object.insert(
"request_id".to_string(),
Value::String(parts.request_id.to_string()),
);
object.insert(
"candidate_id".to_string(),
Value::String(parts.candidate_id.to_string()),
);
object.insert(
"candidate_index".to_string(),
Value::Number(parts.candidate_index.into()),
);
object.insert(
"retry_index".to_string(),
Value::Number(parts.retry_index.into()),
);
object.insert("model".to_string(), Value::String(parts.model.to_string()));
object.insert(
"provider_name".to_string(),
Value::String(parts.provider_name.to_string()),
);
object.insert(
"provider_id".to_string(),
Value::String(parts.provider_id.to_string()),
);
object.insert(
"endpoint_id".to_string(),
Value::String(parts.endpoint_id.to_string()),
);
object.insert(
"key_id".to_string(),
Value::String(parts.key_id.to_string()),
);
object.insert(
"provider_api_format".to_string(),
Value::String(parts.provider_api_format.to_string()),
);
object.insert(
"client_api_format".to_string(),
Value::String(parts.client_api_format.to_string()),
);
object.insert(
"original_headers".to_string(),
serde_json::to_value(crate::ai_pipeline::collect_control_headers(
parts.original_headers,
))
.expect("control headers should serialize"),
);
object.insert(
"original_request_body".to_string(),
crate::ai_pipeline::build_report_context_original_request_echo(parts.original_request_body)
.unwrap_or(Value::Null),
);
object.insert("has_envelope".to_string(), Value::Bool(parts.has_envelope));
object.insert(
"needs_conversion".to_string(),
Value::Bool(parts.needs_conversion),
);
if let Some(key_name) = parts.key_name {
object.insert("key_name".to_string(), Value::String(key_name.to_string()));
}
if let Some(mapped_model) = parts.mapped_model {
object.insert(
"mapped_model".to_string(),
Value::String(mapped_model.to_string()),
);
}
if let Some(upstream_url) = parts.upstream_url {
object.insert(
"upstream_url".to_string(),
Value::String(upstream_url.to_string()),
);
}
if let Some(provider_request_method) = parts.provider_request_method {
object.insert(
"provider_request_method".to_string(),
provider_request_method,
);
}
if let Some(provider_request_headers) = parts.provider_request_headers {
object.insert(
"provider_request_headers".to_string(),
serde_json::to_value(provider_request_headers)
.expect("provider request headers should serialize"),
);
}
object.extend(parts.extra_fields);
Value::Object(object)
}

View File

@@ -0,0 +1,137 @@
use crate::ai_pipeline::planner::common::{
apply_local_candidate_evaluation_progress, apply_local_candidate_terminal_plan_reason,
build_local_runtime_miss_diagnostic,
};
use crate::ai_pipeline::GatewayControlDecision;
use crate::{AppState, LocalExecutionRuntimeMissDiagnostic};
pub(crate) fn set_local_runtime_miss_diagnostic_reason(
state: &AppState,
trace_id: &str,
decision: &GatewayControlDecision,
plan_kind: &str,
requested_model: Option<&str>,
reason: &str,
) {
state.set_local_execution_runtime_miss_diagnostic(
trace_id,
build_local_runtime_miss_diagnostic(decision, plan_kind, requested_model, reason),
);
}
pub(crate) fn build_local_runtime_execution_exhausted_diagnostic(
decision: &GatewayControlDecision,
plan_kind: &str,
requested_model: Option<&str>,
candidate_count: usize,
) -> LocalExecutionRuntimeMissDiagnostic {
let mut diagnostic = build_local_runtime_miss_diagnostic(
decision,
plan_kind,
requested_model,
"execution_runtime_candidates_exhausted",
);
diagnostic.candidate_count = Some(candidate_count);
diagnostic
}
pub(crate) fn set_local_runtime_execution_exhausted_diagnostic(
state: &AppState,
trace_id: &str,
decision: &GatewayControlDecision,
plan_kind: &str,
requested_model: Option<&str>,
candidate_count: usize,
) {
state.set_local_execution_runtime_miss_diagnostic(
trace_id,
build_local_runtime_execution_exhausted_diagnostic(
decision,
plan_kind,
requested_model,
candidate_count,
),
);
}
pub(crate) fn build_local_runtime_candidate_evaluation_diagnostic(
decision: &GatewayControlDecision,
plan_kind: &str,
requested_model: Option<&str>,
candidate_count: usize,
) -> LocalExecutionRuntimeMissDiagnostic {
let mut diagnostic = build_local_runtime_miss_diagnostic(
decision,
plan_kind,
requested_model,
"candidate_evaluation_incomplete",
);
apply_local_candidate_evaluation_progress(&mut diagnostic, candidate_count);
diagnostic
}
pub(crate) fn set_local_runtime_candidate_evaluation_diagnostic(
state: &AppState,
trace_id: &str,
decision: &GatewayControlDecision,
plan_kind: &str,
requested_model: Option<&str>,
candidate_count: usize,
) {
state.set_local_execution_runtime_miss_diagnostic(
trace_id,
build_local_runtime_candidate_evaluation_diagnostic(
decision,
plan_kind,
requested_model,
candidate_count,
),
);
}
pub(crate) fn apply_local_runtime_candidate_evaluation_progress(
state: &AppState,
trace_id: &str,
candidate_count: usize,
) {
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
apply_local_candidate_evaluation_progress(diagnostic, candidate_count);
});
}
pub(crate) fn apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
state: &AppState,
trace_id: &str,
candidate_count: usize,
) {
let preserve_existing_candidate_signal = candidate_count == 0
&& state.local_execution_runtime_miss_diagnostic_has_candidate_signal(trace_id);
if preserve_existing_candidate_signal {
return;
}
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
}
pub(crate) fn apply_local_runtime_candidate_terminal_reason(
state: &AppState,
trace_id: &str,
no_plan_reason: &'static str,
) {
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
apply_local_candidate_terminal_plan_reason(diagnostic, no_plan_reason);
});
}
pub(crate) fn record_local_runtime_candidate_skip_reason(
state: &AppState,
trace_id: &str,
skip_reason: &'static str,
) {
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
*diagnostic
.skip_reasons
.entry(skip_reason.to_string())
.or_insert(0) += 1;
*diagnostic.skipped_candidate_count.get_or_insert(0) += 1;
});
}

View File

@@ -0,0 +1,140 @@
use crate::ai_pipeline::planner::common::RequestedModelFamily;
use crate::ai_pipeline::planner::plan_builders::{
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
};
use crate::ai_pipeline::{
GatewayControlSyncDecisionResponse, LocalGeminiFilesSpec, LocalOpenAiCliSpec,
LocalSameFormatProviderFamily, LocalSameFormatProviderSpec, LocalStandardSourceFamily,
LocalStandardSpec, LocalVideoCreateFamily, LocalVideoCreateSpec,
};
use crate::GatewayError;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct LocalExecutionSurfaceSpecMetadata {
pub(crate) api_format: &'static str,
pub(crate) decision_kind: &'static str,
pub(crate) report_kind: Option<&'static str>,
pub(crate) require_streaming: bool,
pub(crate) requested_model_family: Option<RequestedModelFamily>,
}
pub(crate) fn requested_model_family_for_standard_source(
family: LocalStandardSourceFamily,
) -> RequestedModelFamily {
match family {
LocalStandardSourceFamily::Standard => RequestedModelFamily::Standard,
LocalStandardSourceFamily::Gemini => RequestedModelFamily::Gemini,
}
}
pub(crate) fn local_standard_spec_metadata(
spec: LocalStandardSpec,
) -> LocalExecutionSurfaceSpecMetadata {
LocalExecutionSurfaceSpecMetadata {
api_format: spec.api_format,
decision_kind: spec.decision_kind,
report_kind: Some(spec.report_kind),
require_streaming: spec.require_streaming,
requested_model_family: Some(requested_model_family_for_standard_source(spec.family)),
}
}
pub(crate) fn local_same_format_provider_spec_metadata(
spec: LocalSameFormatProviderSpec,
) -> LocalExecutionSurfaceSpecMetadata {
LocalExecutionSurfaceSpecMetadata {
api_format: spec.api_format,
decision_kind: spec.decision_kind,
report_kind: Some(spec.report_kind),
require_streaming: spec.require_streaming,
requested_model_family: Some(requested_model_family_for_same_format_provider(spec.family)),
}
}
pub(crate) fn local_openai_cli_spec_metadata(
spec: LocalOpenAiCliSpec,
) -> LocalExecutionSurfaceSpecMetadata {
LocalExecutionSurfaceSpecMetadata {
api_format: spec.api_format,
decision_kind: spec.decision_kind,
report_kind: Some(spec.report_kind),
require_streaming: spec.require_streaming,
requested_model_family: None,
}
}
pub(crate) fn local_gemini_files_spec_metadata(
spec: LocalGeminiFilesSpec,
) -> LocalExecutionSurfaceSpecMetadata {
LocalExecutionSurfaceSpecMetadata {
api_format: "gemini:files",
decision_kind: spec.decision_kind,
report_kind: spec.report_kind,
require_streaming: spec.require_streaming,
requested_model_family: None,
}
}
pub(crate) fn local_video_create_spec_metadata(
spec: LocalVideoCreateSpec,
) -> LocalExecutionSurfaceSpecMetadata {
LocalExecutionSurfaceSpecMetadata {
api_format: spec.api_format,
decision_kind: spec.decision_kind,
report_kind: Some(spec.report_kind),
require_streaming: false,
requested_model_family: Some(requested_model_family_for_video_create(spec.family)),
}
}
pub(crate) fn requested_model_family_for_same_format_provider(
family: LocalSameFormatProviderFamily,
) -> RequestedModelFamily {
match family {
LocalSameFormatProviderFamily::Standard => RequestedModelFamily::Standard,
LocalSameFormatProviderFamily::Gemini => RequestedModelFamily::Gemini,
}
}
pub(crate) fn requested_model_family_for_video_create(
family: LocalVideoCreateFamily,
) -> RequestedModelFamily {
match family {
LocalVideoCreateFamily::OpenAi => RequestedModelFamily::Standard,
LocalVideoCreateFamily::Gemini => RequestedModelFamily::Gemini,
}
}
pub(crate) fn build_sync_plan_from_requested_model_family(
family: RequestedModelFamily,
parts: &http::request::Parts,
body_json: &serde_json::Value,
payload: GatewayControlSyncDecisionResponse,
) -> Result<Option<LocalSyncPlanAndReport>, GatewayError> {
match family {
RequestedModelFamily::Standard => {
build_standard_sync_plan_from_decision(parts, body_json, payload)
}
RequestedModelFamily::Gemini => {
build_gemini_sync_plan_from_decision(parts, body_json, payload)
}
}
}
pub(crate) fn build_stream_plan_from_requested_model_family(
family: RequestedModelFamily,
parts: &http::request::Parts,
body_json: &serde_json::Value,
payload: GatewayControlSyncDecisionResponse,
) -> Result<Option<LocalStreamPlanAndReport>, GatewayError> {
match family {
RequestedModelFamily::Standard => {
build_standard_stream_plan_from_decision(parts, body_json, payload, false)
}
RequestedModelFamily::Gemini => {
build_gemini_stream_plan_from_decision(parts, body_json, payload)
}
}
}

View File

@@ -1,4 +1,5 @@
mod decision;
mod request;
mod support;
use tracing::warn;
@@ -7,6 +8,7 @@ use crate::ai_pipeline::planner::plan_builders::{
build_passthrough_stream_plan_from_decision, build_passthrough_sync_plan_from_decision,
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
};
use crate::ai_pipeline::planner::spec_metadata::local_gemini_files_spec_metadata;
use crate::ai_pipeline::GatewayControlDecision;
use crate::ai_pipeline::{
resolve_gemini_files_stream_spec as resolve_stream_spec,
@@ -154,6 +156,7 @@ async fn build_local_sync_plan_and_reports(
decision: &GatewayControlDecision,
spec: LocalGeminiFilesSpec,
) -> Result<Vec<LocalSyncPlanAndReport>, GatewayError> {
let spec_metadata = local_gemini_files_spec_metadata(spec);
let Some(input) = resolve_local_gemini_files_decision_input(state, trace_id, decision).await
else {
return Ok(Vec::new());
@@ -186,7 +189,7 @@ async fn build_local_sync_plan_and_reports(
Err(err) => {
warn!(
trace_id = %trace_id,
decision_kind = spec.decision_kind,
decision_kind = spec_metadata.decision_kind,
error = ?err,
"gateway local gemini files sync decision plan build failed"
);
@@ -204,6 +207,7 @@ async fn build_local_stream_plan_and_reports(
decision: &GatewayControlDecision,
spec: LocalGeminiFilesSpec,
) -> Result<Vec<LocalStreamPlanAndReport>, GatewayError> {
let spec_metadata = local_gemini_files_spec_metadata(spec);
let Some(input) = resolve_local_gemini_files_decision_input(state, trace_id, decision).await
else {
return Ok(Vec::new());
@@ -237,7 +241,7 @@ async fn build_local_stream_plan_and_reports(
Err(err) => {
warn!(
trace_id = %trace_id,
decision_kind = spec.decision_kind,
decision_kind = spec_metadata.decision_kind,
error = ?err,
"gateway local gemini files stream decision plan build failed"
);

View File

@@ -1,33 +1,26 @@
use std::collections::BTreeMap;
use serde_json::json;
use tracing::warn;
use crate::ai_pipeline::contracts::GEMINI_FILES_UPLOAD_PLAN_KIND;
use crate::ai_pipeline::planner::common::{
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
use crate::ai_pipeline::planner::payload_metadata::{
build_local_execution_decision_response, LocalExecutionDecisionResponseParts,
};
use crate::ai_pipeline::transport::auth::{
build_passthrough_headers_with_auth, resolve_local_gemini_auth,
use crate::ai_pipeline::planner::report_context::{
build_local_execution_report_context, LocalExecutionReportContextParts,
};
use crate::ai_pipeline::transport::policy::supports_local_gemini_transport_with_network;
use crate::ai_pipeline::transport::url::build_gemini_files_passthrough_url;
use crate::ai_pipeline::planner::spec_metadata::local_gemini_files_spec_metadata;
use crate::ai_pipeline::transport::{
apply_local_body_rules, apply_local_header_rules, resolve_transport_execution_timeouts,
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_transport_tls_profile,
};
use crate::ai_pipeline::{
collect_control_headers, ConversionMode, ExecutionStrategy, PlannerAppState,
resolve_transport_execution_timeouts, resolve_transport_proxy_snapshot_with_tunnel_affinity,
resolve_transport_tls_profile,
};
use crate::ai_pipeline::{ConversionMode, ExecutionStrategy, PlannerAppState};
use crate::{AppState, GatewayControlSyncDecisionResponse};
use super::request::resolve_local_gemini_files_candidate_payload_parts;
use super::support::{
mark_skipped_local_gemini_files_candidate, LocalGeminiFilesCandidateAttempt,
LocalGeminiFilesDecisionInput, GEMINI_FILES_CANDIDATE_API_FORMAT,
GEMINI_FILES_CLIENT_API_FORMAT,
LocalGeminiFilesCandidateAttempt, LocalGeminiFilesDecisionInput, GEMINI_FILES_CLIENT_API_FORMAT,
};
use super::LocalGeminiFilesSpec;
#[allow(clippy::too_many_arguments)]
pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidate(
state: &AppState,
parts: &http::request::Parts,
@@ -39,268 +32,99 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
attempt: LocalGeminiFilesCandidateAttempt,
spec: LocalGeminiFilesSpec,
) -> Option<GatewayControlSyncDecisionResponse> {
let spec_metadata = local_gemini_files_spec_metadata(spec);
let planner_state = PlannerAppState::new(state);
let resolved = resolve_local_gemini_files_candidate_payload_parts(
state,
parts,
body_json,
body_base64,
body_is_empty,
trace_id,
input,
&attempt,
spec,
)
.await?;
let LocalGeminiFilesCandidateAttempt {
candidate,
eligible,
candidate_index,
candidate_id,
} = attempt;
let transport = match planner_state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
{
Ok(Some(snapshot)) => snapshot,
Ok(None) => {
mark_skipped_local_gemini_files_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"transport_snapshot_missing",
)
.await;
return None;
}
Err(err) => {
warn!(
trace_id = %trace_id,
error = ?err,
"gateway local gemini files provider transport read failed"
);
mark_skipped_local_gemini_files_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"transport_snapshot_read_failed",
)
.await;
return None;
}
};
if !supports_local_gemini_transport_with_network(&transport, GEMINI_FILES_CANDIDATE_API_FORMAT)
{
mark_skipped_local_gemini_files_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"transport_unsupported",
)
.await;
return None;
}
let Some((auth_header, auth_value)) = resolve_local_gemini_auth(&transport) else {
mark_skipped_local_gemini_files_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"transport_auth_unavailable",
)
.await;
return None;
};
let custom_path = transport
.endpoint
.custom_path
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
let passthrough_path = custom_path.unwrap_or(parts.uri.path());
let Some(upstream_url) = build_gemini_files_passthrough_url(
&transport.endpoint.base_url,
passthrough_path,
parts.uri.query(),
) else {
mark_skipped_local_gemini_files_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"upstream_url_missing",
)
.await;
return None;
};
let mut provider_request_body = if spec.decision_kind == GEMINI_FILES_UPLOAD_PLAN_KIND
&& !body_is_empty
&& body_base64.is_none()
{
Some(body_json.clone())
} else {
None
};
let provider_request_body_base64 = if spec.decision_kind == GEMINI_FILES_UPLOAD_PLAN_KIND {
body_base64
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
} else {
None
};
let original_request_body = if let Some(body_bytes_b64) = provider_request_body_base64.clone() {
json!({"body_bytes_b64": body_bytes_b64})
} else if !body_is_empty {
body_json.clone()
} else {
serde_json::Value::Null
};
if provider_request_body_base64.is_some() && transport.endpoint.body_rules.is_some() {
mark_skipped_local_gemini_files_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"transport_body_rules_unsupported_for_binary_upload",
)
.await;
return None;
}
if let Some(body) = provider_request_body.as_mut() {
if !apply_local_body_rules(
body,
transport.endpoint.body_rules.as_ref(),
Some(body_json),
) {
mark_skipped_local_gemini_files_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"transport_body_rules_apply_failed",
)
.await;
return None;
}
}
let mut provider_request_headers = build_passthrough_headers_with_auth(
&parts.headers,
&auth_header,
&auth_value,
&BTreeMap::new(),
);
if !apply_local_header_rules(
&mut provider_request_headers,
transport.endpoint.header_rules.as_ref(),
&[&auth_header, "content-type"],
provider_request_body
.as_ref()
.unwrap_or(&original_request_body),
Some(&original_request_body),
) {
mark_skipped_local_gemini_files_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"transport_header_rules_apply_failed",
)
.await;
return None;
}
let file_name = parts
.uri
.path()
.trim_start_matches("/v1beta/")
.trim()
.to_string();
let candidate = eligible.candidate;
let transport = resolved.transport;
let proxy =
resolve_transport_proxy_snapshot_with_tunnel_affinity(planner_state.app(), &transport)
.await;
let tls_profile = resolve_transport_tls_profile(&transport);
let mut extra_fields = serde_json::Map::new();
extra_fields.insert("file_key_id".to_string(), json!(candidate.key_id));
extra_fields.insert("file_name".to_string(), json!(resolved.file_name));
Some(GatewayControlSyncDecisionResponse {
action: if spec.require_streaming {
EXECUTION_RUNTIME_STREAM_DECISION_ACTION.to_string()
} else {
EXECUTION_RUNTIME_SYNC_DECISION_ACTION.to_string()
Some(build_local_execution_decision_response(
LocalExecutionDecisionResponseParts {
decision_is_stream: spec_metadata.require_streaming,
decision_kind: spec_metadata.decision_kind.to_string(),
execution_strategy: ExecutionStrategy::LocalSameFormat,
conversion_mode: ConversionMode::None,
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: 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: resolved.provider_request_headers,
provider_request_body: resolved.provider_request_body,
provider_request_body_base64: resolved.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,
tls_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(build_local_execution_report_context(
LocalExecutionReportContextParts {
auth_context: &input.auth_context,
request_id: trace_id,
candidate_id: &candidate_id,
candidate_index,
retry_index: 0,
model: "gemini-files",
provider_name: &transport.provider.name,
provider_id: &candidate.provider_id,
endpoint_id: &candidate.endpoint_id,
key_id: &candidate.key_id,
key_name: None,
provider_api_format: GEMINI_FILES_CLIENT_API_FORMAT,
client_api_format: GEMINI_FILES_CLIENT_API_FORMAT,
mapped_model: None,
upstream_url: None,
provider_request_method: None,
provider_request_headers: None,
original_headers: &parts.headers,
original_request_body: &resolved.original_request_body,
has_envelope: false,
needs_conversion: false,
extra_fields,
},
)),
auth_context: input.auth_context.clone(),
},
decision_kind: Some(spec.decision_kind.to_string()),
execution_strategy: Some(ExecutionStrategy::LocalSameFormat.as_str().to_string()),
conversion_mode: Some(ConversionMode::None.as_str().to_string()),
request_id: Some(trace_id.to_string()),
candidate_id: Some(candidate_id.clone()),
provider_name: Some(transport.provider.name.clone()),
provider_id: Some(candidate.provider_id.clone()),
endpoint_id: Some(candidate.endpoint_id.clone()),
key_id: Some(candidate.key_id.clone()),
upstream_base_url: Some(transport.endpoint.base_url.clone()),
upstream_url: Some(upstream_url),
provider_request_method: Some(parts.method.to_string()),
auth_header: Some(auth_header),
auth_value: Some(auth_value),
provider_api_format: Some(GEMINI_FILES_CLIENT_API_FORMAT.to_string()),
client_api_format: Some(GEMINI_FILES_CLIENT_API_FORMAT.to_string()),
provider_contract: Some(GEMINI_FILES_CLIENT_API_FORMAT.to_string()),
client_contract: Some(GEMINI_FILES_CLIENT_API_FORMAT.to_string()),
model_name: Some("gemini-files".to_string()),
mapped_model: Some(candidate.selected_provider_model_name.clone()),
prompt_cache_key: None,
extra_headers: BTreeMap::new(),
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,
tls_profile,
timeouts: resolve_transport_execution_timeouts(&transport),
upstream_is_stream: spec.require_streaming,
report_kind: spec.report_kind.map(ToOwned::to_owned),
report_context: Some(json!({
"user_id": input.auth_context.user_id,
"api_key_id": input.auth_context.api_key_id,
"username": input.auth_context.username,
"api_key_name": input.auth_context.api_key_name,
"request_id": trace_id,
"candidate_id": candidate_id,
"candidate_index": candidate_index,
"retry_index": 0,
"model": "gemini-files",
"provider_name": transport.provider.name,
"provider_id": candidate.provider_id,
"endpoint_id": candidate.endpoint_id,
"key_id": candidate.key_id,
"file_key_id": candidate.key_id,
"file_name": file_name,
"provider_api_format": GEMINI_FILES_CLIENT_API_FORMAT,
"client_api_format": GEMINI_FILES_CLIENT_API_FORMAT,
"original_headers": collect_control_headers(&parts.headers),
"original_request_body": crate::ai_pipeline::build_report_context_original_request_echo(&original_request_body),
"has_envelope": false,
"needs_conversion": false,
})),
auth_context: Some(input.auth_context.clone()),
})
))
}

View File

@@ -0,0 +1,205 @@
use std::collections::BTreeMap;
use serde_json::json;
use crate::ai_pipeline::contracts::GEMINI_FILES_UPLOAD_PLAN_KIND;
use crate::ai_pipeline::planner::spec_metadata::local_gemini_files_spec_metadata;
use crate::ai_pipeline::transport::auth::{
build_passthrough_headers_with_auth, resolve_local_gemini_auth,
};
use crate::ai_pipeline::transport::policy::supports_local_gemini_transport_with_network;
use crate::ai_pipeline::transport::url::build_gemini_files_passthrough_url;
use crate::ai_pipeline::transport::{apply_local_body_rules, apply_local_header_rules};
use crate::ai_pipeline::GatewayProviderTransportSnapshot;
use crate::AppState;
use super::support::{
mark_skipped_local_gemini_files_candidate, LocalGeminiFilesCandidateAttempt,
LocalGeminiFilesDecisionInput, GEMINI_FILES_CANDIDATE_API_FORMAT,
};
use super::LocalGeminiFilesSpec;
pub(super) struct LocalGeminiFilesCandidatePayloadParts {
pub(super) transport: GatewayProviderTransportSnapshot,
pub(super) auth_header: String,
pub(super) auth_value: String,
pub(super) provider_request_headers: BTreeMap<String, String>,
pub(super) provider_request_body: Option<serde_json::Value>,
pub(super) provider_request_body_base64: Option<String>,
pub(super) original_request_body: serde_json::Value,
pub(super) upstream_url: String,
pub(super) file_name: String,
}
#[allow(clippy::too_many_arguments)]
pub(super) async fn resolve_local_gemini_files_candidate_payload_parts(
state: &AppState,
parts: &http::request::Parts,
body_json: &serde_json::Value,
body_base64: Option<&str>,
body_is_empty: bool,
trace_id: &str,
input: &LocalGeminiFilesDecisionInput,
attempt: &LocalGeminiFilesCandidateAttempt,
spec: LocalGeminiFilesSpec,
) -> Option<LocalGeminiFilesCandidatePayloadParts> {
let spec_metadata = local_gemini_files_spec_metadata(spec);
let candidate = &attempt.eligible.candidate;
let transport = &attempt.eligible.transport;
if !supports_local_gemini_transport_with_network(transport, GEMINI_FILES_CANDIDATE_API_FORMAT) {
mark_skipped_local_gemini_files_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_unsupported",
)
.await;
return None;
}
let Some((auth_header, auth_value)) = resolve_local_gemini_auth(transport) else {
mark_skipped_local_gemini_files_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_auth_unavailable",
)
.await;
return None;
};
let custom_path = transport
.endpoint
.custom_path
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
let passthrough_path = custom_path.unwrap_or(parts.uri.path());
let Some(upstream_url) = build_gemini_files_passthrough_url(
&transport.endpoint.base_url,
passthrough_path,
parts.uri.query(),
) else {
mark_skipped_local_gemini_files_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"upstream_url_missing",
)
.await;
return None;
};
let mut provider_request_body = if spec_metadata.decision_kind == GEMINI_FILES_UPLOAD_PLAN_KIND
&& !body_is_empty
&& body_base64.is_none()
{
Some(body_json.clone())
} else {
None
};
let provider_request_body_base64 =
if spec_metadata.decision_kind == GEMINI_FILES_UPLOAD_PLAN_KIND {
body_base64
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
} else {
None
};
let original_request_body = if let Some(body_bytes_b64) = provider_request_body_base64.clone() {
json!({"body_bytes_b64": body_bytes_b64})
} else if !body_is_empty {
body_json.clone()
} else {
serde_json::Value::Null
};
if provider_request_body_base64.is_some() && transport.endpoint.body_rules.is_some() {
mark_skipped_local_gemini_files_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_body_rules_unsupported_for_binary_upload",
)
.await;
return None;
}
if let Some(body) = provider_request_body.as_mut() {
if !apply_local_body_rules(
body,
transport.endpoint.body_rules.as_ref(),
Some(body_json),
) {
mark_skipped_local_gemini_files_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_body_rules_apply_failed",
)
.await;
return None;
}
}
let mut provider_request_headers = build_passthrough_headers_with_auth(
&parts.headers,
&auth_header,
&auth_value,
&BTreeMap::new(),
);
if !apply_local_header_rules(
&mut provider_request_headers,
transport.endpoint.header_rules.as_ref(),
&[&auth_header, "content-type"],
provider_request_body
.as_ref()
.unwrap_or(&original_request_body),
Some(&original_request_body),
) {
mark_skipped_local_gemini_files_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_header_rules_apply_failed",
)
.await;
return None;
}
let file_name = parts
.uri
.path()
.trim_start_matches("/v1beta/")
.trim()
.to_string();
Some(LocalGeminiFilesCandidatePayloadParts {
transport: transport.clone(),
auth_header,
auth_value,
provider_request_headers,
provider_request_body,
provider_request_body_base64,
original_request_body,
upstream_url,
file_name,
})
}

View File

@@ -1,56 +1,57 @@
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use serde_json::json;
use tracing::warn;
use uuid::Uuid;
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::planner::candidate_affinity::{
rank_local_execution_candidates, remember_scheduler_affinity_for_candidate,
use crate::ai_pipeline::planner::candidate_eligibility::filter_and_rank_local_execution_candidates_without_transport_pair_gate;
use crate::ai_pipeline::planner::candidate_materialization::{
mark_skipped_local_execution_candidate,
persist_available_local_execution_candidates_with_context,
persist_skipped_local_execution_candidates_with_context,
remember_first_local_candidate_affinity,
};
use crate::ai_pipeline::planner::candidate_metadata::{
build_local_execution_candidate_metadata, LocalExecutionCandidateMetadataParts,
};
use crate::ai_pipeline::planner::decision_input::{
build_local_authenticated_decision_input, resolve_local_authenticated_decision_input,
};
use crate::ai_pipeline::planner::materialization_policy::{
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
};
use crate::ai_pipeline::PlannerAppState;
use crate::ai_pipeline::{
resolve_local_decision_execution_runtime_auth_context, GatewayControlDecision,
};
use crate::ai_pipeline::{GatewayAuthApiKeySnapshot, PlannerAppState};
use crate::clock::{current_unix_ms, current_unix_secs};
use crate::clock::current_unix_secs;
use crate::{AppState, GatewayError};
pub(super) use crate::ai_pipeline::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalGeminiFilesCandidateAttempt;
pub(super) use crate::ai_pipeline::planner::decision_input::LocalAuthenticatedDecisionInput as LocalGeminiFilesDecisionInput;
pub(super) const GEMINI_FILES_CANDIDATE_API_FORMAT: &str = "gemini:chat";
pub(super) const GEMINI_FILES_CLIENT_API_FORMAT: &str = "gemini:files";
pub(super) const GEMINI_FILES_REQUIRED_CAPABILITY: &str = "gemini_files";
#[derive(Debug, Clone)]
pub(super) struct LocalGeminiFilesDecisionInput {
pub(super) auth_context: ExecutionRuntimeAuthContext,
pub(super) auth_snapshot: GatewayAuthApiKeySnapshot,
pub(super) required_capabilities: Option<serde_json::Value>,
}
#[derive(Debug, Clone)]
pub(super) struct LocalGeminiFilesCandidateAttempt {
pub(super) candidate: SchedulerMinimalCandidateSelectionCandidate,
pub(super) candidate_index: u32,
pub(super) candidate_id: String,
}
pub(super) async fn resolve_local_gemini_files_decision_input(
state: &AppState,
trace_id: &str,
decision: &GatewayControlDecision,
) -> Option<LocalGeminiFilesDecisionInput> {
let planner_state = PlannerAppState::new(state);
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
return None;
};
let auth_snapshot = match planner_state
.read_auth_api_key_snapshot(
&auth_context.user_id,
&auth_context.api_key_id,
current_unix_secs(),
)
.await
let explicit_required_capabilities = json!({ "gemini_files": true });
let resolved_input = match resolve_local_authenticated_decision_input(
state,
auth_context,
None,
Some(&explicit_required_capabilities),
)
.await
{
Ok(Some(snapshot)) => snapshot,
Ok(Some(resolved_input)) => resolved_input,
Ok(None) => return None,
Err(err) => {
warn!(
@@ -62,21 +63,7 @@ pub(super) async fn resolve_local_gemini_files_decision_input(
}
};
let explicit_required_capabilities = json!({ "gemini_files": true });
let required_capabilities = planner_state
.resolve_request_candidate_required_capabilities(
&auth_context.user_id,
&auth_context.api_key_id,
None,
Some(&explicit_required_capabilities),
)
.await;
Some(LocalGeminiFilesDecisionInput {
auth_context,
auth_snapshot,
required_capabilities,
})
Some(build_local_authenticated_decision_input(resolved_input))
}
pub(super) async fn materialize_local_gemini_files_candidate_attempts(
@@ -85,6 +72,11 @@ pub(super) async fn materialize_local_gemini_files_candidate_attempts(
input: &LocalGeminiFilesDecisionInput,
) -> Result<Vec<LocalGeminiFilesCandidateAttempt>, GatewayError> {
let planner_state = PlannerAppState::new(state);
let persistence_policy = build_local_candidate_persistence_policy(
&input.auth_context,
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::GeminiFilesDecision,
);
let candidates = planner_state
.list_selectable_candidates_for_required_capability_without_requested_model(
GEMINI_FILES_CANDIDATE_API_FORMAT,
@@ -94,63 +86,54 @@ pub(super) async fn materialize_local_gemini_files_candidate_attempts(
current_unix_secs(),
)
.await?;
let candidates = rank_local_execution_candidates(
let (candidates, skipped_candidates) =
filter_and_rank_local_execution_candidates_without_transport_pair_gate(
planner_state,
candidates,
GEMINI_FILES_CLIENT_API_FORMAT,
None,
input.required_capabilities.as_ref(),
)
.await;
remember_first_local_candidate_affinity(
planner_state,
candidates,
Some(&input.auth_snapshot),
GEMINI_FILES_CLIENT_API_FORMAT,
input.required_capabilities.as_ref(),
None,
&candidates,
);
let attempts = persist_available_local_execution_candidates_with_context(
planner_state,
trace_id,
persistence_policy.available,
candidates,
|eligible| {
let mut extra_fields = serde_json::Map::new();
extra_fields.insert(
"candidate_api_format".to_string(),
json!(GEMINI_FILES_CANDIDATE_API_FORMAT),
);
Some(build_local_execution_candidate_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: GEMINI_FILES_CLIENT_API_FORMAT,
client_api_format: GEMINI_FILES_CLIENT_API_FORMAT,
extra_fields,
},
))
},
)
.await;
let created_at_unix_ms = current_unix_ms();
let mut attempts = Vec::with_capacity(candidates.len());
let mut affinity_remembered = false;
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
let generated_candidate_id = Uuid::new_v4().to_string();
if !affinity_remembered {
remember_scheduler_affinity_for_candidate(
planner_state,
Some(&input.auth_snapshot),
GEMINI_FILES_CLIENT_API_FORMAT,
&candidate.global_model_name,
&candidate,
);
affinity_remembered = true;
}
let extra_data = json!({
"provider_api_format": GEMINI_FILES_CLIENT_API_FORMAT,
"client_api_format": GEMINI_FILES_CLIENT_API_FORMAT,
"candidate_api_format": GEMINI_FILES_CANDIDATE_API_FORMAT,
"global_model_id": candidate.global_model_id.clone(),
"global_model_name": candidate.global_model_name.clone(),
"model_id": candidate.model_id.clone(),
"selected_provider_model_name": candidate.selected_provider_model_name.clone(),
"mapping_matched_model": candidate.mapping_matched_model.clone(),
"provider_name": candidate.provider_name.clone(),
"key_name": candidate.key_name.clone(),
});
let candidate_id = planner_state
.persist_available_local_candidate(
trace_id,
&input.auth_context.user_id,
&input.auth_context.api_key_id,
&candidate,
candidate_index as u32,
&generated_candidate_id,
input.required_capabilities.as_ref(),
Some(extra_data),
created_at_unix_ms,
"gateway local gemini files request candidate upsert failed",
)
.await;
attempts.push(LocalGeminiFilesCandidateAttempt {
candidate,
candidate_index: candidate_index as u32,
candidate_id,
});
}
persist_skipped_local_execution_candidates_with_context(
state,
trace_id,
persistence_policy.skipped,
attempts.len() as u32,
skipped_candidates,
)
.await;
Ok(attempts)
}
@@ -164,18 +147,19 @@ pub(super) async fn mark_skipped_local_gemini_files_candidate(
candidate_id: &str,
skip_reason: &'static str,
) {
PlannerAppState::new(state)
.persist_skipped_local_candidate(
trace_id,
&input.auth_context.user_id,
&input.auth_context.api_key_id,
candidate,
candidate_index,
candidate_id,
input.required_capabilities.as_ref(),
skip_reason,
current_unix_ms(),
"gateway local gemini files failed to persist skipped candidate",
)
.await;
let persistence_policy = build_local_candidate_persistence_policy(
&input.auth_context,
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::GeminiFilesDecision,
);
mark_skipped_local_execution_candidate(
state,
trace_id,
persistence_policy.skipped,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
}

View File

@@ -1,4 +1,5 @@
mod decision;
mod request;
mod support;
use tracing::warn;
@@ -6,6 +7,7 @@ use tracing::warn;
use crate::ai_pipeline::planner::plan_builders::{
build_passthrough_sync_plan_from_decision, LocalSyncPlanAndReport,
};
use crate::ai_pipeline::planner::spec_metadata::local_video_create_spec_metadata;
use crate::ai_pipeline::GatewayControlDecision;
use crate::ai_pipeline::{
resolve_local_video_sync_spec as resolve_sync_spec, LocalVideoCreateFamily,
@@ -44,6 +46,7 @@ pub(crate) async fn maybe_build_sync_local_video_decision_payload(
let Some(spec) = resolve_sync_spec(plan_kind) else {
return Ok(None);
};
let spec_metadata = local_video_create_spec_metadata(spec);
let Some(input) = resolve_local_video_create_decision_input(
state, parts, trace_id, decision, body_json, spec,
@@ -57,8 +60,8 @@ pub(crate) async fn maybe_build_sync_local_video_decision_payload(
state,
trace_id,
&input,
spec.api_format,
spec.decision_kind,
spec_metadata.api_format,
spec_metadata.decision_kind,
)
.await
else {
@@ -86,6 +89,7 @@ async fn build_local_sync_plan_and_reports(
decision: &GatewayControlDecision,
spec: LocalVideoCreateSpec,
) -> Result<Vec<LocalSyncPlanAndReport>, GatewayError> {
let spec_metadata = local_video_create_spec_metadata(spec);
let Some(input) = resolve_local_video_create_decision_input(
state, parts, trace_id, decision, body_json, spec,
)
@@ -98,8 +102,8 @@ async fn build_local_sync_plan_and_reports(
state,
trace_id,
&input,
spec.api_format,
spec.decision_kind,
spec_metadata.api_format,
spec_metadata.decision_kind,
)
.await
else {
@@ -122,7 +126,7 @@ async fn build_local_sync_plan_and_reports(
Err(err) => {
warn!(
trace_id = %trace_id,
decision_kind = spec.decision_kind,
decision_kind = spec_metadata.decision_kind,
error = ?err,
"gateway local video sync decision plan build failed"
);

View File

@@ -1,33 +1,20 @@
use std::collections::BTreeMap;
use serde_json::{json, Value};
use tracing::warn;
use crate::ai_pipeline::planner::common::EXECUTION_RUNTIME_SYNC_DECISION_ACTION;
use crate::ai_pipeline::transport::auth::{
build_passthrough_headers_with_auth, resolve_local_gemini_auth, resolve_local_openai_chat_auth,
use crate::ai_pipeline::planner::payload_metadata::{
build_local_execution_decision_response, LocalExecutionDecisionResponseParts,
};
use crate::ai_pipeline::transport::policy::{
supports_local_gemini_transport_with_network, supports_local_standard_transport_with_network,
};
use crate::ai_pipeline::transport::url::{
build_gemini_video_predict_long_running_url, build_passthrough_path_url,
use crate::ai_pipeline::planner::report_context::{
build_local_execution_report_context, LocalExecutionReportContextParts,
};
use crate::ai_pipeline::planner::spec_metadata::local_video_create_spec_metadata;
use crate::ai_pipeline::transport::{
apply_local_body_rules, apply_local_header_rules, resolve_transport_execution_timeouts,
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_transport_tls_profile,
};
use crate::ai_pipeline::{
collect_control_headers, ConversionMode, ExecutionStrategy, GatewayProviderTransportSnapshot,
PlannerAppState,
resolve_transport_execution_timeouts, resolve_transport_proxy_snapshot_with_tunnel_affinity,
resolve_transport_tls_profile,
};
use crate::ai_pipeline::{ConversionMode, ExecutionStrategy, PlannerAppState};
use crate::{AppState, GatewayControlSyncDecisionResponse};
use super::support::{
mark_skipped_local_video_candidate, LocalVideoCreateCandidateAttempt,
LocalVideoCreateDecisionInput,
};
use super::{LocalVideoCreateFamily, LocalVideoCreateSpec};
use super::request::resolve_local_video_create_candidate_payload_parts;
use super::support::{LocalVideoCreateCandidateAttempt, LocalVideoCreateDecisionInput};
use super::LocalVideoCreateSpec;
pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidate(
state: &AppState,
@@ -38,297 +25,88 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
attempt: LocalVideoCreateCandidateAttempt,
spec: LocalVideoCreateSpec,
) -> Option<GatewayControlSyncDecisionResponse> {
let spec_metadata = local_video_create_spec_metadata(spec);
let planner_state = PlannerAppState::new(state);
let resolved = resolve_local_video_create_candidate_payload_parts(
state, parts, body_json, trace_id, input, &attempt, spec,
)
.await?;
let LocalVideoCreateCandidateAttempt {
candidate,
eligible,
candidate_index,
candidate_id,
} = attempt;
let transport = match planner_state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
{
Ok(Some(snapshot)) => snapshot,
Ok(None) => {
mark_skipped_local_video_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"transport_snapshot_missing",
)
.await;
return None;
}
Err(err) => {
warn!(
trace_id = %trace_id,
decision_kind = spec.decision_kind,
error = ?err,
"gateway local video decision provider transport read failed"
);
mark_skipped_local_video_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"transport_snapshot_read_failed",
)
.await;
return None;
}
};
let transport_supported = match spec.family {
LocalVideoCreateFamily::OpenAi => {
supports_local_standard_transport_with_network(&transport, spec.api_format)
}
LocalVideoCreateFamily::Gemini => {
supports_local_gemini_transport_with_network(&transport, spec.api_format)
}
};
if !transport_supported {
mark_skipped_local_video_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"transport_unsupported",
)
.await;
return None;
}
let auth = match spec.family {
LocalVideoCreateFamily::OpenAi => resolve_local_openai_chat_auth(&transport),
LocalVideoCreateFamily::Gemini => resolve_local_gemini_auth(&transport),
};
let Some((auth_header, auth_value)) = auth else {
mark_skipped_local_video_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"transport_auth_unavailable",
)
.await;
return None;
};
let mapped_model = candidate.selected_provider_model_name.trim().to_string();
if mapped_model.is_empty() {
mark_skipped_local_video_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"mapped_model_missing",
)
.await;
return None;
}
let Some(upstream_url) =
build_video_upstream_url(parts, &transport, &mapped_model, spec.family)
else {
mark_skipped_local_video_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"upstream_url_missing",
)
.await;
return None;
};
let Some(provider_request_body) = build_provider_request_body(
body_json,
spec.family,
&mapped_model,
transport.endpoint.body_rules.as_ref(),
) else {
mark_skipped_local_video_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"provider_request_body_missing",
)
.await;
return None;
};
let mut provider_request_headers = build_passthrough_headers_with_auth(
&parts.headers,
&auth_header,
&auth_value,
&BTreeMap::new(),
);
if !apply_local_header_rules(
&mut provider_request_headers,
transport.endpoint.header_rules.as_ref(),
&[&auth_header, "content-type"],
&provider_request_body,
Some(body_json),
) {
mark_skipped_local_video_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"transport_header_rules_apply_failed",
)
.await;
return None;
}
let candidate = eligible.candidate;
let transport = resolved.transport;
let proxy =
resolve_transport_proxy_snapshot_with_tunnel_affinity(planner_state.app(), &transport)
.await;
let tls_profile = resolve_transport_tls_profile(&transport);
Some(GatewayControlSyncDecisionResponse {
action: EXECUTION_RUNTIME_SYNC_DECISION_ACTION.to_string(),
decision_kind: Some(spec.decision_kind.to_string()),
execution_strategy: Some(ExecutionStrategy::LocalSameFormat.as_str().to_string()),
conversion_mode: Some(ConversionMode::None.as_str().to_string()),
request_id: Some(trace_id.to_string()),
candidate_id: Some(candidate_id.clone()),
provider_name: Some(transport.provider.name.clone()),
provider_id: Some(candidate.provider_id.clone()),
endpoint_id: Some(candidate.endpoint_id.clone()),
key_id: Some(candidate.key_id.clone()),
upstream_base_url: Some(transport.endpoint.base_url.clone()),
upstream_url: Some(upstream_url),
provider_request_method: Some(parts.method.to_string()),
auth_header: Some(auth_header),
auth_value: Some(auth_value),
provider_api_format: Some(spec.api_format.to_string()),
client_api_format: Some(spec.api_format.to_string()),
provider_contract: Some(spec.api_format.to_string()),
client_contract: Some(spec.api_format.to_string()),
model_name: Some(input.requested_model.clone()),
mapped_model: Some(mapped_model.clone()),
prompt_cache_key: None,
extra_headers: BTreeMap::new(),
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,
tls_profile,
timeouts: resolve_transport_execution_timeouts(&transport),
upstream_is_stream: false,
report_kind: Some(spec.report_kind.to_string()),
report_context: Some(json!({
"user_id": input.auth_context.user_id.clone(),
"api_key_id": input.auth_context.api_key_id.clone(),
"username": input.auth_context.username.clone(),
"api_key_name": input.auth_context.api_key_name.clone(),
"request_id": trace_id,
"candidate_id": candidate_id,
"candidate_index": candidate_index,
"retry_index": 0,
"model": input.requested_model.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(),
"provider_api_format": spec.api_format,
"client_api_format": spec.api_format,
"mapped_model": mapped_model,
"original_headers": collect_control_headers(&parts.headers),
"original_request_body": crate::ai_pipeline::build_report_context_original_request_echo(body_json),
"has_envelope": false,
"needs_conversion": false,
})),
auth_context: Some(input.auth_context.clone()),
})
}
fn build_provider_request_body(
body_json: &serde_json::Value,
family: LocalVideoCreateFamily,
mapped_model: &str,
body_rules: Option<&serde_json::Value>,
) -> Option<serde_json::Value> {
let mut provider_request_body = match family {
LocalVideoCreateFamily::OpenAi => {
let mut provider_request_body = body_json.as_object().cloned().unwrap_or_default();
provider_request_body
.insert("model".to_string(), Value::String(mapped_model.to_string()));
serde_json::Value::Object(provider_request_body)
}
LocalVideoCreateFamily::Gemini => body_json.clone(),
};
if !apply_local_body_rules(&mut provider_request_body, body_rules, Some(body_json)) {
return None;
}
Some(provider_request_body)
}
fn build_video_upstream_url(
parts: &http::request::Parts,
transport: &GatewayProviderTransportSnapshot,
mapped_model: &str,
family: LocalVideoCreateFamily,
) -> Option<String> {
let custom_path = transport
.endpoint
.custom_path
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
if let Some(path) = custom_path {
let blocked_keys = match family {
LocalVideoCreateFamily::OpenAi => &[][..],
LocalVideoCreateFamily::Gemini => &["key"][..],
};
return build_passthrough_path_url(
&transport.endpoint.base_url,
path,
parts.uri.query(),
blocked_keys,
);
}
match family {
LocalVideoCreateFamily::OpenAi => build_passthrough_path_url(
&transport.endpoint.base_url,
parts.uri.path(),
parts.uri.query(),
&[],
),
LocalVideoCreateFamily::Gemini => build_gemini_video_predict_long_running_url(
&transport.endpoint.base_url,
mapped_model,
parts.uri.query(),
),
}
Some(build_local_execution_decision_response(
LocalExecutionDecisionResponseParts {
decision_is_stream: false,
decision_kind: spec_metadata.decision_kind.to_string(),
execution_strategy: ExecutionStrategy::LocalSameFormat,
conversion_mode: ConversionMode::None,
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: spec_metadata.api_format.to_string(),
client_api_format: spec_metadata.api_format.to_string(),
model_name: input.requested_model.clone(),
mapped_model: resolved.mapped_model.clone(),
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: 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,
tls_profile,
timeouts: resolve_transport_execution_timeouts(&transport),
upstream_is_stream: false,
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
report_context: Some(build_local_execution_report_context(
LocalExecutionReportContextParts {
auth_context: &input.auth_context,
request_id: trace_id,
candidate_id: &candidate_id,
candidate_index,
retry_index: 0,
model: &input.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,
provider_api_format: spec_metadata.api_format,
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),
upstream_url: None,
provider_request_method: None,
provider_request_headers: None,
original_headers: &parts.headers,
original_request_body: body_json,
has_envelope: false,
needs_conversion: false,
extra_fields: serde_json::Map::new(),
},
)),
auth_context: input.auth_context.clone(),
},
))
}

View File

@@ -0,0 +1,237 @@
use std::collections::BTreeMap;
use serde_json::Value;
use crate::ai_pipeline::planner::candidate_preparation::resolve_candidate_mapped_model;
use crate::ai_pipeline::planner::spec_metadata::local_video_create_spec_metadata;
use crate::ai_pipeline::transport::auth::{
build_passthrough_headers_with_auth, resolve_local_gemini_auth, resolve_local_openai_chat_auth,
};
use crate::ai_pipeline::transport::policy::{
supports_local_gemini_transport_with_network, supports_local_standard_transport_with_network,
};
use crate::ai_pipeline::transport::url::{
build_gemini_video_predict_long_running_url, build_passthrough_path_url,
};
use crate::ai_pipeline::transport::{apply_local_body_rules, apply_local_header_rules};
use crate::ai_pipeline::GatewayProviderTransportSnapshot;
use crate::AppState;
use super::support::{
mark_skipped_local_video_candidate, LocalVideoCreateCandidateAttempt,
LocalVideoCreateDecisionInput,
};
use super::{LocalVideoCreateFamily, LocalVideoCreateSpec};
pub(super) struct LocalVideoCreateCandidatePayloadParts {
pub(super) transport: GatewayProviderTransportSnapshot,
pub(super) auth_header: String,
pub(super) auth_value: String,
pub(super) mapped_model: String,
pub(super) provider_request_headers: BTreeMap<String, String>,
pub(super) provider_request_body: Value,
pub(super) upstream_url: String,
}
pub(super) async fn resolve_local_video_create_candidate_payload_parts(
state: &AppState,
parts: &http::request::Parts,
body_json: &serde_json::Value,
trace_id: &str,
input: &LocalVideoCreateDecisionInput,
attempt: &LocalVideoCreateCandidateAttempt,
spec: LocalVideoCreateSpec,
) -> Option<LocalVideoCreateCandidatePayloadParts> {
let spec_metadata = local_video_create_spec_metadata(spec);
let candidate = &attempt.eligible.candidate;
let transport = &attempt.eligible.transport;
let transport_supported = match spec.family {
LocalVideoCreateFamily::OpenAi => {
supports_local_standard_transport_with_network(transport, spec_metadata.api_format)
}
LocalVideoCreateFamily::Gemini => {
supports_local_gemini_transport_with_network(transport, spec_metadata.api_format)
}
};
if !transport_supported {
mark_skipped_local_video_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_unsupported",
)
.await;
return None;
}
let auth = match spec.family {
LocalVideoCreateFamily::OpenAi => resolve_local_openai_chat_auth(transport),
LocalVideoCreateFamily::Gemini => resolve_local_gemini_auth(transport),
};
let Some((auth_header, auth_value)) = auth else {
mark_skipped_local_video_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_auth_unavailable",
)
.await;
return None;
};
let mapped_model = match resolve_candidate_mapped_model(candidate) {
Ok(mapped_model) => mapped_model,
Err(skip_reason) => {
mark_skipped_local_video_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
skip_reason,
)
.await;
return None;
}
};
let Some(upstream_url) = build_video_upstream_url(parts, transport, &mapped_model, spec.family)
else {
mark_skipped_local_video_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"upstream_url_missing",
)
.await;
return None;
};
let Some(provider_request_body) = build_provider_request_body(
body_json,
spec.family,
&mapped_model,
transport.endpoint.body_rules.as_ref(),
) else {
mark_skipped_local_video_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"provider_request_body_missing",
)
.await;
return None;
};
let mut provider_request_headers = build_passthrough_headers_with_auth(
&parts.headers,
&auth_header,
&auth_value,
&BTreeMap::new(),
);
if !apply_local_header_rules(
&mut provider_request_headers,
transport.endpoint.header_rules.as_ref(),
&[&auth_header, "content-type"],
&provider_request_body,
Some(body_json),
) {
mark_skipped_local_video_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_header_rules_apply_failed",
)
.await;
return None;
}
Some(LocalVideoCreateCandidatePayloadParts {
transport: transport.clone(),
auth_header,
auth_value,
mapped_model,
provider_request_headers,
provider_request_body,
upstream_url,
})
}
fn build_provider_request_body(
body_json: &serde_json::Value,
family: LocalVideoCreateFamily,
mapped_model: &str,
body_rules: Option<&serde_json::Value>,
) -> Option<serde_json::Value> {
let mut provider_request_body = match family {
LocalVideoCreateFamily::OpenAi => {
let mut provider_request_body = body_json.as_object().cloned().unwrap_or_default();
provider_request_body
.insert("model".to_string(), Value::String(mapped_model.to_string()));
serde_json::Value::Object(provider_request_body)
}
LocalVideoCreateFamily::Gemini => body_json.clone(),
};
if !apply_local_body_rules(&mut provider_request_body, body_rules, Some(body_json)) {
return None;
}
Some(provider_request_body)
}
fn build_video_upstream_url(
parts: &http::request::Parts,
transport: &GatewayProviderTransportSnapshot,
mapped_model: &str,
family: LocalVideoCreateFamily,
) -> Option<String> {
let custom_path = transport
.endpoint
.custom_path
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
if let Some(path) = custom_path {
let blocked_keys = match family {
LocalVideoCreateFamily::OpenAi => &[][..],
LocalVideoCreateFamily::Gemini => &["key"][..],
};
return build_passthrough_path_url(
&transport.endpoint.base_url,
path,
parts.uri.query(),
blocked_keys,
);
}
match family {
LocalVideoCreateFamily::OpenAi => build_passthrough_path_url(
&transport.endpoint.base_url,
parts.uri.path(),
parts.uri.query(),
&[],
),
LocalVideoCreateFamily::Gemini => build_gemini_video_predict_long_running_url(
&transport.endpoint.base_url,
mapped_model,
parts.uri.query(),
),
}
}

View File

@@ -1,34 +1,35 @@
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use serde_json::json;
use tracing::warn;
use uuid::Uuid;
use super::{LocalVideoCreateFamily, LocalVideoCreateSpec};
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::planner::candidate_affinity::{
rank_local_execution_candidates, remember_scheduler_affinity_for_candidate,
use crate::ai_pipeline::planner::candidate_eligibility::filter_and_rank_local_execution_candidates;
use crate::ai_pipeline::planner::candidate_materialization::{
mark_skipped_local_execution_candidate,
persist_available_local_execution_candidates_with_context,
persist_skipped_local_execution_candidates_with_context,
remember_first_local_candidate_affinity,
};
use crate::ai_pipeline::planner::candidate_metadata::{
build_local_execution_candidate_metadata, LocalExecutionCandidateMetadataParts,
};
use crate::ai_pipeline::planner::common::extract_requested_model_from_request;
use crate::ai_pipeline::planner::decision_input::{
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
};
use crate::ai_pipeline::planner::materialization_policy::{
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
};
use crate::ai_pipeline::planner::spec_metadata::local_video_create_spec_metadata;
use crate::ai_pipeline::PlannerAppState;
use crate::ai_pipeline::{
resolve_local_decision_execution_runtime_auth_context, GatewayControlDecision,
};
use crate::ai_pipeline::{GatewayAuthApiKeySnapshot, PlannerAppState};
use crate::clock::{current_unix_ms, current_unix_secs};
use crate::clock::current_unix_secs;
use crate::AppState;
#[derive(Debug, Clone)]
pub(super) struct LocalVideoCreateDecisionInput {
pub(super) auth_context: ExecutionRuntimeAuthContext,
pub(super) requested_model: String,
pub(super) auth_snapshot: GatewayAuthApiKeySnapshot,
pub(super) required_capabilities: Option<serde_json::Value>,
}
#[derive(Debug, Clone)]
pub(super) struct LocalVideoCreateCandidateAttempt {
pub(super) candidate: SchedulerMinimalCandidateSelectionCandidate,
pub(super) candidate_index: u32,
pub(super) candidate_id: String,
}
pub(super) use crate::ai_pipeline::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalVideoCreateCandidateAttempt;
pub(super) use crate::ai_pipeline::planner::decision_input::LocalRequestedModelDecisionInput as LocalVideoCreateDecisionInput;
pub(super) async fn resolve_local_video_create_decision_input(
state: &AppState,
@@ -38,35 +39,33 @@ pub(super) async fn resolve_local_video_create_decision_input(
body_json: &serde_json::Value,
spec: LocalVideoCreateSpec,
) -> Option<LocalVideoCreateDecisionInput> {
let planner_state = PlannerAppState::new(state);
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
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;
};
let requested_model = match spec.family {
LocalVideoCreateFamily::OpenAi => body_json
.get("model")
.and_then(|value| value.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)?,
LocalVideoCreateFamily::Gemini => extract_gemini_video_model_from_path(parts.uri.path())?,
};
let requested_model = extract_requested_model_from_request(
parts,
body_json,
spec_metadata
.requested_model_family
.expect("video specs should declare requested-model family"),
)?;
let auth_snapshot = match planner_state
.read_auth_api_key_snapshot(
&auth_context.user_id,
&auth_context.api_key_id,
current_unix_secs(),
)
.await
let resolved_input = match resolve_local_authenticated_decision_input(
state,
auth_context,
Some(requested_model.as_str()),
None,
)
.await
{
Ok(Some(snapshot)) => snapshot,
Ok(Some(resolved_input)) => resolved_input,
Ok(None) => return None,
Err(err) => {
warn!(
trace_id = %trace_id,
decision_kind = spec.decision_kind,
decision_kind = spec_metadata.decision_kind,
error = ?err,
"gateway local video decision auth snapshot read failed"
);
@@ -74,21 +73,20 @@ pub(super) async fn resolve_local_video_create_decision_input(
}
};
let required_capabilities = planner_state
.resolve_request_candidate_required_capabilities(
&auth_context.user_id,
&auth_context.api_key_id,
Some(requested_model.as_str()),
None,
)
.await;
Some(LocalVideoCreateDecisionInput {
auth_context,
Some(build_local_requested_model_decision_input(
resolved_input,
requested_model,
auth_snapshot,
required_capabilities,
})
))
}
fn resolve_local_video_create_auth_context(
decision: &GatewayControlDecision,
family: LocalVideoCreateFamily,
) -> Option<ExecutionRuntimeAuthContext> {
let auth_context = resolve_local_decision_execution_runtime_auth_context(decision)?;
match family {
LocalVideoCreateFamily::OpenAi | LocalVideoCreateFamily::Gemini => Some(auth_context),
}
}
pub(super) async fn list_local_video_create_candidate_attempts(
@@ -141,62 +139,52 @@ async fn materialize_local_video_create_candidate_attempts(
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
api_format: &str,
) -> Vec<LocalVideoCreateCandidateAttempt> {
let candidates = rank_local_execution_candidates(
let persistence_policy = build_local_candidate_persistence_policy(
&input.auth_context,
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::VideoDecision,
);
let (candidates, skipped_candidates) = filter_and_rank_local_execution_candidates(
state,
candidates,
api_format,
&input.requested_model,
input.required_capabilities.as_ref(),
)
.await;
let created_at_unix_ms = current_unix_ms();
let mut attempts = Vec::with_capacity(candidates.len());
let mut affinity_remembered = false;
remember_first_local_candidate_affinity(
state,
Some(&input.auth_snapshot),
api_format,
Some(&input.requested_model),
&candidates,
);
let attempts = persist_available_local_execution_candidates_with_context(
state,
trace_id,
persistence_policy.available,
candidates,
|eligible| {
Some(build_local_execution_candidate_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: api_format,
client_api_format: api_format,
extra_fields: serde_json::Map::new(),
},
))
},
)
.await;
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
let generated_candidate_id = Uuid::new_v4().to_string();
if !affinity_remembered {
remember_scheduler_affinity_for_candidate(
state,
Some(&input.auth_snapshot),
api_format,
&input.requested_model,
&candidate,
);
affinity_remembered = true;
}
let extra_data = json!({
"provider_api_format": api_format,
"client_api_format": api_format,
"global_model_id": candidate.global_model_id.clone(),
"global_model_name": candidate.global_model_name.clone(),
"model_id": candidate.model_id.clone(),
"selected_provider_model_name": candidate.selected_provider_model_name.clone(),
"mapping_matched_model": candidate.mapping_matched_model.clone(),
"provider_name": candidate.provider_name.clone(),
"key_name": candidate.key_name.clone(),
});
let candidate_id = state
.persist_available_local_candidate(
trace_id,
&input.auth_context.user_id,
&input.auth_context.api_key_id,
&candidate,
candidate_index as u32,
&generated_candidate_id,
input.required_capabilities.as_ref(),
Some(extra_data),
created_at_unix_ms,
"gateway local video decision request candidate upsert failed",
)
.await;
attempts.push(LocalVideoCreateCandidateAttempt {
candidate,
candidate_index: candidate_index as u32,
candidate_id,
});
}
persist_skipped_local_execution_candidates_with_context(
state.app(),
trace_id,
persistence_policy.skipped,
attempts.len() as u32,
skipped_candidates,
)
.await;
attempts
}
@@ -210,27 +198,19 @@ pub(super) async fn mark_skipped_local_video_candidate(
candidate_id: &str,
skip_reason: &'static str,
) {
PlannerAppState::new(state)
.persist_skipped_local_candidate(
trace_id,
&input.auth_context.user_id,
&input.auth_context.api_key_id,
candidate,
candidate_index,
candidate_id,
input.required_capabilities.as_ref(),
skip_reason,
current_unix_ms(),
"gateway local video decision failed to persist skipped candidate",
)
.await;
}
fn extract_gemini_video_model_from_path(path: &str) -> Option<String> {
let suffix = path.strip_prefix("/v1beta/models/")?;
let model = suffix.split(':').next()?.trim();
if model.is_empty() {
return None;
}
Some(model.to_string())
let persistence_policy = build_local_candidate_persistence_policy(
&input.auth_context,
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::VideoDecision,
);
mark_skipped_local_execution_candidate(
state,
trace_id,
persistence_policy.skipped,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
}

View File

@@ -1,66 +1,25 @@
use tracing::warn;
use crate::ai_pipeline::planner::common::extract_requested_model_from_request;
use crate::ai_pipeline::planner::plan_builders::{
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
};
use crate::ai_pipeline::GatewayControlDecision;
use crate::{
AppState, GatewayControlSyncDecisionResponse, GatewayError, LocalExecutionRuntimeMissDiagnostic,
use crate::ai_pipeline::planner::runtime_miss::{
apply_local_runtime_candidate_evaluation_progress,
apply_local_runtime_candidate_terminal_reason, set_local_runtime_miss_diagnostic_reason,
};
use crate::ai_pipeline::planner::spec_metadata::{
build_stream_plan_from_requested_model_family, build_sync_plan_from_requested_model_family,
local_standard_spec_metadata,
};
use crate::ai_pipeline::GatewayControlDecision;
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
use super::candidates::{
materialize_local_standard_candidate_attempts, resolve_local_standard_decision_input,
};
use super::payload::maybe_build_local_standard_decision_payload_for_candidate;
use super::{LocalStandardSourceFamily, LocalStandardSpec};
fn extract_requested_model(
parts: &http::request::Parts,
body_json: &serde_json::Value,
spec: LocalStandardSpec,
) -> Option<String> {
match spec.family {
LocalStandardSourceFamily::Standard => body_json
.get("model")
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
LocalStandardSourceFamily::Gemini => {
let marker = "/models/";
let start = parts.uri.path().find(marker)? + marker.len();
let tail = &parts.uri.path()[start..];
let end = tail.find(':').unwrap_or(tail.len());
let model = tail[..end].trim();
if model.is_empty() {
None
} else {
Some(model.to_string())
}
}
}
}
fn build_local_standard_miss_diagnostic(
decision: &GatewayControlDecision,
spec: LocalStandardSpec,
requested_model: Option<&str>,
reason: &str,
) -> LocalExecutionRuntimeMissDiagnostic {
LocalExecutionRuntimeMissDiagnostic {
reason: reason.to_string(),
route_family: decision.route_family.clone(),
route_kind: decision.route_kind.clone(),
public_path: Some(decision.public_path.clone()),
plan_kind: Some(spec.decision_kind.to_string()),
requested_model: requested_model.map(ToOwned::to_owned),
candidate_count: None,
skipped_candidate_count: None,
skip_reasons: std::collections::BTreeMap::new(),
}
}
use super::LocalStandardSpec;
pub(crate) async fn maybe_build_sync_via_standard_family_payload(
state: &AppState,
@@ -74,6 +33,7 @@ pub(crate) async fn maybe_build_sync_via_standard_family_payload(
let Some(spec) = resolve_sync_spec(plan_kind) else {
return Ok(None);
};
let spec_metadata = local_standard_spec_metadata(spec);
let Some(input) =
resolve_local_standard_decision_input(state, parts, trace_id, decision, body_json, spec)
@@ -82,25 +42,17 @@ pub(crate) async fn maybe_build_sync_via_standard_family_payload(
return Ok(None);
};
state.set_local_execution_runtime_miss_diagnostic(
set_local_runtime_miss_diagnostic_reason(
state,
trace_id,
build_local_standard_miss_diagnostic(
decision,
spec,
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
),
decision,
spec_metadata.decision_kind,
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) =
materialize_local_standard_candidate_attempts(state, trace_id, &input, spec).await?;
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
diagnostic.candidate_count = Some(candidate_count);
diagnostic.reason = if candidate_count == 0 {
"candidate_list_empty".to_string()
} else {
"candidate_evaluation_incomplete".to_string()
};
});
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
for attempt in attempts {
if let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
@@ -112,17 +64,7 @@ pub(crate) async fn maybe_build_sync_via_standard_family_payload(
}
}
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
let candidate_count = diagnostic.candidate_count.unwrap_or(0);
let skipped_candidate_count = diagnostic.skipped_candidate_count.unwrap_or(0);
diagnostic.reason = if candidate_count == 0 {
"candidate_list_empty".to_string()
} else if skipped_candidate_count >= candidate_count {
"all_candidates_skipped".to_string()
} else {
"no_local_sync_plans".to_string()
};
});
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_sync_plans");
Ok(None)
}
@@ -139,6 +81,7 @@ pub(crate) async fn maybe_build_stream_via_standard_family_payload(
let Some(spec) = resolve_stream_spec(plan_kind) else {
return Ok(None);
};
let spec_metadata = local_standard_spec_metadata(spec);
let Some(input) =
resolve_local_standard_decision_input(state, parts, trace_id, decision, body_json, spec)
@@ -147,25 +90,17 @@ pub(crate) async fn maybe_build_stream_via_standard_family_payload(
return Ok(None);
};
state.set_local_execution_runtime_miss_diagnostic(
set_local_runtime_miss_diagnostic_reason(
state,
trace_id,
build_local_standard_miss_diagnostic(
decision,
spec,
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
),
decision,
spec_metadata.decision_kind,
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) =
materialize_local_standard_candidate_attempts(state, trace_id, &input, spec).await?;
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
diagnostic.candidate_count = Some(candidate_count);
diagnostic.reason = if candidate_count == 0 {
"candidate_list_empty".to_string()
} else {
"candidate_evaluation_incomplete".to_string()
};
});
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
for attempt in attempts {
if let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
@@ -177,17 +112,7 @@ pub(crate) async fn maybe_build_stream_via_standard_family_payload(
}
}
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
let candidate_count = diagnostic.candidate_count.unwrap_or(0);
let skipped_candidate_count = diagnostic.skipped_candidate_count.unwrap_or(0);
diagnostic.reason = if candidate_count == 0 {
"candidate_list_empty".to_string()
} else if skipped_candidate_count >= candidate_count {
"all_candidates_skipped".to_string()
} else {
"no_local_stream_plans".to_string()
};
});
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_stream_plans");
Ok(None)
}
@@ -200,40 +125,36 @@ pub(crate) async fn build_local_sync_plan_and_reports(
body_json: &serde_json::Value,
spec: LocalStandardSpec,
) -> Result<Vec<LocalSyncPlanAndReport>, GatewayError> {
let spec_metadata = local_standard_spec_metadata(spec);
let requested_model_family = spec_metadata
.requested_model_family
.expect("standard spec metadata should include requested-model family");
let Some(input) =
resolve_local_standard_decision_input(state, parts, trace_id, decision, body_json, spec)
.await
else {
state.set_local_execution_runtime_miss_diagnostic(
set_local_runtime_miss_diagnostic_reason(
state,
trace_id,
build_local_standard_miss_diagnostic(
decision,
spec,
extract_requested_model(parts, body_json, spec).as_deref(),
"decision_input_unavailable",
),
decision,
spec_metadata.decision_kind,
extract_requested_model_from_request(parts, body_json, requested_model_family)
.as_deref(),
"decision_input_unavailable",
);
return Ok(Vec::new());
};
state.set_local_execution_runtime_miss_diagnostic(
set_local_runtime_miss_diagnostic_reason(
state,
trace_id,
build_local_standard_miss_diagnostic(
decision,
spec,
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
),
decision,
spec_metadata.decision_kind,
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) =
materialize_local_standard_candidate_attempts(state, trace_id, &input, spec).await?;
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
diagnostic.candidate_count = Some(candidate_count);
diagnostic.reason = if candidate_count == 0 {
"candidate_list_empty".to_string()
} else {
"candidate_evaluation_incomplete".to_string()
};
});
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
if candidate_count == 0 {
return Ok(Vec::new());
}
@@ -246,36 +167,26 @@ pub(crate) async fn build_local_sync_plan_and_reports(
else {
continue;
};
let built = match spec.family {
LocalStandardSourceFamily::Standard => {
build_standard_sync_plan_from_decision(parts, body_json, payload)
}
LocalStandardSourceFamily::Gemini => {
build_gemini_sync_plan_from_decision(parts, body_json, payload)
}
};
let built = build_sync_plan_from_requested_model_family(
requested_model_family,
parts,
body_json,
payload,
);
match built {
Ok(Some(value)) => plans.push(value),
Ok(None) => {}
Err(err) => {
warn!(
trace_id = %trace_id,
api_format = spec.api_format,
api_format = spec_metadata.api_format,
error = ?err,
"gateway local standard sync plan build failed"
);
}
}
}
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
let candidate_count = diagnostic.candidate_count.unwrap_or(0);
let skipped_candidate_count = diagnostic.skipped_candidate_count.unwrap_or(0);
diagnostic.reason = if candidate_count > 0 && skipped_candidate_count >= candidate_count {
"all_candidates_skipped".to_string()
} else {
"no_local_sync_plans".to_string()
};
});
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_sync_plans");
Ok(plans)
}
@@ -287,40 +198,36 @@ pub(crate) async fn build_local_stream_plan_and_reports(
body_json: &serde_json::Value,
spec: LocalStandardSpec,
) -> Result<Vec<LocalStreamPlanAndReport>, GatewayError> {
let spec_metadata = local_standard_spec_metadata(spec);
let requested_model_family = spec_metadata
.requested_model_family
.expect("standard spec metadata should include requested-model family");
let Some(input) =
resolve_local_standard_decision_input(state, parts, trace_id, decision, body_json, spec)
.await
else {
state.set_local_execution_runtime_miss_diagnostic(
set_local_runtime_miss_diagnostic_reason(
state,
trace_id,
build_local_standard_miss_diagnostic(
decision,
spec,
extract_requested_model(parts, body_json, spec).as_deref(),
"decision_input_unavailable",
),
decision,
spec_metadata.decision_kind,
extract_requested_model_from_request(parts, body_json, requested_model_family)
.as_deref(),
"decision_input_unavailable",
);
return Ok(Vec::new());
};
state.set_local_execution_runtime_miss_diagnostic(
set_local_runtime_miss_diagnostic_reason(
state,
trace_id,
build_local_standard_miss_diagnostic(
decision,
spec,
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
),
decision,
spec_metadata.decision_kind,
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) =
materialize_local_standard_candidate_attempts(state, trace_id, &input, spec).await?;
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
diagnostic.candidate_count = Some(candidate_count);
diagnostic.reason = if candidate_count == 0 {
"candidate_list_empty".to_string()
} else {
"candidate_evaluation_incomplete".to_string()
};
});
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
if candidate_count == 0 {
return Ok(Vec::new());
}
@@ -333,35 +240,25 @@ pub(crate) async fn build_local_stream_plan_and_reports(
else {
continue;
};
let built = match spec.family {
LocalStandardSourceFamily::Standard => {
build_standard_stream_plan_from_decision(parts, body_json, payload, false)
}
LocalStandardSourceFamily::Gemini => {
build_gemini_stream_plan_from_decision(parts, body_json, payload)
}
};
let built = build_stream_plan_from_requested_model_family(
requested_model_family,
parts,
body_json,
payload,
);
match built {
Ok(Some(value)) => plans.push(value),
Ok(None) => {}
Err(err) => {
warn!(
trace_id = %trace_id,
api_format = spec.api_format,
api_format = spec_metadata.api_format,
error = ?err,
"gateway local standard stream plan build failed"
);
}
}
}
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
let candidate_count = diagnostic.candidate_count.unwrap_or(0);
let skipped_candidate_count = diagnostic.skipped_candidate_count.unwrap_or(0);
diagnostic.reason = if candidate_count > 0 && skipped_candidate_count >= candidate_count {
"all_candidates_skipped".to_string()
} else {
"no_local_stream_plans".to_string()
};
});
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_stream_plans");
Ok(plans)
}

View File

@@ -1,29 +1,35 @@
use std::collections::BTreeSet;
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use serde_json::json;
use tracing::warn;
use uuid::Uuid;
use crate::ai_pipeline::conversion::{
request_candidate_api_formats, request_conversion_kind,
request_conversion_requires_enable_flag, request_pair_allowed_for_transport,
use crate::ai_pipeline::conversion::{request_candidate_api_formats, request_conversion_kind};
use crate::ai_pipeline::planner::candidate_eligibility::filter_and_rank_local_execution_candidates;
use crate::ai_pipeline::planner::candidate_materialization::{
persist_available_local_execution_candidates_with_context,
persist_skipped_local_execution_candidates_with_context,
remember_first_local_candidate_affinity,
};
use crate::ai_pipeline::planner::candidate_affinity::{
rank_local_execution_candidates, remember_scheduler_affinity_for_candidate,
use crate::ai_pipeline::planner::candidate_metadata::{
build_local_execution_candidate_contract_metadata, LocalExecutionCandidateMetadataParts,
};
use crate::ai_pipeline::planner::candidate_source::auth_snapshot_allows_cross_format_candidate;
use crate::ai_pipeline::planner::common::extract_requested_model_from_request;
use crate::ai_pipeline::planner::decision_input::{
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
};
use crate::ai_pipeline::planner::materialization_policy::{
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
};
use crate::ai_pipeline::planner::spec_metadata::local_standard_spec_metadata;
use crate::ai_pipeline::PlannerAppState;
use crate::ai_pipeline::{
resolve_local_decision_execution_runtime_auth_context, ConversionMode, ExecutionStrategy,
GatewayControlDecision,
};
use crate::ai_pipeline::{GatewayAuthApiKeySnapshot, PlannerAppState};
use crate::clock::{current_unix_ms, current_unix_secs};
use crate::{append_execution_contract_fields_to_value, AppState, GatewayError};
use crate::clock::current_unix_secs;
use crate::{AppState, GatewayError};
use super::{
LocalStandardCandidateAttempt, LocalStandardDecisionInput, LocalStandardSourceFamily,
LocalStandardSpec,
};
use super::{LocalStandardCandidateAttempt, LocalStandardDecisionInput, LocalStandardSpec};
pub(super) async fn resolve_local_standard_decision_input(
state: &AppState,
@@ -33,35 +39,33 @@ pub(super) async fn resolve_local_standard_decision_input(
body_json: &serde_json::Value,
spec: LocalStandardSpec,
) -> Option<LocalStandardDecisionInput> {
let planner_state = PlannerAppState::new(state);
let spec_metadata = local_standard_spec_metadata(spec);
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
return None;
};
let requested_model = match spec.family {
LocalStandardSourceFamily::Standard => body_json
.get("model")
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)?,
LocalStandardSourceFamily::Gemini => extract_gemini_model_from_path(parts.uri.path())?,
};
let requested_model = extract_requested_model_from_request(
parts,
body_json,
spec_metadata
.requested_model_family
.expect("standard specs should declare requested-model family"),
)?;
let auth_snapshot = match planner_state
.read_auth_api_key_snapshot(
&auth_context.user_id,
&auth_context.api_key_id,
current_unix_secs(),
)
.await
let resolved_input = match resolve_local_authenticated_decision_input(
state,
auth_context,
Some(requested_model.as_str()),
None,
)
.await
{
Ok(Some(snapshot)) => snapshot,
Ok(Some(resolved_input)) => resolved_input,
Ok(None) => return None,
Err(err) => {
warn!(
trace_id = %trace_id,
api_format = spec.api_format,
api_format = spec_metadata.api_format,
error = ?err,
"gateway local standard decision auth snapshot read failed"
);
@@ -69,21 +73,10 @@ pub(super) async fn resolve_local_standard_decision_input(
}
};
let required_capabilities = planner_state
.resolve_request_candidate_required_capabilities(
&auth_context.user_id,
&auth_context.api_key_id,
Some(requested_model.as_str()),
None,
)
.await;
Some(LocalStandardDecisionInput {
auth_context,
Some(build_local_requested_model_decision_input(
resolved_input,
requested_model,
auth_snapshot,
required_capabilities,
})
))
}
pub(super) async fn materialize_local_standard_candidate_attempts(
@@ -92,13 +85,19 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
input: &LocalStandardDecisionInput,
spec: LocalStandardSpec,
) -> Result<(Vec<LocalStandardCandidateAttempt>, usize), GatewayError> {
let spec_metadata = local_standard_spec_metadata(spec);
let planner_state = PlannerAppState::new(state);
let persistence_policy = build_local_candidate_persistence_policy(
&input.auth_context,
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::StandardDecision,
);
let mut seen_candidates = BTreeSet::new();
let mut candidates = Vec::new();
for candidate_api_format in
request_candidate_api_formats(spec.api_format, spec.require_streaming)
request_candidate_api_formats(spec_metadata.api_format, spec_metadata.require_streaming)
{
let auth_snapshot = if candidate_api_format == spec.api_format {
let auth_snapshot = if candidate_api_format == spec_metadata.api_format {
Some(&input.auth_snapshot)
} else {
None
@@ -107,7 +106,7 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
.list_selectable_candidates(
candidate_api_format,
&input.requested_model,
spec.require_streaming,
spec_metadata.require_streaming,
input.required_capabilities.as_ref(),
auth_snapshot,
current_unix_secs(),
@@ -137,170 +136,66 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
}
}
}
let candidates = rank_local_execution_candidates(
let (candidates, skipped_candidates) = filter_and_rank_local_execution_candidates(
planner_state,
candidates,
spec.api_format,
spec_metadata.api_format,
&input.requested_model,
input.required_capabilities.as_ref(),
)
.await;
let candidate_count = candidates.len();
let candidate_count = candidates.len() + skipped_candidates.len();
let created_at_unix_ms = current_unix_ms();
let mut attempts = Vec::with_capacity(candidates.len());
let mut affinity_remembered = false;
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
let candidate_id = Uuid::new_v4().to_string();
let provider_api_format = candidate.endpoint_api_format.trim().to_ascii_lowercase();
if provider_api_format != spec.api_format {
if let Ok(Some(transport)) = planner_state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
{
if !request_pair_allowed_for_transport(
&transport,
spec.api_format,
provider_api_format.as_str(),
) {
let skip_reason =
if request_conversion_kind(spec.api_format, provider_api_format.as_str())
.is_some()
&& request_conversion_requires_enable_flag(
spec.api_format,
provider_api_format.as_str(),
)
&& !transport.provider.enable_format_conversion
{
"format_conversion_disabled"
} else {
"transport_unsupported"
};
super::payload::mark_skipped_local_standard_candidate(
state,
input,
trace_id,
&candidate,
candidate_index as u32,
&candidate_id,
skip_reason,
)
.await;
continue;
}
}
}
if !affinity_remembered {
remember_scheduler_affinity_for_candidate(
planner_state,
Some(&input.auth_snapshot),
spec.api_format,
&input.requested_model,
&candidate,
);
affinity_remembered = true;
}
let execution_strategy = if provider_api_format == spec.api_format {
ExecutionStrategy::LocalSameFormat
} else {
ExecutionStrategy::LocalCrossFormat
};
let conversion_mode = if crate::ai_pipeline::conversion::request_conversion_kind(
spec.api_format,
provider_api_format.as_str(),
)
.is_some()
{
ConversionMode::Bidirectional
} else {
ConversionMode::None
};
let extra_data = append_execution_contract_fields_to_value(
json!({
"provider_api_format": provider_api_format,
"client_api_format": spec.api_format,
"global_model_id": candidate.global_model_id.clone(),
"global_model_name": candidate.global_model_name.clone(),
"model_id": candidate.model_id.clone(),
"selected_provider_model_name": candidate.selected_provider_model_name.clone(),
"mapping_matched_model": candidate.mapping_matched_model.clone(),
"provider_name": candidate.provider_name.clone(),
"key_name": candidate.key_name.clone(),
}),
execution_strategy,
conversion_mode,
spec.api_format,
candidate.endpoint_api_format.as_str(),
);
remember_first_local_candidate_affinity(
planner_state,
Some(&input.auth_snapshot),
spec_metadata.api_format,
Some(&input.requested_model),
&candidates,
);
let attempts = persist_available_local_execution_candidates_with_context(
planner_state,
trace_id,
persistence_policy.available,
candidates,
|eligible| {
let provider_api_format = eligible.provider_api_format.clone();
let execution_strategy = if provider_api_format == spec_metadata.api_format {
ExecutionStrategy::LocalSameFormat
} else {
ExecutionStrategy::LocalCrossFormat
};
let conversion_mode =
if request_conversion_kind(spec_metadata.api_format, provider_api_format.as_str())
.is_some()
{
ConversionMode::Bidirectional
} else {
ConversionMode::None
};
Some(build_local_execution_candidate_contract_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: provider_api_format.as_str(),
client_api_format: spec_metadata.api_format,
extra_fields: serde_json::Map::new(),
},
execution_strategy,
conversion_mode,
eligible.candidate.endpoint_api_format.as_str(),
))
},
)
.await;
let stored_candidate_id = planner_state
.persist_available_local_candidate(
trace_id,
&input.auth_context.user_id,
&input.auth_context.api_key_id,
&candidate,
candidate_index as u32,
&candidate_id,
input.required_capabilities.as_ref(),
Some(extra_data),
created_at_unix_ms,
"gateway local standard decision request candidate upsert failed",
)
.await;
attempts.push(LocalStandardCandidateAttempt {
candidate,
candidate_index: candidate_index as u32,
candidate_id: stored_candidate_id,
});
}
persist_skipped_local_execution_candidates_with_context(
state,
trace_id,
persistence_policy.skipped,
attempts.len() as u32,
skipped_candidates,
)
.await;
Ok((attempts, candidate_count))
}
fn auth_snapshot_allows_cross_format_candidate(
auth_snapshot: &GatewayAuthApiKeySnapshot,
requested_model: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> bool {
if let Some(allowed_providers) = auth_snapshot.effective_allowed_providers() {
let provider_allowed = allowed_providers.iter().any(|value| {
value
.trim()
.eq_ignore_ascii_case(candidate.provider_id.trim())
|| value
.trim()
.eq_ignore_ascii_case(candidate.provider_name.trim())
});
if !provider_allowed {
return false;
}
}
if let Some(allowed_models) = auth_snapshot.effective_allowed_models() {
let model_allowed = allowed_models
.iter()
.any(|value| value == requested_model || value == &candidate.global_model_name);
if !model_allowed {
return false;
}
}
true
}
fn extract_gemini_model_from_path(path: &str) -> Option<String> {
let marker = "/models/";
let start = path.find(marker)? + marker.len();
let tail = &path[start..];
let end = tail.find(':').unwrap_or(tail.len());
let model = tail[..end].trim();
if model.is_empty() {
None
} else {
Some(model.to_string())
}
}

View File

@@ -1,29 +1,14 @@
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::GatewayAuthApiKeySnapshot;
mod build;
mod candidates;
mod payload;
mod request;
pub(crate) use self::build::{
build_local_stream_plan_and_reports, build_local_sync_plan_and_reports,
maybe_build_stream_via_standard_family_payload, maybe_build_sync_via_standard_family_payload,
};
pub(super) use crate::ai_pipeline::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalStandardCandidateAttempt;
pub(super) use crate::ai_pipeline::planner::decision_input::LocalRequestedModelDecisionInput as LocalStandardDecisionInput;
pub(crate) use crate::ai_pipeline::{
LocalStandardSourceFamily, LocalStandardSourceMode, LocalStandardSpec,
};
#[derive(Debug, Clone)]
pub(super) struct LocalStandardDecisionInput {
pub(super) auth_context: ExecutionRuntimeAuthContext,
pub(super) requested_model: String,
pub(super) auth_snapshot: GatewayAuthApiKeySnapshot,
pub(super) required_capabilities: Option<serde_json::Value>,
}
#[derive(Debug, Clone)]
pub(super) struct LocalStandardCandidateAttempt {
pub(super) candidate: aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate,
pub(super) candidate_index: u32,
pub(super) candidate_id: String,
}

View File

@@ -1,26 +1,25 @@
use std::collections::BTreeMap;
use serde_json::json;
use tracing::warn;
use crate::ai_pipeline::planner::common::force_upstream_streaming_for_provider;
use crate::ai_pipeline::planner::standard::apply_codex_openai_cli_special_headers;
use crate::ai_pipeline::transport::auth::{
build_claude_passthrough_headers, build_openai_passthrough_headers, ensure_upstream_auth_header,
use crate::ai_pipeline::planner::candidate_materialization::mark_skipped_local_execution_candidate;
use crate::ai_pipeline::planner::materialization_policy::{
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
};
use crate::ai_pipeline::planner::payload_metadata::{
build_local_execution_decision_response, LocalExecutionDecisionResponseParts,
};
use crate::ai_pipeline::planner::report_context::{
build_local_execution_report_context, LocalExecutionReportContextParts,
};
use crate::ai_pipeline::planner::spec_metadata::local_standard_spec_metadata;
use crate::ai_pipeline::transport::{
apply_local_header_rules, resolve_transport_execution_timeouts,
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_transport_tls_profile,
resolve_transport_execution_timeouts, resolve_transport_proxy_snapshot_with_tunnel_affinity,
resolve_transport_tls_profile,
};
use crate::ai_pipeline::{collect_control_headers, ConversionMode, ExecutionStrategy};
use crate::ai_pipeline::{LocalResolvedOAuthRequestAuth, PlannerAppState};
use crate::clock::current_unix_ms;
use crate::ai_pipeline::{ConversionMode, ExecutionStrategy};
use crate::{
append_execution_contract_fields_to_value, append_local_failover_policy_to_value, AppState,
GatewayControlSyncDecisionResponse, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
GatewayControlSyncDecisionResponse,
};
use super::request::resolve_local_standard_candidate_payload_parts;
use super::{LocalStandardCandidateAttempt, LocalStandardDecisionInput, LocalStandardSpec};
pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
@@ -32,346 +31,89 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
attempt: LocalStandardCandidateAttempt,
spec: LocalStandardSpec,
) -> Option<GatewayControlSyncDecisionResponse> {
let planner_state = PlannerAppState::new(state);
let spec_metadata = local_standard_spec_metadata(spec);
let LocalStandardCandidateAttempt {
candidate,
eligible,
candidate_index,
candidate_id,
} = attempt;
let provider_api_format = candidate.endpoint_api_format.trim().to_ascii_lowercase();
let Some(conversion_kind) = crate::ai_pipeline::conversion::request_conversion_kind(
spec.api_format,
provider_api_format.as_str(),
) else {
if provider_api_format == spec.api_format {
return None;
}
return None;
};
} = &attempt;
let candidate = &eligible.candidate;
let resolved = resolve_local_standard_candidate_payload_parts(
state, parts, trace_id, body_json, input, &attempt, spec,
)
.await?;
let transport = match planner_state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
{
Ok(Some(snapshot)) => snapshot,
Ok(None) => {
mark_skipped_local_standard_candidate(
Some(build_local_execution_decision_response(
LocalExecutionDecisionResponseParts {
decision_is_stream: spec_metadata.require_streaming,
decision_kind: spec_metadata.decision_kind.to_string(),
execution_strategy: ExecutionStrategy::LocalCrossFormat,
conversion_mode: ConversionMode::Bidirectional,
request_id: trace_id.to_string(),
candidate_id: candidate_id.to_string(),
provider_name: candidate.provider_name.clone(),
provider_id: candidate.provider_id.clone(),
endpoint_id: candidate.endpoint_id.clone(),
key_id: candidate.key_id.clone(),
upstream_base_url: resolved.transport.endpoint.base_url.clone(),
upstream_url: resolved.upstream_url.clone(),
provider_request_method: None,
auth_header: Some(resolved.auth_header.clone()),
auth_value: Some(resolved.auth_value.clone()),
provider_api_format: resolved.provider_api_format.clone(),
client_api_format: spec_metadata.api_format.to_string(),
model_name: input.requested_model.clone(),
mapped_model: resolved.mapped_model.clone(),
prompt_cache_key: None,
provider_request_headers: resolved.provider_request_headers.clone(),
provider_request_body: Some(resolved.provider_request_body.clone()),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
proxy: resolve_transport_proxy_snapshot_with_tunnel_affinity(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"transport_snapshot_missing",
&resolved.transport,
)
.await;
return None;
}
Err(err) => {
warn!(
trace_id = %trace_id,
api_format = spec.api_format,
error = ?err,
"gateway local standard decision provider transport read failed"
);
mark_skipped_local_standard_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"transport_snapshot_read_failed",
)
.await;
return None;
}
};
if !crate::ai_pipeline::conversion::request_pair_allowed_for_transport(
&transport,
spec.api_format,
provider_api_format.as_str(),
) {
let skip_reason = if crate::ai_pipeline::conversion::request_conversion_requires_enable_flag(
spec.api_format,
provider_api_format.as_str(),
) && !transport.provider.enable_format_conversion
{
"format_conversion_disabled"
} else {
"transport_unsupported"
};
mark_skipped_local_standard_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
skip_reason,
)
.await;
return None;
}
if !crate::ai_pipeline::conversion::request_conversion_transport_supported(
&transport,
conversion_kind,
) {
mark_skipped_local_standard_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"transport_unsupported",
)
.await;
return None;
}
let resolved_auth =
crate::ai_pipeline::conversion::request_conversion_direct_auth(&transport, conversion_kind);
let oauth_auth = if resolved_auth.is_none() {
match planner_state
.resolve_local_oauth_request_auth(&transport)
.await
{
Ok(Some(LocalResolvedOAuthRequestAuth::Header { name, value })) => Some((name, value)),
Ok(Some(LocalResolvedOAuthRequestAuth::Kiro(_))) => None,
Ok(None) => None,
Err(err) => {
warn!(
trace_id = %trace_id,
api_format = spec.api_format,
provider_type = %transport.provider.provider_type,
error = ?err,
"gateway local standard oauth auth resolution failed"
);
None
}
}
} else {
None
};
let Some((auth_header, auth_value)) = resolved_auth.or(oauth_auth) else {
mark_skipped_local_standard_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"transport_auth_unavailable",
)
.await;
return None;
};
let mapped_model = candidate.selected_provider_model_name.trim().to_string();
if mapped_model.is_empty() {
mark_skipped_local_standard_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"mapped_model_missing",
)
.await;
return None;
}
let upstream_is_stream = spec.require_streaming
|| force_upstream_streaming_for_provider(
transport.provider.provider_type.as_str(),
provider_api_format.as_str(),
);
let provider_request_body =
match crate::ai_pipeline::planner::standard::build_standard_request_body(
body_json,
spec.api_format,
&mapped_model,
transport.provider.provider_type.as_str(),
provider_api_format.as_str(),
parts.uri.path(),
upstream_is_stream,
transport.endpoint.body_rules.as_ref(),
Some(input.auth_context.api_key_id.as_str()),
) {
Some(body) => body,
None => {
mark_skipped_local_standard_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"provider_request_body_missing",
)
.await;
return None;
}
};
let upstream_url = match crate::ai_pipeline::planner::standard::build_standard_upstream_url(
parts,
&transport,
&mapped_model,
provider_api_format.as_str(),
upstream_is_stream,
) {
Some(url) => url,
None => {
mark_skipped_local_standard_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"upstream_url_missing",
)
.await;
return None;
}
};
let mut provider_request_headers = if provider_api_format.starts_with("claude:") {
build_claude_passthrough_headers(
&parts.headers,
&auth_header,
&auth_value,
&BTreeMap::new(),
Some("application/json"),
)
} else {
build_openai_passthrough_headers(
&parts.headers,
&auth_header,
&auth_value,
&BTreeMap::new(),
Some("application/json"),
)
};
if !apply_local_header_rules(
&mut provider_request_headers,
transport.endpoint.header_rules.as_ref(),
&[&auth_header, "content-type"],
&provider_request_body,
Some(body_json),
) {
mark_skipped_local_standard_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"transport_header_rules_apply_failed",
)
.await;
return None;
}
apply_codex_openai_cli_special_headers(
&mut provider_request_headers,
&provider_request_body,
&parts.headers,
transport.provider.provider_type.as_str(),
provider_api_format.as_str(),
Some(trace_id),
transport.key.decrypted_auth_config.as_deref(),
);
ensure_upstream_auth_header(&mut provider_request_headers, &auth_header, &auth_value);
if upstream_is_stream {
provider_request_headers
.entry("accept".to_string())
.or_insert_with(|| "text/event-stream".to_string());
}
Some(GatewayControlSyncDecisionResponse {
action: if spec.require_streaming {
EXECUTION_RUNTIME_STREAM_DECISION_ACTION.to_string()
} else {
EXECUTION_RUNTIME_SYNC_DECISION_ACTION.to_string()
.await,
tls_profile: resolve_transport_tls_profile(&resolved.transport),
timeouts: resolve_transport_execution_timeouts(&resolved.transport),
upstream_is_stream: resolved.upstream_is_stream,
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
report_context: Some(append_local_failover_policy_to_value(
append_execution_contract_fields_to_value(
build_local_execution_report_context(LocalExecutionReportContextParts {
auth_context: &input.auth_context,
request_id: trace_id,
candidate_id,
candidate_index: *candidate_index,
retry_index: 0,
model: &input.requested_model,
provider_name: &candidate.provider_name,
provider_id: &candidate.provider_id,
endpoint_id: &candidate.endpoint_id,
key_id: &candidate.key_id,
key_name: Some(&candidate.key_name),
provider_api_format: &resolved.provider_api_format,
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),
upstream_url: Some(&resolved.upstream_url),
provider_request_method: Some(serde_json::Value::Null),
provider_request_headers: Some(&resolved.provider_request_headers),
original_headers: &parts.headers,
original_request_body: body_json,
has_envelope: false,
needs_conversion: true,
extra_fields: serde_json::Map::new(),
}),
ExecutionStrategy::LocalCrossFormat,
ConversionMode::Bidirectional,
spec_metadata.api_format,
candidate.endpoint_api_format.as_str(),
),
&resolved.transport,
)),
auth_context: input.auth_context.clone(),
},
decision_kind: Some(spec.decision_kind.to_string()),
execution_strategy: Some(ExecutionStrategy::LocalCrossFormat.as_str().to_string()),
conversion_mode: Some(ConversionMode::Bidirectional.as_str().to_string()),
request_id: Some(trace_id.to_string()),
candidate_id: Some(candidate_id.clone()),
provider_name: Some(candidate.provider_name.clone()),
provider_id: Some(candidate.provider_id.clone()),
endpoint_id: Some(candidate.endpoint_id.clone()),
key_id: Some(candidate.key_id.clone()),
upstream_base_url: Some(transport.endpoint.base_url.clone()),
upstream_url: Some(upstream_url.clone()),
provider_request_method: None,
auth_header: Some(auth_header),
auth_value: Some(auth_value),
provider_api_format: Some(provider_api_format.clone()),
client_api_format: Some(spec.api_format.to_string()),
provider_contract: Some(provider_api_format.clone()),
client_contract: Some(spec.api_format.to_string()),
model_name: Some(input.requested_model.clone()),
mapped_model: Some(mapped_model.clone()),
prompt_cache_key: None,
extra_headers: BTreeMap::new(),
provider_request_headers: provider_request_headers.clone(),
provider_request_body: Some(provider_request_body.clone()),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
proxy: resolve_transport_proxy_snapshot_with_tunnel_affinity(state, &transport).await,
tls_profile: resolve_transport_tls_profile(&transport),
timeouts: resolve_transport_execution_timeouts(&transport),
upstream_is_stream,
report_kind: Some(spec.report_kind.to_string()),
report_context: Some(append_local_failover_policy_to_value(
append_execution_contract_fields_to_value(
json!({
"user_id": input.auth_context.user_id,
"api_key_id": input.auth_context.api_key_id,
"username": input.auth_context.username,
"api_key_name": input.auth_context.api_key_name,
"request_id": trace_id,
"candidate_id": candidate_id,
"candidate_index": candidate_index,
"retry_index": 0,
"model": input.requested_model,
"provider_name": candidate.provider_name,
"provider_id": candidate.provider_id,
"endpoint_id": candidate.endpoint_id,
"key_id": candidate.key_id,
"key_name": candidate.key_name,
"provider_api_format": provider_api_format,
"client_api_format": spec.api_format,
"mapped_model": mapped_model,
"upstream_url": upstream_url,
"provider_request_method": serde_json::Value::Null,
"provider_request_headers": provider_request_headers,
"original_headers": collect_control_headers(&parts.headers),
"original_request_body": crate::ai_pipeline::build_report_context_original_request_echo(body_json),
"has_envelope": false,
"needs_conversion": true,
}),
ExecutionStrategy::LocalCrossFormat,
ConversionMode::Bidirectional,
spec.api_format,
candidate.endpoint_api_format.as_str(),
),
&transport,
)),
auth_context: Some(input.auth_context.clone()),
})
))
}
pub(super) async fn mark_skipped_local_standard_candidate(
@@ -383,25 +125,19 @@ pub(super) async fn mark_skipped_local_standard_candidate(
candidate_id: &str,
skip_reason: &'static str,
) {
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
*diagnostic
.skip_reasons
.entry(skip_reason.to_string())
.or_insert(0) += 1;
*diagnostic.skipped_candidate_count.get_or_insert(0) += 1;
});
PlannerAppState::new(state)
.persist_skipped_local_candidate(
trace_id,
&input.auth_context.user_id,
&input.auth_context.api_key_id,
candidate,
candidate_index,
candidate_id,
input.required_capabilities.as_ref(),
skip_reason,
current_unix_ms(),
"gateway local standard decision failed to persist skipped candidate",
)
.await;
let persistence_policy = build_local_candidate_persistence_policy(
&input.auth_context,
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::StandardDecision,
);
mark_skipped_local_execution_candidate(
state,
trace_id,
persistence_policy.skipped,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
}

View File

@@ -0,0 +1,223 @@
use std::collections::BTreeMap;
use serde_json::Value;
use crate::ai_pipeline::planner::candidate_preparation::{
prepare_header_authenticated_candidate, OauthPreparationContext,
};
use crate::ai_pipeline::planner::common::force_upstream_streaming_for_provider;
use crate::ai_pipeline::planner::spec_metadata::local_standard_spec_metadata;
use crate::ai_pipeline::planner::standard::apply_codex_openai_cli_special_headers;
use crate::ai_pipeline::transport::apply_local_header_rules;
use crate::ai_pipeline::transport::auth::{
build_claude_passthrough_headers, build_openai_passthrough_headers, ensure_upstream_auth_header,
};
use crate::ai_pipeline::GatewayProviderTransportSnapshot;
use crate::AppState;
use super::payload::mark_skipped_local_standard_candidate;
use super::{LocalStandardCandidateAttempt, LocalStandardDecisionInput, LocalStandardSpec};
pub(crate) struct LocalStandardCandidatePayloadParts {
pub(super) auth_header: String,
pub(super) auth_value: String,
pub(super) mapped_model: String,
pub(super) provider_api_format: String,
pub(super) provider_request_body: Value,
pub(super) provider_request_headers: BTreeMap<String, String>,
pub(super) upstream_url: String,
pub(super) upstream_is_stream: bool,
pub(super) transport: GatewayProviderTransportSnapshot,
}
pub(crate) async fn resolve_local_standard_candidate_payload_parts(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
body_json: &serde_json::Value,
input: &LocalStandardDecisionInput,
attempt: &LocalStandardCandidateAttempt,
spec: LocalStandardSpec,
) -> Option<LocalStandardCandidatePayloadParts> {
let spec_metadata = local_standard_spec_metadata(spec);
let planner_state = crate::ai_pipeline::PlannerAppState::new(state);
let candidate = &attempt.eligible.candidate;
let transport = &attempt.eligible.transport;
let provider_api_format = attempt.eligible.provider_api_format.as_str();
let Some(conversion_kind) = crate::ai_pipeline::conversion::request_conversion_kind(
spec_metadata.api_format,
provider_api_format,
) else {
return None;
};
if !crate::ai_pipeline::conversion::request_conversion_transport_supported(
transport,
conversion_kind,
) {
mark_skipped_local_standard_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_unsupported",
)
.await;
return None;
}
let prepared_candidate = match prepare_header_authenticated_candidate(
planner_state,
transport,
candidate,
crate::ai_pipeline::conversion::request_conversion_direct_auth(transport, conversion_kind),
OauthPreparationContext {
trace_id,
api_format: provider_api_format,
operation: "standard_family_cross_format",
},
)
.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 upstream_is_stream = spec_metadata.require_streaming
|| force_upstream_streaming_for_provider(
transport.provider.provider_type.as_str(),
provider_api_format,
);
let provider_request_body =
match crate::ai_pipeline::planner::standard::build_standard_request_body(
body_json,
spec_metadata.api_format,
&prepared_candidate.mapped_model,
transport.provider.provider_type.as_str(),
provider_api_format,
parts.uri.path(),
upstream_is_stream,
transport.endpoint.body_rules.as_ref(),
Some(input.auth_context.api_key_id.as_str()),
) {
Some(body) => body,
None => {
mark_skipped_local_standard_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"provider_request_body_missing",
)
.await;
return None;
}
};
let upstream_url = match crate::ai_pipeline::planner::standard::build_standard_upstream_url(
parts,
transport,
&prepared_candidate.mapped_model,
provider_api_format,
upstream_is_stream,
) {
Some(url) => url,
None => {
mark_skipped_local_standard_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"upstream_url_missing",
)
.await;
return None;
}
};
let mut provider_request_headers = if provider_api_format.starts_with("claude:") {
build_claude_passthrough_headers(
&parts.headers,
&prepared_candidate.auth_header,
&prepared_candidate.auth_value,
&BTreeMap::new(),
Some("application/json"),
)
} else {
build_openai_passthrough_headers(
&parts.headers,
&prepared_candidate.auth_header,
&prepared_candidate.auth_value,
&BTreeMap::new(),
Some("application/json"),
)
};
if !apply_local_header_rules(
&mut provider_request_headers,
transport.endpoint.header_rules.as_ref(),
&[&prepared_candidate.auth_header, "content-type"],
&provider_request_body,
Some(body_json),
) {
mark_skipped_local_standard_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_header_rules_apply_failed",
)
.await;
return None;
}
apply_codex_openai_cli_special_headers(
&mut provider_request_headers,
&provider_request_body,
&parts.headers,
transport.provider.provider_type.as_str(),
provider_api_format,
Some(trace_id),
transport.key.decrypted_auth_config.as_deref(),
);
ensure_upstream_auth_header(
&mut provider_request_headers,
&prepared_candidate.auth_header,
&prepared_candidate.auth_value,
);
if upstream_is_stream {
provider_request_headers
.entry("accept".to_string())
.or_insert_with(|| "text/event-stream".to_string());
}
Some(LocalStandardCandidatePayloadParts {
auth_header: prepared_candidate.auth_header,
auth_value: prepared_candidate.auth_value,
mapped_model: prepared_candidate.mapped_model,
provider_api_format: provider_api_format.to_string(),
provider_request_body,
provider_request_headers,
upstream_url,
upstream_is_stream,
transport: transport.clone(),
})
}

View File

@@ -1,130 +1,12 @@
use tracing::warn;
use crate::ai_pipeline::PlannerAppState;
use crate::{AppState, GatewayControlSyncDecisionResponse};
#[path = "decision/cross_format.rs"]
mod cross_format;
#[path = "decision/same_format.rs"]
mod same_format;
#[path = "decision/payload.rs"]
mod payload;
#[path = "decision/request.rs"]
mod request;
#[path = "decision/support.rs"]
mod support;
use self::cross_format::build_cross_format_local_openai_chat_decision_payload_for_candidate;
use self::same_format::build_same_format_local_openai_chat_decision_payload_for_candidate;
use self::support::mark_skipped_local_openai_chat_candidate;
pub(super) use self::payload::maybe_build_local_openai_chat_decision_payload_for_candidate;
pub(super) use self::support::{
materialize_local_openai_chat_candidate_attempts, LocalOpenAiChatCandidateAttempt,
LocalOpenAiChatDecisionInput,
};
pub(super) async fn maybe_build_local_openai_chat_decision_payload_for_candidate(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
body_json: &serde_json::Value,
input: &LocalOpenAiChatDecisionInput,
attempt: LocalOpenAiChatCandidateAttempt,
decision_kind: &str,
report_kind: &str,
upstream_is_stream: bool,
) -> Option<GatewayControlSyncDecisionResponse> {
let planner_state = PlannerAppState::new(state);
let LocalOpenAiChatCandidateAttempt {
candidate,
candidate_index,
candidate_id,
} = attempt;
let transport = match planner_state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
{
Ok(Some(snapshot)) => snapshot,
Ok(None) => {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"transport_snapshot_missing",
)
.await;
return None;
}
Err(err) => {
warn!(
trace_id = %trace_id,
error = ?err,
"gateway local openai chat decision provider transport read failed"
);
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"transport_snapshot_read_failed",
)
.await;
return None;
}
};
let provider_api_format = transport.endpoint.api_format.trim().to_ascii_lowercase();
match provider_api_format.as_str() {
"openai:chat" => {
build_same_format_local_openai_chat_decision_payload_for_candidate(
state,
parts,
trace_id,
body_json,
input,
&candidate,
candidate_index,
&candidate_id,
decision_kind,
report_kind,
upstream_is_stream,
&transport,
)
.await
}
"claude:chat" | "claude:cli" | "gemini:chat" | "gemini:cli" | "openai:cli" => {
build_cross_format_local_openai_chat_decision_payload_for_candidate(
state,
parts,
trace_id,
body_json,
input,
&candidate,
candidate_index,
&candidate_id,
decision_kind,
upstream_is_stream,
&transport,
provider_api_format.as_str(),
)
.await
}
_ => {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
"transport_unsupported",
)
.await;
None
}
}
}

View File

@@ -1,327 +0,0 @@
use std::collections::BTreeMap;
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use serde_json::json;
use tracing::warn;
use crate::ai_pipeline::collect_control_headers;
use crate::ai_pipeline::conversion::{
request_conversion_direct_auth, request_conversion_kind,
request_conversion_requires_enable_flag, request_conversion_transport_supported,
request_pair_allowed_for_transport,
};
use crate::ai_pipeline::planner::common::OPENAI_CHAT_STREAM_PLAN_KIND;
use crate::ai_pipeline::planner::standard::{
apply_codex_openai_cli_special_headers, build_cross_format_openai_chat_request_body,
build_cross_format_openai_chat_upstream_url,
};
use crate::ai_pipeline::transport::auth::{
build_claude_passthrough_headers, build_openai_passthrough_headers, ensure_upstream_auth_header,
};
use crate::ai_pipeline::transport::{
apply_local_header_rules, resolve_transport_execution_timeouts,
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_transport_tls_profile,
};
use crate::ai_pipeline::{ConversionMode, ExecutionStrategy, PlannerAppState};
use crate::ai_pipeline::{GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth};
use crate::{
append_execution_contract_fields_to_value, append_local_failover_policy_to_value, AppState,
GatewayControlSyncDecisionResponse,
};
use super::support::{mark_skipped_local_openai_chat_candidate, LocalOpenAiChatDecisionInput};
#[allow(clippy::too_many_arguments)]
pub(super) async fn build_cross_format_local_openai_chat_decision_payload_for_candidate(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
body_json: &serde_json::Value,
input: &LocalOpenAiChatDecisionInput,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
candidate_index: u32,
candidate_id: &str,
decision_kind: &str,
upstream_is_stream: bool,
transport: &GatewayProviderTransportSnapshot,
provider_api_format: &str,
) -> Option<GatewayControlSyncDecisionResponse> {
let planner_state = PlannerAppState::new(state);
let provider_api_format = provider_api_format.trim().to_ascii_lowercase();
let Some(conversion_kind) =
request_conversion_kind("openai:chat", provider_api_format.as_str())
else {
return None;
};
if !request_pair_allowed_for_transport(transport, "openai:chat", provider_api_format.as_str()) {
let skip_reason =
if request_conversion_requires_enable_flag("openai:chat", provider_api_format.as_str())
&& !transport.provider.enable_format_conversion
{
"format_conversion_disabled"
} else {
"transport_unsupported"
};
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
return None;
}
if !request_conversion_transport_supported(transport, conversion_kind) {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_unsupported",
)
.await;
return None;
}
let resolve_auth = request_conversion_direct_auth(transport, conversion_kind);
let oauth_auth = if resolve_auth.is_none() {
match planner_state
.resolve_local_oauth_request_auth(transport)
.await
{
Ok(Some(LocalResolvedOAuthRequestAuth::Header { name, value })) => Some((name, value)),
Ok(Some(LocalResolvedOAuthRequestAuth::Kiro(_))) => None,
Ok(None) => None,
Err(err) => {
warn!(
trace_id = %trace_id,
provider_type = %transport.provider.provider_type,
provider_api_format = %provider_api_format,
error = ?err,
"gateway local openai chat cross-format oauth auth resolution failed"
);
None
}
}
} else {
None
};
let Some((auth_header, auth_value)) = resolve_auth.or(oauth_auth) else {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_auth_unavailable",
)
.await;
return None;
};
let mapped_model = candidate.selected_provider_model_name.trim().to_string();
if mapped_model.is_empty() {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"mapped_model_missing",
)
.await;
return None;
}
let Some(provider_request_body) = build_cross_format_openai_chat_request_body(
body_json,
&mapped_model,
transport.provider.provider_type.as_str(),
provider_api_format.as_str(),
upstream_is_stream,
transport.endpoint.body_rules.as_ref(),
Some(input.auth_context.api_key_id.as_str()),
) else {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"provider_request_body_missing",
)
.await;
return None;
};
let Some(upstream_url) = build_cross_format_openai_chat_upstream_url(
parts,
transport,
&mapped_model,
provider_api_format.as_str(),
upstream_is_stream,
) else {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"upstream_url_missing",
)
.await;
return None;
};
let mut provider_request_headers = if provider_api_format.starts_with("claude:") {
build_claude_passthrough_headers(
&parts.headers,
&auth_header,
&auth_value,
&BTreeMap::new(),
Some("application/json"),
)
} else {
build_openai_passthrough_headers(
&parts.headers,
&auth_header,
&auth_value,
&BTreeMap::new(),
Some("application/json"),
)
};
if !apply_local_header_rules(
&mut provider_request_headers,
transport.endpoint.header_rules.as_ref(),
&[&auth_header, "content-type"],
&provider_request_body,
Some(body_json),
) {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_header_rules_apply_failed",
)
.await;
return None;
}
apply_codex_openai_cli_special_headers(
&mut provider_request_headers,
&provider_request_body,
&parts.headers,
transport.provider.provider_type.as_str(),
provider_api_format.as_str(),
Some(trace_id),
transport.key.decrypted_auth_config.as_deref(),
);
ensure_upstream_auth_header(&mut provider_request_headers, &auth_header, &auth_value);
if upstream_is_stream {
provider_request_headers
.entry("accept".to_string())
.or_insert_with(|| "text/event-stream".to_string());
}
let report_kind = if decision_kind == OPENAI_CHAT_STREAM_PLAN_KIND {
"openai_chat_stream_success"
} else {
"openai_chat_sync_finalize"
};
let proxy =
resolve_transport_proxy_snapshot_with_tunnel_affinity(planner_state.app(), transport).await;
let tls_profile = resolve_transport_tls_profile(transport);
let prompt_cache_key = provider_request_body
.get("prompt_cache_key")
.and_then(|value| value.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
Some(GatewayControlSyncDecisionResponse {
action: if upstream_is_stream {
crate::ai_pipeline::planner::common::EXECUTION_RUNTIME_STREAM_DECISION_ACTION
.to_string()
} else {
crate::ai_pipeline::planner::common::EXECUTION_RUNTIME_SYNC_DECISION_ACTION.to_string()
},
decision_kind: Some(decision_kind.to_string()),
execution_strategy: Some(ExecutionStrategy::LocalCrossFormat.as_str().to_string()),
conversion_mode: Some(ConversionMode::Bidirectional.as_str().to_string()),
request_id: Some(trace_id.to_string()),
candidate_id: Some(candidate_id.to_string()),
provider_name: Some(transport.provider.name.clone()),
provider_id: Some(candidate.provider_id.clone()),
endpoint_id: Some(candidate.endpoint_id.clone()),
key_id: Some(candidate.key_id.clone()),
upstream_base_url: Some(transport.endpoint.base_url.clone()),
upstream_url: Some(upstream_url.clone()),
provider_request_method: None,
auth_header: Some(auth_header),
auth_value: Some(auth_value),
provider_api_format: Some(provider_api_format.clone()),
client_api_format: Some("openai:chat".to_string()),
provider_contract: Some(provider_api_format.clone()),
client_contract: Some("openai:chat".to_string()),
model_name: Some(input.requested_model.clone()),
mapped_model: Some(mapped_model.clone()),
prompt_cache_key,
extra_headers: BTreeMap::new(),
provider_request_headers: provider_request_headers.clone(),
provider_request_body: Some(provider_request_body.clone()),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
proxy,
tls_profile,
timeouts: resolve_transport_execution_timeouts(transport),
upstream_is_stream,
report_kind: Some(report_kind.to_string()),
report_context: Some(append_local_failover_policy_to_value(
append_execution_contract_fields_to_value(
json!({
"user_id": input.auth_context.user_id,
"api_key_id": input.auth_context.api_key_id,
"username": input.auth_context.username,
"api_key_name": input.auth_context.api_key_name,
"request_id": trace_id,
"candidate_id": candidate_id,
"candidate_index": candidate_index,
"retry_index": 0,
"model": input.requested_model,
"provider_name": transport.provider.name,
"provider_id": candidate.provider_id,
"endpoint_id": candidate.endpoint_id,
"key_id": candidate.key_id,
"key_name": candidate.key_name,
"provider_api_format": provider_api_format,
"client_api_format": "openai:chat",
"mapped_model": mapped_model,
"upstream_url": upstream_url,
"provider_request_method": serde_json::Value::Null,
"provider_request_headers": provider_request_headers,
"original_headers": collect_control_headers(&parts.headers),
"original_request_body": crate::ai_pipeline::build_report_context_original_request_echo(body_json),
"has_envelope": false,
"needs_conversion": true,
}),
ExecutionStrategy::LocalCrossFormat,
ConversionMode::Bidirectional,
"openai:chat",
provider_api_format.as_str(),
),
transport,
)),
auth_context: Some(input.auth_context.clone()),
})
}

View File

@@ -0,0 +1,134 @@
use crate::ai_pipeline::planner::payload_metadata::{
build_local_execution_decision_response, LocalExecutionDecisionResponseParts,
};
use crate::ai_pipeline::planner::report_context::{
build_local_execution_report_context, LocalExecutionReportContextParts,
};
use crate::ai_pipeline::transport::{
resolve_transport_execution_timeouts, resolve_transport_proxy_snapshot_with_tunnel_affinity,
resolve_transport_tls_profile,
};
use crate::{
append_execution_contract_fields_to_value, append_local_failover_policy_to_value, AppState,
GatewayControlSyncDecisionResponse,
};
use super::request::resolve_local_openai_chat_candidate_payload_parts;
use super::support::{LocalOpenAiChatCandidateAttempt, LocalOpenAiChatDecisionInput};
#[allow(clippy::too_many_arguments)]
pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
body_json: &serde_json::Value,
input: &LocalOpenAiChatDecisionInput,
attempt: LocalOpenAiChatCandidateAttempt,
decision_kind: &str,
report_kind: &str,
upstream_is_stream: bool,
) -> Option<GatewayControlSyncDecisionResponse> {
let LocalOpenAiChatCandidateAttempt {
eligible,
candidate_index,
candidate_id,
} = attempt;
let resolved = resolve_local_openai_chat_candidate_payload_parts(
state,
parts,
trace_id,
body_json,
input,
&eligible,
candidate_index,
&candidate_id,
decision_kind,
report_kind,
upstream_is_stream,
)
.await?;
let candidate = &eligible.candidate;
let prompt_cache_key = resolved
.provider_request_body
.get("prompt_cache_key")
.and_then(|value| value.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let proxy =
resolve_transport_proxy_snapshot_with_tunnel_affinity(state, &resolved.transport).await;
let tls_profile = resolve_transport_tls_profile(&resolved.transport);
let timeouts = resolve_transport_execution_timeouts(&resolved.transport);
Some(build_local_execution_decision_response(
LocalExecutionDecisionResponseParts {
decision_is_stream: upstream_is_stream,
decision_kind: decision_kind.to_string(),
execution_strategy: resolved.execution_strategy,
conversion_mode: resolved.conversion_mode,
request_id: trace_id.to_string(),
candidate_id: candidate_id.clone(),
provider_name: resolved.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: resolved.transport.endpoint.base_url.clone(),
upstream_url: resolved.upstream_url.clone(),
provider_request_method: None,
auth_header: Some(resolved.auth_header.clone()),
auth_value: Some(resolved.auth_value.clone()),
provider_api_format: resolved.provider_api_format.clone(),
client_api_format: "openai:chat".to_string(),
model_name: input.requested_model.clone(),
mapped_model: resolved.mapped_model.clone(),
prompt_cache_key,
provider_request_headers: resolved.provider_request_headers.clone(),
provider_request_body: Some(resolved.provider_request_body.clone()),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
proxy,
tls_profile,
timeouts,
upstream_is_stream,
report_kind: Some(resolved.report_kind.clone()),
report_context: Some(append_local_failover_policy_to_value(
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,
candidate_index,
retry_index: 0,
model: &input.requested_model,
provider_name: &resolved.transport.provider.name,
provider_id: &candidate.provider_id,
endpoint_id: &candidate.endpoint_id,
key_id: &candidate.key_id,
key_name: Some(&candidate.key_name),
provider_api_format: &resolved.provider_api_format,
client_api_format: "openai:chat",
mapped_model: Some(&resolved.mapped_model),
upstream_url: Some(&resolved.upstream_url),
provider_request_method: Some(serde_json::Value::Null),
provider_request_headers: Some(&resolved.provider_request_headers),
original_headers: &parts.headers,
original_request_body: body_json,
has_envelope: false,
needs_conversion: matches!(
resolved.conversion_mode,
crate::ai_pipeline::ConversionMode::Bidirectional
),
extra_fields: serde_json::Map::new(),
}),
resolved.execution_strategy,
resolved.conversion_mode,
"openai:chat",
candidate.endpoint_api_format.as_str(),
),
&resolved.transport,
)),
auth_context: input.auth_context.clone(),
},
))
}

View File

@@ -0,0 +1,375 @@
use std::collections::BTreeMap;
use serde_json::Value;
use crate::ai_pipeline::conversion::{
request_conversion_direct_auth, request_conversion_kind, request_conversion_transport_supported,
};
use crate::ai_pipeline::planner::candidate_eligibility::EligibleLocalExecutionCandidate;
use crate::ai_pipeline::planner::candidate_preparation::{
prepare_header_authenticated_candidate, OauthPreparationContext,
};
use crate::ai_pipeline::planner::common::OPENAI_CHAT_STREAM_PLAN_KIND;
use crate::ai_pipeline::planner::standard::{
apply_codex_openai_cli_special_headers, build_cross_format_openai_chat_request_body,
build_cross_format_openai_chat_upstream_url, build_local_openai_chat_request_body,
build_local_openai_chat_upstream_url,
};
use crate::ai_pipeline::transport::apply_local_header_rules;
use crate::ai_pipeline::transport::auth::{
build_claude_passthrough_headers, build_complete_passthrough_headers_with_auth,
build_openai_passthrough_headers, ensure_upstream_auth_header, resolve_local_openai_chat_auth,
};
use crate::ai_pipeline::transport::policy::supports_local_openai_chat_transport;
use crate::ai_pipeline::{ConversionMode, ExecutionStrategy, GatewayProviderTransportSnapshot};
use crate::AppState;
use super::support::{mark_skipped_local_openai_chat_candidate, LocalOpenAiChatDecisionInput};
pub(crate) struct LocalOpenAiChatCandidatePayloadParts {
pub(super) auth_header: String,
pub(super) auth_value: String,
pub(super) mapped_model: String,
pub(super) provider_api_format: String,
pub(super) provider_request_body: Value,
pub(super) provider_request_headers: BTreeMap<String, String>,
pub(super) upstream_url: String,
pub(super) execution_strategy: ExecutionStrategy,
pub(super) conversion_mode: ConversionMode,
pub(super) report_kind: String,
pub(super) transport: GatewayProviderTransportSnapshot,
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
body_json: &serde_json::Value,
input: &LocalOpenAiChatDecisionInput,
eligible: &EligibleLocalExecutionCandidate,
candidate_index: u32,
candidate_id: &str,
decision_kind: &str,
report_kind: &str,
upstream_is_stream: bool,
) -> Option<LocalOpenAiChatCandidatePayloadParts> {
let planner_state = crate::ai_pipeline::PlannerAppState::new(state);
let candidate = &eligible.candidate;
let provider_api_format = eligible.provider_api_format.as_str();
let transport = &eligible.transport;
if provider_api_format == "openai:chat" {
if !supports_local_openai_chat_transport(transport) {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_unsupported",
)
.await;
return None;
}
let prepared_candidate = match prepare_header_authenticated_candidate(
planner_state,
transport,
candidate,
resolve_local_openai_chat_auth(transport),
OauthPreparationContext {
trace_id,
api_format: "openai:chat",
operation: "openai_chat_same_format",
},
)
.await
{
Ok(prepared) => prepared,
Err(skip_reason) => {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
return None;
}
};
let Some(provider_request_body) = build_local_openai_chat_request_body(
body_json,
&prepared_candidate.mapped_model,
upstream_is_stream,
transport.endpoint.body_rules.as_ref(),
) else {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"provider_request_body_missing",
)
.await;
return None;
};
let Some(upstream_url) = build_local_openai_chat_upstream_url(parts, transport) else {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"upstream_url_missing",
)
.await;
return None;
};
let mut provider_request_headers = build_complete_passthrough_headers_with_auth(
&parts.headers,
&prepared_candidate.auth_header,
&prepared_candidate.auth_value,
&BTreeMap::new(),
Some("application/json"),
);
if !apply_local_header_rules(
&mut provider_request_headers,
transport.endpoint.header_rules.as_ref(),
&[&prepared_candidate.auth_header, "content-type"],
&provider_request_body,
Some(body_json),
) {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_header_rules_apply_failed",
)
.await;
return None;
}
apply_codex_openai_cli_special_headers(
&mut provider_request_headers,
&provider_request_body,
&parts.headers,
transport.provider.provider_type.as_str(),
transport.endpoint.api_format.as_str(),
Some(trace_id),
transport.key.decrypted_auth_config.as_deref(),
);
ensure_upstream_auth_header(
&mut provider_request_headers,
&prepared_candidate.auth_header,
&prepared_candidate.auth_value,
);
if upstream_is_stream {
provider_request_headers
.entry("accept".to_string())
.or_insert_with(|| "text/event-stream".to_string());
}
return Some(LocalOpenAiChatCandidatePayloadParts {
auth_header: prepared_candidate.auth_header,
auth_value: prepared_candidate.auth_value,
mapped_model: prepared_candidate.mapped_model,
provider_api_format: "openai:chat".to_string(),
provider_request_body,
provider_request_headers,
upstream_url,
execution_strategy: ExecutionStrategy::LocalSameFormat,
conversion_mode: ConversionMode::None,
report_kind: report_kind.to_string(),
transport: transport.clone(),
});
}
let provider_api_format = provider_api_format.trim().to_ascii_lowercase();
let Some(conversion_kind) =
request_conversion_kind("openai:chat", provider_api_format.as_str())
else {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_unsupported",
)
.await;
return None;
};
if !request_conversion_transport_supported(transport, conversion_kind) {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_unsupported",
)
.await;
return None;
}
let prepared_candidate = match prepare_header_authenticated_candidate(
planner_state,
transport,
candidate,
request_conversion_direct_auth(transport, conversion_kind),
OauthPreparationContext {
trace_id,
api_format: provider_api_format.as_str(),
operation: "openai_chat_cross_format",
},
)
.await
{
Ok(prepared) => prepared,
Err(skip_reason) => {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
return None;
}
};
let Some(provider_request_body) = build_cross_format_openai_chat_request_body(
body_json,
&prepared_candidate.mapped_model,
transport.provider.provider_type.as_str(),
provider_api_format.as_str(),
upstream_is_stream,
transport.endpoint.body_rules.as_ref(),
Some(input.auth_context.api_key_id.as_str()),
) else {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"provider_request_body_missing",
)
.await;
return None;
};
let Some(upstream_url) = build_cross_format_openai_chat_upstream_url(
parts,
transport,
&prepared_candidate.mapped_model,
provider_api_format.as_str(),
upstream_is_stream,
) else {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"upstream_url_missing",
)
.await;
return None;
};
let mut provider_request_headers = if provider_api_format.starts_with("claude:") {
build_claude_passthrough_headers(
&parts.headers,
&prepared_candidate.auth_header,
&prepared_candidate.auth_value,
&BTreeMap::new(),
Some("application/json"),
)
} else {
build_openai_passthrough_headers(
&parts.headers,
&prepared_candidate.auth_header,
&prepared_candidate.auth_value,
&BTreeMap::new(),
Some("application/json"),
)
};
if !apply_local_header_rules(
&mut provider_request_headers,
transport.endpoint.header_rules.as_ref(),
&[&prepared_candidate.auth_header, "content-type"],
&provider_request_body,
Some(body_json),
) {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_header_rules_apply_failed",
)
.await;
return None;
}
apply_codex_openai_cli_special_headers(
&mut provider_request_headers,
&provider_request_body,
&parts.headers,
transport.provider.provider_type.as_str(),
provider_api_format.as_str(),
Some(trace_id),
transport.key.decrypted_auth_config.as_deref(),
);
ensure_upstream_auth_header(
&mut provider_request_headers,
&prepared_candidate.auth_header,
&prepared_candidate.auth_value,
);
if upstream_is_stream {
provider_request_headers
.entry("accept".to_string())
.or_insert_with(|| "text/event-stream".to_string());
}
let resolved_report_kind = if decision_kind == OPENAI_CHAT_STREAM_PLAN_KIND {
"openai_chat_stream_success".to_string()
} else {
"openai_chat_sync_finalize".to_string()
};
Some(LocalOpenAiChatCandidatePayloadParts {
auth_header: prepared_candidate.auth_header,
auth_value: prepared_candidate.auth_value,
mapped_model: prepared_candidate.mapped_model,
provider_api_format,
provider_request_body,
provider_request_headers,
upstream_url,
execution_strategy: ExecutionStrategy::LocalCrossFormat,
conversion_mode: ConversionMode::Bidirectional,
report_kind: resolved_report_kind,
transport: transport.clone(),
})
}

View File

@@ -1,271 +0,0 @@
use std::collections::BTreeMap;
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use serde_json::json;
use tracing::warn;
use crate::ai_pipeline::planner::standard::{
apply_codex_openai_cli_special_headers, build_local_openai_chat_request_body,
build_local_openai_chat_upstream_url,
};
use crate::ai_pipeline::transport::auth::{
build_complete_passthrough_headers_with_auth, ensure_upstream_auth_header,
resolve_local_openai_chat_auth,
};
use crate::ai_pipeline::transport::policy::supports_local_openai_chat_transport;
use crate::ai_pipeline::transport::{
apply_local_header_rules, resolve_transport_execution_timeouts,
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_transport_tls_profile,
};
use crate::ai_pipeline::{
collect_control_headers, ConversionMode, ExecutionStrategy, PlannerAppState,
};
use crate::ai_pipeline::{GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth};
use crate::{
append_execution_contract_fields_to_value, append_local_failover_policy_to_value, AppState,
GatewayControlSyncDecisionResponse,
};
use super::support::{mark_skipped_local_openai_chat_candidate, LocalOpenAiChatDecisionInput};
#[allow(clippy::too_many_arguments)]
pub(super) async fn build_same_format_local_openai_chat_decision_payload_for_candidate(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
body_json: &serde_json::Value,
input: &LocalOpenAiChatDecisionInput,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
candidate_index: u32,
candidate_id: &str,
decision_kind: &str,
report_kind: &str,
upstream_is_stream: bool,
transport: &GatewayProviderTransportSnapshot,
) -> Option<GatewayControlSyncDecisionResponse> {
let planner_state = PlannerAppState::new(state);
if !supports_local_openai_chat_transport(transport) {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_unsupported",
)
.await;
return None;
}
let oauth_auth = if resolve_local_openai_chat_auth(transport).is_none() {
match planner_state
.resolve_local_oauth_request_auth(transport)
.await
{
Ok(Some(LocalResolvedOAuthRequestAuth::Header { name, value })) => Some((name, value)),
Ok(Some(LocalResolvedOAuthRequestAuth::Kiro(_))) => None,
Ok(None) => None,
Err(err) => {
warn!(
trace_id = %trace_id,
provider_type = %transport.provider.provider_type,
error = ?err,
"gateway local openai chat oauth auth resolution failed"
);
None
}
}
} else {
None
};
let Some((auth_header, auth_value)) = resolve_local_openai_chat_auth(transport).or(oauth_auth)
else {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_auth_unavailable",
)
.await;
return None;
};
let mapped_model = candidate.selected_provider_model_name.trim().to_string();
if mapped_model.is_empty() {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"mapped_model_missing",
)
.await;
return None;
}
let Some(provider_request_body) = build_local_openai_chat_request_body(
body_json,
&mapped_model,
upstream_is_stream,
transport.endpoint.body_rules.as_ref(),
) else {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"provider_request_body_missing",
)
.await;
return None;
};
let Some(upstream_url) = build_local_openai_chat_upstream_url(parts, transport) else {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"upstream_url_missing",
)
.await;
return None;
};
let mut provider_request_headers = build_complete_passthrough_headers_with_auth(
&parts.headers,
&auth_header,
&auth_value,
&BTreeMap::new(),
Some("application/json"),
);
if !apply_local_header_rules(
&mut provider_request_headers,
transport.endpoint.header_rules.as_ref(),
&[&auth_header, "content-type"],
&provider_request_body,
Some(body_json),
) {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_header_rules_apply_failed",
)
.await;
return None;
}
apply_codex_openai_cli_special_headers(
&mut provider_request_headers,
&provider_request_body,
&parts.headers,
transport.provider.provider_type.as_str(),
transport.endpoint.api_format.as_str(),
Some(trace_id),
transport.key.decrypted_auth_config.as_deref(),
);
ensure_upstream_auth_header(&mut provider_request_headers, &auth_header, &auth_value);
if upstream_is_stream {
provider_request_headers
.entry("accept".to_string())
.or_insert_with(|| "text/event-stream".to_string());
}
let proxy =
resolve_transport_proxy_snapshot_with_tunnel_affinity(planner_state.app(), transport).await;
let tls_profile = resolve_transport_tls_profile(transport);
let prompt_cache_key = provider_request_body
.get("prompt_cache_key")
.and_then(|value| value.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
Some(GatewayControlSyncDecisionResponse {
action: if upstream_is_stream {
crate::ai_pipeline::planner::common::EXECUTION_RUNTIME_STREAM_DECISION_ACTION
.to_string()
} else {
crate::ai_pipeline::planner::common::EXECUTION_RUNTIME_SYNC_DECISION_ACTION.to_string()
},
decision_kind: Some(decision_kind.to_string()),
execution_strategy: Some(ExecutionStrategy::LocalSameFormat.as_str().to_string()),
conversion_mode: Some(ConversionMode::None.as_str().to_string()),
request_id: Some(trace_id.to_string()),
candidate_id: Some(candidate_id.to_string()),
provider_name: Some(transport.provider.name.clone()),
provider_id: Some(candidate.provider_id.clone()),
endpoint_id: Some(candidate.endpoint_id.clone()),
key_id: Some(candidate.key_id.clone()),
upstream_base_url: Some(transport.endpoint.base_url.clone()),
upstream_url: Some(upstream_url.clone()),
provider_request_method: None,
auth_header: Some(auth_header),
auth_value: Some(auth_value),
provider_api_format: Some("openai:chat".to_string()),
client_api_format: Some("openai:chat".to_string()),
provider_contract: Some("openai:chat".to_string()),
client_contract: Some("openai:chat".to_string()),
model_name: Some(input.requested_model.clone()),
mapped_model: Some(mapped_model.clone()),
prompt_cache_key,
extra_headers: BTreeMap::new(),
provider_request_headers: provider_request_headers.clone(),
provider_request_body: Some(provider_request_body.clone()),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
proxy,
tls_profile,
timeouts: resolve_transport_execution_timeouts(transport),
upstream_is_stream,
report_kind: Some(report_kind.to_string()),
report_context: Some(append_local_failover_policy_to_value(
append_execution_contract_fields_to_value(
json!({
"user_id": input.auth_context.user_id,
"api_key_id": input.auth_context.api_key_id,
"username": input.auth_context.username,
"api_key_name": input.auth_context.api_key_name,
"request_id": trace_id,
"candidate_id": candidate_id,
"candidate_index": candidate_index,
"retry_index": 0,
"model": input.requested_model,
"provider_name": transport.provider.name,
"provider_id": candidate.provider_id,
"endpoint_id": candidate.endpoint_id,
"key_id": candidate.key_id,
"key_name": candidate.key_name,
"provider_api_format": "openai:chat",
"client_api_format": "openai:chat",
"mapped_model": mapped_model,
"upstream_url": upstream_url,
"provider_request_method": serde_json::Value::Null,
"provider_request_headers": provider_request_headers,
"original_headers": collect_control_headers(&parts.headers),
"original_request_body": crate::ai_pipeline::build_report_context_original_request_echo(body_json),
"has_envelope": false,
"needs_conversion": false,
}),
ExecutionStrategy::LocalSameFormat,
ConversionMode::None,
"openai:chat",
"openai:chat",
),
transport,
)),
auth_context: Some(input.auth_context.clone()),
})
}

View File

@@ -1,34 +1,24 @@
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use serde_json::json;
use uuid::Uuid;
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::conversion::{
request_conversion_kind, request_conversion_requires_enable_flag,
request_pair_allowed_for_transport,
use crate::ai_pipeline::planner::candidate_eligibility::filter_and_rank_local_execution_candidates;
use crate::ai_pipeline::planner::candidate_materialization::{
mark_skipped_local_execution_candidate,
persist_available_local_execution_candidates_with_context,
persist_skipped_local_execution_candidates_with_context,
remember_first_local_candidate_affinity,
};
use crate::ai_pipeline::planner::candidate_affinity::{
rank_local_execution_candidates, remember_scheduler_affinity_for_candidate,
use crate::ai_pipeline::planner::candidate_metadata::{
build_local_execution_candidate_contract_metadata, LocalExecutionCandidateMetadataParts,
};
use crate::ai_pipeline::planner::materialization_policy::{
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
};
use crate::ai_pipeline::GatewayAuthApiKeySnapshot;
use crate::ai_pipeline::{ConversionMode, ExecutionStrategy, PlannerAppState};
use crate::clock::{current_unix_ms, current_unix_secs};
use crate::{append_execution_contract_fields_to_value, AppState};
use crate::AppState;
#[derive(Debug, Clone)]
pub(crate) struct LocalOpenAiChatDecisionInput {
pub(crate) auth_context: ExecutionRuntimeAuthContext,
pub(crate) requested_model: String,
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
pub(crate) required_capabilities: Option<serde_json::Value>,
}
#[derive(Debug, Clone)]
pub(crate) struct LocalOpenAiChatCandidateAttempt {
pub(crate) candidate: SchedulerMinimalCandidateSelectionCandidate,
pub(crate) candidate_index: u32,
pub(crate) candidate_id: String,
}
pub(crate) use crate::ai_pipeline::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalOpenAiChatCandidateAttempt;
pub(crate) use crate::ai_pipeline::planner::decision_input::LocalRequestedModelDecisionInput as LocalOpenAiChatDecisionInput;
pub(crate) async fn mark_skipped_local_openai_chat_candidate(
state: &AppState,
@@ -39,28 +29,22 @@ pub(crate) async fn mark_skipped_local_openai_chat_candidate(
candidate_id: &str,
skip_reason: &'static str,
) {
let planner_state = PlannerAppState::new(state);
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
*diagnostic
.skip_reasons
.entry(skip_reason.to_string())
.or_insert(0) += 1;
*diagnostic.skipped_candidate_count.get_or_insert(0) += 1;
});
planner_state
.persist_skipped_local_candidate(
trace_id,
&input.auth_context.user_id,
&input.auth_context.api_key_id,
candidate,
candidate_index,
candidate_id,
input.required_capabilities.as_ref(),
skip_reason,
current_unix_ms(),
"gateway local openai chat decision failed to persist skipped candidate",
)
.await;
let auth_context: &ExecutionRuntimeAuthContext = &input.auth_context;
let persistence_policy = build_local_candidate_persistence_policy(
auth_context,
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::OpenAiChatDecision,
);
mark_skipped_local_execution_candidate(
state,
trace_id,
persistence_policy.skipped,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
}
pub(crate) async fn materialize_local_openai_chat_candidate_attempts(
@@ -70,118 +54,65 @@ pub(crate) async fn materialize_local_openai_chat_candidate_attempts(
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
) -> Vec<LocalOpenAiChatCandidateAttempt> {
let planner_state = PlannerAppState::new(state);
let candidates = rank_local_execution_candidates(
let auth_context: &ExecutionRuntimeAuthContext = &input.auth_context;
let persistence_policy = build_local_candidate_persistence_policy(
auth_context,
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::OpenAiChatDecision,
);
let (candidates, skipped_candidates) = filter_and_rank_local_execution_candidates(
planner_state,
candidates,
"openai:chat",
&input.requested_model,
input.required_capabilities.as_ref(),
)
.await;
let created_at_unix_ms = current_unix_ms();
let mut attempts = Vec::with_capacity(candidates.len());
let mut affinity_remembered = false;
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
let generated_candidate_id = Uuid::new_v4().to_string();
let provider_api_format = candidate.endpoint_api_format.trim().to_ascii_lowercase();
if provider_api_format != "openai:chat" {
if let Ok(Some(transport)) = planner_state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
remember_first_local_candidate_affinity(
planner_state,
Some(&input.auth_snapshot),
"openai:chat",
Some(&input.requested_model),
&candidates,
);
let attempts = persist_available_local_execution_candidates_with_context(
planner_state,
trace_id,
persistence_policy.available,
candidates,
|eligible| {
let provider_api_format = eligible.provider_api_format.clone();
let (execution_strategy, conversion_mode) = if provider_api_format == "openai:chat" {
(ExecutionStrategy::LocalSameFormat, ConversionMode::None)
} else {
(
ExecutionStrategy::LocalCrossFormat,
ConversionMode::Bidirectional,
)
.await
{
if !request_pair_allowed_for_transport(
&transport,
"openai:chat",
provider_api_format.as_str(),
) {
let skip_reason =
if request_conversion_kind("openai:chat", provider_api_format.as_str())
.is_some()
&& request_conversion_requires_enable_flag(
"openai:chat",
provider_api_format.as_str(),
)
&& !transport.provider.enable_format_conversion
{
"format_conversion_disabled"
} else {
"transport_unsupported"
};
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
&candidate,
candidate_index as u32,
&generated_candidate_id,
skip_reason,
)
.await;
continue;
}
}
}
if !affinity_remembered {
remember_scheduler_affinity_for_candidate(
planner_state,
Some(&input.auth_snapshot),
"openai:chat",
&input.requested_model,
&candidate,
);
affinity_remembered = true;
}
let (execution_strategy, conversion_mode) = if provider_api_format == "openai:chat" {
(ExecutionStrategy::LocalSameFormat, ConversionMode::None)
} else {
(
ExecutionStrategy::LocalCrossFormat,
ConversionMode::Bidirectional,
)
};
let extra_data = append_execution_contract_fields_to_value(
json!({
"provider_api_format": provider_api_format,
"client_api_format": "openai:chat",
"global_model_id": candidate.global_model_id.clone(),
"global_model_name": candidate.global_model_name.clone(),
"model_id": candidate.model_id.clone(),
"selected_provider_model_name": candidate.selected_provider_model_name.clone(),
"mapping_matched_model": candidate.mapping_matched_model.clone(),
"provider_name": candidate.provider_name.clone(),
"key_name": candidate.key_name.clone(),
}),
execution_strategy,
conversion_mode,
"openai:chat",
candidate.endpoint_api_format.trim(),
);
};
Some(build_local_execution_candidate_contract_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: provider_api_format.as_str(),
client_api_format: "openai:chat",
extra_fields: serde_json::Map::new(),
},
execution_strategy,
conversion_mode,
eligible.candidate.endpoint_api_format.trim(),
))
},
)
.await;
let candidate_id = planner_state
.persist_available_local_candidate(
trace_id,
&input.auth_context.user_id,
&input.auth_context.api_key_id,
&candidate,
candidate_index as u32,
&generated_candidate_id,
input.required_capabilities.as_ref(),
Some(extra_data),
created_at_unix_ms,
"gateway local openai chat decision request candidate upsert failed",
)
.await;
attempts.push(LocalOpenAiChatCandidateAttempt {
candidate,
candidate_index: candidate_index as u32,
candidate_id,
});
}
persist_skipped_local_execution_candidates_with_context(
state,
trace_id,
persistence_policy.skipped,
attempts.len() as u32,
skipped_candidates,
)
.await;
attempts
}

View File

@@ -4,10 +4,9 @@ use tracing::warn;
use crate::ai_pipeline::planner::common::{
OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND,
};
use crate::ai_pipeline::planner::runtime_miss::set_local_runtime_execution_exhausted_diagnostic;
use crate::ai_pipeline::GatewayControlDecision;
use crate::{
AppState, GatewayControlSyncDecisionResponse, GatewayError, LocalExecutionRuntimeMissDiagnostic,
};
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
mod decision;
mod plans;
@@ -17,9 +16,9 @@ use self::decision::{
maybe_build_local_openai_chat_decision_payload_for_candidate, LocalOpenAiChatDecisionInput,
};
use self::plans::{
build_local_openai_chat_miss_diagnostic, build_local_openai_chat_stream_plan_and_reports,
build_local_openai_chat_sync_plan_and_reports, list_local_openai_chat_candidates,
resolve_local_openai_chat_decision_input, set_local_openai_chat_miss_diagnostic,
build_local_openai_chat_stream_plan_and_reports, build_local_openai_chat_sync_plan_and_reports,
list_local_openai_chat_candidates, resolve_local_openai_chat_decision_input,
set_local_openai_chat_miss_diagnostic,
};
pub(crate) async fn build_local_openai_chat_sync_plan_and_reports_for_kind(
@@ -70,17 +69,13 @@ pub(crate) fn set_local_openai_chat_execution_exhausted_diagnostic(
model = body_json.get("model").and_then(|value| value.as_str()).unwrap_or(""),
"gateway local openai chat execution exhausted all candidates"
);
state.set_local_execution_runtime_miss_diagnostic(
set_local_runtime_execution_exhausted_diagnostic(
state,
trace_id,
LocalExecutionRuntimeMissDiagnostic {
candidate_count: Some(plan_count),
..build_local_openai_chat_miss_diagnostic(
decision,
plan_kind,
body_json.get("model").and_then(|value| value.as_str()),
"execution_runtime_candidates_exhausted",
)
},
decision,
plan_kind,
body_json.get("model").and_then(|value| value.as_str()),
plan_count,
);
}

View File

@@ -10,9 +10,7 @@ mod stream;
mod sync;
pub(super) use self::candidates::list_local_openai_chat_candidates;
pub(super) use self::diagnostic::{
build_local_openai_chat_miss_diagnostic, set_local_openai_chat_miss_diagnostic,
};
pub(super) use self::diagnostic::set_local_openai_chat_miss_diagnostic;
pub(super) use self::resolve::resolve_local_openai_chat_decision_input;
pub(super) use self::stream::build_local_openai_chat_stream_plan_and_reports;
pub(super) use self::sync::build_local_openai_chat_sync_plan_and_reports;

View File

@@ -4,7 +4,8 @@ use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use super::super::{GatewayError, LocalOpenAiChatDecisionInput};
use crate::ai_pipeline::conversion::request_candidate_api_formats;
use crate::ai_pipeline::{GatewayAuthApiKeySnapshot, PlannerAppState};
use crate::ai_pipeline::planner::candidate_source::auth_snapshot_allows_cross_format_candidate;
use crate::ai_pipeline::PlannerAppState;
use crate::clock::current_unix_secs;
use crate::AppState;
@@ -38,7 +39,7 @@ pub(crate) async fn list_local_openai_chat_candidates(
.await?;
if api_format != "openai:chat" {
candidates.retain(|candidate| {
auth_snapshot_allows_cross_format_openai_chat_candidate(
auth_snapshot_allows_cross_format_candidate(
&input.auth_snapshot,
&input.requested_model,
candidate,
@@ -62,34 +63,3 @@ pub(crate) async fn list_local_openai_chat_candidates(
Ok(combined)
}
fn auth_snapshot_allows_cross_format_openai_chat_candidate(
auth_snapshot: &GatewayAuthApiKeySnapshot,
requested_model: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> bool {
if let Some(allowed_providers) = auth_snapshot.effective_allowed_providers() {
let provider_allowed = allowed_providers.iter().any(|value| {
value
.trim()
.eq_ignore_ascii_case(candidate.provider_id.trim())
|| value
.trim()
.eq_ignore_ascii_case(candidate.provider_name.trim())
});
if !provider_allowed {
return false;
}
}
if let Some(allowed_models) = auth_snapshot.effective_allowed_models() {
let model_allowed = allowed_models
.iter()
.any(|value| value == requested_model || value == &candidate.global_model_name);
if !model_allowed {
return false;
}
}
true
}

View File

@@ -1,27 +1,9 @@
use std::collections::BTreeMap;
use super::super::{GatewayControlDecision, LocalExecutionRuntimeMissDiagnostic};
use super::super::GatewayControlDecision;
use crate::ai_pipeline::planner::runtime_miss::{
set_local_runtime_candidate_evaluation_diagnostic, set_local_runtime_miss_diagnostic_reason,
};
use crate::AppState;
pub(crate) fn build_local_openai_chat_miss_diagnostic(
decision: &GatewayControlDecision,
plan_kind: &str,
requested_model: Option<&str>,
reason: &str,
) -> LocalExecutionRuntimeMissDiagnostic {
LocalExecutionRuntimeMissDiagnostic {
reason: reason.to_string(),
route_family: decision.route_family.clone(),
route_kind: decision.route_kind.clone(),
public_path: Some(decision.public_path.clone()),
plan_kind: Some(plan_kind.to_string()),
requested_model: requested_model.map(ToOwned::to_owned),
candidate_count: None,
skipped_candidate_count: None,
skip_reasons: BTreeMap::new(),
}
}
pub(crate) fn set_local_openai_chat_miss_diagnostic(
state: &AppState,
trace_id: &str,
@@ -30,8 +12,30 @@ pub(crate) fn set_local_openai_chat_miss_diagnostic(
requested_model: Option<&str>,
reason: &str,
) {
state.set_local_execution_runtime_miss_diagnostic(
set_local_runtime_miss_diagnostic_reason(
state,
trace_id,
build_local_openai_chat_miss_diagnostic(decision, plan_kind, requested_model, reason),
decision,
plan_kind,
requested_model,
reason,
);
}
pub(crate) fn set_local_openai_chat_candidate_evaluation_diagnostic(
state: &AppState,
trace_id: &str,
decision: &GatewayControlDecision,
plan_kind: &str,
requested_model: Option<&str>,
candidate_count: usize,
) {
set_local_runtime_candidate_evaluation_diagnostic(
state,
trace_id,
decision,
plan_kind,
requested_model,
candidate_count,
);
}

View File

@@ -2,8 +2,11 @@ use tracing::warn;
use super::super::{GatewayControlDecision, LocalOpenAiChatDecisionInput};
use super::diagnostic::set_local_openai_chat_miss_diagnostic;
use crate::ai_pipeline::{resolve_local_decision_execution_runtime_auth_context, PlannerAppState};
use crate::clock::current_unix_secs;
use crate::ai_pipeline::planner::common::extract_standard_requested_model;
use crate::ai_pipeline::planner::decision_input::{
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
};
use crate::ai_pipeline::resolve_local_decision_execution_runtime_auth_context;
use crate::AppState;
pub(crate) async fn resolve_local_openai_chat_decision_input(
@@ -14,7 +17,6 @@ pub(crate) async fn resolve_local_openai_chat_decision_input(
plan_kind: &str,
record_miss_diagnostic: bool,
) -> Option<LocalOpenAiChatDecisionInput> {
let planner_state = PlannerAppState::new(state);
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
warn!(
trace_id = %trace_id,
@@ -29,20 +31,14 @@ pub(crate) async fn resolve_local_openai_chat_decision_input(
trace_id,
decision,
plan_kind,
body_json.get("model").and_then(|value| value.as_str()),
extract_standard_requested_model(body_json).as_deref(),
"missing_auth_context",
);
}
return None;
};
let Some(requested_model) = body_json
.get("model")
.and_then(|value| value.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
else {
let Some(requested_model) = extract_standard_requested_model(body_json) else {
warn!(
trace_id = %trace_id,
"gateway local openai chat decision skipped: missing_requested_model"
@@ -60,15 +56,15 @@ pub(crate) async fn resolve_local_openai_chat_decision_input(
return None;
};
let auth_snapshot = match planner_state
.read_auth_api_key_snapshot(
&auth_context.user_id,
&auth_context.api_key_id,
current_unix_secs(),
)
.await
let resolved_input = match resolve_local_authenticated_decision_input(
state,
auth_context.clone(),
Some(requested_model.as_str()),
None,
)
.await
{
Ok(Some(snapshot)) => snapshot,
Ok(Some(resolved_input)) => resolved_input,
Ok(None) => {
warn!(
trace_id = %trace_id,
@@ -108,19 +104,8 @@ pub(crate) async fn resolve_local_openai_chat_decision_input(
}
};
let required_capabilities = planner_state
.resolve_request_candidate_required_capabilities(
&auth_context.user_id,
&auth_context.api_key_id,
Some(requested_model.as_str()),
None,
)
.await;
Some(LocalOpenAiChatDecisionInput {
auth_context,
Some(build_local_requested_model_decision_input(
resolved_input,
requested_model,
auth_snapshot,
required_capabilities,
})
))
}

View File

@@ -3,17 +3,18 @@ use tracing::warn;
use super::super::{
materialize_local_openai_chat_candidate_attempts,
maybe_build_local_openai_chat_decision_payload_for_candidate, AppState, GatewayControlDecision,
GatewayError, LocalExecutionRuntimeMissDiagnostic,
GatewayError,
};
use super::candidates::list_local_openai_chat_candidates;
use super::diagnostic::{
build_local_openai_chat_miss_diagnostic, set_local_openai_chat_miss_diagnostic,
set_local_openai_chat_candidate_evaluation_diagnostic, set_local_openai_chat_miss_diagnostic,
};
use super::resolve::resolve_local_openai_chat_decision_input;
use crate::ai_pipeline::planner::common::OPENAI_CHAT_STREAM_PLAN_KIND;
use crate::ai_pipeline::planner::plan_builders::{
build_openai_chat_stream_plan_from_decision, LocalStreamPlanAndReport,
};
use crate::ai_pipeline::planner::runtime_miss::apply_local_runtime_candidate_terminal_reason;
pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
state: &AppState,
@@ -55,31 +56,23 @@ pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
}
};
if candidates.is_empty() {
state.set_local_execution_runtime_miss_diagnostic(
set_local_openai_chat_candidate_evaluation_diagnostic(
state,
trace_id,
LocalExecutionRuntimeMissDiagnostic {
candidate_count: Some(0),
..build_local_openai_chat_miss_diagnostic(
decision,
plan_kind,
Some(input.requested_model.as_str()),
"candidate_list_empty",
)
},
decision,
plan_kind,
Some(input.requested_model.as_str()),
0,
);
return Ok(Vec::new());
}
state.set_local_execution_runtime_miss_diagnostic(
set_local_openai_chat_candidate_evaluation_diagnostic(
state,
trace_id,
LocalExecutionRuntimeMissDiagnostic {
candidate_count: Some(candidates.len()),
..build_local_openai_chat_miss_diagnostic(
decision,
plan_kind,
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
)
},
decision,
plan_kind,
Some(input.requested_model.as_str()),
candidates.len(),
);
let attempts =
@@ -116,15 +109,7 @@ pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
}
}
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
let candidate_count = diagnostic.candidate_count.unwrap_or(0);
let skipped_candidate_count = diagnostic.skipped_candidate_count.unwrap_or(0);
diagnostic.reason = if candidate_count > 0 && skipped_candidate_count >= candidate_count {
"all_candidates_skipped".to_string()
} else {
"no_local_stream_plans".to_string()
};
});
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_stream_plans");
Ok(plans)
}

View File

@@ -3,17 +3,18 @@ use tracing::warn;
use super::super::{
materialize_local_openai_chat_candidate_attempts,
maybe_build_local_openai_chat_decision_payload_for_candidate, AppState, GatewayControlDecision,
GatewayError, LocalExecutionRuntimeMissDiagnostic,
GatewayError,
};
use super::candidates::list_local_openai_chat_candidates;
use super::diagnostic::{
build_local_openai_chat_miss_diagnostic, set_local_openai_chat_miss_diagnostic,
set_local_openai_chat_candidate_evaluation_diagnostic, set_local_openai_chat_miss_diagnostic,
};
use super::resolve::resolve_local_openai_chat_decision_input;
use crate::ai_pipeline::planner::common::OPENAI_CHAT_SYNC_PLAN_KIND;
use crate::ai_pipeline::planner::plan_builders::{
build_openai_chat_sync_plan_from_decision, LocalSyncPlanAndReport,
};
use crate::ai_pipeline::planner::runtime_miss::apply_local_runtime_candidate_terminal_reason;
pub(crate) async fn build_local_openai_chat_sync_plan_and_reports(
state: &AppState,
@@ -55,31 +56,23 @@ pub(crate) async fn build_local_openai_chat_sync_plan_and_reports(
}
};
if candidates.is_empty() {
state.set_local_execution_runtime_miss_diagnostic(
set_local_openai_chat_candidate_evaluation_diagnostic(
state,
trace_id,
LocalExecutionRuntimeMissDiagnostic {
candidate_count: Some(0),
..build_local_openai_chat_miss_diagnostic(
decision,
plan_kind,
Some(input.requested_model.as_str()),
"candidate_list_empty",
)
},
decision,
plan_kind,
Some(input.requested_model.as_str()),
0,
);
return Ok(Vec::new());
}
state.set_local_execution_runtime_miss_diagnostic(
set_local_openai_chat_candidate_evaluation_diagnostic(
state,
trace_id,
LocalExecutionRuntimeMissDiagnostic {
candidate_count: Some(candidates.len()),
..build_local_openai_chat_miss_diagnostic(
decision,
plan_kind,
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
)
},
decision,
plan_kind,
Some(input.requested_model.as_str()),
candidates.len(),
);
let attempts =
@@ -116,15 +109,7 @@ pub(crate) async fn build_local_openai_chat_sync_plan_and_reports(
}
}
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
let candidate_count = diagnostic.candidate_count.unwrap_or(0);
let skipped_candidate_count = diagnostic.skipped_candidate_count.unwrap_or(0);
diagnostic.reason = if candidate_count > 0 && skipped_candidate_count >= candidate_count {
"all_candidates_skipped".to_string()
} else {
"no_local_sync_plans".to_string()
};
});
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_sync_plans");
Ok(plans)
}

View File

@@ -1,9 +1,13 @@
use std::collections::BTreeMap;
use serde_json::json;
use tracing::debug;
use crate::ai_pipeline::collect_control_headers;
use crate::ai_pipeline::planner::payload_metadata::{
build_local_execution_decision_response, LocalExecutionDecisionResponseParts,
};
use crate::ai_pipeline::planner::report_context::{
build_local_execution_report_context, LocalExecutionReportContextParts,
};
use crate::ai_pipeline::planner::spec_metadata::local_openai_cli_spec_metadata;
use crate::ai_pipeline::transport::{
resolve_transport_execution_timeouts, resolve_transport_proxy_snapshot_with_tunnel_affinity,
resolve_transport_tls_profile,
@@ -26,8 +30,9 @@ pub(crate) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
attempt: LocalOpenAiCliCandidateAttempt,
spec: LocalOpenAiCliSpec,
) -> Option<GatewayControlSyncDecisionResponse> {
let spec_metadata = local_openai_cli_spec_metadata(spec);
let LocalOpenAiCliCandidateAttempt {
candidate,
eligible,
candidate_index,
candidate_id,
} = attempt;
@@ -37,12 +42,13 @@ pub(crate) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
trace_id,
body_json,
input,
&candidate,
&eligible,
candidate_index,
&candidate_id,
spec,
)
.await?;
let candidate = &eligible.candidate;
let prompt_cache_key = resolved
.provider_request_body
@@ -55,6 +61,10 @@ pub(crate) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
resolve_transport_proxy_snapshot_with_tunnel_affinity(state, &resolved.transport).await;
let tls_profile = resolve_transport_tls_profile(&resolved.transport);
let timeouts = resolve_transport_execution_timeouts(&resolved.transport);
let mut extra_fields = serde_json::Map::new();
if resolved.is_antigravity {
extra_fields.insert("envelope_name".to_string(), json!("antigravity:v1internal"));
}
debug!(
event_name = "local_openai_cli_decision_payload_built",
@@ -66,10 +76,10 @@ pub(crate) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
provider_id = %candidate.provider_id,
endpoint_id = %candidate.endpoint_id,
key_id = %candidate.key_id,
decision_kind = spec.decision_kind,
decision_kind = spec_metadata.decision_kind,
execution_strategy = resolved.execution_strategy.as_str(),
conversion_mode = resolved.conversion_mode.as_str(),
client_api_format = spec.api_format,
client_api_format = spec_metadata.api_format,
provider_api_format = %resolved.provider_api_format,
request_path = %parts.uri.path(),
request_query = ?parts.uri.query(),
@@ -80,84 +90,74 @@ pub(crate) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
"gateway built local openai cli decision payload"
);
Some(GatewayControlSyncDecisionResponse {
action: if spec.require_streaming {
crate::ai_pipeline::planner::common::EXECUTION_RUNTIME_STREAM_DECISION_ACTION
.to_string()
} else {
crate::ai_pipeline::planner::common::EXECUTION_RUNTIME_SYNC_DECISION_ACTION.to_string()
Some(build_local_execution_decision_response(
LocalExecutionDecisionResponseParts {
decision_is_stream: spec_metadata.require_streaming,
decision_kind: spec_metadata.decision_kind.to_string(),
execution_strategy: resolved.execution_strategy,
conversion_mode: resolved.conversion_mode,
request_id: trace_id.to_string(),
candidate_id: candidate_id.clone(),
provider_name: resolved.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: resolved.transport.endpoint.base_url.clone(),
upstream_url: resolved.upstream_url.clone(),
provider_request_method: None,
auth_header: Some(resolved.auth_header.clone()),
auth_value: Some(resolved.auth_value.clone()),
provider_api_format: resolved.provider_api_format.clone(),
client_api_format: spec_metadata.api_format.to_string(),
model_name: input.requested_model.clone(),
mapped_model: resolved.mapped_model.clone(),
prompt_cache_key,
provider_request_headers: resolved.provider_request_headers.clone(),
provider_request_body: Some(resolved.provider_request_body.clone()),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
proxy,
tls_profile,
timeouts,
upstream_is_stream: resolved.upstream_is_stream,
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
report_context: Some(append_local_failover_policy_to_value(
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,
candidate_index,
retry_index: 0,
model: &input.requested_model,
provider_name: &resolved.transport.provider.name,
provider_id: &candidate.provider_id,
endpoint_id: &candidate.endpoint_id,
key_id: &candidate.key_id,
key_name: Some(&candidate.key_name),
provider_api_format: &resolved.provider_api_format,
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),
upstream_url: Some(&resolved.upstream_url),
provider_request_method: Some(serde_json::Value::Null),
provider_request_headers: Some(&resolved.provider_request_headers),
original_headers: &parts.headers,
original_request_body: body_json,
has_envelope: resolved.is_antigravity,
needs_conversion: matches!(
resolved.conversion_mode,
crate::ai_pipeline::ConversionMode::Bidirectional
),
extra_fields,
}),
resolved.execution_strategy,
resolved.conversion_mode,
spec_metadata.api_format,
candidate.endpoint_api_format.as_str(),
),
&resolved.transport,
)),
auth_context: input.auth_context.clone(),
},
decision_kind: Some(spec.decision_kind.to_string()),
execution_strategy: Some(resolved.execution_strategy.as_str().to_string()),
conversion_mode: Some(resolved.conversion_mode.as_str().to_string()),
request_id: Some(trace_id.to_string()),
candidate_id: Some(candidate_id.clone()),
provider_name: Some(resolved.transport.provider.name.clone()),
provider_id: Some(candidate.provider_id.clone()),
endpoint_id: Some(candidate.endpoint_id.clone()),
key_id: Some(candidate.key_id.clone()),
upstream_base_url: Some(resolved.transport.endpoint.base_url.clone()),
upstream_url: Some(resolved.upstream_url.clone()),
provider_request_method: None,
auth_header: Some(resolved.auth_header.clone()),
auth_value: Some(resolved.auth_value.clone()),
provider_api_format: Some(resolved.provider_api_format.clone()),
client_api_format: Some(spec.api_format.to_string()),
provider_contract: Some(resolved.provider_api_format.clone()),
client_contract: Some(spec.api_format.to_string()),
model_name: Some(input.requested_model.clone()),
mapped_model: Some(resolved.mapped_model.clone()),
prompt_cache_key,
extra_headers: BTreeMap::new(),
provider_request_headers: resolved.provider_request_headers.clone(),
provider_request_body: Some(resolved.provider_request_body.clone()),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
proxy,
tls_profile,
timeouts,
upstream_is_stream: resolved.upstream_is_stream,
report_kind: Some(spec.report_kind.to_string()),
report_context: Some(append_local_failover_policy_to_value(
append_execution_contract_fields_to_value(
json!({
"user_id": input.auth_context.user_id,
"api_key_id": input.auth_context.api_key_id,
"username": input.auth_context.username,
"api_key_name": input.auth_context.api_key_name,
"request_id": trace_id,
"candidate_id": candidate_id,
"candidate_index": candidate_index,
"retry_index": 0,
"model": input.requested_model,
"provider_name": resolved.transport.provider.name,
"provider_id": candidate.provider_id,
"endpoint_id": candidate.endpoint_id,
"key_id": candidate.key_id,
"key_name": candidate.key_name,
"provider_api_format": resolved.provider_api_format,
"client_api_format": spec.api_format,
"mapped_model": resolved.mapped_model,
"upstream_url": resolved.upstream_url,
"provider_request_method": serde_json::Value::Null,
"provider_request_headers": resolved.provider_request_headers,
"original_headers": collect_control_headers(&parts.headers),
"original_request_body": crate::ai_pipeline::build_report_context_original_request_echo(body_json),
"has_envelope": resolved.is_antigravity,
"envelope_name": if resolved.is_antigravity {
Some("antigravity:v1internal")
} else {
None
},
"needs_conversion": matches!(resolved.conversion_mode, crate::ai_pipeline::ConversionMode::Bidirectional),
}),
resolved.execution_strategy,
resolved.conversion_mode,
spec.api_format,
candidate.endpoint_api_format.as_str(),
),
&resolved.transport,
)),
auth_context: Some(input.auth_context.clone()),
})
))
}

View File

@@ -1,15 +1,17 @@
use std::collections::BTreeMap;
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use serde_json::Value;
use tracing::{debug, warn};
use tracing::debug;
use crate::ai_pipeline::conversion::{
request_conversion_direct_auth, request_conversion_kind,
request_conversion_requires_enable_flag, request_conversion_transport_supported,
request_pair_allowed_for_transport,
request_conversion_direct_auth, request_conversion_kind, request_conversion_transport_supported,
};
use crate::ai_pipeline::planner::candidate_eligibility::EligibleLocalExecutionCandidate;
use crate::ai_pipeline::planner::candidate_preparation::{
prepare_header_authenticated_candidate, OauthPreparationContext,
};
use crate::ai_pipeline::planner::common::force_upstream_streaming_for_provider;
use crate::ai_pipeline::planner::spec_metadata::local_openai_cli_spec_metadata;
use crate::ai_pipeline::planner::standard::{
apply_codex_openai_cli_special_headers, build_cross_format_openai_cli_request_body,
build_cross_format_openai_cli_upstream_url, build_local_openai_cli_request_body,
@@ -28,9 +30,7 @@ use crate::ai_pipeline::transport::auth::{
};
use crate::ai_pipeline::transport::policy::supports_local_standard_transport_with_network;
use crate::ai_pipeline::{ConversionMode, ExecutionStrategy};
use crate::ai_pipeline::{
GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth, PlannerAppState,
};
use crate::ai_pipeline::{GatewayProviderTransportSnapshot, PlannerAppState};
use crate::AppState;
use super::support::{mark_skipped_local_openai_cli_candidate, LocalOpenAiCliDecisionInput};
@@ -60,100 +60,31 @@ pub(crate) async fn resolve_local_openai_cli_candidate_payload_parts(
trace_id: &str,
body_json: &serde_json::Value,
input: &LocalOpenAiCliDecisionInput,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
eligible: &EligibleLocalExecutionCandidate,
candidate_index: u32,
candidate_id: &str,
spec: LocalOpenAiCliSpec,
) -> Option<LocalOpenAiCliCandidatePayloadParts> {
let spec_metadata = local_openai_cli_spec_metadata(spec);
let client_api_format = spec_metadata.api_format.trim().to_ascii_lowercase();
let planner_state = PlannerAppState::new(state);
let provider_api_format = candidate.endpoint_api_format.trim().to_ascii_lowercase();
let transport = match planner_state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
{
Ok(Some(snapshot)) => snapshot,
Ok(None) => {
mark_skipped_local_openai_cli_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_snapshot_missing",
)
.await;
return None;
}
Err(err) => {
warn!(
trace_id = %trace_id,
api_format = spec.api_format,
error = ?err,
"gateway local openai cli decision provider transport read failed"
);
mark_skipped_local_openai_cli_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_snapshot_read_failed",
)
.await;
return None;
}
};
let candidate = &eligible.candidate;
let provider_api_format = eligible.provider_api_format.as_str();
let transport = &eligible.transport;
let is_antigravity = transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("antigravity");
let same_format = provider_api_format == spec.api_format.trim().to_ascii_lowercase();
let conversion_kind = request_conversion_kind(spec.api_format, provider_api_format.as_str());
if !same_format
&& !request_pair_allowed_for_transport(
&transport,
spec.api_format,
provider_api_format.as_str(),
)
{
let skip_reason = if conversion_kind.is_some()
&& request_conversion_requires_enable_flag(
spec.api_format,
provider_api_format.as_str(),
)
&& !transport.provider.enable_format_conversion
{
"format_conversion_disabled"
} else {
"transport_unsupported"
};
mark_skipped_local_openai_cli_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
return None;
}
let same_format = provider_api_format == client_api_format;
let conversion_kind = request_conversion_kind(spec_metadata.api_format, provider_api_format);
let transport_supported = if same_format {
supports_local_standard_transport_with_network(&transport, provider_api_format.as_str())
supports_local_standard_transport_with_network(transport, provider_api_format)
} else {
match conversion_kind {
Some(_) if is_antigravity && provider_api_format == "gemini:cli" => true,
Some(kind) => request_conversion_transport_supported(&transport, kind),
Some(kind) => request_conversion_transport_supported(transport, kind),
None => false,
}
};
@@ -171,82 +102,62 @@ pub(crate) async fn resolve_local_openai_cli_candidate_payload_parts(
return None;
}
let resolved_auth = if same_format {
match provider_api_format.as_str() {
"gemini:cli" => resolve_local_gemini_auth(&transport),
let direct_auth = if same_format {
match provider_api_format {
"gemini:cli" => resolve_local_gemini_auth(transport),
"claude:cli" | "openai:cli" | "openai:compact" => {
resolve_local_standard_auth(&transport)
resolve_local_standard_auth(transport)
}
_ => None,
}
} else {
conversion_kind.and_then(|kind| request_conversion_direct_auth(&transport, kind))
conversion_kind.and_then(|kind| request_conversion_direct_auth(transport, kind))
};
let oauth_auth = if resolved_auth.is_none() {
match planner_state
.resolve_local_oauth_request_auth(&transport)
.await
{
Ok(Some(LocalResolvedOAuthRequestAuth::Header { name, value })) => Some((name, value)),
Ok(Some(LocalResolvedOAuthRequestAuth::Kiro(_))) => None,
Ok(None) => None,
Err(err) => {
warn!(
trace_id = %trace_id,
api_format = spec.api_format,
provider_type = %transport.provider.provider_type,
error = ?err,
"gateway local openai cli oauth auth resolution failed"
);
None
}
let prepared_candidate = match prepare_header_authenticated_candidate(
planner_state,
transport,
candidate,
direct_auth,
OauthPreparationContext {
trace_id,
api_format: provider_api_format,
operation: "openai_cli_candidate_request",
},
)
.await
{
Ok(prepared) => prepared,
Err(skip_reason) => {
mark_skipped_local_openai_cli_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
return None;
}
} else {
None
};
let Some((auth_header, auth_value)) = resolved_auth.or(oauth_auth) else {
mark_skipped_local_openai_cli_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_auth_unavailable",
)
.await;
return None;
};
let mapped_model = candidate.selected_provider_model_name.trim().to_string();
if mapped_model.is_empty() {
mark_skipped_local_openai_cli_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"mapped_model_missing",
)
.await;
return None;
}
let auth_header = prepared_candidate.auth_header;
let auth_value = prepared_candidate.auth_value;
let mapped_model = prepared_candidate.mapped_model;
let needs_bidirectional_conversion = !same_format && conversion_kind.is_some();
let upstream_is_stream = spec.require_streaming
let upstream_is_stream = spec_metadata.require_streaming
|| is_antigravity
|| force_upstream_streaming_for_provider(
transport.provider.provider_type.as_str(),
provider_api_format.as_str(),
provider_api_format,
);
let Some(base_provider_request_body) = (if needs_bidirectional_conversion {
build_cross_format_openai_cli_request_body(
body_json,
&mapped_model,
spec.api_format,
provider_api_format.as_str(),
spec_metadata.api_format,
provider_api_format,
upstream_is_stream,
transport.provider.provider_type.as_str(),
transport.endpoint.body_rules.as_ref(),
@@ -258,7 +169,7 @@ pub(crate) async fn resolve_local_openai_cli_candidate_payload_parts(
&mapped_model,
upstream_is_stream,
transport.provider.provider_type.as_str(),
provider_api_format.as_str(),
provider_api_format,
transport.endpoint.body_rules.as_ref(),
Some(input.auth_context.api_key_id.as_str()),
)
@@ -277,7 +188,7 @@ pub(crate) async fn resolve_local_openai_cli_candidate_payload_parts(
};
let antigravity_auth = if is_antigravity {
match classify_local_antigravity_request_support(
&transport,
transport,
&base_provider_request_body,
AntigravityEnvelopeRequestType::Agent,
) {
@@ -329,17 +240,17 @@ pub(crate) async fn resolve_local_openai_cli_candidate_payload_parts(
let Some(upstream_url) = (if needs_bidirectional_conversion {
build_cross_format_openai_cli_upstream_url(
parts,
&transport,
transport,
&mapped_model,
spec.api_format,
provider_api_format.as_str(),
spec_metadata.api_format,
provider_api_format,
upstream_is_stream,
)
} else {
build_local_openai_cli_upstream_url(
parts,
&transport,
provider_api_format.as_str() == "openai:compact",
transport,
provider_api_format == "openai:compact",
)
}) else {
mark_skipped_local_openai_cli_candidate(
@@ -408,7 +319,7 @@ pub(crate) async fn resolve_local_openai_cli_candidate_payload_parts(
&provider_request_body,
&parts.headers,
transport.provider.provider_type.as_str(),
provider_api_format.as_str(),
provider_api_format,
Some(trace_id),
transport.key.decrypted_auth_config.as_deref(),
);
@@ -440,7 +351,7 @@ pub(crate) async fn resolve_local_openai_cli_candidate_payload_parts(
endpoint_id = %candidate.endpoint_id,
key_id = %candidate.key_id,
provider_type = %transport.provider.provider_type,
client_api_format = spec.api_format,
client_api_format = spec_metadata.api_format,
provider_api_format = %provider_api_format,
execution_strategy = execution_strategy.as_str(),
conversion_mode = conversion_mode.as_str(),
@@ -458,7 +369,7 @@ pub(crate) async fn resolve_local_openai_cli_candidate_payload_parts(
auth_header,
auth_value,
mapped_model,
provider_api_format,
provider_api_format: provider_api_format.to_string(),
provider_request_body,
provider_request_headers,
upstream_url,
@@ -467,6 +378,6 @@ pub(crate) async fn resolve_local_openai_cli_candidate_payload_parts(
is_antigravity: is_antigravity
|| antigravity_auth.is_some() && ANTIGRAVITY_ENVELOPE_NAME == "antigravity:v1internal",
upstream_is_stream,
transport,
transport: transport.clone(),
})
}

View File

@@ -1,42 +1,41 @@
use std::collections::BTreeSet;
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use serde_json::json;
use tracing::warn;
use uuid::Uuid;
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::conversion::{
request_candidate_api_formats, request_conversion_kind,
request_conversion_requires_enable_flag, request_pair_allowed_for_transport,
use crate::ai_pipeline::conversion::{request_candidate_api_formats, request_conversion_kind};
use crate::ai_pipeline::planner::candidate_eligibility::filter_and_rank_local_execution_candidates;
use crate::ai_pipeline::planner::candidate_materialization::{
mark_skipped_local_execution_candidate,
persist_available_local_execution_candidates_with_context,
persist_skipped_local_execution_candidates_with_context,
remember_first_local_candidate_affinity,
};
use crate::ai_pipeline::planner::candidate_affinity::{
rank_local_execution_candidates, remember_scheduler_affinity_for_candidate,
use crate::ai_pipeline::planner::candidate_metadata::{
build_local_execution_candidate_contract_metadata, LocalExecutionCandidateMetadataParts,
};
use crate::ai_pipeline::planner::candidate_source::auth_snapshot_allows_cross_format_candidate;
use crate::ai_pipeline::planner::common::extract_standard_requested_model;
use crate::ai_pipeline::planner::decision_input::{
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
};
use crate::ai_pipeline::planner::materialization_policy::{
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
};
use crate::ai_pipeline::planner::spec_metadata::local_openai_cli_spec_metadata;
use crate::ai_pipeline::PlannerAppState;
use crate::ai_pipeline::{
resolve_local_decision_execution_runtime_auth_context, ConversionMode, ExecutionStrategy,
GatewayControlDecision,
};
use crate::ai_pipeline::{GatewayAuthApiKeySnapshot, PlannerAppState};
use crate::clock::{current_unix_ms, current_unix_secs};
use crate::{append_execution_contract_fields_to_value, AppState, GatewayError};
use crate::clock::current_unix_secs;
use crate::{AppState, GatewayError};
use super::LocalOpenAiCliSpec;
#[derive(Debug, Clone)]
pub(crate) struct LocalOpenAiCliDecisionInput {
pub(crate) auth_context: ExecutionRuntimeAuthContext,
pub(crate) requested_model: String,
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
pub(crate) required_capabilities: Option<serde_json::Value>,
}
#[derive(Debug, Clone)]
pub(crate) struct LocalOpenAiCliCandidateAttempt {
pub(crate) candidate: SchedulerMinimalCandidateSelectionCandidate,
pub(crate) candidate_index: u32,
pub(crate) candidate_id: String,
}
pub(crate) use crate::ai_pipeline::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalOpenAiCliCandidateAttempt;
pub(crate) use crate::ai_pipeline::planner::decision_input::LocalRequestedModelDecisionInput as LocalOpenAiCliDecisionInput;
pub(crate) async fn resolve_local_openai_cli_decision_input(
state: &AppState,
@@ -44,27 +43,20 @@ pub(crate) async fn resolve_local_openai_cli_decision_input(
decision: &GatewayControlDecision,
body_json: &serde_json::Value,
) -> Option<LocalOpenAiCliDecisionInput> {
let planner_state = PlannerAppState::new(state);
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
return None;
};
let auth_context: ExecutionRuntimeAuthContext =
resolve_local_decision_execution_runtime_auth_context(decision)?;
let requested_model = body_json
.get("model")
.and_then(|value| value.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)?;
let requested_model = extract_standard_requested_model(body_json)?;
let auth_snapshot = match planner_state
.read_auth_api_key_snapshot(
&auth_context.user_id,
&auth_context.api_key_id,
current_unix_secs(),
)
.await
let resolved_input = match resolve_local_authenticated_decision_input(
state,
auth_context,
Some(requested_model.as_str()),
None,
)
.await
{
Ok(Some(snapshot)) => snapshot,
Ok(Some(resolved_input)) => resolved_input,
Ok(None) => return None,
Err(err) => {
warn!(
@@ -76,21 +68,10 @@ pub(crate) async fn resolve_local_openai_cli_decision_input(
}
};
let required_capabilities = planner_state
.resolve_request_candidate_required_capabilities(
&auth_context.user_id,
&auth_context.api_key_id,
Some(requested_model.as_str()),
None,
)
.await;
Some(LocalOpenAiCliDecisionInput {
auth_context,
Some(build_local_requested_model_decision_input(
resolved_input,
requested_model,
auth_snapshot,
required_capabilities,
})
))
}
pub(crate) async fn materialize_local_openai_cli_candidate_attempts(
@@ -99,13 +80,21 @@ pub(crate) async fn materialize_local_openai_cli_candidate_attempts(
input: &LocalOpenAiCliDecisionInput,
spec: LocalOpenAiCliSpec,
) -> Result<Vec<LocalOpenAiCliCandidateAttempt>, GatewayError> {
let spec_metadata = local_openai_cli_spec_metadata(spec);
let client_api_format = spec_metadata.api_format.trim().to_ascii_lowercase();
let planner_state = PlannerAppState::new(state);
let auth_context: &ExecutionRuntimeAuthContext = &input.auth_context;
let persistence_policy = build_local_candidate_persistence_policy(
auth_context,
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::OpenAiCliDecision,
);
let mut seen_candidates = BTreeSet::new();
let mut candidates = Vec::new();
for candidate_api_format in
request_candidate_api_formats(spec.api_format, spec.require_streaming)
request_candidate_api_formats(spec_metadata.api_format, spec_metadata.require_streaming)
{
let auth_snapshot = if candidate_api_format == spec.api_format {
let auth_snapshot = if candidate_api_format == spec_metadata.api_format {
Some(&input.auth_snapshot)
} else {
None
@@ -114,7 +103,7 @@ pub(crate) async fn materialize_local_openai_cli_candidate_attempts(
.list_selectable_candidates(
candidate_api_format,
&input.requested_model,
spec.require_streaming,
spec_metadata.require_streaming,
input.required_capabilities.as_ref(),
auth_snapshot,
current_unix_secs(),
@@ -122,7 +111,7 @@ pub(crate) async fn materialize_local_openai_cli_candidate_attempts(
.await?;
if auth_snapshot.is_none() {
selected_candidates.retain(|candidate| {
auth_snapshot_allows_cross_format_openai_cli_candidate(
auth_snapshot_allows_cross_format_candidate(
&input.auth_snapshot,
&input.requested_model,
candidate,
@@ -144,157 +133,68 @@ pub(crate) async fn materialize_local_openai_cli_candidate_attempts(
}
}
}
let candidates = rank_local_execution_candidates(
let (candidates, skipped_candidates) = filter_and_rank_local_execution_candidates(
planner_state,
candidates,
spec.api_format,
spec_metadata.api_format,
&input.requested_model,
input.required_capabilities.as_ref(),
)
.await;
let created_at_unix_ms = current_unix_ms();
let mut attempts = Vec::with_capacity(candidates.len());
let mut affinity_remembered = false;
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
let generated_candidate_id = Uuid::new_v4().to_string();
let provider_api_format = candidate.endpoint_api_format.trim().to_ascii_lowercase();
if provider_api_format != spec.api_format.trim().to_ascii_lowercase() {
if let Ok(Some(transport)) = planner_state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
{
if !request_pair_allowed_for_transport(
&transport,
spec.api_format,
provider_api_format.as_str(),
) {
let skip_reason =
if request_conversion_kind(spec.api_format, provider_api_format.as_str())
.is_some()
&& request_conversion_requires_enable_flag(
spec.api_format,
provider_api_format.as_str(),
)
&& !transport.provider.enable_format_conversion
{
"format_conversion_disabled"
} else {
"transport_unsupported"
};
mark_skipped_local_openai_cli_candidate(
state,
input,
trace_id,
&candidate,
candidate_index as u32,
&generated_candidate_id,
skip_reason,
)
.await;
continue;
}
}
}
if !affinity_remembered {
remember_scheduler_affinity_for_candidate(
planner_state,
Some(&input.auth_snapshot),
spec.api_format,
&input.requested_model,
&candidate,
);
affinity_remembered = true;
}
let execution_strategy =
if provider_api_format == spec.api_format.trim().to_ascii_lowercase() {
remember_first_local_candidate_affinity(
planner_state,
Some(&input.auth_snapshot),
spec_metadata.api_format,
Some(&input.requested_model),
&candidates,
);
let attempts = persist_available_local_execution_candidates_with_context(
planner_state,
trace_id,
persistence_policy.available,
candidates,
|eligible| {
let provider_api_format = eligible.provider_api_format.clone();
let execution_strategy = if provider_api_format == client_api_format {
ExecutionStrategy::LocalSameFormat
} else {
ExecutionStrategy::LocalCrossFormat
};
let conversion_mode =
if request_conversion_kind(spec.api_format, provider_api_format.as_str()).is_some() {
ConversionMode::Bidirectional
} else {
ConversionMode::None
};
let extra_data = append_execution_contract_fields_to_value(
json!({
"provider_api_format": provider_api_format,
"client_api_format": spec.api_format,
"global_model_id": candidate.global_model_id.clone(),
"global_model_name": candidate.global_model_name.clone(),
"model_id": candidate.model_id.clone(),
"selected_provider_model_name": candidate.selected_provider_model_name.clone(),
"mapping_matched_model": candidate.mapping_matched_model.clone(),
"provider_name": candidate.provider_name.clone(),
"key_name": candidate.key_name.clone(),
}),
execution_strategy,
conversion_mode,
spec.api_format,
candidate.endpoint_api_format.as_str(),
);
let conversion_mode =
if request_conversion_kind(spec_metadata.api_format, provider_api_format.as_str())
.is_some()
{
ConversionMode::Bidirectional
} else {
ConversionMode::None
};
Some(build_local_execution_candidate_contract_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: provider_api_format.as_str(),
client_api_format: spec_metadata.api_format,
extra_fields: serde_json::Map::new(),
},
execution_strategy,
conversion_mode,
eligible.candidate.endpoint_api_format.as_str(),
))
},
)
.await;
let candidate_id = planner_state
.persist_available_local_candidate(
trace_id,
&input.auth_context.user_id,
&input.auth_context.api_key_id,
&candidate,
candidate_index as u32,
&generated_candidate_id,
input.required_capabilities.as_ref(),
Some(extra_data),
created_at_unix_ms,
"gateway local openai cli decision request candidate upsert failed",
)
.await;
attempts.push(LocalOpenAiCliCandidateAttempt {
candidate,
candidate_index: candidate_index as u32,
candidate_id,
});
}
persist_skipped_local_execution_candidates_with_context(
state,
trace_id,
persistence_policy.skipped,
attempts.len() as u32,
skipped_candidates,
)
.await;
Ok(attempts)
}
fn auth_snapshot_allows_cross_format_openai_cli_candidate(
auth_snapshot: &GatewayAuthApiKeySnapshot,
requested_model: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> bool {
if let Some(allowed_providers) = auth_snapshot.effective_allowed_providers() {
let provider_allowed = allowed_providers.iter().any(|value| {
value
.trim()
.eq_ignore_ascii_case(candidate.provider_id.trim())
|| value
.trim()
.eq_ignore_ascii_case(candidate.provider_name.trim())
});
if !provider_allowed {
return false;
}
}
if let Some(allowed_models) = auth_snapshot.effective_allowed_models() {
let model_allowed = allowed_models
.iter()
.any(|value| value == requested_model || value == &candidate.global_model_name);
if !model_allowed {
return false;
}
}
true
}
pub(crate) async fn mark_skipped_local_openai_cli_candidate(
state: &AppState,
input: &LocalOpenAiCliDecisionInput,
@@ -304,18 +204,20 @@ pub(crate) async fn mark_skipped_local_openai_cli_candidate(
candidate_id: &str,
skip_reason: &'static str,
) {
PlannerAppState::new(state)
.persist_skipped_local_candidate(
trace_id,
&input.auth_context.user_id,
&input.auth_context.api_key_id,
candidate,
candidate_index,
candidate_id,
input.required_capabilities.as_ref(),
skip_reason,
current_unix_ms(),
"gateway local openai cli decision failed to persist skipped candidate",
)
.await;
let auth_context: &ExecutionRuntimeAuthContext = &input.auth_context;
let persistence_policy = build_local_candidate_persistence_policy(
auth_context,
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::OpenAiCliDecision,
);
mark_skipped_local_execution_candidate(
state,
trace_id,
persistence_policy.skipped,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
}

View File

@@ -9,6 +9,7 @@ use crate::ai_pipeline::planner::plan_builders::{
build_openai_cli_stream_plan_from_decision, build_openai_cli_sync_plan_from_decision,
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
};
use crate::ai_pipeline::planner::spec_metadata::local_openai_cli_spec_metadata;
use crate::ai_pipeline::GatewayControlDecision;
pub(crate) use crate::ai_pipeline::{
resolve_openai_cli_stream_spec as resolve_stream_spec,
@@ -24,6 +25,7 @@ pub(super) async fn build_local_sync_plan_and_reports(
body_json: &serde_json::Value,
spec: LocalOpenAiCliSpec,
) -> Result<Vec<LocalSyncPlanAndReport>, GatewayError> {
let spec_metadata = local_openai_cli_spec_metadata(spec);
let Some(input) =
resolve_local_openai_cli_decision_input(state, trace_id, decision, body_json).await
else {
@@ -49,7 +51,7 @@ pub(super) async fn build_local_sync_plan_and_reports(
Err(err) => {
warn!(
trace_id = %trace_id,
api_format = spec.api_format,
api_format = spec_metadata.api_format,
error = ?err,
"gateway local openai cli sync decision plan build failed"
);
@@ -68,6 +70,7 @@ pub(super) async fn build_local_stream_plan_and_reports(
body_json: &serde_json::Value,
spec: LocalOpenAiCliSpec,
) -> Result<Vec<LocalStreamPlanAndReport>, GatewayError> {
let spec_metadata = local_openai_cli_spec_metadata(spec);
let Some(input) =
resolve_local_openai_cli_decision_input(state, trace_id, decision, body_json).await
else {
@@ -93,7 +96,7 @@ pub(super) async fn build_local_stream_plan_and_reports(
Err(err) => {
warn!(
trace_id = %trace_id,
api_format = spec.api_format,
api_format = spec_metadata.api_format,
error = ?err,
"gateway local openai cli stream decision plan build failed"
);

View File

@@ -5,7 +5,10 @@ use axum::http::Uri;
use sha2::{Digest, Sha256};
use url::form_urlencoded;
use crate::headers::{header_value_str, is_json_request};
use crate::{
ai_pipeline::extract_gemini_model_from_path,
headers::{header_value_str, is_json_request},
};
use super::super::GatewayControlDecision;
use super::types::{
@@ -108,20 +111,6 @@ pub(super) fn build_auth_context_cache_key(
))
}
fn extract_gemini_model_from_path(path: &str) -> Option<String> {
let (_, suffix) = path.split_once("/models/")?;
let model = suffix
.split_once(':')
.map(|(value, _)| value)
.unwrap_or(suffix);
let model = model.trim();
if model.is_empty() {
None
} else {
Some(model.to_string())
}
}
fn extract_trusted_auth_headers(headers: &http::HeaderMap) -> Option<GatewayTrustedAuthHeaders> {
if !has_trusted_gateway_marker(headers) {
return None;

View File

@@ -111,6 +111,22 @@ pub(crate) async fn execute_execution_runtime_stream(
state
.usage_runtime
.record_pending(state.data.as_ref(), &lifecycle_seed);
let candidate_started_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status(
state,
&plan,
report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Pending,
status_code: None,
error_type: None,
error_message: None,
latency_ms: None,
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: None,
},
)
.await;
let plan_request_id_for_log = short_request_id(plan.request_id.as_str());
let provider_name = plan.provider_name.as_deref().unwrap_or("-");
let endpoint_id = plan.endpoint_id.as_str();
@@ -142,6 +158,22 @@ pub(crate) async fn execute_execution_runtime_stream(
error = %err,
"gateway in-process stream execution unavailable"
);
let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status(
state,
&plan,
report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: None,
error_type: Some("execution_runtime_unavailable".to_string()),
error_message: Some(format!("{err:?}")),
latency_ms: None,
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
return Ok(None);
}
};
@@ -154,6 +186,7 @@ pub(crate) async fn execute_execution_runtime_stream(
plan_kind,
report_kind,
report_context,
candidate_started_unix_secs,
frame_stream,
)
.await;
@@ -184,6 +217,22 @@ pub(crate) async fn execute_execution_runtime_stream(
error = %err,
"gateway in-process stream execution unavailable"
);
let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status(
state,
&plan,
report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: None,
error_type: Some("execution_runtime_unavailable".to_string()),
error_message: Some(err.to_string()),
latency_ms: None,
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
return Ok(None);
}
};
@@ -196,6 +245,7 @@ pub(crate) async fn execute_execution_runtime_stream(
plan_kind,
report_kind,
report_context,
candidate_started_unix_secs,
frame_stream,
)
.await;
@@ -220,6 +270,22 @@ pub(crate) async fn execute_execution_runtime_stream(
error = ?err,
"gateway remote execution runtime stream unavailable"
);
let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status(
state,
&plan,
report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: None,
error_type: Some("execution_runtime_unavailable".to_string()),
error_message: Some(format!("{err:?}")),
latency_ms: None,
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
return Ok(None);
}
};
@@ -239,7 +305,7 @@ pub(crate) async fn execute_execution_runtime_stream(
response.status()
)),
latency_ms: None,
started_at_unix_ms: Some(terminal_unix_secs),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
@@ -263,6 +329,7 @@ pub(crate) async fn execute_execution_runtime_stream(
plan_kind,
report_kind,
report_context,
candidate_started_unix_secs,
frame_stream,
)
.await;
@@ -375,6 +442,7 @@ async fn execute_stream_from_frame_stream(
plan_kind: &str,
report_kind: Option<String>,
report_context: Option<serde_json::Value>,
candidate_started_unix_secs: u64,
frame_stream: BoxStream<'static, Result<Bytes, IoError>>,
) -> Result<Option<Response<Body>>, GatewayError> {
let request_id = plan.request_id.as_str();
@@ -430,7 +498,7 @@ async fn execute_stream_from_frame_stream(
.to_string(),
),
latency_ms: None,
started_at_unix_ms: Some(terminal_unix_secs),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
@@ -498,7 +566,7 @@ async fn execute_stream_from_frame_stream(
"execution runtime stream returned retryable status {status_code}"
)),
latency_ms: None,
started_at_unix_ms: Some(terminal_unix_secs),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
@@ -539,7 +607,7 @@ async fn execute_stream_from_frame_stream(
"stream decision fell back to control after status {status_code}"
)),
latency_ms: None,
started_at_unix_ms: Some(terminal_unix_secs),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
@@ -576,7 +644,7 @@ async fn execute_stream_from_frame_stream(
"execution runtime stream returned error status {status_code}"
)),
latency_ms: None,
started_at_unix_ms: Some(terminal_unix_secs),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
@@ -899,7 +967,6 @@ async fn execute_stream_from_frame_stream(
}
}
let candidate_started_unix_secs = current_request_candidate_unix_ms();
state.usage_runtime.record_stream_started(
state.data.as_ref(),
&lifecycle_seed,

View File

@@ -107,6 +107,21 @@ pub(crate) async fn execute_execution_runtime_sync(
state
.usage_runtime
.record_pending(state.data.as_ref(), &lifecycle_seed);
record_local_request_candidate_status(
state,
&plan,
report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Pending,
status_code: None,
error_type: None,
error_message: None,
latency_ms: None,
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: None,
},
)
.await;
#[cfg(not(test))]
let result = {
match DirectSyncExecutionRuntime::new()
@@ -129,6 +144,22 @@ pub(crate) async fn execute_execution_runtime_sync(
error = %err,
"gateway in-process sync execution unavailable"
);
let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status(
state,
&plan,
report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: None,
error_type: Some("execution_runtime_unavailable".to_string()),
error_message: Some(err.to_string()),
latency_ms: None,
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
return Ok(None);
}
}
@@ -159,6 +190,22 @@ pub(crate) async fn execute_execution_runtime_sync(
error = %err,
"gateway in-process sync execution unavailable"
);
let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status(
state,
&plan,
report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: None,
error_type: Some("execution_runtime_unavailable".to_string()),
error_message: Some(err.to_string()),
latency_ms: None,
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
return Ok(None);
}
}
@@ -656,6 +703,22 @@ async fn execute_sync_via_remote_execution_runtime(
error = ?err,
"gateway remote execution runtime sync unavailable"
);
let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status(
state,
plan,
report_context,
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: None,
error_type: Some("execution_runtime_unavailable".to_string()),
error_message: Some(format!("{err:?}")),
latency_ms: None,
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
return Ok(RemoteSyncFallbackOutcome::Unavailable);
}
};

File diff suppressed because it is too large Load Diff

View File

@@ -204,7 +204,7 @@ async fn gateway_handles_admin_dashboard_stats_locally_without_proxying_upstream
json!("session-dashboard-stats-admin"),
),
]),
now + chrono::Duration::hours(1),
chrono::Utc::now() + chrono::Duration::hours(1),
);
let session = sample_auth_session(
"admin-auth-1",

View File

@@ -617,6 +617,48 @@ mod tests {
);
}
#[tokio::test]
async fn submit_sync_report_treats_null_error_field_as_success() {
let repository = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
sample_request_candidate("cand-reporting-sync-null-1", "req-reporting-sync-null-1"),
]));
let state = build_test_state(Arc::clone(&repository));
submit_sync_report(
&state,
"trace-reporting-sync-null-1",
GatewaySyncReportRequest {
trace_id: "trace-reporting-sync-null-1".to_string(),
report_kind: "claude_cli_sync_success".to_string(),
report_context: Some(json!({
"request_id": "req-reporting-sync-null-1",
"client_api_format": "claude:cli",
"provider_api_format": "openai:cli"
})),
status_code: 200,
headers: BTreeMap::new(),
body_json: Some(json!({
"id": "resp_1",
"status": "completed",
"error": null
})),
client_body_json: None,
body_base64: None,
telemetry: None,
},
)
.await
.expect("sync report should stay local");
let stored = repository
.list_by_request_id("req-reporting-sync-null-1")
.await
.expect("request candidates should list");
assert_eq!(stored.len(), 1);
assert_eq!(stored[0].status, RequestCandidateStatus::Success);
assert_eq!(stored[0].status_code, Some(200));
}
#[tokio::test]
async fn submit_stream_report_handles_request_id_only_context_locally_when_unique_candidate_exists(
) {