mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-05 00:47:48 +08:00
Merge remote-tracking branch 'origin/pr/613'
# Conflicts: # apps/aether-gateway/src/tests/usage/direct.rs
This commit is contained in:
elky
@@ -154,74 +154,79 @@ async fn gateway_records_usage_for_execution_runtime_sync_when_runtime_enabled()
|
||||
|
||||
#[test]
|
||||
fn gateway_records_pending_usage_before_execution_runtime_sync_result_arrives() {
|
||||
run_async_test_on_large_stack("pending-usage-sync-before-runtime-result", async move {
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let execution_request_started = Arc::new(tokio::sync::Notify::new());
|
||||
let allow_execution_response = Arc::new(tokio::sync::Notify::new());
|
||||
run_async_test_on_large_stack(
|
||||
"gateway_records_pending_usage_before_execution_runtime_sync_result_arrives",
|
||||
gateway_records_pending_usage_before_execution_runtime_sync_result_arrives_impl(),
|
||||
);
|
||||
}
|
||||
|
||||
let upstream = Router::new().route(
|
||||
"/api/internal/gateway/report-sync",
|
||||
any(|_request: Request| async move { Json(json!({"ok": true})) }),
|
||||
);
|
||||
async fn gateway_records_pending_usage_before_execution_runtime_sync_result_arrives_impl() {
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let execution_request_started = Arc::new(tokio::sync::Notify::new());
|
||||
let allow_execution_response = Arc::new(tokio::sync::Notify::new());
|
||||
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any({
|
||||
let upstream = Router::new().route(
|
||||
"/api/internal/gateway/report-sync",
|
||||
any(|_request: Request| async move { Json(json!({"ok": true})) }),
|
||||
);
|
||||
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any({
|
||||
let execution_request_started = Arc::clone(&execution_request_started);
|
||||
let allow_execution_response = Arc::clone(&allow_execution_response);
|
||||
move |_request: Request| {
|
||||
let execution_request_started = Arc::clone(&execution_request_started);
|
||||
let allow_execution_response = Arc::clone(&allow_execution_response);
|
||||
move |_request: Request| {
|
||||
let execution_request_started = Arc::clone(&execution_request_started);
|
||||
let allow_execution_response = Arc::clone(&allow_execution_response);
|
||||
async move {
|
||||
execution_request_started.notify_one();
|
||||
allow_execution_response.notified().await;
|
||||
Json(json!({
|
||||
"request_id": "req-usage-sync-pending-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"id": "chatcmpl-usage-sync-pending-123",
|
||||
"usage": {
|
||||
"input_tokens": 3,
|
||||
"output_tokens": 5,
|
||||
"total_tokens": 8
|
||||
}
|
||||
async move {
|
||||
execution_request_started.notify_one();
|
||||
allow_execution_response.notified().await;
|
||||
Json(json!({
|
||||
"request_id": "req-usage-sync-pending-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"id": "chatcmpl-usage-sync-pending-123",
|
||||
"usage": {
|
||||
"input_tokens": 3,
|
||||
"output_tokens": 5,
|
||||
"total_tokens": 8
|
||||
}
|
||||
},
|
||||
"telemetry": {
|
||||
"elapsed_ms": 45
|
||||
}
|
||||
}))
|
||||
}
|
||||
},
|
||||
"telemetry": {
|
||||
"elapsed_ms": 45
|
||||
}
|
||||
}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-client-openai-usage-sync-pending")),
|
||||
sample_local_openai_auth_snapshot(
|
||||
"api-key-usage-sync-pending-123",
|
||||
"user-usage-sync-pending-123",
|
||||
),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_local_openai_candidate_row(),
|
||||
]));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_local_openai_provider()],
|
||||
vec![sample_local_openai_endpoint()],
|
||||
vec![sample_local_openai_key()],
|
||||
));
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-client-openai-usage-sync-pending")),
|
||||
sample_local_openai_auth_snapshot(
|
||||
"api-key-usage-sync-pending-123",
|
||||
"user-usage-sync-pending-123",
|
||||
),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_local_openai_candidate_row(),
|
||||
]));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_local_openai_provider()],
|
||||
vec![sample_local_openai_endpoint()],
|
||||
vec![sample_local_openai_key()],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) =
|
||||
start_server(execution_runtime).await;
|
||||
let gateway_state =
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state =
|
||||
build_state_with_execution_runtime_override(execution_runtime_url)
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
|
||||
@@ -237,76 +242,74 @@ fn gateway_records_pending_usage_before_execution_runtime_sync_result_arrives()
|
||||
enabled: true,
|
||||
..UsageRuntimeConfig::default()
|
||||
});
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let request_task = tokio::spawn({
|
||||
let gateway_url = gateway_url.clone();
|
||||
async move {
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-client-openai-usage-sync-pending",
|
||||
)
|
||||
.header(TRACE_ID_HEADER, "req-usage-sync-pending-123")
|
||||
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
let status = response.status();
|
||||
let body = response.text().await.expect("body should read");
|
||||
(status, body)
|
||||
}
|
||||
});
|
||||
|
||||
execution_request_started.notified().await;
|
||||
|
||||
let mut pending = None;
|
||||
for _ in 0..50 {
|
||||
pending = usage_repository
|
||||
.find_by_request_id("req-usage-sync-pending-123")
|
||||
let request_task = tokio::spawn({
|
||||
let gateway_url = gateway_url.clone();
|
||||
async move {
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-client-openai-usage-sync-pending",
|
||||
)
|
||||
.header(TRACE_ID_HEADER, "req-usage-sync-pending-123")
|
||||
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("usage lookup should succeed");
|
||||
if pending
|
||||
.as_ref()
|
||||
.is_some_and(|stored| stored.status == "pending")
|
||||
{
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
.expect("request should succeed");
|
||||
let status = response.status();
|
||||
let body = response.text().await.expect("body should read");
|
||||
(status, body)
|
||||
}
|
||||
let pending =
|
||||
pending.expect("pending usage should be recorded before sync result resolves");
|
||||
assert_eq!(pending.status, "pending");
|
||||
assert_eq!(pending.billing_status, "pending");
|
||||
assert_eq!(pending.response_time_ms, None);
|
||||
|
||||
allow_execution_response.notify_one();
|
||||
|
||||
let (status, _body) = request_task.await.expect("request task should join");
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
|
||||
let mut stored = None;
|
||||
for _ in 0..50 {
|
||||
stored = usage_repository
|
||||
.find_by_request_id("req-usage-sync-pending-123")
|
||||
.await
|
||||
.expect("usage lookup should succeed");
|
||||
if stored.as_ref().is_some_and(|row| row.status == "completed") {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
let stored = stored.expect("usage should be finalized");
|
||||
assert_eq!(stored.status, "completed");
|
||||
assert_eq!(stored.response_time_ms, Some(45));
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
});
|
||||
|
||||
execution_request_started.notified().await;
|
||||
|
||||
let mut pending = None;
|
||||
for _ in 0..50 {
|
||||
pending = usage_repository
|
||||
.find_by_request_id("req-usage-sync-pending-123")
|
||||
.await
|
||||
.expect("usage lookup should succeed");
|
||||
if pending
|
||||
.as_ref()
|
||||
.is_some_and(|stored| stored.status == "pending")
|
||||
{
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
let pending = pending.expect("pending usage should be recorded before sync result resolves");
|
||||
assert_eq!(pending.status, "pending");
|
||||
assert_eq!(pending.billing_status, "pending");
|
||||
assert_eq!(pending.response_time_ms, None);
|
||||
|
||||
allow_execution_response.notify_one();
|
||||
|
||||
let (status, _body) = request_task.await.expect("request task should join");
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
|
||||
let mut stored = None;
|
||||
for _ in 0..50 {
|
||||
stored = usage_repository
|
||||
.find_by_request_id("req-usage-sync-pending-123")
|
||||
.await
|
||||
.expect("usage lookup should succeed");
|
||||
if stored.as_ref().is_some_and(|row| row.status == "completed") {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
let stored = stored.expect("usage should be finalized");
|
||||
assert_eq!(stored.status, "completed");
|
||||
assert_eq!(stored.response_time_ms, Some(45));
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -559,8 +559,15 @@ async fn gateway_applies_system_max_request_body_size_to_local_openai_chat_sync_
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_strips_request_and_response_bodies_when_request_record_level_is_base() {
|
||||
#[test]
|
||||
fn gateway_strips_request_and_response_bodies_when_request_record_level_is_base() {
|
||||
run_async_test_on_large_stack(
|
||||
"gateway_strips_request_and_response_bodies_when_request_record_level_is_base",
|
||||
gateway_strips_request_and_response_bodies_when_request_record_level_is_base_impl(),
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_strips_request_and_response_bodies_when_request_record_level_is_base_impl() {
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
|
||||
|
||||
@@ -11,8 +11,36 @@ use super::{
|
||||
};
|
||||
use aether_data::repository::settlement::InMemorySettlementRepository;
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_settles_wallet_for_completed_execution_runtime_sync_usage() {
|
||||
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(8 * 1024 * 1024)
|
||||
.spawn(move || {
|
||||
tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.expect("tokio runtime should build")
|
||||
.block_on(future);
|
||||
})
|
||||
.expect("large-stack test thread should spawn");
|
||||
|
||||
if let Err(payload) = handle.join() {
|
||||
std::panic::resume_unwind(payload);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gateway_settles_wallet_for_completed_execution_runtime_sync_usage() {
|
||||
run_async_test_on_large_stack(
|
||||
"gateway_settles_wallet_for_completed_execution_runtime_sync_usage",
|
||||
gateway_settles_wallet_for_completed_execution_runtime_sync_usage_impl(),
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_settles_wallet_for_completed_execution_runtime_sync_usage_impl() {
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let billing_repository = Arc::new(InMemoryBillingReadRepository::seed(vec![
|
||||
|
||||
Reference in New Issue
Block a user