mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
Merge remote-tracking branch 'upstream/aether-rust-pioneer' into rust
This commit is contained in:
91
.github/workflows/build-hub.yml
vendored
91
.github/workflows/build-hub.yml
vendored
@@ -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
|
||||
@@ -5,8 +5,9 @@ pub(crate) use crate::ai_pipeline::{
|
||||
core_error_default_client_api_format, is_core_error_finalize_kind, LocalCoreSyncErrorKind,
|
||||
};
|
||||
pub(crate) use crate::ai_pipeline::{
|
||||
request_conversion_direct_auth, request_conversion_kind,
|
||||
request_conversion_transport_supported, sync_chat_response_conversion_kind,
|
||||
request_candidate_api_formats, request_conversion_direct_auth, request_conversion_kind,
|
||||
request_conversion_requires_enable_flag, request_conversion_transport_supported,
|
||||
request_pair_allowed_for_transport, sync_chat_response_conversion_kind,
|
||||
sync_cli_response_conversion_kind, RequestConversionKind, SyncChatResponseConversionKind,
|
||||
SyncCliResponseConversionKind,
|
||||
};
|
||||
@@ -14,7 +15,7 @@ pub(crate) use crate::ai_pipeline::{
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
request_conversion_kind, sync_chat_response_conversion_kind,
|
||||
request_candidate_api_formats, request_conversion_kind, sync_chat_response_conversion_kind,
|
||||
sync_cli_response_conversion_kind, RequestConversionKind, SyncChatResponseConversionKind,
|
||||
SyncCliResponseConversionKind,
|
||||
};
|
||||
@@ -22,7 +23,15 @@ mod tests {
|
||||
#[test]
|
||||
fn request_conversion_registry_supports_bidirectional_standard_matrix() {
|
||||
assert_eq!(
|
||||
request_conversion_kind("claude:chat", "openai:chat"),
|
||||
request_conversion_kind("openai:chat", "openai:cli"),
|
||||
Some(RequestConversionKind::ToOpenAIFamilyCli)
|
||||
);
|
||||
assert_eq!(
|
||||
request_conversion_kind("openai:chat", "claude:cli"),
|
||||
Some(RequestConversionKind::ToClaudeStandard)
|
||||
);
|
||||
assert_eq!(
|
||||
request_conversion_kind("openai:cli", "openai:chat"),
|
||||
Some(RequestConversionKind::ToOpenAIChat)
|
||||
);
|
||||
assert_eq!(
|
||||
@@ -30,12 +39,16 @@ mod tests {
|
||||
Some(RequestConversionKind::ToClaudeStandard)
|
||||
);
|
||||
assert_eq!(
|
||||
request_conversion_kind("gemini:cli", "openai:compact"),
|
||||
Some(RequestConversionKind::ToOpenAICompact)
|
||||
request_conversion_kind("openai:compact", "gemini:cli"),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
request_conversion_kind("openai:compact", "gemini:cli"),
|
||||
Some(RequestConversionKind::ToGeminiStandard)
|
||||
request_conversion_kind("gemini:cli", "openai:compact"),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
request_conversion_kind("openai:chat", "openai:compact"),
|
||||
None
|
||||
);
|
||||
assert_eq!(request_conversion_kind("claude:chat", "claude:chat"), None);
|
||||
}
|
||||
@@ -59,12 +72,50 @@ mod tests {
|
||||
Some(SyncCliResponseConversionKind::ToGeminiCli)
|
||||
);
|
||||
assert_eq!(
|
||||
sync_cli_response_conversion_kind("claude:cli", "openai:compact"),
|
||||
sync_cli_response_conversion_kind("claude:chat", "openai:cli"),
|
||||
Some(SyncCliResponseConversionKind::ToOpenAIFamilyCli)
|
||||
);
|
||||
assert_eq!(
|
||||
sync_cli_response_conversion_kind("claude:cli", "openai:compact"),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
sync_cli_response_conversion_kind("openai:compact", "claude:cli"),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
sync_cli_response_conversion_kind("gemini:cli", "claude:cli"),
|
||||
Some(SyncCliResponseConversionKind::ToClaudeCli)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_candidate_registry_excludes_compact_as_cross_format_target() {
|
||||
assert_eq!(
|
||||
request_candidate_api_formats("openai:chat", false),
|
||||
vec![
|
||||
"openai:chat",
|
||||
"openai:cli",
|
||||
"claude:chat",
|
||||
"claude:cli",
|
||||
"gemini:chat",
|
||||
"gemini:cli",
|
||||
]
|
||||
);
|
||||
assert_eq!(
|
||||
request_candidate_api_formats("openai:cli", false),
|
||||
vec![
|
||||
"openai:chat",
|
||||
"openai:cli",
|
||||
"claude:chat",
|
||||
"claude:cli",
|
||||
"gemini:chat",
|
||||
"gemini:cli",
|
||||
]
|
||||
);
|
||||
assert_eq!(
|
||||
request_candidate_api_formats("openai:compact", false),
|
||||
vec!["openai:compact"]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -52,6 +52,130 @@ fn gemini_cli_v1internal_stream_rewriter_unwraps_response_object() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_chat_error_to_openai_cli_stream_rewriter_converts_to_response_failed() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "openai:chat",
|
||||
"client_api_format": "openai:cli",
|
||||
"needs_conversion": true,
|
||||
"mapped_model": "gpt-5.4",
|
||||
});
|
||||
let mut rewriter =
|
||||
maybe_build_local_stream_rewriter(Some(&report_context)).expect("rewriter should exist");
|
||||
let output = rewriter
|
||||
.push_chunk(
|
||||
b"data: {\"error\":{\"message\":\"bad request\",\"type\":\"invalid_request_error\",\"code\":\"invalid_request\"}}\n\n",
|
||||
)
|
||||
.expect("rewrite should succeed");
|
||||
let output_text = utf8(output);
|
||||
assert!(output_text.contains("event: response.failed"));
|
||||
assert!(output_text.contains("\"sequence_number\":1"));
|
||||
assert!(output_text.contains("\"message\":\"bad request\""));
|
||||
assert!(output_text.contains("\"type\":\"invalid_request_error\""));
|
||||
assert!(!output_text.contains("data: [DONE]"));
|
||||
assert!(rewriter.finish().expect("finish should succeed").is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn claude_error_to_openai_cli_stream_rewriter_converts_to_response_failed() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "claude:chat",
|
||||
"client_api_format": "openai:cli",
|
||||
"needs_conversion": true,
|
||||
"mapped_model": "claude-sonnet-4-5",
|
||||
});
|
||||
let mut rewriter =
|
||||
maybe_build_local_stream_rewriter(Some(&report_context)).expect("rewriter should exist");
|
||||
let output = rewriter
|
||||
.push_chunk(
|
||||
concat!(
|
||||
"event: error\n",
|
||||
"data: {\"type\":\"error\",\"error\":{\"type\":\"rate_limit_error\",\"message\":\"slow down\",\"code\":\"rate_limit\"}}\n\n"
|
||||
)
|
||||
.as_bytes(),
|
||||
)
|
||||
.expect("rewrite should succeed");
|
||||
let output_text = utf8(output);
|
||||
assert!(output_text.contains("event: response.failed"));
|
||||
assert!(output_text.contains("\"sequence_number\":1"));
|
||||
assert!(output_text.contains("\"message\":\"slow down\""));
|
||||
assert!(output_text.contains("\"type\":\"rate_limit_error\""));
|
||||
assert!(!output_text.contains("data: [DONE]"));
|
||||
assert!(rewriter.finish().expect("finish should succeed").is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemini_error_to_openai_cli_stream_rewriter_converts_to_response_failed() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "gemini:cli",
|
||||
"client_api_format": "openai:cli",
|
||||
"needs_conversion": true,
|
||||
"mapped_model": "gemini-2.5-pro",
|
||||
});
|
||||
let mut rewriter =
|
||||
maybe_build_local_stream_rewriter(Some(&report_context)).expect("rewriter should exist");
|
||||
let output = rewriter
|
||||
.push_chunk(
|
||||
b"data: {\"error\":{\"code\":429,\"message\":\"quota exceeded\",\"status\":\"RESOURCE_EXHAUSTED\"}}\n\n",
|
||||
)
|
||||
.expect("rewrite should succeed");
|
||||
let output_text = utf8(output);
|
||||
assert!(output_text.contains("event: response.failed"));
|
||||
assert!(output_text.contains("\"sequence_number\":1"));
|
||||
assert!(output_text.contains("\"message\":\"quota exceeded\""));
|
||||
assert!(output_text.contains("\"type\":\"rate_limit_error\""));
|
||||
assert!(!output_text.contains("data: [DONE]"));
|
||||
assert!(rewriter.finish().expect("finish should succeed").is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_chat_error_to_claude_chat_stream_rewriter_uses_error_event_line() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "openai:chat",
|
||||
"client_api_format": "claude:chat",
|
||||
"needs_conversion": true,
|
||||
"mapped_model": "gpt-5.4",
|
||||
});
|
||||
let mut rewriter =
|
||||
maybe_build_local_stream_rewriter(Some(&report_context)).expect("rewriter should exist");
|
||||
let output = rewriter
|
||||
.push_chunk(
|
||||
b"data: {\"error\":{\"message\":\"bad request\",\"type\":\"invalid_request_error\",\"code\":\"invalid_request\"}}\n\n",
|
||||
)
|
||||
.expect("rewrite should succeed");
|
||||
let output_text = utf8(output);
|
||||
assert!(output_text.contains("event: error"));
|
||||
assert!(output_text.contains("\"type\":\"error\""));
|
||||
assert!(output_text.contains("\"message\":\"bad request\""));
|
||||
assert!(output_text.contains("\"code\":\"invalid_request\""));
|
||||
assert!(!output_text.contains("data: [DONE]"));
|
||||
assert!(rewriter.finish().expect("finish should succeed").is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_chat_error_to_gemini_chat_stream_rewriter_keeps_data_only_error() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "openai:chat",
|
||||
"client_api_format": "gemini:chat",
|
||||
"needs_conversion": true,
|
||||
"mapped_model": "gpt-5.4",
|
||||
});
|
||||
let mut rewriter =
|
||||
maybe_build_local_stream_rewriter(Some(&report_context)).expect("rewriter should exist");
|
||||
let output = rewriter
|
||||
.push_chunk(
|
||||
b"data: {\"error\":{\"message\":\"rate limited\",\"type\":\"rate_limit_error\",\"code\":\"rate_limit\"}}\n\n",
|
||||
)
|
||||
.expect("rewrite should succeed");
|
||||
let output_text = utf8(output);
|
||||
assert!(output_text.starts_with("data: {\"error\":"));
|
||||
assert!(!output_text.contains("event: "));
|
||||
assert!(output_text.contains("\"message\":\"rate limited\""));
|
||||
assert!(output_text.contains("\"status\":\"RESOURCE_EXHAUSTED\""));
|
||||
assert!(!output_text.contains("data: [DONE]"));
|
||||
assert!(rewriter.finish().expect("finish should succeed").is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn claude_to_openai_chat_stream_rewriter_converts_text_deltas() {
|
||||
let report_context = json!({
|
||||
@@ -234,6 +358,22 @@ fn openai_cli_to_openai_chat_stream_rewriter_converts_text_deltas_immediately()
|
||||
.expect("rewrite should succeed");
|
||||
let completed_text = String::from_utf8(completed).expect("utf8 should decode");
|
||||
assert!(completed_text.contains("\"finish_reason\":\"stop\""));
|
||||
assert!(
|
||||
completed_text.contains("\"choices\":[]"),
|
||||
"{completed_text}"
|
||||
);
|
||||
assert!(
|
||||
completed_text.contains("\"prompt_tokens\":1"),
|
||||
"{completed_text}"
|
||||
);
|
||||
assert!(
|
||||
completed_text.contains("\"completion_tokens\":2"),
|
||||
"{completed_text}"
|
||||
);
|
||||
assert!(
|
||||
completed_text.contains("\"total_tokens\":3"),
|
||||
"{completed_text}"
|
||||
);
|
||||
assert!(completed_text.contains("data: [DONE]"));
|
||||
assert!(rewriter.finish().expect("finish should succeed").is_empty());
|
||||
}
|
||||
@@ -262,6 +402,13 @@ fn openai_cli_to_openai_chat_stream_rewriter_converts_completed_event_without_bu
|
||||
assert!(output_text.contains("\"role\":\"assistant\""));
|
||||
assert!(output_text.contains("\"content\":\"Hello Codex\""));
|
||||
assert!(output_text.contains("\"finish_reason\":\"stop\""));
|
||||
assert!(output_text.contains("\"choices\":[]"), "{output_text}");
|
||||
assert!(output_text.contains("\"prompt_tokens\":1"), "{output_text}");
|
||||
assert!(
|
||||
output_text.contains("\"completion_tokens\":2"),
|
||||
"{output_text}"
|
||||
);
|
||||
assert!(output_text.contains("\"total_tokens\":3"), "{output_text}");
|
||||
assert!(output_text.contains("data: [DONE]"));
|
||||
assert!(rewriter.finish().expect("finish should succeed").is_empty());
|
||||
}
|
||||
@@ -514,6 +661,7 @@ fn openai_chat_to_claude_chat_stream_rewriter_converts_via_standard_matrix() {
|
||||
assert!(first_text.contains("event: content_block_start"));
|
||||
assert!(first_text.contains("event: content_block_delta"));
|
||||
assert!(first_text.contains("\"text\":\"Hello Claude\""));
|
||||
assert!(first_text.contains("\"usage\":{\"input_tokens\":0,\"output_tokens\":0}"));
|
||||
|
||||
let second = rewriter
|
||||
.push_chunk(
|
||||
@@ -531,6 +679,44 @@ fn openai_chat_to_claude_chat_stream_rewriter_converts_via_standard_matrix() {
|
||||
assert!(rewriter.finish().expect("finish should succeed").is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_chat_to_claude_chat_stream_rewriter_injects_default_usage_when_finish_chunk_lacks_usage()
|
||||
{
|
||||
let report_context = json!({
|
||||
"provider_api_format": "openai:chat",
|
||||
"client_api_format": "claude:chat",
|
||||
"needs_conversion": true,
|
||||
"mapped_model": "gpt-5.4",
|
||||
});
|
||||
let mut rewriter =
|
||||
maybe_build_local_stream_rewriter(Some(&report_context)).expect("rewriter should exist");
|
||||
let first = rewriter
|
||||
.push_chunk(
|
||||
"data: {\"id\":\"chatcmpl_usage_missing_123\",\"object\":\"chat.completion.chunk\",\"created\":1,\"model\":\"gpt-5.4\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"Hello\"},\"finish_reason\":null}]}\n\n"
|
||||
.as_bytes(),
|
||||
)
|
||||
.expect("rewrite should succeed");
|
||||
let first_text = utf8(first);
|
||||
assert!(first_text.contains("event: message_start"));
|
||||
|
||||
let second = rewriter
|
||||
.push_chunk(
|
||||
concat!(
|
||||
"data: {\"id\":\"chatcmpl_usage_missing_123\",\"object\":\"chat.completion.chunk\",\"model\":\"gpt-5.4\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n",
|
||||
"data: [DONE]\n\n"
|
||||
)
|
||||
.as_bytes(),
|
||||
)
|
||||
.expect("rewrite should succeed");
|
||||
let output_text = utf8(second);
|
||||
assert!(output_text.is_empty(), "{output_text}");
|
||||
let final_text = utf8(rewriter.finish().expect("finish should succeed"));
|
||||
assert!(final_text.contains("event: message_delta"));
|
||||
assert!(final_text.contains("\"stop_reason\":\"end_turn\""));
|
||||
assert!(final_text.contains("\"usage\":{\"input_tokens\":0,\"output_tokens\":0}"));
|
||||
assert!(final_text.contains("event: message_stop"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_chat_to_gemini_cli_stream_rewriter_converts_via_standard_matrix() {
|
||||
let report_context = json!({
|
||||
|
||||
@@ -68,6 +68,8 @@ pub(crate) fn build_execution_runtime_auth_context(
|
||||
ExecutionRuntimeAuthContext {
|
||||
user_id: auth_context.user_id.clone(),
|
||||
api_key_id: auth_context.api_key_id.clone(),
|
||||
username: auth_context.username.clone(),
|
||||
api_key_name: auth_context.api_key_name.clone(),
|
||||
balance_remaining: auth_context.balance_remaining,
|
||||
access_allowed: auth_context.access_allowed,
|
||||
}
|
||||
|
||||
@@ -1,9 +1,18 @@
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_pipeline::transport::resolve_transport_proxy_snapshot;
|
||||
use crate::ai_pipeline::PlannerAppState;
|
||||
use crate::AppState;
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use crate::ai_pipeline::{
|
||||
GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, PlannerAppState,
|
||||
};
|
||||
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
|
||||
use crate::scheduler::config::{read_scheduler_ordering_config, SchedulerOrderingConfig};
|
||||
use aether_scheduler_core::{
|
||||
build_scheduler_affinity_cache_key_for_api_key_id, compare_candidates_by_priority_mode,
|
||||
requested_capability_priority_for_candidate, SchedulerAffinityTarget,
|
||||
SchedulerMinimalCandidateSelectionCandidate,
|
||||
};
|
||||
|
||||
const PLANNER_SCHEDULER_AFFINITY_MAX_ENTRIES: usize = 10_000;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
|
||||
enum TunnelOwnerAffinityBucket {
|
||||
@@ -12,6 +21,12 @@ enum TunnelOwnerAffinityBucket {
|
||||
RemoteTunnel = 2,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
struct CandidateExecutionOrdering {
|
||||
tunnel_bucket: TunnelOwnerAffinityBucket,
|
||||
keep_priority_on_conversion: bool,
|
||||
}
|
||||
|
||||
pub(crate) async fn prefer_local_tunnel_owner_candidates(
|
||||
state: PlannerAppState<'_>,
|
||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
@@ -28,11 +43,127 @@ pub(crate) async fn prefer_local_tunnel_owner_candidates(
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(crate) async fn rank_local_execution_candidates(
|
||||
state: PlannerAppState<'_>,
|
||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
client_api_format: &str,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
) -> Vec<SchedulerMinimalCandidateSelectionCandidate> {
|
||||
let normalized_client_api_format = client_api_format.trim().to_ascii_lowercase();
|
||||
let ordering_config = read_scheduler_ordering_config_or_default(state).await;
|
||||
let mut ranked = Vec::with_capacity(candidates.len());
|
||||
|
||||
for (original_index, candidate) in candidates.into_iter().enumerate() {
|
||||
let ordering =
|
||||
resolve_candidate_execution_ordering(state, &candidate, ordering_config).await;
|
||||
let is_same_format = candidate
|
||||
.endpoint_api_format
|
||||
.trim()
|
||||
.eq_ignore_ascii_case(normalized_client_api_format.as_str());
|
||||
let demote_cross_format = !is_same_format && !ordering.keep_priority_on_conversion;
|
||||
let capability_priority =
|
||||
requested_capability_priority_for_candidate(required_capabilities, &candidate);
|
||||
ranked.push((
|
||||
capability_priority.0,
|
||||
capability_priority.1,
|
||||
ordering.tunnel_bucket,
|
||||
demote_cross_format,
|
||||
original_index,
|
||||
candidate,
|
||||
));
|
||||
}
|
||||
|
||||
ranked.sort_by(|left, right| {
|
||||
left.0
|
||||
.cmp(&right.0)
|
||||
.then(left.1.cmp(&right.1))
|
||||
.then(left.2.cmp(&right.2))
|
||||
.then(left.3.cmp(&right.3))
|
||||
.then_with(|| {
|
||||
compare_candidates_by_priority_mode(
|
||||
&left.5,
|
||||
&right.5,
|
||||
ordering_config.priority_mode,
|
||||
None,
|
||||
)
|
||||
})
|
||||
.then(left.4.cmp(&right.4))
|
||||
});
|
||||
|
||||
ranked
|
||||
.into_iter()
|
||||
.map(|(_, _, _, _, _, candidate)| candidate)
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(crate) fn remember_scheduler_affinity_for_candidate(
|
||||
state: PlannerAppState<'_>,
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_api_format: &str,
|
||||
requested_model: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
) {
|
||||
let Some(api_key_id) = auth_snapshot
|
||||
.map(|snapshot| snapshot.api_key_id.trim())
|
||||
.filter(|value| !value.is_empty())
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let Some(cache_key) = build_scheduler_affinity_cache_key_for_api_key_id(
|
||||
api_key_id,
|
||||
client_api_format,
|
||||
requested_model,
|
||||
) else {
|
||||
return;
|
||||
};
|
||||
|
||||
state.app().remember_scheduler_affinity_target(
|
||||
&cache_key,
|
||||
SchedulerAffinityTarget {
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
key_id: candidate.key_id.clone(),
|
||||
},
|
||||
SCHEDULER_AFFINITY_TTL,
|
||||
PLANNER_SCHEDULER_AFFINITY_MAX_ENTRIES,
|
||||
);
|
||||
}
|
||||
|
||||
async fn resolve_candidate_tunnel_owner_affinity(
|
||||
state: PlannerAppState<'_>,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
) -> TunnelOwnerAffinityBucket {
|
||||
let transport = match state
|
||||
let Some(transport) = read_candidate_transport_snapshot(state, candidate).await else {
|
||||
return TunnelOwnerAffinityBucket::Neutral;
|
||||
};
|
||||
|
||||
resolve_tunnel_owner_affinity_from_transport(state, &transport).await
|
||||
}
|
||||
|
||||
async fn resolve_candidate_execution_ordering(
|
||||
state: PlannerAppState<'_>,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
ordering_config: SchedulerOrderingConfig,
|
||||
) -> CandidateExecutionOrdering {
|
||||
let Some(transport) = read_candidate_transport_snapshot(state, candidate).await else {
|
||||
return CandidateExecutionOrdering {
|
||||
tunnel_bucket: TunnelOwnerAffinityBucket::Neutral,
|
||||
keep_priority_on_conversion: ordering_config.keep_priority_on_conversion,
|
||||
};
|
||||
};
|
||||
|
||||
CandidateExecutionOrdering {
|
||||
tunnel_bucket: resolve_tunnel_owner_affinity_from_transport(state, &transport).await,
|
||||
keep_priority_on_conversion: ordering_config.keep_priority_on_conversion
|
||||
|| transport.provider.keep_priority_on_conversion,
|
||||
}
|
||||
}
|
||||
|
||||
async fn read_candidate_transport_snapshot(
|
||||
state: PlannerAppState<'_>,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
) -> Option<GatewayProviderTransportSnapshot> {
|
||||
match state
|
||||
.read_provider_transport_snapshot(
|
||||
&candidate.provider_id,
|
||||
&candidate.endpoint_id,
|
||||
@@ -40,8 +171,8 @@ async fn resolve_candidate_tunnel_owner_affinity(
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(transport)) => transport,
|
||||
Ok(None) => return TunnelOwnerAffinityBucket::Neutral,
|
||||
Ok(Some(transport)) => Some(transport),
|
||||
Ok(None) => None,
|
||||
Err(error) => {
|
||||
warn!(
|
||||
event_name = "candidate_affinity_transport_load_failed",
|
||||
@@ -50,12 +181,17 @@ async fn resolve_candidate_tunnel_owner_affinity(
|
||||
endpoint_id = %candidate.endpoint_id,
|
||||
key_id = %candidate.key_id,
|
||||
error = ?error,
|
||||
"failed to load provider transport while evaluating tunnel owner affinity"
|
||||
"failed to load provider transport while evaluating execution ordering"
|
||||
);
|
||||
return TunnelOwnerAffinityBucket::Neutral;
|
||||
None
|
||||
}
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
async fn resolve_tunnel_owner_affinity_from_transport(
|
||||
state: PlannerAppState<'_>,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> TunnelOwnerAffinityBucket {
|
||||
let Some(proxy) = resolve_transport_proxy_snapshot(&transport) else {
|
||||
return TunnelOwnerAffinityBucket::Neutral;
|
||||
};
|
||||
@@ -99,6 +235,23 @@ async fn resolve_candidate_tunnel_owner_affinity(
|
||||
}
|
||||
}
|
||||
|
||||
async fn read_scheduler_ordering_config_or_default(
|
||||
state: PlannerAppState<'_>,
|
||||
) -> SchedulerOrderingConfig {
|
||||
match read_scheduler_ordering_config(state.app()).await {
|
||||
Ok(config) => config,
|
||||
Err(error) => {
|
||||
warn!(
|
||||
event_name = "planner_scheduler_ordering_config_load_failed",
|
||||
log_type = "event",
|
||||
error = ?error,
|
||||
"failed to load scheduler ordering config while ranking local execution candidates"
|
||||
);
|
||||
SchedulerOrderingConfig::default()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
@@ -110,11 +263,15 @@ mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::{
|
||||
prefer_local_tunnel_owner_candidates, AppState, PlannerAppState,
|
||||
prefer_local_tunnel_owner_candidates, rank_local_execution_candidates,
|
||||
remember_scheduler_affinity_for_candidate, PlannerAppState,
|
||||
SchedulerMinimalCandidateSelectionCandidate,
|
||||
};
|
||||
use crate::data::auth::GatewayAuthApiKeySnapshot;
|
||||
use crate::data::GatewayDataState;
|
||||
use crate::tunnel::TunnelAttachmentRecord;
|
||||
use crate::{scheduler::affinity::SCHEDULER_AFFINITY_TTL, AppState};
|
||||
use aether_data::repository::auth::StoredAuthApiKeySnapshot;
|
||||
|
||||
fn sample_candidate(
|
||||
endpoint_id: &str,
|
||||
@@ -142,22 +299,55 @@ mod tests {
|
||||
}
|
||||
|
||||
fn sample_provider() -> StoredProviderCatalogProvider {
|
||||
sample_provider_with_options("provider-1", false, 0)
|
||||
}
|
||||
|
||||
fn sample_provider_with_options(
|
||||
id: &str,
|
||||
keep_priority_on_conversion: bool,
|
||||
provider_priority: i32,
|
||||
) -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
"provider-1".to_string(),
|
||||
"provider-1".to_string(),
|
||||
id.to_string(),
|
||||
id.to_string(),
|
||||
Some("https://provider.example".to_string()),
|
||||
"custom".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_transport_fields(true, false, false, None, None, None, None, None, None)
|
||||
.with_transport_fields(
|
||||
true,
|
||||
keep_priority_on_conversion,
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.with_routing_fields(provider_priority)
|
||||
}
|
||||
|
||||
fn sample_endpoint(id: &str) -> StoredProviderCatalogEndpoint {
|
||||
sample_endpoint_for_provider("provider-1", id, "openai:chat")
|
||||
}
|
||||
|
||||
fn sample_endpoint_for_provider(
|
||||
provider_id: &str,
|
||||
id: &str,
|
||||
api_format: &str,
|
||||
) -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
id.to_string(),
|
||||
"provider-1".to_string(),
|
||||
"openai:chat".to_string(),
|
||||
Some("openai".to_string()),
|
||||
provider_id.to_string(),
|
||||
api_format.to_string(),
|
||||
Some(
|
||||
api_format
|
||||
.split(':')
|
||||
.next()
|
||||
.unwrap_or(api_format)
|
||||
.to_string(),
|
||||
),
|
||||
Some("chat".to_string()),
|
||||
true,
|
||||
)
|
||||
@@ -176,9 +366,17 @@ mod tests {
|
||||
}
|
||||
|
||||
fn sample_key(id: &str, node_id: &str) -> StoredProviderCatalogKey {
|
||||
sample_key_for_provider("provider-1", id, node_id)
|
||||
}
|
||||
|
||||
fn sample_key_for_provider(
|
||||
provider_id: &str,
|
||||
id: &str,
|
||||
node_id: &str,
|
||||
) -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
id.to_string(),
|
||||
"provider-1".to_string(),
|
||||
provider_id.to_string(),
|
||||
id.to_string(),
|
||||
"api_key".to_string(),
|
||||
None,
|
||||
@@ -214,6 +412,35 @@ mod tests {
|
||||
.as_secs()
|
||||
}
|
||||
|
||||
fn sample_priority_candidate(
|
||||
provider_id: &str,
|
||||
endpoint_id: &str,
|
||||
key_id: &str,
|
||||
endpoint_api_format: &str,
|
||||
key_global_priority_for_format: Option<i32>,
|
||||
provider_priority: i32,
|
||||
) -> SchedulerMinimalCandidateSelectionCandidate {
|
||||
SchedulerMinimalCandidateSelectionCandidate {
|
||||
provider_id: provider_id.to_string(),
|
||||
provider_name: provider_id.to_string(),
|
||||
provider_type: "custom".to_string(),
|
||||
provider_priority,
|
||||
endpoint_id: endpoint_id.to_string(),
|
||||
endpoint_api_format: endpoint_api_format.to_string(),
|
||||
key_id: key_id.to_string(),
|
||||
key_name: key_id.to_string(),
|
||||
key_auth_type: "api_key".to_string(),
|
||||
key_internal_priority: 0,
|
||||
key_global_priority_for_format,
|
||||
key_capabilities: None,
|
||||
model_id: format!("model-{provider_id}"),
|
||||
global_model_id: "global-model-1".to_string(),
|
||||
global_model_name: "gpt-4.1".to_string(),
|
||||
selected_provider_model_name: "gpt-4.1".to_string(),
|
||||
mapping_matched_model: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn prefers_local_tunnel_owner_candidates_before_remote_tunnel_candidates() {
|
||||
let provider_catalog = InMemoryProviderCatalogReadRepository::seed(
|
||||
@@ -300,4 +527,339 @@ mod tests {
|
||||
assert_eq!(reordered[0].endpoint_id, "endpoint-a");
|
||||
assert_eq!(reordered[1].endpoint_id, "endpoint-b");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_execution_ranking_demotes_cross_format_candidates_without_keep_priority() {
|
||||
let provider_catalog = InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![
|
||||
sample_provider_with_options("provider-same", false, 10),
|
||||
sample_provider_with_options("provider-cross", false, 0),
|
||||
],
|
||||
vec![
|
||||
sample_endpoint_for_provider("provider-same", "endpoint-same", "openai:chat"),
|
||||
sample_endpoint_for_provider("provider-cross", "endpoint-cross", "claude:chat"),
|
||||
],
|
||||
vec![
|
||||
sample_key_for_provider("provider-same", "key-same", ""),
|
||||
sample_key_for_provider("provider-cross", "key-cross", ""),
|
||||
],
|
||||
);
|
||||
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
|
||||
std::sync::Arc::new(provider_catalog),
|
||||
"development-key",
|
||||
);
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(data_state);
|
||||
|
||||
let ranked = rank_local_execution_candidates(
|
||||
PlannerAppState::new(&state),
|
||||
vec![
|
||||
sample_priority_candidate(
|
||||
"provider-cross",
|
||||
"endpoint-cross",
|
||||
"key-cross",
|
||||
"claude:chat",
|
||||
Some(0),
|
||||
0,
|
||||
),
|
||||
sample_priority_candidate(
|
||||
"provider-same",
|
||||
"endpoint-same",
|
||||
"key-same",
|
||||
"openai:chat",
|
||||
Some(10),
|
||||
10,
|
||||
),
|
||||
],
|
||||
"openai:chat",
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(ranked[0].endpoint_id, "endpoint-same");
|
||||
assert_eq!(ranked[1].endpoint_id, "endpoint-cross");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_execution_ranking_keeps_cross_format_priority_when_enabled() {
|
||||
let provider_catalog = InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![
|
||||
sample_provider_with_options("provider-same", false, 10),
|
||||
sample_provider_with_options("provider-cross", true, 0),
|
||||
],
|
||||
vec![
|
||||
sample_endpoint_for_provider("provider-same", "endpoint-same", "openai:chat"),
|
||||
sample_endpoint_for_provider("provider-cross", "endpoint-cross", "claude:chat"),
|
||||
],
|
||||
vec![
|
||||
sample_key_for_provider("provider-same", "key-same", ""),
|
||||
sample_key_for_provider("provider-cross", "key-cross", ""),
|
||||
],
|
||||
);
|
||||
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
|
||||
std::sync::Arc::new(provider_catalog),
|
||||
"development-key",
|
||||
);
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(data_state);
|
||||
|
||||
let ranked = rank_local_execution_candidates(
|
||||
PlannerAppState::new(&state),
|
||||
vec![
|
||||
sample_priority_candidate(
|
||||
"provider-cross",
|
||||
"endpoint-cross",
|
||||
"key-cross",
|
||||
"claude:chat",
|
||||
Some(0),
|
||||
0,
|
||||
),
|
||||
sample_priority_candidate(
|
||||
"provider-same",
|
||||
"endpoint-same",
|
||||
"key-same",
|
||||
"openai:chat",
|
||||
Some(10),
|
||||
10,
|
||||
),
|
||||
],
|
||||
"openai:chat",
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(ranked[0].endpoint_id, "endpoint-cross");
|
||||
assert_eq!(ranked[1].endpoint_id, "endpoint-same");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_execution_ranking_keeps_cross_format_priority_when_global_override_is_enabled() {
|
||||
let provider_catalog = InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![
|
||||
sample_provider_with_options("provider-same", false, 10),
|
||||
sample_provider_with_options("provider-cross", false, 0),
|
||||
],
|
||||
vec![
|
||||
sample_endpoint_for_provider("provider-same", "endpoint-same", "openai:chat"),
|
||||
sample_endpoint_for_provider("provider-cross", "endpoint-cross", "claude:chat"),
|
||||
],
|
||||
vec![
|
||||
sample_key_for_provider("provider-same", "key-same", ""),
|
||||
sample_key_for_provider("provider-cross", "key-cross", ""),
|
||||
],
|
||||
);
|
||||
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
|
||||
std::sync::Arc::new(provider_catalog),
|
||||
"development-key",
|
||||
)
|
||||
.with_system_config_values_for_tests(vec![(
|
||||
"keep_priority_on_conversion".to_string(),
|
||||
json!(true),
|
||||
)]);
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(data_state);
|
||||
|
||||
let ranked = rank_local_execution_candidates(
|
||||
PlannerAppState::new(&state),
|
||||
vec![
|
||||
sample_priority_candidate(
|
||||
"provider-cross",
|
||||
"endpoint-cross",
|
||||
"key-cross",
|
||||
"claude:chat",
|
||||
Some(0),
|
||||
0,
|
||||
),
|
||||
sample_priority_candidate(
|
||||
"provider-same",
|
||||
"endpoint-same",
|
||||
"key-same",
|
||||
"openai:chat",
|
||||
Some(10),
|
||||
10,
|
||||
),
|
||||
],
|
||||
"openai:chat",
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(ranked[0].endpoint_id, "endpoint-cross");
|
||||
assert_eq!(ranked[1].endpoint_id, "endpoint-same");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_execution_ranking_uses_provider_priority_mode_when_configured() {
|
||||
let provider_catalog = InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![
|
||||
sample_provider_with_options("provider-provider-first", false, 0),
|
||||
sample_provider_with_options("provider-global-first", false, 10),
|
||||
],
|
||||
vec![
|
||||
sample_endpoint_for_provider(
|
||||
"provider-provider-first",
|
||||
"endpoint-provider-first",
|
||||
"openai:chat",
|
||||
),
|
||||
sample_endpoint_for_provider(
|
||||
"provider-global-first",
|
||||
"endpoint-global-first",
|
||||
"openai:chat",
|
||||
),
|
||||
],
|
||||
vec![
|
||||
sample_key_for_provider("provider-provider-first", "key-provider-first", ""),
|
||||
sample_key_for_provider("provider-global-first", "key-global-first", ""),
|
||||
],
|
||||
);
|
||||
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
|
||||
std::sync::Arc::new(provider_catalog),
|
||||
"development-key",
|
||||
)
|
||||
.with_system_config_values_for_tests(vec![(
|
||||
"provider_priority_mode".to_string(),
|
||||
json!("provider"),
|
||||
)]);
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(data_state);
|
||||
|
||||
let ranked = rank_local_execution_candidates(
|
||||
PlannerAppState::new(&state),
|
||||
vec![
|
||||
sample_priority_candidate(
|
||||
"provider-global-first",
|
||||
"endpoint-global-first",
|
||||
"key-global-first",
|
||||
"openai:chat",
|
||||
Some(0),
|
||||
10,
|
||||
),
|
||||
sample_priority_candidate(
|
||||
"provider-provider-first",
|
||||
"endpoint-provider-first",
|
||||
"key-provider-first",
|
||||
"openai:chat",
|
||||
Some(10),
|
||||
0,
|
||||
),
|
||||
],
|
||||
"openai:chat",
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(ranked[0].endpoint_id, "endpoint-provider-first");
|
||||
assert_eq!(ranked[1].endpoint_id, "endpoint-global-first");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_execution_ranking_prefers_candidates_matching_requested_capabilities() {
|
||||
let provider_catalog = InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![
|
||||
sample_provider_with_options("provider-miss", false, 0),
|
||||
sample_provider_with_options("provider-hit", false, 0),
|
||||
],
|
||||
vec![
|
||||
sample_endpoint_for_provider("provider-miss", "endpoint-miss", "openai:chat"),
|
||||
sample_endpoint_for_provider("provider-hit", "endpoint-hit", "openai:chat"),
|
||||
],
|
||||
vec![
|
||||
sample_key_for_provider("provider-miss", "key-miss", ""),
|
||||
sample_key_for_provider("provider-hit", "key-hit", ""),
|
||||
],
|
||||
);
|
||||
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
|
||||
std::sync::Arc::new(provider_catalog),
|
||||
"development-key",
|
||||
);
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(data_state);
|
||||
|
||||
let mut candidate_miss = sample_priority_candidate(
|
||||
"provider-miss",
|
||||
"endpoint-miss",
|
||||
"key-miss",
|
||||
"openai:chat",
|
||||
Some(0),
|
||||
0,
|
||||
);
|
||||
let mut candidate_hit = sample_priority_candidate(
|
||||
"provider-hit",
|
||||
"endpoint-hit",
|
||||
"key-hit",
|
||||
"openai:chat",
|
||||
Some(0),
|
||||
0,
|
||||
);
|
||||
candidate_miss.key_capabilities = Some(json!({"cache_1h": false}));
|
||||
candidate_hit.key_capabilities = Some(json!({"cache_1h": true}));
|
||||
|
||||
let required_capabilities = json!({"cache_1h": true});
|
||||
let ranked = rank_local_execution_candidates(
|
||||
PlannerAppState::new(&state),
|
||||
vec![candidate_miss, candidate_hit],
|
||||
"openai:chat",
|
||||
Some(&required_capabilities),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(ranked[0].endpoint_id, "endpoint-hit");
|
||||
assert_eq!(ranked[1].endpoint_id, "endpoint-miss");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn remembers_scheduler_affinity_for_candidate_using_requested_model_key() {
|
||||
let state = AppState::new().expect("state should build");
|
||||
let auth_snapshot = GatewayAuthApiKeySnapshot::from_stored(
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
"user-1".to_string(),
|
||||
"alice".to_string(),
|
||||
Some("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");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,7 +2,9 @@ use serde_json::json;
|
||||
use tracing::warn;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
|
||||
use crate::ai_pipeline::planner::candidate_affinity::{
|
||||
rank_local_execution_candidates, remember_scheduler_affinity_for_candidate,
|
||||
};
|
||||
use crate::ai_pipeline::{
|
||||
resolve_local_decision_execution_runtime_auth_context, ConversionMode, ExecutionStrategy,
|
||||
GatewayControlDecision, PlannerAppState,
|
||||
@@ -61,10 +63,20 @@ pub(crate) async fn resolve_local_same_format_provider_decision_input(
|
||||
}
|
||||
};
|
||||
|
||||
let required_capabilities = planner_state
|
||||
.resolve_request_candidate_required_capabilities(
|
||||
&auth_context.user_id,
|
||||
&auth_context.api_key_id,
|
||||
Some(requested_model.as_str()),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
Some(LocalSameFormatProviderDecisionInput {
|
||||
auth_context,
|
||||
requested_model,
|
||||
auth_snapshot,
|
||||
required_capabilities,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -80,16 +92,34 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
|
||||
spec.api_format,
|
||||
&input.requested_model,
|
||||
spec.require_streaming,
|
||||
input.required_capabilities.as_ref(),
|
||||
Some(&input.auth_snapshot),
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await?;
|
||||
let candidates = prefer_local_tunnel_owner_candidates(planner_state, candidates).await;
|
||||
let candidates = rank_local_execution_candidates(
|
||||
planner_state,
|
||||
candidates,
|
||||
spec.api_format,
|
||||
input.required_capabilities.as_ref(),
|
||||
)
|
||||
.await;
|
||||
|
||||
let created_at_unix_secs = current_unix_secs();
|
||||
let created_at_unix_ms = current_unix_secs();
|
||||
let mut attempts = Vec::with_capacity(candidates.len());
|
||||
let mut affinity_remembered = false;
|
||||
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
|
||||
let generated_candidate_id = Uuid::new_v4().to_string();
|
||||
if !affinity_remembered {
|
||||
remember_scheduler_affinity_for_candidate(
|
||||
planner_state,
|
||||
Some(&input.auth_snapshot),
|
||||
spec.api_format,
|
||||
&input.requested_model,
|
||||
&candidate,
|
||||
);
|
||||
affinity_remembered = true;
|
||||
}
|
||||
let extra_data = append_execution_contract_fields_to_value(
|
||||
json!({
|
||||
"provider_api_format": spec.api_format,
|
||||
@@ -116,8 +146,9 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
|
||||
&candidate,
|
||||
candidate_index as u32,
|
||||
&generated_candidate_id,
|
||||
input.required_capabilities.as_ref(),
|
||||
Some(extra_data),
|
||||
created_at_unix_secs,
|
||||
created_at_unix_ms,
|
||||
"gateway local same-format decision request candidate upsert failed",
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -21,6 +21,7 @@ pub(crate) struct LocalSameFormatProviderDecisionInput {
|
||||
pub(crate) auth_context: ExecutionRuntimeAuthContext,
|
||||
pub(crate) requested_model: String,
|
||||
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
pub(crate) required_capabilities: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
|
||||
@@ -265,6 +265,8 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
||||
json!({
|
||||
"user_id": input.auth_context.user_id,
|
||||
"api_key_id": input.auth_context.api_key_id,
|
||||
"username": input.auth_context.username,
|
||||
"api_key_name": input.auth_context.api_key_name,
|
||||
"request_id": trace_id,
|
||||
"candidate_id": candidate_id,
|
||||
"candidate_index": candidate_index,
|
||||
@@ -359,6 +361,7 @@ pub(super) async fn mark_skipped_local_same_format_provider_candidate(
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
input.required_capabilities.as_ref(),
|
||||
skip_reason,
|
||||
current_unix_secs(),
|
||||
"gateway local same-format decision failed to persist skipped candidate",
|
||||
|
||||
@@ -281,6 +281,8 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
|
||||
report_context: Some(json!({
|
||||
"user_id": input.auth_context.user_id,
|
||||
"api_key_id": input.auth_context.api_key_id,
|
||||
"username": input.auth_context.username,
|
||||
"api_key_name": input.auth_context.api_key_name,
|
||||
"request_id": trace_id,
|
||||
"candidate_id": candidate_id,
|
||||
"candidate_index": candidate_index,
|
||||
|
||||
@@ -4,7 +4,9 @@ use tracing::warn;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
||||
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
|
||||
use crate::ai_pipeline::planner::candidate_affinity::{
|
||||
rank_local_execution_candidates, remember_scheduler_affinity_for_candidate,
|
||||
};
|
||||
use crate::ai_pipeline::{
|
||||
resolve_local_decision_execution_runtime_auth_context, GatewayControlDecision,
|
||||
};
|
||||
@@ -20,6 +22,7 @@ pub(super) const GEMINI_FILES_REQUIRED_CAPABILITY: &str = "gemini_files";
|
||||
pub(super) struct LocalGeminiFilesDecisionInput {
|
||||
pub(super) auth_context: ExecutionRuntimeAuthContext,
|
||||
pub(super) auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
pub(super) required_capabilities: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -59,9 +62,20 @@ pub(super) async fn resolve_local_gemini_files_decision_input(
|
||||
}
|
||||
};
|
||||
|
||||
let explicit_required_capabilities = json!({ "gemini_files": true });
|
||||
let required_capabilities = planner_state
|
||||
.resolve_request_candidate_required_capabilities(
|
||||
&auth_context.user_id,
|
||||
&auth_context.api_key_id,
|
||||
None,
|
||||
Some(&explicit_required_capabilities),
|
||||
)
|
||||
.await;
|
||||
|
||||
Some(LocalGeminiFilesDecisionInput {
|
||||
auth_context,
|
||||
auth_snapshot,
|
||||
required_capabilities,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -80,12 +94,29 @@ pub(super) async fn materialize_local_gemini_files_candidate_attempts(
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await?;
|
||||
let candidates = prefer_local_tunnel_owner_candidates(planner_state, candidates).await;
|
||||
let candidates = rank_local_execution_candidates(
|
||||
planner_state,
|
||||
candidates,
|
||||
GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
input.required_capabilities.as_ref(),
|
||||
)
|
||||
.await;
|
||||
|
||||
let created_at_unix_secs = current_unix_secs();
|
||||
let created_at_unix_ms = current_unix_secs();
|
||||
let mut attempts = Vec::with_capacity(candidates.len());
|
||||
let mut affinity_remembered = false;
|
||||
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
|
||||
let generated_candidate_id = Uuid::new_v4().to_string();
|
||||
if !affinity_remembered {
|
||||
remember_scheduler_affinity_for_candidate(
|
||||
planner_state,
|
||||
Some(&input.auth_snapshot),
|
||||
GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
&candidate.global_model_name,
|
||||
&candidate,
|
||||
);
|
||||
affinity_remembered = true;
|
||||
}
|
||||
let extra_data = json!({
|
||||
"provider_api_format": GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
"client_api_format": GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
@@ -107,8 +138,9 @@ pub(super) async fn materialize_local_gemini_files_candidate_attempts(
|
||||
&candidate,
|
||||
candidate_index as u32,
|
||||
&generated_candidate_id,
|
||||
input.required_capabilities.as_ref(),
|
||||
Some(extra_data),
|
||||
created_at_unix_secs,
|
||||
created_at_unix_ms,
|
||||
"gateway local gemini files request candidate upsert failed",
|
||||
)
|
||||
.await;
|
||||
@@ -140,6 +172,7 @@ pub(super) async fn mark_skipped_local_gemini_files_candidate(
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
input.required_capabilities.as_ref(),
|
||||
skip_reason,
|
||||
current_unix_secs(),
|
||||
"gateway local gemini files failed to persist skipped candidate",
|
||||
|
||||
@@ -248,6 +248,8 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
|
||||
report_context: Some(json!({
|
||||
"user_id": input.auth_context.user_id.clone(),
|
||||
"api_key_id": input.auth_context.api_key_id.clone(),
|
||||
"username": input.auth_context.username.clone(),
|
||||
"api_key_name": input.auth_context.api_key_name.clone(),
|
||||
"request_id": trace_id,
|
||||
"candidate_id": candidate_id,
|
||||
"candidate_index": candidate_index,
|
||||
|
||||
@@ -5,7 +5,9 @@ use uuid::Uuid;
|
||||
|
||||
use super::{LocalVideoCreateFamily, LocalVideoCreateSpec};
|
||||
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
||||
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
|
||||
use crate::ai_pipeline::planner::candidate_affinity::{
|
||||
rank_local_execution_candidates, remember_scheduler_affinity_for_candidate,
|
||||
};
|
||||
use crate::ai_pipeline::{
|
||||
resolve_local_decision_execution_runtime_auth_context, GatewayControlDecision,
|
||||
};
|
||||
@@ -18,6 +20,7 @@ pub(super) struct LocalVideoCreateDecisionInput {
|
||||
pub(super) auth_context: ExecutionRuntimeAuthContext,
|
||||
pub(super) requested_model: String,
|
||||
pub(super) auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
pub(super) required_capabilities: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -71,10 +74,20 @@ pub(super) async fn resolve_local_video_create_decision_input(
|
||||
}
|
||||
};
|
||||
|
||||
let required_capabilities = planner_state
|
||||
.resolve_request_candidate_required_capabilities(
|
||||
&auth_context.user_id,
|
||||
&auth_context.api_key_id,
|
||||
Some(requested_model.as_str()),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
Some(LocalVideoCreateDecisionInput {
|
||||
auth_context,
|
||||
requested_model,
|
||||
auth_snapshot,
|
||||
required_capabilities,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -91,6 +104,7 @@ pub(super) async fn list_local_video_create_candidate_attempts(
|
||||
api_format,
|
||||
&input.requested_model,
|
||||
false,
|
||||
input.required_capabilities.as_ref(),
|
||||
Some(&input.auth_snapshot),
|
||||
current_unix_secs(),
|
||||
)
|
||||
@@ -127,12 +141,29 @@ async fn materialize_local_video_create_candidate_attempts(
|
||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
api_format: &str,
|
||||
) -> Vec<LocalVideoCreateCandidateAttempt> {
|
||||
let candidates = prefer_local_tunnel_owner_candidates(state, candidates).await;
|
||||
let created_at_unix_secs = current_unix_secs();
|
||||
let candidates = rank_local_execution_candidates(
|
||||
state,
|
||||
candidates,
|
||||
api_format,
|
||||
input.required_capabilities.as_ref(),
|
||||
)
|
||||
.await;
|
||||
let created_at_unix_ms = current_unix_secs();
|
||||
let mut attempts = Vec::with_capacity(candidates.len());
|
||||
let mut affinity_remembered = false;
|
||||
|
||||
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
|
||||
let generated_candidate_id = Uuid::new_v4().to_string();
|
||||
if !affinity_remembered {
|
||||
remember_scheduler_affinity_for_candidate(
|
||||
state,
|
||||
Some(&input.auth_snapshot),
|
||||
api_format,
|
||||
&input.requested_model,
|
||||
&candidate,
|
||||
);
|
||||
affinity_remembered = true;
|
||||
}
|
||||
let extra_data = json!({
|
||||
"provider_api_format": api_format,
|
||||
"client_api_format": api_format,
|
||||
@@ -153,8 +184,9 @@ async fn materialize_local_video_create_candidate_attempts(
|
||||
&candidate,
|
||||
candidate_index as u32,
|
||||
&generated_candidate_id,
|
||||
input.required_capabilities.as_ref(),
|
||||
Some(extra_data),
|
||||
created_at_unix_secs,
|
||||
created_at_unix_ms,
|
||||
"gateway local video decision request candidate upsert failed",
|
||||
)
|
||||
.await;
|
||||
@@ -186,6 +218,7 @@ pub(super) async fn mark_skipped_local_video_candidate(
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
input.required_capabilities.as_ref(),
|
||||
skip_reason,
|
||||
current_unix_secs(),
|
||||
"gateway local video decision failed to persist skipped candidate",
|
||||
|
||||
@@ -5,7 +5,13 @@ use serde_json::json;
|
||||
use tracing::warn;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
|
||||
use crate::ai_pipeline::conversion::{
|
||||
request_candidate_api_formats, request_conversion_kind,
|
||||
request_conversion_requires_enable_flag, request_pair_allowed_for_transport,
|
||||
};
|
||||
use crate::ai_pipeline::planner::candidate_affinity::{
|
||||
rank_local_execution_candidates, remember_scheduler_affinity_for_candidate,
|
||||
};
|
||||
use crate::ai_pipeline::{
|
||||
resolve_local_decision_execution_runtime_auth_context, ConversionMode, ExecutionStrategy,
|
||||
GatewayControlDecision,
|
||||
@@ -16,7 +22,7 @@ use crate::{append_execution_contract_fields_to_value, AppState, GatewayError};
|
||||
|
||||
use super::{
|
||||
LocalStandardCandidateAttempt, LocalStandardDecisionInput, LocalStandardSourceFamily,
|
||||
LocalStandardSourceMode, LocalStandardSpec,
|
||||
LocalStandardSpec,
|
||||
};
|
||||
|
||||
pub(super) async fn resolve_local_standard_decision_input(
|
||||
@@ -63,10 +69,20 @@ pub(super) async fn resolve_local_standard_decision_input(
|
||||
}
|
||||
};
|
||||
|
||||
let required_capabilities = planner_state
|
||||
.resolve_request_candidate_required_capabilities(
|
||||
&auth_context.user_id,
|
||||
&auth_context.api_key_id,
|
||||
Some(requested_model.as_str()),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
Some(LocalStandardDecisionInput {
|
||||
auth_context,
|
||||
requested_model,
|
||||
auth_snapshot,
|
||||
required_capabilities,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -79,8 +95,10 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let mut seen_candidates = BTreeSet::new();
|
||||
let mut candidates = Vec::new();
|
||||
for candidate_api_format in candidate_api_formats_for_spec(spec) {
|
||||
let auth_snapshot = if *candidate_api_format == spec.api_format {
|
||||
for candidate_api_format in
|
||||
request_candidate_api_formats(spec.api_format, spec.require_streaming)
|
||||
{
|
||||
let auth_snapshot = if candidate_api_format == spec.api_format {
|
||||
Some(&input.auth_snapshot)
|
||||
} else {
|
||||
None
|
||||
@@ -90,6 +108,7 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
|
||||
candidate_api_format,
|
||||
&input.requested_model,
|
||||
spec.require_streaming,
|
||||
input.required_capabilities.as_ref(),
|
||||
auth_snapshot,
|
||||
current_unix_secs(),
|
||||
)
|
||||
@@ -118,13 +137,71 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
|
||||
}
|
||||
}
|
||||
}
|
||||
let candidates = prefer_local_tunnel_owner_candidates(planner_state, candidates).await;
|
||||
let candidates = rank_local_execution_candidates(
|
||||
planner_state,
|
||||
candidates,
|
||||
spec.api_format,
|
||||
input.required_capabilities.as_ref(),
|
||||
)
|
||||
.await;
|
||||
|
||||
let created_at_unix_secs = current_unix_secs();
|
||||
let created_at_unix_ms = current_unix_secs();
|
||||
let mut attempts = Vec::with_capacity(candidates.len());
|
||||
let mut affinity_remembered = false;
|
||||
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
|
||||
let candidate_id = Uuid::new_v4().to_string();
|
||||
let provider_api_format = candidate.endpoint_api_format.trim().to_ascii_lowercase();
|
||||
if provider_api_format != spec.api_format {
|
||||
if let Ok(Some(transport)) = planner_state
|
||||
.read_provider_transport_snapshot(
|
||||
&candidate.provider_id,
|
||||
&candidate.endpoint_id,
|
||||
&candidate.key_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
if !request_pair_allowed_for_transport(
|
||||
&transport,
|
||||
spec.api_format,
|
||||
provider_api_format.as_str(),
|
||||
) {
|
||||
let skip_reason =
|
||||
if request_conversion_kind(spec.api_format, provider_api_format.as_str())
|
||||
.is_some()
|
||||
&& request_conversion_requires_enable_flag(
|
||||
spec.api_format,
|
||||
provider_api_format.as_str(),
|
||||
)
|
||||
&& !transport.provider.enable_format_conversion
|
||||
{
|
||||
"format_conversion_disabled"
|
||||
} else {
|
||||
"transport_unsupported"
|
||||
};
|
||||
super::payload::mark_skipped_local_standard_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index as u32,
|
||||
&candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
if !affinity_remembered {
|
||||
remember_scheduler_affinity_for_candidate(
|
||||
planner_state,
|
||||
Some(&input.auth_snapshot),
|
||||
spec.api_format,
|
||||
&input.requested_model,
|
||||
&candidate,
|
||||
);
|
||||
affinity_remembered = true;
|
||||
}
|
||||
let execution_strategy = if provider_api_format == spec.api_format {
|
||||
ExecutionStrategy::LocalSameFormat
|
||||
} else {
|
||||
@@ -166,8 +243,9 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
|
||||
&candidate,
|
||||
candidate_index as u32,
|
||||
&candidate_id,
|
||||
input.required_capabilities.as_ref(),
|
||||
Some(extra_data),
|
||||
created_at_unix_secs,
|
||||
created_at_unix_ms,
|
||||
"gateway local standard decision request candidate upsert failed",
|
||||
)
|
||||
.await;
|
||||
@@ -213,20 +291,6 @@ fn auth_snapshot_allows_cross_format_candidate(
|
||||
true
|
||||
}
|
||||
|
||||
fn candidate_api_formats_for_spec(spec: LocalStandardSpec) -> &'static [&'static str] {
|
||||
match spec.mode {
|
||||
LocalStandardSourceMode::Chat | LocalStandardSourceMode::Cli => &[
|
||||
"openai:chat",
|
||||
"openai:cli",
|
||||
"openai:compact",
|
||||
"claude:chat",
|
||||
"claude:cli",
|
||||
"gemini:chat",
|
||||
"gemini:cli",
|
||||
],
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_gemini_model_from_path(path: &str) -> Option<String> {
|
||||
let marker = "/models/";
|
||||
let start = path.find(marker)? + marker.len();
|
||||
|
||||
@@ -18,6 +18,7 @@ pub(super) struct LocalStandardDecisionInput {
|
||||
pub(super) auth_context: ExecutionRuntimeAuthContext,
|
||||
pub(super) requested_model: String,
|
||||
pub(super) auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
pub(super) required_capabilities: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
|
||||
@@ -91,6 +91,33 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
|
||||
}
|
||||
};
|
||||
|
||||
if !crate::ai_pipeline::conversion::request_pair_allowed_for_transport(
|
||||
&transport,
|
||||
spec.api_format,
|
||||
provider_api_format.as_str(),
|
||||
) {
|
||||
let skip_reason = if crate::ai_pipeline::conversion::request_conversion_requires_enable_flag(
|
||||
spec.api_format,
|
||||
provider_api_format.as_str(),
|
||||
) && !transport.provider.enable_format_conversion
|
||||
{
|
||||
"format_conversion_disabled"
|
||||
} else {
|
||||
"transport_unsupported"
|
||||
};
|
||||
mark_skipped_local_standard_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
if !crate::ai_pipeline::conversion::request_conversion_transport_supported(
|
||||
&transport,
|
||||
conversion_kind,
|
||||
@@ -301,6 +328,8 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
|
||||
json!({
|
||||
"user_id": input.auth_context.user_id,
|
||||
"api_key_id": input.auth_context.api_key_id,
|
||||
"username": input.auth_context.username,
|
||||
"api_key_name": input.auth_context.api_key_name,
|
||||
"request_id": trace_id,
|
||||
"candidate_id": candidate_id,
|
||||
"candidate_index": candidate_index,
|
||||
@@ -349,6 +378,7 @@ pub(super) async fn mark_skipped_local_standard_candidate(
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
input.required_capabilities.as_ref(),
|
||||
skip_reason,
|
||||
current_unix_secs(),
|
||||
"gateway local standard decision failed to persist skipped candidate",
|
||||
|
||||
@@ -112,11 +112,6 @@ pub(crate) fn build_cross_format_openai_chat_upstream_url(
|
||||
parts.uri.query(),
|
||||
false,
|
||||
)),
|
||||
RequestConversionKind::ToOpenAICompact => Some(build_openai_cli_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
true,
|
||||
)),
|
||||
_ => None,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -10,8 +10,8 @@ use crate::ai_pipeline::transport::antigravity::{
|
||||
};
|
||||
use crate::ai_pipeline::transport::apply_local_body_rules;
|
||||
use crate::ai_pipeline::transport::url::{
|
||||
build_claude_messages_url, build_gemini_content_url, build_openai_cli_url,
|
||||
build_passthrough_path_url,
|
||||
build_claude_messages_url, build_gemini_content_url, build_openai_chat_url,
|
||||
build_openai_cli_url, build_passthrough_path_url,
|
||||
};
|
||||
use crate::ai_pipeline::{
|
||||
build_cross_format_openai_cli_request_body as pipeline_build_cross_format_openai_cli_request_body,
|
||||
@@ -140,16 +140,15 @@ pub(crate) fn build_cross_format_openai_cli_upstream_url(
|
||||
build_passthrough_path_url(&transport.endpoint.base_url, path, parts.uri.query(), &[])
|
||||
}
|
||||
None => match conversion_kind {
|
||||
RequestConversionKind::ToOpenAIChat => Some(build_openai_chat_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
)),
|
||||
RequestConversionKind::ToOpenAIFamilyCli => Some(build_openai_cli_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
false,
|
||||
)),
|
||||
RequestConversionKind::ToOpenAICompact => Some(build_openai_cli_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
true,
|
||||
)),
|
||||
RequestConversionKind::ToClaudeStandard => Some(build_claude_messages_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
@@ -160,7 +159,6 @@ pub(crate) fn build_cross_format_openai_cli_upstream_url(
|
||||
upstream_is_stream,
|
||||
parts.uri.query(),
|
||||
),
|
||||
_ => None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,7 +3,7 @@ use serde_json::json;
|
||||
use super::build_cross_format_openai_cli_request_body;
|
||||
|
||||
#[test]
|
||||
fn builds_openai_family_cross_format_request_body_from_compact_source() {
|
||||
fn builds_openai_chat_cross_format_request_body_from_openai_cli_source() {
|
||||
let body_json = json!({
|
||||
"model": "gpt-5",
|
||||
"input": "hello",
|
||||
@@ -12,18 +12,18 @@ fn builds_openai_family_cross_format_request_body_from_compact_source() {
|
||||
let provider_request_body = build_cross_format_openai_cli_request_body(
|
||||
&body_json,
|
||||
"gpt-5-upstream",
|
||||
"openai:compact",
|
||||
"openai:cli",
|
||||
"openai:chat",
|
||||
false,
|
||||
"openai",
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("compact to openai cli body should build");
|
||||
.expect("openai cli to openai chat body should build");
|
||||
|
||||
assert_eq!(provider_request_body["model"], "gpt-5-upstream");
|
||||
assert_eq!(provider_request_body["input"][0]["type"], "message");
|
||||
assert_eq!(provider_request_body["input"][0]["role"], "user");
|
||||
assert_eq!(provider_request_body["messages"][0]["role"], "user");
|
||||
assert_eq!(provider_request_body["messages"][0]["content"], "hello");
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -96,7 +96,7 @@ pub(super) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
|
||||
)
|
||||
.await
|
||||
}
|
||||
"claude:chat" | "gemini:chat" | "openai:cli" | "openai:compact" => {
|
||||
"claude:chat" | "claude:cli" | "gemini:chat" | "gemini:cli" | "openai:cli" => {
|
||||
build_cross_format_local_openai_chat_decision_payload_for_candidate(
|
||||
state,
|
||||
parts,
|
||||
|
||||
@@ -6,7 +6,9 @@ use tracing::warn;
|
||||
|
||||
use crate::ai_pipeline::collect_control_headers;
|
||||
use crate::ai_pipeline::conversion::{
|
||||
request_conversion_direct_auth, request_conversion_kind, request_conversion_transport_supported,
|
||||
request_conversion_direct_auth, request_conversion_kind,
|
||||
request_conversion_requires_enable_flag, request_conversion_transport_supported,
|
||||
request_pair_allowed_for_transport,
|
||||
};
|
||||
use crate::ai_pipeline::planner::common::OPENAI_CHAT_STREAM_PLAN_KIND;
|
||||
use crate::ai_pipeline::planner::standard::{
|
||||
@@ -50,6 +52,28 @@ pub(super) async fn build_cross_format_local_openai_chat_decision_payload_for_ca
|
||||
else {
|
||||
return None;
|
||||
};
|
||||
if !request_pair_allowed_for_transport(&transport, "openai:chat", provider_api_format.as_str())
|
||||
{
|
||||
let skip_reason =
|
||||
if request_conversion_requires_enable_flag("openai:chat", provider_api_format.as_str())
|
||||
&& !transport.provider.enable_format_conversion
|
||||
{
|
||||
"format_conversion_disabled"
|
||||
} else {
|
||||
"transport_unsupported"
|
||||
};
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
if !request_conversion_transport_supported(transport, conversion_kind) {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
@@ -63,7 +87,6 @@ pub(super) async fn build_cross_format_local_openai_chat_decision_payload_for_ca
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let resolve_auth = request_conversion_direct_auth(transport, conversion_kind);
|
||||
let oauth_auth = if resolve_auth.is_none() {
|
||||
match planner_state
|
||||
@@ -258,6 +281,8 @@ pub(super) async fn build_cross_format_local_openai_chat_decision_payload_for_ca
|
||||
json!({
|
||||
"user_id": input.auth_context.user_id,
|
||||
"api_key_id": input.auth_context.api_key_id,
|
||||
"username": input.auth_context.username,
|
||||
"api_key_name": input.auth_context.api_key_name,
|
||||
"request_id": trace_id,
|
||||
"candidate_id": candidate_id,
|
||||
"candidate_index": candidate_index,
|
||||
|
||||
@@ -233,6 +233,8 @@ pub(super) async fn build_same_format_local_openai_chat_decision_payload_for_can
|
||||
json!({
|
||||
"user_id": input.auth_context.user_id,
|
||||
"api_key_id": input.auth_context.api_key_id,
|
||||
"username": input.auth_context.username,
|
||||
"api_key_name": input.auth_context.api_key_name,
|
||||
"request_id": trace_id,
|
||||
"candidate_id": candidate_id,
|
||||
"candidate_index": candidate_index,
|
||||
|
||||
@@ -3,7 +3,13 @@ use serde_json::json;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
||||
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
|
||||
use crate::ai_pipeline::conversion::{
|
||||
request_conversion_kind, request_conversion_requires_enable_flag,
|
||||
request_pair_allowed_for_transport,
|
||||
};
|
||||
use crate::ai_pipeline::planner::candidate_affinity::{
|
||||
rank_local_execution_candidates, remember_scheduler_affinity_for_candidate,
|
||||
};
|
||||
use crate::ai_pipeline::GatewayAuthApiKeySnapshot;
|
||||
use crate::ai_pipeline::{ConversionMode, ExecutionStrategy, PlannerAppState};
|
||||
use crate::clock::current_unix_secs;
|
||||
@@ -14,6 +20,7 @@ pub(crate) struct LocalOpenAiChatDecisionInput {
|
||||
pub(crate) auth_context: ExecutionRuntimeAuthContext,
|
||||
pub(crate) requested_model: String,
|
||||
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
pub(crate) required_capabilities: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -48,6 +55,7 @@ pub(crate) async fn mark_skipped_local_openai_chat_candidate(
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
input.required_capabilities.as_ref(),
|
||||
skip_reason,
|
||||
current_unix_secs(),
|
||||
"gateway local openai chat decision failed to persist skipped candidate",
|
||||
@@ -62,13 +70,71 @@ pub(crate) async fn materialize_local_openai_chat_candidate_attempts(
|
||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
) -> Vec<LocalOpenAiChatCandidateAttempt> {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let candidates = prefer_local_tunnel_owner_candidates(planner_state, candidates).await;
|
||||
let created_at_unix_secs = current_unix_secs();
|
||||
let candidates = rank_local_execution_candidates(
|
||||
planner_state,
|
||||
candidates,
|
||||
"openai:chat",
|
||||
input.required_capabilities.as_ref(),
|
||||
)
|
||||
.await;
|
||||
let created_at_unix_ms = current_unix_secs();
|
||||
let mut attempts = Vec::with_capacity(candidates.len());
|
||||
let mut affinity_remembered = false;
|
||||
|
||||
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
|
||||
let generated_candidate_id = Uuid::new_v4().to_string();
|
||||
let provider_api_format = candidate.endpoint_api_format.trim().to_ascii_lowercase();
|
||||
if provider_api_format != "openai:chat" {
|
||||
if let Ok(Some(transport)) = planner_state
|
||||
.read_provider_transport_snapshot(
|
||||
&candidate.provider_id,
|
||||
&candidate.endpoint_id,
|
||||
&candidate.key_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
if !request_pair_allowed_for_transport(
|
||||
&transport,
|
||||
"openai:chat",
|
||||
provider_api_format.as_str(),
|
||||
) {
|
||||
let skip_reason =
|
||||
if request_conversion_kind("openai:chat", provider_api_format.as_str())
|
||||
.is_some()
|
||||
&& request_conversion_requires_enable_flag(
|
||||
"openai:chat",
|
||||
provider_api_format.as_str(),
|
||||
)
|
||||
&& !transport.provider.enable_format_conversion
|
||||
{
|
||||
"format_conversion_disabled"
|
||||
} else {
|
||||
"transport_unsupported"
|
||||
};
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index as u32,
|
||||
&generated_candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
if !affinity_remembered {
|
||||
remember_scheduler_affinity_for_candidate(
|
||||
planner_state,
|
||||
Some(&input.auth_snapshot),
|
||||
"openai:chat",
|
||||
&input.requested_model,
|
||||
&candidate,
|
||||
);
|
||||
affinity_remembered = true;
|
||||
}
|
||||
let (execution_strategy, conversion_mode) = if provider_api_format == "openai:chat" {
|
||||
(ExecutionStrategy::LocalSameFormat, ConversionMode::None)
|
||||
} else {
|
||||
@@ -103,8 +169,9 @@ pub(crate) async fn materialize_local_openai_chat_candidate_attempts(
|
||||
&candidate,
|
||||
candidate_index as u32,
|
||||
&generated_candidate_id,
|
||||
input.required_capabilities.as_ref(),
|
||||
Some(extra_data),
|
||||
created_at_unix_secs,
|
||||
created_at_unix_ms,
|
||||
"gateway local openai chat decision request candidate upsert failed",
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -3,6 +3,7 @@ use std::collections::BTreeSet;
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
|
||||
use super::super::{GatewayError, LocalOpenAiChatDecisionInput};
|
||||
use crate::ai_pipeline::conversion::request_candidate_api_formats;
|
||||
use crate::ai_pipeline::{GatewayAuthApiKeySnapshot, PlannerAppState};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::AppState;
|
||||
@@ -17,17 +18,7 @@ pub(crate) async fn list_local_openai_chat_candidates(
|
||||
let mut combined = Vec::new();
|
||||
let mut seen = BTreeSet::new();
|
||||
|
||||
let api_formats = if require_streaming {
|
||||
vec!["openai:chat", "claude:chat", "gemini:chat", "openai:cli"]
|
||||
} else {
|
||||
vec![
|
||||
"openai:chat",
|
||||
"claude:chat",
|
||||
"gemini:chat",
|
||||
"openai:cli",
|
||||
"openai:compact",
|
||||
]
|
||||
};
|
||||
let api_formats = request_candidate_api_formats("openai:chat", require_streaming);
|
||||
|
||||
for api_format in api_formats {
|
||||
let auth_snapshot = if api_format == "openai:chat" {
|
||||
@@ -40,6 +31,7 @@ pub(crate) async fn list_local_openai_chat_candidates(
|
||||
api_format,
|
||||
&input.requested_model,
|
||||
require_streaming,
|
||||
input.required_capabilities.as_ref(),
|
||||
auth_snapshot,
|
||||
now_unix_secs,
|
||||
)
|
||||
|
||||
@@ -108,9 +108,19 @@ pub(crate) async fn resolve_local_openai_chat_decision_input(
|
||||
}
|
||||
};
|
||||
|
||||
let required_capabilities = planner_state
|
||||
.resolve_request_candidate_required_capabilities(
|
||||
&auth_context.user_id,
|
||||
&auth_context.api_key_id,
|
||||
Some(requested_model.as_str()),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
Some(LocalOpenAiChatDecisionInput {
|
||||
auth_context,
|
||||
requested_model,
|
||||
auth_snapshot,
|
||||
required_capabilities,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -95,6 +95,8 @@ pub(crate) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
|
||||
json!({
|
||||
"user_id": input.auth_context.user_id,
|
||||
"api_key_id": input.auth_context.api_key_id,
|
||||
"username": input.auth_context.username,
|
||||
"api_key_name": input.auth_context.api_key_name,
|
||||
"request_id": trace_id,
|
||||
"candidate_id": candidate_id,
|
||||
"candidate_index": candidate_index,
|
||||
|
||||
@@ -5,7 +5,9 @@ use serde_json::Value;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_pipeline::conversion::{
|
||||
request_conversion_direct_auth, request_conversion_kind, request_conversion_transport_supported,
|
||||
request_conversion_direct_auth, request_conversion_kind,
|
||||
request_conversion_requires_enable_flag, request_conversion_transport_supported,
|
||||
request_pair_allowed_for_transport,
|
||||
};
|
||||
use crate::ai_pipeline::planner::common::force_upstream_streaming_for_provider;
|
||||
use crate::ai_pipeline::planner::standard::{
|
||||
@@ -115,6 +117,36 @@ pub(crate) async fn resolve_local_openai_cli_candidate_payload_parts(
|
||||
|
||||
let same_format = provider_api_format == spec.api_format.trim().to_ascii_lowercase();
|
||||
let conversion_kind = request_conversion_kind(spec.api_format, provider_api_format.as_str());
|
||||
if !same_format
|
||||
&& !request_pair_allowed_for_transport(
|
||||
&transport,
|
||||
spec.api_format,
|
||||
provider_api_format.as_str(),
|
||||
)
|
||||
{
|
||||
let skip_reason = if conversion_kind.is_some()
|
||||
&& request_conversion_requires_enable_flag(
|
||||
spec.api_format,
|
||||
provider_api_format.as_str(),
|
||||
)
|
||||
&& !transport.provider.enable_format_conversion
|
||||
{
|
||||
"format_conversion_disabled"
|
||||
} else {
|
||||
"transport_unsupported"
|
||||
};
|
||||
mark_skipped_local_openai_cli_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
let transport_supported = if same_format {
|
||||
supports_local_standard_transport_with_network(&transport, provider_api_format.as_str())
|
||||
} else {
|
||||
|
||||
@@ -6,8 +6,13 @@ use tracing::warn;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
||||
use crate::ai_pipeline::conversion::request_conversion_kind;
|
||||
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
|
||||
use crate::ai_pipeline::conversion::{
|
||||
request_candidate_api_formats, request_conversion_kind,
|
||||
request_conversion_requires_enable_flag, request_pair_allowed_for_transport,
|
||||
};
|
||||
use crate::ai_pipeline::planner::candidate_affinity::{
|
||||
rank_local_execution_candidates, remember_scheduler_affinity_for_candidate,
|
||||
};
|
||||
use crate::ai_pipeline::{
|
||||
resolve_local_decision_execution_runtime_auth_context, ConversionMode, ExecutionStrategy,
|
||||
GatewayControlDecision,
|
||||
@@ -23,6 +28,7 @@ pub(crate) struct LocalOpenAiCliDecisionInput {
|
||||
pub(crate) auth_context: ExecutionRuntimeAuthContext,
|
||||
pub(crate) requested_model: String,
|
||||
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
pub(crate) required_capabilities: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -70,10 +76,20 @@ pub(crate) async fn resolve_local_openai_cli_decision_input(
|
||||
}
|
||||
};
|
||||
|
||||
let required_capabilities = planner_state
|
||||
.resolve_request_candidate_required_capabilities(
|
||||
&auth_context.user_id,
|
||||
&auth_context.api_key_id,
|
||||
Some(requested_model.as_str()),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
Some(LocalOpenAiCliDecisionInput {
|
||||
auth_context,
|
||||
requested_model,
|
||||
auth_snapshot,
|
||||
required_capabilities,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -86,8 +102,10 @@ pub(crate) async fn materialize_local_openai_cli_candidate_attempts(
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let mut seen_candidates = BTreeSet::new();
|
||||
let mut candidates = Vec::new();
|
||||
for candidate_api_format in candidate_api_formats_for_spec(spec) {
|
||||
let auth_snapshot = if *candidate_api_format == spec.api_format {
|
||||
for candidate_api_format in
|
||||
request_candidate_api_formats(spec.api_format, spec.require_streaming)
|
||||
{
|
||||
let auth_snapshot = if candidate_api_format == spec.api_format {
|
||||
Some(&input.auth_snapshot)
|
||||
} else {
|
||||
None
|
||||
@@ -97,6 +115,7 @@ pub(crate) async fn materialize_local_openai_cli_candidate_attempts(
|
||||
candidate_api_format,
|
||||
&input.requested_model,
|
||||
spec.require_streaming,
|
||||
input.required_capabilities.as_ref(),
|
||||
auth_snapshot,
|
||||
current_unix_secs(),
|
||||
)
|
||||
@@ -125,13 +144,71 @@ pub(crate) async fn materialize_local_openai_cli_candidate_attempts(
|
||||
}
|
||||
}
|
||||
}
|
||||
let candidates = prefer_local_tunnel_owner_candidates(planner_state, candidates).await;
|
||||
let candidates = rank_local_execution_candidates(
|
||||
planner_state,
|
||||
candidates,
|
||||
spec.api_format,
|
||||
input.required_capabilities.as_ref(),
|
||||
)
|
||||
.await;
|
||||
|
||||
let created_at_unix_secs = current_unix_secs();
|
||||
let created_at_unix_ms = current_unix_secs();
|
||||
let mut attempts = Vec::with_capacity(candidates.len());
|
||||
let mut affinity_remembered = false;
|
||||
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
|
||||
let generated_candidate_id = Uuid::new_v4().to_string();
|
||||
let provider_api_format = candidate.endpoint_api_format.trim().to_ascii_lowercase();
|
||||
if provider_api_format != spec.api_format.trim().to_ascii_lowercase() {
|
||||
if let Ok(Some(transport)) = planner_state
|
||||
.read_provider_transport_snapshot(
|
||||
&candidate.provider_id,
|
||||
&candidate.endpoint_id,
|
||||
&candidate.key_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
if !request_pair_allowed_for_transport(
|
||||
&transport,
|
||||
spec.api_format,
|
||||
provider_api_format.as_str(),
|
||||
) {
|
||||
let skip_reason =
|
||||
if request_conversion_kind(spec.api_format, provider_api_format.as_str())
|
||||
.is_some()
|
||||
&& request_conversion_requires_enable_flag(
|
||||
spec.api_format,
|
||||
provider_api_format.as_str(),
|
||||
)
|
||||
&& !transport.provider.enable_format_conversion
|
||||
{
|
||||
"format_conversion_disabled"
|
||||
} else {
|
||||
"transport_unsupported"
|
||||
};
|
||||
mark_skipped_local_openai_cli_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index as u32,
|
||||
&generated_candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
if !affinity_remembered {
|
||||
remember_scheduler_affinity_for_candidate(
|
||||
planner_state,
|
||||
Some(&input.auth_snapshot),
|
||||
spec.api_format,
|
||||
&input.requested_model,
|
||||
&candidate,
|
||||
);
|
||||
affinity_remembered = true;
|
||||
}
|
||||
let execution_strategy =
|
||||
if provider_api_format == spec.api_format.trim().to_ascii_lowercase() {
|
||||
ExecutionStrategy::LocalSameFormat
|
||||
@@ -170,8 +247,9 @@ pub(crate) async fn materialize_local_openai_cli_candidate_attempts(
|
||||
&candidate,
|
||||
candidate_index as u32,
|
||||
&generated_candidate_id,
|
||||
input.required_capabilities.as_ref(),
|
||||
Some(extra_data),
|
||||
created_at_unix_secs,
|
||||
created_at_unix_ms,
|
||||
"gateway local openai cli decision request candidate upsert failed",
|
||||
)
|
||||
.await;
|
||||
@@ -217,14 +295,6 @@ fn auth_snapshot_allows_cross_format_openai_cli_candidate(
|
||||
true
|
||||
}
|
||||
|
||||
fn candidate_api_formats_for_spec(spec: LocalOpenAiCliSpec) -> &'static [&'static str] {
|
||||
match spec.api_format {
|
||||
"openai:compact" => &["openai:compact", "openai:cli", "claude:cli", "gemini:cli"],
|
||||
"openai:cli" => &["openai:cli", "claude:cli", "gemini:cli"],
|
||||
_ => &[],
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn mark_skipped_local_openai_cli_candidate(
|
||||
state: &AppState,
|
||||
input: &LocalOpenAiCliDecisionInput,
|
||||
@@ -242,6 +312,7 @@ pub(crate) async fn mark_skipped_local_openai_cli_candidate(
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
input.required_capabilities.as_ref(),
|
||||
skip_reason,
|
||||
current_unix_secs(),
|
||||
"gateway local openai cli decision failed to persist skipped candidate",
|
||||
|
||||
@@ -4,6 +4,23 @@ use serde_json::Value;
|
||||
use super::PlannerAppState;
|
||||
|
||||
impl<'a> PlannerAppState<'a> {
|
||||
pub(crate) async fn resolve_request_candidate_required_capabilities(
|
||||
self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
requested_model: Option<&str>,
|
||||
explicit_required_capabilities: Option<&Value>,
|
||||
) -> Option<Value> {
|
||||
crate::request_candidate_runtime::resolve_request_candidate_required_capabilities(
|
||||
self.app(),
|
||||
user_id,
|
||||
api_key_id,
|
||||
requested_model,
|
||||
explicit_required_capabilities,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn persist_available_local_candidate(
|
||||
self,
|
||||
@@ -13,8 +30,9 @@ impl<'a> PlannerAppState<'a> {
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
required_capabilities: Option<&Value>,
|
||||
extra_data: Option<Value>,
|
||||
created_at_unix_secs: u64,
|
||||
created_at_unix_ms: u64,
|
||||
error_context: &'static str,
|
||||
) -> String {
|
||||
crate::request_candidate_runtime::persist_available_local_candidate(
|
||||
@@ -25,8 +43,9 @@ impl<'a> PlannerAppState<'a> {
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
required_capabilities,
|
||||
extra_data,
|
||||
created_at_unix_secs,
|
||||
created_at_unix_ms,
|
||||
error_context,
|
||||
)
|
||||
.await
|
||||
@@ -41,8 +60,9 @@ impl<'a> PlannerAppState<'a> {
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
required_capabilities: Option<&Value>,
|
||||
skip_reason: &str,
|
||||
finished_at_unix_secs: u64,
|
||||
finished_at_unix_ms: u64,
|
||||
error_context: &'static str,
|
||||
) {
|
||||
crate::request_candidate_runtime::persist_skipped_local_candidate(
|
||||
@@ -53,8 +73,9 @@ impl<'a> PlannerAppState<'a> {
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
required_capabilities,
|
||||
skip_reason,
|
||||
finished_at_unix_secs,
|
||||
finished_at_unix_ms,
|
||||
error_context,
|
||||
)
|
||||
.await
|
||||
|
||||
@@ -9,6 +9,7 @@ impl<'a> PlannerAppState<'a> {
|
||||
api_format: &str,
|
||||
global_model_name: &str,
|
||||
require_streaming: bool,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
|
||||
@@ -18,6 +19,7 @@ impl<'a> PlannerAppState<'a> {
|
||||
api_format,
|
||||
global_model_name,
|
||||
require_streaming,
|
||||
required_capabilities,
|
||||
auth_snapshot,
|
||||
now_unix_secs,
|
||||
)
|
||||
|
||||
@@ -44,18 +44,20 @@ pub(crate) use aether_ai_pipeline::api::{
|
||||
provider_adaptation_descriptor_for_provider_type,
|
||||
provider_adaptation_requires_eventstream_accept,
|
||||
provider_adaptation_should_unwrap_stream_envelope,
|
||||
provider_private_response_allows_sync_finalize, request_conversion_direct_auth,
|
||||
request_conversion_kind, request_conversion_transport_supported, resolve_claude_stream_spec,
|
||||
resolve_claude_sync_spec, resolve_execution_runtime_stream_plan_kind,
|
||||
resolve_execution_runtime_sync_plan_kind, resolve_finalize_stream_rewrite_mode,
|
||||
resolve_gemini_files_stream_spec, resolve_gemini_files_sync_spec, resolve_gemini_stream_spec,
|
||||
resolve_gemini_sync_spec, resolve_local_same_format_stream_spec,
|
||||
resolve_local_same_format_sync_spec, resolve_local_video_sync_spec,
|
||||
resolve_openai_chat_max_tokens, resolve_openai_cli_stream_spec, resolve_openai_cli_sync_spec,
|
||||
stream_body_contains_error_event, supports_stream_scheduler_decision_kind,
|
||||
supports_sync_scheduler_decision_kind, sync_chat_response_conversion_kind,
|
||||
sync_cli_response_conversion_kind, transform_provider_private_stream_line, value_as_u64,
|
||||
CanonicalStreamFrame, ClaudeClientEmitter, ClaudeProviderState, ExecutionRuntimeAuthContext,
|
||||
provider_private_response_allows_sync_finalize, request_candidate_api_formats,
|
||||
request_conversion_direct_auth, request_conversion_kind,
|
||||
request_conversion_requires_enable_flag, request_conversion_transport_supported,
|
||||
request_pair_allowed_for_transport, resolve_claude_stream_spec, resolve_claude_sync_spec,
|
||||
resolve_execution_runtime_stream_plan_kind, resolve_execution_runtime_sync_plan_kind,
|
||||
resolve_finalize_stream_rewrite_mode, resolve_gemini_files_stream_spec,
|
||||
resolve_gemini_files_sync_spec, resolve_gemini_stream_spec, resolve_gemini_sync_spec,
|
||||
resolve_local_same_format_stream_spec, resolve_local_same_format_sync_spec,
|
||||
resolve_local_video_sync_spec, resolve_openai_chat_max_tokens, resolve_openai_cli_stream_spec,
|
||||
resolve_openai_cli_sync_spec, stream_body_contains_error_event,
|
||||
supports_stream_scheduler_decision_kind, supports_sync_scheduler_decision_kind,
|
||||
sync_chat_response_conversion_kind, sync_cli_response_conversion_kind,
|
||||
transform_provider_private_stream_line, value_as_u64, CanonicalStreamFrame,
|
||||
ClaudeClientEmitter, ClaudeProviderState, ExecutionRuntimeAuthContext,
|
||||
FinalizeStreamRewriteMode, GatewayControlPlanRequest, GatewayControlPlanResponse,
|
||||
GatewayControlSyncDecisionResponse, GeminiClientEmitter, GeminiProviderState,
|
||||
LocalCoreSyncErrorKind, LocalGeminiFilesSpec, LocalOpenAiCliSpec,
|
||||
|
||||
@@ -253,7 +253,7 @@ async fn persist_cancelled_video_task(
|
||||
next_poll_at_unix_secs: None,
|
||||
poll_count: task.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,
|
||||
completed_at_unix_secs: Some(now_unix_secs),
|
||||
updated_at_unix_secs: now_unix_secs,
|
||||
|
||||
@@ -272,7 +272,7 @@ fn build_successful_poll_update(
|
||||
record.resolution = task.resolution.clone();
|
||||
record.aspect_ratio = task.aspect_ratio.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.updated_at_unix_secs = now_unix_secs;
|
||||
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,
|
||||
poll_count: task.poll_count,
|
||||
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,
|
||||
completed_at_unix_secs: task.completed_at_unix_secs,
|
||||
updated_at_unix_secs: task.updated_at_unix_secs,
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
use axum::body::Body;
|
||||
use axum::http::{self, Response};
|
||||
use axum::http::{self, Response, StatusCode};
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::control::GatewayControlDecision;
|
||||
@@ -68,12 +68,8 @@ pub(crate) fn emit_admin_audit(
|
||||
)
|
||||
};
|
||||
|
||||
let audit_status = if response.status().is_success() {
|
||||
"completed"
|
||||
} else {
|
||||
"failed"
|
||||
};
|
||||
if response.status().is_success() {
|
||||
let (audit_status, log_level) = classify_admin_audit_response(method, response.status());
|
||||
if log_level == AdminAuditLogLevel::Info {
|
||||
info!(
|
||||
event_name,
|
||||
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 {
|
||||
route_family
|
||||
.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
|
||||
)
|
||||
}
|
||||
|
||||
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)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
4
apps/aether-gateway/src/cache/mod.rs
vendored
4
apps/aether-gateway/src/cache/mod.rs
vendored
@@ -6,4 +6,6 @@ mod scheduler_affinity;
|
||||
pub(crate) use auth_api_key_last_used::AuthApiKeyLastUsedCache;
|
||||
pub(crate) use auth_context::AuthContextCache;
|
||||
pub(crate) use direct_plan_bypass::DirectPlanBypassCache;
|
||||
pub(crate) use scheduler_affinity::{SchedulerAffinityCache, SchedulerAffinityTarget};
|
||||
pub(crate) use scheduler_affinity::{
|
||||
SchedulerAffinityCache, SchedulerAffinitySnapshotEntry, SchedulerAffinityTarget,
|
||||
};
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use aether_cache::ExpiringMap;
|
||||
use aether_cache::{ExpiringMap, ExpiringMapFreshEntry};
|
||||
pub(crate) use aether_scheduler_core::SchedulerAffinityTarget;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
@@ -8,6 +8,13 @@ pub(crate) struct SchedulerAffinityCache {
|
||||
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 {
|
||||
pub(crate) fn get_fresh(
|
||||
&self,
|
||||
@@ -31,4 +38,18 @@ impl SchedulerAffinityCache {
|
||||
pub(crate) fn remove(&self, cache_key: &str) -> Option<SchedulerAffinityTarget> {
|
||||
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()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,3 +6,10 @@ pub(crate) fn current_unix_secs() -> u64 {
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}
|
||||
|
||||
pub(crate) fn current_unix_ms() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_millis() as u64
|
||||
}
|
||||
|
||||
@@ -27,6 +27,10 @@ const AUTH_CONTEXT_CACHE_MAX_ENTRIES: usize = 256;
|
||||
pub(crate) struct GatewayControlAuthContext {
|
||||
pub(crate) user_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) access_allowed: bool,
|
||||
#[serde(skip)]
|
||||
@@ -489,6 +493,8 @@ pub(super) async fn resolve_data_backed_auth_context(
|
||||
return Ok(Some(GatewayControlAuthContext {
|
||||
user_id: String::new(),
|
||||
api_key_id: String::new(),
|
||||
username: None,
|
||||
api_key_name: None,
|
||||
balance_remaining: None,
|
||||
access_allowed: false,
|
||||
user_rate_limit: None,
|
||||
@@ -539,6 +545,8 @@ async fn resolve_trusted_auth_context(
|
||||
return Ok(Some(GatewayControlAuthContext {
|
||||
user_id: trusted_headers.user_id,
|
||||
api_key_id: trusted_headers.api_key_id,
|
||||
username: None,
|
||||
api_key_name: None,
|
||||
balance_remaining: trusted_headers.balance_remaining,
|
||||
access_allowed: false,
|
||||
user_rate_limit: None,
|
||||
@@ -620,6 +628,8 @@ fn build_data_backed_auth_context(
|
||||
};
|
||||
|
||||
GatewayControlAuthContext {
|
||||
username: Some(snapshot.username.clone()),
|
||||
api_key_name: snapshot.api_key_name.clone(),
|
||||
user_id: snapshot.user_id,
|
||||
api_key_id: snapshot.api_key_id,
|
||||
balance_remaining: wallet_remaining.or(balance_remaining),
|
||||
|
||||
@@ -4,9 +4,10 @@ use aether_scheduler_core::{
|
||||
auth_constraints_allow_api_format, build_minimal_candidate_selection,
|
||||
collect_global_model_names_for_required_capability, normalize_api_format,
|
||||
resolve_requested_global_model_name, SchedulerAuthConstraints,
|
||||
SchedulerMinimalCandidateSelectionCandidate,
|
||||
SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode,
|
||||
};
|
||||
use async_trait::async_trait;
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use super::auth::GatewayAuthApiKeySnapshot;
|
||||
|
||||
@@ -62,6 +63,91 @@ pub(crate) async fn read_minimal_candidate_selection(
|
||||
requested_model_name: &str,
|
||||
require_streaming: bool,
|
||||
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> {
|
||||
let normalized_api_format = normalize_api_format(api_format);
|
||||
if normalized_api_format.is_empty() {
|
||||
@@ -81,17 +167,16 @@ pub(crate) async fn read_minimal_candidate_selection(
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
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(
|
||||
rows,
|
||||
&normalized_api_format,
|
||||
requested_model_name,
|
||||
resolved_global_model_name.as_str(),
|
||||
require_streaming,
|
||||
required_capabilities,
|
||||
auth_constraints.as_ref(),
|
||||
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 {
|
||||
SchedulerAuthConstraints {
|
||||
allowed_providers: snapshot
|
||||
|
||||
@@ -36,7 +36,7 @@ mod tests {
|
||||
request_id: &str,
|
||||
candidate_index: i32,
|
||||
status: RequestCandidateStatus,
|
||||
started_at_unix_secs: Option<i64>,
|
||||
started_at_unix_ms: Option<i64>,
|
||||
latency_ms: Option<i32>,
|
||||
status_code: Option<i32>,
|
||||
) -> StoredRequestCandidate {
|
||||
@@ -62,9 +62,9 @@ mod tests {
|
||||
Some(1),
|
||||
None,
|
||||
None,
|
||||
100 + i64::from(candidate_index),
|
||||
started_at_unix_secs,
|
||||
started_at_unix_secs.map(|value| value + 1),
|
||||
(100 + i64::from(candidate_index)) * 1_000,
|
||||
started_at_unix_ms.map(|v| v * 1_000),
|
||||
started_at_unix_ms.map(|value| (value + 1) * 1_000),
|
||||
)
|
||||
.expect("candidate should build")
|
||||
}
|
||||
|
||||
@@ -101,9 +101,9 @@ mod tests {
|
||||
Some(1),
|
||||
None,
|
||||
Some(serde_json::json!({"cache_1h": true})),
|
||||
100,
|
||||
Some(101),
|
||||
Some(102),
|
||||
100_000,
|
||||
Some(101_000),
|
||||
Some(102_000),
|
||||
)
|
||||
.expect("candidate should build")
|
||||
}
|
||||
@@ -174,11 +174,15 @@ mod tests {
|
||||
provider_name: Some("OpenAI".to_string()),
|
||||
provider_website: Some("https://openai.com".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_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("chat".to_string()),
|
||||
provider_key_name: Some("prod-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_is_active: Some(true),
|
||||
}],
|
||||
|
||||
@@ -28,7 +28,9 @@ use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
|
||||
use aether_data_contracts::repository::video_tasks::{
|
||||
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 serde_json::json;
|
||||
|
||||
@@ -108,7 +110,7 @@ async fn data_state_find_uses_configured_read_repository() {
|
||||
next_poll_at_unix_secs: Some(100),
|
||||
poll_count: 0,
|
||||
max_poll_count: 360,
|
||||
created_at_unix_secs: 100,
|
||||
created_at_unix_ms: 100,
|
||||
submitted_at_unix_secs: Some(100),
|
||||
completed_at_unix_secs: None,
|
||||
updated_at_unix_secs: 100,
|
||||
@@ -384,9 +386,9 @@ async fn data_state_reads_decision_trace_with_provider_catalog_metadata() {
|
||||
Some(1),
|
||||
None,
|
||||
Some(serde_json::json!({"cache_1h": true})),
|
||||
100,
|
||||
Some(101),
|
||||
Some(102),
|
||||
100_000,
|
||||
Some(101_000),
|
||||
Some(102_000),
|
||||
)
|
||||
.expect("candidate should build"),
|
||||
]));
|
||||
@@ -472,9 +474,9 @@ async fn data_state_reads_request_audit_bundle_from_multiple_readers() {
|
||||
Some(1),
|
||||
None,
|
||||
Some(serde_json::json!({"cache_1h": true})),
|
||||
100,
|
||||
Some(101),
|
||||
Some(102),
|
||||
100_000,
|
||||
Some(101_000),
|
||||
Some(102_000),
|
||||
)
|
||||
.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",
|
||||
false,
|
||||
None,
|
||||
Some(&auth_constraints),
|
||||
Some(auth_snapshot.api_key_id.as_str()),
|
||||
SchedulerPriorityMode::Provider,
|
||||
)
|
||||
.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),
|
||||
poll_count: 1,
|
||||
max_poll_count: 360,
|
||||
created_at_unix_secs: 100,
|
||||
created_at_unix_ms: 100,
|
||||
submitted_at_unix_secs: Some(100),
|
||||
completed_at_unix_secs: None,
|
||||
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,
|
||||
poll_count: 4,
|
||||
max_poll_count: 360,
|
||||
created_at_unix_secs: 100,
|
||||
created_at_unix_ms: 100,
|
||||
submitted_at_unix_secs: Some(100),
|
||||
completed_at_unix_secs: Some(120),
|
||||
updated_at_unix_secs: 120,
|
||||
@@ -843,7 +847,7 @@ async fn data_state_write_uses_configured_shadow_result_writer() {
|
||||
match_status: ShadowResultMatchStatus::Pending,
|
||||
status_code: Some(200),
|
||||
error_message: None,
|
||||
created_at_unix_secs: 100,
|
||||
created_at_unix_ms: 100,
|
||||
updated_at_unix_secs: 100,
|
||||
})
|
||||
.await
|
||||
@@ -906,7 +910,7 @@ async fn data_state_records_shadow_result_samples_and_merges_match_status() {
|
||||
.expect("second stored result should exist");
|
||||
|
||||
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.request_id.as_deref(), Some("req-1"));
|
||||
}
|
||||
@@ -986,7 +990,7 @@ fn sample_request_candidate(
|
||||
request_id: &str,
|
||||
candidate_index: i32,
|
||||
status: RequestCandidateStatus,
|
||||
started_at_unix_secs: Option<i64>,
|
||||
started_at_unix_ms: Option<i64>,
|
||||
latency_ms: Option<i32>,
|
||||
status_code: Option<i32>,
|
||||
) -> StoredRequestCandidate {
|
||||
@@ -1013,8 +1017,8 @@ fn sample_request_candidate(
|
||||
None,
|
||||
None,
|
||||
100 + i64::from(candidate_index),
|
||||
started_at_unix_secs,
|
||||
started_at_unix_secs.map(|value| value + 1),
|
||||
started_at_unix_ms,
|
||||
started_at_unix_ms.map(|value| value + 1),
|
||||
)
|
||||
.expect("candidate should build")
|
||||
}
|
||||
|
||||
@@ -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
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.and_then(|context| context.get("candidate_index"))
|
||||
.and_then(serde_json::Value::as_u64)
|
||||
.is_some()
|
||||
}
|
||||
|
||||
fn is_retryable_local_upstream_status(status_code: u16) -> bool {
|
||||
status_code == 429 || status_code >= 500
|
||||
}
|
||||
|
||||
pub(crate) fn should_retry_next_local_candidate_sync(
|
||||
plan_kind: &str,
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
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>,
|
||||
result: &ExecutionResult,
|
||||
response_text: Option<&str>,
|
||||
) -> bool {
|
||||
is_local_candidate_attempt(report_context)
|
||||
&& plan_kind == "openai_chat_sync"
|
||||
&& is_retryable_local_upstream_status(result.status_code)
|
||||
matches!(
|
||||
resolve_local_failover_decision(
|
||||
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(
|
||||
@@ -102,14 +160,245 @@ pub(crate) fn resolve_core_sync_error_finalize_report_kind(
|
||||
Some(report_kind.to_string())
|
||||
}
|
||||
|
||||
pub(crate) fn should_retry_next_local_candidate_stream(
|
||||
plan_kind: &str,
|
||||
pub(crate) async fn should_retry_next_local_candidate_stream(
|
||||
state: &AppState,
|
||||
plan: &ExecutionPlan,
|
||||
_plan_kind: &str,
|
||||
report_context: Option<&serde_json::Value>,
|
||||
status_code: u16,
|
||||
response_text: Option<&str>,
|
||||
) -> bool {
|
||||
is_local_candidate_attempt(report_context)
|
||||
&& plan_kind == "openai_chat_stream"
|
||||
&& is_retryable_local_upstream_status(status_code)
|
||||
matches!(
|
||||
resolve_local_failover_decision(state, plan, report_context, status_code, response_text)
|
||||
.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(
|
||||
@@ -172,14 +461,121 @@ pub(crate) fn resolve_core_stream_direct_finalize_report_kind(plan_kind: &str) -
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use aether_contracts::ExecutionResult;
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
|
||||
use super::{
|
||||
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_retry_next_local_candidate_stream,
|
||||
should_retry_next_local_candidate_sync,
|
||||
resolve_core_sync_error_finalize_report_kind, resolve_local_failover_policy,
|
||||
should_fallback_to_control_stream, should_fallback_to_control_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]
|
||||
fn sync_failover_marks_chat_errors() {
|
||||
@@ -220,8 +616,8 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sync_retry_next_candidate_is_local_openai_chat_only() {
|
||||
#[tokio::test]
|
||||
async fn sync_retry_next_candidate_requires_local_candidate_context() {
|
||||
let result = ExecutionResult {
|
||||
request_id: "req-1".to_string(),
|
||||
candidate_id: None,
|
||||
@@ -235,26 +631,57 @@ mod tests {
|
||||
"candidate_index": 0,
|
||||
"retry_index": 0,
|
||||
});
|
||||
let state = build_state_with_provider_config(None);
|
||||
let plan = sample_plan();
|
||||
|
||||
assert!(should_retry_next_local_candidate_sync(
|
||||
"openai_chat_sync",
|
||||
Some(&local_report_context),
|
||||
&result,
|
||||
));
|
||||
assert!(!should_retry_next_local_candidate_sync(
|
||||
"openai_chat_sync",
|
||||
None,
|
||||
&result,
|
||||
));
|
||||
assert!(!should_retry_next_local_candidate_sync(
|
||||
"claude_chat_sync",
|
||||
None,
|
||||
&result,
|
||||
));
|
||||
assert!(
|
||||
should_retry_next_local_candidate_sync(
|
||||
&state,
|
||||
&plan,
|
||||
"openai_chat_sync",
|
||||
Some(&local_report_context),
|
||||
&result,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
);
|
||||
assert!(
|
||||
should_retry_next_local_candidate_sync(
|
||||
&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]
|
||||
fn sync_retry_next_candidate_treats_rate_limit_as_retryable() {
|
||||
#[tokio::test]
|
||||
async fn sync_retry_next_candidate_treats_rate_limit_as_retryable() {
|
||||
let result = ExecutionResult {
|
||||
request_id: "req-1".to_string(),
|
||||
candidate_id: None,
|
||||
@@ -268,49 +695,304 @@ mod tests {
|
||||
"candidate_index": 0,
|
||||
"retry_index": 0,
|
||||
});
|
||||
let state = build_state_with_provider_config(None);
|
||||
let plan = sample_plan();
|
||||
|
||||
assert!(should_retry_next_local_candidate_sync(
|
||||
"openai_chat_sync",
|
||||
Some(&local_report_context),
|
||||
&result,
|
||||
));
|
||||
assert!(
|
||||
should_retry_next_local_candidate_sync(
|
||||
&state,
|
||||
&plan,
|
||||
"openai_chat_sync",
|
||||
Some(&local_report_context),
|
||||
&result,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stream_retry_next_candidate_is_local_openai_chat_only() {
|
||||
#[tokio::test]
|
||||
async fn stream_retry_next_candidate_requires_local_candidate_context() {
|
||||
let local_report_context = serde_json::json!({
|
||||
"candidate_index": 0,
|
||||
"retry_index": 0,
|
||||
});
|
||||
let state = build_state_with_provider_config(None);
|
||||
let plan = sample_plan();
|
||||
|
||||
assert!(should_retry_next_local_candidate_stream(
|
||||
"openai_chat_stream",
|
||||
Some(&local_report_context),
|
||||
502,
|
||||
));
|
||||
assert!(!should_retry_next_local_candidate_stream(
|
||||
"openai_chat_stream",
|
||||
None,
|
||||
502,
|
||||
));
|
||||
assert!(!should_retry_next_local_candidate_stream(
|
||||
"claude_chat_stream",
|
||||
Some(&local_report_context),
|
||||
502,
|
||||
));
|
||||
assert!(
|
||||
should_retry_next_local_candidate_stream(
|
||||
&state,
|
||||
&plan,
|
||||
"openai_chat_stream",
|
||||
Some(&local_report_context),
|
||||
502,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
);
|
||||
assert!(
|
||||
should_retry_next_local_candidate_stream(
|
||||
&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]
|
||||
fn stream_retry_next_candidate_treats_rate_limit_as_retryable() {
|
||||
#[tokio::test]
|
||||
async fn stream_retry_next_candidate_treats_rate_limit_as_retryable() {
|
||||
let local_report_context = serde_json::json!({
|
||||
"candidate_index": 0,
|
||||
"retry_index": 0,
|
||||
});
|
||||
let state = build_state_with_provider_config(None);
|
||||
let plan = sample_plan();
|
||||
|
||||
assert!(should_retry_next_local_candidate_stream(
|
||||
"openai_chat_stream",
|
||||
Some(&local_report_context),
|
||||
429,
|
||||
));
|
||||
assert!(
|
||||
should_retry_next_local_candidate_stream(
|
||||
&state,
|
||||
&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
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -19,11 +19,12 @@ pub(crate) use self::constants::{
|
||||
MAX_ERROR_BODY_BYTES, MAX_STREAM_PREFETCH_BYTES, MAX_STREAM_PREFETCH_FRAMES,
|
||||
};
|
||||
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,
|
||||
should_fallback_to_control_stream, should_fallback_to_control_sync,
|
||||
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::{
|
||||
build_execution_runtime_router, build_execution_runtime_router_with_request_concurrency_limit,
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
use std::collections::VecDeque;
|
||||
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_scheduler_core::SchedulerRequestCandidateStatusUpdate;
|
||||
use async_stream::stream;
|
||||
@@ -32,7 +33,7 @@ use crate::ai_pipeline_api::{
|
||||
use crate::api::response::{
|
||||
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::control::GatewayControlDecision;
|
||||
use crate::execution_runtime::build_direct_execution_frame_stream;
|
||||
@@ -45,9 +46,9 @@ use crate::execution_runtime::transport::{
|
||||
DirectSyncExecutionRuntime, DirectUpstreamStreamExecution,
|
||||
};
|
||||
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,
|
||||
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::log_ids::short_request_id;
|
||||
@@ -165,7 +166,7 @@ pub(crate) async fn execute_execution_runtime_stream(
|
||||
};
|
||||
|
||||
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(
|
||||
state,
|
||||
&plan,
|
||||
@@ -179,8 +180,8 @@ pub(crate) async fn execute_execution_runtime_stream(
|
||||
response.status()
|
||||
)),
|
||||
latency_ms: None,
|
||||
started_at_unix_secs: Some(terminal_unix_secs),
|
||||
finished_at_unix_secs: Some(terminal_unix_secs),
|
||||
started_at_unix_ms: Some(terminal_unix_secs),
|
||||
finished_at_unix_ms: Some(terminal_unix_secs),
|
||||
},
|
||||
)
|
||||
.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(
|
||||
state: &AppState,
|
||||
plan: ExecutionPlan,
|
||||
@@ -237,69 +292,137 @@ async fn execute_stream_from_frame_stream(
|
||||
"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) {
|
||||
let terminal_unix_secs = current_request_candidate_unix_secs();
|
||||
record_local_request_candidate_status(
|
||||
if status_code == 200 {
|
||||
let success_probe_text =
|
||||
probe_local_stream_success_failover_text(&mut buffered_frames, &mut lines).await?;
|
||||
if should_retry_next_local_candidate_stream(
|
||||
state,
|
||||
&plan,
|
||||
plan_kind,
|
||||
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,
|
||||
"gateway local stream decision retrying next candidate after retryable execution runtime status"
|
||||
);
|
||||
return Ok(None);
|
||||
success_probe_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("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 =
|
||||
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 {
|
||||
let error_body = collect_error_body(&mut lines).await?;
|
||||
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
|
||||
.clone()
|
||||
.or_else(|| report_kind.clone())
|
||||
@@ -324,7 +447,7 @@ async fn execute_stream_from_frame_stream(
|
||||
&usage_payload,
|
||||
)
|
||||
.await;
|
||||
let terminal_unix_secs = current_request_candidate_unix_secs();
|
||||
let terminal_unix_secs = current_request_candidate_unix_ms();
|
||||
record_local_request_candidate_status(
|
||||
state,
|
||||
&plan,
|
||||
@@ -337,8 +460,8 @@ async fn execute_stream_from_frame_stream(
|
||||
"execution runtime stream returned error status {status_code}"
|
||||
)),
|
||||
latency_ms: None,
|
||||
started_at_unix_secs: Some(terminal_unix_secs),
|
||||
finished_at_unix_secs: Some(terminal_unix_secs),
|
||||
started_at_unix_ms: Some(terminal_unix_secs),
|
||||
finished_at_unix_ms: Some(terminal_unix_secs),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
@@ -399,7 +522,7 @@ async fn execute_stream_from_frame_stream(
|
||||
while prefetched_chunks.len() < MAX_STREAM_PREFETCH_FRAMES
|
||||
&& 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,
|
||||
Err(err) => {
|
||||
let failure = build_stream_failure_report(
|
||||
@@ -429,14 +552,14 @@ async fn execute_stream_from_frame_stream(
|
||||
};
|
||||
match frame.payload {
|
||||
StreamFramePayload::Data { chunk_b64, text } => {
|
||||
let chunk = if let Some(chunk_b64) = chunk_b64 {
|
||||
match base64::engine::general_purpose::STANDARD.decode(chunk_b64) {
|
||||
Ok(decoded) => decoded,
|
||||
let chunk =
|
||||
match decode_stream_data_chunk(chunk_b64.as_deref(), text.as_deref()) {
|
||||
Ok(chunk) => chunk,
|
||||
Err(err) => {
|
||||
let failure = build_stream_failure_report(
|
||||
"execution_runtime_stream_chunk_decode_error",
|
||||
format!(
|
||||
"failed to decode execution runtime stream chunk: {err}"
|
||||
"failed to decode execution runtime stream chunk: {err:?}"
|
||||
),
|
||||
502,
|
||||
);
|
||||
@@ -456,12 +579,7 @@ async fn execute_stream_from_frame_stream(
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
} else if let Some(text) = text {
|
||||
text.into_bytes()
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
};
|
||||
|
||||
if chunk.is_empty() {
|
||||
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
|
||||
.usage_runtime
|
||||
.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
|
||||
.as_ref()
|
||||
.and_then(|telemetry| telemetry.elapsed_ms),
|
||||
started_at_unix_secs: Some(candidate_started_unix_secs),
|
||||
finished_at_unix_secs: None,
|
||||
started_at_unix_ms: Some(candidate_started_unix_secs),
|
||||
finished_at_unix_ms: None,
|
||||
},
|
||||
)
|
||||
.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_log = short_request_id(request_id);
|
||||
let candidate_id_for_report = candidate_id.map(ToOwned::to_owned);
|
||||
let mut buffered_frames = buffered_frames;
|
||||
tokio::spawn(async move {
|
||||
let mut provider_buffered_body = provider_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 {
|
||||
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,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
@@ -707,9 +826,9 @@ async fn execute_stream_from_frame_stream(
|
||||
};
|
||||
match frame.payload {
|
||||
StreamFramePayload::Data { chunk_b64, text } => {
|
||||
let chunk = if let Some(chunk_b64) = chunk_b64 {
|
||||
match base64::engine::general_purpose::STANDARD.decode(chunk_b64) {
|
||||
Ok(decoded) => decoded,
|
||||
let chunk =
|
||||
match decode_stream_data_chunk(chunk_b64.as_deref(), text.as_deref()) {
|
||||
Ok(chunk) => chunk,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
event_name = "stream_execution_chunk_decode_failed",
|
||||
@@ -717,22 +836,19 @@ async fn execute_stream_from_frame_stream(
|
||||
trace_id = %trace_id_owned,
|
||||
request_id = %request_id_for_report_log,
|
||||
candidate_id = ?candidate_id_for_report.as_deref(),
|
||||
error = %err,
|
||||
error = ?err,
|
||||
"gateway failed to decode execution runtime chunk"
|
||||
);
|
||||
terminal_failure = Some(build_stream_failure_report(
|
||||
"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,
|
||||
));
|
||||
break;
|
||||
}
|
||||
}
|
||||
} else if let Some(text) = text {
|
||||
text.into_bytes()
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
};
|
||||
|
||||
if chunk.is_empty() {
|
||||
continue;
|
||||
@@ -998,8 +1114,8 @@ async fn execute_stream_from_frame_stream(
|
||||
error_type: Some("downstream_disconnect".to_string()),
|
||||
error_message: Some("client disconnected before stream completion".to_string()),
|
||||
latency_ms: telemetry.as_ref().and_then(|value| value.elapsed_ms),
|
||||
started_at_unix_secs: Some(candidate_started_unix_secs_for_report),
|
||||
finished_at_unix_secs: Some(current_request_candidate_unix_secs()),
|
||||
started_at_unix_ms: Some(candidate_started_unix_secs_for_report),
|
||||
finished_at_unix_ms: Some(current_request_candidate_unix_ms()),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
@@ -1055,8 +1171,8 @@ async fn execute_stream_from_frame_stream(
|
||||
error_type: None,
|
||||
error_message: None,
|
||||
latency_ms: telemetry.as_ref().and_then(|value| value.elapsed_ms),
|
||||
started_at_unix_secs: Some(candidate_started_unix_secs_for_report),
|
||||
finished_at_unix_secs: Some(current_request_candidate_unix_secs()),
|
||||
started_at_unix_ms: Some(candidate_started_unix_secs_for_report),
|
||||
finished_at_unix_ms: Some(current_request_candidate_unix_ms()),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -8,7 +8,7 @@ use serde_json::{Map, Value};
|
||||
use tracing::warn;
|
||||
|
||||
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::execution_runtime::submission::{
|
||||
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>,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
failure: &StreamFailureReport,
|
||||
started_at_unix_secs: Option<u64>,
|
||||
started_at_unix_ms: Option<u64>,
|
||||
) {
|
||||
state
|
||||
.usage_runtime
|
||||
.record_sync_terminal(state.data.as_ref(), plan, report_context, payload)
|
||||
.await;
|
||||
let terminal_unix_secs = current_request_candidate_unix_secs();
|
||||
let terminal_unix_secs = current_request_candidate_unix_ms();
|
||||
record_report_request_candidate_status(
|
||||
state,
|
||||
report_context,
|
||||
@@ -136,8 +136,8 @@ async fn record_stream_sync_failure(
|
||||
.telemetry
|
||||
.as_ref()
|
||||
.and_then(|telemetry| telemetry.elapsed_ms),
|
||||
started_at_unix_secs: started_at_unix_secs.or(Some(terminal_unix_secs)),
|
||||
finished_at_unix_secs: Some(terminal_unix_secs),
|
||||
started_at_unix_ms: started_at_unix_ms.or(Some(terminal_unix_secs)),
|
||||
finished_at_unix_ms: Some(terminal_unix_secs),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
@@ -195,7 +195,7 @@ pub(super) async fn submit_midstream_stream_failure(
|
||||
headers: &std::collections::BTreeMap<String, String>,
|
||||
telemetry: Option<ExecutionTelemetry>,
|
||||
buffered_body: &[u8],
|
||||
started_at_unix_secs: u64,
|
||||
started_at_unix_ms: u64,
|
||||
failure: StreamFailureReport,
|
||||
) {
|
||||
let Some(report_kind) =
|
||||
@@ -219,7 +219,7 @@ pub(super) async fn submit_midstream_stream_failure(
|
||||
report_context,
|
||||
&payload,
|
||||
&failure,
|
||||
Some(started_at_unix_secs),
|
||||
Some(started_at_unix_ms),
|
||||
)
|
||||
.await;
|
||||
if let Err(err) = submit_sync_report(state, trace_id, payload).await {
|
||||
|
||||
@@ -297,29 +297,39 @@ fn resolve_local_sync_error_status_code(status_code: u16, body_json: &serde_json
|
||||
return status_code;
|
||||
}
|
||||
|
||||
let Some(error_object) = body_json.get("error").and_then(|value| value.as_object()) else {
|
||||
return 400;
|
||||
};
|
||||
let body_object = body_json.as_object();
|
||||
let error_object = body_object
|
||||
.and_then(|object| object.get("error"))
|
||||
.and_then(|value| value.as_object());
|
||||
|
||||
for key in ["code", "status"] {
|
||||
let Some(value) = error_object.get(key) else {
|
||||
continue;
|
||||
};
|
||||
if let Some(number) = value.as_u64() {
|
||||
let raw_code = first_non_empty_error_text(error_object, body_object, &["code"]);
|
||||
let raw_status = first_non_empty_error_text(error_object, body_object, &["status"]);
|
||||
for numeric_hint in [raw_code.as_deref(), raw_status.as_deref()]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
{
|
||||
if let Ok(number) = numeric_hint.parse::<u16>() {
|
||||
if (400..600).contains(&number) {
|
||||
return number as u16;
|
||||
}
|
||||
}
|
||||
if let Some(text) = value.as_str() {
|
||||
if let Ok(number) = text.parse::<u16>() {
|
||||
if (400..600).contains(&number) {
|
||||
return 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(
|
||||
@@ -436,6 +446,20 @@ fn classify_local_sync_error_kind(
|
||||
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] {
|
||||
loop {
|
||||
while let Some(first) = body.first() {
|
||||
@@ -534,3 +558,128 @@ pub(crate) async fn submit_local_core_error_or_sync_finalize(
|
||||
|
||||
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"
|
||||
}
|
||||
})
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -15,7 +15,7 @@ use crate::ai_pipeline_api::{
|
||||
use crate::api::response::{
|
||||
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::control::GatewayControlDecision;
|
||||
#[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::transport::DirectSyncExecutionRuntime;
|
||||
use crate::execution_runtime::{
|
||||
resolve_core_sync_error_finalize_report_kind, should_fallback_to_control_sync,
|
||||
should_finalize_sync_response, should_retry_next_local_candidate_sync,
|
||||
local_failover_response_text, resolve_core_sync_error_finalize_report_kind,
|
||||
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::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_for_log = short_request_id(plan_request_id);
|
||||
let plan_candidate_id = plan.candidate_id.as_deref();
|
||||
let candidate_started_unix_secs = current_request_candidate_unix_ms();
|
||||
#[cfg(not(test))]
|
||||
let result = {
|
||||
match DirectSyncExecutionRuntime::new()
|
||||
@@ -140,6 +142,7 @@ pub(crate) async fn execute_execution_runtime_sync(
|
||||
plan_request_id,
|
||||
plan_candidate_id,
|
||||
report_context.as_ref(),
|
||||
candidate_started_unix_secs,
|
||||
)
|
||||
.await?;
|
||||
match remote_outcome {
|
||||
@@ -159,8 +162,34 @@ pub(crate) async fn execute_execution_runtime_sync(
|
||||
.telemetry
|
||||
.as_ref()
|
||||
.and_then(|telemetry| telemetry.elapsed_ms);
|
||||
if should_retry_next_local_candidate_sync(plan_kind, report_context.as_ref(), &result) {
|
||||
let terminal_unix_secs = current_request_candidate_unix_secs();
|
||||
let mut headers = result.headers.clone();
|
||||
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(
|
||||
state,
|
||||
&plan,
|
||||
@@ -171,8 +200,8 @@ pub(crate) async fn execute_execution_runtime_sync(
|
||||
error_type: result_error_type.clone(),
|
||||
error_message: result_error_message.clone(),
|
||||
latency_ms: result_latency_ms,
|
||||
started_at_unix_secs: Some(terminal_unix_secs),
|
||||
finished_at_unix_secs: Some(terminal_unix_secs),
|
||||
started_at_unix_ms: Some(candidate_started_unix_secs),
|
||||
finished_at_unix_ms: Some(terminal_unix_secs),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
@@ -191,8 +220,6 @@ pub(crate) async fn execute_execution_runtime_sync(
|
||||
.or(Some(plan_request_id));
|
||||
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 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 explicit_finalize = should_finalize_sync_response(report_kind.as_deref());
|
||||
let mapped_error_finalize_kind =
|
||||
@@ -220,15 +247,17 @@ pub(crate) async fn execute_execution_runtime_sync(
|
||||
mapped_error_finalize_kind.clone()
|
||||
};
|
||||
|
||||
if should_fallback_to_control_sync(
|
||||
plan_kind,
|
||||
&result,
|
||||
body_json.as_ref(),
|
||||
has_body_bytes,
|
||||
explicit_finalize || implicit_finalize.is_some(),
|
||||
mapped_error_finalize_kind.is_some(),
|
||||
) {
|
||||
let terminal_unix_secs = current_request_candidate_unix_secs();
|
||||
if !stop_local_failover
|
||||
&& should_fallback_to_control_sync(
|
||||
plan_kind,
|
||||
&result,
|
||||
body_json.as_ref(),
|
||||
has_body_bytes,
|
||||
explicit_finalize || implicit_finalize.is_some(),
|
||||
mapped_error_finalize_kind.is_some(),
|
||||
)
|
||||
{
|
||||
let terminal_unix_secs = current_request_candidate_unix_ms();
|
||||
record_local_request_candidate_status(
|
||||
state,
|
||||
&plan,
|
||||
@@ -239,8 +268,8 @@ pub(crate) async fn execute_execution_runtime_sync(
|
||||
error_type: result_error_type.clone(),
|
||||
error_message: result_error_message.clone(),
|
||||
latency_ms: result_latency_ms,
|
||||
started_at_unix_secs: Some(terminal_unix_secs),
|
||||
finished_at_unix_secs: Some(terminal_unix_secs),
|
||||
started_at_unix_ms: Some(candidate_started_unix_secs),
|
||||
finished_at_unix_ms: Some(terminal_unix_secs),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
@@ -251,7 +280,7 @@ pub(crate) async fn execute_execution_runtime_sync(
|
||||
.usage_runtime
|
||||
.record_pending(state.data.as_ref(), &plan, report_context.as_ref())
|
||||
.await;
|
||||
let terminal_unix_secs = current_request_candidate_unix_secs();
|
||||
let terminal_unix_secs = current_request_candidate_unix_ms();
|
||||
record_local_request_candidate_status(
|
||||
state,
|
||||
&plan,
|
||||
@@ -266,8 +295,8 @@ pub(crate) async fn execute_execution_runtime_sync(
|
||||
error_type: result_error_type.clone(),
|
||||
error_message: result_error_message.clone(),
|
||||
latency_ms: result_latency_ms,
|
||||
started_at_unix_secs: Some(terminal_unix_secs),
|
||||
finished_at_unix_secs: Some(terminal_unix_secs),
|
||||
started_at_unix_ms: Some(candidate_started_unix_secs),
|
||||
finished_at_unix_ms: Some(terminal_unix_secs),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
@@ -579,6 +608,7 @@ async fn execute_sync_via_remote_execution_runtime(
|
||||
plan_request_id: &str,
|
||||
plan_candidate_id: Option<&str>,
|
||||
report_context: Option<&serde_json::Value>,
|
||||
candidate_started_unix_secs: u64,
|
||||
) -> Result<RemoteSyncFallbackOutcome, GatewayError> {
|
||||
let response = match post_sync_plan_to_remote_execution_runtime(
|
||||
state,
|
||||
@@ -604,7 +634,7 @@ async fn execute_sync_via_remote_execution_runtime(
|
||||
};
|
||||
|
||||
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(
|
||||
state,
|
||||
plan,
|
||||
@@ -618,8 +648,8 @@ async fn execute_sync_via_remote_execution_runtime(
|
||||
response.status()
|
||||
)),
|
||||
latency_ms: None,
|
||||
started_at_unix_secs: Some(terminal_unix_secs),
|
||||
finished_at_unix_secs: Some(terminal_unix_secs),
|
||||
started_at_unix_ms: Some(candidate_started_unix_secs),
|
||||
finished_at_unix_ms: Some(terminal_unix_secs),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -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]
|
||||
fn resolve_local_core_error_response_body_json_parses_body_base64_json_for_cross_format_cli_error()
|
||||
{
|
||||
|
||||
@@ -125,8 +125,8 @@ where
|
||||
error_type: None,
|
||||
error_message: None,
|
||||
latency_ms: None,
|
||||
started_at_unix_secs: None,
|
||||
finished_at_unix_secs: None,
|
||||
started_at_unix_ms: None,
|
||||
finished_at_unix_ms: None,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
@@ -153,8 +153,8 @@ pub(crate) async fn mark_unused_local_candidate_items<T, FPlan, FContext>(
|
||||
error_type: None,
|
||||
error_message: None,
|
||||
latency_ms: None,
|
||||
started_at_unix_secs: None,
|
||||
finished_at_unix_secs: None,
|
||||
started_at_unix_ms: None,
|
||||
finished_at_unix_ms: None,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -53,7 +53,7 @@ pub(super) fn build_admin_billing_collector_payload_from_record(
|
||||
"default_value": record.default_value,
|
||||
"priority": record.priority,
|
||||
"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),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -213,7 +213,7 @@ pub(super) fn build_admin_payment_order_payload(
|
||||
"gateway_order_id": record.gateway_order_id,
|
||||
"gateway_response": record.gateway_response,
|
||||
"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),
|
||||
"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),
|
||||
@@ -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()))?,
|
||||
"error_message": row.try_get::<Option<String>, _>("error_message").map_err(|err| GatewayError::Internal(err.to_string()))?,
|
||||
"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()))?
|
||||
.and_then(|value| u64::try_from(value).ok())
|
||||
.and_then(unix_secs_to_rfc3339),
|
||||
@@ -263,7 +263,7 @@ pub(super) fn build_admin_payment_callback_payload_from_record(
|
||||
"status": record.status,
|
||||
"payload": record.payload,
|
||||
"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),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -53,7 +53,7 @@ fn build_admin_billing_rule_payload_from_record(
|
||||
"variables": record.variables,
|
||||
"dimension_mappings": record.dimension_mappings,
|
||||
"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),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -97,7 +97,7 @@ pub(in super::super) async fn build_admin_wallet_adjust_response(
|
||||
transaction.link_id.as_deref(),
|
||||
transaction.operator_id.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!({
|
||||
"wallet": wallet_payload,
|
||||
|
||||
@@ -84,7 +84,7 @@ pub(in super::super) async fn build_admin_wallet_fail_refund_response(
|
||||
transaction.link_id.as_deref(),
|
||||
transaction.operator_id.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),
|
||||
|
||||
@@ -69,7 +69,7 @@ pub(in super::super) async fn build_admin_wallet_process_refund_response(
|
||||
transaction.link_id.as_deref(),
|
||||
transaction.operator_id.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();
|
||||
|
||||
@@ -87,7 +87,7 @@ pub(in super::super) async fn build_admin_wallet_recharge_response(
|
||||
payment_order.amount_usd,
|
||||
payment_order.payment_method,
|
||||
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
|
||||
.credited_at_unix_secs
|
||||
.and_then(unix_secs_to_rfc3339),
|
||||
|
||||
@@ -69,7 +69,7 @@ pub(in super::super) async fn build_admin_wallet_ledger_response(
|
||||
"operator_name": entry.operator_name,
|
||||
"operator_email": entry.operator_email,
|
||||
"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<_>>();
|
||||
|
||||
@@ -59,7 +59,7 @@ pub(in super::super) async fn build_admin_wallet_list_response(
|
||||
"total_consumed": wallet.total_consumed,
|
||||
"total_refunded": wallet.total_refunded,
|
||||
"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),
|
||||
})
|
||||
})
|
||||
|
||||
@@ -79,7 +79,7 @@ pub(in super::super) async fn build_admin_wallet_refund_requests_response(
|
||||
"requested_by": refund.requested_by,
|
||||
"approved_by": refund.approved_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),
|
||||
"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),
|
||||
|
||||
@@ -80,7 +80,7 @@ pub(in super::super) async fn build_admin_wallet_transactions_response(
|
||||
"operator_name": operator_name,
|
||||
"operator_email": operator_email,
|
||||
"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),
|
||||
}));
|
||||
}
|
||||
|
||||
|
||||
@@ -188,7 +188,7 @@ pub(in super::super) fn build_admin_wallet_refund_payload(
|
||||
"requested_by": refund.requested_by.clone(),
|
||||
"approved_by": refund.approved_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),
|
||||
"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),
|
||||
|
||||
@@ -3,6 +3,7 @@ use crate::handlers::admin::shared::unix_secs_to_rfc3339;
|
||||
use crate::handlers::public::{
|
||||
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_scheduler_core::{is_provider_key_circuit_open, provider_key_health_score};
|
||||
use serde_json::json;
|
||||
@@ -124,28 +125,24 @@ pub(crate) async fn build_admin_endpoint_health_status_payload(
|
||||
total_count: 0,
|
||||
success_count: 0,
|
||||
failed_count: 0,
|
||||
min_created_at_unix_secs: None,
|
||||
max_created_at_unix_secs: None,
|
||||
min_created_at_unix_ms: None,
|
||||
max_created_at_unix_ms: None,
|
||||
});
|
||||
bucket.total_count += row.total_count;
|
||||
bucket.success_count += row.success_count;
|
||||
bucket.failed_count += row.failed_count;
|
||||
bucket.min_created_at_unix_secs = match (
|
||||
bucket.min_created_at_unix_secs,
|
||||
row.min_created_at_unix_secs,
|
||||
) {
|
||||
(Some(left), Some(right)) => Some(left.min(right)),
|
||||
(None, Some(right)) => Some(right),
|
||||
(left, None) => left,
|
||||
};
|
||||
bucket.max_created_at_unix_secs = match (
|
||||
bucket.max_created_at_unix_secs,
|
||||
row.max_created_at_unix_secs,
|
||||
) {
|
||||
(Some(left), Some(right)) => Some(left.max(right)),
|
||||
(None, Some(right)) => Some(right),
|
||||
(left, None) => left,
|
||||
};
|
||||
bucket.min_created_at_unix_ms =
|
||||
match (bucket.min_created_at_unix_ms, row.min_created_at_unix_ms) {
|
||||
(Some(left), Some(right)) => Some(left.min(right)),
|
||||
(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) {
|
||||
(Some(left), Some(right)) => Some(left.max(right)),
|
||||
(None, Some(right)) => Some(right),
|
||||
(left, None) => left,
|
||||
};
|
||||
}
|
||||
|
||||
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),
|
||||
"health_score": health_score,
|
||||
"timeline": timeline,
|
||||
"time_range_start": time_range_start.and_then(unix_secs_to_rfc3339),
|
||||
"time_range_end": time_range_end.or(Some(now_unix_secs)).and_then(unix_secs_to_rfc3339),
|
||||
"time_range_start": time_range_start.and_then(unix_ms_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_keys": total_keys,
|
||||
"active_keys": active_keys_by_format.get(&api_format).map(BTreeSet::len).unwrap_or(0),
|
||||
|
||||
@@ -140,7 +140,7 @@ fn build_admin_gemini_file_mapping_payload(
|
||||
"username": username,
|
||||
"display_name": mapping.display_name,
|
||||
"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),
|
||||
"is_expired": mapping.expires_at_unix_secs <= now_unix_secs,
|
||||
})
|
||||
|
||||
@@ -94,7 +94,7 @@ pub(super) fn build_admin_video_task_list_item(
|
||||
"error_message": task.error_message,
|
||||
"poll_count": task.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),
|
||||
"submitted_at": admin_video_task_timestamp(task.submitted_at_unix_secs),
|
||||
})
|
||||
|
||||
@@ -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(
|
||||
"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(
|
||||
"updated_at".to_string(),
|
||||
|
||||
@@ -49,6 +49,7 @@ pub(super) async fn build_admin_monitoring_cache_stats_response(
|
||||
"affinity_stats": {
|
||||
"storage_type": snapshot.storage_type,
|
||||
"total_affinities": snapshot.total_affinities,
|
||||
"active_affinities": snapshot.total_affinities,
|
||||
"cache_hits": snapshot.cache_hits,
|
||||
"cache_misses": snapshot.cache_misses,
|
||||
"cache_hit_rate": snapshot.cache_hit_rate,
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
use super::cache_types::AdminMonitoringCacheAffinityRecord;
|
||||
use crate::cache::SchedulerAffinityTarget;
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::GatewayError;
|
||||
use std::time::Duration;
|
||||
|
||||
fn parse_admin_monitoring_cache_affinity_key(raw_key: &str) -> Option<(String, String, String)> {
|
||||
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))
|
||||
}
|
||||
|
||||
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(
|
||||
record: &AdminMonitoringCacheAffinityRecord,
|
||||
) -> 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(
|
||||
state: &AdminAppState<'_>,
|
||||
records: &[AdminMonitoringCacheAffinityRecord],
|
||||
|
||||
@@ -8,6 +8,7 @@ use super::super::cache_route_helpers::{
|
||||
admin_monitoring_cache_affinity_unavailable_response,
|
||||
};
|
||||
use super::super::cache_store::{
|
||||
admin_monitoring_has_runtime_scheduler_affinity_entries,
|
||||
list_admin_monitoring_cache_affinity_records_by_affinity_keys,
|
||||
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()
|
||||
&& 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());
|
||||
}
|
||||
|
||||
@@ -11,6 +11,7 @@ use super::super::cache_route_helpers::{
|
||||
admin_monitoring_cache_users_user_identifier_from_path,
|
||||
};
|
||||
use super::super::cache_store::{
|
||||
admin_monitoring_has_runtime_scheduler_affinity_entries,
|
||||
list_admin_monitoring_cache_affinity_records_by_affinity_keys,
|
||||
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()
|
||||
&& 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());
|
||||
}
|
||||
|
||||
@@ -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 crate::handlers::admin::observability::stats::round_to;
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
|
||||
use crate::GatewayError;
|
||||
use aether_data_contracts::repository::usage::UsageAuditListQuery;
|
||||
|
||||
async fn count_admin_monitoring_cache_affinity_entries(state: &AdminAppState<'_>) -> usize {
|
||||
let Some(runner) = state.redis_kv_runner() else {
|
||||
return 0;
|
||||
};
|
||||
let mut connection = match runner.client().get_multiplexed_async_connection().await {
|
||||
Ok(value) => value,
|
||||
Err(_) => return 0,
|
||||
};
|
||||
let pattern = runner.keyspace().key("cache_affinity:*");
|
||||
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
|
||||
list_admin_monitoring_cache_affinity_records(state)
|
||||
.await
|
||||
.map(|items| items.len())
|
||||
.unwrap_or_else(|_| {
|
||||
state
|
||||
.as_ref()
|
||||
.list_scheduler_affinity_entries(SCHEDULER_AFFINITY_TTL)
|
||||
.len()
|
||||
})
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
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(
|
||||
state: &AdminAppState<'_>,
|
||||
affinity_keys: Option<&std::collections::BTreeSet<String>>,
|
||||
) -> Result<Vec<AdminMonitoringCacheAffinityRecord>, GatewayError> {
|
||||
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() {
|
||||
let mut connection = runner
|
||||
@@ -209,14 +208,24 @@ async fn list_admin_monitoring_cache_affinity_records_matching(
|
||||
let patterns = affinity_keys
|
||||
.map(|keys| {
|
||||
keys.iter()
|
||||
.map(|affinity_key| {
|
||||
runner
|
||||
.keyspace()
|
||||
.key(&format!("cache_affinity:{affinity_key}:*"))
|
||||
.flat_map(|affinity_key| {
|
||||
[
|
||||
runner
|
||||
.keyspace()
|
||||
.key(&format!("cache_affinity:{affinity_key}:*")),
|
||||
runner
|
||||
.keyspace()
|
||||
.key(&format!("scheduler_affinity:{affinity_key}:*")),
|
||||
]
|
||||
})
|
||||
.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 {
|
||||
let mut cursor = 0u64;
|
||||
@@ -246,16 +255,18 @@ async fn list_admin_monitoring_cache_affinity_records_matching(
|
||||
let Some(raw_value) = raw_value else {
|
||||
continue;
|
||||
};
|
||||
let Some(record) = admin_monitoring_cache_affinity_record(&key, &raw_value)
|
||||
else {
|
||||
let record = if key.contains("scheduler_affinity:") {
|
||||
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;
|
||||
};
|
||||
if affinity_keys.is_some_and(|keys| !keys.contains(&record.affinity_key)) {
|
||||
continue;
|
||||
}
|
||||
if seen_raw_keys.insert(record.raw_key.clone()) {
|
||||
records.push(record);
|
||||
}
|
||||
push_record(record);
|
||||
}
|
||||
}
|
||||
if next_cursor == 0 {
|
||||
@@ -264,7 +275,6 @@ async fn list_admin_monitoring_cache_affinity_records_matching(
|
||||
cursor = next_cursor;
|
||||
}
|
||||
}
|
||||
return Ok(records);
|
||||
}
|
||||
|
||||
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)) {
|
||||
continue;
|
||||
}
|
||||
if seen_raw_keys.insert(record.raw_key.clone()) {
|
||||
records.push(record);
|
||||
push_record(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)
|
||||
|
||||
@@ -161,7 +161,7 @@ pub(super) async fn build_admin_monitoring_resilience_snapshot(
|
||||
.filter(admin_monitoring_usage_is_error)
|
||||
.collect::<Vec<_>>();
|
||||
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 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_type": error_type,
|
||||
"operation": operation,
|
||||
"timestamp": unix_secs_to_rfc3339(item.created_at_unix_secs),
|
||||
"timestamp": unix_secs_to_rfc3339(item.created_at_unix_ms),
|
||||
"context": {
|
||||
"request_id": item.request_id,
|
||||
"provider_id": item.provider_id,
|
||||
|
||||
@@ -31,7 +31,7 @@ pub(super) fn sample_usage(
|
||||
total_cost_usd: f64,
|
||||
status: &str,
|
||||
status_code: Option<i32>,
|
||||
created_at_unix_secs: i64,
|
||||
created_at_unix_ms: i64,
|
||||
) -> StoredRequestUsageAudit {
|
||||
let is_error = status_code.is_some_and(|value| value >= 400)
|
||||
|| status.trim().eq_ignore_ascii_case("failed")
|
||||
@@ -70,9 +70,9 @@ pub(super) fn sample_usage(
|
||||
Some(30),
|
||||
status.to_string(),
|
||||
"billed".to_string(),
|
||||
created_at_unix_secs,
|
||||
created_at_unix_secs,
|
||||
Some(created_at_unix_secs),
|
||||
created_at_unix_ms,
|
||||
created_at_unix_ms,
|
||||
Some(created_at_unix_ms),
|
||||
)
|
||||
.expect("usage should build")
|
||||
}
|
||||
@@ -82,7 +82,7 @@ pub(super) fn sample_candidate(
|
||||
request_id: &str,
|
||||
candidate_index: i32,
|
||||
status: RequestCandidateStatus,
|
||||
started_at_unix_secs: Option<i64>,
|
||||
started_at_unix_ms: Option<i64>,
|
||||
latency_ms: Option<i32>,
|
||||
status_code: Option<i32>,
|
||||
) -> StoredRequestCandidate {
|
||||
@@ -108,9 +108,9 @@ pub(super) fn sample_candidate(
|
||||
Some(1),
|
||||
None,
|
||||
Some(json!({"cache_1h": true})),
|
||||
100 + i64::from(candidate_index),
|
||||
started_at_unix_secs,
|
||||
started_at_unix_secs.map(|value| value + 1),
|
||||
(100 + i64::from(candidate_index)) * 1_000,
|
||||
started_at_unix_ms.map(|v| v * 1_000),
|
||||
started_at_unix_ms.map(|value| (value + 1) * 1_000),
|
||||
)
|
||||
.expect("candidate should build")
|
||||
}
|
||||
|
||||
@@ -258,6 +258,56 @@ async fn admin_monitoring_resilience_status_returns_local_payload() {
|
||||
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]
|
||||
async fn admin_monitoring_cache_stats_returns_local_payload() {
|
||||
let now = chrono::Utc::now().timestamp();
|
||||
|
||||
@@ -209,6 +209,136 @@ async fn admin_monitoring_cache_affinities_and_affinity_return_local_payload_fro
|
||||
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]
|
||||
async fn admin_monitoring_cache_users_delete_returns_local_payload_from_test_store() {
|
||||
let user_repository = Arc::new(
|
||||
|
||||
@@ -13,7 +13,7 @@ use axum::{
|
||||
body::Body,
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use tracing::warn;
|
||||
use tracing::debug;
|
||||
|
||||
pub(super) async fn build_admin_monitoring_trace_request_response(
|
||||
state: &AdminAppState<'_>,
|
||||
@@ -38,7 +38,7 @@ pub(super) async fn build_admin_monitoring_trace_request_response(
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?
|
||||
else {
|
||||
warn!(
|
||||
debug!(
|
||||
event_name = "admin_monitoring_request_trace_not_found",
|
||||
log_type = "admin_monitoring",
|
||||
request_id = %short_request_id(request_id.as_str()),
|
||||
|
||||
@@ -20,8 +20,8 @@ pub(in super::super) async fn list_recent_completed_usage_for_cache_affinity(
|
||||
.await?;
|
||||
items.retain(|item| item.status == "completed");
|
||||
items.sort_by(|left, right| {
|
||||
left.created_at_unix_secs
|
||||
.cmp(&right.created_at_unix_secs)
|
||||
left.created_at_unix_ms
|
||||
.cmp(&right.created_at_unix_ms)
|
||||
.then_with(|| left.id.cmp(&right.id))
|
||||
});
|
||||
Ok(items)
|
||||
|
||||
@@ -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 {
|
||||
if let Some(previous) = previous_created_at_unix_secs {
|
||||
if let Some(previous) = previous_created_at_unix_ms {
|
||||
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 {
|
||||
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,
|
||||
});
|
||||
if !item.model.trim().is_empty() {
|
||||
@@ -78,7 +78,7 @@ pub(super) async fn build_admin_usage_cache_affinity_interval_timeline_response(
|
||||
.push(point);
|
||||
}
|
||||
}
|
||||
previous_created_at_unix_secs = Some(item.created_at_unix_secs);
|
||||
previous_created_at_unix_ms = Some(item.created_at_unix_ms);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -78,8 +78,8 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
|
||||
.collect();
|
||||
items.sort_by(|left, right| {
|
||||
right
|
||||
.created_at_unix_secs
|
||||
.cmp(&left.created_at_unix_secs)
|
||||
.created_at_unix_ms
|
||||
.cmp(&left.created_at_unix_ms)
|
||||
.then_with(|| left.id.cmp(&right.id))
|
||||
});
|
||||
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| {
|
||||
right
|
||||
.created_at_unix_secs
|
||||
.cmp(&left.created_at_unix_secs)
|
||||
.created_at_unix_ms
|
||||
.cmp(&left.created_at_unix_ms)
|
||||
.then_with(|| left.id.cmp(&right.id))
|
||||
});
|
||||
let total = usage.len();
|
||||
|
||||
@@ -28,9 +28,9 @@ pub(crate) async fn build_admin_provider_endpoints_payload(
|
||||
.unwrap_or_default();
|
||||
endpoints.sort_by(|left, right| {
|
||||
right
|
||||
.created_at_unix_secs
|
||||
.created_at_unix_ms
|
||||
.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))
|
||||
});
|
||||
let keys = state
|
||||
|
||||
@@ -165,7 +165,7 @@ pub(super) async fn handle_admin_provider_oauth_device_authorize(
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
created_at_unix_secs: now_unix_secs,
|
||||
created_at_unix_ms: now_unix_secs,
|
||||
key_id: None,
|
||||
email: None,
|
||||
replaced: false,
|
||||
|
||||
@@ -128,7 +128,7 @@ pub(crate) async fn create_provider_oauth_catalog_key(
|
||||
record.total_response_time_ms = Some(0);
|
||||
record.health_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);
|
||||
state.create_provider_catalog_key(&record).await
|
||||
}
|
||||
|
||||
@@ -803,7 +803,7 @@ pub(super) fn build_admin_pool_key_payload(
|
||||
);
|
||||
payload.insert(
|
||||
"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(
|
||||
"last_used_at".to_string(),
|
||||
|
||||
@@ -177,7 +177,7 @@ pub(crate) async fn build_admin_providers_summary_payload(
|
||||
.is_active
|
||||
.cmp(&left.is_active)
|
||||
.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();
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
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::{
|
||||
RequestCandidateStatus, StoredRequestCandidate,
|
||||
};
|
||||
@@ -82,14 +83,14 @@ pub(crate) async fn build_admin_provider_health_monitor_payload(
|
||||
for candidates in attempts_by_endpoint.values_mut() {
|
||||
candidates.sort_by(|left, right| {
|
||||
right
|
||||
.created_at_unix_secs
|
||||
.cmp(&left.created_at_unix_secs)
|
||||
.created_at_unix_ms
|
||||
.cmp(&left.created_at_unix_ms)
|
||||
.then_with(|| right.id.cmp(&left.id))
|
||||
});
|
||||
candidates.truncate(per_endpoint_limit);
|
||||
candidates.sort_by(|left, right| {
|
||||
request_candidate_event_unix_secs(left)
|
||||
.cmp(&request_candidate_event_unix_secs(right))
|
||||
request_candidate_event_unix_ms(left)
|
||||
.cmp(&request_candidate_event_unix_ms(right))
|
||||
.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
|
||||
.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
|
||||
.into_iter()
|
||||
.filter_map(|candidate| {
|
||||
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_code": candidate.status_code,
|
||||
"latency_ms": candidate.latency_ms,
|
||||
|
||||
@@ -88,7 +88,7 @@ pub(crate) async fn build_admin_providers_payload(
|
||||
"api_key": has_any_key_by_provider.contains(&provider_id).then_some("***"),
|
||||
"priority": provider.provider_priority,
|
||||
"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),
|
||||
})
|
||||
})
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
use crate::handlers::admin::shared::unix_secs_to_rfc3339;
|
||||
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::{
|
||||
RequestCandidateStatus, StoredRequestCandidate,
|
||||
@@ -172,7 +172,7 @@ pub(crate) fn build_admin_provider_summary_value(
|
||||
"endpoint_health_details": endpoint_health_details,
|
||||
"ops_configured": ops_configured,
|
||||
"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),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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));
|
||||
key.health_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);
|
||||
Ok(key)
|
||||
}
|
||||
|
||||
@@ -24,9 +24,9 @@ pub(crate) async fn build_admin_provider_keys_payload(
|
||||
left.internal_priority
|
||||
.cmp(&right.internal_priority)
|
||||
.then_with(|| {
|
||||
left.created_at_unix_secs
|
||||
left.created_at_unix_ms
|
||||
.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))
|
||||
});
|
||||
|
||||
@@ -82,9 +82,9 @@ pub(crate) async fn build_admin_system_export_providers_payload(
|
||||
left.internal_priority
|
||||
.cmp(&right.internal_priority)
|
||||
.then(
|
||||
left.created_at_unix_secs
|
||||
left.created_at_unix_ms
|
||||
.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))
|
||||
});
|
||||
|
||||
@@ -463,7 +463,7 @@ pub(crate) fn build_management_token_payload(
|
||||
"last_used_ip": token.last_used_ip,
|
||||
"usage_count": token.usage_count,
|
||||
"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),
|
||||
});
|
||||
if let Some(user) = user {
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
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 aether_data_contracts::repository::candidates::{
|
||||
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
|
||||
.finished_at_unix_secs
|
||||
.or(candidate.started_at_unix_secs)
|
||||
.unwrap_or(candidate.created_at_unix_secs)
|
||||
.finished_at_unix_ms
|
||||
.or(candidate.started_at_unix_ms)
|
||||
.unwrap_or(candidate.created_at_unix_ms)
|
||||
}
|
||||
|
||||
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,
|
||||
success_count: 0,
|
||||
failed_count: 0,
|
||||
min_created_at_unix_secs: None,
|
||||
max_created_at_unix_secs: None,
|
||||
min_created_at_unix_ms: None,
|
||||
max_created_at_unix_ms: None,
|
||||
});
|
||||
bucket.total_count += row.total_count;
|
||||
bucket.success_count += row.success_count;
|
||||
bucket.failed_count += row.failed_count;
|
||||
bucket.min_created_at_unix_secs = match (
|
||||
bucket.min_created_at_unix_secs,
|
||||
row.min_created_at_unix_secs,
|
||||
) {
|
||||
(Some(left), Some(right)) => Some(left.min(right)),
|
||||
(None, Some(right)) => Some(right),
|
||||
(left, None) => left,
|
||||
};
|
||||
bucket.max_created_at_unix_secs = match (
|
||||
bucket.max_created_at_unix_secs,
|
||||
row.max_created_at_unix_secs,
|
||||
) {
|
||||
(Some(left), Some(right)) => Some(left.max(right)),
|
||||
(None, Some(right)) => Some(right),
|
||||
(left, None) => left,
|
||||
};
|
||||
bucket.min_created_at_unix_ms =
|
||||
match (bucket.min_created_at_unix_ms, row.min_created_at_unix_ms) {
|
||||
(Some(left), Some(right)) => Some(left.min(right)),
|
||||
(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) {
|
||||
(Some(left), Some(right)) => Some(left.max(right)),
|
||||
(None, Some(right)) => Some(right),
|
||||
(left, None) => left,
|
||||
};
|
||||
}
|
||||
|
||||
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| {
|
||||
candidate
|
||||
.finished_at_unix_secs
|
||||
.or(candidate.started_at_unix_secs)
|
||||
.or(Some(candidate.created_at_unix_secs))
|
||||
.finished_at_unix_ms
|
||||
.or(candidate.started_at_unix_ms)
|
||||
.or(Some(candidate.created_at_unix_ms))
|
||||
});
|
||||
let events = attempts
|
||||
.into_iter()
|
||||
.filter_map(|candidate| {
|
||||
let timestamp = candidate
|
||||
.finished_at_unix_secs
|
||||
.or(candidate.started_at_unix_secs)
|
||||
.unwrap_or(candidate.created_at_unix_secs);
|
||||
.finished_at_unix_ms
|
||||
.or(candidate.started_at_unix_ms)
|
||||
.unwrap_or(candidate.created_at_unix_ms);
|
||||
Some(json!({
|
||||
"timestamp": unix_secs_to_rfc3339(timestamp)?,
|
||||
"timestamp": unix_ms_to_rfc3339(timestamp)?,
|
||||
"status": request_candidate_status_label(candidate.status),
|
||||
"status_code": candidate.status_code,
|
||||
"latency_ms": candidate.latency_ms,
|
||||
@@ -490,11 +488,11 @@ pub(crate) async fn build_api_format_health_monitor_payload(
|
||||
"failed_count": failed_count,
|
||||
"skipped_count": skipped_count,
|
||||
"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,
|
||||
"timeline": timeline,
|
||||
"time_range_start": time_range_start.and_then(unix_secs_to_rfc3339),
|
||||
"time_range_end": time_range_end.or(Some(now_unix_secs)).and_then(unix_secs_to_rfc3339),
|
||||
"time_range_start": time_range_start.and_then(unix_ms_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 {
|
||||
format_payload["api_path"] = json!(public_api_format_local_path(&api_format));
|
||||
@@ -536,12 +534,12 @@ pub(crate) fn build_public_health_timeline(
|
||||
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)),
|
||||
(None, Some(right)) => Some(right),
|
||||
(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)),
|
||||
(None, Some(right)) => Some(right),
|
||||
(left, None) => left,
|
||||
@@ -592,3 +590,50 @@ pub(crate) fn api_format_display_name(api_format: &str) -> String {
|
||||
};
|
||||
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")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,7 +10,7 @@ pub(crate) use self::catalog_helpers::{
|
||||
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_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,
|
||||
};
|
||||
pub(crate) use self::system_modules_helpers::{
|
||||
|
||||
@@ -84,7 +84,7 @@ pub(super) fn build_public_announcement_payload(
|
||||
},
|
||||
"start_time": format_optional_unix_datetime(announcement.start_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),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -361,7 +361,7 @@ fn dashboard_usage_local_date(
|
||||
item: &StoredRequestUsageAudit,
|
||||
tz_offset_minutes: i32,
|
||||
) -> 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)?;
|
||||
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,
|
||||
"model": dashboard_non_empty_value(&item.model, "N/A"),
|
||||
"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,
|
||||
})
|
||||
})
|
||||
|
||||
@@ -95,7 +95,7 @@ fn users_me_usage_total_input_context(item: &StoredRequestUsageAudit) -> u64 {
|
||||
|
||||
fn users_me_usage_effective_unix_secs(item: &StoredRequestUsageAudit) -> u64 {
|
||||
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 {
|
||||
@@ -159,7 +159,7 @@ fn build_users_me_usage_record_payload(
|
||||
"first_byte_time_ms": item.first_byte_time_ms,
|
||||
"is_stream": item.is_stream,
|
||||
"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_read_input_tokens": item.cache_read_input_tokens,
|
||||
"status_code": item.status_code,
|
||||
@@ -649,8 +649,8 @@ pub(super) async fn handle_users_me_usage_get(
|
||||
.collect::<Vec<_>>();
|
||||
records.sort_by(|left, right| {
|
||||
right
|
||||
.created_at_unix_secs
|
||||
.cmp(&left.created_at_unix_secs)
|
||||
.created_at_unix_ms
|
||||
.cmp(&left.created_at_unix_ms)
|
||||
.then_with(|| left.id.cmp(&right.id))
|
||||
});
|
||||
let total_record_count = records.len();
|
||||
@@ -735,8 +735,8 @@ pub(super) async fn handle_users_me_usage_active_get(
|
||||
.collect::<Vec<_>>();
|
||||
items.sort_by(|left, right| {
|
||||
right
|
||||
.created_at_unix_secs
|
||||
.cmp(&left.created_at_unix_secs)
|
||||
.created_at_unix_ms
|
||||
.cmp(&left.created_at_unix_ms)
|
||||
.then_with(|| left.id.cmp(&right.id))
|
||||
});
|
||||
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.sort_by(|left, right| {
|
||||
left.created_at_unix_secs
|
||||
.cmp(&right.created_at_unix_secs)
|
||||
left.created_at_unix_ms
|
||||
.cmp(&right.created_at_unix_ms)
|
||||
.then_with(|| left.id.cmp(&right.id))
|
||||
});
|
||||
|
||||
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 {
|
||||
if let Some(previous) = previous_created_at_unix_secs {
|
||||
let interval_minutes =
|
||||
(item.created_at_unix_secs.saturating_sub(previous) as f64) / 60.0;
|
||||
if let Some(previous) = previous_created_at_unix_ms {
|
||||
let interval_minutes = (item.created_at_unix_ms.saturating_sub(previous) as f64) / 60.0;
|
||||
if interval_minutes <= 120.0 {
|
||||
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),
|
||||
"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!({
|
||||
|
||||
@@ -153,7 +153,7 @@ pub(super) fn wallet_transaction_payload_from_record(
|
||||
"link_id": record.link_id.clone(),
|
||||
"operator_id": record.operator_id.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),
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -207,7 +207,7 @@ pub(crate) fn wallet_payment_order_payload_from_row(
|
||||
row: &sqlx::postgres::PgRow,
|
||||
) -> Result<serde_json::Value, GatewayError> {
|
||||
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()))?
|
||||
.and_then(|value| u64::try_from(value).ok())
|
||||
.and_then(unix_secs_to_rfc3339);
|
||||
@@ -280,7 +280,7 @@ fn wallet_payment_order_payload_from_record(
|
||||
record.gateway_order_id.clone(),
|
||||
record.gateway_response.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.credited_at_unix_secs.and_then(unix_secs_to_rfc3339),
|
||||
record.expires_at_unix_secs.and_then(unix_secs_to_rfc3339),
|
||||
|
||||
@@ -98,7 +98,7 @@ fn wallet_refund_payload_from_row(
|
||||
row: &sqlx::postgres::PgRow,
|
||||
) -> Result<serde_json::Value, GatewayError> {
|
||||
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()))?
|
||||
.and_then(|value| u64::try_from(value).ok())
|
||||
.and_then(unix_secs_to_rfc3339);
|
||||
@@ -157,7 +157,7 @@ fn wallet_refund_payload_from_record(
|
||||
"payout_method": record.payout_method,
|
||||
"payout_reference": record.payout_reference,
|
||||
"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),
|
||||
"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),
|
||||
@@ -247,7 +247,7 @@ pub(super) async fn handle_wallet_refunds_list(
|
||||
"payout_method": record.payout_method,
|
||||
"payout_reference": record.payout_reference,
|
||||
"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),
|
||||
"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),
|
||||
|
||||
@@ -156,7 +156,7 @@ pub(crate) async fn build_admin_keys_grouped_by_format_payload(
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
.unwrap_or(false),
|
||||
"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)),
|
||||
}));
|
||||
}
|
||||
|
||||
@@ -598,7 +598,7 @@ pub(crate) fn build_admin_provider_key_response(
|
||||
payload.insert(
|
||||
"created_at".to_string(),
|
||||
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(
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user