refactor: 大规模模块拆分与代码精简,新增 ai-pipeline/data-contracts 独立 crate

- 新增 aether-ai-pipeline 和 aether-data-contracts crate,将 pipeline 逻辑与数据契约从 gateway 中解耦
- 重构 admin handlers:拆分单体模块为 auth/billing/endpoint/features/model/observability/provider/system 等独立子模块
- 合并 chat/cli 重复代码路径:精简 conversion、finalize、planner 中的 sync/chat/cli 分支
- 重构 scheduler/executor/data 层,引入 facade 模式降低模块间耦合
- 移除冗余的 intent 模块,将 plan_fallback/policy/stream_path/sync_path 迁移至 executor
- 前端适配:调整 admin API 调用和 provider 模型测试对话框
This commit is contained in:
fawney19
2026-04-07 02:50:19 +08:00
parent 763ff03a7b
commit 5d96d6673b
732 changed files with 28589 additions and 20662 deletions
@@ -0,0 +1,15 @@
pub(crate) use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::{AppState, GatewayError};
pub(crate) async fn read_auth_api_key_snapshot(
state: &AppState,
user_id: &str,
api_key_id: &str,
now_unix_secs: u64,
) -> Result<Option<GatewayAuthApiKeySnapshot>, GatewayError> {
state
.data
.read_auth_api_key_snapshot(user_id, api_key_id, now_unix_secs)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
@@ -1,8 +1,9 @@
use tracing::warn;
use crate::provider_transport::resolve_transport_proxy_snapshot;
use crate::scheduler::GatewayMinimalCandidateSelectionCandidate;
use crate::ai_pipeline::planner::transport_facade::read_provider_transport_snapshot;
use crate::ai_pipeline::provider_transport_facade::resolve_transport_proxy_snapshot;
use crate::AppState;
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
enum TunnelOwnerAffinityBucket {
@@ -13,8 +14,8 @@ enum TunnelOwnerAffinityBucket {
pub(crate) async fn prefer_local_tunnel_owner_candidates(
state: &AppState,
candidates: Vec<GatewayMinimalCandidateSelectionCandidate>,
) -> Vec<GatewayMinimalCandidateSelectionCandidate> {
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
) -> Vec<SchedulerMinimalCandidateSelectionCandidate> {
let mut ranked = Vec::with_capacity(candidates.len());
for (original_index, candidate) in candidates.into_iter().enumerate() {
let bucket = resolve_candidate_tunnel_owner_affinity(state, &candidate).await;
@@ -29,15 +30,15 @@ pub(crate) async fn prefer_local_tunnel_owner_candidates(
async fn resolve_candidate_tunnel_owner_affinity(
state: &AppState,
candidate: &GatewayMinimalCandidateSelectionCandidate,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> TunnelOwnerAffinityBucket {
let transport = match state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
let transport = match read_provider_transport_snapshot(
state,
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
{
Ok(Some(transport)) => transport,
Ok(None) => return TunnelOwnerAffinityBucket::Neutral,
@@ -101,14 +102,14 @@ async fn resolve_candidate_tunnel_owner_affinity(
mod tests {
use std::time::{SystemTime, UNIX_EPOCH};
use aether_data::repository::provider_catalog::{
InMemoryProviderCatalogReadRepository, StoredProviderCatalogEndpoint,
StoredProviderCatalogKey, StoredProviderCatalogProvider,
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use serde_json::json;
use super::{
prefer_local_tunnel_owner_candidates, AppState, GatewayMinimalCandidateSelectionCandidate,
prefer_local_tunnel_owner_candidates, AppState, SchedulerMinimalCandidateSelectionCandidate,
};
use crate::data::GatewayDataState;
use crate::tunnel::TunnelAttachmentRecord;
@@ -116,8 +117,8 @@ mod tests {
fn sample_candidate(
endpoint_id: &str,
key_id: &str,
) -> GatewayMinimalCandidateSelectionCandidate {
GatewayMinimalCandidateSelectionCandidate {
) -> SchedulerMinimalCandidateSelectionCandidate {
SchedulerMinimalCandidateSelectionCandidate {
provider_id: "provider-1".to_string(),
provider_name: "provider-1".to_string(),
provider_type: "custom".to_string(),
@@ -0,0 +1,60 @@
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use serde_json::Value;
use crate::{AppState, GatewayError};
#[allow(clippy::too_many_arguments)]
pub(crate) async fn persist_available_local_candidate(
state: &AppState,
trace_id: &str,
user_id: &str,
api_key_id: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
candidate_index: u32,
candidate_id: &str,
extra_data: Option<Value>,
created_at_unix_secs: u64,
error_context: &'static str,
) -> String {
crate::request_candidate_runtime::persist_available_local_candidate(
state,
trace_id,
user_id,
api_key_id,
candidate,
candidate_index,
candidate_id,
extra_data,
created_at_unix_secs,
error_context,
)
.await
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn persist_skipped_local_candidate(
state: &AppState,
trace_id: &str,
user_id: &str,
api_key_id: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
candidate_index: u32,
candidate_id: &str,
skip_reason: &str,
finished_at_unix_secs: u64,
error_context: &'static str,
) {
crate::request_candidate_runtime::persist_skipped_local_candidate(
state,
trace_id,
user_id,
api_key_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
finished_at_unix_secs,
error_context,
)
.await
}
@@ -1,7 +1,7 @@
use axum::body::Bytes;
use base64::Engine as _;
pub(crate) use crate::ai_pipeline::contracts::{
use crate::ai_pipeline::control_facade::is_json_request;
pub(crate) use aether_ai_pipeline::contracts::{
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
CLAUDE_CLI_SYNC_PLAN_KIND, EXECUTION_RUNTIME_STREAM_ACTION,
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_ACTION,
@@ -15,25 +15,45 @@ 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::headers::is_json_request;
pub(crate) fn parse_direct_request_body(
parts: &http::request::Parts,
body_bytes: &Bytes,
) -> Option<(serde_json::Value, Option<String>)> {
if is_json_request(&parts.headers) {
if body_bytes.is_empty() {
Some((serde_json::json!({}), None))
} else {
serde_json::from_slice::<serde_json::Value>(body_bytes)
.ok()
.map(|value| (value, None))
}
} else {
Some((
serde_json::json!({}),
(!body_bytes.is_empty())
.then(|| base64::engine::general_purpose::STANDARD.encode(body_bytes)),
))
aether_ai_pipeline::planner::common::parse_direct_request_body(
is_json_request(&parts.headers),
body_bytes.as_ref(),
)
}
pub(crate) fn force_upstream_streaming_for_provider(
provider_type: &str,
provider_api_format: &str,
) -> bool {
provider_type.trim().eq_ignore_ascii_case("codex")
&& provider_api_format
.trim()
.eq_ignore_ascii_case("openai:cli")
}
#[cfg(test)]
mod tests {
use super::force_upstream_streaming_for_provider;
#[test]
fn forces_streaming_for_codex_openai_cli() {
assert!(force_upstream_streaming_for_provider("codex", "openai:cli"));
}
#[test]
fn does_not_force_streaming_for_compact_or_other_provider_types() {
assert!(!force_upstream_streaming_for_provider(
"codex",
"openai:compact"
));
assert!(!force_upstream_streaming_for_provider(
"openai",
"openai:cli"
));
}
}
@@ -1,3 +1,4 @@
use crate::ai_pipeline::control_facade::{GatewayControlAuthContext, GatewayControlDecision};
use crate::ai_pipeline::planner::common::{
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
CLAUDE_CLI_SYNC_PLAN_KIND, EXECUTION_RUNTIME_STREAM_ACTION, EXECUTION_RUNTIME_SYNC_ACTION,
@@ -18,8 +19,6 @@ use crate::ai_pipeline::planner::plan_builders::{
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
};
use crate::control::GatewayControlAuthContext;
use crate::control::GatewayControlDecision;
use crate::{
AppState, GatewayControlPlanResponse, GatewayControlSyncDecisionResponse, GatewayError,
};
@@ -23,7 +23,7 @@ pub(crate) use super::standard::{
maybe_build_sync_local_openai_cli_decision_payload,
maybe_build_sync_local_standard_decision_payload,
};
pub(crate) use crate::scheduler::{
pub(crate) use crate::ai_pipeline::planner::{
resolve_execution_runtime_stream_plan_kind as resolve_stream_plan_kind,
resolve_execution_runtime_sync_plan_kind as resolve_sync_plan_kind,
};
@@ -1,11 +1,11 @@
use std::collections::BTreeMap;
use crate::ai_pipeline::control_facade::GatewayControlDecision;
use crate::ai_pipeline::execution_facade::{ConversionMode, ExecutionStrategy};
use crate::ai_pipeline::planner::common::{
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, OPENAI_VIDEO_CONTENT_PLAN_KIND,
};
use crate::control::GatewayControlDecision;
use crate::execution_runtime::{ConversionMode, ExecutionStrategy};
use crate::scheduler::{
use crate::ai_pipeline::planner::{
is_matching_stream_request, resolve_execution_runtime_stream_plan_kind,
};
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
@@ -2,16 +2,17 @@ use std::collections::BTreeMap;
use url::Url;
use crate::ai_pipeline::control_facade::{
resolve_execution_runtime_auth_context, GatewayControlDecision,
};
use crate::ai_pipeline::execution_facade::{ConversionMode, ExecutionStrategy};
use crate::ai_pipeline::planner::common::{
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_FILES_DELETE_PLAN_KIND,
GEMINI_FILES_GET_PLAN_KIND, GEMINI_FILES_LIST_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
};
use crate::control::resolve_execution_runtime_auth_context;
use crate::control::GatewayControlDecision;
use crate::execution_runtime::{ConversionMode, ExecutionStrategy};
use crate::scheduler::resolve_execution_runtime_sync_plan_kind;
use crate::ai_pipeline::planner::resolve_execution_runtime_sync_plan_kind;
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
pub(crate) async fn maybe_build_sync_decision_payload(
@@ -0,0 +1,16 @@
use aether_contracts::ExecutionPlan;
use serde_json::Value;
use crate::AppState;
pub(crate) async fn mark_unused_local_candidate_items<T, FPlan, FContext>(
state: &AppState,
remaining: Vec<T>,
plan: FPlan,
report_context: FContext,
) where
FPlan: Fn(&T) -> &ExecutionPlan,
FContext: Fn(&T) -> Option<&Value>,
{
crate::executor::mark_unused_local_candidate_items(state, remaining, plan, report_context).await
}
@@ -1,16 +1,28 @@
use crate::ai_pipeline::contracts::{
GatewayControlPlanResponse, GatewayControlSyncDecisionResponse,
};
use crate::control::GatewayControlDecision;
use crate::ai_pipeline::control_facade::GatewayControlDecision;
use crate::{AppState, GatewayError};
pub(crate) mod auth_snapshot_facade;
pub(crate) mod candidate_affinity;
pub(crate) mod candidate_runtime_facade;
pub(crate) mod common;
mod decision;
pub(crate) mod executor_facade;
pub(crate) mod passthrough;
pub(crate) mod plan_builders;
mod route;
pub(crate) mod scheduler_facade;
pub(crate) mod specialized;
pub(crate) mod standard;
pub(crate) mod transport_facade;
pub(crate) use self::route::{
is_matching_stream_request, resolve_execution_runtime_stream_plan_kind,
resolve_execution_runtime_sync_plan_kind, supports_stream_scheduler_decision_kind,
supports_sync_scheduler_decision_kind,
};
pub(crate) async fn maybe_build_sync_decision_payload(
state: &AppState,
@@ -6,4 +6,4 @@ pub(crate) use self::provider::{
maybe_build_stream_local_same_format_provider_decision_payload,
maybe_build_sync_local_same_format_provider_decision_payload,
};
pub(crate) use crate::provider_transport::provider_types::provider_type_supports_local_same_format_transport;
pub(crate) use crate::ai_pipeline::provider_transport_facade::provider_types::provider_type_supports_local_same_format_transport;
@@ -3,11 +3,16 @@ use axum::http::Response;
use std::collections::BTreeMap;
use url::form_urlencoded;
use aether_data::repository::candidates::{RequestCandidateStatus, UpsertRequestCandidateRecord};
use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, UpsertRequestCandidateRecord,
};
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use serde_json::{json, Value};
use tracing::warn;
use uuid::Uuid;
use crate::ai_pipeline::control_facade::{collect_control_headers, GatewayControlDecision};
use crate::ai_pipeline::execution_facade::{ConversionMode, ExecutionStrategy};
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
use crate::ai_pipeline::planner::common::{
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
@@ -15,47 +20,41 @@ use crate::ai_pipeline::planner::common::{
use crate::ai_pipeline::planner::plan_builders::{
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
};
use crate::control::GatewayControlDecision;
use crate::execution_runtime::{ConversionMode, ExecutionStrategy};
use crate::headers::collect_control_headers;
use crate::provider_transport::antigravity::{
use crate::ai_pipeline::provider_transport_facade::antigravity::{
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
build_antigravity_v1internal_url, classify_local_antigravity_request_support,
AntigravityEnvelopeRequestType, AntigravityRequestEnvelopeSupport,
AntigravityRequestSideSupport, AntigravityRequestUrlAction,
};
use crate::provider_transport::auth::{
use crate::ai_pipeline::provider_transport_facade::auth::{
build_openai_passthrough_headers, resolve_local_gemini_auth, resolve_local_standard_auth,
};
use crate::provider_transport::claude_code::{
use crate::ai_pipeline::provider_transport_facade::claude_code::{
build_claude_code_messages_url, build_claude_code_passthrough_headers,
sanitize_claude_code_request_body, supports_local_claude_code_transport_with_network,
};
use crate::provider_transport::kiro::{
use crate::ai_pipeline::provider_transport_facade::kiro::{
build_kiro_generate_assistant_response_url, build_kiro_provider_headers,
build_kiro_provider_request_body, supports_local_kiro_request_transport_with_network,
KIRO_ENVELOPE_NAME,
};
use crate::provider_transport::policy::{
use crate::ai_pipeline::provider_transport_facade::policy::{
supports_local_gemini_transport_with_network, supports_local_standard_transport_with_network,
};
use crate::provider_transport::url::{
use crate::ai_pipeline::provider_transport_facade::url::{
build_claude_messages_url, build_gemini_content_url, build_passthrough_path_url,
};
use crate::provider_transport::vertex::{
use crate::ai_pipeline::provider_transport_facade::vertex::{
build_vertex_api_key_gemini_content_url, resolve_local_vertex_api_key_query_auth,
supports_local_vertex_api_key_gemini_transport_with_network,
};
use crate::provider_transport::{
use crate::ai_pipeline::provider_transport_facade::{
apply_local_body_rules, apply_local_header_rules, build_passthrough_headers,
ensure_upstream_auth_header, resolve_transport_execution_timeouts,
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_transport_tls_profile,
LocalResolvedOAuthRequestAuth,
};
use crate::scheduler::{
current_unix_secs, list_selectable_candidates, record_local_request_candidate_status,
GatewayMinimalCandidateSelectionCandidate,
};
use crate::clock::current_unix_secs;
use crate::{
append_execution_contract_fields_to_value, AppState, GatewayControlSyncDecisionResponse,
GatewayError,
@@ -1,4 +1,4 @@
use crate::control::GatewayControlDecision;
use crate::ai_pipeline::control_facade::GatewayControlDecision;
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
use super::super::plans::{resolve_stream_spec, resolve_sync_spec};
@@ -1,12 +1,14 @@
use aether_data::repository::candidates::{RequestCandidateStatus, UpsertRequestCandidateRecord};
use serde_json::json;
use tracing::warn;
use uuid::Uuid;
use crate::ai_pipeline::control_facade::GatewayControlDecision;
use crate::ai_pipeline::execution_facade::{ConversionMode, ExecutionStrategy};
use crate::ai_pipeline::planner::auth_snapshot_facade::read_auth_api_key_snapshot;
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
use crate::control::GatewayControlDecision;
use crate::execution_runtime::{ConversionMode, ExecutionStrategy};
use crate::scheduler::{current_unix_secs, list_selectable_candidates};
use crate::ai_pipeline::planner::candidate_runtime_facade::persist_available_local_candidate;
use crate::ai_pipeline::planner::scheduler_facade::list_selectable_candidates;
use crate::clock::current_unix_secs;
use crate::{append_execution_contract_fields_to_value, AppState, GatewayError};
use super::types::{
@@ -40,13 +42,13 @@ pub(crate) async fn resolve_local_same_format_provider_decision_input(
}
};
let auth_snapshot = match state
.read_auth_api_key_snapshot(
&auth_context.user_id,
&auth_context.api_key_id,
current_unix_secs(),
)
.await
let auth_snapshot = match read_auth_api_key_snapshot(
state,
&auth_context.user_id,
&auth_context.api_key_id,
current_unix_secs(),
)
.await
{
Ok(Some(snapshot)) => snapshot,
Ok(None) => return None,
@@ -107,47 +109,19 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
spec.api_format,
);
let candidate_id = match state
.upsert_request_candidate(UpsertRequestCandidateRecord {
id: generated_candidate_id.clone(),
request_id: trace_id.to_string(),
user_id: Some(input.auth_context.user_id.clone()),
api_key_id: Some(input.auth_context.api_key_id.clone()),
username: None,
api_key_name: None,
candidate_index: candidate_index as u32,
retry_index: 0,
provider_id: Some(candidate.provider_id.clone()),
endpoint_id: Some(candidate.endpoint_id.clone()),
key_id: Some(candidate.key_id.clone()),
status: RequestCandidateStatus::Available,
skip_reason: None,
is_cached: Some(false),
status_code: None,
error_type: None,
error_message: None,
latency_ms: None,
concurrent_requests: None,
extra_data: Some(extra_data),
required_capabilities: candidate.key_capabilities.clone(),
created_at_unix_secs: Some(created_at_unix_secs),
started_at_unix_secs: None,
finished_at_unix_secs: None,
})
.await
{
Ok(Some(stored)) => stored.id,
Ok(None) => generated_candidate_id.clone(),
Err(err) => {
warn!(
trace_id = %trace_id,
api_format = spec.api_format,
error = ?err,
"gateway local same-format decision request candidate upsert failed"
);
generated_candidate_id.clone()
}
};
let candidate_id = persist_available_local_candidate(
state,
trace_id,
&input.auth_context.user_id,
&input.auth_context.api_key_id,
&candidate,
candidate_index as u32,
&generated_candidate_id,
Some(extra_data),
created_at_unix_secs,
"gateway local same-format decision request candidate upsert failed",
)
.await;
attempts.push(LocalSameFormatProviderCandidateAttempt {
candidate,
@@ -1,43 +1,48 @@
use std::collections::BTreeMap;
use aether_data::repository::candidates::{RequestCandidateStatus, UpsertRequestCandidateRecord};
use serde_json::json;
use tracing::warn;
use crate::execution_runtime::{ConversionMode, ExecutionStrategy};
use crate::headers::collect_control_headers;
use crate::provider_transport::antigravity::{
use crate::ai_pipeline::control_facade::collect_control_headers;
use crate::ai_pipeline::execution_facade::{ConversionMode, ExecutionStrategy};
use crate::ai_pipeline::planner::candidate_runtime_facade::persist_skipped_local_candidate;
use crate::ai_pipeline::planner::transport_facade::{
read_provider_transport_snapshot, resolve_local_oauth_request_auth,
LocalResolvedOAuthRequestAuth,
};
use crate::ai_pipeline::provider_transport_facade::antigravity::{
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
classify_local_antigravity_request_support, AntigravityEnvelopeRequestType,
AntigravityRequestEnvelopeSupport, AntigravityRequestSideSupport,
};
use crate::provider_transport::auth::{
use crate::ai_pipeline::provider_transport_facade::auth::{
build_openai_passthrough_headers, resolve_local_gemini_auth, resolve_local_standard_auth,
};
use crate::provider_transport::claude_code::{
use crate::ai_pipeline::provider_transport_facade::claude_code::{
build_claude_code_passthrough_headers, supports_local_claude_code_transport_with_network,
};
use crate::provider_transport::kiro::{
use crate::ai_pipeline::provider_transport_facade::kiro::{
build_kiro_provider_headers, supports_local_kiro_request_transport_with_network,
KIRO_ENVELOPE_NAME,
};
use crate::provider_transport::policy::{
use crate::ai_pipeline::provider_transport_facade::policy::{
supports_local_gemini_transport_with_network, supports_local_standard_transport_with_network,
};
use crate::provider_transport::vertex::{
use crate::ai_pipeline::provider_transport_facade::vertex::{
resolve_local_vertex_api_key_query_auth,
supports_local_vertex_api_key_gemini_transport_with_network,
};
use crate::provider_transport::{
use crate::ai_pipeline::provider_transport_facade::{
apply_local_header_rules, build_passthrough_headers, ensure_upstream_auth_header,
resolve_transport_execution_timeouts, resolve_transport_proxy_snapshot_with_tunnel_affinity,
resolve_transport_tls_profile, LocalResolvedOAuthRequestAuth,
resolve_transport_tls_profile,
};
use crate::scheduler::{current_unix_secs, GatewayMinimalCandidateSelectionCandidate};
use crate::clock::current_unix_secs;
use crate::{
append_execution_contract_fields_to_value, AppState, GatewayControlSyncDecisionResponse,
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
};
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use super::types::{
LocalSameFormatProviderCandidateAttempt, LocalSameFormatProviderDecisionInput,
@@ -59,13 +64,13 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
candidate_id,
} = attempt;
let transport = match state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
let transport = match read_provider_transport_snapshot(
state,
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
{
Ok(Some(snapshot)) => snapshot,
Ok(None) => {
@@ -169,7 +174,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
&& !is_vertex
&& resolve_local_gemini_auth(&transport).is_none();
let oauth_auth = if should_try_oauth_auth {
match state.resolve_local_oauth_request_auth(&transport).await {
match resolve_local_oauth_request_auth(state, &transport).await {
Ok(Some(LocalResolvedOAuthRequestAuth::Kiro(auth))) => {
Some(LocalResolvedOAuthRequestAuth::Kiro(auth))
}
@@ -528,46 +533,22 @@ pub(super) async fn mark_skipped_local_same_format_provider_candidate(
state: &AppState,
input: &LocalSameFormatProviderDecisionInput,
trace_id: &str,
candidate: &GatewayMinimalCandidateSelectionCandidate,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
candidate_index: u32,
candidate_id: &str,
skip_reason: &'static str,
) {
if let Err(err) = state
.upsert_request_candidate(UpsertRequestCandidateRecord {
id: candidate_id.to_string(),
request_id: trace_id.to_string(),
user_id: Some(input.auth_context.user_id.clone()),
api_key_id: Some(input.auth_context.api_key_id.clone()),
username: None,
api_key_name: None,
candidate_index,
retry_index: 0,
provider_id: Some(candidate.provider_id.clone()),
endpoint_id: Some(candidate.endpoint_id.clone()),
key_id: Some(candidate.key_id.clone()),
status: RequestCandidateStatus::Skipped,
skip_reason: Some(skip_reason.to_string()),
is_cached: Some(false),
status_code: None,
error_type: None,
error_message: None,
latency_ms: None,
concurrent_requests: None,
extra_data: None,
required_capabilities: candidate.key_capabilities.clone(),
created_at_unix_secs: None,
started_at_unix_secs: None,
finished_at_unix_secs: Some(current_unix_secs()),
})
.await
{
warn!(
trace_id = %trace_id,
candidate_id = %candidate_id,
skip_reason,
error = ?err,
"gateway local same-format decision failed to persist skipped candidate"
);
}
persist_skipped_local_candidate(
state,
trace_id,
&input.auth_context.user_id,
&input.auth_context.api_key_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
current_unix_secs(),
"gateway local same-format decision failed to persist skipped candidate",
)
.await;
}
@@ -1,28 +1,20 @@
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum LocalSameFormatProviderFamily {
Standard,
Gemini,
}
use crate::ai_pipeline::control_facade::GatewayControlAuthContext;
use crate::ai_pipeline::planner::auth_snapshot_facade::GatewayAuthApiKeySnapshot;
#[derive(Debug, Clone, Copy)]
pub(crate) struct LocalSameFormatProviderSpec {
pub(crate) api_format: &'static str,
pub(crate) decision_kind: &'static str,
pub(crate) report_kind: &'static str,
pub(crate) family: LocalSameFormatProviderFamily,
pub(crate) require_streaming: bool,
}
pub(crate) use aether_ai_pipeline::planner::passthrough::provider::{
LocalSameFormatProviderFamily, LocalSameFormatProviderSpec,
};
#[derive(Debug, Clone)]
pub(crate) struct LocalSameFormatProviderDecisionInput {
pub(crate) auth_context: crate::control::GatewayControlAuthContext,
pub(crate) auth_context: GatewayControlAuthContext,
pub(crate) requested_model: String,
pub(crate) auth_snapshot: crate::data::auth::GatewayAuthApiKeySnapshot,
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
}
#[derive(Debug, Clone)]
pub(crate) struct LocalSameFormatProviderCandidateAttempt {
pub(crate) candidate: crate::scheduler::GatewayMinimalCandidateSelectionCandidate,
pub(crate) candidate: aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate,
pub(crate) candidate_index: u32,
pub(crate) candidate_id: String,
}
@@ -1,5 +1,9 @@
use tracing::warn;
pub(crate) use aether_ai_pipeline::planner::passthrough::provider::{
resolve_stream_spec, resolve_sync_spec,
};
use super::{
materialize_local_same_format_provider_candidate_attempts,
maybe_build_local_same_format_provider_decision_payload_for_candidate,
@@ -7,84 +11,11 @@ use super::{
GatewayError, LocalSameFormatProviderFamily, LocalSameFormatProviderSpec,
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
};
use crate::ai_pipeline::planner::common::{
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
CLAUDE_CLI_SYNC_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND,
GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND,
};
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,
};
pub(crate) fn resolve_sync_spec(plan_kind: &str) -> Option<LocalSameFormatProviderSpec> {
match plan_kind {
CLAUDE_CHAT_SYNC_PLAN_KIND => Some(LocalSameFormatProviderSpec {
api_format: "claude:chat",
decision_kind: CLAUDE_CHAT_SYNC_PLAN_KIND,
report_kind: "claude_chat_sync_success",
family: LocalSameFormatProviderFamily::Standard,
require_streaming: false,
}),
CLAUDE_CLI_SYNC_PLAN_KIND => Some(LocalSameFormatProviderSpec {
api_format: "claude:cli",
decision_kind: CLAUDE_CLI_SYNC_PLAN_KIND,
report_kind: "claude_cli_sync_success",
family: LocalSameFormatProviderFamily::Standard,
require_streaming: false,
}),
GEMINI_CHAT_SYNC_PLAN_KIND => Some(LocalSameFormatProviderSpec {
api_format: "gemini:chat",
decision_kind: GEMINI_CHAT_SYNC_PLAN_KIND,
report_kind: "gemini_chat_sync_success",
family: LocalSameFormatProviderFamily::Gemini,
require_streaming: false,
}),
GEMINI_CLI_SYNC_PLAN_KIND => Some(LocalSameFormatProviderSpec {
api_format: "gemini:cli",
decision_kind: GEMINI_CLI_SYNC_PLAN_KIND,
report_kind: "gemini_cli_sync_success",
family: LocalSameFormatProviderFamily::Gemini,
require_streaming: false,
}),
_ => None,
}
}
pub(crate) fn resolve_stream_spec(plan_kind: &str) -> Option<LocalSameFormatProviderSpec> {
match plan_kind {
CLAUDE_CHAT_STREAM_PLAN_KIND => Some(LocalSameFormatProviderSpec {
api_format: "claude:chat",
decision_kind: CLAUDE_CHAT_STREAM_PLAN_KIND,
report_kind: "claude_chat_stream_success",
family: LocalSameFormatProviderFamily::Standard,
require_streaming: true,
}),
CLAUDE_CLI_STREAM_PLAN_KIND => Some(LocalSameFormatProviderSpec {
api_format: "claude:cli",
decision_kind: CLAUDE_CLI_STREAM_PLAN_KIND,
report_kind: "claude_cli_stream_success",
family: LocalSameFormatProviderFamily::Standard,
require_streaming: true,
}),
GEMINI_CHAT_STREAM_PLAN_KIND => Some(LocalSameFormatProviderSpec {
api_format: "gemini:chat",
decision_kind: GEMINI_CHAT_STREAM_PLAN_KIND,
report_kind: "gemini_chat_stream_success",
family: LocalSameFormatProviderFamily::Gemini,
require_streaming: true,
}),
GEMINI_CLI_STREAM_PLAN_KIND => Some(LocalSameFormatProviderSpec {
api_format: "gemini:cli",
decision_kind: GEMINI_CLI_STREAM_PLAN_KIND,
report_kind: "gemini_cli_stream_success",
family: LocalSameFormatProviderFamily::Gemini,
require_streaming: true,
}),
_ => None,
}
}
pub(crate) async fn build_local_sync_plan_and_reports(
state: &AppState,
parts: &http::request::Parts,
@@ -3,6 +3,8 @@ use std::collections::BTreeMap;
use serde_json::Value;
use url::form_urlencoded;
use crate::ai_pipeline::planner::transport_facade::GatewayProviderTransportSnapshot;
use super::{
apply_local_body_rules, build_antigravity_v1internal_url, build_claude_code_messages_url,
build_claude_messages_url, build_gemini_content_url,
@@ -18,7 +20,7 @@ pub(super) fn build_same_format_provider_request_body(
spec: LocalSameFormatProviderSpec,
body_rules: Option<&Value>,
upstream_is_stream: bool,
kiro_auth: Option<&crate::provider_transport::kiro::KiroRequestAuth>,
kiro_auth: Option<&crate::ai_pipeline::provider_transport_facade::kiro::KiroRequestAuth>,
is_claude_code: bool,
) -> Option<Value> {
if let Some(kiro_auth) = kiro_auth {
@@ -60,11 +62,11 @@ pub(super) fn build_same_format_provider_request_body(
pub(super) fn build_same_format_upstream_url(
parts: &http::request::Parts,
transport: &crate::provider_transport::GatewayProviderTransportSnapshot,
transport: &GatewayProviderTransportSnapshot,
mapped_model: &str,
spec: LocalSameFormatProviderSpec,
upstream_is_stream: bool,
kiro_auth: Option<&crate::provider_transport::kiro::KiroRequestAuth>,
kiro_auth: Option<&crate::ai_pipeline::provider_transport_facade::kiro::KiroRequestAuth>,
) -> Option<String> {
if let Some(kiro_auth) = kiro_auth {
return build_kiro_generate_assistant_response_url(
@@ -0,0 +1,119 @@
use crate::ai_pipeline::control_facade::GatewayControlDecision;
pub(crate) fn resolve_execution_runtime_stream_plan_kind(
parts: &http::request::Parts,
decision: &GatewayControlDecision,
) -> Option<&'static str> {
aether_ai_pipeline::planner::route::resolve_execution_runtime_stream_plan_kind(
decision.route_class.as_deref(),
decision.route_family.as_deref(),
decision.route_kind.as_deref(),
&parts.method,
parts.uri.path(),
)
}
pub(crate) fn resolve_execution_runtime_sync_plan_kind(
parts: &http::request::Parts,
decision: &GatewayControlDecision,
) -> Option<&'static str> {
aether_ai_pipeline::planner::route::resolve_execution_runtime_sync_plan_kind(
decision.route_class.as_deref(),
decision.route_family.as_deref(),
decision.route_kind.as_deref(),
&parts.method,
parts.uri.path(),
)
}
pub(crate) fn is_matching_stream_request(
plan_kind: &str,
parts: &http::request::Parts,
body_json: &serde_json::Value,
) -> bool {
aether_ai_pipeline::planner::route::is_matching_stream_request(
plan_kind,
parts.uri.path(),
body_json,
)
}
pub(crate) fn supports_sync_scheduler_decision_kind(plan_kind: &str) -> bool {
aether_ai_pipeline::planner::route::supports_sync_scheduler_decision_kind(plan_kind)
}
pub(crate) fn supports_stream_scheduler_decision_kind(plan_kind: &str) -> bool {
aether_ai_pipeline::planner::route::supports_stream_scheduler_decision_kind(plan_kind)
}
#[cfg(test)]
mod tests {
use axum::http::{Method, Request};
use super::{
is_matching_stream_request, resolve_execution_runtime_stream_plan_kind,
resolve_execution_runtime_sync_plan_kind, supports_stream_scheduler_decision_kind,
supports_sync_scheduler_decision_kind,
};
use crate::ai_pipeline::control_facade::GatewayControlDecision;
fn sample_decision(route_family: &str, route_kind: &str) -> GatewayControlDecision {
GatewayControlDecision {
public_path: "/".to_string(),
public_query_string: None,
route_class: Some("ai_public".to_string()),
route_family: Some(route_family.to_string()),
route_kind: Some(route_kind.to_string()),
auth_context: None,
admin_principal: None,
auth_endpoint_signature: None,
execution_runtime_candidate: true,
local_auth_rejection: None,
}
}
#[test]
fn resolves_openai_chat_plan_kinds_via_pipeline_crate() {
let request = Request::builder()
.method(Method::POST)
.uri("/v1/chat/completions")
.body(())
.expect("request should build");
let (parts, _) = request.into_parts();
let decision = sample_decision("openai", "chat");
assert_eq!(
resolve_execution_runtime_sync_plan_kind(&parts, &decision),
Some("openai_chat_sync")
);
assert_eq!(
resolve_execution_runtime_stream_plan_kind(&parts, &decision),
Some("openai_chat_stream")
);
}
#[test]
fn stream_matching_uses_pipeline_route_logic() {
let request = Request::builder()
.method(Method::POST)
.uri("/v1/chat/completions")
.body(())
.expect("request should build");
let (parts, _) = request.into_parts();
assert!(!is_matching_stream_request(
"openai_chat_stream",
&parts,
&serde_json::json!({"stream": false}),
));
assert!(is_matching_stream_request(
"openai_chat_stream",
&parts,
&serde_json::json!({"stream": true}),
));
assert!(supports_sync_scheduler_decision_kind("openai_chat_sync"));
assert!(supports_stream_scheduler_decision_kind(
"openai_chat_stream"
));
}
}
@@ -0,0 +1,44 @@
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use crate::ai_pipeline::planner::auth_snapshot_facade::GatewayAuthApiKeySnapshot;
use crate::{AppState, GatewayError};
pub(crate) async fn list_selectable_candidates(
state: &AppState,
api_format: &str,
global_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
crate::scheduler::candidate::list_selectable_candidates(
state.data.as_ref(),
state,
api_format,
global_model_name,
require_streaming,
auth_snapshot,
now_unix_secs,
)
.await
}
pub(crate) async fn list_selectable_candidates_for_required_capability_without_requested_model(
state: &AppState,
candidate_api_format: &str,
required_capability: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
crate::scheduler::candidate::list_selectable_candidates_for_required_capability_without_requested_model(
state.data.as_ref(),
state,
candidate_api_format,
required_capability,
require_streaming,
auth_snapshot,
now_unix_secs,
)
.await
}
@@ -1,57 +1,60 @@
use std::collections::BTreeMap;
use aether_data::repository::candidates::{RequestCandidateStatus, UpsertRequestCandidateRecord};
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use serde_json::json;
use tracing::warn;
use uuid::Uuid;
use crate::ai_pipeline::control_facade::{
collect_control_headers, GatewayControlAuthContext, GatewayControlDecision,
};
use crate::ai_pipeline::execution_facade::{ConversionMode, ExecutionStrategy};
use crate::ai_pipeline::planner::auth_snapshot_facade::{
read_auth_api_key_snapshot, GatewayAuthApiKeySnapshot,
};
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
use crate::ai_pipeline::planner::candidate_runtime_facade::{
persist_available_local_candidate, persist_skipped_local_candidate,
};
use crate::ai_pipeline::planner::common::{
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
GEMINI_FILES_DELETE_PLAN_KIND, GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND,
GEMINI_FILES_LIST_PLAN_KIND, GEMINI_FILES_UPLOAD_PLAN_KIND,
};
use crate::ai_pipeline::planner::executor_facade::mark_unused_local_candidate_items;
use crate::ai_pipeline::planner::plan_builders::{
build_passthrough_stream_plan_from_decision, build_passthrough_sync_plan_from_decision,
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
};
use crate::control::GatewayControlDecision;
use crate::headers::collect_control_headers;
use crate::provider_transport::auth::{
use crate::ai_pipeline::planner::scheduler_facade::list_selectable_candidates_for_required_capability_without_requested_model;
use crate::ai_pipeline::planner::transport_facade::read_provider_transport_snapshot;
use crate::ai_pipeline::provider_transport_facade::auth::{
build_passthrough_headers_with_auth, resolve_local_gemini_auth,
};
use crate::provider_transport::policy::supports_local_gemini_transport_with_network;
use crate::provider_transport::url::build_gemini_files_passthrough_url;
use crate::provider_transport::{
use crate::ai_pipeline::provider_transport_facade::policy::supports_local_gemini_transport_with_network;
use crate::ai_pipeline::provider_transport_facade::url::build_gemini_files_passthrough_url;
use crate::ai_pipeline::provider_transport_facade::{
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::scheduler::{
current_unix_secs, list_selectable_candidates_for_required_capability_without_requested_model,
record_local_request_candidate_status, GatewayMinimalCandidateSelectionCandidate,
};
use crate::clock::current_unix_secs;
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
use aether_ai_pipeline::contracts::GEMINI_FILES_UPLOAD_PLAN_KIND;
use aether_ai_pipeline::planner::specialized::files::{
resolve_stream_spec, resolve_sync_spec, LocalGeminiFilesSpec,
};
const GEMINI_FILES_CANDIDATE_API_FORMAT: &str = "gemini:chat";
const GEMINI_FILES_CLIENT_API_FORMAT: &str = "gemini:files";
const GEMINI_FILES_REQUIRED_CAPABILITY: &str = "gemini_files";
#[derive(Debug, Clone, Copy)]
struct LocalGeminiFilesSpec {
decision_kind: &'static str,
report_kind: Option<&'static str>,
require_streaming: bool,
}
#[derive(Debug, Clone)]
struct LocalGeminiFilesDecisionInput {
auth_context: crate::control::GatewayControlAuthContext,
auth_snapshot: crate::data::auth::GatewayAuthApiKeySnapshot,
auth_context: GatewayControlAuthContext,
auth_snapshot: GatewayAuthApiKeySnapshot,
}
#[derive(Debug, Clone)]
struct LocalGeminiFilesCandidateAttempt {
candidate: GatewayMinimalCandidateSelectionCandidate,
candidate: SchedulerMinimalCandidateSelectionCandidate,
candidate_index: u32,
candidate_id: String,
}
@@ -181,43 +184,6 @@ pub(crate) async fn maybe_build_stream_local_gemini_files_decision_payload(
Ok(None)
}
fn resolve_sync_spec(plan_kind: &str) -> Option<LocalGeminiFilesSpec> {
match plan_kind {
GEMINI_FILES_UPLOAD_PLAN_KIND => Some(LocalGeminiFilesSpec {
decision_kind: GEMINI_FILES_UPLOAD_PLAN_KIND,
report_kind: Some("gemini_files_store_mapping"),
require_streaming: false,
}),
GEMINI_FILES_LIST_PLAN_KIND => Some(LocalGeminiFilesSpec {
decision_kind: GEMINI_FILES_LIST_PLAN_KIND,
report_kind: Some("gemini_files_store_mapping"),
require_streaming: false,
}),
GEMINI_FILES_GET_PLAN_KIND => Some(LocalGeminiFilesSpec {
decision_kind: GEMINI_FILES_GET_PLAN_KIND,
report_kind: Some("gemini_files_store_mapping"),
require_streaming: false,
}),
GEMINI_FILES_DELETE_PLAN_KIND => Some(LocalGeminiFilesSpec {
decision_kind: GEMINI_FILES_DELETE_PLAN_KIND,
report_kind: Some("gemini_files_delete_mapping"),
require_streaming: false,
}),
_ => None,
}
}
fn resolve_stream_spec(plan_kind: &str) -> Option<LocalGeminiFilesSpec> {
match plan_kind {
GEMINI_FILES_DOWNLOAD_PLAN_KIND => Some(LocalGeminiFilesSpec {
decision_kind: GEMINI_FILES_DOWNLOAD_PLAN_KIND,
report_kind: None,
require_streaming: true,
}),
_ => None,
}
}
async fn build_local_sync_plan_and_reports(
state: &AppState,
parts: &http::request::Parts,
@@ -333,13 +299,13 @@ async fn resolve_local_gemini_files_decision_input(
return None;
};
let auth_snapshot = match state
.read_auth_api_key_snapshot(
&auth_context.user_id,
&auth_context.api_key_id,
current_unix_secs(),
)
.await
let auth_snapshot = match read_auth_api_key_snapshot(
state,
&auth_context.user_id,
&auth_context.api_key_id,
current_unix_secs(),
)
.await
{
Ok(Some(snapshot)) => snapshot,
Ok(None) => return None,
@@ -392,46 +358,19 @@ async fn materialize_local_gemini_files_candidate_attempts(
"key_name": candidate.key_name.clone(),
});
let candidate_id = match state
.upsert_request_candidate(UpsertRequestCandidateRecord {
id: generated_candidate_id.clone(),
request_id: trace_id.to_string(),
user_id: Some(input.auth_context.user_id.clone()),
api_key_id: Some(input.auth_context.api_key_id.clone()),
username: None,
api_key_name: None,
candidate_index: candidate_index as u32,
retry_index: 0,
provider_id: Some(candidate.provider_id.clone()),
endpoint_id: Some(candidate.endpoint_id.clone()),
key_id: Some(candidate.key_id.clone()),
status: RequestCandidateStatus::Available,
skip_reason: None,
is_cached: Some(false),
status_code: None,
error_type: None,
error_message: None,
latency_ms: None,
concurrent_requests: None,
extra_data: Some(extra_data),
required_capabilities: candidate.key_capabilities.clone(),
created_at_unix_secs: Some(created_at_unix_secs),
started_at_unix_secs: None,
finished_at_unix_secs: None,
})
.await
{
Ok(Some(stored)) => stored.id,
Ok(None) => generated_candidate_id.clone(),
Err(err) => {
warn!(
trace_id = %trace_id,
error = ?err,
"gateway local gemini files request candidate upsert failed"
);
generated_candidate_id.clone()
}
};
let candidate_id = persist_available_local_candidate(
state,
trace_id,
&input.auth_context.user_id,
&input.auth_context.api_key_id,
&candidate,
candidate_index as u32,
&generated_candidate_id,
Some(extra_data),
created_at_unix_secs,
"gateway local gemini files request candidate upsert failed",
)
.await;
attempts.push(LocalGeminiFilesCandidateAttempt {
candidate,
@@ -460,13 +399,13 @@ async fn maybe_build_local_gemini_files_decision_payload_for_candidate(
candidate_id,
} = attempt;
let transport = match state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
let transport = match read_provider_transport_snapshot(
state,
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
{
Ok(Some(snapshot)) => snapshot,
Ok(None) => {
@@ -664,16 +603,8 @@ async fn maybe_build_local_gemini_files_decision_payload_for_candidate(
EXECUTION_RUNTIME_SYNC_DECISION_ACTION.to_string()
},
decision_kind: Some(spec.decision_kind.to_string()),
execution_strategy: Some(
crate::execution_runtime::ExecutionStrategy::LocalSameFormat
.as_str()
.to_string(),
),
conversion_mode: Some(
crate::execution_runtime::ConversionMode::None
.as_str()
.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()),
@@ -737,69 +668,37 @@ async fn mark_skipped_local_gemini_files_candidate(
state: &AppState,
input: &LocalGeminiFilesDecisionInput,
trace_id: &str,
candidate: &GatewayMinimalCandidateSelectionCandidate,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
candidate_index: u32,
candidate_id: &str,
skip_reason: &'static str,
) {
if let Err(err) = state
.upsert_request_candidate(UpsertRequestCandidateRecord {
id: candidate_id.to_string(),
request_id: trace_id.to_string(),
user_id: Some(input.auth_context.user_id.clone()),
api_key_id: Some(input.auth_context.api_key_id.clone()),
username: None,
api_key_name: None,
candidate_index,
retry_index: 0,
provider_id: Some(candidate.provider_id.clone()),
endpoint_id: Some(candidate.endpoint_id.clone()),
key_id: Some(candidate.key_id.clone()),
status: RequestCandidateStatus::Skipped,
skip_reason: Some(skip_reason.to_string()),
is_cached: Some(false),
status_code: None,
error_type: None,
error_message: None,
latency_ms: None,
concurrent_requests: None,
extra_data: None,
required_capabilities: candidate.key_capabilities.clone(),
created_at_unix_secs: None,
started_at_unix_secs: None,
finished_at_unix_secs: Some(current_unix_secs()),
})
.await
{
warn!(
trace_id = %trace_id,
candidate_id = %candidate_id,
skip_reason,
error = ?err,
"gateway local gemini files failed to persist skipped candidate"
);
}
persist_skipped_local_candidate(
state,
trace_id,
&input.auth_context.user_id,
&input.auth_context.api_key_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
current_unix_secs(),
"gateway local gemini files failed to persist skipped candidate",
)
.await;
}
async fn mark_unused_local_files_candidates<T>(state: &AppState, remaining: Vec<T>)
where
T: LocalGeminiFilesPlanAndReport,
{
for plan_and_report in remaining {
record_local_request_candidate_status(
state,
plan_and_report.plan(),
plan_and_report.report_context(),
RequestCandidateStatus::Unused,
None,
None,
None,
None,
None,
None,
)
.await;
}
mark_unused_local_candidate_items(
state,
remaining,
|item| item.plan(),
|item| item.report_context(),
)
.await;
}
trait LocalGeminiFilesPlanAndReport {
@@ -1,63 +1,59 @@
use std::collections::BTreeMap;
use aether_data::repository::candidates::{RequestCandidateStatus, UpsertRequestCandidateRecord};
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use serde_json::{json, Value};
use tracing::warn;
use uuid::Uuid;
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
use crate::ai_pipeline::planner::common::{
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND,
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND,
use crate::ai_pipeline::control_facade::{
collect_control_headers, GatewayControlAuthContext, GatewayControlDecision,
};
use crate::ai_pipeline::execution_facade::{ConversionMode, ExecutionStrategy};
use crate::ai_pipeline::planner::auth_snapshot_facade::{
read_auth_api_key_snapshot, GatewayAuthApiKeySnapshot,
};
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
use crate::ai_pipeline::planner::candidate_runtime_facade::{
persist_available_local_candidate, persist_skipped_local_candidate,
};
use crate::ai_pipeline::planner::common::EXECUTION_RUNTIME_SYNC_DECISION_ACTION;
use crate::ai_pipeline::planner::executor_facade::mark_unused_local_candidate_items;
use crate::ai_pipeline::planner::plan_builders::{
build_passthrough_sync_plan_from_decision, LocalSyncPlanAndReport,
};
use crate::control::GatewayControlDecision;
use crate::headers::collect_control_headers;
use crate::provider_transport::auth::{
use crate::ai_pipeline::planner::scheduler_facade::list_selectable_candidates;
use crate::ai_pipeline::planner::transport_facade::{
read_provider_transport_snapshot, GatewayProviderTransportSnapshot,
};
use crate::ai_pipeline::provider_transport_facade::auth::{
build_passthrough_headers_with_auth, resolve_local_gemini_auth, resolve_local_openai_chat_auth,
};
use crate::provider_transport::policy::{
use crate::ai_pipeline::provider_transport_facade::policy::{
supports_local_gemini_transport_with_network, supports_local_standard_transport_with_network,
};
use crate::provider_transport::url::{
use crate::ai_pipeline::provider_transport_facade::url::{
build_gemini_video_predict_long_running_url, build_passthrough_path_url,
};
use crate::provider_transport::{
use crate::ai_pipeline::provider_transport_facade::{
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::scheduler::{
current_unix_secs, list_selectable_candidates, record_local_request_candidate_status,
GatewayMinimalCandidateSelectionCandidate,
};
use crate::clock::current_unix_secs;
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum LocalVideoCreateFamily {
OpenAi,
Gemini,
}
#[derive(Debug, Clone, Copy)]
struct LocalVideoCreateSpec {
api_format: &'static str,
decision_kind: &'static str,
report_kind: &'static str,
family: LocalVideoCreateFamily,
}
use aether_ai_pipeline::planner::specialized::video::{
resolve_sync_spec, LocalVideoCreateFamily, LocalVideoCreateSpec,
};
#[derive(Debug, Clone)]
struct LocalVideoCreateDecisionInput {
auth_context: crate::control::GatewayControlAuthContext,
auth_context: GatewayControlAuthContext,
requested_model: String,
auth_snapshot: crate::data::auth::GatewayAuthApiKeySnapshot,
auth_snapshot: GatewayAuthApiKeySnapshot,
}
#[derive(Debug, Clone)]
struct LocalVideoCreateCandidateAttempt {
candidate: GatewayMinimalCandidateSelectionCandidate,
candidate: SchedulerMinimalCandidateSelectionCandidate,
candidate_index: u32,
candidate_id: String,
}
@@ -141,24 +137,6 @@ pub(crate) async fn maybe_build_sync_local_video_decision_payload(
Ok(None)
}
fn resolve_sync_spec(plan_kind: &str) -> Option<LocalVideoCreateSpec> {
match plan_kind {
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND => Some(LocalVideoCreateSpec {
api_format: "openai:video",
decision_kind: OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND,
report_kind: "openai_video_create_sync_finalize",
family: LocalVideoCreateFamily::OpenAi,
}),
GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND => Some(LocalVideoCreateSpec {
api_format: "gemini:video",
decision_kind: GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND,
report_kind: "gemini_video_create_sync_finalize",
family: LocalVideoCreateFamily::Gemini,
}),
_ => None,
}
}
async fn build_local_sync_plan_and_reports(
state: &AppState,
parts: &http::request::Parts,
@@ -257,13 +235,13 @@ async fn resolve_local_video_create_decision_input(
LocalVideoCreateFamily::Gemini => extract_gemini_video_model_from_path(parts.uri.path())?,
};
let auth_snapshot = match state
.read_auth_api_key_snapshot(
&auth_context.user_id,
&auth_context.api_key_id,
current_unix_secs(),
)
.await
let auth_snapshot = match read_auth_api_key_snapshot(
state,
&auth_context.user_id,
&auth_context.api_key_id,
current_unix_secs(),
)
.await
{
Ok(Some(snapshot)) => snapshot,
Ok(None) => return None,
@@ -299,13 +277,13 @@ async fn maybe_build_local_video_create_decision_payload_for_candidate(
candidate_index,
candidate_id,
} = attempt;
let transport = match state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
let transport = match read_provider_transport_snapshot(
state,
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
{
Ok(Some(snapshot)) => snapshot,
Ok(None) => {
@@ -461,16 +439,8 @@ async fn maybe_build_local_video_create_decision_payload_for_candidate(
Some(GatewayControlSyncDecisionResponse {
action: EXECUTION_RUNTIME_SYNC_DECISION_ACTION.to_string(),
decision_kind: Some(spec.decision_kind.to_string()),
execution_strategy: Some(
crate::execution_runtime::ExecutionStrategy::LocalSameFormat
.as_str()
.to_string(),
),
conversion_mode: Some(
crate::execution_runtime::ConversionMode::None
.as_str()
.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()),
@@ -552,7 +522,7 @@ fn build_provider_request_body(
fn build_video_upstream_url(
parts: &http::request::Parts,
transport: &crate::provider_transport::GatewayProviderTransportSnapshot,
transport: &GatewayProviderTransportSnapshot,
mapped_model: &str,
family: LocalVideoCreateFamily,
) -> Option<String> {
@@ -595,7 +565,7 @@ async fn materialize_local_video_create_candidate_attempts(
state: &AppState,
trace_id: &str,
input: &LocalVideoCreateDecisionInput,
candidates: Vec<GatewayMinimalCandidateSelectionCandidate>,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
api_format: &str,
) -> Vec<LocalVideoCreateCandidateAttempt> {
let candidates = prefer_local_tunnel_owner_candidates(state, candidates).await;
@@ -616,47 +586,19 @@ async fn materialize_local_video_create_candidate_attempts(
"key_name": candidate.key_name.clone(),
});
let candidate_id = match state
.upsert_request_candidate(UpsertRequestCandidateRecord {
id: generated_candidate_id.clone(),
request_id: trace_id.to_string(),
user_id: Some(input.auth_context.user_id.clone()),
api_key_id: Some(input.auth_context.api_key_id.clone()),
username: None,
api_key_name: None,
candidate_index: candidate_index as u32,
retry_index: 0,
provider_id: Some(candidate.provider_id.clone()),
endpoint_id: Some(candidate.endpoint_id.clone()),
key_id: Some(candidate.key_id.clone()),
status: RequestCandidateStatus::Available,
skip_reason: None,
is_cached: Some(false),
status_code: None,
error_type: None,
error_message: None,
latency_ms: None,
concurrent_requests: None,
extra_data: Some(extra_data),
required_capabilities: candidate.key_capabilities.clone(),
created_at_unix_secs: Some(created_at_unix_secs),
started_at_unix_secs: None,
finished_at_unix_secs: None,
})
.await
{
Ok(Some(stored)) => stored.id,
Ok(None) => generated_candidate_id.clone(),
Err(err) => {
warn!(
trace_id = %trace_id,
decision_api_format = api_format,
error = ?err,
"gateway local video decision request candidate upsert failed"
);
generated_candidate_id.clone()
}
};
let candidate_id = persist_available_local_candidate(
state,
trace_id,
&input.auth_context.user_id,
&input.auth_context.api_key_id,
&candidate,
candidate_index as u32,
&generated_candidate_id,
Some(extra_data),
created_at_unix_secs,
"gateway local video decision request candidate upsert failed",
)
.await;
attempts.push(LocalVideoCreateCandidateAttempt {
candidate,
@@ -672,70 +614,37 @@ async fn mark_skipped_local_video_candidate(
state: &AppState,
input: &LocalVideoCreateDecisionInput,
trace_id: &str,
candidate: &GatewayMinimalCandidateSelectionCandidate,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
candidate_index: u32,
candidate_id: &str,
skip_reason: &'static str,
) {
let terminal_unix_secs = current_unix_secs();
if let Err(err) = state
.upsert_request_candidate(UpsertRequestCandidateRecord {
id: candidate_id.to_string(),
request_id: trace_id.to_string(),
user_id: Some(input.auth_context.user_id.clone()),
api_key_id: Some(input.auth_context.api_key_id.clone()),
username: None,
api_key_name: None,
candidate_index,
retry_index: 0,
provider_id: Some(candidate.provider_id.clone()),
endpoint_id: Some(candidate.endpoint_id.clone()),
key_id: Some(candidate.key_id.clone()),
status: RequestCandidateStatus::Skipped,
skip_reason: Some(skip_reason.to_string()),
is_cached: Some(false),
status_code: None,
error_type: None,
error_message: None,
latency_ms: None,
concurrent_requests: None,
extra_data: None,
required_capabilities: candidate.key_capabilities.clone(),
created_at_unix_secs: None,
started_at_unix_secs: None,
finished_at_unix_secs: Some(terminal_unix_secs),
})
.await
{
warn!(
trace_id = %trace_id,
candidate_id = %candidate_id,
skip_reason,
error = ?err,
"gateway local video decision failed to persist skipped candidate"
);
}
persist_skipped_local_candidate(
state,
trace_id,
&input.auth_context.user_id,
&input.auth_context.api_key_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
current_unix_secs(),
"gateway local video decision failed to persist skipped candidate",
)
.await;
}
async fn mark_unused_local_video_candidates(
state: &AppState,
remaining: Vec<LocalSyncPlanAndReport>,
) {
for plan_and_report in remaining {
record_local_request_candidate_status(
state,
&plan_and_report.plan,
plan_and_report.report_context.as_ref(),
RequestCandidateStatus::Unused,
None,
None,
None,
None,
None,
None,
)
.await;
}
mark_unused_local_candidate_items(
state,
remaining,
|item| &item.plan,
|item| item.report_context.as_ref(),
)
.await;
}
fn extract_gemini_video_model_from_path(path: &str) -> Option<String> {
@@ -1,33 +0,0 @@
use crate::ai_pipeline::planner::common::{
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND,
};
use super::super::family::{LocalStandardSourceFamily, LocalStandardSourceMode, LocalStandardSpec};
pub(crate) fn resolve_sync_spec(plan_kind: &str) -> Option<LocalStandardSpec> {
match plan_kind {
CLAUDE_CHAT_SYNC_PLAN_KIND => Some(LocalStandardSpec {
api_format: "claude:chat",
decision_kind: CLAUDE_CHAT_SYNC_PLAN_KIND,
report_kind: "claude_chat_sync_finalize",
family: LocalStandardSourceFamily::Standard,
mode: LocalStandardSourceMode::Chat,
require_streaming: false,
}),
_ => None,
}
}
pub(crate) fn resolve_stream_spec(plan_kind: &str) -> Option<LocalStandardSpec> {
match plan_kind {
CLAUDE_CHAT_STREAM_PLAN_KIND => Some(LocalStandardSpec {
api_format: "claude:chat",
decision_kind: CLAUDE_CHAT_STREAM_PLAN_KIND,
report_kind: "claude_chat_stream_success",
family: LocalStandardSourceFamily::Standard,
mode: LocalStandardSourceMode::Chat,
require_streaming: true,
}),
_ => None,
}
}
@@ -1,33 +0,0 @@
use crate::ai_pipeline::planner::common::{
CLAUDE_CLI_STREAM_PLAN_KIND, CLAUDE_CLI_SYNC_PLAN_KIND,
};
use super::super::family::{LocalStandardSourceFamily, LocalStandardSourceMode, LocalStandardSpec};
pub(crate) fn resolve_sync_spec(plan_kind: &str) -> Option<LocalStandardSpec> {
match plan_kind {
CLAUDE_CLI_SYNC_PLAN_KIND => Some(LocalStandardSpec {
api_format: "claude:cli",
decision_kind: CLAUDE_CLI_SYNC_PLAN_KIND,
report_kind: "claude_cli_sync_finalize",
family: LocalStandardSourceFamily::Standard,
mode: LocalStandardSourceMode::Cli,
require_streaming: false,
}),
_ => None,
}
}
pub(crate) fn resolve_stream_spec(plan_kind: &str) -> Option<LocalStandardSpec> {
match plan_kind {
CLAUDE_CLI_STREAM_PLAN_KIND => Some(LocalStandardSpec {
api_format: "claude:cli",
decision_kind: CLAUDE_CLI_STREAM_PLAN_KIND,
report_kind: "claude_cli_stream_success",
family: LocalStandardSourceFamily::Standard,
mode: LocalStandardSourceMode::Cli,
require_streaming: true,
}),
_ => None,
}
}
@@ -1,13 +1,23 @@
use crate::control::GatewayControlDecision;
use crate::ai_pipeline::control_facade::GatewayControlDecision;
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
use aether_ai_pipeline::planner::standard::claude::{
resolve_stream_spec as resolve_pipeline_stream_spec,
resolve_sync_spec as resolve_pipeline_sync_spec,
};
use super::family::{
maybe_build_stream_via_standard_family_payload, maybe_build_sync_via_standard_family_payload,
};
pub(crate) use crate::ai_pipeline::conversion::request::normalize_claude_request_to_openai_chat_request;
pub(crate) mod chat;
pub(crate) mod cli;
pub(crate) fn resolve_sync_spec(plan_kind: &str) -> Option<super::family::LocalStandardSpec> {
resolve_pipeline_sync_spec(plan_kind)
}
pub(crate) fn resolve_stream_spec(plan_kind: &str) -> Option<super::family::LocalStandardSpec> {
resolve_pipeline_stream_spec(plan_kind)
}
pub(crate) async fn maybe_build_sync_local_claude_decision_payload(
state: &AppState,
@@ -24,9 +34,7 @@ pub(crate) async fn maybe_build_sync_local_claude_decision_payload(
decision,
body_json,
plan_kind,
|plan_kind| {
chat::resolve_sync_spec(plan_kind).or_else(|| cli::resolve_sync_spec(plan_kind))
},
resolve_sync_spec,
)
.await
}
@@ -46,9 +54,7 @@ pub(crate) async fn maybe_build_stream_local_claude_decision_payload(
decision,
body_json,
plan_kind,
|plan_kind| {
chat::resolve_stream_spec(plan_kind).or_else(|| cli::resolve_stream_spec(plan_kind))
},
resolve_stream_spec,
)
.await
}
@@ -0,0 +1,433 @@
use std::collections::BTreeMap;
use std::fmt::Write;
use crate::ai_pipeline::provider_transport_facade::body_rules_handle_path;
use serde_json::{json, Value};
use sha1::{Digest as Sha1Digest, Sha1};
use sha2::{Digest as Sha2Digest, Sha256};
use uuid::Uuid;
const CODEX_PROMPT_CACHE_NAMESPACE_VERSION: &str = "v3";
const UUID_NAMESPACE_OID_BYTES: [u8; 16] = [
0x6b, 0xa7, 0xb8, 0x12, 0x9d, 0xad, 0x11, 0xd1, 0x80, 0xb4, 0x00, 0xc0, 0x4f, 0xd4, 0x30, 0xc8,
];
fn is_codex_openai_cli_request(provider_type: &str, provider_api_format: &str) -> bool {
provider_type.trim().eq_ignore_ascii_case("codex")
&& matches!(
provider_api_format.trim().to_ascii_lowercase().as_str(),
"openai:cli" | "openai:compact"
)
}
fn build_stable_codex_prompt_cache_key(user_api_key_id: &str) -> Option<String> {
let normalized = user_api_key_id.trim();
if normalized.is_empty() {
return None;
}
let namespace = format!(
"aether:codex:prompt-cache:{CODEX_PROMPT_CACHE_NAMESPACE_VERSION}:user:{normalized}"
);
let mut hasher = Sha1::new();
hasher.update(UUID_NAMESPACE_OID_BYTES);
hasher.update(namespace.as_bytes());
let digest = hasher.finalize();
let mut bytes = [0u8; 16];
bytes.copy_from_slice(&digest[..16]);
bytes[6] = (bytes[6] & 0x0f) | 0x50;
bytes[8] = (bytes[8] & 0x3f) | 0x80;
Some(Uuid::from_bytes(bytes).to_string())
}
fn maybe_inject_codex_prompt_cache_key(
provider_request_body: &mut Value,
provider_type: &str,
provider_api_format: &str,
user_api_key_id: Option<&str>,
) {
if !is_codex_openai_cli_request(provider_type, provider_api_format) {
return;
}
let Some(body_object) = provider_request_body.as_object_mut() else {
return;
};
let existing = body_object
.get("prompt_cache_key")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default();
if !existing.is_empty() {
return;
}
let Some(prompt_cache_key) = user_api_key_id.and_then(build_stable_codex_prompt_cache_key)
else {
return;
};
body_object.insert(
"prompt_cache_key".to_string(),
Value::String(prompt_cache_key),
);
}
fn build_short_codex_header_id(seed: &str) -> Option<String> {
let normalized = seed.trim();
if normalized.is_empty() {
return None;
}
let digest = Sha256::digest(normalized.as_bytes());
let mut short_id = String::with_capacity(16);
for byte in digest.iter().take(8) {
let _ = write!(&mut short_id, "{byte:02x}");
}
Some(short_id)
}
fn header_map_has_non_empty_value(headers: &http::HeaderMap, header_name: &str) -> bool {
let target = header_name.trim().to_ascii_lowercase();
if target.is_empty() {
return false;
}
headers.iter().any(|(name, value)| {
if name.as_str().trim().to_ascii_lowercase() != target {
return false;
}
value
.to_str()
.ok()
.map(str::trim)
.map(|value| !value.is_empty())
.unwrap_or(false)
})
}
fn extract_codex_account_id(decrypted_auth_config_raw: Option<&str>) -> Option<String> {
let raw = decrypted_auth_config_raw?.trim();
if raw.is_empty() {
return None;
}
serde_json::from_str::<Value>(raw).ok().and_then(|value| {
value
.get("account_id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
})
}
pub(crate) fn apply_codex_openai_cli_special_headers(
provider_request_headers: &mut BTreeMap<String, String>,
provider_request_body: &Value,
original_headers: &http::HeaderMap,
provider_type: &str,
provider_api_format: &str,
request_id: Option<&str>,
decrypted_auth_config_raw: Option<&str>,
) {
if !is_codex_openai_cli_request(provider_type, provider_api_format) {
return;
}
if let Some(account_id) = extract_codex_account_id(decrypted_auth_config_raw) {
provider_request_headers.insert("chatgpt-account-id".to_string(), account_id);
}
if !provider_request_headers
.get("x-client-request-id")
.map(|value| !value.trim().is_empty())
.unwrap_or(false)
{
if let Some(request_id) = request_id.map(str::trim).filter(|value| !value.is_empty()) {
provider_request_headers
.insert("x-client-request-id".to_string(), request_id.to_string());
}
}
let prompt_cache_key = provider_request_body
.get("prompt_cache_key")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
let Some(short_id) = prompt_cache_key.and_then(build_short_codex_header_id) else {
return;
};
if !header_map_has_non_empty_value(original_headers, "session_id") {
provider_request_headers.insert("session_id".to_string(), short_id.clone());
}
if provider_api_format.trim().to_ascii_lowercase() != "openai:compact"
&& !header_map_has_non_empty_value(original_headers, "conversation_id")
{
provider_request_headers.insert("conversation_id".to_string(), short_id);
}
}
pub(crate) fn apply_codex_openai_cli_special_body_edits(
provider_request_body: &mut Value,
provider_type: &str,
provider_api_format: &str,
body_rules: Option<&Value>,
user_api_key_id: Option<&str>,
) {
if !is_codex_openai_cli_request(provider_type, provider_api_format) {
return;
}
let Some(body_object) = provider_request_body.as_object_mut() else {
return;
};
if !body_rules_handle_path(body_rules, "max_output_tokens") {
body_object.remove("max_output_tokens");
}
if !body_rules_handle_path(body_rules, "temperature") {
body_object.remove("temperature");
}
if !body_rules_handle_path(body_rules, "top_p") {
body_object.remove("top_p");
}
if !body_rules_handle_path(body_rules, "metadata") {
body_object.remove("metadata");
}
if !body_rules_handle_path(body_rules, "store") {
body_object.insert("store".to_string(), json!(false));
}
if !body_rules_handle_path(body_rules, "instructions")
&& !body_object.contains_key("instructions")
{
body_object.insert("instructions".to_string(), json!("You are GPT-5."));
}
maybe_inject_codex_prompt_cache_key(
provider_request_body,
provider_type,
provider_api_format,
user_api_key_id,
);
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use super::{
apply_codex_openai_cli_special_body_edits, apply_codex_openai_cli_special_headers,
};
use http::{HeaderMap, HeaderValue};
use serde_json::json;
#[test]
fn applies_codex_defaults_when_body_rules_do_not_handle_fields() {
let mut body = json!({
"model": "gpt-5",
"max_output_tokens": 128,
"temperature": 0.3,
"top_p": 0.9,
"metadata": {"client": "desktop"},
"store": true
});
apply_codex_openai_cli_special_body_edits(&mut body, "codex", "openai:cli", None, None);
assert!(body.get("max_output_tokens").is_none());
assert!(body.get("temperature").is_none());
assert!(body.get("top_p").is_none());
assert!(body.get("metadata").is_none());
assert_eq!(body["store"], false);
assert_eq!(body["instructions"], "You are GPT-5.");
}
#[test]
fn defers_to_user_body_rules_for_handled_fields() {
let body_rules = json!([
{"action":"set","path":"store","value":true},
{"action":"set","path":"instructions","value":"Keep custom"},
{"action":"set","path":"metadata","value":{"client":"desktop","mode":"custom"}},
{"action":"set","path":"top_p","value":0.5}
]);
let mut body = json!({
"model": "gpt-5",
"max_output_tokens": 128,
"metadata": {"client": "desktop", "mode": "custom"},
"store": true,
"instructions": "Keep custom",
"top_p": 0.5
});
apply_codex_openai_cli_special_body_edits(
&mut body,
"codex",
"openai:compact",
Some(&body_rules),
None,
);
assert!(body.get("max_output_tokens").is_none());
assert_eq!(body["store"], true);
assert_eq!(body["instructions"], "Keep custom");
assert_eq!(body["metadata"]["mode"], "custom");
assert_eq!(body["top_p"], 0.5);
}
#[test]
fn injects_stable_prompt_cache_key_for_codex_requests() {
let mut body = json!({
"model": "gpt-5",
"input": "hello",
});
apply_codex_openai_cli_special_body_edits(
&mut body,
"codex",
"openai:cli",
None,
Some("key-123"),
);
assert_eq!(
body["prompt_cache_key"],
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
);
}
#[test]
fn keeps_existing_prompt_cache_key_for_codex_requests() {
let mut body = json!({
"model": "gpt-5",
"input": "hello",
"prompt_cache_key": "existing-key",
});
apply_codex_openai_cli_special_body_edits(
&mut body,
"codex",
"openai:cli",
None,
Some("key-123"),
);
assert_eq!(body["prompt_cache_key"], "existing-key");
}
#[test]
fn injects_chatgpt_account_id_and_session_headers_for_codex_requests() {
let mut headers = BTreeMap::new();
let body = json!({
"model": "gpt-5",
"prompt_cache_key": "172c39e6-c0a0-5a70-8b63-e0f8e0d185a3",
});
apply_codex_openai_cli_special_headers(
&mut headers,
&body,
&HeaderMap::new(),
"codex",
"openai:cli",
Some("trace-codex-123"),
Some(r#"{"account_id":"acc-123"}"#),
);
assert_eq!(
headers.get("chatgpt-account-id"),
Some(&"acc-123".to_string())
);
assert_eq!(
headers.get("x-client-request-id"),
Some(&"trace-codex-123".to_string())
);
assert_eq!(
headers.get("session_id"),
Some(&"ab5ecce4f0d110fe".to_string())
);
assert_eq!(
headers.get("conversation_id"),
Some(&"ab5ecce4f0d110fe".to_string())
);
}
#[test]
fn respects_existing_codex_request_and_session_headers() {
let mut headers = BTreeMap::new();
headers.insert(
"x-client-request-id".to_string(),
"kept-by-rule-request".to_string(),
);
headers.insert("session_id".to_string(), "kept-by-rule".to_string());
let body = json!({
"model": "gpt-5",
"prompt_cache_key": "172c39e6-c0a0-5a70-8b63-e0f8e0d185a3",
});
let mut original_headers = HeaderMap::new();
original_headers.insert(
"x-client-request-id",
HeaderValue::from_static("user-specified-request"),
);
original_headers.insert(
"session_id",
HeaderValue::from_static("user-specified-session"),
);
original_headers.insert(
"conversation_id",
HeaderValue::from_static("user-specified-conversation"),
);
apply_codex_openai_cli_special_headers(
&mut headers,
&body,
&original_headers,
"codex",
"openai:cli",
Some("trace-codex-123"),
Some(r#"{"account_id":"acc-123"}"#),
);
assert_eq!(
headers.get("x-client-request-id"),
Some(&"kept-by-rule-request".to_string())
);
assert_eq!(headers.get("session_id"), Some(&"kept-by-rule".to_string()));
assert!(headers.get("conversation_id").is_none());
}
#[test]
fn skips_conversation_id_for_compact_codex_requests() {
let mut headers = BTreeMap::new();
let body = json!({
"model": "gpt-5",
"prompt_cache_key": "172c39e6-c0a0-5a70-8b63-e0f8e0d185a3",
});
apply_codex_openai_cli_special_headers(
&mut headers,
&body,
&HeaderMap::new(),
"codex",
"openai:compact",
Some("trace-codex-compact-123"),
Some(r#"{"account_id":"acc-123"}"#),
);
assert_eq!(
headers.get("chatgpt-account-id"),
Some(&"acc-123".to_string())
);
assert_eq!(
headers.get("x-client-request-id"),
Some(&"trace-codex-compact-123".to_string())
);
assert_eq!(
headers.get("session_id"),
Some(&"ab5ecce4f0d110fe".to_string())
);
assert!(headers.get("conversation_id").is_none());
}
}
@@ -1,11 +1,11 @@
use tracing::warn;
use crate::ai_pipeline::control_facade::GatewayControlDecision;
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::control::GatewayControlDecision;
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
use super::candidates::{
@@ -1,16 +1,19 @@
use std::collections::BTreeSet;
use aether_data::repository::candidates::{RequestCandidateStatus, UpsertRequestCandidateRecord};
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use serde_json::json;
use tracing::warn;
use uuid::Uuid;
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
use crate::control::GatewayControlDecision;
use crate::execution_runtime::{ConversionMode, ExecutionStrategy};
use crate::scheduler::{
current_unix_secs, list_selectable_candidates, GatewayMinimalCandidateSelectionCandidate,
use crate::ai_pipeline::control_facade::GatewayControlDecision;
use crate::ai_pipeline::execution_facade::{ConversionMode, ExecutionStrategy};
use crate::ai_pipeline::planner::auth_snapshot_facade::{
read_auth_api_key_snapshot, GatewayAuthApiKeySnapshot,
};
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
use crate::ai_pipeline::planner::candidate_runtime_facade::persist_available_local_candidate;
use crate::ai_pipeline::planner::scheduler_facade::list_selectable_candidates;
use crate::clock::current_unix_secs;
use crate::{append_execution_contract_fields_to_value, AppState, GatewayError};
use super::types::{
@@ -42,13 +45,13 @@ pub(super) async fn resolve_local_standard_decision_input(
LocalStandardSourceFamily::Gemini => extract_gemini_model_from_path(parts.uri.path())?,
};
let auth_snapshot = match state
.read_auth_api_key_snapshot(
&auth_context.user_id,
&auth_context.api_key_id,
current_unix_secs(),
)
.await
let auth_snapshot = match read_auth_api_key_snapshot(
state,
&auth_context.user_id,
&auth_context.api_key_id,
current_unix_secs(),
)
.await
{
Ok(Some(snapshot)) => snapshot,
Ok(None) => return None,
@@ -157,47 +160,19 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
candidate.endpoint_api_format.as_str(),
);
let stored_candidate_id = match state
.upsert_request_candidate(UpsertRequestCandidateRecord {
id: candidate_id.clone(),
request_id: trace_id.to_string(),
user_id: Some(input.auth_context.user_id.clone()),
api_key_id: Some(input.auth_context.api_key_id.clone()),
username: None,
api_key_name: None,
candidate_index: candidate_index as u32,
retry_index: 0,
provider_id: Some(candidate.provider_id.clone()),
endpoint_id: Some(candidate.endpoint_id.clone()),
key_id: Some(candidate.key_id.clone()),
status: RequestCandidateStatus::Available,
skip_reason: None,
is_cached: Some(false),
status_code: None,
error_type: None,
error_message: None,
latency_ms: None,
concurrent_requests: None,
extra_data: Some(extra_data),
required_capabilities: candidate.key_capabilities.clone(),
created_at_unix_secs: Some(created_at_unix_secs),
started_at_unix_secs: None,
finished_at_unix_secs: None,
})
.await
{
Ok(Some(stored)) => stored.id,
Ok(None) => candidate_id.clone(),
Err(err) => {
warn!(
trace_id = %trace_id,
api_format = spec.api_format,
error = ?err,
"gateway local standard decision request candidate upsert failed"
);
candidate_id.clone()
}
};
let stored_candidate_id = persist_available_local_candidate(
state,
trace_id,
&input.auth_context.user_id,
&input.auth_context.api_key_id,
&candidate,
candidate_index as u32,
&candidate_id,
Some(extra_data),
created_at_unix_secs,
"gateway local standard decision request candidate upsert failed",
)
.await;
attempts.push(LocalStandardCandidateAttempt {
candidate,
@@ -210,9 +185,9 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
}
fn auth_snapshot_allows_cross_format_candidate(
auth_snapshot: &crate::data::auth::GatewayAuthApiKeySnapshot,
auth_snapshot: &GatewayAuthApiKeySnapshot,
requested_model: &str,
candidate: &GatewayMinimalCandidateSelectionCandidate,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> bool {
if let Some(allowed_providers) = auth_snapshot.effective_allowed_providers() {
let provider_allowed = allowed_providers.iter().any(|value| {
@@ -1,20 +1,25 @@
use std::collections::BTreeMap;
use aether_data::repository::candidates::{RequestCandidateStatus, UpsertRequestCandidateRecord};
use serde_json::json;
use tracing::warn;
use crate::execution_runtime::{ConversionMode, ExecutionStrategy};
use crate::headers::collect_control_headers;
use crate::provider_transport::auth::{
build_openai_passthrough_headers, ensure_upstream_auth_header,
};
use crate::provider_transport::{
apply_local_header_rules, resolve_transport_execution_timeouts,
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_transport_tls_profile,
use crate::ai_pipeline::control_facade::collect_control_headers;
use crate::ai_pipeline::execution_facade::{ConversionMode, ExecutionStrategy};
use crate::ai_pipeline::planner::candidate_runtime_facade::persist_skipped_local_candidate;
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::planner::transport_facade::{
read_provider_transport_snapshot, resolve_local_oauth_request_auth,
LocalResolvedOAuthRequestAuth,
};
use crate::scheduler::current_unix_secs;
use crate::ai_pipeline::provider_transport_facade::auth::{
build_openai_passthrough_headers, ensure_upstream_auth_header,
};
use crate::ai_pipeline::provider_transport_facade::{
apply_local_header_rules, resolve_transport_execution_timeouts,
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_transport_tls_profile,
};
use crate::clock::current_unix_secs;
use crate::{
append_execution_contract_fields_to_value, AppState, GatewayControlSyncDecisionResponse,
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
@@ -47,13 +52,13 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
return None;
};
let transport = match state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
let transport = match read_provider_transport_snapshot(
state,
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
{
Ok(Some(snapshot)) => snapshot,
Ok(None) => {
@@ -107,12 +112,10 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
return None;
}
let resolved_auth = crate::ai_pipeline::conversion::request_conversion_direct_auth(
&transport,
conversion_kind,
);
let resolved_auth =
crate::ai_pipeline::conversion::request_conversion_direct_auth(&transport, conversion_kind);
let oauth_auth = if resolved_auth.is_none() {
match state.resolve_local_oauth_request_auth(&transport).await {
match resolve_local_oauth_request_auth(state, &transport).await {
Ok(Some(LocalResolvedOAuthRequestAuth::Header { name, value })) => Some((name, value)),
Ok(Some(LocalResolvedOAuthRequestAuth::Kiro(_))) => None,
Ok(None) => None,
@@ -160,15 +163,22 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
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(),
spec.require_streaming,
upstream_is_stream,
transport.endpoint.body_rules.as_ref(),
Some(input.auth_context.api_key_id.as_str()),
) {
Some(body) => body,
None => {
@@ -186,29 +196,28 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
}
};
let upstream_url =
match crate::ai_pipeline::planner::standard::build_standard_upstream_url(
parts,
&transport,
&mapped_model,
provider_api_format.as_str(),
spec.require_streaming,
) {
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 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 = build_openai_passthrough_headers(
&parts.headers,
@@ -236,8 +245,17 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
.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 spec.require_streaming {
if upstream_is_stream {
provider_request_headers
.entry("accept".to_string())
.or_insert_with(|| "text/event-stream".to_string());
@@ -278,7 +296,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
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: spec.require_streaming,
upstream_is_stream,
report_kind: Some(spec.report_kind.to_string()),
report_context: Some(append_execution_contract_fields_to_value(
json!({
@@ -319,46 +337,22 @@ pub(super) async fn mark_skipped_local_standard_candidate(
state: &AppState,
input: &LocalStandardDecisionInput,
trace_id: &str,
candidate: &crate::scheduler::GatewayMinimalCandidateSelectionCandidate,
candidate: &aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate,
candidate_index: u32,
candidate_id: &str,
skip_reason: &'static str,
) {
if let Err(err) = state
.upsert_request_candidate(UpsertRequestCandidateRecord {
id: candidate_id.to_string(),
request_id: trace_id.to_string(),
user_id: Some(input.auth_context.user_id.clone()),
api_key_id: Some(input.auth_context.api_key_id.clone()),
username: None,
api_key_name: None,
candidate_index,
retry_index: 0,
provider_id: Some(candidate.provider_id.clone()),
endpoint_id: Some(candidate.endpoint_id.clone()),
key_id: Some(candidate.key_id.clone()),
status: RequestCandidateStatus::Skipped,
skip_reason: Some(skip_reason.to_string()),
is_cached: Some(false),
status_code: None,
error_type: None,
error_message: None,
latency_ms: None,
concurrent_requests: None,
extra_data: None,
required_capabilities: candidate.key_capabilities.clone(),
created_at_unix_secs: None,
started_at_unix_secs: None,
finished_at_unix_secs: Some(current_unix_secs()),
})
.await
{
warn!(
trace_id = %trace_id,
candidate_id = %candidate_id,
skip_reason,
error = ?err,
"gateway local standard decision failed to persist skipped candidate"
);
}
persist_skipped_local_candidate(
state,
trace_id,
&input.auth_context.user_id,
&input.auth_context.api_key_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
current_unix_secs(),
"gateway local standard decision failed to persist skipped candidate",
)
.await;
}
@@ -1,35 +1,20 @@
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum LocalStandardSourceFamily {
Standard,
Gemini,
}
use crate::ai_pipeline::control_facade::GatewayControlAuthContext;
use crate::ai_pipeline::planner::auth_snapshot_facade::GatewayAuthApiKeySnapshot;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum LocalStandardSourceMode {
Chat,
Cli,
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct LocalStandardSpec {
pub(crate) api_format: &'static str,
pub(crate) decision_kind: &'static str,
pub(crate) report_kind: &'static str,
pub(crate) family: LocalStandardSourceFamily,
pub(crate) mode: LocalStandardSourceMode,
pub(crate) require_streaming: bool,
}
pub(crate) use aether_ai_pipeline::planner::standard::family::{
LocalStandardSourceFamily, LocalStandardSourceMode, LocalStandardSpec,
};
#[derive(Debug, Clone)]
pub(super) struct LocalStandardDecisionInput {
pub(super) auth_context: crate::control::GatewayControlAuthContext,
pub(super) auth_context: GatewayControlAuthContext,
pub(super) requested_model: String,
pub(super) auth_snapshot: crate::data::auth::GatewayAuthApiKeySnapshot,
pub(super) auth_snapshot: GatewayAuthApiKeySnapshot,
}
#[derive(Debug, Clone)]
pub(super) struct LocalStandardCandidateAttempt {
pub(super) candidate: crate::scheduler::GatewayMinimalCandidateSelectionCandidate,
pub(super) candidate: aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate,
pub(super) candidate_index: u32,
pub(super) candidate_id: String,
}
@@ -1,33 +0,0 @@
use crate::ai_pipeline::planner::common::{
GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND,
};
use super::super::family::{LocalStandardSourceFamily, LocalStandardSourceMode, LocalStandardSpec};
pub(crate) fn resolve_sync_spec(plan_kind: &str) -> Option<LocalStandardSpec> {
match plan_kind {
GEMINI_CHAT_SYNC_PLAN_KIND => Some(LocalStandardSpec {
api_format: "gemini:chat",
decision_kind: GEMINI_CHAT_SYNC_PLAN_KIND,
report_kind: "gemini_chat_sync_finalize",
family: LocalStandardSourceFamily::Gemini,
mode: LocalStandardSourceMode::Chat,
require_streaming: false,
}),
_ => None,
}
}
pub(crate) fn resolve_stream_spec(plan_kind: &str) -> Option<LocalStandardSpec> {
match plan_kind {
GEMINI_CHAT_STREAM_PLAN_KIND => Some(LocalStandardSpec {
api_format: "gemini:chat",
decision_kind: GEMINI_CHAT_STREAM_PLAN_KIND,
report_kind: "gemini_chat_stream_success",
family: LocalStandardSourceFamily::Gemini,
mode: LocalStandardSourceMode::Chat,
require_streaming: true,
}),
_ => None,
}
}
@@ -1,33 +0,0 @@
use crate::ai_pipeline::planner::common::{
GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND,
};
use super::super::family::{LocalStandardSourceFamily, LocalStandardSourceMode, LocalStandardSpec};
pub(crate) fn resolve_sync_spec(plan_kind: &str) -> Option<LocalStandardSpec> {
match plan_kind {
GEMINI_CLI_SYNC_PLAN_KIND => Some(LocalStandardSpec {
api_format: "gemini:cli",
decision_kind: GEMINI_CLI_SYNC_PLAN_KIND,
report_kind: "gemini_cli_sync_finalize",
family: LocalStandardSourceFamily::Gemini,
mode: LocalStandardSourceMode::Cli,
require_streaming: false,
}),
_ => None,
}
}
pub(crate) fn resolve_stream_spec(plan_kind: &str) -> Option<LocalStandardSpec> {
match plan_kind {
GEMINI_CLI_STREAM_PLAN_KIND => Some(LocalStandardSpec {
api_format: "gemini:cli",
decision_kind: GEMINI_CLI_STREAM_PLAN_KIND,
report_kind: "gemini_cli_stream_success",
family: LocalStandardSourceFamily::Gemini,
mode: LocalStandardSourceMode::Cli,
require_streaming: true,
}),
_ => None,
}
}
@@ -1,13 +1,23 @@
use crate::control::GatewayControlDecision;
use crate::ai_pipeline::control_facade::GatewayControlDecision;
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
use aether_ai_pipeline::planner::standard::gemini::{
resolve_stream_spec as resolve_pipeline_stream_spec,
resolve_sync_spec as resolve_pipeline_sync_spec,
};
use super::family::{
maybe_build_stream_via_standard_family_payload, maybe_build_sync_via_standard_family_payload,
};
pub(crate) use crate::ai_pipeline::conversion::request::normalize_gemini_request_to_openai_chat_request;
pub(crate) mod chat;
pub(crate) mod cli;
pub(crate) fn resolve_sync_spec(plan_kind: &str) -> Option<super::family::LocalStandardSpec> {
resolve_pipeline_sync_spec(plan_kind)
}
pub(crate) fn resolve_stream_spec(plan_kind: &str) -> Option<super::family::LocalStandardSpec> {
resolve_pipeline_stream_spec(plan_kind)
}
pub(crate) async fn maybe_build_sync_local_gemini_decision_payload(
state: &AppState,
@@ -24,9 +34,7 @@ pub(crate) async fn maybe_build_sync_local_gemini_decision_payload(
decision,
body_json,
plan_kind,
|plan_kind| {
chat::resolve_sync_spec(plan_kind).or_else(|| cli::resolve_sync_spec(plan_kind))
},
resolve_sync_spec,
)
.await
}
@@ -46,9 +54,7 @@ pub(crate) async fn maybe_build_stream_local_gemini_decision_payload(
decision,
body_json,
plan_kind,
|plan_kind| {
chat::resolve_stream_spec(plan_kind).or_else(|| cli::resolve_stream_spec(plan_kind))
},
resolve_stream_spec,
)
.await
}
@@ -4,7 +4,7 @@ use super::{
augment_sync_report_context, generic_decision_missing_exact_provider_request,
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
};
use crate::provider_transport::ensure_upstream_auth_header;
use crate::ai_pipeline::provider_transport_facade::ensure_upstream_auth_header;
use crate::{GatewayControlSyncDecisionResponse, GatewayError};
pub(crate) fn build_gemini_sync_plan_from_decision(
@@ -1,79 +1,59 @@
use std::collections::BTreeMap;
use serde_json::{json, Map, Value};
use super::{
claude::normalize_claude_request_to_openai_chat_request,
gemini::normalize_gemini_request_to_openai_chat_request,
};
use crate::ai_pipeline::conversion::request::{
convert_openai_chat_request_to_claude_request, convert_openai_chat_request_to_gemini_request,
convert_openai_chat_request_to_openai_cli_request,
normalize_openai_cli_request_to_openai_chat_request,
};
use crate::ai_pipeline::conversion::{request_conversion_kind, RequestConversionKind};
use crate::provider_transport::apply_local_body_rules;
use crate::provider_transport::url::{
use super::codex::apply_codex_openai_cli_special_body_edits;
use crate::ai_pipeline::planner::transport_facade::GatewayProviderTransportSnapshot;
use crate::ai_pipeline::provider_transport_facade::apply_local_body_rules;
use crate::ai_pipeline::provider_transport_facade::url::{
build_claude_messages_url, build_gemini_content_url, build_openai_chat_url,
build_openai_cli_url, build_passthrough_path_url,
};
use aether_ai_pipeline::planner::matrix::build_standard_request_body_from_canonical;
pub(crate) use aether_ai_pipeline::planner::standard::normalize_standard_request_to_openai_chat_request;
use serde_json::{json, Value};
pub(crate) fn build_standard_request_body(
body_json: &Value,
client_api_format: &str,
mapped_model: &str,
provider_type: &str,
provider_api_format: &str,
request_path: &str,
upstream_is_stream: bool,
body_rules: Option<&Value>,
user_api_key_id: Option<&str>,
) -> Option<Value> {
let canonical_request = normalize_standard_request_to_openai_chat_request(
body_json,
client_api_format,
request_path,
)?;
let conversion_kind = request_conversion_kind(client_api_format, provider_api_format)?;
let mut provider_request_body = match conversion_kind {
RequestConversionKind::ToOpenAIChat => {
build_openai_chat_request_body(&canonical_request, mapped_model, upstream_is_stream)?
}
RequestConversionKind::ToOpenAIFamilyCli => {
convert_openai_chat_request_to_openai_cli_request(
&canonical_request,
mapped_model,
upstream_is_stream,
false,
)?
}
RequestConversionKind::ToOpenAICompact => {
convert_openai_chat_request_to_openai_cli_request(
&canonical_request,
mapped_model,
false,
true,
)?
}
RequestConversionKind::ToClaudeStandard => convert_openai_chat_request_to_claude_request(
&canonical_request,
mapped_model,
upstream_is_stream,
)?,
RequestConversionKind::ToGeminiStandard => convert_openai_chat_request_to_gemini_request(
&canonical_request,
mapped_model,
upstream_is_stream,
)?,
};
if cfg!(test) {
println!("canonical_request: {canonical_request:#?}");
}
let mut provider_request_body = build_standard_request_body_from_canonical(
&canonical_request,
mapped_model,
provider_api_format,
upstream_is_stream,
)?;
if cfg!(test) {
println!("provider_request_body before rules: {provider_request_body:#?}");
}
if !apply_local_body_rules(&mut provider_request_body, body_rules, Some(body_json)) {
return None;
}
apply_codex_openai_cli_special_body_edits(
&mut provider_request_body,
provider_type,
provider_api_format,
body_rules,
user_api_key_id,
);
Some(provider_request_body)
}
pub(crate) fn build_standard_upstream_url(
parts: &http::request::Parts,
transport: &crate::provider_transport::GatewayProviderTransportSnapshot,
transport: &GatewayProviderTransportSnapshot,
mapped_model: &str,
provider_api_format: &str,
upstream_is_stream: bool,
@@ -119,42 +99,6 @@ pub(crate) fn build_standard_upstream_url(
}
}
pub(crate) fn normalize_standard_request_to_openai_chat_request(
body_json: &Value,
client_api_format: &str,
request_path: &str,
) -> Option<Value> {
match client_api_format.trim().to_ascii_lowercase().as_str() {
"openai:chat" => Some(body_json.clone()),
"openai:cli" | "openai:compact" => {
normalize_openai_cli_request_to_openai_chat_request(body_json)
}
"claude:chat" | "claude:cli" => normalize_claude_request_to_openai_chat_request(body_json),
"gemini:chat" | "gemini:cli" => {
normalize_gemini_request_to_openai_chat_request(body_json, request_path)
}
_ => None,
}
}
fn build_openai_chat_request_body(
body_json: &Value,
mapped_model: &str,
upstream_is_stream: bool,
) -> Option<Value> {
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())),
);
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));
}
Some(Value::Object(provider_request_body))
}
#[cfg(test)]
mod tests {
use super::build_standard_request_body;
@@ -178,10 +122,12 @@ mod tests {
&request,
"claude:chat",
"gpt-5",
"openai",
"openai:chat",
"/v1/messages",
false,
None,
None,
)
.expect("claude chat should convert to openai chat");
@@ -210,10 +156,12 @@ mod tests {
&request,
"gemini:chat",
"claude-sonnet-4-5",
"anthropic",
"claude:chat",
"/v1beta/models/gemini-2.5-pro:generateContent",
false,
None,
None,
)
.expect("gemini chat should convert to claude chat");
@@ -244,10 +192,12 @@ mod tests {
&request,
"claude:cli",
"gemini-2.5-pro",
"google",
"gemini:cli",
"/v1/messages",
false,
None,
None,
)
.expect("claude cli should convert to gemini cli");
@@ -257,4 +207,133 @@ mod tests {
"Need CLI output"
);
}
#[test]
fn builds_openai_cli_request_from_claude_cli_source_with_forced_stream() {
let request = json!({
"model": "claude-sonnet-4-5",
"messages": [
{
"role": "user",
"content": [{"type": "text", "text": "Need OpenAI CLI output"}]
}
],
"max_tokens": 64
});
let converted = build_standard_request_body(
&request,
"claude:cli",
"gpt-5",
"openai",
"openai:cli",
"/v1/messages",
true,
None,
None,
)
.expect("claude cli should convert to openai cli");
assert_eq!(converted["model"], "gpt-5");
assert_eq!(converted["input"][0]["role"], "user");
assert_eq!(converted["input"][0]["content"][0]["type"], "input_text");
assert_eq!(
converted["input"][0]["content"][0]["text"],
"Need OpenAI CLI output"
);
assert_eq!(converted["stream"], true);
}
#[test]
fn strips_metadata_for_codex_openai_cli_requests() {
let request = json!({
"model": "claude-sonnet-4-5",
"metadata": {"trace_id": "abc"},
"messages": [{
"role": "user",
"content": [{"type": "text", "text": "Need OpenAI CLI output"}]
}],
"max_tokens": 64
});
let converted = build_standard_request_body(
&request,
"claude:cli",
"gpt-5",
"codex",
"openai:cli",
"/v1/messages",
true,
None,
None,
)
.expect("claude cli should convert to codex cli");
assert!(converted.get("metadata").is_none());
}
#[test]
fn applies_codex_defaults_unless_body_rules_handle_the_field() {
let request = json!({
"model": "claude-sonnet-4-5",
"metadata": {"trace_id": "abc"},
"messages": [{
"role": "user",
"content": [{"type": "text", "text": "Need OpenAI CLI output"}]
}],
"max_tokens": 64
});
let body_rules = json!([
{"action":"set","path":"store","value":true},
{"action":"set","path":"instructions","value":"Custom instructions"},
{"action":"set","path":"metadata","value":{"trace_id":"keep-me"}}
]);
let converted = build_standard_request_body(
&request,
"claude:cli",
"gpt-5",
"codex",
"openai:cli",
"/v1/messages",
true,
Some(&body_rules),
None,
)
.expect("claude cli should convert to codex cli");
assert_eq!(converted["store"], true);
assert_eq!(converted["instructions"], "Custom instructions");
assert_eq!(converted["metadata"]["trace_id"], "keep-me");
}
#[test]
fn injects_codex_prompt_cache_key_for_standard_requests() {
let request = json!({
"model": "claude-sonnet-4-5",
"messages": [{
"role": "user",
"content": [{"type": "text", "text": "Need OpenAI CLI output"}]
}],
"max_tokens": 64
});
let converted = build_standard_request_body(
&request,
"claude:cli",
"gpt-5",
"codex",
"openai:cli",
"/v1/messages",
true,
None,
Some("key-123"),
)
.expect("claude cli should convert to codex cli");
assert_eq!(
converted["prompt_cache_key"],
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
);
}
}
@@ -3,16 +3,18 @@
//! This groups the standard planning surface in one place:
//! request-side conversion, matrix registry, and decision payload builders.
use crate::control::GatewayControlDecision;
use crate::ai_pipeline::control_facade::GatewayControlDecision;
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
pub(crate) mod claude;
mod codex;
pub(crate) mod family;
pub(crate) mod gemini;
mod matrix;
mod normalize;
pub(crate) mod openai;
pub(crate) use self::codex::apply_codex_openai_cli_special_headers;
pub(crate) use self::matrix::{
build_standard_request_body, build_standard_upstream_url,
normalize_standard_request_to_openai_chat_request,
@@ -3,20 +3,23 @@ use std::collections::BTreeMap;
use serde_json::{json, Map, Value};
use url::form_urlencoded;
use crate::ai_pipeline::conversion::request::{
convert_openai_chat_request_to_claude_request, convert_openai_chat_request_to_gemini_request,
convert_openai_chat_request_to_openai_cli_request, extract_openai_text_content,
normalize_openai_cli_request_to_openai_chat_request, parse_openai_tool_result_content,
};
use super::codex::apply_codex_openai_cli_special_body_edits;
use crate::ai_pipeline::conversion::{request_conversion_kind, RequestConversionKind};
use crate::provider_transport::antigravity::{
use crate::ai_pipeline::planner::transport_facade::GatewayProviderTransportSnapshot;
use crate::ai_pipeline::provider_transport_facade::antigravity::{
build_antigravity_v1internal_url, AntigravityRequestUrlAction,
};
use crate::provider_transport::apply_local_body_rules;
use crate::provider_transport::url::{
use crate::ai_pipeline::provider_transport_facade::apply_local_body_rules;
use crate::ai_pipeline::provider_transport_facade::url::{
build_claude_messages_url, build_gemini_content_url, build_openai_chat_url,
build_openai_cli_url, build_passthrough_path_url,
};
use aether_ai_pipeline::planner::standard::normalize::{
build_cross_format_openai_chat_request_body as pipeline_build_cross_format_openai_chat_request_body,
build_cross_format_openai_cli_request_body as pipeline_build_cross_format_openai_cli_request_body,
build_local_openai_chat_request_body as pipeline_build_local_openai_chat_request_body,
build_local_openai_cli_request_body as pipeline_build_local_openai_cli_request_body,
};
pub(crate) fn build_local_openai_chat_request_body(
body_json: &Value,
@@ -24,17 +27,8 @@ pub(crate) fn build_local_openai_chat_request_body(
upstream_is_stream: bool,
body_rules: Option<&Value>,
) -> Option<Value> {
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())),
);
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));
}
let mut provider_request_body = Value::Object(provider_request_body);
let mut provider_request_body =
pipeline_build_local_openai_chat_request_body(body_json, mapped_model, upstream_is_stream)?;
if !apply_local_body_rules(&mut provider_request_body, body_rules, Some(body_json)) {
return None;
}
@@ -43,7 +37,7 @@ pub(crate) fn build_local_openai_chat_request_body(
pub(crate) fn build_local_openai_chat_upstream_url(
parts: &http::request::Parts,
transport: &crate::provider_transport::GatewayProviderTransportSnapshot,
transport: &GatewayProviderTransportSnapshot,
) -> Option<String> {
let custom_path = transport
.endpoint
@@ -66,45 +60,34 @@ pub(crate) fn build_local_openai_chat_upstream_url(
pub(crate) fn build_cross_format_openai_chat_request_body(
body_json: &Value,
mapped_model: &str,
provider_type: &str,
provider_api_format: &str,
upstream_is_stream: bool,
body_rules: Option<&Value>,
user_api_key_id: Option<&str>,
) -> Option<Value> {
let conversion_kind = request_conversion_kind("openai:chat", provider_api_format)?;
let mut provider_request_body = match conversion_kind {
RequestConversionKind::ToClaudeStandard => convert_openai_chat_request_to_claude_request(
body_json,
mapped_model,
upstream_is_stream,
)?,
RequestConversionKind::ToGeminiStandard => convert_openai_chat_request_to_gemini_request(
body_json,
mapped_model,
upstream_is_stream,
)?,
RequestConversionKind::ToOpenAIFamilyCli => {
convert_openai_chat_request_to_openai_cli_request(
body_json,
mapped_model,
upstream_is_stream,
false,
)?
}
RequestConversionKind::ToOpenAICompact => {
convert_openai_chat_request_to_openai_cli_request(body_json, mapped_model, false, true)?
}
_ => return None,
};
let mut provider_request_body = pipeline_build_cross_format_openai_chat_request_body(
body_json,
mapped_model,
provider_api_format,
upstream_is_stream,
)?;
if !apply_local_body_rules(&mut provider_request_body, body_rules, Some(body_json)) {
return None;
}
apply_codex_openai_cli_special_body_edits(
&mut provider_request_body,
provider_type,
provider_api_format,
body_rules,
user_api_key_id,
);
Some(provider_request_body)
}
pub(crate) fn build_cross_format_openai_chat_upstream_url(
parts: &http::request::Parts,
transport: &crate::provider_transport::GatewayProviderTransportSnapshot,
transport: &GatewayProviderTransportSnapshot,
mapped_model: &str,
provider_api_format: &str,
upstream_is_stream: bool,
@@ -151,22 +134,23 @@ pub(crate) fn build_local_openai_cli_request_body(
body_json: &Value,
mapped_model: &str,
require_streaming: bool,
provider_type: &str,
provider_api_format: &str,
body_rules: Option<&Value>,
user_api_key_id: Option<&str>,
) -> Option<Value> {
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())),
);
provider_request_body.insert("model".to_string(), Value::String(mapped_model.to_string()));
if require_streaming {
provider_request_body.insert("stream".to_string(), Value::Bool(true));
}
let mut provider_request_body = Value::Object(provider_request_body);
let mut provider_request_body =
pipeline_build_local_openai_cli_request_body(body_json, mapped_model, require_streaming)?;
if !apply_local_body_rules(&mut provider_request_body, body_rules, Some(body_json)) {
return None;
}
apply_codex_openai_cli_special_body_edits(
&mut provider_request_body,
provider_type,
provider_api_format,
body_rules,
user_api_key_id,
);
Some(provider_request_body)
}
@@ -176,48 +160,33 @@ pub(crate) fn build_cross_format_openai_cli_request_body(
client_api_format: &str,
provider_api_format: &str,
upstream_is_stream: bool,
provider_type: &str,
body_rules: Option<&Value>,
user_api_key_id: Option<&str>,
) -> Option<Value> {
let chat_like_request = normalize_openai_cli_request_to_openai_chat_request(body_json)?;
let conversion_kind = request_conversion_kind(client_api_format, provider_api_format)?;
let mut provider_request_body = match conversion_kind {
RequestConversionKind::ToOpenAIFamilyCli => {
convert_openai_chat_request_to_openai_cli_request(
&chat_like_request,
mapped_model,
upstream_is_stream,
false,
)?
}
RequestConversionKind::ToOpenAICompact => {
convert_openai_chat_request_to_openai_cli_request(
&chat_like_request,
mapped_model,
false,
true,
)?
}
RequestConversionKind::ToClaudeStandard => convert_openai_chat_request_to_claude_request(
&chat_like_request,
mapped_model,
upstream_is_stream,
)?,
RequestConversionKind::ToGeminiStandard => convert_openai_chat_request_to_gemini_request(
&chat_like_request,
mapped_model,
upstream_is_stream,
)?,
_ => return None,
};
let mut provider_request_body = pipeline_build_cross_format_openai_cli_request_body(
body_json,
mapped_model,
client_api_format,
provider_api_format,
upstream_is_stream,
)?;
if !apply_local_body_rules(&mut provider_request_body, body_rules, Some(body_json)) {
return None;
}
apply_codex_openai_cli_special_body_edits(
&mut provider_request_body,
provider_type,
provider_api_format,
body_rules,
user_api_key_id,
);
Some(provider_request_body)
}
pub(crate) fn build_local_openai_cli_upstream_url(
parts: &http::request::Parts,
transport: &crate::provider_transport::GatewayProviderTransportSnapshot,
transport: &GatewayProviderTransportSnapshot,
compact: bool,
) -> Option<String> {
let custom_path = transport
@@ -241,7 +210,7 @@ pub(crate) fn build_local_openai_cli_upstream_url(
pub(crate) fn build_cross_format_openai_cli_upstream_url(
parts: &http::request::Parts,
transport: &crate::provider_transport::GatewayProviderTransportSnapshot,
transport: &GatewayProviderTransportSnapshot,
mapped_model: &str,
client_api_format: &str,
provider_api_format: &str,
@@ -325,6 +294,8 @@ mod tests {
"openai:compact",
"openai:cli",
false,
"openai",
None,
None,
)
.expect("compact to openai cli body should build");
@@ -333,4 +304,119 @@ mod tests {
assert_eq!(provider_request_body["input"][0]["type"], "message");
assert_eq!(provider_request_body["input"][0]["role"], "user");
}
#[test]
fn strips_metadata_for_codex_openai_cli_requests() {
let body_json = json!({
"model": "claude-sonnet-4-5",
"metadata": {"trace_id": "abc"},
"messages": [{
"role": "user",
"content": [{"type": "text", "text": "hello"}]
}],
});
let provider_request_body = build_cross_format_openai_cli_request_body(
&body_json,
"gpt-5-upstream",
"claude:cli",
"openai:cli",
true,
"codex",
None,
None,
)
.expect("claude cli to codex request should build");
assert!(provider_request_body.get("metadata").is_none());
}
#[test]
fn applies_codex_defaults_unless_body_rules_handle_the_field() {
let body_json = json!({
"model": "claude-sonnet-4-5",
"messages": [{
"role": "user",
"content": [{"type": "text", "text": "hello"}]
}],
"metadata": {"trace_id": "abc"},
"store": true
});
let body_rules = json!([
{"action":"set","path":"store","value":true},
{"action":"set","path":"instructions","value":"Custom instructions"},
{"action":"set","path":"metadata","value":{"trace_id":"keep-me"}}
]);
let provider_request_body = build_cross_format_openai_cli_request_body(
&body_json,
"gpt-5-upstream",
"claude:cli",
"openai:cli",
true,
"codex",
Some(&body_rules),
None,
)
.expect("claude cli to codex request should build");
assert_eq!(provider_request_body["store"], true);
assert_eq!(provider_request_body["instructions"], "Custom instructions");
assert_eq!(provider_request_body["metadata"]["trace_id"], "keep-me");
}
#[test]
fn injects_codex_prompt_cache_key_for_openai_cli_cross_format_requests() {
let body_json = json!({
"model": "claude-sonnet-4-5",
"messages": [{
"role": "user",
"content": [{"type": "text", "text": "hello"}]
}],
});
let provider_request_body = build_cross_format_openai_cli_request_body(
&body_json,
"gpt-5-upstream",
"claude:cli",
"openai:cli",
true,
"codex",
None,
Some("key-123"),
)
.expect("claude cli to codex request should build");
assert_eq!(
provider_request_body["prompt_cache_key"],
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
);
}
#[test]
fn injects_codex_prompt_cache_key_for_openai_chat_cross_format_requests() {
let body_json = json!({
"model": "gpt-5",
"messages": [{
"role": "user",
"content": "hello"
}],
});
let provider_request_body = super::build_cross_format_openai_chat_request_body(
&body_json,
"gpt-5-upstream",
"codex",
"openai:cli",
false,
None,
Some("key-123"),
)
.expect("openai chat to codex request should build");
assert_eq!(
provider_request_body["prompt_cache_key"],
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
);
}
}
@@ -1,55 +1,61 @@
use std::collections::BTreeMap;
use aether_data::repository::candidates::{RequestCandidateStatus, UpsertRequestCandidateRecord};
use serde_json::json;
use tracing::warn;
use uuid::Uuid;
use crate::ai_pipeline::control_facade::{collect_control_headers, GatewayControlAuthContext};
use crate::ai_pipeline::conversion::{
request_conversion_direct_auth, request_conversion_kind, request_conversion_transport_supported,
};
use crate::ai_pipeline::execution_facade::{ConversionMode, ExecutionStrategy};
use crate::ai_pipeline::planner::auth_snapshot_facade::GatewayAuthApiKeySnapshot;
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
use crate::ai_pipeline::planner::candidate_runtime_facade::{
persist_available_local_candidate, persist_skipped_local_candidate,
};
use crate::ai_pipeline::planner::common::{
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
OPENAI_CHAT_STREAM_PLAN_KIND,
};
use crate::ai_pipeline::planner::executor_facade::mark_unused_local_candidate_items;
use crate::ai_pipeline::planner::plan_builders::{
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
};
use crate::execution_runtime::{ConversionMode, ExecutionStrategy};
use crate::headers::collect_control_headers;
use crate::provider_transport::auth::{
use crate::ai_pipeline::planner::transport_facade::{
read_provider_transport_snapshot, resolve_local_oauth_request_auth,
GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
};
use crate::ai_pipeline::provider_transport_facade::auth::{
build_openai_passthrough_headers, ensure_upstream_auth_header, resolve_local_openai_chat_auth,
};
use crate::provider_transport::policy::supports_local_openai_chat_transport;
use crate::provider_transport::{
use crate::ai_pipeline::provider_transport_facade::policy::supports_local_openai_chat_transport;
use crate::ai_pipeline::provider_transport_facade::{
apply_local_header_rules, resolve_transport_execution_timeouts,
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_transport_tls_profile,
LocalResolvedOAuthRequestAuth,
};
use crate::scheduler::{
record_local_request_candidate_status, GatewayMinimalCandidateSelectionCandidate,
};
use crate::{
append_execution_contract_fields_to_value, AppState, GatewayControlSyncDecisionResponse,
};
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use super::plans::current_unix_secs;
use crate::ai_pipeline::planner::standard::{
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,
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,
};
#[derive(Debug, Clone)]
pub(super) struct LocalOpenAiChatDecisionInput {
pub(super) auth_context: crate::control::GatewayControlAuthContext,
pub(super) auth_context: GatewayControlAuthContext,
pub(super) requested_model: String,
pub(super) auth_snapshot: crate::data::auth::GatewayAuthApiKeySnapshot,
pub(super) auth_snapshot: GatewayAuthApiKeySnapshot,
}
#[derive(Debug, Clone)]
pub(super) struct LocalOpenAiChatCandidateAttempt {
pub(super) candidate: GatewayMinimalCandidateSelectionCandidate,
pub(super) candidate: SchedulerMinimalCandidateSelectionCandidate,
pub(super) candidate_index: u32,
pub(super) candidate_id: String,
}
@@ -70,13 +76,13 @@ pub(super) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
candidate_index,
candidate_id,
} = attempt;
let transport = match state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
let transport = match read_provider_transport_snapshot(
state,
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
{
Ok(Some(snapshot)) => snapshot,
Ok(None) => {
@@ -171,13 +177,13 @@ async fn build_same_format_local_openai_chat_decision_payload_for_candidate(
trace_id: &str,
body_json: &serde_json::Value,
input: &LocalOpenAiChatDecisionInput,
candidate: &GatewayMinimalCandidateSelectionCandidate,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
candidate_index: u32,
candidate_id: &str,
decision_kind: &str,
report_kind: &str,
upstream_is_stream: bool,
transport: &crate::provider_transport::GatewayProviderTransportSnapshot,
transport: &GatewayProviderTransportSnapshot,
) -> Option<GatewayControlSyncDecisionResponse> {
if !supports_local_openai_chat_transport(transport) {
mark_skipped_local_openai_chat_candidate(
@@ -194,7 +200,7 @@ async fn build_same_format_local_openai_chat_decision_payload_for_candidate(
}
let oauth_auth = if resolve_local_openai_chat_auth(transport).is_none() {
match state.resolve_local_oauth_request_auth(transport).await {
match resolve_local_oauth_request_auth(state, transport).await {
Ok(Some(LocalResolvedOAuthRequestAuth::Header { name, value })) => Some((name, value)),
Ok(Some(LocalResolvedOAuthRequestAuth::Kiro(_))) => None,
Ok(None) => None,
@@ -300,6 +306,15 @@ async fn build_same_format_local_openai_chat_decision_payload_for_candidate(
.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
@@ -394,12 +409,12 @@ async fn build_cross_format_local_openai_chat_decision_payload_for_candidate(
trace_id: &str,
body_json: &serde_json::Value,
input: &LocalOpenAiChatDecisionInput,
candidate: &GatewayMinimalCandidateSelectionCandidate,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
candidate_index: u32,
candidate_id: &str,
decision_kind: &str,
upstream_is_stream: bool,
transport: &crate::provider_transport::GatewayProviderTransportSnapshot,
transport: &GatewayProviderTransportSnapshot,
provider_api_format: &str,
) -> Option<GatewayControlSyncDecisionResponse> {
let provider_api_format = provider_api_format.trim().to_ascii_lowercase();
@@ -425,7 +440,7 @@ async fn build_cross_format_local_openai_chat_decision_payload_for_candidate(
let resolve_auth = request_conversion_direct_auth(transport, conversion_kind);
let oauth_auth = if resolve_auth.is_none() {
match state.resolve_local_oauth_request_auth(transport).await {
match resolve_local_oauth_request_auth(state, transport).await {
Ok(Some(LocalResolvedOAuthRequestAuth::Header { name, value })) => Some((name, value)),
Ok(Some(LocalResolvedOAuthRequestAuth::Kiro(_))) => None,
Ok(None) => None,
@@ -476,9 +491,11 @@ async fn build_cross_format_local_openai_chat_decision_payload_for_candidate(
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,
@@ -539,6 +556,15 @@ async fn build_cross_format_local_openai_chat_decision_payload_for_candidate(
.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
@@ -636,7 +662,7 @@ pub(super) async fn mark_skipped_local_openai_chat_candidate(
state: &AppState,
input: &LocalOpenAiChatDecisionInput,
trace_id: &str,
candidate: &GatewayMinimalCandidateSelectionCandidate,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
candidate_index: u32,
candidate_id: &str,
skip_reason: &'static str,
@@ -649,50 +675,26 @@ pub(super) async fn mark_skipped_local_openai_chat_candidate(
*diagnostic.skipped_candidate_count.get_or_insert(0) += 1;
});
let terminal_unix_secs = current_unix_secs();
if let Err(err) = state
.upsert_request_candidate(UpsertRequestCandidateRecord {
id: candidate_id.to_string(),
request_id: trace_id.to_string(),
user_id: Some(input.auth_context.user_id.clone()),
api_key_id: Some(input.auth_context.api_key_id.clone()),
username: None,
api_key_name: None,
candidate_index,
retry_index: 0,
provider_id: Some(candidate.provider_id.clone()),
endpoint_id: Some(candidate.endpoint_id.clone()),
key_id: Some(candidate.key_id.clone()),
status: RequestCandidateStatus::Skipped,
skip_reason: Some(skip_reason.to_string()),
is_cached: Some(false),
status_code: None,
error_type: None,
error_message: None,
latency_ms: None,
concurrent_requests: None,
extra_data: None,
required_capabilities: candidate.key_capabilities.clone(),
created_at_unix_secs: None,
started_at_unix_secs: None,
finished_at_unix_secs: Some(terminal_unix_secs),
})
.await
{
warn!(
trace_id = %trace_id,
candidate_id = %candidate_id,
skip_reason,
error = ?err,
"gateway local openai chat decision failed to persist skipped candidate"
);
}
persist_skipped_local_candidate(
state,
trace_id,
&input.auth_context.user_id,
&input.auth_context.api_key_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
terminal_unix_secs,
"gateway local openai chat decision failed to persist skipped candidate",
)
.await;
}
pub(super) async fn materialize_local_openai_chat_candidate_attempts(
state: &AppState,
trace_id: &str,
input: &LocalOpenAiChatDecisionInput,
candidates: Vec<GatewayMinimalCandidateSelectionCandidate>,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
) -> Vec<LocalOpenAiChatCandidateAttempt> {
let candidates = prefer_local_tunnel_owner_candidates(state, candidates).await;
let created_at_unix_secs = current_unix_secs();
@@ -727,46 +729,19 @@ pub(super) async fn materialize_local_openai_chat_candidate_attempts(
candidate.endpoint_api_format.trim(),
);
let candidate_id = match state
.upsert_request_candidate(UpsertRequestCandidateRecord {
id: generated_candidate_id.clone(),
request_id: trace_id.to_string(),
user_id: Some(input.auth_context.user_id.clone()),
api_key_id: Some(input.auth_context.api_key_id.clone()),
username: None,
api_key_name: None,
candidate_index: candidate_index as u32,
retry_index: 0,
provider_id: Some(candidate.provider_id.clone()),
endpoint_id: Some(candidate.endpoint_id.clone()),
key_id: Some(candidate.key_id.clone()),
status: RequestCandidateStatus::Available,
skip_reason: None,
is_cached: Some(false),
status_code: None,
error_type: None,
error_message: None,
latency_ms: None,
concurrent_requests: None,
extra_data: Some(extra_data),
required_capabilities: candidate.key_capabilities.clone(),
created_at_unix_secs: Some(created_at_unix_secs),
started_at_unix_secs: None,
finished_at_unix_secs: None,
})
.await
{
Ok(Some(stored)) => stored.id,
Ok(None) => generated_candidate_id.clone(),
Err(err) => {
warn!(
trace_id = %trace_id,
error = ?err,
"gateway local openai chat decision request candidate upsert failed"
);
generated_candidate_id.clone()
}
};
let candidate_id = persist_available_local_candidate(
state,
trace_id,
&input.auth_context.user_id,
&input.auth_context.api_key_id,
&candidate,
candidate_index as u32,
&generated_candidate_id,
Some(extra_data),
created_at_unix_secs,
"gateway local openai chat decision request candidate upsert failed",
)
.await;
attempts.push(LocalOpenAiChatCandidateAttempt {
candidate,
@@ -782,21 +757,13 @@ pub(super) async fn mark_unused_local_openai_chat_candidates<T>(state: &AppState
where
T: LocalOpenAiChatPlanAndReport,
{
for plan_and_report in remaining {
record_local_request_candidate_status(
state,
plan_and_report.plan(),
plan_and_report.report_context(),
RequestCandidateStatus::Unused,
None,
None,
None,
None,
None,
None,
)
.await;
}
mark_unused_local_candidate_items(
state,
remaining,
|item| item.plan(),
|item| item.report_context(),
)
.await;
}
pub(super) trait LocalOpenAiChatPlanAndReport {
@@ -1,10 +1,10 @@
use serde_json::{Map, Value};
use serde_json::Value;
use tracing::warn;
use crate::ai_pipeline::control_facade::GatewayControlDecision;
use crate::ai_pipeline::planner::common::{
OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND,
};
use crate::control::GatewayControlDecision;
use crate::{
AppState, GatewayControlSyncDecisionResponse, GatewayError, LocalExecutionRuntimeMissDiagnostic,
};
@@ -30,10 +30,7 @@ pub(crate) async fn build_local_openai_chat_sync_plan_and_reports_for_kind(
decision: &GatewayControlDecision,
body_json: &serde_json::Value,
plan_kind: &str,
) -> Result<
Vec<crate::ai_pipeline::planner::plan_builders::LocalSyncPlanAndReport>,
GatewayError,
> {
) -> Result<Vec<crate::ai_pipeline::planner::plan_builders::LocalSyncPlanAndReport>, GatewayError> {
build_local_openai_chat_sync_plan_and_reports(
state, parts, trace_id, decision, body_json, plan_kind,
)
@@ -47,10 +44,8 @@ pub(crate) async fn build_local_openai_chat_stream_plan_and_reports_for_kind(
decision: &GatewayControlDecision,
body_json: &serde_json::Value,
plan_kind: &str,
) -> Result<
Vec<crate::ai_pipeline::planner::plan_builders::LocalStreamPlanAndReport>,
GatewayError,
> {
) -> Result<Vec<crate::ai_pipeline::planner::plan_builders::LocalStreamPlanAndReport>, GatewayError>
{
build_local_openai_chat_stream_plan_and_reports(
state, parts, trace_id, decision, body_json, plan_kind,
)
@@ -205,76 +200,3 @@ pub(crate) async fn maybe_build_stream_local_decision_payload(
Ok(None)
}
pub(crate) fn parse_openai_stop_sequences(stop: Option<&Value>) -> Option<Vec<Value>> {
match stop {
Some(Value::String(value)) if !value.trim().is_empty() => {
Some(vec![Value::String(value.clone())])
}
Some(Value::Array(values)) => Some(
values
.iter()
.filter_map(|value| value.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| Value::String(value.to_string()))
.collect::<Vec<_>>(),
)
.filter(|values| !values.is_empty()),
_ => None,
}
}
pub(crate) fn resolve_openai_chat_max_tokens(request: &Map<String, Value>) -> u64 {
request
.get("max_completion_tokens")
.and_then(value_as_u64)
.or_else(|| request.get("max_tokens").and_then(value_as_u64))
.unwrap_or(4096)
}
pub(crate) fn value_as_u64(value: &Value) -> Option<u64> {
value
.as_u64()
.or_else(|| value.as_i64().and_then(|value| u64::try_from(value).ok()))
}
pub(crate) fn copy_request_number_field(
request: &Map<String, Value>,
target: &mut Map<String, Value>,
key: &str,
) {
copy_request_number_field_as(request, target, key, key);
}
pub(crate) fn copy_request_number_field_as(
request: &Map<String, Value>,
target: &mut Map<String, Value>,
source_key: &str,
target_key: &str,
) {
if let Some(value) = request.get(source_key).cloned() {
if value.is_number() {
target.insert(target_key.to_string(), value);
}
}
}
pub(crate) fn map_openai_reasoning_effort_to_claude_output(value: &str) -> Option<&'static str> {
match value.trim().to_ascii_lowercase().as_str() {
"low" => Some("low"),
"medium" => Some("medium"),
"high" | "xhigh" => Some("high"),
_ => None,
}
}
pub(crate) fn map_openai_reasoning_effort_to_gemini_budget(value: &str) -> Option<u64> {
match value.trim().to_ascii_lowercase().as_str() {
"low" => Some(1024),
"medium" => Some(4096),
"high" => Some(8192),
"xhigh" => Some(16_384),
_ => None,
}
}
@@ -1,6 +1,7 @@
use std::collections::{BTreeMap, BTreeSet};
use std::time::{SystemTime, UNIX_EPOCH};
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use tracing::warn;
use super::{
@@ -8,6 +9,9 @@ use super::{
maybe_build_local_openai_chat_decision_payload_for_candidate, AppState, GatewayControlDecision,
GatewayError, LocalExecutionRuntimeMissDiagnostic, LocalOpenAiChatDecisionInput,
};
use crate::ai_pipeline::planner::auth_snapshot_facade::{
read_auth_api_key_snapshot, GatewayAuthApiKeySnapshot,
};
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
use crate::ai_pipeline::planner::common::{
OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND,
@@ -16,9 +20,7 @@ use crate::ai_pipeline::planner::plan_builders::{
build_openai_chat_stream_plan_from_decision, build_openai_chat_sync_plan_from_decision,
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
};
use crate::scheduler::{
list_selectable_candidates, GatewayMinimalCandidateSelectionCandidate,
};
use crate::ai_pipeline::planner::scheduler_facade::list_selectable_candidates;
pub(super) fn build_local_openai_chat_miss_diagnostic(
decision: &GatewayControlDecision,
@@ -292,7 +294,7 @@ pub(super) async fn list_local_openai_chat_candidates(
state: &AppState,
input: &LocalOpenAiChatDecisionInput,
require_streaming: bool,
) -> Result<Vec<GatewayMinimalCandidateSelectionCandidate>, GatewayError> {
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
let now_unix_secs = current_unix_secs();
let mut combined = Vec::new();
let mut seen = BTreeSet::new();
@@ -352,9 +354,9 @@ pub(super) async fn list_local_openai_chat_candidates(
}
fn auth_snapshot_allows_cross_format_openai_chat_candidate(
auth_snapshot: &crate::data::auth::GatewayAuthApiKeySnapshot,
auth_snapshot: &GatewayAuthApiKeySnapshot,
requested_model: &str,
candidate: &GatewayMinimalCandidateSelectionCandidate,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> bool {
if let Some(allowed_providers) = auth_snapshot.effective_allowed_providers() {
let provider_allowed = allowed_providers.iter().any(|value| {
@@ -438,13 +440,13 @@ pub(super) async fn resolve_local_openai_chat_decision_input(
};
let now_unix_secs = current_unix_secs();
let auth_snapshot = match state
.read_auth_api_key_snapshot(
&auth_context.user_id,
&auth_context.api_key_id,
now_unix_secs,
)
.await
let auth_snapshot = match read_auth_api_key_snapshot(
state,
&auth_context.user_id,
&auth_context.api_key_id,
now_unix_secs,
)
.await
{
Ok(Some(snapshot)) => snapshot,
Ok(None) => {
@@ -1,46 +1,57 @@
use std::collections::{BTreeMap, BTreeSet};
use aether_data::repository::candidates::{RequestCandidateStatus, UpsertRequestCandidateRecord};
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use serde_json::{json, Value};
use tracing::warn;
use uuid::Uuid;
use crate::ai_pipeline::control_facade::{
collect_control_headers, GatewayControlAuthContext, GatewayControlDecision,
};
use crate::ai_pipeline::conversion::{
request_conversion_direct_auth, request_conversion_kind, request_conversion_transport_supported,
};
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
use crate::ai_pipeline::planner::common::{
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
use crate::ai_pipeline::execution_facade::{ConversionMode, ExecutionStrategy};
use crate::ai_pipeline::planner::auth_snapshot_facade::{
read_auth_api_key_snapshot, GatewayAuthApiKeySnapshot,
};
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
use crate::ai_pipeline::planner::candidate_runtime_facade::{
persist_available_local_candidate, persist_skipped_local_candidate,
};
use crate::ai_pipeline::planner::common::{
force_upstream_streaming_for_provider, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
};
use crate::ai_pipeline::planner::executor_facade::mark_unused_local_candidate_items;
use crate::ai_pipeline::planner::plan_builders::{
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
};
use crate::ai_pipeline::planner::scheduler_facade::list_selectable_candidates;
use crate::ai_pipeline::planner::standard::{
build_cross_format_openai_cli_request_body, build_cross_format_openai_cli_upstream_url,
build_local_openai_cli_request_body, build_local_openai_cli_upstream_url,
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,
build_local_openai_cli_upstream_url,
};
use crate::control::GatewayControlDecision;
use crate::execution_runtime::{ConversionMode, ExecutionStrategy};
use crate::headers::collect_control_headers;
use crate::provider_transport::antigravity::{
use crate::ai_pipeline::planner::transport_facade::{
read_provider_transport_snapshot, resolve_local_oauth_request_auth,
LocalResolvedOAuthRequestAuth,
};
use crate::ai_pipeline::provider_transport_facade::antigravity::{
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
classify_local_antigravity_request_support, AntigravityEnvelopeRequestType,
AntigravityRequestEnvelopeSupport, AntigravityRequestSideSupport,
};
use crate::provider_transport::auth::{
use crate::ai_pipeline::provider_transport_facade::auth::{
build_openai_passthrough_headers, ensure_upstream_auth_header, resolve_local_gemini_auth,
resolve_local_standard_auth,
};
use crate::provider_transport::policy::supports_local_standard_transport_with_network;
use crate::provider_transport::{
use crate::ai_pipeline::provider_transport_facade::policy::supports_local_standard_transport_with_network;
use crate::ai_pipeline::provider_transport_facade::{
apply_local_body_rules, apply_local_header_rules, resolve_transport_execution_timeouts,
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_transport_tls_profile,
LocalResolvedOAuthRequestAuth,
};
use crate::scheduler::{
current_unix_secs, list_selectable_candidates, record_local_request_candidate_status,
GatewayMinimalCandidateSelectionCandidate,
};
use crate::clock::current_unix_secs;
use crate::{
append_execution_contract_fields_to_value, AppState, GatewayControlSyncDecisionResponse,
GatewayError,
@@ -48,25 +59,18 @@ use crate::{
const ANTIGRAVITY_ENVELOPE_NAME: &str = "antigravity:v1internal";
#[derive(Debug, Clone, Copy)]
pub(super) struct LocalOpenAiCliSpec {
pub(super) api_format: &'static str,
pub(super) decision_kind: &'static str,
pub(super) report_kind: &'static str,
pub(super) compact: bool,
pub(super) require_streaming: bool,
}
pub(super) use aether_ai_pipeline::planner::standard::openai_cli::LocalOpenAiCliSpec;
#[derive(Debug, Clone)]
pub(super) struct LocalOpenAiCliDecisionInput {
pub(super) auth_context: crate::control::GatewayControlAuthContext,
pub(super) auth_context: GatewayControlAuthContext,
pub(super) requested_model: String,
pub(super) auth_snapshot: crate::data::auth::GatewayAuthApiKeySnapshot,
pub(super) auth_snapshot: GatewayAuthApiKeySnapshot,
}
#[derive(Debug, Clone)]
pub(super) struct LocalOpenAiCliCandidateAttempt {
pub(super) candidate: GatewayMinimalCandidateSelectionCandidate,
pub(super) candidate: SchedulerMinimalCandidateSelectionCandidate,
pub(super) candidate_index: u32,
pub(super) candidate_id: String,
}
@@ -90,13 +94,13 @@ pub(super) async fn resolve_local_openai_cli_decision_input(
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)?;
let auth_snapshot = match state
.read_auth_api_key_snapshot(
&auth_context.user_id,
&auth_context.api_key_id,
current_unix_secs(),
)
.await
let auth_snapshot = match read_auth_api_key_snapshot(
state,
&auth_context.user_id,
&auth_context.api_key_id,
current_unix_secs(),
)
.await
{
Ok(Some(snapshot)) => snapshot,
Ok(None) => return None,
@@ -201,47 +205,19 @@ pub(super) async fn materialize_local_openai_cli_candidate_attempts(
candidate.endpoint_api_format.as_str(),
);
let candidate_id = match state
.upsert_request_candidate(UpsertRequestCandidateRecord {
id: generated_candidate_id.clone(),
request_id: trace_id.to_string(),
user_id: Some(input.auth_context.user_id.clone()),
api_key_id: Some(input.auth_context.api_key_id.clone()),
username: None,
api_key_name: None,
candidate_index: candidate_index as u32,
retry_index: 0,
provider_id: Some(candidate.provider_id.clone()),
endpoint_id: Some(candidate.endpoint_id.clone()),
key_id: Some(candidate.key_id.clone()),
status: RequestCandidateStatus::Available,
skip_reason: None,
is_cached: Some(false),
status_code: None,
error_type: None,
error_message: None,
latency_ms: None,
concurrent_requests: None,
extra_data: Some(extra_data),
required_capabilities: candidate.key_capabilities.clone(),
created_at_unix_secs: Some(created_at_unix_secs),
started_at_unix_secs: None,
finished_at_unix_secs: None,
})
.await
{
Ok(Some(stored)) => stored.id,
Ok(None) => generated_candidate_id.clone(),
Err(err) => {
warn!(
trace_id = %trace_id,
api_format = spec.api_format,
error = ?err,
"gateway local openai cli decision request candidate upsert failed"
);
generated_candidate_id.clone()
}
};
let candidate_id = persist_available_local_candidate(
state,
trace_id,
&input.auth_context.user_id,
&input.auth_context.api_key_id,
&candidate,
candidate_index as u32,
&generated_candidate_id,
Some(extra_data),
created_at_unix_secs,
"gateway local openai cli decision request candidate upsert failed",
)
.await;
attempts.push(LocalOpenAiCliCandidateAttempt {
candidate,
@@ -254,9 +230,9 @@ pub(super) async fn materialize_local_openai_cli_candidate_attempts(
}
fn auth_snapshot_allows_cross_format_openai_cli_candidate(
auth_snapshot: &crate::data::auth::GatewayAuthApiKeySnapshot,
auth_snapshot: &GatewayAuthApiKeySnapshot,
requested_model: &str,
candidate: &GatewayMinimalCandidateSelectionCandidate,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> bool {
if let Some(allowed_providers) = auth_snapshot.effective_allowed_providers() {
let provider_allowed = allowed_providers.iter().any(|value| {
@@ -300,13 +276,13 @@ pub(super) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
} = attempt;
let provider_api_format = candidate.endpoint_api_format.trim().to_ascii_lowercase();
let transport = match state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
let transport = match read_provider_transport_snapshot(
state,
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
{
Ok(Some(snapshot)) => snapshot,
Ok(None) => {
@@ -385,7 +361,7 @@ pub(super) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
conversion_kind.and_then(|kind| request_conversion_direct_auth(&transport, kind))
};
let oauth_auth = if resolved_auth.is_none() {
match state.resolve_local_oauth_request_auth(&transport).await {
match resolve_local_oauth_request_auth(state, &transport).await {
Ok(Some(LocalResolvedOAuthRequestAuth::Header { name, value })) => Some((name, value)),
Ok(Some(LocalResolvedOAuthRequestAuth::Kiro(_))) => None,
Ok(None) => None,
@@ -434,21 +410,32 @@ pub(super) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
}
let needs_bidirectional_conversion = !same_format && conversion_kind.is_some();
let upstream_is_stream = spec.require_streaming
|| is_antigravity
|| force_upstream_streaming_for_provider(
transport.provider.provider_type.as_str(),
provider_api_format.as_str(),
);
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.require_streaming,
upstream_is_stream,
transport.provider.provider_type.as_str(),
transport.endpoint.body_rules.as_ref(),
Some(input.auth_context.api_key_id.as_str()),
)
} else {
build_local_openai_cli_request_body(
body_json,
&mapped_model,
spec.require_streaming,
upstream_is_stream,
transport.provider.provider_type.as_str(),
provider_api_format.as_str(),
transport.endpoint.body_rules.as_ref(),
Some(input.auth_context.api_key_id.as_str()),
)
}) else {
mark_skipped_local_openai_cli_candidate(
@@ -513,7 +500,6 @@ pub(super) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
} else {
base_provider_request_body
};
let upstream_is_stream = spec.require_streaming || is_antigravity;
let Some(upstream_url) = (if needs_bidirectional_conversion {
build_cross_format_openai_cli_upstream_url(
@@ -572,6 +558,15 @@ pub(super) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
.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
@@ -685,69 +680,37 @@ async fn mark_skipped_local_openai_cli_candidate(
state: &AppState,
input: &LocalOpenAiCliDecisionInput,
trace_id: &str,
candidate: &GatewayMinimalCandidateSelectionCandidate,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
candidate_index: u32,
candidate_id: &str,
skip_reason: &'static str,
) {
if let Err(err) = state
.upsert_request_candidate(UpsertRequestCandidateRecord {
id: candidate_id.to_string(),
request_id: trace_id.to_string(),
user_id: Some(input.auth_context.user_id.clone()),
api_key_id: Some(input.auth_context.api_key_id.clone()),
username: None,
api_key_name: None,
candidate_index,
retry_index: 0,
provider_id: Some(candidate.provider_id.clone()),
endpoint_id: Some(candidate.endpoint_id.clone()),
key_id: Some(candidate.key_id.clone()),
status: RequestCandidateStatus::Skipped,
skip_reason: Some(skip_reason.to_string()),
is_cached: Some(false),
status_code: None,
error_type: None,
error_message: None,
latency_ms: None,
concurrent_requests: None,
extra_data: None,
required_capabilities: candidate.key_capabilities.clone(),
created_at_unix_secs: None,
started_at_unix_secs: None,
finished_at_unix_secs: Some(current_unix_secs()),
})
.await
{
warn!(
trace_id = %trace_id,
candidate_id = %candidate_id,
skip_reason,
error = ?err,
"gateway local openai cli decision failed to persist skipped candidate"
);
}
persist_skipped_local_candidate(
state,
trace_id,
&input.auth_context.user_id,
&input.auth_context.api_key_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
current_unix_secs(),
"gateway local openai cli decision failed to persist skipped candidate",
)
.await;
}
pub(super) async fn mark_unused_local_openai_cli_candidates<T>(state: &AppState, remaining: Vec<T>)
where
T: LocalOpenAiCliPlanAndReport,
{
for plan_and_report in remaining {
record_local_request_candidate_status(
state,
plan_and_report.plan(),
plan_and_report.report_context(),
RequestCandidateStatus::Unused,
None,
None,
None,
None,
None,
None,
)
.await;
}
mark_unused_local_candidate_items(
state,
remaining,
|item| item.plan(),
|item| item.report_context(),
)
.await;
}
pub(super) trait LocalOpenAiCliPlanAndReport {
@@ -1,7 +1,7 @@
use crate::ai_pipeline::control_facade::GatewayControlDecision;
use crate::ai_pipeline::planner::plan_builders::{
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
};
use crate::control::GatewayControlDecision;
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
mod decision;
@@ -5,56 +5,15 @@ use super::decision::{
maybe_build_local_openai_cli_decision_payload_for_candidate,
resolve_local_openai_cli_decision_input, LocalOpenAiCliSpec,
};
use crate::ai_pipeline::planner::common::{
OPENAI_CLI_STREAM_PLAN_KIND, OPENAI_CLI_SYNC_PLAN_KIND, OPENAI_COMPACT_STREAM_PLAN_KIND,
OPENAI_COMPACT_SYNC_PLAN_KIND,
};
use crate::ai_pipeline::control_facade::GatewayControlDecision;
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::control::GatewayControlDecision;
use crate::{AppState, GatewayError};
pub(super) fn resolve_sync_spec(plan_kind: &str) -> Option<LocalOpenAiCliSpec> {
match plan_kind {
OPENAI_CLI_SYNC_PLAN_KIND => Some(LocalOpenAiCliSpec {
api_format: "openai:cli",
decision_kind: OPENAI_CLI_SYNC_PLAN_KIND,
report_kind: "openai_cli_sync_success",
compact: false,
require_streaming: false,
}),
OPENAI_COMPACT_SYNC_PLAN_KIND => Some(LocalOpenAiCliSpec {
api_format: "openai:compact",
decision_kind: OPENAI_COMPACT_SYNC_PLAN_KIND,
report_kind: "openai_cli_sync_success",
compact: true,
require_streaming: false,
}),
_ => None,
}
}
pub(super) fn resolve_stream_spec(plan_kind: &str) -> Option<LocalOpenAiCliSpec> {
match plan_kind {
OPENAI_CLI_STREAM_PLAN_KIND => Some(LocalOpenAiCliSpec {
api_format: "openai:cli",
decision_kind: OPENAI_CLI_STREAM_PLAN_KIND,
report_kind: "openai_cli_stream_success",
compact: false,
require_streaming: true,
}),
OPENAI_COMPACT_STREAM_PLAN_KIND => Some(LocalOpenAiCliSpec {
api_format: "openai:compact",
decision_kind: OPENAI_COMPACT_STREAM_PLAN_KIND,
report_kind: "openai_cli_stream_success",
compact: true,
require_streaming: true,
}),
_ => None,
}
}
pub(super) use aether_ai_pipeline::planner::standard::openai_cli::{
resolve_stream_spec, resolve_sync_spec,
};
pub(super) async fn build_local_sync_plan_and_reports(
state: &AppState,
@@ -1,12 +1,14 @@
pub(crate) mod chat;
pub(crate) mod cli;
pub(crate) use chat::{
pub(crate) use aether_ai_pipeline::planner::openai::{
copy_request_number_field, copy_request_number_field_as,
map_openai_reasoning_effort_to_claude_output, map_openai_reasoning_effort_to_gemini_budget,
maybe_build_stream_local_decision_payload, maybe_build_sync_local_decision_payload,
parse_openai_stop_sequences, resolve_openai_chat_max_tokens, value_as_u64,
};
pub(crate) use chat::{
maybe_build_stream_local_decision_payload, maybe_build_sync_local_decision_payload,
};
pub(crate) use cli::{
maybe_build_stream_local_openai_cli_decision_payload,
maybe_build_sync_local_openai_cli_decision_payload,
@@ -6,10 +6,12 @@ use super::{
LocalSyncPlanAndReport,
};
use crate::ai_pipeline::adaptation::surfaces::provider_adaptation_requires_eventstream_accept;
use crate::provider_transport::auth::{
use crate::ai_pipeline::provider_transport_facade::auth::{
build_openai_passthrough_headers, ensure_upstream_auth_header,
};
use crate::provider_transport::url::{build_openai_chat_url, build_openai_cli_url};
use crate::ai_pipeline::provider_transport_facade::url::{
build_openai_chat_url, build_openai_cli_url,
};
pub(crate) fn build_openai_chat_sync_plan_from_decision(
parts: &http::request::Parts,
@@ -3,7 +3,7 @@ use aether_contracts::{ExecutionPlan, RequestBody};
use super::{augment_sync_report_context, LocalStreamPlanAndReport, LocalSyncPlanAndReport};
use crate::ai_pipeline::adaptation::surfaces::provider_adaptation_requires_eventstream_accept;
use crate::ai_pipeline::contracts::generic_decision_missing_exact_provider_request;
use crate::provider_transport::ensure_upstream_auth_header;
use crate::ai_pipeline::provider_transport_facade::ensure_upstream_auth_header;
use crate::{GatewayControlSyncDecisionResponse, GatewayError};
pub(crate) fn build_standard_sync_plan_from_decision(
@@ -0,0 +1,22 @@
pub(crate) use crate::ai_pipeline::provider_transport_facade::{
GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
};
use crate::{AppState, GatewayError};
pub(crate) async fn read_provider_transport_snapshot(
state: &AppState,
provider_id: &str,
endpoint_id: &str,
key_id: &str,
) -> Result<Option<GatewayProviderTransportSnapshot>, GatewayError> {
state
.read_provider_transport_snapshot(provider_id, endpoint_id, key_id)
.await
}
pub(crate) async fn resolve_local_oauth_request_auth(
state: &AppState,
transport: &GatewayProviderTransportSnapshot,
) -> Result<Option<LocalResolvedOAuthRequestAuth>, GatewayError> {
state.resolve_local_oauth_request_auth(transport).await
}