mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-10 19:29:50 +08:00
722 lines
29 KiB
Rust
722 lines
29 KiB
Rust
|
use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine as _};
|
||
|
|
|
||
|
|
use super::*;
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn gateway_executes_gemini_files_upload_via_local_decision_gate_without_python_plan_or_decision(
|
||
|
|
) {
|
||
|
|
#[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",
|
||
|
|
"executor_candidate": true,
|
||
|
|
"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_test_remote_execution_runtime(
|
||
|
|
upstream_url.clone(),
|
||
|
|
execution_runtime_url.clone(),
|
||
|
|
)
|
||
|
|
.with_data_state_for_tests(
|
||
|
|
crate::gateway::gateway_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_without_python_plan_or_decision(
|
||
|
|
) {
|
||
|
|
#[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",
|
||
|
|
"executor_candidate": true,
|
||
|
|
"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_test_remote_execution_runtime(
|
||
|
|
upstream_url.clone(),
|
||
|
|
execution_runtime_url.clone(),
|
||
|
|
)
|
||
|
|
.with_data_state_for_tests(
|
||
|
|
crate::gateway::gateway_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_without_python_plan_or_decision(
|
||
|
|
) {
|
||
|
|
#[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",
|
||
|
|
"executor_candidate": true,
|
||
|
|
"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_test_remote_execution_runtime(
|
||
|
|
upstream_url.clone(),
|
||
|
|
execution_runtime_url.clone(),
|
||
|
|
)
|
||
|
|
.with_data_state_for_tests(
|
||
|
|
crate::gateway::gateway_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();
|
||
|
|
}
|