i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix scheduler affinity candidate selection

Browse Source
This commit is contained in:
fawney19
2026-04-30 16:27:24 +08:00
parent 558abfcfa3
commit 33aa70c22b
24 changed files with 1715 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
apps/aether-gateway/src/ai_pipeline/conversion/mod.rs
View File

@@ -115,6 +115,15 @@ mod tests {
"gemini:generate_content",
]
);
assert_eq!(
request_candidate_api_formats("claude:messages", false),
vec![
"claude:messages",
"openai:chat",
"openai:responses",
"gemini:generate_content",
]
);
assert_eq!(
request_candidate_api_formats("openai:cli", false),
Vec::<&'static str>::new()

182
apps/aether-gateway/src/ai_pipeline/planner/candidate_ranking.rs
View File

@@ -609,8 +609,7 @@ mod tests {
}
#[tokio::test]
async fn fixed_order_local_execution_ranking_keeps_provider_priority_before_format_preference()
{
async fn fixed_order_local_execution_ranking_demotes_cross_format_before_provider_priority() {
let provider_catalog = InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider_with_options("provider-same", false, 10),
@@ -662,8 +661,8 @@ mod tests {
)
.await;
assert_eq!(ranked[0].endpoint_id, "endpoint-cross");
assert_eq!(ranked[1].endpoint_id, "endpoint-same");
assert_eq!(ranked[0].endpoint_id, "endpoint-same");
assert_eq!(ranked[1].endpoint_id, "endpoint-cross");
}
#[tokio::test]
@@ -1405,6 +1404,181 @@ mod tests {
);
}
#[tokio::test]
async fn first_request_same_key_exact_endpoint_beats_cross_format_without_affinity() {
let mut openai_endpoint =
sample_endpoint_for_provider("provider-shared", "endpoint-openai", "openai:chat");
openai_endpoint.format_acceptance_config = Some(json!({
"enabled": true,
"accept_formats": ["claude:messages"],
}));
let provider_catalog = InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider_with_options("provider-shared", false, 0)],
vec![
openai_endpoint,
sample_endpoint_for_provider(
"provider-shared",
"endpoint-claude",
"claude:messages",
),
],
vec![sample_key_for_provider_with_options(
"provider-shared",
"key-shared",
"",
true,
Some(json!(["openai:chat", "claude:messages"])),
None,
)],
);
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, skipped) = resolve_and_rank_local_execution_candidates(
PlannerAppState::new(&state),
vec![
sample_priority_candidate(
"provider-shared",
"endpoint-openai",
"key-shared",
"openai:chat",
Some(0),
0,
),
sample_priority_candidate(
"provider-shared",
"endpoint-claude",
"key-shared",
"claude:messages",
Some(0),
0,
),
],
"claude:messages",
"gpt-4.1",
None,
None,
None,
)
.await;
assert!(skipped.is_empty());
assert_eq!(ranked[0].candidate.endpoint_id, "endpoint-claude");
assert_eq!(ranked[1].candidate.endpoint_id, "endpoint-openai");
assert_eq!(
ranked[1]
.ranking
.as_ref()
.and_then(|ranking| ranking.promoted_by),
None
);
assert_eq!(
ranked[1]
.ranking
.as_ref()
.and_then(|ranking| ranking.demoted_by),
Some(aether_scheduler_core::RANKING_REASON_CROSS_FORMAT)
);
}
#[tokio::test]
async fn cached_affinity_promotes_cross_format_over_same_key_exact_endpoint() {
let mut openai_endpoint =
sample_endpoint_for_provider("provider-shared", "endpoint-openai", "openai:chat");
openai_endpoint.format_acceptance_config = Some(json!({
"enabled": true,
"accept_formats": ["claude:messages"],
}));
let provider_catalog = InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider_with_options("provider-shared", false, 0)],
vec![
openai_endpoint,
sample_endpoint_for_provider(
"provider-shared",
"endpoint-claude",
"claude:messages",
),
],
vec![sample_key_for_provider_with_options(
"provider-shared",
"key-shared",
"",
true,
Some(json!(["openai:chat", "claude:messages"])),
None,
)],
);
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 auth_snapshot = sample_auth_snapshot();
let cached_cross_format = sample_priority_candidate(
"provider-shared",
"endpoint-openai",
"key-shared",
"openai:chat",
Some(0),
0,
);
remember_scheduler_affinity_for_candidate(
PlannerAppState::new(&state),
Some(&auth_snapshot),
"claude:messages",
"gpt-4.1",
&cached_cross_format,
);
let (ranked, skipped) = resolve_and_rank_local_execution_candidates(
PlannerAppState::new(&state),
vec![
cached_cross_format,
sample_priority_candidate(
"provider-shared",
"endpoint-claude",
"key-shared",
"claude:messages",
Some(0),
0,
),
],
"claude:messages",
"gpt-4.1",
Some(&auth_snapshot),
None,
None,
)
.await;
assert!(skipped.is_empty());
assert_eq!(ranked[0].candidate.endpoint_id, "endpoint-openai");
assert_eq!(
ranked[0]
.ranking
.as_ref()
.and_then(|ranking| ranking.promoted_by),
Some(RANKING_REASON_CACHED_AFFINITY)
);
assert_eq!(
ranked[0]
.ranking
.as_ref()
.and_then(|ranking| ranking.demoted_by),
Some(aether_scheduler_core::RANKING_REASON_CROSS_FORMAT)
);
assert_eq!(ranked[1].candidate.endpoint_id, "endpoint-claude");
}
#[tokio::test]
async fn non_pool_key_affinity_does_not_promote_sibling_key_when_cached_key_is_inactive() {
let provider_catalog = InMemoryProviderCatalogReadRepository::seed(

1
apps/aether-gateway/src/ai_pipeline/planner/passthrough/mod.rs
View File

@@ -5,6 +5,7 @@ mod provider;
pub(crate) use self::provider::{
build_local_stream_plan_and_reports as build_local_same_format_stream_plan_and_reports,
build_local_sync_plan_and_reports as build_local_same_format_sync_plan_and_reports,
maybe_build_local_same_format_provider_decision_payload_for_candidate,
maybe_build_stream_local_same_format_provider_decision_payload,
maybe_build_sync_local_same_format_provider_decision_payload,
resolve_same_format_provider_transport_unsupported_reason_for_trace,

307
apps/aether-gateway/src/ai_pipeline/planner/standard/family/payload.rs
View File

@@ -6,6 +6,7 @@ use crate::ai_pipeline::planner::candidate_metadata::build_request_trace_proxy_v
use crate::ai_pipeline::planner::materialization_policy::{
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
};
use crate::ai_pipeline::planner::passthrough::maybe_build_local_same_format_provider_decision_payload_for_candidate;
use crate::ai_pipeline::planner::payload_metadata::{
build_local_execution_decision_response, LocalExecutionDecisionResponseParts,
};
@@ -17,7 +18,10 @@ use crate::ai_pipeline::planner::CandidateFailureDiagnostic;
use crate::ai_pipeline::transport::{
resolve_transport_execution_timeouts, resolve_transport_tls_profile,
};
use crate::ai_pipeline::{ConversionMode, ExecutionStrategy};
use crate::ai_pipeline::{
api_format_alias_matches, resolve_local_same_format_stream_spec,
resolve_local_same_format_sync_spec, ConversionMode, ExecutionStrategy,
};
use crate::{
append_execution_contract_fields_to_value, append_local_failover_policy_to_value, AppState,
GatewayControlSyncDecisionResponse,
@@ -36,6 +40,29 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
spec: LocalStandardSpec,
) -> Option<GatewayControlSyncDecisionResponse> {
let spec_metadata = local_standard_spec_metadata(spec);
if api_format_alias_matches(
&attempt.eligible.provider_api_format,
spec_metadata.api_format,
) {
let same_format_spec = if spec_metadata.require_streaming {
resolve_local_same_format_stream_spec(spec_metadata.decision_kind)
} else {
resolve_local_same_format_sync_spec(spec_metadata.decision_kind)
};
if let Some(same_format_spec) = same_format_spec {
return maybe_build_local_same_format_provider_decision_payload_for_candidate(
state,
parts,
trace_id,
body_json,
input,
attempt,
same_format_spec,
)
.await;
}
}
let LocalStandardCandidateAttempt {
eligible,
candidate_index,
@@ -243,3 +270,281 @@ pub(super) async fn mark_skipped_local_standard_candidate_with_failure_diagnosti
)
.await;
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use aether_provider_transport::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use serde_json::json;
use super::maybe_build_local_standard_decision_payload_for_candidate;
use crate::ai_pipeline::planner::candidate_materialization::LocalExecutionCandidateAttempt;
use crate::ai_pipeline::planner::candidate_resolution::EligibleLocalExecutionCandidate;
use crate::ai_pipeline::planner::decision_input::LocalRequestedModelDecisionInput;
use crate::ai_pipeline::{
ExecutionRuntimeAuthContext, GatewayAuthApiKeySnapshot, LocalStandardSourceFamily,
LocalStandardSourceMode, LocalStandardSpec,
};
use crate::orchestration::LocalExecutionCandidateMetadata;
fn sample_auth_snapshot() -> GatewayAuthApiKeySnapshot {
GatewayAuthApiKeySnapshot {
user_id: "user-1".to_string(),
username: "alice".to_string(),
email: None,
user_role: "user".to_string(),
user_auth_source: "local".to_string(),
user_is_active: true,
user_is_deleted: false,
user_rate_limit: None,
user_allowed_providers: None,
user_allowed_api_formats: None,
user_allowed_models: None,
api_key_id: "api-key-1".to_string(),
api_key_name: Some("default".to_string()),
api_key_is_active: true,
api_key_is_locked: false,
api_key_is_standalone: false,
api_key_rate_limit: None,
api_key_concurrent_limit: None,
api_key_expires_at_unix_secs: None,
api_key_allowed_providers: None,
api_key_allowed_api_formats: None,
api_key_allowed_models: None,
currently_usable: true,
}
}
fn sample_input() -> LocalRequestedModelDecisionInput {
LocalRequestedModelDecisionInput {
auth_context: ExecutionRuntimeAuthContext {
user_id: "user-1".to_string(),
api_key_id: "api-key-1".to_string(),
username: Some("alice".to_string()),
api_key_name: Some("default".to_string()),
balance_remaining: Some(10.0),
access_allowed: true,
api_key_is_standalone: false,
},
requested_model: "claude-sonnet-4-5".to_string(),
auth_snapshot: sample_auth_snapshot(),
required_capabilities: None,
}
}
fn sample_transport(api_format: &str, endpoint_id: &str) -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "provider".to_string(),
provider_type: "custom".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: true,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: None,
},
endpoint: GatewayProviderTransportEndpoint {
id: endpoint_id.to_string(),
provider_id: "provider-1".to_string(),
api_format: api_format.to_string(),
api_family: Some(
api_format
.split_once(':')
.map(|(family, _)| family)
.unwrap_or(api_format)
.to_string(),
),
endpoint_kind: Some("chat".to_string()),
is_active: true,
base_url: "https://api.example.test".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: None,
format_acceptance_config: if api_format == "openai:chat" {
Some(json!({
"enabled": true,
"accept_formats": ["claude:messages"],
}))
} else {
None
},
proxy: None,
},
key: GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "key".to_string(),
auth_type: "api_key".to_string(),
is_active: true,
api_formats: Some(vec![
"claude:messages".to_string(),
"openai:chat".to_string(),
]),
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: Some(json!({
"claude:messages": 1,
"openai:chat": 1,
})),
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
decrypted_api_key: "sk-upstream".to_string(),
decrypted_auth_config: None,
},
}
}
fn sample_candidate(
api_format: &str,
endpoint_id: &str,
) -> SchedulerMinimalCandidateSelectionCandidate {
SchedulerMinimalCandidateSelectionCandidate {
provider_id: "provider-1".to_string(),
provider_name: "provider".to_string(),
provider_type: "custom".to_string(),
provider_priority: 1,
endpoint_id: endpoint_id.to_string(),
endpoint_api_format: api_format.to_string(),
key_id: "key-1".to_string(),
key_name: "key".to_string(),
key_auth_type: "api_key".to_string(),
key_internal_priority: 1,
key_global_priority_for_format: Some(1),
key_capabilities: None,
model_id: format!("model-{endpoint_id}"),
global_model_id: "global-model-1".to_string(),
global_model_name: "claude-sonnet-4-5".to_string(),
selected_provider_model_name: if api_format == "claude:messages" {
"claude-sonnet-4-5-upstream".to_string()
} else {
"gpt-4o-upstream".to_string()
},
mapping_matched_model: None,
}
}
fn sample_attempt(
api_format: &str,
endpoint_id: &str,
candidate_index: u32,
) -> LocalExecutionCandidateAttempt {
LocalExecutionCandidateAttempt {
eligible: EligibleLocalExecutionCandidate {
candidate: sample_candidate(api_format, endpoint_id),
transport: Arc::new(sample_transport(api_format, endpoint_id)),
provider_api_format: api_format.to_string(),
orchestration: LocalExecutionCandidateMetadata::default(),
ranking: None,
},
candidate_index,
retry_index: 0,
candidate_id: format!("candidate-{candidate_index}"),
}
}
fn claude_stream_spec() -> LocalStandardSpec {
LocalStandardSpec {
api_format: "claude:messages",
decision_kind: "claude_chat_stream",
report_kind: "claude_chat_stream_success",
family: LocalStandardSourceFamily::Standard,
mode: LocalStandardSourceMode::Chat,
require_streaming: true,
}
}
#[tokio::test]
async fn standard_family_builds_same_format_candidate_before_cross_format_candidate() {
let state = crate::AppState::new().expect("state should build");
let request = http::Request::builder()
.method("POST")
.uri("/v1/messages?beta=true")
.header(http::header::CONTENT_TYPE, "application/json")
.body(())
.expect("request should build");
let (parts, _) = request.into_parts();
let body_json = json!({
"model": "claude-sonnet-4-5",
"messages": [{"role": "user", "content": "hello"}],
"max_tokens": 32,
"stream": true
});
let input = sample_input();
let payload = maybe_build_local_standard_decision_payload_for_candidate(
&state,
&parts,
"trace-standard-same-format-first",
&body_json,
&input,
sample_attempt("claude:messages", "endpoint-claude", 0),
claude_stream_spec(),
)
.await
.expect("same-format candidate should build a standard-family payload");
assert_eq!(payload.endpoint_id.as_deref(), Some("endpoint-claude"));
assert_eq!(
payload.execution_strategy.as_deref(),
Some("local_same_format")
);
assert_eq!(payload.conversion_mode.as_deref(), Some("none"));
assert_eq!(
payload.provider_api_format.as_deref(),
Some("claude:messages")
);
assert_eq!(
payload.client_api_format.as_deref(),
Some("claude:messages")
);
assert_eq!(
payload
.provider_request_body
.as_ref()
.and_then(|body| body.get("model"))
.and_then(serde_json::Value::as_str),
Some("claude-sonnet-4-5-upstream")
);
let cross_format_payload = maybe_build_local_standard_decision_payload_for_candidate(
&state,
&parts,
"trace-standard-same-format-first",
&body_json,
&input,
sample_attempt("openai:chat", "endpoint-openai-chat", 1),
claude_stream_spec(),
)
.await
.expect("cross-format candidate should still build after the same-format candidate");
assert_eq!(
cross_format_payload.endpoint_id.as_deref(),
Some("endpoint-openai-chat")
);
assert_eq!(
cross_format_payload.execution_strategy.as_deref(),
Some("local_cross_format")
);
assert_eq!(
cross_format_payload.conversion_mode.as_deref(),
Some("bidirectional")
);
}
}

236
apps/aether-gateway/src/handlers/admin/observability/usage/summary_routes.rs
View File

@@ -8,15 +8,25 @@ use aether_admin::observability::usage::{
admin_usage_bad_request_response, admin_usage_data_unavailable_response,
admin_usage_has_fallback, admin_usage_is_failed, admin_usage_matches_search,
admin_usage_matches_username, admin_usage_parse_ids, admin_usage_parse_limit,
admin_usage_parse_offset, build_admin_usage_active_requests_response,
build_admin_usage_records_response, build_admin_usage_summary_stats_response_from_summary,
ADMIN_USAGE_DATA_UNAVAILABLE_DETAIL,
admin_usage_parse_offset, admin_usage_provider_key_name, admin_usage_record_json,
build_admin_usage_active_requests_response, build_admin_usage_records_response,
build_admin_usage_summary_stats_response_from_summary, ADMIN_USAGE_DATA_UNAVAILABLE_DETAIL,
};
use aether_data_contracts::repository::usage::{
StoredRequestUsageAudit, UsageAuditKeywordSearchQuery, UsageAuditListQuery,
UsageAuditSummaryQuery,
use aether_data::repository::users::StoredUserSummary;
use aether_data_contracts::repository::{
candidates::{RequestCandidateStatus, StoredRequestCandidate},
usage::{
StoredRequestUsageAudit, UsageAuditKeywordSearchQuery, UsageAuditListQuery,
UsageAuditSummaryQuery,
},
};
use axum::{body::Body, http, response::Response};
use axum::{
body::Body,
http,
response::{IntoResponse, Response},
Json,
};
use serde_json::json;
use std::collections::{BTreeMap, BTreeSet};
const ADMIN_USAGE_ACTIVE_LIMIT: usize = 50;
@@ -56,11 +66,198 @@ fn apply_admin_usage_status_filter(query: &mut UsageAuditListQuery, status: Opti
"pending" | "streaming" | "completed" | "cancelled" => {
query.statuses = Some(vec![status.to_string()]);
}
"has_fallback" => {}
"has_fallback" | "has_retry" => {}
_ => {}
}
}
#[derive(Clone, Copy, Debug, Default)]
struct AdminUsageAttemptFlags {
has_fallback: bool,
has_retry: bool,
}
fn admin_usage_attempt_status_filter(status: Option<&str>) -> Option<&'static str> {
match status?.trim().to_ascii_lowercase().as_str() {
"has_fallback" => Some("has_fallback"),
"has_retry" => Some("has_retry"),
_ => None,
}
}
fn admin_usage_candidate_failed_before_fallback(candidate: &StoredRequestCandidate) -> bool {
candidate.status.is_attempted(candidate.started_at_unix_ms)
&& (matches!(
candidate.status,
RequestCandidateStatus::Failed | RequestCandidateStatus::Cancelled
) || candidate.status_code.is_some_and(|code| code >= 400))
}
fn admin_usage_candidate_was_retried(candidate: &StoredRequestCandidate) -> bool {
candidate.retry_index > 0 && candidate.status.is_attempted(candidate.started_at_unix_ms)
}
fn admin_usage_final_candidate_index(
item: &StoredRequestUsageAudit,
candidates: &[StoredRequestCandidate],
) -> Option<u64> {
if let Some(candidate_id) = item.routing_candidate_id() {
if let Some(candidate) = candidates
.iter()
.find(|candidate| candidate.id == candidate_id)
{
return Some(u64::from(candidate.candidate_index));
}
}
item.routing_candidate_index().or_else(|| {
candidates
.iter()
.filter(|candidate| candidate.status.is_attempted(candidate.started_at_unix_ms))
.max_by(|left, right| {
left.candidate_index
.cmp(&right.candidate_index)
.then(left.retry_index.cmp(&right.retry_index))
})
.map(|candidate| u64::from(candidate.candidate_index))
})
}
fn admin_usage_attempt_flags_from_candidates(
item: &StoredRequestUsageAudit,
candidates: &[StoredRequestCandidate],
) -> AdminUsageAttemptFlags {
let final_candidate_index = admin_usage_final_candidate_index(item, candidates);
let has_fallback = final_candidate_index.is_some_and(|final_index| {
candidates.iter().any(|candidate| {
u64::from(candidate.candidate_index) < final_index
&& admin_usage_candidate_failed_before_fallback(candidate)
})
});
let has_retry = candidates.iter().any(admin_usage_candidate_was_retried);
AdminUsageAttemptFlags {
has_fallback,
has_retry,
}
}
fn admin_usage_attempt_flags_for_item(
item: &StoredRequestUsageAudit,
flags_by_usage_id: &BTreeMap<String, AdminUsageAttemptFlags>,
request_candidate_reader_available: bool,
) -> AdminUsageAttemptFlags {
flags_by_usage_id.get(&item.id).copied().unwrap_or_else(|| {
if request_candidate_reader_available {
AdminUsageAttemptFlags::default()
} else {
AdminUsageAttemptFlags {
has_fallback: admin_usage_has_fallback(item),
has_retry: false,
}
}
})
}
async fn resolve_admin_usage_attempt_flags_by_usage_id(
state: &AdminAppState<'_>,
items: &[StoredRequestUsageAudit],
) -> Result<BTreeMap<String, AdminUsageAttemptFlags>, GatewayError> {
if !state.has_request_candidate_data_reader() || items.is_empty() {
return Ok(BTreeMap::new());
}
let request_ids = items
.iter()
.map(|item| item.request_id.clone())
.collect::<BTreeSet<_>>();
let mut candidates_by_request_id = BTreeMap::new();
for request_id in request_ids {
candidates_by_request_id.insert(
request_id.clone(),
state
.app()
.read_request_candidates_by_request_id(&request_id)
.await?,
);
}
Ok(items
.iter()
.filter_map(|item| {
let candidates = candidates_by_request_id.get(&item.request_id)?;
Some((
item.id.clone(),
admin_usage_attempt_flags_from_candidates(item, candidates),
))
})
.collect())
}
fn admin_usage_matches_attempt_status(
item: &StoredRequestUsageAudit,
status: &str,
flags_by_usage_id: &BTreeMap<String, AdminUsageAttemptFlags>,
request_candidate_reader_available: bool,
) -> bool {
let flags = admin_usage_attempt_flags_for_item(
item,
flags_by_usage_id,
request_candidate_reader_available,
);
match status {
"has_fallback" => flags.has_fallback,
"has_retry" => flags.has_retry,
_ => true,
}
}
#[allow(clippy::too_many_arguments)]
fn build_admin_usage_records_response_with_attempt_flags(
items: &[StoredRequestUsageAudit],
users_by_id: &BTreeMap<String, StoredUserSummary>,
api_key_names: &BTreeMap<String, String>,
auth_user_reader_available: bool,
auth_api_key_reader_available: bool,
provider_key_names: &BTreeMap<String, String>,
attempt_flags_by_usage_id: &BTreeMap<String, AdminUsageAttemptFlags>,
request_candidate_reader_available: bool,
total: usize,
limit: usize,
offset: usize,
) -> Response<Body> {
let records: Vec<_> = items
.iter()
.map(|item| {
let provider_key_name = admin_usage_provider_key_name(item, provider_key_names);
let mut record = admin_usage_record_json(
item,
users_by_id,
api_key_names,
auth_user_reader_available,
auth_api_key_reader_available,
provider_key_name.as_deref(),
);
let flags = admin_usage_attempt_flags_for_item(
item,
attempt_flags_by_usage_id,
request_candidate_reader_available,
);
record["has_fallback"] = json!(flags.has_fallback);
record["has_retry"] = json!(flags.has_retry);
record
})
.collect();
Json(json!({
"records": records,
"total": total,
"limit": limit,
"offset": offset,
}))
.into_response()
}
fn build_admin_usage_records_query(
created_from_unix_secs: u64,
created_until_unix_secs: u64,
@@ -342,9 +539,8 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
Ok(value) => value,
Err(detail) => return Ok(Some(admin_usage_bad_request_response(detail))),
};
let has_fallback_only = query_param_value(query, "status")
.as_deref()
.is_some_and(|value| value.trim().eq_ignore_ascii_case("has_fallback"));
let attempt_status_filter =
admin_usage_attempt_status_filter(query_param_value(query, "status").as_deref());
let search = query_param_value(query, "search");
let username_filter = query_param_value(query, "username");
let limit = match admin_usage_parse_limit(query) {
@@ -381,7 +577,7 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
let active_username_filter = username_filter
.as_deref()
.filter(|value| !value.trim().is_empty());
let (usage, total) = if has_fallback_only {
let (usage, total) = if let Some(attempt_status) = attempt_status_filter {
let mut usage = state.list_usage_audits(&base_query).await?;
let user_ids: Vec<String> = usage
.iter()
@@ -394,6 +590,9 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
aether_data::repository::users::StoredUserSummary,
> = state.resolve_auth_user_summaries_by_ids(&user_ids).await?;
let api_key_names = admin_usage_api_key_names(state, &usage).await?;
let attempt_flags_by_usage_id =
resolve_admin_usage_attempt_flags_by_usage_id(state, &usage).await?;
let request_candidate_reader_available = state.has_request_candidate_data_reader();
usage.retain(|item| {
admin_usage_matches_search(
@@ -408,7 +607,12 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
active_username_filter,
&users_by_id,
state.has_auth_user_data_reader(),
) && admin_usage_has_fallback(item)
) && admin_usage_matches_attempt_status(
item,
attempt_status,
&attempt_flags_by_usage_id,
request_candidate_reader_available,
)
});
sort_usage_newest_first(&mut usage);
let total = usage.len();
@@ -473,14 +677,18 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
state.resolve_auth_user_summaries_by_ids(&user_ids).await?;
let api_key_names = admin_usage_api_key_names(state, &usage).await?;
let provider_key_names = admin_usage_provider_key_names(state, &usage).await?;
let attempt_flags_by_usage_id =
resolve_admin_usage_attempt_flags_by_usage_id(state, &usage).await?;
return Ok(Some(build_admin_usage_records_response(
return Ok(Some(build_admin_usage_records_response_with_attempt_flags(
&usage,
&users_by_id,
&api_key_names,
state.has_auth_user_data_reader(),
state.has_auth_api_key_data_reader(),
&provider_key_names,
&attempt_flags_by_usage_id,
state.has_request_candidate_data_reader(),
total,
limit,
offset,

135
apps/aether-gateway/src/orchestration/effects.rs
View File

@@ -4,6 +4,7 @@ use aether_admin::provider::quota as admin_provider_quota_pure;
use aether_contracts::{ExecutionPlan, ExecutionTelemetry};
use aether_scheduler_core::{
build_scheduler_affinity_cache_key_for_api_key_id, count_recent_rpm_requests_for_provider_key,
SchedulerAffinityTarget,
};
use aether_usage_runtime::{
build_stream_terminal_usage_outcome, build_sync_terminal_usage_outcome,
@@ -24,6 +25,7 @@ use crate::handlers::shared::provider_pool::{
record_admin_provider_pool_error, record_admin_provider_pool_stream_timeout,
record_admin_provider_pool_success, AdminProviderPoolConfig,
};
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
use crate::AppState;
#[derive(Debug, Clone, Copy)]
@@ -96,6 +98,7 @@ struct PoolFeedbackContext {
}
const ADAPTIVE_RPM_RECENT_CANDIDATE_LIMIT: usize = 512;
const LOCAL_EXECUTION_SCHEDULER_AFFINITY_MAX_ENTRIES: usize = 10_000;
pub(crate) async fn apply_local_execution_effect(
state: &AppState,
@@ -155,6 +158,40 @@ fn local_scheduler_affinity_cache_key(report_context: Option<&Value>) -> Option<
)
}
fn local_scheduler_affinity_target(plan: &ExecutionPlan) -> Option<SchedulerAffinityTarget> {
let provider_id = plan.provider_id.trim();
let endpoint_id = plan.endpoint_id.trim();
let key_id = plan.key_id.trim();
if provider_id.is_empty() || endpoint_id.is_empty() || key_id.is_empty() {
return None;
}
Some(SchedulerAffinityTarget {
provider_id: provider_id.to_string(),
endpoint_id: endpoint_id.to_string(),
key_id: key_id.to_string(),
})
}
fn remember_successful_local_scheduler_affinity(
state: &AppState,
context: LocalExecutionEffectContext<'_>,
) {
let Some(cache_key) = local_scheduler_affinity_cache_key(context.report_context) else {
return;
};
let Some(target) = local_scheduler_affinity_target(context.plan) else {
return;
};
state.remember_scheduler_affinity_target(
&cache_key,
target,
SCHEDULER_AFFINITY_TTL,
LOCAL_EXECUTION_SCHEDULER_AFFINITY_MAX_ENTRIES,
);
}
fn pool_feedback_request_body<'a>(
plan: &'a ExecutionPlan,
report_context: Option<&'a Value>,
@@ -415,6 +452,8 @@ async fn record_health_success_effect(
context: LocalExecutionEffectContext<'_>,
_effect: LocalHealthSuccessEffect,
) {
remember_successful_local_scheduler_affinity(state, context);
let api_format = context.plan.provider_api_format.trim();
if api_format.is_empty() {
return;
@@ -1039,6 +1078,102 @@ mod tests {
.is_some());
}
#[tokio::test]
async fn success_remembers_scheduler_affinity_cache_for_final_candidate() {
let state = AppState::new().expect("gateway state should build");
let plan = sample_plan();
let report_context = json!({
"api_key_id": "api-key-1",
"client_api_format": "openai:chat",
"model": "gpt-5",
});
let cache_key =
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
.expect("scheduler affinity cache key should build");
apply_local_execution_effect(
&state,
LocalExecutionEffectContext {
plan: &plan,
report_context: Some(&report_context),
},
LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect),
)
.await;
assert_eq!(
state.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL),
Some(SchedulerAffinityTarget {
provider_id: "prov-1".to_string(),
endpoint_id: "ep-1".to_string(),
key_id: "key-1".to_string(),
})
);
}
#[tokio::test]
async fn fallback_success_rewarms_scheduler_affinity_after_failed_candidate_invalidates() {
let state = AppState::new().expect("gateway state should build");
let failed_plan = sample_plan();
let mut success_plan = sample_plan();
success_plan.provider_id = "prov-2".to_string();
success_plan.endpoint_id = "ep-2".to_string();
success_plan.key_id = "key-2".to_string();
let report_context = json!({
"api_key_id": "api-key-1",
"client_api_format": "openai:chat",
"model": "gpt-5",
});
let cache_key =
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
.expect("scheduler affinity cache key should build");
state.scheduler_affinity_cache.insert(
cache_key.clone(),
SchedulerAffinityTarget {
provider_id: "prov-1".to_string(),
endpoint_id: "ep-1".to_string(),
key_id: "key-1".to_string(),
},
SCHEDULER_AFFINITY_TTL,
16,
);
apply_local_execution_effect(
&state,
LocalExecutionEffectContext {
plan: &failed_plan,
report_context: Some(&report_context),
},
LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect {
status_code: 429,
classification: LocalFailoverClassification::RetryUpstreamFailure,
}),
)
.await;
assert!(state
.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL)
.is_none());
apply_local_execution_effect(
&state,
LocalExecutionEffectContext {
plan: &success_plan,
report_context: Some(&report_context),
},
LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect),
)
.await;
assert_eq!(
state.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL),
Some(SchedulerAffinityTarget {
provider_id: "prov-2".to_string(),
endpoint_id: "ep-2".to_string(),
key_id: "key-2".to_string(),
})
);
}
#[test]
fn semantic_client_error_does_not_penalize_pool_feedback() {
assert!(!local_candidate_failure_should_record_pool_error(

35
apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/direct.rs
View File

@@ -25,6 +25,30 @@ use aether_data_contracts::repository::provider_catalog::{
};
use sha2::{Digest, Sha256};
const KIRO_CLAUDE_CLI_FINALIZE_TEST_STACK_BYTES: usize = 16 * 1024 * 1024;
fn run_kiro_claude_cli_finalize_test<F, Fut>(test_name: &'static str, make_future: F)
where
F: FnOnce() -> Fut + Send + 'static,
Fut: std::future::Future<Output = ()> + 'static,
{
let handle = std::thread::Builder::new()
.name(test_name.to_string())
.stack_size(KIRO_CLAUDE_CLI_FINALIZE_TEST_STACK_BYTES)
.spawn(move || {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("test runtime should build");
runtime.block_on(make_future());
})
.expect("kiro claude cli finalize test thread should spawn");
if let Err(payload) = handle.join() {
std::panic::resume_unwind(payload);
}
}
#[tokio::test]
async fn gateway_executes_openai_responses_sync_upstream_stream_via_local_finalize_response() {
use base64::Engine as _;
@@ -474,8 +498,15 @@ async fn gateway_executes_openai_responses_sync_upstream_stream_via_local_finali
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_executes_kiro_claude_cli_sync_upstream_stream_via_local_finalize_response() {
#[test]
fn gateway_executes_kiro_claude_cli_sync_upstream_stream_via_local_finalize_response() {
run_kiro_claude_cli_finalize_test(
"gateway_executes_kiro_claude_cli_sync_upstream_stream_via_local_finalize_response",
gateway_executes_kiro_claude_cli_sync_upstream_stream_via_local_finalize_response_impl,
);
}
async fn gateway_executes_kiro_claude_cli_sync_upstream_stream_via_local_finalize_response_impl() {
use base64::Engine as _;
fn crc32(data: &[u8]) -> u32 {

57
apps/aether-gateway/src/tests/ai_execute/finalize_local_provider/claude.rs
View File

@@ -23,8 +23,39 @@ use aether_data_contracts::repository::provider_catalog::{
};
use sha2::{Digest, Sha256};
#[tokio::test]
async fn gateway_executes_claude_chat_sync_same_format_via_local_finalize_response() {
const CLAUDE_PROVIDER_FINALIZE_TEST_STACK_BYTES: usize = 16 * 1024 * 1024;
fn run_claude_provider_finalize_test<F, Fut>(test_name: &'static str, make_future: F)
where
F: FnOnce() -> Fut + Send + 'static,
Fut: std::future::Future<Output = ()> + 'static,
{
let handle = std::thread::Builder::new()
.name(test_name.to_string())
.stack_size(CLAUDE_PROVIDER_FINALIZE_TEST_STACK_BYTES)
.spawn(move || {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("test runtime should build");
runtime.block_on(make_future());
})
.expect("claude provider finalize test thread should spawn");
if let Err(payload) = handle.join() {
std::panic::resume_unwind(payload);
}
}
#[test]
fn gateway_executes_claude_chat_sync_same_format_via_local_finalize_response() {
run_claude_provider_finalize_test(
"gateway_executes_claude_chat_sync_same_format_via_local_finalize_response",
gateway_executes_claude_chat_sync_same_format_via_local_finalize_response_impl,
);
}
async fn gateway_executes_claude_chat_sync_same_format_via_local_finalize_response_impl() {
#[derive(Debug, Clone)]
struct SeenRemoteExecutionRuntimeRequest {
trace_id: String,
@@ -468,8 +499,15 @@ async fn gateway_executes_claude_chat_sync_same_format_via_local_finalize_respon
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_executes_claude_chat_sync_upstream_stream_via_local_finalize_response() {
#[test]
fn gateway_executes_claude_chat_sync_upstream_stream_via_local_finalize_response() {
run_claude_provider_finalize_test(
"gateway_executes_claude_chat_sync_upstream_stream_via_local_finalize_response",
gateway_executes_claude_chat_sync_upstream_stream_via_local_finalize_response_impl,
);
}
async fn gateway_executes_claude_chat_sync_upstream_stream_via_local_finalize_response_impl() {
use base64::Engine as _;
#[derive(Debug, Clone)]
@@ -921,8 +959,15 @@ async fn gateway_executes_claude_chat_sync_upstream_stream_via_local_finalize_re
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_executes_claude_cli_sync_upstream_stream_via_local_finalize_response() {
#[test]
fn gateway_executes_claude_cli_sync_upstream_stream_via_local_finalize_response() {
run_claude_provider_finalize_test(
"gateway_executes_claude_cli_sync_upstream_stream_via_local_finalize_response",
gateway_executes_claude_cli_sync_upstream_stream_via_local_finalize_response_impl,
);
}
async fn gateway_executes_claude_cli_sync_upstream_stream_via_local_finalize_response_impl() {
use base64::Engine as _;
#[derive(Debug, Clone)]

70
apps/aether-gateway/src/tests/ai_execute/finalize_local_provider/gemini.rs
View File

@@ -23,8 +23,39 @@ use aether_data_contracts::repository::provider_catalog::{
};
use sha2::{Digest, Sha256};
#[tokio::test]
async fn gateway_executes_gemini_chat_sync_same_format_via_local_finalize_response() {
const GEMINI_PROVIDER_FINALIZE_TEST_STACK_BYTES: usize = 16 * 1024 * 1024;
fn run_gemini_provider_finalize_test<F, Fut>(test_name: &'static str, make_future: F)
where
F: FnOnce() -> Fut + Send + 'static,
Fut: std::future::Future<Output = ()> + 'static,
{
let handle = std::thread::Builder::new()
.name(test_name.to_string())
.stack_size(GEMINI_PROVIDER_FINALIZE_TEST_STACK_BYTES)
.spawn(move || {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("test runtime should build");
runtime.block_on(make_future());
})
.expect("gemini provider finalize test thread should spawn");
if let Err(payload) = handle.join() {
std::panic::resume_unwind(payload);
}
}
#[test]
fn gateway_executes_gemini_chat_sync_same_format_via_local_finalize_response() {
run_gemini_provider_finalize_test(
"gateway_executes_gemini_chat_sync_same_format_via_local_finalize_response",
gateway_executes_gemini_chat_sync_same_format_via_local_finalize_response_impl,
);
}
async fn gateway_executes_gemini_chat_sync_same_format_via_local_finalize_response_impl() {
#[derive(Debug, Clone)]
struct SeenRemoteExecutionRuntimeRequest {
trace_id: String,
@@ -517,8 +548,15 @@ async fn gateway_executes_gemini_chat_sync_same_format_via_local_finalize_respon
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_executes_gemini_chat_sync_upstream_stream_via_local_finalize_response() {
#[test]
fn gateway_executes_gemini_chat_sync_upstream_stream_via_local_finalize_response() {
run_gemini_provider_finalize_test(
"gateway_executes_gemini_chat_sync_upstream_stream_via_local_finalize_response",
gateway_executes_gemini_chat_sync_upstream_stream_via_local_finalize_response_impl,
);
}
async fn gateway_executes_gemini_chat_sync_upstream_stream_via_local_finalize_response_impl() {
use base64::Engine as _;
#[derive(Debug, Clone)]
@@ -1003,8 +1041,15 @@ async fn gateway_executes_gemini_chat_sync_upstream_stream_via_local_finalize_re
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_executes_gemini_cli_sync_upstream_stream_via_local_finalize_response() {
#[test]
fn gateway_executes_gemini_cli_sync_upstream_stream_via_local_finalize_response() {
run_gemini_provider_finalize_test(
"gateway_executes_gemini_cli_sync_upstream_stream_via_local_finalize_response",
gateway_executes_gemini_cli_sync_upstream_stream_via_local_finalize_response_impl,
);
}
async fn gateway_executes_gemini_cli_sync_upstream_stream_via_local_finalize_response_impl() {
use base64::Engine as _;
#[derive(Debug, Clone)]
@@ -1491,9 +1536,16 @@ async fn gateway_executes_gemini_cli_sync_upstream_stream_via_local_finalize_res
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_executes_antigravity_gemini_cli_sync_upstream_stream_via_local_finalize_response()
{
#[test]
fn gateway_executes_antigravity_gemini_cli_sync_upstream_stream_via_local_finalize_response() {
run_gemini_provider_finalize_test(
"gateway_executes_antigravity_gemini_cli_sync_upstream_stream_via_local_finalize_response",
gateway_executes_antigravity_gemini_cli_sync_upstream_stream_via_local_finalize_response_impl,
);
}
async fn gateway_executes_antigravity_gemini_cli_sync_upstream_stream_via_local_finalize_response_impl(
) {
use base64::Engine as _;
#[derive(Debug, Clone)]

36
apps/aether-gateway/src/tests/ai_execute/sync/claude/claude_code.rs
View File

@@ -9,8 +9,40 @@ use super::{
DEVELOPMENT_ENCRYPTION_KEY, TRACE_ID_HEADER,
};
#[tokio::test]
async fn gateway_executes_claude_code_cli_sync_via_local_decision_gate_with_local_sync_decision() {
const CLAUDE_CODE_CLI_SYNC_TEST_STACK_BYTES: usize = 16 * 1024 * 1024;
fn run_claude_code_cli_sync_test<F, Fut>(test_name: &'static str, make_future: F)
where
F: FnOnce() -> Fut + Send + 'static,
Fut: std::future::Future<Output = ()> + 'static,
{
let handle = std::thread::Builder::new()
.name(test_name.to_string())
.stack_size(CLAUDE_CODE_CLI_SYNC_TEST_STACK_BYTES)
.spawn(move || {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("test runtime should build");
runtime.block_on(make_future());
})
.expect("claude code cli sync test thread should spawn");
if let Err(payload) = handle.join() {
std::panic::resume_unwind(payload);
}
}
#[test]
fn gateway_executes_claude_code_cli_sync_via_local_decision_gate_with_local_sync_decision() {
run_claude_code_cli_sync_test(
"gateway_executes_claude_code_cli_sync_via_local_decision_gate_with_local_sync_decision",
gateway_executes_claude_code_cli_sync_via_local_decision_gate_with_local_sync_decision_impl,
);
}
async fn gateway_executes_claude_code_cli_sync_via_local_decision_gate_with_local_sync_decision_impl(
) {
#[derive(Debug, Clone)]
struct SeenExecutionRuntimeSyncRequest {
trace_id: String,

57
apps/aether-gateway/src/tests/ai_execute/sync/claude/local_chat.rs
View File

@@ -11,8 +11,39 @@ use super::{
TRACE_ID_HEADER,
};
#[tokio::test]
async fn gateway_executes_claude_chat_sync_via_local_decision_gate_with_local_sync_decision() {
const CLAUDE_CHAT_SYNC_TEST_STACK_BYTES: usize = 16 * 1024 * 1024;
fn run_claude_chat_sync_test<F, Fut>(test_name: &'static str, make_future: F)
where
F: FnOnce() -> Fut + Send + 'static,
Fut: std::future::Future<Output = ()> + 'static,
{
let handle = std::thread::Builder::new()
.name(test_name.to_string())
.stack_size(CLAUDE_CHAT_SYNC_TEST_STACK_BYTES)
.spawn(move || {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("test runtime should build");
runtime.block_on(make_future());
})
.expect("claude chat sync test thread should spawn");
if let Err(payload) = handle.join() {
std::panic::resume_unwind(payload);
}
}
#[test]
fn gateway_executes_claude_chat_sync_via_local_decision_gate_with_local_sync_decision() {
run_claude_chat_sync_test(
"gateway_executes_claude_chat_sync_via_local_decision_gate_with_local_sync_decision",
gateway_executes_claude_chat_sync_via_local_decision_gate_with_local_sync_decision_impl,
);
}
async fn gateway_executes_claude_chat_sync_via_local_decision_gate_with_local_sync_decision_impl() {
#[derive(Debug, Clone)]
struct SeenExecutionRuntimeSyncRequest {
trace_id: String,
@@ -452,8 +483,15 @@ async fn gateway_executes_claude_chat_sync_via_local_decision_gate_with_local_sy
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_surfaces_candidate_list_empty_reason_for_claude_chat_runtime_miss() {
#[test]
fn gateway_surfaces_candidate_list_empty_reason_for_claude_chat_runtime_miss() {
run_claude_chat_sync_test(
"gateway_surfaces_candidate_list_empty_reason_for_claude_chat_runtime_miss",
gateway_surfaces_candidate_list_empty_reason_for_claude_chat_runtime_miss_impl,
);
}
async fn gateway_surfaces_candidate_list_empty_reason_for_claude_chat_runtime_miss_impl() {
fn hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());
@@ -577,8 +615,15 @@ async fn gateway_surfaces_candidate_list_empty_reason_for_claude_chat_runtime_mi
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_returns_claude_chat_error_for_local_sync_failure() {
#[test]
fn gateway_returns_claude_chat_error_for_local_sync_failure() {
run_claude_chat_sync_test(
"gateway_returns_claude_chat_error_for_local_sync_failure",
gateway_returns_claude_chat_error_for_local_sync_failure_impl,
);
}
async fn gateway_returns_claude_chat_error_for_local_sync_failure_impl() {
fn hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());

58
apps/aether-gateway/src/tests/ai_execute/sync/claude/local_cli.rs
View File

@@ -11,8 +11,39 @@ use super::{
TRACE_ID_HEADER,
};
#[tokio::test]
async fn gateway_executes_claude_cli_sync_via_local_decision_gate_with_local_sync_decision() {
const CLAUDE_CLI_SYNC_TEST_STACK_BYTES: usize = 16 * 1024 * 1024;
fn run_claude_cli_sync_test<F, Fut>(test_name: &'static str, make_future: F)
where
F: FnOnce() -> Fut + Send + 'static,
Fut: std::future::Future<Output = ()> + 'static,
{
let handle = std::thread::Builder::new()
.name(test_name.to_string())
.stack_size(CLAUDE_CLI_SYNC_TEST_STACK_BYTES)
.spawn(move || {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("test runtime should build");
runtime.block_on(make_future());
})
.expect("claude cli sync test thread should spawn");
if let Err(payload) = handle.join() {
std::panic::resume_unwind(payload);
}
}
#[test]
fn gateway_executes_claude_cli_sync_via_local_decision_gate_with_local_sync_decision() {
run_claude_cli_sync_test(
"gateway_executes_claude_cli_sync_via_local_decision_gate_with_local_sync_decision",
gateway_executes_claude_cli_sync_via_local_decision_gate_with_local_sync_decision_impl,
);
}
async fn gateway_executes_claude_cli_sync_via_local_decision_gate_with_local_sync_decision_impl() {
#[derive(Debug, Clone)]
struct SeenExecutionRuntimeSyncRequest {
trace_id: String,
@@ -460,8 +491,15 @@ async fn gateway_executes_claude_cli_sync_via_local_decision_gate_with_local_syn
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_returns_claude_cli_error_for_local_sync_failure() {
#[test]
fn gateway_returns_claude_cli_error_for_local_sync_failure() {
run_claude_cli_sync_test(
"gateway_returns_claude_cli_error_for_local_sync_failure",
gateway_returns_claude_cli_error_for_local_sync_failure_impl,
);
}
async fn gateway_returns_claude_cli_error_for_local_sync_failure_impl() {
fn hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());
@@ -734,8 +772,16 @@ async fn gateway_returns_claude_cli_error_for_local_sync_failure() {
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_marks_claude_cli_cross_format_runtime_miss_when_format_conversion_is_disabled() {
#[test]
fn gateway_marks_claude_cli_cross_format_runtime_miss_when_format_conversion_is_disabled() {
run_claude_cli_sync_test(
"gateway_marks_claude_cli_cross_format_runtime_miss_when_format_conversion_is_disabled",
gateway_marks_claude_cli_cross_format_runtime_miss_when_format_conversion_is_disabled_impl,
);
}
async fn gateway_marks_claude_cli_cross_format_runtime_miss_when_format_conversion_is_disabled_impl(
) {
fn hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());

81
apps/aether-gateway/src/tests/ai_execute/sync/gemini/cli.rs
View File

@@ -10,8 +10,39 @@ use super::{
TRACE_ID_HEADER,
};
#[tokio::test]
async fn gateway_executes_gemini_cli_sync_via_local_decision_gate_with_local_sync_decision() {
const GEMINI_CLI_SYNC_TEST_STACK_BYTES: usize = 16 * 1024 * 1024;
fn run_gemini_cli_sync_test<F, Fut>(test_name: &'static str, make_future: F)
where
F: FnOnce() -> Fut + Send + 'static,
Fut: std::future::Future<Output = ()> + 'static,
{
let handle = std::thread::Builder::new()
.name(test_name.to_string())
.stack_size(GEMINI_CLI_SYNC_TEST_STACK_BYTES)
.spawn(move || {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("test runtime should build");
runtime.block_on(make_future());
})
.expect("gemini cli sync test thread should spawn");
if let Err(payload) = handle.join() {
std::panic::resume_unwind(payload);
}
}
#[test]
fn gateway_executes_gemini_cli_sync_via_local_decision_gate_with_local_sync_decision() {
run_gemini_cli_sync_test(
"gateway_executes_gemini_cli_sync_via_local_decision_gate_with_local_sync_decision",
gateway_executes_gemini_cli_sync_via_local_decision_gate_with_local_sync_decision_impl,
);
}
async fn gateway_executes_gemini_cli_sync_via_local_decision_gate_with_local_sync_decision_impl() {
#[derive(Debug, Clone)]
struct SeenExecutionRuntimeSyncRequest {
trace_id: String,
@@ -445,8 +476,15 @@ async fn gateway_executes_gemini_cli_sync_via_local_decision_gate_with_local_syn
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_returns_gemini_cli_error_for_local_sync_failure() {
#[test]
fn gateway_returns_gemini_cli_error_for_local_sync_failure() {
run_gemini_cli_sync_test(
"gateway_returns_gemini_cli_error_for_local_sync_failure",
gateway_returns_gemini_cli_error_for_local_sync_failure_impl,
);
}
async fn gateway_returns_gemini_cli_error_for_local_sync_failure_impl() {
fn hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());
@@ -716,8 +754,15 @@ async fn gateway_returns_gemini_cli_error_for_local_sync_failure() {
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_executes_gemini_cli_sync_via_local_decision_gate_after_oauth_refresh() {
#[test]
fn gateway_executes_gemini_cli_sync_via_local_decision_gate_after_oauth_refresh() {
run_gemini_cli_sync_test(
"gateway_executes_gemini_cli_sync_via_local_decision_gate_after_oauth_refresh",
gateway_executes_gemini_cli_sync_via_local_decision_gate_after_oauth_refresh_impl,
);
}
async fn gateway_executes_gemini_cli_sync_via_local_decision_gate_after_oauth_refresh_impl() {
#[derive(Debug, Clone)]
struct SeenExecutionRuntimeSyncRequest {
trace_id: String,
@@ -1220,8 +1265,15 @@ async fn gateway_executes_gemini_cli_sync_via_local_decision_gate_after_oauth_re
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_executes_vertex_ai_gemini_cli_sync_via_local_decision_gate_with_local_sync_decision(
#[test]
fn gateway_executes_vertex_ai_gemini_cli_sync_via_local_decision_gate_with_local_sync_decision() {
run_gemini_cli_sync_test(
"gateway_executes_vertex_ai_gemini_cli_sync_via_local_decision_gate_with_local_sync_decision",
gateway_executes_vertex_ai_gemini_cli_sync_via_local_decision_gate_with_local_sync_decision_impl,
);
}
async fn gateway_executes_vertex_ai_gemini_cli_sync_via_local_decision_gate_with_local_sync_decision_impl(
) {
#[derive(Debug, Clone)]
struct SeenExecutionRuntimeSyncRequest {
@@ -1642,9 +1694,16 @@ async fn gateway_executes_vertex_ai_gemini_cli_sync_via_local_decision_gate_with
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_executes_antigravity_gemini_cli_sync_via_local_decision_gate_after_oauth_refresh()
{
#[test]
fn gateway_executes_antigravity_gemini_cli_sync_via_local_decision_gate_after_oauth_refresh() {
run_gemini_cli_sync_test(
"gateway_executes_antigravity_gemini_cli_sync_via_local_decision_gate_after_oauth_refresh",
gateway_executes_antigravity_gemini_cli_sync_via_local_decision_gate_after_oauth_refresh_impl,
);
}
async fn gateway_executes_antigravity_gemini_cli_sync_via_local_decision_gate_after_oauth_refresh_impl(
) {
use base64::Engine as _;
#[derive(Debug, Clone)]

46
apps/aether-gateway/src/tests/ai_execute/sync/gemini/local_chat.rs
View File

@@ -10,8 +10,39 @@ use super::{
TRACE_ID_HEADER,
};
#[tokio::test]
async fn gateway_executes_gemini_chat_sync_via_local_decision_gate_with_local_sync_decision() {
const GEMINI_CHAT_SYNC_TEST_STACK_BYTES: usize = 16 * 1024 * 1024;
fn run_gemini_chat_sync_test<F, Fut>(test_name: &'static str, make_future: F)
where
F: FnOnce() -> Fut + Send + 'static,
Fut: std::future::Future<Output = ()> + 'static,
{
let handle = std::thread::Builder::new()
.name(test_name.to_string())
.stack_size(GEMINI_CHAT_SYNC_TEST_STACK_BYTES)
.spawn(move || {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("test runtime should build");
runtime.block_on(make_future());
})
.expect("gemini chat sync test thread should spawn");
if let Err(payload) = handle.join() {
std::panic::resume_unwind(payload);
}
}
#[test]
fn gateway_executes_gemini_chat_sync_via_local_decision_gate_with_local_sync_decision() {
run_gemini_chat_sync_test(
"gateway_executes_gemini_chat_sync_via_local_decision_gate_with_local_sync_decision",
gateway_executes_gemini_chat_sync_via_local_decision_gate_with_local_sync_decision_impl,
);
}
async fn gateway_executes_gemini_chat_sync_via_local_decision_gate_with_local_sync_decision_impl() {
#[derive(Debug, Clone)]
struct SeenExecutionRuntimeSyncRequest {
trace_id: String,
@@ -435,8 +466,15 @@ async fn gateway_executes_gemini_chat_sync_via_local_decision_gate_with_local_sy
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_returns_gemini_chat_error_for_local_sync_failure() {
#[test]
fn gateway_returns_gemini_chat_error_for_local_sync_failure() {
run_gemini_chat_sync_test(
"gateway_returns_gemini_chat_error_for_local_sync_failure",
gateway_returns_gemini_chat_error_for_local_sync_failure_impl,
);
}
async fn gateway_returns_gemini_chat_error_for_local_sync_failure_impl() {
fn hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());

225
apps/aether-gateway/src/tests/control/admin/usage.rs
View File

@@ -3,11 +3,15 @@ use std::sync::{Arc, Mutex};
use aether_data::repository::auth::{
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
};
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::usage::InMemoryUsageReadRepository;
use aether_data::repository::users::{
InMemoryUserReadRepository, StoredUserAuthRecord, StoredUserSummary,
};
use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, StoredRequestCandidate,
};
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
use axum::body::{Body, Bytes};
use axum::routing::{any, get, post};
@@ -201,6 +205,43 @@ fn sample_usage_row(
usage
}
fn sample_request_candidate(
id: &str,
request_id: &str,
candidate_index: i32,
retry_index: i32,
status: RequestCandidateStatus,
) -> StoredRequestCandidate {
let attempted = status.is_attempted(None);
StoredRequestCandidate::new(
id.to_string(),
request_id.to_string(),
Some("user-1".to_string()),
Some("key-1".to_string()),
Some("alice".to_string()),
Some("primary".to_string()),
candidate_index,
retry_index,
Some("provider-1".to_string()),
Some(format!("endpoint-{candidate_index}")),
Some(format!("provider-key-{candidate_index}")),
status,
None,
false,
matches!(status, RequestCandidateStatus::Failed).then_some(503),
None,
matches!(status, RequestCandidateStatus::Failed).then(|| "upstream failed".to_string()),
attempted.then_some(50),
None,
None,
None,
1_711_000_000_000 + i64::from(candidate_index) * 10 + i64::from(retry_index),
attempted.then_some(1_711_000_000_000 + i64::from(candidate_index) * 10),
attempted.then_some(1_711_000_000_005 + i64::from(candidate_index) * 10),
)
.expect("request candidate should build")
}
fn sample_user_summary(id: &str, username: &str) -> StoredUserSummary {
StoredUserSummary::new(
id.to_string(),
@@ -1203,6 +1244,190 @@ async fn gateway_filters_admin_usage_records_by_has_fallback_status() {
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_usage_record_attempt_flags_follow_request_candidate_timeline() {
let (_upstream_url, upstream_hits, upstream_handle) =
start_usage_upstream("/api/admin/usage/records").await;
let mut non_fallback_usage = sample_usage_row(
"usage-non-fallback-secondary",
"req-non-fallback-secondary",
Some("user-1"),
Some("key-1"),
Some("primary"),
"OpenAI",
"gpt-5",
"completed",
12,
8,
0.02,
0.02,
DAY_1_UNIX_SECS + 3,
);
non_fallback_usage.candidate_id = Some("cand-non-fallback-success".to_string());
non_fallback_usage.candidate_index = Some(1);
let mut fallback_usage = sample_usage_row(
"usage-real-fallback",
"req-real-fallback",
Some("user-1"),
Some("key-1"),
Some("primary"),
"OpenAI",
"gpt-5",
"completed",
12,
8,
0.02,
0.02,
DAY_1_UNIX_SECS + 2,
);
fallback_usage.candidate_id = Some("cand-fallback-success".to_string());
fallback_usage.candidate_index = Some(1);
let mut retry_usage = sample_usage_row(
"usage-real-retry",
"req-real-retry",
Some("user-1"),
Some("key-1"),
Some("primary"),
"OpenAI",
"gpt-5",
"completed",
12,
8,
0.02,
0.02,
DAY_1_UNIX_SECS + 1,
);
retry_usage.candidate_id = Some("cand-retry-success".to_string());
retry_usage.candidate_index = Some(0);
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
non_fallback_usage,
fallback_usage,
retry_usage,
]));
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
sample_request_candidate(
"cand-non-fallback-success",
"req-non-fallback-secondary",
1,
0,
RequestCandidateStatus::Success,
),
sample_request_candidate(
"cand-non-fallback-unused",
"req-non-fallback-secondary",
2,
0,
RequestCandidateStatus::Unused,
),
sample_request_candidate(
"cand-fallback-failed",
"req-real-fallback",
0,
0,
RequestCandidateStatus::Failed,
),
sample_request_candidate(
"cand-fallback-success",
"req-real-fallback",
1,
0,
RequestCandidateStatus::Success,
),
sample_request_candidate(
"cand-retry-failed",
"req-real-retry",
0,
0,
RequestCandidateStatus::Failed,
),
sample_request_candidate(
"cand-retry-success",
"req-real-retry",
0,
1,
RequestCandidateStatus::Success,
),
]));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
request_candidate_repository,
usage_repository,
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = admin_request(reqwest::Client::new().get(format!(
"{gateway_url}/api/admin/usage/records?start_date=2024-03-21&end_date=2024-03-22&tz_offset_minutes=0&limit=10&offset=0"
)))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
let records = payload["records"]
.as_array()
.expect("records should be array");
let record_by_id = |id: &str| {
records
.iter()
.find(|record| record["id"].as_str() == Some(id))
.expect("record should exist")
};
assert_eq!(
record_by_id("usage-non-fallback-secondary")["has_fallback"],
false
);
assert_eq!(
record_by_id("usage-non-fallback-secondary")["has_retry"],
false
);
assert_eq!(record_by_id("usage-real-fallback")["has_fallback"], true);
assert_eq!(record_by_id("usage-real-fallback")["has_retry"], false);
assert_eq!(record_by_id("usage-real-retry")["has_fallback"], false);
assert_eq!(record_by_id("usage-real-retry")["has_retry"], true);
let fallback_response = admin_request(reqwest::Client::new().get(format!(
"{gateway_url}/api/admin/usage/records?start_date=2024-03-21&end_date=2024-03-22&tz_offset_minutes=0&status=has_fallback&limit=10&offset=0"
)))
.send()
.await
.expect("fallback request should succeed");
assert_eq!(fallback_response.status(), StatusCode::OK);
let fallback_payload: serde_json::Value = fallback_response
.json()
.await
.expect("fallback json body should parse");
assert_eq!(fallback_payload["total"], 1);
assert_eq!(fallback_payload["records"][0]["id"], "usage-real-fallback");
let retry_response = admin_request(reqwest::Client::new().get(format!(
"{gateway_url}/api/admin/usage/records?start_date=2024-03-21&end_date=2024-03-22&tz_offset_minutes=0&status=has_retry&limit=10&offset=0"
)))
.send()
.await
.expect("retry request should succeed");
assert_eq!(retry_response.status(), StatusCode::OK);
let retry_payload: serde_json::Value = retry_response
.json()
.await
.expect("retry json body should parse");
assert_eq!(retry_payload["total"], 1);
assert_eq!(retry_payload["records"][0]["id"], "usage-real-retry");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_usage_records_with_snapshot_first_user_and_api_key_names() {
let (_upstream_url, upstream_hits, upstream_handle) =

87
apps/aether-gateway/src/tests/usage/pricing.rs
View File

@@ -21,6 +21,27 @@ const OUTPUT_PRICE_PER_1M: f64 = 15.0;
const CACHE_CREATION_PRICE_PER_1M: f64 = 3.75;
const CACHE_READ_PRICE_PER_1M: f64 = 0.30;
fn run_async_test_on_large_stack<F>(name: &'static str, future: F)
where
F: std::future::Future<Output = ()> + Send + 'static,
{
let handle = std::thread::Builder::new()
.name(name.to_string())
.stack_size(16 * 1024 * 1024)
.spawn(move || {
tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("tokio runtime should build")
.block_on(future);
})
.expect("large-stack usage pricing test thread should spawn");
if let Err(payload) = handle.join() {
std::panic::resume_unwind(payload);
}
}
#[derive(Clone, Copy)]
struct ProviderSpec {
provider_id: &'static str,
@@ -752,8 +773,15 @@ async fn assert_candidate_success(
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
}
#[tokio::test]
async fn gateway_records_openai_sync_usage_and_pricing_with_cache_tokens() {
#[test]
fn gateway_records_openai_sync_usage_and_pricing_with_cache_tokens() {
run_async_test_on_large_stack(
"gateway_records_openai_sync_usage_and_pricing_with_cache_tokens",
gateway_records_openai_sync_usage_and_pricing_with_cache_tokens_impl(),
);
}
async fn gateway_records_openai_sync_usage_and_pricing_with_cache_tokens_impl() {
let expected = ExpectedUsagePricing {
input_tokens: 120,
billed_input_tokens: 100,
@@ -839,8 +867,15 @@ async fn gateway_records_openai_sync_usage_and_pricing_with_cache_tokens() {
gateway.shutdown();
}
#[tokio::test]
async fn gateway_records_openai_stream_usage_and_pricing_with_cache_tokens() {
#[test]
fn gateway_records_openai_stream_usage_and_pricing_with_cache_tokens() {
run_async_test_on_large_stack(
"gateway_records_openai_stream_usage_and_pricing_with_cache_tokens",
gateway_records_openai_stream_usage_and_pricing_with_cache_tokens_impl(),
);
}
async fn gateway_records_openai_stream_usage_and_pricing_with_cache_tokens_impl() {
let expected = ExpectedUsagePricing {
input_tokens: 240,
billed_input_tokens: 200,
@@ -924,8 +959,15 @@ async fn gateway_records_openai_stream_usage_and_pricing_with_cache_tokens() {
gateway.shutdown();
}
#[tokio::test]
async fn gateway_records_claude_sync_usage_and_pricing_with_cache_breakdown() {
#[test]
fn gateway_records_claude_sync_usage_and_pricing_with_cache_breakdown() {
run_async_test_on_large_stack(
"gateway_records_claude_sync_usage_and_pricing_with_cache_breakdown",
gateway_records_claude_sync_usage_and_pricing_with_cache_breakdown_impl(),
);
}
async fn gateway_records_claude_sync_usage_and_pricing_with_cache_breakdown_impl() {
let expected = ExpectedUsagePricing {
input_tokens: 50,
billed_input_tokens: 50,
@@ -1016,8 +1058,15 @@ async fn gateway_records_claude_sync_usage_and_pricing_with_cache_breakdown() {
gateway.shutdown();
}
#[tokio::test]
async fn gateway_records_claude_stream_usage_and_pricing_with_cache_breakdown() {
#[test]
fn gateway_records_claude_stream_usage_and_pricing_with_cache_breakdown() {
run_async_test_on_large_stack(
"gateway_records_claude_stream_usage_and_pricing_with_cache_breakdown",
gateway_records_claude_stream_usage_and_pricing_with_cache_breakdown_impl(),
);
}
async fn gateway_records_claude_stream_usage_and_pricing_with_cache_breakdown_impl() {
let expected = ExpectedUsagePricing {
input_tokens: 90,
billed_input_tokens: 90,
@@ -1103,8 +1152,15 @@ async fn gateway_records_claude_stream_usage_and_pricing_with_cache_breakdown()
gateway.shutdown();
}
#[tokio::test]
async fn gateway_records_gemini_sync_usage_and_pricing_with_cache_read_tokens() {
#[test]
fn gateway_records_gemini_sync_usage_and_pricing_with_cache_read_tokens() {
run_async_test_on_large_stack(
"gateway_records_gemini_sync_usage_and_pricing_with_cache_read_tokens",
gateway_records_gemini_sync_usage_and_pricing_with_cache_read_tokens_impl(),
);
}
async fn gateway_records_gemini_sync_usage_and_pricing_with_cache_read_tokens_impl() {
let expected = ExpectedUsagePricing {
input_tokens: 70,
billed_input_tokens: 60,
@@ -1186,8 +1242,15 @@ async fn gateway_records_gemini_sync_usage_and_pricing_with_cache_read_tokens()
gateway.shutdown();
}
#[tokio::test]
async fn gateway_records_gemini_stream_usage_and_pricing_with_cache_read_tokens() {
#[test]
fn gateway_records_gemini_stream_usage_and_pricing_with_cache_read_tokens() {
run_async_test_on_large_stack(
"gateway_records_gemini_stream_usage_and_pricing_with_cache_read_tokens",
gateway_records_gemini_stream_usage_and_pricing_with_cache_read_tokens_impl(),
);
}
async fn gateway_records_gemini_stream_usage_and_pricing_with_cache_read_tokens_impl() {
let expected = ExpectedUsagePricing {
input_tokens: 110,
billed_input_tokens: 80,