refactor: 拆分 gateway 单体为独立 crate,新增 systemd 部署方案

将 gateway 内部的 model-fetch、provider-transport、scheduler-core、
usage-runtime、video-tasks-core 模块提取为独立 crate;重构 gateway
内部模块结构(state/router/cache/data/query 等);移除大量遗留模块
文件;新增 systemd 二进制部署骨架及相关文档;更新前端 usage 相关
API 和组件。
This commit is contained in:
fawney19
2026-04-05 20:23:16 +08:00
parent cbc811f6ce
commit 763ff03a7b
777 changed files with 42659 additions and 21469 deletions

View File

@@ -0,0 +1,788 @@
use super::{
hash_api_key, sample_models_candidate_row, unrestricted_models_snapshot,
InMemoryAuthApiKeySnapshotRepository, InMemoryMinimalCandidateSelectionReadRepository,
InMemoryVideoTaskRepository, UpsertVideoTask, VideoTaskLookupKey, VideoTaskReadRepository,
VideoTaskStatus, VideoTaskWriteRepository, DEVELOPMENT_ENCRYPTION_KEY,
};
use crate::tests::{
any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server,
to_bytes, AppState, Arc, Body, Json, Mutex, Request, Router, StatusCode, EXECUTION_PATH_HEADER,
EXECUTION_PATH_LOCAL_AI_PUBLIC,
};
use axum::response::IntoResponse;
fn gemini_operation_status_label(status: VideoTaskStatus) -> &'static str {
match status {
VideoTaskStatus::Pending => "Pending",
VideoTaskStatus::Submitted => "Submitted",
VideoTaskStatus::Queued => "Queued",
VideoTaskStatus::Processing => "Processing",
VideoTaskStatus::Completed => "Completed",
VideoTaskStatus::Failed => "Failed",
VideoTaskStatus::Cancelled => "Cancelled",
VideoTaskStatus::Expired => "Expired",
VideoTaskStatus::Deleted => "Deleted",
}
}
fn sample_gemini_video_task(
id: &str,
short_id: &str,
user_id: &str,
api_key_id: &str,
external_task_id: &str,
status: VideoTaskStatus,
) -> UpsertVideoTask {
let completed = matches!(status, VideoTaskStatus::Completed);
UpsertVideoTask {
id: id.to_string(),
short_id: Some(short_id.to_string()),
request_id: format!("request-{id}"),
user_id: Some(user_id.to_string()),
api_key_id: Some(api_key_id.to_string()),
username: Some(format!("user-{user_id}")),
api_key_name: Some("video-key".to_string()),
external_task_id: Some(external_task_id.to_string()),
provider_id: Some("provider-gemini-video-local-1".to_string()),
endpoint_id: Some("endpoint-gemini-video-local-1".to_string()),
key_id: Some("key-gemini-video-local-1".to_string()),
client_api_format: Some("gemini:video".to_string()),
provider_api_format: Some("gemini:video".to_string()),
format_converted: false,
model: Some("veo-3".to_string()),
prompt: Some("gemini video prompt".to_string()),
original_request_body: Some(json!({"prompt": "gemini video prompt"})),
duration_seconds: Some(8),
resolution: Some("720p".to_string()),
aspect_ratio: Some("16:9".to_string()),
size: Some("720p".to_string()),
status,
progress_percent: if completed { 100 } else { 50 },
progress_message: None,
retry_count: 0,
poll_interval_seconds: 10,
next_poll_at_unix_secs: (!completed).then_some(124),
poll_count: 0,
max_poll_count: 360,
created_at_unix_secs: 123,
submitted_at_unix_secs: Some(123),
completed_at_unix_secs: completed.then_some(124),
updated_at_unix_secs: 124,
error_code: None,
error_message: None,
video_url: None,
request_metadata: Some(json!({
"rust_local_snapshot": {
"Gemini": {
"local_short_id": short_id,
"upstream_operation_name": external_task_id,
"user_id": user_id,
"api_key_id": api_key_id,
"model": "veo-3",
"status": gemini_operation_status_label(status),
"progress_percent": if completed { 100 } else { 50 },
"error_code": null,
"error_message": null,
"metadata": {},
"persistence": {
"request_id": format!("request-{id}"),
"username": format!("user-{user_id}"),
"api_key_name": "video-key",
"client_api_format": "gemini:video",
"provider_api_format": "gemini:video",
"original_request_body": {
"prompt": "gemini video prompt"
},
"format_converted": false
},
"transport": {
"upstream_base_url": "https://generativelanguage.googleapis.com",
"provider_name": "gemini-video",
"provider_id": "provider-gemini-video-local-1",
"endpoint_id": "endpoint-gemini-video-local-1",
"key_id": "key-gemini-video-local-1",
"headers": {
"x-goog-api-key": "sk-upstream-gemini-video",
"content-type": "application/json"
},
"content_type": "application/json",
"model_name": "veo-3-upstream",
"proxy": null,
"tls_profile": null,
"timeouts": null
}
}
}
})),
}
}
#[tokio::test]
async fn gateway_handles_public_openai_models_without_hitting_fallback_probe() {
let fallback_probe_hits = Arc::new(Mutex::new(0usize));
let fallback_probe_hits_clone = Arc::clone(&fallback_probe_hits);
let fallback_probe = Router::new().route(
"/{*path}",
any(move |_request: Request| {
let fallback_probe_hits_inner = Arc::clone(&fallback_probe_hits_clone);
async move {
*fallback_probe_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("proxied"))
}
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-openai-models")),
unrestricted_models_snapshot("key-1", "user-1"),
)]));
let candidate_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_models_candidate_row("provider-openai", "openai", "openai:chat", "gpt-5", 10),
sample_models_candidate_row("provider-openai", "openai", "openai:chat", "gpt-4.1", 10),
]));
let (_unused_fallback_probe_url, fallback_probe_handle) = start_server(fallback_probe).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_minimal_candidate_selection_and_auth_for_tests(
candidate_repository,
auth_repository,
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/v1/models"))
.header("authorization", "Bearer sk-openai-models")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["object"], "list");
assert_eq!(payload["data"][0]["id"], "gpt-4.1");
assert_eq!(payload["data"][1]["id"], "gpt-5");
assert_eq!(payload["data"][0]["owned_by"], "openai");
assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
fallback_probe_handle.abort();
}
#[tokio::test]
async fn gateway_handles_public_claude_models_without_hitting_fallback_probe() {
let fallback_probe_hits = Arc::new(Mutex::new(0usize));
let fallback_probe_hits_clone = Arc::clone(&fallback_probe_hits);
let fallback_probe = Router::new().route(
"/{*path}",
any(move |_request: Request| {
let fallback_probe_hits_inner = Arc::clone(&fallback_probe_hits_clone);
async move {
*fallback_probe_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Json(json!({"proxied": true}))).into_response()
}
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-claude-models")),
unrestricted_models_snapshot("key-claude", "user-claude"),
)]));
let candidate_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_models_candidate_row(
"provider-claude",
"claude",
"claude:chat",
"claude-3-7-sonnet",
10,
),
sample_models_candidate_row(
"provider-claude",
"claude",
"claude:chat",
"claude-3-5-haiku",
10,
),
]));
let (_unused_fallback_probe_url, fallback_probe_handle) = start_server(fallback_probe).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_minimal_candidate_selection_and_auth_for_tests(
candidate_repository,
auth_repository,
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/v1/models?limit=1"))
.header("x-api-key", "sk-claude-models")
.header("anthropic-version", "2023-06-01")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["data"][0]["id"], "claude-3-5-haiku");
assert_eq!(payload["first_id"], "claude-3-5-haiku");
assert_eq!(payload["last_id"], "claude-3-5-haiku");
assert_eq!(payload["has_more"], true);
assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
fallback_probe_handle.abort();
}
#[tokio::test]
async fn gateway_handles_public_gemini_models_without_hitting_fallback_probe() {
let fallback_probe_hits = Arc::new(Mutex::new(0usize));
let fallback_probe_hits_clone = Arc::clone(&fallback_probe_hits);
let fallback_probe = Router::new().route(
"/{*path}",
any(move |_request: Request| {
let fallback_probe_hits_inner = Arc::clone(&fallback_probe_hits_clone);
async move {
*fallback_probe_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Json(json!({"proxied": true}))).into_response()
}
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-gemini-models")),
unrestricted_models_snapshot("key-gemini", "user-gemini"),
)]));
let candidate_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_models_candidate_row(
"provider-gemini",
"gemini",
"gemini:chat",
"gemini-2.5-flash",
10,
),
sample_models_candidate_row(
"provider-gemini",
"gemini",
"gemini:chat",
"gemini-2.5-pro",
10,
),
]));
let (_unused_fallback_probe_url, fallback_probe_handle) = start_server(fallback_probe).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_minimal_candidate_selection_and_auth_for_tests(
candidate_repository,
auth_repository,
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/v1beta/models?pageSize=1&key=sk-gemini-models"
))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["models"][0]["name"], "models/gemini-2.5-flash");
assert_eq!(payload["nextPageToken"], "1");
assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
fallback_probe_handle.abort();
}
#[tokio::test]
async fn gateway_handles_claude_count_tokens_without_hitting_fallback_probe() {
let fallback_probe_hits = Arc::new(Mutex::new(0usize));
let fallback_probe_hits_clone = Arc::clone(&fallback_probe_hits);
let fallback_probe = Router::new().route(
"/{*path}",
any(move |_request: Request| {
let fallback_probe_hits_inner = Arc::clone(&fallback_probe_hits_clone);
async move {
*fallback_probe_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Json(json!({"proxied": true}))).into_response()
}
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-claude-count")),
unrestricted_models_snapshot("key-claude-count", "user-claude-count"),
)]));
let (_unused_fallback_probe_url, fallback_probe_handle) = start_server(fallback_probe).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_auth_api_key_data_reader_for_tests(auth_repository),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/messages/count_tokens"))
.header("x-api-key", "sk-claude-count")
.header("anthropic-version", "2023-06-01")
.body(
serde_json::to_vec(&json!({
"model": "claude-sonnet-4-5",
"system": [{"type": "text", "text": "abcdefghijklmnop"}],
"messages": [
{
"role": "user",
"content": "abcdefghijkl"
},
{
"role": "assistant",
"content": [
{"type": "text", "text": "abcdefgh"},
{"type": "tool_use", "name": "ignored", "input": {"city": "SF"}}
]
}
]
}))
.expect("request body should encode"),
)
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok()),
Some(EXECUTION_PATH_LOCAL_AI_PUBLIC)
);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["input_tokens"], 17);
assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
fallback_probe_handle.abort();
}
#[tokio::test]
async fn gateway_rejects_invalid_claude_count_tokens_payload_without_hitting_fallback_probe() {
let fallback_probe_hits = Arc::new(Mutex::new(0usize));
let fallback_probe_hits_clone = Arc::clone(&fallback_probe_hits);
let fallback_probe = Router::new().route(
"/{*path}",
any(move |_request: Request| {
let fallback_probe_hits_inner = Arc::clone(&fallback_probe_hits_clone);
async move {
*fallback_probe_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Json(json!({"proxied": true}))).into_response()
}
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-claude-count-invalid")),
unrestricted_models_snapshot("key-claude-count-invalid", "user-claude-count-invalid"),
)]));
let (_unused_fallback_probe_url, fallback_probe_handle) = start_server(fallback_probe).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_auth_api_key_data_reader_for_tests(auth_repository),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/messages/count_tokens"))
.header("x-api-key", "sk-claude-count-invalid")
.body(
serde_json::to_vec(&json!({
"model": "claude-sonnet-4-5",
"messages": [{"role": "system", "content": "bad"}]
}))
.expect("request body should encode"),
)
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
assert_eq!(
response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok()),
Some(EXECUTION_PATH_LOCAL_AI_PUBLIC)
);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["detail"], "Invalid token count payload");
assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
fallback_probe_handle.abort();
}
#[tokio::test]
async fn gateway_handles_gemini_operation_detail_without_hitting_fallback_probe() {
let fallback_probe_hits = Arc::new(Mutex::new(0usize));
let fallback_probe_hits_clone = Arc::clone(&fallback_probe_hits);
let fallback_probe = Router::new().route(
"/{*path}",
any(move |_request: Request| {
let fallback_probe_hits_inner = Arc::clone(&fallback_probe_hits_clone);
async move {
*fallback_probe_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Json(json!({"proxied": true}))).into_response()
}
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-gemini-operation-detail")),
unrestricted_models_snapshot(
"key-gemini-operation-detail",
"user-gemini-operation-detail",
),
)]));
let repository = Arc::new(InMemoryVideoTaskRepository::default());
repository
.upsert(sample_gemini_video_task(
"task-gemini-operation-detail",
"opshort123",
"user-gemini-operation-detail",
"key-gemini-operation-detail",
"operations/ext-op-123",
VideoTaskStatus::Completed,
))
.await
.expect("upsert should succeed");
let (_unused_fallback_probe_url, fallback_probe_handle) = start_server(fallback_probe).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_auth_and_video_task_repository_for_tests(
auth_repository,
repository,
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/v1beta/operations/opshort123?key=sk-gemini-operation-detail"
))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok()),
Some(EXECUTION_PATH_LOCAL_AI_PUBLIC)
);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["name"], "models/veo-3/operations/opshort123");
assert_eq!(payload["done"], true);
assert_eq!(
payload["response"]["generateVideoResponse"]["generatedSamples"][0]["video"]["uri"],
"/v1beta/files/aev_opshort123:download?alt=media"
);
assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
fallback_probe_handle.abort();
}
#[tokio::test]
async fn gateway_lists_gemini_operations_without_hitting_fallback_probe() {
let fallback_probe_hits = Arc::new(Mutex::new(0usize));
let fallback_probe_hits_clone = Arc::clone(&fallback_probe_hits);
let fallback_probe = Router::new().route(
"/{*path}",
any(move |_request: Request| {
let fallback_probe_hits_inner = Arc::clone(&fallback_probe_hits_clone);
async move {
*fallback_probe_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Json(json!({"proxied": true}))).into_response()
}
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-gemini-operation-list")),
unrestricted_models_snapshot("key-gemini-operation-list", "user-gemini-operation-list"),
)]));
let repository = Arc::new(InMemoryVideoTaskRepository::default());
repository
.upsert(sample_gemini_video_task(
"task-gemini-operation-list-1",
"opshort-list-1",
"user-gemini-operation-list",
"key-gemini-operation-list",
"operations/ext-list-1",
VideoTaskStatus::Completed,
))
.await
.expect("upsert should succeed");
repository
.upsert(sample_gemini_video_task(
"task-gemini-operation-list-2",
"opshort-list-2",
"user-gemini-operation-list",
"key-gemini-operation-list",
"operations/ext-list-2",
VideoTaskStatus::Processing,
))
.await
.expect("upsert should succeed");
repository
.upsert(sample_gemini_video_task(
"task-gemini-operation-list-other",
"opshort-list-other",
"user-other",
"key-other",
"operations/ext-list-other",
VideoTaskStatus::Completed,
))
.await
.expect("upsert should succeed");
let (_unused_fallback_probe_url, fallback_probe_handle) = start_server(fallback_probe).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_auth_and_video_task_repository_for_tests(
auth_repository,
repository,
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/v1beta/operations?key=sk-gemini-operation-list"
))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok()),
Some(EXECUTION_PATH_LOCAL_AI_PUBLIC)
);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
let operations = payload["operations"]
.as_array()
.expect("operations should be an array");
assert_eq!(operations.len(), 2);
let operation_names = operations
.iter()
.map(|value| {
value["name"]
.as_str()
.expect("operation name should be a string")
.to_string()
})
.collect::<std::collections::BTreeSet<_>>();
assert_eq!(
operation_names,
std::collections::BTreeSet::from([
"models/veo-3/operations/opshort-list-1".to_string(),
"models/veo-3/operations/opshort-list-2".to_string(),
])
);
assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
fallback_probe_handle.abort();
}
#[tokio::test]
async fn gateway_cancels_gemini_operation_without_hitting_fallback_probe() {
#[derive(Debug, Clone, PartialEq, Eq)]
struct SeenExecutionRuntimeSyncRequest {
method: String,
url: String,
api_key: String,
}
let fallback_probe_hits = Arc::new(Mutex::new(0usize));
let fallback_probe_hits_clone = Arc::clone(&fallback_probe_hits);
let fallback_probe = Router::new().route(
"/{*path}",
any(move |_request: Request| {
let fallback_probe_hits_inner = Arc::clone(&fallback_probe_hits_clone);
async move {
*fallback_probe_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Json(json!({"proxied": true}))).into_response()
}
}),
);
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeSyncRequest>));
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |request: Request| {
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
async move {
let (_parts, body) = request.into_parts();
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");
*seen_execution_runtime_inner
.lock()
.expect("mutex should lock") = Some(SeenExecutionRuntimeSyncRequest {
method: payload
.get("method")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
url: payload
.get("url")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
api_key: payload
.get("headers")
.and_then(|value| value.get("x-goog-api-key"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
});
Json(json!({
"request_id": "trace-gemini-operation-cancel",
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {}
},
"telemetry": {
"elapsed_ms": 12
}
}))
}
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-gemini-operation-cancel")),
unrestricted_models_snapshot(
"key-gemini-operation-cancel",
"user-gemini-operation-cancel",
),
)]));
let repository = Arc::new(InMemoryVideoTaskRepository::default());
repository
.upsert(sample_gemini_video_task(
"task-gemini-operation-cancel",
"opshort-cancel",
"user-gemini-operation-cancel",
"key-gemini-operation-cancel",
"operations/ext-op-123",
VideoTaskStatus::Submitted,
))
.await
.expect("upsert should succeed");
let (fallback_probe_url, fallback_probe_handle) = start_server(fallback_probe).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
crate::data::GatewayDataState::with_auth_and_video_task_repository_for_tests(
auth_repository,
Arc::clone(&repository),
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/v1beta/operations/opshort-cancel:cancel"
))
.header("x-goog-api-key", "sk-gemini-operation-cancel")
.header(http::header::CONTENT_TYPE, "application/json")
.body("{}")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok()),
Some(EXECUTION_PATH_LOCAL_AI_PUBLIC)
);
assert_eq!(
response
.json::<serde_json::Value>()
.await
.expect("json body should parse"),
json!({})
);
let seen_execution_runtime_request = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("execution runtime sync should be captured");
assert_eq!(seen_execution_runtime_request.method, "POST");
assert_eq!(
seen_execution_runtime_request.url,
"https://generativelanguage.googleapis.com/v1beta/models/veo-3/operations/ext-op-123:cancel"
);
assert_eq!(
seen_execution_runtime_request.api_key,
"sk-upstream-gemini-video"
);
let stored = repository
.find(VideoTaskLookupKey::Id("task-gemini-operation-cancel"))
.await
.expect("task lookup should succeed")
.expect("task should exist");
assert_eq!(stored.status, VideoTaskStatus::Cancelled);
assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
execution_runtime_handle.abort();
fallback_probe_handle.abort();
}

View File

@@ -0,0 +1,186 @@
use crate::tests::{
any, build_router, start_server, Arc, Body, Mutex, Request, Router, StatusCode, READYZ_PATH,
};
#[tokio::test]
async fn gateway_exposes_readyz_without_proxying_upstream() {
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 gateway = build_router().expect("gateway should build");
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}{READYZ_PATH}"))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["status"], "ready");
assert_eq!(payload["component"], "aether-gateway");
assert_eq!(payload["warmup_status"], "disabled");
assert_eq!(payload["gate_readiness"], false);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_public_health_without_proxying_upstream() {
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 gateway = build_router().expect("gateway should build");
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/health"))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["status"], "healthy");
assert_eq!(payload["database_pool"]["source"], "rust_frontdoor");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_public_service_health_without_proxying_upstream() {
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 gateway = build_router().expect("gateway should build");
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/v1/health"))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["status"], "ok");
assert!(payload["stats"].is_object());
assert!(payload["dependencies"]["database"]["status"].is_string());
assert!(payload["dependencies"]["redis"]["status"].is_string());
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_public_root_without_proxying_upstream() {
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 gateway = build_router().expect("gateway should build");
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/"))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["status"], "running");
assert_eq!(
payload["message"],
"AI Proxy with Modular Architecture v4.0.0"
);
assert_eq!(payload["endpoints"]["health"], "/v1/health");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_public_site_info_without_proxying_upstream() {
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 gateway = build_router().expect("gateway should build");
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/public/site-info"))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["site_name"], "Aether");
assert_eq!(payload["site_subtitle"], "AI Gateway");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,109 @@
use crate::tests::{
any, build_router, start_server, Arc, Body, Mutex, Request, Router, StatusCode,
};
#[tokio::test]
async fn gateway_rejects_oauth_public_providers_as_local_not_found_without_hitting_upstream() {
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 gateway = build_router().expect("gateway should build");
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/oauth/providers"))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::NOT_FOUND);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["error"]["type"], "http_error");
assert_eq!(payload["error"]["message"], "Route not found");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_rejects_oauth_user_bindable_providers_as_local_not_found_without_hitting_upstream()
{
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 gateway = build_router().expect("gateway should build");
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/user/oauth/bindable-providers"))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::NOT_FOUND);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["error"]["type"], "http_error");
assert_eq!(payload["error"]["message"], "Route not found");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_rejects_oauth_user_bind_token_as_local_not_found_without_hitting_upstream() {
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 gateway = build_router().expect("gateway should build");
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/user/oauth/linuxdo/bind-token"))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::NOT_FOUND);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["error"]["type"], "http_error");
assert_eq!(payload["error"]["message"], "Route not found");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}

View File

@@ -0,0 +1,444 @@
use super::super::usage::{
hash_api_key, sample_local_openai_auth_snapshot, sample_local_openai_candidate_row,
sample_local_openai_endpoint, sample_local_openai_key, sample_local_openai_provider,
};
use crate::tests::{
any, build_router_with_execution_runtime_override, build_router_with_state,
build_state_with_execution_runtime_override, json, start_server, AppState, Arc, Body,
FrontdoorCorsConfig, Mutex, Request, Router, StatusCode, FRONTDOOR_MANIFEST_PATH, READYZ_PATH,
};
use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
use aether_data::repository::auth::InMemoryAuthApiKeySnapshotRepository;
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::usage::InMemoryUsageReadRepository;
use axum::Json;
use crate::data::GatewayDataState;
#[tokio::test]
async fn gateway_exposes_frontdoor_manifest_without_proxying_upstream() {
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 gateway = build_router_with_execution_runtime_override("http://127.0.0.1:19091");
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}{FRONTDOOR_MANIFEST_PATH}"))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["component"], "aether-gateway");
assert_eq!(payload["mode"], "compatibility_frontdoor");
assert_eq!(
payload["entrypoints"]["public_manifest"],
FRONTDOOR_MANIFEST_PATH
);
assert_eq!(payload["entrypoints"]["readiness"], READYZ_PATH);
assert_eq!(payload["entrypoints"]["health"], "/_gateway/health");
assert_eq!(
payload["rust_frontdoor"]["capabilities"]["public_proxy_catch_all"],
true
);
let owned_routes = payload["rust_frontdoor"]["owned_route_patterns"]
.as_array()
.expect("owned route patterns should be an array");
assert!(owned_routes
.iter()
.any(|value| value == "/v1/chat/completions"));
assert!(owned_routes.iter().any(|value| value == "/v1/messages"));
assert!(owned_routes
.iter()
.any(|value| value == "/v1/messages/count_tokens"));
assert!(owned_routes.iter().any(|value| value == "/v1/responses"));
assert!(owned_routes
.iter()
.any(|value| value == "/v1/responses/compact"));
assert!(owned_routes.iter().any(|value| value == "/health"));
assert!(owned_routes.iter().any(|value| value == "/v1/health"));
assert!(owned_routes.iter().any(|value| value == "/v1/providers"));
assert!(owned_routes
.iter()
.any(|value| value == "/v1/providers/{path...}"));
assert!(owned_routes
.iter()
.any(|value| value == "/v1/test-connection"));
assert!(owned_routes.iter().any(|value| value == "/test-connection"));
assert!(owned_routes
.iter()
.any(|value| value == "/api/public/providers"));
assert!(owned_routes
.iter()
.any(|value| value == "/api/oauth/providers"));
assert!(owned_routes
.iter()
.any(|value| value == "/api/oauth/{provider_type}/authorize"));
assert!(owned_routes
.iter()
.any(|value| value == "/api/oauth/{provider_type}/callback"));
assert!(owned_routes
.iter()
.any(|value| value == "/api/user/oauth/bindable-providers"));
assert!(owned_routes
.iter()
.any(|value| value == "/api/user/oauth/links"));
assert!(owned_routes
.iter()
.any(|value| value == "/api/user/oauth/{provider_type}/bind-token"));
assert!(owned_routes
.iter()
.any(|value| value == "/api/user/oauth/{provider_type}/bind"));
assert!(owned_routes
.iter()
.any(|value| value == "/api/user/oauth/{provider_type}"));
assert!(owned_routes
.iter()
.any(|value| value == "/api/capabilities"));
assert!(owned_routes
.iter()
.any(|value| value == "/api/public/health/api-formats"));
assert!(owned_routes
.iter()
.any(|value| value == "/api/modules/auth-status"));
assert!(owned_routes
.iter()
.any(|value| value == "/api/internal/gateway/{path...}"));
assert!(owned_routes
.iter()
.any(|value| value == "/api/internal/proxy-tunnel"));
assert!(owned_routes
.iter()
.any(|value| value == "/api/internal/tunnel/heartbeat"));
assert!(owned_routes
.iter()
.any(|value| value == "/api/internal/tunnel/node-status"));
assert!(owned_routes
.iter()
.any(|value| value == "/api/internal/tunnel/relay/{node_id}"));
assert!(owned_routes
.iter()
.any(|value| value == "/api/capabilities/user-configurable"));
assert!(owned_routes
.iter()
.any(|value| value == "/api/capabilities/model/{path...}"));
assert!(owned_routes.iter().any(|value| value == "/v1/models"));
assert!(owned_routes
.iter()
.any(|value| value == "/v1/models/{path...}"));
assert!(owned_routes.iter().any(|value| value == "/v1beta/models"));
assert!(owned_routes
.iter()
.any(|value| value == "/v1beta/models/{path...}"));
assert!(owned_routes
.iter()
.any(|value| value == "/v1beta/models/{model}:generateContent"));
assert!(owned_routes
.iter()
.any(|value| value == "/v1beta/models/{model}:streamGenerateContent"));
assert!(owned_routes
.iter()
.any(|value| value == "/v1beta/models/{model}:predictLongRunning"));
assert!(owned_routes
.iter()
.any(|value| value == "/v1beta/operations/{id}"));
assert!(owned_routes.iter().any(|value| value == "/v1/videos"));
assert!(owned_routes
.iter()
.any(|value| value == "/v1/videos/{path...}"));
assert!(owned_routes
.iter()
.any(|value| value == "/upload/v1beta/files"));
assert!(owned_routes.iter().any(|value| value == "/v1beta/files"));
assert!(owned_routes
.iter()
.any(|value| value == "/v1beta/files/{path...}"));
assert_eq!(
payload["rust_frontdoor"]["internal_gateway"]["status"],
"rust_native_control_plane"
);
assert_eq!(
payload["rust_frontdoor"]["internal_gateway"]["path_prefixes"][0],
"/api/internal/gateway"
);
assert_eq!(payload["features"]["control_api_configured"], true);
assert_eq!(payload["features"]["execution_runtime_configured"], true);
assert!(payload["features"]
.get("remote_executor_configured")
.is_none());
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_reports_local_control_plane_as_configured_without_external_control_config() {
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 gateway = build_router_with_execution_runtime_override("http://127.0.0.1:19091");
let (gateway_url, gateway_handle) = start_server(gateway).await;
let manifest = reqwest::Client::new()
.get(format!("{gateway_url}{FRONTDOOR_MANIFEST_PATH}"))
.send()
.await
.expect("manifest request should succeed");
assert_eq!(manifest.status(), StatusCode::OK);
let manifest_payload: serde_json::Value = manifest.json().await.expect("manifest should parse");
assert_eq!(manifest_payload["features"]["control_api_configured"], true);
assert_eq!(
manifest_payload["features"]["execution_runtime_configured"],
true
);
assert!(manifest_payload["features"]
.get("remote_executor_configured")
.is_none());
let health = reqwest::Client::new()
.get(format!("{gateway_url}/_gateway/health"))
.send()
.await
.expect("health request should succeed");
assert_eq!(health.status(), StatusCode::OK);
let health_payload: serde_json::Value = health.json().await.expect("health should parse");
assert_eq!(health_payload["control_api_enabled"], true);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_reports_execution_runtime_as_configured_without_execution_runtime_override() {
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 gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let manifest = reqwest::Client::new()
.get(format!("{gateway_url}{FRONTDOOR_MANIFEST_PATH}"))
.send()
.await
.expect("manifest request should succeed");
assert_eq!(manifest.status(), StatusCode::OK);
let manifest_payload: serde_json::Value = manifest.json().await.expect("manifest should parse");
assert_eq!(manifest_payload["features"]["control_api_configured"], true);
assert_eq!(
manifest_payload["features"]["execution_runtime_configured"],
true
);
assert!(manifest_payload["features"]
.get("remote_executor_configured")
.is_none());
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_cors_preflight_without_proxying_upstream() {
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 state = AppState::new()
.expect("state should build")
.with_frontdoor_cors_config(
FrontdoorCorsConfig::new(vec!["http://localhost:3000".to_string()], true)
.expect("cors config should build"),
);
let gateway = build_router_with_state(state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.request(
reqwest::Method::OPTIONS,
format!("{gateway_url}/v1/chat/completions"),
)
.header("origin", "http://localhost:3000")
.header("access-control-request-method", "POST")
.header(
"access-control-request-headers",
"authorization,content-type",
)
.send()
.await
.expect("preflight should succeed");
assert_eq!(response.status(), StatusCode::NO_CONTENT);
assert_eq!(
response
.headers()
.get("access-control-allow-origin")
.expect("allow origin header"),
"http://localhost:3000"
);
assert_eq!(
response
.headers()
.get("access-control-allow-credentials")
.expect("allow credentials header"),
"true"
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_adds_cors_headers_to_proxied_responses() {
let execution_runtime_hits = Arc::new(Mutex::new(0usize));
let execution_runtime_hits_clone = Arc::clone(&execution_runtime_hits);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |_request: Request| {
let execution_runtime_hits = Arc::clone(&execution_runtime_hits_clone);
async move {
*execution_runtime_hits.lock().expect("mutex should lock") += 1;
Json(json!({
"request_id": "trace-openai-cors-proxy-123",
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-cors-proxy-123",
"object": "chat.completion",
"model": "gpt-5-upstream",
"choices": []
}
}
}))
}
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-openai-cors")),
sample_local_openai_auth_snapshot("api-key-openai-cors-1", "user-openai-cors-1"),
)]));
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 request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let state = build_state_with_execution_runtime_override(execution_runtime_url)
.with_frontdoor_cors_config(
FrontdoorCorsConfig::new(vec!["http://localhost:3000".to_string()], true)
.expect("cors config should build"),
)
.with_data_state_for_tests(
GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
request_candidate_repository,
usage_repository,
DEVELOPMENT_ENCRYPTION_KEY,
),
);
let gateway = build_router_with_state(state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header("origin", "http://localhost:3000")
.header(http::header::AUTHORIZATION, "Bearer sk-client-openai-cors")
.header(http::header::CONTENT_TYPE, "application/json")
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
.send()
.await
.expect("proxy request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let response_headers = response.headers().clone();
assert_eq!(
response
.json::<serde_json::Value>()
.await
.expect("body should parse")["id"],
"chatcmpl-cors-proxy-123"
);
assert_eq!(
response_headers
.get("access-control-allow-origin")
.expect("allow origin header"),
"http://localhost:3000"
);
assert_eq!(
response_headers
.get("access-control-expose-headers")
.expect("expose headers header"),
"*"
);
assert_eq!(
*execution_runtime_hits.lock().expect("mutex should lock"),
1
);
gateway_handle.abort();
execution_runtime_handle.abort();
}

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,713 @@
use super::{
build_test_auth_token, json, sample_auth_session, sample_auth_user, sample_auth_wallet,
sample_provider, sample_user_usage_audit, start_auth_dashboard_gateway_with_state,
start_auth_gateway_with_builder, start_auth_gateway_with_usage_state, AppState, Arc,
GatewayDataState, InMemoryAuthApiKeySnapshotRepository, InMemoryProviderCatalogReadRepository,
InMemoryUsageReadRepository, InMemoryUserReadRepository, InMemoryWalletRepository, StatusCode,
StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot, StoredUserAuthRecord,
StoredUserExportRow, Utc,
};
fn stable_dashboard_now() -> chrono::DateTime<Utc> {
Utc::now()
.date_naive()
.and_hms_opt(12, 0, 0)
.expect("stable dashboard test time should build")
.and_utc()
}
#[tokio::test]
async fn gateway_handles_dashboard_stats_locally_without_proxying_upstream() {
let now = stable_dashboard_now();
let user = sample_auth_user(now);
let access_token = build_test_auth_token(
"access",
serde_json::Map::from_iter([
("user_id".to_string(), json!(user.id)),
("role".to_string(), json!(user.role)),
(
"created_at".to_string(),
json!(user.created_at.map(|value| value.to_rfc3339())),
),
("session_id".to_string(), json!("session-dashboard-stats")),
]),
chrono::Utc::now() + chrono::Duration::hours(1),
);
let session = sample_auth_session(
"user-auth-1",
"session-dashboard-stats",
"device-dashboard-stats",
"refresh-dashboard-stats",
now,
);
let current_usage = sample_user_usage_audit(
"usage-dashboard-stats-1",
"req-dashboard-stats-1",
"user-auth-1",
"gpt-5",
"openai",
"completed",
now - chrono::Duration::minutes(10),
);
let mut prior_usage = sample_user_usage_audit(
"usage-dashboard-stats-2",
"req-dashboard-stats-2",
"user-auth-1",
"gpt-4.1",
"openai",
"completed",
now - chrono::Duration::days(2),
);
prior_usage.actual_total_cost_usd = 0.75;
let other_usage = sample_user_usage_audit(
"usage-dashboard-stats-3",
"req-dashboard-stats-3",
"user-auth-2",
"claude-3-7",
"claude",
"completed",
now - chrono::Duration::minutes(3),
);
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
current_usage,
prior_usage,
other_usage,
]));
let user_repository = Arc::new(InMemoryUserReadRepository::seed_auth_users(vec![
user.clone()
]));
let wallet_repository = Arc::new(InMemoryWalletRepository::seed(vec![sample_auth_wallet(
"user-auth-1",
now,
)]));
let auth_repository = Arc::new(
InMemoryAuthApiKeySnapshotRepository::seed(
Vec::<(Option<String>, StoredAuthApiKeySnapshot)>::new(),
)
.with_export_records(vec![
StoredAuthApiKeyExportRecord::new(
"user-auth-1".to_string(),
"user-key-1".to_string(),
"hash-user-key-1".to_string(),
None,
Some("primary".to_string()),
Some(json!(["openai"])),
Some(json!(["openai:chat"])),
Some(json!(["gpt-5"])),
Some(60),
Some(5),
None,
true,
None,
false,
5,
1.5,
false,
)
.expect("api key export should build"),
StoredAuthApiKeyExportRecord::new(
"user-auth-1".to_string(),
"user-key-2".to_string(),
"hash-user-key-2".to_string(),
None,
Some("secondary".to_string()),
Some(json!(["openai"])),
Some(json!(["openai:chat"])),
Some(json!(["gpt-4.1"])),
Some(60),
Some(5),
None,
false,
None,
false,
1,
0.5,
false,
)
.expect("api key export should build"),
]),
);
let (gateway_url, upstream_hits, gateway_handle, upstream_handle) =
start_auth_gateway_with_builder(|| {
let data_state = GatewayDataState::with_user_wallet_and_usage_for_tests(
user_repository,
wallet_repository,
usage_repository,
)
.with_auth_api_key_reader(auth_repository);
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data_state)
.with_auth_sessions_for_tests([session])
})
.await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/dashboard/stats?days=3"))
.header("authorization", format!("Bearer {access_token}"))
.header("x-client-device-id", "device-dashboard-stats")
.header("user-agent", "AetherTest/1.0")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["today"]["requests"], 1);
assert_eq!(payload["today"]["tokens"], 150);
assert_eq!(payload["api_keys"]["total"], 2);
assert_eq!(payload["api_keys"]["active"], 1);
assert_eq!(payload["token_breakdown"]["input"], 240);
assert_eq!(payload["token_breakdown"]["output"], 60);
assert_eq!(payload["token_breakdown"]["cache_creation"], 20);
assert_eq!(payload["token_breakdown"]["cache_read"], 30);
assert_eq!(payload["monthly_cost"], json!(2.5));
assert_eq!(payload["stats"].as_array().map(Vec::len), Some(4));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_dashboard_stats_locally_without_proxying_upstream() {
let now = Utc::now();
let admin = StoredUserAuthRecord::new(
"admin-auth-1".to_string(),
Some("admin@example.com".to_string()),
true,
"admin".to_string(),
Some("$2y$10$.OBQfixAECpsb8V/VS3csOMf00x2E/jD/gnud20t6RG0yiQosyOZ2".to_string()),
"admin".to_string(),
"local".to_string(),
None,
None,
None,
true,
false,
Some(now),
Some(now),
)
.expect("admin auth user should build");
let access_token = build_test_auth_token(
"access",
serde_json::Map::from_iter([
("user_id".to_string(), json!(admin.id)),
("role".to_string(), json!(admin.role)),
(
"created_at".to_string(),
json!(admin.created_at.map(|value| value.to_rfc3339())),
),
(
"session_id".to_string(),
json!("session-dashboard-stats-admin"),
),
]),
now + chrono::Duration::hours(1),
);
let session = sample_auth_session(
"admin-auth-1",
"session-dashboard-stats-admin",
"device-dashboard-stats-admin",
"refresh-dashboard-stats-admin",
now,
);
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
sample_user_usage_audit(
"usage-dashboard-admin-1",
"req-dashboard-admin-1",
"user-auth-1",
"gpt-5",
"openai",
"completed",
now - chrono::Duration::minutes(10),
),
sample_user_usage_audit(
"usage-dashboard-admin-2",
"req-dashboard-admin-2",
"user-auth-2",
"claude-3-7",
"claude",
"completed",
now - chrono::Duration::minutes(5),
),
]));
let user_repository = Arc::new(
InMemoryUserReadRepository::seed_auth_users(vec![admin.clone()]).with_export_users(vec![
StoredUserExportRow::new(
"user-auth-1".to_string(),
Some("alice@example.com".to_string()),
true,
"alice".to_string(),
Some("hash".to_string()),
"user".to_string(),
"local".to_string(),
Some(json!(["openai"])),
Some(json!(["openai:chat"])),
Some(json!(["gpt-5"])),
Some(60),
None,
true,
)
.expect("user export row should build"),
StoredUserExportRow::new(
"user-auth-2".to_string(),
Some("bob@example.com".to_string()),
true,
"bob".to_string(),
Some("hash".to_string()),
"user".to_string(),
"local".to_string(),
Some(json!(["anthropic"])),
Some(json!(["anthropic:messages"])),
Some(json!(["claude-3-7"])),
Some(30),
None,
false,
)
.expect("user export row should build"),
]),
);
let wallet_repository = Arc::new(InMemoryWalletRepository::seed(vec![sample_auth_wallet(
"admin-auth-1",
now,
)]));
let auth_repository = Arc::new(
InMemoryAuthApiKeySnapshotRepository::seed(
Vec::<(Option<String>, StoredAuthApiKeySnapshot)>::new(),
)
.with_export_records(vec![
StoredAuthApiKeyExportRecord::new(
"user-auth-1".to_string(),
"user-key-1".to_string(),
"hash-user-key-1".to_string(),
None,
Some("primary".to_string()),
Some(json!(["openai"])),
Some(json!(["openai:chat"])),
Some(json!(["gpt-5"])),
Some(60),
Some(5),
None,
true,
None,
false,
5,
1.5,
false,
)
.expect("api key export should build"),
StoredAuthApiKeyExportRecord::new(
"user-auth-2".to_string(),
"user-key-2".to_string(),
"hash-user-key-2".to_string(),
None,
Some("secondary".to_string()),
Some(json!(["anthropic"])),
Some(json!(["anthropic:messages"])),
Some(json!(["claude-3-7"])),
Some(60),
Some(5),
None,
false,
None,
false,
1,
0.5,
false,
)
.expect("api key export should build"),
StoredAuthApiKeyExportRecord::new(
"admin-auth-1".to_string(),
"standalone-key-1".to_string(),
"hash-standalone-key-1".to_string(),
None,
Some("standalone".to_string()),
Some(json!(["openai"])),
Some(json!(["openai:chat"])),
Some(json!(["gpt-5"])),
Some(60),
Some(5),
None,
true,
None,
false,
3,
0.75,
true,
)
.expect("standalone api key export should build"),
]),
);
let (gateway_url, upstream_hits, gateway_handle, upstream_handle) =
start_auth_gateway_with_builder(|| {
let data_state = GatewayDataState::with_user_wallet_and_usage_for_tests(
user_repository,
wallet_repository,
usage_repository,
)
.with_auth_api_key_reader(auth_repository);
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data_state)
.with_auth_sessions_for_tests([session])
})
.await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/dashboard/stats"))
.header("authorization", format!("Bearer {access_token}"))
.header("x-client-device-id", "device-dashboard-stats-admin")
.header("user-agent", "AetherTest/1.0")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["users"]["total"], 2);
assert_eq!(payload["users"]["active"], 1);
assert_eq!(payload["api_keys"]["total"], 3);
assert_eq!(payload["api_keys"]["active"], 2);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_dashboard_daily_stats_locally_without_proxying_upstream() {
let now = stable_dashboard_now();
let admin = StoredUserAuthRecord::new(
"admin-auth-1".to_string(),
Some("admin@example.com".to_string()),
true,
"admin".to_string(),
Some("$2y$10$.OBQfixAECpsb8V/VS3csOMf00x2E/jD/gnud20t6RG0yiQosyOZ2".to_string()),
"admin".to_string(),
"local".to_string(),
None,
None,
None,
true,
false,
Some(now),
Some(now),
)
.expect("admin auth user should build");
let access_token = build_test_auth_token(
"access",
serde_json::Map::from_iter([
("user_id".to_string(), json!(admin.id)),
("role".to_string(), json!(admin.role)),
(
"created_at".to_string(),
json!(admin.created_at.map(|value| value.to_rfc3339())),
),
(
"session_id".to_string(),
json!("session-dashboard-daily-stats"),
),
]),
chrono::Utc::now() + chrono::Duration::hours(1),
);
let session = sample_auth_session(
"admin-auth-1",
"session-dashboard-daily-stats",
"device-dashboard-daily-stats",
"refresh-dashboard-daily-stats",
now,
);
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
sample_user_usage_audit(
"usage-dashboard-daily-1",
"req-dashboard-daily-1",
"user-auth-1",
"gpt-5",
"openai",
"completed",
now - chrono::Duration::hours(1),
),
sample_user_usage_audit(
"usage-dashboard-daily-2",
"req-dashboard-daily-2",
"user-auth-2",
"claude-3-7",
"claude",
"completed",
now - chrono::Duration::hours(2),
),
sample_user_usage_audit(
"usage-dashboard-daily-3",
"req-dashboard-daily-3",
"user-auth-3",
"gpt-5",
"openai",
"completed",
now - chrono::Duration::days(1) - chrono::Duration::hours(2),
),
]));
let (gateway_url, upstream_hits, gateway_handle, upstream_handle) =
start_auth_gateway_with_builder(|| {
let user_repository = Arc::new(InMemoryUserReadRepository::seed_auth_users(vec![
admin.clone(),
]));
let wallet_repository =
Arc::new(InMemoryWalletRepository::seed(vec![sample_auth_wallet(
"admin-auth-1",
now,
)]));
let data_state = GatewayDataState::with_user_wallet_and_usage_for_tests(
user_repository,
wallet_repository,
usage_repository,
);
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data_state)
.with_auth_sessions_for_tests([session])
})
.await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/dashboard/daily-stats?days=2"))
.header("authorization", format!("Bearer {access_token}"))
.header("x-client-device-id", "device-dashboard-daily-stats")
.header("user-agent", "AetherTest/1.0")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
let daily_stats = payload["daily_stats"]
.as_array()
.expect("daily stats should be array");
assert_eq!(daily_stats.len(), 2);
assert_eq!(
daily_stats[0]["date"],
json!((now - chrono::Duration::days(1)).date_naive().to_string())
);
assert_eq!(daily_stats[0]["requests"], 1);
assert_eq!(daily_stats[0]["unique_providers"], 1);
assert_eq!(daily_stats[1]["date"], json!(now.date_naive().to_string()));
assert_eq!(daily_stats[1]["requests"], 2);
assert_eq!(daily_stats[1]["unique_models"], 2);
assert_eq!(daily_stats[1]["unique_providers"], 2);
assert_eq!(
daily_stats[1]["model_breakdown"].as_array().map(Vec::len),
Some(2)
);
let model_summary = payload["model_summary"]
.as_array()
.expect("model summary should exist");
assert_eq!(model_summary.len(), 2);
assert_eq!(model_summary[0]["model"], "gpt-5");
assert_eq!(model_summary[0]["requests"], 2);
let provider_summary = payload["provider_summary"]
.as_array()
.expect("provider summary should exist");
assert_eq!(provider_summary.len(), 2);
assert_eq!(provider_summary[0]["provider"], "openai");
assert_eq!(provider_summary[0]["requests"], 2);
assert_eq!(provider_summary[1]["provider"], "claude");
assert_eq!(provider_summary[1]["requests"], 1);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_dashboard_recent_requests_locally_without_proxying_upstream() {
let now = Utc::now();
let user = sample_auth_user(now);
let access_token = build_test_auth_token(
"access",
serde_json::Map::from_iter([
("user_id".to_string(), json!(user.id)),
("role".to_string(), json!(user.role)),
(
"created_at".to_string(),
json!(user.created_at.map(|value| value.to_rfc3339())),
),
("session_id".to_string(), json!("session-dashboard-recent")),
]),
now + chrono::Duration::hours(1),
);
let session = sample_auth_session(
"user-auth-1",
"session-dashboard-recent",
"device-dashboard-recent",
"refresh-dashboard-recent",
now,
);
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
sample_user_usage_audit(
"usage-dashboard-1",
"req-dashboard-1",
"user-auth-1",
"gpt-5",
"OpenAI",
"completed",
now - chrono::Duration::minutes(5),
),
sample_user_usage_audit(
"usage-dashboard-2",
"req-dashboard-2",
"user-auth-2",
"claude-3-7",
"Anthropic",
"completed",
now - chrono::Duration::minutes(2),
),
]));
let (gateway_url, upstream_hits, gateway_handle, upstream_handle) =
start_auth_gateway_with_usage_state(
user,
sample_auth_wallet("user-auth-1", now),
[session],
usage_repository,
)
.await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/dashboard/recent-requests?limit=5"
))
.header("authorization", format!("Bearer {access_token}"))
.header("x-client-device-id", "device-dashboard-recent")
.header("user-agent", "AetherTest/1.0")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
let requests = payload["requests"].as_array().expect("array");
assert_eq!(requests.len(), 1);
assert_eq!(requests[0]["id"], "usage-dashboard-1");
assert_eq!(requests[0]["user"], "alice");
assert_eq!(requests[0]["model"], "gpt-5");
assert_eq!(requests[0]["tokens"], 150);
assert_eq!(requests[0]["is_stream"], false);
assert!(requests[0]["time"].as_str().is_some());
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_dashboard_provider_status_locally_without_proxying_upstream() {
let now = Utc::now();
let user = sample_auth_user(now);
let access_token = build_test_auth_token(
"access",
serde_json::Map::from_iter([
("user_id".to_string(), json!(user.id)),
("role".to_string(), json!(user.role)),
(
"created_at".to_string(),
json!(user.created_at.map(|value| value.to_rfc3339())),
),
(
"session_id".to_string(),
json!("session-dashboard-provider-status"),
),
]),
now + chrono::Duration::hours(1),
);
let session = sample_auth_session(
"user-auth-1",
"session-dashboard-provider-status",
"device-dashboard-provider-status",
"refresh-dashboard-provider-status",
now,
);
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
sample_user_usage_audit(
"usage-provider-1",
"req-provider-1",
"user-auth-1",
"gpt-5",
"openai",
"completed",
now - chrono::Duration::hours(1),
),
sample_user_usage_audit(
"usage-provider-2",
"req-provider-2",
"user-auth-2",
"gpt-5",
"openai",
"completed",
now - chrono::Duration::hours(2),
),
sample_user_usage_audit(
"usage-provider-3",
"req-provider-3",
"user-auth-1",
"claude-3-7",
"claude",
"completed",
now - chrono::Duration::hours(3),
),
sample_user_usage_audit(
"usage-provider-4",
"req-provider-4",
"user-auth-1",
"claude-3-7",
"claude",
"completed",
now - chrono::Duration::hours(30),
),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider("provider-openai", "openai", 10),
sample_provider("provider-claude", "claude", 20),
sample_provider("provider-gemini", "gemini", 30),
],
vec![],
vec![],
));
let (gateway_url, upstream_hits, gateway_handle, upstream_handle) =
start_auth_dashboard_gateway_with_state(
user,
sample_auth_wallet("user-auth-1", now),
[session],
usage_repository,
provider_catalog_repository,
)
.await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/dashboard/provider-status"))
.header("authorization", format!("Bearer {access_token}"))
.header("x-client-device-id", "device-dashboard-provider-status")
.header("user-agent", "AetherTest/1.0")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
let providers = payload["providers"].as_array().expect("array");
assert_eq!(providers.len(), 3);
assert_eq!(providers[0]["name"], "openai");
assert_eq!(providers[0]["requests"], 2);
assert_eq!(providers[1]["name"], "claude");
assert_eq!(providers[1]["requests"], 1);
assert_eq!(providers[2]["name"], "gemini");
assert_eq!(providers[2]["requests"], 0);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}