Merge remote-tracking branch 'upstream/aether-rust-pioneer' into rust

This commit is contained in:
AAEE86
2026-04-10 15:08:32 +08:00
255 changed files with 15056 additions and 3205 deletions

View File

@@ -1,91 +0,0 @@
name: Build aether-hub
on:
push:
tags: ['hub-v*']
workflow_dispatch:
permissions:
contents: write
jobs:
build:
name: ${{ matrix.name }}
runs-on: ubuntu-latest
strategy:
fail-fast: false
matrix:
include:
- name: linux-amd64
target: x86_64-unknown-linux-gnu
use_cross: true
- name: linux-arm64
target: aarch64-unknown-linux-gnu
use_cross: true
steps:
- uses: actions/checkout@v5
- name: Install Rust toolchain
uses: dtolnay/rust-toolchain@stable
with:
targets: ${{ matrix.target }}
- name: Rust cache
uses: Swatinem/rust-cache@v2
with:
workspaces: aether-hub -> target
key: ${{ matrix.target }}
- name: Install cross
if: matrix.use_cross
uses: taiki-e/install-action@cross
- name: Build
working-directory: aether-hub
shell: bash
run: |
if [ "${{ matrix.use_cross }}" = "true" ]; then
cross build --release --target ${{ matrix.target }}
else
cargo build --release --target ${{ matrix.target }}
fi
- name: Package
shell: bash
run: |
cd aether-hub/target/${{ matrix.target }}/release
chmod +x aether-hub
tar czf ../../../../aether-hub-${{ matrix.name }}.tar.gz aether-hub
- name: Upload artifact
uses: actions/upload-artifact@v5
with:
name: aether-hub-${{ matrix.name }}
path: aether-hub-*.tar.gz
if-no-files-found: error
release:
needs: build
runs-on: ubuntu-latest
if: startsWith(github.ref, 'refs/tags/')
steps:
- name: Download all artifacts
uses: actions/download-artifact@v5
with:
merge-multiple: true
path: artifacts
- name: Generate checksums
working-directory: artifacts
run: sha256sum aether-hub-* > SHA256SUMS.txt
- name: Create GitHub Release
uses: softprops/action-gh-release@v2
with:
name: "${{ github.ref_name }}"
generate_release_notes: true
files: |
artifacts/aether-hub-*
artifacts/SHA256SUMS.txt
fail_on_unmatched_files: true

View File

@@ -5,8 +5,9 @@ pub(crate) use crate::ai_pipeline::{
core_error_default_client_api_format, is_core_error_finalize_kind, LocalCoreSyncErrorKind, core_error_default_client_api_format, is_core_error_finalize_kind, LocalCoreSyncErrorKind,
}; };
pub(crate) use crate::ai_pipeline::{ pub(crate) use crate::ai_pipeline::{
request_conversion_direct_auth, request_conversion_kind, request_candidate_api_formats, request_conversion_direct_auth, request_conversion_kind,
request_conversion_transport_supported, sync_chat_response_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, sync_cli_response_conversion_kind, RequestConversionKind, SyncChatResponseConversionKind,
SyncCliResponseConversionKind, SyncCliResponseConversionKind,
}; };
@@ -14,7 +15,7 @@ pub(crate) use crate::ai_pipeline::{
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::{ 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, sync_cli_response_conversion_kind, RequestConversionKind, SyncChatResponseConversionKind,
SyncCliResponseConversionKind, SyncCliResponseConversionKind,
}; };
@@ -22,7 +23,15 @@ mod tests {
#[test] #[test]
fn request_conversion_registry_supports_bidirectional_standard_matrix() { fn request_conversion_registry_supports_bidirectional_standard_matrix() {
assert_eq!( 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) Some(RequestConversionKind::ToOpenAIChat)
); );
assert_eq!( assert_eq!(
@@ -30,12 +39,16 @@ mod tests {
Some(RequestConversionKind::ToClaudeStandard) Some(RequestConversionKind::ToClaudeStandard)
); );
assert_eq!( assert_eq!(
request_conversion_kind("gemini:cli", "openai:compact"), request_conversion_kind("openai:compact", "gemini:cli"),
Some(RequestConversionKind::ToOpenAICompact) None
); );
assert_eq!( assert_eq!(
request_conversion_kind("openai:compact", "gemini:cli"), request_conversion_kind("gemini:cli", "openai:compact"),
Some(RequestConversionKind::ToGeminiStandard) None
);
assert_eq!(
request_conversion_kind("openai:chat", "openai:compact"),
None
); );
assert_eq!(request_conversion_kind("claude:chat", "claude:chat"), None); assert_eq!(request_conversion_kind("claude:chat", "claude:chat"), None);
} }
@@ -59,12 +72,50 @@ mod tests {
Some(SyncCliResponseConversionKind::ToGeminiCli) Some(SyncCliResponseConversionKind::ToGeminiCli)
); );
assert_eq!( assert_eq!(
sync_cli_response_conversion_kind("claude:cli", "openai:compact"), sync_cli_response_conversion_kind("claude:chat", "openai:cli"),
Some(SyncCliResponseConversionKind::ToOpenAIFamilyCli) 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!( assert_eq!(
sync_cli_response_conversion_kind("gemini:cli", "claude:cli"), sync_cli_response_conversion_kind("gemini:cli", "claude:cli"),
Some(SyncCliResponseConversionKind::ToClaudeCli) 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"]
);
}
} }

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] #[test]
fn claude_to_openai_chat_stream_rewriter_converts_text_deltas() { fn claude_to_openai_chat_stream_rewriter_converts_text_deltas() {
let report_context = json!({ let report_context = json!({
@@ -234,6 +358,22 @@ fn openai_cli_to_openai_chat_stream_rewriter_converts_text_deltas_immediately()
.expect("rewrite should succeed"); .expect("rewrite should succeed");
let completed_text = String::from_utf8(completed).expect("utf8 should decode"); let completed_text = String::from_utf8(completed).expect("utf8 should decode");
assert!(completed_text.contains("\"finish_reason\":\"stop\"")); 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!(completed_text.contains("data: [DONE]"));
assert!(rewriter.finish().expect("finish should succeed").is_empty()); 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("\"role\":\"assistant\""));
assert!(output_text.contains("\"content\":\"Hello Codex\"")); assert!(output_text.contains("\"content\":\"Hello Codex\""));
assert!(output_text.contains("\"finish_reason\":\"stop\"")); 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!(output_text.contains("data: [DONE]"));
assert!(rewriter.finish().expect("finish should succeed").is_empty()); 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_start"));
assert!(first_text.contains("event: content_block_delta")); assert!(first_text.contains("event: content_block_delta"));
assert!(first_text.contains("\"text\":\"Hello Claude\"")); assert!(first_text.contains("\"text\":\"Hello Claude\""));
assert!(first_text.contains("\"usage\":{\"input_tokens\":0,\"output_tokens\":0}"));
let second = rewriter let second = rewriter
.push_chunk( .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()); 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] #[test]
fn openai_chat_to_gemini_cli_stream_rewriter_converts_via_standard_matrix() { fn openai_chat_to_gemini_cli_stream_rewriter_converts_via_standard_matrix() {
let report_context = json!({ let report_context = json!({

View File

@@ -68,6 +68,8 @@ pub(crate) fn build_execution_runtime_auth_context(
ExecutionRuntimeAuthContext { ExecutionRuntimeAuthContext {
user_id: auth_context.user_id.clone(), user_id: auth_context.user_id.clone(),
api_key_id: auth_context.api_key_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, balance_remaining: auth_context.balance_remaining,
access_allowed: auth_context.access_allowed, access_allowed: auth_context.access_allowed,
} }

View File

@@ -1,9 +1,18 @@
use tracing::warn; use tracing::warn;
use crate::ai_pipeline::transport::resolve_transport_proxy_snapshot; use crate::ai_pipeline::transport::resolve_transport_proxy_snapshot;
use crate::ai_pipeline::PlannerAppState; use crate::ai_pipeline::{
use crate::AppState; GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, PlannerAppState,
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate; };
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)] #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
enum TunnelOwnerAffinityBucket { enum TunnelOwnerAffinityBucket {
@@ -12,6 +21,12 @@ enum TunnelOwnerAffinityBucket {
RemoteTunnel = 2, 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( pub(crate) async fn prefer_local_tunnel_owner_candidates(
state: PlannerAppState<'_>, state: PlannerAppState<'_>,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>, candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
@@ -28,11 +43,127 @@ pub(crate) async fn prefer_local_tunnel_owner_candidates(
.collect() .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( async fn resolve_candidate_tunnel_owner_affinity(
state: PlannerAppState<'_>, state: PlannerAppState<'_>,
candidate: &SchedulerMinimalCandidateSelectionCandidate, candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> TunnelOwnerAffinityBucket { ) -> 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( .read_provider_transport_snapshot(
&candidate.provider_id, &candidate.provider_id,
&candidate.endpoint_id, &candidate.endpoint_id,
@@ -40,8 +171,8 @@ async fn resolve_candidate_tunnel_owner_affinity(
) )
.await .await
{ {
Ok(Some(transport)) => transport, Ok(Some(transport)) => Some(transport),
Ok(None) => return TunnelOwnerAffinityBucket::Neutral, Ok(None) => None,
Err(error) => { Err(error) => {
warn!( warn!(
event_name = "candidate_affinity_transport_load_failed", event_name = "candidate_affinity_transport_load_failed",
@@ -50,12 +181,17 @@ async fn resolve_candidate_tunnel_owner_affinity(
endpoint_id = %candidate.endpoint_id, endpoint_id = %candidate.endpoint_id,
key_id = %candidate.key_id, key_id = %candidate.key_id,
error = ?error, 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 { let Some(proxy) = resolve_transport_proxy_snapshot(&transport) else {
return TunnelOwnerAffinityBucket::Neutral; 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)] #[cfg(test)]
mod tests { mod tests {
use std::time::{SystemTime, UNIX_EPOCH}; use std::time::{SystemTime, UNIX_EPOCH};
@@ -110,11 +263,15 @@ mod tests {
use serde_json::json; use serde_json::json;
use super::{ 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, SchedulerMinimalCandidateSelectionCandidate,
}; };
use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::data::GatewayDataState; use crate::data::GatewayDataState;
use crate::tunnel::TunnelAttachmentRecord; use crate::tunnel::TunnelAttachmentRecord;
use crate::{scheduler::affinity::SCHEDULER_AFFINITY_TTL, AppState};
use aether_data::repository::auth::StoredAuthApiKeySnapshot;
fn sample_candidate( fn sample_candidate(
endpoint_id: &str, endpoint_id: &str,
@@ -142,22 +299,55 @@ mod tests {
} }
fn sample_provider() -> StoredProviderCatalogProvider { 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( StoredProviderCatalogProvider::new(
"provider-1".to_string(), id.to_string(),
"provider-1".to_string(), id.to_string(),
Some("https://provider.example".to_string()), Some("https://provider.example".to_string()),
"custom".to_string(), "custom".to_string(),
) )
.expect("provider should build") .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 { 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( StoredProviderCatalogEndpoint::new(
id.to_string(), id.to_string(),
"provider-1".to_string(), provider_id.to_string(),
"openai:chat".to_string(), api_format.to_string(),
Some("openai".to_string()), Some(
api_format
.split(':')
.next()
.unwrap_or(api_format)
.to_string(),
),
Some("chat".to_string()), Some("chat".to_string()),
true, true,
) )
@@ -176,9 +366,17 @@ mod tests {
} }
fn sample_key(id: &str, node_id: &str) -> StoredProviderCatalogKey { 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( StoredProviderCatalogKey::new(
id.to_string(), id.to_string(),
"provider-1".to_string(), provider_id.to_string(),
id.to_string(), id.to_string(),
"api_key".to_string(), "api_key".to_string(),
None, None,
@@ -214,6 +412,35 @@ mod tests {
.as_secs() .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] #[tokio::test]
async fn prefers_local_tunnel_owner_candidates_before_remote_tunnel_candidates() { async fn prefers_local_tunnel_owner_candidates_before_remote_tunnel_candidates() {
let provider_catalog = InMemoryProviderCatalogReadRepository::seed( let provider_catalog = InMemoryProviderCatalogReadRepository::seed(
@@ -300,4 +527,339 @@ mod tests {
assert_eq!(reordered[0].endpoint_id, "endpoint-a"); assert_eq!(reordered[0].endpoint_id, "endpoint-a");
assert_eq!(reordered[1].endpoint_id, "endpoint-b"); 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("alice@example.com".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");
}
} }

View File

@@ -2,7 +2,9 @@ use serde_json::json;
use tracing::warn; use tracing::warn;
use uuid::Uuid; 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::{ use crate::ai_pipeline::{
resolve_local_decision_execution_runtime_auth_context, ConversionMode, ExecutionStrategy, resolve_local_decision_execution_runtime_auth_context, ConversionMode, ExecutionStrategy,
GatewayControlDecision, PlannerAppState, 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 { Some(LocalSameFormatProviderDecisionInput {
auth_context, auth_context,
requested_model, requested_model,
auth_snapshot, auth_snapshot,
required_capabilities,
}) })
} }
@@ -80,16 +92,34 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
spec.api_format, spec.api_format,
&input.requested_model, &input.requested_model,
spec.require_streaming, spec.require_streaming,
input.required_capabilities.as_ref(),
Some(&input.auth_snapshot), Some(&input.auth_snapshot),
current_unix_secs(), current_unix_secs(),
) )
.await?; .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 attempts = Vec::with_capacity(candidates.len());
let mut affinity_remembered = false;
for (candidate_index, candidate) in candidates.into_iter().enumerate() { for (candidate_index, candidate) in candidates.into_iter().enumerate() {
let generated_candidate_id = Uuid::new_v4().to_string(); 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( let extra_data = append_execution_contract_fields_to_value(
json!({ json!({
"provider_api_format": spec.api_format, "provider_api_format": spec.api_format,
@@ -116,8 +146,9 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
&candidate, &candidate,
candidate_index as u32, candidate_index as u32,
&generated_candidate_id, &generated_candidate_id,
input.required_capabilities.as_ref(),
Some(extra_data), Some(extra_data),
created_at_unix_secs, created_at_unix_ms,
"gateway local same-format decision request candidate upsert failed", "gateway local same-format decision request candidate upsert failed",
) )
.await; .await;

View File

@@ -21,6 +21,7 @@ pub(crate) struct LocalSameFormatProviderDecisionInput {
pub(crate) auth_context: ExecutionRuntimeAuthContext, pub(crate) auth_context: ExecutionRuntimeAuthContext,
pub(crate) requested_model: String, pub(crate) requested_model: String,
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot, pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
pub(crate) required_capabilities: Option<serde_json::Value>,
} }
#[derive(Debug, Clone)] #[derive(Debug, Clone)]

View File

@@ -265,6 +265,8 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
json!({ json!({
"user_id": input.auth_context.user_id, "user_id": input.auth_context.user_id,
"api_key_id": input.auth_context.api_key_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, "request_id": trace_id,
"candidate_id": candidate_id, "candidate_id": candidate_id,
"candidate_index": candidate_index, "candidate_index": candidate_index,
@@ -359,6 +361,7 @@ pub(super) async fn mark_skipped_local_same_format_provider_candidate(
candidate, candidate,
candidate_index, candidate_index,
candidate_id, candidate_id,
input.required_capabilities.as_ref(),
skip_reason, skip_reason,
current_unix_secs(), current_unix_secs(),
"gateway local same-format decision failed to persist skipped candidate", "gateway local same-format decision failed to persist skipped candidate",

View File

@@ -281,6 +281,8 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
report_context: Some(json!({ report_context: Some(json!({
"user_id": input.auth_context.user_id, "user_id": input.auth_context.user_id,
"api_key_id": input.auth_context.api_key_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, "request_id": trace_id,
"candidate_id": candidate_id, "candidate_id": candidate_id,
"candidate_index": candidate_index, "candidate_index": candidate_index,

View File

@@ -4,7 +4,9 @@ use tracing::warn;
use uuid::Uuid; use uuid::Uuid;
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext; 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::{ use crate::ai_pipeline::{
resolve_local_decision_execution_runtime_auth_context, GatewayControlDecision, 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) struct LocalGeminiFilesDecisionInput {
pub(super) auth_context: ExecutionRuntimeAuthContext, pub(super) auth_context: ExecutionRuntimeAuthContext,
pub(super) auth_snapshot: GatewayAuthApiKeySnapshot, pub(super) auth_snapshot: GatewayAuthApiKeySnapshot,
pub(super) required_capabilities: Option<serde_json::Value>,
} }
#[derive(Debug, Clone)] #[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 { Some(LocalGeminiFilesDecisionInput {
auth_context, auth_context,
auth_snapshot, auth_snapshot,
required_capabilities,
}) })
} }
@@ -80,12 +94,29 @@ pub(super) async fn materialize_local_gemini_files_candidate_attempts(
current_unix_secs(), current_unix_secs(),
) )
.await?; .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 attempts = Vec::with_capacity(candidates.len());
let mut affinity_remembered = false;
for (candidate_index, candidate) in candidates.into_iter().enumerate() { for (candidate_index, candidate) in candidates.into_iter().enumerate() {
let generated_candidate_id = Uuid::new_v4().to_string(); 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!({ let extra_data = json!({
"provider_api_format": GEMINI_FILES_CLIENT_API_FORMAT, "provider_api_format": GEMINI_FILES_CLIENT_API_FORMAT,
"client_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,
candidate_index as u32, candidate_index as u32,
&generated_candidate_id, &generated_candidate_id,
input.required_capabilities.as_ref(),
Some(extra_data), Some(extra_data),
created_at_unix_secs, created_at_unix_ms,
"gateway local gemini files request candidate upsert failed", "gateway local gemini files request candidate upsert failed",
) )
.await; .await;
@@ -140,6 +172,7 @@ pub(super) async fn mark_skipped_local_gemini_files_candidate(
candidate, candidate,
candidate_index, candidate_index,
candidate_id, candidate_id,
input.required_capabilities.as_ref(),
skip_reason, skip_reason,
current_unix_secs(), current_unix_secs(),
"gateway local gemini files failed to persist skipped candidate", "gateway local gemini files failed to persist skipped candidate",

View File

@@ -248,6 +248,8 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
report_context: Some(json!({ report_context: Some(json!({
"user_id": input.auth_context.user_id.clone(), "user_id": input.auth_context.user_id.clone(),
"api_key_id": input.auth_context.api_key_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, "request_id": trace_id,
"candidate_id": candidate_id, "candidate_id": candidate_id,
"candidate_index": candidate_index, "candidate_index": candidate_index,

View File

@@ -5,7 +5,9 @@ use uuid::Uuid;
use super::{LocalVideoCreateFamily, LocalVideoCreateSpec}; use super::{LocalVideoCreateFamily, LocalVideoCreateSpec};
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext; 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::{ use crate::ai_pipeline::{
resolve_local_decision_execution_runtime_auth_context, GatewayControlDecision, resolve_local_decision_execution_runtime_auth_context, GatewayControlDecision,
}; };
@@ -18,6 +20,7 @@ pub(super) struct LocalVideoCreateDecisionInput {
pub(super) auth_context: ExecutionRuntimeAuthContext, pub(super) auth_context: ExecutionRuntimeAuthContext,
pub(super) requested_model: String, pub(super) requested_model: String,
pub(super) auth_snapshot: GatewayAuthApiKeySnapshot, pub(super) auth_snapshot: GatewayAuthApiKeySnapshot,
pub(super) required_capabilities: Option<serde_json::Value>,
} }
#[derive(Debug, Clone)] #[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 { Some(LocalVideoCreateDecisionInput {
auth_context, auth_context,
requested_model, requested_model,
auth_snapshot, auth_snapshot,
required_capabilities,
}) })
} }
@@ -91,6 +104,7 @@ pub(super) async fn list_local_video_create_candidate_attempts(
api_format, api_format,
&input.requested_model, &input.requested_model,
false, false,
input.required_capabilities.as_ref(),
Some(&input.auth_snapshot), Some(&input.auth_snapshot),
current_unix_secs(), current_unix_secs(),
) )
@@ -127,12 +141,29 @@ async fn materialize_local_video_create_candidate_attempts(
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>, candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
api_format: &str, api_format: &str,
) -> Vec<LocalVideoCreateCandidateAttempt> { ) -> Vec<LocalVideoCreateCandidateAttempt> {
let candidates = prefer_local_tunnel_owner_candidates(state, candidates).await; let candidates = rank_local_execution_candidates(
let created_at_unix_secs = current_unix_secs(); 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 attempts = Vec::with_capacity(candidates.len());
let mut affinity_remembered = false;
for (candidate_index, candidate) in candidates.into_iter().enumerate() { for (candidate_index, candidate) in candidates.into_iter().enumerate() {
let generated_candidate_id = Uuid::new_v4().to_string(); 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!({ let extra_data = json!({
"provider_api_format": api_format, "provider_api_format": api_format,
"client_api_format": api_format, "client_api_format": api_format,
@@ -153,8 +184,9 @@ async fn materialize_local_video_create_candidate_attempts(
&candidate, &candidate,
candidate_index as u32, candidate_index as u32,
&generated_candidate_id, &generated_candidate_id,
input.required_capabilities.as_ref(),
Some(extra_data), Some(extra_data),
created_at_unix_secs, created_at_unix_ms,
"gateway local video decision request candidate upsert failed", "gateway local video decision request candidate upsert failed",
) )
.await; .await;
@@ -186,6 +218,7 @@ pub(super) async fn mark_skipped_local_video_candidate(
candidate, candidate,
candidate_index, candidate_index,
candidate_id, candidate_id,
input.required_capabilities.as_ref(),
skip_reason, skip_reason,
current_unix_secs(), current_unix_secs(),
"gateway local video decision failed to persist skipped candidate", "gateway local video decision failed to persist skipped candidate",

View File

@@ -5,7 +5,13 @@ use serde_json::json;
use tracing::warn; use tracing::warn;
use uuid::Uuid; 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::{ use crate::ai_pipeline::{
resolve_local_decision_execution_runtime_auth_context, ConversionMode, ExecutionStrategy, resolve_local_decision_execution_runtime_auth_context, ConversionMode, ExecutionStrategy,
GatewayControlDecision, GatewayControlDecision,
@@ -16,7 +22,7 @@ use crate::{append_execution_contract_fields_to_value, AppState, GatewayError};
use super::{ use super::{
LocalStandardCandidateAttempt, LocalStandardDecisionInput, LocalStandardSourceFamily, LocalStandardCandidateAttempt, LocalStandardDecisionInput, LocalStandardSourceFamily,
LocalStandardSourceMode, LocalStandardSpec, LocalStandardSpec,
}; };
pub(super) async fn resolve_local_standard_decision_input( 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 { Some(LocalStandardDecisionInput {
auth_context, auth_context,
requested_model, requested_model,
auth_snapshot, auth_snapshot,
required_capabilities,
}) })
} }
@@ -79,8 +95,10 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
let planner_state = PlannerAppState::new(state); let planner_state = PlannerAppState::new(state);
let mut seen_candidates = BTreeSet::new(); let mut seen_candidates = BTreeSet::new();
let mut candidates = Vec::new(); let mut candidates = Vec::new();
for candidate_api_format in candidate_api_formats_for_spec(spec) { for candidate_api_format in
let auth_snapshot = if *candidate_api_format == spec.api_format { request_candidate_api_formats(spec.api_format, spec.require_streaming)
{
let auth_snapshot = if candidate_api_format == spec.api_format {
Some(&input.auth_snapshot) Some(&input.auth_snapshot)
} else { } else {
None None
@@ -90,6 +108,7 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
candidate_api_format, candidate_api_format,
&input.requested_model, &input.requested_model,
spec.require_streaming, spec.require_streaming,
input.required_capabilities.as_ref(),
auth_snapshot, auth_snapshot,
current_unix_secs(), 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 attempts = Vec::with_capacity(candidates.len());
let mut affinity_remembered = false;
for (candidate_index, candidate) in candidates.into_iter().enumerate() { for (candidate_index, candidate) in candidates.into_iter().enumerate() {
let candidate_id = Uuid::new_v4().to_string(); let candidate_id = Uuid::new_v4().to_string();
let provider_api_format = candidate.endpoint_api_format.trim().to_ascii_lowercase(); 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 { let execution_strategy = if provider_api_format == spec.api_format {
ExecutionStrategy::LocalSameFormat ExecutionStrategy::LocalSameFormat
} else { } else {
@@ -166,8 +243,9 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
&candidate, &candidate,
candidate_index as u32, candidate_index as u32,
&candidate_id, &candidate_id,
input.required_capabilities.as_ref(),
Some(extra_data), Some(extra_data),
created_at_unix_secs, created_at_unix_ms,
"gateway local standard decision request candidate upsert failed", "gateway local standard decision request candidate upsert failed",
) )
.await; .await;
@@ -213,20 +291,6 @@ fn auth_snapshot_allows_cross_format_candidate(
true 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> { fn extract_gemini_model_from_path(path: &str) -> Option<String> {
let marker = "/models/"; let marker = "/models/";
let start = path.find(marker)? + marker.len(); let start = path.find(marker)? + marker.len();

View File

@@ -18,6 +18,7 @@ pub(super) struct LocalStandardDecisionInput {
pub(super) auth_context: ExecutionRuntimeAuthContext, pub(super) auth_context: ExecutionRuntimeAuthContext,
pub(super) requested_model: String, pub(super) requested_model: String,
pub(super) auth_snapshot: GatewayAuthApiKeySnapshot, pub(super) auth_snapshot: GatewayAuthApiKeySnapshot,
pub(super) required_capabilities: Option<serde_json::Value>,
} }
#[derive(Debug, Clone)] #[derive(Debug, Clone)]

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( if !crate::ai_pipeline::conversion::request_conversion_transport_supported(
&transport, &transport,
conversion_kind, conversion_kind,
@@ -301,6 +328,8 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
json!({ json!({
"user_id": input.auth_context.user_id, "user_id": input.auth_context.user_id,
"api_key_id": input.auth_context.api_key_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, "request_id": trace_id,
"candidate_id": candidate_id, "candidate_id": candidate_id,
"candidate_index": candidate_index, "candidate_index": candidate_index,
@@ -349,6 +378,7 @@ pub(super) async fn mark_skipped_local_standard_candidate(
candidate, candidate,
candidate_index, candidate_index,
candidate_id, candidate_id,
input.required_capabilities.as_ref(),
skip_reason, skip_reason,
current_unix_secs(), current_unix_secs(),
"gateway local standard decision failed to persist skipped candidate", "gateway local standard decision failed to persist skipped candidate",

View File

@@ -112,11 +112,6 @@ pub(crate) fn build_cross_format_openai_chat_upstream_url(
parts.uri.query(), parts.uri.query(),
false, false,
)), )),
RequestConversionKind::ToOpenAICompact => Some(build_openai_cli_url(
&transport.endpoint.base_url,
parts.uri.query(),
true,
)),
_ => None, _ => None,
}, },
} }

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::apply_local_body_rules;
use crate::ai_pipeline::transport::url::{ use crate::ai_pipeline::transport::url::{
build_claude_messages_url, build_gemini_content_url, build_openai_cli_url, build_claude_messages_url, build_gemini_content_url, build_openai_chat_url,
build_passthrough_path_url, build_openai_cli_url, build_passthrough_path_url,
}; };
use crate::ai_pipeline::{ use crate::ai_pipeline::{
build_cross_format_openai_cli_request_body as pipeline_build_cross_format_openai_cli_request_body, 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(), &[]) build_passthrough_path_url(&transport.endpoint.base_url, path, parts.uri.query(), &[])
} }
None => match conversion_kind { 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( RequestConversionKind::ToOpenAIFamilyCli => Some(build_openai_cli_url(
&transport.endpoint.base_url, &transport.endpoint.base_url,
parts.uri.query(), parts.uri.query(),
false, false,
)), )),
RequestConversionKind::ToOpenAICompact => Some(build_openai_cli_url(
&transport.endpoint.base_url,
parts.uri.query(),
true,
)),
RequestConversionKind::ToClaudeStandard => Some(build_claude_messages_url( RequestConversionKind::ToClaudeStandard => Some(build_claude_messages_url(
&transport.endpoint.base_url, &transport.endpoint.base_url,
parts.uri.query(), parts.uri.query(),
@@ -160,7 +159,6 @@ pub(crate) fn build_cross_format_openai_cli_upstream_url(
upstream_is_stream, upstream_is_stream,
parts.uri.query(), parts.uri.query(),
), ),
_ => None,
}, },
} }
} }

View File

@@ -3,7 +3,7 @@ use serde_json::json;
use super::build_cross_format_openai_cli_request_body; use super::build_cross_format_openai_cli_request_body;
#[test] #[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!({ let body_json = json!({
"model": "gpt-5", "model": "gpt-5",
"input": "hello", "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( let provider_request_body = build_cross_format_openai_cli_request_body(
&body_json, &body_json,
"gpt-5-upstream", "gpt-5-upstream",
"openai:compact",
"openai:cli", "openai:cli",
"openai:chat",
false, false,
"openai", "openai",
None, None,
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["model"], "gpt-5-upstream");
assert_eq!(provider_request_body["input"][0]["type"], "message"); assert_eq!(provider_request_body["messages"][0]["role"], "user");
assert_eq!(provider_request_body["input"][0]["role"], "user"); assert_eq!(provider_request_body["messages"][0]["content"], "hello");
} }
#[test] #[test]

View File

@@ -96,7 +96,7 @@ pub(super) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
) )
.await .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( build_cross_format_local_openai_chat_decision_payload_for_candidate(
state, state,
parts, parts,

View File

@@ -6,7 +6,9 @@ use tracing::warn;
use crate::ai_pipeline::collect_control_headers; use crate::ai_pipeline::collect_control_headers;
use crate::ai_pipeline::conversion::{ 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::common::OPENAI_CHAT_STREAM_PLAN_KIND;
use crate::ai_pipeline::planner::standard::{ 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 { else {
return None; 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) { if !request_conversion_transport_supported(transport, conversion_kind) {
mark_skipped_local_openai_chat_candidate( mark_skipped_local_openai_chat_candidate(
state, state,
@@ -63,7 +87,6 @@ pub(super) async fn build_cross_format_local_openai_chat_decision_payload_for_ca
.await; .await;
return None; return None;
} }
let resolve_auth = request_conversion_direct_auth(transport, conversion_kind); let resolve_auth = request_conversion_direct_auth(transport, conversion_kind);
let oauth_auth = if resolve_auth.is_none() { let oauth_auth = if resolve_auth.is_none() {
match planner_state match planner_state
@@ -258,6 +281,8 @@ pub(super) async fn build_cross_format_local_openai_chat_decision_payload_for_ca
json!({ json!({
"user_id": input.auth_context.user_id, "user_id": input.auth_context.user_id,
"api_key_id": input.auth_context.api_key_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, "request_id": trace_id,
"candidate_id": candidate_id, "candidate_id": candidate_id,
"candidate_index": candidate_index, "candidate_index": candidate_index,

View File

@@ -233,6 +233,8 @@ pub(super) async fn build_same_format_local_openai_chat_decision_payload_for_can
json!({ json!({
"user_id": input.auth_context.user_id, "user_id": input.auth_context.user_id,
"api_key_id": input.auth_context.api_key_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, "request_id": trace_id,
"candidate_id": candidate_id, "candidate_id": candidate_id,
"candidate_index": candidate_index, "candidate_index": candidate_index,

View File

@@ -3,7 +3,13 @@ use serde_json::json;
use uuid::Uuid; use uuid::Uuid;
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext; 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::GatewayAuthApiKeySnapshot;
use crate::ai_pipeline::{ConversionMode, ExecutionStrategy, PlannerAppState}; use crate::ai_pipeline::{ConversionMode, ExecutionStrategy, PlannerAppState};
use crate::clock::current_unix_secs; use crate::clock::current_unix_secs;
@@ -14,6 +20,7 @@ pub(crate) struct LocalOpenAiChatDecisionInput {
pub(crate) auth_context: ExecutionRuntimeAuthContext, pub(crate) auth_context: ExecutionRuntimeAuthContext,
pub(crate) requested_model: String, pub(crate) requested_model: String,
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot, pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
pub(crate) required_capabilities: Option<serde_json::Value>,
} }
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
@@ -48,6 +55,7 @@ pub(crate) async fn mark_skipped_local_openai_chat_candidate(
candidate, candidate,
candidate_index, candidate_index,
candidate_id, candidate_id,
input.required_capabilities.as_ref(),
skip_reason, skip_reason,
current_unix_secs(), current_unix_secs(),
"gateway local openai chat decision failed to persist skipped candidate", "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>, candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
) -> Vec<LocalOpenAiChatCandidateAttempt> { ) -> Vec<LocalOpenAiChatCandidateAttempt> {
let planner_state = PlannerAppState::new(state); let planner_state = PlannerAppState::new(state);
let candidates = prefer_local_tunnel_owner_candidates(planner_state, candidates).await; let candidates = rank_local_execution_candidates(
let created_at_unix_secs = current_unix_secs(); 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 attempts = Vec::with_capacity(candidates.len());
let mut affinity_remembered = false;
for (candidate_index, candidate) in candidates.into_iter().enumerate() { for (candidate_index, candidate) in candidates.into_iter().enumerate() {
let generated_candidate_id = Uuid::new_v4().to_string(); let generated_candidate_id = Uuid::new_v4().to_string();
let provider_api_format = candidate.endpoint_api_format.trim().to_ascii_lowercase(); 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" { let (execution_strategy, conversion_mode) = if provider_api_format == "openai:chat" {
(ExecutionStrategy::LocalSameFormat, ConversionMode::None) (ExecutionStrategy::LocalSameFormat, ConversionMode::None)
} else { } else {
@@ -103,8 +169,9 @@ pub(crate) async fn materialize_local_openai_chat_candidate_attempts(
&candidate, &candidate,
candidate_index as u32, candidate_index as u32,
&generated_candidate_id, &generated_candidate_id,
input.required_capabilities.as_ref(),
Some(extra_data), Some(extra_data),
created_at_unix_secs, created_at_unix_ms,
"gateway local openai chat decision request candidate upsert failed", "gateway local openai chat decision request candidate upsert failed",
) )
.await; .await;

View File

@@ -3,6 +3,7 @@ use std::collections::BTreeSet;
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate; use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use super::super::{GatewayError, LocalOpenAiChatDecisionInput}; use super::super::{GatewayError, LocalOpenAiChatDecisionInput};
use crate::ai_pipeline::conversion::request_candidate_api_formats;
use crate::ai_pipeline::{GatewayAuthApiKeySnapshot, PlannerAppState}; use crate::ai_pipeline::{GatewayAuthApiKeySnapshot, PlannerAppState};
use crate::clock::current_unix_secs; use crate::clock::current_unix_secs;
use crate::AppState; use crate::AppState;
@@ -17,17 +18,7 @@ pub(crate) async fn list_local_openai_chat_candidates(
let mut combined = Vec::new(); let mut combined = Vec::new();
let mut seen = BTreeSet::new(); let mut seen = BTreeSet::new();
let api_formats = if require_streaming { let api_formats = request_candidate_api_formats("openai:chat", require_streaming);
vec!["openai:chat", "claude:chat", "gemini:chat", "openai:cli"]
} else {
vec![
"openai:chat",
"claude:chat",
"gemini:chat",
"openai:cli",
"openai:compact",
]
};
for api_format in api_formats { for api_format in api_formats {
let auth_snapshot = if api_format == "openai:chat" { let auth_snapshot = if api_format == "openai:chat" {
@@ -40,6 +31,7 @@ pub(crate) async fn list_local_openai_chat_candidates(
api_format, api_format,
&input.requested_model, &input.requested_model,
require_streaming, require_streaming,
input.required_capabilities.as_ref(),
auth_snapshot, auth_snapshot,
now_unix_secs, now_unix_secs,
) )

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 { Some(LocalOpenAiChatDecisionInput {
auth_context, auth_context,
requested_model, requested_model,
auth_snapshot, auth_snapshot,
required_capabilities,
}) })
} }

View File

@@ -95,6 +95,8 @@ pub(crate) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
json!({ json!({
"user_id": input.auth_context.user_id, "user_id": input.auth_context.user_id,
"api_key_id": input.auth_context.api_key_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, "request_id": trace_id,
"candidate_id": candidate_id, "candidate_id": candidate_id,
"candidate_index": candidate_index, "candidate_index": candidate_index,

View File

@@ -5,7 +5,9 @@ use serde_json::Value;
use tracing::warn; use tracing::warn;
use crate::ai_pipeline::conversion::{ 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::common::force_upstream_streaming_for_provider;
use crate::ai_pipeline::planner::standard::{ 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 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()); 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 { let transport_supported = if same_format {
supports_local_standard_transport_with_network(&transport, provider_api_format.as_str()) supports_local_standard_transport_with_network(&transport, provider_api_format.as_str())
} else { } else {

View File

@@ -6,8 +6,13 @@ use tracing::warn;
use uuid::Uuid; use uuid::Uuid;
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext; use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::conversion::request_conversion_kind; use crate::ai_pipeline::conversion::{
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates; 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::{ use crate::ai_pipeline::{
resolve_local_decision_execution_runtime_auth_context, ConversionMode, ExecutionStrategy, resolve_local_decision_execution_runtime_auth_context, ConversionMode, ExecutionStrategy,
GatewayControlDecision, GatewayControlDecision,
@@ -23,6 +28,7 @@ pub(crate) struct LocalOpenAiCliDecisionInput {
pub(crate) auth_context: ExecutionRuntimeAuthContext, pub(crate) auth_context: ExecutionRuntimeAuthContext,
pub(crate) requested_model: String, pub(crate) requested_model: String,
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot, pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
pub(crate) required_capabilities: Option<serde_json::Value>,
} }
#[derive(Debug, Clone)] #[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 { Some(LocalOpenAiCliDecisionInput {
auth_context, auth_context,
requested_model, requested_model,
auth_snapshot, 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 planner_state = PlannerAppState::new(state);
let mut seen_candidates = BTreeSet::new(); let mut seen_candidates = BTreeSet::new();
let mut candidates = Vec::new(); let mut candidates = Vec::new();
for candidate_api_format in candidate_api_formats_for_spec(spec) { for candidate_api_format in
let auth_snapshot = if *candidate_api_format == spec.api_format { request_candidate_api_formats(spec.api_format, spec.require_streaming)
{
let auth_snapshot = if candidate_api_format == spec.api_format {
Some(&input.auth_snapshot) Some(&input.auth_snapshot)
} else { } else {
None None
@@ -97,6 +115,7 @@ pub(crate) async fn materialize_local_openai_cli_candidate_attempts(
candidate_api_format, candidate_api_format,
&input.requested_model, &input.requested_model,
spec.require_streaming, spec.require_streaming,
input.required_capabilities.as_ref(),
auth_snapshot, auth_snapshot,
current_unix_secs(), 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 attempts = Vec::with_capacity(candidates.len());
let mut affinity_remembered = false;
for (candidate_index, candidate) in candidates.into_iter().enumerate() { for (candidate_index, candidate) in candidates.into_iter().enumerate() {
let generated_candidate_id = Uuid::new_v4().to_string(); let generated_candidate_id = Uuid::new_v4().to_string();
let provider_api_format = candidate.endpoint_api_format.trim().to_ascii_lowercase(); 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 = let execution_strategy =
if provider_api_format == spec.api_format.trim().to_ascii_lowercase() { if provider_api_format == spec.api_format.trim().to_ascii_lowercase() {
ExecutionStrategy::LocalSameFormat ExecutionStrategy::LocalSameFormat
@@ -170,8 +247,9 @@ pub(crate) async fn materialize_local_openai_cli_candidate_attempts(
&candidate, &candidate,
candidate_index as u32, candidate_index as u32,
&generated_candidate_id, &generated_candidate_id,
input.required_capabilities.as_ref(),
Some(extra_data), Some(extra_data),
created_at_unix_secs, created_at_unix_ms,
"gateway local openai cli decision request candidate upsert failed", "gateway local openai cli decision request candidate upsert failed",
) )
.await; .await;
@@ -217,14 +295,6 @@ fn auth_snapshot_allows_cross_format_openai_cli_candidate(
true 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( pub(crate) async fn mark_skipped_local_openai_cli_candidate(
state: &AppState, state: &AppState,
input: &LocalOpenAiCliDecisionInput, input: &LocalOpenAiCliDecisionInput,
@@ -242,6 +312,7 @@ pub(crate) async fn mark_skipped_local_openai_cli_candidate(
candidate, candidate,
candidate_index, candidate_index,
candidate_id, candidate_id,
input.required_capabilities.as_ref(),
skip_reason, skip_reason,
current_unix_secs(), current_unix_secs(),
"gateway local openai cli decision failed to persist skipped candidate", "gateway local openai cli decision failed to persist skipped candidate",

View File

@@ -4,6 +4,23 @@ use serde_json::Value;
use super::PlannerAppState; use super::PlannerAppState;
impl<'a> PlannerAppState<'a> { 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)] #[allow(clippy::too_many_arguments)]
pub(crate) async fn persist_available_local_candidate( pub(crate) async fn persist_available_local_candidate(
self, self,
@@ -13,8 +30,9 @@ impl<'a> PlannerAppState<'a> {
candidate: &SchedulerMinimalCandidateSelectionCandidate, candidate: &SchedulerMinimalCandidateSelectionCandidate,
candidate_index: u32, candidate_index: u32,
candidate_id: &str, candidate_id: &str,
required_capabilities: Option<&Value>,
extra_data: Option<Value>, extra_data: Option<Value>,
created_at_unix_secs: u64, created_at_unix_ms: u64,
error_context: &'static str, error_context: &'static str,
) -> String { ) -> String {
crate::request_candidate_runtime::persist_available_local_candidate( crate::request_candidate_runtime::persist_available_local_candidate(
@@ -25,8 +43,9 @@ impl<'a> PlannerAppState<'a> {
candidate, candidate,
candidate_index, candidate_index,
candidate_id, candidate_id,
required_capabilities,
extra_data, extra_data,
created_at_unix_secs, created_at_unix_ms,
error_context, error_context,
) )
.await .await
@@ -41,8 +60,9 @@ impl<'a> PlannerAppState<'a> {
candidate: &SchedulerMinimalCandidateSelectionCandidate, candidate: &SchedulerMinimalCandidateSelectionCandidate,
candidate_index: u32, candidate_index: u32,
candidate_id: &str, candidate_id: &str,
required_capabilities: Option<&Value>,
skip_reason: &str, skip_reason: &str,
finished_at_unix_secs: u64, finished_at_unix_ms: u64,
error_context: &'static str, error_context: &'static str,
) { ) {
crate::request_candidate_runtime::persist_skipped_local_candidate( crate::request_candidate_runtime::persist_skipped_local_candidate(
@@ -53,8 +73,9 @@ impl<'a> PlannerAppState<'a> {
candidate, candidate,
candidate_index, candidate_index,
candidate_id, candidate_id,
required_capabilities,
skip_reason, skip_reason,
finished_at_unix_secs, finished_at_unix_ms,
error_context, error_context,
) )
.await .await

View File

@@ -9,6 +9,7 @@ impl<'a> PlannerAppState<'a> {
api_format: &str, api_format: &str,
global_model_name: &str, global_model_name: &str,
require_streaming: bool, require_streaming: bool,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>, auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64, now_unix_secs: u64,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> { ) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
@@ -18,6 +19,7 @@ impl<'a> PlannerAppState<'a> {
api_format, api_format,
global_model_name, global_model_name,
require_streaming, require_streaming,
required_capabilities,
auth_snapshot, auth_snapshot,
now_unix_secs, now_unix_secs,
) )

View File

@@ -44,18 +44,20 @@ pub(crate) use aether_ai_pipeline::api::{
provider_adaptation_descriptor_for_provider_type, provider_adaptation_descriptor_for_provider_type,
provider_adaptation_requires_eventstream_accept, provider_adaptation_requires_eventstream_accept,
provider_adaptation_should_unwrap_stream_envelope, provider_adaptation_should_unwrap_stream_envelope,
provider_private_response_allows_sync_finalize, request_conversion_direct_auth, provider_private_response_allows_sync_finalize, request_candidate_api_formats,
request_conversion_kind, request_conversion_transport_supported, resolve_claude_stream_spec, request_conversion_direct_auth, request_conversion_kind,
resolve_claude_sync_spec, resolve_execution_runtime_stream_plan_kind, request_conversion_requires_enable_flag, request_conversion_transport_supported,
resolve_execution_runtime_sync_plan_kind, resolve_finalize_stream_rewrite_mode, request_pair_allowed_for_transport, resolve_claude_stream_spec, resolve_claude_sync_spec,
resolve_gemini_files_stream_spec, resolve_gemini_files_sync_spec, resolve_gemini_stream_spec, resolve_execution_runtime_stream_plan_kind, resolve_execution_runtime_sync_plan_kind,
resolve_gemini_sync_spec, resolve_local_same_format_stream_spec, resolve_finalize_stream_rewrite_mode, resolve_gemini_files_stream_spec,
resolve_local_same_format_sync_spec, resolve_local_video_sync_spec, resolve_gemini_files_sync_spec, resolve_gemini_stream_spec, resolve_gemini_sync_spec,
resolve_openai_chat_max_tokens, resolve_openai_cli_stream_spec, resolve_openai_cli_sync_spec, resolve_local_same_format_stream_spec, resolve_local_same_format_sync_spec,
stream_body_contains_error_event, supports_stream_scheduler_decision_kind, resolve_local_video_sync_spec, resolve_openai_chat_max_tokens, resolve_openai_cli_stream_spec,
supports_sync_scheduler_decision_kind, sync_chat_response_conversion_kind, resolve_openai_cli_sync_spec, stream_body_contains_error_event,
sync_cli_response_conversion_kind, transform_provider_private_stream_line, value_as_u64, supports_stream_scheduler_decision_kind, supports_sync_scheduler_decision_kind,
CanonicalStreamFrame, ClaudeClientEmitter, ClaudeProviderState, ExecutionRuntimeAuthContext, 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, FinalizeStreamRewriteMode, GatewayControlPlanRequest, GatewayControlPlanResponse,
GatewayControlSyncDecisionResponse, GeminiClientEmitter, GeminiProviderState, GatewayControlSyncDecisionResponse, GeminiClientEmitter, GeminiProviderState,
LocalCoreSyncErrorKind, LocalGeminiFilesSpec, LocalOpenAiCliSpec, LocalCoreSyncErrorKind, LocalGeminiFilesSpec, LocalOpenAiCliSpec,

View File

@@ -253,7 +253,7 @@ async fn persist_cancelled_video_task(
next_poll_at_unix_secs: None, next_poll_at_unix_secs: None,
poll_count: task.poll_count, poll_count: task.poll_count,
max_poll_count: task.max_poll_count, max_poll_count: task.max_poll_count,
created_at_unix_secs: task.created_at_unix_secs, created_at_unix_ms: task.created_at_unix_ms,
submitted_at_unix_secs: task.submitted_at_unix_secs, submitted_at_unix_secs: task.submitted_at_unix_secs,
completed_at_unix_secs: Some(now_unix_secs), completed_at_unix_secs: Some(now_unix_secs),
updated_at_unix_secs: now_unix_secs, updated_at_unix_secs: now_unix_secs,

View File

@@ -272,7 +272,7 @@ fn build_successful_poll_update(
record.resolution = task.resolution.clone(); record.resolution = task.resolution.clone();
record.aspect_ratio = task.aspect_ratio.clone(); record.aspect_ratio = task.aspect_ratio.clone();
record.size = task.size.clone(); record.size = task.size.clone();
record.created_at_unix_secs = task.created_at_unix_secs; record.created_at_unix_ms = task.created_at_unix_ms;
record.submitted_at_unix_secs = task.submitted_at_unix_secs; record.submitted_at_unix_secs = task.submitted_at_unix_secs;
record.updated_at_unix_secs = now_unix_secs; record.updated_at_unix_secs = now_unix_secs;
record.retry_count = task.retry_count; record.retry_count = task.retry_count;
@@ -379,7 +379,7 @@ fn stored_task_to_upsert(task: &StoredVideoTask) -> UpsertVideoTask {
next_poll_at_unix_secs: task.next_poll_at_unix_secs, next_poll_at_unix_secs: task.next_poll_at_unix_secs,
poll_count: task.poll_count, poll_count: task.poll_count,
max_poll_count: task.max_poll_count.max(1), max_poll_count: task.max_poll_count.max(1),
created_at_unix_secs: task.created_at_unix_secs, created_at_unix_ms: task.created_at_unix_ms,
submitted_at_unix_secs: task.submitted_at_unix_secs, submitted_at_unix_secs: task.submitted_at_unix_secs,
completed_at_unix_secs: task.completed_at_unix_secs, completed_at_unix_secs: task.completed_at_unix_secs,
updated_at_unix_secs: task.updated_at_unix_secs, updated_at_unix_secs: task.updated_at_unix_secs,

View File

@@ -1,5 +1,5 @@
use axum::body::Body; use axum::body::Body;
use axum::http::{self, Response}; use axum::http::{self, Response, StatusCode};
use tracing::{info, warn}; use tracing::{info, warn};
use crate::control::GatewayControlDecision; use crate::control::GatewayControlDecision;
@@ -68,12 +68,8 @@ pub(crate) fn emit_admin_audit(
) )
}; };
let audit_status = if response.status().is_success() { let (audit_status, log_level) = classify_admin_audit_response(method, response.status());
"completed" if log_level == AdminAuditLogLevel::Info {
} else {
"failed"
};
if response.status().is_success() {
info!( info!(
event_name, event_name,
log_type = "audit", log_type = "audit",
@@ -116,6 +112,25 @@ pub(crate) fn emit_admin_audit(
} }
} }
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum AdminAuditLogLevel {
Info,
Warn,
}
fn classify_admin_audit_response(
method: &http::Method,
status: StatusCode,
) -> (&'static str, AdminAuditLogLevel) {
if status.is_success() {
return ("completed", AdminAuditLogLevel::Info);
}
if status == StatusCode::NOT_FOUND && is_admin_read_method(method) {
return ("not_found", AdminAuditLogLevel::Info);
}
("failed", AdminAuditLogLevel::Warn)
}
fn default_target_type(route_family: &str) -> &str { fn default_target_type(route_family: &str) -> &str {
route_family route_family
.strip_suffix("_manage") .strip_suffix("_manage")
@@ -129,3 +144,29 @@ fn is_admin_mutation_method(method: &http::Method) -> bool {
http::Method::POST | http::Method::PUT | http::Method::PATCH | http::Method::DELETE http::Method::POST | http::Method::PUT | http::Method::PATCH | http::Method::DELETE
) )
} }
fn is_admin_read_method(method: &http::Method) -> bool {
matches!(*method, http::Method::GET | http::Method::HEAD)
}
#[cfg(test)]
mod tests {
use super::{classify_admin_audit_response, AdminAuditLogLevel};
use axum::http::{Method, StatusCode};
#[test]
fn classifies_read_not_found_as_info_not_found() {
assert_eq!(
classify_admin_audit_response(&Method::GET, StatusCode::NOT_FOUND),
("not_found", AdminAuditLogLevel::Info)
);
}
#[test]
fn classifies_mutation_not_found_as_warn_failed() {
assert_eq!(
classify_admin_audit_response(&Method::DELETE, StatusCode::NOT_FOUND),
("failed", AdminAuditLogLevel::Warn)
);
}
}

View File

@@ -6,4 +6,6 @@ mod scheduler_affinity;
pub(crate) use auth_api_key_last_used::AuthApiKeyLastUsedCache; pub(crate) use auth_api_key_last_used::AuthApiKeyLastUsedCache;
pub(crate) use auth_context::AuthContextCache; pub(crate) use auth_context::AuthContextCache;
pub(crate) use direct_plan_bypass::DirectPlanBypassCache; pub(crate) use direct_plan_bypass::DirectPlanBypassCache;
pub(crate) use scheduler_affinity::{SchedulerAffinityCache, SchedulerAffinityTarget}; pub(crate) use scheduler_affinity::{
SchedulerAffinityCache, SchedulerAffinitySnapshotEntry, SchedulerAffinityTarget,
};

View File

@@ -1,6 +1,6 @@
use std::time::Duration; use std::time::Duration;
use aether_cache::ExpiringMap; use aether_cache::{ExpiringMap, ExpiringMapFreshEntry};
pub(crate) use aether_scheduler_core::SchedulerAffinityTarget; pub(crate) use aether_scheduler_core::SchedulerAffinityTarget;
#[derive(Debug, Default)] #[derive(Debug, Default)]
@@ -8,6 +8,13 @@ pub(crate) struct SchedulerAffinityCache {
entries: ExpiringMap<String, SchedulerAffinityTarget>, entries: ExpiringMap<String, SchedulerAffinityTarget>,
} }
#[derive(Debug, Clone)]
pub(crate) struct SchedulerAffinitySnapshotEntry {
pub(crate) cache_key: String,
pub(crate) target: SchedulerAffinityTarget,
pub(crate) age: Duration,
}
impl SchedulerAffinityCache { impl SchedulerAffinityCache {
pub(crate) fn get_fresh( pub(crate) fn get_fresh(
&self, &self,
@@ -31,4 +38,18 @@ impl SchedulerAffinityCache {
pub(crate) fn remove(&self, cache_key: &str) -> Option<SchedulerAffinityTarget> { pub(crate) fn remove(&self, cache_key: &str) -> Option<SchedulerAffinityTarget> {
self.entries.remove(&cache_key.to_string()) self.entries.remove(&cache_key.to_string())
} }
pub(crate) fn fresh_entries(&self, ttl: Duration) -> Vec<SchedulerAffinitySnapshotEntry> {
self.entries
.snapshot_fresh(ttl)
.into_iter()
.map(
|ExpiringMapFreshEntry { key, value, age }| SchedulerAffinitySnapshotEntry {
cache_key: key,
target: value,
age,
},
)
.collect()
}
} }

View File

@@ -6,3 +6,10 @@ pub(crate) fn current_unix_secs() -> u64 {
.unwrap_or_default() .unwrap_or_default()
.as_secs() .as_secs()
} }
pub(crate) fn current_unix_ms() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64
}

View File

@@ -27,6 +27,10 @@ const AUTH_CONTEXT_CACHE_MAX_ENTRIES: usize = 256;
pub(crate) struct GatewayControlAuthContext { pub(crate) struct GatewayControlAuthContext {
pub(crate) user_id: String, pub(crate) user_id: String,
pub(crate) api_key_id: String, pub(crate) api_key_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) username: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) api_key_name: Option<String>,
pub(crate) balance_remaining: Option<f64>, pub(crate) balance_remaining: Option<f64>,
pub(crate) access_allowed: bool, pub(crate) access_allowed: bool,
#[serde(skip)] #[serde(skip)]
@@ -489,6 +493,8 @@ pub(super) async fn resolve_data_backed_auth_context(
return Ok(Some(GatewayControlAuthContext { return Ok(Some(GatewayControlAuthContext {
user_id: String::new(), user_id: String::new(),
api_key_id: String::new(), api_key_id: String::new(),
username: None,
api_key_name: None,
balance_remaining: None, balance_remaining: None,
access_allowed: false, access_allowed: false,
user_rate_limit: None, user_rate_limit: None,
@@ -539,6 +545,8 @@ async fn resolve_trusted_auth_context(
return Ok(Some(GatewayControlAuthContext { return Ok(Some(GatewayControlAuthContext {
user_id: trusted_headers.user_id, user_id: trusted_headers.user_id,
api_key_id: trusted_headers.api_key_id, api_key_id: trusted_headers.api_key_id,
username: None,
api_key_name: None,
balance_remaining: trusted_headers.balance_remaining, balance_remaining: trusted_headers.balance_remaining,
access_allowed: false, access_allowed: false,
user_rate_limit: None, user_rate_limit: None,
@@ -620,6 +628,8 @@ fn build_data_backed_auth_context(
}; };
GatewayControlAuthContext { GatewayControlAuthContext {
username: Some(snapshot.username.clone()),
api_key_name: snapshot.api_key_name.clone(),
user_id: snapshot.user_id, user_id: snapshot.user_id,
api_key_id: snapshot.api_key_id, api_key_id: snapshot.api_key_id,
balance_remaining: wallet_remaining.or(balance_remaining), balance_remaining: wallet_remaining.or(balance_remaining),

View File

@@ -4,9 +4,10 @@ use aether_scheduler_core::{
auth_constraints_allow_api_format, build_minimal_candidate_selection, auth_constraints_allow_api_format, build_minimal_candidate_selection,
collect_global_model_names_for_required_capability, normalize_api_format, collect_global_model_names_for_required_capability, normalize_api_format,
resolve_requested_global_model_name, SchedulerAuthConstraints, resolve_requested_global_model_name, SchedulerAuthConstraints,
SchedulerMinimalCandidateSelectionCandidate, SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode,
}; };
use async_trait::async_trait; use async_trait::async_trait;
use std::collections::BTreeSet;
use super::auth::GatewayAuthApiKeySnapshot; use super::auth::GatewayAuthApiKeySnapshot;
@@ -62,6 +63,91 @@ pub(crate) async fn read_minimal_candidate_selection(
requested_model_name: &str, requested_model_name: &str,
require_streaming: bool, require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>, auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
read_minimal_candidate_selection_with_priority_mode(
state,
api_format,
requested_model_name,
require_streaming,
auth_snapshot,
SchedulerPriorityMode::Provider,
)
.await
}
pub(crate) async fn read_minimal_candidate_selection_with_priority_mode(
state: &(impl MinimalCandidateSelectionRowSource + Sync),
api_format: &str,
requested_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
priority_mode: SchedulerPriorityMode,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
read_minimal_candidate_selection_with_priority_mode_and_affinity_key(
state,
api_format,
requested_model_name,
require_streaming,
auth_snapshot,
priority_mode,
auth_snapshot_affinity_key(auth_snapshot),
)
.await
}
pub(crate) async fn read_minimal_candidate_selection_with_priority_mode_and_required_capabilities(
state: &(impl MinimalCandidateSelectionRowSource + Sync),
api_format: &str,
requested_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
priority_mode: SchedulerPriorityMode,
required_capabilities: Option<&serde_json::Value>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
read_minimal_candidate_selection_with_priority_mode_and_affinity_key_and_required_capabilities(
state,
api_format,
requested_model_name,
require_streaming,
auth_snapshot,
priority_mode,
auth_snapshot_affinity_key(auth_snapshot),
required_capabilities,
)
.await
}
pub(crate) async fn read_minimal_candidate_selection_with_priority_mode_and_affinity_key(
state: &(impl MinimalCandidateSelectionRowSource + Sync),
api_format: &str,
requested_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
priority_mode: SchedulerPriorityMode,
affinity_key: Option<&str>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
read_minimal_candidate_selection_with_priority_mode_and_affinity_key_and_required_capabilities(
state,
api_format,
requested_model_name,
require_streaming,
auth_snapshot,
priority_mode,
affinity_key,
None,
)
.await
}
pub(crate) async fn read_minimal_candidate_selection_with_priority_mode_and_affinity_key_and_required_capabilities(
state: &(impl MinimalCandidateSelectionRowSource + Sync),
api_format: &str,
requested_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
priority_mode: SchedulerPriorityMode,
affinity_key: Option<&str>,
required_capabilities: Option<&serde_json::Value>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> { ) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
let normalized_api_format = normalize_api_format(api_format); let normalized_api_format = normalize_api_format(api_format);
if normalized_api_format.is_empty() { if normalized_api_format.is_empty() {
@@ -81,17 +167,16 @@ pub(crate) async fn read_minimal_candidate_selection(
return Ok(Vec::new()); return Ok(Vec::new());
}; };
let auth_constraints = auth_snapshot.map(auth_snapshot_constraints); let auth_constraints = auth_snapshot.map(auth_snapshot_constraints);
let affinity_key = auth_snapshot
.map(|snapshot| snapshot.api_key_id.trim())
.filter(|value| !value.is_empty());
build_minimal_candidate_selection( build_minimal_candidate_selection(
rows, rows,
&normalized_api_format, &normalized_api_format,
requested_model_name, requested_model_name,
resolved_global_model_name.as_str(), resolved_global_model_name.as_str(),
require_streaming, require_streaming,
required_capabilities,
auth_constraints.as_ref(), auth_constraints.as_ref(),
affinity_key, affinity_key,
priority_mode,
) )
} }
@@ -128,6 +213,60 @@ pub(crate) async fn read_global_model_names_for_required_capability(
)) ))
} }
pub(crate) async fn read_global_model_names_for_api_format(
state: &(impl MinimalCandidateSelectionRowSource + Sync),
api_format: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
) -> Result<Vec<String>, DataLayerError> {
let normalized_api_format = normalize_api_format(api_format);
if normalized_api_format.is_empty() {
return Ok(Vec::new());
}
if !auth_constraints_allow_api_format(
auth_snapshot.map(auth_snapshot_constraints).as_ref(),
&normalized_api_format,
) {
return Ok(Vec::new());
}
let rows = state
.read_minimal_candidate_selection_rows_for_api_format(&normalized_api_format)
.await?;
let auth_constraints = auth_snapshot.map(auth_snapshot_constraints);
let mut model_names = BTreeSet::new();
for row in rows {
if require_streaming && !row.supports_streaming() {
continue;
}
if !aether_scheduler_core::auth_constraints_allow_provider(
auth_constraints.as_ref(),
&row.provider_id,
&row.provider_name,
) {
continue;
}
if !aether_scheduler_core::auth_constraints_allow_model(
auth_constraints.as_ref(),
&row.global_model_name,
&row.global_model_name,
) {
continue;
}
model_names.insert(row.global_model_name);
}
Ok(model_names.into_iter().collect())
}
fn auth_snapshot_affinity_key(auth_snapshot: Option<&GatewayAuthApiKeySnapshot>) -> Option<&str> {
auth_snapshot
.map(|snapshot| snapshot.api_key_id.trim())
.filter(|value| !value.is_empty())
}
fn auth_snapshot_constraints(snapshot: &GatewayAuthApiKeySnapshot) -> SchedulerAuthConstraints { fn auth_snapshot_constraints(snapshot: &GatewayAuthApiKeySnapshot) -> SchedulerAuthConstraints {
SchedulerAuthConstraints { SchedulerAuthConstraints {
allowed_providers: snapshot allowed_providers: snapshot

View File

@@ -36,7 +36,7 @@ mod tests {
request_id: &str, request_id: &str,
candidate_index: i32, candidate_index: i32,
status: RequestCandidateStatus, status: RequestCandidateStatus,
started_at_unix_secs: Option<i64>, started_at_unix_ms: Option<i64>,
latency_ms: Option<i32>, latency_ms: Option<i32>,
status_code: Option<i32>, status_code: Option<i32>,
) -> StoredRequestCandidate { ) -> StoredRequestCandidate {
@@ -62,9 +62,9 @@ mod tests {
Some(1), Some(1),
None, None,
None, None,
100 + i64::from(candidate_index), (100 + i64::from(candidate_index)) * 1_000,
started_at_unix_secs, started_at_unix_ms.map(|v| v * 1_000),
started_at_unix_secs.map(|value| value + 1), started_at_unix_ms.map(|value| (value + 1) * 1_000),
) )
.expect("candidate should build") .expect("candidate should build")
} }

View File

@@ -101,9 +101,9 @@ mod tests {
Some(1), Some(1),
None, None,
Some(serde_json::json!({"cache_1h": true})), Some(serde_json::json!({"cache_1h": true})),
100, 100_000,
Some(101), Some(101_000),
Some(102), Some(102_000),
) )
.expect("candidate should build") .expect("candidate should build")
} }
@@ -174,11 +174,15 @@ mod tests {
provider_name: Some("OpenAI".to_string()), provider_name: Some("OpenAI".to_string()),
provider_website: Some("https://openai.com".to_string()), provider_website: Some("https://openai.com".to_string()),
provider_type: Some("custom".to_string()), provider_type: Some("custom".to_string()),
provider_priority: Some(0),
provider_keep_priority_on_conversion: Some(false),
endpoint_api_format: Some("openai:chat".to_string()), endpoint_api_format: Some("openai:chat".to_string()),
endpoint_api_family: Some("openai".to_string()), endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("chat".to_string()), endpoint_kind: Some("chat".to_string()),
provider_key_name: Some("prod-key".to_string()), provider_key_name: Some("prod-key".to_string()),
provider_key_auth_type: Some("api_key".to_string()), provider_key_auth_type: Some("api_key".to_string()),
provider_key_internal_priority: Some(50),
provider_key_global_priority_by_format: None,
provider_key_capabilities: Some(serde_json::json!({"cache_1h": true})), provider_key_capabilities: Some(serde_json::json!({"cache_1h": true})),
provider_key_is_active: Some(true), provider_key_is_active: Some(true),
}], }],

View File

@@ -28,7 +28,9 @@ use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
use aether_data_contracts::repository::video_tasks::{ use aether_data_contracts::repository::video_tasks::{
UpsertVideoTask, VideoTaskLookupKey, VideoTaskStatus, VideoTaskWriteRepository, UpsertVideoTask, VideoTaskLookupKey, VideoTaskStatus, VideoTaskWriteRepository,
}; };
use aether_scheduler_core::{build_minimal_candidate_selection, SchedulerAuthConstraints}; use aether_scheduler_core::{
build_minimal_candidate_selection, SchedulerAuthConstraints, SchedulerPriorityMode,
};
use async_trait::async_trait; use async_trait::async_trait;
use serde_json::json; use serde_json::json;
@@ -108,7 +110,7 @@ async fn data_state_find_uses_configured_read_repository() {
next_poll_at_unix_secs: Some(100), next_poll_at_unix_secs: Some(100),
poll_count: 0, poll_count: 0,
max_poll_count: 360, max_poll_count: 360,
created_at_unix_secs: 100, created_at_unix_ms: 100,
submitted_at_unix_secs: Some(100), submitted_at_unix_secs: Some(100),
completed_at_unix_secs: None, completed_at_unix_secs: None,
updated_at_unix_secs: 100, updated_at_unix_secs: 100,
@@ -384,9 +386,9 @@ async fn data_state_reads_decision_trace_with_provider_catalog_metadata() {
Some(1), Some(1),
None, None,
Some(serde_json::json!({"cache_1h": true})), Some(serde_json::json!({"cache_1h": true})),
100, 100_000,
Some(101), Some(101_000),
Some(102), Some(102_000),
) )
.expect("candidate should build"), .expect("candidate should build"),
])); ]));
@@ -472,9 +474,9 @@ async fn data_state_reads_request_audit_bundle_from_multiple_readers() {
Some(1), Some(1),
None, None,
Some(serde_json::json!({"cache_1h": true})), Some(serde_json::json!({"cache_1h": true})),
100, 100_000,
Some(101), Some(101_000),
Some(102), Some(102_000),
) )
.expect("candidate should build"), .expect("candidate should build"),
])); ]));
@@ -668,8 +670,10 @@ async fn data_state_reads_minimal_candidate_selection_with_auth_filters() {
"gpt-4.1", "gpt-4.1",
"gpt-4.1", "gpt-4.1",
false, false,
None,
Some(&auth_constraints), Some(&auth_constraints),
Some(auth_snapshot.api_key_id.as_str()), Some(auth_snapshot.api_key_id.as_str()),
SchedulerPriorityMode::Provider,
) )
.expect("selection should read"); .expect("selection should read");
@@ -722,7 +726,7 @@ async fn maps_openai_video_task_repository_row_into_read_response() {
next_poll_at_unix_secs: Some(120), next_poll_at_unix_secs: Some(120),
poll_count: 1, poll_count: 1,
max_poll_count: 360, max_poll_count: 360,
created_at_unix_secs: 100, created_at_unix_ms: 100,
submitted_at_unix_secs: Some(100), submitted_at_unix_secs: Some(100),
completed_at_unix_secs: None, completed_at_unix_secs: None,
updated_at_unix_secs: 120, updated_at_unix_secs: 120,
@@ -781,7 +785,7 @@ async fn maps_gemini_video_task_repository_row_into_read_response() {
next_poll_at_unix_secs: None, next_poll_at_unix_secs: None,
poll_count: 4, poll_count: 4,
max_poll_count: 360, max_poll_count: 360,
created_at_unix_secs: 100, created_at_unix_ms: 100,
submitted_at_unix_secs: Some(100), submitted_at_unix_secs: Some(100),
completed_at_unix_secs: Some(120), completed_at_unix_secs: Some(120),
updated_at_unix_secs: 120, updated_at_unix_secs: 120,
@@ -843,7 +847,7 @@ async fn data_state_write_uses_configured_shadow_result_writer() {
match_status: ShadowResultMatchStatus::Pending, match_status: ShadowResultMatchStatus::Pending,
status_code: Some(200), status_code: Some(200),
error_message: None, error_message: None,
created_at_unix_secs: 100, created_at_unix_ms: 100,
updated_at_unix_secs: 100, updated_at_unix_secs: 100,
}) })
.await .await
@@ -906,7 +910,7 @@ async fn data_state_records_shadow_result_samples_and_merges_match_status() {
.expect("second stored result should exist"); .expect("second stored result should exist");
assert_eq!(second.match_status, ShadowResultMatchStatus::Match); assert_eq!(second.match_status, ShadowResultMatchStatus::Match);
assert_eq!(second.created_at_unix_secs, 100); assert_eq!(second.created_at_unix_ms, 100);
assert_eq!(second.updated_at_unix_secs, 200); assert_eq!(second.updated_at_unix_secs, 200);
assert_eq!(second.request_id.as_deref(), Some("req-1")); assert_eq!(second.request_id.as_deref(), Some("req-1"));
} }
@@ -986,7 +990,7 @@ fn sample_request_candidate(
request_id: &str, request_id: &str,
candidate_index: i32, candidate_index: i32,
status: RequestCandidateStatus, status: RequestCandidateStatus,
started_at_unix_secs: Option<i64>, started_at_unix_ms: Option<i64>,
latency_ms: Option<i32>, latency_ms: Option<i32>,
status_code: Option<i32>, status_code: Option<i32>,
) -> StoredRequestCandidate { ) -> StoredRequestCandidate {
@@ -1013,8 +1017,8 @@ fn sample_request_candidate(
None, None,
None, None,
100 + i64::from(candidate_index), 100 + i64::from(candidate_index),
started_at_unix_secs, started_at_unix_ms,
started_at_unix_secs.map(|value| value + 1), started_at_unix_ms.map(|value| value + 1),
) )
.expect("candidate should build") .expect("candidate should build")
} }

View File

@@ -1,25 +1,83 @@
use aether_contracts::ExecutionResult; use std::collections::BTreeSet;
fn is_local_candidate_attempt(report_context: Option<&serde_json::Value>) -> bool { use aether_contracts::{ExecutionPlan, ExecutionResult};
use regex::Regex;
use crate::AppState;
fn local_candidate_index(report_context: Option<&serde_json::Value>) -> Option<u64> {
report_context report_context
.and_then(serde_json::Value::as_object) .and_then(serde_json::Value::as_object)
.and_then(|context| context.get("candidate_index")) .and_then(|context| context.get("candidate_index"))
.and_then(serde_json::Value::as_u64) .and_then(serde_json::Value::as_u64)
.is_some()
} }
fn is_retryable_local_upstream_status(status_code: u16) -> bool { fn is_retryable_local_upstream_status(status_code: u16) -> bool {
status_code == 429 || status_code >= 500 status_code == 429 || status_code >= 500
} }
pub(crate) fn should_retry_next_local_candidate_sync( #[derive(Debug, Clone, Default, PartialEq, Eq)]
plan_kind: &str, struct LocalFailoverPolicy {
max_retries: Option<u64>,
stop_status_codes: BTreeSet<u16>,
continue_status_codes: BTreeSet<u16>,
success_failover_patterns: Vec<LocalFailoverRegexRule>,
error_stop_patterns: Vec<LocalFailoverRegexRule>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct LocalFailoverRegexRule {
pattern: String,
status_codes: BTreeSet<u16>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum LocalFailoverDecision {
UseDefault,
RetryNextCandidate,
StopLocalFailover,
}
pub(crate) async fn should_retry_next_local_candidate_sync(
state: &AppState,
plan: &ExecutionPlan,
_plan_kind: &str,
report_context: Option<&serde_json::Value>, report_context: Option<&serde_json::Value>,
result: &ExecutionResult, result: &ExecutionResult,
response_text: Option<&str>,
) -> bool { ) -> bool {
is_local_candidate_attempt(report_context) matches!(
&& plan_kind == "openai_chat_sync" resolve_local_failover_decision(
&& is_retryable_local_upstream_status(result.status_code) state,
plan,
report_context,
result.status_code,
response_text,
)
.await,
LocalFailoverDecision::RetryNextCandidate
)
}
pub(crate) async fn should_stop_local_candidate_failover_sync(
state: &AppState,
plan: &ExecutionPlan,
_plan_kind: &str,
report_context: Option<&serde_json::Value>,
result: &ExecutionResult,
response_text: Option<&str>,
) -> bool {
matches!(
resolve_local_failover_decision(
state,
plan,
report_context,
result.status_code,
response_text,
)
.await,
LocalFailoverDecision::StopLocalFailover
)
} }
pub(crate) fn should_fallback_to_control_sync( pub(crate) fn should_fallback_to_control_sync(
@@ -102,14 +160,245 @@ pub(crate) fn resolve_core_sync_error_finalize_report_kind(
Some(report_kind.to_string()) Some(report_kind.to_string())
} }
pub(crate) fn should_retry_next_local_candidate_stream( pub(crate) async fn should_retry_next_local_candidate_stream(
plan_kind: &str, state: &AppState,
plan: &ExecutionPlan,
_plan_kind: &str,
report_context: Option<&serde_json::Value>, report_context: Option<&serde_json::Value>,
status_code: u16, status_code: u16,
response_text: Option<&str>,
) -> bool { ) -> bool {
is_local_candidate_attempt(report_context) matches!(
&& plan_kind == "openai_chat_stream" resolve_local_failover_decision(state, plan, report_context, status_code, response_text)
&& is_retryable_local_upstream_status(status_code) .await,
LocalFailoverDecision::RetryNextCandidate
)
}
pub(crate) async fn should_stop_local_candidate_failover_stream(
state: &AppState,
plan: &ExecutionPlan,
_plan_kind: &str,
report_context: Option<&serde_json::Value>,
status_code: u16,
response_text: Option<&str>,
) -> bool {
matches!(
resolve_local_failover_decision(state, plan, report_context, status_code, response_text)
.await,
LocalFailoverDecision::StopLocalFailover
)
}
pub(crate) fn local_failover_response_text(
body_json: Option<&serde_json::Value>,
body_bytes: &[u8],
fallback_text: Option<&str>,
) -> Option<String> {
if let Some(body_json) = body_json {
return serde_json::to_string(body_json).ok();
}
if !body_bytes.is_empty() {
return Some(String::from_utf8_lossy(body_bytes).into_owned());
}
fallback_text
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
async fn resolve_local_failover_decision(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
status_code: u16,
response_text: Option<&str>,
) -> LocalFailoverDecision {
let Some(candidate_index) = local_candidate_index(report_context) else {
return LocalFailoverDecision::UseDefault;
};
let policy = resolve_local_failover_policy(state, plan).await;
let response_text = response_text
.map(str::trim)
.filter(|value| !value.is_empty());
if policy.stop_status_codes.contains(&status_code) {
return LocalFailoverDecision::StopLocalFailover;
}
if status_code >= 400
&& response_text.is_some_and(|text| {
policy
.error_stop_patterns
.iter()
.any(|rule| local_failover_regex_rule_matches(rule, text, status_code))
})
{
return LocalFailoverDecision::StopLocalFailover;
}
if policy
.max_retries
.is_some_and(|max_retries| candidate_index >= max_retries)
{
return LocalFailoverDecision::UseDefault;
}
if status_code == 200
&& response_text.is_some_and(|text| {
policy
.success_failover_patterns
.iter()
.any(|rule| local_failover_regex_rule_matches(rule, text, status_code))
})
{
return LocalFailoverDecision::RetryNextCandidate;
}
if policy.continue_status_codes.contains(&status_code) {
return LocalFailoverDecision::RetryNextCandidate;
}
if is_retryable_local_upstream_status(status_code) {
return LocalFailoverDecision::RetryNextCandidate;
}
LocalFailoverDecision::UseDefault
}
async fn resolve_local_failover_policy(
state: &AppState,
plan: &ExecutionPlan,
) -> LocalFailoverPolicy {
let transport = match state
.read_provider_transport_snapshot(&plan.provider_id, &plan.endpoint_id, &plan.key_id)
.await
{
Ok(Some(transport)) => transport,
Ok(None) | Err(_) => return LocalFailoverPolicy::default(),
};
let rules = transport
.provider
.config
.as_ref()
.and_then(|config| config.get("failover_rules"))
.and_then(serde_json::Value::as_object);
let max_retries = rules
.and_then(|value| value.get("max_retries"))
.and_then(parse_u64_value)
.or_else(|| {
transport
.endpoint
.max_retries
.and_then(|value| u64::try_from(value).ok())
})
.or_else(|| {
transport
.provider
.max_retries
.and_then(|value| u64::try_from(value).ok())
});
LocalFailoverPolicy {
max_retries,
stop_status_codes: rules
.map(|value| {
parse_status_code_set(
value,
&[
"stop_on_status_codes",
"early_stop_status_codes",
"non_retryable_status_codes",
"stop_status_codes",
],
)
})
.unwrap_or_default(),
continue_status_codes: rules
.map(|value| {
parse_status_code_set(
value,
&[
"continue_on_status_codes",
"retryable_status_codes",
"retry_on_status_codes",
"continue_status_codes",
],
)
})
.unwrap_or_default(),
success_failover_patterns: rules
.map(|value| parse_regex_rules(value, "success_failover_patterns"))
.unwrap_or_default(),
error_stop_patterns: rules
.map(|value| parse_regex_rules(value, "error_stop_patterns"))
.unwrap_or_default(),
}
}
fn parse_regex_rules(
rules: &serde_json::Map<String, serde_json::Value>,
key: &str,
) -> Vec<LocalFailoverRegexRule> {
rules
.get(key)
.and_then(serde_json::Value::as_array)
.into_iter()
.flat_map(|items| items.iter())
.filter_map(parse_regex_rule)
.collect()
}
fn parse_regex_rule(value: &serde_json::Value) -> Option<LocalFailoverRegexRule> {
let object = value.as_object()?;
let pattern = object
.get("pattern")
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())?;
Some(LocalFailoverRegexRule {
pattern: pattern.to_string(),
status_codes: object
.get("status_codes")
.and_then(serde_json::Value::as_array)
.into_iter()
.flat_map(|values| values.iter())
.filter_map(|value| parse_u64_value(value).and_then(|value| u16::try_from(value).ok()))
.collect(),
})
}
fn local_failover_regex_rule_matches(
rule: &LocalFailoverRegexRule,
response_text: &str,
status_code: u16,
) -> bool {
if !rule.status_codes.is_empty() && !rule.status_codes.contains(&status_code) {
return false;
}
Regex::new(&rule.pattern)
.ok()
.is_some_and(|regex| regex.is_match(response_text))
}
fn parse_status_code_set(
rules: &serde_json::Map<String, serde_json::Value>,
keys: &[&str],
) -> BTreeSet<u16> {
keys.iter()
.filter_map(|key| rules.get(*key))
.filter_map(serde_json::Value::as_array)
.flat_map(|values| values.iter())
.filter_map(|value| parse_u64_value(value).and_then(|value| u16::try_from(value).ok()))
.collect()
}
fn parse_u64_value(value: &serde_json::Value) -> Option<u64> {
value
.as_u64()
.or_else(|| value.as_i64().and_then(|value| u64::try_from(value).ok()))
} }
pub(crate) fn should_fallback_to_control_stream( pub(crate) fn should_fallback_to_control_stream(
@@ -172,14 +461,121 @@ pub(crate) fn resolve_core_stream_direct_finalize_report_kind(plan_kind: &str) -
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use std::collections::BTreeSet;
use aether_contracts::ExecutionResult; use aether_contracts::ExecutionResult;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use super::{ use super::{
resolve_core_stream_error_finalize_report_kind, resolve_core_stream_error_finalize_report_kind,
resolve_core_sync_error_finalize_report_kind, should_fallback_to_control_stream, resolve_core_sync_error_finalize_report_kind, resolve_local_failover_policy,
should_fallback_to_control_sync, should_retry_next_local_candidate_stream, should_fallback_to_control_stream, should_fallback_to_control_sync,
should_retry_next_local_candidate_sync, should_retry_next_local_candidate_stream, should_retry_next_local_candidate_sync,
should_stop_local_candidate_failover_stream, should_stop_local_candidate_failover_sync,
LocalFailoverPolicy, LocalFailoverRegexRule,
}; };
use crate::data::GatewayDataState;
use crate::AppState;
fn sample_plan() -> aether_contracts::ExecutionPlan {
aether_contracts::ExecutionPlan {
request_id: "req-1".to_string(),
candidate_id: Some("cand-1".to_string()),
provider_name: Some("provider-1".to_string()),
provider_id: "provider-1".to_string(),
endpoint_id: "endpoint-1".to_string(),
key_id: "key-1".to_string(),
method: "POST".to_string(),
url: "https://example.com/v1/chat/completions".to_string(),
headers: Default::default(),
content_type: Some("application/json".to_string()),
content_encoding: None,
body: aether_contracts::RequestBody::from_json(serde_json::json!({"model":"gpt-5"})),
stream: false,
client_api_format: "openai:chat".to_string(),
provider_api_format: "openai:chat".to_string(),
model_name: Some("gpt-5".to_string()),
proxy: None,
tls_profile: None,
timeouts: None,
}
}
fn sample_provider(config: Option<serde_json::Value>) -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-1".to_string(),
"provider-1".to_string(),
Some("https://provider.example".to_string()),
"custom".to_string(),
)
.expect("provider should build")
.with_transport_fields(true, false, false, None, Some(3), None, None, None, config)
}
fn sample_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-1".to_string(),
"provider-1".to_string(),
"openai:chat".to_string(),
Some("openai".to_string()),
Some("chat".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://api.provider.example".to_string(),
None,
None,
Some(2),
None,
None,
None,
None,
)
.expect("endpoint transport should build")
}
fn sample_key() -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
"key-1".to_string(),
"provider-1".to_string(),
"key-1".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["openai:chat"])),
"plain-upstream-key".to_string(),
None,
None,
Some(serde_json::json!({"openai:chat": 1})),
None,
None,
None,
None,
)
.expect("key transport should build")
}
fn build_state_with_provider_config(config: Option<serde_json::Value>) -> AppState {
let provider_catalog = InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider(config)],
vec![sample_endpoint()],
vec![sample_key()],
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
);
AppState::new()
.expect("state should build")
.with_data_state_for_tests(data_state)
}
#[test] #[test]
fn sync_failover_marks_chat_errors() { fn sync_failover_marks_chat_errors() {
@@ -220,8 +616,8 @@ mod tests {
); );
} }
#[test] #[tokio::test]
fn sync_retry_next_candidate_is_local_openai_chat_only() { async fn sync_retry_next_candidate_requires_local_candidate_context() {
let result = ExecutionResult { let result = ExecutionResult {
request_id: "req-1".to_string(), request_id: "req-1".to_string(),
candidate_id: None, candidate_id: None,
@@ -235,26 +631,57 @@ mod tests {
"candidate_index": 0, "candidate_index": 0,
"retry_index": 0, "retry_index": 0,
}); });
let state = build_state_with_provider_config(None);
let plan = sample_plan();
assert!(should_retry_next_local_candidate_sync( assert!(
"openai_chat_sync", should_retry_next_local_candidate_sync(
Some(&local_report_context), &state,
&result, &plan,
)); "openai_chat_sync",
assert!(!should_retry_next_local_candidate_sync( Some(&local_report_context),
"openai_chat_sync", &result,
None, None,
&result, )
)); .await
assert!(!should_retry_next_local_candidate_sync( );
"claude_chat_sync", assert!(
None, should_retry_next_local_candidate_sync(
&result, &state,
)); &plan,
"claude_cli_sync",
Some(&local_report_context),
&result,
None,
)
.await
);
assert!(
!should_retry_next_local_candidate_sync(
&state,
&plan,
"openai_chat_sync",
None,
&result,
None,
)
.await
);
assert!(
!should_retry_next_local_candidate_sync(
&state,
&plan,
"claude_chat_sync",
None,
&result,
None,
)
.await
);
} }
#[test] #[tokio::test]
fn sync_retry_next_candidate_treats_rate_limit_as_retryable() { async fn sync_retry_next_candidate_treats_rate_limit_as_retryable() {
let result = ExecutionResult { let result = ExecutionResult {
request_id: "req-1".to_string(), request_id: "req-1".to_string(),
candidate_id: None, candidate_id: None,
@@ -268,49 +695,304 @@ mod tests {
"candidate_index": 0, "candidate_index": 0,
"retry_index": 0, "retry_index": 0,
}); });
let state = build_state_with_provider_config(None);
let plan = sample_plan();
assert!(should_retry_next_local_candidate_sync( assert!(
"openai_chat_sync", should_retry_next_local_candidate_sync(
Some(&local_report_context), &state,
&result, &plan,
)); "openai_chat_sync",
Some(&local_report_context),
&result,
None,
)
.await
);
} }
#[test] #[tokio::test]
fn stream_retry_next_candidate_is_local_openai_chat_only() { async fn stream_retry_next_candidate_requires_local_candidate_context() {
let local_report_context = serde_json::json!({ let local_report_context = serde_json::json!({
"candidate_index": 0, "candidate_index": 0,
"retry_index": 0, "retry_index": 0,
}); });
let state = build_state_with_provider_config(None);
let plan = sample_plan();
assert!(should_retry_next_local_candidate_stream( assert!(
"openai_chat_stream", should_retry_next_local_candidate_stream(
Some(&local_report_context), &state,
502, &plan,
)); "openai_chat_stream",
assert!(!should_retry_next_local_candidate_stream( Some(&local_report_context),
"openai_chat_stream", 502,
None, None,
502, )
)); .await
assert!(!should_retry_next_local_candidate_stream( );
"claude_chat_stream", assert!(
Some(&local_report_context), should_retry_next_local_candidate_stream(
502, &state,
)); &plan,
"gemini_cli_stream",
Some(&local_report_context),
502,
None,
)
.await
);
assert!(
!should_retry_next_local_candidate_stream(
&state,
&plan,
"openai_chat_stream",
None,
502,
None,
)
.await
);
assert!(
!should_retry_next_local_candidate_stream(
&state,
&plan,
"claude_chat_stream",
None,
502,
None,
)
.await
);
} }
#[test] #[tokio::test]
fn stream_retry_next_candidate_treats_rate_limit_as_retryable() { async fn stream_retry_next_candidate_treats_rate_limit_as_retryable() {
let local_report_context = serde_json::json!({ let local_report_context = serde_json::json!({
"candidate_index": 0, "candidate_index": 0,
"retry_index": 0, "retry_index": 0,
}); });
let state = build_state_with_provider_config(None);
let plan = sample_plan();
assert!(should_retry_next_local_candidate_stream( assert!(
"openai_chat_stream", should_retry_next_local_candidate_stream(
Some(&local_report_context), &state,
429, &plan,
)); "openai_chat_stream",
Some(&local_report_context),
429,
None,
)
.await
);
}
#[test]
fn resolve_local_failover_policy_reads_provider_rules() {
let state = build_state_with_provider_config(Some(serde_json::json!({
"failover_rules": {
"max_retries": 1,
"stop_on_status_codes": [503],
"continue_on_status_codes": [409, 429]
}
})));
let plan = sample_plan();
let runtime = tokio::runtime::Runtime::new().expect("runtime should build");
let policy = runtime.block_on(resolve_local_failover_policy(&state, &plan));
assert_eq!(
policy,
LocalFailoverPolicy {
max_retries: Some(1),
stop_status_codes: [503].into_iter().collect(),
continue_status_codes: [409, 429].into_iter().collect(),
success_failover_patterns: Vec::new(),
error_stop_patterns: Vec::new(),
}
);
}
#[tokio::test]
async fn local_failover_policy_can_stop_retryable_statuses_and_continue_non_retryable_statuses()
{
let state = build_state_with_provider_config(Some(serde_json::json!({
"failover_rules": {
"max_retries": 2,
"stop_on_status_codes": [503],
"continue_on_status_codes": [409]
}
})));
let plan = sample_plan();
let first_candidate = serde_json::json!({
"candidate_index": 0,
"retry_index": 0,
});
let third_candidate = serde_json::json!({
"candidate_index": 2,
"retry_index": 0,
});
assert!(
!should_retry_next_local_candidate_stream(
&state,
&plan,
"openai_chat_stream",
Some(&first_candidate),
503,
None,
)
.await
);
assert!(
should_stop_local_candidate_failover_stream(
&state,
&plan,
"openai_chat_stream",
Some(&first_candidate),
503,
None,
)
.await
);
assert!(
should_retry_next_local_candidate_stream(
&state,
&plan,
"openai_chat_stream",
Some(&first_candidate),
409,
None,
)
.await
);
assert!(
!should_retry_next_local_candidate_stream(
&state,
&plan,
"openai_chat_stream",
Some(&third_candidate),
429,
None,
)
.await
);
}
#[test]
fn resolve_local_failover_policy_reads_regex_rules() {
let state = build_state_with_provider_config(Some(serde_json::json!({
"failover_rules": {
"success_failover_patterns": [
{"pattern": "relay:.*格式错误"}
],
"error_stop_patterns": [
{"pattern": "content_policy_violation", "status_codes": [400, 403]}
]
}
})));
let plan = sample_plan();
let runtime = tokio::runtime::Runtime::new().expect("runtime should build");
let policy = runtime.block_on(resolve_local_failover_policy(&state, &plan));
assert_eq!(
policy.success_failover_patterns,
vec![LocalFailoverRegexRule {
pattern: "relay:.*格式错误".to_string(),
status_codes: BTreeSet::new(),
}]
);
assert_eq!(
policy.error_stop_patterns,
vec![LocalFailoverRegexRule {
pattern: "content_policy_violation".to_string(),
status_codes: [400, 403].into_iter().collect(),
}]
);
}
#[tokio::test]
async fn success_failover_pattern_can_retry_sync_candidate() {
let result = ExecutionResult {
request_id: "req-1".to_string(),
candidate_id: None,
status_code: 200,
headers: Default::default(),
body: None,
telemetry: None,
error: None,
};
let local_report_context = serde_json::json!({
"candidate_index": 0,
"retry_index": 0,
});
let state = build_state_with_provider_config(Some(serde_json::json!({
"failover_rules": {
"success_failover_patterns": [
{"pattern": "relay:.*格式错误"}
]
}
})));
let plan = sample_plan();
assert!(
should_retry_next_local_candidate_sync(
&state,
&plan,
"openai_chat_sync",
Some(&local_report_context),
&result,
Some("{\"error\":\"relay: 返回格式错误\"}"),
)
.await
);
}
#[tokio::test]
async fn error_stop_pattern_can_stop_sync_failover() {
let result = ExecutionResult {
request_id: "req-1".to_string(),
candidate_id: None,
status_code: 400,
headers: Default::default(),
body: None,
telemetry: None,
error: None,
};
let local_report_context = serde_json::json!({
"candidate_index": 0,
"retry_index": 0,
});
let state = build_state_with_provider_config(Some(serde_json::json!({
"failover_rules": {
"error_stop_patterns": [
{"pattern": "content_policy_violation", "status_codes": [400]}
]
}
})));
let plan = sample_plan();
assert!(
should_stop_local_candidate_failover_sync(
&state,
&plan,
"openai_chat_sync",
Some(&local_report_context),
&result,
Some("{\"error\":\"content_policy_violation\"}"),
)
.await
);
assert!(
!should_retry_next_local_candidate_sync(
&state,
&plan,
"openai_chat_sync",
Some(&local_report_context),
&result,
Some("{\"error\":\"content_policy_violation\"}"),
)
.await
);
} }
} }

View File

@@ -19,11 +19,12 @@ pub(crate) use self::constants::{
MAX_ERROR_BODY_BYTES, MAX_STREAM_PREFETCH_BYTES, MAX_STREAM_PREFETCH_FRAMES, MAX_ERROR_BODY_BYTES, MAX_STREAM_PREFETCH_BYTES, MAX_STREAM_PREFETCH_FRAMES,
}; };
pub(crate) use self::fallback::{ pub(crate) use self::fallback::{
resolve_core_stream_direct_finalize_report_kind, local_failover_response_text, resolve_core_stream_direct_finalize_report_kind,
resolve_core_stream_error_finalize_report_kind, resolve_core_sync_error_finalize_report_kind, resolve_core_stream_error_finalize_report_kind, resolve_core_sync_error_finalize_report_kind,
should_fallback_to_control_stream, should_fallback_to_control_sync, should_fallback_to_control_stream, should_fallback_to_control_sync,
should_finalize_sync_response, should_retry_next_local_candidate_stream, should_finalize_sync_response, should_retry_next_local_candidate_stream,
should_retry_next_local_candidate_sync, should_retry_next_local_candidate_sync, should_stop_local_candidate_failover_stream,
should_stop_local_candidate_failover_sync,
}; };
pub use server::{ pub use server::{
build_execution_runtime_router, build_execution_runtime_router_with_request_concurrency_limit, build_execution_runtime_router, build_execution_runtime_router_with_request_concurrency_limit,

View File

@@ -1,6 +1,7 @@
use std::collections::VecDeque;
use std::io::Error as IoError; use std::io::Error as IoError;
use aether_contracts::{ExecutionPlan, ExecutionTelemetry, StreamFramePayload}; use aether_contracts::{ExecutionPlan, ExecutionTelemetry, StreamFrame, StreamFramePayload};
use aether_data_contracts::repository::candidates::RequestCandidateStatus; use aether_data_contracts::repository::candidates::RequestCandidateStatus;
use aether_scheduler_core::SchedulerRequestCandidateStatusUpdate; use aether_scheduler_core::SchedulerRequestCandidateStatusUpdate;
use async_stream::stream; use async_stream::stream;
@@ -32,7 +33,7 @@ use crate::ai_pipeline_api::{
use crate::api::response::{ use crate::api::response::{
attach_control_metadata_headers, build_client_response, build_client_response_from_parts, attach_control_metadata_headers, build_client_response, build_client_response_from_parts,
}; };
use crate::clock::current_unix_secs as current_request_candidate_unix_secs; use crate::clock::current_unix_ms as current_request_candidate_unix_ms;
use crate::constants::{CONTROL_CANDIDATE_ID_HEADER, CONTROL_REQUEST_ID_HEADER}; use crate::constants::{CONTROL_CANDIDATE_ID_HEADER, CONTROL_REQUEST_ID_HEADER};
use crate::control::GatewayControlDecision; use crate::control::GatewayControlDecision;
use crate::execution_runtime::build_direct_execution_frame_stream; use crate::execution_runtime::build_direct_execution_frame_stream;
@@ -45,9 +46,9 @@ use crate::execution_runtime::transport::{
DirectSyncExecutionRuntime, DirectUpstreamStreamExecution, DirectSyncExecutionRuntime, DirectUpstreamStreamExecution,
}; };
use crate::execution_runtime::{ use crate::execution_runtime::{
resolve_core_stream_direct_finalize_report_kind, local_failover_response_text, resolve_core_stream_direct_finalize_report_kind,
resolve_core_stream_error_finalize_report_kind, should_fallback_to_control_stream, resolve_core_stream_error_finalize_report_kind, should_fallback_to_control_stream,
should_retry_next_local_candidate_stream, should_retry_next_local_candidate_stream, should_stop_local_candidate_failover_stream,
}; };
use crate::execution_runtime::{MAX_STREAM_PREFETCH_BYTES, MAX_STREAM_PREFETCH_FRAMES}; use crate::execution_runtime::{MAX_STREAM_PREFETCH_BYTES, MAX_STREAM_PREFETCH_FRAMES};
use crate::log_ids::short_request_id; use crate::log_ids::short_request_id;
@@ -165,7 +166,7 @@ pub(crate) async fn execute_execution_runtime_stream(
}; };
if response.status() != http::StatusCode::OK { if response.status() != http::StatusCode::OK {
let terminal_unix_secs = current_request_candidate_unix_secs(); let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status( record_local_request_candidate_status(
state, state,
&plan, &plan,
@@ -179,8 +180,8 @@ pub(crate) async fn execute_execution_runtime_stream(
response.status() response.status()
)), )),
latency_ms: None, latency_ms: None,
started_at_unix_secs: Some(terminal_unix_secs), started_at_unix_ms: Some(terminal_unix_secs),
finished_at_unix_secs: Some(terminal_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs),
}, },
) )
.await; .await;
@@ -209,6 +210,60 @@ pub(crate) async fn execute_execution_runtime_stream(
} }
} }
fn decode_stream_data_chunk(
chunk_b64: Option<&str>,
text: Option<&str>,
) -> Result<Vec<u8>, GatewayError> {
if let Some(chunk_b64) = chunk_b64 {
return base64::engine::general_purpose::STANDARD
.decode(chunk_b64)
.map_err(|err| GatewayError::Internal(err.to_string()));
}
Ok(text.unwrap_or_default().as_bytes().to_vec())
}
async fn next_stream_frame<R>(
buffered_frames: &mut VecDeque<StreamFrame>,
lines: &mut FramedRead<R, LinesCodec>,
) -> Result<Option<StreamFrame>, GatewayError>
where
R: tokio::io::AsyncRead + Unpin,
{
if let Some(frame) = buffered_frames.pop_front() {
return Ok(Some(frame));
}
read_next_frame(lines).await
}
async fn probe_local_stream_success_failover_text<R>(
buffered_frames: &mut VecDeque<StreamFrame>,
lines: &mut FramedRead<R, LinesCodec>,
) -> Result<Option<String>, GatewayError>
where
R: tokio::io::AsyncRead + Unpin,
{
while let Some(frame) = read_next_frame(lines).await? {
let probe_text = match &frame.payload {
StreamFramePayload::Data { chunk_b64, text } => {
match decode_stream_data_chunk(chunk_b64.as_deref(), text.as_deref()) {
Ok(chunk) if !chunk.is_empty() => {
Some(String::from_utf8_lossy(&chunk).into_owned())
}
Ok(_) | Err(_) => None,
}
}
StreamFramePayload::Error { .. } | StreamFramePayload::Eof { .. } => None,
StreamFramePayload::Headers { .. } | StreamFramePayload::Telemetry { .. } => None,
};
buffered_frames.push_back(frame);
if probe_text.is_some() {
return Ok(probe_text);
}
}
Ok(None)
}
async fn execute_stream_from_frame_stream( async fn execute_stream_from_frame_stream(
state: &AppState, state: &AppState,
plan: ExecutionPlan, plan: ExecutionPlan,
@@ -237,69 +292,137 @@ async fn execute_stream_from_frame_stream(
"execution runtime stream must start with headers frame".to_string(), "execution runtime stream must start with headers frame".to_string(),
)); ));
}; };
let mut buffered_frames = VecDeque::new();
if should_retry_next_local_candidate_stream(plan_kind, report_context.as_ref(), status_code) { if status_code == 200 {
let terminal_unix_secs = current_request_candidate_unix_secs(); let success_probe_text =
record_local_request_candidate_status( probe_local_stream_success_failover_text(&mut buffered_frames, &mut lines).await?;
if should_retry_next_local_candidate_stream(
state, state,
&plan, &plan,
plan_kind,
report_context.as_ref(), report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: Some(status_code),
error_type: Some("retryable_upstream_status".to_string()),
error_message: Some(format!(
"execution runtime stream returned retryable status {status_code}"
)),
latency_ms: None,
started_at_unix_secs: Some(terminal_unix_secs),
finished_at_unix_secs: Some(terminal_unix_secs),
},
)
.await;
warn!(
event_name = "local_stream_candidate_retry_scheduled",
log_type = "event",
trace_id = %trace_id,
request_id = %request_id_for_log,
status_code, status_code,
"gateway local stream decision retrying next candidate after retryable execution runtime status" success_probe_text.as_deref(),
); )
return Ok(None); .await
{
let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status(
state,
&plan,
report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: Some(status_code),
error_type: Some("success_failover_pattern".to_string()),
error_message: Some(
"execution runtime stream matched provider success failover rule"
.to_string(),
),
latency_ms: None,
started_at_unix_ms: Some(terminal_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
warn!(
event_name = "local_stream_candidate_retry_scheduled",
log_type = "event",
trace_id = %trace_id,
request_id = %request_id_for_log,
status_code,
"gateway local stream decision retrying next candidate after success failover rule match"
);
return Ok(None);
}
} }
let stream_error_finalize_kind = let stream_error_finalize_kind =
resolve_core_stream_error_finalize_report_kind(plan_kind, status_code); resolve_core_stream_error_finalize_report_kind(plan_kind, status_code);
if should_fallback_to_control_stream(
plan_kind,
status_code,
stream_error_finalize_kind.is_some(),
) {
let terminal_unix_secs = current_request_candidate_unix_secs();
record_local_request_candidate_status(
state,
&plan,
report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: Some(status_code),
error_type: Some("control_fallback".to_string()),
error_message: Some(format!(
"stream decision fell back to control after status {status_code}"
)),
latency_ms: None,
started_at_unix_secs: Some(terminal_unix_secs),
finished_at_unix_secs: Some(terminal_unix_secs),
},
)
.await;
return Ok(None);
}
if status_code >= 400 { if status_code >= 400 {
let error_body = collect_error_body(&mut lines).await?; let error_body = collect_error_body(&mut lines).await?;
let (body_json, body_base64) = decode_stream_error_body(&headers, &error_body); let (body_json, body_base64) = decode_stream_error_body(&headers, &error_body);
let error_response_text =
local_failover_response_text(body_json.as_ref(), &error_body, None);
let stop_local_failover = should_stop_local_candidate_failover_stream(
state,
&plan,
plan_kind,
report_context.as_ref(),
status_code,
error_response_text.as_deref(),
)
.await;
if !stop_local_failover
&& should_retry_next_local_candidate_stream(
state,
&plan,
plan_kind,
report_context.as_ref(),
status_code,
error_response_text.as_deref(),
)
.await
{
let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status(
state,
&plan,
report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: Some(status_code),
error_type: Some("retryable_upstream_status".to_string()),
error_message: Some(format!(
"execution runtime stream returned retryable status {status_code}"
)),
latency_ms: None,
started_at_unix_ms: Some(terminal_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
warn!(
event_name = "local_stream_candidate_retry_scheduled",
log_type = "event",
trace_id = %trace_id,
request_id = %request_id_for_log,
status_code,
"gateway local stream decision retrying next candidate after retryable execution runtime status"
);
return Ok(None);
}
if !stop_local_failover
&& should_fallback_to_control_stream(
plan_kind,
status_code,
stream_error_finalize_kind.is_some(),
)
{
let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status(
state,
&plan,
report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: Some(status_code),
error_type: Some("control_fallback".to_string()),
error_message: Some(format!(
"stream decision fell back to control after status {status_code}"
)),
latency_ms: None,
started_at_unix_ms: Some(terminal_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
return Ok(None);
}
let usage_report_kind = stream_error_finalize_kind let usage_report_kind = stream_error_finalize_kind
.clone() .clone()
.or_else(|| report_kind.clone()) .or_else(|| report_kind.clone())
@@ -324,7 +447,7 @@ async fn execute_stream_from_frame_stream(
&usage_payload, &usage_payload,
) )
.await; .await;
let terminal_unix_secs = current_request_candidate_unix_secs(); let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status( record_local_request_candidate_status(
state, state,
&plan, &plan,
@@ -337,8 +460,8 @@ async fn execute_stream_from_frame_stream(
"execution runtime stream returned error status {status_code}" "execution runtime stream returned error status {status_code}"
)), )),
latency_ms: None, latency_ms: None,
started_at_unix_secs: Some(terminal_unix_secs), started_at_unix_ms: Some(terminal_unix_secs),
finished_at_unix_secs: Some(terminal_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs),
}, },
) )
.await; .await;
@@ -399,7 +522,7 @@ async fn execute_stream_from_frame_stream(
while prefetched_chunks.len() < MAX_STREAM_PREFETCH_FRAMES while prefetched_chunks.len() < MAX_STREAM_PREFETCH_FRAMES
&& prefetched_inspection_body.len() < MAX_STREAM_PREFETCH_BYTES && prefetched_inspection_body.len() < MAX_STREAM_PREFETCH_BYTES
{ {
let Some(frame) = (match read_next_frame(&mut lines).await { let Some(frame) = (match next_stream_frame(&mut buffered_frames, &mut lines).await {
Ok(frame) => frame, Ok(frame) => frame,
Err(err) => { Err(err) => {
let failure = build_stream_failure_report( let failure = build_stream_failure_report(
@@ -429,14 +552,14 @@ async fn execute_stream_from_frame_stream(
}; };
match frame.payload { match frame.payload {
StreamFramePayload::Data { chunk_b64, text } => { StreamFramePayload::Data { chunk_b64, text } => {
let chunk = if let Some(chunk_b64) = chunk_b64 { let chunk =
match base64::engine::general_purpose::STANDARD.decode(chunk_b64) { match decode_stream_data_chunk(chunk_b64.as_deref(), text.as_deref()) {
Ok(decoded) => decoded, Ok(chunk) => chunk,
Err(err) => { Err(err) => {
let failure = build_stream_failure_report( let failure = build_stream_failure_report(
"execution_runtime_stream_chunk_decode_error", "execution_runtime_stream_chunk_decode_error",
format!( format!(
"failed to decode execution runtime stream chunk: {err}" "failed to decode execution runtime stream chunk: {err:?}"
), ),
502, 502,
); );
@@ -456,12 +579,7 @@ async fn execute_stream_from_frame_stream(
) )
.await; .await;
} }
} };
} else if let Some(text) = text {
text.into_bytes()
} else {
Vec::new()
};
if chunk.is_empty() { if chunk.is_empty() {
continue; continue;
@@ -621,7 +739,7 @@ async fn execute_stream_from_frame_stream(
} }
} }
let candidate_started_unix_secs = current_request_candidate_unix_secs(); let candidate_started_unix_secs = current_request_candidate_unix_ms();
state state
.usage_runtime .usage_runtime
.record_pending(state.data.as_ref(), &plan, report_context.as_ref()) .record_pending(state.data.as_ref(), &plan, report_context.as_ref())
@@ -649,8 +767,8 @@ async fn execute_stream_from_frame_stream(
latency_ms: prefetched_telemetry latency_ms: prefetched_telemetry
.as_ref() .as_ref()
.and_then(|telemetry| telemetry.elapsed_ms), .and_then(|telemetry| telemetry.elapsed_ms),
started_at_unix_secs: Some(candidate_started_unix_secs), started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_secs: None, finished_at_unix_ms: None,
}, },
) )
.await; .await;
@@ -672,6 +790,7 @@ async fn execute_stream_from_frame_stream(
let request_id_for_report = request_id.to_string(); let request_id_for_report = request_id.to_string();
let request_id_for_report_log = short_request_id(request_id); let request_id_for_report_log = short_request_id(request_id);
let candidate_id_for_report = candidate_id.map(ToOwned::to_owned); let candidate_id_for_report = candidate_id.map(ToOwned::to_owned);
let mut buffered_frames = buffered_frames;
tokio::spawn(async move { tokio::spawn(async move {
let mut provider_buffered_body = provider_prefetched_body_for_report; let mut provider_buffered_body = provider_prefetched_body_for_report;
let mut buffered_body = prefetched_body_for_report; let mut buffered_body = prefetched_body_for_report;
@@ -682,7 +801,7 @@ async fn execute_stream_from_frame_stream(
if !reached_eof { if !reached_eof {
loop { loop {
let next_frame = match read_next_frame(&mut lines).await { let next_frame = match next_stream_frame(&mut buffered_frames, &mut lines).await {
Ok(frame) => frame, Ok(frame) => frame,
Err(err) => { Err(err) => {
warn!( warn!(
@@ -707,9 +826,9 @@ async fn execute_stream_from_frame_stream(
}; };
match frame.payload { match frame.payload {
StreamFramePayload::Data { chunk_b64, text } => { StreamFramePayload::Data { chunk_b64, text } => {
let chunk = if let Some(chunk_b64) = chunk_b64 { let chunk =
match base64::engine::general_purpose::STANDARD.decode(chunk_b64) { match decode_stream_data_chunk(chunk_b64.as_deref(), text.as_deref()) {
Ok(decoded) => decoded, Ok(chunk) => chunk,
Err(err) => { Err(err) => {
warn!( warn!(
event_name = "stream_execution_chunk_decode_failed", event_name = "stream_execution_chunk_decode_failed",
@@ -717,22 +836,19 @@ async fn execute_stream_from_frame_stream(
trace_id = %trace_id_owned, trace_id = %trace_id_owned,
request_id = %request_id_for_report_log, request_id = %request_id_for_report_log,
candidate_id = ?candidate_id_for_report.as_deref(), candidate_id = ?candidate_id_for_report.as_deref(),
error = %err, error = ?err,
"gateway failed to decode execution runtime chunk" "gateway failed to decode execution runtime chunk"
); );
terminal_failure = Some(build_stream_failure_report( terminal_failure = Some(build_stream_failure_report(
"execution_runtime_stream_chunk_decode_error", "execution_runtime_stream_chunk_decode_error",
format!("failed to decode execution runtime stream chunk: {err}"), format!(
"failed to decode execution runtime stream chunk: {err:?}"
),
502, 502,
)); ));
break; break;
} }
} };
} else if let Some(text) = text {
text.into_bytes()
} else {
Vec::new()
};
if chunk.is_empty() { if chunk.is_empty() {
continue; continue;
@@ -998,8 +1114,8 @@ async fn execute_stream_from_frame_stream(
error_type: Some("downstream_disconnect".to_string()), error_type: Some("downstream_disconnect".to_string()),
error_message: Some("client disconnected before stream completion".to_string()), error_message: Some("client disconnected before stream completion".to_string()),
latency_ms: telemetry.as_ref().and_then(|value| value.elapsed_ms), latency_ms: telemetry.as_ref().and_then(|value| value.elapsed_ms),
started_at_unix_secs: Some(candidate_started_unix_secs_for_report), started_at_unix_ms: Some(candidate_started_unix_secs_for_report),
finished_at_unix_secs: Some(current_request_candidate_unix_secs()), finished_at_unix_ms: Some(current_request_candidate_unix_ms()),
}, },
) )
.await; .await;
@@ -1055,8 +1171,8 @@ async fn execute_stream_from_frame_stream(
error_type: None, error_type: None,
error_message: None, error_message: None,
latency_ms: telemetry.as_ref().and_then(|value| value.elapsed_ms), latency_ms: telemetry.as_ref().and_then(|value| value.elapsed_ms),
started_at_unix_secs: Some(candidate_started_unix_secs_for_report), started_at_unix_ms: Some(candidate_started_unix_secs_for_report),
finished_at_unix_secs: Some(current_request_candidate_unix_secs()), finished_at_unix_ms: Some(current_request_candidate_unix_ms()),
}, },
) )
.await; .await;

View File

@@ -8,7 +8,7 @@ use serde_json::{Map, Value};
use tracing::warn; use tracing::warn;
use crate::api::response::attach_control_metadata_headers; use crate::api::response::attach_control_metadata_headers;
use crate::clock::current_unix_secs as current_request_candidate_unix_secs; use crate::clock::current_unix_ms as current_request_candidate_unix_ms;
use crate::control::GatewayControlDecision; use crate::control::GatewayControlDecision;
use crate::execution_runtime::submission::{ use crate::execution_runtime::submission::{
resolve_core_error_background_report_kind, submit_local_core_error_or_sync_finalize, resolve_core_error_background_report_kind, submit_local_core_error_or_sync_finalize,
@@ -117,13 +117,13 @@ async fn record_stream_sync_failure(
report_context: Option<&Value>, report_context: Option<&Value>,
payload: &GatewaySyncReportRequest, payload: &GatewaySyncReportRequest,
failure: &StreamFailureReport, failure: &StreamFailureReport,
started_at_unix_secs: Option<u64>, started_at_unix_ms: Option<u64>,
) { ) {
state state
.usage_runtime .usage_runtime
.record_sync_terminal(state.data.as_ref(), plan, report_context, payload) .record_sync_terminal(state.data.as_ref(), plan, report_context, payload)
.await; .await;
let terminal_unix_secs = current_request_candidate_unix_secs(); let terminal_unix_secs = current_request_candidate_unix_ms();
record_report_request_candidate_status( record_report_request_candidate_status(
state, state,
report_context, report_context,
@@ -136,8 +136,8 @@ async fn record_stream_sync_failure(
.telemetry .telemetry
.as_ref() .as_ref()
.and_then(|telemetry| telemetry.elapsed_ms), .and_then(|telemetry| telemetry.elapsed_ms),
started_at_unix_secs: started_at_unix_secs.or(Some(terminal_unix_secs)), started_at_unix_ms: started_at_unix_ms.or(Some(terminal_unix_secs)),
finished_at_unix_secs: Some(terminal_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs),
}, },
) )
.await; .await;
@@ -195,7 +195,7 @@ pub(super) async fn submit_midstream_stream_failure(
headers: &std::collections::BTreeMap<String, String>, headers: &std::collections::BTreeMap<String, String>,
telemetry: Option<ExecutionTelemetry>, telemetry: Option<ExecutionTelemetry>,
buffered_body: &[u8], buffered_body: &[u8],
started_at_unix_secs: u64, started_at_unix_ms: u64,
failure: StreamFailureReport, failure: StreamFailureReport,
) { ) {
let Some(report_kind) = let Some(report_kind) =
@@ -219,7 +219,7 @@ pub(super) async fn submit_midstream_stream_failure(
report_context, report_context,
&payload, &payload,
&failure, &failure,
Some(started_at_unix_secs), Some(started_at_unix_ms),
) )
.await; .await;
if let Err(err) = submit_sync_report(state, trace_id, payload).await { if let Err(err) = submit_sync_report(state, trace_id, payload).await {

View File

@@ -297,29 +297,39 @@ fn resolve_local_sync_error_status_code(status_code: u16, body_json: &serde_json
return status_code; return status_code;
} }
let Some(error_object) = body_json.get("error").and_then(|value| value.as_object()) else { let body_object = body_json.as_object();
return 400; let error_object = body_object
}; .and_then(|object| object.get("error"))
.and_then(|value| value.as_object());
for key in ["code", "status"] { let raw_code = first_non_empty_error_text(error_object, body_object, &["code"]);
let Some(value) = error_object.get(key) else { let raw_status = first_non_empty_error_text(error_object, body_object, &["status"]);
continue; for numeric_hint in [raw_code.as_deref(), raw_status.as_deref()]
}; .into_iter()
if let Some(number) = value.as_u64() { .flatten()
{
if let Ok(number) = numeric_hint.parse::<u16>() {
if (400..600).contains(&number) { if (400..600).contains(&number) {
return number as u16; return number;
}
}
if let Some(text) = value.as_str() {
if let Ok(number) = text.parse::<u16>() {
if (400..600).contains(&number) {
return number;
}
} }
} }
} }
400 let raw_type = first_non_empty_error_text(error_object, body_object, &["type", "__type"]);
let message = first_non_empty_error_text(
error_object,
body_object,
&["message", "detail", "reason", "status", "type", "__type"],
)
.unwrap_or_else(|| "HTTP 400".to_string());
let kind = classify_local_sync_error_kind(
status_code,
raw_type.as_deref(),
raw_status.as_deref(),
raw_code.as_deref(),
message.as_str(),
);
default_status_code_for_local_sync_error_kind(kind)
} }
fn extract_local_sync_error_details( fn extract_local_sync_error_details(
@@ -436,6 +446,20 @@ fn classify_local_sync_error_kind(
LocalCoreSyncErrorKind::InvalidRequest LocalCoreSyncErrorKind::InvalidRequest
} }
fn default_status_code_for_local_sync_error_kind(kind: LocalCoreSyncErrorKind) -> u16 {
match kind {
LocalCoreSyncErrorKind::InvalidRequest | LocalCoreSyncErrorKind::ContextLengthExceeded => {
400
}
LocalCoreSyncErrorKind::Authentication => 401,
LocalCoreSyncErrorKind::PermissionDenied => 403,
LocalCoreSyncErrorKind::NotFound => 404,
LocalCoreSyncErrorKind::RateLimit => 429,
LocalCoreSyncErrorKind::Overloaded => 503,
LocalCoreSyncErrorKind::ServerError => 500,
}
}
pub(crate) fn strip_utf8_bom_and_ws(mut body: &[u8]) -> &[u8] { pub(crate) fn strip_utf8_bom_and_ws(mut body: &[u8]) -> &[u8] {
loop { loop {
while let Some(first) = body.first() { while let Some(first) = body.first() {
@@ -534,3 +558,128 @@ pub(crate) async fn submit_local_core_error_or_sync_finalize(
Ok(response) Ok(response)
} }
#[cfg(test)]
mod tests {
use axum::body::to_bytes;
use serde_json::json;
use super::maybe_build_local_core_error_response;
use crate::control::GatewayControlDecision;
use crate::usage::GatewaySyncReportRequest;
fn test_decision() -> GatewayControlDecision {
GatewayControlDecision::synthetic(
"/v1/chat/completions",
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
Some("openai:chat".to_string()),
)
.with_execution_runtime_candidate(true)
}
fn core_finalize_payload(
report_kind: &str,
client_api_format: &str,
provider_api_format: &str,
status_code: u16,
body_json: serde_json::Value,
) -> GatewaySyncReportRequest {
GatewaySyncReportRequest {
trace_id: "trace-core-error-status-123".to_string(),
report_kind: report_kind.to_string(),
report_context: Some(json!({
"client_api_format": client_api_format,
"provider_api_format": provider_api_format,
})),
status_code,
headers: Default::default(),
body_json: Some(body_json),
client_body_json: None,
body_base64: None,
telemetry: None,
}
}
#[tokio::test]
async fn maybe_build_local_core_error_response_infers_status_from_semantic_error_type() {
let payload = core_finalize_payload(
"openai_chat_sync_finalize",
"openai:chat",
"claude:chat",
200,
json!({
"type": "error",
"error": {
"type": "rate_limit_error",
"message": "slow down"
}
}),
);
let response = maybe_build_local_core_error_response(
"trace-sync-status-type",
&test_decision(),
&payload,
)
.expect("response build should not error")
.expect("response should exist");
assert_eq!(response.status(), http::StatusCode::TOO_MANY_REQUESTS);
assert_eq!(
serde_json::from_slice::<serde_json::Value>(
&to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("body should decode"),
json!({
"error": {
"message": "slow down",
"type": "rate_limit_error"
}
})
);
}
#[tokio::test]
async fn maybe_build_local_core_error_response_infers_status_from_gemini_status_text() {
let payload = core_finalize_payload(
"gemini_chat_sync_finalize",
"gemini:chat",
"gemini:chat",
200,
json!({
"error": {
"message": "quota reached",
"status": "RESOURCE_EXHAUSTED"
}
}),
);
let response = maybe_build_local_core_error_response(
"trace-sync-status-gemini",
&test_decision(),
&payload,
)
.expect("response build should not error")
.expect("response should exist");
assert_eq!(response.status(), http::StatusCode::TOO_MANY_REQUESTS);
assert_eq!(
serde_json::from_slice::<serde_json::Value>(
&to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("body should decode"),
json!({
"error": {
"message": "quota reached",
"status": "RESOURCE_EXHAUSTED"
}
})
);
}
}

View File

@@ -15,7 +15,7 @@ use crate::ai_pipeline_api::{
use crate::api::response::{ use crate::api::response::{
attach_control_metadata_headers, build_client_response, build_client_response_from_parts, attach_control_metadata_headers, build_client_response, build_client_response_from_parts,
}; };
use crate::clock::current_unix_secs as current_request_candidate_unix_secs; use crate::clock::current_unix_ms as current_request_candidate_unix_ms;
use crate::constants::{CONTROL_CANDIDATE_ID_HEADER, CONTROL_REQUEST_ID_HEADER}; use crate::constants::{CONTROL_CANDIDATE_ID_HEADER, CONTROL_REQUEST_ID_HEADER};
use crate::control::GatewayControlDecision; use crate::control::GatewayControlDecision;
#[cfg(test)] #[cfg(test)]
@@ -23,8 +23,9 @@ use crate::execution_runtime::remote_compat::post_sync_plan_to_remote_execution_
use crate::execution_runtime::submission::submit_local_core_error_or_sync_finalize; use crate::execution_runtime::submission::submit_local_core_error_or_sync_finalize;
use crate::execution_runtime::transport::DirectSyncExecutionRuntime; use crate::execution_runtime::transport::DirectSyncExecutionRuntime;
use crate::execution_runtime::{ use crate::execution_runtime::{
resolve_core_sync_error_finalize_report_kind, should_fallback_to_control_sync, local_failover_response_text, resolve_core_sync_error_finalize_report_kind,
should_finalize_sync_response, should_retry_next_local_candidate_sync, should_fallback_to_control_sync, should_finalize_sync_response,
should_retry_next_local_candidate_sync, should_stop_local_candidate_failover_sync,
}; };
use crate::log_ids::short_request_id; use crate::log_ids::short_request_id;
use crate::request_candidate_runtime::{ use crate::request_candidate_runtime::{
@@ -85,6 +86,7 @@ pub(crate) async fn execute_execution_runtime_sync(
let plan_request_id = plan.request_id.as_str(); let plan_request_id = plan.request_id.as_str();
let plan_request_id_for_log = short_request_id(plan_request_id); let plan_request_id_for_log = short_request_id(plan_request_id);
let plan_candidate_id = plan.candidate_id.as_deref(); let plan_candidate_id = plan.candidate_id.as_deref();
let candidate_started_unix_secs = current_request_candidate_unix_ms();
#[cfg(not(test))] #[cfg(not(test))]
let result = { let result = {
match DirectSyncExecutionRuntime::new() match DirectSyncExecutionRuntime::new()
@@ -140,6 +142,7 @@ pub(crate) async fn execute_execution_runtime_sync(
plan_request_id, plan_request_id,
plan_candidate_id, plan_candidate_id,
report_context.as_ref(), report_context.as_ref(),
candidate_started_unix_secs,
) )
.await?; .await?;
match remote_outcome { match remote_outcome {
@@ -159,8 +162,34 @@ pub(crate) async fn execute_execution_runtime_sync(
.telemetry .telemetry
.as_ref() .as_ref()
.and_then(|telemetry| telemetry.elapsed_ms); .and_then(|telemetry| telemetry.elapsed_ms);
if should_retry_next_local_candidate_sync(plan_kind, report_context.as_ref(), &result) { let mut headers = result.headers.clone();
let terminal_unix_secs = current_request_candidate_unix_secs(); let (body_bytes, body_json, body_base64) = decode_execution_result_body(&result, &mut headers)?;
let local_failover_response_text = local_failover_response_text(
body_json.as_ref(),
&body_bytes,
result.error.as_ref().map(|error| error.message.as_str()),
);
let stop_local_failover = should_stop_local_candidate_failover_sync(
state,
&plan,
plan_kind,
report_context.as_ref(),
&result,
local_failover_response_text.as_deref(),
)
.await;
if should_retry_next_local_candidate_sync(
state,
&plan,
plan_kind,
report_context.as_ref(),
&result,
local_failover_response_text.as_deref(),
)
.await
&& !stop_local_failover
{
let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status( record_local_request_candidate_status(
state, state,
&plan, &plan,
@@ -171,8 +200,8 @@ pub(crate) async fn execute_execution_runtime_sync(
error_type: result_error_type.clone(), error_type: result_error_type.clone(),
error_message: result_error_message.clone(), error_message: result_error_message.clone(),
latency_ms: result_latency_ms, latency_ms: result_latency_ms,
started_at_unix_secs: Some(terminal_unix_secs), started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_secs: Some(terminal_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs),
}, },
) )
.await; .await;
@@ -191,8 +220,6 @@ pub(crate) async fn execute_execution_runtime_sync(
.or(Some(plan_request_id)); .or(Some(plan_request_id));
let request_id_for_log = short_request_id(request_id.unwrap_or("-")); let request_id_for_log = short_request_id(request_id.unwrap_or("-"));
let candidate_id = result.candidate_id.as_deref().or(plan_candidate_id); let candidate_id = result.candidate_id.as_deref().or(plan_candidate_id);
let mut headers = result.headers.clone();
let (body_bytes, body_json, body_base64) = decode_execution_result_body(&result, &mut headers)?;
let has_body_bytes = body_base64.is_some(); let has_body_bytes = body_base64.is_some();
let explicit_finalize = should_finalize_sync_response(report_kind.as_deref()); let explicit_finalize = should_finalize_sync_response(report_kind.as_deref());
let mapped_error_finalize_kind = let mapped_error_finalize_kind =
@@ -220,15 +247,17 @@ pub(crate) async fn execute_execution_runtime_sync(
mapped_error_finalize_kind.clone() mapped_error_finalize_kind.clone()
}; };
if should_fallback_to_control_sync( if !stop_local_failover
plan_kind, && should_fallback_to_control_sync(
&result, plan_kind,
body_json.as_ref(), &result,
has_body_bytes, body_json.as_ref(),
explicit_finalize || implicit_finalize.is_some(), has_body_bytes,
mapped_error_finalize_kind.is_some(), explicit_finalize || implicit_finalize.is_some(),
) { mapped_error_finalize_kind.is_some(),
let terminal_unix_secs = current_request_candidate_unix_secs(); )
{
let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status( record_local_request_candidate_status(
state, state,
&plan, &plan,
@@ -239,8 +268,8 @@ pub(crate) async fn execute_execution_runtime_sync(
error_type: result_error_type.clone(), error_type: result_error_type.clone(),
error_message: result_error_message.clone(), error_message: result_error_message.clone(),
latency_ms: result_latency_ms, latency_ms: result_latency_ms,
started_at_unix_secs: Some(terminal_unix_secs), started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_secs: Some(terminal_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs),
}, },
) )
.await; .await;
@@ -251,7 +280,7 @@ pub(crate) async fn execute_execution_runtime_sync(
.usage_runtime .usage_runtime
.record_pending(state.data.as_ref(), &plan, report_context.as_ref()) .record_pending(state.data.as_ref(), &plan, report_context.as_ref())
.await; .await;
let terminal_unix_secs = current_request_candidate_unix_secs(); let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status( record_local_request_candidate_status(
state, state,
&plan, &plan,
@@ -266,8 +295,8 @@ pub(crate) async fn execute_execution_runtime_sync(
error_type: result_error_type.clone(), error_type: result_error_type.clone(),
error_message: result_error_message.clone(), error_message: result_error_message.clone(),
latency_ms: result_latency_ms, latency_ms: result_latency_ms,
started_at_unix_secs: Some(terminal_unix_secs), started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_secs: Some(terminal_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs),
}, },
) )
.await; .await;
@@ -579,6 +608,7 @@ async fn execute_sync_via_remote_execution_runtime(
plan_request_id: &str, plan_request_id: &str,
plan_candidate_id: Option<&str>, plan_candidate_id: Option<&str>,
report_context: Option<&serde_json::Value>, report_context: Option<&serde_json::Value>,
candidate_started_unix_secs: u64,
) -> Result<RemoteSyncFallbackOutcome, GatewayError> { ) -> Result<RemoteSyncFallbackOutcome, GatewayError> {
let response = match post_sync_plan_to_remote_execution_runtime( let response = match post_sync_plan_to_remote_execution_runtime(
state, state,
@@ -604,7 +634,7 @@ async fn execute_sync_via_remote_execution_runtime(
}; };
if response.status() != http::StatusCode::OK { if response.status() != http::StatusCode::OK {
let terminal_unix_secs = current_request_candidate_unix_secs(); let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status( record_local_request_candidate_status(
state, state,
plan, plan,
@@ -618,8 +648,8 @@ async fn execute_sync_via_remote_execution_runtime(
response.status() response.status()
)), )),
latency_ms: None, latency_ms: None,
started_at_unix_secs: Some(terminal_unix_secs), started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_secs: Some(terminal_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs),
}, },
) )
.await; .await;

View File

@@ -228,6 +228,160 @@ fn build_best_effort_local_core_error_body_converts_claude_cli_error_to_openai_c
); );
} }
#[test]
fn build_best_effort_local_core_error_body_converts_sync_errors_across_standard_families() {
let cases = vec![
(
"claude chat -> openai chat",
core_finalize_payload(
"openai_chat_sync_finalize",
"openai:chat",
"claude:chat",
429,
json!({
"type": "error",
"error": {
"type": "rate_limit_error",
"message": "slow down"
}
}),
),
json!({
"error": {
"message": "slow down",
"type": "rate_limit_error"
}
}),
),
(
"openai chat -> claude chat",
core_finalize_payload(
"claude_chat_sync_finalize",
"claude:chat",
"openai:chat",
404,
json!({
"error": {
"message": "missing model",
"type": "not_found_error",
"code": "model_missing"
}
}),
),
json!({
"type": "error",
"error": {
"message": "missing model",
"type": "not_found_error",
"code": "model_missing"
}
}),
),
(
"openai chat -> gemini chat",
core_finalize_payload(
"gemini_chat_sync_finalize",
"gemini:chat",
"openai:chat",
401,
json!({
"error": {
"message": "bad auth",
"type": "authentication_error"
}
}),
),
json!({
"error": {
"code": 401,
"message": "bad auth",
"status": "UNAUTHENTICATED"
}
}),
),
(
"gemini cli -> openai cli",
core_finalize_payload(
"openai_cli_sync_finalize",
"openai:cli",
"gemini:cli",
429,
json!({
"error": {
"message": "quota hit",
"status": "RESOURCE_EXHAUSTED",
"code": 429
}
}),
),
json!({
"error": {
"message": "quota hit",
"type": "rate_limit_error",
"code": "429"
}
}),
),
(
"gemini cli -> claude cli",
core_finalize_payload(
"claude_cli_sync_finalize",
"claude:cli",
"gemini:cli",
503,
json!({
"error": {
"message": "backend busy",
"status": "UNAVAILABLE"
}
}),
),
json!({
"type": "error",
"error": {
"message": "backend busy",
"type": "api_error",
"code": "UNAVAILABLE"
}
}),
),
(
"claude cli -> gemini cli",
core_finalize_payload(
"gemini_cli_sync_finalize",
"gemini:cli",
"claude:cli",
404,
json!({
"type": "error",
"error": {
"type": "not_found_error",
"message": "resource missing"
}
}),
),
json!({
"error": {
"code": 404,
"message": "resource missing",
"status": "NOT_FOUND"
}
}),
),
];
for (label, payload, expected) in cases {
let converted = build_best_effort_local_core_error_body(
&payload,
payload.body_json.as_ref().expect("body_json should exist"),
)
.expect("conversion should not error")
.expect("conversion should produce a client error body");
assert_eq!(converted, expected, "unexpected conversion for {label}");
}
}
#[test] #[test]
fn resolve_local_core_error_response_body_json_parses_body_base64_json_for_cross_format_cli_error() fn resolve_local_core_error_response_body_json_parses_body_base64_json_for_cross_format_cli_error()
{ {

View File

@@ -125,8 +125,8 @@ where
error_type: None, error_type: None,
error_message: None, error_message: None,
latency_ms: None, latency_ms: None,
started_at_unix_secs: None, started_at_unix_ms: None,
finished_at_unix_secs: None, finished_at_unix_ms: None,
}, },
) )
.await; .await;
@@ -153,8 +153,8 @@ pub(crate) async fn mark_unused_local_candidate_items<T, FPlan, FContext>(
error_type: None, error_type: None,
error_message: None, error_message: None,
latency_ms: None, latency_ms: None,
started_at_unix_secs: None, started_at_unix_ms: None,
finished_at_unix_secs: None, finished_at_unix_ms: None,
}, },
) )
.await; .await;

View File

@@ -53,7 +53,7 @@ pub(super) fn build_admin_billing_collector_payload_from_record(
"default_value": record.default_value, "default_value": record.default_value,
"priority": record.priority, "priority": record.priority,
"is_enabled": record.is_enabled, "is_enabled": record.is_enabled,
"created_at": unix_secs_to_rfc3339(record.created_at_unix_secs), "created_at": unix_secs_to_rfc3339(record.created_at_unix_ms),
"updated_at": unix_secs_to_rfc3339(record.updated_at_unix_secs), "updated_at": unix_secs_to_rfc3339(record.updated_at_unix_secs),
}) })
} }

View File

@@ -213,7 +213,7 @@ pub(super) fn build_admin_payment_order_payload(
"gateway_order_id": record.gateway_order_id, "gateway_order_id": record.gateway_order_id,
"gateway_response": record.gateway_response, "gateway_response": record.gateway_response,
"status": admin_payment_effective_status(&record.status, record.expires_at_unix_secs), "status": admin_payment_effective_status(&record.status, record.expires_at_unix_secs),
"created_at": unix_secs_to_rfc3339(record.created_at_unix_secs), "created_at": unix_secs_to_rfc3339(record.created_at_unix_ms),
"paid_at": record.paid_at_unix_secs.and_then(unix_secs_to_rfc3339), "paid_at": record.paid_at_unix_secs.and_then(unix_secs_to_rfc3339),
"credited_at": record.credited_at_unix_secs.and_then(unix_secs_to_rfc3339), "credited_at": record.credited_at_unix_secs.and_then(unix_secs_to_rfc3339),
"expires_at": record.expires_at_unix_secs.and_then(unix_secs_to_rfc3339), "expires_at": record.expires_at_unix_secs.and_then(unix_secs_to_rfc3339),
@@ -236,7 +236,7 @@ pub(super) fn build_admin_payment_callback_payload(
"payload": row.try_get::<Option<serde_json::Value>, _>("payload").map_err(|err| GatewayError::Internal(err.to_string()))?, "payload": row.try_get::<Option<serde_json::Value>, _>("payload").map_err(|err| GatewayError::Internal(err.to_string()))?,
"error_message": row.try_get::<Option<String>, _>("error_message").map_err(|err| GatewayError::Internal(err.to_string()))?, "error_message": row.try_get::<Option<String>, _>("error_message").map_err(|err| GatewayError::Internal(err.to_string()))?,
"created_at": row "created_at": row
.try_get::<Option<i64>, _>("created_at_unix_secs") .try_get::<Option<i64>, _>("created_at_unix_ms")
.map_err(|err| GatewayError::Internal(err.to_string()))? .map_err(|err| GatewayError::Internal(err.to_string()))?
.and_then(|value| u64::try_from(value).ok()) .and_then(|value| u64::try_from(value).ok())
.and_then(unix_secs_to_rfc3339), .and_then(unix_secs_to_rfc3339),
@@ -263,7 +263,7 @@ pub(super) fn build_admin_payment_callback_payload_from_record(
"status": record.status, "status": record.status,
"payload": record.payload, "payload": record.payload,
"error_message": record.error_message, "error_message": record.error_message,
"created_at": unix_secs_to_rfc3339(record.created_at_unix_secs), "created_at": unix_secs_to_rfc3339(record.created_at_unix_ms),
"processed_at": record.processed_at_unix_secs.and_then(unix_secs_to_rfc3339), "processed_at": record.processed_at_unix_secs.and_then(unix_secs_to_rfc3339),
}) })
} }

View File

@@ -53,7 +53,7 @@ fn build_admin_billing_rule_payload_from_record(
"variables": record.variables, "variables": record.variables,
"dimension_mappings": record.dimension_mappings, "dimension_mappings": record.dimension_mappings,
"is_enabled": record.is_enabled, "is_enabled": record.is_enabled,
"created_at": unix_secs_to_rfc3339(record.created_at_unix_secs), "created_at": unix_secs_to_rfc3339(record.created_at_unix_ms),
"updated_at": unix_secs_to_rfc3339(record.updated_at_unix_secs), "updated_at": unix_secs_to_rfc3339(record.updated_at_unix_secs),
}) })
} }

View File

@@ -97,7 +97,7 @@ pub(in super::super) async fn build_admin_wallet_adjust_response(
transaction.link_id.as_deref(), transaction.link_id.as_deref(),
transaction.operator_id.as_deref(), transaction.operator_id.as_deref(),
transaction.description.as_deref(), transaction.description.as_deref(),
unix_secs_to_rfc3339(transaction.created_at_unix_secs), unix_secs_to_rfc3339(transaction.created_at_unix_ms),
); );
let response = Json(json!({ let response = Json(json!({
"wallet": wallet_payload, "wallet": wallet_payload,

View File

@@ -84,7 +84,7 @@ pub(in super::super) async fn build_admin_wallet_fail_refund_response(
transaction.link_id.as_deref(), transaction.link_id.as_deref(),
transaction.operator_id.as_deref(), transaction.operator_id.as_deref(),
transaction.description.as_deref(), transaction.description.as_deref(),
unix_secs_to_rfc3339(transaction.created_at_unix_secs), unix_secs_to_rfc3339(transaction.created_at_unix_ms),
) )
}) })
.unwrap_or(serde_json::Value::Null), .unwrap_or(serde_json::Value::Null),

View File

@@ -69,7 +69,7 @@ pub(in super::super) async fn build_admin_wallet_process_refund_response(
transaction.link_id.as_deref(), transaction.link_id.as_deref(),
transaction.operator_id.as_deref(), transaction.operator_id.as_deref(),
transaction.description.as_deref(), transaction.description.as_deref(),
unix_secs_to_rfc3339(transaction.created_at_unix_secs), unix_secs_to_rfc3339(transaction.created_at_unix_ms),
), ),
})) }))
.into_response(); .into_response();

View File

@@ -87,7 +87,7 @@ pub(in super::super) async fn build_admin_wallet_recharge_response(
payment_order.amount_usd, payment_order.amount_usd,
payment_order.payment_method, payment_order.payment_method,
payment_order.status, payment_order.status,
unix_secs_to_rfc3339(payment_order.created_at_unix_secs), unix_secs_to_rfc3339(payment_order.created_at_unix_ms),
payment_order payment_order
.credited_at_unix_secs .credited_at_unix_secs
.and_then(unix_secs_to_rfc3339), .and_then(unix_secs_to_rfc3339),

View File

@@ -69,7 +69,7 @@ pub(in super::super) async fn build_admin_wallet_ledger_response(
"operator_name": entry.operator_name, "operator_name": entry.operator_name,
"operator_email": entry.operator_email, "operator_email": entry.operator_email,
"description": entry.description, "description": entry.description,
"created_at": entry.created_at_unix_secs.and_then(unix_secs_to_rfc3339), "created_at": entry.created_at_unix_ms.and_then(unix_secs_to_rfc3339),
}) })
}) })
.collect::<Vec<_>>(); .collect::<Vec<_>>();

View File

@@ -59,7 +59,7 @@ pub(in super::super) async fn build_admin_wallet_list_response(
"total_consumed": wallet.total_consumed, "total_consumed": wallet.total_consumed,
"total_refunded": wallet.total_refunded, "total_refunded": wallet.total_refunded,
"total_adjusted": wallet.total_adjusted, "total_adjusted": wallet.total_adjusted,
"created_at": wallet.created_at_unix_secs.and_then(unix_secs_to_rfc3339), "created_at": wallet.created_at_unix_ms.and_then(unix_secs_to_rfc3339),
"updated_at": wallet.updated_at_unix_secs.and_then(unix_secs_to_rfc3339), "updated_at": wallet.updated_at_unix_secs.and_then(unix_secs_to_rfc3339),
}) })
}) })

View File

@@ -79,7 +79,7 @@ pub(in super::super) async fn build_admin_wallet_refund_requests_response(
"requested_by": refund.requested_by, "requested_by": refund.requested_by,
"approved_by": refund.approved_by, "approved_by": refund.approved_by,
"processed_by": refund.processed_by, "processed_by": refund.processed_by,
"created_at": refund.created_at_unix_secs.and_then(unix_secs_to_rfc3339), "created_at": refund.created_at_unix_ms.and_then(unix_secs_to_rfc3339),
"updated_at": refund.updated_at_unix_secs.and_then(unix_secs_to_rfc3339), "updated_at": refund.updated_at_unix_secs.and_then(unix_secs_to_rfc3339),
"processed_at": refund.processed_at_unix_secs.and_then(unix_secs_to_rfc3339), "processed_at": refund.processed_at_unix_secs.and_then(unix_secs_to_rfc3339),
"completed_at": refund.completed_at_unix_secs.and_then(unix_secs_to_rfc3339), "completed_at": refund.completed_at_unix_secs.and_then(unix_secs_to_rfc3339),

View File

@@ -80,7 +80,7 @@ pub(in super::super) async fn build_admin_wallet_transactions_response(
"operator_name": operator_name, "operator_name": operator_name,
"operator_email": operator_email, "operator_email": operator_email,
"description": transaction.description, "description": transaction.description,
"created_at": transaction.created_at_unix_secs.and_then(unix_secs_to_rfc3339), "created_at": transaction.created_at_unix_ms.and_then(unix_secs_to_rfc3339),
})); }));
} }

View File

@@ -188,7 +188,7 @@ pub(in super::super) fn build_admin_wallet_refund_payload(
"requested_by": refund.requested_by.clone(), "requested_by": refund.requested_by.clone(),
"approved_by": refund.approved_by.clone(), "approved_by": refund.approved_by.clone(),
"processed_by": refund.processed_by.clone(), "processed_by": refund.processed_by.clone(),
"created_at": unix_secs_to_rfc3339(refund.created_at_unix_secs), "created_at": unix_secs_to_rfc3339(refund.created_at_unix_ms),
"updated_at": unix_secs_to_rfc3339(refund.updated_at_unix_secs), "updated_at": unix_secs_to_rfc3339(refund.updated_at_unix_secs),
"processed_at": refund.processed_at_unix_secs.and_then(unix_secs_to_rfc3339), "processed_at": refund.processed_at_unix_secs.and_then(unix_secs_to_rfc3339),
"completed_at": refund.completed_at_unix_secs.and_then(unix_secs_to_rfc3339), "completed_at": refund.completed_at_unix_secs.and_then(unix_secs_to_rfc3339),

View File

@@ -3,6 +3,7 @@ use crate::handlers::admin::shared::unix_secs_to_rfc3339;
use crate::handlers::public::{ use crate::handlers::public::{
api_format_display_name, build_public_health_timeline, provider_key_api_formats, api_format_display_name, build_public_health_timeline, provider_key_api_formats,
}; };
use crate::handlers::shared::unix_ms_to_rfc3339;
use aether_data_contracts::repository::candidates::PublicHealthTimelineBucket; use aether_data_contracts::repository::candidates::PublicHealthTimelineBucket;
use aether_scheduler_core::{is_provider_key_circuit_open, provider_key_health_score}; use aether_scheduler_core::{is_provider_key_circuit_open, provider_key_health_score};
use serde_json::json; use serde_json::json;
@@ -124,28 +125,24 @@ pub(crate) async fn build_admin_endpoint_health_status_payload(
total_count: 0, total_count: 0,
success_count: 0, success_count: 0,
failed_count: 0, failed_count: 0,
min_created_at_unix_secs: None, min_created_at_unix_ms: None,
max_created_at_unix_secs: None, max_created_at_unix_ms: None,
}); });
bucket.total_count += row.total_count; bucket.total_count += row.total_count;
bucket.success_count += row.success_count; bucket.success_count += row.success_count;
bucket.failed_count += row.failed_count; bucket.failed_count += row.failed_count;
bucket.min_created_at_unix_secs = match ( bucket.min_created_at_unix_ms =
bucket.min_created_at_unix_secs, match (bucket.min_created_at_unix_ms, row.min_created_at_unix_ms) {
row.min_created_at_unix_secs, (Some(left), Some(right)) => Some(left.min(right)),
) { (None, Some(right)) => Some(right),
(Some(left), Some(right)) => Some(left.min(right)), (left, None) => left,
(None, Some(right)) => Some(right), };
(left, None) => left, bucket.max_created_at_unix_ms =
}; match (bucket.max_created_at_unix_ms, row.max_created_at_unix_ms) {
bucket.max_created_at_unix_secs = match ( (Some(left), Some(right)) => Some(left.max(right)),
bucket.max_created_at_unix_secs, (None, Some(right)) => Some(right),
row.max_created_at_unix_secs, (left, None) => left,
) { };
(Some(left), Some(right)) => Some(left.max(right)),
(None, Some(right)) => Some(right),
(left, None) => left,
};
} }
let mut payload = endpoint_ids_by_format let mut payload = endpoint_ids_by_format
@@ -182,8 +179,8 @@ pub(crate) async fn build_admin_endpoint_health_status_payload(
"display_name": api_format_display_name(&api_format), "display_name": api_format_display_name(&api_format),
"health_score": health_score, "health_score": health_score,
"timeline": timeline, "timeline": timeline,
"time_range_start": time_range_start.and_then(unix_secs_to_rfc3339), "time_range_start": time_range_start.and_then(unix_ms_to_rfc3339),
"time_range_end": time_range_end.or(Some(now_unix_secs)).and_then(unix_secs_to_rfc3339), "time_range_end": time_range_end.map(|ms| unix_ms_to_rfc3339(ms)).unwrap_or_else(|| unix_secs_to_rfc3339(now_unix_secs)),
"total_endpoints": endpoint_ids.len(), "total_endpoints": endpoint_ids.len(),
"total_keys": total_keys, "total_keys": total_keys,
"active_keys": active_keys_by_format.get(&api_format).map(BTreeSet::len).unwrap_or(0), "active_keys": active_keys_by_format.get(&api_format).map(BTreeSet::len).unwrap_or(0),

View File

@@ -140,7 +140,7 @@ fn build_admin_gemini_file_mapping_payload(
"username": username, "username": username,
"display_name": mapping.display_name, "display_name": mapping.display_name,
"mime_type": mapping.mime_type, "mime_type": mapping.mime_type,
"created_at": unix_secs_to_rfc3339(mapping.created_at_unix_secs), "created_at": unix_secs_to_rfc3339(mapping.created_at_unix_ms),
"expires_at": unix_secs_to_rfc3339(mapping.expires_at_unix_secs), "expires_at": unix_secs_to_rfc3339(mapping.expires_at_unix_secs),
"is_expired": mapping.expires_at_unix_secs <= now_unix_secs, "is_expired": mapping.expires_at_unix_secs <= now_unix_secs,
}) })

View File

@@ -94,7 +94,7 @@ pub(super) fn build_admin_video_task_list_item(
"error_message": task.error_message, "error_message": task.error_message,
"poll_count": task.poll_count, "poll_count": task.poll_count,
"max_poll_count": task.max_poll_count, "max_poll_count": task.max_poll_count,
"created_at": admin_video_task_timestamp(Some(task.created_at_unix_secs)), "created_at": admin_video_task_timestamp(Some(task.created_at_unix_ms)),
"completed_at": admin_video_task_timestamp(task.completed_at_unix_secs), "completed_at": admin_video_task_timestamp(task.completed_at_unix_secs),
"submitted_at": admin_video_task_timestamp(task.submitted_at_unix_secs), "submitted_at": admin_video_task_timestamp(task.submitted_at_unix_secs),
}) })

View File

@@ -272,7 +272,7 @@ pub(super) async fn maybe_build_local_admin_video_tasks_response(
payload.insert("max_poll_count".to_string(), json!(task.max_poll_count)); payload.insert("max_poll_count".to_string(), json!(task.max_poll_count));
payload.insert( payload.insert(
"created_at".to_string(), "created_at".to_string(),
json!(admin_video_task_timestamp(Some(task.created_at_unix_secs))), json!(admin_video_task_timestamp(Some(task.created_at_unix_ms))),
); );
payload.insert( payload.insert(
"updated_at".to_string(), "updated_at".to_string(),

View File

@@ -49,6 +49,7 @@ pub(super) async fn build_admin_monitoring_cache_stats_response(
"affinity_stats": { "affinity_stats": {
"storage_type": snapshot.storage_type, "storage_type": snapshot.storage_type,
"total_affinities": snapshot.total_affinities, "total_affinities": snapshot.total_affinities,
"active_affinities": snapshot.total_affinities,
"cache_hits": snapshot.cache_hits, "cache_hits": snapshot.cache_hits,
"cache_misses": snapshot.cache_misses, "cache_misses": snapshot.cache_misses,
"cache_hit_rate": snapshot.cache_hit_rate, "cache_hit_rate": snapshot.cache_hit_rate,

View File

@@ -1,6 +1,8 @@
use super::cache_types::AdminMonitoringCacheAffinityRecord; use super::cache_types::AdminMonitoringCacheAffinityRecord;
use crate::cache::SchedulerAffinityTarget;
use crate::handlers::admin::request::AdminAppState; use crate::handlers::admin::request::AdminAppState;
use crate::GatewayError; use crate::GatewayError;
use std::time::Duration;
fn parse_admin_monitoring_cache_affinity_key(raw_key: &str) -> Option<(String, String, String)> { fn parse_admin_monitoring_cache_affinity_key(raw_key: &str) -> Option<(String, String, String)> {
let parts = raw_key.split(':').collect::<Vec<_>>(); let parts = raw_key.split(':').collect::<Vec<_>>();
@@ -25,6 +27,48 @@ fn parse_admin_monitoring_cache_affinity_key(raw_key: &str) -> Option<(String, S
Some((affinity_key.to_string(), api_format, model_name)) Some((affinity_key.to_string(), api_format, model_name))
} }
fn parse_admin_monitoring_scheduler_affinity_key(
raw_key: &str,
) -> Option<(String, String, String)> {
let parts = raw_key.split(':').collect::<Vec<_>>();
let start = parts
.iter()
.position(|segment| *segment == "scheduler_affinity")?;
let affinity_key = parts.get(start + 1)?.trim();
if affinity_key.is_empty() {
return None;
}
let remaining = parts.get(start + 2..)?;
if remaining.len() < 2 {
return None;
}
let (api_format, model_name_parts) = if remaining.len() == 2 {
(remaining[0].trim().to_string(), &remaining[1..])
} else {
(
format!("{}:{}", remaining[0].trim(), remaining[1].trim()),
&remaining[2..],
)
};
if api_format.trim().is_empty() {
return None;
}
let model_name = model_name_parts
.iter()
.map(|segment| segment.trim())
.filter(|segment| !segment.is_empty())
.collect::<Vec<_>>()
.join(":");
if model_name.is_empty() {
return None;
}
Some((affinity_key.to_string(), api_format, model_name))
}
pub(super) fn admin_monitoring_scheduler_affinity_cache_key( pub(super) fn admin_monitoring_scheduler_affinity_cache_key(
record: &AdminMonitoringCacheAffinityRecord, record: &AdminMonitoringCacheAffinityRecord,
) -> Option<String> { ) -> Option<String> {
@@ -92,6 +136,93 @@ pub(super) fn admin_monitoring_cache_affinity_record(
}) })
} }
pub(super) fn admin_monitoring_scheduler_affinity_record(
cache_key: &str,
target: &SchedulerAffinityTarget,
age: Duration,
ttl: Duration,
now_unix_secs: u64,
) -> Option<AdminMonitoringCacheAffinityRecord> {
let (affinity_key, api_format, model_name) =
parse_admin_monitoring_scheduler_affinity_key(cache_key)?;
let age_secs = age.as_secs();
let created_at = now_unix_secs.saturating_sub(age_secs);
let expire_at = created_at.saturating_add(ttl.as_secs());
Some(AdminMonitoringCacheAffinityRecord {
raw_key: cache_key.to_string(),
affinity_key,
api_format,
model_name,
provider_id: Some(target.provider_id.clone()),
endpoint_id: Some(target.endpoint_id.clone()),
key_id: Some(target.key_id.clone()),
created_at: Some(serde_json::json!(created_at)),
expire_at: Some(serde_json::json!(expire_at)),
request_count: 0,
})
}
pub(super) fn admin_monitoring_scheduler_affinity_record_from_raw(
raw_key: &str,
raw_value: &str,
) -> Option<AdminMonitoringCacheAffinityRecord> {
let payload = serde_json::from_str::<serde_json::Value>(raw_value).ok()?;
let object = payload.as_object()?;
let (affinity_key, parsed_api_format, parsed_model_name) =
parse_admin_monitoring_scheduler_affinity_key(raw_key)?;
let api_format = object
.get("api_format")
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(parsed_api_format.as_str())
.to_string();
let model_name = object
.get("model_name")
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(parsed_model_name.as_str())
.to_string();
let request_count = object
.get("request_count")
.and_then(|value| {
value
.as_u64()
.or_else(|| value.as_i64().and_then(|number| u64::try_from(number).ok()))
})
.unwrap_or(0);
Some(AdminMonitoringCacheAffinityRecord {
raw_key: raw_key.to_string(),
affinity_key,
api_format,
model_name,
provider_id: object
.get("provider_id")
.and_then(serde_json::Value::as_str)
.map(ToOwned::to_owned),
endpoint_id: object
.get("endpoint_id")
.and_then(serde_json::Value::as_str)
.map(ToOwned::to_owned),
key_id: object
.get("key_id")
.and_then(serde_json::Value::as_str)
.map(ToOwned::to_owned),
created_at: object.get("created_at").cloned(),
expire_at: object.get("expire_at").cloned(),
request_count,
})
}
pub(super) fn admin_monitoring_cache_affinity_record_identity(
record: &AdminMonitoringCacheAffinityRecord,
) -> String {
admin_monitoring_scheduler_affinity_cache_key(record).unwrap_or_else(|| record.raw_key.clone())
}
pub(super) fn clear_admin_monitoring_scheduler_affinity_entries( pub(super) fn clear_admin_monitoring_scheduler_affinity_entries(
state: &AdminAppState<'_>, state: &AdminAppState<'_>,
records: &[AdminMonitoringCacheAffinityRecord], records: &[AdminMonitoringCacheAffinityRecord],

View File

@@ -8,6 +8,7 @@ use super::super::cache_route_helpers::{
admin_monitoring_cache_affinity_unavailable_response, admin_monitoring_cache_affinity_unavailable_response,
}; };
use super::super::cache_store::{ use super::super::cache_store::{
admin_monitoring_has_runtime_scheduler_affinity_entries,
list_admin_monitoring_cache_affinity_records_by_affinity_keys, list_admin_monitoring_cache_affinity_records_by_affinity_keys,
load_admin_monitoring_cache_affinity_entries_for_tests, load_admin_monitoring_cache_affinity_entries_for_tests,
}; };
@@ -33,6 +34,7 @@ pub(in super::super) async fn build_admin_monitoring_cache_affinity_delete_respo
if state.redis_kv_runner().is_none() if state.redis_kv_runner().is_none()
&& load_admin_monitoring_cache_affinity_entries_for_tests(state).is_empty() && load_admin_monitoring_cache_affinity_entries_for_tests(state).is_empty()
&& !admin_monitoring_has_runtime_scheduler_affinity_entries(state)
{ {
return Ok(admin_monitoring_cache_affinity_unavailable_response()); return Ok(admin_monitoring_cache_affinity_unavailable_response());
} }

View File

@@ -11,6 +11,7 @@ use super::super::cache_route_helpers::{
admin_monitoring_cache_users_user_identifier_from_path, admin_monitoring_cache_users_user_identifier_from_path,
}; };
use super::super::cache_store::{ use super::super::cache_store::{
admin_monitoring_has_runtime_scheduler_affinity_entries,
list_admin_monitoring_cache_affinity_records_by_affinity_keys, list_admin_monitoring_cache_affinity_records_by_affinity_keys,
load_admin_monitoring_cache_affinity_entries_for_tests, load_admin_monitoring_cache_affinity_entries_for_tests,
}; };
@@ -37,6 +38,7 @@ pub(in super::super) async fn build_admin_monitoring_cache_users_delete_response
if state.redis_kv_runner().is_none() if state.redis_kv_runner().is_none()
&& load_admin_monitoring_cache_affinity_entries_for_tests(state).is_empty() && load_admin_monitoring_cache_affinity_entries_for_tests(state).is_empty()
&& !admin_monitoring_has_runtime_scheduler_affinity_entries(state)
{ {
return Ok(admin_monitoring_cache_affinity_unavailable_response()); return Ok(admin_monitoring_cache_affinity_unavailable_response());
} }

View File

@@ -1,41 +1,25 @@
use super::cache_affinity::admin_monitoring_cache_affinity_record; use super::cache_affinity::{
admin_monitoring_cache_affinity_record, admin_monitoring_cache_affinity_record_identity,
admin_monitoring_scheduler_affinity_record,
admin_monitoring_scheduler_affinity_record_from_raw,
};
use super::cache_types::{AdminMonitoringCacheAffinityRecord, AdminMonitoringCacheSnapshot}; use super::cache_types::{AdminMonitoringCacheAffinityRecord, AdminMonitoringCacheSnapshot};
use crate::handlers::admin::observability::stats::round_to; use crate::handlers::admin::observability::stats::round_to;
use crate::handlers::admin::request::AdminAppState; use crate::handlers::admin::request::AdminAppState;
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
use crate::GatewayError; use crate::GatewayError;
use aether_data_contracts::repository::usage::UsageAuditListQuery; use aether_data_contracts::repository::usage::UsageAuditListQuery;
async fn count_admin_monitoring_cache_affinity_entries(state: &AdminAppState<'_>) -> usize { async fn count_admin_monitoring_cache_affinity_entries(state: &AdminAppState<'_>) -> usize {
let Some(runner) = state.redis_kv_runner() else { list_admin_monitoring_cache_affinity_records(state)
return 0; .await
}; .map(|items| items.len())
let mut connection = match runner.client().get_multiplexed_async_connection().await { .unwrap_or_else(|_| {
Ok(value) => value, state
Err(_) => return 0, .as_ref()
}; .list_scheduler_affinity_entries(SCHEDULER_AFFINITY_TTL)
let pattern = runner.keyspace().key("cache_affinity:*"); .len()
let mut cursor = 0u64; })
let mut total = 0usize;
loop {
let (next_cursor, keys) = match redis::cmd("SCAN")
.arg(cursor)
.arg("MATCH")
.arg(&pattern)
.arg("COUNT")
.arg(200)
.query_async::<(u64, Vec<String>)>(&mut connection)
.await
{
Ok(value) => value,
Err(_) => return total,
};
total += keys.len();
if next_cursor == 0 {
break;
}
cursor = next_cursor;
}
total
} }
async fn scan_admin_monitoring_namespaced_keys( async fn scan_admin_monitoring_namespaced_keys(
@@ -191,12 +175,27 @@ pub(super) async fn list_admin_monitoring_cache_affinity_records_by_affinity_key
list_admin_monitoring_cache_affinity_records_matching(state, Some(affinity_keys)).await list_admin_monitoring_cache_affinity_records_matching(state, Some(affinity_keys)).await
} }
pub(super) fn admin_monitoring_has_runtime_scheduler_affinity_entries(
state: &AdminAppState<'_>,
) -> bool {
!state
.as_ref()
.list_scheduler_affinity_entries(SCHEDULER_AFFINITY_TTL)
.is_empty()
}
async fn list_admin_monitoring_cache_affinity_records_matching( async fn list_admin_monitoring_cache_affinity_records_matching(
state: &AdminAppState<'_>, state: &AdminAppState<'_>,
affinity_keys: Option<&std::collections::BTreeSet<String>>, affinity_keys: Option<&std::collections::BTreeSet<String>>,
) -> Result<Vec<AdminMonitoringCacheAffinityRecord>, GatewayError> { ) -> Result<Vec<AdminMonitoringCacheAffinityRecord>, GatewayError> {
let mut records = Vec::new(); let mut records = Vec::new();
let mut seen_raw_keys = std::collections::BTreeSet::new(); let mut seen_record_ids = std::collections::BTreeSet::new();
let mut push_record = |record: AdminMonitoringCacheAffinityRecord| {
if seen_record_ids.insert(admin_monitoring_cache_affinity_record_identity(&record)) {
records.push(record);
}
};
if let Some(runner) = state.redis_kv_runner() { if let Some(runner) = state.redis_kv_runner() {
let mut connection = runner let mut connection = runner
@@ -209,14 +208,24 @@ async fn list_admin_monitoring_cache_affinity_records_matching(
let patterns = affinity_keys let patterns = affinity_keys
.map(|keys| { .map(|keys| {
keys.iter() keys.iter()
.map(|affinity_key| { .flat_map(|affinity_key| {
runner [
.keyspace() runner
.key(&format!("cache_affinity:{affinity_key}:*")) .keyspace()
.key(&format!("cache_affinity:{affinity_key}:*")),
runner
.keyspace()
.key(&format!("scheduler_affinity:{affinity_key}:*")),
]
}) })
.collect::<Vec<_>>() .collect::<Vec<_>>()
}) })
.unwrap_or_else(|| vec![runner.keyspace().key("cache_affinity:*")]); .unwrap_or_else(|| {
vec![
runner.keyspace().key("cache_affinity:*"),
runner.keyspace().key("scheduler_affinity:*"),
]
});
for pattern in patterns { for pattern in patterns {
let mut cursor = 0u64; let mut cursor = 0u64;
@@ -246,16 +255,18 @@ async fn list_admin_monitoring_cache_affinity_records_matching(
let Some(raw_value) = raw_value else { let Some(raw_value) = raw_value else {
continue; continue;
}; };
let Some(record) = admin_monitoring_cache_affinity_record(&key, &raw_value) let record = if key.contains("scheduler_affinity:") {
else { admin_monitoring_scheduler_affinity_record_from_raw(&key, &raw_value)
} else {
admin_monitoring_cache_affinity_record(&key, &raw_value)
};
let Some(record) = record else {
continue; continue;
}; };
if affinity_keys.is_some_and(|keys| !keys.contains(&record.affinity_key)) { if affinity_keys.is_some_and(|keys| !keys.contains(&record.affinity_key)) {
continue; continue;
} }
if seen_raw_keys.insert(record.raw_key.clone()) { push_record(record);
records.push(record);
}
} }
} }
if next_cursor == 0 { if next_cursor == 0 {
@@ -264,7 +275,6 @@ async fn list_admin_monitoring_cache_affinity_records_matching(
cursor = next_cursor; cursor = next_cursor;
} }
} }
return Ok(records);
} }
for (key, raw_value) in load_admin_monitoring_cache_affinity_entries_for_tests(state) { for (key, raw_value) in load_admin_monitoring_cache_affinity_entries_for_tests(state) {
@@ -274,9 +284,27 @@ async fn list_admin_monitoring_cache_affinity_records_matching(
if affinity_keys.is_some_and(|keys| !keys.contains(&record.affinity_key)) { if affinity_keys.is_some_and(|keys| !keys.contains(&record.affinity_key)) {
continue; continue;
} }
if seen_raw_keys.insert(record.raw_key.clone()) { push_record(record);
records.push(record); }
let now_unix_secs = chrono::Utc::now().timestamp().max(0) as u64;
for entry in state
.as_ref()
.list_scheduler_affinity_entries(SCHEDULER_AFFINITY_TTL)
{
let Some(record) = admin_monitoring_scheduler_affinity_record(
&entry.cache_key,
&entry.target,
entry.age,
SCHEDULER_AFFINITY_TTL,
now_unix_secs,
) else {
continue;
};
if affinity_keys.is_some_and(|keys| !keys.contains(&record.affinity_key)) {
continue;
} }
push_record(record);
} }
Ok(records) Ok(records)

View File

@@ -161,7 +161,7 @@ pub(super) async fn build_admin_monitoring_resilience_snapshot(
.filter(admin_monitoring_usage_is_error) .filter(admin_monitoring_usage_is_error)
.collect::<Vec<_>>(); .collect::<Vec<_>>();
recent_usage_errors recent_usage_errors
.sort_by(|left, right| right.created_at_unix_secs.cmp(&left.created_at_unix_secs)); .sort_by(|left, right| right.created_at_unix_ms.cmp(&left.created_at_unix_ms));
let total_errors = recent_usage_errors.len(); let total_errors = recent_usage_errors.len();
let mut error_breakdown = BTreeMap::<String, usize>::new(); let mut error_breakdown = BTreeMap::<String, usize>::new();
@@ -200,7 +200,7 @@ pub(super) async fn build_admin_monitoring_resilience_snapshot(
"error_id": item.id, "error_id": item.id,
"error_type": error_type, "error_type": error_type,
"operation": operation, "operation": operation,
"timestamp": unix_secs_to_rfc3339(item.created_at_unix_secs), "timestamp": unix_secs_to_rfc3339(item.created_at_unix_ms),
"context": { "context": {
"request_id": item.request_id, "request_id": item.request_id,
"provider_id": item.provider_id, "provider_id": item.provider_id,

View File

@@ -31,7 +31,7 @@ pub(super) fn sample_usage(
total_cost_usd: f64, total_cost_usd: f64,
status: &str, status: &str,
status_code: Option<i32>, status_code: Option<i32>,
created_at_unix_secs: i64, created_at_unix_ms: i64,
) -> StoredRequestUsageAudit { ) -> StoredRequestUsageAudit {
let is_error = status_code.is_some_and(|value| value >= 400) let is_error = status_code.is_some_and(|value| value >= 400)
|| status.trim().eq_ignore_ascii_case("failed") || status.trim().eq_ignore_ascii_case("failed")
@@ -70,9 +70,9 @@ pub(super) fn sample_usage(
Some(30), Some(30),
status.to_string(), status.to_string(),
"billed".to_string(), "billed".to_string(),
created_at_unix_secs, created_at_unix_ms,
created_at_unix_secs, created_at_unix_ms,
Some(created_at_unix_secs), Some(created_at_unix_ms),
) )
.expect("usage should build") .expect("usage should build")
} }
@@ -82,7 +82,7 @@ pub(super) fn sample_candidate(
request_id: &str, request_id: &str,
candidate_index: i32, candidate_index: i32,
status: RequestCandidateStatus, status: RequestCandidateStatus,
started_at_unix_secs: Option<i64>, started_at_unix_ms: Option<i64>,
latency_ms: Option<i32>, latency_ms: Option<i32>,
status_code: Option<i32>, status_code: Option<i32>,
) -> StoredRequestCandidate { ) -> StoredRequestCandidate {
@@ -108,9 +108,9 @@ pub(super) fn sample_candidate(
Some(1), Some(1),
None, None,
Some(json!({"cache_1h": true})), Some(json!({"cache_1h": true})),
100 + i64::from(candidate_index), (100 + i64::from(candidate_index)) * 1_000,
started_at_unix_secs, started_at_unix_ms.map(|v| v * 1_000),
started_at_unix_secs.map(|value| value + 1), started_at_unix_ms.map(|value| (value + 1) * 1_000),
) )
.expect("candidate should build") .expect("candidate should build")
} }

View File

@@ -258,6 +258,56 @@ async fn admin_monitoring_resilience_status_returns_local_payload() {
assert!(payload["timestamp"].as_str().is_some()); assert!(payload["timestamp"].as_str().is_some());
} }
#[tokio::test]
async fn admin_monitoring_cache_stats_count_runtime_scheduler_affinities() {
let state = AppState::new().expect("state should build");
let affinity_cache_key =
aether_scheduler_core::build_scheduler_affinity_cache_key_for_api_key_id(
"user-key-1",
"openai:chat",
"model-alpha",
)
.expect("scheduler affinity cache key should build");
state.scheduler_affinity_cache.insert(
affinity_cache_key,
crate::cache::SchedulerAffinityTarget {
provider_id: "provider-1".to_string(),
endpoint_id: "endpoint-1".to_string(),
key_id: "provider-key-1".to_string(),
},
crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL,
128,
);
let response = local_monitoring_response(
&state,
&request_context(http::Method::GET, "/api/admin/monitoring/cache/stats"),
)
.await
.expect("handler should not error")
.expect("monitoring route should be handled locally");
assert_eq!(response.status(), http::StatusCode::OK);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read");
let payload: serde_json::Value = serde_json::from_slice(&body).expect("json body should parse");
assert_eq!(payload["status"], json!("ok"));
assert_eq!(payload["data"]["total_affinities"], json!(1));
assert_eq!(
payload["data"]["affinity_stats"]["total_affinities"],
json!(1)
);
assert_eq!(
payload["data"]["affinity_stats"]["active_affinities"],
json!(1)
);
assert_eq!(
payload["data"]["affinity_stats"]["storage_type"],
json!("memory")
);
}
#[tokio::test] #[tokio::test]
async fn admin_monitoring_cache_stats_returns_local_payload() { async fn admin_monitoring_cache_stats_returns_local_payload() {
let now = chrono::Utc::now().timestamp(); let now = chrono::Utc::now().timestamp();

View File

@@ -209,6 +209,136 @@ async fn admin_monitoring_cache_affinities_and_affinity_return_local_payload_fro
assert_eq!(detail_payload["total_endpoints"], json!(1)); assert_eq!(detail_payload["total_endpoints"], json!(1));
} }
#[tokio::test]
async fn admin_monitoring_cache_affinities_and_delete_use_runtime_scheduler_affinity_cache() {
let user_repository = Arc::new(
InMemoryUserReadRepository::seed_auth_users(vec![sample_monitoring_auth_user("user-1")])
.with_export_users(vec![sample_monitoring_export_user("user-1")]),
);
let auth_repository = Arc::new(
InMemoryAuthApiKeySnapshotRepository::default().with_export_records(vec![
sample_monitoring_export_api_key("user-1", "user-key-1"),
]),
);
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider()],
vec![sample_monitoring_catalog_endpoint()],
vec![sample_monitoring_catalog_key()],
));
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_provider_catalog_reader_for_tests(provider_catalog)
.with_user_reader(user_repository)
.with_auth_api_key_reader(auth_repository),
);
let affinity_cache_key =
aether_scheduler_core::build_scheduler_affinity_cache_key_for_api_key_id(
"user-key-1",
"openai:chat",
"model-alpha",
)
.expect("scheduler affinity cache key should build");
state.scheduler_affinity_cache.insert(
affinity_cache_key.clone(),
crate::cache::SchedulerAffinityTarget {
provider_id: "provider-1".to_string(),
endpoint_id: "endpoint-1".to_string(),
key_id: "provider-key-1".to_string(),
},
crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL,
128,
);
let list_context = request_context(
http::Method::GET,
"/api/admin/monitoring/cache/affinities?keyword=alice&limit=20&offset=0",
);
let list_response = local_monitoring_response(&state, &list_context)
.await
.expect("handler should not error")
.expect("route should be handled locally");
assert_eq!(list_response.status(), http::StatusCode::OK);
let list_body = to_bytes(list_response.into_body(), usize::MAX)
.await
.expect("body should read");
let list_payload: serde_json::Value =
serde_json::from_slice(&list_body).expect("json body should parse");
assert_eq!(list_payload["status"], json!("ok"));
assert_eq!(list_payload["data"]["meta"]["total"], json!(1));
assert_eq!(list_payload["data"]["matched_user_id"], json!("user-1"));
assert_eq!(
list_payload["data"]["items"][0]["affinity_key"],
json!("user-key-1")
);
assert_eq!(
list_payload["data"]["items"][0]["api_format"],
json!("openai:chat")
);
assert_eq!(
list_payload["data"]["items"][0]["provider_name"],
json!("OpenAI")
);
assert_eq!(
list_payload["data"]["items"][0]["endpoint_url"],
json!("https://api.openai.example/v1")
);
assert_eq!(list_payload["data"]["items"][0]["request_count"], json!(0));
assert!(list_payload["data"]["items"][0]["expire_at"]
.as_u64()
.is_some_and(|value| value > 0));
let detail_context = request_context(
http::Method::GET,
"/api/admin/monitoring/cache/affinity/alice",
);
let detail_response = local_monitoring_response(&state, &detail_context)
.await
.expect("handler should not error")
.expect("route should be handled locally");
assert_eq!(detail_response.status(), http::StatusCode::OK);
let detail_body = to_bytes(detail_response.into_body(), usize::MAX)
.await
.expect("body should read");
let detail_payload: serde_json::Value =
serde_json::from_slice(&detail_body).expect("json body should parse");
assert_eq!(detail_payload["status"], json!("ok"));
assert_eq!(
detail_payload["affinities"][0]["api_format"],
json!("openai:chat")
);
assert_eq!(detail_payload["total_endpoints"], json!(1));
let delete_response = local_monitoring_response(
&state,
&request_context(
http::Method::DELETE,
"/api/admin/monitoring/cache/affinity/user-key-1/endpoint-1/model-alpha/openai:chat",
),
)
.await
.expect("handler should not error")
.expect("route should be handled locally");
assert_eq!(delete_response.status(), http::StatusCode::OK);
let delete_body = to_bytes(delete_response.into_body(), usize::MAX)
.await
.expect("body should read");
let delete_payload: serde_json::Value =
serde_json::from_slice(&delete_body).expect("json body should parse");
assert_eq!(
delete_payload["message"],
json!("已清除缓存亲和性: Alice Key")
);
assert_eq!(delete_payload["affinity_key"], json!("user-key-1"));
assert_eq!(
state.read_scheduler_affinity_target(
&affinity_cache_key,
crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL,
),
None
);
}
#[tokio::test] #[tokio::test]
async fn admin_monitoring_cache_users_delete_returns_local_payload_from_test_store() { async fn admin_monitoring_cache_users_delete_returns_local_payload_from_test_store() {
let user_repository = Arc::new( let user_repository = Arc::new(

View File

@@ -13,7 +13,7 @@ use axum::{
body::Body, body::Body,
response::{IntoResponse, Response}, response::{IntoResponse, Response},
}; };
use tracing::warn; use tracing::debug;
pub(super) async fn build_admin_monitoring_trace_request_response( pub(super) async fn build_admin_monitoring_trace_request_response(
state: &AdminAppState<'_>, state: &AdminAppState<'_>,
@@ -38,7 +38,7 @@ pub(super) async fn build_admin_monitoring_trace_request_response(
.await .await
.map_err(|err| GatewayError::Internal(err.to_string()))? .map_err(|err| GatewayError::Internal(err.to_string()))?
else { else {
warn!( debug!(
event_name = "admin_monitoring_request_trace_not_found", event_name = "admin_monitoring_request_trace_not_found",
log_type = "admin_monitoring", log_type = "admin_monitoring",
request_id = %short_request_id(request_id.as_str()), request_id = %short_request_id(request_id.as_str()),

View File

@@ -20,8 +20,8 @@ pub(in super::super) async fn list_recent_completed_usage_for_cache_affinity(
.await?; .await?;
items.retain(|item| item.status == "completed"); items.retain(|item| item.status == "completed");
items.sort_by(|left, right| { items.sort_by(|left, right| {
left.created_at_unix_secs left.created_at_unix_ms
.cmp(&right.created_at_unix_secs) .cmp(&right.created_at_unix_ms)
.then_with(|| left.id.cmp(&right.id)) .then_with(|| left.id.cmp(&right.id))
}); });
Ok(items) Ok(items)

View File

@@ -50,14 +50,14 @@ pub(super) async fn build_admin_usage_cache_affinity_interval_timeline_response(
} }
} }
let mut previous_created_at_unix_secs = None; let mut previous_created_at_unix_ms = None;
for item in items { for item in items {
if let Some(previous) = previous_created_at_unix_secs { if let Some(previous) = previous_created_at_unix_ms {
let interval_minutes = let interval_minutes =
item.created_at_unix_secs.saturating_sub(previous) as f64 / 60.0; item.created_at_unix_ms.saturating_sub(previous) as f64 / 60.0;
if interval_minutes <= 120.0 { if interval_minutes <= 120.0 {
let mut point = json!({ let mut point = json!({
"x": unix_secs_to_rfc3339(item.created_at_unix_secs), "x": unix_secs_to_rfc3339(item.created_at_unix_ms),
"y": ((interval_minutes * 100.0).round()) / 100.0, "y": ((interval_minutes * 100.0).round()) / 100.0,
}); });
if !item.model.trim().is_empty() { if !item.model.trim().is_empty() {
@@ -78,7 +78,7 @@ pub(super) async fn build_admin_usage_cache_affinity_interval_timeline_response(
.push(point); .push(point);
} }
} }
previous_created_at_unix_secs = Some(item.created_at_unix_secs); previous_created_at_unix_ms = Some(item.created_at_unix_ms);
} }
} }

View File

@@ -78,8 +78,8 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
.collect(); .collect();
items.sort_by(|left, right| { items.sort_by(|left, right| {
right right
.created_at_unix_secs .created_at_unix_ms
.cmp(&left.created_at_unix_secs) .cmp(&left.created_at_unix_ms)
.then_with(|| left.id.cmp(&right.id)) .then_with(|| left.id.cmp(&right.id))
}); });
if requested_ids.is_none() && items.len() > 50 { if requested_ids.is_none() && items.len() > 50 {
@@ -147,8 +147,8 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
}); });
usage.sort_by(|left, right| { usage.sort_by(|left, right| {
right right
.created_at_unix_secs .created_at_unix_ms
.cmp(&left.created_at_unix_secs) .cmp(&left.created_at_unix_ms)
.then_with(|| left.id.cmp(&right.id)) .then_with(|| left.id.cmp(&right.id))
}); });
let total = usage.len(); let total = usage.len();

View File

@@ -28,9 +28,9 @@ pub(crate) async fn build_admin_provider_endpoints_payload(
.unwrap_or_default(); .unwrap_or_default();
endpoints.sort_by(|left, right| { endpoints.sort_by(|left, right| {
right right
.created_at_unix_secs .created_at_unix_ms
.unwrap_or_default() .unwrap_or_default()
.cmp(&left.created_at_unix_secs.unwrap_or_default()) .cmp(&left.created_at_unix_ms.unwrap_or_default())
.then_with(|| left.id.cmp(&right.id)) .then_with(|| left.id.cmp(&right.id))
}); });
let keys = state let keys = state

View File

@@ -165,7 +165,7 @@ pub(super) async fn handle_admin_provider_oauth_device_authorize(
.map(str::trim) .map(str::trim)
.filter(|value| !value.is_empty()) .filter(|value| !value.is_empty())
.map(ToOwned::to_owned), .map(ToOwned::to_owned),
created_at_unix_secs: now_unix_secs, created_at_unix_ms: now_unix_secs,
key_id: None, key_id: None,
email: None, email: None,
replaced: false, replaced: false,

View File

@@ -128,7 +128,7 @@ pub(crate) async fn create_provider_oauth_catalog_key(
record.total_response_time_ms = Some(0); record.total_response_time_ms = Some(0);
record.health_by_format = Some(json!({})); record.health_by_format = Some(json!({}));
record.circuit_breaker_by_format = Some(json!({})); record.circuit_breaker_by_format = Some(json!({}));
record.created_at_unix_secs = Some(now_unix_secs); record.created_at_unix_ms = Some(now_unix_secs);
record.updated_at_unix_secs = Some(now_unix_secs); record.updated_at_unix_secs = Some(now_unix_secs);
state.create_provider_catalog_key(&record).await state.create_provider_catalog_key(&record).await
} }

View File

@@ -803,7 +803,7 @@ pub(super) fn build_admin_pool_key_payload(
); );
payload.insert( payload.insert(
"created_at".to_string(), "created_at".to_string(),
json!(key.created_at_unix_secs.and_then(unix_secs_to_rfc3339)), json!(key.created_at_unix_ms.and_then(unix_secs_to_rfc3339)),
); );
payload.insert( payload.insert(
"last_used_at".to_string(), "last_used_at".to_string(),

View File

@@ -177,7 +177,7 @@ pub(crate) async fn build_admin_providers_summary_payload(
.is_active .is_active
.cmp(&left.is_active) .cmp(&left.is_active)
.then_with(|| left.provider_priority.cmp(&right.provider_priority)) .then_with(|| left.provider_priority.cmp(&right.provider_priority))
.then_with(|| left.created_at_unix_secs.cmp(&right.created_at_unix_secs)) .then_with(|| left.created_at_unix_ms.cmp(&right.created_at_unix_ms))
}); });
let total = providers.len(); let total = providers.len();

View File

@@ -1,6 +1,7 @@
use crate::handlers::admin::request::AdminAppState; use crate::handlers::admin::request::AdminAppState;
use crate::handlers::admin::shared::unix_secs_to_rfc3339; use crate::handlers::admin::shared::unix_secs_to_rfc3339;
use crate::handlers::public::{request_candidate_event_unix_secs, request_candidate_status_label}; use crate::handlers::public::{request_candidate_event_unix_ms, request_candidate_status_label};
use crate::handlers::shared::unix_ms_to_rfc3339;
use aether_data_contracts::repository::candidates::{ use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, StoredRequestCandidate, RequestCandidateStatus, StoredRequestCandidate,
}; };
@@ -82,14 +83,14 @@ pub(crate) async fn build_admin_provider_health_monitor_payload(
for candidates in attempts_by_endpoint.values_mut() { for candidates in attempts_by_endpoint.values_mut() {
candidates.sort_by(|left, right| { candidates.sort_by(|left, right| {
right right
.created_at_unix_secs .created_at_unix_ms
.cmp(&left.created_at_unix_secs) .cmp(&left.created_at_unix_ms)
.then_with(|| right.id.cmp(&left.id)) .then_with(|| right.id.cmp(&left.id))
}); });
candidates.truncate(per_endpoint_limit); candidates.truncate(per_endpoint_limit);
candidates.sort_by(|left, right| { candidates.sort_by(|left, right| {
request_candidate_event_unix_secs(left) request_candidate_event_unix_ms(left)
.cmp(&request_candidate_event_unix_secs(right)) .cmp(&request_candidate_event_unix_ms(right))
.then_with(|| left.id.cmp(&right.id)) .then_with(|| left.id.cmp(&right.id))
}); });
} }
@@ -118,12 +119,12 @@ pub(crate) async fn build_admin_provider_health_monitor_payload(
}; };
let last_event_at = candidates let last_event_at = candidates
.last() .last()
.and_then(|candidate| unix_secs_to_rfc3339(request_candidate_event_unix_secs(candidate))); .and_then(|candidate| unix_ms_to_rfc3339(request_candidate_event_unix_ms(candidate)));
let events = candidates let events = candidates
.into_iter() .into_iter()
.filter_map(|candidate| { .filter_map(|candidate| {
Some(json!({ Some(json!({
"timestamp": unix_secs_to_rfc3339(request_candidate_event_unix_secs(&candidate))?, "timestamp": unix_ms_to_rfc3339(request_candidate_event_unix_ms(&candidate))?,
"status": request_candidate_status_label(candidate.status), "status": request_candidate_status_label(candidate.status),
"status_code": candidate.status_code, "status_code": candidate.status_code,
"latency_ms": candidate.latency_ms, "latency_ms": candidate.latency_ms,

View File

@@ -88,7 +88,7 @@ pub(crate) async fn build_admin_providers_payload(
"api_key": has_any_key_by_provider.contains(&provider_id).then_some("***"), "api_key": has_any_key_by_provider.contains(&provider_id).then_some("***"),
"priority": provider.provider_priority, "priority": provider.provider_priority,
"is_active": provider.is_active, "is_active": provider.is_active,
"created_at": provider.created_at_unix_secs.and_then(unix_secs_to_rfc3339), "created_at": provider.created_at_unix_ms.and_then(unix_secs_to_rfc3339),
"updated_at": provider.updated_at_unix_secs.and_then(unix_secs_to_rfc3339), "updated_at": provider.updated_at_unix_secs.and_then(unix_secs_to_rfc3339),
}) })
}) })

View File

@@ -1,6 +1,6 @@
use crate::handlers::admin::shared::unix_secs_to_rfc3339; use crate::handlers::admin::shared::unix_secs_to_rfc3339;
use crate::handlers::public::{ use crate::handlers::public::{
provider_key_api_formats, request_candidate_event_unix_secs, request_candidate_status_label, provider_key_api_formats, request_candidate_event_unix_ms, request_candidate_status_label,
}; };
use aether_data_contracts::repository::candidates::{ use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, StoredRequestCandidate, RequestCandidateStatus, StoredRequestCandidate,
@@ -172,7 +172,7 @@ pub(crate) fn build_admin_provider_summary_value(
"endpoint_health_details": endpoint_health_details, "endpoint_health_details": endpoint_health_details,
"ops_configured": ops_configured, "ops_configured": ops_configured,
"ops_architecture_id": ops_architecture_id, "ops_architecture_id": ops_architecture_id,
"created_at": endpoint_timestamp_or_now(provider.created_at_unix_secs, now_unix_secs), "created_at": endpoint_timestamp_or_now(provider.created_at_unix_ms, now_unix_secs),
"updated_at": endpoint_timestamp_or_now(provider.updated_at_unix_secs, now_unix_secs), "updated_at": endpoint_timestamp_or_now(provider.updated_at_unix_secs, now_unix_secs),
}) })
} }

View File

@@ -176,7 +176,7 @@ pub(crate) async fn build_admin_create_provider_key_record(
normalize_string_list(payload.model_exclude_patterns).map(|value| json!(value)); normalize_string_list(payload.model_exclude_patterns).map(|value| json!(value));
key.health_by_format = Some(json!({})); key.health_by_format = Some(json!({}));
key.circuit_breaker_by_format = Some(json!({})); key.circuit_breaker_by_format = Some(json!({}));
key.created_at_unix_secs = Some(now_unix_secs); key.created_at_unix_ms = Some(now_unix_secs);
key.updated_at_unix_secs = Some(now_unix_secs); key.updated_at_unix_secs = Some(now_unix_secs);
Ok(key) Ok(key)
} }

View File

@@ -24,9 +24,9 @@ pub(crate) async fn build_admin_provider_keys_payload(
left.internal_priority left.internal_priority
.cmp(&right.internal_priority) .cmp(&right.internal_priority)
.then_with(|| { .then_with(|| {
left.created_at_unix_secs left.created_at_unix_ms
.unwrap_or_default() .unwrap_or_default()
.cmp(&right.created_at_unix_secs.unwrap_or_default()) .cmp(&right.created_at_unix_ms.unwrap_or_default())
}) })
.then_with(|| left.id.cmp(&right.id)) .then_with(|| left.id.cmp(&right.id))
}); });

View File

@@ -82,9 +82,9 @@ pub(crate) async fn build_admin_system_export_providers_payload(
left.internal_priority left.internal_priority
.cmp(&right.internal_priority) .cmp(&right.internal_priority)
.then( .then(
left.created_at_unix_secs left.created_at_unix_ms
.unwrap_or(0) .unwrap_or(0)
.cmp(&right.created_at_unix_secs.unwrap_or(0)), .cmp(&right.created_at_unix_ms.unwrap_or(0)),
) )
.then(left.id.cmp(&right.id)) .then(left.id.cmp(&right.id))
}); });

View File

@@ -463,7 +463,7 @@ pub(crate) fn build_management_token_payload(
"last_used_ip": token.last_used_ip, "last_used_ip": token.last_used_ip,
"usage_count": token.usage_count, "usage_count": token.usage_count,
"is_active": token.is_active, "is_active": token.is_active,
"created_at": token.created_at_unix_secs.and_then(unix_secs_to_rfc3339), "created_at": token.created_at_unix_ms.and_then(unix_secs_to_rfc3339),
"updated_at": token.updated_at_unix_secs.and_then(unix_secs_to_rfc3339), "updated_at": token.updated_at_unix_secs.and_then(unix_secs_to_rfc3339),
}); });
if let Some(user) = user { if let Some(user) = user {

View File

@@ -1,5 +1,7 @@
use crate::api::ai::public_api_format_local_path; use crate::api::ai::public_api_format_local_path;
use crate::handlers::shared::{query_param_optional_bool, query_param_value, unix_secs_to_rfc3339}; use crate::handlers::shared::{
query_param_optional_bool, query_param_value, unix_ms_to_rfc3339, unix_secs_to_rfc3339,
};
use crate::AppState; use crate::AppState;
use aether_data_contracts::repository::candidates::{ use aether_data_contracts::repository::candidates::{
PublicHealthTimelineBucket, RequestCandidateStatus, StoredRequestCandidate, PublicHealthTimelineBucket, RequestCandidateStatus, StoredRequestCandidate,
@@ -25,11 +27,11 @@ pub(crate) fn request_candidate_status_label(status: RequestCandidateStatus) ->
} }
} }
pub(crate) fn request_candidate_event_unix_secs(candidate: &StoredRequestCandidate) -> u64 { pub(crate) fn request_candidate_event_unix_ms(candidate: &StoredRequestCandidate) -> u64 {
candidate candidate
.finished_at_unix_secs .finished_at_unix_ms
.or(candidate.started_at_unix_secs) .or(candidate.started_at_unix_ms)
.unwrap_or(candidate.created_at_unix_secs) .unwrap_or(candidate.created_at_unix_ms)
} }
pub(crate) fn normalize_admin_base_url(base_url: &str) -> Result<String, String> { pub(crate) fn normalize_admin_base_url(base_url: &str) -> Result<String, String> {
@@ -410,28 +412,24 @@ pub(crate) async fn build_api_format_health_monitor_payload(
total_count: 0, total_count: 0,
success_count: 0, success_count: 0,
failed_count: 0, failed_count: 0,
min_created_at_unix_secs: None, min_created_at_unix_ms: None,
max_created_at_unix_secs: None, max_created_at_unix_ms: None,
}); });
bucket.total_count += row.total_count; bucket.total_count += row.total_count;
bucket.success_count += row.success_count; bucket.success_count += row.success_count;
bucket.failed_count += row.failed_count; bucket.failed_count += row.failed_count;
bucket.min_created_at_unix_secs = match ( bucket.min_created_at_unix_ms =
bucket.min_created_at_unix_secs, match (bucket.min_created_at_unix_ms, row.min_created_at_unix_ms) {
row.min_created_at_unix_secs, (Some(left), Some(right)) => Some(left.min(right)),
) { (None, Some(right)) => Some(right),
(Some(left), Some(right)) => Some(left.min(right)), (left, None) => left,
(None, Some(right)) => Some(right), };
(left, None) => left, bucket.max_created_at_unix_ms =
}; match (bucket.max_created_at_unix_ms, row.max_created_at_unix_ms) {
bucket.max_created_at_unix_secs = match ( (Some(left), Some(right)) => Some(left.max(right)),
bucket.max_created_at_unix_secs, (None, Some(right)) => Some(right),
row.max_created_at_unix_secs, (left, None) => left,
) { };
(Some(left), Some(right)) => Some(left.max(right)),
(None, Some(right)) => Some(right),
(left, None) => left,
};
} }
let mut formats = Vec::new(); let mut formats = Vec::new();
@@ -456,19 +454,19 @@ pub(crate) async fn build_api_format_health_monitor_payload(
}; };
let last_event_at = attempts.first().and_then(|candidate| { let last_event_at = attempts.first().and_then(|candidate| {
candidate candidate
.finished_at_unix_secs .finished_at_unix_ms
.or(candidate.started_at_unix_secs) .or(candidate.started_at_unix_ms)
.or(Some(candidate.created_at_unix_secs)) .or(Some(candidate.created_at_unix_ms))
}); });
let events = attempts let events = attempts
.into_iter() .into_iter()
.filter_map(|candidate| { .filter_map(|candidate| {
let timestamp = candidate let timestamp = candidate
.finished_at_unix_secs .finished_at_unix_ms
.or(candidate.started_at_unix_secs) .or(candidate.started_at_unix_ms)
.unwrap_or(candidate.created_at_unix_secs); .unwrap_or(candidate.created_at_unix_ms);
Some(json!({ Some(json!({
"timestamp": unix_secs_to_rfc3339(timestamp)?, "timestamp": unix_ms_to_rfc3339(timestamp)?,
"status": request_candidate_status_label(candidate.status), "status": request_candidate_status_label(candidate.status),
"status_code": candidate.status_code, "status_code": candidate.status_code,
"latency_ms": candidate.latency_ms, "latency_ms": candidate.latency_ms,
@@ -490,11 +488,11 @@ pub(crate) async fn build_api_format_health_monitor_payload(
"failed_count": failed_count, "failed_count": failed_count,
"skipped_count": skipped_count, "skipped_count": skipped_count,
"success_rate": success_rate, "success_rate": success_rate,
"last_event_at": last_event_at.and_then(unix_secs_to_rfc3339), "last_event_at": last_event_at.and_then(unix_ms_to_rfc3339),
"events": events, "events": events,
"timeline": timeline, "timeline": timeline,
"time_range_start": time_range_start.and_then(unix_secs_to_rfc3339), "time_range_start": time_range_start.and_then(unix_ms_to_rfc3339),
"time_range_end": time_range_end.or(Some(now_unix_secs)).and_then(unix_secs_to_rfc3339), "time_range_end": time_range_end.map(|ms| unix_ms_to_rfc3339(ms)).unwrap_or_else(|| unix_secs_to_rfc3339(now_unix_secs)),
}); });
if options.include_api_path { if options.include_api_path {
format_payload["api_path"] = json!(public_api_format_local_path(&api_format)); format_payload["api_path"] = json!(public_api_format_local_path(&api_format));
@@ -536,12 +534,12 @@ pub(crate) fn build_public_health_timeline(
continue; continue;
} }
earliest_time = match (earliest_time, bucket.min_created_at_unix_secs) { earliest_time = match (earliest_time, bucket.min_created_at_unix_ms) {
(Some(left), Some(right)) => Some(left.min(right)), (Some(left), Some(right)) => Some(left.min(right)),
(None, Some(right)) => Some(right), (None, Some(right)) => Some(right),
(left, None) => left, (left, None) => left,
}; };
latest_time = match (latest_time, bucket.max_created_at_unix_secs) { latest_time = match (latest_time, bucket.max_created_at_unix_ms) {
(Some(left), Some(right)) => Some(left.max(right)), (Some(left), Some(right)) => Some(left.max(right)),
(None, Some(right)) => Some(right), (None, Some(right)) => Some(right),
(left, None) => left, (left, None) => left,
@@ -592,3 +590,50 @@ pub(crate) fn api_format_display_name(api_format: &str) -> String {
}; };
format!("{family_label} {kind_label}") format!("{family_label} {kind_label}")
} }
#[cfg(test)]
mod tests {
use super::request_candidate_event_unix_ms;
use crate::handlers::shared::unix_ms_to_rfc3339;
use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, StoredRequestCandidate,
};
#[test]
fn request_candidate_event_timestamp_uses_millisecond_precision() {
let candidate = StoredRequestCandidate::new(
"cand-1".to_string(),
"req-1".to_string(),
None,
None,
None,
None,
0,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("key-1".to_string()),
RequestCandidateStatus::Success,
None,
false,
Some(200),
None,
None,
Some(42),
Some(1),
None,
None,
1_700_000_000_000,
Some(1_700_000_000_111),
Some(1_700_000_000_123),
)
.expect("candidate should build");
let event_unix_ms = request_candidate_event_unix_ms(&candidate);
assert_eq!(event_unix_ms, 1_700_000_000_123);
assert_eq!(
unix_ms_to_rfc3339(event_unix_ms).as_deref(),
Some("2023-11-14T22:13:20.123Z")
);
}
}

View File

@@ -10,7 +10,7 @@ pub(crate) use self::catalog_helpers::{
admin_requested_force_stream, api_format_display_name, build_api_format_health_monitor_payload, admin_requested_force_stream, api_format_display_name, build_api_format_health_monitor_payload,
build_public_catalog_models_payload, build_public_catalog_search_models_payload, build_public_catalog_models_payload, build_public_catalog_search_models_payload,
build_public_health_timeline, build_public_providers_payload, normalize_admin_base_url, build_public_health_timeline, build_public_providers_payload, normalize_admin_base_url,
provider_key_api_formats, request_candidate_event_unix_secs, request_candidate_status_label, provider_key_api_formats, request_candidate_event_unix_ms, request_candidate_status_label,
ApiFormatHealthMonitorOptions, ApiFormatHealthMonitorOptions,
}; };
pub(crate) use self::system_modules_helpers::{ pub(crate) use self::system_modules_helpers::{

View File

@@ -84,7 +84,7 @@ pub(super) fn build_public_announcement_payload(
}, },
"start_time": format_optional_unix_datetime(announcement.start_time_unix_secs), "start_time": format_optional_unix_datetime(announcement.start_time_unix_secs),
"end_time": format_optional_unix_datetime(announcement.end_time_unix_secs), "end_time": format_optional_unix_datetime(announcement.end_time_unix_secs),
"created_at": format_required_unix_datetime(announcement.created_at_unix_secs), "created_at": format_required_unix_datetime(announcement.created_at_unix_ms),
"updated_at": format_required_unix_datetime(announcement.updated_at_unix_secs), "updated_at": format_required_unix_datetime(announcement.updated_at_unix_secs),
}) })
} }

View File

@@ -361,7 +361,7 @@ fn dashboard_usage_local_date(
item: &StoredRequestUsageAudit, item: &StoredRequestUsageAudit,
tz_offset_minutes: i32, tz_offset_minutes: i32,
) -> Option<chrono::NaiveDate> { ) -> Option<chrono::NaiveDate> {
let timestamp = i64::try_from(item.created_at_unix_secs).ok()?; let timestamp = i64::try_from(item.created_at_unix_ms).ok()?;
let datetime = chrono::DateTime::<chrono::Utc>::from_timestamp(timestamp, 0)?; let datetime = chrono::DateTime::<chrono::Utc>::from_timestamp(timestamp, 0)?;
Some((datetime + chrono::Duration::minutes(i64::from(tz_offset_minutes))).date_naive()) Some((datetime + chrono::Duration::minutes(i64::from(tz_offset_minutes))).date_naive())
} }
@@ -1008,7 +1008,7 @@ pub(super) async fn handle_dashboard_recent_requests_get(
"user": username, "user": username,
"model": dashboard_non_empty_value(&item.model, "N/A"), "model": dashboard_non_empty_value(&item.model, "N/A"),
"tokens": item.total_tokens, "tokens": item.total_tokens,
"time": dashboard_format_time_hhmm(item.created_at_unix_secs), "time": dashboard_format_time_hhmm(item.created_at_unix_ms),
"is_stream": item.is_stream, "is_stream": item.is_stream,
}) })
}) })

View File

@@ -95,7 +95,7 @@ fn users_me_usage_total_input_context(item: &StoredRequestUsageAudit) -> u64 {
fn users_me_usage_effective_unix_secs(item: &StoredRequestUsageAudit) -> u64 { fn users_me_usage_effective_unix_secs(item: &StoredRequestUsageAudit) -> u64 {
item.finalized_at_unix_secs item.finalized_at_unix_secs
.unwrap_or(item.created_at_unix_secs) .unwrap_or(item.created_at_unix_ms)
} }
fn users_me_usage_cache_hit_rate(total_input_context: u64, cache_read_tokens: u64) -> f64 { fn users_me_usage_cache_hit_rate(total_input_context: u64, cache_read_tokens: u64) -> f64 {
@@ -159,7 +159,7 @@ fn build_users_me_usage_record_payload(
"first_byte_time_ms": item.first_byte_time_ms, "first_byte_time_ms": item.first_byte_time_ms,
"is_stream": item.is_stream, "is_stream": item.is_stream,
"status": item.status, "status": item.status,
"created_at": unix_secs_to_rfc3339(item.created_at_unix_secs), "created_at": unix_secs_to_rfc3339(item.created_at_unix_ms),
"cache_creation_input_tokens": item.cache_creation_input_tokens, "cache_creation_input_tokens": item.cache_creation_input_tokens,
"cache_read_input_tokens": item.cache_read_input_tokens, "cache_read_input_tokens": item.cache_read_input_tokens,
"status_code": item.status_code, "status_code": item.status_code,
@@ -649,8 +649,8 @@ pub(super) async fn handle_users_me_usage_get(
.collect::<Vec<_>>(); .collect::<Vec<_>>();
records.sort_by(|left, right| { records.sort_by(|left, right| {
right right
.created_at_unix_secs .created_at_unix_ms
.cmp(&left.created_at_unix_secs) .cmp(&left.created_at_unix_ms)
.then_with(|| left.id.cmp(&right.id)) .then_with(|| left.id.cmp(&right.id))
}); });
let total_record_count = records.len(); let total_record_count = records.len();
@@ -735,8 +735,8 @@ pub(super) async fn handle_users_me_usage_active_get(
.collect::<Vec<_>>(); .collect::<Vec<_>>();
items.sort_by(|left, right| { items.sort_by(|left, right| {
right right
.created_at_unix_secs .created_at_unix_ms
.cmp(&left.created_at_unix_secs) .cmp(&left.created_at_unix_ms)
.then_with(|| left.id.cmp(&right.id)) .then_with(|| left.id.cmp(&right.id))
}); });
if ids.is_none() && items.len() > 50 { if ids.is_none() && items.len() > 50 {
@@ -798,20 +798,19 @@ pub(super) async fn handle_users_me_usage_interval_timeline_get(
}; };
items.retain(|item| item.status == "completed"); items.retain(|item| item.status == "completed");
items.sort_by(|left, right| { items.sort_by(|left, right| {
left.created_at_unix_secs left.created_at_unix_ms
.cmp(&right.created_at_unix_secs) .cmp(&right.created_at_unix_ms)
.then_with(|| left.id.cmp(&right.id)) .then_with(|| left.id.cmp(&right.id))
}); });
let mut points = Vec::new(); let mut points = Vec::new();
let mut previous_created_at_unix_secs = None; let mut previous_created_at_unix_ms = None;
for item in items { for item in items {
if let Some(previous) = previous_created_at_unix_secs { if let Some(previous) = previous_created_at_unix_ms {
let interval_minutes = let interval_minutes = (item.created_at_unix_ms.saturating_sub(previous) as f64) / 60.0;
(item.created_at_unix_secs.saturating_sub(previous) as f64) / 60.0;
if interval_minutes <= 120.0 { if interval_minutes <= 120.0 {
points.push(json!({ points.push(json!({
"x": unix_secs_to_rfc3339(item.created_at_unix_secs), "x": unix_secs_to_rfc3339(item.created_at_unix_ms),
"y": round_to(interval_minutes, 2), "y": round_to(interval_minutes, 2),
"model": item.model, "model": item.model,
})); }));
@@ -820,7 +819,7 @@ pub(super) async fn handle_users_me_usage_interval_timeline_get(
} }
} }
} }
previous_created_at_unix_secs = Some(item.created_at_unix_secs); previous_created_at_unix_ms = Some(item.created_at_unix_ms);
} }
Json(json!({ Json(json!({

View File

@@ -153,7 +153,7 @@ pub(super) fn wallet_transaction_payload_from_record(
"link_id": record.link_id.clone(), "link_id": record.link_id.clone(),
"operator_id": record.operator_id.clone(), "operator_id": record.operator_id.clone(),
"description": record.description.clone(), "description": record.description.clone(),
"created_at": record.created_at_unix_secs.and_then(unix_secs_to_rfc3339), "created_at": record.created_at_unix_ms.and_then(unix_secs_to_rfc3339),
}) })
} }

View File

@@ -207,7 +207,7 @@ pub(crate) fn wallet_payment_order_payload_from_row(
row: &sqlx::postgres::PgRow, row: &sqlx::postgres::PgRow,
) -> Result<serde_json::Value, GatewayError> { ) -> Result<serde_json::Value, GatewayError> {
let created_at = row let created_at = row
.try_get::<Option<i64>, _>("created_at_unix_secs") .try_get::<Option<i64>, _>("created_at_unix_ms")
.map_err(|err| GatewayError::Internal(err.to_string()))? .map_err(|err| GatewayError::Internal(err.to_string()))?
.and_then(|value| u64::try_from(value).ok()) .and_then(|value| u64::try_from(value).ok())
.and_then(unix_secs_to_rfc3339); .and_then(unix_secs_to_rfc3339);
@@ -280,7 +280,7 @@ fn wallet_payment_order_payload_from_record(
record.gateway_order_id.clone(), record.gateway_order_id.clone(),
record.gateway_response.clone(), record.gateway_response.clone(),
record.status.clone(), record.status.clone(),
Some(unix_secs_to_rfc3339(record.created_at_unix_secs)).flatten(), Some(unix_secs_to_rfc3339(record.created_at_unix_ms)).flatten(),
record.paid_at_unix_secs.and_then(unix_secs_to_rfc3339), record.paid_at_unix_secs.and_then(unix_secs_to_rfc3339),
record.credited_at_unix_secs.and_then(unix_secs_to_rfc3339), record.credited_at_unix_secs.and_then(unix_secs_to_rfc3339),
record.expires_at_unix_secs.and_then(unix_secs_to_rfc3339), record.expires_at_unix_secs.and_then(unix_secs_to_rfc3339),

View File

@@ -98,7 +98,7 @@ fn wallet_refund_payload_from_row(
row: &sqlx::postgres::PgRow, row: &sqlx::postgres::PgRow,
) -> Result<serde_json::Value, GatewayError> { ) -> Result<serde_json::Value, GatewayError> {
let created_at = row let created_at = row
.try_get::<Option<i64>, _>("created_at_unix_secs") .try_get::<Option<i64>, _>("created_at_unix_ms")
.map_err(|err| GatewayError::Internal(err.to_string()))? .map_err(|err| GatewayError::Internal(err.to_string()))?
.and_then(|value| u64::try_from(value).ok()) .and_then(|value| u64::try_from(value).ok())
.and_then(unix_secs_to_rfc3339); .and_then(unix_secs_to_rfc3339);
@@ -157,7 +157,7 @@ fn wallet_refund_payload_from_record(
"payout_method": record.payout_method, "payout_method": record.payout_method,
"payout_reference": record.payout_reference, "payout_reference": record.payout_reference,
"payout_proof": record.payout_proof, "payout_proof": record.payout_proof,
"created_at": unix_secs_to_rfc3339(record.created_at_unix_secs), "created_at": unix_secs_to_rfc3339(record.created_at_unix_ms),
"updated_at": unix_secs_to_rfc3339(record.updated_at_unix_secs), "updated_at": unix_secs_to_rfc3339(record.updated_at_unix_secs),
"processed_at": record.processed_at_unix_secs.and_then(unix_secs_to_rfc3339), "processed_at": record.processed_at_unix_secs.and_then(unix_secs_to_rfc3339),
"completed_at": record.completed_at_unix_secs.and_then(unix_secs_to_rfc3339), "completed_at": record.completed_at_unix_secs.and_then(unix_secs_to_rfc3339),
@@ -247,7 +247,7 @@ pub(super) async fn handle_wallet_refunds_list(
"payout_method": record.payout_method, "payout_method": record.payout_method,
"payout_reference": record.payout_reference, "payout_reference": record.payout_reference,
"payout_proof": record.payout_proof, "payout_proof": record.payout_proof,
"created_at": unix_secs_to_rfc3339(record.created_at_unix_secs), "created_at": unix_secs_to_rfc3339(record.created_at_unix_ms),
"updated_at": unix_secs_to_rfc3339(record.updated_at_unix_secs), "updated_at": unix_secs_to_rfc3339(record.updated_at_unix_secs),
"processed_at": record.processed_at_unix_secs.and_then(unix_secs_to_rfc3339), "processed_at": record.processed_at_unix_secs.and_then(unix_secs_to_rfc3339),
"completed_at": record.completed_at_unix_secs.and_then(unix_secs_to_rfc3339), "completed_at": record.completed_at_unix_secs.and_then(unix_secs_to_rfc3339),

View File

@@ -156,7 +156,7 @@ pub(crate) async fn build_admin_keys_grouped_by_format_payload(
.and_then(serde_json::Value::as_bool) .and_then(serde_json::Value::as_bool)
.unwrap_or(false), .unwrap_or(false),
"last_used_at": key.last_used_at_unix_secs.and_then(unix_secs_to_rfc3339), "last_used_at": key.last_used_at_unix_secs.and_then(unix_secs_to_rfc3339),
"created_at": unix_secs_to_rfc3339(key.created_at_unix_secs.unwrap_or(now_unix_secs)), "created_at": unix_secs_to_rfc3339(key.created_at_unix_ms.unwrap_or(now_unix_secs)),
"updated_at": unix_secs_to_rfc3339(key.updated_at_unix_secs.unwrap_or(now_unix_secs)), "updated_at": unix_secs_to_rfc3339(key.updated_at_unix_secs.unwrap_or(now_unix_secs)),
})); }));
} }

View File

@@ -598,7 +598,7 @@ pub(crate) fn build_admin_provider_key_response(
payload.insert( payload.insert(
"created_at".to_string(), "created_at".to_string(),
json!(unix_secs_to_rfc3339( json!(unix_secs_to_rfc3339(
key.created_at_unix_secs.unwrap_or(now_unix_secs) key.created_at_unix_ms.unwrap_or(now_unix_secs)
)), )),
); );
payload.insert( payload.insert(

Some files were not shown because too many files have changed in this diff Show More