fix scheduler affinity candidate selection

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

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 {

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)]

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)]

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,

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());

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());

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)]

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());

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) =

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,