mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-08 10:27:46 +08:00
Merge remote-tracking branch 'upstream/main'
This commit is contained in:
@@ -24,8 +24,39 @@ use aether_data_contracts::repository::provider_catalog::{
|
||||
};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_executes_openai_chat_sync_upstream_stream_via_local_finalize_response() {
|
||||
const OPENAI_CHAT_FINALIZE_TEST_STACK_BYTES: usize = 16 * 1024 * 1024;
|
||||
|
||||
fn run_openai_chat_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(OPENAI_CHAT_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("openai chat finalize test thread should spawn");
|
||||
|
||||
if let Err(payload) = handle.join() {
|
||||
std::panic::resume_unwind(payload);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gateway_executes_openai_chat_sync_upstream_stream_via_local_finalize_response() {
|
||||
run_openai_chat_finalize_test(
|
||||
"gateway_executes_openai_chat_sync_upstream_stream_via_local_finalize_response",
|
||||
gateway_executes_openai_chat_sync_upstream_stream_via_local_finalize_response_impl,
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_executes_openai_chat_sync_upstream_stream_via_local_finalize_response_impl() {
|
||||
use base64::Engine as _;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -521,8 +552,16 @@ async fn gateway_executes_openai_chat_sync_upstream_stream_via_local_finalize_re
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_executes_openai_chat_cross_format_upstream_stream_via_local_finalize_response() {
|
||||
#[test]
|
||||
fn gateway_executes_openai_chat_cross_format_upstream_stream_via_local_finalize_response() {
|
||||
run_openai_chat_finalize_test(
|
||||
"gateway_executes_openai_chat_cross_format_upstream_stream_via_local_finalize_response",
|
||||
gateway_executes_openai_chat_cross_format_upstream_stream_via_local_finalize_response_impl,
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_executes_openai_chat_cross_format_upstream_stream_via_local_finalize_response_impl(
|
||||
) {
|
||||
use base64::Engine as _;
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenRemoteExecutionRuntimeRequest {
|
||||
@@ -955,8 +994,16 @@ async fn gateway_executes_openai_chat_cross_format_upstream_stream_via_local_fin
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_executes_openai_chat_cross_format_tool_use_upstream_stream_via_local_finalize_response(
|
||||
#[test]
|
||||
fn gateway_executes_openai_chat_cross_format_tool_use_upstream_stream_via_local_finalize_response()
|
||||
{
|
||||
run_openai_chat_finalize_test(
|
||||
"gateway_executes_openai_chat_cross_format_tool_use_upstream_stream_via_local_finalize_response",
|
||||
gateway_executes_openai_chat_cross_format_tool_use_upstream_stream_via_local_finalize_response_impl,
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_executes_openai_chat_cross_format_tool_use_upstream_stream_via_local_finalize_response_impl(
|
||||
) {
|
||||
use base64::Engine as _;
|
||||
|
||||
@@ -1782,8 +1829,15 @@ async fn gateway_skips_openai_chat_antigravity_cross_format_sync_candidate_as_tr
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_executes_openai_chat_cross_format_claude_upstream_sync_via_local_finalize_response(
|
||||
#[test]
|
||||
fn gateway_executes_openai_chat_cross_format_claude_upstream_sync_via_local_finalize_response() {
|
||||
run_openai_chat_finalize_test(
|
||||
"gateway_executes_openai_chat_cross_format_claude_upstream_sync_via_local_finalize_response",
|
||||
gateway_executes_openai_chat_cross_format_claude_upstream_sync_via_local_finalize_response_impl,
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_executes_openai_chat_cross_format_claude_upstream_sync_via_local_finalize_response_impl(
|
||||
) {
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
@@ -2134,8 +2188,15 @@ async fn gateway_executes_openai_chat_cross_format_claude_upstream_sync_via_loca
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_executes_openai_chat_cross_format_gemini_upstream_sync_via_local_finalize_response(
|
||||
#[test]
|
||||
fn gateway_executes_openai_chat_cross_format_gemini_upstream_sync_via_local_finalize_response() {
|
||||
run_openai_chat_finalize_test(
|
||||
"gateway_executes_openai_chat_cross_format_gemini_upstream_sync_via_local_finalize_response",
|
||||
gateway_executes_openai_chat_cross_format_gemini_upstream_sync_via_local_finalize_response_impl,
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_executes_openai_chat_cross_format_gemini_upstream_sync_via_local_finalize_response_impl(
|
||||
) {
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
|
||||
@@ -8,6 +8,9 @@ use super::{
|
||||
StoredProviderCatalogKey, StoredProviderCatalogProvider, StoredProviderModelMapping,
|
||||
DEVELOPMENT_ENCRYPTION_KEY, TRACE_ID_HEADER,
|
||||
};
|
||||
use aether_data::repository::usage::InMemoryUsageReadRepository;
|
||||
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UsageReadRepository};
|
||||
use aether_usage_runtime::UsageRuntimeConfig;
|
||||
|
||||
const KIRO_CLAUDE_CLI_SYNC_TEST_STACK_BYTES: usize = 16 * 1024 * 1024;
|
||||
|
||||
@@ -33,6 +36,30 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
async fn wait_for_completed_usage<T>(repository: &T, request_id: &str) -> StoredRequestUsageAudit
|
||||
where
|
||||
T: UsageReadRepository + ?Sized,
|
||||
{
|
||||
let timeout = std::time::Duration::from_secs(60);
|
||||
let deadline = tokio::time::Instant::now() + timeout;
|
||||
loop {
|
||||
if let Some(usage) = repository
|
||||
.find_by_request_id(request_id)
|
||||
.await
|
||||
.expect("usage should read")
|
||||
{
|
||||
if usage.status == "completed" {
|
||||
return usage;
|
||||
}
|
||||
}
|
||||
assert!(
|
||||
tokio::time::Instant::now() < deadline,
|
||||
"usage {request_id} should complete within {timeout:?}"
|
||||
);
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gateway_executes_kiro_claude_cli_sync_via_local_provider_catalog_candidate() {
|
||||
run_kiro_claude_cli_sync_test(
|
||||
@@ -194,7 +221,7 @@ async fn gateway_executes_kiro_claude_cli_sync_via_local_provider_catalog_candid
|
||||
Some(serde_json::json!({"url":"http://provider-proxy.internal:8080"})),
|
||||
Some(20.0),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"kiro": {"simulated_cache_enabled": true}})),
|
||||
)
|
||||
}
|
||||
|
||||
@@ -354,14 +381,15 @@ async fn gateway_executes_kiro_claude_cli_sync_via_local_provider_catalog_candid
|
||||
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
let payload: serde_json::Value =
|
||||
serde_json::from_slice(&raw_body).expect("execution runtime payload should parse");
|
||||
let trace_id = parts
|
||||
.headers
|
||||
.get(TRACE_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
*seen_execution_runtime_inner.lock().expect("mutex should lock") =
|
||||
Some(SeenExecutionRuntimeSyncRequest {
|
||||
trace_id: parts
|
||||
.headers
|
||||
.get(TRACE_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
trace_id: trace_id.clone(),
|
||||
url: payload
|
||||
.get("url")
|
||||
.and_then(|value| value.as_str())
|
||||
@@ -453,7 +481,7 @@ async fn gateway_executes_kiro_claude_cli_sync_via_local_provider_catalog_candid
|
||||
.concat();
|
||||
|
||||
Json(json!({
|
||||
"request_id": "trace-kiro-cli-local-sync-123",
|
||||
"request_id": trace_id,
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/vnd.amazon.eventstream"
|
||||
@@ -481,6 +509,7 @@ async fn gateway_executes_kiro_claude_cli_sync_via_local_provider_catalog_candid
|
||||
sample_candidate_row(),
|
||||
]));
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider_catalog_provider()],
|
||||
vec![sample_provider_catalog_endpoint()],
|
||||
@@ -491,34 +520,51 @@ async fn gateway_executes_kiro_claude_cli_sync_via_local_provider_catalog_candid
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
Arc::clone(&usage_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
);
|
||||
)
|
||||
.with_usage_runtime_for_tests(UsageRuntimeConfig {
|
||||
enabled: true,
|
||||
..UsageRuntimeConfig::default()
|
||||
});
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/messages"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-client-kiro-cli-local-sync",
|
||||
)
|
||||
.header(TRACE_ID_HEADER, "trace-kiro-cli-local-sync-123")
|
||||
.body(
|
||||
"{\"model\":\"claude-sonnet-4\",\"messages\":[{\"role\":\"user\",\"content\":\"hello kiro\"}],\"thinking\":{\"type\":\"enabled\",\"budget_tokens\":64}}",
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
async fn send_kiro_request(
|
||||
gateway_url: &str,
|
||||
trace_id: &str,
|
||||
body: String,
|
||||
) -> (StatusCode, String) {
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/messages"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-client-kiro-cli-local-sync",
|
||||
)
|
||||
.header(TRACE_ID_HEADER, trace_id)
|
||||
.body(body)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
let status = response.status();
|
||||
let response_body = response.text().await.expect("body should read");
|
||||
let status = response.status();
|
||||
let response_body = response.text().await.expect("body should read");
|
||||
(status, response_body)
|
||||
}
|
||||
|
||||
let (status, response_body) = send_kiro_request(
|
||||
&gateway_url,
|
||||
"trace-kiro-cli-local-sync-123",
|
||||
"{\"model\":\"claude-sonnet-4\",\"messages\":[{\"role\":\"user\",\"content\":\"hello kiro\"}],\"thinking\":{\"type\":\"enabled\",\"budget_tokens\":64}}".to_string(),
|
||||
)
|
||||
.await;
|
||||
assert!(
|
||||
status == StatusCode::OK,
|
||||
"unexpected status={status} body={response_body} decision_hits={} plan_hits={} public_hits={}",
|
||||
@@ -598,6 +644,58 @@ async fn gateway_executes_kiro_claude_cli_sync_via_local_provider_catalog_candid
|
||||
"report-sync should stay local when request candidate persistence is available"
|
||||
);
|
||||
|
||||
let cacheable_request_body = serde_json::json!({
|
||||
"model": "claude-sonnet-4",
|
||||
"system": [{
|
||||
"type": "text",
|
||||
"text": format!("sync cacheable prompt {}", "cacheable prompt chunk ".repeat(300)),
|
||||
"cache_control": {"type": "ephemeral"}
|
||||
}],
|
||||
"messages": [{"role": "user", "content": "reuse this Kiro prompt"}]
|
||||
})
|
||||
.to_string();
|
||||
let (first_cache_status, first_cache_body) = send_kiro_request(
|
||||
&gateway_url,
|
||||
"trace-kiro-cli-local-sync-cache-1",
|
||||
cacheable_request_body.clone(),
|
||||
)
|
||||
.await;
|
||||
assert!(
|
||||
first_cache_status == StatusCode::OK,
|
||||
"unexpected first cache status={first_cache_status} body={first_cache_body}"
|
||||
);
|
||||
let first_usage = wait_for_completed_usage(
|
||||
usage_repository.as_ref(),
|
||||
"trace-kiro-cli-local-sync-cache-1",
|
||||
)
|
||||
.await;
|
||||
assert!(
|
||||
first_usage.cache_creation_input_tokens > 0,
|
||||
"first Kiro sync cacheable request should create simulated cache"
|
||||
);
|
||||
assert_eq!(first_usage.cache_read_input_tokens, 0);
|
||||
|
||||
let (second_cache_status, second_cache_body) = send_kiro_request(
|
||||
&gateway_url,
|
||||
"trace-kiro-cli-local-sync-cache-2",
|
||||
cacheable_request_body,
|
||||
)
|
||||
.await;
|
||||
assert!(
|
||||
second_cache_status == StatusCode::OK,
|
||||
"unexpected second cache status={second_cache_status} body={second_cache_body}"
|
||||
);
|
||||
let second_usage = wait_for_completed_usage(
|
||||
usage_repository.as_ref(),
|
||||
"trace-kiro-cli-local-sync-cache-2",
|
||||
)
|
||||
.await;
|
||||
assert!(
|
||||
second_usage.cache_read_input_tokens > 0,
|
||||
"second Kiro sync cacheable request should read simulated cache"
|
||||
);
|
||||
assert_eq!(second_usage.cache_creation_input_tokens, 0);
|
||||
|
||||
assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
@@ -26,8 +26,40 @@ use aether_data_contracts::repository::provider_catalog::{
|
||||
use base64::Engine as _;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_executes_openai_responses_sync_via_local_decision_gate_with_local_sync_decision() {
|
||||
const CLI_SYNC_TEST_STACK_BYTES: usize = 16 * 1024 * 1024;
|
||||
|
||||
fn run_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(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("cli sync test thread should spawn");
|
||||
|
||||
if let Err(payload) = handle.join() {
|
||||
std::panic::resume_unwind(payload);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gateway_executes_openai_responses_sync_via_local_decision_gate_with_local_sync_decision() {
|
||||
run_cli_sync_test(
|
||||
"gateway_executes_openai_responses_sync_via_local_decision_gate_with_local_sync_decision",
|
||||
gateway_executes_openai_responses_sync_via_local_decision_gate_with_local_sync_decision_impl,
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_executes_openai_responses_sync_via_local_decision_gate_with_local_sync_decision_impl(
|
||||
) {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeSyncRequest {
|
||||
trace_id: String,
|
||||
@@ -546,8 +578,15 @@ async fn gateway_executes_openai_responses_sync_via_local_decision_gate_with_loc
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_waits_for_api_key_concurrency_slot_then_executes_openai_responses_sync() {
|
||||
#[test]
|
||||
fn gateway_waits_for_api_key_concurrency_slot_then_executes_openai_responses_sync() {
|
||||
run_cli_sync_test(
|
||||
"gateway_waits_for_api_key_concurrency_slot_then_executes_openai_responses_sync",
|
||||
gateway_waits_for_api_key_concurrency_slot_then_executes_openai_responses_sync_impl,
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_waits_for_api_key_concurrency_slot_then_executes_openai_responses_sync_impl() {
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
@@ -929,8 +968,16 @@ async fn gateway_waits_for_api_key_concurrency_slot_then_executes_openai_respons
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_returns_concurrency_limited_after_wait_budget_expires_for_openai_responses_sync() {
|
||||
#[test]
|
||||
fn gateway_returns_concurrency_limited_after_wait_budget_expires_for_openai_responses_sync() {
|
||||
run_cli_sync_test(
|
||||
"gateway_returns_concurrency_limited_after_wait_budget_expires_for_openai_responses_sync",
|
||||
gateway_returns_concurrency_limited_after_wait_budget_expires_for_openai_responses_sync_impl,
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_returns_concurrency_limited_after_wait_budget_expires_for_openai_responses_sync_impl(
|
||||
) {
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
@@ -1278,8 +1325,15 @@ async fn gateway_returns_concurrency_limited_after_wait_budget_expires_for_opena
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_returns_openai_responses_error_for_local_sync_failure() {
|
||||
#[test]
|
||||
fn gateway_returns_openai_responses_error_for_local_sync_failure() {
|
||||
run_cli_sync_test(
|
||||
"gateway_returns_openai_responses_error_for_local_sync_failure",
|
||||
gateway_returns_openai_responses_error_for_local_sync_failure_impl,
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_returns_openai_responses_error_for_local_sync_failure_impl() {
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
@@ -1571,8 +1625,16 @@ async fn gateway_returns_openai_responses_error_for_local_sync_failure() {
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_returns_openai_responses_error_for_local_cross_format_gemini_cli_sync_failure() {
|
||||
#[test]
|
||||
fn gateway_returns_openai_responses_error_for_local_cross_format_gemini_cli_sync_failure() {
|
||||
run_cli_sync_test(
|
||||
"gateway_returns_openai_responses_error_for_local_cross_format_gemini_cli_sync_failure",
|
||||
gateway_returns_openai_responses_error_for_local_cross_format_gemini_cli_sync_failure_impl,
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_returns_openai_responses_error_for_local_cross_format_gemini_cli_sync_failure_impl(
|
||||
) {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeSyncRequest {
|
||||
trace_id: String,
|
||||
@@ -1972,8 +2034,15 @@ async fn gateway_returns_openai_responses_error_for_local_cross_format_gemini_cl
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_returns_openai_responses_error_for_local_cross_format_claude_sync_failure() {
|
||||
#[test]
|
||||
fn gateway_returns_openai_responses_error_for_local_cross_format_claude_sync_failure() {
|
||||
run_cli_sync_test(
|
||||
"gateway_returns_openai_responses_error_for_local_cross_format_claude_sync_failure",
|
||||
gateway_returns_openai_responses_error_for_local_cross_format_claude_sync_failure_impl,
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_returns_openai_responses_error_for_local_cross_format_claude_sync_failure_impl() {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeSyncRequest {
|
||||
trace_id: String,
|
||||
@@ -2349,8 +2418,16 @@ async fn gateway_returns_openai_responses_error_for_local_cross_format_claude_sy
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_returns_openai_responses_error_for_local_cross_format_claude_chat_sync_failure() {
|
||||
#[test]
|
||||
fn gateway_returns_openai_responses_error_for_local_cross_format_claude_chat_sync_failure() {
|
||||
run_cli_sync_test(
|
||||
"gateway_returns_openai_responses_error_for_local_cross_format_claude_chat_sync_failure",
|
||||
gateway_returns_openai_responses_error_for_local_cross_format_claude_chat_sync_failure_impl,
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_returns_openai_responses_error_for_local_cross_format_claude_chat_sync_failure_impl(
|
||||
) {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeSyncRequest {
|
||||
trace_id: String,
|
||||
@@ -2729,8 +2806,16 @@ async fn gateway_returns_openai_responses_error_for_local_cross_format_claude_ch
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_returns_openai_responses_error_for_local_cross_format_gemini_chat_sync_failure() {
|
||||
#[test]
|
||||
fn gateway_returns_openai_responses_error_for_local_cross_format_gemini_chat_sync_failure() {
|
||||
run_cli_sync_test(
|
||||
"gateway_returns_openai_responses_error_for_local_cross_format_gemini_chat_sync_failure",
|
||||
gateway_returns_openai_responses_error_for_local_cross_format_gemini_chat_sync_failure_impl,
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_returns_openai_responses_error_for_local_cross_format_gemini_chat_sync_failure_impl(
|
||||
) {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeSyncRequest {
|
||||
trace_id: String,
|
||||
@@ -3109,8 +3194,15 @@ async fn gateway_returns_openai_responses_error_for_local_cross_format_gemini_ch
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_executes_codex_cli_sync_via_local_decision_gate_after_oauth_refresh() {
|
||||
#[test]
|
||||
fn gateway_executes_codex_cli_sync_via_local_decision_gate_after_oauth_refresh() {
|
||||
run_cli_sync_test(
|
||||
"gateway_executes_codex_cli_sync_via_local_decision_gate_after_oauth_refresh",
|
||||
gateway_executes_codex_cli_sync_via_local_decision_gate_after_oauth_refresh_impl,
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_executes_codex_cli_sync_via_local_decision_gate_after_oauth_refresh_impl() {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeSyncRequest {
|
||||
trace_id: String,
|
||||
|
||||
@@ -226,19 +226,24 @@ impl StandardFormat {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ai_execute_openai_responses_pii_redaction_round_trip_same_format() {
|
||||
let (response_json, seen) = run_redaction_format_case(RedactionFormatCase {
|
||||
test_id: "openai-responses-pii-redaction-same-format",
|
||||
trace_id: "trace-openai-responses-pii-redaction-same-format",
|
||||
client_format: StandardFormat::OpenAiResponses,
|
||||
provider_format: StandardFormat::OpenAiResponses,
|
||||
})
|
||||
.await;
|
||||
#[test]
|
||||
fn ai_execute_openai_responses_pii_redaction_round_trip_same_format() {
|
||||
run_async_test_on_large_stack(
|
||||
"ai_execute_openai_responses_pii_redaction_round_trip_same_format",
|
||||
async {
|
||||
let (response_json, seen) = run_redaction_format_case(RedactionFormatCase {
|
||||
test_id: "openai-responses-pii-redaction-same-format",
|
||||
trace_id: "trace-openai-responses-pii-redaction-same-format",
|
||||
client_format: StandardFormat::OpenAiResponses,
|
||||
provider_format: StandardFormat::OpenAiResponses,
|
||||
})
|
||||
.await;
|
||||
|
||||
assert_provider_request_redacted(&seen, StandardFormat::OpenAiResponses);
|
||||
assert!(seen.body.get("input").is_some());
|
||||
assert_restored_response(&response_json, StandardFormat::OpenAiResponses);
|
||||
assert_provider_request_redacted(&seen, StandardFormat::OpenAiResponses);
|
||||
assert!(seen.body.get("input").is_some());
|
||||
assert_restored_response(&response_json, StandardFormat::OpenAiResponses);
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -261,42 +266,52 @@ fn ai_execute_claude_messages_pii_redaction_round_trip_same_format() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ai_execute_openai_chat_pii_redaction_before_claude_conversion() {
|
||||
let (response_json, seen) = run_redaction_format_case(RedactionFormatCase {
|
||||
test_id: "openai-chat-pii-redaction-before-claude-conversion",
|
||||
trace_id: "trace-openai-chat-pii-redaction-before-claude-conversion",
|
||||
client_format: StandardFormat::OpenAiChat,
|
||||
provider_format: StandardFormat::ClaudeMessages,
|
||||
})
|
||||
.await;
|
||||
#[test]
|
||||
fn ai_execute_openai_chat_pii_redaction_before_claude_conversion() {
|
||||
run_async_test_on_large_stack(
|
||||
"ai_execute_openai_chat_pii_redaction_before_claude_conversion",
|
||||
async {
|
||||
let (response_json, seen) = run_redaction_format_case(RedactionFormatCase {
|
||||
test_id: "openai-chat-pii-redaction-before-claude-conversion",
|
||||
trace_id: "trace-openai-chat-pii-redaction-before-claude-conversion",
|
||||
client_format: StandardFormat::OpenAiChat,
|
||||
provider_format: StandardFormat::ClaudeMessages,
|
||||
})
|
||||
.await;
|
||||
|
||||
assert_provider_request_redacted(&seen, StandardFormat::ClaudeMessages);
|
||||
assert!(seen.body.get("messages").is_some());
|
||||
assert_eq!(
|
||||
seen.body["model"],
|
||||
StandardFormat::ClaudeMessages.provider_model()
|
||||
assert_provider_request_redacted(&seen, StandardFormat::ClaudeMessages);
|
||||
assert!(seen.body.get("messages").is_some());
|
||||
assert_eq!(
|
||||
seen.body["model"],
|
||||
StandardFormat::ClaudeMessages.provider_model()
|
||||
);
|
||||
assert_restored_response(&response_json, StandardFormat::OpenAiChat);
|
||||
},
|
||||
);
|
||||
assert_restored_response(&response_json, StandardFormat::OpenAiChat);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ai_execute_openai_responses_pii_redaction_before_claude_conversion() {
|
||||
let (response_json, seen) = run_redaction_format_case(RedactionFormatCase {
|
||||
test_id: "openai-responses-pii-redaction-before-claude-conversion",
|
||||
trace_id: "trace-openai-responses-pii-redaction-before-claude-conversion",
|
||||
client_format: StandardFormat::OpenAiResponses,
|
||||
provider_format: StandardFormat::ClaudeMessages,
|
||||
})
|
||||
.await;
|
||||
#[test]
|
||||
fn ai_execute_openai_responses_pii_redaction_before_claude_conversion() {
|
||||
run_async_test_on_large_stack(
|
||||
"ai_execute_openai_responses_pii_redaction_before_claude_conversion",
|
||||
async {
|
||||
let (response_json, seen) = run_redaction_format_case(RedactionFormatCase {
|
||||
test_id: "openai-responses-pii-redaction-before-claude-conversion",
|
||||
trace_id: "trace-openai-responses-pii-redaction-before-claude-conversion",
|
||||
client_format: StandardFormat::OpenAiResponses,
|
||||
provider_format: StandardFormat::ClaudeMessages,
|
||||
})
|
||||
.await;
|
||||
|
||||
assert_provider_request_redacted(&seen, StandardFormat::ClaudeMessages);
|
||||
assert!(seen.body.get("messages").is_some());
|
||||
assert_eq!(
|
||||
seen.body["model"],
|
||||
StandardFormat::ClaudeMessages.provider_model()
|
||||
assert_provider_request_redacted(&seen, StandardFormat::ClaudeMessages);
|
||||
assert!(seen.body.get("messages").is_some());
|
||||
assert_eq!(
|
||||
seen.body["model"],
|
||||
StandardFormat::ClaudeMessages.provider_model()
|
||||
);
|
||||
assert_restored_response(&response_json, StandardFormat::OpenAiResponses);
|
||||
},
|
||||
);
|
||||
assert_restored_response(&response_json, StandardFormat::OpenAiResponses);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -28,7 +28,7 @@ fn admin_system_build_version_contract_uses_explicit_local_build_arg() {
|
||||
let deploy = read_workspace_file("deploy.sh");
|
||||
for pattern in [
|
||||
"detect_build_version()",
|
||||
"git describe --tags --always --dirty",
|
||||
"git describe --tags --match 'v[0-9]*' --always --dirty",
|
||||
"AETHER_BUILD_VERSION=\"${AETHER_BUILD_VERSION:-$(detect_build_version)}\"",
|
||||
"--build-arg \"AETHER_BUILD_VERSION=$AETHER_BUILD_VERSION\"",
|
||||
">>> AETHER_BUILD_VERSION",
|
||||
@@ -43,6 +43,8 @@ fn admin_system_build_version_contract_uses_explicit_local_build_arg() {
|
||||
for pattern in [
|
||||
"process.env.AETHER_BUILD_VERSION",
|
||||
"process.env.AETHER_VERSION",
|
||||
"git describe --tags --match \"v[0-9]*\" --always --dirty",
|
||||
"trimmed.startsWith('tunnel-v')",
|
||||
] {
|
||||
assert!(
|
||||
vite_config.contains(pattern),
|
||||
@@ -60,6 +62,18 @@ fn admin_system_build_version_contract_uses_explicit_local_build_arg() {
|
||||
"api/core.rs should expose build version pattern {pattern}"
|
||||
);
|
||||
}
|
||||
|
||||
let build_rs = read_workspace_file("apps/aether-gateway/build.rs");
|
||||
for pattern in [
|
||||
"\"--match\"",
|
||||
"\"v[0-9]*\"",
|
||||
"trimmed.starts_with(\"tunnel-v\")",
|
||||
] {
|
||||
assert!(
|
||||
build_rs.contains(pattern),
|
||||
"apps/aether-gateway/build.rs should ignore tunnel release tags for gateway version pattern {pattern}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -1452,6 +1452,14 @@ async fn gateway_handles_admin_system_api_formats_locally_with_trusted_admin_pri
|
||||
assert!(formats.iter().any(|item| item["value"] == "jina:embedding"));
|
||||
assert!(formats.iter().any(|item| item["value"] == "jina:rerank"));
|
||||
assert!(formats.iter().any(|item| item["value"] == "gemini:video"));
|
||||
let aliyun_embedding = formats
|
||||
.iter()
|
||||
.find(|item| item["value"] == "aliyun:multimodal_embedding")
|
||||
.expect("aliyun multimodal embedding format should exist");
|
||||
assert_eq!(
|
||||
aliyun_embedding["default_path"],
|
||||
"/api/v1/services/embeddings/multimodal-embedding/multimodal-embedding"
|
||||
);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
|
||||
@@ -179,6 +179,129 @@ fn vertex_gemini_embedding_success_state(execution_runtime_url: String) -> AppSt
|
||||
.with_data_state_for_tests(data_state)
|
||||
}
|
||||
|
||||
fn aliyun_embedding_success_state(execution_runtime_url: String) -> AppState {
|
||||
let mut snapshot = sample_currently_usable_auth_snapshot(
|
||||
"key-aliyun-embedding-success",
|
||||
"user-aliyun-embedding-success",
|
||||
);
|
||||
snapshot.user_allowed_providers = None;
|
||||
snapshot.api_key_allowed_providers = None;
|
||||
snapshot.user_allowed_api_formats = Some(vec!["openai:embedding".to_string()]);
|
||||
snapshot.api_key_allowed_api_formats = Some(vec!["openai:embedding".to_string()]);
|
||||
snapshot.user_allowed_models = Some(vec!["qwen3-vl-embedding".to_string()]);
|
||||
snapshot.api_key_allowed_models = Some(vec!["qwen3-vl-embedding".to_string()]);
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-aliyun-embedding-success")),
|
||||
snapshot,
|
||||
)]));
|
||||
let candidate_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
aliyun_embedding_candidate_row(),
|
||||
]));
|
||||
let mut provider = sample_provider("provider-aliyun-embedding", "Aliyun DashScope", 1);
|
||||
provider.provider_type = "aliyun".to_string();
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![provider],
|
||||
vec![sample_endpoint(
|
||||
"endpoint-aliyun-embedding",
|
||||
"provider-aliyun-embedding",
|
||||
"aliyun:multimodal_embedding",
|
||||
"https://dashscope.aliyuncs.com",
|
||||
)],
|
||||
vec![sample_key(
|
||||
"key-upstream-aliyun-embedding",
|
||||
"provider-aliyun-embedding",
|
||||
"aliyun:multimodal_embedding",
|
||||
"sk-upstream-aliyun-embedding",
|
||||
)],
|
||||
));
|
||||
let data_state =
|
||||
GatewayDataState::with_provider_catalog_and_minimal_candidate_selection_for_tests(
|
||||
provider_catalog_repository,
|
||||
candidate_repository,
|
||||
)
|
||||
.with_auth_api_key_reader(auth_repository)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY);
|
||||
|
||||
build_state_with_execution_runtime_override(execution_runtime_url)
|
||||
.with_data_state_for_tests(data_state)
|
||||
}
|
||||
|
||||
fn mixed_embedding_success_state(execution_runtime_url: String) -> AppState {
|
||||
let mut snapshot = sample_currently_usable_auth_snapshot(
|
||||
"key-mixed-embedding-success",
|
||||
"user-mixed-embedding-success",
|
||||
);
|
||||
snapshot.user_allowed_providers = None;
|
||||
snapshot.api_key_allowed_providers = None;
|
||||
snapshot.user_allowed_api_formats = Some(vec!["openai:embedding".to_string()]);
|
||||
snapshot.api_key_allowed_api_formats = Some(vec!["openai:embedding".to_string()]);
|
||||
snapshot.user_allowed_models = Some(vec!["qwen3-vl-embedding".to_string()]);
|
||||
snapshot.api_key_allowed_models = Some(vec!["qwen3-vl-embedding".to_string()]);
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-mixed-embedding-success")),
|
||||
snapshot,
|
||||
)]));
|
||||
|
||||
let mut openai_candidate = embedding_candidate_row();
|
||||
openai_candidate.model_id = "model-openai-qwen-vl-embedding".to_string();
|
||||
openai_candidate.global_model_id = "global-qwen3-vl-embedding".to_string();
|
||||
openai_candidate.global_model_name = "qwen3-vl-embedding".to_string();
|
||||
openai_candidate.model_provider_model_name = "openai-qwen-fallback".to_string();
|
||||
let candidate_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
openai_candidate,
|
||||
aliyun_embedding_candidate_row(),
|
||||
]));
|
||||
|
||||
let mut aliyun_provider = sample_provider("provider-aliyun-embedding", "Aliyun DashScope", 1);
|
||||
aliyun_provider.provider_type = "aliyun".to_string();
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![
|
||||
sample_provider("provider-embedding", "OpenAI Embeddings", 1),
|
||||
aliyun_provider,
|
||||
],
|
||||
vec![
|
||||
sample_endpoint(
|
||||
"endpoint-embedding",
|
||||
"provider-embedding",
|
||||
"openai:embedding",
|
||||
"https://api.openai.example",
|
||||
),
|
||||
sample_endpoint(
|
||||
"endpoint-aliyun-embedding",
|
||||
"provider-aliyun-embedding",
|
||||
"aliyun:multimodal_embedding",
|
||||
"https://dashscope.aliyuncs.com",
|
||||
),
|
||||
],
|
||||
vec![
|
||||
sample_key(
|
||||
"key-upstream-embedding",
|
||||
"provider-embedding",
|
||||
"openai:embedding",
|
||||
"sk-upstream-embedding",
|
||||
),
|
||||
sample_key(
|
||||
"key-upstream-aliyun-embedding",
|
||||
"provider-aliyun-embedding",
|
||||
"aliyun:multimodal_embedding",
|
||||
"sk-upstream-aliyun-embedding",
|
||||
),
|
||||
],
|
||||
));
|
||||
let data_state =
|
||||
GatewayDataState::with_provider_catalog_and_minimal_candidate_selection_for_tests(
|
||||
provider_catalog_repository,
|
||||
candidate_repository,
|
||||
)
|
||||
.with_auth_api_key_reader(auth_repository)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY);
|
||||
|
||||
build_state_with_execution_runtime_override(execution_runtime_url)
|
||||
.with_data_state_for_tests(data_state)
|
||||
}
|
||||
|
||||
fn gemini_embedding_conversion_execution_runtime() -> Router {
|
||||
Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
@@ -189,6 +312,29 @@ fn gemini_embedding_conversion_execution_runtime() -> Router {
|
||||
)
|
||||
}
|
||||
|
||||
fn aliyun_embedding_conversion_execution_runtime(
|
||||
expected_contents: serde_json::Value,
|
||||
expected_parameters: Option<serde_json::Value>,
|
||||
) -> Router {
|
||||
let expected_contents = Arc::new(expected_contents);
|
||||
let expected_parameters = Arc::new(expected_parameters);
|
||||
Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(move |Json(plan): Json<ExecutionPlan>| {
|
||||
let expected_contents = Arc::clone(&expected_contents);
|
||||
let expected_parameters = Arc::clone(&expected_parameters);
|
||||
async move {
|
||||
assert_openai_to_aliyun_embedding_execution_plan(
|
||||
&plan,
|
||||
&expected_contents,
|
||||
expected_parameters.as_ref().as_ref(),
|
||||
);
|
||||
Json(aliyun_embedding_execution_result(&plan))
|
||||
}
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
fn vertex_gemini_embedding_conversion_execution_runtime() -> Router {
|
||||
Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
@@ -300,6 +446,40 @@ fn vertex_gemini_embedding_candidate_row() -> StoredMinimalCandidateSelectionRow
|
||||
row
|
||||
}
|
||||
|
||||
fn aliyun_embedding_candidate_row() -> StoredMinimalCandidateSelectionRow {
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: "provider-aliyun-embedding".to_string(),
|
||||
provider_name: "Aliyun DashScope".to_string(),
|
||||
provider_type: "aliyun".to_string(),
|
||||
provider_priority: 1,
|
||||
provider_is_active: true,
|
||||
endpoint_id: "endpoint-aliyun-embedding".to_string(),
|
||||
endpoint_api_format: "aliyun:multimodal_embedding".to_string(),
|
||||
endpoint_api_family: Some("aliyun".to_string()),
|
||||
endpoint_kind: Some("multimodal_embedding".to_string()),
|
||||
endpoint_is_active: true,
|
||||
key_id: "key-upstream-aliyun-embedding".to_string(),
|
||||
key_name: "default".to_string(),
|
||||
key_auth_type: "api_key".to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec!["aliyun:multimodal_embedding".to_string()]),
|
||||
key_allowed_models: None,
|
||||
key_capabilities: None,
|
||||
key_internal_priority: 50,
|
||||
key_global_priority_by_format: None,
|
||||
model_id: "model-qwen3-vl-embedding".to_string(),
|
||||
global_model_id: "global-qwen3-vl-embedding".to_string(),
|
||||
global_model_name: "qwen3-vl-embedding".to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(false),
|
||||
model_provider_model_name: "qwen3-vl-embedding".to_string(),
|
||||
model_provider_model_mappings: None,
|
||||
model_supports_streaming: Some(false),
|
||||
model_is_active: true,
|
||||
model_is_available: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn assert_embedding_execution_plan(plan: &ExecutionPlan) {
|
||||
assert_eq!(plan.client_api_format, "openai:embedding");
|
||||
assert_eq!(plan.provider_api_format, "openai:embedding");
|
||||
@@ -311,6 +491,34 @@ fn assert_embedding_execution_plan(plan: &ExecutionPlan) {
|
||||
assert!(body.get("input").is_some());
|
||||
}
|
||||
|
||||
fn assert_openai_to_aliyun_embedding_execution_plan(
|
||||
plan: &ExecutionPlan,
|
||||
expected_contents: &serde_json::Value,
|
||||
expected_parameters: Option<&serde_json::Value>,
|
||||
) {
|
||||
assert_eq!(plan.client_api_format, "openai:embedding");
|
||||
assert_eq!(plan.provider_api_format, "aliyun:multimodal_embedding");
|
||||
assert_eq!(plan.method, "POST");
|
||||
assert_eq!(
|
||||
plan.url,
|
||||
"https://dashscope.aliyuncs.com/api/v1/services/embeddings/multimodal-embedding/multimodal-embedding"
|
||||
);
|
||||
assert_eq!(
|
||||
plan.headers.get("authorization").map(String::as_str),
|
||||
Some("Bearer sk-upstream-aliyun-embedding")
|
||||
);
|
||||
assert_eq!(plan.model_name.as_deref(), Some("qwen3-vl-embedding"));
|
||||
assert!(!plan.stream);
|
||||
let body = plan.body.json_body.as_ref().expect("json request body");
|
||||
assert_eq!(body["model"], "qwen3-vl-embedding");
|
||||
assert_eq!(&body["input"]["contents"], expected_contents);
|
||||
match expected_parameters {
|
||||
Some(expected) => assert_eq!(&body["parameters"], expected),
|
||||
None => assert!(body.get("parameters").is_none()),
|
||||
}
|
||||
assert!(body.get("messages").is_none());
|
||||
}
|
||||
|
||||
fn assert_openai_to_gemini_embedding_execution_plan(plan: &ExecutionPlan) {
|
||||
assert_eq!(plan.client_api_format, "openai:embedding");
|
||||
assert_eq!(plan.provider_api_format, "gemini:embedding");
|
||||
@@ -503,6 +711,41 @@ fn gemini_batch_embedding_execution_result(plan: &ExecutionPlan) -> ExecutionRes
|
||||
}
|
||||
}
|
||||
|
||||
fn aliyun_embedding_execution_result(plan: &ExecutionPlan) -> ExecutionResult {
|
||||
ExecutionResult {
|
||||
request_id: plan.request_id.clone(),
|
||||
candidate_id: plan.candidate_id.clone(),
|
||||
status_code: 200,
|
||||
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
|
||||
body: Some(ResponseBody {
|
||||
json_body: Some(json!({
|
||||
"output": {
|
||||
"embeddings": [
|
||||
{
|
||||
"index": 0,
|
||||
"embedding": [0.1, 0.2, 0.3],
|
||||
"type": "fusion"
|
||||
}
|
||||
]
|
||||
},
|
||||
"usage": {
|
||||
"input_tokens": 432,
|
||||
"input_tokens_details": {
|
||||
"image_tokens": 402,
|
||||
"text_tokens": 30
|
||||
},
|
||||
"output_tokens": 1,
|
||||
"total_tokens": 433
|
||||
},
|
||||
"request_id": "aliyun-request-1"
|
||||
})),
|
||||
body_bytes_b64: None,
|
||||
}),
|
||||
telemetry: None,
|
||||
error: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn embeddings_route_accepts_openai_payload() {
|
||||
let (execution_runtime_url, execution_runtime_handle) =
|
||||
@@ -714,6 +957,183 @@ async fn embeddings_route_converts_openai_batch_payload_to_gemini_batch_endpoint
|
||||
execution_runtime_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn embeddings_route_converts_text_payload_to_aliyun_embedding_provider() {
|
||||
let (execution_runtime_url, execution_runtime_handle) =
|
||||
start_server(aliyun_embedding_conversion_execution_runtime(
|
||||
json!([{ "text": "hello" }]),
|
||||
Some(json!({ "dimension": 1024 })),
|
||||
))
|
||||
.await;
|
||||
let gateway = build_router_with_state(aliyun_embedding_success_state(execution_runtime_url));
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/embeddings"))
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-aliyun-embedding-success",
|
||||
)
|
||||
.json(&json!({
|
||||
"model": "qwen3-vl-embedding",
|
||||
"input": "hello",
|
||||
"dimensions": 1024
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(CONTROL_ENDPOINT_SIGNATURE_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("openai:embedding")
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(payload["object"], "list");
|
||||
assert_eq!(payload["request_id"], "aliyun-request-1");
|
||||
assert_eq!(payload["model"], "qwen3-vl-embedding");
|
||||
assert_eq!(payload["data"][0]["object"], "embedding");
|
||||
assert_eq!(payload["data"][0]["embedding"], json!([0.1, 0.2, 0.3]));
|
||||
assert_eq!(payload["data"][0]["type"], "fusion");
|
||||
assert_eq!(payload["usage"]["prompt_tokens"], json!(432));
|
||||
assert_eq!(payload["usage"]["completion_tokens"], json!(1));
|
||||
assert_eq!(payload["usage"]["total_tokens"], json!(433));
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn embeddings_route_converts_multimodal_payload_to_aliyun_embedding_provider() {
|
||||
let expected_contents = json!([
|
||||
{ "text": "white running shoes" },
|
||||
{ "image": "https://example.com/shoe.png" },
|
||||
{ "multi_images": ["https://example.com/a.png", "https://example.com/b.png"] }
|
||||
]);
|
||||
let (execution_runtime_url, execution_runtime_handle) =
|
||||
start_server(aliyun_embedding_conversion_execution_runtime(
|
||||
expected_contents.clone(),
|
||||
Some(json!({ "res_level": 2, "max_video_frames": 64 })),
|
||||
))
|
||||
.await;
|
||||
let gateway = build_router_with_state(aliyun_embedding_success_state(execution_runtime_url));
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/embeddings"))
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-aliyun-embedding-success",
|
||||
)
|
||||
.json(&json!({
|
||||
"model": "qwen3-vl-embedding",
|
||||
"input": expected_contents,
|
||||
"parameters": {
|
||||
"res_level": 2,
|
||||
"max_video_frames": 64
|
||||
}
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(payload["data"][0]["embedding"], json!([0.1, 0.2, 0.3]));
|
||||
assert_eq!(payload["data"][0]["type"], "fusion");
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn embeddings_route_skips_openai_candidate_for_multimodal_payload() {
|
||||
let expected_contents = json!([
|
||||
{ "text": "white running shoes" },
|
||||
{ "image": "https://example.com/shoe.png" }
|
||||
]);
|
||||
let (execution_runtime_url, execution_runtime_handle) =
|
||||
start_server(aliyun_embedding_conversion_execution_runtime(
|
||||
expected_contents.clone(),
|
||||
Some(json!({ "enable_fusion": true })),
|
||||
))
|
||||
.await;
|
||||
let gateway = build_router_with_state(mixed_embedding_success_state(execution_runtime_url));
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/embeddings"))
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-mixed-embedding-success",
|
||||
)
|
||||
.json(&json!({
|
||||
"model": "qwen3-vl-embedding",
|
||||
"input": expected_contents,
|
||||
"parameters": {
|
||||
"enable_fusion": true
|
||||
}
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(payload["data"][0]["embedding"], json!([0.1, 0.2, 0.3]));
|
||||
assert_eq!(payload["data"][0]["type"], "fusion");
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn embeddings_route_converts_fusion_payload_to_aliyun_embedding_provider() {
|
||||
let expected_contents = json!([
|
||||
{
|
||||
"text": "white running shoes",
|
||||
"image": "https://example.com/shoe.png"
|
||||
},
|
||||
{ "video": "https://example.com/demo.mp4" }
|
||||
]);
|
||||
let (execution_runtime_url, execution_runtime_handle) =
|
||||
start_server(aliyun_embedding_conversion_execution_runtime(
|
||||
expected_contents.clone(),
|
||||
Some(json!({ "enable_fusion": true })),
|
||||
))
|
||||
.await;
|
||||
let gateway = build_router_with_state(aliyun_embedding_success_state(execution_runtime_url));
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/embeddings"))
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-aliyun-embedding-success",
|
||||
)
|
||||
.json(&json!({
|
||||
"model": "qwen3-vl-embedding",
|
||||
"input": expected_contents,
|
||||
"parameters": {
|
||||
"enable_fusion": true
|
||||
}
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(payload["data"][0]["embedding"], json!([0.1, 0.2, 0.3]));
|
||||
assert_eq!(payload["data"][0]["type"], "fusion");
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gemini_embed_content_route_uses_native_gemini_embedding_provider() {
|
||||
let (execution_runtime_url, execution_runtime_handle) =
|
||||
@@ -836,6 +1256,14 @@ async fn embeddings_route_rejects_invalid_local_payloads() {
|
||||
r#"{"model":"text-embedding-3-small","input":[[1],[]]}"#,
|
||||
"Embedding request input is required",
|
||||
),
|
||||
(
|
||||
r#"{"model":"text-embedding-3-small","input":[{}]}"#,
|
||||
"Embedding request input is required",
|
||||
),
|
||||
(
|
||||
r#"{"model":"text-embedding-3-small","input":[{"image":" "} ]}"#,
|
||||
"Embedding request input is required",
|
||||
),
|
||||
(
|
||||
r#"{"model":"text-embedding-3-small","input":"hello","stream":true}"#,
|
||||
"Embedding requests do not support streaming",
|
||||
|
||||
@@ -24,6 +24,7 @@ use aether_data::repository::auth::{
|
||||
use aether_data::repository::auth_modules::{
|
||||
InMemoryAuthModuleReadRepository, StoredLdapModuleConfig, StoredOAuthProviderModuleConfig,
|
||||
};
|
||||
use aether_data::repository::billing::InMemoryBillingReadRepository;
|
||||
use aether_data::repository::management_tokens::{
|
||||
InMemoryManagementTokenRepository, StoredManagementToken, StoredManagementTokenUserSummary,
|
||||
StoredManagementTokenWithUser,
|
||||
@@ -36,6 +37,10 @@ use aether_data::repository::users::{
|
||||
use aether_data::repository::wallet::{
|
||||
InMemoryWalletRepository, StoredWalletSnapshot, WalletWriteRepository,
|
||||
};
|
||||
use aether_data_contracts::repository::billing::{
|
||||
AdminBillingMutationOutcome, BillingPlanWriteInput, BillingReadRepository,
|
||||
PaymentGatewayConfigWriteInput,
|
||||
};
|
||||
use aether_data_contracts::repository::global_models::StoredProviderActiveGlobalModel;
|
||||
use aether_data_contracts::repository::provider_catalog::ProviderCatalogReadRepository;
|
||||
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UsageRepository};
|
||||
@@ -3901,6 +3906,209 @@ async fn gateway_creates_wallet_recharge_orders_locally_without_proxying_upstrea
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_reuses_pending_billing_plan_checkout_order_without_proxying_upstream() {
|
||||
let now = Utc::now();
|
||||
let user = StoredUserAuthRecord::new(
|
||||
"user-billing-checkout-reuse".to_string(),
|
||||
Some("[email protected]".to_string()),
|
||||
true,
|
||||
"billing_checkout_reuse_user".to_string(),
|
||||
Some("$2y$10$.OBQfixAECpsb8V/VS3csOMf00x2E/jD/gnud20t6RG0yiQosyOZ2".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
Some(json!(["openai"])),
|
||||
Some(json!(["openai:chat"])),
|
||||
Some(json!(["gpt-5"])),
|
||||
true,
|
||||
false,
|
||||
Some(now),
|
||||
Some(now),
|
||||
)
|
||||
.expect("auth user should build");
|
||||
let wallet = StoredWalletSnapshot::new(
|
||||
"wallet-billing-checkout-reuse".to_string(),
|
||||
Some(user.id.clone()),
|
||||
None,
|
||||
12.5,
|
||||
3.0,
|
||||
"finite".to_string(),
|
||||
"USD".to_string(),
|
||||
"active".to_string(),
|
||||
20.0,
|
||||
4.5,
|
||||
0.0,
|
||||
0.0,
|
||||
now.timestamp(),
|
||||
)
|
||||
.expect("wallet should build");
|
||||
let access_token = build_test_auth_token(
|
||||
"access",
|
||||
serde_json::Map::from_iter([
|
||||
("user_id".to_string(), json!(user.id.clone())),
|
||||
("role".to_string(), json!(user.role.clone())),
|
||||
(
|
||||
"created_at".to_string(),
|
||||
json!(user.created_at.map(|value| value.to_rfc3339())),
|
||||
),
|
||||
(
|
||||
"session_id".to_string(),
|
||||
json!("session-billing-checkout-reuse"),
|
||||
),
|
||||
]),
|
||||
now + chrono::Duration::hours(1),
|
||||
);
|
||||
let billing_repository = Arc::new(InMemoryBillingReadRepository::seed(Vec::new()));
|
||||
let encrypted_key = encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "epay-secret")
|
||||
.expect("merchant key should encrypt");
|
||||
let AdminBillingMutationOutcome::Applied(_) = billing_repository
|
||||
.upsert_payment_gateway_config(&PaymentGatewayConfigWriteInput {
|
||||
provider: "epay".to_string(),
|
||||
enabled: true,
|
||||
endpoint_url: "https://pay.example.com/".to_string(),
|
||||
callback_base_url: Some("https://app.example.com".to_string()),
|
||||
merchant_id: "merchant-1".to_string(),
|
||||
merchant_key_encrypted: Some(encrypted_key),
|
||||
preserve_existing_secret: false,
|
||||
pay_currency: "CNY".to_string(),
|
||||
usd_exchange_rate: 7.25,
|
||||
min_recharge_usd: 1.0,
|
||||
channels_json: json!([
|
||||
{
|
||||
"channel": "alipay",
|
||||
"display_name": "支付宝"
|
||||
}
|
||||
]),
|
||||
})
|
||||
.await
|
||||
.expect("gateway config should create")
|
||||
else {
|
||||
panic!("gateway config should apply");
|
||||
};
|
||||
let plan = match billing_repository
|
||||
.create_billing_plan(&BillingPlanWriteInput {
|
||||
title: "每日额度月卡".to_string(),
|
||||
description: Some("测试套餐".to_string()),
|
||||
price_amount: 100.0,
|
||||
price_currency: "CNY".to_string(),
|
||||
duration_unit: "month".to_string(),
|
||||
duration_value: 1,
|
||||
enabled: true,
|
||||
sort_order: 1,
|
||||
max_active_per_user: 1,
|
||||
purchase_limit_scope: "active_period".to_string(),
|
||||
entitlements_json: json!([
|
||||
{
|
||||
"type": "daily_quota",
|
||||
"daily_quota_usd": 50.0,
|
||||
"reset_timezone": "Asia/Shanghai",
|
||||
"allow_wallet_overage": false
|
||||
}
|
||||
]),
|
||||
})
|
||||
.await
|
||||
.expect("billing plan should create")
|
||||
{
|
||||
AdminBillingMutationOutcome::Applied(plan) => plan,
|
||||
other => panic!("billing plan should apply, got {other:?}"),
|
||||
};
|
||||
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/{*path}",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("proxied"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (_upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let user_repository = Arc::new(InMemoryUserReadRepository::seed_auth_users(vec![user]));
|
||||
let wallet_repository = Arc::new(InMemoryWalletRepository::seed(vec![wallet]));
|
||||
let state = AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_user_billing_and_wallet_for_tests(
|
||||
user_repository,
|
||||
billing_repository,
|
||||
wallet_repository,
|
||||
))
|
||||
.with_auth_sessions_for_tests([sample_auth_session(
|
||||
"user-billing-checkout-reuse",
|
||||
"session-billing-checkout-reuse",
|
||||
"device-billing-checkout-reuse",
|
||||
"refresh-token-placeholder",
|
||||
now,
|
||||
)]);
|
||||
let gateway = build_router_with_state(state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let checkout_body = json!({
|
||||
"payment_provider": "epay",
|
||||
"payment_method": "epay",
|
||||
"payment_channel": "alipay",
|
||||
});
|
||||
let first_response = client
|
||||
.post(format!(
|
||||
"{gateway_url}/api/billing/plans/{}/checkout",
|
||||
plan.id
|
||||
))
|
||||
.header("authorization", format!("Bearer {access_token}"))
|
||||
.header("x-client-device-id", "device-billing-checkout-reuse")
|
||||
.header("user-agent", "AetherTest/1.0")
|
||||
.json(&checkout_body)
|
||||
.send()
|
||||
.await
|
||||
.expect("first checkout request should succeed");
|
||||
assert_eq!(first_response.status(), StatusCode::OK);
|
||||
let first_payload: serde_json::Value = first_response
|
||||
.json()
|
||||
.await
|
||||
.expect("first checkout json should parse");
|
||||
let first_order_id = first_payload["order"]["id"]
|
||||
.as_str()
|
||||
.expect("first order id should exist")
|
||||
.to_string();
|
||||
assert_eq!(first_payload["order"]["status"], "pending");
|
||||
assert_eq!(first_payload["order"]["product_id"], plan.id);
|
||||
assert_eq!(
|
||||
first_payload["reused_pending_order"],
|
||||
serde_json::Value::Null
|
||||
);
|
||||
|
||||
let second_response = client
|
||||
.post(format!(
|
||||
"{gateway_url}/api/billing/plans/{}/checkout",
|
||||
plan.id
|
||||
))
|
||||
.header("authorization", format!("Bearer {access_token}"))
|
||||
.header("x-client-device-id", "device-billing-checkout-reuse")
|
||||
.header("user-agent", "AetherTest/1.0")
|
||||
.json(&checkout_body)
|
||||
.send()
|
||||
.await
|
||||
.expect("second checkout request should succeed");
|
||||
assert_eq!(second_response.status(), StatusCode::OK);
|
||||
let second_payload: serde_json::Value = second_response
|
||||
.json()
|
||||
.await
|
||||
.expect("second checkout json should parse");
|
||||
assert_eq!(second_payload["order"]["id"], first_order_id);
|
||||
assert_eq!(second_payload["reused_pending_order"], true);
|
||||
assert_eq!(
|
||||
second_payload["payment_instructions"],
|
||||
first_payload["payment_instructions"]
|
||||
);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_creates_wallet_refunds_locally_without_proxying_upstream() {
|
||||
let now = Utc::now();
|
||||
|
||||
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user