i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-04 16:37:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge upstream main into feat/500-api-key-ip-whitelist

Browse Source
This commit is contained in:
RWDai
2026-05-20 10:26:56 +08:00
parent d19bf71343 ed75ae6d56
commit bdc848b19e
501 changed files with 47013 additions and 3667 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/tests/ai_execute/stream/image.rs
+449
View File
@@ -13,6 +13,7 @@ use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadReposi
use aether_data_contracts::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};
use aether_data_contracts::repository::candidates::RequestCandidateReadRepository;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
@@ -741,3 +742,451 @@ async fn gateway_bridges_codex_image_sync_json_to_streaming_image_sse() {
execution_runtime_handle.abort();
refresh_handle.abort();
}
#[derive(Debug, Clone)]
struct SeenImageBridgeExecutionPlan {
trace_id: String,
client_api_format: String,
provider_api_format: String,
url: String,
plan_stream: bool,
auth_header: String,
chatgpt_web_marker: String,
body_json: serde_json::Value,
}
fn image_bridge_hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());
format!("{:x}", hasher.finalize())
}
fn image_bridge_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
user_id.to_string(),
"alice".to_string(),
Some("[email protected]".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
None,
Some(serde_json::json!([
"openai:chat",
"openai:responses",
"openai:image"
])),
Some(serde_json::json!(["gpt-image-2"])),
api_key_id.to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(4_102_444_800_i64),
None,
Some(serde_json::json!([
"openai:chat",
"openai:responses",
"openai:image"
])),
Some(serde_json::json!(["gpt-image-2"])),
)
.expect("auth snapshot should build")
}
fn image_bridge_candidate_row(
prefix: &str,
provider_name: &str,
provider_type: &str,
) -> StoredMinimalCandidateSelectionRow {
let key_auth_type = if provider_type == "chatgpt_web" {
"bearer"
} else {
"api_key"
};
StoredMinimalCandidateSelectionRow {
provider_id: format!("provider-{prefix}"),
provider_name: provider_name.to_string(),
provider_type: provider_type.to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: format!("endpoint-{prefix}"),
endpoint_api_format: "openai:image".to_string(),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("image".to_string()),
endpoint_is_active: true,
key_id: format!("key-{prefix}"),
key_name: "prod".to_string(),
key_auth_type: key_auth_type.to_string(),
key_is_active: true,
key_api_formats: Some(vec!["openai:image".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 5,
key_global_priority_by_format: Some(serde_json::json!({"openai:image": 1})),
model_id: format!("model-{prefix}"),
global_model_id: format!("global-model-{prefix}"),
global_model_name: "gpt-image-2".to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(false),
model_provider_model_name: "gpt-image-2".to_string(),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: "gpt-image-2".to_string(),
priority: 1,
api_formats: Some(vec!["openai:image".to_string()]),
endpoint_ids: None,
}]),
model_supports_streaming: Some(false),
model_is_active: true,
model_is_available: true,
}
}
fn image_bridge_provider_catalog_provider(
prefix: &str,
provider_name: &str,
provider_type: &str,
base_url: &str,
) -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
format!("provider-{prefix}"),
provider_name.to_string(),
Some(base_url.to_string()),
provider_type.to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
false,
None,
Some(2),
None,
Some(20.0),
None,
None,
)
}
fn image_bridge_provider_catalog_endpoint(
prefix: &str,
base_url: &str,
) -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
format!("endpoint-{prefix}"),
format!("provider-{prefix}"),
"openai:image".to_string(),
Some("openai".to_string()),
Some("image".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
base_url.to_string(),
None,
None,
Some(2),
None,
None,
None,
None,
)
.expect("endpoint transport should build")
}
fn image_bridge_provider_catalog_key(
prefix: &str,
provider_type: &str,
) -> StoredProviderCatalogKey {
let auth_type = if provider_type == "chatgpt_web" {
"bearer"
} else {
"api_key"
};
StoredProviderCatalogKey::new(
format!("key-{prefix}"),
format!("provider-{prefix}"),
"prod".to_string(),
auth_type.to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["openai:image"])),
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-upstream-image-bridge")
.expect("api key should encrypt"),
None,
None,
Some(serde_json::json!({"openai:image": 1})),
None,
None,
None,
None,
)
.expect("key transport should build")
}
async fn start_image_bridge_gateway(
prefix: &str,
provider_name: &str,
provider_type: &str,
base_url: &str,
execution_runtime_url: String,
) -> (
String,
tokio::task::JoinHandle<()>,
String,
Arc<InMemoryRequestCandidateRepository>,
) {
let client_api_key = format!("sk-client-{prefix}");
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(image_bridge_hash_api_key(&client_api_key)),
image_bridge_auth_snapshot(&format!("api-key-{prefix}"), &format!("user-{prefix}")),
)]));
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
image_bridge_candidate_row(prefix, provider_name, provider_type),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![image_bridge_provider_catalog_provider(
prefix,
provider_name,
provider_type,
base_url,
)],
vec![image_bridge_provider_catalog_endpoint(prefix, base_url)],
vec![image_bridge_provider_catalog_key(prefix, provider_type)],
));
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url)
.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;
(
gateway_url,
gateway_handle,
client_api_key,
request_candidate_repository,
)
}
fn capture_image_bridge_execution_plan(
parts: http::request::Parts,
payload: serde_json::Value,
) -> SeenImageBridgeExecutionPlan {
SeenImageBridgeExecutionPlan {
trace_id: parts
.headers
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
client_api_format: payload
.get("client_api_format")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
provider_api_format: payload
.get("provider_api_format")
.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(),
plan_stream: payload
.get("stream")
.and_then(|value| value.as_bool())
.unwrap_or(false),
auth_header: payload
.get("headers")
.and_then(|value| value.get("authorization"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
chatgpt_web_marker: payload
.get("headers")
.and_then(|value| value.get("x-aether-chatgpt-web-image"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
body_json: payload
.get("body")
.and_then(|value| value.get("json_body"))
.cloned()
.unwrap_or(serde_json::Value::Null),
}
}
fn image_bridge_execution_runtime(
seen_execution_plan: Arc<Mutex<Option<SeenImageBridgeExecutionPlan>>>,
) -> Router {
Router::new().route(
"/v1/execute/stream",
any(move |request: Request| {
let seen_execution_plan_inner = Arc::clone(&seen_execution_plan);
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_plan_inner.lock().expect("mutex should lock") =
Some(capture_image_bridge_execution_plan(parts, payload));
let frames = concat!(
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.output_item.done\\ndata: {\\\"type\\\":\\\"response.output_item.done\\\",\\\"output_index\\\":0,\\\"item\\\":{\\\"id\\\":\\\"ig_bridge_123\\\",\\\"type\\\":\\\"image_generation_call\\\",\\\"result\\\":\\\"aGVsbG8=\\\",\\\"output_format\\\":\\\"png\\\"}}\\n\\n\"}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.completed\\ndata: {\\\"type\\\":\\\"response.completed\\\",\\\"response\\\":{\\\"id\\\":\\\"resp_bridge_123\\\",\\\"object\\\":\\\"response\\\",\\\"model\\\":\\\"gpt-image-2\\\",\\\"status\\\":\\\"completed\\\",\\\"output\\\":[]}}\\n\\n\"}}\n",
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
);
let mut response = http::Response::builder()
.status(StatusCode::OK)
.body(Body::from(frames))
.expect("response should build");
response.headers_mut().insert(
http::header::CONTENT_TYPE,
http::HeaderValue::from_static("application/x-ndjson"),
);
response
}
}),
)
}
#[tokio::test]
async fn gateway_routes_openai_chat_stream_image_intent_to_openai_image_plan_without_streaming_support(
) {
let seen_execution_plan = Arc::new(Mutex::new(None::<SeenImageBridgeExecutionPlan>));
let execution_runtime = image_bridge_execution_runtime(Arc::clone(&seen_execution_plan));
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let (gateway_url, gateway_handle, client_api_key, request_candidate_repository) =
start_image_bridge_gateway(
"chat-stream-image-bridge",
"image-provider",
"custom",
"https://images.example.com",
execution_runtime_url,
)
.await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(http::header::AUTHORIZATION, format!("Bearer {client_api_key}"))
.header(TRACE_ID_HEADER, "trace-chat-stream-image-bridge-123")
.body(
r#"{"model":"gpt-image-2","messages":[{"role":"user","content":"Draw a city made of glass"}],"stream":true,"size":"1024x1024"}"#,
)
.send()
.await
.expect("request should succeed");
let status = response.status();
let response_text = response.text().await.expect("body should read");
let stored_candidates = request_candidate_repository
.list_by_request_id("trace-chat-stream-image-bridge-123")
.await
.expect("request candidates should read");
assert_eq!(
status,
StatusCode::OK,
"{response_text}\n{stored_candidates:#?}"
);
assert!(response_text.contains("\"object\":\"chat.completion.chunk\""));
assert!(response_text.contains("![generated image](data:image/png;base64,aGVsbG8=)"));
assert!(response_text.contains("data: [DONE]"));
assert!(!response_text.contains("image_generation.completed"));
let seen_plan = seen_execution_plan
.lock()
.expect("mutex should lock")
.clone()
.expect("execution plan should be captured");
assert_eq!(seen_plan.trace_id, "trace-chat-stream-image-bridge-123");
assert_eq!(seen_plan.client_api_format, "openai:chat");
assert_eq!(seen_plan.provider_api_format, "openai:image");
assert_eq!(seen_plan.url, "https://images.example.com/v1/responses");
assert!(seen_plan.plan_stream);
assert_eq!(seen_plan.auth_header, "Bearer sk-upstream-image-bridge");
assert_eq!(seen_plan.chatgpt_web_marker, "");
assert_eq!(seen_plan.body_json["model"], "gpt-image-2");
assert_eq!(seen_plan.body_json["stream"], true);
assert_eq!(
seen_plan.body_json["input"][0]["content"],
"Draw a city made of glass"
);
assert_eq!(seen_plan.body_json["tools"][0]["type"], "image_generation");
assert_eq!(seen_plan.body_json["tools"][0]["size"], "1024x1024");
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_routes_openai_responses_stream_image_intent_to_openai_image_plan_without_streaming_support(
) {
let seen_execution_plan = Arc::new(Mutex::new(None::<SeenImageBridgeExecutionPlan>));
let execution_runtime = image_bridge_execution_runtime(Arc::clone(&seen_execution_plan));
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let (gateway_url, gateway_handle, client_api_key, _request_candidate_repository) =
start_image_bridge_gateway(
"responses-stream-image-bridge",
"image-provider",
"custom",
"https://images.example.com",
execution_runtime_url,
)
.await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/responses"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(http::header::AUTHORIZATION, format!("Bearer {client_api_key}"))
.header(TRACE_ID_HEADER, "trace-responses-stream-image-bridge-123")
.body(
r#"{"model":"gpt-image-2","input":"Draw a mountain observatory","tools":[{"type":"image_generation","size":"1024x1024"}],"stream":true}"#,
)
.send()
.await
.expect("request should succeed");
let status = response.status();
let response_text = response.text().await.expect("body should read");
assert_eq!(status, StatusCode::OK, "{response_text}");
assert!(response_text.contains("response.output_item.done"));
assert!(response_text.contains("image_generation_call"));
let seen_plan = seen_execution_plan
.lock()
.expect("mutex should lock")
.clone()
.expect("execution plan should be captured");
assert_eq!(
seen_plan.trace_id,
"trace-responses-stream-image-bridge-123"
);
assert_eq!(seen_plan.client_api_format, "openai:responses");
assert_eq!(seen_plan.provider_api_format, "openai:image");
assert_eq!(seen_plan.url, "https://images.example.com/v1/responses");
assert!(seen_plan.plan_stream);
assert_eq!(seen_plan.auth_header, "Bearer sk-upstream-image-bridge");
assert_eq!(seen_plan.body_json["stream"], true);
assert_eq!(seen_plan.body_json["input"], "Draw a mountain observatory");
assert_eq!(seen_plan.body_json["tools"][0]["type"], "image_generation");
assert_eq!(seen_plan.body_json["tools"][0]["size"], "1024x1024");
gateway_handle.abort();
execution_runtime_handle.abort();
}
apps/aether-gateway/src/tests/ai_execute/sync/gemini/cli.rs
+21 -3
View File
@@ -349,7 +349,13 @@ async fn gateway_executes_gemini_cli_sync_via_local_decision_gate_with_local_syn
},
"body": {
"json_body": {
"candidates": [],
"candidates": [{
"content": {
"role": "model",
"parts": [{"text": "Hello from Gemini CLI"}]
},
"finishReason": "STOP"
}],
"usageMetadata": {
"promptTokenCount": 1,
"candidatesTokenCount": 2,
@@ -1117,7 +1123,13 @@ async fn gateway_executes_gemini_cli_sync_via_local_decision_gate_after_oauth_re
},
"body": {
"json_body": {
"candidates": [],
"candidates": [{
"content": {
"role": "model",
"parts": [{"text": "Hello from Gemini CLI"}]
},
"finishReason": "STOP"
}],
"usageMetadata": {
"promptTokenCount": 1,
"candidatesTokenCount": 2,
@@ -1595,7 +1607,13 @@ async fn gateway_executes_vertex_ai_gemini_cli_sync_via_local_decision_gate_with
},
"body": {
"json_body": {
"candidates": [],
"candidates": [{
"content": {
"role": "model",
"parts": [{"text": "Hello from Vertex Gemini CLI"}]
},
"finishReason": "STOP"
}],
"usageMetadata": {
"promptTokenCount": 1,
"candidatesTokenCount": 2,
apps/aether-gateway/src/tests/ai_execute/sync/gemini/local_chat.rs
+7 -1
View File
@@ -349,7 +349,13 @@ async fn gateway_executes_gemini_chat_sync_via_local_decision_gate_with_local_sy
},
"body": {
"json_body": {
"candidates": [],
"candidates": [{
"content": {
"role": "model",
"parts": [{"text": "Hello from Gemini"}]
},
"finishReason": "STOP"
}],
"usageMetadata": {
"promptTokenCount": 1,
"candidatesTokenCount": 2,
apps/aether-gateway/src/tests/architecture/runtime_and_security.rs
+1 -2
View File
@@ -406,7 +406,6 @@ fn scheduler_candidate_runtime_paths_depend_on_scheduler_core_and_state_trait()
"fn candidate_provider_key_health_bucket(",
"fn candidate_provider_key_health_score(",
"count_recent_active_requests_for_provider(",
"is_candidate_in_recent_failure_cooldown(",
"provider_key_health_score(",
"provider_key_rpm_allows_request_since(",
"read_recent_request_candidates(128)",
@@ -1397,7 +1396,7 @@ fn usage_reporting_does_not_log_raw_report_context() {
#[test]
fn proxy_registration_client_does_not_log_raw_management_response_body() {
let source = read_workspace_file("apps/aether-proxy/src/registration/client.rs");
let source = read_workspace_file("apps/aether-tunnel/src/registration/client.rs");
assert!(
!source.contains("error!(body = %text"),
"registration/client.rs should not log raw management response bodies"
apps/aether-gateway/src/tests/control/admin/endpoints/keys.rs
+101
View File
@@ -118,6 +118,107 @@ async fn gateway_handles_admin_provider_keys_locally_with_trusted_admin_principa
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_provider_keys_expose_circuit_breaker_and_recover_clears_it() {
let key = sample_key("key-1", "provider-1", "openai:chat", "sk-test-a").with_health_fields(
Some(json!({"openai:chat": {
"health_score": 0.2,
"consecutive_failures": 8,
"last_failure_at": "2026-03-26T12:00:00+00:00"
}})),
Some(json!({"openai:chat": {
"open": true,
"open_at": "2026-03-26T12:00:00+00:00",
"reason": "consecutive_failures_8",
"next_probe_at": "2026-03-26T12:01:00+00:00",
"next_probe_at_unix_secs": 1774526460u64,
"probe_interval_minutes": 1,
"max_probe_interval_minutes": 32,
"half_open_until": null,
"half_open_successes": 0,
"half_open_failures": 0
}})),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-1", "openai", 10)],
vec![sample_endpoint(
"endpoint-1",
"provider-1",
"openai:chat",
"https://example.com/v1",
)],
vec![key],
));
let gateway_state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone(
&provider_catalog_repository,
)),
);
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let response = client
.get(format!(
"{gateway_url}/api/admin/endpoints/providers/provider-1/keys?skip=0&limit=50"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.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[0]["circuit_breaker_open"], true);
assert_eq!(
payload[0]["circuit_breaker_by_format"]["openai:chat"]["reason"],
"consecutive_failures_8"
);
assert_eq!(
payload[0]["circuit_breaker_by_format"]["openai:chat"]["probe_interval_minutes"],
1
);
assert!(
payload[0]["circuit_breaker_by_format"]["openai:chat"]["next_probe_at_unix_secs"]
.as_u64()
.is_some()
);
let recover_response = client
.patch(format!(
"{gateway_url}/api/admin/endpoints/health/keys/key-1"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("recover request should succeed");
assert_eq!(recover_response.status(), StatusCode::OK);
let response = client
.get(format!(
"{gateway_url}/api/admin/endpoints/providers/provider-1/keys?skip=0&limit=50"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.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[0]["circuit_breaker_open"], false);
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_keys_page_locally_with_total() {
let upstream_hits = Arc::new(Mutex::new(0usize));
apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs
+1 -1
View File
@@ -1224,7 +1224,7 @@ async fn gateway_refresh_quota_reconciles_unsupported_fixed_provider_endpoints_b
(
"provider-vertex-ai-reconcile",
"vertex_ai",
2usize,
3usize,
"gemini:generate_content",
"https://aiplatform.googleapis.com",
"Vertex AI 暂不支持自动刷新额度",
apps/aether-gateway/src/tests/control/admin/endpoints/routes.rs
+115 -1
View File
@@ -148,6 +148,120 @@ async fn gateway_handles_admin_provider_endpoints_locally_with_trusted_admin_pri
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_counts_fixed_provider_oauth_keys_for_inherited_endpoint_formats() {
let mut codex_provider = sample_provider("provider-codex", "codex", 10);
codex_provider.provider_type = "codex".to_string();
let mut chatgpt_web_provider = sample_provider("provider-chatgpt-web", "chatgpt_web", 20);
chatgpt_web_provider.provider_type = "chatgpt_web".to_string();
let mut codex_key = sample_key(
"key-codex-oauth",
"provider-codex",
"openai:responses:compact",
"oauth-token",
);
codex_key.auth_type = "oauth".to_string();
codex_key.api_formats = Some(json!(["legacy:mismatch"]));
let mut chatgpt_web_key = sample_key(
"key-chatgpt-web-oauth",
"provider-chatgpt-web",
"openai:image",
"oauth-token",
);
chatgpt_web_key.auth_type = "oauth".to_string();
chatgpt_web_key.api_formats = Some(json!(["legacy:mismatch"]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![codex_provider, chatgpt_web_provider],
vec![
sample_endpoint(
"endpoint-codex-compact",
"provider-codex",
"openai:responses:compact",
"https://chatgpt.com/backend-api/codex",
),
sample_endpoint(
"endpoint-codex-image",
"provider-codex",
"openai:image",
"https://chatgpt.com/backend-api/codex",
),
sample_endpoint(
"endpoint-chatgpt-web-image",
"provider-chatgpt-web",
"openai:image",
"https://chatgpt.com",
),
],
vec![codex_key, chatgpt_web_key],
));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let codex_response = client
.get(format!(
"{gateway_url}/api/admin/endpoints/providers/provider-codex/endpoints?skip=0&limit=50"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
assert_eq!(codex_response.status(), StatusCode::OK);
let codex_payload: serde_json::Value = codex_response.json().await.expect("json should parse");
let codex_items = codex_payload
.as_array()
.expect("payload should be an array");
for api_format in ["openai:responses:compact", "openai:image"] {
let endpoint = codex_items
.iter()
.find(|item| item["api_format"] == api_format)
.expect("endpoint should exist");
assert_eq!(endpoint["total_keys"], 1);
assert_eq!(endpoint["active_keys"], 1);
}
let chatgpt_web_response = client
.get(format!(
"{gateway_url}/api/admin/endpoints/providers/provider-chatgpt-web/endpoints?skip=0&limit=50"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
assert_eq!(chatgpt_web_response.status(), StatusCode::OK);
let chatgpt_web_payload: serde_json::Value = chatgpt_web_response
.json()
.await
.expect("json should parse");
let chatgpt_web_items = chatgpt_web_payload
.as_array()
.expect("payload should be an array");
let chatgpt_web_image = chatgpt_web_items
.iter()
.find(|item| item["api_format"] == "openai:image")
.expect("image endpoint should exist");
assert_eq!(chatgpt_web_image["total_keys"], 1);
assert_eq!(chatgpt_web_image["active_keys"], 1);
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_counts_keys_with_null_api_formats_for_each_fixed_provider_endpoint() {
let upstream_hits = Arc::new(Mutex::new(0usize));
@@ -195,7 +309,7 @@ async fn gateway_counts_keys_with_null_api_formats_for_each_fixed_provider_endpo
vec![inherited_key],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (_, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
apps/aether-gateway/src/tests/control/admin/models/global.rs
+136
View File
@@ -936,6 +936,142 @@ async fn gateway_handles_admin_global_model_routing_locally_with_trusted_admin_p
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_global_model_routing_counts_image_provider_keys_by_provider_model_name() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/models/global/global-gpt-image/routing",
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("unexpected upstream hit"))
}
}),
);
let mut image_provider = sample_provider("provider-image", "image", 10);
image_provider.provider_type = "chatgpt_web".to_string();
let grok_provider = sample_provider("provider-grok", "grok2api", 20);
let mut image_key = sample_key(
"key-image-routing",
"provider-image",
"legacy:mismatch",
"sk-image-routing-1234",
);
image_key.name = "image-account".to_string();
image_key.auth_type = "oauth".to_string();
image_key.allowed_models = Some(json!(["gpt-image-2"]));
let mut grok_key = sample_key(
"key-grok-routing",
"provider-grok",
"openai:chat",
"sk-grok-routing-5678",
);
grok_key.name = "all".to_string();
grok_key.allowed_models = Some(json!(["gpt-image-2"]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![image_provider, grok_provider],
vec![
sample_endpoint(
"endpoint-image",
"provider-image",
"openai:image",
"https://chatgpt.example",
),
sample_endpoint(
"endpoint-grok-chat",
"provider-grok",
"openai:chat",
"https://grok.example",
),
],
vec![image_key, grok_key],
));
let global_model_repository = Arc::new(
InMemoryGlobalModelReadRepository::seed(Vec::new())
.with_admin_global_models(vec![sample_admin_global_model(
"global-gpt-image",
"GPT-Image-2",
"GPT-Image-2",
)])
.with_admin_provider_models(vec![
sample_admin_provider_model(
"model-image-gpt-image",
"provider-image",
"global-gpt-image",
"gpt-image-2",
),
sample_admin_provider_model(
"model-grok-gpt-image",
"provider-grok",
"global-gpt-image",
"gpt-image-2",
),
]),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)
.with_global_model_repository_for_tests(global_model_repository),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/models/global/global-gpt-image/routing"
))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.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["global_model_name"], "GPT-Image-2");
assert_eq!(payload["total_providers"], 2);
assert_eq!(payload["active_providers"], 2);
let providers = payload["providers"].as_array().expect("providers array");
assert_eq!(providers.len(), 2);
let image_endpoints = providers[0]["endpoints"]
.as_array()
.expect("image endpoints array");
assert_eq!(providers[0]["id"], "provider-image");
assert_eq!(image_endpoints[0]["api_format"], "openai:image");
assert_eq!(image_endpoints[0]["total_keys"], 1);
assert_eq!(image_endpoints[0]["active_keys"], 1);
assert_eq!(image_endpoints[0]["keys"][0]["name"], "image-account");
let grok_endpoints = providers[1]["endpoints"]
.as_array()
.expect("grok endpoints array");
assert_eq!(providers[1]["id"], "provider-grok");
assert_eq!(grok_endpoints[0]["api_format"], "openai:chat");
assert_eq!(grok_endpoints[0]["total_keys"], 1);
assert_eq!(grok_endpoints[0]["active_keys"], 1);
assert_eq!(grok_endpoints[0]["keys"][0]["name"], "all");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_creates_admin_global_model_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
apps/aether-gateway/src/tests/control/admin/models/provider.rs
+1
View File
@@ -94,6 +94,7 @@ async fn gateway_handles_admin_provider_models_locally_with_trusted_admin_princi
assert_eq!(items[0]["effective_input_price"], 3.0);
assert_eq!(items[0]["effective_output_price"], 15.0);
assert_eq!(items[0]["effective_supports_streaming"], true);
assert!(items[0]["model_test_capabilities"]["openai:image"].is_null());
assert_eq!(items[0]["created_at"], "2024-03-21T05:46:40Z");
assert_eq!(items[0]["updated_at"], "2024-03-21T05:48:20Z");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
apps/aether-gateway/src/tests/control/admin/monitoring.rs
+2
View File
@@ -273,6 +273,8 @@ async fn gateway_handles_admin_monitoring_system_status_locally_with_trusted_adm
assert_eq!(payload["tunnel"]["nodes"], json!(0));
assert_eq!(payload["tunnel"]["active_streams"], json!(0));
assert_eq!(payload["recent_errors"], json!(1));
assert_eq!(payload["usage_counter"]["status"], json!("idle"));
assert_eq!(payload["usage_counter"]["outbox_pending_rows"], json!(0));
assert!(payload["timestamp"].as_str().is_some());
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
apps/aether-gateway/src/tests/control/admin/provider_query.rs
+278
View File
@@ -1256,6 +1256,146 @@ async fn gateway_handles_admin_provider_query_gemini_embedding_model_test() {
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_vertex_gemini_embedding_model_test() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-vertex-ai");
assert_eq!(plan.endpoint_id, "endpoint-vertex-gemini-embedding");
assert_eq!(plan.key_id, "key-vertex-gemini-embedding");
assert_eq!(plan.client_api_format, "openai:embedding");
assert_eq!(plan.provider_api_format, "gemini:embedding");
assert_eq!(
plan.url,
"https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-embedding-2:predict?key=sk-vertex-gemini-embedding"
);
assert_eq!(plan.model_name.as_deref(), Some("gemini-embedding-2"));
assert!(!plan.stream);
let body = plan.body.json_body.as_ref().expect("json body");
assert!(
body.get("model").is_none(),
"Vertex predict carries the model in the URL path; the test body must not repeat it"
);
assert_eq!(
body["instances"][0]["content"],
json!("This is a test embedding input.")
);
assert!(body.get("content").is_none());
assert!(body.get("requests").is_none());
assert!(
body.get("stream").is_none(),
"gemini embedding provider body must not carry stream"
);
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"predictions": [
{
"embeddings": {
"values": [0.1, 0.2, 0.3]
}
}
],
"deployedModelId": "gemini-embedding-2"
}
},
"telemetry": {
"elapsed_ms": 27
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-vertex-ai", "Vertex AI", 10);
provider.provider_type = "vertex_ai".to_string();
let mut key = sample_key(
"key-vertex-gemini-embedding",
"provider-vertex-ai",
"gemini:embedding",
"sk-vertex-gemini-embedding",
);
key.allowed_models = Some(json!(["gemini-embedding-2"]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-vertex-gemini-embedding",
"provider-vertex-ai",
"gemini:embedding",
"https://aiplatform.googleapis.com",
)],
vec![key],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-vertex-ai",
"model": "gemini-embedding-2",
"api_format": "gemini:embedding",
"endpoint_id": "endpoint-vertex-gemini-embedding",
"request_body": {
"model": "gemini-embedding-2",
"input": "This is a test embedding input."
}
}))
.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["success"], json!(true));
assert_eq!(payload["error"], serde_json::Value::Null);
assert_eq!(payload["attempts"][0]["status"], json!("success"));
assert_eq!(
payload["attempts"][0]["request_body"]["instances"][0]["content"],
json!("This is a test embedding input.")
);
assert_eq!(
payload["attempts"][0]["endpoint_product"],
json!("Vertex AI")
);
assert_eq!(
payload["attempts"][0]["endpoint_variant"],
json!("vertex_native")
);
assert_eq!(payload["attempts"][0]["endpoint_action"], json!("predict"));
assert_eq!(
payload["attempts"][0]["endpoint_batch_strategy"],
json!("single_instance")
);
assert!(
payload["attempts"][0]["request_body"]
.get("model")
.is_none(),
"attempt debug payload must expose the exact Vertex body without a duplicate model"
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_jina_embedding_model_test() {
let execution_runtime = Router::new().route(
@@ -1967,6 +2107,144 @@ async fn gateway_streams_codex_openai_responses_upstream_for_admin_pool_model_te
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_routes_grok_responses_admin_pool_model_test_through_grok_runtime() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-grok");
assert_eq!(plan.endpoint_id, "endpoint-grok-responses");
assert_eq!(plan.key_id, "key-grok-oauth");
assert_eq!(plan.client_api_format, "openai:responses");
assert_eq!(plan.provider_api_format, "openai:responses");
assert_eq!(plan.url, "https://grok.com/rest/app-chat/conversations/new");
assert_eq!(plan.model_name.as_deref(), Some("grok-4.20-fast"));
assert!(plan.stream, "Grok model test should request a stream");
assert_eq!(
plan.headers
.get(aether_provider_transport::GROK_INTERNAL_HEADER)
.map(String::as_str),
Some("1")
);
assert_eq!(
plan.headers.get("cookie").map(String::as_str),
Some("sso=grok-sso; sso-rw=grok-rw")
);
let body = plan.body.json_body.as_ref().expect("json body");
assert_eq!(body["model"], json!("grok-4.20-fast"));
assert_eq!(body["input"], json!("Hello! This is a test message."));
assert_eq!(
body["messages"][0]["content"],
json!("stale chat-shaped frontend body")
);
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "resp-grok-model-test",
"model": "grok-4.20-fast",
"output_text": "ok"
}
},
"telemetry": {
"elapsed_ms": 18
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-grok", "Grok", 10);
provider.provider_type = "grok".to_string();
provider.config = Some(json!({"pool_advanced": {}}));
let mut key = sample_key(
"key-grok-oauth",
"provider-grok",
"openai:responses",
"__placeholder__",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
aether_crypto::encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{
"provider_type":"grok",
"sso_token":"grok-sso",
"sso_rw_token":"grok-rw"
}"#,
)
.expect("auth config should encrypt"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-grok-responses",
"provider-grok",
"openai:responses",
"https://grok.com",
)],
vec![key],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-query/test-model-failover"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-grok",
"mode": "pool",
"model": "grok-4.20-fast",
"failover_models": ["grok-4.20-fast"],
"api_format": "openai:responses",
"endpoint_id": "endpoint-grok-responses",
"request_body": {
"model": "grok-4.20-fast",
"messages": [{
"role": "user",
"content": "stale chat-shaped frontend body"
}],
"stream": true
}
}))
.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["success"], json!(true));
assert_eq!(payload["attempts"][0]["status"], json!("success"));
assert_eq!(
payload["attempts"][0]["request_body"]["message"],
json!("Hello! This is a test message.")
);
assert_eq!(
payload["attempts"][0]["request_headers"][aether_provider_transport::GROK_INTERNAL_HEADER],
json!("1")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_uses_pool_scheduler_order_for_admin_pool_model_test() {
let execution_runtime = Router::new().route(
apps/aether-gateway/src/tests/control/admin/providers.rs
+2
View File
@@ -241,6 +241,7 @@ async fn gateway_handles_admin_provider_summary_locally_with_trusted_admin_princ
"pool_advanced": {"enabled": true},
"failover_rules": {"strategy": "ordered"},
"chat_pii_redaction": {"enabled": true},
"kiro": {"simulated_cache_enabled": true},
"provider_ops": {"architecture_id": "anyrouter"}
})),
);
@@ -353,6 +354,7 @@ async fn gateway_handles_admin_provider_summary_locally_with_trusted_admin_princ
assert_eq!(payload["ops_configured"], true);
assert_eq!(payload["ops_architecture_id"], "anyrouter");
assert_eq!(payload["chat_pii_redaction"], json!({"enabled": true}));
assert_eq!(payload["kiro_simulated_cache_enabled"], true);
assert_eq!(payload["created_at"], "2024-03-21T05:46:40Z");
assert_eq!(payload["updated_at"], "2024-03-21T05:48:20Z");
assert_eq!(
apps/aether-gateway/src/tests/control/admin/proxy_nodes.rs
+1
View File
@@ -1224,6 +1224,7 @@ async fn gateway_tests_connected_tunnel_proxy_nodes_with_active_probe() {
proxy_tx,
proxy_close_tx,
16,
2,
)));
let gateway = build_router_with_state(state);
apps/aether-gateway/src/tests/control/admin/stats.rs
+4
View File
@@ -963,6 +963,8 @@ async fn gateway_handles_admin_stats_provider_performance_locally_with_trusted_a
assert_eq!(payload["summary"]["p99_response_time_ms"], 1000);
assert_eq!(payload["summary"]["response_time_sample_count"], 11);
assert_eq!(payload["summary"]["slow_request_count"], 0);
assert_eq!(payload["usage_counter"]["status"], json!("idle"));
assert_eq!(payload["usage_counter"]["outbox_pending_rows"], json!(0));
assert_eq!(payload["providers"].as_array().map(Vec::len), Some(2));
assert_eq!(payload["providers"][0]["provider_id"], "provider-1");
@@ -1038,6 +1040,8 @@ async fn gateway_returns_empty_admin_stats_provider_performance_without_usage_re
);
assert_eq!(payload["providers"].as_array().map(Vec::len), Some(0));
assert_eq!(payload["timeline"].as_array().map(Vec::len), Some(0));
assert_eq!(payload["usage_counter"]["status"], json!("idle"));
assert_eq!(payload["usage_counter"]["outbox_pending_rows"], json!(0));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
apps/aether-gateway/src/tests/control/admin/system.rs
+10
View File
@@ -296,6 +296,8 @@ async fn gateway_handles_admin_system_stats_locally_with_trusted_admin_principal
assert_eq!(payload["providers"]["active"], json!(1));
assert_eq!(payload["api_keys"], json!(0));
assert_eq!(payload["requests"], json!(0));
assert_eq!(payload["usage_counter"]["status"], json!("idle"));
assert_eq!(payload["usage_counter"]["outbox_pending_rows"], json!(0));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
@@ -1152,6 +1154,14 @@ async fn gateway_handles_admin_system_api_formats_locally_with_trusted_admin_pri
.expect("formats should be an array");
assert_eq!(formats[0]["value"], "openai:chat");
assert_eq!(formats[0]["default_path"], "/v1/chat/completions");
let gemini_embedding = formats
.iter()
.find(|item| item["value"] == "gemini:embedding")
.expect("gemini embedding format should exist");
assert_eq!(
gemini_embedding["default_path"],
"/v1beta/models/{model}:{action}"
);
assert!(formats
.iter()
.any(|item| item["value"] == "openai:embedding"));
apps/aether-gateway/src/tests/control/admin/users.rs
+115
View File
@@ -567,6 +567,121 @@ async fn gateway_allows_default_user_group_access_policy_updates() {
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_allows_removing_default_group_members_when_other_group_remains() {
let upstream = Router::new().fallback(any(|_request: Request| async {
(StatusCode::OK, Body::from("unexpected upstream hit"))
}));
let user_repository = Arc::new(
InMemoryUserReadRepository::seed_auth_users(vec![
sample_admin_user_with_role("admin-1", "admin", "[email protected]", "admin"),
sample_admin_user_with_role("user-2", "user", "[email protected]", "bob"),
sample_admin_user_with_role("user-3", "user", "[email protected]", "carol"),
])
.with_export_users(vec![
sample_admin_export_user_with("admin", true, "admin-1", "[email protected]", "admin"),
sample_admin_export_user_with("user", true, "user-2", "[email protected]", "bob"),
sample_admin_export_user_with("user", true, "user-3", "[email protected]", "carol"),
]),
);
let default_group = user_repository
.create_user_group(UpsertUserGroupRecord {
name: "Default".to_string(),
description: None,
priority: 0,
allowed_providers: None,
allowed_providers_mode: "unrestricted".to_string(),
allowed_api_formats: None,
allowed_api_formats_mode: "unrestricted".to_string(),
allowed_models: None,
allowed_models_mode: "unrestricted".to_string(),
rate_limit: None,
rate_limit_mode: "system".to_string(),
})
.await
.expect("default group should create")
.expect("default group should exist");
let team_group = user_repository
.create_user_group(UpsertUserGroupRecord {
name: "Team".to_string(),
description: None,
priority: 0,
allowed_providers: None,
allowed_providers_mode: "unrestricted".to_string(),
allowed_api_formats: None,
allowed_api_formats_mode: "unrestricted".to_string(),
allowed_models: None,
allowed_models_mode: "unrestricted".to_string(),
rate_limit: None,
rate_limit_mode: "system".to_string(),
})
.await
.expect("team group should create")
.expect("team group should exist");
user_repository
.add_user_to_group(&team_group.id, "user-2")
.await
.expect("team membership should create");
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_user_reader_for_tests(user_repository.clone())
.with_system_config_values_for_tests(vec![(
crate::constants::DEFAULT_USER_GROUP_CONFIG_KEY.to_string(),
json!(default_group.id),
)]),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
user_repository
.add_user_to_group(&default_group.id, "user-2")
.await
.expect("default membership should create");
user_repository
.add_user_to_group(&default_group.id, "user-3")
.await
.expect("default membership should create");
let remove_user_with_other_group = client
.put(format!(
"{gateway_url}/api/admin/user-groups/{}/members",
default_group.id
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({ "user_ids": ["user-3"] }))
.send()
.await
.expect("request should succeed");
assert_eq!(remove_user_with_other_group.status(), StatusCode::OK);
let reject_groupless_user = client
.put(format!(
"{gateway_url}/api/admin/user-groups/{}/members",
default_group.id
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({ "user_ids": [] }))
.send()
.await
.expect("request should succeed");
assert_eq!(reject_groupless_user.status(), StatusCode::BAD_REQUEST);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_resolves_admin_user_batch_selection_locally() {
let upstream_hits = Arc::new(Mutex::new(0usize));
apps/aether-gateway/src/tests/control/helpers.rs
+1
View File
@@ -436,6 +436,7 @@ pub(super) fn sample_oauth_provider_config(provider_type: &str) -> StoredOAuthPr
Some(vec!["openid".to_string()]),
Some(json!({"email": "email"})),
Some(json!({"team": true})),
None,
true,
)
}
apps/aether-gateway/src/tests/control/proxy/embeddings.rs
+553
View File
@@ -78,6 +78,147 @@ fn embedding_execution_runtime() -> Router {
)
}
fn gemini_embedding_success_state(
execution_runtime_url: String,
client_api_format: &str,
) -> AppState {
let mut snapshot = sample_currently_usable_auth_snapshot(
"key-gemini-embedding-success",
"user-gemini-embedding-success",
);
snapshot.user_allowed_providers = None;
snapshot.api_key_allowed_providers = None;
snapshot.user_allowed_api_formats = Some(vec![client_api_format.to_string()]);
snapshot.api_key_allowed_api_formats = Some(vec![client_api_format.to_string()]);
snapshot.user_allowed_models = Some(vec!["gemini-embedding-2-preview".to_string()]);
snapshot.api_key_allowed_models = Some(vec!["gemini-embedding-2-preview".to_string()]);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-gemini-embedding-success")),
snapshot,
)]));
let candidate_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
gemini_embedding_candidate_row(),
]));
let mut provider = sample_provider("provider-gemini-embedding", "Gemini Embeddings", 1);
provider.provider_type = "gemini".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-gemini-embedding",
"provider-gemini-embedding",
"gemini:embedding",
"https://generativelanguage.googleapis.com/v1beta",
)],
vec![sample_key(
"key-upstream-gemini-embedding",
"provider-gemini-embedding",
"gemini:embedding",
"sk-upstream-gemini-embedding",
)],
));
let data_state =
GatewayDataState::with_provider_catalog_and_minimal_candidate_selection_for_tests(
provider_catalog_repository,
candidate_repository,
)
.with_auth_api_key_reader(auth_repository)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY);
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(data_state)
}
fn vertex_gemini_embedding_success_state(execution_runtime_url: String) -> AppState {
let mut snapshot = sample_currently_usable_auth_snapshot(
"key-vertex-gemini-embedding-success",
"user-vertex-gemini-embedding-success",
);
snapshot.user_allowed_providers = None;
snapshot.api_key_allowed_providers = Some(vec!["openai".to_string(), "vertex_ai".to_string()]);
snapshot.user_allowed_api_formats = Some(vec!["openai:embedding".to_string()]);
snapshot.api_key_allowed_api_formats = Some(vec!["openai:embedding".to_string()]);
snapshot.user_allowed_models = Some(vec!["gemini-embedding-2-preview".to_string()]);
snapshot.api_key_allowed_models = Some(vec!["gemini-embedding-2-preview".to_string()]);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-vertex-gemini-embedding-success")),
snapshot,
)]));
let candidate_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
vertex_gemini_embedding_candidate_row(),
]));
let mut provider = sample_provider("provider-vertex-gemini-embedding", "Vertex AI", 1);
provider.provider_type = "vertex_ai".to_string();
let mut key = sample_key(
"key-upstream-vertex-gemini-embedding",
"provider-vertex-gemini-embedding",
"gemini:embedding",
"sk-upstream-vertex-gemini-embedding",
);
key.allowed_models = Some(json!(["gemini-embedding-2"]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-vertex-gemini-embedding",
"provider-vertex-gemini-embedding",
"gemini:embedding",
"https://aiplatform.googleapis.com",
)],
vec![key],
));
let data_state =
GatewayDataState::with_provider_catalog_and_minimal_candidate_selection_for_tests(
provider_catalog_repository,
candidate_repository,
)
.with_auth_api_key_reader(auth_repository)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY);
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(data_state)
}
fn gemini_embedding_conversion_execution_runtime() -> Router {
Router::new().route(
"/v1/execute/sync",
any(|Json(plan): Json<ExecutionPlan>| async move {
assert_openai_to_gemini_embedding_execution_plan(&plan);
Json(gemini_embedding_execution_result(&plan))
}),
)
}
fn vertex_gemini_embedding_conversion_execution_runtime() -> Router {
Router::new().route(
"/v1/execute/sync",
any(|Json(plan): Json<ExecutionPlan>| async move {
assert_openai_to_vertex_gemini_embedding_execution_plan(&plan);
Json(vertex_gemini_embedding_execution_result(&plan))
}),
)
}
fn gemini_embedding_batch_conversion_execution_runtime() -> Router {
Router::new().route(
"/v1/execute/sync",
any(|Json(plan): Json<ExecutionPlan>| async move {
assert_openai_to_gemini_batch_embedding_execution_plan(&plan);
Json(gemini_batch_embedding_execution_result(&plan))
}),
)
}
fn gemini_embedding_native_execution_runtime() -> Router {
Router::new().route(
"/v1/execute/sync",
any(|Json(plan): Json<ExecutionPlan>| async move {
assert_native_gemini_embedding_execution_plan(&plan);
Json(gemini_embedding_execution_result(&plan))
}),
)
}
fn embedding_candidate_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-embedding".to_string(),
@@ -112,6 +253,53 @@ fn embedding_candidate_row() -> StoredMinimalCandidateSelectionRow {
}
}
fn gemini_embedding_candidate_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-gemini-embedding".to_string(),
provider_name: "Gemini Embeddings".to_string(),
provider_type: "gemini".to_string(),
provider_priority: 1,
provider_is_active: true,
endpoint_id: "endpoint-gemini-embedding".to_string(),
endpoint_api_format: "gemini:embedding".to_string(),
endpoint_api_family: Some("gemini".to_string()),
endpoint_kind: Some("embedding".to_string()),
endpoint_is_active: true,
key_id: "key-upstream-gemini-embedding".to_string(),
key_name: "default".to_string(),
key_auth_type: "api_key".to_string(),
key_is_active: true,
key_api_formats: Some(vec!["gemini:embedding".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 50,
key_global_priority_by_format: None,
model_id: "model-gemini-embedding-preview".to_string(),
global_model_id: "global-gemini-embedding-preview".to_string(),
global_model_name: "gemini-embedding-2-preview".to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(false),
model_provider_model_name: "gemini-embedding-2-preview".to_string(),
model_provider_model_mappings: None,
model_supports_streaming: Some(false),
model_is_active: true,
model_is_available: true,
}
}
fn vertex_gemini_embedding_candidate_row() -> StoredMinimalCandidateSelectionRow {
let mut row = gemini_embedding_candidate_row();
row.provider_id = "provider-vertex-gemini-embedding".to_string();
row.provider_name = "Vertex AI".to_string();
row.provider_type = "vertex_ai".to_string();
row.endpoint_id = "endpoint-vertex-gemini-embedding".to_string();
row.key_id = "key-upstream-vertex-gemini-embedding".to_string();
row.key_name = "default".to_string();
row.key_allowed_models = Some(vec!["gemini-embedding-2".to_string()]);
row.model_provider_model_name = "gemini-embedding-2".to_string();
row
}
fn assert_embedding_execution_plan(plan: &ExecutionPlan) {
assert_eq!(plan.client_api_format, "openai:embedding");
assert_eq!(plan.provider_api_format, "openai:embedding");
@@ -123,6 +311,103 @@ fn assert_embedding_execution_plan(plan: &ExecutionPlan) {
assert!(body.get("input").is_some());
}
fn assert_openai_to_gemini_embedding_execution_plan(plan: &ExecutionPlan) {
assert_eq!(plan.client_api_format, "openai:embedding");
assert_eq!(plan.provider_api_format, "gemini:embedding");
assert_eq!(plan.method, "POST");
assert_eq!(
plan.url,
"https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2-preview:embedContent"
);
assert_eq!(
plan.headers.get("x-goog-api-key").map(String::as_str),
Some("sk-upstream-gemini-embedding")
);
assert_eq!(
plan.model_name.as_deref(),
Some("gemini-embedding-2-preview")
);
assert!(!plan.stream);
let body = plan.body.json_body.as_ref().expect("json request body");
assert_eq!(body["model"], "gemini-embedding-2-preview");
assert_eq!(body["content"]["parts"][0]["text"], "hello");
assert!(body.get("input").is_none());
assert!(body.get("messages").is_none());
}
fn assert_openai_to_vertex_gemini_embedding_execution_plan(plan: &ExecutionPlan) {
assert_eq!(plan.provider_id, "provider-vertex-gemini-embedding");
assert_eq!(plan.client_api_format, "openai:embedding");
assert_eq!(plan.provider_api_format, "gemini:embedding");
assert_eq!(plan.method, "POST");
assert_eq!(
plan.url,
"https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-embedding-2:predict?key=sk-upstream-vertex-gemini-embedding"
);
assert_eq!(
plan.model_name.as_deref(),
Some("gemini-embedding-2-preview")
);
assert!(!plan.stream);
let body = plan.body.json_body.as_ref().expect("json request body");
assert!(
body.get("model").is_none(),
"Vertex predict carries the model in the path; the body must not repeat it"
);
assert_eq!(body["instances"][0]["content"], "hello");
assert!(body.get("content").is_none());
assert!(body.get("input").is_none());
assert!(body.get("messages").is_none());
}
fn assert_openai_to_gemini_batch_embedding_execution_plan(plan: &ExecutionPlan) {
assert_eq!(plan.client_api_format, "openai:embedding");
assert_eq!(plan.provider_api_format, "gemini:embedding");
assert_eq!(plan.method, "POST");
assert_eq!(
plan.url,
"https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2-preview:batchEmbedContents"
);
assert_eq!(
plan.headers.get("x-goog-api-key").map(String::as_str),
Some("sk-upstream-gemini-embedding")
);
assert!(!plan.stream);
let body = plan.body.json_body.as_ref().expect("json request body");
assert!(body.get("model").is_none());
let requests = body["requests"].as_array().expect("batch requests");
assert_eq!(requests.len(), 2);
assert_eq!(requests[0]["model"], "models/gemini-embedding-2-preview");
assert_eq!(requests[0]["content"]["parts"][0]["text"], "hello");
assert_eq!(requests[1]["model"], "models/gemini-embedding-2-preview");
assert_eq!(requests[1]["content"]["parts"][0]["text"], "world");
assert!(body.get("input").is_none());
assert!(body.get("messages").is_none());
}
fn assert_native_gemini_embedding_execution_plan(plan: &ExecutionPlan) {
assert_eq!(plan.client_api_format, "gemini:embedding");
assert_eq!(plan.provider_api_format, "gemini:embedding");
assert_eq!(plan.method, "POST");
assert_eq!(
plan.url,
"https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2-preview:embedContent"
);
assert_eq!(
plan.headers.get("x-goog-api-key").map(String::as_str),
Some("sk-upstream-gemini-embedding")
);
assert_eq!(
plan.model_name.as_deref(),
Some("gemini-embedding-2-preview")
);
assert!(!plan.stream);
let body = plan.body.json_body.as_ref().expect("json request body");
assert_eq!(body["content"]["parts"][0]["text"], "hello");
assert!(body.get("input").is_none());
assert!(body.get("messages").is_none());
}
fn embedding_execution_result(plan: &ExecutionPlan) -> ExecutionResult {
ExecutionResult {
request_id: plan.request_id.clone(),
@@ -145,6 +430,79 @@ fn embedding_execution_result(plan: &ExecutionPlan) -> ExecutionResult {
}
}
fn gemini_embedding_execution_result(plan: &ExecutionPlan) -> ExecutionResult {
ExecutionResult {
request_id: plan.request_id.clone(),
candidate_id: plan.candidate_id.clone(),
status_code: 200,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
body: Some(ResponseBody {
json_body: Some(json!({
"model": "gemini-embedding-2-preview",
"embedding": {
"values": [0.1, 0.2, 0.3]
},
"usageMetadata": {
"promptTokenCount": 4,
"totalTokenCount": 4
}
})),
body_bytes_b64: None,
}),
telemetry: None,
error: None,
}
}
fn vertex_gemini_embedding_execution_result(plan: &ExecutionPlan) -> ExecutionResult {
ExecutionResult {
request_id: plan.request_id.clone(),
candidate_id: plan.candidate_id.clone(),
status_code: 200,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
body: Some(ResponseBody {
json_body: Some(json!({
"predictions": [
{
"embeddings": {
"values": [0.1, 0.2, 0.3]
}
}
],
"deployedModelId": "gemini-embedding-2"
})),
body_bytes_b64: None,
}),
telemetry: None,
error: None,
}
}
fn gemini_batch_embedding_execution_result(plan: &ExecutionPlan) -> ExecutionResult {
ExecutionResult {
request_id: plan.request_id.clone(),
candidate_id: plan.candidate_id.clone(),
status_code: 200,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
body: Some(ResponseBody {
json_body: Some(json!({
"model": "gemini-embedding-2-preview",
"embeddings": [
{"values": [0.1, 0.2, 0.3]},
{"values": [0.4, 0.5, 0.6]}
],
"usageMetadata": {
"promptTokenCount": 8,
"totalTokenCount": 8
}
})),
body_bytes_b64: None,
}),
telemetry: None,
error: None,
}
}
#[tokio::test]
async fn embeddings_route_accepts_openai_payload() {
let (execution_runtime_url, execution_runtime_handle) =
@@ -215,6 +573,201 @@ async fn embeddings_route_accepts_openai_payload() {
execution_runtime_handle.abort();
}
#[tokio::test]
async fn embeddings_route_converts_openai_payload_to_gemini_embedding_provider() {
let (execution_runtime_url, execution_runtime_handle) =
start_server(gemini_embedding_conversion_execution_runtime()).await;
let gateway = build_router_with_state(gemini_embedding_success_state(
execution_runtime_url,
"openai:embedding",
));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/embeddings"))
.header(
http::header::AUTHORIZATION,
"Bearer sk-gemini-embedding-success",
)
.json(&json!({
"model": "gemini-embedding-2-preview",
"input": "hello"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get(CONTROL_ENDPOINT_SIGNATURE_HEADER)
.and_then(|value| value.to_str().ok()),
Some("openai:embedding")
);
assert_eq!(
response
.headers()
.get(CONTROL_EXECUTION_RUNTIME_HEADER)
.and_then(|value| value.to_str().ok()),
Some("true")
);
let payload: serde_json::Value = response.json().await.expect("body should parse");
assert_eq!(payload["object"], "list");
assert_eq!(payload["model"], "gemini-embedding-2-preview");
assert_eq!(payload["data"][0]["object"], "embedding");
assert_eq!(payload["data"][0]["embedding"], json!([0.1, 0.2, 0.3]));
assert_eq!(payload["usage"]["prompt_tokens"], json!(4));
assert_eq!(payload["usage"]["total_tokens"], json!(4));
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn embeddings_route_converts_openai_payload_to_vertex_gemini_embedding_provider() {
let (execution_runtime_url, execution_runtime_handle) =
start_server(vertex_gemini_embedding_conversion_execution_runtime()).await;
let gateway =
build_router_with_state(vertex_gemini_embedding_success_state(execution_runtime_url));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/embeddings"))
.header(
http::header::AUTHORIZATION,
"Bearer sk-vertex-gemini-embedding-success",
)
.json(&json!({
"model": "gemini-embedding-2-preview",
"input": "hello"
}))
.send()
.await
.expect("request should succeed");
let endpoint_signature = response
.headers()
.get(CONTROL_ENDPOINT_SIGNATURE_HEADER)
.and_then(|value| value.to_str().ok())
.map(str::to_string);
let status = response.status();
let body_text = response.text().await.expect("body should read");
assert_eq!(
status,
StatusCode::OK,
"unexpected response body: {body_text}"
);
assert_eq!(endpoint_signature.as_deref(), Some("openai:embedding"));
let payload: serde_json::Value = serde_json::from_str(&body_text).expect("body should parse");
assert_eq!(payload["object"], "list");
assert_eq!(payload["model"], "gemini-embedding-2");
assert_eq!(payload["data"][0]["embedding"], json!([0.1, 0.2, 0.3]));
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn embeddings_route_converts_openai_batch_payload_to_gemini_batch_endpoint() {
let (execution_runtime_url, execution_runtime_handle) =
start_server(gemini_embedding_batch_conversion_execution_runtime()).await;
let gateway = build_router_with_state(gemini_embedding_success_state(
execution_runtime_url,
"openai:embedding",
));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/embeddings"))
.header(
http::header::AUTHORIZATION,
"Bearer sk-gemini-embedding-success",
)
.json(&json!({
"model": "gemini-embedding-2-preview",
"input": ["hello", "world"]
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get(CONTROL_ENDPOINT_SIGNATURE_HEADER)
.and_then(|value| value.to_str().ok()),
Some("openai:embedding")
);
let payload: serde_json::Value = response.json().await.expect("body should parse");
assert_eq!(payload["object"], "list");
assert_eq!(payload["data"].as_array().map(Vec::len), Some(2));
assert_eq!(payload["data"][0]["index"], json!(0));
assert_eq!(payload["data"][0]["embedding"], json!([0.1, 0.2, 0.3]));
assert_eq!(payload["data"][1]["index"], json!(1));
assert_eq!(payload["data"][1]["embedding"], json!([0.4, 0.5, 0.6]));
assert_eq!(payload["usage"]["prompt_tokens"], json!(8));
assert_eq!(payload["usage"]["total_tokens"], json!(8));
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gemini_embed_content_route_uses_native_gemini_embedding_provider() {
let (execution_runtime_url, execution_runtime_handle) =
start_server(gemini_embedding_native_execution_runtime()).await;
let gateway = build_router_with_state(gemini_embedding_success_state(
execution_runtime_url,
"gemini:embedding",
));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/v1beta/models/gemini-embedding-2-preview:embedContent"
))
.header("x-goog-api-key", "sk-gemini-embedding-success")
.json(&json!({
"content": {
"parts": [{"text": "hello"}]
}
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get(CONTROL_ROUTE_FAMILY_HEADER)
.and_then(|value| value.to_str().ok()),
Some("gemini")
);
assert_eq!(
response
.headers()
.get(CONTROL_ROUTE_KIND_HEADER)
.and_then(|value| value.to_str().ok()),
Some("embedding")
);
assert_eq!(
response
.headers()
.get(CONTROL_ENDPOINT_SIGNATURE_HEADER)
.and_then(|value| value.to_str().ok()),
Some("gemini:embedding")
);
let payload: serde_json::Value = response.json().await.expect("body should parse");
assert_eq!(payload["embedding"]["values"], json!([0.1, 0.2, 0.3]));
assert_eq!(payload["model"], "gemini-embedding-2-preview");
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn embeddings_route_accepts_all_canonical_input_shapes() {
let (execution_runtime_url, execution_runtime_handle) =
apps/aether-gateway/src/tests/frontdoor/ai.rs
+13 -15
View File
@@ -4,6 +4,7 @@ use super::{
InMemoryVideoTaskRepository, UpsertVideoTask, VideoTaskLookupKey, VideoTaskReadRepository,
VideoTaskStatus, VideoTaskWriteRepository, DEVELOPMENT_ENCRYPTION_KEY,
};
use crate::image_capabilities::openai_image_gateway_max_generation_count;
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,
@@ -582,7 +583,7 @@ async fn gateway_does_not_locally_reject_image_model_name_on_chat_completions()
}
#[tokio::test]
async fn gateway_rejects_image_request_with_n_greater_than_one_without_hitting_fallback_probe() {
async fn gateway_rejects_image_request_with_n_greater_than_four_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(
@@ -615,9 +616,9 @@ async fn gateway_rejects_image_request_with_n_greater_than_one_without_hitting_f
.header(http::header::CONTENT_TYPE, "application/json")
.body(
serde_json::to_vec(&json!({
"model": "gpt-image-2",
"model": "grok-imagine-image-lite",
"prompt": "draw",
"n": 2,
"n": 5,
"response_format": "b64_json"
}))
.expect("request body should encode"),
@@ -635,7 +636,13 @@ async fn gateway_rejects_image_request_with_n_greater_than_one_without_hitting_f
Some(EXECUTION_PATH_LOCAL_AI_PUBLIC)
);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["detail"], "当前 Codex 图片反代仅支持 n=1");
assert_eq!(
payload["detail"],
format!(
"当前图片反代仅支持 n=1..{}",
openai_image_gateway_max_generation_count()
)
);
assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
@@ -643,7 +650,7 @@ async fn gateway_rejects_image_request_with_n_greater_than_one_without_hitting_f
}
#[tokio::test]
async fn gateway_rejects_variation_request_without_image_without_hitting_fallback_probe() {
async fn gateway_does_not_mount_image_variation_route_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(
@@ -685,16 +692,7 @@ async fn gateway_rejects_variation_request_without_image_without_hitting_fallbac
.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"], "图片变体请求需要 image 文件");
assert_eq!(response.status(), StatusCode::NOT_FOUND);
assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
apps/aether-gateway/src/tests/frontdoor/public_support.rs
+293
View File
@@ -69,6 +69,7 @@ async fn gateway_handles_public_announcements_list_without_proxying_upstream() {
5,
true,
true,
false,
Some("admin-1".to_string()),
Some("admin".to_string()),
None,
@@ -85,6 +86,7 @@ async fn gateway_handles_public_announcements_list_without_proxying_upstream() {
3,
true,
false,
false,
Some("admin-2".to_string()),
Some("ops".to_string()),
None,
@@ -101,6 +103,7 @@ async fn gateway_handles_public_announcements_list_without_proxying_upstream() {
100,
false,
true,
false,
Some("admin-3".to_string()),
Some("root".to_string()),
None,
@@ -170,6 +173,7 @@ async fn gateway_handles_public_active_announcements_without_proxying_upstream()
50,
true,
false,
false,
Some("admin-1".to_string()),
Some("admin".to_string()),
Some((now.saturating_sub(60)) as i64),
@@ -186,6 +190,7 @@ async fn gateway_handles_public_active_announcements_without_proxying_upstream()
10,
true,
false,
false,
Some("admin-2".to_string()),
Some("ops".to_string()),
Some((now.saturating_add(3600)) as i64),
@@ -251,6 +256,7 @@ async fn gateway_handles_public_announcement_detail_without_proxying_upstream()
10,
true,
true,
false,
Some("admin-1".to_string()),
Some("admin".to_string()),
Some(1_711_000_000),
@@ -390,6 +396,7 @@ async fn gateway_updates_announcement_locally_with_trusted_admin_principal() {
10,
true,
true,
false,
Some("admin-1".to_string()),
Some("admin".to_string()),
None,
@@ -479,6 +486,7 @@ async fn gateway_deletes_announcement_locally_with_trusted_admin_principal() {
10,
true,
true,
false,
Some("admin-1".to_string()),
Some("admin".to_string()),
None,
@@ -561,6 +569,7 @@ async fn gateway_returns_service_unavailable_for_admin_announcement_writes_witho
10,
true,
true,
false,
Some("admin-1".to_string()),
Some("admin".to_string()),
None,
@@ -1436,6 +1445,22 @@ async fn gateway_handles_auth_registration_settings_without_proxying_upstream()
"turnstile_secret_key".to_string(),
json!("secret-private-key"),
),
(
"registration_privacy_policy_enabled".to_string(),
json!(true),
),
(
"registration_privacy_policy_format".to_string(),
json!("html"),
),
(
"registration_privacy_policy_content".to_string(),
json!("<p>Policy</p>"),
),
(
"registration_privacy_policy_version".to_string(),
json!("2026-05-16"),
),
]);
let (upstream_url, upstream_handle) = start_server(upstream).await;
@@ -1464,6 +1489,12 @@ async fn gateway_handles_auth_registration_settings_without_proxying_upstream()
"turnstile_enabled": true,
"turnstile_site_key": "site-public-key",
"turnstile_required_actions": ["send_verification_code", "register"],
"privacy_policy": {
"enabled": true,
"format": "html",
"content": "<p>Policy</p>",
"version": "2026-05-16",
},
})
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
@@ -1920,6 +1951,10 @@ async fn gateway_handles_public_test_connection_without_hitting_fallback_probe()
);
assert_eq!(body_json["model"], "gpt-5");
assert_eq!(body_json["messages"][0]["content"], "Health check");
assert!(
body_json.get("max_tokens").is_none(),
"public OpenAI-compatible test connection must not force a tiny max_tokens value"
);
Json(json!({"id": "resp_local_test"})).into_response()
}
}),
@@ -1978,6 +2013,91 @@ async fn gateway_handles_public_test_connection_without_hitting_fallback_probe()
provider_handle.abort();
}
#[tokio::test]
async fn gateway_gemini_test_connection_does_not_force_low_max_output_tokens() {
let provider_hits = Arc::new(Mutex::new(0usize));
let provider_hits_clone = Arc::clone(&provider_hits);
let provider = Router::new().route(
"/{*path}",
any(move |request: Request| {
let provider_hits_inner = Arc::clone(&provider_hits_clone);
async move {
*provider_hits_inner.lock().expect("mutex should lock") += 1;
let body = to_bytes(request.into_body(), usize::MAX)
.await
.expect("body should read");
let body_json: serde_json::Value =
serde_json::from_slice(&body).expect("json body should parse");
assert_eq!(body_json["contents"][0]["parts"][0]["text"], "Health check");
assert!(
body_json
.get("generationConfig")
.and_then(|config| config.get("maxOutputTokens"))
.is_none(),
"Gemini test connection must not force a tiny maxOutputTokens value"
);
Json(json!({
"candidates": [{
"content": {
"role": "model",
"parts": [{"text": "ok"}]
},
"finishReason": "STOP"
}],
"responseId": "gemini_test_connection_ok"
}))
.into_response()
}
}),
);
let (provider_url, provider_handle) = start_server(provider).await;
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-gemini", "google", 10)],
vec![sample_endpoint(
"endpoint-gemini",
"provider-gemini",
"gemini:generate_content",
&provider_url,
)],
vec![sample_key(
"key-gemini",
"provider-gemini",
"gemini:generate_content",
"google-api-key",
)],
));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY,
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/v1/test-connection?provider=provider-gemini&model=gemini-3-flash-preview&api_format=gemini:generate_content"
))
.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"], "success");
assert_eq!(payload["provider_id"], "provider-gemini");
assert_eq!(payload["endpoint_id"], "endpoint-gemini");
assert_eq!(payload["api_format"], "gemini:generate_content");
assert_eq!(*provider_hits.lock().expect("mutex should lock"), 1);
gateway_handle.abort();
provider_handle.abort();
}
async fn assert_public_support_route_returns_local_503(
method: reqwest::Method,
path: &str,
@@ -2839,6 +2959,7 @@ async fn gateway_reads_announcement_unread_count_locally_without_proxying_upstre
10,
true,
false,
false,
Some("admin-1".to_string()),
Some("admin".to_string()),
None,
@@ -2855,6 +2976,7 @@ async fn gateway_reads_announcement_unread_count_locally_without_proxying_upstre
8,
true,
true,
false,
Some("admin-1".to_string()),
Some("admin".to_string()),
None,
@@ -2871,6 +2993,7 @@ async fn gateway_reads_announcement_unread_count_locally_without_proxying_upstre
6,
false,
false,
false,
Some("admin-1".to_string()),
Some("admin".to_string()),
None,
@@ -2917,6 +3040,122 @@ async fn gateway_reads_announcement_unread_count_locally_without_proxying_upstre
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_lists_required_unread_announcements_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-announcement-required-1"),
),
]),
now + chrono::Duration::hours(1),
);
let announcement_repository = Arc::new(InMemoryAnnouncementReadRepository::seed_with_reads(
vec![
StoredAnnouncement::new(
"announcement-required".to_string(),
"必读公告".to_string(),
"需要确认".to_string(),
"important".to_string(),
20,
true,
false,
true,
Some("admin-1".to_string()),
Some("admin".to_string()),
None,
None,
now.timestamp(),
now.timestamp(),
)
.expect("announcement should build"),
StoredAnnouncement::new(
"announcement-normal".to_string(),
"普通公告".to_string(),
"不需要弹窗".to_string(),
"info".to_string(),
10,
true,
false,
false,
Some("admin-1".to_string()),
Some("admin".to_string()),
None,
None,
now.timestamp(),
now.timestamp(),
)
.expect("announcement should build"),
StoredAnnouncement::new(
"announcement-read-required".to_string(),
"已读必读公告".to_string(),
"已经确认".to_string(),
"warning".to_string(),
8,
true,
false,
true,
Some("admin-1".to_string()),
Some("admin".to_string()),
None,
None,
now.timestamp(),
now.timestamp(),
)
.expect("announcement should build"),
],
[(
"user-auth-1".to_string(),
"announcement-read-required".to_string(),
)],
));
let (gateway_url, upstream_hits, gateway_handle, upstream_handle) =
start_auth_announcement_gateway_with_state(
user,
sample_auth_wallet("user-auth-1", now),
[sample_auth_session(
"user-auth-1",
"session-announcement-required-1",
"device-announcement-required-1",
"refresh-token-placeholder",
now,
)],
announcement_repository,
)
.await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/announcements/users/me/required-unread"
))
.header("authorization", format!("Bearer {access_token}"))
.header("x-client-device-id", "device-announcement-required-1")
.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["total"], 1);
assert_eq!(payload["items"][0]["id"], "announcement-required");
assert_eq!(payload["items"][0]["requires_ack"], true);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_marks_announcement_read_status_locally_without_proxying_upstream() {
let now = Utc::now();
@@ -2946,6 +3185,7 @@ async fn gateway_marks_announcement_read_status_locally_without_proxying_upstrea
20,
true,
true,
false,
Some("admin-1".to_string()),
Some("admin".to_string()),
None,
@@ -3040,6 +3280,7 @@ async fn gateway_marks_all_announcements_read_locally_without_proxying_upstream(
10,
true,
false,
false,
Some("admin-1".to_string()),
Some("admin".to_string()),
None,
@@ -3056,6 +3297,7 @@ async fn gateway_marks_all_announcements_read_locally_without_proxying_upstream(
8,
false,
false,
false,
Some("admin-1".to_string()),
Some("admin".to_string()),
None,
@@ -3072,6 +3314,7 @@ async fn gateway_marks_all_announcements_read_locally_without_proxying_upstream(
6,
true,
true,
false,
Some("admin-1".to_string()),
Some("admin".to_string()),
None,
@@ -3163,6 +3406,7 @@ async fn gateway_handles_announcement_user_routes_with_trailing_slash_locally()
10,
true,
false,
false,
Some("admin-1".to_string()),
Some("admin".to_string()),
None,
@@ -3323,6 +3567,7 @@ async fn gateway_rejects_invalid_nested_announcement_paths_as_local_not_found_wi
10,
true,
false,
false,
Some("admin-1".to_string()),
Some("admin".to_string()),
None,
@@ -7978,6 +8223,54 @@ async fn gateway_handles_auth_register_locally_without_proxying_upstream() {
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_rejects_auth_register_without_current_privacy_policy_acceptance() {
let (gateway_url, upstream_hits, gateway_handle, upstream_handle) =
start_auth_gateway_with_builder(|| {
let data_state = crate::data::GatewayDataState::disabled()
.with_system_config_values_for_tests(vec![
("enable_registration".to_string(), json!(true)),
("require_email_verification".to_string(), json!(true)),
("smtp_host".to_string(), json!("smtp.example.com")),
("smtp_from_email".to_string(), json!("[email protected]")),
(
"registration_privacy_policy_enabled".to_string(),
json!(true),
),
(
"registration_privacy_policy_version".to_string(),
json!("2026-05-16"),
),
]);
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data_state)
.with_auth_email_verified_for_tests("[email protected]")
})
.await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/auth/register"))
.json(&json!({
"email": "[email protected]",
"username": "alice",
"password": "secret123",
"privacy_policy_accepted": true,
"privacy_policy_version": "old-version",
}))
.send()
.await
.expect("register request should succeed");
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["detail"], "请先阅读并同意当前版本的隐私政策");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
async fn start_turnstile_siteverify_server(
response_payload: serde_json::Value,
status: StatusCode,
apps/aether-gateway/src/tests/usage/pricing.rs
+8 -5
View File
@@ -113,10 +113,7 @@ struct ExpectedUsagePricing {
impl ExpectedUsagePricing {
fn total_tokens(self) -> u64 {
self.input_tokens
.saturating_add(self.output_tokens)
.saturating_add(self.cache_creation_tokens)
.saturating_add(self.cache_read_tokens)
self.input_tokens.saturating_add(self.output_tokens)
}
fn cache_creation_uncategorized_tokens(self) -> u64 {
@@ -1180,7 +1177,13 @@ async fn gateway_records_gemini_sync_usage_and_pricing_with_cache_read_tokens_im
},
"body": {
"json_body": {
"candidates": [],
"candidates": [{
"content": {
"role": "model",
"parts": [{"text": "Hello from Gemini"}]
},
"finishReason": "STOP"
}],
"usageMetadata": {
"promptTokenCount": expected.input_tokens,
"candidatesTokenCount": expected.output_tokens,