mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-10 11:19: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,714 @@
|
||||
use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine as _};
|
||||
|
||||
use super::{
|
||||
any, build_router_with_state, build_state_with_execution_runtime_override, hash_api_key, json,
|
||||
sample_auth_snapshot, sample_files_candidate_row, sample_files_provider_catalog_endpoint,
|
||||
sample_files_provider_catalog_key, sample_files_provider_catalog_provider, start_server,
|
||||
to_bytes, Arc, Body, InMemoryAuthApiKeySnapshotRepository,
|
||||
InMemoryMinimalCandidateSelectionReadRepository, InMemoryProviderCatalogReadRepository,
|
||||
InMemoryRequestCandidateRepository, Json, Mutex, Request, RequestCandidateReadRepository,
|
||||
RequestCandidateStatus, Router, StatusCode, DEVELOPMENT_ENCRYPTION_KEY, TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_executes_gemini_files_upload_via_local_decision_gate_with_local_planning_only() {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeSyncRequest {
|
||||
method: String,
|
||||
url: String,
|
||||
auth_header_value: String,
|
||||
content_type: String,
|
||||
body_bytes_b64: String,
|
||||
endpoint_tag: String,
|
||||
proxy_node_id: String,
|
||||
tls_profile: String,
|
||||
}
|
||||
|
||||
let decision_hits = Arc::new(Mutex::new(0usize));
|
||||
let decision_hits_clone = Arc::clone(&decision_hits);
|
||||
let plan_hits = Arc::new(Mutex::new(0usize));
|
||||
let plan_hits_clone = Arc::clone(&plan_hits);
|
||||
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeSyncRequest>));
|
||||
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
|
||||
let report_hits = Arc::new(Mutex::new(0usize));
|
||||
let report_hits_clone = Arc::clone(&report_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/resolve",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"action": "proxy_public",
|
||||
"route_class": "ai_public",
|
||||
"route_family": "gemini",
|
||||
"route_kind": "files",
|
||||
"auth_endpoint_signature": "gemini:chat",
|
||||
"execution_runtime_candidate": true,
|
||||
"auth_context": {
|
||||
"user_id": "user-files-upload-local-123",
|
||||
"api_key_id": "key-files-upload-local-123",
|
||||
"access_allowed": true
|
||||
},
|
||||
"public_path": "/upload/v1beta/files"
|
||||
}))
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/decision-sync",
|
||||
any(move |_request: Request| {
|
||||
let decision_hits_inner = Arc::clone(&decision_hits_clone);
|
||||
async move {
|
||||
*decision_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({
|
||||
"action": "execution_runtime_sync_decision",
|
||||
"decision_kind": "unexpected_remote_decision"
|
||||
}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/plan-sync",
|
||||
any(move |_request: Request| {
|
||||
let plan_hits_inner = Arc::clone(&plan_hits_clone);
|
||||
async move {
|
||||
*plan_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({
|
||||
"action": "execution_runtime_sync",
|
||||
"plan_kind": "unexpected_plan_fallback"
|
||||
}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/report-sync",
|
||||
any(move |_request: Request| {
|
||||
let report_hits_inner = Arc::clone(&report_hits_clone);
|
||||
async move {
|
||||
*report_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"ok": true}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/upload/v1beta/files",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
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(),
|
||||
auth_header_value: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-goog-api-key"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
content_type: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("content-type"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
body_bytes_b64: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("body_bytes_b64"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
endpoint_tag: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-endpoint-tag"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
proxy_node_id: payload
|
||||
.get("proxy")
|
||||
.and_then(|value| value.get("node_id"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
tls_profile: payload
|
||||
.get("tls_profile")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
});
|
||||
Json(json!({
|
||||
"request_id": "trace-gemini-files-upload-local-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"file": {
|
||||
"name": "files/uploaded-local-123"
|
||||
}
|
||||
}
|
||||
},
|
||||
"telemetry": {
|
||||
"elapsed_ms": 19
|
||||
}
|
||||
}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("client-files-upload-local-key")),
|
||||
sample_auth_snapshot("key-files-upload-local-123", "user-files-upload-local-123"),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_files_candidate_row(),
|
||||
]));
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let mut provider = sample_files_provider_catalog_provider();
|
||||
provider.proxy = Some(serde_json::json!({"url":"http://provider-proxy.internal:8080"}));
|
||||
let mut endpoint = sample_files_provider_catalog_endpoint();
|
||||
endpoint.custom_path = Some("/custom/upload/v1beta/files".to_string());
|
||||
endpoint.header_rules = Some(
|
||||
serde_json::json!([{"action":"set","key":"x-endpoint-tag","value":"gemini-files-upload-local"}]),
|
||||
);
|
||||
let mut key = sample_files_provider_catalog_key();
|
||||
key.proxy = Some(
|
||||
serde_json::json!({"enabled": true, "node_id":"proxy-node-gemini-files-upload-local"}),
|
||||
);
|
||||
key.fingerprint = Some(serde_json::json!({"tls_profile":"chrome_136"}));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![provider],
|
||||
vec![endpoint],
|
||||
vec![key],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state =
|
||||
build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{gateway_url}/upload/v1beta/files?uploadType=resumable&key=client-files-upload-local-key"
|
||||
))
|
||||
.header("x-goog-api-key", "client-header-key")
|
||||
.header(http::header::CONTENT_TYPE, "application/octet-stream")
|
||||
.header(TRACE_ID_HEADER, "trace-gemini-files-upload-local-123")
|
||||
.body("upload-local-bytes")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response
|
||||
.json::<serde_json::Value>()
|
||||
.await
|
||||
.expect("json should parse"),
|
||||
json!({"file": {"name": "files/uploaded-local-123"}})
|
||||
);
|
||||
|
||||
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/custom/upload/v1beta/files?uploadType=resumable"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.auth_header_value,
|
||||
"sk-upstream-gemini-files"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.content_type,
|
||||
"application/octet-stream"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.endpoint_tag,
|
||||
"gemini-files-upload-local"
|
||||
);
|
||||
assert_eq!(
|
||||
BASE64_STANDARD
|
||||
.decode(seen_execution_runtime_request.body_bytes_b64)
|
||||
.expect("execution runtime body should decode"),
|
||||
b"upload-local-bytes"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.proxy_node_id,
|
||||
"proxy-node-gemini-files-upload-local"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.tls_profile, "chrome_136");
|
||||
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-gemini-files-upload-local-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert_eq!(stored_candidates.len(), 1);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
|
||||
|
||||
assert_eq!(*report_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_executes_gemini_files_list_via_local_decision_gate_with_local_planning_only() {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeSyncRequest {
|
||||
method: String,
|
||||
url: String,
|
||||
auth_header_value: String,
|
||||
}
|
||||
|
||||
let decision_hits = Arc::new(Mutex::new(0usize));
|
||||
let decision_hits_clone = Arc::clone(&decision_hits);
|
||||
let plan_hits = Arc::new(Mutex::new(0usize));
|
||||
let plan_hits_clone = Arc::clone(&plan_hits);
|
||||
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeSyncRequest>));
|
||||
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
|
||||
let report_hits = Arc::new(Mutex::new(0usize));
|
||||
let report_hits_clone = Arc::clone(&report_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/resolve",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"action": "proxy_public",
|
||||
"route_class": "ai_public",
|
||||
"route_family": "gemini",
|
||||
"route_kind": "files",
|
||||
"auth_endpoint_signature": "gemini:chat",
|
||||
"execution_runtime_candidate": true,
|
||||
"auth_context": {
|
||||
"user_id": "user-files-list-local-123",
|
||||
"api_key_id": "key-files-list-local-123",
|
||||
"access_allowed": true
|
||||
},
|
||||
"public_path": "/v1beta/files"
|
||||
}))
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/decision-sync",
|
||||
any(move |_request: Request| {
|
||||
let decision_hits_inner = Arc::clone(&decision_hits_clone);
|
||||
async move {
|
||||
*decision_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"action": "proxy_public"}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/plan-sync",
|
||||
any(move |_request: Request| {
|
||||
let plan_hits_inner = Arc::clone(&plan_hits_clone);
|
||||
async move {
|
||||
*plan_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"action": "proxy_public"}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/report-sync",
|
||||
any(move |_request: Request| {
|
||||
let report_hits_inner = Arc::clone(&report_hits_clone);
|
||||
async move {
|
||||
*report_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
StatusCode::NO_CONTENT
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1beta/files",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
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(),
|
||||
auth_header_value: 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-files-list-local-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"files": []
|
||||
}
|
||||
},
|
||||
"telemetry": {
|
||||
"elapsed_ms": 14
|
||||
}
|
||||
}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("client-files-list-local-key")),
|
||||
sample_auth_snapshot("key-files-list-local-123", "user-files-list-local-123"),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_files_candidate_row(),
|
||||
]));
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_files_provider_catalog_provider()],
|
||||
vec![sample_files_provider_catalog_endpoint()],
|
||||
vec![sample_files_provider_catalog_key()],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state =
|
||||
build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/v1beta/files?pageSize=20&key=client-files-list-local-key"
|
||||
))
|
||||
.header("x-goog-api-key", "client-header-key")
|
||||
.header(TRACE_ID_HEADER, "trace-gemini-files-list-local-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response.text().await.expect("body should read"),
|
||||
"{\"files\":[]}"
|
||||
);
|
||||
|
||||
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, "GET");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.url,
|
||||
"https://generativelanguage.googleapis.com/v1beta/files?pageSize=20"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.auth_header_value,
|
||||
"sk-upstream-gemini-files"
|
||||
);
|
||||
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-gemini-files-list-local-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert_eq!(stored_candidates.len(), 1);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
|
||||
|
||||
assert_eq!(*report_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_executes_gemini_files_delete_via_local_decision_gate_with_local_planning_only() {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeSyncRequest {
|
||||
method: String,
|
||||
url: String,
|
||||
auth_header_value: String,
|
||||
}
|
||||
|
||||
let decision_hits = Arc::new(Mutex::new(0usize));
|
||||
let decision_hits_clone = Arc::clone(&decision_hits);
|
||||
let plan_hits = Arc::new(Mutex::new(0usize));
|
||||
let plan_hits_clone = Arc::clone(&plan_hits);
|
||||
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeSyncRequest>));
|
||||
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
|
||||
let report_hits = Arc::new(Mutex::new(0usize));
|
||||
let report_hits_clone = Arc::clone(&report_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/resolve",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"action": "proxy_public",
|
||||
"route_class": "ai_public",
|
||||
"route_family": "gemini",
|
||||
"route_kind": "files",
|
||||
"auth_endpoint_signature": "gemini:chat",
|
||||
"execution_runtime_candidate": true,
|
||||
"auth_context": {
|
||||
"user_id": "user-files-delete-local-123",
|
||||
"api_key_id": "key-files-delete-local-123",
|
||||
"access_allowed": true
|
||||
},
|
||||
"public_path": "/v1beta/files/files/abc-123"
|
||||
}))
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/decision-sync",
|
||||
any(move |_request: Request| {
|
||||
let decision_hits_inner = Arc::clone(&decision_hits_clone);
|
||||
async move {
|
||||
*decision_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"action": "proxy_public"}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/plan-sync",
|
||||
any(move |_request: Request| {
|
||||
let plan_hits_inner = Arc::clone(&plan_hits_clone);
|
||||
async move {
|
||||
*plan_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"action": "proxy_public"}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/report-sync",
|
||||
any(move |_request: Request| {
|
||||
let report_hits_inner = Arc::clone(&report_hits_clone);
|
||||
async move {
|
||||
*report_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
StatusCode::NO_CONTENT
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1beta/files/files/abc-123",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
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(),
|
||||
auth_header_value: 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-files-delete-local-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {}
|
||||
},
|
||||
"telemetry": {
|
||||
"elapsed_ms": 17
|
||||
}
|
||||
}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("client-files-delete-local-key")),
|
||||
sample_auth_snapshot("key-files-delete-local-123", "user-files-delete-local-123"),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_files_candidate_row(),
|
||||
]));
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_files_provider_catalog_provider()],
|
||||
vec![sample_files_provider_catalog_endpoint()],
|
||||
vec![sample_files_provider_catalog_key()],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state =
|
||||
build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.delete(format!(
|
||||
"{gateway_url}/v1beta/files/files/abc-123?key=client-files-delete-local-key"
|
||||
))
|
||||
.header("x-goog-api-key", "client-header-key")
|
||||
.header(TRACE_ID_HEADER, "trace-gemini-files-delete-local-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response
|
||||
.json::<serde_json::Value>()
|
||||
.await
|
||||
.expect("json 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, "DELETE");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.url,
|
||||
"https://generativelanguage.googleapis.com/v1beta/files/files/abc-123"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.auth_header_value,
|
||||
"sk-upstream-gemini-files"
|
||||
);
|
||||
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-gemini-files-delete-local-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert_eq!(stored_candidates.len(), 1);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
|
||||
|
||||
assert_eq!(*report_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
Reference in New Issue
Block a user