i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-11 03:39:49 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: 抽离 AI pipeline 与调度共享能力逻辑

Browse Source
This commit is contained in:
fawney19
2026-04-10 01:46:14 +08:00
parent b901a6ffc7
commit 5014e2f5fd
255 changed files with 15057 additions and 3115 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/ai_pipeline/conversion/mod.rs
+60 -9
View File
@@ -5,8 +5,9 @@ pub(crate) use crate::ai_pipeline::{
core_error_default_client_api_format, is_core_error_finalize_kind, LocalCoreSyncErrorKind,
};
pub(crate) use crate::ai_pipeline::{
request_conversion_direct_auth, request_conversion_kind,
request_conversion_transport_supported, sync_chat_response_conversion_kind,
request_candidate_api_formats, request_conversion_direct_auth, request_conversion_kind,
request_conversion_requires_enable_flag, request_conversion_transport_supported,
request_pair_allowed_for_transport, sync_chat_response_conversion_kind,
sync_cli_response_conversion_kind, RequestConversionKind, SyncChatResponseConversionKind,
SyncCliResponseConversionKind,
};
@@ -14,7 +15,7 @@ pub(crate) use crate::ai_pipeline::{
#[cfg(test)]
mod tests {
use super::{
request_conversion_kind, sync_chat_response_conversion_kind,
request_candidate_api_formats, request_conversion_kind, sync_chat_response_conversion_kind,
sync_cli_response_conversion_kind, RequestConversionKind, SyncChatResponseConversionKind,
SyncCliResponseConversionKind,
};
@@ -22,7 +23,15 @@ mod tests {
#[test]
fn request_conversion_registry_supports_bidirectional_standard_matrix() {
assert_eq!(
request_conversion_kind("claude:chat", "openai:chat"),
request_conversion_kind("openai:chat", "openai:cli"),
Some(RequestConversionKind::ToOpenAIFamilyCli)
);
assert_eq!(
request_conversion_kind("openai:chat", "claude:cli"),
Some(RequestConversionKind::ToClaudeStandard)
);
assert_eq!(
request_conversion_kind("openai:cli", "openai:chat"),
Some(RequestConversionKind::ToOpenAIChat)
);
assert_eq!(
@@ -30,12 +39,16 @@ mod tests {
Some(RequestConversionKind::ToClaudeStandard)
);
assert_eq!(
request_conversion_kind("gemini:cli", "openai:compact"),
Some(RequestConversionKind::ToOpenAICompact)
request_conversion_kind("openai:compact", "gemini:cli"),
None
);
assert_eq!(
request_conversion_kind("openai:compact", "gemini:cli"),
Some(RequestConversionKind::ToGeminiStandard)
request_conversion_kind("gemini:cli", "openai:compact"),
None
);
assert_eq!(
request_conversion_kind("openai:chat", "openai:compact"),
None
);
assert_eq!(request_conversion_kind("claude:chat", "claude:chat"), None);
}
@@ -59,12 +72,50 @@ mod tests {
Some(SyncCliResponseConversionKind::ToGeminiCli)
);
assert_eq!(
sync_cli_response_conversion_kind("claude:cli", "openai:compact"),
sync_cli_response_conversion_kind("claude:chat", "openai:cli"),
Some(SyncCliResponseConversionKind::ToOpenAIFamilyCli)
);
assert_eq!(
sync_cli_response_conversion_kind("claude:cli", "openai:compact"),
None
);
assert_eq!(
sync_cli_response_conversion_kind("openai:compact", "claude:cli"),
None
);
assert_eq!(
sync_cli_response_conversion_kind("gemini:cli", "claude:cli"),
Some(SyncCliResponseConversionKind::ToClaudeCli)
);
}
#[test]
fn request_candidate_registry_excludes_compact_as_cross_format_target() {
assert_eq!(
request_candidate_api_formats("openai:chat", false),
vec![
"openai:chat",
"openai:cli",
"claude:chat",
"claude:cli",
"gemini:chat",
"gemini:cli",
]
);
assert_eq!(
request_candidate_api_formats("openai:cli", false),
vec![
"openai:chat",
"openai:cli",
"claude:chat",
"claude:cli",
"gemini:chat",
"gemini:cli",
]
);
assert_eq!(
request_candidate_api_formats("openai:compact", false),
vec!["openai:compact"]
);
}
}
apps/aether-gateway/src/ai_pipeline/finalize/tests_stream.rs
+186
View File
@@ -52,6 +52,130 @@ fn gemini_cli_v1internal_stream_rewriter_unwraps_response_object() {
);
}
#[test]
fn openai_chat_error_to_openai_cli_stream_rewriter_converts_to_response_failed() {
let report_context = json!({
"provider_api_format": "openai:chat",
"client_api_format": "openai:cli",
"needs_conversion": true,
"mapped_model": "gpt-5.4",
});
let mut rewriter =
maybe_build_local_stream_rewriter(Some(&report_context)).expect("rewriter should exist");
let output = rewriter
.push_chunk(
b"data: {\"error\":{\"message\":\"bad request\",\"type\":\"invalid_request_error\",\"code\":\"invalid_request\"}}\n\n",
)
.expect("rewrite should succeed");
let output_text = utf8(output);
assert!(output_text.contains("event: response.failed"));
assert!(output_text.contains("\"sequence_number\":1"));
assert!(output_text.contains("\"message\":\"bad request\""));
assert!(output_text.contains("\"type\":\"invalid_request_error\""));
assert!(!output_text.contains("data: [DONE]"));
assert!(rewriter.finish().expect("finish should succeed").is_empty());
}
#[test]
fn claude_error_to_openai_cli_stream_rewriter_converts_to_response_failed() {
let report_context = json!({
"provider_api_format": "claude:chat",
"client_api_format": "openai:cli",
"needs_conversion": true,
"mapped_model": "claude-sonnet-4-5",
});
let mut rewriter =
maybe_build_local_stream_rewriter(Some(&report_context)).expect("rewriter should exist");
let output = rewriter
.push_chunk(
concat!(
"event: error\n",
"data: {\"type\":\"error\",\"error\":{\"type\":\"rate_limit_error\",\"message\":\"slow down\",\"code\":\"rate_limit\"}}\n\n"
)
.as_bytes(),
)
.expect("rewrite should succeed");
let output_text = utf8(output);
assert!(output_text.contains("event: response.failed"));
assert!(output_text.contains("\"sequence_number\":1"));
assert!(output_text.contains("\"message\":\"slow down\""));
assert!(output_text.contains("\"type\":\"rate_limit_error\""));
assert!(!output_text.contains("data: [DONE]"));
assert!(rewriter.finish().expect("finish should succeed").is_empty());
}
#[test]
fn gemini_error_to_openai_cli_stream_rewriter_converts_to_response_failed() {
let report_context = json!({
"provider_api_format": "gemini:cli",
"client_api_format": "openai:cli",
"needs_conversion": true,
"mapped_model": "gemini-2.5-pro",
});
let mut rewriter =
maybe_build_local_stream_rewriter(Some(&report_context)).expect("rewriter should exist");
let output = rewriter
.push_chunk(
b"data: {\"error\":{\"code\":429,\"message\":\"quota exceeded\",\"status\":\"RESOURCE_EXHAUSTED\"}}\n\n",
)
.expect("rewrite should succeed");
let output_text = utf8(output);
assert!(output_text.contains("event: response.failed"));
assert!(output_text.contains("\"sequence_number\":1"));
assert!(output_text.contains("\"message\":\"quota exceeded\""));
assert!(output_text.contains("\"type\":\"rate_limit_error\""));
assert!(!output_text.contains("data: [DONE]"));
assert!(rewriter.finish().expect("finish should succeed").is_empty());
}
#[test]
fn openai_chat_error_to_claude_chat_stream_rewriter_uses_error_event_line() {
let report_context = json!({
"provider_api_format": "openai:chat",
"client_api_format": "claude:chat",
"needs_conversion": true,
"mapped_model": "gpt-5.4",
});
let mut rewriter =
maybe_build_local_stream_rewriter(Some(&report_context)).expect("rewriter should exist");
let output = rewriter
.push_chunk(
b"data: {\"error\":{\"message\":\"bad request\",\"type\":\"invalid_request_error\",\"code\":\"invalid_request\"}}\n\n",
)
.expect("rewrite should succeed");
let output_text = utf8(output);
assert!(output_text.contains("event: error"));
assert!(output_text.contains("\"type\":\"error\""));
assert!(output_text.contains("\"message\":\"bad request\""));
assert!(output_text.contains("\"code\":\"invalid_request\""));
assert!(!output_text.contains("data: [DONE]"));
assert!(rewriter.finish().expect("finish should succeed").is_empty());
}
#[test]
fn openai_chat_error_to_gemini_chat_stream_rewriter_keeps_data_only_error() {
let report_context = json!({
"provider_api_format": "openai:chat",
"client_api_format": "gemini:chat",
"needs_conversion": true,
"mapped_model": "gpt-5.4",
});
let mut rewriter =
maybe_build_local_stream_rewriter(Some(&report_context)).expect("rewriter should exist");
let output = rewriter
.push_chunk(
b"data: {\"error\":{\"message\":\"rate limited\",\"type\":\"rate_limit_error\",\"code\":\"rate_limit\"}}\n\n",
)
.expect("rewrite should succeed");
let output_text = utf8(output);
assert!(output_text.starts_with("data: {\"error\":"));
assert!(!output_text.contains("event: "));
assert!(output_text.contains("\"message\":\"rate limited\""));
assert!(output_text.contains("\"status\":\"RESOURCE_EXHAUSTED\""));
assert!(!output_text.contains("data: [DONE]"));
assert!(rewriter.finish().expect("finish should succeed").is_empty());
}
#[test]
fn claude_to_openai_chat_stream_rewriter_converts_text_deltas() {
let report_context = json!({
@@ -234,6 +358,22 @@ fn openai_cli_to_openai_chat_stream_rewriter_converts_text_deltas_immediately()
.expect("rewrite should succeed");
let completed_text = String::from_utf8(completed).expect("utf8 should decode");
assert!(completed_text.contains("\"finish_reason\":\"stop\""));
assert!(
completed_text.contains("\"choices\":[]"),
"{completed_text}"
);
assert!(
completed_text.contains("\"prompt_tokens\":1"),
"{completed_text}"
);
assert!(
completed_text.contains("\"completion_tokens\":2"),
"{completed_text}"
);
assert!(
completed_text.contains("\"total_tokens\":3"),
"{completed_text}"
);
assert!(completed_text.contains("data: [DONE]"));
assert!(rewriter.finish().expect("finish should succeed").is_empty());
}
@@ -262,6 +402,13 @@ fn openai_cli_to_openai_chat_stream_rewriter_converts_completed_event_without_bu
assert!(output_text.contains("\"role\":\"assistant\""));
assert!(output_text.contains("\"content\":\"Hello Codex\""));
assert!(output_text.contains("\"finish_reason\":\"stop\""));
assert!(output_text.contains("\"choices\":[]"), "{output_text}");
assert!(output_text.contains("\"prompt_tokens\":1"), "{output_text}");
assert!(
output_text.contains("\"completion_tokens\":2"),
"{output_text}"
);
assert!(output_text.contains("\"total_tokens\":3"), "{output_text}");
assert!(output_text.contains("data: [DONE]"));
assert!(rewriter.finish().expect("finish should succeed").is_empty());
}
@@ -514,6 +661,7 @@ fn openai_chat_to_claude_chat_stream_rewriter_converts_via_standard_matrix() {
assert!(first_text.contains("event: content_block_start"));
assert!(first_text.contains("event: content_block_delta"));
assert!(first_text.contains("\"text\":\"Hello Claude\""));
assert!(first_text.contains("\"usage\":{\"input_tokens\":0,\"output_tokens\":0}"));
let second = rewriter
.push_chunk(
@@ -531,6 +679,44 @@ fn openai_chat_to_claude_chat_stream_rewriter_converts_via_standard_matrix() {
assert!(rewriter.finish().expect("finish should succeed").is_empty());
}
#[test]
fn openai_chat_to_claude_chat_stream_rewriter_injects_default_usage_when_finish_chunk_lacks_usage()
{
let report_context = json!({
"provider_api_format": "openai:chat",
"client_api_format": "claude:chat",
"needs_conversion": true,
"mapped_model": "gpt-5.4",
});
let mut rewriter =
maybe_build_local_stream_rewriter(Some(&report_context)).expect("rewriter should exist");
let first = rewriter
.push_chunk(
"data: {\"id\":\"chatcmpl_usage_missing_123\",\"object\":\"chat.completion.chunk\",\"created\":1,\"model\":\"gpt-5.4\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"Hello\"},\"finish_reason\":null}]}\n\n"
.as_bytes(),
)
.expect("rewrite should succeed");
let first_text = utf8(first);
assert!(first_text.contains("event: message_start"));
let second = rewriter
.push_chunk(
concat!(
"data: {\"id\":\"chatcmpl_usage_missing_123\",\"object\":\"chat.completion.chunk\",\"model\":\"gpt-5.4\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n",
"data: [DONE]\n\n"
)
.as_bytes(),
)
.expect("rewrite should succeed");
let output_text = utf8(second);
assert!(output_text.is_empty(), "{output_text}");
let final_text = utf8(rewriter.finish().expect("finish should succeed"));
assert!(final_text.contains("event: message_delta"));
assert!(final_text.contains("\"stop_reason\":\"end_turn\""));
assert!(final_text.contains("\"usage\":{\"input_tokens\":0,\"output_tokens\":0}"));
assert!(final_text.contains("event: message_stop"));
}
#[test]
fn openai_chat_to_gemini_cli_stream_rewriter_converts_via_standard_matrix() {
let report_context = json!({
apps/aether-gateway/src/ai_pipeline/mod.rs
+2
View File
@@ -68,6 +68,8 @@ pub(crate) fn build_execution_runtime_auth_context(
ExecutionRuntimeAuthContext {
user_id: auth_context.user_id.clone(),
api_key_id: auth_context.api_key_id.clone(),
username: auth_context.username.clone(),
api_key_name: auth_context.api_key_name.clone(),
balance_remaining: auth_context.balance_remaining,
access_allowed: auth_context.access_allowed,
}
apps/aether-gateway/src/ai_pipeline/planner/candidate_affinity.rs
+579 -17
View File
@@ -1,9 +1,18 @@
use tracing::warn;
use crate::ai_pipeline::transport::resolve_transport_proxy_snapshot;
use crate::ai_pipeline::PlannerAppState;
use crate::AppState;
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use crate::ai_pipeline::{
GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, PlannerAppState,
};
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
use crate::scheduler::config::{read_scheduler_ordering_config, SchedulerOrderingConfig};
use aether_scheduler_core::{
build_scheduler_affinity_cache_key_for_api_key_id, compare_candidates_by_priority_mode,
requested_capability_priority_for_candidate, SchedulerAffinityTarget,
SchedulerMinimalCandidateSelectionCandidate,
};
const PLANNER_SCHEDULER_AFFINITY_MAX_ENTRIES: usize = 10_000;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
enum TunnelOwnerAffinityBucket {
@@ -12,6 +21,12 @@ enum TunnelOwnerAffinityBucket {
RemoteTunnel = 2,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct CandidateExecutionOrdering {
tunnel_bucket: TunnelOwnerAffinityBucket,
keep_priority_on_conversion: bool,
}
pub(crate) async fn prefer_local_tunnel_owner_candidates(
state: PlannerAppState<'_>,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
@@ -28,11 +43,127 @@ pub(crate) async fn prefer_local_tunnel_owner_candidates(
.collect()
}
pub(crate) async fn rank_local_execution_candidates(
state: PlannerAppState<'_>,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
client_api_format: &str,
required_capabilities: Option<&serde_json::Value>,
) -> Vec<SchedulerMinimalCandidateSelectionCandidate> {
let normalized_client_api_format = client_api_format.trim().to_ascii_lowercase();
let ordering_config = read_scheduler_ordering_config_or_default(state).await;
let mut ranked = Vec::with_capacity(candidates.len());
for (original_index, candidate) in candidates.into_iter().enumerate() {
let ordering =
resolve_candidate_execution_ordering(state, &candidate, ordering_config).await;
let is_same_format = candidate
.endpoint_api_format
.trim()
.eq_ignore_ascii_case(normalized_client_api_format.as_str());
let demote_cross_format = !is_same_format && !ordering.keep_priority_on_conversion;
let capability_priority =
requested_capability_priority_for_candidate(required_capabilities, &candidate);
ranked.push((
capability_priority.0,
capability_priority.1,
ordering.tunnel_bucket,
demote_cross_format,
original_index,
candidate,
));
}
ranked.sort_by(|left, right| {
left.0
.cmp(&right.0)
.then(left.1.cmp(&right.1))
.then(left.2.cmp(&right.2))
.then(left.3.cmp(&right.3))
.then_with(|| {
compare_candidates_by_priority_mode(
&left.5,
&right.5,
ordering_config.priority_mode,
None,
)
})
.then(left.4.cmp(&right.4))
});
ranked
.into_iter()
.map(|(_, _, _, _, _, candidate)| candidate)
.collect()
}
pub(crate) fn remember_scheduler_affinity_for_candidate(
state: PlannerAppState<'_>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_api_format: &str,
requested_model: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) {
let Some(api_key_id) = auth_snapshot
.map(|snapshot| snapshot.api_key_id.trim())
.filter(|value| !value.is_empty())
else {
return;
};
let Some(cache_key) = build_scheduler_affinity_cache_key_for_api_key_id(
api_key_id,
client_api_format,
requested_model,
) else {
return;
};
state.app().remember_scheduler_affinity_target(
&cache_key,
SchedulerAffinityTarget {
provider_id: candidate.provider_id.clone(),
endpoint_id: candidate.endpoint_id.clone(),
key_id: candidate.key_id.clone(),
},
SCHEDULER_AFFINITY_TTL,
PLANNER_SCHEDULER_AFFINITY_MAX_ENTRIES,
);
}
async fn resolve_candidate_tunnel_owner_affinity(
state: PlannerAppState<'_>,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> TunnelOwnerAffinityBucket {
let transport = match state
let Some(transport) = read_candidate_transport_snapshot(state, candidate).await else {
return TunnelOwnerAffinityBucket::Neutral;
};
resolve_tunnel_owner_affinity_from_transport(state, &transport).await
}
async fn resolve_candidate_execution_ordering(
state: PlannerAppState<'_>,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
ordering_config: SchedulerOrderingConfig,
) -> CandidateExecutionOrdering {
let Some(transport) = read_candidate_transport_snapshot(state, candidate).await else {
return CandidateExecutionOrdering {
tunnel_bucket: TunnelOwnerAffinityBucket::Neutral,
keep_priority_on_conversion: ordering_config.keep_priority_on_conversion,
};
};
CandidateExecutionOrdering {
tunnel_bucket: resolve_tunnel_owner_affinity_from_transport(state, &transport).await,
keep_priority_on_conversion: ordering_config.keep_priority_on_conversion
|| transport.provider.keep_priority_on_conversion,
}
}
async fn read_candidate_transport_snapshot(
state: PlannerAppState<'_>,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> Option<GatewayProviderTransportSnapshot> {
match state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
@@ -40,8 +171,8 @@ async fn resolve_candidate_tunnel_owner_affinity(
)
.await
{
Ok(Some(transport)) => transport,
Ok(None) => return TunnelOwnerAffinityBucket::Neutral,
Ok(Some(transport)) => Some(transport),
Ok(None) => None,
Err(error) => {
warn!(
event_name = "candidate_affinity_transport_load_failed",
@@ -50,12 +181,17 @@ async fn resolve_candidate_tunnel_owner_affinity(
endpoint_id = %candidate.endpoint_id,
key_id = %candidate.key_id,
error = ?error,
"failed to load provider transport while evaluating tunnel owner affinity"
"failed to load provider transport while evaluating execution ordering"
);
return TunnelOwnerAffinityBucket::Neutral;
None
}
};
}
}
async fn resolve_tunnel_owner_affinity_from_transport(
state: PlannerAppState<'_>,
transport: &GatewayProviderTransportSnapshot,
) -> TunnelOwnerAffinityBucket {
let Some(proxy) = resolve_transport_proxy_snapshot(&transport) else {
return TunnelOwnerAffinityBucket::Neutral;
};
@@ -99,6 +235,23 @@ async fn resolve_candidate_tunnel_owner_affinity(
}
}
async fn read_scheduler_ordering_config_or_default(
state: PlannerAppState<'_>,
) -> SchedulerOrderingConfig {
match read_scheduler_ordering_config(state.app()).await {
Ok(config) => config,
Err(error) => {
warn!(
event_name = "planner_scheduler_ordering_config_load_failed",
log_type = "event",
error = ?error,
"failed to load scheduler ordering config while ranking local execution candidates"
);
SchedulerOrderingConfig::default()
}
}
}
#[cfg(test)]
mod tests {
use std::time::{SystemTime, UNIX_EPOCH};
@@ -110,11 +263,15 @@ mod tests {
use serde_json::json;
use super::{
prefer_local_tunnel_owner_candidates, AppState, PlannerAppState,
prefer_local_tunnel_owner_candidates, rank_local_execution_candidates,
remember_scheduler_affinity_for_candidate, PlannerAppState,
SchedulerMinimalCandidateSelectionCandidate,
};
use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::data::GatewayDataState;
use crate::tunnel::TunnelAttachmentRecord;
use crate::{scheduler::affinity::SCHEDULER_AFFINITY_TTL, AppState};
use aether_data::repository::auth::StoredAuthApiKeySnapshot;
fn sample_candidate(
endpoint_id: &str,
@@ -142,22 +299,55 @@ mod tests {
}
fn sample_provider() -> StoredProviderCatalogProvider {
sample_provider_with_options("provider-1", false, 0)
}
fn sample_provider_with_options(
id: &str,
keep_priority_on_conversion: bool,
provider_priority: i32,
) -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-1".to_string(),
"provider-1".to_string(),
id.to_string(),
id.to_string(),
Some("https://provider.example".to_string()),
"custom".to_string(),
)
.expect("provider should build")
.with_transport_fields(true, false, false, None, None, None, None, None, None)
.with_transport_fields(
true,
keep_priority_on_conversion,
false,
None,
None,
None,
None,
None,
None,
)
.with_routing_fields(provider_priority)
}
fn sample_endpoint(id: &str) -> StoredProviderCatalogEndpoint {
sample_endpoint_for_provider("provider-1", id, "openai:chat")
}
fn sample_endpoint_for_provider(
provider_id: &str,
id: &str,
api_format: &str,
) -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
id.to_string(),
"provider-1".to_string(),
"openai:chat".to_string(),
Some("openai".to_string()),
provider_id.to_string(),
api_format.to_string(),
Some(
api_format
.split(':')
.next()
.unwrap_or(api_format)
.to_string(),
),
Some("chat".to_string()),
true,
)
@@ -176,9 +366,17 @@ mod tests {
}
fn sample_key(id: &str, node_id: &str) -> StoredProviderCatalogKey {
sample_key_for_provider("provider-1", id, node_id)
}
fn sample_key_for_provider(
provider_id: &str,
id: &str,
node_id: &str,
) -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
id.to_string(),
"provider-1".to_string(),
provider_id.to_string(),
id.to_string(),
"api_key".to_string(),
None,
@@ -214,6 +412,35 @@ mod tests {
.as_secs()
}
fn sample_priority_candidate(
provider_id: &str,
endpoint_id: &str,
key_id: &str,
endpoint_api_format: &str,
key_global_priority_for_format: Option<i32>,
provider_priority: i32,
) -> SchedulerMinimalCandidateSelectionCandidate {
SchedulerMinimalCandidateSelectionCandidate {
provider_id: provider_id.to_string(),
provider_name: provider_id.to_string(),
provider_type: "custom".to_string(),
provider_priority,
endpoint_id: endpoint_id.to_string(),
endpoint_api_format: endpoint_api_format.to_string(),
key_id: key_id.to_string(),
key_name: key_id.to_string(),
key_auth_type: "api_key".to_string(),
key_internal_priority: 0,
key_global_priority_for_format,
key_capabilities: None,
model_id: format!("model-{provider_id}"),
global_model_id: "global-model-1".to_string(),
global_model_name: "gpt-4.1".to_string(),
selected_provider_model_name: "gpt-4.1".to_string(),
mapping_matched_model: None,
}
}
#[tokio::test]
async fn prefers_local_tunnel_owner_candidates_before_remote_tunnel_candidates() {
let provider_catalog = InMemoryProviderCatalogReadRepository::seed(
@@ -300,4 +527,339 @@ mod tests {
assert_eq!(reordered[0].endpoint_id, "endpoint-a");
assert_eq!(reordered[1].endpoint_id, "endpoint-b");
}
#[tokio::test]
async fn local_execution_ranking_demotes_cross_format_candidates_without_keep_priority() {
let provider_catalog = InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider_with_options("provider-same", false, 10),
sample_provider_with_options("provider-cross", false, 0),
],
vec![
sample_endpoint_for_provider("provider-same", "endpoint-same", "openai:chat"),
sample_endpoint_for_provider("provider-cross", "endpoint-cross", "claude:chat"),
],
vec![
sample_key_for_provider("provider-same", "key-same", ""),
sample_key_for_provider("provider-cross", "key-cross", ""),
],
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
);
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data_state);
let ranked = rank_local_execution_candidates(
PlannerAppState::new(&state),
vec![
sample_priority_candidate(
"provider-cross",
"endpoint-cross",
"key-cross",
"claude:chat",
Some(0),
0,
),
sample_priority_candidate(
"provider-same",
"endpoint-same",
"key-same",
"openai:chat",
Some(10),
10,
),
],
"openai:chat",
None,
)
.await;
assert_eq!(ranked[0].endpoint_id, "endpoint-same");
assert_eq!(ranked[1].endpoint_id, "endpoint-cross");
}
#[tokio::test]
async fn local_execution_ranking_keeps_cross_format_priority_when_enabled() {
let provider_catalog = InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider_with_options("provider-same", false, 10),
sample_provider_with_options("provider-cross", true, 0),
],
vec![
sample_endpoint_for_provider("provider-same", "endpoint-same", "openai:chat"),
sample_endpoint_for_provider("provider-cross", "endpoint-cross", "claude:chat"),
],
vec![
sample_key_for_provider("provider-same", "key-same", ""),
sample_key_for_provider("provider-cross", "key-cross", ""),
],
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
);
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data_state);
let ranked = rank_local_execution_candidates(
PlannerAppState::new(&state),
vec![
sample_priority_candidate(
"provider-cross",
"endpoint-cross",
"key-cross",
"claude:chat",
Some(0),
0,
),
sample_priority_candidate(
"provider-same",
"endpoint-same",
"key-same",
"openai:chat",
Some(10),
10,
),
],
"openai:chat",
None,
)
.await;
assert_eq!(ranked[0].endpoint_id, "endpoint-cross");
assert_eq!(ranked[1].endpoint_id, "endpoint-same");
}
#[tokio::test]
async fn local_execution_ranking_keeps_cross_format_priority_when_global_override_is_enabled() {
let provider_catalog = InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider_with_options("provider-same", false, 10),
sample_provider_with_options("provider-cross", false, 0),
],
vec![
sample_endpoint_for_provider("provider-same", "endpoint-same", "openai:chat"),
sample_endpoint_for_provider("provider-cross", "endpoint-cross", "claude:chat"),
],
vec![
sample_key_for_provider("provider-same", "key-same", ""),
sample_key_for_provider("provider-cross", "key-cross", ""),
],
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
)
.with_system_config_values_for_tests(vec![(
"keep_priority_on_conversion".to_string(),
json!(true),
)]);
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data_state);
let ranked = rank_local_execution_candidates(
PlannerAppState::new(&state),
vec![
sample_priority_candidate(
"provider-cross",
"endpoint-cross",
"key-cross",
"claude:chat",
Some(0),
0,
),
sample_priority_candidate(
"provider-same",
"endpoint-same",
"key-same",
"openai:chat",
Some(10),
10,
),
],
"openai:chat",
None,
)
.await;
assert_eq!(ranked[0].endpoint_id, "endpoint-cross");
assert_eq!(ranked[1].endpoint_id, "endpoint-same");
}
#[tokio::test]
async fn local_execution_ranking_uses_provider_priority_mode_when_configured() {
let provider_catalog = InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider_with_options("provider-provider-first", false, 0),
sample_provider_with_options("provider-global-first", false, 10),
],
vec![
sample_endpoint_for_provider(
"provider-provider-first",
"endpoint-provider-first",
"openai:chat",
),
sample_endpoint_for_provider(
"provider-global-first",
"endpoint-global-first",
"openai:chat",
),
],
vec![
sample_key_for_provider("provider-provider-first", "key-provider-first", ""),
sample_key_for_provider("provider-global-first", "key-global-first", ""),
],
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
)
.with_system_config_values_for_tests(vec![(
"provider_priority_mode".to_string(),
json!("provider"),
)]);
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data_state);
let ranked = rank_local_execution_candidates(
PlannerAppState::new(&state),
vec![
sample_priority_candidate(
"provider-global-first",
"endpoint-global-first",
"key-global-first",
"openai:chat",
Some(0),
10,
),
sample_priority_candidate(
"provider-provider-first",
"endpoint-provider-first",
"key-provider-first",
"openai:chat",
Some(10),
0,
),
],
"openai:chat",
None,
)
.await;
assert_eq!(ranked[0].endpoint_id, "endpoint-provider-first");
assert_eq!(ranked[1].endpoint_id, "endpoint-global-first");
}
#[tokio::test]
async fn local_execution_ranking_prefers_candidates_matching_requested_capabilities() {
let provider_catalog = InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider_with_options("provider-miss", false, 0),
sample_provider_with_options("provider-hit", false, 0),
],
vec![
sample_endpoint_for_provider("provider-miss", "endpoint-miss", "openai:chat"),
sample_endpoint_for_provider("provider-hit", "endpoint-hit", "openai:chat"),
],
vec![
sample_key_for_provider("provider-miss", "key-miss", ""),
sample_key_for_provider("provider-hit", "key-hit", ""),
],
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
);
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data_state);
let mut candidate_miss = sample_priority_candidate(
"provider-miss",
"endpoint-miss",
"key-miss",
"openai:chat",
Some(0),
0,
);
let mut candidate_hit = sample_priority_candidate(
"provider-hit",
"endpoint-hit",
"key-hit",
"openai:chat",
Some(0),
0,
);
candidate_miss.key_capabilities = Some(json!({"cache_1h": false}));
candidate_hit.key_capabilities = Some(json!({"cache_1h": true}));
let required_capabilities = json!({"cache_1h": true});
let ranked = rank_local_execution_candidates(
PlannerAppState::new(&state),
vec![candidate_miss, candidate_hit],
"openai:chat",
Some(&required_capabilities),
)
.await;
assert_eq!(ranked[0].endpoint_id, "endpoint-hit");
assert_eq!(ranked[1].endpoint_id, "endpoint-miss");
}
#[tokio::test]
async fn remembers_scheduler_affinity_for_candidate_using_requested_model_key() {
let state = AppState::new().expect("state should build");
let auth_snapshot = GatewayAuthApiKeySnapshot::from_stored(
StoredAuthApiKeySnapshot::new(
"user-1".to_string(),
"alice".to_string(),
Some("[email protected]".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
None,
None,
None,
"api-key-1".to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(4_102_444_800),
None,
None,
None,
)
.expect("stored auth snapshot should build"),
current_unix_secs(),
);
let candidate = sample_candidate("endpoint-1", "key-1");
remember_scheduler_affinity_for_candidate(
PlannerAppState::new(&state),
Some(&auth_snapshot),
"openai:chat",
"gpt-5",
&candidate,
);
let remembered = state
.read_scheduler_affinity_target(
"scheduler_affinity:api-key-1:openai:chat:gpt-5",
SCHEDULER_AFFINITY_TTL,
)
.expect("affinity target should be cached");
assert_eq!(remembered.provider_id, "provider-1");
assert_eq!(remembered.endpoint_id, "endpoint-1");
assert_eq!(remembered.key_id, "key-1");
}
}
apps/aether-gateway/src/ai_pipeline/planner/passthrough/provider/family/candidates.rs
+35 -4
View File
@@ -2,7 +2,9 @@ use serde_json::json;
use tracing::warn;
use uuid::Uuid;
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
use crate::ai_pipeline::planner::candidate_affinity::{
rank_local_execution_candidates, remember_scheduler_affinity_for_candidate,
};
use crate::ai_pipeline::{
resolve_local_decision_execution_runtime_auth_context, ConversionMode, ExecutionStrategy,
GatewayControlDecision, PlannerAppState,
@@ -61,10 +63,20 @@ pub(crate) async fn resolve_local_same_format_provider_decision_input(
}
};
let required_capabilities = planner_state
.resolve_request_candidate_required_capabilities(
&auth_context.user_id,
&auth_context.api_key_id,
Some(requested_model.as_str()),
None,
)
.await;
Some(LocalSameFormatProviderDecisionInput {
auth_context,
requested_model,
auth_snapshot,
required_capabilities,
})
}
@@ -80,16 +92,34 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
spec.api_format,
&input.requested_model,
spec.require_streaming,
input.required_capabilities.as_ref(),
Some(&input.auth_snapshot),
current_unix_secs(),
)
.await?;
let candidates = prefer_local_tunnel_owner_candidates(planner_state, candidates).await;
let candidates = rank_local_execution_candidates(
planner_state,
candidates,
spec.api_format,
input.required_capabilities.as_ref(),
)
.await;
let created_at_unix_secs = current_unix_secs();
let created_at_unix_ms = current_unix_secs();
let mut attempts = Vec::with_capacity(candidates.len());
let mut affinity_remembered = false;
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
let generated_candidate_id = Uuid::new_v4().to_string();
if !affinity_remembered {
remember_scheduler_affinity_for_candidate(
planner_state,
Some(&input.auth_snapshot),
spec.api_format,
&input.requested_model,
&candidate,
);
affinity_remembered = true;
}
let extra_data = append_execution_contract_fields_to_value(
json!({
"provider_api_format": spec.api_format,
@@ -116,8 +146,9 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
&candidate,
candidate_index as u32,
&generated_candidate_id,
input.required_capabilities.as_ref(),
Some(extra_data),
created_at_unix_secs,
created_at_unix_ms,
"gateway local same-format decision request candidate upsert failed",
)
.await;
apps/aether-gateway/src/ai_pipeline/planner/passthrough/provider/family/mod.rs
+1
View File
@@ -21,6 +21,7 @@ pub(crate) struct LocalSameFormatProviderDecisionInput {
pub(crate) auth_context: ExecutionRuntimeAuthContext,
pub(crate) requested_model: String,
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
pub(crate) required_capabilities: Option<serde_json::Value>,
}
#[derive(Debug, Clone)]
apps/aether-gateway/src/ai_pipeline/planner/passthrough/provider/family/payload.rs
+3
View File
@@ -265,6 +265,8 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
json!({
"user_id": input.auth_context.user_id,
"api_key_id": input.auth_context.api_key_id,
"username": input.auth_context.username,
"api_key_name": input.auth_context.api_key_name,
"request_id": trace_id,
"candidate_id": candidate_id,
"candidate_index": candidate_index,
@@ -359,6 +361,7 @@ pub(super) async fn mark_skipped_local_same_format_provider_candidate(
candidate,
candidate_index,
candidate_id,
input.required_capabilities.as_ref(),
skip_reason,
current_unix_secs(),
"gateway local same-format decision failed to persist skipped candidate",
apps/aether-gateway/src/ai_pipeline/planner/specialized/files/decision.rs
+2
View File
@@ -281,6 +281,8 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
report_context: Some(json!({
"user_id": input.auth_context.user_id,
"api_key_id": input.auth_context.api_key_id,
"username": input.auth_context.username,
"api_key_name": input.auth_context.api_key_name,
"request_id": trace_id,
"candidate_id": candidate_id,
"candidate_index": candidate_index,
apps/aether-gateway/src/ai_pipeline/planner/specialized/files/support.rs
+37 -4
View File
@@ -4,7 +4,9 @@ use tracing::warn;
use uuid::Uuid;
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
use crate::ai_pipeline::planner::candidate_affinity::{
rank_local_execution_candidates, remember_scheduler_affinity_for_candidate,
};
use crate::ai_pipeline::{
resolve_local_decision_execution_runtime_auth_context, GatewayControlDecision,
};
@@ -20,6 +22,7 @@ pub(super) const GEMINI_FILES_REQUIRED_CAPABILITY: &str = "gemini_files";
pub(super) struct LocalGeminiFilesDecisionInput {
pub(super) auth_context: ExecutionRuntimeAuthContext,
pub(super) auth_snapshot: GatewayAuthApiKeySnapshot,
pub(super) required_capabilities: Option<serde_json::Value>,
}
#[derive(Debug, Clone)]
@@ -59,9 +62,20 @@ pub(super) async fn resolve_local_gemini_files_decision_input(
}
};
let explicit_required_capabilities = json!({ "gemini_files": true });
let required_capabilities = planner_state
.resolve_request_candidate_required_capabilities(
&auth_context.user_id,
&auth_context.api_key_id,
None,
Some(&explicit_required_capabilities),
)
.await;
Some(LocalGeminiFilesDecisionInput {
auth_context,
auth_snapshot,
required_capabilities,
})
}
@@ -80,12 +94,29 @@ pub(super) async fn materialize_local_gemini_files_candidate_attempts(
current_unix_secs(),
)
.await?;
let candidates = prefer_local_tunnel_owner_candidates(planner_state, candidates).await;
let candidates = rank_local_execution_candidates(
planner_state,
candidates,
GEMINI_FILES_CLIENT_API_FORMAT,
input.required_capabilities.as_ref(),
)
.await;
let created_at_unix_secs = current_unix_secs();
let created_at_unix_ms = current_unix_secs();
let mut attempts = Vec::with_capacity(candidates.len());
let mut affinity_remembered = false;
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
let generated_candidate_id = Uuid::new_v4().to_string();
if !affinity_remembered {
remember_scheduler_affinity_for_candidate(
planner_state,
Some(&input.auth_snapshot),
GEMINI_FILES_CLIENT_API_FORMAT,
&candidate.global_model_name,
&candidate,
);
affinity_remembered = true;
}
let extra_data = json!({
"provider_api_format": GEMINI_FILES_CLIENT_API_FORMAT,
"client_api_format": GEMINI_FILES_CLIENT_API_FORMAT,
@@ -107,8 +138,9 @@ pub(super) async fn materialize_local_gemini_files_candidate_attempts(
&candidate,
candidate_index as u32,
&generated_candidate_id,
input.required_capabilities.as_ref(),
Some(extra_data),
created_at_unix_secs,
created_at_unix_ms,
"gateway local gemini files request candidate upsert failed",
)
.await;
@@ -140,6 +172,7 @@ pub(super) async fn mark_skipped_local_gemini_files_candidate(
candidate,
candidate_index,
candidate_id,
input.required_capabilities.as_ref(),
skip_reason,
current_unix_secs(),
"gateway local gemini files failed to persist skipped candidate",
apps/aether-gateway/src/ai_pipeline/planner/specialized/video/decision.rs
+2
View File
@@ -248,6 +248,8 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
report_context: Some(json!({
"user_id": input.auth_context.user_id.clone(),
"api_key_id": input.auth_context.api_key_id.clone(),
"username": input.auth_context.username.clone(),
"api_key_name": input.auth_context.api_key_name.clone(),
"request_id": trace_id,
"candidate_id": candidate_id,
"candidate_index": candidate_index,
apps/aether-gateway/src/ai_pipeline/planner/specialized/video/support.rs
+37 -4
View File
@@ -5,7 +5,9 @@ use uuid::Uuid;
use super::{LocalVideoCreateFamily, LocalVideoCreateSpec};
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
use crate::ai_pipeline::planner::candidate_affinity::{
rank_local_execution_candidates, remember_scheduler_affinity_for_candidate,
};
use crate::ai_pipeline::{
resolve_local_decision_execution_runtime_auth_context, GatewayControlDecision,
};
@@ -18,6 +20,7 @@ pub(super) struct LocalVideoCreateDecisionInput {
pub(super) auth_context: ExecutionRuntimeAuthContext,
pub(super) requested_model: String,
pub(super) auth_snapshot: GatewayAuthApiKeySnapshot,
pub(super) required_capabilities: Option<serde_json::Value>,
}
#[derive(Debug, Clone)]
@@ -71,10 +74,20 @@ pub(super) async fn resolve_local_video_create_decision_input(
}
};
let required_capabilities = planner_state
.resolve_request_candidate_required_capabilities(
&auth_context.user_id,
&auth_context.api_key_id,
Some(requested_model.as_str()),
None,
)
.await;
Some(LocalVideoCreateDecisionInput {
auth_context,
requested_model,
auth_snapshot,
required_capabilities,
})
}
@@ -91,6 +104,7 @@ pub(super) async fn list_local_video_create_candidate_attempts(
api_format,
&input.requested_model,
false,
input.required_capabilities.as_ref(),
Some(&input.auth_snapshot),
current_unix_secs(),
)
@@ -127,12 +141,29 @@ async fn materialize_local_video_create_candidate_attempts(
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
api_format: &str,
) -> Vec<LocalVideoCreateCandidateAttempt> {
let candidates = prefer_local_tunnel_owner_candidates(state, candidates).await;
let created_at_unix_secs = current_unix_secs();
let candidates = rank_local_execution_candidates(
state,
candidates,
api_format,
input.required_capabilities.as_ref(),
)
.await;
let created_at_unix_ms = current_unix_secs();
let mut attempts = Vec::with_capacity(candidates.len());
let mut affinity_remembered = false;
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
let generated_candidate_id = Uuid::new_v4().to_string();
if !affinity_remembered {
remember_scheduler_affinity_for_candidate(
state,
Some(&input.auth_snapshot),
api_format,
&input.requested_model,
&candidate,
);
affinity_remembered = true;
}
let extra_data = json!({
"provider_api_format": api_format,
"client_api_format": api_format,
@@ -153,8 +184,9 @@ async fn materialize_local_video_create_candidate_attempts(
&candidate,
candidate_index as u32,
&generated_candidate_id,
input.required_capabilities.as_ref(),
Some(extra_data),
created_at_unix_secs,
created_at_unix_ms,
"gateway local video decision request candidate upsert failed",
)
.await;
@@ -186,6 +218,7 @@ pub(super) async fn mark_skipped_local_video_candidate(
candidate,
candidate_index,
candidate_id,
input.required_capabilities.as_ref(),
skip_reason,
current_unix_secs(),
"gateway local video decision failed to persist skipped candidate",
apps/aether-gateway/src/ai_pipeline/planner/standard/family/candidates.rs
+85 -21
View File
@@ -5,7 +5,13 @@ use serde_json::json;
use tracing::warn;
use uuid::Uuid;
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
use crate::ai_pipeline::conversion::{
request_candidate_api_formats, request_conversion_kind,
request_conversion_requires_enable_flag, request_pair_allowed_for_transport,
};
use crate::ai_pipeline::planner::candidate_affinity::{
rank_local_execution_candidates, remember_scheduler_affinity_for_candidate,
};
use crate::ai_pipeline::{
resolve_local_decision_execution_runtime_auth_context, ConversionMode, ExecutionStrategy,
GatewayControlDecision,
@@ -16,7 +22,7 @@ use crate::{append_execution_contract_fields_to_value, AppState, GatewayError};
use super::{
LocalStandardCandidateAttempt, LocalStandardDecisionInput, LocalStandardSourceFamily,
LocalStandardSourceMode, LocalStandardSpec,
LocalStandardSpec,
};
pub(super) async fn resolve_local_standard_decision_input(
@@ -63,10 +69,20 @@ pub(super) async fn resolve_local_standard_decision_input(
}
};
let required_capabilities = planner_state
.resolve_request_candidate_required_capabilities(
&auth_context.user_id,
&auth_context.api_key_id,
Some(requested_model.as_str()),
None,
)
.await;
Some(LocalStandardDecisionInput {
auth_context,
requested_model,
auth_snapshot,
required_capabilities,
})
}
@@ -79,8 +95,10 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
let planner_state = PlannerAppState::new(state);
let mut seen_candidates = BTreeSet::new();
let mut candidates = Vec::new();
for candidate_api_format in candidate_api_formats_for_spec(spec) {
let auth_snapshot = if *candidate_api_format == spec.api_format {
for candidate_api_format in
request_candidate_api_formats(spec.api_format, spec.require_streaming)
{
let auth_snapshot = if candidate_api_format == spec.api_format {
Some(&input.auth_snapshot)
} else {
None
@@ -90,6 +108,7 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
candidate_api_format,
&input.requested_model,
spec.require_streaming,
input.required_capabilities.as_ref(),
auth_snapshot,
current_unix_secs(),
)
@@ -118,13 +137,71 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
}
}
}
let candidates = prefer_local_tunnel_owner_candidates(planner_state, candidates).await;
let candidates = rank_local_execution_candidates(
planner_state,
candidates,
spec.api_format,
input.required_capabilities.as_ref(),
)
.await;
let created_at_unix_secs = current_unix_secs();
let created_at_unix_ms = current_unix_secs();
let mut attempts = Vec::with_capacity(candidates.len());
let mut affinity_remembered = false;
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
let candidate_id = Uuid::new_v4().to_string();
let provider_api_format = candidate.endpoint_api_format.trim().to_ascii_lowercase();
if provider_api_format != spec.api_format {
if let Ok(Some(transport)) = planner_state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
{
if !request_pair_allowed_for_transport(
&transport,
spec.api_format,
provider_api_format.as_str(),
) {
let skip_reason =
if request_conversion_kind(spec.api_format, provider_api_format.as_str())
.is_some()
&& request_conversion_requires_enable_flag(
spec.api_format,
provider_api_format.as_str(),
)
&& !transport.provider.enable_format_conversion
{
"format_conversion_disabled"
} else {
"transport_unsupported"
};
super::payload::mark_skipped_local_standard_candidate(
state,
input,
trace_id,
&candidate,
candidate_index as u32,
&candidate_id,
skip_reason,
)
.await;
continue;
}
}
}
if !affinity_remembered {
remember_scheduler_affinity_for_candidate(
planner_state,
Some(&input.auth_snapshot),
spec.api_format,
&input.requested_model,
&candidate,
);
affinity_remembered = true;
}
let execution_strategy = if provider_api_format == spec.api_format {
ExecutionStrategy::LocalSameFormat
} else {
@@ -166,8 +243,9 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
&candidate,
candidate_index as u32,
&candidate_id,
input.required_capabilities.as_ref(),
Some(extra_data),
created_at_unix_secs,
created_at_unix_ms,
"gateway local standard decision request candidate upsert failed",
)
.await;
@@ -213,20 +291,6 @@ fn auth_snapshot_allows_cross_format_candidate(
true
}
fn candidate_api_formats_for_spec(spec: LocalStandardSpec) -> &'static [&'static str] {
match spec.mode {
LocalStandardSourceMode::Chat | LocalStandardSourceMode::Cli => &[
"openai:chat",
"openai:cli",
"openai:compact",
"claude:chat",
"claude:cli",
"gemini:chat",
"gemini:cli",
],
}
}
fn extract_gemini_model_from_path(path: &str) -> Option<String> {
let marker = "/models/";
let start = path.find(marker)? + marker.len();
apps/aether-gateway/src/ai_pipeline/planner/standard/family/mod.rs
+1
View File
@@ -18,6 +18,7 @@ pub(super) struct LocalStandardDecisionInput {
pub(super) auth_context: ExecutionRuntimeAuthContext,
pub(super) requested_model: String,
pub(super) auth_snapshot: GatewayAuthApiKeySnapshot,
pub(super) required_capabilities: Option<serde_json::Value>,
}
#[derive(Debug, Clone)]
apps/aether-gateway/src/ai_pipeline/planner/standard/family/payload.rs
+30
View File
@@ -91,6 +91,33 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
}
};
if !crate::ai_pipeline::conversion::request_pair_allowed_for_transport(
&transport,
spec.api_format,
provider_api_format.as_str(),
) {
let skip_reason = if crate::ai_pipeline::conversion::request_conversion_requires_enable_flag(
spec.api_format,
provider_api_format.as_str(),
) && !transport.provider.enable_format_conversion
{
"format_conversion_disabled"
} else {
"transport_unsupported"
};
mark_skipped_local_standard_candidate(
state,
input,
trace_id,
&candidate,
candidate_index,
&candidate_id,
skip_reason,
)
.await;
return None;
}
if !crate::ai_pipeline::conversion::request_conversion_transport_supported(
&transport,
conversion_kind,
@@ -301,6 +328,8 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
json!({
"user_id": input.auth_context.user_id,
"api_key_id": input.auth_context.api_key_id,
"username": input.auth_context.username,
"api_key_name": input.auth_context.api_key_name,
"request_id": trace_id,
"candidate_id": candidate_id,
"candidate_index": candidate_index,
@@ -349,6 +378,7 @@ pub(super) async fn mark_skipped_local_standard_candidate(
candidate,
candidate_index,
candidate_id,
input.required_capabilities.as_ref(),
skip_reason,
current_unix_secs(),
"gateway local standard decision failed to persist skipped candidate",
apps/aether-gateway/src/ai_pipeline/planner/standard/normalize/chat.rs
-5
View File
@@ -112,11 +112,6 @@ pub(crate) fn build_cross_format_openai_chat_upstream_url(
parts.uri.query(),
false,
)),
RequestConversionKind::ToOpenAICompact => Some(build_openai_cli_url(
&transport.endpoint.base_url,
parts.uri.query(),
true,
)),
_ => None,
},
}
apps/aether-gateway/src/ai_pipeline/planner/standard/normalize/cli.rs
+6 -8
View File
@@ -10,8 +10,8 @@ use crate::ai_pipeline::transport::antigravity::{
};
use crate::ai_pipeline::transport::apply_local_body_rules;
use crate::ai_pipeline::transport::url::{
build_claude_messages_url, build_gemini_content_url, build_openai_cli_url,
build_passthrough_path_url,
build_claude_messages_url, build_gemini_content_url, build_openai_chat_url,
build_openai_cli_url, build_passthrough_path_url,
};
use crate::ai_pipeline::{
build_cross_format_openai_cli_request_body as pipeline_build_cross_format_openai_cli_request_body,
@@ -140,16 +140,15 @@ pub(crate) fn build_cross_format_openai_cli_upstream_url(
build_passthrough_path_url(&transport.endpoint.base_url, path, parts.uri.query(), &[])
}
None => match conversion_kind {
RequestConversionKind::ToOpenAIChat => Some(build_openai_chat_url(
&transport.endpoint.base_url,
parts.uri.query(),
)),
RequestConversionKind::ToOpenAIFamilyCli => Some(build_openai_cli_url(
&transport.endpoint.base_url,
parts.uri.query(),
false,
)),
RequestConversionKind::ToOpenAICompact => Some(build_openai_cli_url(
&transport.endpoint.base_url,
parts.uri.query(),
true,
)),
RequestConversionKind::ToClaudeStandard => Some(build_claude_messages_url(
&transport.endpoint.base_url,
parts.uri.query(),
@@ -160,7 +159,6 @@ pub(crate) fn build_cross_format_openai_cli_upstream_url(
upstream_is_stream,
parts.uri.query(),
),
_ => None,
},
}
}
apps/aether-gateway/src/ai_pipeline/planner/standard/normalize/tests.rs
+5 -5
View File
@@ -3,7 +3,7 @@ use serde_json::json;
use super::build_cross_format_openai_cli_request_body;
#[test]
fn builds_openai_family_cross_format_request_body_from_compact_source() {
fn builds_openai_chat_cross_format_request_body_from_openai_cli_source() {
let body_json = json!({
"model": "gpt-5",
"input": "hello",
@@ -12,18 +12,18 @@ fn builds_openai_family_cross_format_request_body_from_compact_source() {
let provider_request_body = build_cross_format_openai_cli_request_body(
&body_json,
"gpt-5-upstream",
"openai:compact",
"openai:cli",
"openai:chat",
false,
"openai",
None,
None,
)
.expect("compact to openai cli body should build");
.expect("openai cli to openai chat body should build");
assert_eq!(provider_request_body["model"], "gpt-5-upstream");
assert_eq!(provider_request_body["input"][0]["type"], "message");
assert_eq!(provider_request_body["input"][0]["role"], "user");
assert_eq!(provider_request_body["messages"][0]["role"], "user");
assert_eq!(provider_request_body["messages"][0]["content"], "hello");
}
#[test]
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/chat/decision.rs
+1 -1
View File
@@ -96,7 +96,7 @@ pub(super) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
)
.await
}
"claude:chat" | "gemini:chat" | "openai:cli" | "openai:compact" => {
"claude:chat" | "claude:cli" | "gemini:chat" | "gemini:cli" | "openai:cli" => {
build_cross_format_local_openai_chat_decision_payload_for_candidate(
state,
parts,
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/chat/decision/cross_format.rs
+27 -2
View File
@@ -6,7 +6,9 @@ use tracing::warn;
use crate::ai_pipeline::collect_control_headers;
use crate::ai_pipeline::conversion::{
request_conversion_direct_auth, request_conversion_kind, request_conversion_transport_supported,
request_conversion_direct_auth, request_conversion_kind,
request_conversion_requires_enable_flag, request_conversion_transport_supported,
request_pair_allowed_for_transport,
};
use crate::ai_pipeline::planner::common::OPENAI_CHAT_STREAM_PLAN_KIND;
use crate::ai_pipeline::planner::standard::{
@@ -50,6 +52,28 @@ pub(super) async fn build_cross_format_local_openai_chat_decision_payload_for_ca
else {
return None;
};
if !request_pair_allowed_for_transport(&transport, "openai:chat", provider_api_format.as_str())
{
let skip_reason =
if request_conversion_requires_enable_flag("openai:chat", provider_api_format.as_str())
&& !transport.provider.enable_format_conversion
{
"format_conversion_disabled"
} else {
"transport_unsupported"
};
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
return None;
}
if !request_conversion_transport_supported(transport, conversion_kind) {
mark_skipped_local_openai_chat_candidate(
state,
@@ -63,7 +87,6 @@ pub(super) async fn build_cross_format_local_openai_chat_decision_payload_for_ca
.await;
return None;
}
let resolve_auth = request_conversion_direct_auth(transport, conversion_kind);
let oauth_auth = if resolve_auth.is_none() {
match planner_state
@@ -258,6 +281,8 @@ pub(super) async fn build_cross_format_local_openai_chat_decision_payload_for_ca
json!({
"user_id": input.auth_context.user_id,
"api_key_id": input.auth_context.api_key_id,
"username": input.auth_context.username,
"api_key_name": input.auth_context.api_key_name,
"request_id": trace_id,
"candidate_id": candidate_id,
"candidate_index": candidate_index,
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/chat/decision/same_format.rs
+2
View File
@@ -233,6 +233,8 @@ pub(super) async fn build_same_format_local_openai_chat_decision_payload_for_can
json!({
"user_id": input.auth_context.user_id,
"api_key_id": input.auth_context.api_key_id,
"username": input.auth_context.username,
"api_key_name": input.auth_context.api_key_name,
"request_id": trace_id,
"candidate_id": candidate_id,
"candidate_index": candidate_index,
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/chat/decision/support.rs
+71 -4
View File
@@ -3,7 +3,13 @@ use serde_json::json;
use uuid::Uuid;
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
use crate::ai_pipeline::conversion::{
request_conversion_kind, request_conversion_requires_enable_flag,
request_pair_allowed_for_transport,
};
use crate::ai_pipeline::planner::candidate_affinity::{
rank_local_execution_candidates, remember_scheduler_affinity_for_candidate,
};
use crate::ai_pipeline::GatewayAuthApiKeySnapshot;
use crate::ai_pipeline::{ConversionMode, ExecutionStrategy, PlannerAppState};
use crate::clock::current_unix_secs;
@@ -14,6 +20,7 @@ pub(crate) struct LocalOpenAiChatDecisionInput {
pub(crate) auth_context: ExecutionRuntimeAuthContext,
pub(crate) requested_model: String,
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
pub(crate) required_capabilities: Option<serde_json::Value>,
}
#[derive(Debug, Clone)]
@@ -48,6 +55,7 @@ pub(crate) async fn mark_skipped_local_openai_chat_candidate(
candidate,
candidate_index,
candidate_id,
input.required_capabilities.as_ref(),
skip_reason,
current_unix_secs(),
"gateway local openai chat decision failed to persist skipped candidate",
@@ -62,13 +70,71 @@ pub(crate) async fn materialize_local_openai_chat_candidate_attempts(
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
) -> Vec<LocalOpenAiChatCandidateAttempt> {
let planner_state = PlannerAppState::new(state);
let candidates = prefer_local_tunnel_owner_candidates(planner_state, candidates).await;
let created_at_unix_secs = current_unix_secs();
let candidates = rank_local_execution_candidates(
planner_state,
candidates,
"openai:chat",
input.required_capabilities.as_ref(),
)
.await;
let created_at_unix_ms = current_unix_secs();
let mut attempts = Vec::with_capacity(candidates.len());
let mut affinity_remembered = false;
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
let generated_candidate_id = Uuid::new_v4().to_string();
let provider_api_format = candidate.endpoint_api_format.trim().to_ascii_lowercase();
if provider_api_format != "openai:chat" {
if let Ok(Some(transport)) = planner_state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
{
if !request_pair_allowed_for_transport(
&transport,
"openai:chat",
provider_api_format.as_str(),
) {
let skip_reason =
if request_conversion_kind("openai:chat", provider_api_format.as_str())
.is_some()
&& request_conversion_requires_enable_flag(
"openai:chat",
provider_api_format.as_str(),
)
&& !transport.provider.enable_format_conversion
{
"format_conversion_disabled"
} else {
"transport_unsupported"
};
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
&candidate,
candidate_index as u32,
&generated_candidate_id,
skip_reason,
)
.await;
continue;
}
}
}
if !affinity_remembered {
remember_scheduler_affinity_for_candidate(
planner_state,
Some(&input.auth_snapshot),
"openai:chat",
&input.requested_model,
&candidate,
);
affinity_remembered = true;
}
let (execution_strategy, conversion_mode) = if provider_api_format == "openai:chat" {
(ExecutionStrategy::LocalSameFormat, ConversionMode::None)
} else {
@@ -103,8 +169,9 @@ pub(crate) async fn materialize_local_openai_chat_candidate_attempts(
&candidate,
candidate_index as u32,
&generated_candidate_id,
input.required_capabilities.as_ref(),
Some(extra_data),
created_at_unix_secs,
created_at_unix_ms,
"gateway local openai chat decision request candidate upsert failed",
)
.await;
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/chat/plans/candidates.rs
+3 -11
View File
@@ -3,6 +3,7 @@ use std::collections::BTreeSet;
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use super::super::{GatewayError, LocalOpenAiChatDecisionInput};
use crate::ai_pipeline::conversion::request_candidate_api_formats;
use crate::ai_pipeline::{GatewayAuthApiKeySnapshot, PlannerAppState};
use crate::clock::current_unix_secs;
use crate::AppState;
@@ -17,17 +18,7 @@ pub(crate) async fn list_local_openai_chat_candidates(
let mut combined = Vec::new();
let mut seen = BTreeSet::new();
let api_formats = if require_streaming {
vec!["openai:chat", "claude:chat", "gemini:chat", "openai:cli"]
} else {
vec![
"openai:chat",
"claude:chat",
"gemini:chat",
"openai:cli",
"openai:compact",
]
};
let api_formats = request_candidate_api_formats("openai:chat", require_streaming);
for api_format in api_formats {
let auth_snapshot = if api_format == "openai:chat" {
@@ -40,6 +31,7 @@ pub(crate) async fn list_local_openai_chat_candidates(
api_format,
&input.requested_model,
require_streaming,
input.required_capabilities.as_ref(),
auth_snapshot,
now_unix_secs,
)
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/chat/plans/resolve.rs
+10
View File
@@ -108,9 +108,19 @@ pub(crate) async fn resolve_local_openai_chat_decision_input(
}
};
let required_capabilities = planner_state
.resolve_request_candidate_required_capabilities(
&auth_context.user_id,
&auth_context.api_key_id,
Some(requested_model.as_str()),
None,
)
.await;
Some(LocalOpenAiChatDecisionInput {
auth_context,
requested_model,
auth_snapshot,
required_capabilities,
})
}
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/cli/decision/payload.rs
+2
View File
@@ -95,6 +95,8 @@ pub(crate) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
json!({
"user_id": input.auth_context.user_id,
"api_key_id": input.auth_context.api_key_id,
"username": input.auth_context.username,
"api_key_name": input.auth_context.api_key_name,
"request_id": trace_id,
"candidate_id": candidate_id,
"candidate_index": candidate_index,
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/cli/decision/request.rs
+33 -1
View File
@@ -5,7 +5,9 @@ use serde_json::Value;
use tracing::warn;
use crate::ai_pipeline::conversion::{
request_conversion_direct_auth, request_conversion_kind, request_conversion_transport_supported,
request_conversion_direct_auth, request_conversion_kind,
request_conversion_requires_enable_flag, request_conversion_transport_supported,
request_pair_allowed_for_transport,
};
use crate::ai_pipeline::planner::common::force_upstream_streaming_for_provider;
use crate::ai_pipeline::planner::standard::{
@@ -115,6 +117,36 @@ pub(crate) async fn resolve_local_openai_cli_candidate_payload_parts(
let same_format = provider_api_format == spec.api_format.trim().to_ascii_lowercase();
let conversion_kind = request_conversion_kind(spec.api_format, provider_api_format.as_str());
if !same_format
&& !request_pair_allowed_for_transport(
&transport,
spec.api_format,
provider_api_format.as_str(),
)
{
let skip_reason = if conversion_kind.is_some()
&& request_conversion_requires_enable_flag(
spec.api_format,
provider_api_format.as_str(),
)
&& !transport.provider.enable_format_conversion
{
"format_conversion_disabled"
} else {
"transport_unsupported"
};
mark_skipped_local_openai_cli_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
return None;
}
let transport_supported = if same_format {
supports_local_standard_transport_with_network(&transport, provider_api_format.as_str())
} else {
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/cli/decision/support.rs
+86 -15
View File
@@ -6,8 +6,13 @@ use tracing::warn;
use uuid::Uuid;
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::conversion::request_conversion_kind;
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
use crate::ai_pipeline::conversion::{
request_candidate_api_formats, request_conversion_kind,
request_conversion_requires_enable_flag, request_pair_allowed_for_transport,
};
use crate::ai_pipeline::planner::candidate_affinity::{
rank_local_execution_candidates, remember_scheduler_affinity_for_candidate,
};
use crate::ai_pipeline::{
resolve_local_decision_execution_runtime_auth_context, ConversionMode, ExecutionStrategy,
GatewayControlDecision,
@@ -23,6 +28,7 @@ pub(crate) struct LocalOpenAiCliDecisionInput {
pub(crate) auth_context: ExecutionRuntimeAuthContext,
pub(crate) requested_model: String,
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
pub(crate) required_capabilities: Option<serde_json::Value>,
}
#[derive(Debug, Clone)]
@@ -70,10 +76,20 @@ pub(crate) async fn resolve_local_openai_cli_decision_input(
}
};
let required_capabilities = planner_state
.resolve_request_candidate_required_capabilities(
&auth_context.user_id,
&auth_context.api_key_id,
Some(requested_model.as_str()),
None,
)
.await;
Some(LocalOpenAiCliDecisionInput {
auth_context,
requested_model,
auth_snapshot,
required_capabilities,
})
}
@@ -86,8 +102,10 @@ pub(crate) async fn materialize_local_openai_cli_candidate_attempts(
let planner_state = PlannerAppState::new(state);
let mut seen_candidates = BTreeSet::new();
let mut candidates = Vec::new();
for candidate_api_format in candidate_api_formats_for_spec(spec) {
let auth_snapshot = if *candidate_api_format == spec.api_format {
for candidate_api_format in
request_candidate_api_formats(spec.api_format, spec.require_streaming)
{
let auth_snapshot = if candidate_api_format == spec.api_format {
Some(&input.auth_snapshot)
} else {
None
@@ -97,6 +115,7 @@ pub(crate) async fn materialize_local_openai_cli_candidate_attempts(
candidate_api_format,
&input.requested_model,
spec.require_streaming,
input.required_capabilities.as_ref(),
auth_snapshot,
current_unix_secs(),
)
@@ -125,13 +144,71 @@ pub(crate) async fn materialize_local_openai_cli_candidate_attempts(
}
}
}
let candidates = prefer_local_tunnel_owner_candidates(planner_state, candidates).await;
let candidates = rank_local_execution_candidates(
planner_state,
candidates,
spec.api_format,
input.required_capabilities.as_ref(),
)
.await;
let created_at_unix_secs = current_unix_secs();
let created_at_unix_ms = current_unix_secs();
let mut attempts = Vec::with_capacity(candidates.len());
let mut affinity_remembered = false;
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
let generated_candidate_id = Uuid::new_v4().to_string();
let provider_api_format = candidate.endpoint_api_format.trim().to_ascii_lowercase();
if provider_api_format != spec.api_format.trim().to_ascii_lowercase() {
if let Ok(Some(transport)) = planner_state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
{
if !request_pair_allowed_for_transport(
&transport,
spec.api_format,
provider_api_format.as_str(),
) {
let skip_reason =
if request_conversion_kind(spec.api_format, provider_api_format.as_str())
.is_some()
&& request_conversion_requires_enable_flag(
spec.api_format,
provider_api_format.as_str(),
)
&& !transport.provider.enable_format_conversion
{
"format_conversion_disabled"
} else {
"transport_unsupported"
};
mark_skipped_local_openai_cli_candidate(
state,
input,
trace_id,
&candidate,
candidate_index as u32,
&generated_candidate_id,
skip_reason,
)
.await;
continue;
}
}
}
if !affinity_remembered {
remember_scheduler_affinity_for_candidate(
planner_state,
Some(&input.auth_snapshot),
spec.api_format,
&input.requested_model,
&candidate,
);
affinity_remembered = true;
}
let execution_strategy =
if provider_api_format == spec.api_format.trim().to_ascii_lowercase() {
ExecutionStrategy::LocalSameFormat
@@ -170,8 +247,9 @@ pub(crate) async fn materialize_local_openai_cli_candidate_attempts(
&candidate,
candidate_index as u32,
&generated_candidate_id,
input.required_capabilities.as_ref(),
Some(extra_data),
created_at_unix_secs,
created_at_unix_ms,
"gateway local openai cli decision request candidate upsert failed",
)
.await;
@@ -217,14 +295,6 @@ fn auth_snapshot_allows_cross_format_openai_cli_candidate(
true
}
fn candidate_api_formats_for_spec(spec: LocalOpenAiCliSpec) -> &'static [&'static str] {
match spec.api_format {
"openai:compact" => &["openai:compact", "openai:cli", "claude:cli", "gemini:cli"],
"openai:cli" => &["openai:cli", "claude:cli", "gemini:cli"],
_ => &[],
}
}
pub(crate) async fn mark_skipped_local_openai_cli_candidate(
state: &AppState,
input: &LocalOpenAiCliDecisionInput,
@@ -242,6 +312,7 @@ pub(crate) async fn mark_skipped_local_openai_cli_candidate(
candidate,
candidate_index,
candidate_id,
input.required_capabilities.as_ref(),
skip_reason,
current_unix_secs(),
"gateway local openai cli decision failed to persist skipped candidate",
apps/aether-gateway/src/ai_pipeline/planner/state/candidate_runtime.rs
+25 -4
View File
@@ -4,6 +4,23 @@ use serde_json::Value;
use super::PlannerAppState;
impl<'a> PlannerAppState<'a> {
pub(crate) async fn resolve_request_candidate_required_capabilities(
self,
user_id: &str,
api_key_id: &str,
requested_model: Option<&str>,
explicit_required_capabilities: Option<&Value>,
) -> Option<Value> {
crate::request_candidate_runtime::resolve_request_candidate_required_capabilities(
self.app(),
user_id,
api_key_id,
requested_model,
explicit_required_capabilities,
)
.await
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn persist_available_local_candidate(
self,
@@ -13,8 +30,9 @@ impl<'a> PlannerAppState<'a> {
candidate: &SchedulerMinimalCandidateSelectionCandidate,
candidate_index: u32,
candidate_id: &str,
required_capabilities: Option<&Value>,
extra_data: Option<Value>,
created_at_unix_secs: u64,
created_at_unix_ms: u64,
error_context: &'static str,
) -> String {
crate::request_candidate_runtime::persist_available_local_candidate(
@@ -25,8 +43,9 @@ impl<'a> PlannerAppState<'a> {
candidate,
candidate_index,
candidate_id,
required_capabilities,
extra_data,
created_at_unix_secs,
created_at_unix_ms,
error_context,
)
.await
@@ -41,8 +60,9 @@ impl<'a> PlannerAppState<'a> {
candidate: &SchedulerMinimalCandidateSelectionCandidate,
candidate_index: u32,
candidate_id: &str,
required_capabilities: Option<&Value>,
skip_reason: &str,
finished_at_unix_secs: u64,
finished_at_unix_ms: u64,
error_context: &'static str,
) {
crate::request_candidate_runtime::persist_skipped_local_candidate(
@@ -53,8 +73,9 @@ impl<'a> PlannerAppState<'a> {
candidate,
candidate_index,
candidate_id,
required_capabilities,
skip_reason,
finished_at_unix_secs,
finished_at_unix_ms,
error_context,
)
.await
apps/aether-gateway/src/ai_pipeline/planner/state/scheduler.rs
+2
View File
@@ -9,6 +9,7 @@ impl<'a> PlannerAppState<'a> {
api_format: &str,
global_model_name: &str,
require_streaming: bool,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
@@ -18,6 +19,7 @@ impl<'a> PlannerAppState<'a> {
api_format,
global_model_name,
require_streaming,
required_capabilities,
auth_snapshot,
now_unix_secs,
)
apps/aether-gateway/src/ai_pipeline/pure/mod.rs
+14 -12
View File
@@ -44,18 +44,20 @@ pub(crate) use aether_ai_pipeline::api::{
provider_adaptation_descriptor_for_provider_type,
provider_adaptation_requires_eventstream_accept,
provider_adaptation_should_unwrap_stream_envelope,
provider_private_response_allows_sync_finalize, request_conversion_direct_auth,
request_conversion_kind, request_conversion_transport_supported, resolve_claude_stream_spec,
resolve_claude_sync_spec, resolve_execution_runtime_stream_plan_kind,
resolve_execution_runtime_sync_plan_kind, resolve_finalize_stream_rewrite_mode,
resolve_gemini_files_stream_spec, resolve_gemini_files_sync_spec, resolve_gemini_stream_spec,
resolve_gemini_sync_spec, resolve_local_same_format_stream_spec,
resolve_local_same_format_sync_spec, resolve_local_video_sync_spec,
resolve_openai_chat_max_tokens, resolve_openai_cli_stream_spec, resolve_openai_cli_sync_spec,
stream_body_contains_error_event, supports_stream_scheduler_decision_kind,
supports_sync_scheduler_decision_kind, sync_chat_response_conversion_kind,
sync_cli_response_conversion_kind, transform_provider_private_stream_line, value_as_u64,
CanonicalStreamFrame, ClaudeClientEmitter, ClaudeProviderState, ExecutionRuntimeAuthContext,
provider_private_response_allows_sync_finalize, request_candidate_api_formats,
request_conversion_direct_auth, request_conversion_kind,
request_conversion_requires_enable_flag, request_conversion_transport_supported,
request_pair_allowed_for_transport, resolve_claude_stream_spec, resolve_claude_sync_spec,
resolve_execution_runtime_stream_plan_kind, resolve_execution_runtime_sync_plan_kind,
resolve_finalize_stream_rewrite_mode, resolve_gemini_files_stream_spec,
resolve_gemini_files_sync_spec, resolve_gemini_stream_spec, resolve_gemini_sync_spec,
resolve_local_same_format_stream_spec, resolve_local_same_format_sync_spec,
resolve_local_video_sync_spec, resolve_openai_chat_max_tokens, resolve_openai_cli_stream_spec,
resolve_openai_cli_sync_spec, stream_body_contains_error_event,
supports_stream_scheduler_decision_kind, supports_sync_scheduler_decision_kind,
sync_chat_response_conversion_kind, sync_cli_response_conversion_kind,
transform_provider_private_stream_line, value_as_u64, CanonicalStreamFrame,
ClaudeClientEmitter, ClaudeProviderState, ExecutionRuntimeAuthContext,
FinalizeStreamRewriteMode, GatewayControlPlanRequest, GatewayControlPlanResponse,
GatewayControlSyncDecisionResponse, GeminiClientEmitter, GeminiProviderState,
LocalCoreSyncErrorKind, LocalGeminiFilesSpec, LocalOpenAiCliSpec,