mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-10 03:09:50 +08:00
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
@@ -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();
|
||||
}
|
||||
Reference in New Issue
Block a user