mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-13 04:29:50 +08:00
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:
@@ -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};
|
||||
|
||||
+26
-52
@@ -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,
|
||||
|
||||
+39
-58
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user