mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
refactor(gateway): 重构 ai pipeline 规划链路
This commit is contained in:
@@ -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"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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");
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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,
|
||||
)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
}))
|
||||
}
|
||||
@@ -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,
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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";
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
})
|
||||
}
|
||||
@@ -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,
|
||||
})
|
||||
}
|
||||
@@ -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),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
})
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
152
apps/aether-gateway/src/ai_pipeline/planner/report_context.rs
Normal file
152
apps/aether-gateway/src/ai_pipeline/planner/report_context.rs
Normal 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)
|
||||
}
|
||||
137
apps/aether-gateway/src/ai_pipeline/planner/runtime_miss.rs
Normal file
137
apps/aether-gateway/src/ai_pipeline/planner/runtime_miss.rs
Normal 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;
|
||||
});
|
||||
}
|
||||
140
apps/aether-gateway/src/ai_pipeline/planner/spec_metadata.rs
Normal file
140
apps/aether-gateway/src/ai_pipeline/planner/spec_metadata.rs
Normal 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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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"
|
||||
);
|
||||
|
||||
@@ -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()),
|
||||
})
|
||||
))
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
})
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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"
|
||||
);
|
||||
|
||||
@@ -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(),
|
||||
},
|
||||
))
|
||||
}
|
||||
|
||||
@@ -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(),
|
||||
),
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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(),
|
||||
})
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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()),
|
||||
})
|
||||
}
|
||||
@@ -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(),
|
||||
},
|
||||
))
|
||||
}
|
||||
@@ -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(),
|
||||
})
|
||||
}
|
||||
@@ -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()),
|
||||
})
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
})
|
||||
))
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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()),
|
||||
})
|
||||
))
|
||||
}
|
||||
|
||||
@@ -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(),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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"
|
||||
);
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
@@ -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",
|
||||
|
||||
@@ -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(
|
||||
) {
|
||||
|
||||
Reference in New Issue
Block a user