Improve gateway transport and usage runtime

This commit is contained in:
elky
2026-06-25 22:36:27 +08:00
parent d336d1a7fa
commit 6f00e9fc67
112 changed files with 12456 additions and 1387 deletions
@@ -20,6 +20,30 @@ use crate::constants::{
};
use aether_data_contracts::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
const EMBEDDING_PROXY_TEST_STACK_BYTES: usize = 16 * 1024 * 1024;
fn run_embedding_proxy_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(EMBEDDING_PROXY_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("embedding proxy test thread should spawn");
if let Err(payload) = handle.join() {
std::panic::resume_unwind(payload);
}
}
fn embedding_success_state(execution_runtime_url: String) -> AppState {
let mut snapshot =
sample_currently_usable_auth_snapshot("key-embedding-success", "user-embedding-success");
@@ -746,8 +770,15 @@ fn aliyun_embedding_execution_result(plan: &ExecutionPlan) -> ExecutionResult {
}
}
#[tokio::test]
async fn embeddings_route_accepts_openai_payload() {
#[test]
fn embeddings_route_accepts_openai_payload() {
run_embedding_proxy_test(
"embeddings_route_accepts_openai_payload",
embeddings_route_accepts_openai_payload_impl,
);
}
async fn embeddings_route_accepts_openai_payload_impl() {
let (execution_runtime_url, execution_runtime_handle) =
start_server(embedding_execution_runtime()).await;
let gateway = build_router_with_state(embedding_success_state(execution_runtime_url));
@@ -816,8 +847,15 @@ async fn embeddings_route_accepts_openai_payload() {
execution_runtime_handle.abort();
}
#[tokio::test]
async fn embeddings_route_converts_openai_payload_to_gemini_embedding_provider() {
#[test]
fn embeddings_route_converts_openai_payload_to_gemini_embedding_provider() {
run_embedding_proxy_test(
"embeddings_route_converts_openai_payload_to_gemini_embedding_provider",
embeddings_route_converts_openai_payload_to_gemini_embedding_provider_impl,
);
}
async fn embeddings_route_converts_openai_payload_to_gemini_embedding_provider_impl() {
let (execution_runtime_url, execution_runtime_handle) =
start_server(gemini_embedding_conversion_execution_runtime()).await;
let gateway = build_router_with_state(gemini_embedding_success_state(
@@ -867,8 +905,15 @@ async fn embeddings_route_converts_openai_payload_to_gemini_embedding_provider()
execution_runtime_handle.abort();
}
#[tokio::test]
async fn embeddings_route_converts_openai_payload_to_vertex_gemini_embedding_provider() {
#[test]
fn embeddings_route_converts_openai_payload_to_vertex_gemini_embedding_provider() {
run_embedding_proxy_test(
"embeddings_route_converts_openai_payload_to_vertex_gemini_embedding_provider",
embeddings_route_converts_openai_payload_to_vertex_gemini_embedding_provider_impl,
);
}
async fn embeddings_route_converts_openai_payload_to_vertex_gemini_embedding_provider_impl() {
let (execution_runtime_url, execution_runtime_handle) =
start_server(vertex_gemini_embedding_conversion_execution_runtime()).await;
let gateway =
@@ -911,8 +956,15 @@ async fn embeddings_route_converts_openai_payload_to_vertex_gemini_embedding_pro
execution_runtime_handle.abort();
}
#[tokio::test]
async fn embeddings_route_converts_openai_batch_payload_to_gemini_batch_endpoint() {
#[test]
fn embeddings_route_converts_openai_batch_payload_to_gemini_batch_endpoint() {
run_embedding_proxy_test(
"embeddings_route_converts_openai_batch_payload_to_gemini_batch_endpoint",
embeddings_route_converts_openai_batch_payload_to_gemini_batch_endpoint_impl,
);
}
async fn embeddings_route_converts_openai_batch_payload_to_gemini_batch_endpoint_impl() {
let (execution_runtime_url, execution_runtime_handle) =
start_server(gemini_embedding_batch_conversion_execution_runtime()).await;
let gateway = build_router_with_state(gemini_embedding_success_state(
@@ -957,8 +1009,15 @@ 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() {
#[test]
fn embeddings_route_converts_text_payload_to_aliyun_embedding_provider() {
run_embedding_proxy_test(
"embeddings_route_converts_text_payload_to_aliyun_embedding_provider",
embeddings_route_converts_text_payload_to_aliyun_embedding_provider_impl,
);
}
async fn embeddings_route_converts_text_payload_to_aliyun_embedding_provider_impl() {
let (execution_runtime_url, execution_runtime_handle) =
start_server(aliyun_embedding_conversion_execution_runtime(
json!([{ "text": "hello" }]),
@@ -1006,8 +1065,15 @@ async fn embeddings_route_converts_text_payload_to_aliyun_embedding_provider() {
execution_runtime_handle.abort();
}
#[tokio::test]
async fn embeddings_route_converts_multimodal_payload_to_aliyun_embedding_provider() {
#[test]
fn embeddings_route_converts_multimodal_payload_to_aliyun_embedding_provider() {
run_embedding_proxy_test(
"embeddings_route_converts_multimodal_payload_to_aliyun_embedding_provider",
embeddings_route_converts_multimodal_payload_to_aliyun_embedding_provider_impl,
);
}
async fn embeddings_route_converts_multimodal_payload_to_aliyun_embedding_provider_impl() {
let expected_contents = json!([
{ "text": "white running shoes" },
{ "image": "https://example.com/shoe.png" },
@@ -1049,8 +1115,15 @@ async fn embeddings_route_converts_multimodal_payload_to_aliyun_embedding_provid
execution_runtime_handle.abort();
}
#[tokio::test]
async fn embeddings_route_skips_openai_candidate_for_multimodal_payload() {
#[test]
fn embeddings_route_skips_openai_candidate_for_multimodal_payload() {
run_embedding_proxy_test(
"embeddings_route_skips_openai_candidate_for_multimodal_payload",
embeddings_route_skips_openai_candidate_for_multimodal_payload_impl,
);
}
async fn embeddings_route_skips_openai_candidate_for_multimodal_payload_impl() {
let expected_contents = json!([
{ "text": "white running shoes" },
{ "image": "https://example.com/shoe.png" }
@@ -1090,8 +1163,15 @@ async fn embeddings_route_skips_openai_candidate_for_multimodal_payload() {
execution_runtime_handle.abort();
}
#[tokio::test]
async fn embeddings_route_converts_fusion_payload_to_aliyun_embedding_provider() {
#[test]
fn embeddings_route_converts_fusion_payload_to_aliyun_embedding_provider() {
run_embedding_proxy_test(
"embeddings_route_converts_fusion_payload_to_aliyun_embedding_provider",
embeddings_route_converts_fusion_payload_to_aliyun_embedding_provider_impl,
);
}
async fn embeddings_route_converts_fusion_payload_to_aliyun_embedding_provider_impl() {
let expected_contents = json!([
{
"text": "white running shoes",
@@ -1134,8 +1214,15 @@ async fn embeddings_route_converts_fusion_payload_to_aliyun_embedding_provider()
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gemini_embed_content_route_uses_native_gemini_embedding_provider() {
#[test]
fn gemini_embed_content_route_uses_native_gemini_embedding_provider() {
run_embedding_proxy_test(
"gemini_embed_content_route_uses_native_gemini_embedding_provider",
gemini_embed_content_route_uses_native_gemini_embedding_provider_impl,
);
}
async fn gemini_embed_content_route_uses_native_gemini_embedding_provider_impl() {
let (execution_runtime_url, execution_runtime_handle) =
start_server(gemini_embedding_native_execution_runtime()).await;
let gateway = build_router_with_state(gemini_embedding_success_state(
@@ -1188,8 +1275,15 @@ async fn gemini_embed_content_route_uses_native_gemini_embedding_provider() {
execution_runtime_handle.abort();
}
#[tokio::test]
async fn embeddings_route_accepts_all_canonical_input_shapes() {
#[test]
fn embeddings_route_accepts_all_canonical_input_shapes() {
run_embedding_proxy_test(
"embeddings_route_accepts_all_canonical_input_shapes",
embeddings_route_accepts_all_canonical_input_shapes_impl,
);
}
async fn embeddings_route_accepts_all_canonical_input_shapes_impl() {
let (execution_runtime_url, execution_runtime_handle) =
start_server(embedding_execution_runtime()).await;
let gateway = build_router_with_state(embedding_success_state(execution_runtime_url));
@@ -1229,8 +1323,15 @@ async fn embeddings_route_accepts_all_canonical_input_shapes() {
execution_runtime_handle.abort();
}
#[tokio::test]
async fn embeddings_route_rejects_invalid_local_payloads() {
#[test]
fn embeddings_route_rejects_invalid_local_payloads() {
run_embedding_proxy_test(
"embeddings_route_rejects_invalid_local_payloads",
embeddings_route_rejects_invalid_local_payloads_impl,
);
}
async fn embeddings_route_rejects_invalid_local_payloads_impl() {
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
@@ -1294,8 +1395,15 @@ async fn embeddings_route_rejects_invalid_local_payloads() {
gateway_handle.abort();
}
#[tokio::test]
async fn embeddings_route_rejects_non_json_content_type() {
#[test]
fn embeddings_route_rejects_non_json_content_type() {
run_embedding_proxy_test(
"embeddings_route_rejects_non_json_content_type",
embeddings_route_rejects_non_json_content_type_impl,
);
}
async fn embeddings_route_rejects_non_json_content_type_impl() {
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
@@ -1317,8 +1425,15 @@ async fn embeddings_route_rejects_non_json_content_type() {
gateway_handle.abort();
}
#[tokio::test]
async fn embeddings_route_rejects_chat_only_model() {
#[test]
fn embeddings_route_rejects_chat_only_model() {
run_embedding_proxy_test(
"embeddings_route_rejects_chat_only_model",
embeddings_route_rejects_chat_only_model_impl,
);
}
async fn embeddings_route_rejects_chat_only_model_impl() {
let mut snapshot = sample_currently_usable_auth_snapshot("key-embedding-1", "user-embedding-1");
snapshot.user_allowed_api_formats = Some(vec!["openai:embedding".to_string()]);
snapshot.api_key_allowed_api_formats = Some(vec!["openai:embedding".to_string()]);
@@ -1366,8 +1481,15 @@ async fn embeddings_route_rejects_chat_only_model() {
gateway_handle.abort();
}
#[tokio::test]
async fn embeddings_route_rejects_chat_only_api_format() {
#[test]
fn embeddings_route_rejects_chat_only_api_format() {
run_embedding_proxy_test(
"embeddings_route_rejects_chat_only_api_format",
embeddings_route_rejects_chat_only_api_format_impl,
);
}
async fn embeddings_route_rejects_chat_only_api_format_impl() {
let mut snapshot = sample_currently_usable_auth_snapshot("key-embedding-2", "user-embedding-2");
snapshot.user_allowed_api_formats = Some(vec!["openai:chat".to_string()]);
snapshot.api_key_allowed_api_formats = Some(vec!["openai:chat".to_string()]);
@@ -20,6 +20,30 @@ use crate::constants::{
EXECUTION_PATH_LOCAL_AUTH_DENIED,
};
const RERANK_PROXY_TEST_STACK_BYTES: usize = 16 * 1024 * 1024;
fn run_rerank_proxy_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(RERANK_PROXY_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("rerank proxy test thread should spawn");
if let Err(payload) = handle.join() {
std::panic::resume_unwind(payload);
}
}
fn rerank_success_state(execution_runtime_url: String) -> AppState {
let mut snapshot =
sample_currently_usable_auth_snapshot("key-rerank-success", "user-rerank-success");
@@ -142,8 +166,15 @@ fn rerank_execution_result(plan: &ExecutionPlan) -> ExecutionResult {
}
}
#[tokio::test]
async fn rerank_route_accepts_openai_payload() {
#[test]
fn rerank_route_accepts_openai_payload() {
run_rerank_proxy_test(
"rerank_route_accepts_openai_payload",
rerank_route_accepts_openai_payload_impl,
);
}
async fn rerank_route_accepts_openai_payload_impl() {
let (execution_runtime_url, execution_runtime_handle) =
start_server(rerank_execution_runtime()).await;
let gateway = build_router_with_state(rerank_success_state(execution_runtime_url));
@@ -200,8 +231,15 @@ async fn rerank_route_accepts_openai_payload() {
execution_runtime_handle.abort();
}
#[tokio::test]
async fn rerank_route_rejects_invalid_local_payloads() {
#[test]
fn rerank_route_rejects_invalid_local_payloads() {
run_rerank_proxy_test(
"rerank_route_rejects_invalid_local_payloads",
rerank_route_rejects_invalid_local_payloads_impl,
);
}
async fn rerank_route_rejects_invalid_local_payloads_impl() {
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
@@ -257,8 +295,15 @@ async fn rerank_route_rejects_invalid_local_payloads() {
gateway_handle.abort();
}
#[tokio::test]
async fn rerank_route_rejects_non_json_content_type() {
#[test]
fn rerank_route_rejects_non_json_content_type() {
run_rerank_proxy_test(
"rerank_route_rejects_non_json_content_type",
rerank_route_rejects_non_json_content_type_impl,
);
}
async fn rerank_route_rejects_non_json_content_type_impl() {
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
@@ -280,8 +325,15 @@ async fn rerank_route_rejects_non_json_content_type() {
gateway_handle.abort();
}
#[tokio::test]
async fn rerank_route_rejects_chat_only_api_format() {
#[test]
fn rerank_route_rejects_chat_only_api_format() {
run_rerank_proxy_test(
"rerank_route_rejects_chat_only_api_format",
rerank_route_rejects_chat_only_api_format_impl,
);
}
async fn rerank_route_rejects_chat_only_api_format_impl() {
let mut snapshot = sample_currently_usable_auth_snapshot("key-rerank-2", "user-rerank-2");
snapshot.user_allowed_api_formats = Some(vec!["openai:chat".to_string()]);
snapshot.api_key_allowed_api_formats = Some(vec!["openai:chat".to_string()]);